Updates

agent.py

@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Annotated
+from typing import Annotated, Dict, Optional
 
 import pandas as pd
 from langchain_community.document_loaders import WikipediaLoader, YoutubeLoader
@@ -23,40 +23,65 @@ from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode, tools_condition
 
 
+# Custom exception for tool errors
+class ToolExecutionError(Exception):
+    """Custom exception for tool execution errors"""
+
+    pass
+
+
 @tool("search_web_sources")
-def search_web_sources(query: Annotated[str, "Search query string"]) -> dict:
+def search_web_sources(query: Annotated[str, "Search query string"]) -> Dict[str, str]:
     """Performs a web search and returns up to 3 formatted documents with content and source."""
+    try:
+        if not os.environ.get("TAVILY_API_KEY"):
+            raise EnvironmentError(
+                "TAVILY_API_KEY is not set in environment variables."
+            )
 
-    if not os.environ.get("TAVILY_API_KEY"):
-        raise EnvironmentError("TAVILY_API_KEY is not set in environment variables.")
-    search_docs = TavilySearch(max_results=3).invoke({"query": query})
-    formatted = "\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 = TavilySearch(max_results=3).invoke({"query": query})
+        if not search_docs:
+            return {"web_results": "No results found for the given query."}
+
+        formatted = "\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}
+    except Exception as e:
+        return {"web_results": f"Error during web search: {str(e)}"}
 
 
 @tool
-def search_wikipedia(query: str) -> dict:
+def search_wikipedia(query: str) -> Dict[str, str]:
     """Search Wikipedia using LangChain's loader and return the first document summary."""
     try:
+        # Input validation
+        if not query or not isinstance(query, str):
+            return {
+                "wiki_results": "Invalid query provided. Please provide a valid search term."
+            }
+
         loader = WikipediaLoader(query=query, lang="en", load_max_docs=2)
         docs = loader.load()
-        if docs:
-            formatted_docs = "---".join(
-                [
-                    f'<WikipediaArticle title="{query}">{doc.page_content}</WikipediaArticle>'
-                    for doc in docs
-                ]
-            )
-            return {"wiki_results": formatted_docs}
-        else:
-            return {"wiki_results": "No content found."}
+
+        if not docs:
+            return {"wiki_results": f"No Wikipedia articles found for query: {query}"}
+
+        formatted_docs = "---".join(
+            [
+                f'<WikipediaArticle title="{query}">{doc.page_content}</WikipediaArticle>'
+                for doc in docs
+            ]
+        )
+        return {"wiki_results": formatted_docs}
     except Exception as e:
-        return {"wiki_results": f"Error fetching Wikipedia article: {e}"}
+        error_msg = str(e)
+        if "Page id" in error_msg and "not found" in error_msg:
+            return {"wiki_results": f"No Wikipedia article found for: {query}"}
+        return {"wiki_results": f"Error searching Wikipedia: {error_msg}"}
 
 
 @tool
@@ -97,17 +122,35 @@ def run_python_code(code: str) -> str:
 # --- System Prompt ---
 system_prompt = SystemMessage(
     content="""
-You are a helpful and precise assistant. You will receive a question and optionally access tools to help answer it.
-
-Your job is to think step-by-step, clearly report your thoughts, and conclude with a formatted response.
+You are a helpful and precise assistant with access to several tools. You will receive questions and use tools appropriately to find answers.
+
+When using tools:
+1. Format tool calls correctly using the tool's exact name and required parameters
+2. Validate inputs before making tool calls
+3. Handle tool responses appropriately, checking for errors
+4. If a tool fails, try an alternative approach or provide a clear error message
+
+Available tools:
+- search_web_sources: Search web for information (requires query parameter)
+- search_wikipedia: Search Wikipedia articles (requires query parameter)
+- extract_youtube_transcript: Get transcript from YouTube videos (requires video_url parameter)
+- run_python_code: Execute Python code (requires code parameter)
+
+Think step-by-step:
+1. Understand the question
+2. Choose appropriate tool(s)
+3. Format tool calls correctly
+4. Process tool responses
+5. Formulate final answer
 
 Use this format strictly:
 FINAL ANSWER: [your concise answer here]
 
 Rules for your answer:
-- If the answer is a number, write only the number (no commas, units, or symbols unless asked).
-- If it's a string, avoid articles (a, an, the), don't abbreviate, and use plain text digits.
-- If a list, follow the rules above for each element and separate with a comma and single space (e.g., "apple, orange, banana").
+- If the answer is a number, write only the number (no commas, units, or symbols unless asked)
+- If it's a string, avoid articles (a, an, the), don't abbreviate, and use plain text digits
+- If a list, follow the rules above for each element and separate with a comma and single space (e.g., "apple, orange, banana")
+- If there's an error, start with "Error:" followed by a clear explanation
 
 Your response must always begin with: FINAL ANSWER:
 """
@@ -124,18 +167,39 @@ def build_agent_graph(provider: str = "groq"):
         run_python_code,
     ]
 
-    # Instantiate LLM
-    os.environ["GROQ_API_KEY"]
-    llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+    # Instantiate LLM with proper error handling
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    if not groq_api_key:
+        raise EnvironmentError("GROQ_API_KEY environment variable is not set")
+
+    try:
+        from pydantic import SecretStr
+
+        llm = ChatGroq(
+            model="qwen-qwq-32b", temperature=0, api_key=SecretStr(groq_api_key)
+        )
+    except Exception as e:
+        raise Exception(f"Failed to initialize Groq LLM: {str(e)}")
 
     # Bind tools to the LLM
     llm_with_tools = llm.bind_tools(tools)
 
     # Assistant: reasoning step that plans next action
     def assistant_node(state: MessagesState) -> dict:
-        messages = state["messages"]
-        response = llm_with_tools.invoke(messages)
-        return {"messages": response}
+        try:
+            messages = state["messages"]
+            response = llm_with_tools.invoke(messages)
+
+            # Validate response format
+            if not response or not isinstance(
+                response, (AIMessage, HumanMessage, SystemMessage)
+            ):
+                raise ValueError("Invalid response format from LLM")
+
+            return {"messages": response}
+        except Exception as e:
+            error_msg = f"Error in assistant node: {str(e)}"
+            return {"messages": AIMessage(content=f"FINAL ANSWER: {error_msg}")}
 
     # Stubbed retriever node for future integration
     def retriever_node(state: MessagesState):
@@ -156,17 +220,51 @@ def build_agent_graph(provider: str = "groq"):
     # ToolNode wrapper for actual tool use
     tool_node = ToolNode(tools)
 
-    # Define the graph with ReAct loop
+    # Define error handling node
+    def error_handler_node(state: MessagesState) -> dict:
+        """Handle errors in the graph execution"""
+        error_msg = state.get("error", "Unknown error occurred")
+        return {
+            "messages": AIMessage(content=f"FINAL ANSWER: Error occurred: {error_msg}")
+        }
+
+    # Define the graph with ReAct loop and error handling
     builder = StateGraph(MessagesState)
     builder.add_node("assistant", RunnableLambda(assistant_node))
     builder.add_node("tools", tool_node)
     builder.add_node("retriever", RunnableLambda(retriever_node))
+    builder.add_node("error_handler", RunnableLambda(error_handler_node))
 
     builder.set_entry_point("assistant")
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
     builder.add_edge("assistant", END)
 
+    # Add error handling edges
+    def route_by_error(state: MessagesState):
+        """Route to error handler if error is present, otherwise continue normal flow"""
+        if "error" in state:
+            return "error_handler"
+        return None
+
+    builder.add_conditional_edges(
+        "assistant",
+        route_by_error,
+        {
+            "error_handler": "error_handler",
+        },
+    )
+
+    builder.add_conditional_edges(
+        "tools",
+        route_by_error,
+        {
+            "error_handler": "error_handler",
+        },
+    )
+
+    builder.add_edge("error_handler", END)
+
     graph = builder.compile()
 
     # Optional: test entrypoint to run the graph manually

requirements.txt

@@ -1,14 +1,15 @@
-gradio
-requests
-langchain
-langchain-core
-langchain-community
-langchain_huggingface
-langchain-groq
-langchain-experimental
-langchain-tavily
-langgraph
-tavily-python
-wikipedia
+gradio>=4.0.0
+requests>=2.31.0
+langchain>=0.1.0
+langchain-core>=0.1.0
+langchain-community>=0.0.10
+langchain_huggingface>=0.0.10
+langchain-groq>=0.0.5
+langchain-experimental>=0.0.40
+langchain-tavily>=0.0.5
+langgraph>=0.0.15
+tavily-python>=0.3.0
+wikipedia>=1.4.0
 youtube-transcript-api==0.6.3
-pytube
+pytube>=15.0.0
+pydantic>=2.0.0