Update app.py

app.py

@@ -49,6 +49,21 @@ import arxiv
 from transformers import pipeline as hf_pipeline # Renamed to avoid clash with main pipeline
 from youtube_transcript_api import YouTubeTranscriptApi
 
+from typing import List, Literal, TypedDict
+
+import torch
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.tools import BaseTool
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from langchain.chains.Youtubeing import load_qa_chain
+from langchain_community.llms import HuggingFacePipeline
+from langchain.prompts import ChatPromptTemplate, SystemMessage
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langgraph.graph import END, StateGraph
+
 # --- Helper function for python_execution tool ---
 def indent_code(code: str, indent: str = "    ") -> str:
     """Indents multi-line code for execution within a function."""
@@ -203,9 +218,9 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
             reasoning = response_content[reasoning_idx + len("Reasoning:"):action_idx].strip()
             if reasoning.startswith('"') and reasoning.endswith('"'):
                 reasoning = reasoning[1:-1]
-        elif reasoning_idx != -1: 
-             reasoning = response_content[reasoning_idx + len("Reasoning:"):].strip()
-             if reasoning.startswith('"') and reasoning.endswith('"'):
+        elif reasoning_idx != -1:
+            reasoning = response_content[reasoning_idx + len("Reasoning:"):].strip()
+            if reasoning.startswith('"') and reasoning.endswith('"'):
                 reasoning = reasoning[1:-1]
 
         # Attempt to find Action and Action Input
@@ -217,27 +232,25 @@ def parse_agent_response(response_content: str) -> tuple[str, str, str]:
                 action_input = response_content[action_input_idx + len("Action Input:"):].strip()
             else:
                 action = response_content[action_idx + len("Action:"):].strip()
-            
+
             if action.startswith('"') and action.endswith('"'):
                 action = action[1:-1]
             if action_input.startswith('"') and action_input.endswith('"'):
                 action_input = action_input[1:-1]
 
         # Final cleanup for any trailing JSON artifacts if heuristic grabs too much
-        action = action.split('"', 1)[0].strip() 
-        action_input = action_input.split('"', 1)[0].strip() 
+        action = action.split('"', 1)[0].strip()
+        action_input = action_input.split('"', 1)[0].strip()
 
         return reasoning, action, action_input
 
-
 # --- Graph Nodes ---
-
 def should_continue(state: AgentState) -> str:
     """
     Determines if the agent should continue reasoning, use a tool, or end.
     """
     print(f"DEBUG: Entering should_continue. Current context: {state.get('context', {})}")
-    
+
     if state.get("final_answer") is not None:
         print("DEBUG: should_continue -> END (Final Answer set in state)")
         return "end"
@@ -249,12 +262,8 @@ def should_continue(state: AgentState) -> str:
     print("DEBUG: should_continue -> REASON (Default to reasoning)")
     return "reason"
 
-
 # ====== NEW IMPORTS ======
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS  # Use FAISS instead of Chroma
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document
+# Already included at the top.
 
 # ====== DOCUMENT PROCESSING SETUP ======
 def create_vector_store():
@@ -265,25 +274,23 @@ def create_vector_store():
         Document(page_content="Python is a popular programming language created by Guido van Rossum.", metadata={"source": "tech"}),
         Document(page_content="The Eiffel Tower is located in Paris, France.", metadata={"source": "landmarks"}),
     ]
-    
+
     # Initialize embedding model
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-    
+
     # Split documents into chunks
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=500,  # Smaller chunks for better precision
         chunk_overlap=100
     )
     chunks = text_splitter.split_documents(documents)
-    
+
     # Create FAISS vector store
     return FAISS.from_documents(
         documents=chunks,
         embedding=embeddings
     )
 
-
-
 def reasoning_node(state: AgentState) -> AgentState:
     """
     Node for the agent to analyze the question, determine next steps,
@@ -297,24 +304,24 @@ def reasoning_node(state: AgentState) -> AgentState:
     state.setdefault("iterations", 0)
     state.setdefault("current_task", "Understand the question and plan the next step.")
     state.setdefault("current_thoughts", "")
-    
-    state["context"].pop("pending_action", None) 
+
+    state["context"].pop("pending_action", None)
 
     # --- Initialize local HuggingFacePipeline ---
     # Using Mistral-7B-Instruct-v0.2 for better agent performance
-    model_name = "mistralai/Mistral-7B-Instruct-v0.2" 
-    
+    model_name = "mistralai/Mistral-7B-Instruct-v0.2"
+
     print(f"DEBUG: Loading local model: {model_name}...")
-    
+
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
+
     # Load model with optimal settings for GPU if available, else CPU
     # Use bfloat16 for GPUs that support it (NVIDIA Ampere architecture and newer)
     # else float16 for older GPUs or float32 for CPU/fallback.
     # device_map="auto" intelligently distributes the model across available devices.
     model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, 
+        model_name,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         device_map="auto"
     )
 
@@ -342,19 +349,19 @@ def reasoning_node(state: AgentState) -> AgentState:
     # Initialize vector store if not present
     if "vector_store" not in state["context"]:
         state["context"]["vector_store"] = create_vector_store()
-    
+
     vector_store = state["context"]["vector_store"]
-    
+
     # Perform retrieval
     relevant_docs = vector_store.similarity_search(
-        state["question"], 
+        state["question"],
         k=3  # Retrieve top 3 most relevant chunks
     )
-    
+
     # Format context for LLM
     rag_context = "\n\n[Relevant Knowledge]\n"
     rag_context += "\n---\n".join([doc.page_content for doc in relevant_docs])
-    
+
     # ====== MODIFIED PROMPT ======
     # Add RAG context to system prompt
     system_prompt = (
@@ -370,33 +377,43 @@ def reasoning_node(state: AgentState) -> AgentState:
         "- 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_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"
-        "**Retrieved Context:**\n{rag_context}\n\n"  # ADDED RAG CONTEXT
+        "**To provide a final answer, use the Action 'final answer' with the complete answer in 'Action Input'.**\n\n" # Added explicit instruction
+        "**Retrieved Context:**\n{rag_context}\n\n"
         "**Current Context:**\n{context}\n\n"
         "**Previous Reasoning Steps:**\n{reasoning}\n\n"
         "**Current Task:** {current_task}\n"
         "**Current Thoughts:** {current_thoughts}\n\n"
-        # ... [rest of prompt remains same] ...
+        "**Question:** {question}\n\n"
+        "**Expected JSON Output Format:**\n"
+        "```json\n"
+        "{\n"
+        "  \"Reasoning\": \"Your reasoning process to decide the next step, including why a tool is chosen or how an answer is derived.\",\n"
+        "  \"Action\": \"The name of the tool to use (e.g., duckduckgo_search, final answer), or 'No Action' if no tool is needed yet.\",\n"
+        "  \"Action Input\": \"The input for the tool (e.g., 'What is the capital of France?', 'The final answer is Paris.').\"\n"
+        "}\n"
+        "```\n"
+        "Ensure your response is ONLY valid JSON and strictly follows this format."
     )
 
     prompt = ChatPromptTemplate.from_messages([
         SystemMessage(content=system_prompt),
-        *state["history"] 
+        *state["history"]
     ])
 
     formatted_messages = prompt.format_messages(
-        rag_context=rag_context,  # ADD THIS ARGUMENT
+        rag_context=rag_context,
         context=state["context"],
         reasoning=state["reasoning"],
         question=state["question"],
         current_task=state["current_task"],
         current_thoughts=state["current_thoughts"]
     )
-    
+
     # Use tokenizer's chat template for optimal formatting with chat models
     try:
         full_input_string = tokenizer.apply_chat_template(
-            formatted_messages, 
-            tokenize=False, 
+            formatted_messages,
+            tokenize=False,
             add_generation_prompt=True # Adds the assistant's turn start token
         )
     except Exception as e:
@@ -406,47 +423,50 @@ def reasoning_node(state: AgentState) -> AgentState:
     def call_with_retry_local(inputs, retries=3): # Reduced retries for local models as network isn't primary issue
         for attempt in range(retries):
             try:
-                response_text = llm.invoke(inputs) 
-                
+                response_text = llm.invoke(inputs)
+
                 # Strip the prompt from the generated text
                 if response_text.startswith(inputs):
                     content = response_text[len(inputs):].strip()
                 else:
-                    content = response_text.strip() 
+                    content = response_text.strip()
 
                 print(f"DEBUG: RAW LOCAL LLM Response (Attempt {attempt+1}):\n---\n{content}\n---")
-                
+
                 # Attempt to parse to validate structure
-                json.loads(content) 
-                
+                json.loads(content)
+
                 return AIMessage(content=content)
             except json.JSONDecodeError as e:
                 print(f"[Retry {attempt+1}/{retries}] Local LLM returned invalid JSON. Error: {e}. Retrying...")
                 print(f"Invalid JSON content (partial): {content[:200]}...")
                 state["history"].append(AIMessage(content=f"[Parsing Error] The previous LLM output was not valid JSON. Expected format: ```json{{\"Reasoning\": \"...\", \"Action\": \"...\", \"Action Input\": \"...\"}}```. Please ensure your response is ONLY valid JSON and strictly follows the format. Error: {e}"))
-                time.sleep(5) 
-            except Exception as e: 
+                time.sleep(5)
+            except Exception as e:
                 print(f"[Retry {attempt+1}/{retries}] An unexpected error occurred during local LLM call: {e}.")
                 state["history"].append(AIMessage(content=f"[Local LLM Error] Failed to get a response from the local LLM: {e}. Trying again."))
-                time.sleep(10) 
+                time.sleep(10)
+
         raise RuntimeError("Failed after multiple retries due to local Hugging Face model issues or invalid JSON.")
 
-    response = call_with_retry_local(full_input_string) 
+    response = call_with_retry_local(full_input_string)
 
     content = response.content
     reasoning, action, action_input = parse_agent_response(content)
-    
+
     print(f"DEBUG: Parsed Action: '{action}', Action Input: '{action_input[:100]}...'")
 
-    if isinstance(response, AIMessage) and content == response.content: 
-        state["history"].append(AIMessage(content=content)) 
-        
+    if isinstance(response, AIMessage) and content == response.content:
+        state["history"].append(AIMessage(content=content))
+
     state["reasoning"] += f"\nStep {state['iterations'] + 1}: {reasoning}"
     state["iterations"] += 1
-    state["current_thoughts"] = reasoning 
+    state["current_thoughts"] = reasoning
 
-    if "final answer" in action.lower():
-        state["final_answer"] = action_input 
+    # --- FIX: Set final_answer directly if the action is "final answer" ---
+    if action.lower() == "final answer":
+        state["final_answer"] = action_input
+        print(f"DEBUG: Final answer set in state: {state['final_answer']}")
     else:
         state["context"]["pending_action"] = {
             "tool": action,
@@ -454,11 +474,9 @@ def reasoning_node(state: AgentState) -> AgentState:
         }
         state["history"].append(AIMessage(content=f"Agent decided to use tool: {action} with input: {action_input}"))
 
-
     print(f"DEBUG: Exiting reasoning_node. New history length: {len(state['history'])}")
     return state
 
-
 def tool_node(state: AgentState) -> AgentState:
     """
     Node for executing the chosen tool and returning its output.
