playground model name (#5001)

Arize-ai · Oct 15, 2024 · 716b1c7 · 716b1c7
1 parent 23958bd
commit 716b1c7
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/phoenix/server/api/subscriptions.py b/src/phoenix/server/api/subscriptions.py
@@ -107,6 +107,7 @@ async def chat_completion(
             attributes=dict(
                 chain(
                     _llm_span_kind(),
+                    _llm_model_name(input.model.name),
                     _input_value_and_mime_type(input),
                     _llm_input_messages(input.messages),
                 )
@@ -200,6 +201,10 @@ def _llm_span_kind() -> Iterator[Tuple[str, Any]]:
     yield OPENINFERENCE_SPAN_KIND, LLM
 
 
+def _llm_model_name(model_name: str) -> Iterator[Tuple[str, Any]]:
+    yield LLM_MODEL_NAME, model_name
+
+
 def _input_value_and_mime_type(input: ChatCompletionInput) -> Iterator[Tuple[str, Any]]:
     yield INPUT_MIME_TYPE, JSON
     yield INPUT_VALUE, json.dumps(asdict(input), cls=GraphQLInputJSONEncoder)
@@ -261,6 +266,7 @@ def default(self, obj: Any) -> Any:
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
 LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
 LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
+LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
 
 MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
 MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE