research_engine.py

"""
Real deep research engine.

Uses:
  * DashScope (Qwen) for decomposition, KG extraction, summarisation.
  * Serper for web search.

Exposes a single generator function `run_research(question, prior_nodes, prior_edges)`
that yields tuples ``(status, nodes, edges)`` after each iteration so the UI can
grow the knowledge graph progressively, exactly like the mock loop did.

Env vars (set as HF Space Secrets):
    DASHSCOPE_API_KEY   — required
    SERPER_API_KEY      — required
    DASHSCOPE_MODEL     — optional, default "qwen-plus"
"""

from __future__ import annotations

import json
import logging
import os
import re
import uuid
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Generator, Iterable

import requests

logger = logging.getLogger("research_engine")

DASHSCOPE_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"
SERPER_URL = "https://google.serper.dev/search"
DEFAULT_MODEL = os.getenv("DASHSCOPE_MODEL", "qwen-plus")

ENTITY_TYPES = ("company", "material", "product", "process", "region", "risk")


# ---------- small helpers ---------------------------------------------------

def _slug(name: str) -> str:
    s = re.sub(r"[^a-zA-Z0-9]+", "_", name.strip().lower()).strip("_")
    return s[:60] or "node"


def _keys_configured() -> bool:
    return bool(os.getenv("DASHSCOPE_API_KEY") and os.getenv("SERPER_API_KEY"))


# ---------- LLM call --------------------------------------------------------

def _call_llm(prompt: str, system: str = "", max_tokens: int = 1024, temperature: float = 0.2) -> str:
    key = os.getenv("DASHSCOPE_API_KEY")
    if not key:
        raise RuntimeError("DASHSCOPE_API_KEY is not set")
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    r = requests.post(
        DASHSCOPE_BASE,
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        json={"model": DEFAULT_MODEL, "messages": messages, "max_tokens": max_tokens, "temperature": temperature},
        timeout=60,
    )
    r.raise_for_status()
    return r.json()["choices"][0]["message"]["content"]


def _parse_json(text: str) -> Any:
    """Extract a JSON blob from an LLM response (may be wrapped in markdown)."""
    text = text.strip()
    m = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
    if m:
        text = m.group(1)
    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        text = text[start:end + 1]
    return json.loads(text)


# ---------- web search ------------------------------------------------------

def _serper_search(query: str, num: int = 5) -> list[dict]:
    key = os.getenv("SERPER_API_KEY")
    if not key:
        raise RuntimeError("SERPER_API_KEY is not set")
    try:
        r = requests.post(
            SERPER_URL,
            headers={"X-API-KEY": key, "Content-Type": "application/json"},
            json={"q": query, "num": num},
            timeout=30,
        )
        r.raise_for_status()
        data = r.json()
        results = []
        for item in data.get("organic", [])[:num]:
            results.append({
                "title":   item.get("title", ""),
                "url":     item.get("link", ""),
                "snippet": item.get("snippet", ""),
            })
        return results
    except Exception as e:  # network / credit / auth errors
        logger.warning("Serper error: %s", e)
        return []


# ---------- research steps --------------------------------------------------

DECOMPOSE_SYS = (
    "You are a supply-chain research planner. Given a question, produce 3–5 "
    "focused sub-queries that, run as web searches, would together gather the "
    "evidence needed to answer the question. Return only a JSON list of "
    "strings, no commentary."
)


def _decompose(question: str) -> list[str]:
    text = _call_llm(
        prompt=f"Question: {question}\n\nReturn a JSON array of 3–5 sub-queries.",
        system=DECOMPOSE_SYS,
        max_tokens=400,
    )
    try:
        out = json.loads(text) if text.strip().startswith("[") else _parse_json(text)
        if isinstance(out, list):
            return [str(x) for x in out[:5]]
    except Exception as e:
        logger.warning("decompose parse failed: %s", e)
    return [question]


EXTRACT_SYS = (
    "You are a supply-chain knowledge-graph extractor. From the provided "
    "numbered search snippets, extract entities and relationships relevant "
    "to the question. Entity types MUST be one of: company, material, "
    "product, process, region, risk.\n\n"
    'Return ONLY valid JSON with this schema:\n'
    '{"nodes":[{"id":"<slug>","label":"<display>","type":"<entity_type>",'
    '"confidence":0.0-1.0,"description":"<one-sentence fact>",'
    '"refs":[<snippet-numbers-that-support-this-entity>]}],\n'
    '"edges":[{"from":"<id>","to":"<id>","label":"<short_relation>",'
    '"refs":[<snippet-numbers>]}]}\n'
    "Use lower-case snake_case ids. Every entity MUST include the `refs` "
    "array citing which numbered snippets (e.g. [1,3]) support it. Only add "
    "entities the snippets explicitly mention. At most 10 new nodes and 15 "
    "new edges per call."
)


