ai-contract-platform/backend/app/api/chat.py

import json
from typing import Any, Dict, Iterable, List, Mapping, Optional

from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.api.reviews import require_current_user
from app.clients.shared_ai_platform import SharedAIPlatformClient, get_shared_ai_platform_client
from app.db import get_db_session
from app.schemas.chat import ChatStreamRequest
from app.services.legal_article_service import LegalArticleService
from app.services.prompt_runtime_service import PromptRuntimeService
from app.services.review_service import ReviewService
from app.services.tenant_service import TenantService

router = APIRouter(prefix="/chat", tags=["chat"])
ROUTE_PERMISSION_SCOPE = "chat"
ROUTE_PERMISSION_ACTIONS = {
    "stream": "chat:stream",
}


def get_chat_review_service(db: Session = Depends(get_db_session)) -> ReviewService:
    return ReviewService(db)


def get_chat_legal_article_service(db: Session = Depends(get_db_session)) -> LegalArticleService:
    return LegalArticleService(db)


def get_chat_prompt_runtime_service(db: Session = Depends(get_db_session)) -> PromptRuntimeService:
    return PromptRuntimeService(db)


def get_tenant_service(db: Session = Depends(get_db_session)) -> TenantService:
    return TenantService(db)


@router.post("/stream")
def stream_chat(
    payload: ChatStreamRequest,
    shared_ai_platform_client: SharedAIPlatformClient = Depends(get_shared_ai_platform_client),
    tenant_service: TenantService = Depends(get_tenant_service),
    current_user: dict = Depends(require_current_user),
    review_service: Optional[ReviewService] = Depends(get_chat_review_service),
    legal_article_service: Optional[LegalArticleService] = Depends(get_chat_legal_article_service),
    prompt_runtime_service: Optional[PromptRuntimeService] = Depends(get_chat_prompt_runtime_service),
):
    tenant_service.assert_billing_allows(current_user, "chat")
    review_context = (
        payload.review_context.model_dump(mode="json", by_alias=True, exclude_none=True)
        if payload.review_context
        else None
    )
    review_context = _enrich_review_context_with_legal_articles(current_user, review_context, review_service)
    runtime_legal_sources = _search_runtime_legal_article_sources(
        current_user,
        payload.message,
        legal_article_service,
    )
    review_context = _enrich_review_context_with_runtime_legal_sources(review_context, runtime_legal_sources)
    prompt_runtime = _resolve_chat_prompt_runtime(current_user, review_context, prompt_runtime_service)
    review_context = _enrich_review_context_with_prompt_runtime(review_context, prompt_runtime, prompt_runtime_service)
    template_key = prompt_runtime.get("template_key") if prompt_runtime and prompt_runtime.get("available") else None
    if prompt_runtime_service is not None:
        prompt_runtime_service.record_run_link(
            current_user,
            prompt_runtime,
            run_type="legal_chat",
            run_id=payload.conversation_id or "new_conversation",
            resource_type="conversation",
            resource_id=payload.conversation_id,
            trace={
                "message_length": len(payload.message or ""),
                "has_review_context": bool(review_context),
                "review_id": (review_context or {}).get("reviewId") or (review_context or {}).get("review_id"),
                "template_key": template_key,
            },
        )
    local_legal_sources = _dedupe_sources([*_build_legal_article_sources(review_context), *runtime_legal_sources])
    return StreamingResponse(
        _merge_local_sources_into_stream(
            shared_ai_platform_client.stream_answer(
                current_user,
                payload.conversation_id,
                payload.message,
                review_context,
                template_key=template_key,
            ),
            local_legal_sources,
        ),
        media_type="text/event-stream",
    )


def _resolve_chat_prompt_runtime(
    current_user: Dict[str, Any],
    review_context: Optional[Dict[str, Any]],
    prompt_runtime_service: Optional[PromptRuntimeService],
) -> Dict[str, Any]:
    if prompt_runtime_service is None:
        return {"available": False, "prompt_type": "legal_chat", "prompt_version": "legal_chat.shared_ai_stream.v1", "source": "local_default"}
    has_review_context = bool(review_context and ((review_context.get("reviewId") or review_context.get("review_id"))))
    return prompt_runtime_service.resolve_published_prompt(
        current_user,
        "legal_chat",
        preferred_keys=[
            "legal_chat_review" if has_review_context else None,
            "legal_chat_default",
        ],
        fallback_prompt_version="legal_chat.shared_ai_stream.v1",
    )


def _enrich_review_context_with_prompt_runtime(
    review_context: Optional[Dict[str, Any]],
    prompt_runtime: Optional[Dict[str, Any]],
    prompt_runtime_service: Optional[PromptRuntimeService],
) -> Optional[Dict[str, Any]]:
    if prompt_runtime_service is None or not prompt_runtime:
        return review_context
    enriched = dict(review_context or {})
    prompt_trace = prompt_runtime_service.to_trace(prompt_runtime)
    enriched["promptRuntime"] = prompt_trace
    enriched["prompt_runtime"] = prompt_trace
    metadata = dict(enriched.get("metadata") or {})
    metadata["promptRuntime"] = prompt_trace
    metadata["promptVersion"] = prompt_trace.get("prompt_version")
    metadata["promptKey"] = prompt_trace.get("prompt_key")
    enriched["metadata"] = metadata
    return enriched


