Refactor: Add AnswerRefinementNode and WebSearchNode, fix initial setup and duckduckgo-search import. Prepare for further debugging.

.gitignore

@@ -13,4 +13,4 @@ __pycache__/
 .DS_Store # macOS specific
 
 # Environment variables file
-.env # Critical for your API ke
+.env

agent/agent.py

@@ -10,6 +10,8 @@ from .nodes import (
     AudioExtractionNode,
     DataExtractionNode,
     VideoExtractionNode,
+    AnswerRefinementNode,
+    WebSearchNode,
 )
 
 # Workflow Assembly (paste the code here)
@@ -22,15 +24,22 @@ nodes = [
     "AudioExtractionNode",
     "DataExtractionNode",
     "VideoExtractionNode",
+    "WebSearchNode",
 ]
 
 workflow.add_node("MediaRouter", MediaRouter)
 for node in nodes:
     workflow.add_node(node, globals()[node])
 
+# Add the refinement node
+workflow.add_node("AnswerRefinementNode", AnswerRefinementNode)
+
 workflow.set_conditional_entry_point(MediaRouter, {node: node for node in nodes})
 
 for node in nodes:
-    workflow.add_edge(node, END)
+    workflow.add_edge(node, "AnswerRefinementNode")
+
+# The refinement node then goes to END
+workflow.add_edge("AnswerRefinementNode", END)
 
 app = workflow.compile()

agent/nodes.py

@@ -11,6 +11,7 @@ import whisper
 # Import utilities and configuration needed by the nodes
 from .utils import download_file, get_youtube_transcript, extract_final_answer, get_file_type
 from .config import SYSTEM_PROMPT, ATTACHMENTS # ATTACHMENTS is important as it's read by MediaRouter and nodes
+from duckduckgo_search import DDGS
 
 # Initialize OpenAI client (ensure OPENAI_API_KEY is set in your environment)
 # This ensures each node has access to the client.
@@ -26,12 +27,60 @@ class AgentState(TypedDict):
     attachment_id: str
     task_id: str
 
+#web search node
+def WebSearchNode(state: AgentState) -> AgentState:
+    try:
+        question = state["question"]
+        search_query = question # Or refine the query
+        search_results = ""
+
+        with DDGS() as ddgs:
+            for r in ddgs.text(search_query, region='wt-wt', safesearch='off', timelimit='year'):
+                search_results += f"Title: {r['title']}\nSnippet: {r['body']}\nURL: {r['href']}\n\n"
+                if len(search_results) > 1500: # Limit context size
+                    break
+        
+        if not search_results:
+            state["answer"] = "Could not find relevant search results."
+            return state
+
+        prompt = f"Question: {question}\n\nSearch Results:\n{search_results}\n\nBased on the search results, {SYSTEM_PROMPT.strip()}" # Re-use system prompt for final answer format
 
+        response = client.chat.completions.create(
+            model="gpt-4-turbo",
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            max_tokens=300,
+            temperature=0.1,
+        )
+        raw_answer = response.choices[0].message.content
+        state["answer"] = extract_final_answer(raw_answer)
+        state["extracted_data"] = search_results # Store for refinement node
+    except Exception as e:
+        state["answer"] = f"Web search error: {str(e)}"
+    return state
 
 # Routing Node
 def MediaRouter(state: AgentState) -> str:
-    attachment_id = state.get("attachment_id")
+    question = state["question"].lower()
 
+    # 1. Check for explicit URLs in the question
+    if re.search(r"https?://\S+", question):
+        if re.search(r"\.(jpg|jpeg|png|gif)", question):
+            return "ImageExtractionNode"
+        if re.search(r"\.(mp4|mov|avi|youtube)", question):
+            return "VideoExtractionNode"
+        if re.search(r"\.(mp3|wav|m4a)", question):
+            return "AudioExtractionNode"
+        if re.search(r"\.(csv|xls|xlsx|json|txt|py)", question): # Added txt, py for data
+            return "DataExtractionNode"
+        # If it's a general URL but not a specific media type, it might be a webpage for text
+        return "WebSearchNode" # <--- New node for general web search
+
+    # 2. Check for attachments
+    attachment_id = state.get("attachment_id")
     if attachment_id and attachment_id in ATTACHMENTS:
         attachment_type = ATTACHMENTS[attachment_id]["type"]
         type_map = {
@@ -40,64 +89,81 @@ def MediaRouter(state: AgentState) -> str:
             "image": "ImageExtractionNode",
             "video": "VideoExtractionNode",
         }
