Update app.py

app.py

@@ -31,15 +31,15 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 GROQ_TEXT_MODEL = os.getenv("GROQ_TEXT_MODEL", "llama-3.1-8b-instant")
 GROQ_FINAL_MODEL = os.getenv("GROQ_FINAL_MODEL", "llama-3.1-8b-instant")
 GROQ_STRONG_MODEL = os.getenv("GROQ_STRONG_MODEL", "openai/gpt-oss-20b")
-GROQ_RESEARCH_MODEL = os.getenv("GROQ_RESEARCH_MODEL", "llama-3.1-8b-instant")
+GROQ_RESEARCH_MODEL = os.getenv("GROQ_RESEARCH_MODEL", GROQ_STRONG_MODEL)
 GROQ_VISION_MODEL = os.getenv("GROQ_VISION_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
 GROQ_AUDIO_MODEL = os.getenv("GROQ_AUDIO_MODEL", "whisper-large-v3-turbo")
 
 GAIA_DIR = os.getenv("GAIA_DIR", "./data/gaia")
 ALLOW_CODE_EXECUTION = 1
 
-MAX_CONTEXT_CHARS = 18_000
-MAX_SEARCH_CONTEXT_CHARS = 12_000
+MAX_CONTEXT_CHARS = 24_000
+MAX_SEARCH_CONTEXT_CHARS = 20_000
 
 
 def get_groq_client() -> Groq:
@@ -196,6 +196,12 @@ def analyze_image(task_id: str, question: str = "") -> str:
     b64 = base64.standard_b64encode(data).decode("utf-8")
     mime = _image_mime(data, ct)
     prompt = question or "Describe this image. Extract all visible text, numbers, symbols, and key details."
+    if "chess" in prompt.lower():
+        prompt = (
+            f"{prompt}\n\n"
+            "This is a chess task. Identify the board coordinates, side to move, relevant pieces, checks, "
+            "mate threats, and the best move. Return the move in standard chess notation if possible."
+        )
 
     try:
         client = get_groq_client()
@@ -211,7 +217,7 @@ def analyze_image(task_id: str, question: str = "") -> str:
                 }
             ],
             temperature=0,
-            max_tokens=384,
+            max_tokens=768,
         )
         return resp.choices[0].message.content.strip()
     except Exception as e:
@@ -318,13 +324,56 @@ def read_spreadsheet_context(path: Path) -> str:
     return truncate_text("\n".join(parts), MAX_CONTEXT_CHARS)
 
 
+def build_spreadsheet_summary(path: Path) -> str:
+    parts: list[str] = []
+    xls = pd.ExcelFile(path)
+
+    for sheet_name in xls.sheet_names:
+        df = pd.read_excel(path, sheet_name=sheet_name)
+        if df.empty:
+            continue
+
+        work = df.copy()
+        work.columns = [str(col).strip() for col in work.columns]
+        numeric_cols = [
+            col
+            for col in work.columns
+            if pd.api.types.is_numeric_dtype(work[col])
+        ]
+        categorical_cols = [
+            col
+            for col in work.columns
+            if col not in numeric_cols and work[col].nunique(dropna=True) <= 40
+        ]
+
+        parts.append(f"Sheet: {sheet_name}")
+        if numeric_cols:
+            totals = work[numeric_cols].sum(numeric_only=True).sort_values(ascending=False)
+            parts.append("Numeric column totals:")
+            parts.append(totals.to_string())
+
+        for category_col in categorical_cols[:6]:
+            if not numeric_cols:
+                break
+            grouped = work.groupby(category_col, dropna=False)[numeric_cols].sum(numeric_only=True)
+            if grouped.empty:
+                continue
+            parts.append(f"Totals grouped by {category_col}:")
+            parts.append(grouped.head(40).to_csv())
+
+        if len("\n".join(parts)) > MAX_CONTEXT_CHARS:
+            break
+
+    return truncate_text("\n".join(parts), MAX_CONTEXT_CHARS)
+
+
 def read_code_context(path: Path) -> str:
     source = path.read_text(encoding="utf-8", errors="replace")
     parts = [f"Code file: {path.name}", "--- Source code ---", source]
     return truncate_text("\n".join(parts), MAX_CONTEXT_CHARS)
 
 
-def run_python_file(path: Path, timeout_seconds: int = 8) -> str:
+def run_python_file(path: Path, timeout_seconds: int = 45) -> str:
     if not ALLOW_CODE_EXECUTION:
         return "execution skipped"
     if path.suffix.lower() != ".py":
@@ -397,6 +446,15 @@ def fetch_url_text(url: str, limit: int = 8000) -> str:
         return f"[fetch error: {type(e).__name__}: {e}]"
 
 
+def likely_relevant_url(url: str) -> bool:
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"}:
+        return False
+    if any(skip in parsed.netloc for skip in ["youtube.com", "youtu.be", "facebook.com", "x.com"]):
+        return False
+    return True
+
+
 def ddg_search(query: str, max_results: int = 5) -> list[dict[str, str]]:
     try:
         results = DDGS().text(query, max_results=max_results)
