add search agent

.gitignore

@@ -1,2 +1,4 @@
 .vscode
 .venv
+*.env
+__pycache__/

agents/agent.py

@@ -1,14 +1,10 @@
-import os
 from typing import TypedDict, List, Dict, Any, Optional
 from langgraph.graph import StateGraph, START, END
 from langchain_anthropic import ChatAnthropic
 from langchain_groq import ChatGroq
 from langchain_core.messages import HumanMessage
 import getpass
-
-if "ANTHROPIC_API_KEY" not in os.environ:
-    os.environ["ANTHROPIC_API_KEY"] = getpass.getpass("Enter your Anthropic API key: ")
-    
+ 
 class EmailState(TypedDict):
     # The email being processed
     email: Dict[str, Any]  # Contains subject, sender, body, etc.

agents/search_agent.py

@@ -0,0 +1,14 @@
+from graphs.search import build_workflow
+from langchain_core.messages import HumanMessage
+class SearchAgent:
+    def __init__(self):
+        print("SearchAgent initialized.")
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+        workflow = build_workflow()
+        messages = [HumanMessage(question)]
+        messages = workflow.invoke({
+            "messages":messages
+        })
+        
+        return messages["messages"][-1].content

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from  models.basic_agent import BasicAgent
+from  agents.search_agent import SearchAgent
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -31,7 +31,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent()
+        agent = SearchAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None

graphs/search.py

@@ -0,0 +1,51 @@
+from models.models import  groq_model
+from tools.search import arxiv_search, web_search, wikipedia_search
+from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import HumanMessage
+
+tools = [
+    arxiv_search,
+    web_search,
+    wikipedia_search
+]
+
+tool_node = ToolNode(tools)
+bound_model = groq_model.bind_tools(tools)
+
+# Define the function that calls the model
+def call_model(state: MessagesState):
+    response = bound_model.invoke(state["messages"])
+    # We return a list, because this will get added to the existing list
+    return {"messages": response}
+
+def should_continue(state:MessagesState):
+    last_message = state["messages"][-1]
+    
+    if not last_message.tool_calls:
+        return END
+    
+    return "action"
+
+def build_workflow():
+    """
+    Build search  workflow
+    """
+    workflow = StateGraph(MessagesState)
+    workflow.add_node("agent", call_model)
+    workflow.add_node("action",tool_node)
+
+    workflow.add_edge(START,"agent")
+    workflow.add_conditional_edges("agent", should_continue)
+    workflow.add_edge("action", "agent")
+    return workflow.compile()
+
+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_workflow()
+    # Run the graph
+    messages = [HumanMessage(content=question)]
+    messages = graph.invoke({"messages": messages})
+    for m in messages["messages"]:
+        m.pretty_print()

models/models.py

@@ -0,0 +1,15 @@
+from langchain_anthropic import ChatAnthropic
+from langchain_groq import ChatGroq
+from dotenv import load_dotenv
+
+load_dotenv()
+ 
+anthropic_model = ChatAnthropic(
+    model="claude-3-5-haiku-latest", 
+    temperature=0
+)
+
+groq_model = ChatGroq(
+    model="meta-llama/llama-4-maverick-17b-128e-instruct",
+    temperature=0
+)

tools/search.py

@@ -0,0 +1,56 @@
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from dotenv import load_dotenv
+
+load_dotenv()
+
+@tool
+def wikipedia_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 1 result.
+    Args:
+        query: The search query."""
+    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 web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 1 results.
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=1).invoke(input=query)
+    
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc["url"]}""/>\n{doc["content"]}\n</Document>'
+            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 {"arxiv_results": formatted_search_docs}