update open ai python sdk (#879)

This commit is contained in:
sreedharsreeram 2026-04-25 01:43:14 +00:00
parent 589edd3437
commit 0104a2572c
4 changed files with 129 additions and 99 deletions

View file

@ -44,11 +44,12 @@ async def main():
# Wrap with Supermemory middleware
openai_with_memory = with_supermemory(
openai,
container_tag="user-123", # Unique identifier for user's memories
options=OpenAIMiddlewareOptions(
mode="full", # "profile", "query", or "full"
verbose=True, # Enable logging
add_memory="always" # Automatically save conversations
OpenAIMiddlewareOptions(
container_tag="user-123", # Required: unique identifier for user's memories
custom_id="chat-123", # Required: groups messages into documents
mode="full", # "profile", "query", or "full"
verbose=True, # Enable logging
add_memory="always" # Automatically save conversations (default)
)
)
@ -122,7 +123,13 @@ from supermemory_openai import with_supermemory
# Sync client
openai = OpenAI(api_key="your-openai-api-key")
openai_with_memory = with_supermemory(openai, "user-123")
openai_with_memory = with_supermemory(
openai,
OpenAIMiddlewareOptions(
container_tag="user-123",
custom_id="session-456"
)
)
# Works the same way
response = openai_with_memory.chat.completions.create(
@ -136,13 +143,21 @@ response = openai_with_memory.chat.completions.create(
**Background Task Management**: When `add_memory="always"`, memory storage happens in background tasks. Use context managers or manual cleanup to ensure tasks complete:
```python
from supermemory_openai import with_supermemory, OpenAIMiddlewareOptions
# Async context manager (recommended)
async with with_supermemory(openai, "user-123") as client:
async with with_supermemory(
openai,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456")
) as client:
response = await client.chat.completions.create(...)
# Background tasks automatically waited for on exit
# Manual cleanup
client = with_supermemory(openai, "user-123")
client = with_supermemory(
openai,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456")
)
response = await client.chat.completions.create(...)
await client.wait_for_background_tasks() # Ensure memory is saved
```
@ -159,8 +174,7 @@ Injects all static and dynamic profile memories into every request. Best for mai
```python
openai_with_memory = with_supermemory(
openai,
"user-123",
OpenAIMiddlewareOptions(mode="profile")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456", mode="profile")
)
```
@ -170,8 +184,7 @@ Only searches for memories relevant to the current user message. More efficient
```python
openai_with_memory = with_supermemory(
openai,
"user-123",
OpenAIMiddlewareOptions(mode="query")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456", mode="query")
)
```
@ -181,8 +194,7 @@ Combines both profile and query modes - includes all profile memories plus relev
```python
openai_with_memory = with_supermemory(
openai,
"user-123",
OpenAIMiddlewareOptions(mode="full")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456", mode="full")
)
```
@ -191,11 +203,11 @@ openai_with_memory = with_supermemory(
Control when conversations are automatically saved as memories:
```python
# Always save conversations as memories
OpenAIMiddlewareOptions(add_memory="always")
# Always save conversations as memories (default in v2.0.0+)
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456", add_memory="always")
# Never save conversations (default)
OpenAIMiddlewareOptions(add_memory="never")
# Never save conversations
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456", add_memory="never")
```
### Complete Configuration Example
@ -205,12 +217,12 @@ from supermemory_openai import with_supermemory, OpenAIMiddlewareOptions
openai_with_memory = with_supermemory(
openai_client,
container_tag="user-123",
options=OpenAIMiddlewareOptions(
conversation_id="chat-session-456", # Group messages into conversations
verbose=True, # Enable detailed logging
mode="full", # Use both profile and query
add_memory="always" # Auto-save conversations
OpenAIMiddlewareOptions(
container_tag="user-123", # Required: unique user/container identifier
custom_id="chat-session-456", # Required: groups messages into documents
verbose=True, # Enable detailed logging
mode="full", # Use both profile and query
add_memory="always" # Auto-save conversations (default)
)
)
```
@ -291,14 +303,12 @@ Wraps an OpenAI client with automatic memory injection middleware.
```python
def with_supermemory(
openai_client: Union[OpenAI, AsyncOpenAI],
container_tag: str,
options: Optional[OpenAIMiddlewareOptions] = None
options: OpenAIMiddlewareOptions
) -> Union[OpenAI, AsyncOpenAI]
```
**Parameters:**
- `openai_client`: OpenAI or AsyncOpenAI client instance
- `container_tag`: Unique identifier for memory storage (e.g., user ID)
- `options`: Configuration options (see `OpenAIMiddlewareOptions`)
#### `OpenAIMiddlewareOptions`
@ -308,10 +318,11 @@ Configuration dataclass for middleware behavior.
```python
@dataclass
class OpenAIMiddlewareOptions:
conversation_id: Optional[str] = None # Group messages into conversations
container_tag: str # Required: unique identifier for memory storage
custom_id: str # Required: groups messages into documents
verbose: bool = False # Enable detailed logging
mode: Literal["profile", "query", "full"] = "profile" # Memory injection mode
add_memory: Literal["always", "never"] = "never" # Auto-save behavior
add_memory: Literal["always", "never"] = "always" # Auto-save behavior
```
### SupermemoryTools
@ -341,6 +352,7 @@ The package provides specific exception types for better error handling:
```python
from supermemory_openai import (
with_supermemory,
OpenAIMiddlewareOptions,
SupermemoryConfigurationError,
SupermemoryAPIError,
SupermemoryNetworkError,
@ -349,7 +361,10 @@ from supermemory_openai import (
try:
# This will raise SupermemoryConfigurationError if API key is missing
client = with_supermemory(openai_client, "user-123")
client = with_supermemory(
openai_client,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="session-456")
)
response = await client.chat.completions.create(
messages=[{"role": "user", "content": "Hello"}],

View file

@ -1,6 +1,7 @@
"""Supermemory middleware for OpenAI clients."""
import asyncio
import inspect
import os
from dataclasses import dataclass
from typing import Any, Literal, Optional, Union, cast
@ -32,10 +33,11 @@ from .utils import (
class OpenAIMiddlewareOptions:
"""Configuration options for OpenAI middleware."""
conversation_id: Optional[str] = None
container_tag: str # Required: identifies the user/container
custom_id: str # Required: groups messages into the same document
verbose: bool = False
mode: Literal["profile", "query", "full"] = "profile"
add_memory: Literal["always", "never"] = "never"
add_memory: Literal["always", "never"] = "always"
class SupermemoryProfileSearch:
@ -225,11 +227,11 @@ async def add_memory_tool(
add_params["custom_id"] = custom_id
# Handle both sync and async supermemory clients
try:
response = await client.add(**add_params)
except TypeError:
# If it's not awaitable, call it synchronously
response = client.add(**add_params)
result = client.add(**add_params)
if inspect.isawaitable(result):
response = await result
else:
response = result
logger.info(
"Memory saved successfully",
@ -262,12 +264,11 @@ class SupermemoryOpenAIWrapper:
def __init__(
self,
openai_client: Union[OpenAI, AsyncOpenAI],
container_tag: str,
options: Optional[OpenAIMiddlewareOptions] = None,
options: OpenAIMiddlewareOptions,
):
self._client: Union[OpenAI, AsyncOpenAI] = openai_client
self._container_tag: str = container_tag
self._options: OpenAIMiddlewareOptions = options or OpenAIMiddlewareOptions()
self._container_tag: str = options.container_tag
self._options: OpenAIMiddlewareOptions = options
self._logger: Logger = create_logger(self._options.verbose)
# Track background tasks to ensure they complete
@ -336,12 +337,12 @@ class SupermemoryOpenAIWrapper:
if user_message and user_message.strip():
content = (
get_conversation_content(messages)
if self._options.conversation_id
if self._options.custom_id
else user_message
)
custom_id = (
f"conversation:{self._options.conversation_id}"
if self._options.conversation_id
f"conversation:{self._options.custom_id}"
if self._options.custom_id
else None
)
@ -399,7 +400,7 @@ class SupermemoryOpenAIWrapper:
"Starting memory search",
{
"container_tag": self._container_tag,
"conversation_id": self._options.conversation_id,
"conversation_id": self._options.custom_id,
"mode": self._options.mode,
},
)
@ -430,12 +431,12 @@ class SupermemoryOpenAIWrapper:
if user_message and user_message.strip():
content = (
get_conversation_content(messages)
if self._options.conversation_id
if self._options.custom_id
else user_message
)
custom_id = (
f"conversation:{self._options.conversation_id}"
if self._options.conversation_id
f"conversation:{self._options.custom_id}"
if self._options.custom_id
else None
)
@ -483,7 +484,7 @@ class SupermemoryOpenAIWrapper:
"Starting memory search",
{
"container_tag": self._container_tag,
"conversation_id": self._options.conversation_id,
"conversation_id": self._options.custom_id,
"mode": self._options.mode,
},
)
@ -617,8 +618,7 @@ class SupermemoryOpenAIWrapper:
def with_supermemory(
openai_client: Union[OpenAI, AsyncOpenAI],
container_tag: str,
options: Optional[OpenAIMiddlewareOptions] = None,
options: OpenAIMiddlewareOptions,
) -> Union[OpenAI, AsyncOpenAI]:
"""
Wraps an OpenAI client with SuperMemory middleware to automatically inject relevant memories
@ -630,8 +630,7 @@ def with_supermemory(
Args:
openai_client: The OpenAI client to wrap with SuperMemory middleware
container_tag: The container tag/identifier for memory search (e.g., user ID, project ID)
options: Optional configuration options for the middleware
options: Configuration options for the middleware (container_tag and custom_id are required)
Returns:
An OpenAI client with SuperMemory middleware injected
@ -645,9 +644,9 @@ def with_supermemory(
openai = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
openai_with_supermemory = with_supermemory(
openai,
"user-123",
OpenAIMiddlewareOptions(
conversation_id="conversation-456",
container_tag="user-123",
custom_id="conversation-456",
mode="full",
add_memory="always"
)
@ -666,6 +665,6 @@ def with_supermemory(
ValueError: When SUPERMEMORY_API_KEY environment variable is not set
Exception: When supermemory API request fails
"""
wrapper = SupermemoryOpenAIWrapper(openai_client, container_tag, options)
wrapper = SupermemoryOpenAIWrapper(openai_client, options)
# Return the wrapper, which delegates all attributes to the original client
return cast(Union[OpenAI, AsyncOpenAI], wrapper)

View file

@ -37,8 +37,9 @@ async def test_async_middleware():
# Wrap with Supermemory middleware
openai_with_memory = with_supermemory(
openai_client,
container_tag="test-user-123",
options=OpenAIMiddlewareOptions(
OpenAIMiddlewareOptions(
container_tag="test-user-123",
custom_id="test-integration",
mode="profile",
verbose=True,
add_memory="never" # Don't save test messages
@ -87,8 +88,9 @@ def test_sync_middleware():
# Wrap with Supermemory middleware
openai_with_memory = with_supermemory(
openai_client,
container_tag="test-user-sync-123",
options=OpenAIMiddlewareOptions(
OpenAIMiddlewareOptions(
container_tag="test-user-sync-123",
custom_id="test-integration-sync",
mode="profile",
verbose=True
)
@ -126,7 +128,7 @@ def test_error_handling():
openai_client = OpenAI(api_key="fake-key")
# This should raise SupermemoryConfigurationError
with_supermemory(openai_client, "test-user")
with_supermemory(openai_client, OpenAIMiddlewareOptions(container_tag="test-user", custom_id="test-conv"))
print("❌ Should have raised SupermemoryConfigurationError")
@ -156,8 +158,9 @@ def test_background_tasks():
# Wrap with memory storage enabled
wrapped_client = with_supermemory(
openai_client,
container_tag="test-background-tasks",
options=OpenAIMiddlewareOptions(
OpenAIMiddlewareOptions(
container_tag="test-background-tasks",
custom_id="test-background",
add_memory="always",
verbose=True
)

View file

@ -103,26 +103,32 @@ class TestMiddlewareInitialization:
def test_with_supermemory_basic(self, mock_openai_client):
"""Test basic middleware initialization."""
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
wrapped_client = with_supermemory(mock_openai_client, "user-123")
options = OpenAIMiddlewareOptions(
container_tag="user-123",
custom_id="test-conv"
)
wrapped_client = with_supermemory(mock_openai_client, options)
assert isinstance(wrapped_client, SupermemoryOpenAIWrapper)
assert wrapped_client._container_tag == "user-123"
assert wrapped_client._options.mode == "profile"
assert wrapped_client._options.verbose is False
assert wrapped_client._options.add_memory == "always" # New default
def test_with_supermemory_with_options(self, mock_openai_client):
"""Test middleware initialization with options."""
options = OpenAIMiddlewareOptions(
conversation_id="conv-456",
container_tag="user-123",
custom_id="conv-456",
verbose=True,
mode="full",
add_memory="always"
)
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
wrapped_client = with_supermemory(mock_openai_client, "user-123", options)
wrapped_client = with_supermemory(mock_openai_client, options)
assert wrapped_client._options.conversation_id == "conv-456"
assert wrapped_client._options.custom_id == "conv-456"
assert wrapped_client._options.verbose is True
assert wrapped_client._options.mode == "full"
assert wrapped_client._options.add_memory == "always"
@ -132,15 +138,23 @@ class TestMiddlewareInitialization:
from supermemory_openai.exceptions import SupermemoryConfigurationError
with patch.dict(os.environ, {}, clear=True):
options = OpenAIMiddlewareOptions(
container_tag="user-123",
custom_id="test-conv"
)
with pytest.raises(SupermemoryConfigurationError, match="SUPERMEMORY_API_KEY"):
with_supermemory(mock_openai_client, "user-123")
with_supermemory(mock_openai_client, options)
def test_wrapper_delegates_attributes(self, mock_openai_client):
"""Test that wrapper delegates attributes to wrapped client."""
mock_openai_client.models = Mock()
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
wrapped_client = with_supermemory(mock_openai_client, "user-123")
options = OpenAIMiddlewareOptions(
container_tag="user-123",
custom_id="test-conv"
)
wrapped_client = with_supermemory(mock_openai_client, options)
# Should delegate to the original client
assert wrapped_client.models is mock_openai_client.models
@ -165,8 +179,7 @@ class TestMemoryInjection:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(mode="profile")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", mode="profile")
)
messages = [
@ -207,8 +220,7 @@ class TestMemoryInjection:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(mode="query")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", mode="query")
)
messages = [
@ -241,8 +253,7 @@ class TestMemoryInjection:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(mode="full")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", mode="full")
)
messages = [
@ -278,7 +289,10 @@ class TestMemoryInjection:
mock_search.return_value.profile = mock_supermemory_response["profile"]
mock_search.return_value.search_results = mock_supermemory_response["searchResults"]
wrapped_client = with_supermemory(mock_async_openai_client, "user-123")
wrapped_client = with_supermemory(
mock_async_openai_client,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv")
)
messages = [
{"role": "system", "content": "You are a helpful assistant."},
@ -324,8 +338,7 @@ class TestMemoryStorage:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
)
messages = [
@ -359,8 +372,7 @@ class TestMemoryStorage:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="never")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="never")
)
await wrapped_client.chat.completions.create(
@ -386,7 +398,10 @@ class TestSyncAsyncCompatibility:
mock_search.return_value.profile = {"static": [], "dynamic": []}
mock_search.return_value.search_results = {"results": []}
wrapped_client = with_supermemory(mock_openai_client, "user-123")
wrapped_client = with_supermemory(
mock_openai_client,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv")
)
# This should work for sync clients too
wrapped_client.chat.completions.create(
@ -410,7 +425,10 @@ class TestSyncAsyncCompatibility:
mock_search.return_value.profile = {"static": [], "dynamic": []}
mock_search.return_value.search_results = {"results": []}
wrapped_client = with_supermemory(mock_openai_client, "user-123")
wrapped_client = with_supermemory(
mock_openai_client,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv")
)
# This should work even when called from async context
result = wrapped_client.chat.completions.create(
@ -441,8 +459,7 @@ class TestSyncAsyncCompatibility:
wrapped_client = with_supermemory(
mock_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
)
# Should not raise exception, should continue with main request
@ -470,7 +487,10 @@ class TestErrorHandling:
with patch("supermemory_openai.middleware.supermemory_profile_search") as mock_search:
mock_search.side_effect = Exception("API Error")
wrapped_client = with_supermemory(mock_async_openai_client, "user-123")
wrapped_client = with_supermemory(
mock_async_openai_client,
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv")
)
# Should not raise exception, should fall back gracefully
with pytest.raises(Exception):
@ -490,8 +510,7 @@ class TestErrorHandling:
with patch.dict(os.environ, {"SUPERMEMORY_API_KEY": "test-key"}):
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(mode="query")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", mode="query")
)
messages = [
@ -529,8 +548,7 @@ class TestLogging:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(verbose=True)
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", verbose=True)
)
await wrapped_client.chat.completions.create(
@ -558,8 +576,7 @@ class TestLogging:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(verbose=False)
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", verbose=False)
)
await wrapped_client.chat.completions.create(
@ -597,8 +614,7 @@ class TestBackgroundTaskManagement:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
)
# Make a request that should create a background task
@ -644,8 +660,7 @@ class TestBackgroundTaskManagement:
# Use async context manager
async with with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
) as wrapped_client:
await wrapped_client.chat.completions.create(
model="gpt-4",
@ -680,8 +695,7 @@ class TestBackgroundTaskManagement:
wrapped_client = with_supermemory(
mock_async_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
)
await wrapped_client.chat.completions.create(
@ -714,8 +728,7 @@ class TestBackgroundTaskManagement:
# Use sync context manager
with with_supermemory(
mock_openai_client,
"user-123",
OpenAIMiddlewareOptions(add_memory="always")
OpenAIMiddlewareOptions(container_tag="user-123", custom_id="test-conv", add_memory="always")
) as wrapped_client:
wrapped_client.chat.completions.create(
model="gpt-4",