add context attributes

Author: axiomofjoy

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

@@ -9,6 +9,7 @@
     Any,
     Dict,
     Iterator,
+    List,
     Literal,
     Optional,
     Tuple,
@@ -40,6 +41,7 @@
     TokenCount,
     Tool,
 )
+from .helpers import safe_json_dumps
 
 pydantic: Optional[ModuleType]
 try:
@@ -51,6 +53,46 @@
     from _typeshed import DataclassInstance
 
 
+def get_context_attributes(
+    *,
+    session_id: Optional[str] = None,
+    user_id: Optional[str] = None,
+    metadata: Optional[Union[str, Dict[str, Any]]] = None,
+    tags: Optional[List[str]] = None,
+) -> Dict[str, AttributeValue]:
+    context_attributes: Dict[str, AttributeValue] = {}
+    if session_id is not None:
+        context_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))
+    if metadata is not None:
+        context_attributes.update(get_metadata_attributes(metadata=metadata))
+    if tags is not None:
+        context_attributes.update(get_tag_attributes(tags=tags))
+    return context_attributes
+
+
+def get_session_attributes(*, session_id: str) -> Dict[str, AttributeValue]:
+    return {SESSION_ID: session_id}
+
+
+def get_tag_attributes(*, tags: List[str]) -> Dict[str, AttributeValue]:
+    return {TAG_TAGS: tags}
+
+
+def get_metadata_attributes(*, metadata: Union[str, Dict[str, Any]]) -> Dict[str, AttributeValue]:
+    serialized_metadata: str
+    if isinstance(metadata, str):
+        serialized_metadata = metadata
+    else:
+        serialized_metadata = safe_json_dumps(metadata)
+    return {METADATA: serialized_metadata}
+
+
+def get_user_id(*, user_id: str) -> Dict[str, AttributeValue]:
+    return {USER_ID: user_id}
+
+
 def get_span_kind_attributes(kind: "OpenInferenceSpanKind", /) -> Dict[str, AttributeValue]:
     normalized_kind = _normalize_openinference_span_kind(kind)
     return {
@@ -402,12 +444,16 @@ def get_llm_tool_attributes(
 LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT
 LLM_TOKEN_COUNT_TOTAL = SpanAttributes.LLM_TOKEN_COUNT_TOTAL
 LLM_TOOLS = SpanAttributes.LLM_TOOLS
+METADATA = SpanAttributes.METADATA
 OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
 OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
+SESSION_ID = SpanAttributes.SESSION_ID
+TAG_TAGS = SpanAttributes.TAG_TAGS
 TOOL_DESCRIPTION = SpanAttributes.TOOL_DESCRIPTION
 TOOL_NAME = SpanAttributes.TOOL_NAME
 TOOL_PARAMETERS = SpanAttributes.TOOL_PARAMETERS
+USER_ID = SpanAttributes.USER_ID
 
 
 # tool attributes