feat: 优化

.gitignore

@@ -1,3 +1,4 @@
 .venv
 .env
 **/__pycache__
+chroma_db

agent.py

@@ -1,97 +1,233 @@
-"""agent.py
-High-level Agent wrapper used by `app.py`.
-
-Current implementation leverages:
-• Gemini 1.5 Flash via `ChatGoogleGenerativeAI` (requires `GOOGLE_API_KEY` env).
-• The LangChain `zero-shot-react-description` agent (simple & robust).
-• Tools defined in `tools.create_tools`.
-
-Later we can migrate the control loop to LangGraph, but this version already
-provides a working agent that meets the API expectations (callable returning a
-plain string answer).
-"""
-from __future__ import annotations
-
 import os
-import re
-from typing import List
-
 from dotenv import load_dotenv
-from langchain.agents import AgentType, initialize_agent
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_core.callbacks import BaseCallbackHandler
-
-from tools import create_tools  # noqa: E402
-
-# -----------------------------------------------------------------------------
-# Callback for minimal logging (optional)
-# -----------------------------------------------------------------------------
-class PrintCallback(BaseCallbackHandler):
-    """Simple callback that prints agent thoughts for debugging."""
-
-    def on_llm_new_token(self, token: str, **kwargs):  # noqa: D401
-        print(token, end="", flush=True)
-
-
-# -----------------------------------------------------------------------------
-# Helper to strip template markers from final answer
-# -----------------------------------------------------------------------------
-_SYSTEM_PROMPT = (
-    "You are a general AI assistant. I will ask you a question. "
-    "Report your thoughts, and finish your answer with the following template: "
-    "FINAL ANSWER: [YOUR 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 above rules depending of whether the element to be put in the list is "
-    "a number or a string."
-)
-_FINAL_PATTERN = re.compile(r"FINAL ANSWER:\s*(.*)", re.IGNORECASE | re.DOTALL)
-
-
-class Agent:
-    """High-level callable Agent used by `app.py`."""
-
-    def __init__(self, *, temperature: float = 0.0):
-        # Ensure env vars are loaded
-        load_dotenv()
-
-        api_key = os.getenv("GOOGLE_API_KEY")
-        if not api_key:
-            raise EnvironmentError("GOOGLE_API_KEY not found in environment or .env file.")
-
-        # Initialise LLM
-        self.llm = ChatGoogleGenerativeAI(
-            model="gemini-2.5-flash",
-            temperature=temperature,
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from langchain_community.vectorstores import Chroma
+from langchain_core.documents import Document
+import shutil
+import pandas as pd # Ny import för pandas
+import json # För att parsa metadata-kolumnen
+
+load_dotenv()
+
+# Tools:
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"wiki_results": formatted_search_docs}
+
+@tool
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
+
+@tool
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+# load the system prompt from the file
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+
+# Retrieval
+CHROMA_DIR = "./chroma_db"
+CSV_PATH = "./supabase_docs.csv"
+EMBED_MODEL = "sentence-transformers/all-mpnet-base-v2"
+_SIMILARITY_THRESHOLD = 0.2  # lower distance means more similar
+
+embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
+
+if os.path.exists(CHROMA_DIR):
+    print(f"Loading existing ChromaDB from {CHROMA_DIR}")
+    vector_store = Chroma(
+        persist_directory=CHROMA_DIR,
+        embedding_function=embeddings,
+    )
+else:
+    print(f"Creating new ChromaDB at {CHROMA_DIR}, and loading documents from {CSV_PATH}")
+    if os.path.exists(CHROMA_DIR):
+        shutil.rmtree(CHROMA_DIR)
+    os.makedirs(CHROMA_DIR)
+
+    if not os.path.exists(CSV_PATH):
+        raise FileNotFoundError(f"CSV file {CSV_PATH} does not exist")
+    
+    df = pd.read_csv(CSV_PATH)
+    documents = []
+    for i, row in df.iterrows():
+        content = row["content"]
+        
+        question_part = content.split("Final answer :")[0].strip()
+        final_answer_part = content.split("Final answer :")[-1].strip() if "Final answer :" in content else ""
+        
+        try:
+            metadata = json.loads(row["metadata"].replace("'", '"'))
+        except json.JSONDecodeError:
+            metadata = {}
+        
+        metadata["final_answer"] = final_answer_part
+
+        documents.append(Document(page_content=question_part, metadata=metadata))
+
+    if not documents:
+        print("No documents loaded from CSV. ChromaDB will be empty.")
+
+        vector_store = Chroma(
+            persist_directory=CHROMA_DIR,
+            embedding_function=embeddings
         )
