osmosis-ai 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of osmosis-ai might be problematic. Click here for more details.
- osmosis_ai/__init__.py +4 -1
- osmosis_ai/consts.py +2 -16
- osmosis_ai/utils.py +52 -87
- osmosis_ai-0.1.9.dist-info/METADATA +143 -0
- osmosis_ai-0.1.9.dist-info/RECORD +8 -0
- {osmosis_ai-0.1.7.dist-info → osmosis_ai-0.1.9.dist-info}/WHEEL +1 -1
- {osmosis_ai-0.1.7.dist-info → osmosis_ai-0.1.9.dist-info}/licenses/LICENSE +1 -1
- osmosis_ai/adapters/__init__.py +0 -9
- osmosis_ai/adapters/anthropic.py +0 -502
- osmosis_ai/adapters/langchain.py +0 -674
- osmosis_ai/adapters/langchain_anthropic.py +0 -338
- osmosis_ai/adapters/langchain_openai.py +0 -596
- osmosis_ai/adapters/openai.py +0 -900
- osmosis_ai/logger.py +0 -77
- osmosis_ai-0.1.7.dist-info/METADATA +0 -281
- osmosis_ai-0.1.7.dist-info/RECORD +0 -15
- {osmosis_ai-0.1.7.dist-info → osmosis_ai-0.1.9.dist-info}/top_level.txt +0 -0
osmosis_ai/adapters/openai.py
DELETED
|
@@ -1,900 +0,0 @@
|
|
|
1
|
-
"""
|
|
2
|
-
OpenAI adapter for Osmosis
|
|
3
|
-
|
|
4
|
-
This module provides monkey patching for the OpenAI Python client.
|
|
5
|
-
"""
|
|
6
|
-
|
|
7
|
-
import functools
|
|
8
|
-
import inspect
|
|
9
|
-
import sys
|
|
10
|
-
|
|
11
|
-
from osmosis_ai import utils
|
|
12
|
-
from osmosis_ai.utils import send_to_osmosis
|
|
13
|
-
from osmosis_ai.logger import logger
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
def wrap_openai() -> None:
|
|
17
|
-
"""
|
|
18
|
-
Monkey patch OpenAI's client to send all prompts and responses to OSMOSIS.
|
|
19
|
-
|
|
20
|
-
This function should be called before creating any OpenAI client instances.
|
|
21
|
-
"""
|
|
22
|
-
try:
|
|
23
|
-
import openai
|
|
24
|
-
except ImportError:
|
|
25
|
-
logger.debug("openai package is not installed.")
|
|
26
|
-
return
|
|
27
|
-
|
|
28
|
-
# Try to detect which version of the OpenAI client is installed
|
|
29
|
-
try:
|
|
30
|
-
from openai import OpenAI
|
|
31
|
-
|
|
32
|
-
# Check for v2 client first
|
|
33
|
-
try:
|
|
34
|
-
import openai.version
|
|
35
|
-
|
|
36
|
-
if openai.version.__version__.startswith("2."):
|
|
37
|
-
_ai_openai_v2()
|
|
38
|
-
return
|
|
39
|
-
except (ImportError, AttributeError):
|
|
40
|
-
pass
|
|
41
|
-
|
|
42
|
-
# Fall back to v1 client
|
|
43
|
-
_ai_openai_v1()
|
|
44
|
-
except (ImportError, AttributeError):
|
|
45
|
-
# Fall back to legacy client
|
|
46
|
-
_ai_openai_legacy()
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
def _ai_openai_v2() -> None:
|
|
50
|
-
"""Monkey patch the OpenAI v2 client."""
|
|
51
|
-
import openai
|
|
52
|
-
|
|
53
|
-
try:
|
|
54
|
-
# Get the OpenAI class
|
|
55
|
-
from openai import OpenAI
|
|
56
|
-
|
|
57
|
-
# Debug: Print all available OpenAI modules
|
|
58
|
-
logger.debug(f"OpenAI modules: {dir(openai)}")
|
|
59
|
-
|
|
60
|
-
# Try to import AsyncOpenAI
|
|
61
|
-
try:
|
|
62
|
-
from openai import AsyncOpenAI
|
|
63
|
-
|
|
64
|
-
logger.info(f"Successfully imported AsyncOpenAI")
|
|
65
|
-
has_async_client = True
|
|
66
|
-
except ImportError:
|
|
67
|
-
logger.warning(f"Failed to import AsyncOpenAI")
|
|
68
|
-
has_async_client = False
|
|
69
|
-
|
|
70
|
-
# Store the original __init__ method for OpenAI
|
|
71
|
-
original_init = OpenAI.__init__
|
|
72
|
-
|
|
73
|
-
@functools.wraps(original_init)
|
|
74
|
-
def wrapped_init(self, *args, **kwargs):
|
|
75
|
-
# Call the original __init__
|
|
76
|
-
original_init(self, *args, **kwargs)
|
|
77
|
-
|
|
78
|
-
# Debug: Print client structure
|
|
79
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
80
|
-
logger.debug(
|
|
81
|
-
f"Client chat completions methods: {dir(self.chat.completions)}"
|
|
82
|
-
)
|
|
83
|
-
if hasattr(self.chat.completions, "create"):
|
|
84
|
-
logger.debug(
|
|
85
|
-
f"create is a coro: {inspect.iscoroutinefunction(self.chat.completions.create)}"
|
|
86
|
-
)
|
|
87
|
-
|
|
88
|
-
# Now wrap the client's chat.completions.create and completions.create methods
|
|
89
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
90
|
-
original_chat_create = self.chat.completions.create
|
|
91
|
-
if not hasattr(original_chat_create, "_osmosis_aiped"):
|
|
92
|
-
|
|
93
|
-
@functools.wraps(original_chat_create)
|
|
94
|
-
def wrapped_chat_create(*args, **kwargs):
|
|
95
|
-
# Check if streaming is enabled
|
|
96
|
-
is_streaming = kwargs.get("stream", False)
|
|
97
|
-
|
|
98
|
-
if is_streaming:
|
|
99
|
-
# For streaming, we need to wrap the iterator
|
|
100
|
-
stream = original_chat_create(*args, **kwargs)
|
|
101
|
-
|
|
102
|
-
if utils.enabled:
|
|
103
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
104
|
-
def wrapped_stream():
|
|
105
|
-
chunks = []
|
|
106
|
-
for chunk in stream:
|
|
107
|
-
chunks.append(
|
|
108
|
-
chunk.model_dump()
|
|
109
|
-
if hasattr(chunk, "model_dump")
|
|
110
|
-
else chunk
|
|
111
|
-
)
|
|
112
|
-
yield chunk
|
|
113
|
-
|
|
114
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
115
|
-
if utils.enabled:
|
|
116
|
-
send_to_osmosis(
|
|
117
|
-
query=kwargs,
|
|
118
|
-
response={"streaming_chunks": chunks},
|
|
119
|
-
status=200,
|
|
120
|
-
)
|
|
121
|
-
|
|
122
|
-
return wrapped_stream()
|
|
123
|
-
else:
|
|
124
|
-
return stream
|
|
125
|
-
else:
|
|
126
|
-
# For non-streaming, handle normally
|
|
127
|
-
response = original_chat_create(*args, **kwargs)
|
|
128
|
-
|
|
129
|
-
if utils.enabled:
|
|
130
|
-
send_to_osmosis(
|
|
131
|
-
query=kwargs,
|
|
132
|
-
response=(
|
|
133
|
-
response.model_dump()
|
|
134
|
-
if hasattr(response, "model_dump")
|
|
135
|
-
else response
|
|
136
|
-
),
|
|
137
|
-
status=200,
|
|
138
|
-
)
|
|
139
|
-
|
|
140
|
-
return response
|
|
141
|
-
|
|
142
|
-
wrapped_chat_create._osmosis_aiped = True
|
|
143
|
-
self.chat.completions.create = wrapped_chat_create
|
|
144
|
-
|
|
145
|
-
if hasattr(self, "completions"):
|
|
146
|
-
original_completions_create = self.completions.create
|
|
147
|
-
if not hasattr(original_completions_create, "_osmosis_aiped"):
|
|
148
|
-
|
|
149
|
-
@functools.wraps(original_completions_create)
|
|
150
|
-
def wrapped_completions_create(*args, **kwargs):
|
|
151
|
-
# Check if streaming is enabled
|
|
152
|
-
is_streaming = kwargs.get("stream", False)
|
|
153
|
-
|
|
154
|
-
if is_streaming:
|
|
155
|
-
# For streaming, we need to wrap the iterator
|
|
156
|
-
stream = original_completions_create(*args, **kwargs)
|
|
157
|
-
|
|
158
|
-
if utils.enabled:
|
|
159
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
160
|
-
def wrapped_stream():
|
|
161
|
-
chunks = []
|
|
162
|
-
for chunk in stream:
|
|
163
|
-
chunks.append(
|
|
164
|
-
chunk.model_dump()
|
|
165
|
-
if hasattr(chunk, "model_dump")
|
|
166
|
-
else chunk
|
|
167
|
-
)
|
|
168
|
-
yield chunk
|
|
169
|
-
|
|
170
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
171
|
-
if utils.enabled:
|
|
172
|
-
send_to_osmosis(
|
|
173
|
-
query=kwargs,
|
|
174
|
-
response={"streaming_chunks": chunks},
|
|
175
|
-
status=200,
|
|
176
|
-
)
|
|
177
|
-
|
|
178
|
-
return wrapped_stream()
|
|
179
|
-
else:
|
|
180
|
-
return stream
|
|
181
|
-
else:
|
|
182
|
-
# For non-streaming, handle normally
|
|
183
|
-
response = original_completions_create(*args, **kwargs)
|
|
184
|
-
|
|
185
|
-
if utils.enabled:
|
|
186
|
-
send_to_osmosis(
|
|
187
|
-
query=kwargs,
|
|
188
|
-
response=(
|
|
189
|
-
response.model_dump()
|
|
190
|
-
if hasattr(response, "model_dump")
|
|
191
|
-
else response
|
|
192
|
-
),
|
|
193
|
-
status=200,
|
|
194
|
-
)
|
|
195
|
-
|
|
196
|
-
return response
|
|
197
|
-
|
|
198
|
-
wrapped_completions_create._osmosis_aiped = True
|
|
199
|
-
self.completions.create = wrapped_completions_create
|
|
200
|
-
|
|
201
|
-
# Wrap async methods
|
|
202
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
203
|
-
if hasattr(self.chat.completions, "acreate"):
|
|
204
|
-
logger.debug(f"Found acreate in chat.completions")
|
|
205
|
-
original_achat_create = self.chat.completions.acreate
|
|
206
|
-
if not hasattr(original_achat_create, "_osmosis_aiped"):
|
|
207
|
-
|
|
208
|
-
@functools.wraps(original_achat_create)
|
|
209
|
-
async def wrapped_achat_create(*args, **kwargs):
|
|
210
|
-
# Check if streaming is enabled
|
|
211
|
-
is_streaming = kwargs.get("stream", False)
|
|
212
|
-
|
|
213
|
-
if is_streaming:
|
|
214
|
-
# For streaming, we need to wrap the async iterator
|
|
215
|
-
stream = await original_achat_create(*args, **kwargs)
|
|
216
|
-
|
|
217
|
-
if utils.enabled:
|
|
218
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
219
|
-
async def wrapped_stream():
|
|
220
|
-
chunks = []
|
|
221
|
-
async for chunk in stream:
|
|
222
|
-
chunks.append(
|
|
223
|
-
chunk.model_dump()
|
|
224
|
-
if hasattr(chunk, "model_dump")
|
|
225
|
-
else chunk
|
|
226
|
-
)
|
|
227
|
-
yield chunk
|
|
228
|
-
|
|
229
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
230
|
-
if utils.enabled:
|
|
231
|
-
send_to_osmosis(
|
|
232
|
-
query=kwargs,
|
|
233
|
-
response={"streaming_chunks": chunks},
|
|
234
|
-
status=200,
|
|
235
|
-
)
|
|
236
|
-
|
|
237
|
-
return wrapped_stream()
|
|
238
|
-
else:
|
|
239
|
-
return stream
|
|
240
|
-
else:
|
|
241
|
-
# For non-streaming, handle normally
|
|
242
|
-
response = await original_achat_create(*args, **kwargs)
|
|
243
|
-
|
|
244
|
-
if utils.enabled:
|
|
245
|
-
send_to_osmosis(
|
|
246
|
-
query=kwargs,
|
|
247
|
-
response=(
|
|
248
|
-
response.model_dump()
|
|
249
|
-
if hasattr(response, "model_dump")
|
|
250
|
-
else response
|
|
251
|
-
),
|
|
252
|
-
status=200,
|
|
253
|
-
)
|
|
254
|
-
|
|
255
|
-
return response
|
|
256
|
-
|
|
257
|
-
wrapped_achat_create._osmosis_aiped = True
|
|
258
|
-
self.chat.completions.acreate = wrapped_achat_create
|
|
259
|
-
else:
|
|
260
|
-
logger.debug(f"acreate not found in chat.completions")
|
|
261
|
-
|
|
262
|
-
if hasattr(self, "completions"):
|
|
263
|
-
if hasattr(self.completions, "acreate"):
|
|
264
|
-
original_acompletions_create = self.completions.acreate
|
|
265
|
-
if not hasattr(original_acompletions_create, "_osmosis_aiped"):
|
|
266
|
-
|
|
267
|
-
@functools.wraps(original_acompletions_create)
|
|
268
|
-
async def wrapped_acompletions_create(*args, **kwargs):
|
|
269
|
-
# Check if streaming is enabled
|
|
270
|
-
is_streaming = kwargs.get("stream", False)
|
|
271
|
-
|
|
272
|
-
if is_streaming:
|
|
273
|
-
# For streaming, we need to wrap the async iterator
|
|
274
|
-
stream = await original_acompletions_create(
|
|
275
|
-
*args, **kwargs
|
|
276
|
-
)
|
|
277
|
-
|
|
278
|
-
if utils.enabled:
|
|
279
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
280
|
-
async def wrapped_stream():
|
|
281
|
-
chunks = []
|
|
282
|
-
async for chunk in stream:
|
|
283
|
-
chunks.append(
|
|
284
|
-
chunk.model_dump()
|
|
285
|
-
if hasattr(chunk, "model_dump")
|
|
286
|
-
else chunk
|
|
287
|
-
)
|
|
288
|
-
yield chunk
|
|
289
|
-
|
|
290
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
291
|
-
if utils.enabled:
|
|
292
|
-
send_to_osmosis(
|
|
293
|
-
query=kwargs,
|
|
294
|
-
response={"streaming_chunks": chunks},
|
|
295
|
-
status=200,
|
|
296
|
-
)
|
|
297
|
-
|
|
298
|
-
return wrapped_stream()
|
|
299
|
-
else:
|
|
300
|
-
return stream
|
|
301
|
-
else:
|
|
302
|
-
# For non-streaming, handle normally
|
|
303
|
-
response = await original_acompletions_create(
|
|
304
|
-
*args, **kwargs
|
|
305
|
-
)
|
|
306
|
-
|
|
307
|
-
if utils.enabled:
|
|
308
|
-
send_to_osmosis(
|
|
309
|
-
query=kwargs,
|
|
310
|
-
response=(
|
|
311
|
-
response.model_dump()
|
|
312
|
-
if hasattr(response, "model_dump")
|
|
313
|
-
else response
|
|
314
|
-
),
|
|
315
|
-
status=200,
|
|
316
|
-
)
|
|
317
|
-
|
|
318
|
-
return response
|
|
319
|
-
|
|
320
|
-
wrapped_acompletions_create._osmosis_aiped = True
|
|
321
|
-
self.completions.acreate = wrapped_acompletions_create
|
|
322
|
-
|
|
323
|
-
wrapped_init._osmosis_aiped = True
|
|
324
|
-
OpenAI.__init__ = wrapped_init
|
|
325
|
-
|
|
326
|
-
# Also wrap AsyncOpenAI if it exists
|
|
327
|
-
if has_async_client:
|
|
328
|
-
logger.debug(f"Wrapping AsyncOpenAI __init__")
|
|
329
|
-
original_async_init = AsyncOpenAI.__init__
|
|
330
|
-
|
|
331
|
-
@functools.wraps(original_async_init)
|
|
332
|
-
def wrapped_async_init(self, *args, **kwargs):
|
|
333
|
-
# Call the original __init__
|
|
334
|
-
original_async_init(self, *args, **kwargs)
|
|
335
|
-
|
|
336
|
-
# Debug: Print AsyncOpenAI client structure
|
|
337
|
-
logger.debug(f"AsyncOpenAI client structure:")
|
|
338
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
339
|
-
logger.debug(
|
|
340
|
-
f"AsyncOpenAI chat completions methods: {dir(self.chat.completions)}"
|
|
341
|
-
)
|
|
342
|
-
if hasattr(self.chat.completions, "create"):
|
|
343
|
-
logger.debug(
|
|
344
|
-
f"create is a coro: {inspect.iscoroutinefunction(self.chat.completions.create)}"
|
|
345
|
-
)
|
|
346
|
-
|
|
347
|
-
# Wrap the async client's methods
|
|
348
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
349
|
-
original_achat_create = self.chat.completions.create
|
|
350
|
-
if not hasattr(original_achat_create, "_osmosis_aiped"):
|
|
351
|
-
logger.debug(f"Wrapping AsyncOpenAI chat.completions.create")
|
|
352
|
-
|
|
353
|
-
@functools.wraps(original_achat_create)
|
|
354
|
-
async def wrapped_achat_create(*args, **kwargs):
|
|
355
|
-
logger.debug(f"AsyncOpenAI wrapped create called")
|
|
356
|
-
response = await original_achat_create(*args, **kwargs)
|
|
357
|
-
|
|
358
|
-
if utils.enabled:
|
|
359
|
-
send_to_osmosis(
|
|
360
|
-
query=kwargs,
|
|
361
|
-
response=(
|
|
362
|
-
response.model_dump()
|
|
363
|
-
if hasattr(response, "model_dump")
|
|
364
|
-
else response
|
|
365
|
-
),
|
|
366
|
-
status=200,
|
|
367
|
-
)
|
|
368
|
-
|
|
369
|
-
return response
|
|
370
|
-
|
|
371
|
-
wrapped_achat_create._osmosis_aiped = True
|
|
372
|
-
self.chat.completions.create = wrapped_achat_create
|
|
373
|
-
|
|
374
|
-
if hasattr(self, "completions"):
|
|
375
|
-
original_acompletions_create = self.completions.create
|
|
376
|
-
if not hasattr(original_acompletions_create, "_osmosis_aiped"):
|
|
377
|
-
|
|
378
|
-
@functools.wraps(original_acompletions_create)
|
|
379
|
-
async def wrapped_acompletions_create(*args, **kwargs):
|
|
380
|
-
response = await original_acompletions_create(
|
|
381
|
-
*args, **kwargs
|
|
382
|
-
)
|
|
383
|
-
|
|
384
|
-
if utils.enabled:
|
|
385
|
-
send_to_osmosis(
|
|
386
|
-
query=kwargs,
|
|
387
|
-
response=(
|
|
388
|
-
response.model_dump()
|
|
389
|
-
if hasattr(response, "model_dump")
|
|
390
|
-
else response
|
|
391
|
-
),
|
|
392
|
-
status=200,
|
|
393
|
-
)
|
|
394
|
-
|
|
395
|
-
return response
|
|
396
|
-
|
|
397
|
-
wrapped_acompletions_create._osmosis_aiped = True
|
|
398
|
-
self.completions.create = wrapped_acompletions_create
|
|
399
|
-
|
|
400
|
-
wrapped_async_init._osmosis_aiped = True
|
|
401
|
-
AsyncOpenAI.__init__ = wrapped_async_init
|
|
402
|
-
|
|
403
|
-
logger.info("OpenAI v2 client has been wrapped by osmosis-ai.")
|
|
404
|
-
except (ImportError, AttributeError) as e:
|
|
405
|
-
logger.error(f"Failed to wrap OpenAI v2 client: {e}")
|
|
406
|
-
|
|
407
|
-
|
|
408
|
-
def _ai_openai_v1() -> None:
|
|
409
|
-
"""Monkey patch the OpenAI v1 client."""
|
|
410
|
-
from openai import OpenAI
|
|
411
|
-
from openai.resources.chat import completions
|
|
412
|
-
from openai.resources import completions as text_completions
|
|
413
|
-
|
|
414
|
-
# Print package structure to debug
|
|
415
|
-
logger.debug(f"OpenAI package structure in v1 wrapper:")
|
|
416
|
-
try:
|
|
417
|
-
import openai
|
|
418
|
-
|
|
419
|
-
logger.debug(f"OpenAI modules: {dir(openai)}")
|
|
420
|
-
except Exception as e:
|
|
421
|
-
logger.error(f"Error inspecting openai module: {e}")
|
|
422
|
-
|
|
423
|
-
# Try to import AsyncOpenAI
|
|
424
|
-
try:
|
|
425
|
-
from openai import AsyncOpenAI
|
|
426
|
-
|
|
427
|
-
logger.info(f"Successfully imported AsyncOpenAI in v1 wrapper")
|
|
428
|
-
has_async_client = True
|
|
429
|
-
except ImportError:
|
|
430
|
-
logger.warning(f"Failed to import AsyncOpenAI in v1 wrapper")
|
|
431
|
-
has_async_client = False
|
|
432
|
-
|
|
433
|
-
# Print available methods in completions module
|
|
434
|
-
logger.debug(f"Available methods in completions: {dir(completions.Completions)}")
|
|
435
|
-
|
|
436
|
-
# Patch the chat completions create method
|
|
437
|
-
original_chat_create = completions.Completions.create
|
|
438
|
-
if not hasattr(original_chat_create, "_osmosis_aiped"):
|
|
439
|
-
|
|
440
|
-
@functools.wraps(original_chat_create)
|
|
441
|
-
def wrapped_chat_create(self, *args, **kwargs):
|
|
442
|
-
# Check if streaming is enabled
|
|
443
|
-
is_streaming = kwargs.get("stream", False)
|
|
444
|
-
|
|
445
|
-
if is_streaming:
|
|
446
|
-
# For streaming, we need to wrap the iterator
|
|
447
|
-
stream = original_chat_create(self, *args, **kwargs)
|
|
448
|
-
|
|
449
|
-
if utils.enabled:
|
|
450
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
451
|
-
def wrapped_stream():
|
|
452
|
-
chunks = []
|
|
453
|
-
for chunk in stream:
|
|
454
|
-
chunks.append(
|
|
455
|
-
chunk.model_dump()
|
|
456
|
-
if hasattr(chunk, "model_dump")
|
|
457
|
-
else chunk
|
|
458
|
-
)
|
|
459
|
-
yield chunk
|
|
460
|
-
|
|
461
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
462
|
-
if utils.enabled:
|
|
463
|
-
send_to_osmosis(
|
|
464
|
-
query=kwargs,
|
|
465
|
-
response={"streaming_chunks": chunks},
|
|
466
|
-
status=200,
|
|
467
|
-
)
|
|
468
|
-
|
|
469
|
-
return wrapped_stream()
|
|
470
|
-
else:
|
|
471
|
-
return stream
|
|
472
|
-
else:
|
|
473
|
-
# For non-streaming, handle normally
|
|
474
|
-
response = original_chat_create(self, *args, **kwargs)
|
|
475
|
-
|
|
476
|
-
if utils.enabled:
|
|
477
|
-
send_to_osmosis(
|
|
478
|
-
query=kwargs,
|
|
479
|
-
response=(
|
|
480
|
-
response.model_dump()
|
|
481
|
-
if hasattr(response, "model_dump")
|
|
482
|
-
else response
|
|
483
|
-
),
|
|
484
|
-
status=200,
|
|
485
|
-
)
|
|
486
|
-
|
|
487
|
-
return response
|
|
488
|
-
|
|
489
|
-
wrapped_chat_create._osmosis_aiped = True
|
|
490
|
-
completions.Completions.create = wrapped_chat_create
|
|
491
|
-
|
|
492
|
-
# Patch the completions create method
|
|
493
|
-
original_completions_create = text_completions.Completions.create
|
|
494
|
-
if not hasattr(original_completions_create, "_osmosis_aiped"):
|
|
495
|
-
|
|
496
|
-
@functools.wraps(original_completions_create)
|
|
497
|
-
def wrapped_completions_create(self, *args, **kwargs):
|
|
498
|
-
# Check if streaming is enabled
|
|
499
|
-
is_streaming = kwargs.get("stream", False)
|
|
500
|
-
|
|
501
|
-
if is_streaming:
|
|
502
|
-
# For streaming, we need to wrap the iterator
|
|
503
|
-
stream = original_completions_create(self, *args, **kwargs)
|
|
504
|
-
|
|
505
|
-
if utils.enabled:
|
|
506
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
507
|
-
def wrapped_stream():
|
|
508
|
-
chunks = []
|
|
509
|
-
for chunk in stream:
|
|
510
|
-
chunks.append(
|
|
511
|
-
chunk.model_dump()
|
|
512
|
-
if hasattr(chunk, "model_dump")
|
|
513
|
-
else chunk
|
|
514
|
-
)
|
|
515
|
-
yield chunk
|
|
516
|
-
|
|
517
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
518
|
-
if utils.enabled:
|
|
519
|
-
send_to_osmosis(
|
|
520
|
-
query=kwargs,
|
|
521
|
-
response={"streaming_chunks": chunks},
|
|
522
|
-
status=200,
|
|
523
|
-
)
|
|
524
|
-
|
|
525
|
-
return wrapped_stream()
|
|
526
|
-
else:
|
|
527
|
-
return stream
|
|
528
|
-
else:
|
|
529
|
-
# For non-streaming, handle normally
|
|
530
|
-
response = original_completions_create(self, *args, **kwargs)
|
|
531
|
-
|
|
532
|
-
if utils.enabled:
|
|
533
|
-
send_to_osmosis(
|
|
534
|
-
query=kwargs,
|
|
535
|
-
response=(
|
|
536
|
-
response.model_dump()
|
|
537
|
-
if hasattr(response, "model_dump")
|
|
538
|
-
else response
|
|
539
|
-
),
|
|
540
|
-
status=200,
|
|
541
|
-
)
|
|
542
|
-
|
|
543
|
-
return response
|
|
544
|
-
|
|
545
|
-
wrapped_completions_create._osmosis_aiped = True
|
|
546
|
-
text_completions.Completions.create = wrapped_completions_create
|
|
547
|
-
|
|
548
|
-
# Find and wrap async methods
|
|
549
|
-
for module in [completions, text_completions]:
|
|
550
|
-
for name, method in inspect.getmembers(module.Completions):
|
|
551
|
-
if (
|
|
552
|
-
name.startswith("a")
|
|
553
|
-
and name.endswith("create")
|
|
554
|
-
and inspect.iscoroutinefunction(method)
|
|
555
|
-
and not hasattr(method, "_osmosis_aiped")
|
|
556
|
-
):
|
|
557
|
-
|
|
558
|
-
logger.debug(f"Found async method {name} in {module.__name__}")
|
|
559
|
-
original_method = method
|
|
560
|
-
|
|
561
|
-
@functools.wraps(original_method)
|
|
562
|
-
async def wrapped_async_method(self, *args, **kwargs):
|
|
563
|
-
# Check if streaming is enabled
|
|
564
|
-
is_streaming = kwargs.get("stream", False)
|
|
565
|
-
|
|
566
|
-
if is_streaming:
|
|
567
|
-
# For streaming, we need to wrap the async iterator
|
|
568
|
-
stream = await original_method(self, *args, **kwargs)
|
|
569
|
-
|
|
570
|
-
if utils.enabled:
|
|
571
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
572
|
-
async def wrapped_stream():
|
|
573
|
-
chunks = []
|
|
574
|
-
async for chunk in stream:
|
|
575
|
-
chunks.append(
|
|
576
|
-
chunk.model_dump()
|
|
577
|
-
if hasattr(chunk, "model_dump")
|
|
578
|
-
else chunk
|
|
579
|
-
)
|
|
580
|
-
yield chunk
|
|
581
|
-
|
|
582
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
583
|
-
if utils.enabled:
|
|
584
|
-
send_to_osmosis(
|
|
585
|
-
query=kwargs,
|
|
586
|
-
response={"streaming_chunks": chunks},
|
|
587
|
-
status=200,
|
|
588
|
-
)
|
|
589
|
-
|
|
590
|
-
return wrapped_stream()
|
|
591
|
-
else:
|
|
592
|
-
return stream
|
|
593
|
-
else:
|
|
594
|
-
# For non-streaming, handle normally
|
|
595
|
-
response = await original_method(self, *args, **kwargs)
|
|
596
|
-
|
|
597
|
-
if utils.enabled:
|
|
598
|
-
send_to_osmosis(
|
|
599
|
-
query=kwargs,
|
|
600
|
-
response=(
|
|
601
|
-
response.model_dump()
|
|
602
|
-
if hasattr(response, "model_dump")
|
|
603
|
-
else response
|
|
604
|
-
),
|
|
605
|
-
status=200,
|
|
606
|
-
)
|
|
607
|
-
|
|
608
|
-
return response
|
|
609
|
-
|
|
610
|
-
wrapped_async_method._osmosis_aiped = True
|
|
611
|
-
setattr(module.Completions, name, wrapped_async_method)
|
|
612
|
-
|
|
613
|
-
# Explicitly wrap AsyncOpenAI if it exists
|
|
614
|
-
if has_async_client:
|
|
615
|
-
logger.debug(f"Wrapping AsyncOpenAI __init__ in v1 wrapper")
|
|
616
|
-
original_async_init = AsyncOpenAI.__init__
|
|
617
|
-
|
|
618
|
-
@functools.wraps(original_async_init)
|
|
619
|
-
def wrapped_async_init(self, *args, **kwargs):
|
|
620
|
-
# Call the original __init__
|
|
621
|
-
original_async_init(self, *args, **kwargs)
|
|
622
|
-
|
|
623
|
-
# Debug: Print AsyncOpenAI client structure
|
|
624
|
-
logger.debug(f"AsyncOpenAI client structure in v1:")
|
|
625
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
626
|
-
logger.debug(
|
|
627
|
-
f"AsyncOpenAI chat completions methods: {dir(self.chat.completions)}"
|
|
628
|
-
)
|
|
629
|
-
if hasattr(self.chat.completions, "create"):
|
|
630
|
-
logger.debug(
|
|
631
|
-
f"create is a coro: {inspect.iscoroutinefunction(self.chat.completions.create)}"
|
|
632
|
-
)
|
|
633
|
-
|
|
634
|
-
# Now wrap the async client's methods
|
|
635
|
-
if hasattr(self, "chat") and hasattr(self.chat, "completions"):
|
|
636
|
-
original_achat_create = self.chat.completions.create
|
|
637
|
-
if not hasattr(original_achat_create, "_osmosis_aiped"):
|
|
638
|
-
logger.debug(f"Wrapping AsyncOpenAI chat.completions.create in v1")
|
|
639
|
-
|
|
640
|
-
@functools.wraps(original_achat_create)
|
|
641
|
-
async def wrapped_achat_create(*args, **kwargs):
|
|
642
|
-
logger.debug(f"AsyncOpenAI v1 wrapped create called")
|
|
643
|
-
|
|
644
|
-
# Check if streaming is enabled
|
|
645
|
-
is_streaming = kwargs.get("stream", False)
|
|
646
|
-
|
|
647
|
-
if is_streaming:
|
|
648
|
-
# For streaming, we need to wrap the async iterator
|
|
649
|
-
stream = await original_achat_create(*args, **kwargs)
|
|
650
|
-
|
|
651
|
-
if utils.enabled:
|
|
652
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
653
|
-
async def wrapped_stream():
|
|
654
|
-
chunks = []
|
|
655
|
-
async for chunk in stream:
|
|
656
|
-
chunks.append(
|
|
657
|
-
chunk.model_dump()
|
|
658
|
-
if hasattr(chunk, "model_dump")
|
|
659
|
-
else chunk
|
|
660
|
-
)
|
|
661
|
-
yield chunk
|
|
662
|
-
|
|
663
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
664
|
-
if utils.enabled:
|
|
665
|
-
send_to_osmosis(
|
|
666
|
-
query=kwargs,
|
|
667
|
-
response={"streaming_chunks": chunks},
|
|
668
|
-
status=200,
|
|
669
|
-
)
|
|
670
|
-
|
|
671
|
-
return wrapped_stream()
|
|
672
|
-
else:
|
|
673
|
-
return stream
|
|
674
|
-
else:
|
|
675
|
-
# For non-streaming, handle normally
|
|
676
|
-
response = await original_achat_create(*args, **kwargs)
|
|
677
|
-
|
|
678
|
-
if utils.enabled:
|
|
679
|
-
send_to_osmosis(
|
|
680
|
-
query=kwargs,
|
|
681
|
-
response=(
|
|
682
|
-
response.model_dump()
|
|
683
|
-
if hasattr(response, "model_dump")
|
|
684
|
-
else response
|
|
685
|
-
),
|
|
686
|
-
status=200,
|
|
687
|
-
)
|
|
688
|
-
|
|
689
|
-
return response
|
|
690
|
-
|
|
691
|
-
wrapped_achat_create._osmosis_aiped = True
|
|
692
|
-
self.chat.completions.create = wrapped_achat_create
|
|
693
|
-
|
|
694
|
-
if hasattr(self, "completions"):
|
|
695
|
-
original_acompletions_create = self.completions.create
|
|
696
|
-
if not hasattr(original_acompletions_create, "_osmosis_aiped"):
|
|
697
|
-
|
|
698
|
-
@functools.wraps(original_acompletions_create)
|
|
699
|
-
async def wrapped_acompletions_create(*args, **kwargs):
|
|
700
|
-
# Check if streaming is enabled
|
|
701
|
-
is_streaming = kwargs.get("stream", False)
|
|
702
|
-
|
|
703
|
-
if is_streaming:
|
|
704
|
-
# For streaming, we need to wrap the async iterator
|
|
705
|
-
stream = await original_acompletions_create(*args, **kwargs)
|
|
706
|
-
|
|
707
|
-
if utils.enabled:
|
|
708
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
709
|
-
async def wrapped_stream():
|
|
710
|
-
chunks = []
|
|
711
|
-
async for chunk in stream:
|
|
712
|
-
chunks.append(
|
|
713
|
-
chunk.model_dump()
|
|
714
|
-
if hasattr(chunk, "model_dump")
|
|
715
|
-
else chunk
|
|
716
|
-
)
|
|
717
|
-
yield chunk
|
|
718
|
-
|
|
719
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
720
|
-
if utils.enabled:
|
|
721
|
-
send_to_osmosis(
|
|
722
|
-
query=kwargs,
|
|
723
|
-
response={"streaming_chunks": chunks},
|
|
724
|
-
status=200,
|
|
725
|
-
)
|
|
726
|
-
|
|
727
|
-
return wrapped_stream()
|
|
728
|
-
else:
|
|
729
|
-
return stream
|
|
730
|
-
else:
|
|
731
|
-
# For non-streaming, handle normally
|
|
732
|
-
response = await original_acompletions_create(
|
|
733
|
-
*args, **kwargs
|
|
734
|
-
)
|
|
735
|
-
|
|
736
|
-
if utils.enabled:
|
|
737
|
-
send_to_osmosis(
|
|
738
|
-
query=kwargs,
|
|
739
|
-
response=(
|
|
740
|
-
response.model_dump()
|
|
741
|
-
if hasattr(response, "model_dump")
|
|
742
|
-
else response
|
|
743
|
-
),
|
|
744
|
-
status=200,
|
|
745
|
-
)
|
|
746
|
-
|
|
747
|
-
return response
|
|
748
|
-
|
|
749
|
-
wrapped_acompletions_create._osmosis_aiped = True
|
|
750
|
-
self.completions.create = wrapped_acompletions_create
|
|
751
|
-
|
|
752
|
-
wrapped_async_init._osmosis_aiped = True
|
|
753
|
-
AsyncOpenAI.__init__ = wrapped_async_init
|
|
754
|
-
|
|
755
|
-
logger.info("OpenAI v1 client has been wrapped by osmosis-ai.")
|
|
756
|
-
|
|
757
|
-
|
|
758
|
-
def _ai_openai_legacy() -> None:
|
|
759
|
-
"""Monkey patch the legacy OpenAI client."""
|
|
760
|
-
import openai
|
|
761
|
-
|
|
762
|
-
# Patch the Completion.create method
|
|
763
|
-
original_completion_create = openai.Completion.create
|
|
764
|
-
if not hasattr(original_completion_create, "_osmosis_aiped"):
|
|
765
|
-
|
|
766
|
-
@functools.wraps(original_completion_create)
|
|
767
|
-
def wrapped_completion_create(*args, **kwargs):
|
|
768
|
-
# Check if streaming is enabled
|
|
769
|
-
is_streaming = kwargs.get("stream", False)
|
|
770
|
-
|
|
771
|
-
if is_streaming:
|
|
772
|
-
# For streaming, we need to wrap the iterator
|
|
773
|
-
stream = original_completion_create(*args, **kwargs)
|
|
774
|
-
|
|
775
|
-
if utils.enabled:
|
|
776
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
777
|
-
def wrapped_stream():
|
|
778
|
-
chunks = []
|
|
779
|
-
for chunk in stream:
|
|
780
|
-
chunks.append(chunk)
|
|
781
|
-
yield chunk
|
|
782
|
-
|
|
783
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
784
|
-
if utils.enabled:
|
|
785
|
-
send_to_osmosis(
|
|
786
|
-
query=kwargs,
|
|
787
|
-
response={"streaming_chunks": chunks},
|
|
788
|
-
status=200,
|
|
789
|
-
)
|
|
790
|
-
|
|
791
|
-
return wrapped_stream()
|
|
792
|
-
else:
|
|
793
|
-
return stream
|
|
794
|
-
else:
|
|
795
|
-
# For non-streaming, handle normally
|
|
796
|
-
response = original_completion_create(*args, **kwargs)
|
|
797
|
-
|
|
798
|
-
if utils.enabled:
|
|
799
|
-
send_to_osmosis(query=kwargs, response=response, status=200)
|
|
800
|
-
|
|
801
|
-
return response
|
|
802
|
-
|
|
803
|
-
wrapped_completion_create._osmosis_aiped = True
|
|
804
|
-
openai.Completion.create = wrapped_completion_create
|
|
805
|
-
|
|
806
|
-
# Patch the ChatCompletion.create method
|
|
807
|
-
if hasattr(openai, "ChatCompletion"):
|
|
808
|
-
original_chat_create = openai.ChatCompletion.create
|
|
809
|
-
if not hasattr(original_chat_create, "_osmosis_aiped"):
|
|
810
|
-
|
|
811
|
-
@functools.wraps(original_chat_create)
|
|
812
|
-
def wrapped_chat_create(*args, **kwargs):
|
|
813
|
-
# Check if streaming is enabled
|
|
814
|
-
is_streaming = kwargs.get("stream", False)
|
|
815
|
-
|
|
816
|
-
if is_streaming:
|
|
817
|
-
# For streaming, we need to wrap the iterator
|
|
818
|
-
stream = original_chat_create(*args, **kwargs)
|
|
819
|
-
|
|
820
|
-
if utils.enabled:
|
|
821
|
-
# Create a new wrapped iterator that sends data to OSMOSIS
|
|
822
|
-
def wrapped_stream():
|
|
823
|
-
chunks = []
|
|
824
|
-
for chunk in stream:
|
|
825
|
-
chunks.append(chunk)
|
|
826
|
-
yield chunk
|
|
827
|
-
|
|
828
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
829
|
-
if utils.enabled:
|
|
830
|
-
send_to_osmosis(
|
|
831
|
-
query=kwargs,
|
|
832
|
-
response={"streaming_chunks": chunks},
|
|
833
|
-
status=200,
|
|
834
|
-
)
|
|
835
|
-
|
|
836
|
-
return wrapped_stream()
|
|
837
|
-
else:
|
|
838
|
-
return stream
|
|
839
|
-
else:
|
|
840
|
-
# For non-streaming, handle normally
|
|
841
|
-
response = original_chat_create(*args, **kwargs)
|
|
842
|
-
|
|
843
|
-
if utils.enabled:
|
|
844
|
-
send_to_osmosis(query=kwargs, response=response, status=200)
|
|
845
|
-
|
|
846
|
-
return response
|
|
847
|
-
|
|
848
|
-
wrapped_chat_create._osmosis_aiped = True
|
|
849
|
-
openai.ChatCompletion.create = wrapped_chat_create
|
|
850
|
-
|
|
851
|
-
# Patch the async methods if they exist
|
|
852
|
-
for obj in [openai.Completion, getattr(openai, "ChatCompletion", None)]:
|
|
853
|
-
if obj is None:
|
|
854
|
-
continue
|
|
855
|
-
|
|
856
|
-
if hasattr(obj, "acreate"):
|
|
857
|
-
original_acreate = obj.acreate
|
|
858
|
-
if not hasattr(original_acreate, "_osmosis_aiped"):
|
|
859
|
-
|
|
860
|
-
@functools.wraps(original_acreate)
|
|
861
|
-
async def wrapped_acreate(*args, **kwargs):
|
|
862
|
-
# Check if streaming is enabled
|
|
863
|
-
is_streaming = kwargs.get("stream", False)
|
|
864
|
-
|
|
865
|
-
if is_streaming:
|
|
866
|
-
# For streaming, we need to wrap the async iterator
|
|
867
|
-
stream = await original_acreate(*args, **kwargs)
|
|
868
|
-
|
|
869
|
-
if utils.enabled:
|
|
870
|
-
# Create a new wrapped async iterator that sends data to OSMOSIS
|
|
871
|
-
async def wrapped_stream():
|
|
872
|
-
chunks = []
|
|
873
|
-
async for chunk in stream:
|
|
874
|
-
chunks.append(chunk)
|
|
875
|
-
yield chunk
|
|
876
|
-
|
|
877
|
-
# After collecting all chunks, send them to OSMOSIS
|
|
878
|
-
if utils.enabled:
|
|
879
|
-
send_to_osmosis(
|
|
880
|
-
query=kwargs,
|
|
881
|
-
response={"streaming_chunks": chunks},
|
|
882
|
-
status=200,
|
|
883
|
-
)
|
|
884
|
-
|
|
885
|
-
return wrapped_stream()
|
|
886
|
-
else:
|
|
887
|
-
return stream
|
|
888
|
-
else:
|
|
889
|
-
# For non-streaming, handle normally
|
|
890
|
-
response = await original_acreate(*args, **kwargs)
|
|
891
|
-
|
|
892
|
-
if utils.enabled:
|
|
893
|
-
send_to_osmosis(query=kwargs, response=response, status=200)
|
|
894
|
-
|
|
895
|
-
return response
|
|
896
|
-
|
|
897
|
-
wrapped_acreate._osmosis_aiped = True
|
|
898
|
-
obj.acreate = wrapped_acreate
|
|
899
|
-
|
|
900
|
-
logger.info("OpenAI legacy client has been wrapped by osmosis-ai.")
|