turefinem

agent.py

@@ -25,14 +25,20 @@ 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
+from langchain_core.messages import (
+    AnyMessage,
+    SystemMessage,
+    HumanMessage,
+    ToolMessage,
+    BaseMessage,
+)
 import operator
 import json
 from langchain_community.llms import LlamaCpp
 from huggingface_hub import hf_hub_download
 
 # (Keep Constants as is)
-# --- Constants -- - 
+# --- Constants -- -
 default_api_url = "https://agents-course-unit4-scoring.hf.space"
 
 
@@ -60,12 +66,12 @@ def python_interpreter(code: str) -> str:
 def file_reader_tool(url: str) -> str:
     """
     Reads the content of a file from a URL and returns the entire content as a string.
-    It can handle Excel (.xlsx, .xls) and CSV (.csv) files. 
+    It can handle Excel (.xlsx, .xls) and CSV (.csv) files.
     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'
+            "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, headers=headers)
@@ -120,7 +126,7 @@ def youtube_transcript_tool(url: str) -> str:
         ytt_api = YouTubeTranscriptApi()
         try:
             transcript_list = ytt_api.list(video_id)
-            transcript = transcript_list.find_transcript(['en']).fetch()
+            transcript = transcript_list.find_transcript(["en"]).fetch()
             transcript_text = " ".join([item["text"] for item in transcript])
             return transcript_text
         except NoTranscriptFound:
@@ -136,7 +142,7 @@ def web_reader_tool(url: str) -> str:
     """
     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'
+            "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()
@@ -225,7 +231,9 @@ def image_analyzer_tool(url: str) -> str:
     try:
         print(f"Analyzing image from URL: {url}")
         # Initialize the image captioning pipeline
-        captioner = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
+        captioner = pipeline(
+            "image-to-text", model="Salesforce/blip-image-captioning-large"
+        )
 
         # Analyze the image
         caption = captioner(url)
@@ -233,12 +241,14 @@ def image_analyzer_tool(url: str) -> str:
     except Exception as e:
         return f"Error analyzing image: {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 ------
 class BasicAgent:
@@ -261,9 +271,10 @@ class BasicAgent:
 
         self.llm = LlamaCpp(
             model_path=model_path,
-            n_ctx=4096,        # Context window size
-            n_gpu_layers=0,    # Set to 0 to use CPU only
-            verbose=True,      # For debugging
+            n_ctx=4096,  # Context window size
+            n_gpu_layers=0,  # Set to 0 to use CPU only
+            verbose=True,  # For debugging
+            stop=["<|end|>"],
         )
 
         tools = [
@@ -280,7 +291,7 @@ class BasicAgent:
             Tool(
                 name="File Reader from URL",
                 func=file_reader_tool,
-                description="Reads the content of a file from a URL and returns the first 5 rows as a string. It can handle Excel (.xlsx, .xls) and CSV (.csv) files. The URL must point directly to the file.",
+                description="Reads the content of a file from a URL and returns the entire content as a string. It can handle Excel (.xlsx, .xls) and CSV (.csv) files. The URL must point directly to the file.",
             ),
             WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
             Tool(
@@ -309,12 +320,12 @@ class BasicAgent:
                 description="Analyzes an image from a URL and returns a description of its content. The URL must point directly to an image file (e.g., .jpg, .png).",
             ),
         ]
-        prompt_template_str = '''<|system|>
+        prompt_template_str = """<|system|>
 You are a helpful and concise assistant. You have access to the following tools:
 
 {tools}
 
-To use a tool, respond with a JSON blob with "action" and "action_input" keys.
+To use a tool, respond with a JSON blob with "action" and "action_input" keys, inside a ```json markdown block. For the Python Interpreter, you can use a ```python markdown block.
 
 If you have the final answer, respond with:
 FINAL ANSWER: [your answer]
@@ -348,7 +359,7 @@ FINAL ANSWER: 3<|end|>
 
 Question: {input}<|end|>
 <|assistant|>
-'''
+"""
         prompt = PromptTemplate.from_template(prompt_template_str).partial(
             tools=render_text_description(tools),
             tool_names=", ".join([t.name for t in tools]),
@@ -381,13 +392,13 @@ Question: {input}<|end|>
         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:
@@ -421,14 +432,16 @@ Question: {input}<|end|>
 
     def _parse_agent_output(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
         if "FINAL ANSWER:" in llm_output:
+            answer = llm_output.split("FINAL ANSWER:")[-1].strip()
+            answer = answer.split("\n")[0]
             return AgentFinish(
-                return_values={"output": llm_output.split("FINAL ANSWER:")[-1].strip()},
+                return_values={"output": answer},
                 log=llm_output,
             )
 
         # Regex to find ```json ... ``` or raw JSON
-        match = re.search(r'''(?:```json\n)?({.*?})(?:\n```)?''', llm_output, re.DOTALL)
-        
+        match = re.search(r"""(?:```json\n)?({.*?})(?:\n```)?""", llm_output, re.DOTALL)
+
         if match:
             json_str = match.group(1).strip()
             try:
@@ -443,7 +456,7 @@ Question: {input}<|end|>
                 pass  # Fall through to the next check if JSON is invalid
 
         # Regex to find python code
-        match = re.search(r'''```(?:python\n)?(.*?)```''', llm_output, re.DOTALL)
+        match = re.search(r"""```(?:python\n)?(.*?)```""", llm_output, re.DOTALL)
         if match:
             code = match.group(1).strip()
             return AgentAction(
@@ -452,15 +465,18 @@ Question: {input}<|end|>
                 log=llm_output,
             )
 
-        return AgentFinish(return_values={"output": f"Could not parse LLM output: {llm_output}"}, log=llm_output)
+        return AgentFinish(
+            return_values={"output": f"Could not parse LLM output: {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, "chat_history": []})
-            analysis = result['agent_outcome'].return_values["output"]
+            analysis = result["agent_outcome"].return_values["output"]
             print(f"Agent returning analysis: {analysis}")
             return analysis
         except Exception as e:
             print(f"Error during Langchain invocation: {e}")
-            return f"Error analyzing question: {e}"
+            return f"Error analyzing question: {e}"