Got 45%

agent.py

@@ -3,11 +3,14 @@ import warnings
 from typing import Annotated, TypedDict
 
 from dotenv import load_dotenv
+from langchain_chroma import Chroma
 from langchain_community.cache import SQLiteCache
 from langchain_core.globals import set_llm_cache
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.system import SystemMessage
 from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph.message import add_messages
 from langgraph.graph.state import END, START, CompiledStateGraph, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
@@ -22,6 +25,11 @@ load_dotenv()
 # set_llm_cache(InMemoryCache())
 set_llm_cache(SQLiteCache(database_path=".langchain_cache.db"))
 
+# Initialize RAG vector store
+CHROMA_PATH = "./chroma_gaia_db"
+EMBEDDINGS = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+VECTOR_STORE = Chroma(persist_directory=CHROMA_PATH, embedding_function=EMBEDDINGS)
+
 
 class AgentState(TypedDict):
     """State passed between nods in the graph"""
@@ -38,25 +46,22 @@ def load_system_prompt() -> SystemMessage:
 SYSTEM_PROMPT: SystemMessage = load_system_prompt()
 
 
-
 class GaiaAgent:
     """
     A LangGraph agent for Gaia questions
     """
 
-    def __init__(self, model: str = "gpt-4o", temperature: float = 0.0):
+    def __init__(self, model: str, temperature: float):
         """Initialize the agent with a specific model"""
-        from langchain_openai import ChatOpenAI
+        import asyncio
 
         from tools import get_tools
 
-        self.tools = get_tools()
+        self.tools = asyncio.run(get_tools())
 
         if model.startswith("glm"):
             api_key = SecretStr(secret_value=os.getenv("ZAI_API_KEY", ""))
-            api_base = "https://api.z.ai/api/paas/v4/"
-            if os.getenv("ZAI_USE_CODING_PLAN", "f") == "t":
-                api_base = "https://api.z.ai/api/coding/paas/v4/"
+            api_base = "https://api.z.ai/api/coding/paas/v4/"
         else:
             api_key = SecretStr(secret_value=os.getenv("OPENAI_API_KEY") or "")
             api_base = None
@@ -75,30 +80,44 @@ class GaiaAgent:
 
         graph = StateGraph(AgentState)
 
-        # graph.add_node("retriever", self._retriever_node)
         graph.add_node("agent", self._agent_node)
         graph.add_node("tools", ToolNode(self.tools))
-        # graph.add_node("tools", self._tools_node)
 
-        # graph.add_edge(START, "retriever")
-        # graph.add_edge("retriever", "agent")
         graph.add_edge(START, "agent")
         graph.add_conditional_edges("agent", tools_condition)
         graph.add_edge("tools", "agent")
 
-        return graph.compile()
+        memory = MemorySaver()
+        return graph.compile(checkpointer=memory)
 
     def _retriever_node(self, state: AgentState) -> AgentState:
-        """Retrieve similar questions from vector store."""
-        question = state["messages"][0].content
+        """Retrieve similar questions and inject solving strategy into the question."""
+        original_question = state["messages"][0].content
 
-        similar_docs = VECTOR_STORE.similarity_search(question, k=1)
+        similar_docs = VECTOR_STORE.similarity_search(original_question, k=1)
 
         if similar_docs:
