format tools

agents/search_agent.py

@@ -17,6 +17,8 @@ class SearchAgent:
         state = workflow.invoke({
             "messages":messages,
         }, config={"callbacks": [langfuse_handler]})
+        
+        print(state["external_information"])
         return state["answer"]
     
 if __name__ == "__main__":

graphs/evaluation.py

@@ -1,5 +1,5 @@
 from models.models import  groq_model, anthropic_model
-from tools.search import arxiv_search, taivily_search, serper_search
+from tools.search import  taivily_search, serper_search
 from langgraph.graph import StateGraph, START, END, MessagesState
 from langchain_core.messages import HumanMessage, SystemMessage
 from typing import List, TypedDict
@@ -7,7 +7,7 @@ from langgraph.prebuilt import ToolNode
 
 
 tools = [
-    taivily_search,
+    #taivily_search,
     serper_search,
 ]
 
@@ -24,6 +24,9 @@ bound_model_llama = groq_model.bind_tools(tools)
 bound_model_antrhropic = anthropic_model.bind_tools(tools)
 
 def call_node(state: EvaluationState):
+    """
+    This node call the model with the question and the tools
+    """
     question = state["messages"][-1].content
     state["question"] = question
     response = bound_model_llama.invoke(state["messages"])
@@ -61,21 +64,11 @@ def map_answer(state: EvaluationState):
     Map the answer to the final answer
     """
     answer = anthropic_model.invoke("Map the answer, I want only the number, string or list. ANSWER:"+ state["answer"])
-    print(answer.content)
+    
     return {
         "answer": answer.content
     }
 
-def map_tool_answer(state: EvaluationState):
-    """
-    Map the tool answer to the final answer
-    """
-    last_message = state["messages"][-1]
-    state["external_information"] = last_message.content
-    
-    return state
-    
-
 def build_workflow():
     """
     Build search  workflow
@@ -83,16 +76,12 @@ def build_workflow():
     workflow = StateGraph(EvaluationState)
     workflow.add_node("agent", call_node)
     workflow.add_node("action", tool_node)
-    workflow.add_node("map_tool_answer", map_tool_answer)
     workflow.add_node("parse_response", parse_response)
     workflow.add_node("map_answer", map_answer)
-    """ workflow.add_node("action",tool_node)
-    workflow.add_node("answer", parse_response) """
 
     workflow.add_edge(START,"agent")
     workflow.add_edge("agent", "action")
-    workflow.add_edge("action", "map_tool_answer")
-    workflow.add_edge("map_tool_answer", "parse_response")
+    workflow.add_edge("action", "parse_response")
     workflow.add_edge("parse_response", "map_answer")
     workflow.add_edge("map_answer", END)
     

tools/search.py

@@ -1,32 +1,18 @@
 from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.utilities import GoogleSerperAPIWrapper
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
 from dotenv import load_dotenv
+from typing import Annotated
+from typing_extensions import Annotated
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.runnables import RunnableConfig
+from langgraph.types import Command
+from langchain_core.messages import ToolMessage
 
 load_dotenv()
 
 @tool
-def wikipedia_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 1 result.
-    Args:
-        query: The search query."""
-    query = "Mercedes Sosa"
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-            if "Closed-ended question" not in doc.metadata.get("title", "")
-        ]
-    )
-    
-    return {"wiki_results": formatted_search_docs}
-
-@tool
-def taivily_search(query: str) -> str:
+def taivily_search(query: str, tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):
     """Tavily is a search engine optimized for LLMs, aimed at efficient, quick and persistent search results. Tavily take care of all the burden of searching, scraping, filtering and extracting the most relevant information from online sources.
     Args:
         query: The search query."""
@@ -35,39 +21,35 @@ def taivily_search(query: str) -> str:
     
     formatted_search_docs = "\n\n---\n\n".join(
         [
-            f'<Document source="{doc["url"]}""/>\n{doc["content"]}\n</Document>'
+            f"""# Taivily \n source="{doc["url"]}"" \n ## Content \n {doc["content"]} \n---"""
             for doc in search_docs
         ]
     )
     
-    return {"web_results": formatted_search_docs}
-
-
-@tool
-def arxiv_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    print(search_docs)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["Title"]}""/>\n{doc.page_content[:1000]}\n</Document>'
-            for doc in search_docs
-        ]
+    return Command(
+        update={
+            "external_information": f"{config.get('external_information', '')}\n\n---\n\n{formatted_search_docs}",
+            "messages": [ToolMessage(content=formatted_search_docs, tool_call_id=tool_call_id)]
+        }
     )
-    return {"arxiv_results": formatted_search_docs}
 
 @tool
-def serper_search(query: str) -> str:
+def serper_search(query: str, tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig) -> Command:
     """
-    Search Google for a query and return maximum 2 result. 
-    Args: query: The search query.
+    lightning-fast Google search results in 1-2 seconds.
+    Args: 
+        query: The search query.
     """
-    search_docs = GoogleSerperAPIWrapper(k=2)
+    search_docs = GoogleSerperAPIWrapper(k=8)
     result = search_docs.run(query)
     
-    return {"google_results": result}
+    formatted_result = f"""# Google Search Results \n\n## Results \n\n{result}"""
+    
+    return Command(
+        update={
+            "external_information": f"{config.get('external_information', '')}\n\n---\n\n{formatted_result}",
+            "messages": [ToolMessage(content=formatted_result, tool_call_id=tool_call_id)]
+        }
+    )
+
 
-   
-