default google

agent.py

@@ -16,7 +16,7 @@ from tools import *
 load_dotenv()
 
 
-def buildAgent(provider="groq"):
+def buildAgent(provider="google"):
     # load the system prompt from the file
     with open("system_prompt.txt", "r", encoding="utf-8") as f:
         system_prompt = f.read()
@@ -32,6 +32,14 @@ def buildAgent(provider="groq"):
         )
     elif provider == "groq":
         llm = ChatGroq(model="qwen-qwq-32b")
+    elif provider == "google":
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash-exp")
+    elif provider == "openrouter":
+        llm = ChatOpenAI(
+            base_url="https://openrouter.ai/api/v1",
+            model="google/gemini-2.0-flash-001",
+            api_key=os.getenv("OPENROUTER_API_KEY"),
+        )
     else:
         raise ValueError("Invalid provider. Choose 'groq' or 'huggingface'.")
 
@@ -61,7 +69,7 @@ def buildAgent(provider="groq"):
     def retriever(state: MessagesState):
         """Retriever node"""
         # Handle the case when no similar questions are found
-        return {"messages": [sys_msg] + state["messages"]}
+        return {"messages": state["messages"]}
 
     ## The graph
     builder = StateGraph(MessagesState)
@@ -83,11 +91,12 @@ def buildAgent(provider="groq"):
 
 
 if __name__ == "__main__":
-    random_question_url = "https://agents-course-unit4-scoring.hf.space/random-question"
-    response = requests.get(random_question_url, timeout=15)
-    questions_data = response.json()
-    question = questions_data.get("question")
-    graph = buildAgent(provider="groq")
+    # random_question_url = "https://agents-course-unit4-scoring.hf.space/random-question"
+    # response = requests.get(random_question_url, timeout=15)
+    # questions_data = response.json()
+    # question = questions_data.get("question")
+    question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
+    graph = buildAgent(provider="google")
     messages = [HumanMessage(content=question)]
     print(messages)
     messages = graph.invoke({"messages": messages})