-            example_msg = HumanMessage(
-                content=f"Here is a similar question and answer for reference:\n\n{similar_docs[0].page_content}"
+            doc = similar_docs[0]
+            steps = (
+                doc.page_content.split("Steps to solve:")[-1]
+                .split("Tools needed:")[0]
+                .strip()
             )
-            return {"messages": [SYSTEM_PROMPT] + state["messages"] + [example_msg]}
+            tools = doc.metadata.get("tools", "")
+
+            # Build enhanced question with strategy
+            enhanced_question = f"""{original_question}
+
+---
+Strategy (from similar solved question):
+{steps}
+
+Tools needed: {tools}
+
+Follow a similar approach to solve the question above."""
+
+            enhanced_msg = HumanMessage(content=enhanced_question)
+            return {"messages": [SYSTEM_PROMPT, enhanced_msg]}
 
         return {"messages": [SYSTEM_PROMPT] + state["messages"]}
 
@@ -115,7 +134,7 @@ class GaiaAgent:
 
         return result
 
-    def __call__(self, question: str) -> str:
+    async def __call__(self, question: str) -> str:
         """
         Run the agent on a given question and return the answer
 
@@ -135,7 +154,11 @@ class GaiaAgent:
         }
 
         try:
-            final_state = self.graph.invoke(initial_state, {"recursion_limit": 50})
+            import uuid
+
+            thread_id = str(uuid.uuid4())
+            config = {"configurable": {"thread_id": thread_id}, "recursion_limit": 50}
+            final_state = await self.graph.ainvoke(initial_state, config)
 
             last_message = final_state["messages"][-1]
 
@@ -200,11 +223,43 @@ class GaiaAgent:
     def _agent_node(self, state: AgentState) -> AgentState:
         """The main agent node that processes messages and generates responses"""
         messages = state["messages"]
+
+        # Debug: show message count
+        print(f"\n[AGENT] Message count: {len(messages)}")
+
+        # Prepend system prompt if not already there
+        if not messages or not isinstance(messages[0], SystemMessage):
+            messages = [SYSTEM_PROMPT] + messages
+
+        # Print the full prompt/messages
+        print("[AGENT] === MESSAGES ===")
+        for i, msg in enumerate(messages):
+            msg_type = type(msg).__name__
+            content = (
+                str(msg.content)[:500] if hasattr(msg, "content") else str(msg)[:500]
+            )
+            print(f"  [{i}] {msg_type}: {content}...")
+        print("[AGENT] === END MESSAGES ===\n")
+
         response = self.llm.invoke(messages)
+
+        # Log what the agent is doing
+        if hasattr(response, "tool_calls") and response.tool_calls:
+            print(
+                f"[AGENT] Calling tools: {[tc['name'] for tc in response.tool_calls]}"
+            )
+        else:
+            content = (
+                str(response.content)[:200]
+                if hasattr(response, "content")
+                else str(response)[:200]
+            )
+            print(f"[AGENT] Final response: {content}...")
+
         return {"messages": [response]}
 
 
 # model="o3-mini"
 MODEL = "glm-4.7"
 
-BasicAgent = GaiaAgent(model=MODEL, temperature=0.0)
+BasicAgent = GaiaAgent(model=MODEL, temperature=1.0)

app.py

@@ -11,7 +11,7 @@ from agent import BasicAgent
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
-def run_and_submit_all(profile: gr.OAuthProfile | None):
+async def run_and_submit_all(profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
     and displays the results.
@@ -72,7 +72,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
             print(f"Skipping item with missing task_id or question: {item}")
             continue
         try:
-            submitted_answer = agent(question_text)
+            submitted_answer = await agent(question_text)
             print(f"Task ID: {task_id}")
             print(f"Answer: {submitted_answer}")
             print("-" * 40)
@@ -183,7 +183,8 @@ with gr.Blocks() as demo:
 
     run_button.click(fn=run_and_submit_all, outputs=[status_output, results_table])
 
-if __name__ == "__main__":
+
+async def main():
     print("\n" + "-" * 30 + " App Starting " + "-" * 30)
     # Check for SPACE_HOST and SPACE_ID at startup for information
     space_host_startup = os.getenv("SPACE_HOST")
@@ -210,3 +211,9 @@ if __name__ == "__main__":
 
     print("Launching Gradio Interface for Basic Agent Evaluation...")
     demo.launch(debug=True, share=False)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

build_rag_index.py

@@ -0,0 +1,104 @@
+"""Build RAG index from GAIA validation dataset with Annotator Metadata."""
+
+import json
+from pathlib import Path
+
+from datasets import load_dataset
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain_huggingface import HuggingFaceEmbeddings
+
+CHROMA_PATH = "./chroma_gaia_db"
+
+
+def build_index():
+    """Load GAIA validation set and index questions with metadata."""
+    print("Loading GAIA dataset...")
+    ds = load_dataset("gaia-benchmark/GAIA", "2023_all", split="validation")
+
+    print(f"Found {len(ds)} examples")
+
+    # Create documents from dataset
+    documents = []
+    for item in ds:
+        question = item.get("Question", "")
+        answer = item.get("Final answer", "")
+        level = item.get("Level", "")
+        task_id = item.get("task_id", "")
+        metadata_raw = item.get("Annotator Metadata", {})
+
+        # Parse annotator metadata
+        if isinstance(metadata_raw, str):
+            try:
+                metadata_raw = json.loads(metadata_raw)
+            except json.JSONDecodeError:
+                metadata_raw = {}
+
+        steps = metadata_raw.get("Steps", "")
+        tools = metadata_raw.get("Tools", "")
+        num_steps = metadata_raw.get("Number of steps", "")
+
+        # Build document content with question, answer, and reasoning
+        content = f"""Question: {question}
+
+Final Answer: {answer}
+
+Steps to solve:
+{steps}
+
+Tools needed: {tools}"""
+
+        doc = Document(
+            page_content=content,
+            metadata={
+                "task_id": task_id,
+                "question": question,
+                "answer": answer,
+                "level": str(level),
+                "num_steps": str(num_steps),
+                "tools": tools,
+            },
+        )
+        documents.append(doc)
+
+    print(f"Created {len(documents)} documents")
+
+    # Initialize embeddings
+    print("Initializing embeddings...")
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-mpnet-base-v2"
+    )
+
+    # Clear existing index if present
+    chroma_path = Path(CHROMA_PATH)
+    if chroma_path.exists():
+        import shutil
+
+        shutil.rmtree(chroma_path)
+        print("Cleared existing index")
+
+    # Create and persist vector store
+    print("Building vector store...")
+    vectorstore = Chroma.from_documents(
+        documents=documents,
+        embedding=embeddings,
+        persist_directory=CHROMA_PATH,
+    )
+
+    print(f"Indexed {len(documents)} documents to {CHROMA_PATH}")
+
+    # Test retrieval
+    print("\nTesting retrieval...")
+    test_query = (
+        "How many studio albums did Mercedes Sosa release between 2000 and 2009?"
+    )
+    results = vectorstore.similarity_search(test_query, k=2)
+    print(f"Query: {test_query}")
+    for i, doc in enumerate(results):
+        print(f"\n--- Result {i+1} ---")
+        print(f"Question: {doc.metadata.get('question', '')[:100]}...")
+        print(f"Answer: {doc.metadata.get('answer', '')}")
+
+
+if __name__ == "__main__":
+    build_index()

chroma_gaia_db/.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.sqlite3 filter=lfs diff=lfs merge=lfs -text

chroma_gaia_db/99bb1417-fe53-457a-8b1f-42a54fb4c17c/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b92e8e451752ee2cb1c2e5bba20ff2aa94ba02b270bcfc3f8f6efffb8b948333
+size 321200

chroma_gaia_db/99bb1417-fe53-457a-8b1f-42a54fb4c17c/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03cb3ac86f3e5bcb15e88b9bf99f760ec6b33e31d64a699e129b49868db6d733
+size 100

chroma_gaia_db/99bb1417-fe53-457a-8b1f-42a54fb4c17c/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a12e561363385e9dfeeab326368731c030ed4b374e7f5897ac819159d2884c5
+size 400

chroma_gaia_db/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7524be1035bb3131badbf5dc2828251aa5df02a6248d4675e7966bc8eef2ddd
+size 2830336

pyproject.toml

@@ -29,6 +29,15 @@ dependencies = [
     "ddgs>=9.10.0",
     "sentence-transformers>=5.2.0",
     "typer>=0.9.0",
+    "httpx>=0.28.1",
+    "pyjwt>=2.10.1",
+    "openpyxl>=3.1.5",
+    "python-docx>=1.2.0",
+    "python-pptx>=1.0.2",
+    "langchain-chroma>=1.1.0",
+    "zai>=0.0.2",
+    "zai-sdk>=0.2.0",
+    "langchain-mcp-adapters>=0.2.1",
 ]
 
 [dependency-groups]

system_prompt.txt

@@ -1,38 +1,17 @@
-You are a precise data extraction assistant for the GAIA benchmark.
+You are a precise assistant for the GAIA benchmark.
 
-## Your Task
-Answer factual questions using tools. Extract answers VERBATIM from sources.
+## 工作流程 (workflow)
+1. 首先: call `get_solving_strategy` with your question
+2. follow the strategy steps using appropriate tools
+3. when you find the answer, call `submit_answer` immediately
 
-## Rules (必须严格遵守)
-1. Use EXACT wording from sources. Copy text verbatim, do not paraphrase or shorten.
-2. For lists: sort items alphabetically, separate with comma and space.
-3. If source says "freshly squeezed lemon juice", write "freshly squeezed lemon juice" (NOT "lemon juice").
-4. If source says "granulated sugar", write "granulated sugar" (NOT "sugar").
-5. If source says "pure vanilla extract", write "pure vanilla extract" (NOT "vanilla extract").
-6. If source says "ripe strawberries", write "ripe strawberries" (NOT "strawberries").
-7. Use tools to find information. Do not guess.
-8. If information is not found, respond: "I don't know"
-
-## Search Tips
-- For discography/albums questions: search "[Artist] discography" to find the full album list
-- For counting items in a date range: list each item with its year, then count
-- If wiki_search doesn't have enough detail, use web_search or jina_search
-- Read the full Wikipedia page with jina_read if needed
+## rules (必须严格遵守)
+1. use exact wording from sources. do not paraphrase or shorten.
+2. for lists: sort items alphabetically, separate with comma and space.
+3. use tools to find information. do not guess.
+4. when you have the answer, call `submit_answer` immediately. 不要继续搜索。
+5. if information is not found, keep trying different tools and approaches.
 
 ## Output Format (严格按照此格式)
 Write ONLY this, nothing else:
 FINAL ANSWER: [your answer here]
-
-## Examples
-
-Question: What are the filling ingredients?
-Source: "You'll need ripe strawberries, granulated sugar, and freshly squeezed lemon juice"
-FINAL ANSWER: freshly squeezed lemon juice, granulated sugar, ripe strawberries
-
-Question: What ingredients are in the sauce?
-Source: "Mix pure vanilla extract with heavy whipping cream"
-FINAL ANSWER: heavy whipping cream, pure vanilla extract
-
-Question: Who wrote the book?
-Source: "The novel was written by Jane Smith in 1995"
-FINAL ANSWER: Jane Smith

test_bench.py

@@ -243,7 +243,9 @@ def run_test_bench(
             question_with_file = q.question
             if q.file_path:
                 question_with_file += f"\n\nFile path: {q.file_path}"
-            actual = agent(question_with_file)
+            import asyncio
+
+            actual = asyncio.run(agent(question_with_file))
         except Exception as e:
             actual = f"ERROR: {e}"
 

tools.py

@@ -1,10 +1,77 @@
 import os
 from typing import List
 
+from langchain_chroma import Chroma
 from langchain_core.documents.base import Document
 from langchain_core.tools import tool
 from langchain_core.tools.base import ArgsSchema
+from langchain_huggingface import HuggingFaceEmbeddings
 from pydantic import SecretStr
+from sqlalchemy.sql.selectable import ForUpdateParameter
+
+# Initialize RAG vector store for strategy retrieval
+CHROMA_PATH = "./chroma_gaia_db"
+_embeddings = None
+_vector_store = None
+
+
+def _get_vector_store():
+    """Lazy load vector store."""
+    global _embeddings, _vector_store
+    if _vector_store is None:
+        _embeddings = HuggingFaceEmbeddings(
+            model_name="sentence-transformers/all-mpnet-base-v2"
+        )
+        _vector_store = Chroma(
+            persist_directory=CHROMA_PATH, embedding_function=_embeddings
+        )
+    return _vector_store
+
+
+@tool
+def get_solving_strategy(question: str) -> str:
+    """Search for similar solved questions and get the solving strategy.
+    Use this FIRST to understand how to approach a problem before using other tools.
+
+    Args:
+        question: The question you need to solve."""
+    print(f"\n[GET_SOLVING_STRATEGY] Searching for: {question[:80]}...")
+    try:
+        vector_store = _get_vector_store()
+        similar_docs = vector_store.similarity_search(question, k=1)
+        print(f"[GET_SOLVING_STRATEGY] Found {len(similar_docs)} similar questions")
+
+        if similar_docs:
+            doc = similar_docs[0]
+            steps = (
+                doc.page_content.split("Steps to solve:")[-1]
+                .split("Tools needed:")[0]
+                .strip()
+            )
+            tools_raw = doc.metadata.get("tools", "")
+            # Clean up tools format - replace inline numbers with newlines
+            tools = tools_raw.replace("\n", "\n- ").strip()
+            if tools and not tools.startswith("-"):
+                tools = "- " + tools
+
+            set_current_strategy(steps)
+
+            return f"""Similar question found!
+
+## Strategy to solve (按此策略执行):
+{steps}
+
+## Rules (必须严格遵守):
+1. Use EXACT wording from sources. Do not paraphrase or shorten.
+2. For lists: sort items alphabetically, separate with comma and space.
+3. Use tools to find information. Do not guess.
+4. When you find the answer, call `submit_answer` immediately. 不要继续搜索。
+
+"""
+        else:
+            return "No similar questions found. Use your best judgment."
+    except Exception as e:
+        return f"Error searching for strategy: {e}"
 
 
 def _get_llm():
@@ -25,55 +92,102 @@ def _get_llm():
         return ChatOpenAI(model="gpt-4o-mini", temperature=0)
 
 
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
+def _fetch_url_with_tables(url: str) -> str:
+    """Fetch URL content including tables using Jina reader."""
+    import requests
 
-    Args:
-        query: The search query."""
-    from langchain_community.document_loaders import WikipediaLoader
+    try:
+        # Use Jina to get full page content including tables
+        api_key = os.getenv("JINA_API_KEY", "")
+        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
 
-    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 formatted_search_docs
+        response = requests.get(f"https://r.jina.ai/{url}", headers=headers, timeout=30)
+        return response.text
+    except Exception:
+        return ""
 
 
 @tool
-def zai_web_search(query: str) -> str:
-    """Search the web using Z.AI web search API.
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return relevant content including tables.
 
     Args:
         query: The search query."""
-    import requests
-
-    api_key = os.getenv("ZAI_API_KEY", "")
-
-    response = requests.post(
-        "https://api.z.ai/api/coding/paas/v4/web_search",
-        headers={
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json",
-        },
-        json={"search_engine": "search-prime", "search_query": query, "count": 5},
-        timeout=30,
-    )
+    import wikipedia
 
-    data = response.json()
-    results = data.get("search_result", [])
-
-    parts = []
-    for r in results:
-        title = r.get("title", "")
-        content = r.get("content", "")
-        link = r.get("link", "")
-        parts.append(f'<Document source="{link}">\n{title}\n{content}\n</Document>')
-
-    return "\n\n---\n\n".join(parts) if parts else "No results found."
+    try:
+        # Search for pages
+        search_results = wikipedia.search(query, results=3)
+        if not search_results:
+            return "No Wikipedia results found."
+
+        formatted_parts = []
+        for title in search_results[:2]:
+            try:
+                page = wikipedia.page(title, auto_suggest=False)
+                url = page.url
+
+                # Fetch the page via Jina to get full content including tables
+                content = _fetch_url_with_tables(url)
+
+                if not content:
+                    # Fallback to wikipedia API content
+                    content = page.content
+
+                # Use smart section extraction
+                extracted = _extract_relevant_content(content, query)
+                formatted_parts.append(
+                    f'<Document source="{url}" title="{title}">\n{extracted}\n</Document>'
+                )
+            except (wikipedia.DisambiguationError, wikipedia.PageError):
+                continue
+            except Exception:
+                continue
+
+        return (
+            "\n\n---\n\n".join(formatted_parts)
+            if formatted_parts
+            else "No results found."
+        )
+    except Exception as e:
+        return f"Wikipedia search error: {e}"
+
+
+_zai_mcp_tools = None
+
+
+async def _get_zai_mcp_tools():
+    """Lazy load Z.AI MCP tools."""
+    global _zai_mcp_tools
+    if _zai_mcp_tools is None:
+        from langchain_mcp_adapters.client import MultiServerMCPClient
+
+        api_key = os.getenv("ZAI_API_KEY", "")
+        client = MultiServerMCPClient(
+            {
+                "web-search": {
+                    "transport": "streamable_http",
+                    "url": "https://api.z.ai/api/mcp/web_search_prime/mcp",
+                    "headers": {"Authorization": f"Bearer {api_key}"},
+                },
+                "web-reader": {
+                    "transport": "streamable_http",
+                    "url": "https://api.z.ai/api/mcp/web_reader/mcp",
+                    "headers": {"Authorization": f"Bearer {api_key}"},
+                },
+                "zai-mcp": {
+                    "transport": "stdio",
+                    "command": "npx",
+                    "args": ["-y", "@z_ai/mcp-server"],
+                    "env": {
+                        "Z_AI_API_KEY": api_key,
+                        "Z_AI_MODE": "ZAI",
+                    },
+                },
+            }
+        )
+        _zai_mcp_tools = await client.get_tools()
+    return _zai_mcp_tools
 
 
 @tool
@@ -91,54 +205,123 @@ def jina_search(query: str) -> str:
     return response.text
 
 
-@tool
-def jina_read(url: str) -> str:
-    """Read a webpage using Jina AI and return clean markdown.
+def _extract_section_by_marker(
+    content: str, section_marker: str, context_lines: int = 50
+) -> str:
+    """Extract a section starting from a marker found in strategy steps.
+
+    This is the SMART extraction - uses strategy steps like "scrolled down to Studio albums"
+    to find the exact section we need.
+    """
+    import re
+
+    lines = content.split("\n")
+    marker_lower = section_marker.lower().strip()
+
+    print(f"[EXTRACT_SECTION] Looking for section marker: '{section_marker}'")
+
+    # Find the line containing the section marker
+    start_idx = None
+    for i, line in enumerate(lines):
+        if marker_lower in line.lower():
+            start_idx = i
+            print(f"[EXTRACT_SECTION] Found marker at line {i}: {line[:80]}")
+            break
+
+    if start_idx is None:
+        # Try partial matching (e.g., "Studio albums" might be "Studio Albums" or "Discography")
+        for i, line in enumerate(lines):
+            # Check if most words from marker are in line
+            marker_words = [
+                w for w in re.findall(r"\b\w+\b", marker_lower) if len(w) > 2
+            ]
+            line_lower = line.lower()
+            matches = sum(1 for w in marker_words if w in line_lower)
+            if matches >= len(marker_words) * 0.6:  # 60% match threshold
+                start_idx = i
+                print(f"[EXTRACT_SECTION] Found partial match at line {i}: {line[:80]}")
+                break
+
+    if start_idx is None:
+        print(f"[EXTRACT_SECTION] Section marker not found")
+        return ""
+
+    # Extract from marker line + context_lines after it
+    end_idx = min(start_idx + context_lines, len(lines))
+    section = "\n".join(lines[start_idx:end_idx])
+
+    print(f"[EXTRACT_SECTION] Extracted {end_idx - start_idx} lines from section")
+    return section
+
+
+def _parse_section_markers_from_strategy(strategy: str) -> list:
+    """Parse strategy steps to extract section markers.
+
+    Looks for phrases like:
+    - "scrolled down to Studio albums" -> "Studio albums"
+    - "found the Discography section" -> "Discography"
+    - "went to Studio albums" -> "Studio albums"
+    """
+    import re
+
+    markers = []
+
+    # Patterns that indicate a section name
+    patterns = [
+        r'scrolled?\s+(?:down\s+)?to\s+["\']?([^"\',.]+)["\']?',  # scrolled down to X
+        r'went\s+to\s+(?:the\s+)?["\']?([^"\',.]+)["\']?\s+section',  # went to X section
+        r'found\s+(?:the\s+)?["\']?([^"\',.]+)["\']?\s+section',  # found X section
+        r'clicked\s+on\s+["\']?([^"\',.]+)["\']?',  # clicked on X
+        r'looked\s+(?:at|under)\s+["\']?([^"\',.]+)["\']?',  # looked at/under X
+        r'(?:in|under)\s+(?:the\s+)?["\']?([^"\',.]+)["\']?\s+section',  # in/under X section
+    ]
 
-    Args:
-        url: The URL to read."""
-    import requests
+    for pattern in patterns:
+        matches = re.findall(pattern, strategy.lower())
+        for match in matches:
+            cleaned = match.strip()
+            if cleaned and len(cleaned) > 2 and len(cleaned) < 50:
+                markers.append(cleaned)
 
-    api_key = os.getenv("JINA_API_KEY", "")
-    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+    # Also look for quoted section names
+    quoted = re.findall(r'"([^"]+)"', strategy)
+    for q in quoted:
+        if len(q) > 2 and len(q) < 50 and q.lower() not in ["wikipedia", "google"]:
+            markers.append(q)
 
-    response = requests.get(f"https://r.jina.ai/{url}", headers=headers, timeout=30)
-    return response.text[:10000]
+    print(f"[PARSE_STRATEGY] Extracted section markers: {markers}")
+    return markers
+
+
+# Global variable to store current strategy for smart extraction
+_current_strategy = None
+
+
+def set_current_strategy(strategy: str):
+    """Store the current strategy for use by content extraction."""
+    global _current_strategy
+    _current_strategy = strategy
+    print(f"[STRATEGY] Updated current strategy")
 
 
 @tool
-def zai_web_reader(url: str) -> str:
-    """Read a webpage using Z.AI MCP web reader and return clean markdown.
+def jina_read(url: str, question: str = "") -> str:
+    """Read a webpage and extract content relevant to the question.
 
     Args:
-        url: The URL to read."""
+        url: The URL to read.
+        question: The question to extract relevant info for."""
     import requests
 
-    api_key = os.getenv("ZAI_API_KEY", "")
-
-    response = requests.post(
-        "https://api.z.ai/api/mcp/web_reader/mcp",
-        headers={
-            "Authorization": f"Bearer {api_key}",
-            "Content-Type": "application/json",
-        },
-        json={
-            "method": "tools/call",
-            "params": {
-                "name": "webReader",
-                "arguments": {"url": url},
-            },
-        },
-        timeout=60,
-    )
+    api_key = os.getenv("JINA_API_KEY", "")
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
 
-    data = response.json()
-    if "result" in data and "content" in data["result"]:
-        content = data["result"]["content"]
-        if isinstance(content, list) and len(content) > 0:
-            return content[0].get("text", str(content))[:15000]
-        return str(content)[:15000]
-    return f"Error: {data}"
+    response = requests.get(f"https://r.jina.ai/{url}", headers=headers, timeout=30)
+    content = response.text
+
+    # Use smart extraction with strategy section markers
+    if question:
+        return content[:10000]
 
 
 @tool
@@ -205,20 +388,177 @@ def analyze_text(text: str, question: str) -> str:
     return response.content
 
 
-def get_tools() -> list:
+@tool
+def read_excel(file_path: str) -> str:
+    """Read and extract data from an Excel file (.xlsx, .xls).
+
+    Args:
+        file_path: Path to the Excel file."""
+    import pandas as pd
+
+    try:
+        # Read all sheets
+        xlsx = pd.ExcelFile(file_path)
+        results = []
+        for sheet_name in xlsx.sheet_names:
+            df = pd.read_excel(xlsx, sheet_name=sheet_name)
+            results.append(f"=== Sheet: {sheet_name} ===\n{df.to_string()}")
+        return "\n\n".join(results)[:15000]
+    except Exception as e:
+        return f"Error reading Excel: {e}"
+
+
+@tool
+def read_csv(file_path: str) -> str:
+    """Read and extract data from a CSV file.
+
+    Args:
+        file_path: Path to the CSV file."""
+    import pandas as pd
+
+    try:
+        df = pd.read_csv(file_path)
+        return df.to_string()[:15000]
+    except Exception as e:
+        return f"Error reading CSV: {e}"
+
+
+@tool
+def read_docx(file_path: str) -> str:
+    """Read and extract text from a Word document (.docx).
+
+    Args:
+        file_path: Path to the Word document."""
+    try:
+        from docx import Document
+
+        doc = Document(file_path)
+        text = "\n".join([para.text for para in doc.paragraphs])
+        return text[:15000]
+    except Exception as e:
+        return f"Error reading Word doc: {e}"
+
+
+@tool
+def read_pptx(file_path: str) -> str:
+    """Read and extract text from a PowerPoint presentation (.pptx).
+
+    Args:
+        file_path: Path to the PowerPoint file."""
+    try:
+        from pptx import Presentation
+
+        prs = Presentation(file_path)
+        text_parts = []
+        for slide_num, slide in enumerate(prs.slides, 1):
+            slide_text = [f"=== Slide {slide_num} ==="]
+            for shape in slide.shapes:
+                if hasattr(shape, "text"):
+                    slide_text.append(shape.text)
+            text_parts.append("\n".join(slide_text))
+        return "\n\n".join(text_parts)[:15000]
+    except Exception as e:
+        return f"Error reading PowerPoint: {e}"
+
+
+@tool
+def extract_zip(file_path: str) -> str:
+    """Extract a zip file and list its contents.
+
+    Args:
+        file_path: Path to the zip file."""
+    import zipfile
+    from pathlib import Path
+
+    try:
+        extract_dir = Path(file_path).parent / Path(file_path).stem
+        extract_dir.mkdir(exist_ok=True)
+
+        with zipfile.ZipFile(file_path, "r") as zip_ref:
+            zip_ref.extractall(extract_dir)
+            file_list = zip_ref.namelist()
+
+        return f"Extracted to: {extract_dir}\nContents:\n" + "\n".join(file_list)
+    except Exception as e:
+        return f"Error extracting zip: {e}"
+
+
+@tool
+def analyze_image(file_path: str, question: str) -> str:
+    """Analyze an image and answer a question about it using vision model.
+
+    Args:
+        file_path: Path to the image file (png, jpg, etc.)
+        question: Question to answer about the image."""
+    import base64
+
+    from langchain_openai import ChatOpenAI
+
+    try:
+        with open(file_path, "rb") as f:
+            image_data = base64.b64encode(f.read()).decode("utf-8")
+
+        # Determine mime type
+        ext = file_path.lower().split(".")[-1]
+        mime_type = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg"}.get(
+            ext, "image/png"
+        )
+
+        # Use GPT-4o for vision
+        llm = ChatOpenAI(model="gpt-4o", temperature=0)
+        response = llm.invoke(
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": question},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:{mime_type};base64,{image_data}"
+                            },
+                        },
+                    ],
+                }
+            ]
+        )
+        return response.content
+    except Exception as e:
+        return f"Error analyzing image: {e}"
+
+
+@tool
+def submit_answer(answer: str) -> str:
+    """Submit your final answer. Use this when you have found the answer.
+
+    Args:
+        answer: The final answer to submit."""
+    print(f"[SUBMIT_ANSWER] {answer}")
+    return f"FINAL ANSWER: {answer}"
+
+
+async def get_tools() -> list:
     """Retrieve the list of available tools for the agent."""
-    return [
-        wiki_search,
-        web_search,
-        jina_search,
-        jina_read,
+    base_tools = [
+        get_solving_strategy,  # Use FIRST to get approach
+        submit_answer,
+        # wiki_search,
         download_file,
         read_pdf,
+        read_excel,
+        read_csv,
+        read_docx,
+        read_pptx,
+        extract_zip,
+        analyze_image,
         py_calc_tool,
         youtube_transcript_tool,
         transcribe_audio,
         arxiv_search,
     ]
+    # Add Z.AI MCP tools (webSearchPrime, webReader)
+    zai_tools = await _get_zai_mcp_tools()
+    return base_tools + zai_tools
 
 
 @tool