Update agent.py

app.py

@@ -3,8 +3,7 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
-from langchain_core.messages import HumanMessage
-from typing import TypedDict, Annotated, Sequence, Dict, Any, List
+from typing import TypedDict, Annotated, Sequence, Dict, Any, List, Optional
 from langchain_core.messages import BaseMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
@@ -15,6 +14,9 @@ from langchain_community.utilities import WikipediaAPIWrapper
 from langchain.agents import create_tool_calling_agent, AgentExecutor
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 import operator
+from langchain_experimental.utilities import PythonREPL
+from functools import wraps
+import logging
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -24,44 +26,97 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 
 
+
+# --- Configure logging ---
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# --- Constants ---
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+DEFAULT_MODEL = "gpt-3.5-turbo"
+MAX_RESPONSE_LENGTH = 2000  # Prevent overly long responses
+
+def handle_errors(func):
+    """Decorator to handle common errors in agent operations."""
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.error(f"Error in {func.__name__}: {str(e)}")
+            return {"error": str(e)}
+    return wrapper
+
 class AgentState(TypedDict):
     messages: Annotated[Sequence[BaseMessage], operator.add]
     sender: str
 
 @tool
 def wikipedia_search(query: str) -> str:
-    """Search Wikipedia for information."""
-    return WikipediaAPIWrapper().run(query)
+    """Search Wikipedia for information. Useful for historical facts, scientific concepts, and general knowledge."""
+    try:
+        return WikipediaAPIWrapper().run(query)[:MAX_RESPONSE_LENGTH]
+    except Exception as e:
+        return f"Wikipedia search failed: {str(e)}"
 
 @tool
 def web_search(query: str, num_results: int = 3) -> list:
-    """Search the web for current information."""
-    return DuckDuckGoSearchResults(num_results=num_results).run(query)
+    """Search the web for current information. Useful for news, recent events, and up-to-date data."""
+    try:
+        results = DuckDuckGoSearchResults(num_results=num_results).run(query)
+        return [str(r)[:500] for r in results][:num_results]  # Limit result size
+    except Exception as e:
+        return [f"Web search failed: {str(e)}"]
 
 @tool
 def calculate(expression: str) -> str:
-    """Evaluate mathematical expressions."""
-    from langchain_experimental.utilities import PythonREPL
-    python_repl = PythonREPL()
-    return python_repl.run(expression)
+    """Evaluate mathematical expressions. Supports basic arithmetic and complex formulas."""
+    try:
+        python_repl = PythonREPL()
+        result = python_repl.run(expression)
+        return str(result)[:100]  # Limit numeric output length
+    except Exception as e:
+        return f"Calculation failed: {str(e)}"
 
 class BasicAgent:
-    """A langgraph agent."""
-    def __init__(self):
-        print("BasicAgent initialized.")
-        self.graph = build_graph()
-
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        # Wrap the question in a HumanMessage from langchain_core
-        messages = [HumanMessage(content=question)]
-        messages = self.graph.invoke({"messages": messages})
-        answer = messages['messages'][-1].content
-        return answer[14:]
-
+    """An enhanced LangGraph agent with better error handling and response processing."""
     
+    def __init__(self, model_name: str = DEFAULT_MODEL, temperature: float = 0.7):
+        """Initialize the agent with tools and workflow."""
+        self.model_name = model_name
+        self.temperature = temperature
+        self.tools = [wikipedia_search, web_search, calculate]
+        self.llm = ChatOpenAI(model=model_name, temperature=temperature)
+        self.agent_executor = self._build_agent_executor()
+        self.workflow = self._build_workflow()
+        logger.info(f"AdvancedAgent initialized with model: {model_name}")
+
+    def _build_agent_executor(self) -> AgentExecutor:
+        """Build the agent executor with proper prompt and tools."""
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", self._get_system_prompt()),
+            MessagesPlaceholder(variable_name="messages"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ])
+        agent = create_tool_calling_agent(self.llm, self.tools, prompt)
+        return AgentExecutor(
+            agent=agent,
+            tools=self.tools,
+            verbose=True,
+            handle_parsing_errors=True
+        )
+
+    def _get_system_prompt(self) -> str:
+        """Return a comprehensive system prompt for the agent."""
+        return """You are an advanced AI assistant with access to tools. Follow these rules:
+1. Be precise and factual
+2. Use tools when needed
+3. Cite your sources
+4. Break complex problems into steps
+5. Admit when you don't know something"""
+
     def _build_workflow(self) -> StateGraph:
-        """Build and return the compiled workflow"""
+        """Build and compile the agent workflow."""
         workflow = StateGraph(AgentState)
         
         workflow.add_node("agent", self._run_agent)
@@ -76,19 +131,25 @@ class BasicAgent:
         workflow.add_edge("tools", "agent")
         
         return workflow.compile()
-    
+
+    @handle_errors
     def _run_agent(self, state: AgentState) -> Dict[str, Any]:
-        """Execute the agent"""
+        """Execute the agent with error handling."""
         response = self.agent_executor.invoke({"messages": state["messages"]})
         return {"messages": [response["output"]]}
-    
+
     def _should_continue(self, state: AgentState) -> str:
-        """Determine if the workflow should continue"""
+        """Determine if the workflow should continue based on tool calls."""
         last_message = state["messages"][-1]
         return "continue" if last_message.additional_kwargs.get("tool_calls") else "end"
-    
+
+    @handle_errors
     def __call__(self, query: str) -> Dict[str, Any]:
-        """Process a user query"""
+        """Process a user query and return a structured response."""
+        if not query or len(query.strip()) == 0:
+            return {"error": "Empty query provided"}
+            
+        logger.info(f"Processing query: {query[:50]}...")
         state = AgentState(messages=[HumanMessage(content=query)], sender="user")
         
         for output in self.workflow.stream(state):
@@ -96,29 +157,36 @@ class BasicAgent:
                 if key == "messages":
                     for message in value:
                         if isinstance(message, BaseMessage):
+                            response = message.content[:MAX_RESPONSE_LENGTH]
                             return {
-                                "response": message.content,
+                                "response": response,
                                 "sources": self._extract_sources(state["messages"]),
-                                "steps": self._extract_steps(state["messages"])
+                                "steps": self._extract_steps(state["messages"]),
+                                "model": self.model_name
                             }
         return {"response": "No response generated", "sources": [], "steps": []}
-    
+
     def _extract_sources(self, messages: Sequence[BaseMessage]) -> List[str]:
-        """Extract sources from tool messages"""
-        return [
-            f"{msg.additional_kwargs.get('name', 'unknown')}: {msg.content}"
-            for msg in messages
-            if hasattr(msg, 'additional_kwargs') and 'name' in msg.additional_kwargs
-        ]
-    
+        """Extract and format sources from tool messages."""
+        sources = []
+        for msg in messages:
+            if hasattr(msg, 'additional_kwargs') and 'name' in msg.additional_kwargs:
+                source_name = msg.additional_kwargs.get('name', 'unknown')
+                content = str(msg.content)[:200]  # Truncate long content
+                sources.append(f"{source_name}: {content}")
+        return sources
+
     def _extract_steps(self, messages: Sequence[BaseMessage]) -> List[str]:
-        """Extract reasoning steps"""
+        """Extract and format the reasoning steps."""
         steps = []
         for msg in messages:
             if hasattr(msg, 'additional_kwargs') and 'tool_calls' in msg.additional_kwargs:
                 for call in msg.additional_kwargs['tool_calls']:
-                    steps.append(f"Used {call['function']['name']}: {call['function']['arguments']}")
+                    tool_name = call['function']['name']
+                    args = call['function']['arguments'][:100]  # Truncate long args
+                    steps.append(f"Used {tool_name} with args: {args}")
         return steps
+
         
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """