langgraph way

app.py

@@ -3,10 +3,13 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+import re
 from langchain_huggingface import HuggingFacePipeline
 from langchain_core.prompts import PromptTemplate
-from langchain.agents import AgentExecutor, create_react_agent
-from langchain import hub
+from langchain.agents.format_scratchpad import format_log_to_str
+from langchain_core.runnables import RunnablePassthrough
+from langchain.tools.render import render_text_description
+from langchain_core.agents import AgentAction, AgentFinish
 from langchain.tools import Tool
 from langchain_community.tools import DuckDuckGoSearchResults, WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
@@ -19,6 +22,12 @@ import torch
 from transformers import pipeline
 import av
 import torchaudio
+from langgraph.graph import StateGraph, END, START
+from langgraph.prebuilt import ToolNode
+from typing import TypedDict, Annotated, List, Union
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage, ToolMessage, BaseMessage
+import operator
+import json
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -53,8 +62,11 @@ def file_reader_tool(url: str) -> str:
     The URL must point directly to the file.
     """
     try:
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
+        }
         # Download the file content
-        response = requests.get(url)
+        response = requests.get(url, headers=headers)
         response.raise_for_status()  # Raise an exception for bad status codes
 
         # Get the content type to guess the file type
@@ -114,10 +126,15 @@ def web_reader_tool(url: str) -> str:
     It fetches the HTML and parses it to extract clean text.
     """
     try:
-        response = requests.get(url)
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
+        }
+        response = requests.get(url, headers=headers)
         response.raise_for_status()
         soup = BeautifulSoup(response.content, "html.parser")
         return soup.get_text()
+    except requests.exceptions.HTTPError as e:
+        return f"Error: Could not download the file from the URL. The server returned a {e.response.status_code} error. Please try a different URL or a different tool."
     except requests.exceptions.RequestException as e:
         return f"Error: Could not download the file from the URL: {e}"
     except Exception as e:
@@ -190,6 +207,11 @@ def video_analyzer_tool(url: str) -> str:
     except Exception as e:
         return f"Error analyzing video: {e}"
 
+class AgentState(TypedDict):
+    input: str
+    chat_history: list[BaseMessage]
+    agent_outcome: Union[AgentAction, AgentFinish, None]
+    intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
@@ -203,13 +225,11 @@ class BasicAgent:
                 "Warning: HUGGING_FACE_HUB_TOKEN secret not found. This will fail for gated models."
             )
 
-        # It's better to initialize the model here to avoid reloading it for every call.
-        # Using a smaller model for faster response times.
-        self.llm = HuggingFacePipeline.from_model_id(
-            model_id="google/gemma-2-2b-it",
-            task="text-generation",
-            model_kwargs={"dtype": "float32"},
-            pipeline_kwargs={"max_new_tokens": 256},
+        self.pipeline = pipeline(
+            "text-generation",
+            model="google/gemma-2-2b-it",
+            torch_dtype=torch.float32,
+            device="mps",
         )
 
         tools = [
@@ -250,90 +270,135 @@ class BasicAgent:
                 description="Analyzes the audio track of a video from a URL by transcribing it into text. The URL must point directly to a video file (e.g., .mp4).",
             ),
         ]
-
-        # A more direct, instruction-based prompt with high-quality examples.
-        gemma_react_prompt = '''<start_of_turn>user
-You are a helpful agent. Your goal is to answer the user's question.
-You have access to the following tools:
+        prompt_template_str = """<start_of_turn>user
+You are a helpful assistant tasked with answering questions. You have access to the following tools:
 
 {tools}
 
-To answer the question, you MUST follow these steps:
-1.  **Thought:** First, create a step-by-step plan to answer the question. The plan should be a numbered list. Then, start executing the first step of the plan.
-2.  **Action:** Choose one of the available tools for your 'Action'. Your 'Action' must be one of [{tool_names}].
-3.  **Action Input:** Provide the input for that tool. The input should be a simple string, not a Python expression.
-4.  **Observation:** Wait for the result of the tool.
-5.  **Thought:** Analyze the result and the current plan. If a tool fails, you MUST try to recover by using a different tool or a different input. If you get stuck, you MUST try to search for the information using DuckDuckGo Search Results. If you find a relevant URL, you can use the Web Page Reader tool to read its content. If the plan needs to be revised, state the new plan. Then, proceed with the next step.
-6.  Repeat steps 2-5 until you have enough information to answer the question.
-7.  **Thought:** Once you have the final answer, you MUST use the following template: `FINAL ANSWER: [YOUR FINAL ANSWER]`.
-
-**IMPORTANT RULES for the FINAL ANSWER:**
-*   YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
-*   If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
-*   If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
-*   If you are asked for a comma separated list, Apply the rules above for each element (number or string), ensure there is exactly one space after each comma.
-*   Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
-
-Do not make up observations.
-
-Here are some examples:
-
-Question: In the YouTube video at `https://www.youtube.com/watch?v=example`, the speaker mentions a specific battle. Use the transcript to find the name of the battle, then find out which year it took place.
-
-Thought:
-My plan is:
-1. Get the transcript of the YouTube video.
-2. Find the name of the battle mentioned in the transcript.
-3. Use Wikipedia to find the year the battle took place.
-I will now execute step 1.
-Action: YouTube Transcript Reader
-Action Input: https://www.youtube.com/watch?v=example
-Observation: ...in the video today we are talking about the famous Battle of Hastings...
-Thought: I have completed step 1 and 2 of my plan. The transcript mentions the "Battle of Hastings". Now I need to execute step 3 and find the year it took place. Wikipedia is best for historical events.
-Action: Wikipedia
-Action Input: Battle of Hastings
-Observation: Page: Battle of Hastings. Summary: The Battle of Hastings was fought on 14 October 1066...
-Thought: I have found the year of the battle, 1066. I have completed all steps of my plan.
-FINAL ANSWER: Battle of Hastings, 1066
-
-Question: How many studio albums did Mercedes Sosa release between 2000 and 2009?
-
-Thought:
-My plan is:
-1. Search for Mercedes Sosa's discography.
-2. From the search results, find a reliable source like Wikipedia or a music database.
-3. Read the content of the source.
-4. Count the number of studio albums released between 2000 and 2009.
-I will now execute step 1.
-Action: DuckDuckGo Search Results
-Action Input: Mercedes Sosa discography
-Observation: [snippet: Mercedes Sosa discography. Studio albums. Year, Album, Label. 1959, La voz de la zafra, RCA Victor. ... 2000, Acústico, ... 2005, Corazón libre, ... 2009, Cantora, ... link: https://en.wikipedia.org/wiki/Mercedes_Sosa_discography], ...
-Thought: The search results show a Wikipedia page that seems to contain the discography. I will read the content of that page.
-Action: Web Page Reader
-Action Input: https://en.wikipedia.org/wiki/Mercedes_Sosa_discography
-Observation: ...Studio albums... 2000 Acústico ... 2005 Corazón libre ... 2009 Cantora ...
-Thought: I have the list of albums. Now I need to count the ones between 2000 and 2009. I see "Acústico" (2000), "Corazón libre" (2005), and "Cantora" (2009). That's 3 albums.
-FINAL ANSWER: 3
-
-Begin!
-
-Question: {input}
-Thought:{agent_scratchpad}<end_of_turn>
+To use a tool, you must respond with a JSON blob containing "action" and "action_input" keys.
+The "action" must be one of [{tool_names}].
+
+If you have the final answer, you must respond with:
+FINAL ANSWER: [your answer]
+
+The user has provided the following constraints for the final answer:
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, Apply the rules above for each element (number or string), ensure there is exactly one space after each comma.
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
+
+Here is the conversation history:
+{agent_scratchpad}
+
+Now, answer this question:
+Question: {input}<end_of_turn>
 <start_of_turn>model
-'''
-        prompt = PromptTemplate.from_template(gemma_react_prompt)
+"""
+        prompt = PromptTemplate.from_template(prompt_template_str).partial(
+            tools=render_text_description(tools),
+            tool_names=", ".join([t.name for t in tools]),
+        )
 
-        agent = create_react_agent(self.llm, tools, prompt)
+        def format_scratchpad(intermediate_steps):
+            log = []
+            for action, observation in intermediate_steps:
+                log.append(action.log)
+                log.append(f"Observation: {observation}")
+            return "\n".join(log)
+
+        def llm_wrapper(prompt_value):
+            prompt_str = prompt_value.to_string()
+            outputs = self.pipeline(
+                prompt_str,
+                max_new_tokens=256,
+                return_full_text=False,
+            )
+            return outputs[0]["generated_text"]
 
-        self.agent_executor = AgentExecutor(
-            agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
+        agent_runnable = (
+            RunnablePassthrough.assign(
+                agent_scratchpad=lambda x: format_scratchpad(x["intermediate_steps"])
+            )
+            | prompt
+            | llm_wrapper
+            | self._parse_agent_output
         )
 
+        def run_agent(state):
+            agent_outcome = agent_runnable.invoke(state)
+            return {"agent_outcome": agent_outcome}
+
+        def execute_tools(state):
+            agent_action = state["agent_outcome"]
+            tool_name = agent_action.tool
+            
+            tool_to_use = None
+            for tool in tools:
+                if tool.name == tool_name:
+                    tool_to_use = tool
+                    break
+            
+            if not tool_to_use:
+                output = f"Error: Tool '{tool_name}' not found."
+            else:
+                try:
+                    output = tool_to_use.run(agent_action.tool_input)
+                except Exception as e:
+                    output = f"Error executing tool '{tool_name}': {e}"
+
+            return {"intermediate_steps": [(agent_action, str(output))]}
+
+        def should_continue(state):
+            if isinstance(state["agent_outcome"], AgentFinish):
+                return "end"
+            else:
+                return "continue"
+
+        workflow = StateGraph(AgentState)
+        workflow.add_node("agent", run_agent)
+        workflow.add_node("action", execute_tools)
+        workflow.set_entry_point("agent")
+        workflow.add_conditional_edges(
+            "agent",
+            should_continue,
+            {
+                "continue": "action",
+                "end": END,
+            },
+        )
+        workflow.add_edge("action", "agent")
+        self.agent_executor = workflow.compile()
+
+    def _parse_agent_output(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
+        if "FINAL ANSWER:" in llm_output:
+            return AgentFinish(
+                return_values={"output": llm_output.split("FINAL ANSWER:")[-1].strip()},
+                log=llm_output,
+            )
+        
+        try:
+            llm_output_stripped = llm_output.strip()
+            if llm_output_stripped.startswith("```json"):
+                llm_output_stripped = llm_output_stripped[7:-3].strip()
+            
+            response = json.loads(llm_output_stripped)
+
+            if "action" in response and "action_input" in response:
+                return AgentAction(
+                    tool=response["action"],
+                    tool_input=response["action_input"],
+                    log=llm_output,
+                )
+            else:
+                # Malformed JSON with missing keys
+                return AgentFinish(return_values={"output": f"Malformed JSON. Missing keys.\nLLM Output:\n{llm_output}"}, log=llm_output)
+
+        except json.JSONDecodeError as e:
+            return AgentFinish(return_values={"output": f"JSON Parse Error: {e}\nLLM Output:\n{llm_output}"}, log=llm_output)
+
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         try:
-            result = self.agent_executor.invoke({"input": question})
-            analysis = result["output"]
+            result = self.agent_executor.invoke({"input": question, "chat_history": []})
+            analysis = result['agent_outcome'].return_values["output"]
             print(f"Agent returning analysis: {analysis}")
             return analysis
         except Exception as e: