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

phoenix/server/api/mutations/chat_mutations.py

@@ -1,9 +1,9 @@
-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
@@ -18,14 +18,19 @@ 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.helpers.playground_spans import (
     input_value_and_mime_type,
@@ -35,17 +40,28 @@ from phoenix.server.api.helpers.playground_spans import (
     llm_span_kind,
     llm_tools,
 )
-from phoenix.server.api.input_types.ChatCompletionInput import ChatCompletionInput
+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 (
@@ -79,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 = [
             (
@@ -122,7 +309,6 @@ class ChatCompletionMutationMixin:
                 llm_input_messages(messages),
                 llm_invocation_parameters(invocation_parameters),
                 input_value_and_mime_type(input),
-                **llm_client.attributes,
             )
         )
 
@@ -167,6 +353,7 @@ 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(
@@ -306,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,
             )
 
 

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,

phoenix/server/api/types/ExperimentRun.py

@@ -1,7 +1,9 @@
 from datetime import datetime
-from typing import Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
+from sqlalchemy import select
+from sqlalchemy.orm import load_only
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
@@ -20,6 +22,9 @@ from phoenix.server.api.types.pagination import (
 )
 from phoenix.server.api.types.Trace import Trace
 
+if TYPE_CHECKING:
+    from phoenix.server.api.types.DatasetExample import DatasetExample
+
 
 @strawberry.type
 class ExperimentRun(Node):
@@ -62,6 +67,38 @@ class ExperimentRun(Node):
         trace_rowid, project_rowid = trace
         return Trace(id_attr=trace_rowid, trace_id=self.trace_id, project_rowid=project_rowid)
 
+    @strawberry.field
+    async def example(
+        self, info: Info
+    ) -> Annotated[
+        "DatasetExample", strawberry.lazy("phoenix.server.api.types.DatasetExample")
+    ]:  # use lazy types to avoid circular import: https://strawberry.rocks/docs/types/lazy
+        from phoenix.server.api.types.DatasetExample import DatasetExample
+
+        async with info.context.db() as session:
+            assert (
+                result := await session.execute(
+                    select(models.DatasetExample, models.Experiment.dataset_version_id)
+                    .select_from(models.ExperimentRun)
+                    .join(
+                        models.DatasetExample,
+                        models.DatasetExample.id == models.ExperimentRun.dataset_example_id,
+                    )
+                    .join(
+                        models.Experiment,
+                        models.Experiment.id == models.ExperimentRun.experiment_id,
+                    )
+                    .where(models.ExperimentRun.id == self.id_attr)
+                    .options(load_only(models.DatasetExample.id, models.DatasetExample.created_at))
+                )
+            ) is not None
+            example, version_id = result.first()
+        return DatasetExample(
+            id_attr=example.id,
+            created_at=example.created_at,
+            version_id=version_id,
+        )
+
 
 def to_gql_experiment_run(run: models.ExperimentRun) -> ExperimentRun:
     """

phoenix/server/api/types/GenerativeProvider.py

@@ -8,6 +8,7 @@ class GenerativeProviderKey(Enum):
     OPENAI = "OpenAI"
     ANTHROPIC = "Anthropic"
     AZURE_OPENAI = "Azure OpenAI"
+    GEMINI = "Google AI Studio"
 
 
 @strawberry.type
@@ -24,7 +25,7 @@ class GenerativeProvider:
 
         default_client = PLAYGROUND_CLIENT_REGISTRY.get_client(self.key, PROVIDER_DEFAULT)
         if default_client:
-            return default_client.dependencies()
+            return [dependency.name for dependency in default_client.dependencies()]
         return []
 
     @strawberry.field

phoenix/server/app.py

@@ -253,7 +253,7 @@ class Static(StaticFiles):
 
 
 class RequestOriginHostnameValidator(BaseHTTPMiddleware):
-    def __init__(self, trusted_hostnames: list[str], *args: Any, **kwargs: Any) -> None:
+    def __init__(self, *args: Any, trusted_hostnames: list[str], **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._trusted_hostnames = trusted_hostnames
 
@@ -767,7 +767,12 @@ def create_app(
     middlewares.extend(user_fastapi_middlewares())
     if origins := get_env_csrf_trusted_origins():
         trusted_hostnames = [h for o in origins if o and (h := urlparse(o).hostname)]
-        middlewares.append(Middleware(RequestOriginHostnameValidator, trusted_hostnames))
+        middlewares.append(
+            Middleware(
+                RequestOriginHostnameValidator,
+                trusted_hostnames=trusted_hostnames,
+            )
+        )
     elif email_sender or oauth2_client_configs:
         logger.warning(
             "CSRF protection can be enabled by listing trusted origins via "

phoenix/server/static/.vite/manifest.json

@@ -1,32 +1,32 @@
 {
-  "_components-MllbfxfJ.js": {
-    "file": "assets/components-MllbfxfJ.js",
+  "_components-DU-8CYbi.js": {
+    "file": "assets/components-DU-8CYbi.js",
     "name": "components",
     "imports": [
-      "_vendor-BEuNhfwH.js",
-      "_vendor-arizeai-Bskhzyjm.js",
-      "_pages-BHfC6jnL.js",
-      "_vendor-three-DwGkEfCM.js",
-      "_vendor-codemirror-DLlXCf0x.js"
+      "_vendor-D04tenE6.js",
+      "_vendor-arizeai-D3NxMQw0.js",
+      "_vendor-codemirror-XTiZSlqq.js",
+      "_pages-t09OI1rC.js",
+      "_vendor-three-DwGkEfCM.js"
     ]
   },
-  "_pages-BHfC6jnL.js": {
-    "file": "assets/pages-BHfC6jnL.js",
+  "_pages-t09OI1rC.js": {
+    "file": "assets/pages-t09OI1rC.js",
     "name": "pages",
     "imports": [
-      "_vendor-BEuNhfwH.js",
-      "_vendor-arizeai-Bskhzyjm.js",
-      "_components-MllbfxfJ.js",
-      "_vendor-recharts-CRqhvLYg.js",
-      "_vendor-codemirror-DLlXCf0x.js"
+      "_vendor-D04tenE6.js",
+      "_vendor-arizeai-D3NxMQw0.js",
+      "_components-DU-8CYbi.js",
+      "_vendor-recharts-p0L0neVs.js",
+      "_vendor-codemirror-XTiZSlqq.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-BEuNhfwH.js": {
-    "file": "assets/vendor-BEuNhfwH.js",
+  "_vendor-D04tenE6.js": {
+    "file": "assets/vendor-D04tenE6.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-Bskhzyjm.js": {
-    "file": "assets/vendor-arizeai-Bskhzyjm.js",
+  "_vendor-arizeai-D3NxMQw0.js": {
+    "file": "assets/vendor-arizeai-D3NxMQw0.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-BEuNhfwH.js"
+      "_vendor-D04tenE6.js"
     ]
   },
-  "_vendor-codemirror-DLlXCf0x.js": {
-    "file": "assets/vendor-codemirror-DLlXCf0x.js",
+  "_vendor-codemirror-XTiZSlqq.js": {
+    "file": "assets/vendor-codemirror-XTiZSlqq.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-BEuNhfwH.js"
+      "_vendor-D04tenE6.js"
     ]
   },
-  "_vendor-recharts-CRqhvLYg.js": {
-    "file": "assets/vendor-recharts-CRqhvLYg.js",
+  "_vendor-recharts-p0L0neVs.js": {
+    "file": "assets/vendor-recharts-p0L0neVs.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-BEuNhfwH.js"
+      "_vendor-D04tenE6.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-BVO2YcT1.js",
+    "file": "assets/index-D9E16vvV.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-BEuNhfwH.js",
-      "_vendor-arizeai-Bskhzyjm.js",
-      "_pages-BHfC6jnL.js",
-      "_components-MllbfxfJ.js",
+      "_vendor-D04tenE6.js",
+      "_vendor-arizeai-D3NxMQw0.js",
+      "_pages-t09OI1rC.js",
+      "_components-DU-8CYbi.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-CRqhvLYg.js",
-      "_vendor-codemirror-DLlXCf0x.js"
+      "_vendor-recharts-p0L0neVs.js",
+      "_vendor-codemirror-XTiZSlqq.js"
     ]
   }
 }