-
-        # Aggregate tools
-        self.tools = create_tools()
-
-        # Build agent executor (Zero-Shot ReAct)
-        self.agent_executor = initialize_agent(
-            self.tools,
-            self.llm,
-            agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=False,
-            handle_parsing_errors=True,
-            system_message=_SYSTEM_PROMPT,
-            callbacks=[PrintCallback()],
+    else:
+        vector_store = Chroma.from_documents(
+            documents=documents,
+            embedding=embeddings,
+            persist_directory=CHROMA_DIR,
         )
+        vector_store.persist()
+        print(f"ChromaDB initialized and persisted with {len(documents)} documents from CSV.")
+        
+
+# Retriever tool
+retriever_tool = create_retriever_tool(
+    retriever = vector_store.as_retriever(),
+    name = "Question_Search",
+    description = "A tool to retrieve similar questions from a vector store. The retrieved document's metadata contains the 'final_answer' to the question."
+)
 
-    # ------------------------------------------------------------------
-    # Public API
-    # ------------------------------------------------------------------
-    def __call__(self, question: str) -> str:  # noqa: D401
-        """Return the agent's answer as a plain string (no prefix)."""
-        print(f"Agent received question: {question[:80]}…")
-        raw_answer: str = self.agent_executor.run(question)
-        # Post-process to remove leading template if present
-        match = _FINAL_PATTERN.search(raw_answer)
-        answer = match.group(1).strip() if match else raw_answer.strip()
-        print(f"Agent final answer: {answer}")
-        return answer
+# Agent
+
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    web_search,
+    arvix_search,
+    retriever_tool,
+]
+
+def build_graph_agent():
+    llm = ChatGoogleGenerativeAI(
+        model_name="gemini-1.5-flash",
+        temperature=0.0,
+    )
+
+    llm_with_tools = llm.bind_tools(tools)
+
+    def assistant(state: MessagesState):
+        return {
+            "messages": [llm_with_tools.invoke(state["messages"])],
+        }
+    
+    def retriever(state: MessagesState):
+        query = state["messages"][-1].content
+        similar_docs = vector_store.similarity_search(query, k=3)
+
+        if similar_docs:
+            similar_doc = similar_docs[0]
+            if "final_answer" in similar_doc.metadata and similar_doc.metadata["final_answer"]:
+                answer = similar_doc.metadata["final_answer"]
+            elif "Final answer :" in similar_doc.page_content:
+                answer = similar_doc.page_content.split("Final answer :")[-1].strip()
+            else:
+                answer = similar_doc.page_content.strip()
+
+            return {"messages": [AIMessage(content=answer)]}
+        else:
+            return {"messages": [AIMessage(content="No similar questions found in the knowledge base.")]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
+    builder.set_entry_point("retriever")
+    builder.set_finish_point("retriever")
+    
+    return builder.compile()

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import requests
 import inspect
 from dotenv import load_dotenv
-from agent import Agent
+from agent import build_graph_agent
 import pandas as pd
 import time
 
@@ -16,6 +16,7 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
+        self.graph = build_graph_agent()
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         fixed_answer = "This is a default answer."
@@ -44,7 +45,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
         load_dotenv()
-        agent = Agent()
+        agent = BasicAgent()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -90,8 +91,8 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
         except Exception as e:
              print(f"Error running agent on task {task_id}: {e}")
              results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
-        # wait 60s before next question
-        time.sleep(60)
+        # wait 20s before next question
+        time.sleep(10)
 
     if not answers_payload:
         print("Agent did not produce any answers to submit.")

pyproject.toml

@@ -5,14 +5,17 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "chromadb>=1.0.15",
     "ddgs>=9.0.0",
     "duckduckgo-search>=8.1.1",
     "gradio[oauth]>=5.36.2",
     "langchain>=0.3.26",
+    "langchain-chroma>=0.2.4",
     "langchain-community>=0.3.27",
     "langchain-experimental>=0.3.4",
     "langchain-google-genai>=2.1.7",
     "langchain-huggingface>=0.3.0",
     "langgraph>=0.5.2",
     "requests>=2.32.4",
+    "sentence-transformers>=5.0.0",
 ]

system_prompt.txt

@@ -0,0 +1,18 @@
+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.
+
+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.
+
+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.
+