osmosis-ai 0.1.6__py3-none-any.whl → 0.1.7__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.

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