@@ -416,7 +474,7 @@ def ddg_search(query: str, max_results: int = 5) -> list[dict[str, str]]:
 
 
 def build_research_queries(question: str, base_query: str) -> list[str]:
-    queries: list[str] = [base_query, question]
+    queries: list[str] = []
     quoted_phrases = re.findall(r'["“]([^"”]{3,120})["”]', question)
     if quoted_phrases:
         queries.append(" ".join(f'"{phrase}"' for phrase in quoted_phrases[:4]))
@@ -436,7 +494,18 @@ def build_research_queries(question: str, base_query: str) -> list[str]:
 
     years = re.findall(r"\b(?:19|20)\d{2}\b", question)
     if years:
-        queries.append(f"{base_query} {' '.join(years[:4])}")
+        year_text = " ".join(years[:4])
+        if capitalized_terms:
+            queries.append(f"{' '.join(capitalized_terms[:3])} {year_text}")
+        queries.append(f"{base_query} {year_text}")
+
+    q_lower = question.lower()
+    if "wikipedia" in q_lower and capitalized_terms:
+        queries.append(f"site:en.wikipedia.org {' '.join(capitalized_terms[:4])}")
+    if "wikipedia" in q_lower:
+        queries.append(f"site:en.wikipedia.org {base_query}")
+
+    queries += [base_query, question]
 
     deduped: list[str] = []
     for query in queries:
@@ -446,6 +515,32 @@ def build_research_queries(question: str, base_query: str) -> list[str]:
     return deduped[:5]
 
 
+def build_additional_research_queries(question: str, previous_queries: list[str]) -> list[str]:
+    messages = [
+        SystemMessage(content=(
+            "Create 3 concise web search queries for answering the task. "
+            "Prefer exact entity names, dates, source names, and required answer type. "
+            "Return one query per line, no numbering."
+        )),
+        HumanMessage(content=question),
+    ]
+    try:
+        llm = make_chat_model(GROQ_FINAL_MODEL, max_tokens=160)
+        raw = llm.invoke(messages).content
+    except Exception as e:
+        print(f"[query expansion warning] {type(e).__name__}: {e}")
+        return []
+
+    queries: list[str] = []
+    previous = {q.lower() for q in previous_queries}
+    for line in raw.splitlines():
+        query = clean_answer(re.sub(r"^\s*[-*\d.)]+\s*", "", line))
+        query = re.sub(r"\s+", " ", query).strip()
+        if query and query.lower() not in previous:
+            queries.append(query)
+    return queries[:3]
+
+
 def fetch_youtube_timedtext(video_id: str) -> str:
     urls = [
         f"https://video.google.com/timedtext?lang=en&v={video_id}",
@@ -470,12 +565,17 @@ def fetch_youtube_timedtext(video_id: str) -> str:
 
 def build_youtube_context(question: str, video_id: str | None) -> str:
     queries: list[str] = []
+    quoted_phrases = re.findall(r'["“]([^"”]{3,120})["”]', question)
     if video_id:
         queries += [
             f'"{video_id}" transcript',
             f'"{video_id}" subtitles',
             f'"{video_id}"',
         ]
+        for phrase in quoted_phrases[:3]:
+            queries.append(f'"{video_id}" "{phrase}"')
+    if quoted_phrases:
+        queries.append(" ".join(f'"{phrase}"' for phrase in quoted_phrases[:3]))
 
     queries.append(question)
 
@@ -516,31 +616,66 @@ def build_youtube_context(question: str, video_id: str | None) -> str:
 def build_research_context(question: str, base_query: str) -> str:
     parts = [f"Question: {question}", f"Primary query: {base_query}"]
     seen_urls: set[str] = set()
+    queries = build_research_queries(question, base_query)
 
-    for query in build_research_queries(question, base_query):
+    for query in queries:
         parts.append(f"\n=== Search query: {query} ===")
-        results = ddg_search(query, max_results=5)
+        results = ddg_search(query, max_results=6)
         if not results:
             parts.append(safe_tool_run(web_search_tool, query, limit=2000))
             continue
 
+        fetched_count = 0
         for i, result in enumerate(results, 1):
             url = result["url"]
             title = result["title"]
             body = result["body"]
             parts.append(f"[{i}] {title}\nURL: {url}\nSnippet: {body}")
 
-            parsed = urlparse(url)
             if not url or url in seen_urls:
                 continue
-            if parsed.scheme not in {"http", "https"}:
+            if not likely_relevant_url(url):
                 continue
-            if any(skip in parsed.netloc for skip in ["youtube.com", "youtu.be", "facebook.com", "x.com"]):
+            if fetched_count >= 2:
                 continue
 
             seen_urls.add(url)
             fetched = fetch_url_text(url, limit=5000)
             if fetched and not fetched.startswith("[fetch error"):
+                fetched_count += 1
+                parts.append(f"Fetched text from {url}:\n{fetched}")
+            if len("\n".join(parts)) > MAX_SEARCH_CONTEXT_CHARS:
+                return truncate_text("\n".join(parts), MAX_SEARCH_CONTEXT_CHARS)
+
+    return truncate_text("\n".join(parts), MAX_SEARCH_CONTEXT_CHARS)
+
+
+def extend_research_context(question: str, context: str, used_query: str) -> str:
+    extra_queries = build_additional_research_queries(question, [used_query])
+    if not extra_queries:
+        return context
+
+    parts = [context, "\n=== Additional focused searches ==="]
+    seen_urls = set(re.findall(r"URL: (https?://\S+)", context))
+    for query in extra_queries:
+        parts.append(f"\n=== Search query: {query} ===")
+        results = ddg_search(query, max_results=6)
+        if not results:
+            parts.append(safe_tool_run(web_search_tool, query, limit=2000))
+            continue
+
+        fetched_count = 0
+        for i, result in enumerate(results, 1):
+            url = result["url"]
+            parts.append(f"[{i}] {result['title']}\nURL: {url}\nSnippet: {result['body']}")
+            if not url or url in seen_urls or not likely_relevant_url(url):
+                continue
+            if fetched_count >= 2:
+                continue
+            seen_urls.add(url)
+            fetched = fetch_url_text(url, limit=5000)
+            if fetched and not fetched.startswith("[fetch error"):
+                fetched_count += 1
                 parts.append(f"Fetched text from {url}:\n{fetched}")
             if len("\n".join(parts)) > MAX_SEARCH_CONTEXT_CHARS:
                 return truncate_text("\n".join(parts), MAX_SEARCH_CONTEXT_CHARS)
@@ -690,8 +825,6 @@ def direct_question(question: str) -> bool:
         "grocery list",
         "shopping list",
         "given this table",
-        "alphabetical order",
-        "sort",
         "opposite of",
         "reverse",
         "what is the final numeric output",
@@ -731,7 +864,8 @@ class BasicAgent:
     def __init__(self):
         self.answer_llm = make_chat_model(GROQ_TEXT_MODEL, max_tokens=256)
         self.final_llm = make_chat_model(GROQ_FINAL_MODEL, max_tokens=48)
-        self.research_llm = make_chat_model(GROQ_RESEARCH_MODEL, max_tokens=384)
+        self.strong_llm = make_chat_model(GROQ_STRONG_MODEL, max_tokens=512)
+        self.research_llm = make_chat_model(GROQ_RESEARCH_MODEL, max_tokens=512)
 
         self.graph = self.build_graph()
 
@@ -844,7 +978,7 @@ class BasicAgent:
             question=question,
             context=context,
             context_label="Image analysis",
-            llm=self.answer_llm,
+            llm=self.strong_llm,
         )
         return {"context": context, "raw_answer": raw_answer}
 
