Update app.py

app.py

@@ -16,7 +16,6 @@ from huggingface_hub import InferenceClient
 import io
 import mimetypes
 import base64
-from decision_maker import DecisionMaker, ToolType
 
 # -------------------------
 # Environment & constants
@@ -95,98 +94,43 @@ class AgentState(TypedDict):
     search_query: Annotated[str, override]
     task_id: Annotated[str, override]
     logs: Annotated[Dict[str, Any], merge_dicts]
-    code_blocks: Annotated[List[Dict[str, str]], list.__add__]
-    attachment_data: Annotated[Dict[str, bytes], merge_dicts]  # Store downloaded file data
+    attachment_data: Annotated[Dict[str, bytes], merge_dicts]
 
 # -------------------------
 # BasicAgent implementation
 # -------------------------
 
-def quick_code_stats(src: str) -> Dict[str, Any]:
-    """Lightweight code analysis using AST."""
-    try:
-        tree = ast.parse(src)
-        funcs = [n.name for n in tree.body if isinstance(n, ast.FunctionDef)]
-        classes = [n.name for n in tree.body if isinstance(n, ast.ClassDef)]
-        imports = []
-        for node in ast.walk(tree):
-            if isinstance(node, ast.Import):
-                imports.extend(n.name for n in node.names)
-            elif isinstance(node, ast.ImportFrom):
-                imports.append(f"{node.module}.{node.names[0].name}")
-        
-        return {
-            "functions": funcs,
-            "classes": classes,
-            "imports": imports,
-            "lines": len(src.splitlines())
-        }
-    except Exception as e:
-        return {
-            "error": str(e),
-            "lines": len(src.splitlines())
-        }
-
 class BasicAgent:
     def __init__(self):
         if not OPENAI_API_KEY:
             raise EnvironmentError("OPENAI_API_KEY not set")
         self.llm = OpenAI(api_key=OPENAI_API_KEY)
-        self.decision_maker = DecisionMaker()
         self.workflow = self._build_workflow()
 
-    # ---- Low‑level LLM call
     def _call_llm(self, prompt: str, max_tokens: int = 256) -> str:
-        # List of models to try in order
-        models = ["gpt-4.1", "gpt-4o-mini", "gpt-3.5-turbo"]
-        
-        for model in models:
-            try:
-                print(f"\nTrying model: {model}")
-                resp = self.llm.chat.completions.create(
-                    model=model,
-                    messages=[
-                        {"role": "system", "content": "You are a careful reasoning assistant."},
-                        {"role": "user", "content": prompt},
-                    ],
-                    temperature=0.3,
-                    max_tokens=max_tokens,
-                )
-                return resp.choices[0].message.content.strip()
-            except Exception as e:
-                print(f"Error with {model}: {str(e)}")
-                if model == models[-1]:  # If this was the last model
-                    print(f"All models failed. Last error: {str(e)}")
-                    print(f"Prompt that caused error:\n{prompt}")
-                    raise
-                print(f"Falling back to next model...")
-                continue
-
-    # ---- Workflow nodes
+        try:
+            resp = self.llm.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[
+                    {"role": "system", "content": "You are a careful reasoning assistant."},
+                    {"role": "user", "content": prompt},
+                ],
+                temperature=0.3,
+                max_tokens=max_tokens,
+            )
+            return resp.choices[0].message.content.strip()
+        except Exception as e:
+            print(f"\nLLM Error: {str(e)}")
+            raise
+
     def _analyze_question(self, state: AgentState) -> AgentState:
-        # First, analyze the request using the decision maker
-        request_analysis = self.decision_maker.analyze_request(state["question"])
-        state["logs"]["request_analysis"] = request_analysis
-        
-        # Check for code-related content
-        if "code" in request_analysis["intent"]:
-            # Extract code blocks from the question
-            code_blocks = self._extract_code_blocks(state["question"])
-            if code_blocks:
-                state["current_step"] = "code_analysis"
-                state["code_blocks"] = code_blocks
-                return state
-        
         # Check for file attachments in the question
         try:
             question_data = json.loads(state["question"])
             if "file_url" in question_data:
                 file_url = question_data["file_url"]
