osmosis-ai 0.1.6__py3-none-any.whl → 0.1.8__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 +44 -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 +52 -20
  11. {osmosis_ai-0.1.6.dist-info → osmosis_ai-0.1.8.dist-info}/METADATA +1 -1
  12. osmosis_ai-0.1.8.dist-info/RECORD +15 -0
  13. {osmosis_ai-0.1.6.dist-info → osmosis_ai-0.1.8.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.8.dist-info}/licenses/LICENSE +0 -0
  16. {osmosis_ai-0.1.6.dist-info → osmosis_ai-0.1.8.dist-info}/top_level.txt +0 -0
osmosis_ai/adapters/anthropic.py CHANGED
@@ -11,12 +11,13 @@ from osmosis_ai.utils import send_to_osmosis
11
11
  from osmosis_ai import utils
12
12
  from osmosis_ai.logger import logger
13
13
 
14
+
14
15
  def wrap_anthropic() -> None:
15
16
  """
16
17
  Monkey patch Anthropic's client to send all prompts and responses to OSMOSIS.
17
-
18
+
18
19
  This function should be called before creating any Anthropic client instances.
19
-
20
+
20
21
  Features supported:
21
22
  - Basic message completions
22
23
  - Async message completions
@@ -29,13 +30,13 @@ def wrap_anthropic() -> None:
29
30
  except ImportError:
30
31
  logger.debug("anthropic package is not installed.")
31
32
  return
32
-
33
+
33
34
  logger.info(f"Wrapping Anthropic client, package version: {anthropic.__version__}")
34
-
35
+
35
36
  # Check which version of Anthropic SDK we're dealing with
36
37
  # v1.x has a "resources" attribute, v0.x doesn't
37
- is_v1 = hasattr(anthropic, 'resources')
38
-
38
+ is_v1 = hasattr(anthropic, "resources")
39
+
39
40
  if is_v1:
40
41
  # Handle newer v1.x SDK
41
42
  logger.info("Detected Anthropic SDK v1.x")
@@ -44,162 +45,219 @@ def wrap_anthropic() -> None:
44
45
  # Handle older v0.x SDK
45
46
  logger.info("Detected Anthropic SDK v0.x")
46
47
  _ai_anthropic_v0(anthropic)
47
-
48
+
48
49
  logger.info("Anthropic client has been wrapped by osmosis-ai.")
49
50
 
51
+
50
52
  def _ai_anthropic_v1(anthropic_module):
51
53
  """Handle wrapping for newer v1.x SDK with resources structure"""
52
54
  try:
53
55
  # Get the resources.messages module and class
54
56
  messages_module = anthropic_module.resources.messages
55
57
  messages_class = messages_module.Messages
56
-
58
+
57
59
  logger.info(f"Found Anthropic messages class: {messages_class}")
58
60
 
59
61
  # Patch the Messages.create method
60
62
  original_messages_create = messages_class.create
61
63
  logger.info(f"Original create method: {original_messages_create}")
62
-
64
+
63
65
  if not hasattr(original_messages_create, "_osmosis_aiped"):
66
+
64
67
  @functools.wraps(original_messages_create)
65
68
  def wrapped_messages_create(self, *args, **kwargs):
66
- logger.debug(f"Wrapped create called with args: {args}, kwargs: {kwargs}")
69
+ logger.debug(
70
+ f"Wrapped create called with args: {args}, kwargs: {kwargs}"
71
+ )
67
72
  try:
68
73
  # Check if the call includes tool use parameters
69
74
  has_tools = "tools" in kwargs
70
75
  if has_tools:
71
- logger.debug(f"Tool use detected with {len(kwargs['tools'])} tools")
72
-
76
+ logger.debug(
77
+ f"Tool use detected with {len(kwargs['tools'])} tools"
78
+ )
79
+
73
80
  # Check if this is a tool response message
74
81
  if "messages" in kwargs:
75
82
  for message in kwargs["messages"]:
76
- if message.get("role") == "user" and isinstance(message.get("content"), list):
83
+ if message.get("role") == "user" and isinstance(
84
+ message.get("content"), list
85
+ ):
77
86
  for content_item in message["content"]:
78
- if isinstance(content_item, dict) and content_item.get("type") == "tool_response":
79
- logger.debug("Tool response detected in messages")
87
+ if (
88
+ isinstance(content_item, dict)
89
+ and content_item.get("type") == "tool_response"
90
+ ):
91
+ logger.debug(
92
+ "Tool response detected in messages"
93
+ )
80
94
  break
81
-
95
+
82
96
  response = original_messages_create(self, *args, **kwargs)
83
-
97
+
84
98
  if utils.enabled:
85
99
  logger.debug("Sending success to OSMOSIS (success)")
86
100
  send_to_osmosis(
87
101
  query=kwargs,
88
- response=response.model_dump() if hasattr(response, 'model_dump') else response,
89
- status=200
102
+ response=(
103
+ response.model_dump()
104
+ if hasattr(response, "model_dump")
105
+ else response
106
+ ),
107
+ status=200,
90
108
  )
91
-
109
+
92
110
  return response
93
111
  except Exception as e:
94
112
  logger.error(f"Error in wrapped create: {e}")
95
113
  if utils.enabled:
96
114
  error_response = {"error": str(e)}
97
115
  send_to_osmosis(
98
- query=kwargs,
99
- response=error_response,
100
- status=400
116
+ query=kwargs, response=error_response, status=400
101
117
  )
102
118
  logger.debug("Sending error to OSMOSIS (success)")
103
119
  raise # Re-raise the exception
104
-
120
+
105
121
  wrapped_messages_create._osmosis_aiped = True
106
122
  messages_class.create = wrapped_messages_create
107
123
  logger.info("Successfully wrapped Messages.create method")
108
-
124
+
109
125
  # Directly wrap the AsyncAnthropic client
110
126
  try:
111
127
  # Get the AsyncAnthropic class
112
128
  AsyncAnthropicClass = anthropic_module.AsyncAnthropic
113
129
  logger.info(f"Found AsyncAnthropic class: {AsyncAnthropicClass}")
114
-
130
+
115
131
  # Store the original __init__ to keep track of created instances
116
132
  original_async_init = AsyncAnthropicClass.__init__
117
-
133
+
118
134
  if not hasattr(original_async_init, "_osmosis_aiped"):
135
+
119
136
  @functools.wraps(original_async_init)
120
137
  def wrapped_async_init(self, *args, **kwargs):
121
138
  # Call the original init
122
139
  result = original_async_init(self, *args, **kwargs)
123
-
124
- logger.info("Wrapping new AsyncAnthropic instance's messages.create method")
125
-
140
+
141
+ logger.info(
142
+ "Wrapping new AsyncAnthropic instance's messages.create method"
143
+ )
144
+
126
145
  # Get the messages client from this instance
127
146
  async_messages = self.messages
128
-
147
+
129
148
  # Store and patch the create method if not already wrapped
130
- if hasattr(async_messages, "create") and not hasattr(async_messages.create, "_osmosis_aiped"):
149
+ if hasattr(async_messages, "create") and not hasattr(
150
+ async_messages.create, "_osmosis_aiped"
151
+ ):
131
152
  original_async_messages_create = async_messages.create
132
-
153
+
133
154
  @functools.wraps(original_async_messages_create)
134
155
  async def wrapped_async_messages_create(*args, **kwargs):
135
- logger.debug(f"AsyncAnthropic.messages.create called with args: {args}, kwargs: {kwargs}")
156
+ logger.debug(
157
+ f"AsyncAnthropic.messages.create called with args: {args}, kwargs: {kwargs}"
158
+ )
136
159
  try:
137
- response = await original_async_messages_create(*args, **kwargs)
138
-
160
+ response = await original_async_messages_create(
161
+ *args, **kwargs
162
+ )
163
+
139
164
  if utils.enabled:
140
- logger.debug("Sending AsyncAnthropic response to OSMOSIS (success)")
165
+ logger.debug(
166
+ "Sending AsyncAnthropic response to OSMOSIS (success)"
167
+ )
141
168
  send_to_osmosis(
142
169
  query=kwargs,
143
- response=response.model_dump() if hasattr(response, 'model_dump') else response,
144
- status=200
170
+ response=(
171
+ response.model_dump()
172
+ if hasattr(response, "model_dump")
173
+ else response
174
+ ),
175
+ status=200,
145
176
  )
146
-
177
+
147
178
  return response
148
179
  except Exception as e:
149
- logger.error(f"Error in wrapped AsyncAnthropic.messages.create: {e}")
180
+ logger.error(
181
+ f"Error in wrapped AsyncAnthropic.messages.create: {e}"
182
+ )
150
183
  if utils.enabled:
151
- logger.debug("Sending AsyncAnthropic error to OSMOSIS")
184
+ logger.debug(
185
+ "Sending AsyncAnthropic error to OSMOSIS"
186
+ )
152
187
  error_response = {"error": str(e)}
153
188
  send_to_osmosis(
154
189
  query=kwargs,
155
190
  response=error_response,
156
- status=400
191
+ status=400,
157
192
  )
158
193
  raise # Re-raise the exception
159
-
194
+
160
195
  wrapped_async_messages_create._osmosis_aiped = True
161
196
  async_messages.create = wrapped_async_messages_create
162
- logger.info("Successfully wrapped AsyncAnthropic.messages.create method")
163
-
197
+ logger.info(
198
+ "Successfully wrapped AsyncAnthropic.messages.create method"
199
+ )
200
+
164
201
  return result
165
-
202
+
166
203
  wrapped_async_init._osmosis_aiped = True
167
204
  AsyncAnthropicClass.__init__ = wrapped_async_init
168
- logger.info("Successfully wrapped AsyncAnthropic.__init__ to patch message methods on new instances")
205
+ logger.info(
206
+ "Successfully wrapped AsyncAnthropic.__init__ to patch message methods on new instances"
207
+ )
169
208
  except (ImportError, AttributeError) as e:
170
- logger.warning(f"AsyncAnthropic class not found or has unexpected structure: {e}")
171
-
209
+ logger.warning(
210
+ f"AsyncAnthropic class not found or has unexpected structure: {e}"
211
+ )
212
+
172
213
  # For compatibility, still try to patch the old-style acreate method if it exists
173
214
  if hasattr(messages_class, "acreate"):
174
215
  original_acreate = messages_class.acreate
175
216
  if not hasattr(original_acreate, "_osmosis_aiped"):
217
+
176
218
  @functools.wraps(original_acreate)
177
219
  async def wrapped_acreate(self, *args, **kwargs):
178
- logger.debug(f"Wrapped async create called with args: {args}, kwargs: {kwargs}")
220
+ logger.debug(
221
+ f"Wrapped async create called with args: {args}, kwargs: {kwargs}"
222
+ )
179
223
  try:
180
224
  # Check if the async call includes tool use parameters
181
225
  has_tools = "tools" in kwargs
182
226
  if has_tools:
183
- logger.debug(f"Async tool use detected with {len(kwargs['tools'])} tools")
184
-
227
+ logger.debug(
228
+ f"Async tool use detected with {len(kwargs['tools'])} tools"
229
+ )
230
+
185
231
  if "messages" in kwargs:
186
232
  for message in kwargs["messages"]:
187
- if message.get("role") == "user" and isinstance(message.get("content"), list):
233
+ if message.get("role") == "user" and isinstance(
234
+ message.get("content"), list
235
+ ):
188
236
  for content_item in message["content"]:
189
- if isinstance(content_item, dict) and content_item.get("type") == "tool_response":
190
- logger.debug("Async tool response detected in messages")
237
+ if (
238
+ isinstance(content_item, dict)
239
+ and content_item.get("type")
240
+ == "tool_response"
241
+ ):
242
+ logger.debug(
243
+ "Async tool response detected in messages"
244
+ )
191
245
  break
192
-
246
+
193
247
  response = await original_acreate(self, *args, **kwargs)
194
-
248
+
195
249
  if utils.enabled:
196
250
  logger.debug("Sending async response to OSMOSIS (success)")
197
251
  send_to_osmosis(
198
252
  query=kwargs,
199
- response=response.model_dump() if hasattr(response, 'model_dump') else response,
200
- status=200
253
+ response=(
254
+ response.model_dump()
255
+ if hasattr(response, "model_dump")
256
+ else response
257
+ ),
258
+ status=200,
201
259
  )
202
-
260
+
203
261
  return response
204
262
  except Exception as e:
205
263
  logger.error(f"Error in wrapped async create: {e}")
@@ -207,12 +265,10 @@ def _ai_anthropic_v1(anthropic_module):
207
265
  logger.debug("Sending async error to OSMOSIS")
208
266
  error_response = {"error": str(e)}
209
267
  send_to_osmosis(
210
- query=kwargs,
211
- response=error_response,
212
- status=400
268
+ query=kwargs, response=error_response, status=400
213
269
  )
214
270
  raise # Re-raise the exception
215
-
271
+
216
272
  wrapped_acreate._osmosis_aiped = True
217
273
  messages_class.acreate = wrapped_acreate
218
274
  logger.info("Successfully wrapped Messages.acreate method")
@@ -223,56 +279,76 @@ def _ai_anthropic_v1(anthropic_module):
223
279
  try:
224
280
  completions_module = anthropic_module.resources.completions
225
281
  completions_class = completions_module.Completions
226
-
282
+
227
283
  original_completions_create = completions_class.create
228
284
  if not hasattr(original_completions_create, "_osmosis_aiped"):
285
+
229
286
  @functools.wraps(original_completions_create)
230
287
  def wrapped_completions_create(self, *args, **kwargs):
231
288
  response = original_completions_create(self, *args, **kwargs)
232
-
289
+
233
290
  if utils.enabled:
234
291
  send_to_osmosis(
235
292
  query=kwargs,
236
- response=response.model_dump() if hasattr(response, 'model_dump') else response,
237
- status=200
293
+ response=(
294
+ response.model_dump()
295
+ if hasattr(response, "model_dump")
296
+ else response
297
+ ),
298
+ status=200,
238
299
  )
239
-
300
+
240
301
  return response
241
-
302
+
242
303
  wrapped_completions_create._osmosis_aiped = True
243
304
  completions_class.create = wrapped_completions_create
244
-
305
+
245
306
  # Patch the async create method if it exists
246
307
  if hasattr(completions_class, "acreate"):
247
308
  original_completions_acreate = completions_class.acreate
248
309
  if not hasattr(original_completions_acreate, "_osmosis_aiped"):
310
+
249
311
  @functools.wraps(original_completions_acreate)
250
312
  async def wrapped_completions_acreate(self, *args, **kwargs):
251
- logger.debug(f"Wrapped Completions async create called with args: {args}, kwargs: {kwargs}")
313
+ logger.debug(
314
+ f"Wrapped Completions async create called with args: {args}, kwargs: {kwargs}"
315
+ )
252
316
  try:
253
- response = await original_completions_acreate(self, *args, **kwargs)
254
-
317
+ response = await original_completions_acreate(
318
+ self, *args, **kwargs
319
+ )
320
+
255
321
  if utils.enabled:
256
- logger.debug("Sending Completions async response to OSMOSIS (success)")
322
+ logger.debug(
323
+ "Sending Completions async response to OSMOSIS (success)"
324
+ )
257
325
  send_to_osmosis(
258
326
  query=kwargs,
259
- response=response.model_dump() if hasattr(response, 'model_dump') else response,
260
- status=200
327
+ response=(
328
+ response.model_dump()
329
+ if hasattr(response, "model_dump")
330
+ else response
331
+ ),
332
+ status=200,
261
333
  )
262
-
334
+
263
335
  return response
264
336
  except Exception as e:
265
- logger.error(f"Error in wrapped Completions async create: {e}")
337
+ logger.error(
338
+ f"Error in wrapped Completions async create: {e}"
339
+ )
266
340
  if utils.enabled:
267
- logger.debug("Sending Completions async error to OSMOSIS")
341
+ logger.debug(
342
+ "Sending Completions async error to OSMOSIS"
343
+ )
268
344
  error_response = {"error": str(e)}
269
345
  send_to_osmosis(
270
346
  query=kwargs,
271
347
  response=error_response,
272
- status=400
348
+ status=400,
273
349
  )
274
350
  raise # Re-raise the exception
275
-
351
+
276
352
  wrapped_completions_acreate._osmosis_aiped = True
277
353
  completions_class.acreate = wrapped_completions_acreate
278
354
  logger.info("Successfully wrapped Completions.acreate method")
@@ -284,119 +360,121 @@ def _ai_anthropic_v1(anthropic_module):
284
360
  logger.info("Completions.create already wrapped")
285
361
  except (ImportError, AttributeError) as e:
286
362
  # Completions module may not exist in this version
287
- logger.warning(f"Completions module not found or has an unexpected structure: {e}")
363
+ logger.warning(
364
+ f"Completions module not found or has an unexpected structure: {e}"
365
+ )
288
366
  except Exception as e:
289
367
  logger.error(f"Error wrapping Anthropic v1.x client: {e}")
290
368
 
369
+
291
370
  def _ai_anthropic_v0(anthropic_module):
292
371
  """Handle wrapping for older v0.x SDK without resources structure"""
293
372
  try:
294
373
  # Get the main Anthropic class
295
374
  AnthropicClass = anthropic_module.Anthropic
296
375
  logger.info(f"Found Anthropic class: {AnthropicClass}")
297
-
376
+
298
377
  # Patch the create_completion method for v0.x
299
378
  if hasattr(AnthropicClass, "complete"):
300
379
  original_complete = AnthropicClass.complete
301
-
380
+
302
381
  if not hasattr(original_complete, "_osmosis_aiped"):
382
+
303
383
  @functools.wraps(original_complete)
304
384
  def wrapped_complete(self, *args, **kwargs):
305
- logger.debug(f"Wrapped complete called with args: {args}, kwargs: {kwargs}")
385
+ logger.debug(
386
+ f"Wrapped complete called with args: {args}, kwargs: {kwargs}"
387
+ )
306
388
  try:
307
389
  response = original_complete(self, *args, **kwargs)
308
-
390
+
309
391
  if utils.enabled:
310
392
  logger.debug("Sending success to OSMOSIS (success)")
311
- send_to_osmosis(
312
- query=kwargs,
313
- response=response,
314
- status=200
315
- )
316
-
393
+ send_to_osmosis(query=kwargs, response=response, status=200)
394
+
317
395
  return response
318
396
  except Exception as e:
319
397
  logger.error(f"Error in wrapped complete: {e}")
320
398
  if utils.enabled:
321
399
  error_response = {"error": str(e)}
322
400
  send_to_osmosis(
323
- query=kwargs,
324
- response=error_response,
325
- status=400
401
+ query=kwargs, response=error_response, status=400
326
402
  )
327
403
  logger.debug("Sending error to OSMOSIS (success)")
328
404
  raise # Re-raise the exception
329
-
405
+
330
406
  wrapped_complete._osmosis_aiped = True
331
407
  AnthropicClass.complete = wrapped_complete
332
408
  logger.info("Successfully wrapped Anthropic.complete method")
333
-
409
+
334
410
  # Patch the create_message method for v0.x if it exists
335
411
  if hasattr(AnthropicClass, "messages"):
336
412
  logger.info("Found messages client on Anthropic class")
337
-
413
+
338
414
  if hasattr(AnthropicClass.messages, "create"):
339
415
  original_messages_create = AnthropicClass.messages.create
340
-
416
+
341
417
  if not hasattr(original_messages_create, "_osmosis_aiped"):
418
+
342
419
  @functools.wraps(original_messages_create)
343
420
  def wrapped_messages_create(self, *args, **kwargs):
344
- logger.debug(f"Wrapped messages.create called with args: {args}, kwargs: {kwargs}")
421
+ logger.debug(
422
+ f"Wrapped messages.create called with args: {args}, kwargs: {kwargs}"
423
+ )
345
424
  try:
346
425
  response = original_messages_create(self, *args, **kwargs)
347
-
426
+
348
427
  if utils.enabled:
349
428
  logger.debug("Sending success to OSMOSIS (success)")
350
429
  send_to_osmosis(
351
- query=kwargs,
352
- response=response,
353
- status=200
430
+ query=kwargs, response=response, status=200
354
431
  )
355
-
432
+
356
433
  return response
357
434
  except Exception as e:
358
435
  logger.error(f"Error in wrapped messages.create: {e}")
359
436
  if utils.enabled:
360
437
  error_response = {"error": str(e)}
361
438
  send_to_osmosis(
362
- query=kwargs,
363
- response=error_response,
364
- status=400
439
+ query=kwargs, response=error_response, status=400
365
440
  )
366
441
  logger.debug("Sending error to OSMOSIS (success)")
367
442
  raise # Re-raise the exception
368
-
443
+
369
444
  wrapped_messages_create._osmosis_aiped = True
370
445
  AnthropicClass.messages.create = wrapped_messages_create
371
446
  logger.info("Successfully wrapped Anthropic.messages.create method")
372
-
447
+
373
448
  # Also try to patch instance methods by monkeypatching the __init__
374
449
  original_init = AnthropicClass.__init__
375
-
450
+
376
451
  if not hasattr(original_init, "_osmosis_aiped"):
452
+
377
453
  @functools.wraps(original_init)
378
454
  def wrapped_init(self, *args, **kwargs):
379
455
  # Call the original init
380
456
  result = original_init(self, *args, **kwargs)
381
-
457
+
382
458
  # Wrap the instance methods if they exist
383
- if hasattr(self, "complete") and not hasattr(self.complete, "_osmosis_aiped"):
459
+ if hasattr(self, "complete") and not hasattr(
460
+ self.complete, "_osmosis_aiped"
461
+ ):
384
462
  original_instance_complete = self.complete
385
-
463
+
386
464
  @functools.wraps(original_instance_complete)
387
465
  def wrapped_instance_complete(*args, **kwargs):
388
- logger.debug(f"Instance complete called with args: {args}, kwargs: {kwargs}")
466
+ logger.debug(
467
+ f"Instance complete called with args: {args}, kwargs: {kwargs}"
468
+ )
389
469
  try:
390
470
  response = original_instance_complete(*args, **kwargs)
391
-
471
+
392
472
  if utils.enabled:
393
473
  logger.debug("Sending success to OSMOSIS (success)")
394
474
  send_to_osmosis(
395
- query=kwargs,
396
- response=response,
397
- status=200
475
+ query=kwargs, response=response, status=200
398
476
  )
399
-
477
+
400
478
  return response
401
479
  except Exception as e:
402
480
  logger.error(f"Error in wrapped instance complete: {e}")
@@ -405,18 +483,20 @@ def _ai_anthropic_v0(anthropic_module):
405
483
  send_to_osmosis(
406
484
  query=kwargs,
407
485
  response=error_response,
408
- status=400
486
+ status=400,
409
487
  )
410
488
  raise
411
-
489
+
412
490
  wrapped_instance_complete._osmosis_aiped = True
413
491
  self.complete = wrapped_instance_complete
414
-
492
+
415
493
  return result
416
-
494
+
417
495
  wrapped_init._osmosis_aiped = True
418
496
  AnthropicClass.__init__ = wrapped_init
419
- logger.info("Successfully wrapped Anthropic.__init__ to patch instance methods")
420
-
497
+ logger.info(
498
+ "Successfully wrapped Anthropic.__init__ to patch instance methods"
499
+ )
500
+
421
501
  except Exception as e:
422
- logger.error(f"Error wrapping Anthropic v0.x client: {e}")
502
+ logger.error(f"Error wrapping Anthropic v0.x client: {e}")