fenra 0.1.0__py3-none-any.whl

fenra/integrations/openai/__init__.py

@@ -0,0 +1,904 @@
+"""OpenAI integration with auto-instrumentation via monkey patching."""
+
+import functools
+import logging
+from typing import Any
+
+from fenra._context import get_context
+from fenra._core import enqueue_transaction
+
+logger = logging.getLogger(__name__)
+
+_patched_chat_sync = False
+_patched_chat_async = False
+_patched_responses_sync = False
+_patched_responses_async = False
+_patched_images_sync = False
+_patched_images_async = False
+_patched_beta_parse_sync = False
+_patched_beta_parse_async = False
+
+# Store original methods for unpatching
+_original_chat_create: Any = None
+_original_chat_create_async: Any = None
+_original_responses_create: Any = None
+_original_responses_create_async: Any = None
+_original_images_generate: Any = None
+_original_images_generate_async: Any = None
+_original_beta_parse: Any = None
+_original_beta_parse_async: Any = None
+
+def _is_openai_endpoint(client: Any) -> bool:
+    """
+    Check if client is actually hitting OpenAI, not other OpenAI-compatible endpoints.
+    """
+
+    base_url = str(getattr(client, "base_url", ""))
+    # OpenAI domains
+    if "api.openai.com" in base_url:
+        return True
+    # Azure OpenAI
+    if "openai.azure.com" in base_url:
+        return True
+    if "generativelanguage.googleapis.com" in base_url:
+        return False
+    # Default to False for custom models that use OpenAI-compatible endpoints
+    return False
+
+
+def _track_chat_completion_response(
+    response: Any,
+    model: str,
+    context: dict[str, Any],
+) -> None:
+    """Extract usage from Chat Completions response and queue transaction with full context."""
+    if not hasattr(response, "usage") or response.usage is None:
+        return
+
+    usage = response.usage
+
+    metrics: dict[str, Any] = {
+        "input_tokens": getattr(usage, "prompt_tokens", 0),
+        "output_tokens": getattr(usage, "completion_tokens", 0),
+        "total_tokens": getattr(usage, "total_tokens", 0),
+    }
+
+    # Include reasoning tokens if present (o1/o3 models)
+    if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
+        details = usage.completion_tokens_details
+        if hasattr(details, "reasoning_tokens") and details.reasoning_tokens:
+            metrics["reasoning_tokens"] = details.reasoning_tokens
+
+    # Include cached tokens if present
+    if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
+        details = usage.prompt_tokens_details
+        if hasattr(details, "cached_tokens") and details.cached_tokens:
+            metrics["cached_tokens"] = details.cached_tokens
+
+    transaction = {
+        "provider": "openai",
+        "model": model,
+        "usage": [
+            {
+                "type": "tokens",
+                "metrics": metrics,
+            }
+        ],
+        "context": context,  # Full merged context from set_context() calls
+    }
+
+    try:
+        transaction["provider_usage_raw"] = {
+            "prompt_tokens": getattr(usage, "prompt_tokens", None),
+            "completion_tokens": getattr(usage, "completion_tokens", None),
+            "total_tokens": getattr(usage, "total_tokens", None),
+        }
+    except Exception:
+        pass
+
+    enqueue_transaction(transaction)
+
+
+def _track_responses_api_response(
+    response: Any,
+    model: str,
+    context: dict[str, Any],
+    tools: list[dict[str, Any]] | None = None,
+) -> None:
+    """Extract usage from Responses API response and queue transaction with full context."""
+    if not hasattr(response, "usage") or response.usage is None:
+        return
+
+    usage = response.usage
+
+    # Build token metrics (Responses API uses different field names)
+    metrics: dict[str, Any] = {
+        "input_tokens": getattr(usage, "input_tokens", 0),
+        "output_tokens": getattr(usage, "output_tokens", 0),
+        "total_tokens": getattr(usage, "total_tokens", 0),
+    }
+
+    # Include reasoning tokens if present (o1/o3 models)
+    if hasattr(usage, "output_tokens_details") and usage.output_tokens_details:
+        details = usage.output_tokens_details
+        if hasattr(details, "reasoning_tokens") and details.reasoning_tokens:
+            metrics["reasoning_tokens"] = details.reasoning_tokens
+
+    # Include cached tokens if present
+    if hasattr(usage, "input_tokens_details") and usage.input_tokens_details:
+        details = usage.input_tokens_details
+        if hasattr(details, "cached_tokens") and details.cached_tokens:
+            metrics["cached_tokens"] = details.cached_tokens
+
+    usage_entries: list[dict[str, Any]] = [
+        {
+            "type": "tokens",
+            "metrics": metrics,
+        }
+    ]
+
+    # Add requests usage when web_search tool is used
+    if tools and any(t.get("type") in ("web_search", "web_search_preview") for t in tools):
+        usage_entries.append(
+            {
+                "type": "requests",
+                "metrics": {"count": 1, "request_type": "web_search"},
+            }
+        )
+
+    transaction: dict[str, Any] = {
+        "provider": "openai",
+        "model": model,
+        "usage": usage_entries,
+        "context": context,
+    }
+
+    # Include raw usage for debugging if available
+    try:
+        transaction["provider_usage_raw"] = {
+            "input_tokens": getattr(usage, "input_tokens", None),
+            "output_tokens": getattr(usage, "output_tokens", None),
+            "total_tokens": getattr(usage, "total_tokens", None),
+        }
+    except Exception:
+        pass
+
+    enqueue_transaction(transaction)
+
+
+def _track_images_response(
+    model: str,
+    n: int,
+    size: str,
+    quality: str,
+    context: dict[str, Any],
+    response: Any = None,
+) -> None:
+    """Track image generation usage and queue transaction with full context.
+
+    Prefers quality/size/n from the API response when available (reflects actual
+    usage); otherwise uses request kwargs.
+    """
+    # Prefer response when available (API returns what was actually used)
+    if response is not None:
+        quality = getattr(response, "quality", None) or quality
+        size = getattr(response, "size", None) or size
+        if hasattr(response, "data") and response.data is not None:
+            n = len(response.data)
+
+    # Parse size to get pixels (e.g., "1024x1024" -> 1024)
+    size_px: int | None = None
+    if size and "x" in str(size):
+        try:
+            size_px = int(str(size).split("x")[0])
+        except (ValueError, TypeError):
+            pass
+
+    metrics: dict[str, Any] = {
+        "generated": n,
+        "size_px": size_px,
+        "quality": quality or None,
+    }
+    # Omit None values for API payload
+    metrics = {k: v for k, v in metrics.items() if v is not None}
+
+    transaction = {
+        "provider": "openai",
+        "model": model,
+        "usage": [
+            {
+                "type": "images",
+                "metrics": metrics,
+            }
+        ],
+        "context": context,
+    }
+    enqueue_transaction(transaction)
+
+
+def _track_openai_stream(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    is_responses_api: bool = False,
+    tools: list[dict[str, Any]] | None = None,
+) -> Any:
+    """
+    Wrap a streaming response to track usage from the final chunk.
+
+    Returns a wrapped stream that tracks usage when the stream is exhausted.
+    """
+    last_chunk = None
+
+    def wrapped_stream() -> Any:
+        nonlocal last_chunk
+        try:
+            for chunk in stream:
+                last_chunk = chunk
+                yield chunk
+        finally:
+            # After stream is exhausted, track usage from final chunk
+            if last_chunk and hasattr(last_chunk, "usage") and last_chunk.usage:
+                try:
+                    if is_responses_api:
+                        _track_responses_api_response(
+                            last_chunk, model, context, tools=tools
+                        )
+                    else:
+                        _track_chat_completion_response(last_chunk, model, context)
+                except Exception as e:
+                    logger.error(
+                        f"Error tracking OpenAI streaming usage: {e}", exc_info=True
+                    )
+
+    return wrapped_stream()
+
+
+async def _track_openai_stream_async(
+    stream: Any,
+    model: str,
+    context: dict[str, Any],
+    is_responses_api: bool = False,
+    tools: list[dict[str, Any]] | None = None,
+) -> Any:
+    """
+    Wrap an async streaming response to track usage from the final chunk.
+
+    Returns a wrapped async stream that tracks usage when the stream is exhausted.
+    """
+    last_chunk = None
+
+    try:
+        async for chunk in stream:
+            last_chunk = chunk
+            yield chunk
+    finally:
+        # After stream is exhausted, track usage from final chunk
+        if last_chunk and hasattr(last_chunk, "usage") and last_chunk.usage:
+            try:
+                if is_responses_api:
+                    _track_responses_api_response(
+                        last_chunk, model, context, tools=tools
+                    )
+                else:
+                    _track_chat_completion_response(last_chunk, model, context)
+            except Exception as e:
+                logger.error(
+                    f"Error tracking OpenAI streaming usage: {e}", exc_info=True
+                )
+
+
+# ============================================================================
+# Chat Completions API Patching
+# ============================================================================
+
+def patch_openai() -> None:
+    """
+    Patch OpenAI Chat Completions client to auto-track usage.
+
+    This patches the synchronous `Completions.create` method.
+    """
+    global _patched_chat_sync, _original_chat_create
+
+    if _patched_chat_sync:
+        logger.debug("OpenAI Chat Completions sync already patched, skipping")
+        return
+
+    try:
+        from openai.resources.chat.completions import Completions
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_create = Completions.create
+    _original_chat_create = original_create
+
+    @functools.wraps(original_create)
+    def patched_create(self: Completions, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return original_create(self, *args, **kwargs)
+
+        is_streaming = kwargs.get("stream", False)
+
+        response = original_create(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        try:
+            if is_streaming:
+                return _track_openai_stream(response, model, context, is_responses_api=False)
+            else:
+                _track_chat_completion_response(response, model, context)
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI usage: {e}", exc_info=True)
+
+        return response
+
+    Completions.create = patched_create  # type: ignore[assignment]
+    _patched_chat_sync = True
+    logger.info("OpenAI Chat Completions SDK patched for auto-instrumentation")
+
+
+def patch_openai_async() -> None:
+    """
+    Patch OpenAI Chat Completions async client to auto-track usage.
+
+    This patches the asynchronous `AsyncCompletions.create` method.
+    """
+    global _patched_chat_async, _original_chat_create_async
+
+    if _patched_chat_async:
+        logger.debug("OpenAI Chat Completions async already patched, skipping")
+        return
+
+    try:
+        from openai.resources.chat.completions import AsyncCompletions
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_create = AsyncCompletions.create
+    _original_chat_create_async = original_create
+
+    @functools.wraps(original_create)
+    async def patched_create(
+        self: AsyncCompletions, *args: Any, **kwargs: Any
+    ) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return await original_create(self, *args, **kwargs)
+
+        is_streaming = kwargs.get("stream", False)
+
+        response = await original_create(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        try:
+            if is_streaming:
+                return _track_openai_stream_async(response, model, context, is_responses_api=False)
+            else:
+                _track_chat_completion_response(response, model, context)
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI usage: {e}", exc_info=True)
+
+        return response
+
+    AsyncCompletions.create = patched_create  # type: ignore[assignment]
+    _patched_chat_async = True
+    logger.info("OpenAI Chat Completions async SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Responses API Patching
+# ============================================================================
+
+
+def patch_openai_responses() -> None:
+    """
+    Patch OpenAI Responses API client to auto-track usage.
+
+    This patches the synchronous `Responses.create` method.
+    """
+    global _patched_responses_sync, _original_responses_create
+
+    if _patched_responses_sync:
+        logger.debug("OpenAI Responses API sync already patched, skipping")
+        return
+
+    try:
+        from openai.resources.responses import Responses
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_create = Responses.create
+    _original_responses_create = original_create
+
+    @functools.wraps(original_create)
+    def patched_create(self: Responses, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return original_create(self, *args, **kwargs)
+
+        is_streaming = kwargs.get("stream", False)
+
+        response = original_create(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        tools = kwargs.get("tools")
+        try:
+            if is_streaming:
+                return _track_openai_stream(
+                    response, model, context, is_responses_api=True, tools=tools
+                )
+            else:
+                _track_responses_api_response(response, model, context, tools=tools)
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI Responses API usage: {e}", exc_info=True)
+
+        return response
+
+    Responses.create = patched_create  # type: ignore[assignment]
+    _patched_responses_sync = True
+    logger.info("OpenAI Responses API SDK patched for auto-instrumentation")
+
+
+def patch_openai_responses_async() -> None:
+    """
+    Patch OpenAI Responses API async client to auto-track usage.
+
+    This patches the asynchronous `AsyncResponses.create` method.
+    """
+    global _patched_responses_async, _original_responses_create_async
+
+    if _patched_responses_async:
+        logger.debug("OpenAI Responses API async already patched, skipping")
+        return
+
+    try:
+        from openai.resources.responses import AsyncResponses
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_create = AsyncResponses.create
+    _original_responses_create_async = original_create
+
+    @functools.wraps(original_create)
+    async def patched_create(
+        self: AsyncResponses, *args: Any, **kwargs: Any
+    ) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return await original_create(self, *args, **kwargs)
+
+        is_streaming = kwargs.get("stream", False)
+
+        response = await original_create(self, *args, **kwargs)
+
+        context = get_context()
+
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        tools = kwargs.get("tools")
+        try:
+            if is_streaming:
+                return _track_openai_stream_async(
+                    response, model, context, is_responses_api=True, tools=tools
+                )
+            else:
+                _track_responses_api_response(response, model, context, tools=tools)
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI Responses API usage: {e}", exc_info=True)
+
+        return response
+
+    AsyncResponses.create = patched_create  # type: ignore[assignment]
+    _patched_responses_async = True
+    logger.info("OpenAI Responses API async SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Images API Patching
+# ============================================================================
+
+
+def patch_openai_images() -> None:
+    """
+    Patch OpenAI Images client to auto-track usage.
+
+    This patches the synchronous `Images.generate` method.
+    """
+    global _patched_images_sync, _original_images_generate
+
+    if _patched_images_sync:
+        logger.debug("OpenAI Images sync already patched, skipping")
+        return
+
+    try:
+        from openai.resources.images import Images
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_generate = Images.generate
+    _original_images_generate = original_generate
+
+    @functools.wraps(original_generate)
+    def patched_generate(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return original_generate(self, *args, **kwargs)
+
+        response = original_generate(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model", "unknown")
+        n = kwargs.get("n", 1)
+        size = kwargs.get("size", "")
+        quality = kwargs.get("quality", "")
+
+        try:
+            _track_images_response(
+                model=model,
+                n=n,
+                size=size,
+                quality=quality,
+                context=context,
+                response=response,
+            )
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI Images usage: {e}", exc_info=True)
+
+        return response
+
+    Images.generate = patched_generate  # type: ignore[assignment]
+    _patched_images_sync = True
+    logger.info("OpenAI Images SDK patched for auto-instrumentation")
+
+
+def patch_openai_images_async() -> None:
+    """
+    Patch OpenAI Images async client to auto-track usage.
+
+    This patches the asynchronous `AsyncImages.generate` method.
+    """
+    global _patched_images_async, _original_images_generate_async
+
+    if _patched_images_async:
+        logger.debug("OpenAI Images async already patched, skipping")
+        return
+
+    try:
+        from openai.resources.images import AsyncImages
+    except ImportError:
+        logger.warning(
+            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
+        )
+        return
+
+    original_generate = AsyncImages.generate
+    _original_images_generate_async = original_generate
+
+    @functools.wraps(original_generate)
+    async def patched_generate(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return await original_generate(self, *args, **kwargs)
+
+        response = await original_generate(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model", "unknown")
+        n = kwargs.get("n", 1)
+        size = kwargs.get("size", "")
+        quality = kwargs.get("quality", "")
+
+        try:
+            _track_images_response(
+                model=model,
+                n=n,
+                size=size,
+                quality=quality,
+                context=context,
+                response=response,
+            )
+        except Exception as e:
+            logger.error(f"Error tracking OpenAI Images usage: {e}", exc_info=True)
+
+        return response
+
+    AsyncImages.generate = patched_generate  # type: ignore[assignment]
+    _patched_images_async = True
+    logger.info("OpenAI Images async SDK patched for auto-instrumentation")
+
+
+# ============================================================================
+# Beta Chat Completions Parse Patching
+# ============================================================================
+
+
+def patch_openai_beta_parse() -> None:
+    """
+    Patch OpenAI Beta Chat Completions parse to auto-track usage.
+
+    This patches the synchronous `beta.chat.completions.parse` method.
+    Response format is identical to Chat Completions, so we reuse
+    _track_chat_completion_response.
+    """
+    global _patched_beta_parse_sync, _original_beta_parse
+
+    if _patched_beta_parse_sync:
+        logger.debug("OpenAI Beta parse sync already patched, skipping")
+        return
+
+    try:
+        from openai.resources.beta.chat.completions import Completions as BetaCompletions
+    except ImportError:
+        logger.debug(
+            "OpenAI Beta chat completions not available (optional). "
+            "Skip patching beta parse."
+        )
+        return
+
+    if not hasattr(BetaCompletions, "parse"):
+        logger.debug("OpenAI Beta Completions has no parse method, skipping")
+        return
+
+    original_parse = BetaCompletions.parse
+    _original_beta_parse = original_parse
+
+    @functools.wraps(original_parse)
+    def patched_parse(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return original_parse(self, *args, **kwargs)
+
+        response = original_parse(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        try:
+            _track_chat_completion_response(response, model, context)
+        except Exception as e:
+            logger.error(
+                f"Error tracking OpenAI Beta parse usage: {e}", exc_info=True
+            )
+
+        return response
+
+    BetaCompletions.parse = patched_parse  # type: ignore[assignment]
+    _patched_beta_parse_sync = True
+    logger.info("OpenAI Beta chat.completions.parse SDK patched for auto-instrumentation")
+
+
+def patch_openai_beta_parse_async() -> None:
+    """
+    Patch OpenAI Beta Chat Completions async parse to auto-track usage.
+
+    This patches the asynchronous `beta.chat.completions.parse` method.
+    """
+    global _patched_beta_parse_async, _original_beta_parse_async
+
+    if _patched_beta_parse_async:
+        logger.debug("OpenAI Beta parse async already patched, skipping")
+        return
+
+    try:
+        from openai.resources.beta.chat.completions import (
+            AsyncCompletions as AsyncBetaCompletions,
+        )
+    except ImportError:
+        logger.debug(
+            "OpenAI Beta async chat completions not available (optional). "
+            "Skip patching beta parse async."
+        )
+        return
+
+    if not hasattr(AsyncBetaCompletions, "parse"):
+        logger.debug("OpenAI AsyncBeta Completions has no parse method, skipping")
+        return
+
+    original_parse = AsyncBetaCompletions.parse
+    _original_beta_parse_async = original_parse
+
+    @functools.wraps(original_parse)
+    async def patched_parse(self: Any, *args: Any, **kwargs: Any) -> Any:
+        if not _is_openai_endpoint(self._client):
+            return await original_parse(self, *args, **kwargs)
+
+        response = await original_parse(self, *args, **kwargs)
+
+        context = get_context()
+        model = kwargs.get("model") or getattr(response, "model", "unknown")
+
+        try:
+            _track_chat_completion_response(response, model, context)
+        except Exception as e:
+            logger.error(
+                f"Error tracking OpenAI Beta parse usage: {e}", exc_info=True
+            )
+
+        return response
+
+    AsyncBetaCompletions.parse = patched_parse  # type: ignore[assignment]
+    _patched_beta_parse_async = True
+    logger.info(
+        "OpenAI Beta chat.completions.parse async SDK patched for auto-instrumentation"
+    )
+
+
+# ============================================================================
+# Unpatch Functions
+# ============================================================================
+
+
+def unpatch_openai() -> None:
+    """Restore original OpenAI Chat Completions create method."""
+    global _patched_chat_sync, _original_chat_create
+
+    if not _patched_chat_sync or _original_chat_create is None:
+        return
+
+    try:
+        from openai.resources.chat.completions import Completions
+
+        Completions.create = _original_chat_create  # type: ignore[assignment]
+        _patched_chat_sync = False
+        _original_chat_create = None
+        logger.info("OpenAI Chat Completions SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_async() -> None:
+    """Restore original OpenAI Chat Completions async create method."""
+    global _patched_chat_async, _original_chat_create_async
+
+    if not _patched_chat_async or _original_chat_create_async is None:
+        return
+
+    try:
+        from openai.resources.chat.completions import AsyncCompletions
+
+        AsyncCompletions.create = _original_chat_create_async  # type: ignore[assignment]
+        _patched_chat_async = False
+        _original_chat_create_async = None
+        logger.info("OpenAI Chat Completions async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_responses() -> None:
+    """Restore original OpenAI Responses API create method."""
+    global _patched_responses_sync, _original_responses_create
+
+    if not _patched_responses_sync or _original_responses_create is None:
+        return
+
+    try:
+        from openai.resources.responses import Responses
+
+        Responses.create = _original_responses_create  # type: ignore[assignment]
+        _patched_responses_sync = False
+        _original_responses_create = None
+        logger.info("OpenAI Responses API SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_responses_async() -> None:
+    """Restore original OpenAI Responses API async create method."""
+    global _patched_responses_async, _original_responses_create_async
+
+    if not _patched_responses_async or _original_responses_create_async is None:
+        return
+
+    try:
+        from openai.resources.responses import AsyncResponses
+
+        AsyncResponses.create = _original_responses_create_async  # type: ignore[assignment]
+        _patched_responses_async = False
+        _original_responses_create_async = None
+        logger.info("OpenAI Responses API async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_images() -> None:
+    """Restore original OpenAI Images generate method."""
+    global _patched_images_sync, _original_images_generate
+
+    if not _patched_images_sync or _original_images_generate is None:
+        return
+
+    try:
+        from openai.resources.images import Images
+
+        Images.generate = _original_images_generate  # type: ignore[assignment]
+        _patched_images_sync = False
+        _original_images_generate = None
+        logger.info("OpenAI Images SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_images_async() -> None:
+    """Restore original OpenAI Images async generate method."""
+    global _patched_images_async, _original_images_generate_async
+
+    if not _patched_images_async or _original_images_generate_async is None:
+        return
+
+    try:
+        from openai.resources.images import AsyncImages
+
+        AsyncImages.generate = _original_images_generate_async  # type: ignore[assignment]
+        _patched_images_async = False
+        _original_images_generate_async = None
+        logger.info("OpenAI Images async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_beta_parse() -> None:
+    """Restore original OpenAI Beta Chat Completions parse method."""
+    global _patched_beta_parse_sync, _original_beta_parse
+
+    if not _patched_beta_parse_sync or _original_beta_parse is None:
+        return
+
+    try:
+        from openai.resources.beta.chat.completions import Completions as BetaCompletions
+
+        BetaCompletions.parse = _original_beta_parse  # type: ignore[assignment]
+        _patched_beta_parse_sync = False
+        _original_beta_parse = None
+        logger.info("OpenAI Beta chat.completions.parse SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_beta_parse_async() -> None:
+    """Restore original OpenAI Beta Chat Completions async parse method."""
+    global _patched_beta_parse_async, _original_beta_parse_async
+
+    if not _patched_beta_parse_async or _original_beta_parse_async is None:
+        return
+
+    try:
+        from openai.resources.beta.chat.completions import (
+            AsyncCompletions as AsyncBetaCompletions,
+        )
+
+        AsyncBetaCompletions.parse = _original_beta_parse_async  # type: ignore[assignment]
+        _patched_beta_parse_async = False
+        _original_beta_parse_async = None
+        logger.info("OpenAI Beta chat.completions.parse async SDK unpatched")
+    except ImportError:
+        pass
+
+
+def unpatch_openai_all() -> None:
+    """Restore all original OpenAI SDK methods."""
+    unpatch_openai()
+    unpatch_openai_async()
+    unpatch_openai_responses()
+    unpatch_openai_responses_async()
+    unpatch_openai_images()
+    unpatch_openai_images_async()
+    unpatch_openai_beta_parse()
+    unpatch_openai_beta_parse_async()