Update app.py

app.py

@@ -33,10 +33,12 @@ from typing import TypedDict, List, Union, Any, Dict, Optional
 from langchain.schema import HumanMessage, AIMessage, SystemMessage
 from langchain.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph, END
-from langchain_core.tools import BaseTool, Tool
+from langchain_community.llms import HuggingFacePipeline 
+
+# Corrected Tool import: Use 'tool' (lowercase)
+from langchain_core.tools import BaseTool, tool # <--- CHANGED HERE
 
 # Hugging Face local model imports
-from langchain_community.llms import HuggingFacePipeline 
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch 
 
@@ -53,7 +55,7 @@ def indent_code(code: str, indent: str = "    ") -> str:
     return "\n".join(indent + line for line in code.splitlines())
 
 # --- Tool Definitions ---
-@Tool
+@tool # <--- CHANGED HERE
 def duckduckgo_search(query: str) -> str:
     """Search web using DuckDuckGo. Returns top 3 results."""
     print(f"DEBUG: duckduckgo_search called with: {query}")
@@ -66,7 +68,7 @@ def duckduckgo_search(query: str) -> str:
     except Exception as e:
         return f"Error performing DuckDuckGo search: {str(e)}"
 
-@Tool
+@tool # <--- CHANGED HERE
 def wikipedia_search(query: str) -> str:
     """Get Wikipedia summaries. Returns first 3 sentences."""
     print(f"DEBUG: wikipedia_search called with: {query}")
@@ -79,7 +81,7 @@ def wikipedia_search(query: str) -> str:
     except Exception as e:
         return f"Error performing Wikipedia search: {str(e)}"
 
-@Tool
+@tool # <--- CHANGED HERE
 def arxiv_search(query: str) -> str:
     """Search academic papers on arXiv. Returns top 3 results."""
     print(f"DEBUG: arxiv_search called with: {query}")
@@ -98,7 +100,7 @@ def arxiv_search(query: str) -> str:
     except Exception as e:
         return f"Error performing ArXiv search: {str(e)}"
 
-@Tool
+@tool # <--- CHANGED HERE
 def document_qa(input_str: str) -> str:
     """Answer questions from documents. Input format: 'document_text||question'"""
     print(f"DEBUG: document_qa called with: {input_str}")
@@ -107,15 +109,15 @@ def document_qa(input_str: str) -> str:
             return "Invalid format. Input must be: 'document_text||question'"
         
         context, question = input_str.split('||', 1)
-        # Load QA model on first call or ensure it's loaded once globally
-        # This makes it a bit slow on first call, but avoids re-loading.
-        # For production, consider loading this outside the tool function.
+        # Load QA model on first call or ensure it's loaded once globally.
+        # It's better to load once in __init__ for BasicAgent if possible,
+        # but this lazy loading prevents initial heavy load if tool is not used.
         qa_model = hf_pipeline('question-answering', model='deepset/roberta-base-squad2')
         return qa_model(question=question, context=context)['answer']
     except Exception as e:
         return f"Error answering question from document: {str(e)}"
 
