Support overriding default model provider in instrument_openai

logfire/_internal/integrations/llm_providers/llm_provider.py

@@ -26,6 +26,7 @@ def instrument_llm_provider(
     get_endpoint_config_fn: Callable[[Any], EndpointConfig],
     on_response_fn: Callable[[Any, LogfireSpan], Any],
     is_async_client_fn: Callable[[type[Any]], bool],
+    model_provider: str | None = None,
 ) -> AbstractContextManager[None]:
     """Instruments the provided `client` (or clients) with `logfire`.
 
@@ -93,6 +94,8 @@ def _instrumentation_setup(*args: Any, **kwargs: Any) -> Any:
                 return None, None, kwargs
 
             span_data['async'] = is_async
+            if model_provider is not None:
+                span_data['overridden_model_provider'] = model_provider
 
             if kwargs.get('stream') and stream_state_cls:
                 stream_cls = kwargs['stream_cls']

logfire/_internal/integrations/llm_providers/openai.py

@@ -183,7 +183,9 @@ def on_response(response: ResponseT, span: LogfireSpan) -> ResponseT:
         on_response(response.parse(), span)  # type: ignore
         return cast('ResponseT', response)
 
-    span.set_attribute('gen_ai.system', 'openai')
+    model_provider: str = cast(str, (span.attributes or {}).get('overridden_model_provider', "openai"))
+
+    span.set_attribute('gen_ai.system', model_provider)
 
     if isinstance(response_model := getattr(response, 'model', None), str):
         span.set_attribute('gen_ai.response.model', response_model)
@@ -194,12 +196,12 @@ def on_response(response: ResponseT, span: LogfireSpan) -> ResponseT:
             response_data = response.model_dump()  # type: ignore
             usage_data = extract_usage(
                 response_data,
-                provider_id='openai',
+                provider_id=model_provider,
                 api_flavor='responses' if isinstance(response, Response) else 'chat',
             )
             span.set_attribute(
                 'operation.cost',
-                float(calc_price(usage_data.usage, model_ref=response_model, provider_id='openai').total_price),
+                float(calc_price(usage_data.usage, model_ref=response_model, provider_id=model_provider).total_price),
             )
         except Exception:
             pass

logfire/_internal/main.py

@@ -1171,6 +1171,7 @@ def instrument_openai(
         | None = None,
         *,
         suppress_other_instrumentation: bool = True,
+        model_provider: str = "openai",
     ) -> AbstractContextManager[None]:
         """Instrument an OpenAI client so that spans are automatically created for each request.
 
@@ -1237,6 +1238,7 @@ def instrument_openai(
             get_endpoint_config,
             on_response,
             is_async_client,
+            model_provider,
         )
 
     def instrument_openai_agents(self) -> None: