feat(tools): add more tool to extend the functionaily of jarvis

.gitignore

@@ -41,6 +41,7 @@ coverage.xml
 *.py,cover
 .tox/
 .pytest_cache/
+cache/
 
 # Logs and temporary files
 *.log

README.md

@@ -74,6 +74,10 @@ jarvis_gaia_agent/
   - `SERPAPI_API_KEY`: SERPAPI key for web searches.
   - `OPENWEATHERMAP_API_KEY`: OpenWeatherMap key for weather queries.
   - `SPACE_ID`: `onisj/jarvis_gaia_agent`.
+- Install dependencies:
+  ```bash
+  pip install -r requirements.txt
+  ```
 
 ## Setup and Local Testing
 

app.py

@@ -3,6 +3,7 @@ import json
 import logging
 import asyncio
 import aiohttp
+import ssl
 import nest_asyncio
 import requests
 import pandas as pd
@@ -10,18 +11,25 @@ from typing import Dict, Any, List
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.messages import SystemMessage, HumanMessage
 from langgraph.graph import StateGraph, END
+import torch
 from sentence_transformers import SentenceTransformer
 import gradio as gr
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import together
-from state import JARVISState
-from tools import (
-    search_tool, multi_hop_search_tool, file_parser_tool, image_parser_tool,
-    calculator_tool, document_retriever_tool, duckduckgo_search_tool,
-    weather_info_tool, hub_stats_tool, guest_info_retriever_tool
-)
+from state import JARVISState, validate_state, reset_state
+from tools.answer_generator import generate_answer, preprocess_question
+from tools.file_fetcher import fetch_task_file
+from tools.search import search_tool, multi_hop_search_tool
+from tools.file_parser import file_parser_tool
+from tools.image_parser import image_parser_tool
+from tools.calculator import calculator_tool
+from tools.document_retriever import document_retriever_tool
+from tools.duckduckgo_search import duckduckgo_search_tool
+from tools.weather_info import weather_info_tool
+from tools.hub_stats import hub_stats_tool
+from tools.guest_info import guest_info_retriever_tool
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -33,10 +41,10 @@ nest_asyncio.apply()
 # Load environment variables
 load_dotenv()
 SPACE_ID = os.getenv("SPACE_ID", "onisj/jarvis_gaia_agent")
-GAIA_API_URL = "https://agents-course-unit4-scoring.hf.space"
-GAIA_FILE_URL = f"{GAIA_API_URL}/files/"
+GAIA_API_URL = "https://agents-course-unit4-api-1.hf.space/api"
 TOGETHER_API_KEY = os.getenv("TOGETHER_API_KEY")
 HF_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+OPENWEATHERMAP_API_KEY = os.getenv("OPENWEATHERMAP_API_KEY")
 
 # Verify environment variables
 if not SPACE_ID:
@@ -45,6 +53,8 @@ if not HF_API_TOKEN:
     raise ValueError("HUGGINGFACEHUB_API_TOKEN not set")
 if not TOGETHER_API_KEY:
     raise ValueError("TOGETHER_API_KEY not set")
+if not OPENWEATHERMAP_API_KEY:
+    logger.warning("OPENWEATHERMAP_API_KEY not set; weather_info_tool may fail")
 logger.info(f"SPACE_ID: {SPACE_ID}")
 
 # Model configuration
@@ -56,23 +66,20 @@ HF_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
 
 # Initialize LLM clients
 def initialize_llm():
-    # Try Together AI models
     for model in TOGETHER_MODELS:
         try:
             together.api_key = TOGETHER_API_KEY
             client = together.Together()
-            # Test the model
             response = client.chat.completions.create(
                 model=model,
                 messages=[{"role": "user", "content": "Test"}],
                 max_tokens=10
             )
             logger.info(f"Initialized Together AI model: {model}")
-            return client, "together"
+            return client, "together", model
         except Exception as e:
             logger.warning(f"Failed to initialize Together AI model {model}: {e}")
     
-    # Fallback to Hugging Face Inference API
     try:
         client = InferenceClient(
             model=HF_MODEL,
@@ -80,381 +87,355 @@ def initialize_llm():
             timeout=30
         )
         logger.info(f"Initialized Hugging Face Inference API model: {HF_MODEL}")
-        return client, "hf_api"
+        return client, "hf_api", HF_MODEL
     except Exception as e:
         logger.warning(f"Failed to initialize HF Inference API: {e}")
     
-    # Fallback to local Hugging Face model
     try:
         tokenizer = AutoTokenizer.from_pretrained(HF_MODEL, token=HF_API_TOKEN)
         model = AutoModelForCausalLM.from_pretrained(HF_MODEL, token=HF_API_TOKEN, device_map="auto")
         logger.info(f"Initialized local Hugging Face model: {HF_MODEL}")
-        return (model, tokenizer), "hf_local"
+        return (model, tokenizer), "hf_local", HF_MODEL
     except Exception as e:
         logger.error(f"Failed to initialize local HF model: {e}")
         raise Exception("No LLM could be initialized")
 
-llm_client, llm_type = initialize_llm()
+llm_client, llm_type, llm_model = initialize_llm()
 
 # Initialize embedder
+_embedder = None
+
+def get_embedder():
+    global _embedder
+    if _embedder is None:
+        try:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            _embedder = SentenceTransformer(
+                "all-MiniLM-L6-v2",
+                device=device,
+                cache_folder="./cache"
+            )
+            logger.info(f"SentenceTransformer initialized on {device.upper()}")
+        except Exception as e:
+            logger.error(f"Failed to initialize SentenceTransformer: {e}")
+            raise RuntimeError(f"Embedder initialization failed: {e}")
+    return _embedder
+
 try:
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
-    logger.info("Sentence transformer initialized")
+    embedder = get_embedder()
 except Exception as e:
     logger.error(f"Failed to initialize embedder: {e}")
     embedder = None
 
-# Download file with local fallback
-async def download_file(task_id: str, ext: str) -> str | None:
-    try:
-        url = f"{GAIA_FILE_URL}{task_id}.{ext}"
-        async with aiohttp.ClientSession() as session:
-            async with session.get(url, timeout=10) as resp:
-                logger.info(f"GAIA API test for task {task_id} with .{ext}: HTTP {resp.status}")
-                if resp.status == 200:
-                    os.makedirs("temp", exist_ok=True)
-                    file_path = f"temp/{task_id}.{ext}"
-                    with open(file_path, "wb") as f:
-                        f.write(await resp.read())
-                    return file_path
-    except Exception as e:
-        logger.warning(f"File download failed for {task_id}.{ext}: {e}")
-    local_path = f"temp/{task_id}.{ext}"
-    if os.path.exists(local_path):
-        logger.info(f"Using local file: {local_path}")
-        return local_path
-    return None
+# Log device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Using device: {device}")
+
+# HTTP session with SSL handling
+async def create_http_session():
+    ssl_context = ssl.create_default_context()
+    ssl_context.check_hostname = False
+    ssl_context.verify_mode = ssl.CERT_NONE
+    return aiohttp.ClientSession(
+        connector=aiohttp.TCPConnector(ssl=ssl_context),
+        timeout=aiohttp.ClientTimeout(total=30)
+    )
+
+# Tool registration
+tools = {
+    "search_tool": search_tool,
+    "multi_hop_search_tool": multi_hop_search_tool,
+    "file_parser_tool": file_parser_tool,
+    "image_parser_tool": image_parser_tool,
+    "calculator_tool": calculator_tool,
+    "document_retriever_tool": document_retriever_tool,
+    "duckduckgo_search_tool": duckduckgo_search_tool,
+    "weather_info_tool": weather_info_tool,
+    "hub_stats_tool": hub_stats_tool,
+    "guest_info_retriever_tool": guest_info_retriever_tool,
+}
 
 # Parse question to select tools
 async def parse_question(state: JARVISState) -> JARVISState:
+    """
+    Parse the question to select appropriate tools using LLM with retries, preprocess the question, and integrate file-based tools.
+    
+    Args:
+        state (JARVISState): The input state containing task_id, question.
+    
+    Returns:
+        JARVISState: Updated state with selected tools_needed and metadata.
+    """
+    state = validate_state(state)
+    task_id = state["task_id"]
+    question = state["question"]
+    
+    logger.info(f"Task {task_id} Parsing question: {question}")
     try:
-        question = state["question"]
-        task_id = state["task_id"]
+        # Preprocess question
+        processed_question = await preprocess_question(question)
+        if processed_question != question:
+            logger.info(f"Task {task_id} Preprocessed question: {processed_question}")
+            state["question"] = processed_question
+            question = processed_question
+        
+        # Default to search_tool
         tools_needed = ["search_tool"]
-
+        
+        # LLM-based tool selection
         if llm_client:
             prompt = ChatPromptTemplate.from_messages([
                 SystemMessage(content="""Select tools from: ['search_tool', 'multi_hop_search_tool', 'file_parser_tool', 'image_parser_tool', 'calculator_tool', 'document_retriever_tool', 'duckduckgo_search_tool', 'weather_info_tool', 'hub_stats_tool', 'guest_info_retriever_tool'].
-                Return JSON list, e.g., ["search_tool", "file_parser_tool"].
+
+                Return a JSON list of all relevant tools, e.g., ["search_tool", "duckduckgo_search_tool"].
+
                 Rules:
-                - Always include "search_tool" unless purely computational.
-                - Use "multi_hop_search_tool" for complex queries (over 20 words or requiring multiple steps).
-                - Use "file_parser_tool" for data, tables, or Excel.
-                - Use "image_parser_tool" for images/videos.
-                - Use "calculator_tool" for math calculations.
-                - Use "document_retriever_tool" for documents/PDFs.
-                - Use "duckduckgo_search_tool" for additional search capability.
-                - Use "weather_info_tool" for weather-related queries.
-                - Use "hub_stats_tool" for Hugging Face Hub queries.
-                - Use "guest_info_retriever_tool" for guest-related queries.
+                - Include "search_tool" for web-based questions unless purely computational or file-based.
+                - Include "multi_hop_search_tool" for questions with >20 words or requiring multiple steps.
+                - Include "file_parser_tool" for 'data', 'table', 'excel', 'csv', 'txt', 'mp3', or file extensions.
+                - Include "image_parser_tool" for 'image', 'video', 'picture', or 'painting'.
+                - Include "calculator_tool" for 'calculate', 'math', 'sum', 'average', 'total', or numerical operations.
+                - Include "document_retriever_tool" for 'document', 'pdf', 'report', or 'paper'.
+                - Include "duckduckgo_search_tool" for 'search', 'wikipedia', 'online', or general knowledge.
+                - Include "weather_info_tool" for 'weather', 'temperature', or 'forecast'.
+                - Include "hub_stats_tool" for 'model', 'huggingface', or 'dataset'.
+                - Include "guest_info_retriever_tool" for 'guest', 'name', 'relation', or 'person'.
+                - Select multiple tools if the question spans multiple domains (e.g., web and file).
                 - Output ONLY valid JSON."""),
                 HumanMessage(content=f"Query: {question}")
             ])
