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/langchain_anthropic.py CHANGED
@@ -11,10 +11,11 @@ from osmosis_ai import utils
11
11
  from osmosis_ai.utils import send_to_osmosis
12
12
  from osmosis_ai.logger import logger
13
13
 
14
+
14
15
  def wrap_langchain_anthropic() -> None:
15
16
  """
16
17
  Monkey patch langchain-anthropic's models to send all prompts and responses to OSMOSIS.
17
-
18
+
18
19
  This function should be called before using any langchain-anthropic models.
19
20
  """
20
21
  try:
@@ -22,53 +23,70 @@ def wrap_langchain_anthropic() -> None:
22
23
  except ImportError:
23
24
  logger.debug("langchain-anthropic package is not installed.")
24
25
  return
25
-
26
+
26
27
  _patch_anthropic_chat_models()
27
-
28
+
28
29
  logger.info("langchain-anthropic has been wrapped by osmosis-ai.")
29
30
 
31
+
30
32
  def _patch_anthropic_chat_models() -> None:
31
33
  """Patch langchain-anthropic chat model classes to send data to OSMOSIS."""
32
34
  try:
33
35
  # Try to import ChatAnthropic class
34
36
  try:
35
37
  from langchain_anthropic import ChatAnthropic
38
+
36
39
  logger.info("Successfully imported ChatAnthropic from langchain_anthropic")
37
40
  except ImportError:
38
41
  # Handle older versions if needed
39
42
  try:
40
43
  from langchain.chat_models.anthropic import ChatAnthropic
44
+
41
45
  logger.info("Found ChatAnthropic in langchain.chat_models.anthropic")
42
46
  except ImportError:
43
- logger.warning("Could not find ChatAnthropic class in any expected location.")
47
+ logger.warning(
48
+ "Could not find ChatAnthropic class in any expected location."
49
+ )
44
50
  return
45
-
51
+
46
52
  # Log available methods on ChatAnthropic for debugging
47
- chat_methods = [method for method in dir(ChatAnthropic) if not method.startswith('__') or method in ['_generate', '_agenerate', '_call', '_acall']]
53
+ chat_methods = [
54
+ method
55
+ for method in dir(ChatAnthropic)
56
+ if not method.startswith("__")
57
+ or method in ["_generate", "_agenerate", "_call", "_acall"]
58
+ ]
48
59
  logger.info(f"Found the following methods on ChatAnthropic: {chat_methods}")
49
-
60
+
50
61
  # Try to get the model attribute name - could be model or model_name
51
62
  model_attr = None
52
- for attr in ['model', 'model_name']:
63
+ for attr in ["model", "model_name"]:
53
64
  if hasattr(ChatAnthropic, attr):
54
65
  model_attr = attr
55
66
  logger.info(f"ChatAnthropic uses '{attr}' attribute for model name")
56
67
  break
57
-
68
+
58
69
  if not model_attr:
59
- model_attr = 'model' # Default to 'model' if we can't determine it
60
- logger.info(f"Could not determine model attribute name, defaulting to '{model_attr}'")
61
-
70
+ model_attr = "model" # Default to 'model' if we can't determine it
71
+ logger.info(
72
+ f"Could not determine model attribute name, defaulting to '{model_attr}'"
73
+ )
74
+
62
75
  # Patch the _generate method if it exists
63
76
  if hasattr(ChatAnthropic, "_generate"):
64
77
  original_generate = ChatAnthropic._generate
65
-
78
+
66
79
  if not hasattr(original_generate, "_osmosis_aiped"):
80
+
67
81
  @functools.wraps(original_generate)
68
- def wrapped_generate(self, messages, stop=None, run_manager=None, **kwargs):
82
+ def wrapped_generate(
83
+ self, messages, stop=None, run_manager=None, **kwargs
84
+ ):
69
85
  # Get the response
70
- response = original_generate(self, messages, stop=stop, run_manager=run_manager, **kwargs)
71
-
86
+ response = original_generate(
87
+ self, messages, stop=stop, run_manager=run_manager, **kwargs
88
+ )
89
+
72
90
  # Send to OSMOSIS if enabled
73
91
  if utils.enabled:
