adjust tools and llm

agent.py

@@ -2,16 +2,30 @@ import os
 from typing import Dict, List, Sequence, TypedDict, cast
 
 from dotenv import load_dotenv
-from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader, WikipediaLoader
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.tools import tool
-from langchain_groq import ChatGroq
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_tavily import TavilySearch
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
+from pydantic import BaseModel
 
 
-@tool
+class WebSearchInput(BaseModel):
+    query: str
+
+
+class WikipediaSearchInput(BaseModel):
+    query: str
+
+
+class ArxivSearchInput(BaseModel):
+    query: str
+
+
+@tool(args_schema=WebSearchInput)
 def search_web(query: str) -> Dict[str, str]:
     """Search the web using Tavily and return relevant results."""
     try:
@@ -34,7 +48,7 @@ def search_web(query: str) -> Dict[str, str]:
         return {"error": f"Error searching web: {str(e)}"}
 
 
-@tool
+@tool(args_schema=WikipediaSearchInput)
 def search_wikipedia(query: str) -> Dict[str, str]:
     """Search Wikipedia using LangChain's loader and return the first document summary."""
     try:
@@ -50,6 +64,21 @@ def search_wikipedia(query: str) -> Dict[str, str]:
         return {"error": f"Error searching Wikipedia: {str(e)}"}
 
 
+@tool(args_schema=ArxivSearchInput)
+def arxiv_search(query: str) -> Dict[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()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return {"arxiv_results": formatted_search_docs}
+
+
 # System prompt
 system_prompt = SystemMessage(
     content="""You are a helpful and precise assistant. When answering questions:
@@ -75,6 +104,12 @@ Available tools:
   Returns: {"web_results": "results"} or {"error": "error message"}
   Best for: Recent events, current information, diverse sources
   Error handling: If no results found, try more specific search terms
+  
+- arxiv_search: Search scholarly papers on arXiv  
+  Input: {"query": "topic or keywords"}  
+  Returns: {"arxiv_results": "paper summaries with title, authors, abstract"} or {"error": "error message"}  
+  Best for: Academic research, recent papers in science and technology  
+  Error handling: If no results, simplify keywords or broaden the topic
 
 Tool usage strategy:
 1. For historical/factual queries:
@@ -87,12 +122,17 @@ Tool usage strategy:
    - If no results, try more specific terms
    - Cross-reference with Wikipedia if needed
 
-3. For complex queries:
-   - Use both tools to gather comprehensive info
+3. For academic/scientific questions:
+   - Use arxiv_search to find recent papers
+   - Summarize key findings, topics, or citations
+   - Cross-check with web or Wikipedia if needed
+
+4. For complex queries:
+   - Use all tools to gather comprehensive info
    - Compare and verify information
    - Note any discrepancies in your answer
 
-4. When both tools fail:
+5. Whenall tools fail:
    - Try different phrasings
    - Break complex queries into simpler parts
    - Be transparent about limitations in your answer
@@ -105,28 +145,40 @@ Your final answer must:
 5. Admit uncertainty when information is unclear"""
 )
 
+# Initialize tools
+tools = [search_wikipedia, search_web, arxiv_search]
+
 
-def build_agent_graph(provider: str = "groq"):
+def build_agent_graph(provider: str = "gemini"):
     """Build the graph"""
-    # Initialize tools
-    tools = [search_wikipedia, search_web]
 
-    # Initialize LLM with error handling
+    # Initialize LLM class
     try:
-        from pydantic import SecretStr
-
-        groq_api_key = os.getenv("GROQ_API_KEY")
-        if not groq_api_key:
-            raise EnvironmentError("GROQ_API_KEY environment variable is not set")
-
-        llm = ChatGroq(
-            model="qwen-qwq-32b",
-            temperature=0,
-            api_key=SecretStr(groq_api_key),
-        )
+        gemini_api_key = os.getenv("GEMINI_API_KEY")
+
+        if provider == "gemini":
+            chat_model = ChatGoogleGenerativeAI(
+                model="gemini-2.5-pro",
+                temperature=1.0,
+                max_retries=2,
+                google_api_key=gemini_api_key,
+            )
+        elif provider == "huggingface":
+            llm = HuggingFaceEndpoint(
+                repo_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+                task="text-generation",
+                max_new_tokens=1024,
+                do_sample=False,
+                repetition_penalty=1.03,
+                temperature=0,
+            )
+            chat_model = ChatHuggingFace(llm=llm, verbose=True)
+        else:
+            raise ValueError("Invalid provider.")
     except Exception as e:
         raise Exception(f"Failed to initialize LLM: {str(e)}")
-    llm_with_tools = llm.bind_tools(tools)
+
+    llm_with_tools = chat_model.bind_tools(tools)
 
     # Create nodes
     def assistant(state: MessagesState) -> Dict[str, List[AIMessage]]:
@@ -169,28 +221,31 @@ def test_agent():
     print("=" * 50)
 
     # Check environment variables
-    if not os.getenv("GROQ_API_KEY"):
-        print("\nError: GROQ_API_KEY not set")
+    if not os.getenv("HUGGINGFACEHUB_API_TOKEN"):
+        print("\nError: HUGGINGFACEHUB_API_TOKEN not set")
+        return
+    if not os.getenv("GEMINI_API_KEY"):
+        print("\nError: GEMINI_API_KEY not set")
         return
     if not os.getenv("TAVILY_API_KEY"):
         print("\nWarning: TAVILY_API_KEY not set - web search will be unavailable")
 
     print("\nInitializing agent...")
     try:
-        graph = build_agent_graph()
+        graph = build_agent_graph(provider="gemini")
         print("Agent initialized successfully")
     except Exception as e:
         print(f"Failed to initialize agent: {str(e)}")
         return
 
     # Test a single question
-    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?"
+    question = "What is the surname of the equine veterinarian mentioned in 1.E Exercises from the chemistry materials licensed by Marisa Alviar-Agnew & Henry Agnew under the CK-12 license in LibreText's Introductory Chemistry materials as compiled 08/21/2023?"
     print("\nTesting question:", question)
     print("-" * 50)
 
     try:
         # Create messages state
-        messages = [system_prompt, HumanMessage(content=question)]
+        messages = [HumanMessage(content=question)]
 
         # Run agent
         print("\nWaiting for response...")

app.py

@@ -42,7 +42,8 @@ class BasicAgent:
             try:
                 if not self.graph:
                     raise ValueError("Agent graph not initialized")
-                result = self.graph.invoke(input_state)
+                # result = self.graph.invoke(input_state)
+                result = self.graph.invoke({"messages": [human_msg]})
             except Exception as e:
                 print(f"Graph invocation error: {str(e)}")
                 return f"Error: Failed to process request - {str(e)}"

requirements.txt

@@ -10,4 +10,8 @@ requests
 pydantic
 lxml
 youtube-transcript-api==0.6.3
-pytube>=15.0.0
+pytube>=15.0.0
+langchain_huggingface
+langchain-google-genai
+pymupdf
+arxiv