-            try:
-                if llm_type == "hf_local":
-                    model, tokenizer = llm_client
-                    inputs = tokenizer.apply_chat_template(
-                        [{"role": "system", "content": prompt[0].content}, {"role": "user", "content": prompt[1].content}],
-                        return_tensors="pt"
-                    ).to(model.device)
-                    outputs = model.generate(inputs, max_new_tokens=512, temperature=0.7)
-                    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                    tools_needed = json.loads(response.strip())
-                elif llm_type == "together":
-                    response = llm_client.chat.completions.create(
-                        model=llm_client.model,
-                        messages=[
-                            {"role": "system", "content": prompt[0].content},
-                            {"role": "user", "content": prompt[1].content}
-                        ],
-                        max_tokens=512,
-                        temperature=0.7
-                    )
-                    tools_needed = json.loads(response.choices[0].message.content.strip())
-                else:  # hf_api
-                    response = llm_client.chat.completions.create(
-                        model=HF_MODEL,
-                        messages=[
-                            {"role": "system", "content": prompt[0].content},
-                            {"role": "user", "content": prompt[1].content}
-                        ],
-                        max_tokens=512,
-                        temperature=0.7
-                    )
-                    tools_needed = json.loads(response.choices[0].message.content.strip())
-                
-                valid_tools = {
-                    "search_tool", "multi_hop_search_tool", "file_parser_tool", "image_parser_tool",
-                    "calculator_tool", "document_retriever_tool", "duckduckgo_search_tool",
-                    "weather_info_tool", "hub_stats_tool", "guest_info_retriever_tool"
-                }
-                tools_needed = [tool for tool in tools_needed if tool in valid_tools]
-            except Exception as e:
-                logger.warning(f"Task {task_id} tool selection failed: {e}")
-                state["error"] = f"Tool selection failed: {str(e)}"
-
-        # Keyword-based fallback
-        question_lower = question.lower()
-        if any(word in question_lower for word in ["image", "video", "picture"]):
-            tools_needed.append("image_parser_tool")
-        if any(word in question_lower for word in ["data", "table", "excel", ".txt", ".csv", ".xlsx"]):
-            tools_needed.append("file_parser_tool")
-        if any(word in question_lower for word in ["calculate", "math", "sum", "average", "total"]):
-            tools_needed.append("calculator_tool")
-        if any(word in question_lower for word in ["document", "pdf", "report", "menu"]):
-            tools_needed.append("document_retriever_tool")
-        if any(word in question_lower for word in ["weather", "temperature"]):
-            tools_needed.append("weather_info_tool")
-        if any(word in question_lower for word in ["model", "huggingface", "dataset"]):
-            tools_needed.append("hub_stats_tool")
-        if any(word in question_lower for word in ["guest", "name", "relation", "person"]):
-            tools_needed.append("guest_info_retriever_tool")
-        if len(question.split()) > 20 or "multiple" in question_lower:
-            tools_needed.append("multi_hop_search_tool")
-        if any(word in question_lower for word in ["search", "wikipedia", "online"]):
-            tools_needed.append("duckduckgo_search_tool")
-
-        # Check file availability
-        for ext in ["txt", "csv", "xlsx", "jpg", "pdf"]:
-            file_path = await download_file(task_id, ext)
-            if file_path:
-                if ext in ["txt", "csv", "xlsx"] and "file_parser_tool" not in tools_needed:
+            messages = prompt.format_messages()
+            
+            for attempt in range(3):  # Retry up to 3 times
+                try:
+                    formatted_messages = [
+                        {"role": "system" if isinstance(m, SystemMessage) else "user", "content": m.content}
+                        for m in messages
+                    ]
+                    if llm_type == "hf_local":
+                        model, tokenizer = llm_client
+                        inputs = tokenizer.apply_chat_template(
+                            formatted_messages,
+                            return_tensors="pt"
+                        ).to(model.device)
+                        outputs = model.generate(inputs, max_new_tokens=100, temperature=0.5)
+                        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+                    elif llm_type == "together":
+                        response = llm_client.chat.completions.create(
+                            model=llm_model,
+                            messages=formatted_messages,
+                            max_tokens=100,
+                            temperature=0.5
+                        )
+                        response = response.choices[0].message.content.strip()
+                    else:  # hf_api
+                        response = llm_client.chat.completions.create(
+                            messages=formatted_messages,
+                            max_tokens=100,
+                            temperature=0.5
+                        )
+                        response = response.choices[0].message.content.strip()
+                    
+                    logger.info(f"Task {task_id} LLM tool selection response: {response}")
+                    try:
+                        tools_needed = json.loads(response)
+                        if isinstance(tools_needed, list) and all(isinstance(t, str) and t in tools for t in tools_needed):
+                            break  # Valid response, exit retry loop
+                        else:
+                            raise ValueError("Invalid tool list format")
+                    except json.JSONDecodeError as e:
+                        logger.warning(f"Task {task_id}: Invalid JSON (attempt {attempt + 1}): {e}")
+                        if attempt == 2:
+                            tools_needed = ["search_tool"]  # Fallback after retries
+                except Exception as e:
+                    logger.warning(f"Task {task_id} Tool selection failed (attempt {attempt + 1}): {e}")
+                    if attempt == 2:
+                        tools_needed = ["search_tool"]  # Fallback after retries
+        
+        # Fallback to keyword-based selection if LLM fails
+        if tools_needed == ["search_tool"] and not any(kw in question.lower() for kw in ["calculate", "math", "image", "document", "file", "weather", "guest", "model"]):
+            question_lower = question.lower()
+            if any(kw in question_lower for kw in ["excel", "csv", "mp3", "data", "table", "xlsx"]):
+                tools_needed.append("file_parser_tool")
+            if any(kw in question_lower for kw in ["image", "video", "picture", "painting"]):
+                tools_needed.append("image_parser_tool")
+            if any(kw in question_lower for kw in ["calculate", "math", "sum", "average", "total"]):
+                tools_needed.append("calculator_tool")
+            if any(kw in question_lower for kw in ["document", "pdf", "report", "paper"]):
+                tools_needed.append("document_retriever_tool")
+            if any(kw in question_lower for kw in ["search", "wikipedia", "online"]):
+                tools_needed.append("duckduckgo_search_tool")
+            if any(kw in question_lower for kw in ["weather", "temperature", "forecast"]):
+                tools_needed.append("weather_info_tool")
+            if any(kw in question_lower for kw in ["model", "huggingface", "dataset"]):
+                tools_needed.append("hub_stats_tool")
+            if any(kw in question_lower for kw in ["guest", "name", "relation", "person"]):
+                tools_needed.append("guest_info_retriever_tool")
+            if len(question.split()) > 20 or "multiple" in question_lower:
+                tools_needed.append("multi_hop_search_tool")
+        
+        # Integrate file-based tools
+        file_results = await fetch_task_file(task_id, question)
+        for ext, content in file_results.items():
+            if content:
+                os.makedirs("temp", exist_ok=True)
+                file_path = f"temp/{task_id}.{ext}"
+                with open(file_path, "wb") as f:
+                    f.write(content)
+                state["metadata"] = state.get("metadata", {}) | {"file_ext": ext, "file_path": file_path}
+                if ext in ["txt", "csv", "xlsx", "mp3"] and "file_parser_tool" not in tools_needed:
                     tools_needed.append("file_parser_tool")
-                if ext == "jpg" and "image_parser_tool" not in tools_needed:
+                elif ext in ["jpg", "png"] and "image_parser_tool" not in tools_needed:
                     tools_needed.append("image_parser_tool")
-                if ext == "pdf" and "document_retriever_tool" not in tools_needed:
+                elif ext == "pdf" and "document_retriever_tool" not in tools_needed:
                     tools_needed.append("document_retriever_tool")
-                state["metadata"] = state.get("metadata", {}) | {"file_ext": ext, "file_path": file_path}
-                break
-
-        state["tools_needed"] = list(set(tools_needed))
-        logger.info(f"Task {task_id}: Selected tools: {tools_needed}")
+        
+        state["tools_needed"] = list(set(tools_needed))  # Remove duplicates
+        logger.info(f"Task {task_id} Selected tools: {state['tools_needed']}")
         return state
     except Exception as e:
-        logger.error(f"Error parsing task {task_id}: {e}")
+        logger.error(f"Task {task_id} Tool selection failed: {e}")
         state["error"] = f"Parse question failed: {str(e)}"
         state["tools_needed"] = ["search_tool"]
         return state
 
 # Tool dispatcher
 async def tool_dispatcher(state: JARVISState) -> JARVISState:
+    state = validate_state(state)
     try:
-        updated_state = state.copy()
-        file_type = "jpg" if "image" in state["question"].lower() else "txt"
-        if any(word in state["question"].lower() for word in ["menu", "report"]):
-            file_type = "pdf"
-        elif "data" in state["question"].lower():
-            file_type = "xlsx"
-
-        for tool in updated_state["tools_needed"]:
-            try:
-                if tool == "search_tool":
-                    result = search_tool(updated_state["question"])
-                    updated_state["web_results"].extend([str(r) for r in result])
-                elif tool == "multi_hop_search_tool":
-                    result = await multi_hop_search_tool.ainvoke({"query": updated_state["question"], "steps": 3, "llm_client": llm_client, "llm_type": llm_type})
-                    updated_state["multi_hop_results"].extend([r["content"] for r in result])
-                    await asyncio.sleep(2)
-                elif tool == "file_parser_tool":
-                    for ext in ["txt", "csv", "xlsx"]:
-                        file_path = await download_file(updated_state["task_id"], ext)
-                        if file_path:
-                            result = file_parser_tool(file_path)
-                            updated_state["file_results"] = str(result)
-                            break
-                elif tool == "image_parser_tool":
-                    file_path = await download_file(updated_state["task_id"], "jpg")
-                    if file_path:
-                        result = image_parser_tool(file_path)
-                        updated_state["image_results"] = str(result)
-                elif tool == "calculator_tool":
-                    result = calculator_tool(updated_state["question"])
-                    updated_state["calculation_results"] = str(result)
-                elif tool == "document_retriever_tool":
-                    file_path = await download_file(updated_state["task_id"], "pdf")
-                    if file_path:
-                        result = document_retriever_tool({"task_id": updated_state["task_id"], "query": updated_state["question"], "file_type": "pdf"})
-                        updated_state["document_results"] = str(result)
-                elif tool == "duckduckgo_search_tool":
-                    result = duckduckgo_search_tool(updated_state["question"])
-                    updated_state["web_results"].append(str(result))
-                elif tool == "weather_info_tool":
-                    location = updated_state["question"].split("weather in ")[1].split()[0] if "weather in" in updated_state["question"].lower() else "Unknown"
-                    result = weather_info_tool({"location": location})
-                    updated_state["web_results"].append(str(result))
-                elif tool == "hub_stats_tool":
-                    author = updated_state["question"].split("by ")[1].split()[0] if "by" in updated_state["question"].lower() else "Unknown"
-                    result = hub_stats_tool({"author": author})
-                    updated_state["web_results"].append(str(result))
-                elif tool == "guest_info_retriever_tool":
-                    query = updated_state["question"].split("about ")[1] if "about" in updated_state["question"].lower() else updated_state["question"]
-                    result = guest_info_retriever_tool({"query": query})
-                    updated_state["web_results"].append(str(result))
-                updated_state["metadata"] = updated_state.get("metadata", {}) | {f"{tool}_executed": True}
-            except Exception as e:
-                logger.warning(f"Error in tool {tool} for task {updated_state['task_id']}: {str(e)}")
-                updated_state["error"] = f"Tool {tool} failed: {str(e)}"
-                updated_state["metadata"] = updated_state.get("metadata", {}) | {f"{tool}_error": str(e)}
-
-        logger.info(f"Task {updated_state['task_id']}: Tool results: {updated_state}")
-        return updated_state
-    except Exception as e:
-        logger.error(f"Tool dispatch failed for task {state['task_id']}: {e}")
-        updated_state["error"] = f"Tool dispatch failed: {str(e)}"
-        return updated_state
-
-# Reasoning
-async def reasoning(state: JARVISState) -> Dict[str, Any]:
-    try:
-        prompt = ChatPromptTemplate.from_messages([
-            SystemMessage(content="""Provide ONLY the exact answer (e.g., '90', 'HUE'). For USD, use two decimal places (e.g., '1234.00'). For lists, use comma-separated values (e.g., 'Smith, Lee'). For IOC codes, use three-letter codes (e.g., 'ARG'). No explanations or conversational text."""),
-            HumanMessage(content="""Task: {task_id}
-Question: {question}
-Web results: {web_results}
-Multi-hop results: {multi_hop_results}
-File results: {file_results}
-Image results: {image_results}
-Calculation results: {calculation_results}
-Document results: {document_results}""")
-        ])
-        messages = [
-            {"role": "system", "content": prompt[0].content},
-            {"role": "user", "content": prompt[1].content.format(
-                task_id=state["task_id"],
-                question=state["question"],
-                web_results="\n".join(state["web_results"]),
-                multi_hop_results="\n".join(state["multi_hop_results"]),
-                file_results=state["file_results"],
-                image_results=state["image_results"],
-                calculation_results=state["calculation_results"],
-                document_results=state["document_results"]
-            )}
-        ]
-        for attempt in range(3):
+        task_id = state["task_id"]
+        question = state["question"]
+        tools_needed = state["tools_needed"]
+        
+        for tool_name in tools_needed:
             try:
