Update agent.py

agent.py

@@ -1,224 +1,116 @@
+import os
+import requests
 from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from langchain_community.document_loaders import WikipediaLoader
 from llama_index.core.tools.types import ToolMetadata
 from llama_index.core.schema import Document
 from llama_index.core.tools import FunctionTool
-from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
 from llama_index.core.agent.workflow import AgentWorkflow
-import os
-import requests
 
 hf_token = os.getenv("HF_TOKEN")
 
-# List of models to try in order
-model_list = [
-    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # Your current working model
-    "microsoft/phi-3-mini-128k-instruct",   # You mentioned this works
-    "google/gemma-2b-it",                  # Backup option
-    "gpt2"                                 # Final fallback
+# List of models to try in order\ nmodel_list = [
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "microsoft/phi-3-mini-128k-instruct",
+    "google/gemma-2b-it",
+    "gpt2"
 ]
 
-# Try to create an LLM with the first model
 current_model_index = 0
 llm = HuggingFaceInferenceAPI(
     model_name=model_list[current_model_index],
-    token=os.getenv("HF_TOKEN"),
+    token=hf_token,
 )
 
-@tool
+# Numerical operation functions (plain definitions)
 def multiply(a: int, b: int) -> int:
-    """Multiply two numbers.
-    Args:
-        a: first int
-        b: second int
-    """
+    """Multiply two numbers."""
     return a * b
 
-@tool
+# Changed: removed @tool decorator above each function
+
 def add(a: int, b: int) -> int:
-    """Add two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
+    """Add two numbers."""
     return a + b
 
-@tool
 def subtract(a: int, b: int) -> int:
-    """Subtract two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
+    """Subtract two numbers."""
     return a - b
 
