Update agent.py

agent.py

@@ -1,253 +1,242 @@
+# ----------------------------------------------------------
+# Section 0: Imports
+# ----------------------------------------------------------
 import json
 import os
 import pickle
 import re
+import subprocess
+import textwrap
+import base64
 from datetime import datetime, timedelta
 from io import BytesIO
 from pathlib import Path
 from typing import List
 
+# Third-party libraries
 import requests
 from cachetools import TTLCache
+from PIL import Image
+
+# LangChain and associated libraries
 from langchain.schema import Document
+from langchain.tools.retriever import create_retriever_tool
 from langchain_community.vectorstores import FAISS
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader # Added loaders
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.messages import SystemMessage, HumanMessage, AIMessageChunk
+from langchain_groq import ChatGroq
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFaceEndpoint, ChatHuggingFace
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_core.tools import tool
-from dotenv import load_dotenv
 
+# Environment variable loading
+from dotenv import load_dotenv
 load_dotenv()
 
 # ----------------------------------------------------------
-# 0. Constants
+# Section 1: Constants and Configuration
 # ----------------------------------------------------------
 JSONL_PATH   = Path("metadata.jsonl")
 FAISS_CACHE  = Path("faiss_index.pkl")
 EMBED_MODEL  = "sentence-transformers/all-mpnet-base-v2"
-RETRIEVER_K  = 5
-CACHE_TTL    = 600
-CACHE        = TTLCache(maxsize=256, ttl=CACHE_TTL)
-
-# ----------------------------------------------------------
-# 1. Build / load FAISS retriever
-# ----------------------------------------------------------
-embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
-
-if FAISS_CACHE.exists():
-    with open(FAISS_CACHE, "rb") as f:
-        vector_store = pickle.load(f)
-else:
-    if not JSONL_PATH.exists():
-        raise FileNotFoundError("metadata.jsonl not found")
-    docs = []
-    with open(JSONL_PATH, "rt", encoding="utf-8") as f:
-        for line in f:
-            rec = json.loads(line)
-            content = f"Question: {rec['Question']}\n\nFinal answer: {rec['Final answer']}"
-            docs.append(Document(page_content=content, metadata={"source": rec["task_id"]}))
-    vector_store = FAISS.from_documents(docs, embeddings)
-    with open(FAISS_CACHE, "wb") as f:
-        pickle.dump(vector_store, f)
-
-retriever = vector_store.as_retriever(search_kwargs={"k": RETRIEVER_K})
-
-# ----------------------------------------------------------
-# 2. Caching helper
-# ----------------------------------------------------------
-def cached_get(key: str, fetch_fn):
-    if key in CACHE:
-        return CACHE[key]
-    val = fetch_fn()
-    CACHE[key] = val
-    return val
-
-# ----------------------------------------------------------
-# 3. Tools
-# ----------------------------------------------------------
-@tool
-def python_repl(code: str) -> str:
-    """Execute Python code and return stdout/stderr."""
-    import subprocess, textwrap
-    code = textwrap.dedent(code).strip()
-    try:
-        result = subprocess.run(
-            ["python", "-c", code],
-            capture_output=True,
-            text=True,
-            timeout=5,
-        )
-        return result.stdout if not result.stderr else f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
-    except subprocess.TimeoutExpired:
-        return "Execution timed out (>5s)."
-
-@tool
-def describe_image(image_source: str) -> str:
-    """Describe an image from local path or URL with Gemini vision."""
-    import base64
-    from PIL import Image
-
-    if image_source.startswith("http"):
-        img = Image.open(BytesIO(requests.get(image_source, timeout=10).content))
-    else:
-        img = Image.open(image_source)
-
-    buffered = BytesIO()
-    img.convert("RGB").save(buffered, format="JPEG")
-    b64 = base64.b64encode(buffered.getvalue()).decode()
-
-    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    msg = HumanMessage(
-        content=[
-            {"type": "text", "text": "Describe this image in detail."},
-            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
+RETRIEVER_K  = 5  # Number of similar documents to retrieve
+CACHE_TTL    = 600 # Cache API calls for 10 minutes
+# Global cache object for API calls
+API_CACHE = TTLCache(maxsize=256, ttl=CACHE_TTL)
+
+# ----------------------------------------------------------
+# Section 2: The Agent Class
+# ----------------------------------------------------------
+class MyAgent:
+    """
+    Encapsulates the agent's state, including LLMs, retriever, and tools.
+    This class-based approach ensures clean management of dependencies.
+    """
+
+    def __init__(self, provider: str = "google"):
+        """
+        Initializes the agent, setting up LLMs and the FAISS retriever.
+        Args:
+            provider (str): The LLM provider to use ('google', 'groq', 'huggingface').
+        """
+        print(f"Initializing agent with provider: {provider}")
+        
+        self.llm = self._build_llm(provider)
+        self.vision_llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
+        self.retriever = self._get_retriever()
+        
+    def _get_retriever(self):
+        """Builds or loads the FAISS retriever from a local cache."""
+        embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
+
+        if FAISS_CACHE.exists():
+            print(f"Loading FAISS index from cache: {FAISS_CACHE}")
+            with open(FAISS_CACHE, "rb") as f:
+                vector_store = pickle.load(f)
+        else:
+            print("FAISS cache not found. Building new index from metadata.jsonl...")
+            if not JSONL_PATH.exists():
+                raise FileNotFoundError(f"{JSONL_PATH} not found. Cannot build vector store.")
+            docs = [Document(page_content=f"Question: {rec['Question']}\n\nFinal answer: {rec['Final answer']}", metadata={"source": rec["task_id"]}) for rec in (json.loads(line) for line in open(JSONL_PATH, "rt", encoding="utf-8"))]
+            if not docs: raise ValueError("No documents found in metadata.jsonl.")
+            vector_store = FAISS.from_documents(docs, embeddings)
+            with open(FAISS_CACHE, "wb") as f: pickle.dump(vector_store, f)
+            print(f"FAISS index built and saved to cache: {FAISS_CACHE}")
+
+        return vector_store.as_retriever(search_kwargs={"k": RETRIEVER_K})
+
+    def _build_llm(self, provider: str):
+        """Helper to build the main text-based LLM based on the chosen provider."""
+        if provider == "google": return ChatGoogleGenerativeAI(model="gemini-1.5-pro-latest", temperature=0)
+        elif provider == "groq": return ChatGroq(model_name="llama3-70b-8192", temperature=0)
+        elif provider == "huggingface": return ChatHuggingFace(llm=HuggingFaceEndpoint(repo_id="Qwen/Qwen2.5-Coder-32B-Instruct", temperature=0))
+        else: raise ValueError("Provider must be 'google', 'groq', or 'huggingface'")
+
+    def _cached_get(self, key: str, fetch_fn):
+        """Helper for caching API calls."""
+        if key in API_CACHE: return API_CACHE[key]
+        val = fetch_fn()
+        API_CACHE[key] = val
+        return val
+
+    # --- Tool Definitions as Class Methods ---
+
+    @tool
+    def python_repl(self, code: str) -> str:
+        """Executes a string of Python code and returns the stdout/stderr."""
+        code = textwrap.dedent(code).strip()
+        try:
+            result = subprocess.run(["python", "-c", code], capture_output=True, text=True, timeout=10, check=False)
+            if result.returncode == 0: return f"Execution successful.\nSTDOUT:\n```\n{result.stdout}\n```"
+            else: return f"Execution failed.\nSTDOUT:\n```\n{result.stdout}\n```\nSTDERR:\n```\n{result.stderr}\n```"
+        except subprocess.TimeoutExpired: return "Execution timed out (>10s)."
+
+    @tool
+    def describe_image(self, image_source: str) -> str:
+        """Describes an image from a local file path or a URL using Gemini vision."""
+        try:
+            if image_source.startswith("http"):
+                img = Image.open(BytesIO(requests.get(image_source, timeout=10).content))
+            else:
+                img = Image.open(image_source)
+            buffered = BytesIO()
+            img.convert("RGB").save(buffered, format="JPEG")
+            b64_string = base64.b64encode(buffered.getvalue()).decode()
+            msg = HumanMessage(content=[{"type": "text", "text": "Describe this image in detail."}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64_string}"}}])
+            return self.vision_llm.invoke([msg]).content
+        except Exception as e: return f"Error processing image: {e}"
+
+    @tool
+    def web_search(self, query: str) -> str:
+        """Performs a web search using Tavily and returns a compilation of results."""
+        key = f"web:{query}"
+        results = self._cached_get(key, lambda: TavilySearchResults(max_results=5).invoke(query))
+        return "\n\n---\n\n".join([f"Source: {res['url']}\nContent: {res['content']}" for res in results])
+
+    @tool
+    def wiki_search(self, query: str) -> str:
+        """Searches Wikipedia and returns the top 2 results."""
+        key = f"wiki:{query}"
+        docs = self._cached_get(key, lambda: WikipediaLoader(query=query, load_max_docs=2, doc_content_chars_max=2000).load())
+        return "\n\n---\n\n".join([f"Source: {d.metadata['source']}\n\n{d.page_content}" for d in docs])
+
+    @tool
+    def arxiv_search(self, query: str) -> str:
+        """Searches Arxiv for scientific papers and returns the top 2 results."""
+        key = f"arxiv:{query}"
+        docs = self._cached_get(key, lambda: ArxivLoader(query=query, load_max_docs=2).load())
+        return "\n\n---\n\n".join([f"Source: {d.metadata['source']}\nPublished: {d.metadata['Published']}\nTitle: {d.metadata['Title']}\n\nSummary:\n{d.page_content}" for d in docs])
+
+    def get_tools(self) -> list:
+        """Returns a list of all tools available to the agent."""
+        tools_list = [
+            self.python_repl,
+            self.describe_image,
+            self.web_search,
+            self.wiki_search,
+            self.arxiv_search,
         ]
-    )
-    return llm.invoke([msg]).content
-
-@tool
-def web_search(query: str) -> str:
-    """Smart web search with 3 keyword variants, cached."""
-    from langchain_community.tools.tavily_search import TavilySearchResults
-
-    keywords = [query, query.replace(" ", " OR "), f'"{query}"']
-    seen = set()
-    results = []
-    for kw in keywords:
-        key = f"web:{kw}"
-        snippets = cached_get(
-            key,
-            lambda: TavilySearchResults(max_results=3, include_raw_content=True).invoke(kw),
+        retriever_tool = create_retriever_tool(
+            retriever=self.retriever,
+            name="retrieve_examples",
+            description="Retrieve solved questions and answers similar to the user's query.",
         )
-        for s in snippets:
-            if s["url"] not in seen:
-                seen.add(s["url"])
-                results.append(s["content"][:2000])
-            if len(results) >= 5:
-                break
-    return "\n\n---\n\n".join(results)
-
-@tool
-def wiki_search(query: str) -> str:
-    from langchain_community.document_loaders import WikipediaLoader
-    key = f"wiki:{query}"
-    docs = cached_get(
-        key,
-        lambda: WikipediaLoader(query=query, load_max_docs=2).load(),
-    )
-    return "\n\n---\n\n".join(
-        f'<Document source="{d.metadata.get("source", "")}">\n{d.page_content}\n</Document>'
-        for d in docs
-    )
-
-@tool
-def arxiv_search(query: str) -> str:
-    from langchain_community.document_loaders import ArxivLoader
-    key = f"arxiv:{query}"
-    docs = cached_get(
-        key,
-        lambda: ArxivLoader(query=query, load_max_docs=2).load(),
-    )
-    return "\n\n---\n\n".join(
-        f'<Document source="{d.metadata.get("source", "")}">\n{d.page_content[:2000]}...\n</Document>'
-        for d in docs
-    )
+        tools_list.append(retriever_tool)
+        return tools_list
 
 # ----------------------------------------------------------
-# 4. System prompt
+# Section 3: System Prompt
 # ----------------------------------------------------------
 SYSTEM_PROMPT = (
-    """You are a helpful assistant tasked with answering questions using a set of tools.
-
-Your final answer must strictly follow this format:
-FINAL ANSWER: [ANSWER]
-
-Only write the answer in that exact format. Do not explain anything. Do not include any other text.
+    """You are a helpful and expert assistant designed to answer questions accurately and concisely.
 
-If you are provided with a similar question and its final answer, and the current question is **exactly the same**, then simply return the same final answer without using any tools.
+**Instructions:**
+1.  **Analyze the Question:** Carefully understand what is being asked.
+2.  **Use Tools:** You have a set of tools to find information. Use them logically.
+3.  **Synthesize the Answer:** Based on the information from the tools, formulate your final answer.
+4.  **Format the Output:** Your final response MUST be in the following format and nothing else:
+    
+    FINAL ANSWER: [Your concise and accurate answer here]
 
-Only use tools if the current question is different from the similar one.
-
-Examples:
-- FINAL ANSWER: FunkMonk
-- FINAL ANSWER: Paris
-- FINAL ANSWER: 128
-If you do not follow this format exactly, your response will be considered incorrect.
+If the `retrieve_examples` tool provides an answer to an identical question, use that answer. Otherwise, use your tools to find the correct answer for the current question.
 """
 )
 
 # ----------------------------------------------------------
-# 5. Manual LangGraph construction
-# ----------------------------------------------------------
-tools_list = [python_repl, describe_image, web_search, wiki_search, arxiv_search]
-
-# retriever tool
-from langchain.tools.retriever import create_retriever_tool
-tools_list.append(
-    create_retriever_tool(
-        retriever=retriever,
-        name="retrieve_examples",
-        description="Retrieve up to 5 solved questions similar to the user query.",
-    )
-)
-
-# ----------------------------------------------------------
-# provider switcher
-# ----------------------------------------------------------
-def build_llm(provider: str = "groq"):
-    if provider == "google":
-        return ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "groq":
-        return ChatGroq(model="llama-3.3-70b-versatile", temperature=0)
-    elif provider == "huggingface":
-        return ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
-                temperature=0,
-            )
-        )
-    else:
-        raise ValueError("provider must be 'google', 'groq', or 'huggingface'")
-
-llm = build_llm("google")  # or "groq", "huggingface"
-llm_with_tools = llm.bind_tools(tools_list)
-
-def assistant(state: MessagesState):
-    """LLM node that can call tools."""
-    return {"messages": [llm_with_tools.invoke(state["messages"])]}
-
-def retriever_node(state: MessagesState):
-    """First node: fetch examples and prepend them."""
-    user_query = state["messages"][-1].content
-    docs = retriever.invoke(user_query)
-    if docs:
-        example_text = "\n\n---\n\n".join(d.page_content for d in docs)
-        example_msg = HumanMessage(
-            content=f"Here are {len(docs)} similar solved examples:\n\n{example_text}"
-        )
-        return {"messages": [SYSTEM_PROMPT] + state["messages"] + [example_msg]}
-    return {"messages": [SYSTEM_PROMPT] + state["messages"]}
-
-builder = StateGraph(MessagesState)
-builder.add_node("retriever", retriever_node)
-builder.add_node("assistant", assistant)
-builder.add_node("tools", ToolNode(tools_list))
-builder.add_edge(START, "retriever")
-builder.add_edge("retriever", "assistant")
-builder.add_conditional_edges("assistant", tools_condition)
-builder.add_edge("tools", "assistant")
-
-agent = builder.compile()
-
-# ----------------------------------------------------------
-# 6. Quick streaming test
+# Section 4: Factory Function for Agent Executor
+# ----------------------------------------------------------
+def create_agent_executor(provider: str = "google"):
+    """Factory function to create and compile the LangGraph agent executor."""
+    my_agent_instance = MyAgent(provider=provider)
+    tools_list = my_agent_instance.get_tools()
+    llm_with_tools = my_agent_instance.llm.bind_tools(tools_list)
+
+    def retriever_node(state: MessagesState):
+        """First node: retrieves examples and prepends them to the message history."""
+        user_query = state["messages"][-1].content
+        docs = my_agent_instance.retriever.invoke(user_query)
+        messages = [SystemMessage(content=SYSTEM_PROMPT)]
+        if docs:
+            example_text = "\n\n---\n\n".join(d.page_content for d in docs)
+            example_msg = AIMessage(content=f"I have found {len(docs)} similar solved examples:\n\n{example_text}", name="ExampleRetriever")
+            messages.append(example_msg)
+        messages.extend(state["messages"])
+        return {"messages": messages}
+
+    def assistant_node(state: MessagesState):
+        """Main assistant node: calls the LLM with the current state to decide the next action."""
+        result = llm_with_tools.invoke(state["messages"])
+        return {"messages": [result]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever_node)
+    builder.add_node("assistant", assistant_node)
+    builder.add_node("tools", ToolNode(tools_list))
+    
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition, {"tools": "tools", "__end__": "__end__"})
+    builder.add_edge("tools", "assistant")
+
+    agent_executor = builder.compile()
+    print("Agent Executor created successfully.")
+    return agent_executor
+
+# ----------------------------------------------------------
+# Section 5: Direct Execution Block for Testing
 # ----------------------------------------------------------
 if __name__ == "__main__":
-    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
-    print("Agent thinking …")
-    for chunk in agent.stream({"messages": [("user", question)]}, stream_mode="values"):
-        last = chunk["messages"][-1]
-        if hasattr(last, "content"):
-            print(last.content, end="", flush=True)
+    """direct testing of the agent's logic."""
+    print("--- Running Agent in Test Mode ---")
+    agent = create_agent_executor(provider="google")
+    question = "According to wikipedia, what is the main difference between a lama and an alpaca?"
+    print(f"\nTest Question: {question}\n\n--- Agent Thinking... ---\n")
+    
+    for chunk in agent.stream({"messages": [("user", question)]}):
+        for key, value in chunk.items():
+            if value['messages']:
+                message = value['messages'][-1]
+                if message.content: print(f"--- Node: {key} ---\n{message.content}\n")