1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
|
#!/usr/bin/env python3
"""
Integration test for Supermemory OpenAI middleware.
This script demonstrates how to test the middleware with real API calls.
Set your API keys as environment variables to run this test.
"""
import asyncio
import os
from openai import AsyncOpenAI, OpenAI
from supermemory_openai import (
with_supermemory,
OpenAIMiddlewareOptions,
SupermemoryConfigurationError,
SupermemoryAPIError,
)
async def test_async_middleware():
"""Test async middleware functionality."""
print("๐ Testing Async Middleware...")
try:
# Check for required environment variables
if not os.getenv("OPENAI_API_KEY"):
print("โ OPENAI_API_KEY not set - skipping OpenAI test")
return
if not os.getenv("SUPERMEMORY_API_KEY"):
print("โ SUPERMEMORY_API_KEY not set - skipping Supermemory test")
return
# Create OpenAI client
openai_client = AsyncOpenAI()
# Wrap with Supermemory middleware
openai_with_memory = with_supermemory(
openai_client,
container_tag="test-user-123",
options=OpenAIMiddlewareOptions(
mode="profile",
verbose=True,
add_memory="never" # Don't save test messages
)
)
# Test context manager
async with openai_with_memory as client:
print("โ
Context manager works")
# Make a test request
response = await client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "Hello! This is a test message."}
],
max_tokens=50
)
print(f"โ
API call successful: {response.choices[0].message.content[:50]}...")
except SupermemoryConfigurationError as e:
print(f"โ ๏ธ Configuration error: {e}")
except SupermemoryAPIError as e:
print(f"โ ๏ธ Supermemory API error: {e}")
except Exception as e:
print(f"โ Unexpected error: {e}")
def test_sync_middleware():
"""Test sync middleware functionality."""
print("\n๐ Testing Sync Middleware...")
try:
if not os.getenv("OPENAI_API_KEY"):
print("โ OPENAI_API_KEY not set - skipping OpenAI test")
return
if not os.getenv("SUPERMEMORY_API_KEY"):
print("โ SUPERMEMORY_API_KEY not set - skipping Supermemory test")
return
# Create sync OpenAI client
openai_client = OpenAI()
# Wrap with Supermemory middleware
openai_with_memory = with_supermemory(
openai_client,
container_tag="test-user-sync-123",
options=OpenAIMiddlewareOptions(
mode="profile",
verbose=True
)
)
# Test context manager
with openai_with_memory as client:
print("โ
Sync context manager works")
# Make a test request
response = client.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "user", "content": "This is a sync test message."}
],
max_tokens=50
)
print(f"โ
Sync API call successful: {response.choices[0].message.content[:50]}...")
except SupermemoryConfigurationError as e:
print(f"โ ๏ธ Configuration error: {e}")
except SupermemoryAPIError as e:
print(f"โ ๏ธ Supermemory API error: {e}")
except Exception as e:
print(f"โ Unexpected error: {e}")
def test_error_handling():
"""Test error handling without API keys."""
print("\n๐ Testing Error Handling...")
try:
# Test with missing API key
openai_client = OpenAI(api_key="fake-key")
# This should raise SupermemoryConfigurationError
with_supermemory(openai_client, "test-user")
print("โ Should have raised SupermemoryConfigurationError")
except SupermemoryConfigurationError as e:
print(f"โ
Correctly caught configuration error: {e}")
except Exception as e:
print(f"โ Wrong exception type: {type(e).__name__}: {e}")
def test_background_tasks():
"""Test background task management."""
print("\n๐ Testing Background Task Management...")
try:
if not os.getenv("SUPERMEMORY_API_KEY"):
print("โ SUPERMEMORY_API_KEY not set - skipping background task test")
return
# Create a fake OpenAI client for testing
from unittest.mock import Mock, AsyncMock
openai_client = Mock()
openai_client.chat = Mock()
openai_client.chat.completions = Mock()
openai_client.chat.completions.create = AsyncMock(return_value=Mock())
# Wrap with memory storage enabled
wrapped_client = with_supermemory(
openai_client,
container_tag="test-background-tasks",
options=OpenAIMiddlewareOptions(
add_memory="always",
verbose=True
)
)
print(f"โ
Background tasks tracking: {len(wrapped_client._background_tasks)} tasks")
except Exception as e:
print(f"โ Background task test error: {e}")
async def main():
"""Run all tests."""
print("๐งช Supermemory OpenAI Middleware Integration Tests")
print("=" * 60)
# Test async middleware
await test_async_middleware()
# Test sync middleware
test_sync_middleware()
# Test error handling
test_error_handling()
# Test background tasks
test_background_tasks()
print("\n" + "=" * 60)
print("๐ Integration tests completed!")
print("\n๐ก To run with real API calls, set these environment variables:")
print(" export OPENAI_API_KEY='your-openai-key'")
print(" export SUPERMEMORY_API_KEY='your-supermemory-key'")
if __name__ == "__main__":
asyncio.run(main())
|