PMFlow/backend/app/services/external_query_service/planner.py

from __future__ import annotations

import json
import re
from typing import Any, Callable

from ...config.ai_prompts import PROMPT_EXTERNAL_QUERY_PLAN
from .schemas import ExternalQueryPlan


class ExternalQueryPlanner:
    FUNCTION_NAME = "external_query_plan"

    def __init__(
        self,
        *,
        llm_service: Any,
        db: Any,
        tenant_id: int,
        query_builder: Callable[[str], list[str]] | None = None,
    ) -> None:
        self.llm_service = llm_service
        self.db = db
        self.tenant_id = int(tenant_id or 1)
        self.query_builder = query_builder

    async def build_plan(
        self,
        *,
        question: str,
        query_semantics: dict[str, Any] | None = None,
    ) -> tuple[ExternalQueryPlan, dict[str, Any]]:
        fallback_plan = self._build_fallback_plan(
            question=question,
            query_semantics=query_semantics,
        )
        generator = getattr(self.llm_service, "generate_structured_object", None)
        if not callable(generator):
            return fallback_plan, {"planner": "heuristic", "fallback_taken": True}

        prompt = PROMPT_EXTERNAL_QUERY_PLAN.format(
            question=str(question or "").strip()[:1200],
            query_semantics=json.dumps(query_semantics or {}, ensure_ascii=False)[:4000],
            fallback_queries=json.dumps(fallback_plan.search_queries, ensure_ascii=False),
        ).strip()
        try:
            structured = await generator(
                prompt=prompt,
                db=self.db,
                tenant_id=self.tenant_id,
                function_name=self.FUNCTION_NAME,
                use_prompt_profile=False,
                temperature=0.0,
                response_mode="fast",
            )
            payload = getattr(structured, "payload", None)
            if payload is None and hasattr(structured, "model_dump"):
                payload = structured.model_dump()
            if not isinstance(payload, dict):
                raise ValueError("external_query_plan_payload_not_dict")
            merged = fallback_plan.model_dump()
            for key, value in payload.items():
                if value in (None, "", [], {}):
                    continue
                merged[key] = value
            plan = ExternalQueryPlan.model_validate(merged)
            if not plan.search_queries:
                plan.search_queries = list(fallback_plan.search_queries)
            if not plan.preferred_domains:
                plan.preferred_domains = list(fallback_plan.preferred_domains)
            if not plan.topic:
                plan.topic = fallback_plan.topic
            return plan, {"planner": "llm_structured", "fallback_taken": False}
        except Exception as exc:
            return fallback_plan, {
                "planner": "heuristic",
                "fallback_taken": True,
                "reason": str(exc)[:200],
            }

    def _build_fallback_plan(
        self,
        *,
        question: str,
        query_semantics: dict[str, Any] | None = None,
    ) -> ExternalQueryPlan:
        text = re.sub(r"\s+", " ", str(question or "")).strip()
        semantics = query_semantics if isinstance(query_semantics, dict) else {}
        query_profile = semantics.get("query_profile") if isinstance(semantics.get("query_profile"), dict) else {}
        object_mentions = [
            str(item or "").strip()
            for item in list(semantics.get("object_mentions") or [])
            if str(item or "").strip()
        ]

        query_type = str(semantics.get("query_type") or "").strip().lower()
        if query_type == "compare":
            answer_contract = "compare_with_sources"
        elif bool(query_profile.get("needs_freshness")):
            answer_contract = "summary_with_sources"
        else:
            answer_contract = "direct_fact"

        topic = text[:120] or "外部查询"
        search_queries = self._build_search_queries(text)
        preferred_domains = self._infer_preferred_domains(text)
        subject_entities = object_mentions[:6]
        location_hint = None
        time_range_hint = self._infer_time_range_hint(text)
        requires_multi_source = answer_contract == "compare_with_sources"
        clarify_on_ambiguity = False
        result_limit = 3 if requires_multi_source or bool(query_profile.get("needs_freshness")) else 2

        return ExternalQueryPlan(
            plan_version="v1",
            query_kind="generic_web" if text else "unknown",
            answer_contract=answer_contract,
            freshness_required=bool(query_profile.get("needs_freshness")),
            topic=topic,
            subject_entities=subject_entities,
            location_hint=location_hint,
            time_range_hint=time_range_hint,
            preferred_domains=preferred_domains,
            search_queries=search_queries,
            result_limit=result_limit,
            requires_multi_source=requires_multi_source,
            clarify_on_ambiguity=clarify_on_ambiguity,
        )

    def _build_search_queries(self, question: str) -> list[str]:
        if callable(self.query_builder):
            try:
                queries = self.query_builder(question)
                normalized = [str(item or "").strip() for item in list(queries or []) if str(item or "").strip()]
                if normalized:
                    return normalized[:5]
            except Exception:
                pass
        text = re.sub(r"\s+", " ", str(question or "")).strip()
        return [text] if text else []

    @staticmethod
    def _infer_preferred_domains(question: str) -> list[str]:
        normalized = str(question or "").lower()
        domains: list[str] = []
        if any(token in normalized for token in ("codex", "gpt", "chatgpt", "openai")):
            domains.extend(["openai.com", "developers.openai.com"])
        if any(token in normalized for token in ("copilot", "github")):
            domains.extend(["github.com", "docs.github.com"])
        uniq: list[str] = []
        for domain in domains:
            if domain not in uniq:
                uniq.append(domain)
        return uniq[:4]

    @staticmethod
    def _infer_time_range_hint(question: str) -> str | None:
        text = str(question or "")
        for pattern in (r"未来\d+天", r"未来一周", r"今天", r"截至今天", r"截至现在", r"目前", r"最新"):
            match = re.search(pattern, text)
            if match:
                return str(match.group(0) or "").strip() or None
        return None