-                if llm_type == "hf_local":
-                    model, tokenizer = llm_client
-                    inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
-                    outputs = model.generate(inputs, max_new_tokens=512, temperature=0.7)
-                    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                elif llm_type == "together":
-                    response = llm_client.chat.completions.create(
-                        model=llm_client.model,
-                        messages=messages,
-                        max_tokens=512,
-                        temperature=0.7
-                    )
-                    answer = response.choices[0].message.content.strip()
-                else:  # hf_api
-                    response = llm_client.chat.completions.create(
-                        model=HF_MODEL,
-                        messages=messages,
-                        max_tokens=512,
-                        temperature=0.7
-                    )
-                    answer = response.choices[0].message.content.strip()
-                
-                # Format answer
-                if "USD" in state["question"].lower():
-                    try:
-                        answer = f"{float(answer):.2f}"
-                    except ValueError:
-                        pass
-                if "before and after" in state["question"].lower():
-                    answer = answer.replace(" and ", ", ")
-                if "IOC code" in state["question"].lower():
-                    answer = answer.upper()[:3]
+                if tool_name == "search_tool":
+                    result = await tools["search_tool"].ainvoke({"query": question})
+                    state["web_results"].extend([str(r) for r in result] if result else ["No results from search_tool"])
+                elif tool_name == "multi_hop_search_tool":
+                    result = await tools["multi_hop_search_tool"].ainvoke({
+                        "query": question,
+                        "steps": 3,
+                        "llm_client": llm_client,
+                        "llm_type": llm_type,
+                        "llm_model": llm_model
+                    })
+                    state["multi_hop_results"].extend([r["content"] if isinstance(r, dict) else str(r) for r in result] if result else ["No results from multi_hop_search_tool"])
+                elif tool_name == "file_parser_tool":
+                    file_path = state["metadata"].get("file_path")
+                    file_ext = state["metadata"].get("file_ext")
+                    if file_path and os.path.exists(file_path) and file_ext:
+                        result = await tools["file_parser_tool"].ainvoke({
+                            "task_id": task_id,
+                            "file_type": file_ext,
+                            "file_path": file_path,
+                            "query": question
+                        })
+                        state["file_results"] = str(result) if result else "No file results"
+                    else:
+                        state["file_results"] = "No file available"
+                elif tool_name == "image_parser_tool":
+                    file_path = state["metadata"].get("file_path")
+                    if file_path and os.path.exists(file_path) and file_path.split('.')[-1] in ["jpg", "png"]:
+                        result = await tools["image_parser_tool"].ainvoke({"task_id": task_id, "file_path": file_path})
+                        state["image_results"] = str(result) if result else "No image results"
+                    else:
+                        state["image_results"] = "No image available"
+                elif tool_name == "calculator_tool":
+                    result = await tools["calculator_tool"].ainvoke({"expression": question})
+                    state["calculation_results"] = str(result) if result else "No calculation results"
+                elif tool_name == "document_retriever_tool":
+                    file_path = state["metadata"].get("file_path")
+                    if file_path and os.path.exists(file_path) and file_path.split('.')[-1] == "pdf":
+                        result = await tools["document_retriever_tool"].ainvoke({
+                            "task_id": task_id,
+                            "query": question,
+                            "file_path": file_path
+                        })
+                        state["document_results"] = str(result) if result else "No document results"
+                    else:
+                        state["document_results"] = "No document available"
+                elif tool_name == "duckduckgo_search_tool":
+                    result = await tools["duckduckgo_search_tool"].ainvoke({
+                        "query": question,
+                        "original_query": question,
+                        "embedder": embedder
+                    })
+                    state["web_results"].extend(result if isinstance(result, list) else [str(result)] if result else ["No results from duckduckgo_search_tool"])
+                elif tool_name == "weather_info_tool":
+                    location = question.split()[-1] if "weather" in question.lower() else "Unknown"
+                    result = await tools["weather_info_tool"].ainvoke({"location": location})
+                    state["web_results"].append(str(result) if result else "No weather results")
+                elif tool_name == "hub_stats_tool":
+                    author = question.split("by ")[1].split()[0] if "by" in question.lower() else "Unknown"
+                    result = await tools["hub_stats_tool"].ainvoke({"author": author})
+                    state["web_results"].append(str(result) if result else "No hub stats results")
+                elif tool_name == "guest_info_retriever_tool":
+                    result = await tools["guest_info_retriever_tool"].ainvoke({"query": question})
+                    state["web_results"].append(str(result) if result else "No guest info results")
                 
-                logger.info(f"Task {state['task_id']}: Answer: {answer}")
-                return {"answer": answer}
+                state["metadata"] = state.get("metadata", {}) | {f"{tool_name}_executed": True}
+                logger.info(f"Task {task_id}: Executed {tool_name}")
             except Exception as e:
-                logger.warning(f"LLM retry {attempt + 1}/3 for task {state['task_id']}: {e}")
-                await asyncio.sleep(2)
-        state["error"] = "LLM failed after retries"
-        return {"answer": "Error: LLM failed after retries"}
+                logger.warning(f"Tool {tool_name} failed for task {task_id}: {e}")
+                state["metadata"] = state.get("metadata", {}) | {f"{tool_name}_error": str(e)}
+        
+        # Ensure results are populated
+        state["web_results"] = state.get("web_results", ["No web results found"])
+        state["file_results"] = state.get("file_results", "No file results found")
+        state["image_results"] = state.get("image_results", "No image results found")
+        state["document_results"] = state.get("document_results", "No document results found")
+        state["calculation_results"] = state.get("calculation_results", "No calculation results found")
+        
+        state["answer"] = await generate_answer(
+            task_id=task_id,
+            question=question,
+            search_results=state.get("web_results", []) + [
+                r["content"] if isinstance(r, dict) else str(r) for r in state.get("multi_hop_results", [])
+            ],
+            file_results=state.get("file_results", "") + state.get("document_results", "") + state.get("image_results", "") + state.get("calculation_results", ""),
+            llm_client=llm_client
+        )
+        
+        logger.info(f"Task {task_id}: Generated answer: {state['answer']}")
+        return state
     except Exception as e:
-        logger.error(f"Reasoning failed for task {state['task_id']}: {e}")
-        state["error"] = f"Reasoning failed: {str(e)}"
-        return {"answer": f"Error: {str(e)}"}
-
-# Router
-def router(state: JARVISState) -> str:
-    if state["tools_needed"]:
-        return "tool_dispatcher"
-    return "reasoning"
+        logger.error(f"Tool dispatch failed: {e}")
+        state["error"] = f"Tool dispatch failed: {e}"
+        return state
 
 # Define StateGraph
 workflow = StateGraph(JARVISState)
-workflow.add_node("parse", parse_question)
+workflow.add_node("parse_question", parse_question)
 workflow.add_node("tool_dispatcher", tool_dispatcher)
-workflow.add_node("reasoning", reasoning)
-workflow.set_entry_point("parse")
-workflow.add_conditional_edges(
-    "parse",
-    router,
-    {
-        "tool_dispatcher": "tool_dispatcher",
-        "reasoning": "reasoning"
-    }
-)
-workflow.add_edge("tool_dispatcher", "reasoning")
-workflow.add_edge("reasoning", END)
+workflow.set_entry_point("parse_question")
+workflow.add_edge("parse_question", "tool_dispatcher")
+workflow.add_edge("tool_dispatcher", END)
 graph = workflow.compile()
 
 # Agent class
 class JARVISAgent:
     def __init__(self):
-        self.state = JARVISState(
-            task_id="",
-            question="",
-            tools_needed=[],
-            web_results=[],
-            file_results="",
-            image_results="",
-            calculation_results="",
-            document_results="",
-            multi_hop_results=[],
-            messages=[],
-            answer="",
-            results_table=[],
-            status_output="",
-            error=None,
-            metadata={}
-        )
+        self.state = reset_state(task_id="init", question="Agent initialized")
+        self.state["results_table"] = []  # Initialize as empty list
         logger.info("JARVISAgent initialized.")
 
     async def process_question(self, task_id: str, question: str) -> str:
-        state = JARVISState(
-            task_id=task_id,
-            question=question,
-            tools_needed=["search_tool"],
-            web_results=[],
-            file_results="",
-            image_results="",
-            calculation_results="",
-            document_results="",
-            multi_hop_results=[],
-            messages=[HumanMessage(content=question)],
-            answer="",
-            results_table=[],
-            status_output="",
-            error=None,
-            metadata={}
-        )
+        state = reset_state(task_id=task_id, question=question)
         try:
             result = await graph.ainvoke(state)
-            answer = result["answer"] or "Unknown"
-            logger.info(f"Task {task_id}: Final answer: {answer}")
-            self.state.results_table.append({"Task ID": task_id, "Question": question, "Answer": answer})
-            self.state.metadata = self.state.get("metadata", {}) | {"last_task": task_id, "answer": answer}
+            answer = result.get("answer", "Unknown")
+            logger.info(f"Task {task_id} Final answer: {answer}")
+            self.state["results_table"].append({"Task ID": task_id, "Question": question, "Answer": answer})
+            self.state["metadata"] = {"last_task_id": task_id, "answer": answer}
             return answer
         except Exception as e:
             logger.error(f"Error processing task {task_id}: {e}")
-            self.state.results_table.append({"Task ID": task_id, "Question": question, "Answer": f"Error: {e}"})
-            self.state.error = f"Task {task_id} failed: {str(e)}"
+            self.state["results_table"].append({"Task ID": task_id, "Question": question, "Answer": f"Error: {e}"})
+            self.state["error"] = f"Task {task_id} failed: {str(e)}"
             return f"Error: {str(e)}"
         finally:
-            for ext in ["txt", "csv", "xlsx", "jpg", "pdf"]:
+            for ext in ["txt", "csv", "xlsx", "mp3", "jpg", "png", "pdf"]:
                 file_path = f"temp/{task_id}.{ext}"
                 if os.path.exists(file_path):
                     try:
@@ -466,25 +447,26 @@ class JARVISAgent:
     async def process_all_questions(self, profile: gr.OAuthProfile | None):
         if not profile:
             logger.error("User not logged in.")