@@ -476,11 +494,11 @@ 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: 
+    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}") 
+        print(f"ERROR: {error_message}")
         state["history"].append(AIMessage(content=error_message))
-        state["context"].pop("pending_action", None)
+        state["context"].pop("pending_action", None) # Clear invalid pending action
         return state
 
     available_tools = state.get("tools", [])
@@ -493,27 +511,26 @@ 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: 
+            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)}"
             print(f"ERROR: {tool_output}")
 
     state["history"].append(AIMessage(content=f"[{tool_name} output]\n{tool_output}"))
-    
+
     print(f"DEBUG: Exiting tool_node. Tool output added to history. New history length: {len(state['history'])}")
     return state
 
-
 # ====== Agent Graph ======
 def create_agent_workflow(tools: List[BaseTool]): # Use BaseTool for consistency
     workflow = StateGraph(AgentState)
-    
+
     workflow.add_node("reason", reasoning_node)
     workflow.add_node("action", tool_node)
-    
+
     workflow.set_entry_point("reason")
-    
+
     workflow.add_conditional_edges(
         "reason",
         should_continue,
@@ -523,13 +540,12 @@ def create_agent_workflow(tools: List[BaseTool]): # Use BaseTool for consistency
             "end": END
         }
     )
-    
+
     workflow.add_edge("action", "reason")
-    
+
     app = workflow.compile()
     return app
 
-
 # ====== Agent Interface ======
 class BasicAgent:
     def __init__(self):
@@ -540,13 +556,13 @@ class BasicAgent:
             arxiv_search,
             document_qa,
             python_execution,
-            VideoTranscriptionTool() 
+            VideoTranscriptionTool()
         ]
-        
+
         # Pre-initialize RAG vector store
         self.vector_store = create_vector_store()
         self.workflow = create_agent_workflow(self.tools)
-    
+
     def __call__(self, question: str) -> str:
         print(f"\n--- Agent received question: {question[:50]}{'...' if len(question) > 50 else ''} ---")
 
@@ -561,23 +577,34 @@ class BasicAgent:
             "final_answer": None,
             "current_task": "Understand the question and plan the next step.",
             "current_thoughts": "",
-            "tools": self.tools 
+            "tools": self.tools
         }
 
-        final_state = self.workflow.invoke(state)
-
-        if final_state.get("final_answer") is not None:
-            answer = final_state["final_answer"]
-            print(f"--- Agent returning FINAL ANSWER: {answer} ---")
-            return answer
-        else:
-            print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
-            if final_state["history"]:
-                last_message = final_state["history"][-1].content
-                print(f"Last message in history: {last_message}")
-                return f"Agent could not fully answer. Last message: {last_message}"
+        # The invoke method returns an iterator, so we need to consume it to get the final state
+        # LangGraph's invoke will run until the graph reaches an END node.
+        try:
+            # Running the graph
+            # Note: For simple single-path graphs, `invoke` often gives the final state directly.
+            # For more complex graphs, streaming with `stream` and then getting the final state
+            # might be more appropriate if you need intermediate steps.
+            final_state = self.workflow.invoke(state)
+
+            if final_state.get("final_answer") is not None:
+                answer = final_state["final_answer"]
+                print(f"--- Agent returning FINAL ANSWER: {answer} ---")
+                return answer
             else:
-                raise ValueError("Agent finished without providing a final answer and no history messages.")
+                print(f"--- ERROR: Agent finished without setting 'final_answer' for question: {question} ---")
+                if final_state["history"]:
+                    last_message = final_state["history"][-1].content
+                    print(f"Last message in history: {last_message}")
+                    return f"Agent could not fully answer. Last message: {last_message}"
+                else:
+                    raise ValueError("Agent finished without providing a final answer and no history messages.")
+        except Exception as e:
+            print(f"--- FATAL ERROR during agent execution: {e} ---")
+            # In case of an unexpected error, return a helpful message
+            return f"An unexpected error occurred during agent execution: {str(e)}"