-@Tool
+@tool # <--- CHANGED HERE
 def python_execution(code: str) -> str:
     """Execute Python code and return output.
     The code should assign its final result to a variable named 'result'.
@@ -126,8 +128,10 @@ def python_execution(code: str) -> str:
         # Create isolated environment
         env = {}
         # Wrap code in a function to isolate scope and capture 'result'
-        exec(f"def __exec_fn__():\n{indent_code(code)}\nresult = __exec_fn__()", globals(), env)
-        return str(env.get('result', 'No explicit result assigned to "result" variable.'))
+        # The exec function is used carefully here. In a production environment,
+        # consider a more robust and secure sandbox (e.g., Docker, dedicated service).
+        exec(f"def __exec_fn__():\n{indent_code(code)}\n_result_value = __exec_fn__()", globals(), env)
+        return str(env.get('_result_value', 'No explicit result assigned to "_result_value" variable.'))
     except Exception as e:
         return f"Python execution error: {str(e)}"
 
@@ -139,10 +143,10 @@ class VideoTranscriptionTool(BaseTool):
         print(f"DEBUG: transcript_video called with: {url_or_id}")
         video_id = None
         # Basic parsing for common YouTube URL formats
-        if "youtu.be/" in url_or_id:
-            video_id = url_or_id.split("youtu.be/")[1].split("?")[0]
-        elif "youtube.com/watch?v=" in url_or_id:
+        if "youtube.com/watch?v=" in url_or_id:
             video_id = url_or_id.split("v=")[1].split("&")[0]
+        elif "youtu.be/" in url_or_id:
+            video_id = url_or_id.split("youtu.be/")[1].split("?")[0]
         elif len(url_or_id.strip()) == 11 and not ("http://" in url_or_id or "https://" in url_or_id):
             video_id = url_or_id.strip() # Assume it's just the ID
 
@@ -151,12 +155,9 @@ class VideoTranscriptionTool(BaseTool):
 
         try:
             transcription = YouTubeTranscriptApi.get_transcript(video_id)
-            # You can add include_timestamps logic here if needed,
-            # but for simplicity, let's just return the text.
             return " ".join([part['text'] for part in transcription])
 
         except Exception as e:
-            # Catch specific errors for better messages, e.g., NoTranscriptFound
             return f"Error fetching transcript for video ID '{video_id}': {str(e)}. It might not have an English transcript, or the video is unavailable."
 
     def _arun(self, *args, **kwargs):
@@ -172,7 +173,7 @@ class AgentState(TypedDict):
     final_answer: Union[str, float, int, None]
     current_task: str
     current_thoughts: str
-    tools: List[Tool] # Make sure tools are passed via state
+    tools: List[BaseTool] # Make sure tools are passed via state, using BaseTool type
 
 # --- Utility Functions ---
 def parse_agent_response(response_content: str) -> tuple[str, str, str]:
@@ -290,8 +291,8 @@ def reasoning_node(state: AgentState) -> AgentState:
         max_new_tokens=1024, # Increased max_new_tokens for potentially longer JSON
         temperature=0.1,    # Keep low for factual, tool-use tasks
         do_sample=True,     # Allow some sampling
-        top_p=0.9,
-        repetition_penalty=1.1, # Help avoid repetitive output
+        top_p=0.9,          # Nucleus sampling
+        repetition_penalty=1.1, # Avoid repetition
         # device_map handled by model loading
     )
 
@@ -314,7 +315,7 @@ def reasoning_node(state: AgentState) -> AgentState:
         "- Use **wikipedia_search** for encyclopedic information, historical context, or detailed topics. Provide a concise search term.\n"
         "- Use **arxiv_search** for scientific papers, research, or cutting-edge technical information. Provide a concise search query.\n"
         "- Use **document_qa** when the question explicitly refers to a specific document or when you have content to query. Input format: 'document_text||question'.\n"
-        "- Use **python_execution** for complex calculations, data manipulation, or logical operations that cannot be done with simple reasoning. Always provide the full Python code, ensuring it's valid and executable, and assign the final result to a variable named 'result' if applicable.\n"
+        "- Use **python_execution** for complex calculations, data manipulation, or logical operations that cannot be done with simple reasoning. Always provide the full Python code, ensuring it's valid and executable, and assign the final result to a variable named '_result_value' (e.g., '_result_value = 1 + 1').\n"
         "- Use **transcript_video** for any question involving video or audio content (e.g., YouTube). Provide the full YouTube URL or video ID.\n\n"
         "**Current Context:**\n{context}\n\n"
         "**Previous Reasoning Steps:**\n{reasoning}\n\n"
@@ -430,7 +431,7 @@ def tool_node(state: AgentState) -> AgentState:
     tool_name = tool_call_dict.get("tool")
     tool_input = tool_call_dict.get("input")
 
-    if not tool_name or tool_input is None: # Check for None as empty string "" might be valid input for some tools
+    if not tool_name or tool_input is None: 
         error_message = f"[Tool Error] Invalid action request from LLM: Tool name '{tool_name}' or input '{tool_input}' was empty or None. LLM needs to provide valid 'Action' and 'Action Input'."
         print(f"ERROR: {error_message}") 
         state["history"].append(AIMessage(content=error_message))
@@ -438,7 +439,6 @@ def tool_node(state: AgentState) -> AgentState:
         return state
 
     available_tools = state.get("tools", [])
-    # Use Tool.name to match, which is what @Tool decorator sets
     tool_fn = next((t for t in available_tools if t.name == tool_name), None)
 
     if tool_fn is None:
@@ -448,7 +448,7 @@ def tool_node(state: AgentState) -> AgentState:
         try:
             print(f"DEBUG: Invoking tool '{tool_name}' with input: '{tool_input[:100]}...'")
             tool_output = tool_fn.run(tool_input)
-            if not tool_output and tool_output is not False: # Ensure 'False' or 0 are not treated as empty
+            if not tool_output and tool_output is not False: 
                 tool_output = f"[{tool_name} output] No specific result found for '{tool_input}'. The tool might have returned an empty response."
         except Exception as e:
             tool_output = f"[Tool Error] An error occurred while running '{tool_name}': {str(e)}"
@@ -461,7 +461,7 @@ def tool_node(state: AgentState) -> AgentState:
 
 
 # ====== Agent Graph ======
-def create_agent_workflow(tools: List[Tool]):
+def create_agent_workflow(tools: List[BaseTool]): # Use BaseTool for consistency
     workflow = StateGraph(AgentState)
     
     workflow.add_node("reason", reasoning_node)
@@ -495,7 +495,7 @@ class BasicAgent:
             arxiv_search,
             document_qa,
             python_execution,
-            VideoTranscriptionTool() # Instantiate the class-based tool
+            VideoTranscriptionTool() 
         ]
         self.workflow = create_agent_workflow(self.tools)
 
@@ -511,10 +511,9 @@ class BasicAgent:
             "final_answer": None,
             "current_task": "Understand the question and plan the next step.",
             "current_thoughts": "",
-            "tools": self.tools # Pass tools through state
+            "tools": self.tools 
         }
 
-        # The invoke method returns the final state after execution
         final_state = self.workflow.invoke(state)
 
         if final_state.get("final_answer") is not None: