arize-phoenix 5.3.1__py3-none-any.whl → 5.5.0__py3-none-any.whl

{arize_phoenix-5.3.1.dist-info → arize_phoenix-5.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 5.3.1
+Version: 5.5.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -56,6 +56,7 @@ Requires-Dist: uvicorn
 Requires-Dist: websockets
 Requires-Dist: wrapt
 Provides-Extra: container
+Requires-Dist: anthropic; extra == 'container'
 Requires-Dist: fast-hdbscan>=0.2.0; extra == 'container'
 Requires-Dist: numba>=0.60.0; extra == 'container'
 Requires-Dist: openai>=1.0.0; extra == 'container'

{arize_phoenix-5.3.1.dist-info → arize_phoenix-5.5.0.dist-info}/RECORD

@@ -6,7 +6,7 @@ phoenix/exceptions.py,sha256=n2L2KKuecrdflB9MsCdAYCiSEvGJptIsfRkXMoJle7A,169
 phoenix/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 phoenix/services.py,sha256=aTxhcOA1pZHB6U-B3TEcp6fqDF5oT0xCUvEUNMZVTUQ,5175
 phoenix/settings.py,sha256=ht-0oN-sMV6SPXrk7Tu1EZlngpAYkGNLYPhO8DyrdQI,661
-phoenix/version.py,sha256=XehO8pb4GklFJT4padrwyHGgMITY1M6KU3ARht53Mw4,22
+phoenix/version.py,sha256=zFTHldBmR5ReiC3uSZ8VkZOEirtsq_l6QbUJYRBHlTs,22
 phoenix/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/core/embedding_dimension.py,sha256=zKGbcvwOXgLf-yrJBpQyKtd-LEOPRKHnUToyAU8Owis,87
 phoenix/core/model.py,sha256=km_a--PBHOuA337ClRw9xqhOHhrUT6Rl9pz_zV0JYkQ,4843
@@ -94,7 +94,7 @@ phoenix/server/api/exceptions.py,sha256=TA0JuY2YRnj35qGuMSQ8d0ToHum9gWm9W--3fSKH
 phoenix/server/api/interceptor.py,sha256=ykDnoC_apUd-llVli3m1CW18kNSIgjz2qZ6m5JmPDu8,1294
 phoenix/server/api/queries.py,sha256=sE_w4X9sz0QJBGNQbOMgiZiKpR24jvxCpv6AQeGeE9Y,27153
 phoenix/server/api/schema.py,sha256=JKxcZd0UBX6VnVic9tnKAzehJ2fGYNWAu6Gr1tR1PMI,637
-phoenix/server/api/subscriptions.py,sha256=l5t8f23J_3znUS6gFrK3XOPkkX9BN5dHHMteASWbhVk,17083
+phoenix/server/api/subscriptions.py,sha256=hxIj8WRNk1m1aXVW7PDlpPt5TwRGNbjp2yxibRZRqrc,21972
 phoenix/server/api/utils.py,sha256=Kl47G-1A7QKTDrc75BU2QK6HupsG6MWuXxy351FOfKQ,858
 phoenix/server/api/dataloaders/__init__.py,sha256=jNYvfXjnZzgA2HWTG7AZdqWGla3ZysBUDUei8Zkz6N8,3290
 phoenix/server/api/dataloaders/annotation_summaries.py,sha256=Wv8AORZoGd5TJ4Y-em8iqJu87AMpZP7lWOTr-SML-x8,5560
@@ -264,15 +264,15 @@ phoenix/server/static/apple-touch-icon-76x76.png,sha256=CT_xT12I0u2i0WU8JzBZBuOQ
 phoenix/server/static/apple-touch-icon.png,sha256=fOfpjqGpWYbJ0eAurKsyoZP1EAs6ZVooBJ_SGk2ZkDs,3801
 phoenix/server/static/favicon.ico,sha256=bY0vvCKRftemZfPShwZtE93DiiQdaYaozkPGwNFr6H8,34494
 phoenix/server/static/modernizr.js,sha256=mvK-XtkNqjOral-QvzoqsyOMECXIMu5BQwSVN_wcU9c,2564
-phoenix/server/static/.vite/manifest.json,sha256=SieG3iZksdblyz3-x7sbYAjWMF6-abUBPpCo2fI-M6M,1929
-phoenix/server/static/assets/components-CeBEsPFA.js,sha256=vQY-UcFt8MuIlB-NCAw02xsIt8Tv26aLTlYCC8GwFQQ,272882
-phoenix/server/static/assets/index-CRTTIb3y.js,sha256=_HHh0kyNV3oT1-vYiwkXC25jInuLGhLMyr62OLhRNS8,7162
-phoenix/server/static/assets/pages-vHjpYRiN.js,sha256=WEnxplw-UwjrizzaVl3c6umZ_mEq4gl8sTo5vijmSLE,586223
-phoenix/server/static/assets/vendor-D-NIjePD.js,sha256=tba6ip4z7SOnhlPkBkyt9CUJiDDzZwi6IRKCnbwHvuA,10898294
+phoenix/server/static/.vite/manifest.json,sha256=ifveuWhOubKDOfkJLSeHrUMz3kwtdip_8ctJ4KPRfP8,1929
+phoenix/server/static/assets/components-DOjuTDx9.js,sha256=JXFIP0SiBdyEvvh53dKsdoi-6o1lOUFY_IIMsGBwMxg,272352
+phoenix/server/static/assets/index-C15MiAe0.js,sha256=pM97A7RT-WTM1F2TqKVvP8VgNWGHC7OehcBqq-sjzR4,7162
+phoenix/server/static/assets/pages-uwuWR0d3.js,sha256=GWoWNnx_8-q56uUQWbBRM_fZv8zxw1ZZXvWbog8Q0-0,591038
+phoenix/server/static/assets/vendor-6IcPAw_j.js,sha256=kA0afCcdEfjdTlLbngra3925MiR8wVgYF5d7XBBViXQ,10898278
 phoenix/server/static/assets/vendor-DxkFTwjz.css,sha256=nZrkr0u6NNElFGvpWHk9GTHeGoibCXCli1bE7mXZGZg,1816
-phoenix/server/static/assets/vendor-arizeai-DoY5jUTO.js,sha256=RkuJgULcP6CtTAhyLSzI9bNRiJnq84MwkYCF6MQ2sCw,306507
-phoenix/server/static/assets/vendor-codemirror-CIhY_nEU.js,sha256=6g23qgUXzCj6VuoJC9Uhq7Y9qb3_WN3u767NLDhfGkM,392703
-phoenix/server/static/assets/vendor-recharts-Dgcm35Jq.js,sha256=9ifuAYGnEKtBtbTV9sFe95rP4e1NR9j8DOVEb-hRKzM,282859
+phoenix/server/static/assets/vendor-arizeai-uC2sozJ6.js,sha256=o5Su5D9_jBumfVpDXwuQC7Hzpkze4LlvVA9t0Pvh1pk,306512
+phoenix/server/static/assets/vendor-codemirror-DVE2_WBr.js,sha256=AxNgfQZIek1HfON1zHobnrzHzgBeLHWu6GCltpvct9U,392705
+phoenix/server/static/assets/vendor-recharts-DwrexFA4.js,sha256=QSkNeG0tfn_fdrjYf-dRBXVygFjusBsyxzQ-lepdH8w,282859
 phoenix/server/static/assets/vendor-three-DwGkEfCM.js,sha256=0D12ZgKzfKCTSdSTKJBFR2RZO_xxeMXrqDp0AszZqHY,620972
 phoenix/server/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/server/templates/index.html,sha256=ULNlp0-agU8pH4wCIapE1goqnwkbctAoQRgrfwH_irg,4379
@@ -286,7 +286,7 @@ phoenix/trace/attributes.py,sha256=B_OrzVaxZwFkrAFXZyicYoIti1UdUysURsvUS2GyW1U,1
 phoenix/trace/errors.py,sha256=wB1z8qdPckngdfU-TORToekvg3344oNFAA83_hC2yFY,180
 phoenix/trace/evaluation_conventions.py,sha256=t8jydM3U0-T5YpiQKRJ3tWdWGlHtzKyttYdw-ddvPOk,1048
 phoenix/trace/exporter.py,sha256=bUXh8fjJIbHurrnt4bAm-cCWqUN5FqNsIc8DZzzklkQ,4695
-phoenix/trace/fixtures.py,sha256=qKYXz33MD7FYkSni7ccyFjXMV5iub7serp1zMGqh1P8,18982
+phoenix/trace/fixtures.py,sha256=fa2yMcjX4sshT88Ec4ce0EpSngfZQ9vXVcssNbcyYAU,19202
 phoenix/trace/otel.py,sha256=c3Iy9sUTYBliwYfChZIQ9_6xTti1BijY7t_4s2tqbyw,10028
 phoenix/trace/projects.py,sha256=2BwlNjFE-uwpqYtCu5YyBiYZk9wRPpM13vh3-Cv7GkA,2157
 phoenix/trace/schemas.py,sha256=HpWSyzec0yDHEQXEDuwyLbhpvKrqkGps8BJqGiIFj8Y,5978
@@ -313,9 +313,9 @@ phoenix/utilities/project.py,sha256=8IJuMM4yUMoooPi37sictGj8Etu9rGmq6RFtc9848cQ,
 phoenix/utilities/re.py,sha256=nr_B0txj_7CXc45953X6vr2KCRSWMuaXJSEkL8s8Sjc,2036
 phoenix/utilities/span_store.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 phoenix/utilities/template_formatters.py,sha256=zWY6gjAMmspYqlmn-ZjKeY5SExylwePk7a1BUdDxVGI,2236
-arize_phoenix-5.3.1.dist-info/METADATA,sha256=8vSIv2Shr5LH4dkS6cshYjQos1ERiCIaQcLc6RZryds,12052
-arize_phoenix-5.3.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-arize_phoenix-5.3.1.dist-info/entry_points.txt,sha256=Pgpn8Upxx9P8z8joPXZWl2LlnAlGc3gcQoVchb06X1Q,94
-arize_phoenix-5.3.1.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
-arize_phoenix-5.3.1.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
-arize_phoenix-5.3.1.dist-info/RECORD,,
+arize_phoenix-5.5.0.dist-info/METADATA,sha256=F32DnF7HFmK_M0NKDsD5mXY91llrxN1Php04_nmzo40,12099
+arize_phoenix-5.5.0.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+arize_phoenix-5.5.0.dist-info/entry_points.txt,sha256=Pgpn8Upxx9P8z8joPXZWl2LlnAlGc3gcQoVchb06X1Q,94
+arize_phoenix-5.5.0.dist-info/licenses/IP_NOTICE,sha256=JBqyyCYYxGDfzQ0TtsQgjts41IJoa-hiwDrBjCb9gHM,469
+arize_phoenix-5.5.0.dist-info/licenses/LICENSE,sha256=HFkW9REuMOkvKRACuwLPT0hRydHb3zNg-fdFt94td18,3794
+arize_phoenix-5.5.0.dist-info/RECORD,,

phoenix/server/api/subscriptions.py

@@ -1,4 +1,5 @@
 import json
+from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import fields
 from datetime import datetime
@@ -9,6 +10,7 @@ from typing import (
     Annotated,
     Any,
     AsyncIterator,
+    Callable,
     DefaultDict,
     Dict,
     Iterable,
@@ -16,6 +18,7 @@ from typing import (
     List,
     Optional,
     Tuple,
+    Type,
     Union,
 )
 
@@ -45,6 +48,7 @@ from phoenix.server.api.input_types.ChatCompletionMessageInput import ChatComple
 from phoenix.server.api.input_types.InvocationParameters import InvocationParameters
 from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
 from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.dml_event import SpanInsertEvent
 from phoenix.trace.attributes import unflatten
 from phoenix.utilities.json import jsonify
@@ -55,14 +59,13 @@ from phoenix.utilities.template_formatters import (
 )
 
 if TYPE_CHECKING:
+    from anthropic.types import MessageParam
     from openai.types import CompletionUsage
-    from openai.types.chat import (
-        ChatCompletionMessageParam,
-    )
+    from openai.types.chat import ChatCompletionMessageParam
 
 PLAYGROUND_PROJECT_NAME = "playground"
 
-ToolCallIndex: TypeAlias = int
+ToolCallID: TypeAlias = str
 
 
 @strawberry.enum
@@ -94,8 +97,14 @@ class ToolCallChunk:
     function: FunctionCallChunk
 
 
-ChatCompletionChunk: TypeAlias = Annotated[
-    Union[TextChunk, ToolCallChunk], strawberry.union("ChatCompletionChunk")
+@strawberry.type
+class FinishedChatCompletion:
+    span: Span
+
+
+ChatCompletionSubscriptionPayload: TypeAlias = Annotated[
+    Union[TextChunk, ToolCallChunk, FinishedChatCompletion],
+    strawberry.union("ChatCompletionSubscriptionPayload"),
 ]
 
 
@@ -120,39 +129,202 @@ class ChatCompletionInput:
     api_key: Optional[str] = strawberry.field(default=None)
 
 
-def to_openai_chat_completion_param(
-    role: ChatCompletionMessageRole, content: JSONScalarType
-) -> "ChatCompletionMessageParam":
-    from openai.types.chat import (
-        ChatCompletionAssistantMessageParam,
-        ChatCompletionSystemMessageParam,
-        ChatCompletionUserMessageParam,
-    )
+PLAYGROUND_STREAMING_CLIENT_REGISTRY: Dict[
+    GenerativeProviderKey, Type["PlaygroundStreamingClient"]
+] = {}
 
-    if role is ChatCompletionMessageRole.USER:
-        return ChatCompletionUserMessageParam(
-            {
-                "content": content,
-                "role": "user",
-            }
-        )
-    if role is ChatCompletionMessageRole.SYSTEM:
-        return ChatCompletionSystemMessageParam(
-            {
-                "content": content,
-                "role": "system",
-            }
+
+def register_llm_client(
+    provider_key: GenerativeProviderKey,
+) -> Callable[[Type["PlaygroundStreamingClient"]], Type["PlaygroundStreamingClient"]]:
+    def decorator(cls: Type["PlaygroundStreamingClient"]) -> Type["PlaygroundStreamingClient"]:
+        PLAYGROUND_STREAMING_CLIENT_REGISTRY[provider_key] = cls
+        return cls
+
+    return decorator
+
+
+class PlaygroundStreamingClient(ABC):
+    def __init__(self, model: GenerativeModelInput, api_key: Optional[str] = None) -> None: ...
+
+    @abstractmethod
+    async def chat_completion_create(
+        self,
+        messages: List[Tuple[ChatCompletionMessageRole, str]],
+        tools: List[JSONScalarType],
+        **invocation_parameters: Any,
+    ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+        # a yield statement is needed to satisfy the type-checker
+        # https://mypy.readthedocs.io/en/stable/more_types.html#asynchronous-iterators
+        yield TextChunk(content="")
+
+    @property
+    @abstractmethod
+    def attributes(self) -> Dict[str, Any]: ...
+
+
+@register_llm_client(GenerativeProviderKey.OPENAI)
+class OpenAIStreamingClient(PlaygroundStreamingClient):
+    def __init__(self, model: GenerativeModelInput, api_key: Optional[str] = None) -> None:
+        from openai import AsyncOpenAI
+
+        self.client = AsyncOpenAI(api_key=api_key)
+        self.model_name = model.name
+        self._attributes: Dict[str, Any] = {}
+
+    async def chat_completion_create(
+        self,
+        messages: List[Tuple[ChatCompletionMessageRole, str]],
+        tools: List[JSONScalarType],
+        **invocation_parameters: Any,
+    ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+        from openai import NOT_GIVEN
+        from openai.types.chat import ChatCompletionStreamOptionsParam
+
+        # Convert standard messages to OpenAI messages
+        openai_messages = [self.to_openai_chat_completion_param(*message) for message in messages]
+        tool_call_ids: Dict[int, str] = {}
+        token_usage: Optional["CompletionUsage"] = None
+        async for chunk in await self.client.chat.completions.create(
+            messages=openai_messages,
+            model=self.model_name,
+            stream=True,
+            stream_options=ChatCompletionStreamOptionsParam(include_usage=True),
+            tools=tools or NOT_GIVEN,
+            **invocation_parameters,
+        ):
+            if (usage := chunk.usage) is not None:
+                token_usage = usage
+                continue
+            choice = chunk.choices[0]
+            delta = choice.delta
+            if choice.finish_reason is None:
+                if isinstance(chunk_content := delta.content, str):
+                    text_chunk = TextChunk(content=chunk_content)
+                    yield text_chunk
+                if (tool_calls := delta.tool_calls) is not None:
+                    for tool_call_index, tool_call in enumerate(tool_calls):
+                        tool_call_id = (
+                            tool_call.id
+                            if tool_call.id is not None
+                            else tool_call_ids[tool_call_index]
+                        )
+                        tool_call_ids[tool_call_index] = tool_call_id
+                        if (function := tool_call.function) is not None:
+                            tool_call_chunk = ToolCallChunk(
+                                id=tool_call_id,
+                                function=FunctionCallChunk(
+                                    name=function.name or "",
+                                    arguments=function.arguments or "",
+                                ),
+                            )
+                            yield tool_call_chunk
+        if token_usage is not None:
+            self._attributes.update(_llm_token_counts(token_usage))
+
+    def to_openai_chat_completion_param(
+        self, role: ChatCompletionMessageRole, content: JSONScalarType
+    ) -> "ChatCompletionMessageParam":
+        from openai.types.chat import (
+            ChatCompletionAssistantMessageParam,
+            ChatCompletionSystemMessageParam,
+            ChatCompletionUserMessageParam,
         )
-    if role is ChatCompletionMessageRole.AI:
-        return ChatCompletionAssistantMessageParam(
-            {
-                "content": content,
-                "role": "assistant",
-            }
+
+        if role is ChatCompletionMessageRole.USER:
+            return ChatCompletionUserMessageParam(
+                {
+                    "content": content,
+                    "role": "user",
+                }
+            )
+        if role is ChatCompletionMessageRole.SYSTEM:
+            return ChatCompletionSystemMessageParam(
+                {
+                    "content": content,
+                    "role": "system",
+                }
+            )
+        if role is ChatCompletionMessageRole.AI:
+            return ChatCompletionAssistantMessageParam(
+                {
+                    "content": content,
+                    "role": "assistant",
+                }
+            )
+        if role is ChatCompletionMessageRole.TOOL:
+            raise NotImplementedError
+        assert_never(role)
+
+    @property
+    def attributes(self) -> Dict[str, Any]:
+        return self._attributes
+
+
+@register_llm_client(GenerativeProviderKey.AZURE_OPENAI)
+class AzureOpenAIStreamingClient(OpenAIStreamingClient):
+    def __init__(self, model: GenerativeModelInput, api_key: Optional[str] = None):
+        from openai import AsyncAzureOpenAI
+
+        if model.endpoint is None or model.api_version is None:
+            raise ValueError("endpoint and api_version are required for Azure OpenAI models")
+        self.client = AsyncAzureOpenAI(
+            api_key=api_key,
+            azure_endpoint=model.endpoint,
+            api_version=model.api_version,
         )
-    if role is ChatCompletionMessageRole.TOOL:
-        raise NotImplementedError
-    assert_never(role)
+
+
+@register_llm_client(GenerativeProviderKey.ANTHROPIC)
+class AnthropicStreamingClient(PlaygroundStreamingClient):
+    def __init__(self, model: GenerativeModelInput, api_key: Optional[str] = None) -> None:
+        import anthropic
+
+        self.client = anthropic.AsyncAnthropic(api_key=api_key)
+        self.model_name = model.name
+
+    async def chat_completion_create(
+        self,
+        messages: List[Tuple[ChatCompletionMessageRole, str]],
+        tools: List[JSONScalarType],
+        **invocation_parameters: Any,
+    ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+        anthropic_messages, system_prompt = self._build_anthropic_messages(messages)
+
+        anthropic_params = {
+            "messages": anthropic_messages,
+            "model": self.model_name,
+            "system": system_prompt,
+            "max_tokens": 1024,
+            **invocation_parameters,
+        }
+
+        async with self.client.messages.stream(**anthropic_params) as stream:
+            async for text in stream.text_stream:
+                yield TextChunk(content=text)
+
+    def _build_anthropic_messages(
+        self, messages: List[Tuple[ChatCompletionMessageRole, str]]
+    ) -> Tuple[List["MessageParam"], str]:
+        anthropic_messages: List["MessageParam"] = []
+        system_prompt = ""
+        for role, content in messages:
+            if role == ChatCompletionMessageRole.USER:
+                anthropic_messages.append({"role": "user", "content": content})
+            elif role == ChatCompletionMessageRole.AI:
+                anthropic_messages.append({"role": "assistant", "content": content})
+            elif role == ChatCompletionMessageRole.SYSTEM:
+                system_prompt += content + "\n"
+            elif role == ChatCompletionMessageRole.TOOL:
+                raise NotImplementedError
+            else:
+                assert_never(role)
+
+        return anthropic_messages, system_prompt
+
+    @property
+    def attributes(self) -> Dict[str, Any]:
+        return dict()
 
 
 @strawberry.type
@@ -160,31 +332,21 @@ class Subscription:
     @strawberry.subscription
     async def chat_completion(
         self, info: Info[Context, None], input: ChatCompletionInput
-    ) -> AsyncIterator[ChatCompletionChunk]:
-        from openai import NOT_GIVEN, AsyncAzureOpenAI, AsyncOpenAI
-        from openai.types.chat import ChatCompletionStreamOptionsParam
+    ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
+        # Determine which LLM client to use based on provider_key
+        provider_key = input.model.provider_key
+        llm_client_class = PLAYGROUND_STREAMING_CLIENT_REGISTRY.get(provider_key)
+        if llm_client_class is None:
+            raise ValueError(f"No LLM client registered for provider '{provider_key}'")
 
-        client: Union[AsyncAzureOpenAI, AsyncOpenAI]
-
-        if input.model.provider_key == GenerativeProviderKey.AZURE_OPENAI:
-            if input.model.endpoint is None or input.model.api_version is None:
-                raise ValueError("endpoint and api_version are required for Azure OpenAI models")
-            client = AsyncAzureOpenAI(
-                api_key=input.api_key,
-                azure_endpoint=input.model.endpoint,
-                api_version=input.model.api_version,
-            )
-        else:
-            client = AsyncOpenAI(api_key=input.api_key)
+        llm_client = llm_client_class(model=input.model, api_key=input.api_key)
 
-        invocation_parameters = jsonify(input.invocation_parameters)
+        messages = [(message.role, message.content) for message in input.messages]
 
-        messages: List[Tuple[ChatCompletionMessageRole, str]] = [
-            (message.role, message.content) for message in input.messages
-        ]
         if template_options := input.template:
             messages = list(_formatted_messages(messages, template_options))
-        openai_messages = [to_openai_chat_completion_param(*message) for message in messages]
+
+        invocation_parameters = jsonify(input.invocation_parameters)
 
         in_memory_span_exporter = InMemorySpanExporter()
         tracer_provider = TracerProvider()
@@ -193,6 +355,7 @@ class Subscription:
         )
         tracer = tracer_provider.get_tracer(__name__)
         span_name = "ChatCompletion"
+
         with tracer.start_span(
             span_name,
             attributes=dict(
@@ -208,52 +371,29 @@ class Subscription:
         ) as span:
             response_chunks = []
             text_chunks: List[TextChunk] = []
-            tool_call_chunks: DefaultDict[ToolCallIndex, List[ToolCallChunk]] = defaultdict(list)
-            role: Optional[str] = None
-            token_usage: Optional[CompletionUsage] = None
-            async for chunk in await client.chat.completions.create(
-                messages=openai_messages,
-                model=input.model.name,
-                stream=True,
-                tools=input.tools or NOT_GIVEN,
-                stream_options=ChatCompletionStreamOptionsParam(include_usage=True),
+            tool_call_chunks: DefaultDict[ToolCallID, List[ToolCallChunk]] = defaultdict(list)
+
+            async for chunk in llm_client.chat_completion_create(
+                messages=messages,
+                tools=input.tools or [],
                 **invocation_parameters,
             ):
                 response_chunks.append(chunk)
-                if (usage := chunk.usage) is not None:
-                    token_usage = usage
-                    continue
-                choice = chunk.choices[0]
-                delta = choice.delta
-                if role is None:
-                    role = delta.role
-                if choice.finish_reason is None:
-                    if isinstance(chunk_content := delta.content, str):
-                        text_chunk = TextChunk(content=chunk_content)
-                        yield text_chunk
-                        text_chunks.append(text_chunk)
-                    if (tool_calls := delta.tool_calls) is not None:
-                        for tool_call_index, tool_call in enumerate(tool_calls):
-                            if (function := tool_call.function) is not None:
-                                if (tool_call_id := tool_call.id) is None:
-                                    first_tool_call_chunk = tool_call_chunks[tool_call_index][0]
-                                    tool_call_id = first_tool_call_chunk.id
-                                tool_call_chunk = ToolCallChunk(
-                                    id=tool_call_id,
-                                    function=FunctionCallChunk(
-                                        name=function.name or "",
-                                        arguments=function.arguments or "",
-                                    ),
-                                )
-                                yield tool_call_chunk
-                                tool_call_chunks[tool_call_index].append(tool_call_chunk)
+                if isinstance(chunk, TextChunk):
+                    yield chunk
+                    text_chunks.append(chunk)
+                elif isinstance(chunk, ToolCallChunk):
+                    yield chunk
+                    tool_call_chunks[chunk.id].append(chunk)
+
             span.set_status(StatusCode.OK)
-            assert role is not None
+            llm_client_attributes = llm_client.attributes
+
             span.set_attributes(
                 dict(
                     chain(
                         _output_value_and_mime_type(response_chunks),
-                        _llm_token_counts(token_usage) if token_usage is not None else [],
+                        llm_client_attributes.items(),
                         _llm_output_messages(text_chunks, tool_call_chunks),
                     )
                 )
@@ -265,8 +405,8 @@ class Subscription:
         assert (attributes := finished_span.attributes) is not None
         start_time = _datetime(epoch_nanoseconds=finished_span.start_time)
         end_time = _datetime(epoch_nanoseconds=finished_span.end_time)
-        prompt_tokens = token_usage.prompt_tokens if token_usage is not None else 0
-        completion_tokens = token_usage.completion_tokens if token_usage is not None else 0
+        prompt_tokens = llm_client_attributes.get(LLM_TOKEN_COUNT_PROMPT, 0)
+        completion_tokens = llm_client_attributes.get(LLM_TOKEN_COUNT_COMPLETION, 0)
         trace_id = _hex(finished_span.context.trace_id)
         span_id = _hex(finished_span.context.span_id)
         status = finished_span.status
@@ -284,36 +424,35 @@ class Subscription:
                         description="Traces from prompt playground",
                     )
                 )
-            trace_rowid = await session.scalar(
-                insert(models.Trace)
-                .returning(models.Trace.id)
-                .values(
-                    project_rowid=playground_project_id,
-                    trace_id=trace_id,
-                    start_time=start_time,
-                    end_time=end_time,
-                )
+            playground_trace = models.Trace(
+                project_rowid=playground_project_id,
+                trace_id=trace_id,
+                start_time=start_time,
+                end_time=end_time,
             )
-            await session.execute(
-                insert(models.Span).values(
-                    trace_rowid=trace_rowid,
-                    span_id=span_id,
-                    parent_id=None,
-                    name=span_name,
-                    span_kind=LLM,
-                    start_time=start_time,
-                    end_time=end_time,
-                    attributes=unflatten(attributes.items()),
-                    events=finished_span.events,
-                    status_code=status.status_code.name,
-                    status_message=status.description or "",
-                    cumulative_error_count=int(not status.is_ok),
-                    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,
-                )
+            playground_span = models.Span(
+                trace_rowid=playground_trace.id,
+                span_id=span_id,
+                parent_id=None,
+                name=span_name,
+                span_kind=LLM,
+                start_time=start_time,
+                end_time=end_time,
+                attributes=unflatten(attributes.items()),
+                events=finished_span.events,
+                status_code=status.status_code.name,
+                status_message=status.description or "",
+                cumulative_error_count=int(not status.is_ok),
+                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=playground_trace,
             )
+            session.add(playground_trace)
+            session.add(playground_span)
+            await session.flush()
+            yield FinishedChatCompletion(span=to_gql_span(playground_span))
         info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
 
 
@@ -361,7 +500,7 @@ def _llm_input_messages(
 
 def _llm_output_messages(
     text_chunks: List[TextChunk],
-    tool_call_chunks: DefaultDict[ToolCallIndex, List[ToolCallChunk]],
+    tool_call_chunks: DefaultDict[ToolCallID, List[ToolCallChunk]],
 ) -> Iterator[Tuple[str, Any]]:
     yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"
     if content := "".join(chunk.content for chunk in text_chunks):