diff --git a/packages/openai-sdk-python/src/supermemory_openai/middleware.py b/packages/openai-sdk-python/src/supermemory_openai/middleware.py index 7c83b90c..f1c4aab8 100644 --- a/packages/openai-sdk-python/src/supermemory_openai/middleware.py +++ b/packages/openai-sdk-python/src/supermemory_openai/middleware.py @@ -48,6 +48,68 @@ class SupermemoryProfileSearch: self.search_results: dict[str, Any] = data.get("searchResults", {}) +class _ResourceFacade: + """Delegate an SDK resource while overriding selected attributes.""" + + def __init__( + self, + resource: Any, + lazy_overrides: Optional[dict[str, Any]] = None, + **overrides: Any, + ) -> None: + self._resource = resource + self._lazy_overrides = lazy_overrides or {} + for name, value in overrides.items(): + setattr(self, name, value) + + def __getattr__(self, name: str) -> Any: + factory = self._lazy_overrides.get(name) + if factory is not None: + value = factory() + setattr(self, name, value) + return value + return getattr(self._resource, name) + + +class _AsyncMemoryResponseContextManager: + """Apply async middleware before entering a streaming response context.""" + + def __init__( + self, + wrapper: "SupermemoryOpenAIWrapper", + original_create: Any, + kwargs: dict[str, Any], + ) -> None: + self._wrapper = wrapper + self._original_create = original_create + self._kwargs = kwargs + self._response_context: Any = None + + async def __aenter__(self) -> Any: + async def enter_original(**kwargs: Any) -> Any: + self._response_context = self._original_create(**kwargs) + return await self._response_context.__aenter__() + + return await self._wrapper._create_with_memory_async( + enter_original, + **self._kwargs, + ) + + async def __aexit__( + self, + exc_type: Any, + exc_value: Any, + traceback: Any, + ) -> Any: + if self._response_context is None: + return None + return await self._response_context.__aexit__( + exc_type, + exc_value, + traceback, + ) + + async def supermemory_profile_search( container_tag: str, query_text: str, @@ -199,9 +261,11 @@ async def add_system_prompt( if system_prompt_exists: logger.debug("Added memories to existing system prompt") return [ - {**msg, "content": f"{msg.get('content', '')} \n {memories}"} - if msg.get("role") == "system" - else msg + ( + {**msg, "content": f"{msg.get('content', '')} \n {memories}"} + if msg.get("role") == "system" + else msg + ) for msg in messages ] @@ -269,10 +333,21 @@ class SupermemoryOpenAIWrapper: openai_client: Union[OpenAI, AsyncOpenAI], options: OpenAIMiddlewareOptions, ): - self._client: Union[OpenAI, AsyncOpenAI] = openai_client + self._client: Union[OpenAI, AsyncOpenAI] = getattr( + openai_client, + "__supermemory_openai_base_client__", + openai_client, + ) + # A stable attribute also lets wrappers from another module copy or hot reload + # recover the pristine client instead of nesting middleware. + self.__supermemory_openai_base_client__ = self._client self._container_tag: str = options.container_tag self._options: OpenAIMiddlewareOptions = options self._logger: Logger = create_logger(self._options.verbose) + base_create = self._client.chat.completions.create + self._is_async_client = isinstance( + self._client, AsyncOpenAI + ) or inspect.iscoroutinefunction(inspect.unwrap(base_create)) # Track background tasks to ensure they complete self._background_tasks: set[asyncio.Task] = set() @@ -293,8 +368,8 @@ class SupermemoryOpenAIWrapper: f"Failed to initialize Supermemory client: {e}", e ) - # Wrap the chat completions create method - self._wrap_chat_completions() + # Expose isolated resource facades without mutating the supplied client. + self.chat = self._create_chat_facade() def _get_api_key(self) -> str: """Get Supermemory API key from environment.""" @@ -307,25 +382,114 @@ class SupermemoryOpenAIWrapper: ) return api_key - def _wrap_chat_completions(self) -> None: - """Wrap the chat completions create method with memory injection.""" - original_create = self._client.chat.completions.create + def _create_chat_facade(self) -> _ResourceFacade: + """Create isolated chat/completions facades with memory injection.""" + completions_resource = self._client.chat.completions + completions = _ResourceFacade( + completions_resource, + lazy_overrides={ + "with_raw_response": lambda: self._create_completion_variant_facade( + "with_raw_response" + ), + "with_streaming_response": lambda: self._create_completion_variant_facade( + "with_streaming_response" + ), + }, + create=self._create_completion_method(completions_resource.create), + ) + return _ResourceFacade( + self._client.chat, + lazy_overrides={ + "with_raw_response": lambda: self._create_chat_variant_facade( + "with_raw_response" + ), + "with_streaming_response": lambda: self._create_chat_variant_facade( + "with_streaming_response" + ), + }, + completions=completions, + ) - if asyncio.iscoroutinefunction(original_create): + def _create_completion_variant_facade(self, name: str) -> _ResourceFacade: + """Preserve raw and streaming response behavior on isolated facades.""" + resource = getattr(self._client.chat.completions, name) + return _ResourceFacade( + resource, + create=self._create_completion_method( + resource.create, + streaming_response=name == "with_streaming_response", + ), + ) - async def create_with_memory( + def _create_chat_variant_facade(self, name: str) -> _ResourceFacade: + """Wrap completions reached through a chat response variant.""" + chat_resource = getattr(self._client.chat, name) + completions_resource = chat_resource.completions + return _ResourceFacade( + chat_resource, + completions=_ResourceFacade( + completions_resource, + create=self._create_completion_method( + completions_resource.create, + streaming_response=name == "with_streaming_response", + ), + ), + ) + + def _create_client_variant_facade(self, name: str) -> _ResourceFacade: + """Wrap completions reached through a client response variant.""" + client_resource = getattr(self._client, name) + chat_resource = client_resource.chat + completions_resource = chat_resource.completions + return _ResourceFacade( + client_resource, + chat=_ResourceFacade( + chat_resource, + completions=_ResourceFacade( + completions_resource, + create=self._create_completion_method( + completions_resource.create, + streaming_response=name == "with_streaming_response", + ), + ), + ), + ) + + def _create_completion_method( + self, + original_create: Any, + *, + streaming_response: bool = False, + ) -> Any: + """Wrap one completion create implementation with memory injection.""" + if self._is_async_client and streaming_response: + + def create_streaming_with_memory( + **kwargs: Any, + ) -> _AsyncMemoryResponseContextManager: + return _AsyncMemoryResponseContextManager( + self, + original_create, + kwargs, + ) + + return create_streaming_with_memory + + if self._is_async_client: + + async def create_async_with_memory( **kwargs: Any, ) -> Any: return await self._create_with_memory_async(original_create, **kwargs) - else: - def create_with_memory( - **kwargs: Any, - ) -> Any: - return self._create_with_memory_sync(original_create, **kwargs) + return create_async_with_memory - # Replace the create method with our wrapper - setattr(self._client.chat.completions, "create", create_with_memory) + def create_sync_with_memory( + **kwargs: Any, + ) -> Any: + return self._create_with_memory_sync(original_create, **kwargs) + + return create_sync_with_memory async def _create_with_memory_async( self, @@ -558,7 +722,9 @@ class SupermemoryOpenAIWrapper: f"Background tasks did not complete within {timeout}s timeout" ) # Cancel remaining tasks - tasks_to_cancel = [task for task in self._background_tasks if not task.done()] + tasks_to_cancel = [ + task for task in self._background_tasks if not task.done() + ] for task in tasks_to_cancel: task.cancel() @@ -616,6 +782,10 @@ class SupermemoryOpenAIWrapper: def __getattr__(self, name: str) -> Any: """Delegate all other attributes to the wrapped client.""" + if name in {"with_raw_response", "with_streaming_response"}: + value = self._create_client_variant_facade(name) + setattr(self, name, value) + return value return getattr(self._client, name) diff --git a/packages/openai-sdk-python/tests/test_client_isolation.py b/packages/openai-sdk-python/tests/test_client_isolation.py new file mode 100644 index 00000000..a0a07ace --- /dev/null +++ b/packages/openai-sdk-python/tests/test_client_isolation.py @@ -0,0 +1,546 @@ +"""Regression tests for shared OpenAI client middleware isolation.""" + +from __future__ import annotations + +import asyncio +import gc +import inspect +import json +import os +import warnings +from types import SimpleNamespace +from typing import Any, Generator, Literal, Optional +from unittest.mock import AsyncMock, Mock, patch + +import httpx +import pytest +from openai import AsyncOpenAI + +from supermemory_openai import OpenAIMiddlewareOptions, with_supermemory + + +def middleware_options( + container_tag: str, + add_memory: Literal["always", "never"] = "never", +) -> OpenAIMiddlewareOptions: + return OpenAIMiddlewareOptions( + container_tag=container_tag, + custom_id=f"thread-{container_tag}", + add_memory=add_memory, + ) + + +def create_sync_client() -> tuple[Any, Mock]: + create = Mock(return_value={"id": "chat-response"}) + completions = SimpleNamespace(create=create, marker="completions-marker") + chat = SimpleNamespace(completions=completions, marker="chat-marker") + return SimpleNamespace(chat=chat), create + + +def create_async_client() -> tuple[Any, AsyncMock]: + create = AsyncMock(return_value={"id": "chat-response"}) + completions = SimpleNamespace(create=create) + chat = SimpleNamespace(completions=completions) + return SimpleNamespace(chat=chat), create + + +class RawResponse: + def __init__(self, label: str) -> None: + self.label = label + + def parse(self) -> str: + return f"parsed-{self.label}" + + +class SyncStreamContext: + def __init__(self, label: str) -> None: + self.label = label + + def __enter__(self) -> "SyncStreamContext": + return self + + def __exit__(self, *args: Any) -> None: + return None + + +class AsyncStreamContext: + def __init__(self, label: str) -> None: + self.label = label + + async def __aenter__(self) -> "AsyncStreamContext": + return self + + async def __aexit__(self, *args: Any) -> None: + return None + + +def attach_response_variants( + client: Any, + raw_create_factory: Any, + streaming_create_factory: Any, +) -> dict[str, Any]: + calls: dict[str, Any] = {} + + def completion_resource(label: str, factory: Any) -> Any: + create = factory(label) + calls[label] = create + return SimpleNamespace(create=create) + + client.chat.completions.with_raw_response = completion_resource( + "completions-raw", raw_create_factory + ) + client.chat.with_raw_response = SimpleNamespace( + completions=completion_resource("chat-raw", raw_create_factory) + ) + client.with_raw_response = SimpleNamespace( + chat=SimpleNamespace( + completions=completion_resource("client-raw", raw_create_factory) + ) + ) + client.chat.completions.with_streaming_response = completion_resource( + "completions-stream", streaming_create_factory + ) + client.chat.with_streaming_response = SimpleNamespace( + completions=completion_resource("chat-stream", streaming_create_factory) + ) + client.with_streaming_response = SimpleNamespace( + chat=SimpleNamespace( + completions=completion_resource("client-stream", streaming_create_factory) + ) + ) + return calls + + +def response_variant_creates(client: Any, name: str) -> list[Any]: + return [ + getattr(client.chat.completions, name).create, + getattr(client.chat, name).completions.create, + getattr(client, name).chat.completions.create, + ] + + +REAL_ASYNC_COMPLETION_PATHS = ( + "normal", + "completions.raw", + "chat.raw", + "client.raw", + "completions.streaming", + "chat.streaming", + "client.streaming", +) + + +def real_async_completion_create(client: Any, path: str) -> Any: + if path == "normal": + return client.chat.completions.create + if path == "completions.raw": + return client.chat.completions.with_raw_response.create + if path == "chat.raw": + return client.chat.with_raw_response.completions.create + if path == "client.raw": + return client.with_raw_response.chat.completions.create + if path == "completions.streaming": + return client.chat.completions.with_streaming_response.create + if path == "chat.streaming": + return client.chat.with_streaming_response.completions.create + if path == "client.streaming": + return client.with_streaming_response.chat.completions.create + raise AssertionError(f"Unknown completion path: {path}") + + +@pytest.fixture(autouse=True) # type: ignore[untyped-decorator] +def supermemory_api_key() -> Generator[None, None, None]: + with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}): + yield + + +def test_shared_sync_client_does_not_stack_tenant_middleware() -> None: + base_client, original_create = create_sync_client() + lookups: list[str] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + side_effect=fake_prompt, + ): + tenant_a: Any = with_supermemory( + base_client, + middleware_options("tenant-a"), + ) + tenant_b: Any = with_supermemory( + base_client, + middleware_options("tenant-b"), + ) + + assert base_client.chat.completions.create is original_create + assert tenant_a.chat is not base_client.chat + assert tenant_b.chat.marker == "chat-marker" + assert tenant_b.chat.completions.marker == "completions-marker" + + tenant_b.chat.completions.create( + model="gpt-test", + messages=[{"role": "user", "content": "private"}], + ) + + assert lookups == ["tenant-b"] + messages = original_create.call_args.kwargs["messages"] + assert messages[0]["content"] == "secret-tenant-b" + assert all("tenant-a" not in str(message) for message in messages) + + base_client.chat.completions.create( + model="gpt-test", + messages=[{"role": "user", "content": "unwrapped"}], + ) + assert lookups == ["tenant-b"] + + +def test_rewrapping_a_facade_recovers_the_pristine_client() -> None: + base_client, original_create = create_sync_client() + lookups: list[str] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + side_effect=fake_prompt, + ): + tenant_a: Any = with_supermemory( + base_client, + middleware_options("tenant-a"), + ) + tenant_b: Any = with_supermemory( + tenant_a, + middleware_options("tenant-b"), + ) + + tenant_b.chat.completions.create( + model="gpt-test", + messages=[{"role": "user", "content": "private"}], + ) + + assert lookups == ["tenant-b"] + assert original_create.call_count == 1 + assert base_client.chat.completions.create is original_create + + +def test_raw_and_streaming_response_facades_remain_memory_aware() -> None: + base_client, _ = create_sync_client() + calls = attach_response_variants( + base_client, + lambda label: Mock(return_value=RawResponse(label)), + lambda label: Mock(return_value=SyncStreamContext(label)), + ) + lookups: list[str] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + side_effect=fake_prompt, + ): + tenant_b: Any = with_supermemory( + base_client, + middleware_options("tenant-b"), + ) + + for create in response_variant_creates(tenant_b, "with_raw_response"): + response = create( + model="gpt-test", + messages=[{"role": "user", "content": "raw"}], + ) + assert response.parse().startswith("parsed-") + + for create in response_variant_creates(tenant_b, "with_streaming_response"): + with create( + model="gpt-test", + messages=[{"role": "user", "content": "streaming"}], + ) as stream: + assert stream.label.endswith("stream") + + assert lookups == ["tenant-b"] * 6 + for create in calls.values(): + assert "secret-tenant-b" in str(create.call_args.kwargs["messages"]) + + +def test_async_raw_and_streaming_response_prefixes_remain_memory_aware() -> None: + base_client, _ = create_async_client() + calls = attach_response_variants( + base_client, + lambda label: AsyncMock(return_value=RawResponse(label)), + lambda label: Mock(return_value=AsyncStreamContext(label)), + ) + lookups: list[str] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + async def call_raw(create: Any) -> None: + response = await create( + model="gpt-test", + messages=[{"role": "user", "content": "raw"}], + ) + assert response.parse().startswith("parsed-") + + async def consume_stream(stream_context: Any) -> None: + async with stream_context as stream: + assert stream.label.endswith("stream") + + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + side_effect=fake_prompt, + ): + tenant_b: Any = with_supermemory( + base_client, + middleware_options("tenant-b"), + ) + + for create in response_variant_creates(tenant_b, "with_raw_response"): + asyncio.run(call_raw(create)) + + for create in response_variant_creates(tenant_b, "with_streaming_response"): + stream_context = create( + model="gpt-test", + messages=[{"role": "user", "content": "streaming"}], + ) + asyncio.run(consume_stream(stream_context)) + + assert lookups == ["tenant-b"] * 6 + for create in calls.values(): + assert "secret-tenant-b" in str(create.call_args.kwargs["messages"]) + + +@pytest.mark.parametrize("path", REAL_ASYNC_COMPLETION_PATHS) +@pytest.mark.asyncio +async def test_real_async_openai_paths_use_async_middleware(path: str) -> None: + lookups: list[str] = [] + writes: list[tuple[str, Optional[str], str]] = [] + sent_messages: list[list[Any]] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + async def fake_add_memory( + client: Any, + container_tag: str, + content: str, + custom_id: Optional[str], + logger: Any, + ) -> None: + writes.append((container_tag, custom_id, content)) + + def handle_request(request: httpx.Request) -> httpx.Response: + body = json.loads(request.content) + sent_messages.append(body["messages"]) + return httpx.Response( + 200, + request=request, + headers={"content-type": "application/json"}, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 0, + "model": "gpt-test", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + }, + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) + base_client = AsyncOpenAI(api_key="openai-test", http_client=http_client) + + try: + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + new=fake_prompt, + ), patch( + "supermemory_openai.middleware.add_memory_tool", + new=fake_add_memory, + ), warnings.catch_warnings( + record=True + ) as caught: + warnings.simplefilter("always", RuntimeWarning) + wrapped: Any = with_supermemory( + base_client, + middleware_options("tenant-real", add_memory="always"), + ) + create = real_async_completion_create(wrapped, path) + kwargs = { + "model": "gpt-test", + "messages": [{"role": "user", "content": "private message"}], + } + + if path == "normal": + response = await create(**kwargs) + assert response.id == "chatcmpl-test" + elif path.endswith(".raw"): + raw_response = await create(**kwargs) + assert raw_response.parse().id == "chatcmpl-test" + else: + response_context = create(**kwargs) + assert not inspect.isawaitable(response_context) + async with response_context as streaming_response: + assert streaming_response.status_code == 200 + + await wrapped.wait_for_background_tasks() + await asyncio.sleep(0) + gc.collect() + + runtime_warnings = [ + warning + for warning in caught + if issubclass(warning.category, RuntimeWarning) + ] + + assert lookups == ["tenant-real"] + assert writes == [ + ( + "tenant-real", + "conversation:thread-tenant-real", + "User: private message", + ) + ] + assert len(sent_messages) == 1 + assert sent_messages[0][0]["content"] == "secret-tenant-real" + assert runtime_warnings == [] + finally: + await base_client.close() + + +@pytest.mark.asyncio # type: ignore[untyped-decorator] +async def test_shared_async_client_saves_only_for_selected_tenant() -> None: + base_client, original_create = create_async_client() + lookups: list[str] = [] + writes: list[tuple[str, Optional[str], str]] = [] + + async def fake_prompt( + messages: list[Any], + container_tag: str, + logger: Any, + mode: Any, + api_key: str, + ) -> list[Any]: + lookups.append(container_tag) + return [ + {"role": "system", "content": f"secret-{container_tag}"}, + *messages, + ] + + async def fake_add_memory( + client: Any, + container_tag: str, + content: str, + custom_id: Optional[str], + logger: Any, + ) -> None: + writes.append((container_tag, custom_id, content)) + + with patch( + "supermemory_openai.middleware.supermemory.Supermemory", + return_value=Mock(), + ), patch( + "supermemory_openai.middleware.add_system_prompt", + side_effect=fake_prompt, + ), patch( + "supermemory_openai.middleware.add_memory_tool", + side_effect=fake_add_memory, + ): + tenant_a: Any = with_supermemory( + base_client, + middleware_options("tenant-a", add_memory="always"), + ) + tenant_b: Any = with_supermemory( + base_client, + middleware_options("tenant-b", add_memory="always"), + ) + + await tenant_b.chat.completions.create( + model="gpt-test", + messages=[{"role": "user", "content": "private tenant B message"}], + ) + await tenant_a.wait_for_background_tasks() + await tenant_b.wait_for_background_tasks() + + assert base_client.chat.completions.create is original_create + assert lookups == ["tenant-b"] + assert writes == [ + ( + "tenant-b", + "conversation:thread-tenant-b", + "User: private tenant B message", + ) + ] + assert original_create.call_count == 1