embedding 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 (
+    EmbeddingAttributes,
     ImageAttributes,
     MessageAttributes,
     MessageContentAttributes,
@@ -33,6 +34,7 @@
 )
 
 from ._types import (
+    Embedding,
     Message,
     OpenInferenceLLMProvider,
     OpenInferenceLLMSystem,
@@ -53,6 +55,25 @@
     from _typeshed import DataclassInstance
 
 
+def get_embedding_attributes(
+    *,
+    model_name: Optional[str] = None,
+    embeddings: Optional[List[Embedding]] = None,
+) -> Dict[str, Any]:
+    embedding_attributes: Dict[str, AttributeValue] = {}
+    if model_name is not None:
+        embedding_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
+            if (vector := embedding.get("vector")) is not None:
+                key = f"{EMBEDDING_EMBEDDINGS}.{index}.{EMBEDDING_VECTOR}"
+                embedding_attributes[key] = vector
+    return embedding_attributes
+
+
 def get_context_attributes(
     *,
     session_id: Optional[str] = None,
@@ -413,6 +434,10 @@ def get_llm_tool_attributes(
     return attributes
 
 
+# embedding attributes
+EMBEDDING_TEXT = EmbeddingAttributes.EMBEDDING_TEXT
+EMBEDDING_VECTOR = EmbeddingAttributes.EMBEDDING_VECTOR
+
 # image attributes
 IMAGE_URL = ImageAttributes.IMAGE_URL
 
@@ -432,6 +457,8 @@ def get_llm_tool_attributes(
 
 
 # span attributes
+EMBEDDING_EMBEDDINGS = SpanAttributes.EMBEDDING_EMBEDDINGS
+EMBEDDING_MODEL_NAME = SpanAttributes.EMBEDDING_MODEL_NAME
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
 INPUT_VALUE = SpanAttributes.INPUT_VALUE
 LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES

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

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Any, Dict, Literal, TypedDict, Union
+from typing import Any, Dict, List, Literal, TypedDict, Union
 
 from typing_extensions import Required, TypeAlias
 
@@ -76,3 +76,8 @@ class TokenCount(TypedDict, total=False):
 
 class Tool(TypedDict, total=False):
     json_schema: Required[Union[str, Dict[str, Any]]]
+
+
+class Embedding(TypedDict, total=False):
+    text: str
+    vector: List[float]