74
92
  # Create payload
@@ -76,19 +94,27 @@ def _patch_anthropic_chat_models() -> None:
76
94
  payload = {
77
95
  "model_type": "ChatAnthropic",
78
96
  "model_name": model_name,
79
- "messages": [str(msg) for msg in messages], # Convert to strings for serialization
80
- "response": str(response), # Convert to string since it may not be serializable
81
- "kwargs": {"stop": stop, **kwargs}
97
+ "messages": [
98
+ str(msg) for msg in messages
99
+ ], # Convert to strings for serialization
100
+ "response": str(
101
+ response
102
+ ), # Convert to string since it may not be serializable
103
+ "kwargs": {"stop": stop, **kwargs},
82
104
  }
83
-
105
+
84
106
  send_to_osmosis(
85
- query={"type": "langchain_anthropic_generate", "messages": [str(msg) for msg in messages], "model": model_name},
107
+ query={
108
+ "type": "langchain_anthropic_generate",
109
+ "messages": [str(msg) for msg in messages],
110
+ "model": model_name,
111
+ },
86
112
  response=payload,
87
- status=200
113
+ status=200,
88
114
  )
89
-
115
+
90
116
  return response
91
-
117
+
92
118
  wrapped_generate._osmosis_aiped = True
93
119
  ChatAnthropic._generate = wrapped_generate
94
120
  logger.info("Successfully wrapped ChatAnthropic._generate method")
@@ -96,17 +122,22 @@ def _patch_anthropic_chat_models() -> None:
96
122
  logger.info("ChatAnthropic._generate already wrapped.")
97
123
  else:
98
124
  logger.info("ChatAnthropic does not have a _generate method, skipping.")
99
-
125
+
100
126
  # Patch the _agenerate method if it exists
101
127
  if hasattr(ChatAnthropic, "_agenerate"):
102
128
  original_agenerate = ChatAnthropic._agenerate
103
-
129
+
104
130
  if not hasattr(original_agenerate, "_osmosis_aiped"):
131
+
105
132
  @functools.wraps(original_agenerate)
106
- async def wrapped_agenerate(self, messages, stop=None, run_manager=None, **kwargs):
133
+ async def wrapped_agenerate(
134
+ self, messages, stop=None, run_manager=None, **kwargs
135
+ ):
107
136
  # Get the response
108
- response = await original_agenerate(self, messages, stop=stop, run_manager=run_manager, **kwargs)
109
-
137
+ response = await original_agenerate(
138
+ self, messages, stop=stop, run_manager=run_manager, **kwargs
139
+ )
140
+
110
141
  # Send to OSMOSIS if enabled
111
142
  if utils.enabled:
112
143
  # Create payload
@@ -114,19 +145,27 @@ def _patch_anthropic_chat_models() -> None:
114
145
  payload = {
115
146
  "model_type": "ChatAnthropic",
116
147
  "model_name": model_name,
117
- "messages": [str(msg) for msg in messages], # Convert to strings for serialization
118
- "response": str(response), # Convert to string since it may not be serializable
119
- "kwargs": {"stop": stop, **kwargs}
148
+ "messages": [
149
+ str(msg) for msg in messages
150
+ ], # Convert to strings for serialization
151
+ "response": str(
152
+ response
153
+ ), # Convert to string since it may not be serializable
154
+ "kwargs": {"stop": stop, **kwargs},
120
155
  }
121
-
156
+
122
157
  send_to_osmosis(
123
- query={"type": "langchain_anthropic_agenerate", "messages": [str(msg) for msg in messages], "model": model_name},
158
+ query={
159
+ "type": "langchain_anthropic_agenerate",
160
+ "messages": [str(msg) for msg in messages],
161
+ "model": model_name,
162
+ },
124
163
  response=payload,
125
- status=200
164
+ status=200,
126
165
  )
127
-
166
+
128
167
  return response
129
-
168
+
130
169
  wrapped_agenerate._osmosis_aiped = True
131
170
  ChatAnthropic._agenerate = wrapped_agenerate
132
171
  logger.info("Successfully wrapped ChatAnthropic._agenerate method")
@@ -134,18 +173,21 @@ def _patch_anthropic_chat_models() -> None:
134
173
  logger.info("ChatAnthropic._agenerate already wrapped.")
135
174
  else:
136
175
  logger.info("ChatAnthropic does not have a _agenerate method, skipping.")
137
-
176
+
138
177
  # Patch _call method if it exists (used in newer versions)
139
178
  if hasattr(ChatAnthropic, "_call"):
140
179
  original_call = ChatAnthropic._call
141
-
180
+
142
181
  if not hasattr(original_call, "_osmosis_aiped"):
182
+
143
183
  @functools.wraps(original_call)
144
184
  def wrapped_call(self, messages, stop=None, run_manager=None, **kwargs):
145
185
  try:
146
186
  # Get the response
147
- response = original_call(self, messages, stop=stop, run_manager=run_manager, **kwargs)
148
-
187
+ response = original_call(
188
+ self, messages, stop=stop, run_manager=run_manager, **kwargs
189
+ )
190
+
149
191
  # Send to OSMOSIS if enabled
150
192
  if utils.enabled:
151
193
  # Create payload
@@ -153,24 +195,32 @@ def _patch_anthropic_chat_models() -> None:
153
195
  payload = {
154
196
  "model_type": "ChatAnthropic",
155
197
  "model_name": model_name,
156
- "messages": [str(msg) for msg in messages], # Convert to strings for serialization
198
+ "messages": [
199
+ str(msg) for msg in messages
200
+ ], # Convert to strings for serialization
157
201
  "response": str(response),
158
- "kwargs": {"stop": stop, **kwargs}
202
+ "kwargs": {"stop": stop, **kwargs},
159
203
  }
160
-
204
+
161
205
  send_to_osmosis(
162
- query={"type": "langchain_anthropic_call", "messages": [str(msg) for msg in messages], "model": model_name},
206
+ query={
207
+ "type": "langchain_anthropic_call",
208
+ "messages": [str(msg) for msg in messages],
209
+ "model": model_name,
210
+ },
163
211
  response=payload,
164
- status=200
212
+ status=200,
165
213
  )
166
-
214
+
167
215
  return response
168
216
  except TypeError as e:
169
217
  # Handle parameter mismatch gracefully
170
- logger.warning(f"TypeError in wrapped _call: {e}, trying without run_manager")
218
+ logger.warning(
219
+ f"TypeError in wrapped _call: {e}, trying without run_manager"
220
+ )
171
221
  # Try calling without run_manager (older versions)
172
222
  response = original_call(self, messages, stop=stop, **kwargs)
173
-
223
+
174
224
  # Send to OSMOSIS if enabled
175
225
  if utils.enabled:
176
226
  model_name = getattr(self, model_attr, "unknown_model")
@@ -179,17 +229,21 @@ def _patch_anthropic_chat_models() -> None:
179
229
  "model_name": model_name,
180
230
  "messages": [str(msg) for msg in messages],
181
231
  "response": str(response),
182
- "kwargs": {"stop": stop, **kwargs}
232
+ "kwargs": {"stop": stop, **kwargs},
183
233
  }
184
-
234
+
185
235
  send_to_osmosis(
186
- query={"type": "langchain_anthropic_call_fallback", "messages": [str(msg) for msg in messages], "model": model_name},
236
+ query={
237
+ "type": "langchain_anthropic_call_fallback",
238
+ "messages": [str(msg) for msg in messages],
239
+ "model": model_name,
240
+ },
187
241
  response=payload,
188
- status=200
242
+ status=200,
189
243
  )
190
-
244
+
191
245
  return response
192
-
246
+
193
247
  wrapped_call._osmosis_aiped = True
194
248
  ChatAnthropic._call = wrapped_call
195
249
  logger.info("Successfully wrapped ChatAnthropic._call method")
@@ -197,18 +251,23 @@ def _patch_anthropic_chat_models() -> None:
197
251
  logger.info("ChatAnthropic._call already wrapped.")
198
252
  else:
199
253
  logger.info("ChatAnthropic does not have a _call method, skipping.")
200
-
254
+
201
255
  # Patch _acall method if it exists
202
256
  if hasattr(ChatAnthropic, "_acall"):
203
257
  original_acall = ChatAnthropic._acall
204
-
258
+
205
259
  if not hasattr(original_acall, "_osmosis_aiped"):
260
+
206
261
  @functools.wraps(original_acall)
207
- async def wrapped_acall(self, messages, stop=None, run_manager=None, **kwargs):
262
+ async def wrapped_acall(
263
+ self, messages, stop=None, run_manager=None, **kwargs
264
+ ):
208
265
  try:
209
266
  # Get the response
210
- response = await original_acall(self, messages, stop=stop, run_manager=run_manager, **kwargs)
211
-
267
+ response = await original_acall(
268
+ self, messages, stop=stop, run_manager=run_manager, **kwargs
269
+ )
270
+
212
271
  # Send to OSMOSIS if enabled
213
272
  if utils.enabled:
214
273
  # Create payload
@@ -216,43 +275,57 @@ def _patch_anthropic_chat_models() -> None:
216
275
  payload = {
217
276
  "model_type": "ChatAnthropic",
218
277
  "model_name": model_name,
219
- "messages": [str(msg) for msg in messages], # Convert to strings for serialization
278
+ "messages": [
279
+ str(msg) for msg in messages
280
+ ], # Convert to strings for serialization
220
281
  "response": str(response),
221
- "kwargs": {"stop": stop, **kwargs}
282
+ "kwargs": {"stop": stop, **kwargs},
222
283
  }
223
-
284
+
224
285
  send_to_osmosis(
225
- query={"type": "langchain_anthropic_acall", "messages": [str(msg) for msg in messages], "model": model_name},
286
+ query={
287
+ "type": "langchain_anthropic_acall",
288
+ "messages": [str(msg) for msg in messages],
289
+ "model": model_name,
290
+ },
226
291
  response=payload,
227
- status=200
292
+ status=200,
228
293
  )
229
-
294
+
230
295
  return response
231
296
  except TypeError as e:
232
297
  # Handle parameter mismatch gracefully
233
- logger.warning(f"TypeError in wrapped _acall: {e}, trying without run_manager")
298
+ logger.warning(
299
+ f"TypeError in wrapped _acall: {e}, trying without run_manager"
300
+ )
234
301
  # Try calling without run_manager (older versions)
235
- response = await original_acall(self, messages, stop=stop, **kwargs)
236
-
302
+ response = await original_acall(
303
+ self, messages, stop=stop, **kwargs
304
+ )
305
+
237
306
  # Send to OSMOSIS if enabled
238
307
  if utils.enabled:
239
308
  model_name = getattr(self, model_attr, "unknown_model")
240
309
  payload = {
241
310
  "model_type": "ChatAnthropic",
242
- "model_name": model_name,
311
+ "model_name": model_name,
243
312
  "messages": [str(msg) for msg in messages],
244
313
  "response": str(response),
245
- "kwargs": {"stop": stop, **kwargs}
314
+ "kwargs": {"stop": stop, **kwargs},
246
315
  }
247
-
316
+
248
317
  send_to_osmosis(
249
- query={"type": "langchain_anthropic_acall_fallback", "messages": [str(msg) for msg in messages], "model": model_name},
318
+ query={
319
+ "type": "langchain_anthropic_acall_fallback",
320
+ "messages": [str(msg) for msg in messages],
321
+ "model": model_name,
322
+ },
250
323
  response=payload,
251
- status=200
324
+ status=200,
252
325
  )
253
-
326
+
254
327
  return response
255
-
328
+
256
329
  wrapped_acall._osmosis_aiped = True
257
330
  ChatAnthropic._acall = wrapped_acall
258
331
  logger.info("Successfully wrapped ChatAnthropic._acall method")
@@ -260,6 +333,6 @@ def _patch_anthropic_chat_models() -> None:
260
333
  logger.info("ChatAnthropic._acall already wrapped.")
261
334
  else:
262
335
  logger.info("ChatAnthropic does not have a _acall method, skipping.")
263
-
336
+
264
337
  except Exception as e:
265
- logger.error(f"Failed to patch langchain-anthropic chat model classes: {e}")
338
+ logger.error(f"Failed to patch langchain-anthropic chat model classes: {e}")