diff --git a/src/ra/clients/openai.py b/src/ra/clients/openai.py index 2f700f42..f78fefcf 100644 --- a/src/ra/clients/openai.py +++ b/src/ra/clients/openai.py @@ -6,6 +6,15 @@ from dotenv import load_dotenv from openai.types.chat import ChatCompletion +# Use Langfuse-wrapped OpenAI client when available for automatic tracing +_LangfuseOpenAI = None +_LangfuseAsyncOpenAI = None +if os.environ.get("LANGFUSE_SECRET_KEY"): + try: + from langfuse.openai import OpenAI as _LangfuseOpenAI, AsyncOpenAI as _LangfuseAsyncOpenAI + except ImportError: + pass + from ra.clients.base_lm import BaseLM from ra.core.types import ModelUsageSummary, UsageSummary @@ -51,7 +60,8 @@ def __init__( extra_headers = ( _OPENROUTER_HEADERS if base_url == "https://openrouter.ai/api/v1" else None ) - self.client = openai.OpenAI( + _ClientCls = _LangfuseOpenAI or openai.OpenAI + self.client = _ClientCls( api_key=api_key, base_url=base_url, default_headers=extra_headers, @@ -121,7 +131,8 @@ async def acompletion( if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL: extra_body["usage"] = {"include": True} - async with openai.AsyncOpenAI( + _AsyncCls = _LangfuseAsyncOpenAI or openai.AsyncOpenAI + async with _AsyncCls( **self._async_client_kwargs, ) as client: response = await client.chat.completions.create(