-            self.state.status_output = "Please Login to Hugging Face."
-            return pd.DataFrame(self.state.results_table), self.state.status_output
+            self.state["status_output"] = "Please Login to Hugging Face."
+            return pd.DataFrame(self.state["results_table"]), self.state["status_output"]
 
-        username = f"{profile.username}"
+        username = profile.username
         logger.info(f"User logged in: {username}")
         questions_url = f"{GAIA_API_URL}/questions"
         submit_url = f"{GAIA_API_URL}/submit"
         agent_code = f"https://huggingface.co/spaces/{SPACE_ID}/tree/main"
 
         try:
-            response = requests.get(questions_url, timeout=15)
-            response.raise_for_status()
-            questions = response.json()
+            async with await create_http_session() as session:
+                async with session.get(questions_url) as response:
+                    response.raise_for_status()
+                    questions = await response.json()
             logger.info(f"Fetched {len(questions)} questions.")
         except Exception as e:
             logger.error(f"Error fetching questions: {e}")
-            self.state.status_output = f"Error fetching questions: {e}"
-            self.state.error = f"Fetch questions failed: {str(e)}"
-            return pd.DataFrame(self.state.results_table), self.state.status_output
+            self.state["status_output"] = f"Error fetching questions: {e}"
+            self.state["error"] = f"Fetch questions failed: {str(e)}"
+            return pd.DataFrame(self.state["results_table"]), self.state["status_output"]
 
         answers_payload = []
         for item in questions:
@@ -498,33 +480,34 @@ class JARVISAgent:
 
         if not answers_payload:
             logger.error("No answers generated.")
-            self.state.status_output = "No answers to submit."
-            self.state.error = "No answers generated"
-            return pd.DataFrame(self.state.results_table), self.state.status_output
+            self.state["status_output"] = "No answers to submit."
+            self.state["error"] = "No answers generated"
+            return pd.DataFrame(self.state["results_table"]), self.state["status_output"]
 
         submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
         try:
-            response = requests.post(submit_url, json=submission_data, timeout=120)
-            response.raise_for_status()
-            result_data = response.json()
-            self.state.status_output = (
+            async with await create_http_session() as session:
+                async with session.post(submit_url, json=submission_data) as response:
+                    response.raise_for_status()
+                    result_data = await response.json()
+            self.state["status_output"] = (
                 f"Submission Successful!\n"
                 f"User: {result_data.get('username')}\n"
                 f"Overall Score: {result_data.get('score', 'N/A')}% "
                 f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
                 f"Message: {result_data.get('message', 'No message received.')}"
             )
-            self.state.metadata = self.state.get("metadata", {}) | {"submission_score": result_data.get('score', 'N/A')}
+            self.state["metadata"] = self.state.get("metadata", {}) | {"submission_score": result_data.get('score', 'N/A')}
         except Exception as e:
             logger.error(f"Submission failed: {e}")
-            self.state.status_output = f"Submission Failed: {e}"
-            self.state.error = f"Submission failed: {str(e)}"
+            self.state["status_output"] = f"Submission Failed: {e}"
+            self.state["error"] = f"Submission failed: {str(e)}"
 
-        return pd.DataFrame(self.state.results_table), self.state.status_output
+        return pd.DataFrame(self.state["results_table"] if self.state["results_table"] else [], columns=["Task ID", "Question", "Answer"]), self.state["status_output"]
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Evolved JARVIS GAIA Agent")
+    gr.Markdown("# JARVIS GAIA Agent")
     gr.Markdown(
         """
         **Instructions:**
@@ -539,7 +522,6 @@ with gr.Blocks() as demo:
     )
     with gr.Row():
         gr.LoginButton(value="Login to Hugging Face")
-        # Removed gr.LogoutButton due to deprecation
     run_button = gr.Button("Run Evaluation & Submit All Answers")
     status_output = gr.Textbox(label="Run Status / Submission Result", lines=5, interactive=False)
     results_table = gr.DataFrame(label="Questions and Answers", wrap=True, headers=["Task ID", "Question", "Answer"])

project_structure.txt

@@ -0,0 +1,22 @@
+.
+├── app.py
+├── dockerfile
+├── README.md
+├── requirements.txt
+├── retriever.py
+├── state.py
+└── tools
+    ├── __init__.py
+    ├── answer_generator.py
+    ├── calculator.py
+    ├── document_retriever.py
+    ├── duckduckgo_search.py
+    ├── file_fetcher.py
+    ├── file_parser.py
+    ├── guest_info.py
+    ├── hub_stats.py
+    ├── image_parser.py
+    ├── search.py
+    └── weather_info.py
+
+3 directories, 18 files

requirements.txt

@@ -20,8 +20,11 @@ transformers
 asyncio
 serpapi
 duckduckgo-search
-torch
+torch==2.2.2
 together
 google-search-results
 beautifulsoup4
-gradio[oauth]
+gradio[oauth]
+nlkt
+speechrecognition
+rank_bm25

retriever.py

@@ -1,25 +1,109 @@
-import datasets
-from langchain.docstore.document import Document
-from langchain_community.retrievers import BM25Retriever
-from smolagents import Tool
+import json
+import os
+import logging
+import torch
+from typing import List
+from langchain_core.documents import Document
+from sentence_transformers import SentenceTransformer
+try:
+    from datasets import load_dataset
+except ImportError:
+    load_dataset = None
 
-def load_guest_dataset():
+logger = logging.getLogger(__name__)
+
+def get_device():
+    """
+    Determine the appropriate device for PyTorch.
+    
+    Returns:
+        str: Device name ('cuda', 'mps', or 'cpu').
+    """
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+def load_guest_dataset(dataset_path: str = "agents-course/unit3-invitees") -> List[Document]:
+    """
+    Load guest dataset from a local JSON file or Hugging Face dataset.
+    
+    Args:
+        dataset_path (str): Path to local JSON file or Hugging Face dataset name.
+    
+    Returns:
+        List[Document]: List of Document objects with guest information.
+    """
     try:
-        guest_dataset = datasets.load_dataset("agents-course/unit3-invitees", split="train")
+        # Try loading from Hugging Face dataset if datasets library is available
+        if load_dataset and not os.path.exists(dataset_path):
+            logger.info(f"Attempting to load Hugging Face dataset: {dataset_path}")
+            guest_dataset = load_dataset(dataset_path, split="train")
+            docs = [
+                Document(
+                    page_content="\n".join([
+                        f"Name: {guest['name']}",
+                        f"Relation: {guest['relation']}",
+                        f"Description: {guest['description']}",
+                        f"Email: {guest['email']}"
+                    ]),
+                    metadata={
+                        "name": guest["name"],
+                        "relation": guest["relation"],
+                        "description": guest["description"],
+                        "email": guest["email"]
+                    }
+                )
+                for guest in guest_dataset
+            ]
+            logger.info(f"Loaded {len(docs)} guests from Hugging Face dataset")
+            return docs
+
+        # Try loading from local JSON file
+        if os.path.exists(dataset_path):
+            logger.info(f"Loading guest dataset from local path: {dataset_path}")
+            with open(dataset_path, 'r') as f:
+                guests = json.load(f)
+            docs = [
+                Document(
+                    page_content=guest.get('description', ''),
+                    metadata={
+                        'name': guest.get('name', ''),
+                        'relation': guest.get('relation', ''),
+                        'description': guest.get('description', ''),
+                        'email': guest.get('email', '')  # Optional email field
+                    }
+                )
+                for guest in guests
+            ]
+            logger.info(f"Loaded {len(docs)} guests from local JSON")
+            return docs
+
+        # Fallback to mock dataset if both fail
+        logger.warning(f"Dataset not found at {dataset_path}, using mock dataset")
         docs = [
             Document(
                 page_content="\n".join([
-                    f"Name: {guest['name']}",
-                    f"Relation: {guest['relation']}",
-                    f"Description: {guest['description']}",
-                    f"Email: {guest['email']}"
+                    "Name: Dr. Nikola Tesla",
+                    "Relation: old friend from university days",
+                    "Description: Dr. Nikola Tesla is an old friend from your university days. He's recently patented a new wireless energy transmission system.",
+                    "Email: [email protected]"
                 ]),
-                metadata={"name": guest["name"]}
+                metadata={
+                    "name": "Dr. Nikola Tesla",
+                    "relation": "old friend from university days",
+                    "description": "Dr. Nikola Tesla is an old friend from your university days. He's recently patented a new wireless energy transmission system.",
+                    "email": "[email protected]"
+                }
             )
-            for guest in guest_dataset
         ]
+        logger.info("Loaded mock dataset with 1 guest")
+        return docs
+
     except Exception as e:
-        # Fallback mock dataset
+        logger.error(f"Failed to load guest dataset: {e}")
+        # Return mock dataset as final fallback
         docs = [
             Document(
                 page_content="\n".join([
@@ -28,7 +112,73 @@ def load_guest_dataset():
                     "Description: Dr. Nikola Tesla is an old friend from your university days. He's recently patented a new wireless energy transmission system.",
                     "Email: [email protected]"
                 ]),
-                metadata={"name": "Dr. Nikola Tesla"}
+                metadata={
+                    "name": "Dr. Nikola Tesla",
+                    "relation": "old friend from university days",
+                    "description": "Dr. Nikola Tesla is an old friend from your university days. He's recently patented a new wireless energy transmission system.",
+                    "email": "[email protected]"
+                }
             )
         ]
-    return docs
+        logger.info("Loaded mock dataset with 1 guest due to error")
+        return docs
+
+class BM25Retriever:
+    """
+    A retriever class using SentenceTransformer for embedding-based search.
+    """
+    def __init__(self, dataset_path: str):
+        """
+        Initialize the retriever with a SentenceTransformer model.
+        
+        Args:
+            dataset_path (str): Path to the dataset for retrieval.
+        
+        Raises:
+            Exception: If embedder initialization fails.
+        """
+        try:
+            self.model = SentenceTransformer("all-MiniLM-L6-v2", device=get_device())
+            self.dataset_path = dataset_path
+            logger.info("Initialized SentenceTransformer")
+        except Exception as e:
+            logger.error(f"Failed to initialize embedder: {e}")
+            raise
+
+    def search(self, query: str) -> List[dict]:
+        """
+        Search the dataset for relevant guest information.
+        
+        Args:
+            query (str): Search query (e.g., guest name or relation).
+        
+        Returns:
+            List[dict]: List of matching guest metadata dictionaries.
+        """
+        try:
+            # Load dataset
+            docs = load_guest_dataset(self.dataset_path)
+            if not docs:
+                logger.warning("No documents available for search")
+                return []
+
+            # Convert documents to text for BM25 (using metadata for consistency)
+            texts = [f"{doc.metadata['name']} {doc.metadata['relation']} {doc.metadata['description']}" for doc in docs]
+            from langchain_community.retrievers import BM25Retriever
+            retriever = BM25Retriever.from_texts(texts)
+            retriever.k = 3  # Limit to top 3 results
+
+            # Perform search
+            results = retriever.invoke(query)
+            # Map results back to original metadata
+            matches = [
+                docs[i].metadata
+                for i in range(len(docs))
+                if any(f"{docs[i].metadata['name']} {docs[i].metadata['relation']} {docs[i].metadata['description']}" in r.page_content for r in results)
+            ]
+            logger.info(f"Found {len(matches)} matches for query: {query}")
+            return matches[:3]  # Return top 3 matches
+
+        except Exception as e:
+            logger.error(f"Search failed for query '{query}': {e}")
+            return []

state.py

@@ -1,5 +1,8 @@
-from typing import TypedDict, List, Dict, Optional, Any
+from typing import TypedDict, List, Dict, Optional, Any, Union
 from langchain_core.messages import BaseMessage
+import logging
+
+logger = logging.getLogger(__name__)
 
 class JARVISState(TypedDict):
     """
@@ -10,11 +13,11 @@ class JARVISState(TypedDict):
         question: The question text to be answered.
         tools_needed: List of tool names to be used for the task.
         web_results: List of web search results (e.g., from SERPAPI, DuckDuckGo).
-        file_results: Parsed content from text, CSV, or Excel files.
+        file_results: Parsed content from text, CSV, Excel, or audio files.
         image_results: OCR or description results from image files.
         calculation_results: Results from mathematical calculations.
-        document_results: Extracted content from PDF documents.
-        multi_hop_results: Results from iterative multi-hop searches.
+        document_results: Extracted content from PDF or text documents.
+        multi_hop_results: Results from iterative multi-hop searches (supports strings or dicts).
         messages: List of messages for LLM context (e.g., user prompts, system instructions).
         answer: Final answer for the task, formatted for GAIA submission.
         results_table: List of task results for Gradio display (Task ID, Question, Answer).
@@ -30,10 +33,84 @@ class JARVISState(TypedDict):
     image_results: str
     calculation_results: str
     document_results: str
-    multi_hop_results: List[str]
+    multi_hop_results: List[Union[str, Dict[str, Any]]]
     messages: List[BaseMessage]
     answer: str
     results_table: List[Dict[str, str]]
     status_output: str
     error: Optional[str]
-    metadata: Optional[Dict[str, Any]]
+    metadata: Optional[Dict[str, Any]]
+
+def validate_state(state: JARVISState) -> JARVISState:
+    """
+    Validate and initialize JARVISState fields.
+    
+    Args:
+        state: Input state dictionary.
+    
+    Returns:
+        Validated and initialized state.
+    """
+    try:
+        if not state.get("task_id"):
+            logger.error("task_id is required")
+            raise ValueError("task_id is required")
+        if not state.get("question"):
+            logger.error("question is required")
+            raise ValueError("question is required")
+        
+        # Initialize default values if missing
+        defaults = {
+            "tools_needed": ["search_tool"],
+            "web_results": [],
+            "file_results": "",
+            "image_results": "",
+            "calculation_results": "",
+            "document_results": "",
+            "multi_hop_results": [],
+            "messages": [],
+            "answer": "",
+            "results_table": [],
+            "status_output": "",
+            "error": None,
+            "metadata": {}
+        }
+        for key, default in defaults.items():
+            if key not in state or state[key] is None:
+                state[key] = default
+        
+        logger.debug(f"Validated state for task {state['task_id']}")
+        return state
+    except Exception as e:
+        logger.error(f"State validation failed: {e}")
+        raise
+
+def reset_state(task_id: str, question: str) -> JARVISState:
+    """
+    Create a fresh JARVISState for a new task.
+    
+    Args:
+        task_id: Task identifier.
+        question: Question text.
+    
+    Returns:
+        Initialized JARVISState.
+    """
+    state = JARVISState(
+        task_id=task_id,
+        question=question,
+        tools_needed=["search_tool"],
+        web_results=[],
+        file_results="",
+        image_results="",
+        calculation_results="",
+        document_results="",
+        multi_hop_results=[],
+        messages=[],
+        answer="",
+        results_table=[],
+        status_output="",
+        error=None,
+        metadata={}
+    )
+    return validate_state(state)

test.py

@@ -1,10 +1,233 @@
-from serpapi import GoogleSearch
+import asyncio
+import os
+import logging
+import tempfile
+from pathlib import Path
+from app import JARVISAgent, llm_client, llm_type, llm_model, embedder
+from tools.search import search_tool, multi_hop_search_tool
+from tools.file_parser import file_parser_tool
+from tools.image_parser import image_parser_tool
+from tools.calculator import calculator_tool
+from tools.document_retriever import document_retriever_tool
+from tools.duckduckgo_search import duckduckgo_search_tool
+from tools.weather_info import weather_info_tool
+from tools.hub_stats import hub_stats_tool
+from tools.guest_info import guest_info_retriever_tool
+from tools.file_fetcher import fetch_task_file
+from tools.answer_generator import preprocess_question, filter_results
+from state import validate_state, reset_state, JARVISState
 
-params = {
-  "q": "drop shipping",
-  "api_key": "e44c79583cac0e507fee32d564f190b7290a313d886edd5ba5fccc93df932733"
-}
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-search = GoogleSearch(params)
-results = search.get_dict()
-ai_overview = results["ai_overview"]
+async def test_tools():
+    """Test all tools."""
+    logger.info("Testing Search Tool (SerpAPI)...")
+    try:
+        if not os.getenv("SERPAPI_API_KEY"):
+            logger.warning("Search Warning: SERPAPI_API_KEY not set")
+        else:
+            result = await search_tool.ainvoke({"query": "What is the capital of France?"})
+            logger.info(f"Search Result: {result}")
+    except Exception as e:
+        logger.error(f"Search Error: {e}")
+
+    logger.info("Testing Multi-Hop Search Tool...")
+    try:
+        result = await multi_hop_search_tool.ainvoke({
+            "query": "What is the population of France's capital?",
+            "steps": 2,
+            "llm_client": llm_client,
+            "llm_type": llm_type,
+            "llm_model": llm_model
+        })
+        logger.info(f"Multi-Hop Search Result: {result}")
+    except Exception as e:
+        logger.error(f"Multi-Hop Search Error: {e}")
+
+    logger.info("Testing DuckDuckGo Search Tool...")
+    try:
+        result = await duckduckgo_search_tool.ainvoke({
+            "query": "What is the capital of France?",
+            "original_query": "What is the capital of France?",
+            "embedder": embedder
+        })
+        logger.info(f"DuckDuckGo Result: {result}")
+    except Exception as e:
+        logger.error(f"DuckDuckGo Error: {e}")
+
+    logger.info("Testing Weather Info Tool...")
+    try:
+        if not os.getenv("OPENWEATHERMAP_API_KEY"):
+            logger.warning("Weather Warning: OPENWEATHERMAP_API_KEY not set")
+        else:
+            result = await weather_info_tool.ainvoke({"location": "London"})
+            logger.info(f"Weather Result: {result}")
+    except Exception as e:
+        logger.error(f"Weather Error: {e}")
+
+    logger.info("Testing Document Retriever Tool...")
+    try:
+        from PyPDF2 import PdfWriter
+        with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
+            writer = PdfWriter()
+            from PyPDF2.generic import NameObject, create_string_object
+            page = writer.add_blank_page(width=72, height=72)
+            page[NameObject("/Contents")] = create_string_object("Sample document content for testing.")
+            writer.write(tmp)
+            tmp_path = tmp.name
+        result = await document_retriever_tool.ainvoke({
+            "task_id": "test_task",
+            "query": "Sample question",
+            "file_path": tmp_path
+        })
+        logger.info(f"Document Retriever Result: {result}")
+        os.unlink(tmp_path)
+    except Exception as e:
+        logger.error(f"Document Retriever Error: {e}")
+
+    logger.info("Testing Image Parser Tool...")
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+            # Create a minimal PNG (1x1 pixel)
+            from PIL import Image
+            img = Image.new('RGB', (1, 1), color='white')
+            img.save(tmp.name, 'PNG')
+            tmp_path = tmp.name
+        result = await image_parser_tool.ainvoke({"task_id": "test_task", "file_path": tmp_path})
+        logger.info(f"Image Parser Result: {result}")
+        os.unlink(tmp_path)
+    except Exception as e:
+        logger.error(f"Image Parser Error: {e}")
+
+    logger.info("Testing File Parser Tool...")
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
+            tmp.write(b"Sample text file content")
+            tmp_path = tmp.name
+        result = await file_parser_tool.ainvoke({
+            "task_id": "test_task",
+            "file_type": "txt",
+            "file_path": tmp_path,
+            "query": "What is in the file?"
+        })
+        logger.info(f"File Parser Result: {result}")
+        os.unlink(tmp_path)
+    except Exception as e:
+        logger.error(f"File Parser Error: {e}")
+
+    logger.info("Testing Calculator Tool...")
+    try:
+        result = await calculator_tool.ainvoke({"expression": "2 + 2"})
+        logger.info(f"Calculator Result: {result}")
+    except Exception as e:
+        logger.error(f"Calculator Error: {e}")
+
+    logger.info("Testing Hub Stats Tool...")
+    try:
+        if not os.getenv("HUGGINGFACEHUB_API_TOKEN"):
+            logger.warning("Hub Stats Warning: HUGGINGFACEHUB_API_TOKEN not set")
+        else:
+            result = await hub_stats_tool.ainvoke({"author": "meta-llama"})
+            logger.info(f"Hub Stats Result: {result}")
+    except Exception as e:
+        logger.error(f"Hub Stats Error: {e}")
+
+    logger.info("Testing Guest Info Retriever Tool...")
+    try:
+        result = await guest_info_retriever_tool.ainvoke({"query": "Who is the guest named John?"})
+        logger.info(f"Guest Info Result: {result}")
+    except Exception as e:
+        logger.error(f"Guest Info Error: {e}")
+
+async def test_file_fetcher():
+    """Test file fetcher."""
+    logger.info("Testing File Fetcher...")
+    try:
+        result = await fetch_task_file("8e867cd7-cff9-4e6c-867a-ff5ddc2550be", "Sample question with data")
+        logger.info(f"File Fetcher Result: {result}")
+    except Exception as e:
+        logger.error(f"File Fetcher Error: {e}")
+
+async def test_answer_generator():
+    """Test answer generator functions."""
+    logger.info("Testing Preprocess Question...")
+    try:
+        result = await preprocess_question("What's the weather in Paris?")
+        logger.info(f"Preprocess Question Result: {result}")
+    except Exception as e:
+        logger.error(f"Preprocess Question Error: {e}")
+
+    logger.info("Testing Filter Results...")
+    try:
+        results = ["Paris is the capital of France.", "Florida is a state.", "Paris is in Texas."]
+        filtered = filter_results(results, "What is the capital of France?")
+        logger.info(f"Filter Results: {filtered}")
+    except Exception as e:
+        logger.error(f"Filter Results Error: {e}")
+
+async def test_state_management():
+    """Test state management functions."""
+    logger.info("Testing Reset State...")
+    try:
+        state = reset_state("test_task", "What is the capital of France?")
+        logger.info(f"Reset State Result: {state}")
+    except Exception as e:
+        logger.error(f"Reset State Error: {e}")
+
+    logger.info("Testing Validate State...")
+    try:
+        invalid_state = {"task_id": "", "question": ""}
+        validate_state(invalid_state)
+        logger.error("Validate State should have failed")
+    except ValueError as e:
+        logger.info(f"Validate State Error (expected): {e}")
+
+    try:
+        valid_state = reset_state("test_task", "Sample question")
+        validated = validate_state(valid_state)
+        logger.info(f"Validate State Result: {validated}")
+    except Exception as e:
+        logger.error(f"Validate State Error: {e}")
+
+async def test_agent():
+    """Test JARVISAgent with various cases."""
+    logger.info("Testing JARVISAgent (Simple Question)...")
+    try:
+        agent = JARVISAgent()
+        answer = await agent.process_question("test_task", "What is the capital of France?")
+        logger.info(f"JARVISAgent Answer: {answer}")
+    except Exception as e:
+        logger.error(f"JARVISAgent Error: {e}")
+
+    logger.info("Testing JARVISAgent (Edge Case: Empty Question)...")
+    try:
+        agent = JARVISAgent()
+        answer = await agent.process_question("test_task", "")
+        logger.info(f"JARVISAgent Empty Question Answer: {answer}")
+    except Exception as e:
+        logger.info(f"JARVISAgent Empty Question Error (expected): {e}")
+
+async def main():
+    required_envs = [
+        "HUGGINGFACEHUB_API_TOKEN",
+        "TOGETHER_API_KEY",
+        "OPENWEATHERMAP_API_KEY",
+        "SERPAPI_API_KEY"
+    ]
+    for env in required_envs:
+        if not os.getenv(env):
+            logger.warning(f"{env} not set, some tools may fail")
+
+    await test_tools()
+    await test_file_fetcher()
+    await test_answer_generator()
+    await test_state_management()
+    await test_agent()
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except Exception as e:
+        logger.error(f"Test script failed: {e}")

tools/__init__.py

@@ -6,4 +6,6 @@ from .document_retriever import document_retriever_tool
 from .duckduckgo_search import duckduckgo_search_tool
 from .weather_info import weather_info_tool
 from .hub_stats import hub_stats_tool
-from .guest_info import guest_info_retriever_tool
+from .guest_info import guest_info_retriever_tool
+from .file_fetcher import fetch_task_file
+from .answer_generator import generate_answer, preprocess_question#, filter_results, get_embedder

tools/answer_generator.py

@@ -0,0 +1,129 @@
+import nltk
+import logging
+import numpy as np
+from typing import List, Any
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.messages import SystemMessage, HumanMessage
+from sentence_transformers import SentenceTransformer
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Download NLTK data
+try:
+    nltk.download('punkt', quiet=True)
+    nltk.download('stopwords', quiet=True)
+except Exception as e:
+    logger.warning(f"NLTK data download failed: {e}")
+
+# Global embedder
+_embedder = None
+
+def get_embedder():
+    global _embedder
+    if _embedder is None:
+        try:
+            _embedder = SentenceTransformer(
+                "all-MiniLM-L6-v2",
+                device="cpu",
+                cache_folder="./cache"
+            )
+            logger.info("SentenceTransformer initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize SentenceTransformer: {e}")
+            raise RuntimeError(f"Embedder initialization failed: {e}")
+    return _embedder
+
+def filter_results(search_results: List[str], question: str) -> List[str]:
+    try:
+        if not search_results or not question:
+            return search_results
+        
+        embedder = get_embedder()
+        question_embedding = embedder.encode([question], convert_to_numpy=True)
+        result_embeddings = embedder.encode(search_results, convert_to_numpy=True)
+        
+        similarities = np.dot(result_embeddings, question_embedding.T).flatten()
+        filtered_results = [
+            search_results[i] for i in range(len(search_results))
+            if similarities[i] > 0.5 and search_results[i].strip()
+        ]
+        
+        return filtered_results if filtered_results else search_results[:3]
+    except Exception as e:
+        logger.warning(f"Result filtering failed: {e}")
+        return search_results[:3]
+
+async def preprocess_question(question: str) -> str:
+    """Preprocess the question to clean and standardize it."""
+    try:
+        question = question.strip().lower()
+        if not question.endswith("?"):
+            question += "?"
+        logger.debug(f"Preprocessed question: {question}")
+        return question
+    except Exception as e:
+        logger.error(f"Error preprocessing question: {e}")
+        return question
+
+async def generate_answer(
+    task_id: str,
+    question: str,
+    search_results: List[str],
+    file_results: str,
+    llm_client: Any
+) -> str:
+    """Generate an answer using LLM with search and file results."""
+    try:
+        if not search_results:
+            search_results = ["No search results available."]
+        if not file_results:
+            file_results = "No file results available."
+
+        context = "\n".join([str(r) for r in search_results]) + "\n" + file_results
+        prompt = ChatPromptTemplate.from_messages([
+            SystemMessage(content="""You are an assistant answering questions using provided context. 
+            - Use ONLY the context to formulate a concise, accurate answer.
+            - If the context is insufficient, state: 'Insufficient information to answer.'
+            - Do NOT generate or assume information beyond the context.
+            - Return a single, clear sentence or phrase as the answer."""),
+            HumanMessage(content=f"Context: {context}\nQuestion: {question}")
+        ])
+        
+        messages = [
+            {"role": "system", "content": prompt[0].content},
+            {"role": "user", "content": prompt[1].content}
+        ]
+        
+        if isinstance(llm_client, tuple):  # hf_local
+            model, tokenizer = llm_client
+            inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
+            outputs = model.generate(inputs, max_new_tokens=100, temperature=0.7)
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        elif hasattr(llm_client, "chat"):  # together
+            response = llm_client.chat.completions.create(
+                model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
+                messages=messages,
+                max_tokens=100,
+                temperature=0.7,
+                top_p=0.9,
+                frequency_penalty=0.5
+            )
+            response = response.choices[0].message.content.strip()
+        else:  # hf_api
+            response = llm_client.chat.completions.create(
+                messages=messages,
+                max_tokens=100,
+                temperature=0.7
+            )
+            response = response.choices[0].message.content.strip()
+        
+        answer = response.strip()
+        if not answer or answer.lower() == "none":
+            answer = "Insufficient information to answer."
+        logger.info(f"Task {task_id}: Generated answer: {answer}")
+        return answer
+    except Exception as e:
+        logger.error(f"Task {task_id}: Answer generation failed: {e}")
+        return "Error generating answer."

tools/calculator.py

@@ -1,15 +1,35 @@
-from langchain_core.tools import tool
-from sympy import sympify
 import logging
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
 
 logger = logging.getLogger(__name__)
 
-@tool
-async def calculator_tool(expression: str) -> str:
-    """Evaluate a mathematical expression."""
+class CalculatorInput(BaseModel):
+    expression: str = Field(description="Mathematical expression to evaluate")
+
+async def calculator_func(expression: str) -> str:
+    """
+    Evaluate a mathematical expression and return the result as a string.
+    
+    Args:
+        expression (str): Mathematical expression (e.g., '2 + 2').
+    
+    Returns:
+        str: Result of the expression.
+    """
     try:
-        result = sympify(expression)
+        logger.info(f"Evaluating expression: {expression}")
+        result = eval(expression, {"__builtins__": {}}, {})  # Safe eval
+        if isinstance(result, float):
+            return f"{result:.2f}" if "USD" in expression else str(result)
         return str(result)
     except Exception as e:
-        logger.error(f"Error evaluating expression '{expression}': {e}")
-        return f"Error: {str(e)}"
+        logger.error(f"Calculator error: {e}")
+        return f"Error: {e}"
+
+calculator_tool = StructuredTool.from_function(
+    func=calculator_func,
+    name="calculator_tool",
+    args_schema=CalculatorInput,
+    coroutine=calculator_func
+)

tools/document_retriever.py

@@ -1,30 +1,47 @@
-from langchain_core.tools import tool
-from langchain_community.document_loaders import TextLoader, CSVLoader, PyPDFLoader
 import logging
 import os
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+from typing import Optional
 
 logger = logging.getLogger(__name__)
 
-@tool
-async def document_retriever_tool(task_id: str, query: str, file_type: str) -> str:
-    """Retrieve content from a document."""
+class DocumentRetrieverInput(BaseModel):
+    task_id: str = Field(description="Task identifier")
+    query: str = Field(description="Search query")
+    file_path: Optional[str] = Field(description="Path to document file", default=None)
+
+async def document_retriever_func(task_id: str, query: str, file_path: Optional[str] = None) -> str:
+    """
+    Retrieve content from documents for a given task and query.
+    
+    Args:
+        task_id (str): Task identifier.
+        query (str): Search query.
+        file_path (Optional[str]): Path to document file.
+    
+    Returns:
+        str: Retrieved document content or error message.
+    """
     try:
-        file_path = f"temp_{task_id}.{file_type}"
-        if not os.path.exists(file_path):
-            logger.warning(f"Document not found: {file_path}")
-            return "Document not found"
-        
-        if file_type == "txt":
-            loader = TextLoader(file_path)
-        elif file_type == "csv":
-            loader = CSVLoader(file_path)
-        elif file_type == "pdf":
-            loader = PyPDFLoader(file_path)
-        else:
-            return f"Unsupported file type: {file_type}"
-        
-        docs = loader.load()
-        return "\n".join(doc.page_content for doc in docs)
+        if file_path and os.path.exists(file_path):
+            logger.info(f"Retrieving document from {file_path} for task {task_id}")
+            if file_path.endswith('.pdf'):
+                import PyPDF2
+                with open(file_path, 'rb') as f:
+                    reader = PyPDF2.PdfReader(f)
+                    text = "".join(page.extract_text() or "" for page in reader.pages)
+                return text[:500] if text else "No text extracted"
+            return "Unsupported file format"
+        logger.warning(f"No valid documents found for task {task_id}")
+        return "Document not found"
     except Exception as e:
         logger.error(f"Error retrieving document for task {task_id}: {e}")
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}"
+
+document_retriever_tool = StructuredTool.from_function(
+    func=document_retriever_func,
+    name="document_retriever_tool",
+    args_schema=DocumentRetrieverInput,
+    coroutine=document_retriever_func
+)

tools/duckduckgo_search.py

@@ -1,6 +1,99 @@
-from smolagents import Tool, DuckDuckGoSearchTool
 import logging
+import os
+import asyncio
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+from typing import Optional, List
+from duckduckgo_search import DDGS
+from serpapi import GoogleSearch
 
 logger = logging.getLogger(__name__)
 
-duckduckgo_search_tool = DuckDuckGoSearchTool()
+class DuckDuckGoSearchInput(BaseModel):
+    query: str = Field(description="Search query")
+    original_query: str = Field(description="Original query for context")
+    embedder: Optional[object] = Field(description="SentenceTransformer embedder", default=None)
+
+async def duckduckgo_search_func(query: str, original_query: str, embedder: Optional[object] = None) -> List[str]:
+    """
+    Perform a DuckDuckGo search with retries and fall back to SerpAPI if needed.
+    
+    Args:
+        query (str): Search query.
+        original_query (str): Original query for context.
+        embedder (Optional[object]): SentenceTransformer for result filtering.
+    
+    Returns:
+        List[str]: List of search result snippets.
+    """
+    async def try_duckduckgo(query: str, max_retries: int = 3) -> List[str]:
+        for attempt in range(max_retries):
+            try:
+                logger.info(f"DuckDuckGo search attempt {attempt + 1} for query: {query}")
+                with DDGS() as ddgs:
+                    results = [r['body'] for r in ddgs.text(query, max_results=5)]
+                return results
+            except Exception as e:
+                if "Ratelimit" in str(e) and attempt < max_retries - 1:
+                    wait_time = 2 ** attempt  # Exponential backoff: 1s, 2s, 4s
+                    logger.warning(f"DuckDuckGo rate limit hit, retrying in {wait_time}s: {e}")
+                    await asyncio.sleep(wait_time)
+                else:
+                    logger.error(f"DuckDuckGo search failed for query '{query}': {e}")
+                    raise e
+        return []
+
+    async def try_serpapi(query: str, max_retries: int = 3) -> List[str]:
+        if not os.getenv("SERPAPI_API_KEY"):
+            logger.warning("SERPAPI_API_KEY not set, cannot use SerpAPI fallback")
+            return []
+        for attempt in range(max_retries):
+            try:
+                logger.info(f"SerpAPI search attempt {attempt + 1} for query: {query}")
+                params = {
+                    "q": query,
+                    "api_key": os.getenv("SERPAPI_API_KEY"),
+                    "num": 5
+                }
+                search = GoogleSearch(params)
+                results = search.get_dict().get("organic_results", [])
+                return [result.get("snippet", "") for result in results if "snippet" in result]
+            except Exception as e:
+                if attempt < max_retries - 1:
+                    wait_time = 2 ** attempt  # Exponential backoff: 1s, 2s, 4s
+                    logger.warning(f"SerpAPI search failed, retrying in {wait_time}s: {e}")
+                    await asyncio.sleep(wait_time)
+                else:
+                    logger.error(f"SerpAPI search failed for query '{query}': {e}")
+                    return []
+
+    try:
+        # Try DuckDuckGo with retries
+        logger.info(f"Executing DuckDuckGo search for query: {query}")
+        results = await try_duckduckgo(query)
+        
+        # Fall back to SerpAPI if DuckDuckGo fails
+        if not results:
+            logger.info(f"DuckDuckGo returned no results, falling back to SerpAPI for query: {query}")
+            results = await try_serpapi(query)
+        
+        # Rank results if embedder is provided
+        if embedder and results:
+            from sentence_transformers import util
+            query_embedding = embedder.encode(original_query, convert_to_tensor=True)
+            result_embeddings = embedder.encode(results, convert_to_tensor=True)
+            scores = util.cos_sim(query_embedding, result_embeddings)[0]
+            ranked_results = [results[i] for i in scores.argsort(descending=True)]
+            return ranked_results[:3]
+        
+        return results[:3] if results else []
+    except Exception as e:
+        logger.error(f"Search failed for query '{query}': {e}")
+        return []
+
+duckduckgo_search_tool = StructuredTool.from_function(
+    func=duckduckgo_search_func,
+    name="duckduckgo_search_tool",
+    args_schema=DuckDuckGoSearchInput,
+    coroutine=duckduckgo_search_func
+)

tools/file_fetcher.py

@@ -0,0 +1,42 @@
+import os
+import ssl
+import aiohttp
+import logging
+from typing import Dict
+from urllib.parse import urljoin
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+async def fetch_task_file(task_id: str, question: str) -> Dict[str, bytes]:
+    """
+    Fetch a file associated with a task from the GAIA API.
+    Returns a dictionary of file extensions to content.
+    """
+    results = {}
+    base_url = "https://gaia-benchmark-api.hf.space/files/"  # Updated URL
+    extensions = ["xlsx", "csv", "pdf", "txt", "mp3", "jpg", "png"]
+
+    ssl_context = ssl.create_default_context()
+    ssl_context.check_hostname = False
+    ssl_context.verify_mode = ssl.CERT_NONE
+
+    async with aiohttp.ClientSession(
+        connector=aiohttp.TCPConnector(ssl_context=ssl_context),
+        timeout=aiohttp.ClientTimeout(total=30)
+    ) as session:
+        for ext in extensions:
+            file_url = urljoin(base_url, f"{task_id}/{task_id}.{ext}")
+            try:
+                async with session.get(file_url) as response:
+                    if response.status == 200:
+                        content = await response.read()
+                        results[ext] = content
+                        logger.info(f"Fetched {ext} for task {task_id}")
+                    else:
+                        logger.warning(f"No {ext} for task {task_id}: HTTP {response.status}")
+            except Exception as e:
+                logger.warning(f"Error fetching {ext} for task {task_id}: {str(e)}")
+
+    return results

tools/file_parser.py

@@ -1,36 +1,112 @@
-from langchain_core.tools import tool
-import pandas as pd
-import PyPDF2
 import logging
 import os
+import pandas as pd
+import PyPDF2
+import speech_recognition as sr
+import re
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+from typing import Optional
 
 logger = logging.getLogger(__name__)
 
-@tool
-async def file_parser_tool(task_id: str, file_type: str) -> str:
-    """Parse a file based on task_id and file_type."""
+class FileParserInput(BaseModel):
+    task_id: str = Field(description="Task identifier")
+    file_type: str = Field(description="File extension (e.g., pdf, csv)")
+    file_path: str = Field(description="Path to the file")
+    query: Optional[str] = Field(description="Query related to the file", default=None)
+
+async def file_parser_func(task_id: str, file_type: str, file_path: str, query: Optional[str] = None) -> str:
+    """
+    Parse a file based on task_id, file_type, file_path, and query context.
+    
+    Args:
+        task_id (str): Task identifier.
+        file_type (str): File extension (e.g., 'xlsx', 'mp3', 'pdf').
+        file_path (str): Path to the file.
+        query (Optional[str]): Question context to guide parsing (e.g., for specific data extraction).
+    
+    Returns:
+        str: Parsed content or error message.
+    """
     try:
-        file_path = f"temp_{task_id}.{file_type}"
         if not os.path.exists(file_path):
             logger.warning(f"File not found: {file_path}")
             return "File not found"
         
-        if file_type == "csv":
+        logger.info(f"Parsing file: {file_path} for task {task_id}")
+
+        if file_type in ["xlsx", "xls"]:
+            df = pd.read_excel(file_path, engine="openpyxl")
+            if query and ("sum" in query.lower() or "total" in query.lower()):
+                numerical_cols = df.select_dtypes(include=['float64', 'int64']).columns
+                if numerical_cols.empty:
+                    return "No numerical data found"
+                if "food" in query.lower():
+                    food_rows = df[df.apply(lambda x: "food" in str(x).lower(), axis=1)]
+                    if not food_rows.empty and numerical_cols[0] in food_rows:
+                        total = food_rows[numerical_cols[0]].sum()
+                        return f"{total:.2f}"
+                total = df[numerical_cols[0]].sum()
+                return f"{total:.2f}"
+            return df.to_string(index=False)
+
+        elif file_type == "csv":
             df = pd.read_csv(file_path)
-            return df.to_string()
-        elif file_type == "txt":
-            with open(file_path, "r", encoding="utf-8") as f:
-                return f.read()
+            if query and ("sum" in query.lower() or "total" in query.lower()):
+                numerical_cols = df.select_dtypes(include=['float64', 'int64']).columns
+                if numerical_cols.empty:
+                    return "No numerical data found"
+                total = df[numerical_cols[0]].sum()
+                return f"{total:.2f}"
+            return df.to_string(index=False)
+
         elif file_type == "pdf":
             with open(file_path, "rb") as f:
                 reader = PyPDF2.PdfReader(f)
-                text = "".join(page.extract_text() for page in reader.pages)
+                text = "".join(page.extract_text() or "" for page in reader.pages)
+                if query and "page number" in query.lower():
+                    pages = re.findall(r'\b\d+\b', text)
+                    return ", ".join(sorted(pages, key=int)) if pages else "No page numbers found"
+                return text.strip() or "No text extracted"
+
+        elif file_type == "txt":
+            with open(file_path, "r", encoding="utf-8") as f:
+                text = f.read()
+                if query and "page number" in query.lower():
+                    pages = re.findall(r'\b\d+\b', text)
+                    return ", ".join(sorted(pages, key=int)) if pages else "No page numbers found"
+                return text.strip()
+
+        elif file_type == "mp3":
+            recognizer = sr.Recognizer()
+            with sr.AudioFile(file_path) as source:
+                audio = recognizer.record(source)
+            try:
+                text = recognizer.recognize_google(audio)
+                logger.debug(f"Transcribed audio: {text}")
+                if query and "page number" in query.lower():
+                    pages = re.findall(r'\b\d+\b', text)
+                    return ", ".join(sorted(pages, key=int)) if pages else "No page numbers provided"
                 return text
-        elif file_type in ["xlsx", "xls"]:
-            df = pd.read_excel(file_path, engine="openpyxl")
-            return df.to_string()
+            except sr.UnknownValueError:
+                logger.error("Could not understand audio")
+                return "No text transcribed from audio"
+            except Exception as e:
+                logger.error(f"Audio parsing failed: {e}")
+                return "Error transcribing audio"
+
         else:
+            logger.warning(f"Unsupported file type: {file_type}")
             return f"Unsupported file type: {file_type}"
+
     except Exception as e:
         logger.error(f"Error parsing file for task {task_id}: {e}")
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}"
+
+file_parser_tool = StructuredTool.from_function(
+    func=file_parser_func,
+    name="file_parser_tool",
+    args_schema=FileParserInput,
+    coroutine=file_parser_func
+)

tools/guest_info.py

@@ -1,20 +1,47 @@
-from langchain_core.tools import tool
-from retriever import load_guest_dataset
 import logging
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+from datasets import load_dataset
+from rank_bm25 import BM25Okapi
 
 logger = logging.getLogger(__name__)
 
-@tool
-async def guest_info_retriever_tool(query: str) -> str:
-    """Retrieve detailed information about gala guests based on their name or relation."""
+class GuestInfoInput(BaseModel):
+    query: str = Field(description="Query about guest information")
+
+async def guest_info_func(query: str) -> str:
+    """
+    Retrieve guest information based on a query.
+    
+    Args:
+        query (str): Query about guest information.
+    
+    Returns:
+        str: Guest information or error message.
+    """
     try:
-        docs = load_guest_dataset()
-        from langchain_community.retrievers import BM25Retriever
-        retriever = BM25Retriever.from_documents(docs)
-        results = retriever.get_relevant_documents(query)
-        if results:
-            return "\n\n".join([doc.page_content for doc in results[:3]])
-        return "No matching guest information found."
+        logger.info(f"Retrieving guest info for query: {query}")
+        dataset = load_dataset("agents-course/unit3-invitees", split="train")
+        logger.info(f"Loaded {len(dataset)} guests from Hugging Face dataset")
+        
+        documents = [f"{row['name']} {row['relation']}" for row in dataset]
+        tokenized_docs = [doc.lower().split() for doc in documents]
+        bm25 = BM25Okapi(tokenized_docs)
+        
+        tokenized_query = query.lower().split()
+        scores = bm25.get_scores(tokenized_query)
+        best_idx = scores.argmax()
+        
+        if scores[best_idx] > 0:
+            return f"Guest: {dataset[best_idx]['name']}, Relation: {dataset[best_idx]['relation']}"
+        return "No matching guest found"
     except Exception as e:
         logger.error(f"Error retrieving guest info for query '{query}': {e}")
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}"
+
+guest_info_retriever_tool = StructuredTool.from_function(
+    func=guest_info_func,
+    name="guest_info_retriever_tool",
+    args_schema=GuestInfoInput,
+    coroutine=guest_info_func
+)

