cavisson-pythonagent 0.0.1__py3-none-any.whl

pythonagent/agent/probes/openai/openai.py

@@ -0,0 +1,797 @@
+from ..base import BaseInterceptor, ExitCallInterceptor
+from ..span import LLMSpan, EmbeddingSpan
+
+from pythonagent.utils import get_validated_duration, get_current_timestamp_in_us, build_host_name_str, split_fp_topo
+import logging
+
+"""
+Instrumentation wrappers for the OpenAI Python SDK (v1.x).
+
+Covers:
+  - Sync chat completions  (non-streaming + streaming)
+  - Async chat completions (non-streaming + streaming)
+  - Legacy (non-chat) completions, sync + async
+  - Embeddings, sync + async
+"""
+
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+def _safe_usage_metrics(usage):
+    """Build a metrics dict from a usage object, guarding every field for None."""
+    if usage is None:
+        return {
+            "prompt_tokens": None,
+            "completion_tokens": None,
+            "total_tokens": None,
+            "prompt_tokens_details": {"prompt_audio_tokens": None, "prompt_cached_tokens": None},
+            "completion_tokens_details": {
+                "completion_accepted_prediction_tokens": None,
+                "completion_rejected_prediction_tokens": None,
+                "completion_audio_tokens": None,
+                "completion_reasoning_tokens": None,
+            },
+        }
+    ptd = getattr(usage, "prompt_tokens_details", None)
+    ctd = getattr(usage, "completion_tokens_details", None)
+    return {
+        "prompt_tokens": getattr(usage, "prompt_tokens", None),
+        "completion_tokens": getattr(usage, "completion_tokens", None),
+        "total_tokens": getattr(usage, "total_tokens", None),
+        "prompt_tokens_details": {
+            "prompt_audio_tokens": getattr(ptd, "audio_tokens", None) if ptd else None,
+            "prompt_cached_tokens": getattr(ptd, "cached_tokens", None) if ptd else None,
+        },
+        "completion_tokens_details": {
+            "completion_accepted_prediction_tokens": getattr(ctd, "accepted_prediction_tokens", None) if ctd else None,
+            "completion_rejected_prediction_tokens": getattr(ctd, "rejected_prediction_tokens", None) if ctd else None,
+            "completion_audio_tokens": getattr(ctd, "audio_tokens", None) if ctd else None,
+            "completion_reasoning_tokens": getattr(ctd, "reasoning_tokens", None) if ctd else None,
+        },
+    }
+
+
+def _get_messages_input(args, kwargs):
+    """Return messages list, handling both positional and keyword argument (Gap 5)."""
+    return kwargs.get("messages") or (args[0] if args else [])
+
+
+def _get_model_name(args, kwargs):
+    """Best-effort extraction of the requested model name before the call is made
+    (ported from the old single-entry-point file's _get_model_name)."""
+    model_name = kwargs.get("model")
+    if not model_name and args:
+        model_name = getattr(args[0], "model", None) or getattr(args[0], "model_name", None)
+    return model_name or "default"
+
+
+def _build_chat_output(result):
+    """
+    Extract output from all choices of a ChatCompletion response.
+    Returns a single dict for n=1, a list for n>1 (Gap 8).
+    """
+    choices = getattr(result, "choices", None) or []
+    outputs = []
+    for choice in choices:
+        msg = getattr(choice, "message", None)
+        if msg is None:
+            continue
+        outputs.append({
+            "content": getattr(msg, "content", None),
+            "role": getattr(msg, "role", None),
+            "tool_calls": str(getattr(msg, "tool_calls", None)),
+            "completion_id": getattr(result, "id", None),
+            "finish_reason": getattr(choice, "finish_reason", None),
+            "refusal": getattr(msg, "refusal", None),
+            "annotations": getattr(msg, "annotations", None),
+            "audio": getattr(msg, "audio", None),
+            "function_call": str(getattr(msg, "function_call", None)),
+        })
+    if len(outputs) == 1:
+        return outputs[0]
+    return outputs if outputs else None
+
+
+def _accumulate_stream_chunks(chunks):
+    """
+    Merge ChatCompletionChunk list into aggregated output fields.
+    Returns (content_by_index, model_name, completion_id, finish_reason, usage).
+    """
+    content_by_index = {}
+    model_name = "default"
+    completion_id = None
+    finish_reason = None
+    usage = None
+
+    for chunk in chunks:
+        model_name = getattr(chunk, "model", None) or model_name
+        completion_id = getattr(chunk, "id", None) or completion_id
+        # usage arrives on the last chunk when stream_options={"include_usage": True}
+        chunk_usage = getattr(chunk, "usage", None)
+        if chunk_usage is not None:
+            usage = chunk_usage
+        for choice in (getattr(chunk, "choices", None) or []):
+            idx = getattr(choice, "index", 0)
+            delta = getattr(choice, "delta", None)
+            if delta:
+                text = getattr(delta, "content", None)
+                if text:
+                    content_by_index[idx] = content_by_index.get(idx, "") + text
+            fr = getattr(choice, "finish_reason", None)
+            if fr:
+                finish_reason = fr
+
+    return content_by_index, model_name, completion_id, finish_reason, usage
+
+
+def _build_stream_output(content_by_index, completion_id, finish_reason):
+    """Build output dict(s) from accumulated streaming content."""
+    indices = sorted(content_by_index.keys())
+    if not indices:
+        return {"content": None, "completion_id": completion_id, "finish_reason": finish_reason,
+                "role": "assistant"}
+    outputs = [
+        {
+            "content": content_by_index.get(i),
+            "completion_id": completion_id,
+            "finish_reason": finish_reason,
+            "role": "assistant",
+        }
+        for i in indices
+    ]
+    return outputs[0] if len(outputs) == 1 else outputs
+
+
+def _split_fp(context):
+    """Safely split fp_topo from context; returns 7-tuple of Nones on failure."""
+    try:
+        if context:
+            return split_fp_topo(context.fp_topo)
+    except Exception:
+        pass
+    return None, None, None, None, None, None, None
+
+
+def _record_llm_span(interceptor, bt, fqm, start_time, duration, context,
+                     args, kwargs, model_name, output, usage):
+    """Create and send an LLMSpan if nf_enabled."""
+    if not interceptor.agent.nf_enabled:
+        return
+    metrics = _safe_usage_metrics(usage)
+    test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = _split_fp(context)
+    meta = {
+        "span_kind": "llm",
+        "input": _get_messages_input(args, kwargs),
+        "output": output,
+        "model_name": model_name,
+        "model_provider": "OpenAI",
+        "metadata": kwargs.get("metadata", None),
+    }
+    llm_span = LLMSpan(
+        fqm, start_time, duration, "ok", meta, metrics,
+        test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp
+    )
+    if interceptor.agent.logger.isEnabledFor(logging.INFO):
+        interceptor.agent.logger.info("openai llm sending llm span output={}".format(output))
+    interceptor.agent.dump_log_message(bt, llm_span.create_span_json())
+
+
+def _record_call_success(interceptor, bt, exit_call, start_time, fqm, context,
+                          username, args, kwargs, model_name, output, usage, label="OpenAI LLM"):
+    """Compute duration, call db_call_end, and send LLMSpan."""
+    end_time = get_current_timestamp_in_us()
+    duration = get_validated_duration(start_time, end_time, label)
+    host_name_str = build_host_name_str("NA", "NA", "NA", model_name, username)
+    if context:
+        context.host_str = host_name_str
+        context.status_code = 200
+        context.db_callout_time = 0
+    if interceptor.agent.logger.isEnabledFor(logging.INFO):
+        interceptor.agent.logger.info(
+            "openai llm db call end bt={} exit_call={} duration={} fqm={} model={}".format(
+                bt, exit_call, duration, fqm, model_name
+            )
+        )
+    interceptor.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+    _record_llm_span(interceptor, bt, fqm, start_time, duration, context,
+                     args, kwargs, model_name, output, usage)
+
+
+def _record_call_error(interceptor, bt, exit_call, start_time, fqm, context, username,
+                        label="OpenAI LLM"):
+    """Call db_call_end on the exception path with status 500."""
+    end_time = get_current_timestamp_in_us()
+    duration = get_validated_duration(start_time, end_time, label)
+    host_name_str = build_host_name_str("NA", "NA", "LLM", "OpenAI_default", username)
+    if context:
+        context.host_str = host_name_str
+        context.status_code = 500
+        context.db_callout_time = 0
+    if interceptor.agent.logger.isEnabledFor(logging.INFO):
+        interceptor.agent.logger.info(
+            "openai llm db call end (error) bt={} exit_call={} fqm={}".format(bt, exit_call, fqm)
+        )
+    interceptor.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+
+
+# ---------------------------------------------------------------------------
+# Stream wrappers
+# ---------------------------------------------------------------------------
+
+class SyncStreamWrapper:
+    """
+    Transparent wrapper around openai.Stream[ChatCompletionChunk].
+    Accumulates chunks and fires on_finish after the last chunk is consumed
+    or the iterator is closed early (via finally). (Gap 2)
+    """
+    def __init__(self, stream, on_finish):
+        self._stream = stream
+        self._on_finish = on_finish
+        self._chunks = []
+        self._finished = False
+
+    def _fire_finish(self):
+        if not self._finished:
+            self._finished = True
+            self._on_finish(self._chunks)
+
+    def __iter__(self):
+        try:
+            for chunk in self._stream:
+                self._chunks.append(chunk)
+                yield chunk
+        finally:
+            self._fire_finish()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self._fire_finish()
+        aexit = getattr(self._stream, "__exit__", None)
+        if aexit:
+            return aexit(*args)
+
+    def __getattr__(self, name):
+        return getattr(self._stream, name)
+
+
+class AsyncStreamWrapper:
+    """
+    Transparent wrapper around openai.AsyncStream[ChatCompletionChunk].
+    Fires on_finish (synchronous SDK call) after the last chunk or early close. (Gap 3)
+    """
+    def __init__(self, stream, on_finish):
+        self._stream = stream
+        self._on_finish = on_finish
+        self._chunks = []
+        self._finished = False
+        self._iterator = None
+
+    def _fire_finish(self):
+        if not self._finished:
+            self._finished = True
+            self._on_finish(self._chunks)
+
+    def __aiter__(self):
+        self._iterator = self._stream.__aiter__()
+        return self
+
+    async def __anext__(self):
+        try:
+            chunk = await self._iterator.__anext__()
+            self._chunks.append(chunk)
+            return chunk
+        except StopAsyncIteration:
+            self._fire_finish()
+            raise
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        self._fire_finish()
+        aexit = getattr(self._stream, "__aexit__", None)
+        if aexit:
+            await aexit(*args)
+
+    def __getattr__(self, name):
+        return getattr(self._stream, name)
+
+
+# ---------------------------------------------------------------------------
+# Internal: build streaming on_finish closure
+# ---------------------------------------------------------------------------
+
+def _make_stream_on_finish(interceptor, bt, exit_call, start_time, fqm,
+                            context, username, args, kwargs):
+    """Return a closure suitable for SyncStreamWrapper / AsyncStreamWrapper."""
+    def on_finish(chunks):
+        try:
+            content_by_idx, model_name, completion_id, finish_reason, usage = \
+                _accumulate_stream_chunks(chunks)
+            output = _build_stream_output(content_by_idx, completion_id, finish_reason)
+            _record_call_success(
+                interceptor, bt, exit_call, start_time, fqm, context,
+                username, args, kwargs, model_name, output, usage
+            )
+        except Exception as e:
+            interceptor.agent.logger.exception(
+                "openai stream on_finish error: {}".format(e)
+            )
+    return on_finish
+
+
+# ---------------------------------------------------------------------------
+# Chat Completions — Sync (existing, with all gap fixes applied)
+# ---------------------------------------------------------------------------
+
+class OpenAIInterceptor(ExitCallInterceptor):
+    def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "LLM", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info(
+                    "openai llm db call begin bt={} start_time={}".format(bt, start_time)
+                )
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai llm db call begin: {}".format(e))
+
+        try:
+            result = create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai llm create error: {}".format(e))
+            try:
+                _record_call_error(self, bt, exit_call, start_time, fqm, context, username)
+            except Exception as inner:
+                self.agent.logger.exception("openai llm db call end (error): {}".format(inner))
+            raise
+
+        # Streaming path — wrap stream, defer telemetry until consumed (Gap 2)
+        if kwargs.get("stream", False):
+            return SyncStreamWrapper(
+                result,
+                _make_stream_on_finish(self, bt, exit_call, start_time, fqm,
+                                       context, username, args, kwargs)
+            )
+
+        # Non-streaming path
+        try:
+            model_name = getattr(result, "model", None) or getattr(result, "model_name", None) or "default"
+            output = _build_chat_output(result)
+            usage = getattr(result, "usage", None)
+            _record_call_success(
+                self, bt, exit_call, start_time, fqm, context,
+                username, args, kwargs, model_name, output, usage
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai llm db call end: {}".format(e))
+
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Chat Completions — Async (Gap 1 + Gap 3)
+# ---------------------------------------------------------------------------
+
+class AsyncOpenAIInterceptor(ExitCallInterceptor):
+    async def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "LLM", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info(
+                    "openai async llm db call begin bt={} start_time={}".format(bt, start_time)
+                )
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai async llm db call begin: {}".format(e))
+
+        try:
+            result = await create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai async llm create error: {}".format(e))
+            try:
+                _record_call_error(self, bt, exit_call, start_time, fqm, context, username)
+            except Exception as inner:
+                self.agent.logger.exception("openai async llm db call end (error): {}".format(inner))
+            raise
+
+        # Async streaming path (Gap 3)
+        if kwargs.get("stream", False):
+            return AsyncStreamWrapper(
+                result,
+                _make_stream_on_finish(self, bt, exit_call, start_time, fqm,
+                                       context, username, args, kwargs)
+            )
+
+        # Non-streaming path
+        try:
+            model_name = getattr(result, "model", None) or getattr(result, "model_name", None) or "default"
+            output = _build_chat_output(result)
+            usage = getattr(result, "usage", None)
+            _record_call_success(
+                self, bt, exit_call, start_time, fqm, context,
+                username, args, kwargs, model_name, output, usage
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai async llm db call end: {}".format(e))
+
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Legacy (non-chat) Completions — Sync (Gap 6)
+# ---------------------------------------------------------------------------
+
+def _build_legacy_output(result):
+    """Extract output from legacy Completion choices (result.choices[n].text)."""
+    choices = getattr(result, "choices", None) or []
+    outputs = [
+        {
+            "text": getattr(c, "text", None),
+            "finish_reason": getattr(c, "finish_reason", None),
+            "completion_id": getattr(result, "id", None),
+        }
+        for c in choices
+    ]
+    if len(outputs) == 1:
+        return outputs[0]
+    return outputs if outputs else None
+
+
+class OpenAILegacyCompletionsInterceptor(ExitCallInterceptor):
+    def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "LLM", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai legacy completions db call begin: {}".format(e))
+
+        try:
+            result = create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai legacy completions create error: {}".format(e))
+            try:
+                _record_call_error(self, bt, exit_call, start_time, fqm, context, username,
+                                   label="OpenAI Legacy Completions")
+            except Exception as inner:
+                self.agent.logger.exception("openai legacy completions db call end (error): {}".format(inner))
+            raise
+
+        if kwargs.get("stream", False):
+            return SyncStreamWrapper(
+                result,
+                _make_stream_on_finish(self, bt, exit_call, start_time, fqm,
+                                       context, username, args, kwargs)
+            )
+
+        try:
+            model_name = getattr(result, "model", None) or "default"
+            output = _build_legacy_output(result)
+            usage = getattr(result, "usage", None)
+            _record_call_success(
+                self, bt, exit_call, start_time, fqm, context,
+                username, args, kwargs, model_name, output, usage,
+                label="OpenAI Legacy Completions"
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai legacy completions db call end: {}".format(e))
+
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Legacy Completions — Async (Gap 6)
+# ---------------------------------------------------------------------------
+
+class AsyncOpenAILegacyCompletionsInterceptor(ExitCallInterceptor):
+    async def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "LLM", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai async legacy completions db call begin: {}".format(e))
+
+        try:
+            result = await create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai async legacy completions create error: {}".format(e))
+            try:
+                _record_call_error(self, bt, exit_call, start_time, fqm, context, username,
+                                   label="OpenAI Legacy Completions")
+            except Exception as inner:
+                self.agent.logger.exception("openai async legacy completions db call end (error): {}".format(inner))
+            raise
+
+        if kwargs.get("stream", False):
+            return AsyncStreamWrapper(
+                result,
+                _make_stream_on_finish(self, bt, exit_call, start_time, fqm,
+                                       context, username, args, kwargs)
+            )
+
+        try:
+            model_name = getattr(result, "model", None) or "default"
+            output = _build_legacy_output(result)
+            usage = getattr(result, "usage", None)
+            _record_call_success(
+                self, bt, exit_call, start_time, fqm, context,
+                username, args, kwargs, model_name, output, usage,
+                label="OpenAI Legacy Completions"
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai async legacy completions db call end: {}".format(e))
+
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Embeddings — Sync + Async (Gap 7)
+# ---------------------------------------------------------------------------
+
+def _build_embedding_output(result):
+    """Build output list from EmbeddingResponse.data."""
+    data = getattr(result, "data", None) or []
+    vectors = [
+        {
+            "index": getattr(d, "index", i),
+            "embedding_length": len(getattr(d, "embedding", None) or []),
+            "embedding_sample": (getattr(d, "embedding", None) or [])[:5],
+        }
+        for i, d in enumerate(data)
+    ]
+    return vectors[0] if len(vectors) == 1 else (vectors or None)
+
+
+def _record_embedding_span(interceptor, bt, exit_call, start_time, fqm,
+                            context, username, kwargs, model_name, output, usage):
+    """db_call_end + EmbeddingSpan."""
+    end_time = get_current_timestamp_in_us()
+    duration = get_validated_duration(start_time, end_time, "OpenAI Embeddings")
+    host_name_str = build_host_name_str("NA", "NA", "NA", model_name, username)
+    if context:
+        context.host_str = host_name_str
+        context.status_code = 200
+        context.db_callout_time = 0
+    interceptor.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+
+    if not interceptor.agent.nf_enabled:
+        return
+    test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = _split_fp(context)
+    meta = {
+        "span_kind": "embedding",
+        "input": kwargs.get("input") or kwargs.get("inputs"),
+        "output": output,
+        "model_name": model_name,
+        "model_provider": "OpenAI",
+        "metadata": kwargs.get("metadata", None),
+    }
+    embedding_span = EmbeddingSpan(
+        fqm, start_time, duration, "ok", meta,
+        test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp
+    )
+    # Attach token metrics (EmbeddingSpan inherits self.metrics = {} from BaseSpan)
+    embedding_span.metrics = _safe_usage_metrics(usage)
+    if interceptor.agent.logger.isEnabledFor(logging.INFO):
+        interceptor.agent.logger.info(
+            "openai embeddings sending span output={}".format(output)
+        )
+    interceptor.agent.dump_log_message(bt, embedding_span.create_span_json())
+
+
+class OpenAIEmbeddingsInterceptor(ExitCallInterceptor):
+    def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "Embedding", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai embeddings db call begin: {}".format(e))
+
+        try:
+            result = create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai embeddings create error: {}".format(e))
+            try:
+                end_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, end_time, "OpenAI Embeddings")
+                host_name_str = build_host_name_str("NA", "NA", "Embedding", "OpenAI_default", username)
+                if context:
+                    context.host_str = host_name_str
+                    context.status_code = 500
+                self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+            except Exception as inner:
+                self.agent.logger.exception("openai embeddings db call end (error): {}".format(inner))
+            raise
+
+        try:
+            model_name = getattr(result, "model", None) or "default"
+            output = _build_embedding_output(result)
+            usage = getattr(result, "usage", None)
+            _record_embedding_span(
+                self, bt, exit_call, start_time, fqm, context,
+                username, kwargs, model_name, output, usage
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai embeddings db call end: {}".format(e))
+
+        return result
+
+
+class AsyncOpenAIEmbeddingsInterceptor(ExitCallInterceptor):
+    async def wrap_create(self, create, *args, **kwargs):
+        start_time, exit_call, bt, fqm, context, username = 0, 0, None, None, None, None
+        try:
+            bt = self.bt
+            username = "NA"
+            host, port, query, query_params = "NA", "NA", "NA", None
+            model_name = _get_model_name(args, kwargs)
+            host_name_str = build_host_name_str(host, port, "Embedding", "OpenAI_" + model_name, username)
+            fqm = "{}.{}".format(create.__module__, create.__qualname__)
+            context = self.agent.get_transaction_context()
+            if context:
+                context.entry_point_fqm = fqm
+                context.host_str = host_name_str
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time, fqm, host_name_str)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("openai async embeddings db call begin: {}".format(e))
+
+        try:
+            result = await create(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("openai async embeddings create error: {}".format(e))
+            try:
+                end_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, end_time, "OpenAI Embeddings")
+                host_name_str = build_host_name_str("NA", "NA", "Embedding", "OpenAI_default", username)
+                if context:
+                    context.host_str = host_name_str
+                    context.status_code = 500
+                self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+            except Exception as inner:
+                self.agent.logger.exception("openai async embeddings db call end (error): {}".format(inner))
+            raise
+
+        try:
+            model_name = getattr(result, "model", None) or "default"
+            output = _build_embedding_output(result)
+            usage = getattr(result, "usage", None)
+            _record_embedding_span(
+                self, bt, exit_call, start_time, fqm, context,
+                username, kwargs, model_name, output, usage
+            )
+        except Exception as e:
+            self.agent.logger.exception("openai async embeddings db call end: {}".format(e))
+
+        return result
+
+
+# ---------------------------------------------------------------------------
+# Registration
+# ---------------------------------------------------------------------------
+
+def intercept_openai(agent, mod):
+    """
+    Patch all relevant OpenAI SDK entry points.
+    Each registration is wrapped independently so a missing class on an older
+    SDK version does not prevent the others from being instrumented.
+    """
+    agent.logger.warning("Instrument module: openai{}".format(", mod: {}".format(mod) if mod else ""))
+    # 1. Sync chat completions (fixed: null-safety, positional messages, streaming, all choices)
+    try:
+        OpenAIInterceptor(agent, mod.resources.chat.completions.Completions) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai sync chat completions intercept failed: {}".format(e))
+
+    # 2. Async chat completions (Gap 1 + Gap 3)
+    try:
+        AsyncOpenAIInterceptor(agent, mod.resources.chat.completions.AsyncCompletions) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai async chat completions intercept failed: {}".format(e))
+
+    # 3. Legacy completions — sync (Gap 6)
+    try:
+        OpenAILegacyCompletionsInterceptor(agent, mod.resources.completions.Completions) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai legacy sync completions intercept failed: {}".format(e))
+
+    # 4. Legacy completions — async (Gap 6)
+    try:
+        AsyncOpenAILegacyCompletionsInterceptor(agent, mod.resources.completions.AsyncCompletions) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai legacy async completions intercept failed: {}".format(e))
+
+    # 5. Embeddings — sync (Gap 7)
+    try:
+        OpenAIEmbeddingsInterceptor(agent, mod.resources.embeddings.Embeddings) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai sync embeddings intercept failed: {}".format(e))
+    return OpenAIInterceptor(agent, mod.resources.chat.completions.Completions).attach('create', patched_method_name='wrap_create')
+
+    # 6. Embeddings — async (Gap 7)
+    try:
+        AsyncOpenAIEmbeddingsInterceptor(agent, mod.resources.embeddings.AsyncEmbeddings) \
+            .attach('create', patched_method_name='wrap_create')
+    except Exception as e:
+        agent.logger.exception("openai async embeddings intercept failed: {}".format(e))