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

phoenix/server/api/mutations/chat_mutations.py

@@ -1,11 +1,12 @@
-import json
-from dataclasses import asdict
+import asyncio
+from dataclasses import asdict, field
 from datetime import datetime, timezone
 from itertools import chain
 from traceback import format_exc
-from typing import Any, Iterable, Iterator, List, Optional
+from typing import Any, Iterable, Iterator, List, Optional, Union
 
 import strawberry
+from openinference.instrumentation import safe_json_dumps
 from openinference.semconv.trace import (
     MessageAttributes,
     OpenInferenceMimeTypeValues,
@@ -17,27 +18,52 @@ from openinference.semconv.trace import (
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator as DefaultOTelIDGenerator
 from opentelemetry.trace import StatusCode
 from sqlalchemy import insert, select
+from strawberry.relay import GlobalID
 from strawberry.types import Info
 from typing_extensions import assert_never
 
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
+from phoenix.db.helpers import get_dataset_example_revisions
 from phoenix.server.api.context import Context
-from phoenix.server.api.exceptions import BadRequest
-from phoenix.server.api.helpers.playground_clients import initialize_playground_clients
+from phoenix.server.api.exceptions import BadRequest, NotFound
+from phoenix.server.api.helpers.playground_clients import (
+    PlaygroundStreamingClient,
+    initialize_playground_clients,
+)
 from phoenix.server.api.helpers.playground_registry import PLAYGROUND_CLIENT_REGISTRY
-from phoenix.server.api.input_types.ChatCompletionInput import ChatCompletionInput
+from phoenix.server.api.helpers.playground_spans import (
+    input_value_and_mime_type,
+    llm_input_messages,
+    llm_invocation_parameters,
+    llm_model_name,
+    llm_span_kind,
+    llm_tools,
+)
+from phoenix.server.api.input_types.ChatCompletionInput import (
+    ChatCompletionInput,
+    ChatCompletionOverDatasetInput,
+)
 from phoenix.server.api.input_types.TemplateOptions import TemplateOptions
+from phoenix.server.api.subscriptions import (
+    _default_playground_experiment_description,
+    _default_playground_experiment_metadata,
+    _default_playground_experiment_name,
+)
 from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
 from phoenix.server.api.types.ChatCompletionSubscriptionPayload import (
     TextChunk,
     ToolCallChunk,
 )
+from phoenix.server.api.types.Dataset import Dataset
+from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.TemplateLanguage import TemplateLanguage
 from phoenix.server.dml_event import SpanInsertEvent
-from phoenix.trace.attributes import unflatten
+from phoenix.trace.attributes import get_attribute_value, unflatten
 from phoenix.trace.schemas import SpanException
+from phoenix.utilities.json import jsonify
 from phoenix.utilities.template_formatters import (
     FStringTemplateFormatter,
     MustacheTemplateFormatter,
@@ -69,21 +95,192 @@ class ChatCompletionMutationPayload:
     error_message: Optional[str]
 
 
+@strawberry.type
+class ChatCompletionMutationError:
+    message: str
+
+
+@strawberry.type
+class ChatCompletionOverDatasetMutationExamplePayload:
+    dataset_example_id: GlobalID
+    experiment_run_id: GlobalID
+    result: Union[ChatCompletionMutationPayload, ChatCompletionMutationError]
+
+
+@strawberry.type
+class ChatCompletionOverDatasetMutationPayload:
+    dataset_id: GlobalID
+    dataset_version_id: GlobalID
+    experiment_id: GlobalID
+    examples: list[ChatCompletionOverDatasetMutationExamplePayload] = field(default_factory=list)
+
+
 @strawberry.type
 class ChatCompletionMutationMixin:
     @strawberry.mutation
+    @classmethod
+    async def chat_completion_over_dataset(
+        cls,
+        info: Info[Context, None],
+        input: ChatCompletionOverDatasetInput,
+    ) -> ChatCompletionOverDatasetMutationPayload:
+        provider_key = input.model.provider_key
+        llm_client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(provider_key, input.model.name)
+        if llm_client_class is None:
+            raise BadRequest(f"No LLM client registered for provider '{provider_key}'")
+        llm_client = llm_client_class(
+            model=input.model,
+            api_key=input.api_key,
+        )
+        dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
+        dataset_version_id = (
+            from_global_id_with_expected_type(
+                global_id=input.dataset_version_id, expected_type_name=DatasetVersion.__name__
+            )
+            if input.dataset_version_id
+            else None
+        )
+        async with info.context.db() as session:
+            dataset = await session.scalar(select(models.Dataset).filter_by(id=dataset_id))
+            if dataset is None:
+                raise NotFound("Dataset not found")
+            if dataset_version_id is None:
+                resolved_version_id = await session.scalar(
+                    select(models.DatasetVersion.id)
+                    .filter_by(dataset_id=dataset_id)
+                    .order_by(models.DatasetVersion.id.desc())
+                    .limit(1)
+                )
+                if resolved_version_id is None:
+                    raise NotFound("No versions found for the given dataset")
+            else:
+                resolved_version_id = dataset_version_id
+            revisions = [
+                revision
+                async for revision in await session.stream_scalars(
+                    get_dataset_example_revisions(resolved_version_id)
+                )
+            ]
+            if not revisions:
+                raise NotFound("No examples found for the given dataset and version")
+            experiment = models.Experiment(
+                dataset_id=from_global_id_with_expected_type(input.dataset_id, Dataset.__name__),
+                dataset_version_id=resolved_version_id,
+                name=input.experiment_name or _default_playground_experiment_name(),
+                description=input.experiment_description
+                or _default_playground_experiment_description(dataset_name=dataset.name),
+                repetitions=1,
+                metadata_=input.experiment_metadata
+                or _default_playground_experiment_metadata(
+                    dataset_name=dataset.name,
+                    dataset_id=input.dataset_id,
+                    version_id=GlobalID(DatasetVersion.__name__, str(resolved_version_id)),
+                ),
+                project_name=PLAYGROUND_PROJECT_NAME,
+            )
+            session.add(experiment)
+            await session.flush()
+
+        start_time = datetime.now(timezone.utc)
+        results = await asyncio.gather(
+            *(
+                cls._chat_completion(
+                    info,
+                    llm_client,
+                    ChatCompletionInput(
+                        model=input.model,
+                        api_key=input.api_key,
+                        messages=input.messages,
+                        tools=input.tools,
+                        invocation_parameters=input.invocation_parameters,
+                        template=TemplateOptions(
+                            language=input.template_language,
+                            variables=revision.input,
+                        ),
+                    ),
+                )
+                for revision in revisions
+            ),
+            return_exceptions=True,
+        )
+
+        payload = ChatCompletionOverDatasetMutationPayload(
+            dataset_id=GlobalID(models.Dataset.__name__, str(dataset.id)),
+            dataset_version_id=GlobalID(DatasetVersion.__name__, str(resolved_version_id)),
+            experiment_id=GlobalID(models.Experiment.__name__, str(experiment.id)),
+        )
+        experiment_runs = []
+        for revision, result in zip(revisions, results):
+            if isinstance(result, BaseException):
+                experiment_run = models.ExperimentRun(
+                    experiment_id=experiment.id,
+                    dataset_example_id=revision.dataset_example_id,
+                    output={},
+                    repetition_number=1,
+                    start_time=start_time,
+                    end_time=start_time,
+                    error=str(result),
+                )
+            else:
+                db_span = result.span.db_span
+                experiment_run = models.ExperimentRun(
+                    experiment_id=experiment.id,
+                    dataset_example_id=revision.dataset_example_id,
+                    trace_id=str(result.span.context.trace_id),
+                    output=models.ExperimentRunOutput(
+                        task_output=get_attribute_value(db_span.attributes, LLM_OUTPUT_MESSAGES),
+                    ),
+                    prompt_token_count=db_span.cumulative_llm_token_count_prompt,
+                    completion_token_count=db_span.cumulative_llm_token_count_completion,
+                    repetition_number=1,
+                    start_time=result.span.start_time,
+                    end_time=result.span.end_time,
+                    error=str(result.error_message) if result.error_message else None,
+                )
+            experiment_runs.append(experiment_run)
+
+        async with info.context.db() as session:
+            session.add_all(experiment_runs)
+            await session.flush()
+
+        for revision, experiment_run, result in zip(revisions, experiment_runs, results):
+            dataset_example_id = GlobalID(
+                models.DatasetExample.__name__, str(revision.dataset_example_id)
+            )
+            experiment_run_id = GlobalID(models.ExperimentRun.__name__, str(experiment_run.id))
+            example_payload = ChatCompletionOverDatasetMutationExamplePayload(
+                dataset_example_id=dataset_example_id,
+                experiment_run_id=experiment_run_id,
+                result=result
+                if isinstance(result, ChatCompletionMutationPayload)
+                else ChatCompletionMutationError(message=str(result)),
+            )
+            payload.examples.append(example_payload)
+        return payload
+
+    @strawberry.mutation
+    @classmethod
     async def chat_completion(
-        self, info: Info[Context, None], input: ChatCompletionInput
+        cls, info: Info[Context, None], input: ChatCompletionInput
     ) -> ChatCompletionMutationPayload:
         provider_key = input.model.provider_key
         llm_client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(provider_key, input.model.name)
         if llm_client_class is None:
             raise BadRequest(f"No LLM client registered for provider '{provider_key}'")
-        attributes: dict[str, Any] = {}
         llm_client = llm_client_class(
             model=input.model,
             api_key=input.api_key,
         )
+        return await cls._chat_completion(info, llm_client, input)
+
+    @classmethod
+    async def _chat_completion(
+        cls,
+        info: Info[Context, None],
+        llm_client: PlaygroundStreamingClient,
+        input: ChatCompletionInput,
+    ) -> ChatCompletionMutationPayload:
+        attributes: dict[str, Any] = {}
 
         messages = [
             (
@@ -94,7 +291,6 @@ class ChatCompletionMutationMixin:
             )
             for message in input.messages
         ]
-
         if template_options := input.template:
             messages = list(_formatted_messages(messages, template_options))
 
@@ -103,17 +299,16 @@ class ChatCompletionMutationMixin:
         )
 
         text_content = ""
-        tool_calls = []
+        tool_calls: dict[str, ChatCompletionToolCall] = {}
         events = []
         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_client.attributes,
+                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),
             )
         )
 
@@ -128,14 +323,16 @@ class ChatCompletionMutationMixin:
                 if isinstance(chunk, TextChunk):
                     text_content += chunk.content
                 elif isinstance(chunk, ToolCallChunk):
-                    tool_call = ChatCompletionToolCall(
-                        id=chunk.id,
-                        function=ChatCompletionFunctionCall(
-                            name=chunk.function.name,
-                            arguments=chunk.function.arguments,
-                        ),
-                    )
-                    tool_calls.append(tool_call)
+                    if chunk.id not in tool_calls:
+                        tool_calls[chunk.id] = ChatCompletionToolCall(
+                            id=chunk.id,
+                            function=ChatCompletionFunctionCall(
+                                name=chunk.function.name,
+                                arguments=chunk.function.arguments,
+                            ),
+                        )
+                    else:
+                        tool_calls[chunk.id].function.arguments += chunk.function.arguments
                 else:
                     assert_never(chunk)
         except Exception as e:
@@ -156,10 +353,11 @@ class ChatCompletionMutationMixin:
         else:
             end_time = normalize_datetime(dt=local_now(), tz=timezone.utc)
 
+        attributes.update(llm_client.attributes)
         if text_content or tool_calls:
             attributes.update(
                 chain(
-                    _output_value_and_mime_type({"text": text_content, "tool_calls": tool_calls}),
+                    _output_value_and_mime_type(text_content, tool_calls),
                     _llm_output_messages(text_content, tool_calls),
                 )
             )
@@ -225,7 +423,7 @@ class ChatCompletionMutationMixin:
         else:
             return ChatCompletionMutationPayload(
                 content=text_content if text_content else None,
-                tool_calls=tool_calls,
+                tool_calls=list(tool_calls.values()),
                 span=gql_span,
                 error_message=None,
             )
@@ -264,61 +462,30 @@ def _template_formatter(template_language: TemplateLanguage) -> TemplateFormatte
     assert_never(template_language)
 
 
-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 _llm_invocation_parameters(invocation_parameters: dict[str, Any]) -> Iterator[tuple[str, Any]]:
-    yield LLM_INVOCATION_PARAMETERS, json.dumps(invocation_parameters)
-
-
-def _llm_tools(tools: List[Any]) -> 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: ChatCompletionInput) -> Iterator[tuple[str, Any]]:
-    input_data = input.__dict__.copy()
-    input_data.pop("api_key", None)
-    yield INPUT_MIME_TYPE, JSON
-    yield INPUT_VALUE, json.dumps(input_data)
-
-
-def _output_value_and_mime_type(output: Any) -> Iterator[tuple[str, Any]]:
-    yield OUTPUT_MIME_TYPE, JSON
-    yield OUTPUT_VALUE, json.dumps(output)
-
-
-def _llm_input_messages(
-    messages: Iterable[ChatCompletionMessage],
+def _output_value_and_mime_type(
+    text: str, tool_calls: dict[str, ChatCompletionToolCall]
 ) -> Iterator[tuple[str, Any]]:
-    for i, (role, content, _tool_call_id, tool_calls) in enumerate(messages):
-        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_ROLE}", role.value.lower()
-        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_CONTENT}", content
-        if tool_calls:
-            for tool_call_index, tool_call in enumerate(tool_calls):
-                yield (
-                    f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
-                    tool_call["function"]["name"],
-                )
-                if arguments := tool_call["function"]["arguments"]:
-                    yield (
-                        f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                        json.dumps(arguments),
-                    )
+    if text and tool_calls:
+        yield OUTPUT_MIME_TYPE, JSON
+        yield (
+            OUTPUT_VALUE,
+            safe_json_dumps({"content": text, "tool_calls": jsonify(list(tool_calls.values()))}),
+        )
+    elif tool_calls:
+        yield OUTPUT_MIME_TYPE, JSON
+        yield OUTPUT_VALUE, safe_json_dumps(jsonify(list(tool_calls.values())))
+    elif text:
+        yield OUTPUT_MIME_TYPE, TEXT
+        yield OUTPUT_VALUE, text
 
 
 def _llm_output_messages(
-    text_content: str, tool_calls: List[ChatCompletionToolCall]
+    text_content: str, tool_calls: dict[str, ChatCompletionToolCall]
 ) -> Iterator[tuple[str, Any]]:
     yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"
     if text_content:
         yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", text_content
-    for tool_call_index, tool_call in enumerate(tool_calls):
+    for tool_call_index, tool_call in enumerate(tool_calls.values()):
         yield (
             f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
             tool_call.function.name,
@@ -326,7 +493,7 @@ def _llm_output_messages(
         if arguments := tool_call.function.arguments:
             yield (
                 f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
-                json.dumps(arguments),
+                arguments,
             )
 
 
@@ -347,6 +514,7 @@ def _serialize_event(event: SpanException) -> dict[str, Any]:
 
 
 JSON = OpenInferenceMimeTypeValues.JSON.value
+TEXT = OpenInferenceMimeTypeValues.TEXT.value
 LLM = OpenInferenceSpanKindValues.LLM.value
 
 OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND

phoenix/server/api/queries.py

@@ -48,6 +48,7 @@ from phoenix.server.api.input_types.DatasetSort import DatasetSort
 from phoenix.server.api.input_types.InvocationParameters import (
     InvocationParameter,
 )
+from phoenix.server.api.subscriptions import PLAYGROUND_PROJECT_NAME
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
@@ -237,7 +238,10 @@ class Query:
             select(models.Project)
             .outerjoin(
                 models.Experiment,
-                models.Project.name == models.Experiment.project_name,
+                and_(
+                    models.Project.name == models.Experiment.project_name,
+                    models.Experiment.project_name != PLAYGROUND_PROJECT_NAME,
+                ),
             )
             .where(models.Experiment.project_name.is_(None))
             .order_by(models.Project.id)

phoenix/server/api/routers/v1/spans.py

@@ -1,8 +1,11 @@
+from asyncio import get_running_loop
 from collections.abc import AsyncIterator
 from datetime import datetime, timezone
+from secrets import token_urlsafe
 from typing import Any, Literal, Optional
 
-from fastapi import APIRouter, HTTPException, Query
+import pandas as pd
+from fastapi import APIRouter, Header, HTTPException, Query
 from pydantic import Field
 from sqlalchemy import select
 from starlette.requests import Request
@@ -19,6 +22,7 @@ from phoenix.db.insertion.types import Precursors
 from phoenix.server.api.routers.utils import df_to_bytes
 from phoenix.server.dml_event import SpanAnnotationInsertEvent
 from phoenix.trace.dsl import SpanQuery as SpanQuery_
+from phoenix.utilities.json import encode_df_as_json_string
 
 from .pydantic_compat import V1RoutesBaseModel
 from .utils import RequestBody, ResponseBody, add_errors_to_responses
@@ -72,6 +76,7 @@ class QuerySpansRequestBody(V1RoutesBaseModel):
 async def query_spans_handler(
     request: Request,
     request_body: QuerySpansRequestBody,
+    accept: Optional[str] = Header(None),
     project_name: Optional[str] = Query(
         default=None, description="The project name to get evaluations from"
     ),
@@ -116,6 +121,13 @@ async def query_spans_handler(
     if not results:
         raise HTTPException(status_code=HTTP_404_NOT_FOUND)
 
+    if accept == "application/json":
+        boundary_token = token_urlsafe(64)
+        return StreamingResponse(
+            content=_json_multipart(results, boundary_token),
+            media_type=f"multipart/mixed; boundary={boundary_token}",
+        )
+
     async def content() -> AsyncIterator[bytes]:
         for result in results:
             yield df_to_bytes(result)
@@ -126,6 +138,18 @@ async def query_spans_handler(
     )
 
 
+async def _json_multipart(
+    results: list[pd.DataFrame],
+    boundary_token: str,
+) -> AsyncIterator[str]:
+    for df in results:
+        yield f"--{boundary_token}\r\n"
+        yield "Content-Type: application/json\r\n\r\n"
+        yield await get_running_loop().run_in_executor(None, encode_df_as_json_string, df)
+        yield "\r\n"
+    yield f"--{boundary_token}--\r\n"
+
+
 @router.get("/spans", include_in_schema=False, deprecated=True)
 async def get_spans_handler(
     request: Request,