arize-phoenix 5.7.0__py3-none-any.whl → 5.9.0__py3-none-any.whl

phoenix/server/api/helpers/playground_spans.py

@@ -26,7 +26,6 @@ from openinference.semconv.trace import (
 )
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator as DefaultOTelIDGenerator
 from opentelemetry.trace import StatusCode
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.scalars import JSON as JSONScalarType
 from typing_extensions import Self, TypeAlias, assert_never
 
@@ -41,7 +40,7 @@ from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
     TextChunk,
     ToolCallChunk,
 )
-from phoenix.trace.attributes import unflatten
+from phoenix.trace.attributes import get_attribute_value, unflatten
 from phoenix.trace.schemas import (
     SpanEvent,
     SpanException,
@@ -56,7 +55,8 @@ ToolCallID: TypeAlias = str
 
 class streaming_llm_span:
     """
-    Creates an LLM span for a streaming chat completion.
+    A context manager that records OpenInference attributes for streaming chat
+    completion LLM spans.
     """
 
     def __init__(
@@ -71,24 +71,23 @@ class streaming_llm_span:
         self._attributes: dict[str, Any] = attributes if attributes is not None else {}
         self._attributes.update(
             chain(
-                _llm_span_kind(),
-                _llm_model_name(input.model.name),
-                _llm_tools(input.tools or []),
-                _llm_input_messages(messages),
-                _llm_invocation_parameters(invocation_parameters),
-                _input_value_and_mime_type(input),
+                llm_span_kind(),
+                llm_model_name(input.model.name),
+                llm_tools(input.tools or []),
+                llm_input_messages(messages),
+                llm_invocation_parameters(invocation_parameters),
+                input_value_and_mime_type(input),
             )
         )
         self._events: list[SpanEvent] = []
-        self._start_time: datetime
-        self._end_time: datetime
-        self._response_chunks: list[Union[TextChunk, ToolCallChunk]] = []
+        self._start_time: Optional[datetime] = None
+        self._end_time: Optional[datetime] = None
         self._text_chunks: list[TextChunk] = []
         self._tool_call_chunks: defaultdict[ToolCallID, list[ToolCallChunk]] = defaultdict(list)
-        self._status_code: StatusCode
-        self._status_message: str
-        self._db_span: models.Span
-        self._db_trace: models.Trace
+        self._status_code: StatusCode = StatusCode.UNSET
+        self._status_message: Optional[str] = None
+        self._trace_id = _generate_trace_id()
+        self._span_id = _generate_span_id()
 
     async def __aenter__(self) -> Self:
         self._start_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
@@ -102,7 +101,6 @@ class streaming_llm_span:
     ) -> bool:
         self._end_time = cast(datetime, normalize_datetime(dt=local_now(), tz=timezone.utc))
         self._status_code = StatusCode.OK
-        self._status_message = ""
         if exc_type is not None:
             self._status_code = StatusCode.ERROR
             self._status_message = str(exc_value)
@@ -115,10 +113,10 @@ class streaming_llm_span:
                     exception_stacktrace=format_exc(),
                 )
             )
-        if self._response_chunks:
+        if self._text_chunks or self._tool_call_chunks:
             self._attributes.update(
                 chain(
-                    _output_value_and_mime_type(self._response_chunks),
+                    _output_value_and_mime_type(self._text_chunks, self._tool_call_chunks),
                     _llm_output_messages(self._text_chunks, self._tool_call_chunks),
                 )
             )
@@ -127,46 +125,7 @@ class streaming_llm_span:
     def set_attributes(self, attributes: Mapping[str, Any]) -> None:
         self._attributes.update(attributes)
 
-    def add_to_session(
-        self,
-        session: AsyncSession,
-        project_id: int,
-    ) -> models.Span:
-        prompt_tokens = self._attributes.get(LLM_TOKEN_COUNT_PROMPT, 0)
-        completion_tokens = self._attributes.get(LLM_TOKEN_COUNT_COMPLETION, 0)
-        trace_id = _generate_trace_id()
-        span_id = _generate_span_id()
-        self._db_trace = models.Trace(
-            project_rowid=project_id,
-            trace_id=trace_id,
-            start_time=self._start_time,
-            end_time=self._end_time,
-        )
-        self._db_span = models.Span(
-            trace_rowid=self._db_trace.id,
-            span_id=span_id,
-            parent_id=None,
-            name="ChatCompletion",
-            span_kind=LLM,
-            start_time=self._start_time,
-            end_time=self._end_time,
-            attributes=unflatten(self._attributes.items()),
-            events=[_serialize_event(event) for event in self._events],
-            status_code=self._status_code.name,
-            status_message=self._status_message,
-            cumulative_error_count=int(self._status_code is StatusCode.ERROR),
-            cumulative_llm_token_count_prompt=prompt_tokens,
-            cumulative_llm_token_count_completion=completion_tokens,
-            llm_token_count_prompt=prompt_tokens,
-            llm_token_count_completion=completion_tokens,
-            trace=self._db_trace,
-        )
-        session.add(self._db_trace)
-        session.add(self._db_span)
-        return self._db_span
-
     def add_response_chunk(self, chunk: Union[TextChunk, ToolCallChunk]) -> None:
-        self._response_chunks.append(chunk)
         if isinstance(chunk, TextChunk):
             self._text_chunks.append(chunk)
         elif isinstance(chunk, ToolCallChunk):
@@ -174,48 +133,128 @@ class streaming_llm_span:
         else:
             assert_never(chunk)
 
+    @property
+    def span_id(self) -> str:
+        return self._span_id
+
+    @property
+    def trace_id(self) -> str:
+        return self._trace_id
+
     @property
     def start_time(self) -> datetime:
-        return self._db_span.start_time
+        if self._start_time is None:
+            raise ValueError("Cannot access start time before the context manager is entered")
+        return self._start_time
 
     @property
     def end_time(self) -> datetime:
-        return self._db_span.end_time
+        if self._end_time is None:
+            raise ValueError("Cannot access end time before the context manager is exited")
+        return self._end_time
 
     @property
-    def error_message(self) -> Optional[str]:
-        return self._status_message if self._status_code is StatusCode.ERROR else None
+    def status_code(self) -> StatusCode:
+        return self._status_code
 
     @property
-    def trace_id(self) -> str:
-        return self._db_trace.trace_id
+    def status_message(self) -> Optional[str]:
+        if self._status_code is StatusCode.UNSET:
+            raise ValueError("Cannot access status message before the context manager is exited")
+        return self._status_message
 
     @property
-    def attributes(self) -> dict[str, Any]:
-        return self._db_span.attributes
-
+    def events(self) -> list[SpanEvent]:
+        return self._events
 
-def _llm_span_kind() -> Iterator[tuple[str, Any]]:
+    @property
+    def attributes(self) -> dict[str, Any]:
+        return unflatten(self._attributes.items())
+
+
+def get_db_trace(span: streaming_llm_span, project_id: int) -> models.Trace:
+    return models.Trace(
+        project_rowid=project_id,
+        trace_id=span.trace_id,
+        start_time=span.start_time,
+        end_time=span.end_time,
+    )
+
+
+def get_db_span(
+    span: streaming_llm_span,
+    db_trace: models.Trace,
+) -> models.Span:
+    prompt_tokens = get_attribute_value(span.attributes, LLM_TOKEN_COUNT_PROMPT) or 0
+    completion_tokens = get_attribute_value(span.attributes, LLM_TOKEN_COUNT_COMPLETION) or 0
+    return models.Span(
+        trace_rowid=db_trace.id,
+        span_id=span.span_id,
+        parent_id=None,
+        name="ChatCompletion",
+        span_kind=LLM,
+        start_time=span.start_time,
+        end_time=span.end_time,
+        attributes=span.attributes,
+        events=[_serialize_event(event) for event in span.events],
+        status_code=span.status_code.name,
+        status_message=span.status_message or "",
+        cumulative_error_count=int(span.status_code is StatusCode.ERROR),
+        cumulative_llm_token_count_prompt=prompt_tokens,
+        cumulative_llm_token_count_completion=completion_tokens,
+        llm_token_count_prompt=prompt_tokens,
+        llm_token_count_completion=completion_tokens,
+        trace=db_trace,
+    )
+
+
+def get_db_experiment_run(
+    db_span: models.Span,
+    db_trace: models.Trace,
+    *,
+    experiment_id: int,
+    example_id: int,
+) -> models.ExperimentRun:
+    return models.ExperimentRun(
+        experiment_id=experiment_id,
+        dataset_example_id=example_id,
+        trace_id=db_trace.trace_id,
+        output=models.ExperimentRunOutput(
+            task_output=get_attribute_value(db_span.attributes, LLM_OUTPUT_MESSAGES),
+        ),
+        repetition_number=1,
+        start_time=db_span.start_time,
+        end_time=db_span.end_time,
+        error=db_span.status_message or None,
+        prompt_token_count=get_attribute_value(db_span.attributes, LLM_TOKEN_COUNT_PROMPT),
+        completion_token_count=get_attribute_value(db_span.attributes, LLM_TOKEN_COUNT_COMPLETION),
+        trace=db_trace,
+    )
+
+
+def llm_span_kind() -> Iterator[tuple[str, Any]]:
     yield OPENINFERENCE_SPAN_KIND, LLM
 
 
-def _llm_model_name(model_name: str) -> Iterator[tuple[str, Any]]:
+def llm_model_name(model_name: str) -> Iterator[tuple[str, Any]]:
     yield LLM_MODEL_NAME, model_name
 
 
-def _llm_invocation_parameters(
+def llm_invocation_parameters(
     invocation_parameters: Mapping[str, Any],
 ) -> Iterator[tuple[str, Any]]:
     if invocation_parameters:
         yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(invocation_parameters)
 
 
-def _llm_tools(tools: list[JSONScalarType]) -> Iterator[tuple[str, Any]]:
+def llm_tools(tools: list[JSONScalarType]) -> Iterator[tuple[str, Any]]:
     for tool_index, tool in enumerate(tools):
         yield f"{LLM_TOOLS}.{tool_index}.{TOOL_JSON_SCHEMA}", json.dumps(tool)
 
 
-def _input_value_and_mime_type(input: Any) -> Iterator[tuple[str, Any]]:
+def input_value_and_mime_type(
+    input: Union[ChatCompletionInput, ChatCompletionOverDatasetInput],
+) -> Iterator[tuple[str, Any]]:
     assert (api_key := "api_key") in (input_data := jsonify(input))
     disallowed_keys = {"api_key", "invocation_parameters"}
     input_data = {k: v for k, v in input_data.items() if k not in disallowed_keys}
@@ -224,12 +263,69 @@ def _input_value_and_mime_type(input: Any) -> Iterator[tuple[str, Any]]:
     yield INPUT_VALUE, safe_json_dumps(input_data)
 
 
-def _output_value_and_mime_type(output: Any) -> Iterator[tuple[str, Any]]:
-    yield OUTPUT_MIME_TYPE, JSON
-    yield OUTPUT_VALUE, safe_json_dumps(jsonify(output))
+def _merge_tool_call_chunks(
+    chunks_by_id: defaultdict[str, list[ToolCallChunk]],
+) -> list[dict[str, Any]]:
+    merged_tool_calls = []
+
+    for tool_id, chunks in chunks_by_id.items():
+        if not chunks:
+            continue
+        first_chunk = chunks[0]
+        if not first_chunk:
+            continue
+
+        if not hasattr(first_chunk, "function") or not hasattr(first_chunk.function, "name"):
+            continue
+        # Combine all argument chunks
+        merged_arguments = "".join(
+            chunk.function.arguments
+            for chunk in chunks
+            if chunk and hasattr(chunk, "function") and hasattr(chunk.function, "arguments")
+        )
+
+        merged_tool_calls.append(
+            {
+                "id": tool_id,
+                # Only the first chunk has the tool name
+                "function": {
+                    "name": first_chunk.function.name,
+                    "arguments": merged_arguments or "{}",
+                },
+            }
+        )
+
+    return merged_tool_calls
+
+
+def _output_value_and_mime_type(
+    text_chunks: list[TextChunk],
+    tool_call_chunks: defaultdict[ToolCallID, list[ToolCallChunk]],
+) -> Iterator[tuple[str, Any]]:
+    content = "".join(chunk.content for chunk in text_chunks)
+    merged_tool_calls = _merge_tool_call_chunks(tool_call_chunks)
+    if content and merged_tool_calls:
+        yield OUTPUT_MIME_TYPE, JSON
+        yield (
+            OUTPUT_VALUE,
+            safe_json_dumps(
+                {
+                    "content": content,
+                    "tool_calls": jsonify(
+                        merged_tool_calls,
+                    ),
+                }
+            ),
+        )
+    elif merged_tool_calls:
+        yield OUTPUT_MIME_TYPE, JSON
+        yield OUTPUT_VALUE, safe_json_dumps(jsonify(merged_tool_calls))
+    elif content:
+        yield OUTPUT_MIME_TYPE, TEXT
+        yield OUTPUT_VALUE, content
 
 
-def _llm_input_messages(
+def llm_input_messages(
     messages: Iterable[
         tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[JSONScalarType]]]
     ],
@@ -299,6 +395,7 @@ def _serialize_event(event: SpanEvent) -> dict[str, Any]:
 
 
 JSON = OpenInferenceMimeTypeValues.JSON.value
+TEXT = OpenInferenceMimeTypeValues.TEXT.value
 
 LLM = OpenInferenceSpanKindValues.LLM.value
 

phoenix/server/api/input_types/InvocationParameters.py

@@ -47,25 +47,24 @@ class InvocationParameterBase:
     canonical_name: Optional[CanonicalParameterName] = None
     label: str
     required: bool = False
-    hidden: bool = False
 
 
 @strawberry.type
 class IntInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_int
-    default_value: Optional[int] = UNSET
+    default_value: Optional[int] = None
 
 
 @strawberry.type
 class FloatInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_float
-    default_value: Optional[float] = UNSET
+    default_value: Optional[float] = None
 
 
 @strawberry.type
 class BoundedFloatInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_float
-    default_value: Optional[float] = UNSET
+    default_value: Optional[float] = None
     min_value: float
     max_value: float
 
@@ -73,25 +72,25 @@ class BoundedFloatInvocationParameter(InvocationParameterBase):
 @strawberry.type
 class StringInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_string
-    default_value: Optional[str] = UNSET
+    default_value: Optional[str] = None
 
 
 @strawberry.type
 class JSONInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_json
-    default_value: Optional[JSON] = UNSET
+    default_value: Optional[JSON] = None
 
 
 @strawberry.type
 class StringListInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_string_list
-    default_value: Optional[list[str]] = UNSET
+    default_value: Optional[list[str]] = None
 
 
 @strawberry.type
 class BooleanInvocationParameter(InvocationParameterBase):
     invocation_input_field: InvocationInputField = InvocationInputField.value_bool
-    default_value: Optional[bool] = UNSET
+    default_value: Optional[bool] = None
 
 
 def extract_parameter(