agent/tests/unit_tests/rag/test_docs_prompt_layer.py

from __future__ import annotations

from app.modules.agent.intent_router_v2 import IntentRouterV2
from app.modules.agent.llm.prompt_loader import PromptLoader
from app.modules.agent.runtime.docs_qa_pipeline import DocsQAPipelineRunner
from app.modules.agent.runtime.docs_qa_pipeline.openapi_postprocessor import OpenAPIPostprocessor
from app.modules.agent.runtime.docs_qa_pipeline.prompt_payload_builder import DocsPromptPayloadBuilder
from app.modules.agent.runtime.steps.generation import RuntimePromptSelector
from tests.docs_qa_eval.fixture_adapter import InMemoryDocsRetrievalAdapter
from tests.unit_tests.rag.intent_router_testkit import repo_context


class FakeLlm:
    def __init__(self, response: str) -> None:
        self.response = response
        self.calls: list[tuple[str, str]] = []

    def generate(self, prompt_name: str, user_input: str, *, log_context: str | None = None) -> str:
        self.calls.append((prompt_name, user_input))
        return self.response


def test_docs_prompt_registry_contains_new_prompts() -> None:
    loader = PromptLoader()

    for name in (
        "docs_explain_answer",
        "docs_general_answer",
        "docs_openapi_answer",
        "docs_openapi_fragment_answer",
    ):
        assert loader.load(name)


def test_prompt_selector_uses_docs_prompts_only() -> None:
    selector = RuntimePromptSelector()

    assert selector.select(intent="DOCUMENTATION_EXPLAIN", sub_intent="COMPONENT_EXPLAIN", answer_mode="normal") == "docs_explain_answer"
    assert selector.select(intent="GENERAL_QA", sub_intent="GENERIC_QA", answer_mode="degraded") == "docs_general_answer"
    assert selector.select(intent="OPENAPI_GENERATION", sub_intent="OPENAPI_METHOD_GENERATE", answer_mode="normal") == "docs_openapi_answer"
    assert selector.select(intent="OPENAPI_GENERATION", sub_intent="OPENAPI_FRAGMENT_GENERATE", answer_mode="normal") == "docs_openapi_fragment_answer"


def test_docs_prompt_payload_contains_required_contract() -> None:
    builder = DocsPromptPayloadBuilder()
    from app.modules.agent.runtime.docs_qa_pipeline.models import DocsEvidenceBundle, OpenAPIResult

    payload = builder.build(
        question="Объясни billing",
        intent="DOCUMENTATION_EXPLAIN",
        sub_intent="COMPONENT_EXPLAIN",
        evidence_bundle=DocsEvidenceBundle(
            intent="DOCUMENTATION_EXPLAIN",
            sub_intent="COMPONENT_EXPLAIN",
            documents=[{"title": "Billing"}],
            facts=[{"content": "Handles payments"}],
            relations=[{"title": "Billing -> Orders"}],
        ),
        api_contract=OpenAPIResult(path="/orders", method="post"),
    )

    assert '"question": "Объясни billing"' in payload
    assert '"intent": "DOCUMENTATION_EXPLAIN"' in payload
    assert '"sub_intent": "COMPONENT_EXPLAIN"' in payload
    assert '"documents"' in payload
    assert '"facts"' in payload
    assert '"relations"' in payload
    assert '"api_contract"' in payload


def test_openapi_postprocessor_requires_paths_for_full_spec() -> None:
    validator = OpenAPIPostprocessor()

    assert validator.validate("paths:\n  /orders:\n    post:\n      responses: {}", require_paths=True)[0] is True
    assert validator.validate("type: object\nproperties:\n  id: {}", require_paths=True)[0] is False
    assert validator.validate("type: object\nproperties:\n  id: {}", require_paths=False)[0] is True


def test_docs_pipeline_uses_llm_prompt_and_validates_openapi() -> None:
    rows = [
        {
            "layer": "D2_FACT_INDEX",
            "path": "docs/api/orders-create.md",
            "title": "POST /orders",
            "content": "Create order",
            "metadata": {
                "endpoint": "/orders",
                "http_method": "post",
                "request_schema": {"type": "object", "properties": {"customer_id": {}}},
                "response_schema": {"type": "object", "properties": {"order_id": {}}},
            },
        }
    ]
    llm = FakeLlm("not yaml")
    runner = DocsQAPipelineRunner(
        IntentRouterV2(),
        InMemoryDocsRetrievalAdapter(rows),
        repo_context=repo_context(),
        llm=llm,
    )

    result = runner.run("Сгенерируй openapi spec для создания заказа", "docs-session")

    assert llm.calls
    assert llm.calls[0][0] == "docs_openapi_answer"
    assert result.prompt_name == "docs_openapi_answer"
    assert result.output_valid is False
    assert "paths:" in result.answer
    assert result.diagnostics.prompt_used == "docs_openapi_answer"
    assert result.diagnostics.llm_mode == "yaml"


def test_pre_llm_mode_skips_llm_calls() -> None:
    rows = [
        {
            "layer": "D2_FACT_INDEX",
            "path": "docs/api/orders-create.md",
            "title": "POST /orders",
            "content": "Create order",
            "metadata": {
                "endpoint": "/orders",
                "http_method": "post",
                "request_schema": {"type": "object", "properties": {"customer_id": {}}},
            },
        }
    ]
    llm = FakeLlm("unused")
    runner = DocsQAPipelineRunner(
        IntentRouterV2(),
        InMemoryDocsRetrievalAdapter(rows),
        repo_context=repo_context(),
        llm=llm,
    )

    result = runner.run("Сгенерируй openapi для /orders", "docs-session", mode="pre_llm_only")

    assert llm.calls == []
    assert result.answer
    assert result.mode == "pre_llm_only"
    assert result.diagnostics.gate_decision == "partial"
    assert result.answer_mode == "ready_partial"
    assert "paths:" in result.answer
    assert result.llm_request["prompt_name"] == "docs_openapi_answer"
    assert "user_prompt" in result.llm_request
    assert "system_prompt" in result.llm_request
    assert result.diagnostics.prompt["prompt_name"] == "docs_openapi_answer"