Intial Commit

Author: Nisarg38

examples/interceptor.py

@@ -0,0 +1,75 @@
+from typing import TypeVar, Any
+import time
+from typing_extensions import override
+from openai import OpenAI
+from openai._interceptor import Interceptor, InterceptorRequest, InterceptorResponse
+from dotenv import load_dotenv
+
+load_dotenv()
+
+T = TypeVar("T")
+
+# Define a custom logging interceptor
+class LoggingInterceptor(Interceptor):
+    @override
+    def before_request(self, request: InterceptorRequest) -> InterceptorRequest:
+        print(f"Request: {request.method} {request.url}")
+        print(f"Headers: {request.headers}")
+        if request.body:
+            print(f"Body: {request.body}")
+        return request
+
+    @override
+    def after_response(self, response: InterceptorResponse[Any]) -> InterceptorResponse[Any]:
+        print(f"Response Status: {response.status_code}")
+        print(f"Response Headers: {response.headers}")
+        print(f"Response Body: {response.body}")
+        return response
+
+# Define an interceptor that implements retry logic with exponential backoff
+class RetryInterceptor(Interceptor):
+    def __init__(self, max_retries: int = 3, initial_delay: float = 1.0):
+        self.max_retries = max_retries
+        self.initial_delay = initial_delay
+        self.current_retry = 0
+
+    @override
+    def before_request(self, request: InterceptorRequest) -> InterceptorRequest:
+        return request
+
+    @override
+    def after_response(self, response: InterceptorResponse[Any]) -> InterceptorResponse[Any]:
+        # If response is successful or we've exhausted retries, return as is
+        if response.status_code < 500 or self.current_retry >= self.max_retries:
+            self.current_retry = 0  # Reset for next request
+            return response
+
+        # Calculate delay with exponential backoff
+        delay = self.initial_delay * (2 ** self.current_retry)
+        print(f"Request failed with status {response.status_code}. Retrying in {delay} seconds...")
+        time.sleep(delay)
+        
+        self.current_retry += 1
+        # Trigger a retry by raising an exception
+        raise Exception(f"Retrying request (attempt {self.current_retry}/{self.max_retries})")
+
+# Initialize the OpenAI client and add interceptors
+if __name__ == "__main__":
+    # Create the interceptor chain
+    logging_interceptor = LoggingInterceptor()
+    retry_interceptor = RetryInterceptor(max_retries=3, initial_delay=1.0)
+
+    # Create client with interceptors
+    client = OpenAI(
+        interceptors=[logging_interceptor, retry_interceptor]
+    )
+
+    # Make a request using the client
+    response = client.chat.completions.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "Tell me about error handling and retries in software systems."}],
+        max_tokens=100,
+    )
+
+    # Output the final response
+    print("Final Response:", response)

src/openai/_base_client.py

@@ -89,6 +89,8 @@
 )
 from ._legacy_response import LegacyAPIResponse
 
+from ._interceptor import InterceptorChain, InterceptorRequest, InterceptorResponse, Interceptor 
+
 log: logging.Logger = logging.getLogger(__name__)
 log.addFilter(SensitiveHeadersFilter())
 
@@ -339,6 +341,8 @@ class BaseClient(Generic[_HttpxClientT, _DefaultStreamT]):
     _strict_response_validation: bool
     _idempotency_header: str | None
     _default_stream_cls: type[_DefaultStreamT] | None = None
+    _interceptor_chain: InterceptorChain
+
 
     def __init__(
         self,
@@ -353,6 +357,8 @@ def __init__(
         proxies: ProxiesTypes | None,
         custom_headers: Mapping[str, str] | None = None,
         custom_query: Mapping[str, object] | None = None,
+        interceptors: list[Interceptor] | None = None,
+
     ) -> None:
         self._version = version
         self._base_url = self._enforce_trailing_slash(URL(base_url))
@@ -372,7 +378,9 @@ def __init__(
                 "max_retries cannot be None. If you want to disable retries, pass `0`; if you want unlimited retries, pass `math.inf` or a very high number; if you want the default behavior, pass `openai.DEFAULT_MAX_RETRIES`"
             )
 
-    def _enforce_trailing_slash(self, url: URL) -> URL:
+
+        self._interceptor_chain = InterceptorChain(interceptors)
+     def _enforce_trailing_slash(self, url: URL) -> URL:
         if url.raw_path.endswith(b"/"):
             return url
         return url.copy_with(raw_path=url.raw_path + b"/")
@@ -463,8 +471,27 @@ def _build_request(
             else:
                 raise RuntimeError(f"Unexpected JSON data type, {type(json_data)}, cannot merge with `extra_body`")
 
+        # Build base headers and params
         headers = self._build_headers(options, retries_taken=retries_taken)
         params = _merge_mappings(self.default_query, options.params)
+        prepared_url = self._prepare_url(options.url)
+
+        # Execute request interceptors
+        interceptor_request = InterceptorRequest(
+            method=options.method,
+            url=str(prepared_url),
+            headers=dict(headers),
+            params=dict(params),
+            body=json_data,
+        )
+        interceptor_request = self._interceptor_chain.execute_before_request(interceptor_request)
+
+        # Apply interceptor modifications
+        headers = httpx.Headers(interceptor_request.headers)
+        params = interceptor_request.params or {}
+        json_data = interceptor_request.body
+        prepared_url = URL(interceptor_request.url)
+
         content_type = headers.get("Content-Type")
         files = options.files
 
@@ -506,7 +533,7 @@ def _build_request(
         return self._client.build_request(  # pyright: ignore[reportUnknownMemberType]
             headers=headers,
             timeout=self.timeout if isinstance(options.timeout, NotGiven) else options.timeout,
-            method=options.method,
+            method=interceptor_request.method,
             url=prepared_url,
             # the `Query` type that we use is incompatible with qs'
             # `Params` type as it needs to be typed as `Mapping[str, object]`
@@ -582,6 +609,22 @@ def _process_response_data(
             return cast(ResponseT, data)
 
         try:
+            # Create InterceptorResponse and execute interceptors
+            interceptor_response = InterceptorResponse(
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                body=data,
+                request=InterceptorRequest(
+                    method=response.request.method,
+                    url=str(response.request.url),
+                    headers=dict(response.request.headers),
+                    body=response.request._content if response.request._content else None,
+                ),
+                raw_response=response,
+            )
+            interceptor_response = self._interceptor_chain.execute_after_response(interceptor_response)
+            data = interceptor_response.body
+
             if inspect.isclass(cast_to) and issubclass(cast_to, ModelBuilderProtocol):
                 return cast(ResponseT, cast_to.build(response=response, data=data))
 
@@ -796,6 +839,9 @@ def __init__(
         custom_headers: Mapping[str, str] | None = None,
         custom_query: Mapping[str, object] | None = None,
         _strict_response_validation: bool,
+
+        interceptors: list[Interceptor] | None = None,
+
     ) -> None:
         kwargs: dict[str, Any] = {}
         if limits is not None:
@@ -859,6 +905,9 @@ def __init__(
             custom_query=custom_query,
             custom_headers=custom_headers,
             _strict_response_validation=_strict_response_validation,
+
+            interceptors=interceptors,
+
         )
         self._client = http_client or SyncHttpxClientWrapper(
             base_url=base_url,
@@ -1382,6 +1431,8 @@ def __init__(
         http_client: httpx.AsyncClient | None = None,
         custom_headers: Mapping[str, str] | None = None,
         custom_query: Mapping[str, object] | None = None,
+
+        interceptors: list[Interceptor] | None = None,
     ) -> None:
         kwargs: dict[str, Any] = {}
         if limits is not None:
@@ -1445,6 +1496,9 @@ def __init__(
             custom_query=custom_query,
             custom_headers=custom_headers,
             _strict_response_validation=_strict_response_validation,
+
+            interceptors=interceptors,
+
         )
         self._client = http_client or AsyncHttpxClientWrapper(
             base_url=base_url,

src/openai/_client.py

@@ -7,6 +7,7 @@
 from typing_extensions import Self, override
 
 import httpx
+from openai._interceptor import Interceptor
 
 from . import _exceptions
 from ._qs import Querystring
@@ -96,7 +97,8 @@ def __init__(
         # outlining your use-case to help us decide if it should be
         # part of our public interface in the future.
         _strict_response_validation: bool = False,
-    ) -> None:
+        interceptors: list[Interceptor] | None = None,
+ ) -> None:
         """Construct a new synchronous openai client instance.
 
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
@@ -136,7 +138,8 @@ def __init__(
             custom_headers=default_headers,
             custom_query=default_query,
             _strict_response_validation=_strict_response_validation,
-        )
+            interceptors=interceptors,  
+)
 
         self._default_stream_cls = Stream
 
@@ -192,6 +195,7 @@ def copy(
         set_default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
         set_default_query: Mapping[str, object] | None = None,
+        interceptors: list[Interceptor] | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -227,6 +231,7 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            interceptors=interceptors,
             **_extra_kwargs,
         )
 
@@ -323,7 +328,8 @@ def __init__(
         # outlining your use-case to help us decide if it should be
         # part of our public interface in the future.
         _strict_response_validation: bool = False,
-    ) -> None:
+        interceptors: list[Interceptor] | None = None,
+ ) -> None:
         """Construct a new async openai client instance.
 
         This automatically infers the following arguments from their corresponding environment variables if they are not provided:
@@ -363,6 +369,7 @@ def __init__(
             custom_headers=default_headers,
             custom_query=default_query,
             _strict_response_validation=_strict_response_validation,
+            interceptors=interceptors,
         )
 
         self._default_stream_cls = AsyncStream
@@ -419,6 +426,7 @@ def copy(
         set_default_headers: Mapping[str, str] | None = None,
         default_query: Mapping[str, object] | None = None,
         set_default_query: Mapping[str, object] | None = None,
+        interceptors: list[Interceptor] | None = None,
         _extra_kwargs: Mapping[str, Any] = {},
     ) -> Self:
         """
@@ -454,6 +462,7 @@ def copy(
             max_retries=max_retries if is_given(max_retries) else self.max_retries,
             default_headers=headers,
             default_query=params,
+            interceptors=interceptors,
             **_extra_kwargs,
         )
 

src/openai/_interceptor.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Dict, Optional, TypeVar, Generic, Any, Union
+
+import httpx
+
+from ._types import Body
+
+T = TypeVar("T")
+
+@dataclass
+class InterceptorRequest:
+    """Container for request data that can be modified by interceptors"""
+    method: str
+    url: str
+    headers: Dict[str, str]
+    params: Optional[Dict[str, Any]] = None
+    body: Optional[Union[Body, bytes]] = None
+
+@dataclass
+class InterceptorResponse(Generic[T]):
+    """Container for response data that can be processed by interceptors"""
+    status_code: int
+    headers: Dict[str, str]
+    body: T
+    request: InterceptorRequest
+    raw_response: httpx.Response
+
+class Interceptor(ABC):
+    """Base class for implementing request/response interceptors"""
+
+    @abstractmethod
+    def before_request(self, request: InterceptorRequest) -> InterceptorRequest:
+        """Process and optionally modify the request before it is sent.
+        
+        Args:
+            request: The request to process
+            
+        Returns:
+            The processed request
+        """
+        pass
+
+    @abstractmethod
+    def after_response(self, response: InterceptorResponse[T]) -> InterceptorResponse[T]:
+        """Process and optionally modify the response after it is received.
+        
+        Args:
+            response: The response to process
+            
+        Returns:
+            The processed response
+        """
+        pass
+
+class InterceptorChain:
+    """Manages a chain of interceptors that process requests/responses in sequence"""
+
+    def __init__(self, interceptors: Optional[list[Interceptor]] = None):
+        self._interceptors = interceptors or []
+
+    def add_interceptor(self, interceptor: Interceptor) -> None:
+        """Add an interceptor to the chain"""
+        self._interceptors.append(interceptor)
+
+    def execute_before_request(self, request: InterceptorRequest) -> InterceptorRequest:
+        """Execute all interceptors' before_request methods in sequence"""
+        print("\n=== Intercepted Request ===")
+        print(f"Method: {request.method}")
+        print(f"URL: {request.url}")
+        print(f"Headers: {request.headers}")
+        if request.params:
+            print(f"Query Params: {request.params}")
+        if request.body:
+            print(f"Request Body: {request.body}")
+        print("========================\n")
+        
+        current_request = request
+        for interceptor in self._interceptors:
+            try:
+                current_request = interceptor.before_request(current_request)
+            except Exception as e:
+                # Log error but continue processing
+                print(f"Error in interceptor {interceptor.__class__.__name__}: {e}")
+        return current_request
+
+    def execute_after_response(self, response: InterceptorResponse[T]) -> InterceptorResponse[T]:
+        """Execute all interceptors' after_response methods in sequence"""
+        print("\n=== Intercepted Response ===")
+        print(f"Status Code: {response.status_code}")
+        print(f"Headers: {response.headers}")
+        print(f"Response Body: {response.body}")
+        print("=========================\n")
+        
+        current_response = response
+        for interceptor in self._interceptors:
+            try:
+                current_response = interceptor.after_response(current_response)
+            except Exception as e:
+                # Log error but continue processing
+                print(f"Error in interceptor {interceptor.__class__.__name__}: {e}")
+        return current_response