Upload 5 files

agents/__init__.py

@@ -0,0 +1,13 @@
+from agents.supervisor import (
+    build_chatbot_graph,
+    route_after_analysis,
+    route_after_data_extraction,
+    run_user_query,
+)
+
+__all__ = [
+    "build_chatbot_graph",
+    "route_after_analysis",
+    "route_after_data_extraction",
+    "run_user_query",
+]

agents/data_extraction.py

@@ -0,0 +1,105 @@
+"""
+Data extraction agent: builds the graph node that extracts metrics data
+based on user queries and filters.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import pandas as pd
+
+from tools import data_extractor, filter_extractor
+from tools.state import GraphState
+
+_DEFAULT_FILTER_COLUMNS = [
+    "Region", "Period", "Cluster", "Product", "Calculation_Type", "TA Market", "Class", "Level",
+]
+
+
+def _tool_call(tool_or_fn, **kwargs):
+    """Invoke a tool or callable, handling both raw functions and tool wrappers."""
+    callable_obj = getattr(tool_or_fn, "func", tool_or_fn)
+    return callable_obj(**kwargs)
+
+
+def _candidate_values(df: pd.DataFrame, max_values: int = 100) -> dict[str, list[str]]:
+    values: dict[str, list[str]] = {}
+    for column in df.columns:
+        unique_values = df[column].dropna().astype(str).str.strip().unique().tolist()
+        values[column] = sorted([v for v in unique_values if v])[:max_values]
+    return values
+
+
+def _validate_filters(filters: dict, columns: list[str]) -> dict:
+    allowed = set(columns)
+    validated = {}
+    for key, value in (filters or {}).items():
+        if key not in allowed:
+            continue
+        if isinstance(value, list):
+            cleaned = [str(v).strip() for v in value if str(v).strip()]
+            if cleaned:
+                validated[key] = cleaned
+    return validated
+
+
+def build_data_extraction_node(
+    metrics_df: pd.DataFrame | None = None,
+    llm_invoke: Callable[[str], object] | None = None,
+    db_query_fn: Callable[[dict], pd.DataFrame] | None = None,
+    column_values: dict[str, list[str]] | None = None,
+):
+    """
+    Factory that returns a data extraction node for the LangGraph.
+
+    SQL mode (preferred):
+        Pass ``db_query_fn`` and ``column_values``. Filters are extracted from the
+        user message, then a targeted SQL query is executed.
+
+    In-memory mode (legacy):
+        Pass ``metrics_df`` (the full metrics DataFrame). Filters are applied in
+        memory via pandas.
+    """
+    if db_query_fn is not None:
+        available_columns = list(column_values.keys()) if column_values else _DEFAULT_FILTER_COLUMNS
+        value_map: dict[str, list[str]] = column_values or {}
+    else:
+        if metrics_df is None:
+            raise ValueError("Either metrics_df or db_query_fn must be provided.")
+        available_columns = metrics_df.columns.tolist()
+        value_map = _candidate_values(metrics_df)
+
+    def data_extraction_node(state: GraphState) -> dict:
+        user_query = state.get("user_query", "")
+        prior_filters = _validate_filters(state.get("filters", {}), available_columns)
+        conversation_history = state.get("conversation_history", []) or []
+
+        raw_filters = _tool_call(
+            filter_extractor,
+            user_query=user_query,
+            available_columns=available_columns,
+            column_values=value_map,
+            llm=llm_invoke,
+            prior_filters=prior_filters,
+            conversation_history=conversation_history,
+        )
+
+        filters = _validate_filters(raw_filters, available_columns)
+
+        if db_query_fn is not None:
+            fetched_df = db_query_fn(filters)
+            rows = fetched_df.to_dict(orient="records") if not fetched_df.empty else []
+        else:
+            rows = _tool_call(data_extractor, filters=filters, dataframe=metrics_df)
+
+        error_message = state.get("error_message", "")
+        if not rows:
+            error_message = "I couldn't find any data matching your criteria."
+
+        return {
+            "filters": filters,
+            "extracted_data": rows,
+            "error_message": error_message,
+        }
+
+    return data_extraction_node

agents/data_knowledge.py

@@ -0,0 +1,34 @@
+"""
+Data knowledge agent: builds the graph node that retrieves parameter/configuration
+information based on user queries (e.g. market mappings, cluster definitions).
+"""
+from __future__ import annotations
+
+from tools import parameter_reader
+from tools.state import GraphState
+
+
+def _tool_call(tool_or_fn, **kwargs):
+    callable_obj = getattr(tool_or_fn, "func", tool_or_fn)
+    return callable_obj(**kwargs)
+
+
+def build_data_knowledge_node():
+    """
+    Factory that returns a data knowledge node for the graph.
+    The node uses the parameter reader to fetch relevant parameter info for the user query.
+    """
+    def data_knowledge_node(state: GraphState) -> dict:
+        user_query = state.get("user_query", "")
+        rows = _tool_call(parameter_reader, user_query=user_query)
+
+        error_message = state.get("error_message", "")
+        if not rows and not error_message:
+            error_message = "I couldn't find relevant parameter information for your request."
+
+        return {
+            "parameter_data": rows,
+            "error_message": error_message,
+        }
+
+    return data_knowledge_node

agents/query_analyzer.py

@@ -0,0 +1,89 @@
+"""
+Query analyzer agent: classifies user intent (data retrieval, parameter info, or out-of-scope)
+to route the chatbot to the appropriate handler.
+"""
+from __future__ import annotations
+
+import json
+from typing import Callable
+
+from tools.state import GraphState
+
+
+def _format_history(history: list[dict], max_turns: int = 6) -> str:
+    if not history:
+        return ""
+    clipped = history[-max_turns:]
+    lines = []
+    for msg in clipped:
+        role = str(msg.get("role", "user")).strip().lower()
+        content = str(msg.get("content", "")).strip()
+        if content:
+            lines.append(f"{role}: {content}")
+    return "\n".join(lines)
+
+
+def build_query_analyzer_node(
+    llm_json_call: Callable[[str], dict] | None = None,
+):
+    """
+    Factory that returns a query analyzer node for the graph.
+    The node classifies user intent via LLM into: data_retrieval, parameter_info,
+    both, or out_of_scope.
+    """
+    def query_analyzer_node(state: GraphState) -> dict:
+        user_query = state.get("user_query", "")
+        intent, is_valid, error = "out_of_scope", False, ""
+        prior_filters = state.get("filters", {}) or {}
+        history_text = _format_history(state.get("conversation_history", []) or [])
+
+        if llm_json_call is not None:
+            prompt = f"""
+You are a query classifier for a pharma market metrics chatbot.
+Classify user queries into one intent:
+- data_retrieval: user wants metrics/data (volume, share, growth, etc.)
+- parameter_info: user asks about parameters (cluster mapping, regions, etc.)
+- both: user asks for data AND parameter information at the same time
+- out_of_scope: unrelated topics (weather, general knowledge, etc.)
+
+Return strict JSON only with keys:
+intent, is_valid, error_message, confidence
+
+CRITICAL: Set is_valid=true for data_retrieval when the user asks a data question,
+including when they want to CHANGE filters (e.g. "same but for Spain", "what about Germany",
+"do the same in LATAM"). Do NOT set is_valid=false just because the user requests
+different country/region/period than previously applied - the filter extraction step
+will handle that. Only set is_valid=false for genuinely out-of-scope questions.
+
+User query: {user_query}
+Previously applied filters: {json.dumps(prior_filters, ensure_ascii=True)}
+Recent conversation:
+{history_text}
+"""
+            try:
+                result = llm_json_call(prompt)
+                if isinstance(result, str):
+                    result = json.loads(result)
+                intent = result.get("intent", intent)
+                is_valid = bool(result.get("is_valid", is_valid))
+                error = result.get("error_message", error) or error
+            except Exception:
+                pass
+
+        if intent not in {"data_retrieval", "parameter_info", "both", "out_of_scope"}:
+            intent = "out_of_scope"
+            is_valid = False
+
+        if intent == "out_of_scope" and not error:
+            error = (
+                "This question is outside my knowledge domain. I can help with data "
+                "queries and parameter information about the demo pharma market."
+            )
+
+        return {
+            "query_intent": intent,
+            "is_valid": is_valid,
+            "error_message": error,
+        }
+
+    return query_analyzer_node

agents/supervisor.py

@@ -0,0 +1,369 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Callable
+
+import pandas as pd
+
+from agents.data_extraction import build_data_extraction_node
+from agents.data_knowledge import build_data_knowledge_node
+from agents.query_analyzer import build_query_analyzer_node
+from tools.leading_country_tools import calculate_leading_country, wants_leading_country
+from tools.state import GraphState, create_initial_state
+
+
+LOG_PREFIX = "[Supervisor]"
+
+
+def _log(msg: str) -> None:
+    print(f"{LOG_PREFIX} {msg}")
+
+
+def _wrap_node(name: str, node_fn: Callable[[GraphState], dict]) -> Callable[[GraphState], dict]:
+    """Wrap a graph node to log entry/exit and a one-line summary to the terminal."""
+
+    def wrapped(state: GraphState) -> dict:
+        _log(f"→ Entering node: {name}")
+        result = node_fn(state)
+        summary_parts = []
+        if result.get("error_message"):
+            err = (result["error_message"] or "")[:60]
+            summary_parts.append(f"error={err!r}...")
+        if "extracted_data" in result:
+            n = len(result.get("extracted_data") or [])
+            summary_parts.append(f"extracted_data={n} rows")
+        if "parameter_data" in result:
+            n = len(result.get("parameter_data") or [])
+            summary_parts.append(f"parameter_data={n} rows")
+        if "leading_country_result" in result:
+            n = len(result.get("leading_country_result") or [])
+            summary_parts.append(f"leading_country_result={n} rows")
+        if "query_intent" in result:
+            summary_parts.append(f"intent={result['query_intent']!r} is_valid={result.get('is_valid', '?')}")
+        if "filters" in result and result["filters"]:
+            summary_parts.append(f"filters={list(result['filters'].keys())}")
+        _log(f"← Exiting node: {name}" + (f" ({', '.join(summary_parts)})" if summary_parts else ""))
+        return result
+
+    return wrapped
+
+
+def _invoke_llm_json(client: Any, prompt: str) -> dict:
+    response = client.chat.completions.create(
+        model=client._default_deployment,  # type: ignore[attr-defined]
+        messages=[
+            {"role": "system", "content": "Return strict JSON only."},
+            {"role": "user", "content": prompt},
+        ],
+        temperature=0,
+    )
+    content = response.choices[0].message.content or "{}"
+    try:
+        return json.loads(content)
+    except Exception:
+        match = re.search(r"\{.*\}", content, flags=re.DOTALL)
+        if not match:
+            return {}
+        try:
+            return json.loads(match.group(0))
+        except Exception:
+            return {}
+
+
+def _invoke_llm_text(client: Any, prompt: str):
+    response = client.chat.completions.create(
+        model=client._default_deployment,  # type: ignore[attr-defined]
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0,
+    )
+    return response.choices[0].message.content or ""
+
+
+def _format_rows(rows: list[dict], max_rows: int = 10) -> str:
+    if not rows:
+        return ""
+    df = pd.DataFrame(rows[:max_rows])
+    return df.to_csv(index=False)
+
+
+def _deterministic_final_response(state: GraphState) -> str:
+    if state.get("error_message"):
+        _log("generate_response: returning error_message as final_response")
+        return state["error_message"]
+
+    sections: list[str] = []
+    extracted_data = state.get("extracted_data", [])
+    leading_country_result = state.get("leading_country_result", [])
+    parameter_data = state.get("parameter_data", [])
+
+    if extracted_data:
+        sections.append(f"Found {len(extracted_data)} matching data rows.")
+        sections.append("Sample data:")
+        sections.append(f"```csv\n{_format_rows(extracted_data)}\n```")
+
+    if parameter_data:
+        sections.append(f"Found {len(parameter_data)} relevant parameter rows.")
+        sections.append("Reference data:")
+        sections.append(f"```csv\n{_format_rows(parameter_data)}\n```")
+
+    if leading_country_result:
+        sections.append(f"Calculated {len(leading_country_result)} leading-country result rows.")
+        sections.append("Leading-country calculations:")
+        sections.append(f"```csv\n{_format_rows(leading_country_result)}\n```")
+
+    if not sections:
+        _log("generate_response: no data → asking user to clarify")
+        return "Could you please clarify what you're looking for?"
+
+    sections.append("If you want, I can narrow this further by region, period, product, or market.")
+    return "\n\n".join(sections)
+
+
+def _format_history(history: list[dict], max_turns: int = 6) -> str:
+    if not history:
+        return ""
+    clipped = history[-max_turns:]
+    lines: list[str] = []
+    for msg in clipped:
+        role = str(msg.get("role", "user")).strip().lower()
+        content = str(msg.get("content", "")).strip()
+        if not content:
+            continue
+        lines.append(f"{role}: {content}")
+    return "\n".join(lines)
+
+
+def build_generate_response_node(
+    llm_text_call: Callable[[str], object] | None = None,
+):
+    def generate_response_node(state: GraphState) -> dict:
+        fallback_response = _deterministic_final_response(state)
+
+        if llm_text_call is None:
+            _log("generate_response: llm unavailable, using deterministic response")
+            return {"final_response": fallback_response}
+
+        extracted_data = state.get("extracted_data", [])
+        leading_country_result = state.get("leading_country_result", [])
+        parameter_data = state.get("parameter_data", [])
+        filters = state.get("filters", {}) or {}
+        conversation_history = state.get("conversation_history", []) or []
+        user_query = state.get("user_query", "")
+
+        prompt = f"""
+You are the supervisor response writer for a pharma market metrics assistant.
+Write the final answer to the user based only on the context below.
+
+Rules:
+- Use only facts from the provided context. Do not invent numbers, metrics, or mappings.
+- Keep the answer concise and business-friendly.
+- Mention applied filters (Region, Period, Calculation_Type, Cluster) briefly.
+- When rank columns are present (Current_Volume_Rank, Current_Value_Rank), use them directly
+  to determine ordering — the lowest rank number is the best-performing product.
+- If the data contains the requested information, answer directly and confidently. Do NOT ask
+  clarifying questions when the data is sufficient to answer.
+- Only ask a clarifying question if critical information is genuinely missing (e.g. no data returned).
+- If Leading-country calculations are present, use them as the source of truth.
+- Do not say absolute contribution is unavailable when Current_Value/Current_Volume and Growth fields
+  are available; the leading-country calculation reconstructs deltas from those fields.
+
+Context:
+User query: {user_query}
+Detected intent: {state.get("query_intent", "out_of_scope")}
+Error message: {state.get("error_message", "")}
+Applied filters: {json.dumps(filters, ensure_ascii=True)}
+Extracted data row count: {len(extracted_data)}
+Extracted data:
+{_format_rows(extracted_data, max_rows=40)}
+Leading-country calculation row count: {len(leading_country_result)}
+Leading-country calculations:
+{_format_rows(leading_country_result, max_rows=40)}
+Parameter row count: {len(parameter_data)}
+Parameter data sample:
+{_format_rows(parameter_data, max_rows=100)}
+Recent conversation history:
+{_format_history(conversation_history)}
+
+Return plain text only. No markdown code fences.
+"""
+        try:
+            generated = llm_text_call(prompt)
+            text = getattr(generated, "content", str(generated)).strip()
+            if text:
+                _log(
+                    "generate_response: llm synthesis success "
+                    f"(extracted={len(extracted_data)}, parameter={len(parameter_data)} rows)"
+                )
+                return {"final_response": text}
+        except Exception as exc:
+            _log(f"generate_response: llm synthesis failed, using fallback ({exc})")
+
+        _log(
+            "generate_response: using deterministic fallback "
+            f"(extracted={len(extracted_data)}, parameter={len(parameter_data)} rows)"
+        )
+        return {"final_response": fallback_response}
+
+    return generate_response_node
+
+
+def build_leading_country_node(
+    db_query_fn: Any | None = None,
+    metrics_df: pd.DataFrame | None = None,
+):
+    def leading_country_node(state: GraphState) -> dict:
+        results = calculate_leading_country(
+            user_query=state.get("user_query", ""),
+            filters=state.get("filters", {}) or {},
+            extracted_rows=state.get("extracted_data", []) or [],
+            db_query_fn=db_query_fn,
+            metrics_df=metrics_df,
+        )
+        return {"leading_country_result": results}
+
+    return leading_country_node
+
+
+def route_after_analysis(state: GraphState) -> str:
+    is_valid = state.get("is_valid", False)
+    intent = state.get("query_intent", "out_of_scope")
+    if not is_valid:
+        _log(f"route_after_analysis: is_valid=False → generate_response")
+        return "generate_response"
+    if intent == "data_retrieval":
+        _log(f"route_after_analysis: intent={intent!r} → data_extraction")
+        return "data_extraction"
+    if intent == "parameter_info":
+        _log(f"route_after_analysis: intent={intent!r} → data_knowledge")
+        return "data_knowledge"
+    if intent == "both":
+        _log(f"route_after_analysis: intent={intent!r} → data_extraction (then data_knowledge)")
+        return "data_extraction"
+    _log(f"route_after_analysis: intent={intent!r} → generate_response")
+    return "generate_response"
+
+
+def route_after_data_extraction(state: GraphState) -> str:
+    if wants_leading_country(state.get("user_query", "")):
+        _log("route_after_data_extraction: leading-country query → leading_country")
+        return "leading_country"
+    if state.get("query_intent") == "both":
+        _log("route_after_data_extraction: intent=both → data_knowledge")
+        return "data_knowledge"
+    _log("route_after_data_extraction: → generate_response")
+    return "generate_response"
+
+
+def route_after_leading_country(state: GraphState) -> str:
+    if state.get("query_intent") == "both":
+        _log("route_after_leading_country: intent=both → data_knowledge")
+        return "data_knowledge"
+    _log("route_after_leading_country: → generate_response")
+    return "generate_response"
+
+
+def build_chatbot_graph(
+    metrics_df: pd.DataFrame | None = None,
+    azure_openai_client: Any | None = None,
+    db_query_fn: Any | None = None,
+    column_values: dict | None = None,
+):
+    """
+    Build and compile the LangGraph chatbot.
+
+    ``azure_openai_client`` is named for backward compatibility with the original
+    codebase but actually accepts ANY OpenAI-compatible chat client that exposes
+    ``client.chat.completions.create(...)`` — including
+    ``huggingface_hub.InferenceClient``.
+    """
+    try:
+        from langgraph.graph import END, StateGraph
+    except Exception as exc:
+        raise ImportError(
+            "LangGraph is required. Install with `pip install langgraph`."
+        ) from exc
+
+    llm_json = None
+    llm_text = None
+    if azure_openai_client is not None:
+        llm_json = lambda prompt: _invoke_llm_json(azure_openai_client, prompt)
+        llm_text = lambda prompt: _invoke_llm_text(azure_openai_client, prompt)
+
+    workflow = StateGraph(GraphState)
+
+    workflow.add_node(
+        "query_analyzer",
+        _wrap_node("query_analyzer", build_query_analyzer_node(llm_json_call=llm_json)),
+    )
+    workflow.add_node(
+        "data_extraction",
+        _wrap_node(
+            "data_extraction",
+            build_data_extraction_node(
+                metrics_df=metrics_df,
+                llm_invoke=llm_text,
+                db_query_fn=db_query_fn,
+                column_values=column_values,
+            ),
+        ),
+    )
+    workflow.add_node(
+        "data_knowledge",
+        _wrap_node("data_knowledge", build_data_knowledge_node()),
+    )
+    workflow.add_node(
+        "leading_country",
+        _wrap_node(
+            "leading_country",
+            build_leading_country_node(db_query_fn=db_query_fn, metrics_df=metrics_df),
+        ),
+    )
+    workflow.add_node(
+        "generate_response",
+        _wrap_node("generate_response", build_generate_response_node(llm_text_call=llm_text)),
+    )
+
+    workflow.set_entry_point("query_analyzer")
+    workflow.add_conditional_edges(
+        "query_analyzer",
+        route_after_analysis,
+        {
+            "data_extraction": "data_extraction",
+            "data_knowledge": "data_knowledge",
+            "generate_response": "generate_response",
+        },
+    )
+    workflow.add_conditional_edges(
+        "data_extraction",
+        route_after_data_extraction,
+        {
+            "leading_country": "leading_country",
+            "data_knowledge": "data_knowledge",
+            "generate_response": "generate_response",
+        },
+    )
+    workflow.add_conditional_edges(
+        "leading_country",
+        route_after_leading_country,
+        {
+            "data_knowledge": "data_knowledge",
+            "generate_response": "generate_response",
+        },
+    )
+    workflow.add_edge("data_knowledge", "generate_response")
+    workflow.add_edge("generate_response", END)
+
+    return workflow.compile()
+
+
+def run_user_query(
+    app: Any,
+    user_query: str,
+    conversation_history: list[dict] | None = None,
+) -> GraphState:
+    _log(f"User query: {user_query[:80]}{'...' if len(user_query) > 80 else ''}")
+    state = create_initial_state(user_query=user_query, conversation_history=conversation_history)
+    result = app.invoke(state)
+    _log(f"Done. final_response length={len(result.get('final_response', '') or '')} chars")
+    return result