-                # Download the file
                 file_data = self._download_file(file_url)
-                # Store in state
                 state["attachment_data"] = file_data
-                # Detect type and set appropriate step
                 file_type = self._detect_file_type(file_url)
                 if file_type == "video":
                     state["current_step"] = "video"
@@ -196,12 +140,12 @@ class BasicAgent:
                     state["current_step"] = "sheet"
                 return state
         except (json.JSONDecodeError, KeyError):
-            pass  # Not a JSON question or no file_url
+            pass
         except Exception as e:
             print(f"\nFile handling error: {str(e)}")
             state["current_step"] = "answer"
             return state
-        
+
         # Regular text question analysis
         prompt = (
             "Decide if this question needs web search. Respond with a Python dict:\n"
@@ -215,53 +159,10 @@ class BasicAgent:
             state["search_query"] = decision.get("search_query", state["question"])
         except Exception as e:
             print(f"\nLLM Error in question analysis: {str(e)}")
-            print(f"Raw response: {raw}")
             state["needs_search"] = True
             state["search_query"] = state["question"]
-            decision = {"parse_error": raw}
             
         state["current_step"] = "search" if state["needs_search"] else "answer"
-        state["history"].append({"step": "analyze", "output": decision})
-        return state
-
-    def _extract_code_blocks(self, text: str) -> List[Dict[str, str]]:
-        """Extract code blocks from text using markdown-style code blocks."""
-        code_blocks = []
-        pattern = r"```(\w+)?\n(.*?)```"
-        matches = re.finditer(pattern, text, re.DOTALL)
-        
-        for match in matches:
-            language = match.group(1) or "python"
-            code = match.group(2).strip()
-            code_blocks.append({
-                "language": language,
-                "code": code
-            })
-            
-        return code_blocks
-
-    def _code_analysis_node(self, state: AgentState) -> AgentState:
-        """Handle code analysis requests."""
-        try:
-            results = []
-            for block in state["code_blocks"]:
-                # Analyze code using the lightweight analyzer
-                analysis = quick_code_stats(block["code"])
-                
-                # Format the results
-                result = {
-                    "language": block["language"],
-                    "analysis": analysis
-                }
-                results.append(result)
-            
-            state["history"].append({"step": "code_analysis", "output": results})
-            state["current_step"] = "answer"
-            
-        except Exception as e:
-            state["logs"]["code_analysis_error"] = str(e)
-            state["current_step"] = "answer"
-            
         return state
 
     def _detect_file_type(self, url: str) -> str:
@@ -320,14 +221,14 @@ class BasicAgent:
 
     def _perform_search(self, state: AgentState) -> AgentState:
         try:
-            results = simple_search(state["search_query"], max_results=6)  # Reduced to 6 results
+            results = simple_search(state["search_query"], max_results=6)
             print("\nSearch Results:")
             for i, s in enumerate(results, 1):
                 print(f"[{i}] {s[:120]}…")
             
             if not results:
                 print("Warning: No search results found")
-                state["needs_search"] = True  # Try reformulating query
+                state["needs_search"] = True
             else:
                 state["needs_search"] = False
                 
@@ -338,34 +239,16 @@ class BasicAgent:
             state["needs_search"] = True
             state["history"].append({"step": "search", "error": str(e)})
             
-        state["current_step"] = "recheck"
+        state["current_step"] = "answer"
         return state
 
-    def _re_evaluate(self, state: AgentState) -> AgentState:
-        """If search returned nothing, reformulate a shorter query."""
-        if state["needs_search"]:
-            state["search_query"] = tighten(state["question"])
-            state["current_step"] = "search"
-        else:
-            state["current_step"] = "answer"
-        return state
-
-    def _extract_boxed_answer(self, text: str) -> str:
-        """Extract answer from boxed format or return original text if no box found."""
-        # Look for text between [box] and [/box] tags
-        box_match = re.search(r'\[box\](.*?)\[/box\]', text, re.DOTALL)
-        if box_match:
-            return box_match.group(1).strip()
-        return text.strip()
-
     def _generate_answer(self, state: AgentState) -> AgentState:
         # Collect all relevant tool outputs
         materials = []
         for h in state["history"]:
-            if h["step"] in {"search", "video", "image", "sheet", "code_analysis"}:
+            if h["step"] in {"search", "video", "image", "sheet"}:
                 materials.append(json.dumps(h.get("output") or h.get("results"), indent=2))
         
-        # Join all materials with proper formatting
         search_block = "\n".join(materials) if materials else "No artefacts available."
         
         prompt = f"""
@@ -383,55 +266,39 @@ Write ANSWER: <answer> on its own line.
             raw = self._call_llm(prompt, 300)
             answer = raw.split("ANSWER:")[-1].strip()
             
-            # Validate answer
             if not answer:
-                print("\nLLM Warning: Empty answer received")
-                print(f"Raw response: {raw}")
                 answer = "I cannot provide a definitive answer at this time."
-            elif any(k in answer.lower() for k in ["i cannot find", "sorry"]):
-                print("\nLLM Warning: LLM indicated it couldn't find an answer")
-                print(f"Raw response: {raw}")
-                answer = "Based on the available information, I cannot provide a complete answer."
             elif "ANSWER:" not in raw:
-                print("\nLLM Warning: Response missing ANSWER: prefix")
-                print(f"Raw response: {raw}")
                 answer = "I cannot provide a definitive answer at this time."
             
             state["final_answer"] = answer
-            state["history"].append({"step": "answer", "output": raw})
             state["current_step"] = "done"
             
         except Exception as e:
             print(f"\nLLM Error in answer generation: {str(e)}")
-            print(f"Question: {state['question']}")
-            print(f"Materials:\n{search_block}")
             state["final_answer"] = "I encountered an error while generating the answer."
             state["current_step"] = "done"
             
         return state
 
-    # ---- Build LangGraph workflow
     def _build_workflow(self) -> Graph:
         sg = StateGraph(state_schema=AgentState)
         
-        # Add all nodes
+        # Add nodes
         sg.add_node("analyze", self._analyze_question)
         sg.add_node("search", self._perform_search)
-        sg.add_node("recheck", self._re_evaluate)
         sg.add_node("answer", self._generate_answer)
         sg.add_node("image", self._image_node)
         sg.add_node("video", self._video_node)
         sg.add_node("sheet", self._sheet_node)
-        sg.add_node("code_analysis", self._code_analysis_node)
 
         # Add edges
         sg.add_edge("analyze", "search")
         sg.add_edge("analyze", "answer")
-        sg.add_edge("search", "recheck")
+        sg.add_edge("search", "answer")
         sg.add_edge("image", "answer")
         sg.add_edge("video", "answer")
         sg.add_edge("sheet", "answer")
-        sg.add_edge("code_analysis", "answer")
 
         def router(state: AgentState):
             return state["current_step"]
@@ -441,19 +308,13 @@ Write ANSWER: <answer> on its own line.
             "answer": "answer",
             "image": "image",
             "video": "video",
-            "sheet": "sheet",
-            "code_analysis": "code_analysis"
-        })
-        sg.add_conditional_edges("recheck", router, {
-            "search": "search", 
-            "answer": "answer"
+            "sheet": "sheet"
         })
         
         sg.set_entry_point("analyze")
         sg.set_finish_point("answer")
         return sg.compile()
 
-    # ---- Public call
     def __call__(self, question: str, task_id: str = "unknown") -> str:
         state: AgentState = {
             "question": question,
@@ -464,7 +325,6 @@ Write ANSWER: <answer> on its own line.
             "search_query": "",
             "task_id": task_id,
             "logs": {},
-            "code_blocks": [],
             "attachment_data": {}
         }
         final_state = self.workflow.invoke(state)
@@ -558,7 +418,6 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
                 "search_query": "",
                 "task_id": task_id,
                 "logs": {},
-                "code_blocks": [],
                 "attachment_data": {}
             }