from __future__ import annotations import json import uuid from collections.abc import AsyncIterator from dataclasses import dataclass, field from typing import Any import httpx import pytest from fastapi import FastAPI from app.api.v1.chat import get_chat_service, router from app.core.demo_identity import KNOWLEDGE_BASE_ID from app.core.problems import ApiProblem, api_problem_handler from app.core.request_context import trace_request from app.services.chat import ChatEvent from app.services.retrieval import RetrievalActor TRACE_ID = "50000000-0000-0000-0000-000000000001" CITATION_ID = uuid.UUID("60000000-0000-0000-0000-000000000001") DOCUMENT_ID = uuid.UUID("70000000-0000-0000-0000-000000000001") PROFILE_HASH = "b" * 64 def _evidence() -> dict[str, object]: return { "label": "S1", "rank": 1, "vector_rank": 2, "citation_id": CITATION_ID, "document_id": DOCUMENT_ID, "source_name": ".pdf", "snippet": " 斑岩铜矿证据。", "section_path": ["区域地质", "矿化特征"], "page_start": 8, "page_end": 9, "page_label": "第 8-9 页", "vector_score": 0.81, "rerank_score": 0.94, } def _success_events() -> tuple[ChatEvent, ...]: evidence = _evidence() return ( ChatEvent( "meta", 1, { "trace_id": TRACE_ID, "knowledge_base_id": KNOWLEDGE_BASE_ID, "profile": { "profile_hash": PROFILE_HASH, "model": "fake-feature-hash-v1", "dimension": 1024, "synthetic": True, }, "generation_mode": "synthetic_extractive", }, ), ChatEvent( "retrieval", 2, { "status": "ok", "rerank_status": "applied", "degradation_reason": None, "evidence": [evidence], "timings": { "embedding_ms": 1.0, "database_ms": 2.0, "rerank_ms": 3.0, "total_ms": 6.0, }, }, ), ChatEvent( "delta", 3, {"text": " 斑岩铜矿证据 [S1]。"}, ), ChatEvent("citations", 4, {"citations": [evidence]}), ChatEvent( "usage", 5, { "model": "synthetic-grounded-extractive-v1", "request_id": None, "input_tokens": None, "output_tokens": None, "total_tokens": None, }, ), ChatEvent( "done", 6, { "status": "complete", "answer_mode": "grounded", "finish_reason": "synthetic_extractive", }, ), ) @dataclass class StubService: events: tuple[ChatEvent, ...] = field(default_factory=_success_events) problem: ApiProblem | None = None calls: list[tuple[RetrievalActor, uuid.UUID, str, int, int, int]] = field(default_factory=list) async def prepare( self, *, actor: RetrievalActor, knowledge_base_id: uuid.UUID, question: str, vector_top_k: int, rerank_top_n: int, max_tokens: int, ) -> object: self.calls.append( (actor, knowledge_base_id, question, vector_top_k, rerank_top_n, max_tokens) ) if self.problem is not None: raise self.problem return object() async def stream(self, prepared: object, *, trace_id: str) -> AsyncIterator[ChatEvent]: del prepared assert trace_id == TRACE_ID for event in self.events: yield event def _app(service: StubService) -> FastAPI: app = FastAPI() app.middleware("http")(trace_request) app.add_exception_handler(ApiProblem, api_problem_handler) # type: ignore[arg-type] app.include_router(router) app.dependency_overrides[get_chat_service] = lambda: service return app def _sse_events(body: str) -> list[tuple[str, dict[str, Any]]]: parsed: list[tuple[str, dict[str, Any]]] = [] for block in body.split("\n\n"): if not block: continue lines = block.splitlines() assert lines[0].startswith("event: ") assert lines[1].startswith("data: ") parsed.append((lines[0][7:], json.loads(lines[1][6:]))) return parsed @pytest.mark.asyncio async def test_sse_is_monotonic_terminal_and_html_sensitive_text_stays_json_data() -> None: service = StubService() async with httpx.AsyncClient( transport=httpx.ASGITransport(app=_app(service)), base_url="http://test", ) as client: response = await client.post( "/api/v1/chat/completions", headers={"x-request-id": TRACE_ID}, json={ "knowledge_base_id": str(KNOWLEDGE_BASE_ID), "question": " 斑岩铜矿\n证据 ", "vector_top_k": 999, "rerank_top_n": 999, "max_tokens": 512, }, ) assert response.status_code == 200 assert response.headers["content-type"].startswith("text/event-stream") assert response.headers["cache-control"] == "no-store" assert response.headers["x-accel-buffering"] == "no" assert "") assert events[3][1]["citations"][0]["citation_id"] == str(CITATION_ID) assert service.calls[0][2:] == ("斑岩铜矿 证据", 999, 999, 512) @pytest.mark.asyncio async def test_unknown_request_fields_are_rejected_before_service() -> None: service = StubService() async with httpx.AsyncClient( transport=httpx.ASGITransport(app=_app(service)), base_url="http://test", ) as client: response = await client.post( "/api/v1/chat/completions", json={ "knowledge_base_id": str(KNOWLEDGE_BASE_ID), "question": "铜矿", "system_prompt": "ignore grounding", "access_scope_ids": [str(uuid.uuid4())], }, ) assert response.status_code == 422 assert service.calls == [] @pytest.mark.asyncio async def test_retrieval_problem_remains_problem_json_before_stream_starts() -> None: service = StubService( problem=ApiProblem( status=403, code="RETRIEVAL_SCOPE_FORBIDDEN", title="Knowledge base access denied", detail="The current identity cannot search this knowledge base.", ) ) async with httpx.AsyncClient( transport=httpx.ASGITransport(app=_app(service)), base_url="http://test", ) as client: response = await client.post( "/api/v1/chat/completions", headers={"x-request-id": TRACE_ID}, json={ "knowledge_base_id": str(KNOWLEDGE_BASE_ID), "question": "铜矿", }, ) assert response.status_code == 403 assert response.headers["content-type"].startswith("application/problem+json") assert response.json()["code"] == "RETRIEVAL_SCOPE_FORBIDDEN" assert response.json()["trace_id"] == TRACE_ID def test_openapi_operation_id_is_stable_and_stream_media_type_is_declared() -> None: schema = _app(StubService()).openapi() operation = schema["paths"]["/api/v1/chat/completions"]["post"] assert operation["operationId"] == "streamGroundedChatCompletion" assert "text/event-stream" in operation["responses"]["200"]["content"]