agent.py

import json
import os
from pathlib import Path
import html
from io import StringIO
import re
import subprocess
import sys
from urllib.parse import parse_qs, unquote, urlparse
import pandas as pd
from typing import TypedDict

import requests
from langgraph.graph import END, START, StateGraph
from huggingface_hub import InferenceClient
from openai import OpenAI
from gaia import CourseAPIClient, TaskRecord



class AgentState(TypedDict, total=False):
    # input level
    task: TaskRecord
    question: str
    task_id: str
    file_name: str
    attachment_path: str
    task_type: str
    level: str

    # analysis part
    analysis: str
    needs_tool: bool
    candidate_tools: list[str]
    constraints: str

    # execution part
    selected_tool: str
    tool_input: dict[str, str]
    tool_output: str
    observations: list[str]
    tool_call_count: int
    last_tool_status: str

    # control part
    should_continue: bool
    max_tool_calls: int
    stop_reason: str

    # answer part
    draft_answer: str
    validation_passed: bool
    validation_notes: str
    final_answer: str


SYSTEM_PROMPT = """You are solving a benchmark task from the GAIA-style agent course.

Rules:
- Solve the task carefully and keep intermediate reasoning private.
- If a local attachment path is provided, use it as part of your analysis.
- Return only the final answer text with no prefix and no explanation.
- Prefer exact, concise answers because scoring uses exact match.
"""


