PMFlow/backend/app/services/model/model_context_bundle/mixin_build.py

"""Bundle build and capability summary methods."""

from __future__ import annotations

from .shared import (
    Any,
    AsyncSession,
    CONTEXT_BUNDLE_KEYS,
    CONTEXT_BUNDLE_MARKER,
    ModelRoutingService,
    TARGET_FUNCTION_CONTEXT,
    TaskBriefService,
    datetime,
    timezone,
)
from ...project_knowledge_graph_service import ProjectKnowledgeGraphService


class ModelContextBundleBuildMixin:
    """Build context bundles and capability summaries."""

    @classmethod
    async def build_context_bundle(
        cls,
        db: AsyncSession,
        *,
        tenant_id: int,
        function_name: str,
    ) -> dict[str, Any] | None:
        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
    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(),
        }