tools/hub_stats.py

@@ -1,17 +1,54 @@
-from langchain_core.tools import tool
-from huggingface_hub import list_models
+import aiohttp
+import ssl
 import logging
+from langchain_core.tools import tool
+from tenacity import retry, stop_after_attempt, wait_exponential
+from typing import Optional
+import json
+import os
 
 logger = logging.getLogger(__name__)
 
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10))
+async def fetch_hf_models(author: str) -> Optional[dict]:
+    url = f"https://huggingface.co/api/models?author={author}&sort=downloads&direction=-1&limit=1"
+    ssl_context = ssl.create_default_context()
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url, ssl=ssl_context) as response:
+                response.raise_for_status()
+                return await response.json()
+    except aiohttp.ClientError as e:
+        logger.error(f"Failed to fetch models for {author}: {e}")
+        raise
+
 @tool
 async def hub_stats_tool(author: str) -> str:
-    """Fetch the most downloaded model from a specific author on Hugging Face Hub."""
+    """
+    Fetch the most downloaded model from a specific author on Hugging Face Hub.
+    
+    Args:
+        author (str): Hugging Face author username.
+    
+    Returns:
+        str: Model information or error message.
+    """
     try:
-        models = list(list_models(author=author, sort="downloads", direction=-1, limit=1))
-        if models:
+        # Check local cache
+        cache_file = f"temp/hf_cache_{author}.json"
+        if os.path.exists(cache_file):
+            with open(cache_file, "r") as f:
+                models = json.load(f)
+            logger.debug(f"Loaded cached models for {author}")
+        else:
+            models = await fetch_hf_models(author)
+            os.makedirs("temp", exist_ok=True)
+            with open(cache_file, "w") as f:
+                json.dump(models, f)
+        
+        if models and isinstance(models, list) and models:
             model = models[0]
