add tavily search

agent.py

@@ -12,6 +12,7 @@ from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.tools import Tool
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_experimental.utilities import PythonREPL
 import assemblyai as aai
 
@@ -77,14 +78,29 @@ def wiki_search(query: str) -> str:
         ])
     return {"wiki_results": formatted_search_docs}
 
+# @tool
+# def web_search(query: str) -> str:
+#     """Search DuckDuckGo for a query and return maximum 5 results.
+    
+#     Args:
+#         query: The search query."""
+#     search_docs = DuckDuckGoSearchResults(max_results=5).invoke(query)
+#     return {"web_results": search_docs}
 @tool
 def web_search(query: str) -> str:
-    """Search DuckDuckGo for a query and return maximum 5 results.
+    """Search Tavily for a query and return maximum 3 results.
     
     Args:
         query: The search query."""
-    search_docs = DuckDuckGoSearchResults(max_results=5).invoke(query)
-    return {"web_results": search_docs}
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
+
+
     
 
 system_prompt =  "You are a helpful assistant"
@@ -98,13 +114,32 @@ tools = [
     transcribe_audio
 ]
 
-llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+llm = ChatGroq(model="qwen-qwq-32b", temperature=0.7)
 llm_with_tools = llm.bind_tools(tools)
 
 def assistant(state: MessagesState):
     """Assistant node"""
     return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
+def final_answer(question, answer):
+  sys_prompt = """
+  You are an assistant. You help to extract the final answer based given an answer
+  of a question.
+
+  Example:
+  **QUESTION:** what is 1+1?
+  **ANSWER:** the result is 2
+  **OUTPUT:** 2
+  """
+  conversation = f"Extract final answer from\QUESTION:{question}\ANSWER:{answer}\OUTPUT:"
+  print(f"conversation: {conversation}")
+  messages = [
+    {"role": "system", "content": sys_prompt},
+    {"role": "user", "content": conversation},
+  ]
+  response = llm.invoke(messages)
+  return response.content.replace("**OUTPUT:**","").replace("**ANSWER:**","")
+
 builder = StateGraph(MessagesState)
 builder.add_node("assistant", assistant)
 builder.add_node("tools", ToolNode(tools))
@@ -120,11 +155,14 @@ graph = builder.compile()
 def get_answer(query):
   messages = [HumanMessage(content=query)]
   results = graph.invoke({"messages": messages})
-  return results["messages"][-1].content
+  return final_answer(query, results["messages"][-1].content)
 
 if __name__ == "__main__":
     question = "In the video https://www.youtube.com/watch?v=L1vXCYZAYYM, what is the highest number of bird species to be on camera simultaneously?"
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
+    question = "Hi, I was out sick from my classes on Friday, so I'm trying to figure out what I need to study for my Calculus mid-term next week. My friend from class sent me an audio recording of Professor Willowbrook giving out the recommended reading for the test, but my headphones are broken :(\n\nCould you please listen to the recording for me and tell me the page numbers I'm supposed to go over? I've attached a file called Homework.mp3 that has the recording. Please provide just the page numbers as a comma-delimited list. And please provide the list in ascending order."
+    question = "What is the first name of the only Malko Competition recipient from the 20th Century (after 1977) whose nationality on record is a country that no longer exists?"
+    # getmessages = [HumanMessage(content=question)]
+    # messages = graph.invoke({"messages": messages})
+    # for m in messages["messages"]:
+    #     m.pretty_print()
+    print(f"FINAL ANSWER: {get_answer(question)}")

requirements.txt

@@ -7,4 +7,5 @@ langchain-community
 wikipedia
 duckduckgo-search
 langchain-experimental
-assemblyai
+assemblyai
+langchain-tavily