From c75f0b646a049dbdd13975dfb4abfae40d332355 Mon Sep 17 00:00:00 2001 From: yuqei Date: Tue, 23 Jun 2026 00:09:48 +0800 Subject: [PATCH] feat: enrich ai run audit traces --- ...-platform-embedding-pgvector-2026-06-22.md | 7 ++ .../packages/ai-openapi/openapi.yaml | 21 ++++ .../services/ai-platform-api/README.md | 4 + .../20260622_0005_add_ai_run_trace_fields.py | 39 +++++++ .../src/yuqei_ai_platform_api/config.py | 2 + .../src/yuqei_ai_platform_api/main.py | 8 ++ .../src/yuqei_ai_platform_api/models.py | 8 ++ .../src/yuqei_ai_platform_api/providers.py | 104 +++++++++++++++++- .../src/yuqei_ai_platform_api/repository.py | 24 ++++ .../tests/test_platform_data.py | 14 +++ .../tests/test_provider_flow.py | 42 ++++++- .../src/yuqei_sdk/ai_platform.py | 8 ++ yuqei-ai-sdk-python/tests/test_ai_platform.py | 16 +++ 13 files changed, 287 insertions(+), 10 deletions(-) create mode 100644 yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0005_add_ai_run_trace_fields.py diff --git a/project-docs/V2-ai-platform-embedding-pgvector-2026-06-22.md b/project-docs/V2-ai-platform-embedding-pgvector-2026-06-22.md index 2dff35ff..19791aa9 100644 --- a/project-docs/V2-ai-platform-embedding-pgvector-2026-06-22.md +++ b/project-docs/V2-ai-platform-embedding-pgvector-2026-06-22.md @@ -51,6 +51,12 @@ - Python SDK 新增 `explain_knowledge_search(...)`。 - JSON fallback 的向量分改为 cosine 口径，便于和 pgvector cosine 距离保持一致。 +8. AI run trace 审计增强。 + - `AiRunAudit` 新增 provider、model、provider_status、provider_error。 + - 记录 prompt_tokens、completion_tokens、total_tokens。 + - 支持按 `AI_PLATFORM_PROVIDER_PROMPT_TOKEN_COST_PER_1K` 和 `AI_PLATFORM_PROVIDER_COMPLETION_TOKEN_COST_PER_1K` 估算费用。 + - OpenAPI 和 Python SDK 同步新增审计字段。 + ## 生产配置建议 ```powershell @@ -81,3 +87,4 @@ python -m alembic upgrade head - pgvector 可用时：优先使用数据库向量距离排序。 - pgvector 不可用时：回退 JSON embedding，不影响法律问答主链路。 - explain 接口可说明命中依据和过滤原因，用于法律问答可信度展示和检索调参。 +- AI run 审计可用于排查模型错误、统计 token、估算成本和后续质量评测。 diff --git a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml index ff5ed664..9b3708c7 100644 --- a/yuqei-ai-platform/packages/ai-openapi/openapi.yaml +++ b/yuqei-ai-platform/packages/ai-openapi/openapi.yaml @@ -1166,6 +1166,18 @@ components: prompt_template_id: type: string nullable: true + provider: + type: string + nullable: true + model: + type: string + nullable: true + provider_status: + type: string + nullable: true + provider_error: + type: string + nullable: true question: type: string nullable: true @@ -1188,6 +1200,15 @@ components: type: array items: $ref: "#/components/schemas/RetrievalSource" + prompt_tokens: + type: integer + completion_tokens: + type: integer + total_tokens: + type: integer + estimated_cost: + type: number + format: float latency_ms: type: integer status: diff --git a/yuqei-ai-platform/services/ai-platform-api/README.md b/yuqei-ai-platform/services/ai-platform-api/README.md index 1be64740..152cfe0a 100644 --- a/yuqei-ai-platform/services/ai-platform-api/README.md +++ b/yuqei-ai-platform/services/ai-platform-api/README.md @@ -32,6 +32,8 @@ $env:AI_PLATFORM_PROVIDER_CALL_MODE="enabled" $env:AI_PLATFORM_PROVIDER_API_KEY="your-api-key" $env:AI_PLATFORM_PROVIDER_BASE_URL="https://api.deepseek.com/v1" $env:AI_PLATFORM_PROVIDER_FALLBACK_ENABLED="true" +$env:AI_PLATFORM_PROVIDER_PROMPT_TOKEN_COST_PER_1K="0" +$env:AI_PLATFORM_PROVIDER_COMPLETION_TOKEN_COST_PER_1K="0" ``` Embedding defaults to deterministic `local-hash-v1` so tests and offline @@ -78,6 +80,8 @@ Core MVP endpoints: Knowledge chunk search returns keyword, vector, and hybrid scores. Legal QA uses hybrid retrieval by default and records retrieved chunk ids plus structured source metadata in AI run audits. +AI run audits also record provider/model, provider status/error, token usage, +estimated cost, and latency for operational traceability. Use `GET /api/v1/knowledge-chunks/search/explain` to inspect normalized query terms, keyword/vector/hybrid scores, inclusion rank, and filter reasons for the diff --git a/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0005_add_ai_run_trace_fields.py b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0005_add_ai_run_trace_fields.py new file mode 100644 index 00000000..aba14231 --- /dev/null +++ b/yuqei-ai-platform/services/ai-platform-api/migrations/versions/20260622_0005_add_ai_run_trace_fields.py @@ -0,0 +1,39 @@ +"""add ai run trace fields + +Revision ID: 20260622_0005 +Revises: 20260622_0004 +Create Date: 2026-06-22 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "20260622_0005" +down_revision: Union[str, None] = "20260622_0004" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column("ai_run_audits", sa.Column("provider", sa.String(length=64), nullable=True)) + op.add_column("ai_run_audits", sa.Column("model", sa.String(length=128), nullable=True)) + op.add_column("ai_run_audits", sa.Column("provider_status", sa.String(length=32), nullable=True)) + op.add_column("ai_run_audits", sa.Column("provider_error", sa.Text(), nullable=True)) + op.add_column("ai_run_audits", sa.Column("prompt_tokens", sa.Integer(), nullable=False, server_default="0")) + op.add_column("ai_run_audits", sa.Column("completion_tokens", sa.Integer(), nullable=False, server_default="0")) + op.add_column("ai_run_audits", sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0")) + op.add_column("ai_run_audits", sa.Column("estimated_cost", sa.Float(), nullable=False, server_default="0")) + + +def downgrade() -> None: + op.drop_column("ai_run_audits", "estimated_cost") + op.drop_column("ai_run_audits", "total_tokens") + op.drop_column("ai_run_audits", "completion_tokens") + op.drop_column("ai_run_audits", "prompt_tokens") + op.drop_column("ai_run_audits", "provider_error") + op.drop_column("ai_run_audits", "provider_status") + op.drop_column("ai_run_audits", "model") + op.drop_column("ai_run_audits", "provider") diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/config.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/config.py index 8a1e6bb9..d7fd687c 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/config.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/config.py @@ -20,6 +20,8 @@ class AiPlatformSettings(BaseSettings): provider_base_url: str | None = None provider_request_timeout_seconds: float = 30.0 provider_fallback_enabled: bool = True + provider_prompt_token_cost_per_1k: float = 0.0 + provider_completion_token_cost_per_1k: float = 0.0 embedding_provider: Literal["local-hash", "openai-compatible"] = "local-hash" embedding_model: str = "local-hash-v1" embedding_dimension: int = 64 diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py index d32f3333..b6b21fa4 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/main.py @@ -255,6 +255,10 @@ def create_app( operation="legal_qa", provider_config_id=provider_config.id if provider_config else None, prompt_template_id=prompt_template.id if prompt_template else None, + provider=provider_response.provider or (provider_config.provider if provider_config else None), + model=provider_response.model or (provider_config.model if provider_config else None), + provider_status=provider_response.status, + provider_error=provider_response.error_message, question=request.question, answer=response.answer, citations_count=len(response.citations), @@ -262,6 +266,10 @@ def create_app( retrieval_strategy=request.retrieval_strategy, retrieved_chunk_ids=[source.chunk_id for source in retrieval_sources if source.chunk_id], retrieved_sources=retrieval_sources, + prompt_tokens=provider_response.prompt_tokens, + completion_tokens=provider_response.completion_tokens, + total_tokens=provider_response.total_tokens, + estimated_cost=provider_response.estimated_cost, latency_ms=max(0, int((perf_counter() - started_at) * 1000)), status=audit_status, ) diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py index 754f7e00..5b71fdfd 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/models.py @@ -107,6 +107,10 @@ class AiRunAuditModel(Base): operation: Mapped[str] = mapped_column(String(64), index=True) provider_config_id: Mapped[str | None] = mapped_column(String(64), nullable=True) prompt_template_id: Mapped[str | None] = mapped_column(String(64), nullable=True) + provider: Mapped[str | None] = mapped_column(String(64), nullable=True) + model: Mapped[str | None] = mapped_column(String(128), nullable=True) + provider_status: Mapped[str | None] = mapped_column(String(32), nullable=True) + provider_error: Mapped[str | None] = mapped_column(Text, nullable=True) question: Mapped[str | None] = mapped_column(Text, nullable=True) answer: Mapped[str | None] = mapped_column(Text, nullable=True) citations_count: Mapped[int] = mapped_column(Integer, default=0) @@ -114,6 +118,10 @@ class AiRunAuditModel(Base): retrieval_strategy: Mapped[str | None] = mapped_column(String(32), nullable=True) retrieved_chunk_ids_json: Mapped[str] = mapped_column(Text, default="[]") retrieved_sources_json: Mapped[str] = mapped_column(Text, default="[]") + prompt_tokens: Mapped[int] = mapped_column(Integer, default=0) + completion_tokens: Mapped[int] = mapped_column(Integer, default=0) + total_tokens: Mapped[int] = mapped_column(Integer, default=0) + estimated_cost: Mapped[float] = mapped_column(Float, default=0.0) latency_ms: Mapped[int] = mapped_column(Integer, default=0) status: Mapped[str] = mapped_column(String(32), default="succeeded") created_at: Mapped[datetime] = mapped_column( diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/providers.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/providers.py index 19dc7c7a..160e2595 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/providers.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/providers.py @@ -20,6 +20,12 @@ class ProviderResponse: answer: str status: str = "succeeded" error_message: str | None = None + provider: str | None = None + model: str | None = None + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + estimated_cost: float = 0.0 class ModelProviderAdapter(Protocol): @@ -32,6 +38,8 @@ class DisabledProviderAdapter: answer="", status="skipped", error_message="Provider calls are disabled.", + provider=request.provider_config.provider if request.provider_config else None, + model=request.provider_config.model if request.provider_config else None, ) @@ -52,11 +60,23 @@ class OpenAICompatibleProviderAdapter: api_key = self._resolve_api_key(config) if not api_key: - return ProviderResponse(answer="", status="failed", error_message="Missing provider API key.") + return ProviderResponse( + answer="", + status="failed", + error_message="Missing provider API key.", + provider=config.provider, + model=config.model, + ) base_url = (config.base_url or self._settings.provider_base_url or "").rstrip("/") if not base_url: - return ProviderResponse(answer="", status="failed", error_message="Missing provider base URL.") + return ProviderResponse( + answer="", + status="failed", + error_message="Missing provider base URL.", + provider=config.provider, + model=config.model, + ) try: with httpx.Client( @@ -84,18 +104,66 @@ class OpenAICompatibleProviderAdapter: ) response.raise_for_status() except httpx.HTTPError as exc: - return ProviderResponse(answer="", status="failed", error_message=str(exc)) + return ProviderResponse( + answer="", + status="failed", + error_message=str(exc), + provider=config.provider, + model=config.model, + ) - answer = _extract_answer(response.json()) + try: + payload = response.json() + except ValueError: + return ProviderResponse( + answer="", + status="failed", + error_message="Provider response is not valid JSON.", + provider=config.provider, + model=config.model, + ) + if not isinstance(payload, dict): + return ProviderResponse( + answer="", + status="failed", + error_message="Provider response JSON is not an object.", + provider=config.provider, + model=config.model, + ) + answer = _extract_answer(payload) + usage = _extract_usage(payload) if not answer: - return ProviderResponse(answer="", status="failed", error_message="Provider response has no answer.") - return ProviderResponse(answer=answer) + return ProviderResponse( + answer="", + status="failed", + error_message="Provider response has no answer.", + provider=config.provider, + model=str(payload.get("model") or config.model) if isinstance(payload, dict) else config.model, + prompt_tokens=usage["prompt_tokens"], + completion_tokens=usage["completion_tokens"], + total_tokens=usage["total_tokens"], + estimated_cost=self._estimate_cost(usage), + ) + return ProviderResponse( + answer=answer, + provider=config.provider, + model=str(payload.get("model") or config.model) if isinstance(payload, dict) else config.model, + prompt_tokens=usage["prompt_tokens"], + completion_tokens=usage["completion_tokens"], + total_tokens=usage["total_tokens"], + estimated_cost=self._estimate_cost(usage), + ) def _resolve_api_key(self, config: ProviderConfig) -> str | None: if config.provider in self._settings.provider_api_keys: return self._settings.provider_api_keys[config.provider] return self._settings.provider_api_key + def _estimate_cost(self, usage: dict[str, int]) -> float: + prompt_cost = usage["prompt_tokens"] / 1000 * self._settings.provider_prompt_token_cost_per_1k + completion_cost = usage["completion_tokens"] / 1000 * self._settings.provider_completion_token_cost_per_1k + return round(prompt_cost + completion_cost, 8) + def build_provider_adapter(settings: AiPlatformSettings) -> ModelProviderAdapter: if settings.provider_call_mode == "enabled": @@ -117,3 +185,27 @@ def _extract_answer(payload: dict[str, Any]) -> str: if isinstance(text, str): return text.strip() return "" + + +def _extract_usage(payload: dict[str, Any]) -> dict[str, int]: + usage = payload.get("usage") + if not isinstance(usage, dict): + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + prompt_tokens = _coerce_usage_int(usage.get("prompt_tokens")) + completion_tokens = _coerce_usage_int(usage.get("completion_tokens")) + total_tokens = _coerce_usage_int(usage.get("total_tokens")) or prompt_tokens + completion_tokens + return { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + } + + +def _coerce_usage_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return max(value, 0) + if isinstance(value, float): + return max(int(value), 0) + return 0 diff --git a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py index 291f707f..fc13c1ce 100644 --- a/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py +++ b/yuqei-ai-platform/services/ai-platform-api/src/yuqei_ai_platform_api/repository.py @@ -195,6 +195,10 @@ class AiRunAudit(BaseModel): operation: str provider_config_id: str | None = None prompt_template_id: str | None = None + provider: str | None = None + model: str | None = None + provider_status: str | None = None + provider_error: str | None = None question: str | None = None answer: str | None = None citations_count: int = 0 @@ -202,6 +206,10 @@ class AiRunAudit(BaseModel): retrieval_strategy: str | None = None retrieved_chunk_ids: list[str] = Field(default_factory=list) retrieved_sources: list[RetrievalSource] = Field(default_factory=list) + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + estimated_cost: float = 0.0 latency_ms: int = 0 status: str = "succeeded" created_at: datetime = Field(default_factory=lambda: datetime.now(UTC)) @@ -934,6 +942,10 @@ class SqlAlchemyAiPlatformRepository: operation=audit.operation, provider_config_id=audit.provider_config_id, prompt_template_id=audit.prompt_template_id, + provider=audit.provider, + model=audit.model, + provider_status=audit.provider_status, + provider_error=audit.provider_error, question=audit.question, answer=audit.answer, citations_count=audit.citations_count, @@ -943,6 +955,10 @@ class SqlAlchemyAiPlatformRepository: retrieved_sources_json=_encode_json_list( [source.model_dump() for source in audit.retrieved_sources] ), + prompt_tokens=audit.prompt_tokens, + completion_tokens=audit.completion_tokens, + total_tokens=audit.total_tokens, + estimated_cost=audit.estimated_cost, latency_ms=audit.latency_ms, status=audit.status, created_at=audit.created_at, @@ -1152,6 +1168,10 @@ def _ai_run_audit_from_model(row: AiRunAuditModel) -> AiRunAudit: operation=row.operation, provider_config_id=row.provider_config_id, prompt_template_id=row.prompt_template_id, + provider=row.provider, + model=row.model, + provider_status=row.provider_status, + provider_error=row.provider_error, question=row.question, answer=row.answer, citations_count=row.citations_count, @@ -1159,6 +1179,10 @@ def _ai_run_audit_from_model(row: AiRunAuditModel) -> AiRunAudit: retrieval_strategy=row.retrieval_strategy, retrieved_chunk_ids=_decode_json_list(row.retrieved_chunk_ids_json), retrieved_sources=[RetrievalSource.model_validate(item) for item in _decode_json_list(row.retrieved_sources_json)], + prompt_tokens=row.prompt_tokens, + completion_tokens=row.completion_tokens, + total_tokens=row.total_tokens, + estimated_cost=row.estimated_cost, latency_ms=row.latency_ms, status=row.status, created_at=_ensure_utc(row.created_at), diff --git a/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py b/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py index d82d7b9b..dac989e1 100644 --- a/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py +++ b/yuqei-ai-platform/services/ai-platform-api/tests/test_platform_data.py @@ -374,6 +374,9 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: operation="legal_qa", provider_config_id=provider.id, prompt_template_id=prompt.id, + provider="deepseek", + model="deepseek-chat", + provider_status="succeeded", question="How should lease deposit refund be handled?", answer="The contract should clearly state deposit refund conditions.", citations_count=1, @@ -397,6 +400,10 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: "matched_terms": ["lease"], } ], + prompt_tokens=12, + completion_tokens=8, + total_tokens=20, + estimated_cost=0.00028, latency_ms=12, status="succeeded", ) @@ -418,5 +425,12 @@ def test_sqlalchemy_repository_persists_platform_data(tmp_path) -> None: assert chunk_result.chunk.embedding_vector audit = reloaded.list_ai_run_audits()[0] assert audit.id == "run-test-001" + assert audit.provider == "deepseek" + assert audit.model == "deepseek-chat" + assert audit.provider_status == "succeeded" + assert audit.prompt_tokens == 12 + assert audit.completion_tokens == 8 + assert audit.total_tokens == 20 + assert audit.estimated_cost == 0.00028 assert audit.retrieved_chunk_ids == ["chunk-test-001"] assert audit.retrieved_sources[0].reference == "Article 703" diff --git a/yuqei-ai-platform/services/ai-platform-api/tests/test_provider_flow.py b/yuqei-ai-platform/services/ai-platform-api/tests/test_provider_flow.py index dc8a5888..72ce02f9 100644 --- a/yuqei-ai-platform/services/ai-platform-api/tests/test_provider_flow.py +++ b/yuqei-ai-platform/services/ai-platform-api/tests/test_provider_flow.py @@ -31,13 +31,15 @@ def test_legal_qa_uses_provider_answer_when_enabled(monkeypatch) -> None: return httpx.Response( 200, json={ + "model": "legal-pro", "choices": [ { "message": { "content": "模型回答：依据检索材料，应先核对合同期限和试用期上限。" } } - ] + ], + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, }, ) @@ -51,6 +53,8 @@ def test_legal_qa_uses_provider_answer_when_enabled(monkeypatch) -> None: environment="test", provider_call_mode="enabled", provider_api_key="test-key", + provider_prompt_token_cost_per_1k=0.001, + provider_completion_token_cost_per_1k=0.002, ), repository=repository, ) @@ -63,7 +67,16 @@ def test_legal_qa_uses_provider_answer_when_enabled(monkeypatch) -> None: assert payload["answer"].startswith("模型回答") audit_response = client.get("/api/v1/audit/ai-runs") - assert audit_response.json()[0]["status"] == "succeeded" + audit = audit_response.json()[0] + assert audit["status"] == "succeeded" + assert audit["provider"] == "openai-compatible" + assert audit["model"] == "legal-pro" + assert audit["provider_status"] == "succeeded" + assert audit["provider_error"] is None + assert audit["prompt_tokens"] == 100 + assert audit["completion_tokens"] == 50 + assert audit["total_tokens"] == 150 + assert audit["estimated_cost"] == 0.0002 def test_provider_failure_falls_back_to_rule_answer_and_audit_status() -> None: @@ -94,7 +107,13 @@ def test_provider_failure_falls_back_to_rule_answer_and_audit_status() -> None: assert response.status_code == 200 assert response.json()["citations"] - assert client.get("/api/v1/audit/ai-runs").json()[0]["status"] == "fallback" + audit = client.get("/api/v1/audit/ai-runs").json()[0] + assert audit["status"] == "fallback" + assert audit["provider"] == "openai-compatible" + assert audit["model"] == "legal-pro" + assert audit["provider_status"] == "failed" + assert audit["provider_error"] == "Missing provider API key." + assert audit["total_tokens"] == 0 def test_openai_compatible_provider_extracts_chat_answer() -> None: @@ -111,13 +130,22 @@ def test_openai_compatible_provider_extracts_chat_answer() -> None: ) def handler(_request: httpx.Request) -> httpx.Response: - return httpx.Response(200, json={"choices": [{"message": {"content": "OK"}}]}) + return httpx.Response( + 200, + json={ + "model": "deepseek-chat", + "choices": [{"message": {"content": "OK"}}], + "usage": {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20}, + }, + ) adapter = OpenAICompatibleProviderAdapter( AiPlatformSettings( environment="test", provider_call_mode="enabled", provider_api_key="test-key", + provider_prompt_token_cost_per_1k=0.01, + provider_completion_token_cost_per_1k=0.02, ), transport=httpx.MockTransport(handler), ) @@ -126,6 +154,12 @@ def test_openai_compatible_provider_extracts_chat_answer() -> None: assert response.status == "succeeded" assert response.answer == "OK" + assert response.provider == "deepseek" + assert response.model == "deepseek-chat" + assert response.prompt_tokens == 12 + assert response.completion_tokens == 8 + assert response.total_tokens == 20 + assert response.estimated_cost == 0.00028 def test_render_legal_qa_prompt_includes_question_context_and_citations() -> None: diff --git a/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py b/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py index 1931dfe2..398448bd 100644 --- a/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py +++ b/yuqei-ai-sdk-python/src/yuqei_sdk/ai_platform.py @@ -193,6 +193,10 @@ class AiRunAudit(BaseModel): operation: str provider_config_id: str | None = None prompt_template_id: str | None = None + provider: str | None = None + model: str | None = None + provider_status: str | None = None + provider_error: str | None = None question: str | None = None answer: str | None = None citations_count: int = 0 @@ -200,6 +204,10 @@ class AiRunAudit(BaseModel): retrieval_strategy: str | None = None retrieved_chunk_ids: list[str] = Field(default_factory=list) retrieved_sources: list[RetrievalSource] = Field(default_factory=list) + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + estimated_cost: float = 0.0 latency_ms: int = 0 status: str created_at: str diff --git a/yuqei-ai-sdk-python/tests/test_ai_platform.py b/yuqei-ai-sdk-python/tests/test_ai_platform.py index 57251fce..d3795f6b 100644 --- a/yuqei-ai-sdk-python/tests/test_ai_platform.py +++ b/yuqei-ai-sdk-python/tests/test_ai_platform.py @@ -349,6 +349,10 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "operation": "legal_qa", "provider_config_id": "provider-1", "prompt_template_id": "prompt-1", + "provider": "deepseek", + "model": "deepseek-chat", + "provider_status": "succeeded", + "provider_error": None, "question": "How should lease deposit refund be handled?", "answer": "The contract should clearly state deposit refund conditions.", "citations_count": 1, @@ -372,6 +376,10 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: "matched_terms": ["lease"], } ], + "prompt_tokens": 12, + "completion_tokens": 8, + "total_tokens": 20, + "estimated_cost": 0.00028, "latency_ms": 12, "status": "succeeded", "created_at": "2026-06-22T00:00:00Z", @@ -431,6 +439,14 @@ def test_ai_platform_client_covers_admin_data_endpoints() -> None: assert explain.results[0].rank == 1 audit = client.list_ai_run_audits(limit=5)[0] assert audit.operation == "legal_qa" + assert audit.provider == "deepseek" + assert audit.model == "deepseek-chat" + assert audit.provider_status == "succeeded" + assert audit.provider_error is None + assert audit.prompt_tokens == 12 + assert audit.completion_tokens == 8 + assert audit.total_tokens == 20 + assert audit.estimated_cost == 0.00028 assert audit.retrieved_chunk_ids == ["chunk-1"] assert audit.retrieved_sources[0].reference == "Article 703" assert "/api/v1/audit/ai-runs" in seen_paths