def _extract_kg(question: str, snippets: list[dict], existing_ids: set[str], ref_offset: int = 0) -> dict:
    """Return ``{nodes, edges, refs}`` where each node/edge carries a list of
    global citation indices (1-based) pointing back into a shared ref2url map.
    """
    if not snippets:
        return {"nodes": [], "edges": []}
    snippet_text = "\n\n".join(
        f"[{i + 1 + ref_offset}] {s['title']}\n{s['snippet']}" for i, s in enumerate(snippets)
    )[:6000]
    existing = ", ".join(sorted(existing_ids)[:40]) or "(none yet)"
    prompt = (
        f"Question: {question}\n\n"
        f"Existing KG node ids (avoid duplicates, reuse where applicable): {existing}\n\n"
        f"Numbered search snippets (cite these in the refs fields):\n{snippet_text}\n\n"
        "Extract the KG JSON now."
    )
    try:
        text = _call_llm(prompt, system=EXTRACT_SYS, max_tokens=1500, temperature=0.15)
        data = _parse_json(text)
    except Exception as e:
        logger.warning("extract parse failed: %s", e)
        return {"nodes": [], "edges": []}

    nodes, edges = [], []
    valid_ids = set(existing_ids)
    snippet_count = len(snippets)

    def _clean_refs(raw) -> list[int]:
        out: list[int] = []
        if isinstance(raw, list):
            for r in raw:
                try:
                    n = int(r)
                except Exception:
                    continue
                # Refs from the LLM are 1-indexed relative to ref_offset.
                if 1 + ref_offset <= n <= ref_offset + snippet_count:
                    out.append(n)
        return out

    for n in data.get("nodes", []):
        if not isinstance(n, dict):
            continue
        nid = _slug(str(n.get("id") or n.get("label", "")))
        label = str(n.get("label") or nid).strip()
        ntype = str(n.get("type", "")).strip().lower()
        if ntype not in ENTITY_TYPES:
            ntype = "process"
        if not nid or not label:
            continue
        try:
            conf = float(n.get("confidence", 0.75))
        except Exception:
            conf = 0.75
        nodes.append({
            "id": nid,
            "label": label,
            "type": ntype,
            "confidence": max(0.0, min(1.0, conf)),
            "description": str(n.get("description", ""))[:280],
            "refs": _clean_refs(n.get("refs")),
        })
        valid_ids.add(nid)

    for e in data.get("edges", []):
        if not isinstance(e, dict):
            continue
        src = _slug(str(e.get("from", "")))
        dst = _slug(str(e.get("to", "")))
        lbl = str(e.get("label", "")).strip()[:30]
        if src and dst and src in valid_ids and dst in valid_ids and lbl:
            edges.append({"from": src, "to": dst, "label": lbl, "refs": _clean_refs(e.get("refs"))})

    return {"nodes": nodes, "edges": edges}


# ---------- top-level generator --------------------------------------------

ANSWER_SYS = (
    "You are a supply-chain research assistant. Using ONLY the knowledge "
    "graph and sources provided, write a concise natural-language answer "
    "to the user's question (3–6 short paragraphs or a bulleted list). "
    "Cite sources inline as [n] whenever you reference a specific fact. "
    "If the graph lacks enough information to answer fully, say so honestly "
    "and indicate what further research would be needed. Do not invent "
    "entities or facts that are not in the graph."
)


def generate_kg_answer(question: str, nodes: list[dict], edges: list[dict], ref2url: dict) -> str:
    """Summarise the KG into a natural-language answer, with inline citations."""
    if not nodes:
        return ""
    # Compact textual rendering of the KG for the LLM context.
    node_lines = []
    id_to_label: dict[str, str] = {}
    for n in nodes[:80]:
        label = n.get("label", n["id"])
        id_to_label[n["id"]] = label
        desc = n.get("description", "")
        refs = ",".join(str(r) for r in (n.get("refs") or [])[:6])
        node_lines.append(
            f"- {label} ({n.get('type','entity')}, conf {n.get('confidence',0.75):.2f})"
            + (f" [{refs}]" if refs else "")
            + (f" — {desc}" if desc else "")
        )
    edge_lines = []
    for e in edges[:150]:
        src = id_to_label.get(e["from"], e["from"])
        dst = id_to_label.get(e["to"], e["to"])
        refs = ",".join(str(r) for r in (e.get("refs") or [])[:4])
        edge_lines.append(f"- {src} — {e.get('label','related')} → {dst}" + (f" [{refs}]" if refs else ""))
    kg_text = (
        "ENTITIES:\n" + "\n".join(node_lines) +
        ("\n\nRELATIONSHIPS:\n" + "\n".join(edge_lines) if edge_lines else "")
    )[:8000]

    prompt = (
        f"Question: {question}\n\n"
        f"Knowledge graph (extracted from web research):\n{kg_text}\n\n"
        "Answer the question based strictly on this graph. Use inline [n] "
        "citations that refer to source numbers. Be concise and direct."
    )
    try:
        return _call_llm(prompt, system=ANSWER_SYS, max_tokens=900, temperature=0.25).strip()
    except Exception as e:
        logger.warning("answer generation failed: %s", e)
        return ""


def _ref2url(nodes: list[dict], edges: list[dict]) -> dict:
    """Reconstruct ref2url mapping from accumulated _ref2url_store on nodes.

    We keep the mapping separately via module-level state or the passed-in
    state store; here callers are expected to read it from ``nodes[0]["_ref2url"]``
    if they persisted it there, otherwise pass an empty dict and rely on refs.
    """
    # Nodes carry refs list only; ref2url must be passed around separately.
    # This helper is a placeholder for backward compat.
    return {}


def run_research(
    question: str,
    prior_nodes: list[dict] | None = None,
    prior_edges: list[dict] | None = None,
    ref2url: dict[str, dict] | None = None,
) -> Generator[tuple[str, list[dict], list[dict]], None, None]:
    """Yield (status, nodes, edges) after each progress step.

    If ``ref2url`` is supplied, new URLs from this round are appended to it
    in-place so callers can persist citations across rounds.
    """
    nodes: list[dict] = list(prior_nodes or [])
    edges: list[dict] = list(prior_edges or [])
    seen = {n["id"] for n in nodes}
    ref2url = ref2url if ref2url is not None else {}
    # Map url → ref index to deduplicate citations across rounds.
    url_to_ref = {v["url"]: int(k) for k, v in ref2url.items() if isinstance(v, dict) and v.get("url")}
    next_ref = max([int(k) for k in ref2url.keys()], default=0) + 1

    if not _keys_configured():
        yield "keys_missing", nodes, edges
        return

    yield "🔬 Planning research strategy…", nodes, edges
    sub_queries = _decompose(question)

    yield f"🔍 Running {len(sub_queries)} parallel web searches…", nodes, edges
    with ThreadPoolExecutor(max_workers=min(5, len(sub_queries))) as ex:
        futures = {ex.submit(_serper_search, sq, 5): sq for sq in sub_queries}
        all_results: list[tuple[str, list[dict]]] = []
        for fut, sq in futures.items():
            try:
                results = fut.result(timeout=45)
            except Exception as exc:
                logger.warning("search failed for %s: %s", sq, exc)
                results = []
            all_results.append((sq, results))

    for i, (sq, results) in enumerate(all_results, 1):
        if not results:
            yield f"⚠️ No results for sub-query {i}/{len(all_results)}", nodes, edges
            continue

        # Assign a global ref index per search result (dedupe by URL).
        local_refs: list[int] = []  # ref index per result (1-based in snippet text)
        for r in results:
            url = r.get("url") or ""
            if not url:
                local_refs.append(next_ref)
                ref2url[str(next_ref)] = {"url": "", "title": r.get("title", "")}
                next_ref += 1
                continue
            if url in url_to_ref:
                local_refs.append(url_to_ref[url])
            else:
                url_to_ref[url] = next_ref
                ref2url[str(next_ref)] = {"url": url, "title": r.get("title", "")}
                local_refs.append(next_ref)
                next_ref += 1

        status = f"🧠 Extracting entities from sub-query {i}/{len(all_results)}: *{sq[:80]}*"
        yield status, nodes, edges

        # Snippets are passed with their *global* ref indices so the LLM can
        # cite them back in the `refs` field.
        indexed_snippets = [
            {"title": r.get("title", ""), "url": r.get("url", ""), "snippet": r.get("snippet", ""), "ref": local_refs[j]}
            for j, r in enumerate(results)
        ]
        ref_min = min(local_refs) - 1 if local_refs else 0
        new_kg = _extract_kg(question, indexed_snippets, seen, ref_offset=ref_min)

        # Rewrite LLM-cited refs (1..N relative to offset) into global indices.
        def _local_to_global(lst: list[int]) -> list[int]:
            out: list[int] = []
            for n in lst:
                idx = n - ref_min - 1
                if 0 <= idx < len(local_refs):
                    out.append(local_refs[idx])
            return out

        added = 0
        for n in new_kg["nodes"]:
            n["refs"] = _local_to_global(n.get("refs", []))
            if n["id"] not in seen:
                nodes.append(n)
                seen.add(n["id"])
                added += 1
        for e in new_kg["edges"]:
            e["refs"] = _local_to_global(e.get("refs", []))
            edges.append(e)

        yield (
            f"✨ Added {added} nodes · {len(new_kg['edges'])} edges from sub-query {i}/{len(all_results)}",
            nodes,
            edges,
        )

    yield "✅ Research complete", nodes, edges