def _enrich_review_context_with_legal_articles(
    current_user: Dict[str, Any],
    review_context: Optional[Dict[str, Any]],
    review_service: Optional[ReviewService],
) -> Optional[Dict[str, Any]]:
    if not review_context or review_service is None:
        return review_context

    review_id = review_context.get("reviewId") or review_context.get("review_id")
    if not review_id:
        return review_context

    try:
        result = review_service.get_review_result(current_user, str(review_id))
    except Exception:
        return review_context
    if not result:
        return review_context

    refs = _collect_legal_article_refs(result.get("risk_items") or [])
    if not refs:
        return review_context

    enriched = dict(review_context)
    enriched["legalArticleRefs"] = refs
    enriched["legal_article_refs"] = refs
    metadata = dict(enriched.get("metadata") or {})
    metadata["legalArticleRefCount"] = len(refs)
    metadata["legalArticleSources"] = [
        {
            "id": ref.get("id"),
            "title": _build_legal_article_source_title(ref),
            "detail_path": ref.get("detail_path"),
        }
        for ref in refs
    ]
    enriched["metadata"] = metadata
    return enriched


def _search_runtime_legal_article_sources(
    current_user: Dict[str, Any],
    message: str,
    legal_article_service: Optional[LegalArticleService],
) -> List[Dict[str, Any]]:
    if legal_article_service is None:
        return []
    try:
        return legal_article_service.search_runtime_sources(current_user, message, limit=5)
    except Exception:
        return []


def _enrich_review_context_with_runtime_legal_sources(
    review_context: Optional[Dict[str, Any]],
    runtime_sources: List[Dict[str, Any]],
) -> Optional[Dict[str, Any]]:
    if not runtime_sources:
        return review_context
    enriched = dict(review_context or {})
    enriched["legalArticleSearchSources"] = runtime_sources
    enriched["legal_article_search_sources"] = runtime_sources
    metadata = dict(enriched.get("metadata") or {})
    metadata["legalArticleSearchSourceCount"] = len(runtime_sources)
    metadata["legalArticleSearchSources"] = [
        {
            "id": source.get("source_id"),
            "title": source.get("title"),
            "detail_path": source.get("detail_path"),
        }
        for source in runtime_sources
    ]
    enriched["metadata"] = metadata
    return enriched


def _collect_legal_article_refs(risk_items: Iterable[Mapping[str, Any]]) -> List[Dict[str, Any]]:
    refs: List[Dict[str, Any]] = []
    seen: set[str] = set()
    for item in risk_items:
        for ref in item.get("legal_article_refs") or []:
            if not isinstance(ref, Mapping):
                continue
            ref_id = str(ref.get("id") or "")
            key = ref_id or f"{ref.get('source_title')}::{ref.get('article_no')}::{ref.get('detail_path')}"
            if not key or key in seen:
                continue
            seen.add(key)
            refs.append(dict(ref))
    return refs


def _build_legal_article_sources(review_context: Optional[Mapping[str, Any]]) -> List[Dict[str, Any]]:
    if not review_context:
        return []
    refs = review_context.get("legalArticleRefs") or review_context.get("legal_article_refs") or []
    if not isinstance(refs, list):
        return []
    sources: List[Dict[str, Any]] = []
    for ref in refs:
        if not isinstance(ref, Mapping):
            continue
        source_id = ref.get("id") or ref.get("source_id")
        title = _build_legal_article_source_title(ref)
        content = ref.get("article_excerpt") or ref.get("match_reason") or ref.get("detail_path") or ""
        sources.append(
            {
                "title": title,
                "content": content,
                "score": 1.0,
                "source_type": "legal_article",
                "source_id": source_id,
                "detail_path": ref.get("detail_path"),
                "article_no": ref.get("article_no"),
                "authority": ref.get("authority"),
                "knowledge_base_id": ref.get("knowledge_base_id"),
                "knowledge_document_id": ref.get("knowledge_document_id"),
                "metadata": {
                    "source_id": ref.get("source_id"),
                    "source_title": ref.get("source_title"),
                    "match_reason": ref.get("match_reason"),
                    "tags": ref.get("tags") or [],
                },
            }
        )
    return _dedupe_sources(sources)


def _build_legal_article_source_title(ref: Mapping[str, Any]) -> str:
    source_title = str(ref.get("source_title") or "法条全文")
    article_no = str(ref.get("article_no") or "")
    return f"{source_title} {article_no}".strip()


def _merge_local_sources_into_stream(upstream: Iterable[str], local_sources: List[Dict[str, Any]]) -> Iterable[str]:
    if not local_sources:
        yield from upstream
        return

    emitted_sources = False
    for chunk in upstream:
        event = _parse_sse_event(chunk)
        if event and event.get("type") == "sources":
            merged = _dedupe_sources([*(event.get("sources") or []), *local_sources])
            emitted_sources = True
            yield _format_sse_event({"type": "sources", "sources": merged})
            continue
        if event and event.get("type") == "end" and not emitted_sources:
            emitted_sources = True
            yield _format_sse_event({"type": "sources", "sources": local_sources})
        yield chunk


def _parse_sse_event(chunk: str) -> Optional[Dict[str, Any]]:
    for line in chunk.splitlines():
        if not line.startswith("data: "):
            continue
        try:
            payload = json.loads(line[6:])
        except json.JSONDecodeError:
            return None
        return payload if isinstance(payload, dict) else None
    return None


def _format_sse_event(payload: Dict[str, Any]) -> str:
    return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"


def _dedupe_sources(sources: Iterable[Mapping[str, Any]]) -> List[Dict[str, Any]]:
    deduped: List[Dict[str, Any]] = []
    seen: set[str] = set()
    for source in sources:
        if not isinstance(source, Mapping):
            continue
        key = str(source.get("source_id") or source.get("detail_path") or f"{source.get('title')}::{source.get('content')}")
        if key in seen:
            continue
        seen.add(key)
        deduped.append(dict(source))
    return deduped