From c5a0bde3ecec31a6baf86b16ab61b8dcd31ef432 Mon Sep 17 00:00:00 2001
From: YBoy <ashleewunsch1226@gmail.com>
Date: Tue, 7 Apr 2026 07:57:22 +0200
Subject: [PATCH] refactor(api): type aliyun trace utils with TypedDict and
 tighten return types (#34642)

---
 api/core/ops/aliyun_trace/utils.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/api/core/ops/aliyun_trace/utils.py b/api/core/ops/aliyun_trace/utils.py
index d8e105d6a32..aa35ac74c27 100644
--- a/api/core/ops/aliyun_trace/utils.py
+++ b/api/core/ops/aliyun_trace/utils.py
@@ -1,6 +1,6 @@
 import json
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, TypedDict
 
 from graphon.entities import WorkflowNodeExecution
 from graphon.enums import WorkflowNodeExecutionStatus
@@ -56,10 +56,22 @@ def create_links_from_trace_id(trace_id: str | None) -> list[Link]:
     return links
 
 
-def extract_retrieval_documents(documents: list[Document]) -> list[dict[str, Any]]:
-    documents_data = []
+class RetrievalDocumentMetadataDict(TypedDict):
+    dataset_id: Any
+    doc_id: Any
+    document_id: Any
+
+
+class RetrievalDocumentDict(TypedDict):
+    content: str
+    metadata: RetrievalDocumentMetadataDict
+    score: Any
+
+
+def extract_retrieval_documents(documents: list[Document]) -> list[RetrievalDocumentDict]:
+    documents_data: list[RetrievalDocumentDict] = []
     for document in documents:
-        document_data = {
+        document_data: RetrievalDocumentDict = {
             "content": document.page_content,
             "metadata": {
                 "dataset_id": document.metadata.get("dataset_id"),
@@ -83,7 +95,7 @@ def create_common_span_attributes(
     framework: str = DEFAULT_FRAMEWORK_NAME,
     inputs: str = "",
     outputs: str = "",
-) -> dict[str, Any]:
+) -> dict[str, str]:
     return {
         GEN_AI_SESSION_ID: session_id,
         GEN_AI_USER_ID: user_id,