retriever attributes

Author: axiomofjoy

python/openinference-instrumentation/src/openinference/instrumentation/_attributes.py

@@ -20,6 +20,7 @@
 from typing_extensions import TypeGuard
 
 from openinference.semconv.trace import (
+    DocumentAttributes,
     EmbeddingAttributes,
     ImageAttributes,
     MessageAttributes,
@@ -34,6 +35,7 @@
 )
 
 from ._types import (
+    Document,
     Embedding,
     Message,
     OpenInferenceLLMProvider,
@@ -55,23 +57,48 @@
     from _typeshed import DataclassInstance
 
 
+def get_retriever_attributes(*, documents: List[Document]) -> Dict[str, AttributeValue]:
+    attributes: Dict[str, AttributeValue] = {}
+    for index, document in enumerate(documents):
+        if not isinstance(document, dict):
+            continue
+        if (content := document.get("content")) is not None:
+            key = f"{RETRIEVAL_DOCUMENTS}.{index}.{DOCUMENT_CONTENT}"
+            attributes[key] = content
+        if (document_id := document.get("id")) is not None:
+            key = f"{RETRIEVAL_DOCUMENTS}.{index}.{DOCUMENT_ID}"
+            attributes[key] = document_id
+        if (metadata := document.get("metadata")) is not None:
+            key = f"{RETRIEVAL_DOCUMENTS}.{index}.{DOCUMENT_METADATA}"
+            serialized_metadata: str
+            if isinstance(metadata, str):
+                serialized_metadata = metadata
+            else:
+                serialized_metadata = safe_json_dumps(metadata)
+            attributes[key] = serialized_metadata
+        if (score := document.get("score")) is not None:
+            key = f"{RETRIEVAL_DOCUMENTS}.{index}.{DOCUMENT_SCORE}"
+            attributes[key] = score
+    return attributes
+
+
 def get_embedding_attributes(
     *,
     model_name: Optional[str] = None,
     embeddings: Optional[List[Embedding]] = None,
-) -> Dict[str, Any]:
-    embedding_attributes: Dict[str, AttributeValue] = {}
+) -> Dict[str, AttributeValue]:
+    attributes: Dict[str, AttributeValue] = {}
     if model_name is not None:
-        embedding_attributes[EMBEDDING_MODEL_NAME] = model_name
+        attributes[EMBEDDING_MODEL_NAME] = model_name
     if isinstance(embeddings, list):
         for index, embedding in enumerate(embeddings):
             if (text := embedding.get("text")) is not None:
                 key = f"{EMBEDDING_EMBEDDINGS}.{index}.{EMBEDDING_TEXT}"
-                embedding_attributes[key] = text
+                attributes[key] = text
             if (vector := embedding.get("vector")) is not None:
                 key = f"{EMBEDDING_EMBEDDINGS}.{index}.{EMBEDDING_VECTOR}"
-                embedding_attributes[key] = vector
-    return embedding_attributes
+                attributes[key] = vector
+    return attributes
 
 
 def get_context_attributes(
@@ -81,16 +108,16 @@ def get_context_attributes(
     metadata: Optional[Union[str, Dict[str, Any]]] = None,
     tags: Optional[List[str]] = None,
 ) -> Dict[str, AttributeValue]:
-    context_attributes: Dict[str, AttributeValue] = {}
+    attributes: Dict[str, AttributeValue] = {}
     if session_id is not None:
-        context_attributes.update(get_session_attributes(session_id=session_id))
+        attributes.update(get_session_attributes(session_id=session_id))
     if user_id is not None:
-        context_attributes.update(get_user_id(user_id=user_id))
+        attributes.update(get_user_id(user_id=user_id))
     if metadata is not None:
-        context_attributes.update(get_metadata_attributes(metadata=metadata))
+        attributes.update(get_metadata_attributes(metadata=metadata))
     if tags is not None:
-        context_attributes.update(get_tag_attributes(tags=tags))
-    return context_attributes
+        attributes.update(get_tag_attributes(tags=tags))
+    return attributes
 
 
 def get_session_attributes(*, session_id: str) -> Dict[str, AttributeValue]:
@@ -434,6 +461,12 @@ def get_llm_tool_attributes(
     return attributes
 
 
+# document attributes
+DOCUMENT_CONTENT = DocumentAttributes.DOCUMENT_CONTENT
+DOCUMENT_ID = DocumentAttributes.DOCUMENT_ID
+DOCUMENT_METADATA = DocumentAttributes.DOCUMENT_METADATA
+DOCUMENT_SCORE = DocumentAttributes.DOCUMENT_SCORE
+
 # embedding attributes
 EMBEDDING_TEXT = EmbeddingAttributes.EMBEDDING_TEXT
 EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
@@ -475,6 +508,7 @@ def get_llm_tool_attributes(
 OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
 OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
+RETRIEVAL_DOCUMENTS = SpanAttributes.RETRIEVAL_DOCUMENTS
 SESSION_ID = SpanAttributes.SESSION_ID
 TAG_TAGS = SpanAttributes.TAG_TAGS
 TOOL_DESCRIPTION = SpanAttributes.TOOL_DESCRIPTION

python/openinference-instrumentation/src/openinference/instrumentation/_types.py

@@ -81,3 +81,10 @@ class Tool(TypedDict, total=False):
 class Embedding(TypedDict, total=False):
     text: str
     vector: List[float]
+
+
+class Document(TypedDict, total=False):
+    content: str
+    id: Union[str, int]
+    metadata: Union[str, Dict[str, Any]]
+    score: float