Update agent.py

agent.py

@@ -21,7 +21,6 @@ load_dotenv()
 @tool
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers.
-
     Args:
         a: first int
         b: second int
@@ -152,7 +151,7 @@ tools = [
 ]
 
 # Build graph function
-def build_graph(provider: str = "google"):
+def build_graph(provider: str = "groq"):
     """Build the graph"""
     # Load environment variables from .env file
     if provider == "google":
@@ -180,36 +179,12 @@ def build_graph(provider: str = "google"):
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
     
     def retriever(state: MessagesState):
-        current_question = state["messages"][-1].content
-        similar_doc = vector_store.similarity_search(current_question, k=1)[0]
-        similar_question = similar_doc.metadata["Question"]
-        similar_answer = similar_doc.page_content
-
-        # Wenn gleich: direkt beantworten
-        if current_question.strip() == similar_question.strip():
-            return {
-                "messages": state["messages"] + [
-                    sys_msg,
-                    HumanMessage(content=f"FINAL ANSWER: {similar_answer}")
-                ],
-                "should_end": True  # überspringt Tool-Nutzung etc.
-            }
-
-        # Andernfalls Beispiel einbauen
+        """Retriever node"""
+        similar_question = vector_store.similarity_search(state["messages"][0].content)
         example_msg = HumanMessage(
-            content=f"""
-            Here I provide a similar question and answer for reference:
-
-            Question: {similar_question}
-            Final answer: {similar_answer}
-
-        If the current question is the same, return the same final answer.
-        """
+            content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
         )
-        return {
-            "messages": [sys_msg] + state["messages"] + [example_msg],
-            "should_end": False  # weitere Bearbeitung notwendig
-        }
+        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
@@ -225,14 +200,3 @@ def build_graph(provider: str = "google"):
 
     # Compile graph
     return builder.compile()
-
-# test
-if __name__ == "__main__":
-    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
-    # Build the graph
-    graph = build_graph(provider="groq")
-    # Run the graph
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()