tavily search agent

graph.py

@@ -25,24 +25,11 @@ qdrant_research_chain = (
     )
 
 # Web Search Agent Pieces
-tavily_tool = TavilySearchResults(max_results=5)
-web_search_chain = (
-    {
-        "topic": itemgetter("topic"),
-        "qdrant_results": itemgetter("qdrant_results"),
-    }
-    | prompts.search_query_prompt
-    | models.gpt4o_mini 
-    | StrOutputParser()
-    | tavily_tool
-    | {
-        "topic": itemgetter("topic"),
-        "qdrant_results": itemgetter("qdrant_results"),
-        "search_results": RunnablePassthrough()
-    }
-    | prompts.summarize_prompt
-    | models.gpt4o_mini 
-    | StrOutputParser()
+tavily_tool = TavilySearchResults(max_results=3)
+query_chain = ( prompts.search_query_prompt | models.gpt4o_mini | StrOutputParser() )
+tavily_simple = ({"tav_results": tavily_tool} | prompts.tavily_prompt | models.gpt4o_mini | StrOutputParser())
+tavily_chain = (
+    {"query": query_chain} | tavily_simple
 )
 
 def query_qdrant(state: State) -> State:
@@ -65,7 +52,7 @@ def web_search(state: State) -> State:
     qdrant_results = state["research_data"].get("qdrant_results", "No previous results available.")
     
     # Run the web search chain
-    result = web_search_chain.invoke({
+    result = tavily_chain.invoke({
         "topic": topic,
         "qdrant_results": qdrant_results
     })

prompts.py

@@ -48,18 +48,16 @@ search_query_prompt = ChatPromptTemplate.from_template(
     """
 )
 
-# Create a prompt for summarizing the search results
-summarize_prompt = ChatPromptTemplate.from_template(
-    """Summarize the following search results, focusing on information that is complementary to what we already know from our database. Include sources for each piece of information:
+tavily_prompt = ChatPromptTemplate.from_template(
+    """Summarize the query results into the following format:
 
-    Topic: {topic}
-    
-    Information from our database:
-    {qdrant_results}
-    
-    Search results:
-    {search_results}
+    Query results: {tav_results}
     
-    Complementary summary with sources:
+    Output: 
+    Search Result 1 source: ...
+    Search Result 1 summary: ...
+    Search Result 2 source: ...
+    Search Result 2 summary: ...
+    ...
     """
 )