-@tool
-def divide(a: int, b: int) -> float:  # Changed to float for accurate division
-    """Divide two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
+def divide(a: int, b: int) -> float:
+    """Divide two numbers, raises error on zero divisor."""
     if b == 0:
         raise ValueError("Cannot divide by zero.")
     return a / b
 
-@tool
 def modulus(a: int, b: int) -> int:
-    """Get the modulus of two numbers.
-    
-    Args:
-        a: first int
-        b: second int
-    """
+    """Get the modulus of two numbers."""
     return a % b
-    
-@tool
+
+# Web search tool function (plain definition)
 def web_search(query: str) -> list:
-    """Search Tavily for a query and return a maximum 3 results as LlamaIndex Documents.
-    Args:
-        query: The search query.
-    """
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    
-    llama_docs = []
-    for doc in search_docs:
-        source = doc.metadata.get("source", "")
-        page = doc.metadata.get("page", "")
-        content = doc.page_content
-        full_metadata = {
-            "source": source,
-            "page": page
-        }
-        llama_docs.append(Document(text=content, metadata=full_metadata))
-
-    return llama_docs
-
-@tool
+    """Search Tavily for a query and return up to 3 results."""
+    results = TavilySearchResults(max_results=3).invoke(query=query)
+    docs = []
+    for r in results:
+        meta = {"source": r.metadata.get("source", ""), "page": r.metadata.get("page", "")}
+        docs.append(Document(text=r.page_content, metadata=meta))
+    return docs
+
+# Wikipedia search tool function (plain definition)
 def wiki_search(query: str) -> list:
-    """Search Wikipedia for a query and return maximum 2 results as LlamaIndex Documents.
-    Args:
-        query: The search query.
-    """
-    # Perform Wikipedia search using WikipediaLoader
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    
-    # Create a list of LlamaIndex Document objects
-    llama_docs = []
-    for doc in search_docs:
-        source = doc.metadata.get("source", "")
-        page = doc.metadata.get("page", "")
-        content = doc.page_content
-        full_metadata = {
-            "source": source,
-            "page": page
-        }
-        llama_docs.append(Document(text=content, metadata=full_metadata))
-
-    return llama_docs
-
-
-web_search = FunctionTool(
+    """Search Wikipedia for a query and return up to 2 results."""
+    results = WikipediaLoader(query=query, load_max_docs=2).load()
+    docs = []
+    for r in results:
+        meta = {"source": r.metadata.get("source", ""), "page": r.metadata.get("page", "")}
+        docs.append(Document(text=r.page_content, metadata=meta))
+    return docs
+
+# Wrap functions into FunctionTool instances\ +
+web_search_tool = FunctionTool(
     web_search,
     metadata=ToolMetadata(name="web_search", description="Tavily 3-hit search")
 )
-
-wiki_search = FunctionTool(
+wiki_search_tool = FunctionTool(
     wiki_search,
     metadata=ToolMetadata(name="wiki_search", description="Wikipedia 2-hit search")
 )
 
-multiply_tool = FunctionTool(
-    multiply,
-    metadata=ToolMetadata(name="multiply", description="Multiply two numbers.")
-)
-add_tool = FunctionTool(
-    add,
-    metadata=ToolMetadata(name="add", description="Add two numbers.")
-)
-subtract_tool = FunctionTool(
-    subtract,
-    metadata=ToolMetadata(name="subtract", description="Subtract two numbers.")
-)
-divide_tool = FunctionTool(
-    divide,
-    metadata=ToolMetadata(name="divide", description="Divide two numbers.")
-)
-modulus_tool = FunctionTool(
-    modulus,
-    metadata=ToolMetadata(name="modulus", description="Modulus operation on two numbers.")
-)
+multiply_tool = FunctionTool(multiply, metadata=ToolMetadata(name="multiply", description="Multiply two numbers."))
+add_tool      = FunctionTool(add,      metadata=ToolMetadata(name="add",      description="Add two numbers."))
+subtract_tool = FunctionTool(subtract, metadata=ToolMetadata(name="subtract", description="Subtract two numbers."))
+divide_tool   = FunctionTool(divide,   metadata=ToolMetadata(name="divide",   description="Divide two numbers."))
+modulus_tool  = FunctionTool(modulus,  metadata=ToolMetadata(name="modulus",  description="Modulus operation on two numbers."))
 
-# Then include all tools in the list
+# Aggregate all tools
 tools = [
-    web_search,
-    wiki_search,
+    web_search_tool,
+    wiki_search_tool,
     multiply_tool,
     add_tool,
     subtract_tool,
     divide_tool,
-    modulus_tool
+    modulus_tool,
 ]
 
-# Create the agent with the first model
+# Initialize agent
 agent = AgentWorkflow.from_tools_or_functions(tools, llm=llm)
 
-# Simple function to try a different model when needed
-def try_next_model():
+# Improved fallback logic with loop
+# Changed: simplified run_with_fallback to iterate through models
+
+def run_with_fallback(query: str):
     global current_model_index, llm, agent
-    
-    # Move to the next model if available
-    if current_model_index < len(model_list) - 1:
-        current_model_index += 1
-        next_model = model_list[current_model_index]
-        print(f"Switching to backup model: {next_model}")
-        
-        # Create a new LLM with the next model
-        llm = HuggingFaceInferenceAPI(
-            model_name=next_model,
-            token=os.getenv("HF_TOKEN"),
-        )
-        
-        # Create a new agent with the new model
-        agent = AgentWorkflow.from_tools_or_functions(tools, llm=llm)
-        return True
-    else:
-        print("No more models to try!")
-        return False
-
-# Function to run the agent with fallback
-def run_with_fallback(query):
-    global agent
-    
-    # Try to run the query with the current model
-    try:
-        return agent.run(query)
-    except Exception as e:
-        print(f"Error running agent: {e}")
-        
-        # Try next model
-        if try_next_model():
-            # Try the query with the new model
-            try:
-                return agent.run(query)
-            except Exception as e2:
-                print(f"Error with fallback model: {e2}")
-                
-                # One more attempt with another model
-                if try_next_model():
-                    try:
-                        return agent.run(query)
-                    except Exception as e3:
-                        print(f"Error with second fallback model: {e3}")
-        
-        # If all attempts fail, return an error message
-        return "Sorry, I encountered technical issues with all available models."
-
-# Usage example:
-# result = run_with_fallback("What is 42 * 5?")
-# print(result)
+    for _ in model_list:
+        try:
+            return agent.run(query)
+        except Exception as e:
+            print(f"Error with model {model_list[current_model_index]}: {e}")
+            if not try_next_model():
+                break
+    return "Sorry, encountered issues with all models."
+
+# try_next_model remains unchanged