-            return f"The most downloaded model by {author} is {model.id} with {model.downloads:,} downloads."
+            return f"The most downloaded model by {author} is {model['id']} with {model.get('downloads', 0):,} downloads."
         return f"No models found for author {author}."
     except Exception as e:
         logger.error(f"Error fetching models for {author}: {e}")

tools/image_parser.py

@@ -1,25 +1,43 @@
-from langchain_core.tools import tool
-import easyocr
 import logging
 import os
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+import easyocr
 
 logger = logging.getLogger(__name__)
-reader = easyocr.Reader(['en'])
 
-@tool
-async def image_parser_tool(file_path: str, task: str = "describe", match_query: str = "") -> str:
-    """Parse text from an image."""
+class ImageParserInput(BaseModel):
+    task_id: str = Field(description="Task identifier")
+    file_path: str = Field(description="Path to the image file")
+
+async def image_parser_func(task_id: str, file_path: str) -> str:
+    """
+    Parse text from an image file using OCR.
+    
+    Args:
+        task_id (str): Task identifier.
+        file_path (str): Path to the image file.
+    
+    Returns:
+        str: Extracted text or error message.
+    """
     try:
         if not os.path.exists(file_path):
-            logger.warning(f"Image not found: {file_path}")
-            return "Image not found"
+            logger.warning(f"Image file not found: {file_path}")
+            return "Image file not found"
         
-        results = reader.readtext(file_path)
-        text = " ".join(result[1] for result in results)
-        
-        if task == "match" and match_query:
-            return str(match_query.lower() in text.lower())
-        return text
+        logger.info(f"Parsing image: {file_path} for task {task_id}")
+        reader = easyocr.Reader(['en'], model_storage_directory='./cache')
+        result = reader.readtext(file_path, detail=0)
+        text = " ".join(result).strip()
+        return text if text else "No text extracted from image"
     except Exception as e:
-        logger.error(f"Error parsing image {file_path}: {e}")
-        return f"Error: {str(e)}"
+        logger.error(f"Error parsing image for task {task_id}: {e}")
+        return f"Error: {str(e)}"
+
+image_parser_tool = StructuredTool.from_function(
+    func=image_parser_func,
+    name="image_parser_tool",
+    args_schema=ImageParserInput,
+    coroutine=image_parser_func
+)

tools/search.py

@@ -1,106 +1,103 @@
+import logging
 import os
-import json
-import asyncio
-# from serpapi import GoogleSearch
-from google_search_results import GoogleSearch
-from langchain.tools import Tool
-from typing import List, Dict, Any
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import StructuredTool
+from pydantic import BaseModel, Field
+from typing import Optional, List
+from serpapi import GoogleSearch
 
-def search_tool(query: str) -> List[str]:
+logger = logging.getLogger(__name__)
+
+class SearchInput(BaseModel):
+    query: str = Field(description="Search query")
+
+async def search_func(query: str) -> List[str]:
     """
-    Perform a web search using SERPAPI with retries.
+    Perform a web search using SerpAPI and return relevant snippets.
     
     Args:
-        query: Search query string.
+        query (str): The search query to execute.
     
     Returns:
-        List of search result snippets.
-    
-    Raises:
-        Exception: If search fails after retries.
+        List[str]: A list of search result snippets.
     """
-    params = {
-        "q": query,
-        "api_key": os.getenv("SERPAPI_API_KEY"),
-        "num": 5,
-    }
-    
-    for attempt in range(3):
-        try:
-            search = GoogleSearch(params, timeout=30)
-            results = search.get_dict()
-            organic_results = results.get("organic_results", [])
-            return [r.get("snippet", "") for r in organic_results]
-        except Exception as e:
-            print(f"INFO - SERPAPI retry {attempt + 1}/3 due to: {e}")
-            asyncio.sleep(2)
-    
-    raise Exception("SERPAPI failed after retries")
+    try:
+        logger.info(f"Executing SerpAPI search for query: {query}")
+        params = {
+            "q": query,
+            "api_key": os.getenv("SERPAPI_API_KEY"),
+            "num": 10
+        }
+        search = GoogleSearch(params)
+        results = search.get_dict().get("organic_results", [])
+        return [result.get("snippet", "") for result in results if "snippet" in result]
+    except Exception as e:
+        logger.error(f"SerpAPI search failed for query '{query}': {e}")
+        return []
 
-async def multi_hop_search_tool(query: str, steps: int = 3, llm_client: Any = None, llm_type: str = None) -> List[Dict[str, str]]:
+search_tool = StructuredTool.from_function(
+    func=search_func,
+    name="search_tool",
+    args_schema=SearchInput,
+    coroutine=search_func
+)
+
+class MultiHopSearchInput(BaseModel):
+    query: str = Field(description="Multi-hop search query")
+    steps: int = Field(description="Number of search steps", ge=1, le=3)
+    llm_client: Optional[object] = Field(description="LLM client", default=None)
+    llm_type: Optional[str] = Field(description="LLM type", default="together")
+    llm_model: Optional[str] = Field(description="LLM model", default="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free")
+
+async def multi_hop_search_func(query: str, steps: int, llm_client: Optional[object] = None, llm_type: Optional[str] = "together", llm_model: Optional[str] = "meta-llama/Llama-3.3-70B-Instruct-Turbo-Free") -> List[str]:
     """
-    Perform iterative web searches for complex queries, refining the query using an LLM.
+    Perform a multi-hop web search using SerpAPI with iterative query refinement.
     
     Args:
-        query: Initial search query.
-        steps: Number of search iterations.
-        llm_client: LLM client for query refinement.
-        llm_type: Type of LLM client ("together", "hf_api", or "hf_local").
+        query (str): The initial multi-hop search query.
+        steps (int): Number of search steps to perform (1 to 3).
+        llm_client (Optional[object]): LLM client for query refinement.
+        llm_type (Optional[str]): Type of LLM (e.g., 'together').
+        llm_model (Optional[str]): LLM model name.
     
     Returns:
-        List of dictionaries containing search result content.
+        List[str]: A list of search result snippets from all steps.
     """
-    results = []
-    current_query = query
-    
-    for step in range(steps):
-        try:
-            # Perform search
-            search_results = search_tool(current_query)
-            results.extend([{"content": str(r)} for r in search_results])
+    try:
+        logger.info(f"Executing multi-hop search for query: {query}, steps: {steps}")
+        results = []
+        current_query = query
+        
+        for step in range(steps):
+            logger.info(f"Multi-hop step {step + 1}: {current_query}")
+            step_results = await search_func(current_query)
+            results.extend(step_results)
             
-            # Refine query using LLM if available
-            if llm_client and step < steps - 1:
-                prompt = ChatPromptTemplate.from_messages([
-                    SystemMessage(content="""Refine the following query to dig deeper into the topic, focusing on missing details or related aspects. Return ONLY the refined query as plain text, no explanations."""),
-                    HumanMessage(content=f"Original query: {current_query}\nPrevious results: {json.dumps(search_results[:2], indent=2)}")
-                ])
+            if step < steps - 1 and llm_client:
+                prompt = f"Given the query '{current_query}' and results: {step_results[:3]}, generate a follow-up search query to refine or expand the search."
                 messages = [
-                    {"role": "system", "content": prompt[0].content},
-                    {"role": "user", "content": prompt[1].content}
+                    {"role": "system", "content": "Generate a single search query as a string."},
+                    {"role": "user", "content": prompt}
                 ]
-                
-                try:
-                    if llm_type == "hf_local":
-                        model, tokenizer = llm_client
-                        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("mps")
-                        outputs = model.generate(inputs, max_new_tokens=100, temperature=0.7)
-                        refined_query = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-                    else:
-                        response = llm_client.chat.completions.create(
-                            model=llm_client.model if llm_type == "together" else "meta-llama/Llama-3.2-1B-Instruct",
-                            messages=messages,
-                            max_tokens=100,
-                            temperature=0.7
-                        )
-                        refined_query = response.choices[0].message.content.strip()
-                    
-                    current_query = refined_query if refined_query else f"more details on {current_query}"
-                except Exception as e:
-                    print(f"INFO - Query refinement failed at step {step + 1}: {e}")
-                    current_query = f"more details on {current_query}"
-            
-            await asyncio.sleep(1)  # Rate limit
-        except Exception as e:
-            print(f"INFO - Multi-hop search step {step + 1} failed: {e}")
-            break
-    
-    return results
+                if llm_type == "together":
+                    response = llm_client.chat.completions.create(
+                        model=llm_model,
+                        messages=messages,
+                        max_tokens=50,
+                        temperature=0.7
+                    )
+                    current_query = response.choices[0].message.content.strip()
+                else:
+                    logger.warning("LLM not configured for multi-hop refinement")
+                    break
+        
+        return results[:5] if results else ["No results found"]
+    except Exception as e:
+        logger.error(f"Multi-hop search failed for query '{query}': {e}")
+        return [f"Error: {str(e)}"]
 
-multi_hop_search_tool = Tool.from_function(
-    func=multi_hop_search_tool,
+multi_hop_search_tool = StructuredTool.from_function(
+    func=multi_hop_search_func,
     name="multi_hop_search_tool",
-    description="Performs iterative web searches for complex queries, refining the query with an LLM."
+    args_schema=MultiHopSearchInput,
+    coroutine=multi_hop_search_func
 )

tools/weather_info.py

@@ -1,23 +1,50 @@
-from langchain_core.tools import tool
-import requests
+import aiohttp
+import ssl
 import logging
 import os
+from langchain_core.tools import tool
+from tenacity import retry, stop_after_attempt, wait_exponential
 from dotenv import load_dotenv
 
 logger = logging.getLogger(__name__)
 load_dotenv()
 
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10))
+async def fetch_weather(location: str, api_key: str) -> dict:
+    url = f"http://api.openweathermap.org/data/2.5/weather?q={location}&appid={api_key}&units=metric"
+    ssl_context = ssl.create_default_context()
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url, ssl=ssl_context) as response:
+                response.raise_for_status()
+                return await response.json()
+    except aiohttp.ClientError as e:
+        logger.error(f"Failed to fetch weather for {location}: {e}")
+        raise
+
 @tool
-async def weather_info_tool(location: str) -> str:
-    """Fetch real weather information for a given location."""
+async def weather_info_tool(location: str, query_type: str = "current") -> str:
+    """
+    Fetch weather information for a given location.
+    
+    Args:
+        location (str): City or location name.
+        query_type (str): Type of weather query ('current', 'forecast'; default: 'current').
+    
+    Returns:
+        str: Weather information or error message.
+    """
     try:
         api_key = os.getenv("OPENWEATHERMAP_API_KEY")
         if not api_key:
             logger.error("OPENWEATHERMAP_API_KEY not set")
             return "Weather unavailable: API key missing"
 
-        url = f"http://api.openweathermap.org/data/2.5/weather?q={location}&appid={api_key}&units=metric"
-        response = requests.get(url).json()
+        if query_type != "current":
+            logger.warning(f"Query type '{query_type}' not supported; using current weather")
+            query_type = "current"
+
+        response = await fetch_weather(location, api_key)
         if response.get("cod") == 200:
             condition = response["weather"][0]["description"]
             temp = response["main"]["temp"]