class LangGraphBenchmarkAgent:
    def __init__(self, api_url: str):
        self.client = CourseAPIClient(api_url=api_url)
        self.attachments_dir = Path("downloads")
        self.model = InferenceClient(
                                     api_key=os.getenv("HF_TOKEN"),
                                )
        self.model_name = os.getenv("HF_MODEL","Qwen/Qwen2.5-7B-Instruct")
        self.openai_compatible_client = None
        self.openai_compatible_model = os.getenv("OPENAI_COMPATIBLE_MODEL", "")
        if os.getenv("LLM_PROVIDER", "").lower() == "openai_compatible":
            self.openai_compatible_client = OpenAI(
                api_key=os.getenv("OPENAI_COMPATIBLE_API_KEY"),
                base_url=os.getenv("OPENAI_COMPATIBLE_BASE_URL"),
            )
        self.graph = self._build_graph()

    def _call_llm(self, messages: list[dict[str, str]], system_fallback: str = SYSTEM_PROMPT) -> str:
        if self.openai_compatible_client:
            response = self.openai_compatible_client.chat.completions.create(
                model=self.openai_compatible_model,
                messages=messages,
            )
        else:
            response = self.model.chat.completions.create(
                model=self.model_name,
                messages=messages,
            )
        return response.choices[0].message.content or ""

    def solve(self, task: TaskRecord) -> str:
        initial_state: AgentState = {
            "task": task,
            "max_tool_calls": 3,
            "tool_call_count": 0,
            "observations": [],
        }
        result = self.graph.invoke(initial_state)
        return result["final_answer"].strip()
    
    def solve_debug(self, task: TaskRecord) -> AgentState:
        initial_state: AgentState = {
            "task": task,
            "max_tool_calls": 3,
            "tool_call_count": 0,
            "observations": [],
        }
        result = self.graph.invoke(initial_state)
        return result

    def _build_graph(self):
        graph = StateGraph(AgentState)
        graph.add_node("prepare_task", self._prepare_task)
        graph.add_node("analyze_task", self._analyze_task)
        graph.add_node("select_tool", self._select_tool)
        graph.add_node("draft_answer", self._draft_answer)
        graph.add_node("finalize_answer", self._finalize_answer)
        graph.add_node("run_tool", self._run_tool)
        graph.add_node("validate_answer", self._validate_answer)
        graph.add_edge(START, "prepare_task")
        graph.add_edge("prepare_task", "analyze_task")
        graph.add_conditional_edges("analyze_task", self._judge_tool_usage,{
            "select_tool": "select_tool",
            "draft_answer": "draft_answer",
        },)
        graph.add_edge("select_tool", "run_tool")
        graph.add_edge("run_tool", "draft_answer")
        graph.add_edge("draft_answer", "validate_answer")
        graph.add_edge("validate_answer", "finalize_answer")
        graph.add_edge("finalize_answer", END)
        return graph.compile()

    def _prepare_task(self, state: AgentState) -> AgentState:
        task = state["task"]
        attachment_path = self.client.fetch_attachment(task, self.attachments_dir)
        return {
            "question": task.question,
            "task_id": task.task_id,
            "level": task.level,
            "file_name": task.file_name or "",
            "attachment_path": str(attachment_path) if attachment_path else "",
            "task_type": self._classify_task(task),
        }

    def _analyze_task(self, state: AgentState) -> AgentState:
        question = state.get("question", "")
        task_type = state.get("task_type", "")
        attachment_path = state.get("attachment_path", "")
        file_name = state.get("file_name", "")

        file_task_type = {
            "code_question",
            "document_question",
            "image_question",
            "table_question",
            "structured_data_question",
            "file_question",
            "image_question",
        }

        web_task_types = {
            "web_question",
            "lookup_or_count_question",
            "video_question",
        }

        if "comma separated" in question.lower():
            constraints = "Return a comma separated answer with no explanation."
        else:
            constraints = "Return only the final answer with no explanation."
        if attachment_path or file_name or task_type in file_task_type:
            return {
                "analysis": f"This is a {task_type} task that needs local attachment inspection.",
                "needs_tool": True,
                "candidate_tools": ["question_lookup"],
                "constraints": constraints
            }
        
        if task_type in web_task_types:
            return {
                "analysis": f"This is a {task_type} task that needs external lookup.",
                "needs_tool": True,
                "candidate_tools": ["web_lookup"],
                "constraints": constraints,
            }
        return {
            "analysis":f"This is a {task_type} task that can be solved directly without tools.",
            "needs_tool": False,
            "candidate_tools": [],
            "constraints": constraints,
        }

    def _judge_tool_usage(self, state: AgentState) -> str:
        if state.get("needs_tool", False):
            return "select_tool"
        return "draft_answer"

    def _select_tool(self, state: AgentState) -> AgentState:
        select_tool=(
            state.get("candidate_tools",[])[0]
            if state.get("candidate_tools", []) else ""
        )
        return {
            "selected_tool": select_tool,
            "tool_input":{
                "question": state.get('question', ''),
                "task_type": state.get('task_type', ''),
                "attachment_path": state.get('attachment_path', ''),
            },
        }
    
    def _run_tool(self, state: AgentState) -> AgentState:
        tool_name = state.get("selected_tool", "")
        tool_input = state.get("tool_input", {})
        attachment_path = tool_input.get("attachment_path", "")
        if tool_name == "question_lookup":
            if attachment_path and Path(attachment_path).exists() and Path(attachment_path).suffix.lower() == ".py":
                try:
                    completed = subprocess.run(
                        [sys.executable, attachment_path],
                        capture_output=True,
                        text=True,
                        timeout=30,
                    )
                    output = (
                        f"attachment_path: {tool_input.get('attachment_path', '')}\n"
                        f"task_type: {tool_input.get('task_type', '')}\n"
                        f"python_returncode: {completed.returncode}\n"
                        f"python_stdout: {completed.stdout.strip()}\n"
                        f"python_stderr: {completed.stderr.strip()}"
                    )
                    last_tool_status = "success" if completed.returncode == 0 else "failure"
                except subprocess.TimeoutExpired as exc:
                    source = Path(attachment_path).read_text(encoding="utf-8", errors="ignore")
                    output = (
                        f"attachment_path: {tool_input.get('attachment_path', '')}\n"
                        f"task_type: {tool_input.get('task_type', '')}\n"
                        f"python_status: timeout after {exc.timeout} seconds\n"
                        f"python_stdout_before_timeout: {(exc.stdout or '').strip()}\n"
                        f"python_stderr_before_timeout: {(exc.stderr or '').strip()}\n"
                        f"python_source:\n{source}"
                    )
                    last_tool_status = "failure"

            elif attachment_path and Path(attachment_path).exists() and Path(attachment_path).suffix.lower() in {".csv", ".xlsx", ".xls"}:
                    path = Path(attachment_path)
                    if path.suffix.lower() == ".csv":
                        df = pd.read_csv(path)
                    else:
                        df = pd.read_excel(path)

                    output = (
                        f"attachment_path: {attachment_path}\n"
                        f"task_type: {tool_input.get('task_type', '')}\n"
                        f"table_shape: {df.shape[0]} rows x {df.shape[1]} columns\n"
                        f"table_columns: {list(df.columns)}\n"
                        f"table_preview:\n{df.head(20).to_string(index=False)}"
                    )
                    last_tool_status = "success"

            elif attachment_path and Path(attachment_path).exists() and Path(attachment_path).suffix.lower() in {".txt", ".md"}:
                with open(attachment_path, "r", encoding="utf-8", errors="ignore") as f:
                    content = f.read()
                output = (
                    f"attachment_path: {tool_input.get('attachment_path', '')}\n"
                    f"task_type: {tool_input.get('task_type', '')}\n"
                    f"attachment_content: {content}"
                )
                last_tool_status = "success"
            else:
                output = (
                    f"attachment_path: {tool_input.get('attachment_path', '')}\n"
                    f"task_type: {tool_input.get('task_type', '')}\n"
                    "attachment_status: no readable local attachment found"
                )
                last_tool_status = "failure"
            tool_call_count = state.get("tool_call_count", 0) + 1
        elif tool_name == "web_lookup":
            question = tool_input.get("question", "")
            output, last_tool_status = self._run_web_lookup(question)
            tool_call_count = state.get("tool_call_count", 0) + 1
        else:
            output = f"No tool executed. Unknown tool: {tool_name}"
            last_tool_status = "failure"
            tool_call_count = state.get("tool_call_count", 0) + 1
        observations = state.get("observations", []) + [output]
        return {"tool_output": output, 
                "last_tool_status": last_tool_status, 
                "tool_call_count": tool_call_count,
                "observations": observations
            }

    def _run_web_lookup(self, question: str) -> tuple[str, str]:
        url = self._extract_first_url(question)
        if url:
            metadata = self._describe_url(url)
            page_text = self._fetch_url_text(url)
            if page_text:
                return (
                    f"lookup_mode: direct_url\n"
                    f"{metadata}\n"
                    f"page_text:\n{page_text}",
                    "success",
                )
            return (
                f"lookup_mode: direct_url\n"
                f"{metadata}\n"
                "page_text: unavailable",
                "failure",
            )

        search_queries = self._build_search_queries(question)
        ddg_text = self._search_duckduckgo(search_queries)
        web_results_text = self._search_duckduckgo_html(search_queries)
        wiki_text = self._search_wikipedia(search_queries)
        evidence_parts = [part for part in [ddg_text, web_results_text, wiki_text] if part]
        if evidence_parts:
            return (
                "lookup_mode: web_search\n"
                f"search_target: {question}\n"
                f"search_queries: {search_queries}\n"
                + "\n\n".join(evidence_parts),
                "success",
            )

        return (
            "lookup_mode: web_search\n"
            f"search_target: {question}\n"
            "search_status: no evidence found",
            "failure",
        )

    def _build_search_queries(self, question: str) -> list[str]:
        rule_queries = self._build_rule_search_queries(question)
        model_queries = self._rewrite_search_queries_with_model(question)
        return list(dict.fromkeys(model_queries + rule_queries))[:3]

    @staticmethod
    def _build_rule_search_queries(question: str) -> list[str]:
        queries = []
        question_lower = question.lower()
        cleaned = re.sub(r"(?i)you can use.*?wikipedia\.?", "", question).strip()
        cleaned = re.sub(r"\s+", " ", cleaned)
        if cleaned:
            queries.append(cleaned)

        quoted_phrases = re.findall(r'"([^"]+)"', question)
        capitalized_phrases = re.findall(
            r"\b[A-Z][A-Za-z0-9'’.-]*(?:\s+[A-Z][A-Za-z0-9'’.-]*)*",
            question,
        )
        years = re.findall(r"\b(?:19|20)\d{2}\b", question)
        source_terms = [
            term
            for term in [
                "Wikipedia",
                "LibreTexts",
                "CK-12",
                "Universe Today",
                "YouTube",
            ]
            if term.lower() in question_lower
        ]
        important_terms = re.findall(
            r"\b(?:nominated|promoted|published|surname|veterinarian|pitchers|walks|at bats|athletes|olympics|featured article|dinosaur)\b",
            question_lower,
        )

        stop_phrases = {
            "How",
            "Who",
            "What",
            "Where",
            "When",
            "Which",
            "You",
            "Give",
            "Return",
            "English Wikipedia",
        }
        compressed_parts = quoted_phrases + capitalized_phrases + years + source_terms + important_terms
        compressed_parts = [
            part
            for part in compressed_parts
            if part.strip() and part.strip() not in stop_phrases
        ]
        compressed_query = " ".join(dict.fromkeys(part.strip() for part in compressed_parts if part.strip()))
        if compressed_query:
            queries.append(compressed_query)

        if "wikipedia" in question_lower:
            site_query = f"{compressed_query or cleaned or question} site:en.wikipedia.org"
            queries.append(site_query)

        return list(dict.fromkeys(query for query in queries if query))[:3]

    def _rewrite_search_queries_with_model(self, question: str) -> list[str]:
        if os.getenv("ENABLE_LLM_QUERY_REWRITE", "").lower() not in {"1", "true", "yes"}:
            return []

        prompt = (
            "Rewrite the user question into 3 concise web search queries.\n"
            "Keep names, dates, source names, and exact quoted phrases.\n"
            "Return a JSON list of strings only.\n\n"
            f"Question: {question}"
        )
        try:
            content = self._call_llm(
                [
                    {"role": "system", "content": "You write concise search queries."},
                    {"role": "user", "content": prompt},
                ]
            )
            parsed = json.loads(content)
        except Exception:
            return []

        if not isinstance(parsed, list):
            return []
        return [str(item).strip() for item in parsed if str(item).strip()][:3]

    @staticmethod
    def _extract_first_url(text: str) -> str:
        match = re.search(r"(https?://[^\s]+)", text)
        return match.group(0).rstrip(".,)") if match else ""

    @staticmethod
    def _describe_url(url: str) -> str:
        video_id_match = re.search(r"(?:v=|/)([0-9A-Za-z_-]{11}).*", url)
        if "youtube.com" in url or "youtu.be" in url:
            return (
                f"url: {url}\n"
                "platform: youtube\n"
                "resource_type: video\n"
                f"video_id: {video_id_match.group(1) if video_id_match else 'unknown'}"
            )
        return (
            f"url: {url}\n"
            "platform: unknown_web\n"
            "resource_type: webpage"
        )

    def _fetch_url_text(self, url: str) -> str:
        if "youtube.com" in url or "youtu.be" in url:
            return ""

        try:
            response = requests.get(
                url,
                headers={"User-Agent": "Mozilla/5.0"},
                timeout=20,
            )
            response.raise_for_status()
        except Exception:
            return ""

        content_type = response.headers.get("content-type", "")
        if "text" not in content_type and "html" not in content_type and "json" not in content_type:
            return ""
        return self._clean_text(response.text, max_chars=5000)

    def _search_duckduckgo(self, queries: list[str]) -> str:
        lines = []
        for query in queries:
            try:
                response = requests.get(
                    "https://api.duckduckgo.com/",
                    params={
                        "q": query,
                        "format": "json",
                        "no_html": 1,
                        "skip_disambig": 1,
                    },
                    headers={"User-Agent": "Mozilla/5.0"},
                    timeout=20,
                )
                response.raise_for_status()
                data = response.json()
            except Exception:
                continue

            if data.get("Answer"):
                lines.append(f"query: {query}")
                lines.append(f"answer: {data['Answer']}")
            if data.get("AbstractText"):
                lines.append(f"query: {query}")
                lines.append(f"abstract: {data['AbstractText']}")
            if data.get("AbstractURL"):
                lines.append(f"source: {data['AbstractURL']}")

            related = data.get("RelatedTopics", [])
            for item in related[:3]:
                if "Text" in item:
                    lines.append(f"query: {query}")
                    lines.append(f"related: {item['Text']}")
                    if item.get("FirstURL"):
                        lines.append(f"related_url: {item['FirstURL']}")

        if not lines:
            return ""
        return "duckduckgo_results:\n" + "\n".join(lines[:12])

    def _search_duckduckgo_html(self, queries: list[str]) -> str:
        lines = []
        for query in queries:
            try:
                response = requests.get(
                    "https://html.duckduckgo.com/html/",
                    params={"q": query},
                    headers={"User-Agent": "Mozilla/5.0"},
                    timeout=20,
                )
                response.raise_for_status()
            except Exception:
                continue

            result_blocks = re.findall(
                r'<a rel="nofollow" class="result__a" href="(.*?)">(.*?)</a>.*?'
                r'<a class="result__snippet".*?>(.*?)</a>',
                response.text,
                flags=re.DOTALL,
            )
            for index, (url, title, snippet) in enumerate(result_blocks[:3]):
                normalized_url = self._normalize_search_url(url)
                lines.append(f"query: {query}")
                lines.append(f"title: {self._clean_text(title, max_chars=200)}")
                lines.append(f"url: {normalized_url}")
                lines.append(f"snippet: {self._clean_text(snippet, max_chars=500)}")
                if index < 2:
                    page_text = self._fetch_url_text(normalized_url)
                    if page_text:
                        relevant_snippets = self._extract_relevant_snippets(query, page_text)
                        if relevant_snippets:
                            lines.append(f"relevant_page_snippets:\n{relevant_snippets}")
                        lines.append(f"page_text: {page_text}")
            if lines:
                break

        if not lines:
            return ""
        return "web_search_results:\n" + "\n".join(lines)

    @staticmethod
    def _normalize_search_url(url: str) -> str:
        url = html.unescape(url)
        if url.startswith("//"):
            url = "https:" + url
        parsed = urlparse(url)
        query = parse_qs(parsed.query)
        if "uddg" in query and query["uddg"]:
            return unquote(query["uddg"][0])
        return url

    @staticmethod
    def _extract_relevant_snippets(question: str, text: str) -> str:
        keywords = {
            word
            for word in re.findall(r"[A-Za-z0-9]+", question.lower())
            if len(word) > 3
            and word
            not in {
                "what",
                "when",
                "where",
                "which",
                "with",
                "that",
                "this",
                "from",
                "were",
                "have",
                "does",
                "about",
                "only",
                "give",
                "answer",
            }
        }
        if not keywords:
            return ""

        keyword_roots = {
            word[:6]
            for word in keywords
            if len(word) > 6
        }
        chunks = re.split(r"(?<=[.!?])\s+|\n+", text)

        lower_text = text.lower()
        for search_term in keywords | keyword_roots:
            for match_index, match in enumerate(re.finditer(re.escape(search_term), lower_text)):
                if match_index >= 3:
                    break
                start = max(0, match.start() - 250)
                end = min(len(text), match.end() + 250)
                chunks.append(text[start:end])

        field_snippets = []
        scored_chunks = []
        for chunk in chunks:
            cleaned = chunk.strip()
            if len(cleaned) < 40:
                continue
            chunk_words = set(re.findall(r"[A-Za-z0-9]+", cleaned.lower()))
            chunk_roots = {
                word[:6]
                for word in chunk_words
                if len(word) > 6
            }
            field_matches = re.findall(
                r"\b([A-Za-z][A-Za-z0-9 ]{2,40}?)(?:\(s\))?\s*:\s*([^.\n]{1,120})",
                cleaned,
            )
            field_terms = set()
            for field_name, field_value in field_matches:
                field_name_words = set(re.findall(r"[A-Za-z0-9]+", field_name.lower()))
                if field_name_words & {"wikipedia", "toggle", "menu", "navigation", "contents"}:
                    continue
                field_terms.update(
                    word[:6]
                    for word in field_name_words
                    if len(word) > 3
                )
                field_name_roots = {
                    word[:6]
                    for word in field_name_words
                    if len(word) > 6
                }
                if (keywords & field_name_words) or (keyword_roots & field_name_roots):
                    field_snippets.append(f"{field_name.strip()}: {field_value.strip()}")
            field_bonus = 5 * len(keyword_roots & field_terms)
            score = (
                (2 * len(keywords & chunk_words))
                + len(keyword_roots & chunk_roots)
                + field_bonus
            )
            if score:
                scored_chunks.append((score, len(cleaned), cleaned))

        scored_chunks.sort(key=lambda item: (-item[0], item[1]))
        field_snippets = list(dict.fromkeys(field_snippets))
        snippets = field_snippets[:3] + [chunk for _, _, chunk in scored_chunks[:3]]
        snippets = list(dict.fromkeys(snippets))[:3]
        return "\n---\n".join(snippets)

    def _search_wikipedia(self, queries: list[str]) -> str:
        lines = []
        for query in queries:
            try:
                response = requests.get(
                    "https://en.wikipedia.org/w/api.php",
                    params={
                        "action": "query",
                        "list": "search",
                        "srsearch": query,
                        "format": "json",
                        "srlimit": 3,
                    },
                    headers={"User-Agent": "Mozilla/5.0"},
                    timeout=20,
                )
                response.raise_for_status()
                data = response.json()
            except Exception:
                continue

            results = data.get("query", {}).get("search", [])
            query_terms = {
                term
                for term in re.findall(r"[A-Za-zÀ-ÖØ-öø-ÿ]+", query.lower())
                if len(term) > 2
            }
            for result in results:
                title = result.get("title", "")
                title_terms = set(re.findall(r"[A-Za-zÀ-ÖØ-öø-ÿ]+", title.lower()))
                if query_terms and not (query_terms & title_terms):
                    continue
                snippet = self._clean_text(result.get("snippet", ""), max_chars=500)
                if title:
                    lines.append(f"query: {query}")
                    lines.append(f"title: {title}")
                if snippet:
                    lines.append(f"snippet: {snippet}")
                if title and len(lines) < 10:
                    extract = self._fetch_wikipedia_extract(title)
                    if extract:
                        lines.append(f"page_extract: {extract}")
                    tables = self._fetch_wikipedia_tables(title)
                    if tables:
                        lines.append(f"page_tables:\n{tables}")
            if lines:
                break
        if not lines:
            return ""
        return "wikipedia_search_results:\n" + "\n".join(lines)

    def _fetch_wikipedia_extract(self, title: str) -> str:
        try:
            response = requests.get(
                "https://en.wikipedia.org/w/api.php",
                params={
                    "action": "query",
                    "prop": "extracts",
                    "explaintext": 1,
                    "titles": title,
                    "format": "json",
                },
                headers={"User-Agent": "Mozilla/5.0"},
                timeout=20,
            )
            response.raise_for_status()
            data = response.json()
        except Exception:
            return ""

        pages = data.get("query", {}).get("pages", {})
        for page in pages.values():
            extract = page.get("extract", "")
            if extract:
                return self._clean_text(extract, max_chars=10000)
        return ""

    def _fetch_wikipedia_tables(self, title: str) -> str:
        url_title = title.replace(" ", "_")
        try:
            response = requests.get(
                f"https://en.wikipedia.org/wiki/{url_title}",
                headers={"User-Agent": "Mozilla/5.0"},
                timeout=20,
            )
            response.raise_for_status()
            tables = pd.read_html(StringIO(response.text))
        except Exception:
            return ""

        table_outputs = []
        for index, df in enumerate(tables[:4]):
            table_outputs.append(
                f"table_index: {index}\n"
                f"table_shape: {df.shape[0]} rows x {df.shape[1]} columns\n"
                f"table_columns: {list(df.columns)}\n"
                f"table_preview:\n{df.head(50).to_string(index=False)}"
            )
        return "\n\n".join(table_outputs)

    @staticmethod
    def _clean_text(text: str, max_chars: int = 4000) -> str:
        text = re.sub(r"(?is)<script.*?>.*?</script>", " ", text)
        text = re.sub(r"(?is)<style.*?>.*?</style>", " ", text)
        text = re.sub(r"(?s)<[^>]+>", " ", text)
        text = html.unescape(text)
        text = re.sub(r"\s+", " ", text).strip()
        return text[:max_chars]

    def _draft_answer(self, state: AgentState) -> AgentState:
        operation_table_answer = self._answer_operation_table_question(
            state.get("question", "")
        )
        if operation_table_answer:
            return {"draft_answer": operation_table_answer}

        tool_output = state.get("tool_output", "")
        if state.get("task_type") == "code_question" and "python_stdout:" in tool_output:
            stdout = tool_output.split("python_stdout:", 1)[1].split("python_stderr:", 1)[0].strip()
            lines = [line.strip() for line in stdout.splitlines() if line.strip()]
            if lines:
                return {"draft_answer": lines[-1]}

        if state.get("task_type") == "table_question":
            answer = self._answer_table_question(
                state.get("question", ""),
                state.get("attachment_path", ""),
            )
            if answer:
                return {"draft_answer": answer}

        primary_evidence = self._extract_primary_evidence(tool_output)
        prompt = (
            f"Question: {state.get('question', '')}\n"
            f"Attachment path: {state.get('attachment_path', '') or 'None'}\n"
            f"Task type: {state.get('task_type', '')}\n"
            f"Analysis: {state.get('analysis', '')}\n\n"
            f"Primary evidence:\n{primary_evidence or 'None'}\n\n"
            f"selected tool: {state.get('selected_tool', '')}\n"
            f"Tool output: {state.get('tool_output', '')}\n"
            f"Observations: {'; '.join(state.get('observations', []))}\n\n"
            f"Constraints: {state.get('constraints', '')}\n\n"
            "Return the best final answer now. Output only the answer.\n"
            "Use Primary evidence first when it is available.\n"
            "Use the full Tool output only to resolve missing context or confirm the Primary evidence.\n"
            "Obey all formatting constraints and instructions provided in the constraints section.\n"
            "Do not include any item unless it satisfies every explicit condition in the question.\n"
            "Only use items from the question when the question provides a list.\n"
            "Exclude items that violate the constraints.\n"
            "When evidence contains multiple candidate facts, choose the fact whose relationship and requested attribute match the question exactly.\n"
            "Do not answer with a nearby fact only because it appears near the same entity, date, or event.\n"
            "Pay attention to whether the question asks for a name, surname, date, count, code, location, title, author, ranking, before/after item, or category.\n"
            "Examples: if asked for a country code, do not answer with the full country name; if asked for a surname, do not answer with a full name; if asked for the item before or after a target item, do not answer with the target item itself; if asked for a count, return the count rather than the list being counted.\n"
            "Return the final answer text with no explanation and no prefix."
        )
        content = self._call_llm(
            [
                {"role":"system", "content": SYSTEM_PROMPT},
                {"role":"user", "content": prompt},
            ]
        )
        return {"draft_answer": content}

    @staticmethod
    def _extract_primary_evidence(tool_output: str) -> str:
        sections = []
        for part in tool_output.split("relevant_page_snippets:\n")[1:]:
            section = part.split("\npage_text:", 1)[0].strip()
            if section:
                sections.append(section)
        return "\n\n".join(sections)[:6000]

    @staticmethod
    def _answer_operation_table_question(question: str) -> str:
        question_lower = question.lower()
        if "not commutative" not in question_lower and "commutative" not in question_lower:
            return ""

        table_lines = [
            line.strip()
            for line in question.splitlines()
            if line.strip().startswith("|") and line.strip().endswith("|")
        ]
        if len(table_lines) < 3:
            return ""

        rows = [
            [cell.strip() for cell in line.strip().strip("|").split("|")]
            for line in table_lines
        ]
        rows = [
            row
            for row in rows
            if row and not all(set(cell) <= {"-"} for cell in row if cell)
        ]
        if len(rows) < 2 or len(rows[0]) < 2:
            return ""

        headers = rows[0][1:]
        operation_table = {}
        for row in rows[1:]:
            if len(row) != len(headers) + 1:
                return ""
            row_label = row[0]
            operation_table[row_label] = dict(zip(headers, row[1:]))

        if set(headers) - set(operation_table):
            return ""

        counterexample_elements = set()
        for left in headers:
            for right in headers:
                left_result = operation_table[left].get(right)
                right_result = operation_table[right].get(left)
                if left_result != right_result:
                    counterexample_elements.update([left, right])

        if not counterexample_elements:
            return ""
        return ", ".join(sorted(counterexample_elements))

    def _answer_table_question(self, question: str, attachment_path: str) -> str:
        if not attachment_path:
            return ""

        path = Path(attachment_path)
        if not path.exists() or path.suffix.lower() not in {".csv", ".xlsx", ".xls"}:
            return ""

        if path.suffix.lower() == ".csv":
            df = pd.read_csv(path)
        else:
            df = pd.read_excel(path)

        if df.empty:
            return ""

        question_lower = question.lower()
        if any(word in question_lower for word in ["sum", "total"]):
            operation = "sum"
        elif any(word in question_lower for word in ["how many", "count", "number of"]):
            operation = "count"
        elif any(word in question_lower for word in ["highest", "most", "max", "maximum"]):
            operation = "max"
        elif any(word in question_lower for word in ["lowest", "least", "min", "minimum"]):
            operation = "min"
        else:
            return ""

        if operation == "count":
            return str(len(df))

        numeric_columns = list(df.select_dtypes(include="number").columns)
        if not numeric_columns:
            return ""

        target_column = ""
        for column in numeric_columns:
            if str(column).lower() in question_lower:
                target_column = column
                break
        if not target_column:
            target_column = numeric_columns[-1]

        force_two_decimals = (
            "two decimal" in question_lower
            or "usd" in question_lower
            or "dollar" in question_lower
        )

        if operation == "sum":
            if "food" in question_lower and (
                "not including drinks" in question_lower
                or "excluding drinks" in question_lower
                or "without drinks" in question_lower
            ):
                drink_markers = {
                    "drink",
                    "drinks",
                    "beverage",
                    "beverages",
                    "soda",
                    "coffee",
                    "tea",
                    "water",
                    "juice",
                    "beer",
                    "wine",
                }
                food_columns = [
                    column
                    for column in numeric_columns
                    if not any(marker in str(column).lower() for marker in drink_markers)
                ]
                if food_columns:
                    return self._format_scalar_answer(
                        df[food_columns].sum(numeric_only=True).sum(),
                        force_two_decimals=force_two_decimals,
                    )

            return self._format_scalar_answer(
                df[target_column].sum(),
                force_two_decimals=force_two_decimals,
            )
        if operation == "max":
            best_row = df.loc[df[target_column].idxmax()]
            return self._format_table_row_answer(best_row)
        if operation == "min":
            best_row = df.loc[df[target_column].idxmin()]
            return self._format_table_row_answer(best_row)

        return ""

    @staticmethod
    def _format_scalar_answer(value, force_two_decimals: bool = False) -> str:
        try:
            number = float(value)
        except (TypeError, ValueError):
            return str(value)

        if force_two_decimals:
            return f"{number:.2f}"
        if number.is_integer():
            return str(int(number))
        return f"{number:.2f}"

    @staticmethod
    def _format_table_row_answer(row) -> str:
        preferred_columns = [
            "name",
            "item",
            "product",
            "menu item",
            "country",
            "team",
            "player",
            "person",
        ]
        lower_to_original = {str(column).lower(): column for column in row.index}
        for column in preferred_columns:
            if column in lower_to_original:
                return str(row[lower_to_original[column]])
        return str(row.iloc[0])
    
    def _validate_answer(self, state: AgentState) -> AgentState:
        # This is a placeholder for answer validation logic.
        # In a real implementation, you would check if the draft answer meets the constraints
        # and possibly use additional tools or logic to validate it.
        draft = state.get("draft_answer", "").strip()
        validation_passed = bool(draft)  # Simple check: answer is not empty
        validation_notes = "Answer is valid." if validation_passed else "Answer is empty."
        final_answer = draft.splitlines()[0].strip() if draft else ""
        return {
            "validation_passed": validation_passed,
            "validation_notes": validation_notes,
            "final_answer": final_answer
        }

    def _finalize_answer(self, state: AgentState) -> AgentState:
        draft = state.get("draft_answer", "").strip()
        final_answer = draft.splitlines()[0].strip() if draft else ""
        return {"final_answer": final_answer}

    @staticmethod
    def _classify_task(task: TaskRecord) -> str:
        if task.file_name:
            suffix = Path(task.file_name).suffix.lower()
            if suffix in {".png", ".jpg", ".jpeg", ".gif", ".webp"}:
                return "image_question"
            if suffix in {".csv", ".xlsx", ".xls"}:
                return "table_question"
            if suffix in {".pdf", ".txt", ".md", ".docx"}:
                return "document_question"
            if suffix in {".py", ".java", ".cpp", ".js"}:
                return "code_question"
            if suffix in {".json", ".xml", ".yaml", ".yml"}:
                return "structured_data_question"
            if suffix in {".mp4", ".avi", ".mov", ".mkv"}:
                return "video_question"
            return "file_question"

        question = task.question.lower()
        if "http://" in question or "www" in question or "https://" in question:
            return "web_question"
        lookup_markers = [
            "how many",
            "what year",
            "who ",
            "which ",
            "what country",
            "what is the surname",
            "as of",
            "january",
            "february",
            "march",
            "april",
            "may",
            "june",
            "july",
            "august",
            "september",
            "october",
            "november",
            "december",
        ]
        if any(marker in question for marker in lookup_markers):
            return "lookup_or_count_question"
        if "calculate" in question or "sum" in question:
            return "calculation_question"
        return "text_question"