@@ -866,23 +1000,16 @@ class BasicAgent:
         path = state["local_path"]
         question = state["question"]
 
-        xls = pd.ExcelFile(path)
-        parts = []
-
-        for sheet in xls.sheet_names:
-            df = pd.read_excel(path, sheet_name=sheet)
-            parts.append(f"Sheet: {sheet}")
-            parts.append(f"Columns: {list(df.columns)}")
-            parts.append(f"Shape: {df.shape}")
-            parts.append(df.head(20).to_csv(index=False))
-
-        context = "\n\n".join(parts)
+        context = read_spreadsheet_context(Path(path))
+        summary = build_spreadsheet_summary(Path(path))
+        if summary:
+            context = f"{context}\n\n--- Computed spreadsheet summary ---\n{summary}"
 
         raw_answer = self.answer_from_context(
             question=question,
-            context=context[:12000],
-            context_label="spreadsheet preview",
-            llm=self.answer_llm,
+            context=context,
+            context_label="Spreadsheet data and computed summary",
+            llm=self.strong_llm,
         )
 
         return {"context": context, "raw_answer": raw_answer}
@@ -909,7 +1036,7 @@ class BasicAgent:
             question=question,
             context=context,
             context_label="Code and execution result",
-            llm=self.answer_llm,
+            llm=self.strong_llm,
         )
         return {"context": context, "raw_answer": raw_answer}
 
@@ -957,6 +1084,16 @@ class BasicAgent:
             llm=self.research_llm,
         )
 
+        if is_bad_answer(raw_answer):
+            context = extend_research_context(question, context, query)
+            print(f"[research extended context len] {len(context)}")
+            raw_answer = self.answer_from_context(
+                question=question,
+                context=context,
+                context_label="Extended web research results",
+                llm=self.strong_llm,
+            )
+
         print(f"[research raw_answer] {repr(raw_answer[:500])}")
         return {"context": context, "raw_answer": raw_answer}
     
@@ -973,6 +1110,15 @@ class BasicAgent:
             llm=self.research_llm,
         )
 
+        if is_bad_answer(raw_answer):
+            context = extend_research_context(question, context, question)
+            raw_answer = self.answer_from_context(
+                question=question,
+                context=context,
+                context_label="Extended YouTube/web transcript search results",
+                llm=self.strong_llm,
+            )
+
         return {
             "context": context,
             "raw_answer": raw_answer,