Show retrieved knowledge chunks in audit logs

ai_knowledge_base_v2/apps/backend/app/services/ai_request_log_service.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import json
+from collections.abc import Sequence
+
 from sqlalchemy.orm import Session
 
 from app.models.logs import AiRequestLog
+from app.services.rag_service import RetrievedChunk
 
 
 class AiRequestLogService:
@@ -20,6 +24,7 @@ class AiRequestLogService:
         input_token: int,
         output_token: int,
         cost_ms: int,
+        retrieved_chunks: Sequence[RetrievedChunk] | None = None,
     ) -> None:
         db.add(
             AiRequestLog(
@@ -30,6 +35,7 @@ class AiRequestLogService:
                 prompt=prompt,
                 knowledge_ids=knowledge_ids,
                 retrieve_count=retrieve_count,
+                retrieved_chunks=_dump_retrieved_chunks(retrieved_chunks),
                 input_token=input_token,
                 output_token=output_token,
                 total_token=input_token + output_token,
@@ -51,6 +57,7 @@ class AiRequestLogService:
         retrieve_count: int,
         cost_ms: int,
         error_message: str,
+        retrieved_chunks: Sequence[RetrievedChunk] | None = None,
     ) -> None:
         db.add(
             AiRequestLog(
@@ -61,8 +68,26 @@ class AiRequestLogService:
                 prompt=prompt,
                 knowledge_ids=knowledge_ids,
                 retrieve_count=retrieve_count,
+                retrieved_chunks=_dump_retrieved_chunks(retrieved_chunks),
                 cost_ms=cost_ms,
                 status="FAILED",
                 error_message=error_message,
             )
         )
+
+
+def _dump_retrieved_chunks(chunks: Sequence[RetrievedChunk] | None) -> str | None:
+    if not chunks:
+        return None
+    payload = [
+        {
+            "index": index,
+            "knowledgeId": chunk.knowledge_id,
+            "knowledgeName": chunk.knowledge_name,
+            "title": chunk.title,
+            "content": chunk.content,
+            "sourceUrl": chunk.source_url,
+        }
+        for index, chunk in enumerate(chunks, start=1)
+    ]
+    return json.dumps(payload, ensure_ascii=False)

ai_knowledge_base_v2/apps/backend/app/services/chat_service.py

@@ -87,6 +87,7 @@ class ChatService:
         db.flush()
 
         started_at = perf_counter()
+        rag_result = None
         try:
             # 获取历史消息（不含刚插入的 user_message，它还没 flush id）
             history = list(
@@ -121,11 +122,12 @@ class ChatService:
                 message_id=user_message.id,
                 user_id=user.id,
                 model_name=None,
-                prompt=normalized_question,
-                knowledge_ids=None,
-                retrieve_count=0,
+                prompt=rag_result.prompt if rag_result is not None else normalized_question,
+                knowledge_ids=rag_result.knowledge_ids if rag_result is not None else None,
+                retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                 cost_ms=cost_ms,
                 error_message=str(exc),
+                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
             )
             db.commit()
             raise HTTPException(
@@ -168,6 +170,7 @@ class ChatService:
             input_token=completion.input_token,
             output_token=completion.output_token,
             cost_ms=cost_ms,
+            retrieved_chunks=rag_result.chunks,
         )
         db.commit()
         return completion.answer

ai_knowledge_base_v2/apps/backend/app/services/chat_stream_service.py

@@ -80,6 +80,7 @@ class ChatStreamService:
                 retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                 cost_ms=cost_ms,
                 error_message=str(exc),
+                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
             )
             db.commit()
             raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc)) from exc
@@ -98,6 +99,7 @@ class ChatStreamService:
                 retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                 cost_ms=cost_ms,
                 error_message="模型未返回有效内容",
+                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
             )
             db.commit()
             raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="模型未返回有效内容")
@@ -137,6 +139,7 @@ class ChatStreamService:
             input_token=model_response.input_token if model_response is not None else 0,
             output_token=_rough_token_count(answer),
             cost_ms=cost_ms,
+            retrieved_chunks=rag_result.chunks if rag_result is not None else None,
         )
         db.commit()
 
@@ -205,6 +208,7 @@ class ChatStreamService:
                 retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                 cost_ms=cost_ms,
                 error_message=str(exc),
+                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
             )
             db.commit()
             raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail=str(exc)) from exc
@@ -223,6 +227,7 @@ class ChatStreamService:
                 retrieve_count=len(rag_result.chunks) if rag_result is not None else 0,
                 cost_ms=cost_ms,
                 error_message="模型未返回有效内容",
+                retrieved_chunks=rag_result.chunks if rag_result is not None else None,
             )
             db.commit()
             raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="模型未返回有效内容")
@@ -295,6 +300,7 @@ def _write_success(
         input_token=model_response.input_token if model_response is not None else 0,
         output_token=_rough_token_count(answer),
         cost_ms=cost_ms,
+        retrieved_chunks=rag_result.chunks if rag_result is not None else None,
     )
     db.commit()