philincloud commited on
Commit
2345f27
·
verified ·
1 Parent(s): 7f064dc

Update langgraph_agent.py

Browse files
Files changed (1) hide show
  1. langgraph_agent.py +10 -7
langgraph_agent.py CHANGED
@@ -66,7 +66,7 @@ def arvix_search(query: str) -> dict:
66
  )
67
  return {"arvix_results": formatted}
68
 
69
- OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
70
  HF_SPACE_TOKEN = os.getenv("HF_SPACE_TOKEN")
71
 
72
 
@@ -82,13 +82,16 @@ sys_msg = SystemMessage(content=system_prompt)
82
 
83
 
84
  def build_graph(provider: str = "openai"):
85
- """Build the LangGraph agent with chosen LLM (default: OpenAI)."""
86
  if provider == "openai":
87
- llm = ChatOpenAI(
88
- model_name="o4-mini-2025-04-16",
89
- openai_api_key=OPENAI_API_KEY,
90
- # no temperature override here
91
- )
 
 
 
92
  elif provider == "huggingface":
93
  llm = ChatHuggingFace(
94
  llm=HuggingFaceEndpoint(
 
66
  )
67
  return {"arvix_results": formatted}
68
 
69
+ API_KEY = os.getenv("GEMINI_API_KEY")
70
  HF_SPACE_TOKEN = os.getenv("HF_SPACE_TOKEN")
71
 
72
 
 
82
 
83
 
84
  def build_graph(provider: str = "openai"):
85
+ """Build the LangGraph agent with chosen LLM (default: Gemini)."""
86
  if provider == "openai":
87
+ lllm = ChatGoogleGenerativeAI(
88
+ model= "gemini-2.5-pro-preview-05-06",
89
+ temperature=1.0,
90
+ max_retries=2,
91
+ google_api_key=api_key,
92
+ )
93
+
94
+
95
  elif provider == "huggingface":
96
  llm = ChatHuggingFace(
97
  llm=HuggingFaceEndpoint(