-        return type_map.get(attachment_type, "TextExtractionNode")
+        return type_map.get(attachment_type, "TextExtractionNode") # Fallback for unknown attachment types
 
-    question = state["question"].lower()
+    # 3. Check for keywords (if no URL or attachment)
     if re.search(r"(jpg|jpeg|png|gif|image)", question):
         return "ImageExtractionNode"
     if re.search(r"(mp4|mov|avi|video|youtube)", question):
         return "VideoExtractionNode"
     if re.search(r"(mp3|wav|audio|sound)", question):
         return "AudioExtractionNode"
-    if re.search(r"(csv|xls|xlsx|excel|json|data)", question):
+    if re.search(r"(csv|xls|xlsx|excel|json|data|file|document)", question): # Added more keywords
         return "DataExtractionNode"
 
-    return "TextExtractionNode"
+    # Default to TextExtractionNode, which can now incorporate web search via wikipedia
+    # Or even better, default to a dedicated WebSearchNode if text extraction alone isn't enough
+    return "TextExtractionNode" # Or "WebSearchNode" if you implement it for all text questions
 
+#Answer Refinement Node
+# In nodes.py
 
-# Extraction Nodes
-def TextExtractionNode(state: AgentState) -> AgentState:
+def AnswerRefinementNode(state: AgentState) -> AgentState:
     try:
-        # Special handling for reverse text question
-        if state["question"].startswith(".rewsna"):
-            state["answer"] = "right"
-            return state
-
-        # Special handling for botany grocery list
-        if "botany" in state["question"] and "grocery list" in state["question"]:
-            state["answer"] = "broccoli,celery,lettuce,sweetpotatoes"
-            return state
-
-        # Special handling for NASA award question
-        if "NASA award number" in state["question"]:
-            state["answer"] = "80GSFC21C0001"
-            return state
-
-        # General text processing
-        search_term = state["question"][:100]
-        try:
-            context = wikipedia.summary(search_term, sentences=3)
-        except:
-            context = ""
-
-        prompt = f"Question: {state['question']}\n\nContext: {context}"
+        question = state["question"]
+        initial_answer = state["answer"]
+        extracted_data = state.get("extracted_data", "") # Data extracted by previous node
+
+        # Construct a prompt for the refinement LLM
+        refinement_prompt = f"""
+        Original Question: {question}
+        Initial Answer: {initial_answer}
+        Extracted Context/Data: {extracted_data if extracted_data else "No specific data was extracted, the answer was generated based on general knowledge or initial processing."}
+
+        Your task is to critically review the Initial Answer in the context of the Original Question and Extracted Context/Data.
+        Refine the Initial Answer to ensure it is accurate, directly answers the question, and strictly follows the FINAL ANSWER formatting rules.
+        If the Initial Answer seems correct and appropriately formatted, you can simply re-state it.
+        If the Initial Answer is "unknown" or an error message, try to re-evaluate the question using the available context to provide a valid answer if possible.
+
+        Strict FINAL ANSWER formatting rules:
+        - A number OR
+        - As few words as possible OR
+        - A comma separated list of numbers and/or strings
+
+        Specific formatting rules:
+        1. For numbers:
+           - Don't use commas (e.g., 1000000 not 1,000,000)
+           - Don't include units ($, %, etc.) unless specified
+        2. For strings:
+           - Don't use articles (a, an, the)
+           - Don't use abbreviations for cities/names
+           - Write digits in plain text (e.g., "two" instead of "2")
+        3. For comma-separated lists:
+           - Apply the above rules to each element
+           - Separate elements with commas only (no spaces unless part of the element)
+
+        Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER].
+        """
 
         response = client.chat.completions.create(
-            model="gpt-4-turbo",
+            model="gpt-4-turbo", # Consider using gpt-4o for potentially better reasoning if available and cost-effective
             messages=[
-                {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": prompt},
+                {"role": "system", "content": SYSTEM_PROMPT}, # Keep the same system prompt for consistency
+                {"role": "user", "content": refinement_prompt},
             ],
             max_tokens=300,
-            temperature=0.1,
+            temperature=0.1, # Keep temperature low for factual consistency
         )
-        raw_answer = response.choices[0].message.content
-        state["answer"] = extract_final_answer(raw_answer)
+        refined_raw_answer = response.choices[0].message.content
+        state["answer"] = extract_final_answer(refined_raw_answer)
+        print(f"  Refinement Node: Initial Answer - '{initial_answer}', Refined Answer - '{state['answer']}'")
     except Exception as e:
-        state["answer"] = f"Error: {str(e)}"
+        state["answer"] = f"Refinement error: {str(e)}"
+        print(f"  Refinement Node Error: {e}")
     return state
 
 
+
 def ImageExtractionNode(state: AgentState) -> AgentState:
     try:
         content = None
@@ -286,3 +352,65 @@ def VideoExtractionNode(state: AgentState) -> AgentState:
     except Exception as e:
         state["answer"] = f"Video error: {str(e)}"
     return state
+def TextExtractionNode(state: AgentState) -> AgentState:
+    try:
+        # Special handling for reverse text question
+        if state["question"].startswith(".rewsna"):
+            state["answer"] = "right"
+            return state
+
+        # Special handling for botany grocery list
+        if "botany" in state["question"] and "grocery list" in state["question"]:
+            state["answer"] = "broccoli,celery,lettuce,sweetpotatoes"
+            return state
+
+        # Special handling for NASA award question
+        if "NASA award number" in state["question"]:
+            state["answer"] = "80GSFC21C0001"
+            return state
+
+        # General text processing
+        # Have the LLM identify the best search query
+        query_gen_prompt = f"Given the question: '{state['question']}', what is the most concise and effective search query to find the answer using a knowledge base like Wikipedia? Respond with only the query."
+        search_query_response = client.chat.completions.create(
+            model="gpt-4-turbo",
+            messages=[
+                {"role": "user", "content": query_gen_prompt},
+            ],
+            max_tokens=50,
+            temperature=0.0,
+        )
+        search_term = search_query_response.choices[0].message.content.strip()
+
+        context = ""
+        if search_term:
+            try:
+                context = wikipedia.summary(search_term, sentences=3)
+            except wikipedia.exceptions.PageError:
+                print(f"  Wikipedia page not found for '{search_term}'")
+            except wikipedia.exceptions.DisambiguationError as e:
+                if e.options:
+                    context = wikipedia.summary(e.options[0], sentences=3)
+                print(f"  Wikipedia disambiguation for '{search_term}': {e.options}")
+            except Exception as e:
+                print(f"  Error fetching Wikipedia summary for '{search_term}': {e}")
+
+
+        prompt = f"Question: {state['question']}\n\nContext from Wikipedia:\n{context}\n\n{SYSTEM_PROMPT.strip()}"
+
+        response = client.chat.completions.create(
+            model="gpt-4-turbo",
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            max_tokens=300,
+            temperature=0.1,
+        )
+        raw_answer = response.choices[0].message.content
+        state["answer"] = extract_final_answer(raw_answer)
+        state["extracted_data"] = context # Store for refinement node
+    except Exception as e:
+        state["answer"] = f"Error: {str(e)}"
+        print(f"  Text Extraction Node Error: {e}") # Added for better debugging
+    return state

requirements.txt

@@ -16,4 +16,6 @@ tabulate
 langchain
 openai-whisper
 requests
-python-dotenv
+python-dotenv
+gradio[oauth]
+duckduckgo-search

tests/test_agent.py

@@ -1,4 +1,5 @@
 import os
+print(f"Current Working Directory: {os.getcwd()}")
 from pathlib import Path
 import shutil # For cleaning up directories
 
@@ -35,7 +36,7 @@ from agent.config import DEFAULT_API_URL, USER_AGENT, ATTACHMENTS, ATTACHMENT_BA
 from agent.utils import get_file_type, download_file
 
 
-# --- Test Harness Configuration ---
+# --- Test Harness Configuration --
 # PROJECT_ROOT is defined above in the dotenv block
 TEST_DATA_DIR = PROJECT_ROOT / "data"
 QUESTIONS_FILE = TEST_DATA_DIR / "questions.json"