436 lines
17 KiB
Python
436 lines
17 KiB
Python
"""Model context bundle service.
|
||
|
||
统一提供三类上下文：
|
||
1. 知识图谱摘要（KG overview）
|
||
2. 任务模板（task briefs）
|
||
3. 契约同步状态（contract sync）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from datetime import datetime, timezone
|
||
from typing import Any, Optional
|
||
from uuid import uuid4
|
||
|
||
from sqlalchemy import and_, select
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..models import AICallLog
|
||
from .model_routing_service import ModelRoutingService
|
||
from .project_knowledge_graph import ProjectKnowledgeGraphService
|
||
from .task_brief_service import (
|
||
HIGH_CONTRACT_TEMPLATE_ID,
|
||
HIGH_TEMPLATE_ID,
|
||
LOW_TEMPLATE_ID,
|
||
TaskBriefService,
|
||
)
|
||
|
||
CONTEXT_BUNDLE_MARKER = "PMFLOW_CONTEXT_BUNDLE_V1"
|
||
CONTEXT_BUNDLE_KEYS = ("knowledge_graph", "task_template", "contract_sync")
|
||
|
||
TARGET_FUNCTION_CONTEXT: dict[str, dict[str, Any]] = {
|
||
"classify_intent_ai_pmo_kernel": {
|
||
"tier": "low",
|
||
"template_ids": [LOW_TEMPLATE_ID],
|
||
},
|
||
"signal_planner_analyze": {
|
||
"tier": "low",
|
||
"template_ids": [LOW_TEMPLATE_ID],
|
||
},
|
||
"classify_intent_complex": {
|
||
"tier": "high",
|
||
"template_ids": [HIGH_TEMPLATE_ID, HIGH_CONTRACT_TEMPLATE_ID],
|
||
},
|
||
}
|
||
|
||
|
||
class ModelContextBundleService:
|
||
@staticmethod
|
||
def _normalize_function_name(function_name: str) -> str:
|
||
return str(function_name or "").strip().lower()
|
||
|
||
@classmethod
|
||
def target_functions(cls) -> list[str]:
|
||
return list(TARGET_FUNCTION_CONTEXT.keys())
|
||
|
||
@classmethod
|
||
def should_inject_context(cls, function_name: str) -> bool:
|
||
fn = cls._normalize_function_name(function_name)
|
||
return fn in TARGET_FUNCTION_CONTEXT
|
||
|
||
@staticmethod
|
||
def _template_snapshot(template: Optional[dict[str, Any]], *, include_content: bool = True) -> Optional[dict[str, Any]]:
|
||
if not isinstance(template, dict):
|
||
return None
|
||
content = str(template.get("content") or "")
|
||
payload: dict[str, Any] = {
|
||
"id": str(template.get("id") or ""),
|
||
"name": str(template.get("name") or ""),
|
||
"description": str(template.get("description") or ""),
|
||
"updated_at": str(template.get("updated_at") or ""),
|
||
}
|
||
if include_content:
|
||
payload["content"] = content[:2400]
|
||
return payload
|
||
|
||
@staticmethod
|
||
def _kg_snapshot(overview: dict[str, Any]) -> dict[str, Any]:
|
||
latest = overview.get("latest_sync_run") if isinstance(overview.get("latest_sync_run"), dict) else {}
|
||
return {
|
||
"entity_total": int(overview.get("entity_total") or 0),
|
||
"relation_total": int(overview.get("relation_total") or 0),
|
||
"evidence_total": int(overview.get("evidence_total") or 0),
|
||
"pending_events": int(overview.get("pending_events") or 0),
|
||
"latest_sync_run": {
|
||
"id": latest.get("id"),
|
||
"run_type": str(latest.get("run_type") or ""),
|
||
"status": str(latest.get("status") or ""),
|
||
"started_at": latest.get("started_at"),
|
||
"ended_at": latest.get("ended_at"),
|
||
}
|
||
if latest
|
||
else None,
|
||
}
|
||
|
||
@staticmethod
|
||
def _contract_snapshot(contract_sync: dict[str, Any]) -> dict[str, Any]:
|
||
files = contract_sync.get("files") if isinstance(contract_sync.get("files"), list) else []
|
||
normalized_files: list[dict[str, Any]] = []
|
||
missing_files: list[str] = []
|
||
for item in files:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
exists = bool(item.get("exists"))
|
||
file_path = str(item.get("file") or "")
|
||
if file_path and not exists:
|
||
missing_files.append(file_path)
|
||
normalized_files.append(
|
||
{
|
||
"file": file_path,
|
||
"exists": exists,
|
||
"updated_at": item.get("updated_at"),
|
||
}
|
||
)
|
||
return {
|
||
"all_synced": bool(contract_sync.get("all_synced")),
|
||
"missing_files": missing_files,
|
||
"files": normalized_files,
|
||
}
|
||
|
||
@classmethod
|
||
async def build_context_bundle(
|
||
cls,
|
||
db: AsyncSession,
|
||
*,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
) -> Optional[dict[str, Any]]:
|
||
fn = cls._normalize_function_name(function_name)
|
||
target = TARGET_FUNCTION_CONTEXT.get(fn)
|
||
if target is None:
|
||
return None
|
||
|
||
kg_service = ProjectKnowledgeGraphService(db, tenant_id=int(tenant_id))
|
||
overview = await kg_service.get_overview()
|
||
templates = await TaskBriefService.list_templates(db, tenant_id=int(tenant_id))
|
||
template_map = {
|
||
str(item.get("id") or ""): item
|
||
for item in templates
|
||
if isinstance(item, dict) and str(item.get("id") or "").strip()
|
||
}
|
||
configured_template_ids = [str(x or "").strip() for x in target.get("template_ids", []) if str(x or "").strip()]
|
||
primary_template_id = configured_template_ids[0] if configured_template_ids else ""
|
||
additional_ids = configured_template_ids[1:] if len(configured_template_ids) > 1 else []
|
||
missing_template_ids = [tid for tid in configured_template_ids if tid not in template_map]
|
||
contract_sync = await TaskBriefService.get_contract_sync_status(db, tenant_id=int(tenant_id))
|
||
|
||
bundle = {
|
||
"marker": CONTEXT_BUNDLE_MARKER,
|
||
"function_name": fn,
|
||
"tier": str(target.get("tier") or ""),
|
||
"knowledge_graph": cls._kg_snapshot(overview),
|
||
"task_template": {
|
||
"primary": cls._template_snapshot(template_map.get(primary_template_id), include_content=True),
|
||
"additional": [
|
||
snap
|
||
for snap in (
|
||
cls._template_snapshot(template_map.get(template_id), include_content=False)
|
||
for template_id in additional_ids
|
||
)
|
||
if snap is not None
|
||
],
|
||
"template_ids": configured_template_ids,
|
||
"missing_template_ids": missing_template_ids,
|
||
},
|
||
"contract_sync": cls._contract_snapshot(contract_sync),
|
||
}
|
||
return bundle
|
||
|
||
@classmethod
|
||
def render_bundle_marker_text(cls, bundle: dict[str, Any]) -> str:
|
||
payload = dict(bundle) if isinstance(bundle, dict) else {}
|
||
marker = f"[{CONTEXT_BUNDLE_MARKER}]"
|
||
return f"{marker}\n{json.dumps(payload, ensure_ascii=False, sort_keys=True)}"
|
||
|
||
@classmethod
|
||
async def inject_prompt_with_bundle(
|
||
cls,
|
||
db: AsyncSession,
|
||
*,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
prompt: str,
|
||
) -> str:
|
||
raw_prompt = str(prompt or "")
|
||
if not cls.should_inject_context(function_name):
|
||
return raw_prompt
|
||
marker = f"[{CONTEXT_BUNDLE_MARKER}]"
|
||
if marker in raw_prompt:
|
||
return raw_prompt
|
||
bundle = await cls.build_context_bundle(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
)
|
||
if bundle is None:
|
||
return raw_prompt
|
||
bundle_text = cls.render_bundle_marker_text(bundle)
|
||
return f"{bundle_text}\n\n{raw_prompt}".strip()
|
||
|
||
@classmethod
|
||
async def build_capability_summary(cls, db: AsyncSession, *, tenant_id: int) -> dict[str, Any]:
|
||
kg_service = ProjectKnowledgeGraphService(db, tenant_id=int(tenant_id))
|
||
overview = await kg_service.get_overview()
|
||
templates = await TaskBriefService.list_templates(db, tenant_id=int(tenant_id))
|
||
template_map = {
|
||
str(item.get("id") or ""): item
|
||
for item in templates
|
||
if isinstance(item, dict) and str(item.get("id") or "").strip()
|
||
}
|
||
required_template_ids = sorted(
|
||
{
|
||
template_id
|
||
for spec in TARGET_FUNCTION_CONTEXT.values()
|
||
for template_id in spec.get("template_ids", [])
|
||
if str(template_id or "").strip()
|
||
}
|
||
)
|
||
present_template_ids = sorted(template_map.keys())
|
||
missing_template_ids = [template_id for template_id in required_template_ids if template_id not in template_map]
|
||
contract_sync_raw = await TaskBriefService.get_contract_sync_status(db, tenant_id=int(tenant_id))
|
||
contract_sync = cls._contract_snapshot(contract_sync_raw)
|
||
|
||
function_matrix: list[dict[str, Any]] = []
|
||
for function_name in cls.target_functions():
|
||
spec = TARGET_FUNCTION_CONTEXT[function_name]
|
||
routing = await ModelRoutingService.get_effective_routing(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
)
|
||
function_matrix.append(
|
||
{
|
||
"function_name": function_name,
|
||
"tier": str(spec.get("tier") or ""),
|
||
"context_sections": list(CONTEXT_BUNDLE_KEYS),
|
||
"template_ids": [str(tid) for tid in spec.get("template_ids", []) if str(tid or "").strip()],
|
||
"model_routing": {
|
||
"provider": str(routing.get("provider") or ""),
|
||
"model": str(routing.get("model") or ""),
|
||
"tier": str(routing.get("tier") or ""),
|
||
"source": str(routing.get("source") or ""),
|
||
"api_env": str(routing.get("api_env") or ""),
|
||
},
|
||
}
|
||
)
|
||
|
||
required_templates = [
|
||
cls._template_snapshot(template_map.get(template_id), include_content=False) or {"id": template_id}
|
||
for template_id in required_template_ids
|
||
]
|
||
|
||
return {
|
||
"marker": CONTEXT_BUNDLE_MARKER,
|
||
"knowledge_graph": cls._kg_snapshot(overview),
|
||
"task_templates": {
|
||
"total": len(templates),
|
||
"required_template_ids": required_template_ids,
|
||
"present_template_ids": present_template_ids,
|
||
"missing_template_ids": missing_template_ids,
|
||
"required_templates": required_templates,
|
||
},
|
||
"contract_sync": contract_sync,
|
||
"function_model_template_matrix": function_matrix,
|
||
"coverage": {
|
||
"template_ready": len(missing_template_ids) == 0,
|
||
"contract_ready": bool(contract_sync.get("all_synced")),
|
||
"target_function_count": len(function_matrix),
|
||
},
|
||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||
}
|
||
|
||
@classmethod
|
||
def _canary_prompt_for_function(cls, *, function_name: str, probe_id: str) -> str:
|
||
fn = cls._normalize_function_name(function_name)
|
||
if fn in {"classify_intent_ai_pmo_kernel", "classify_intent_complex"}:
|
||
return (
|
||
f"这是 PMFlow 上下文注入探针，probe_id={probe_id}。"
|
||
"请只返回 JSON："
|
||
'{"intent":"query.read_only","confidence":0.91,"reasoning":"probe_context_ok"}'
|
||
)
|
||
if fn == "signal_planner_analyze":
|
||
return (
|
||
f"这是 PMFlow 上下文注入探针，probe_id={probe_id}。"
|
||
"请只返回 JSON："
|
||
'{"summary":"probe","priority":"medium","execution_plan":[],"risks":[],"confidence":0.77}'
|
||
)
|
||
return f"probe_id={probe_id}"
|
||
|
||
@classmethod
|
||
async def _trigger_canary_call(
|
||
cls,
|
||
db: AsyncSession,
|
||
*,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
probe_id: str,
|
||
) -> dict[str, Any]:
|
||
prompt = cls._canary_prompt_for_function(function_name=function_name, probe_id=probe_id)
|
||
try:
|
||
from .llm import llm_service # 延迟导入，避免循环依赖
|
||
|
||
result = await llm_service.generate_structured_object(
|
||
prompt,
|
||
db=db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
use_prompt_profile=False,
|
||
temperature=0.0,
|
||
max_tokens=220,
|
||
response_mode="fast",
|
||
)
|
||
return {
|
||
"called": True,
|
||
"error": None,
|
||
"error_code": str(result.error_code or ""),
|
||
"status": str(result.error_code or "success"),
|
||
"provider": str(result.provider or ""),
|
||
"model": str(result.model or ""),
|
||
}
|
||
except Exception as exc:
|
||
return {
|
||
"called": False,
|
||
"error": str(exc)[:500],
|
||
"error_code": "canary_call_failed",
|
||
"status": "failed",
|
||
"provider": "",
|
||
"model": "",
|
||
}
|
||
|
||
@classmethod
|
||
async def _find_prompt_evidence(
|
||
cls,
|
||
db: AsyncSession,
|
||
*,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
probe_id: Optional[str],
|
||
) -> Optional[AICallLog]:
|
||
stmt = select(AICallLog).where(
|
||
and_(
|
||
AICallLog.tenant_id == int(tenant_id),
|
||
AICallLog.function_name == str(function_name),
|
||
AICallLog.prompt.like(f"%[{CONTEXT_BUNDLE_MARKER}]%"),
|
||
)
|
||
)
|
||
if probe_id:
|
||
stmt = stmt.where(AICallLog.prompt.like(f"%{probe_id}%"))
|
||
stmt = stmt.order_by(AICallLog.id.desc()).limit(1)
|
||
return (await db.execute(stmt)).scalar_one_or_none()
|
||
|
||
@classmethod
|
||
async def verify_context_consumption(
|
||
cls,
|
||
db: AsyncSession,
|
||
*,
|
||
tenant_id: int,
|
||
run_canary: bool = True,
|
||
) -> dict[str, Any]:
|
||
probe_id = uuid4().hex[:12] if run_canary else ""
|
||
function_results: list[dict[str, Any]] = []
|
||
|
||
for function_name in cls.target_functions():
|
||
target = TARGET_FUNCTION_CONTEXT.get(function_name) or {}
|
||
canary = {"called": False, "status": "skipped", "error": None, "error_code": "", "provider": "", "model": ""}
|
||
if run_canary:
|
||
canary = await cls._trigger_canary_call(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
probe_id=probe_id,
|
||
)
|
||
|
||
evidence = await cls._find_prompt_evidence(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
probe_id=probe_id if run_canary else None,
|
||
)
|
||
prompt_text = str(evidence.prompt or "") if evidence is not None else ""
|
||
checks = {
|
||
"marker": f"[{CONTEXT_BUNDLE_MARKER}]" in prompt_text,
|
||
"knowledge_graph": '"knowledge_graph"' in prompt_text,
|
||
"task_template": '"task_template"' in prompt_text,
|
||
"contract_sync": '"contract_sync"' in prompt_text,
|
||
}
|
||
route = await ModelRoutingService.get_effective_routing(
|
||
db,
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
)
|
||
function_results.append(
|
||
{
|
||
"function_name": function_name,
|
||
"tier": str(target.get("tier") or ""),
|
||
"template_ids": [str(x) for x in target.get("template_ids", []) if str(x or "").strip()],
|
||
"routing_snapshot": {
|
||
"provider": str(route.get("provider") or ""),
|
||
"model": str(route.get("model") or ""),
|
||
"tier": str(route.get("tier") or ""),
|
||
"source": str(route.get("source") or ""),
|
||
},
|
||
"canary": canary,
|
||
"checks": checks,
|
||
"pass": all(checks.values()),
|
||
"log_evidence": (
|
||
{
|
||
"id": int(evidence.id),
|
||
"status": str(evidence.status or ""),
|
||
"created_at": evidence.created_at.isoformat() if evidence.created_at else None,
|
||
"prompt_excerpt": prompt_text[:800],
|
||
}
|
||
if evidence is not None
|
||
else None
|
||
),
|
||
}
|
||
)
|
||
|
||
tier_result: dict[str, dict[str, Any]] = {"low": {"pass": False, "functions": []}, "high": {"pass": False, "functions": []}}
|
||
for tier_name in ("low", "high"):
|
||
tier_functions = [row for row in function_results if str(row.get("tier") or "") == tier_name]
|
||
tier_result[tier_name] = {
|
||
"pass": bool(tier_functions) and all(bool(row.get("pass")) for row in tier_functions),
|
||
"functions": tier_functions,
|
||
}
|
||
|
||
overall_pass = bool(tier_result["low"]["pass"]) and bool(tier_result["high"]["pass"])
|
||
return {
|
||
"marker": CONTEXT_BUNDLE_MARKER,
|
||
"run_canary": bool(run_canary),
|
||
"probe_id": probe_id or None,
|
||
"overall_pass": overall_pass,
|
||
"tiers": tier_result,
|
||
"functions": function_results,
|
||
}
|