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