Skip to content

Commit 2e7c9a1

Browse files
JakeTrockJakeTrock
committed
hide logs
1 parent c24da01 commit 2e7c9a1

File tree

7 files changed

+536
-955
lines changed

7 files changed

+536
-955
lines changed

‎osmosis_wrap/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
from . import utils
3838
from . import consts
3939
from .adapters import anthropic, openai
40+
from .logger import logger, reconfigure_logger
4041

4142
# Re-export configuration flags for easy access
4243
enabled = utils.enabled
@@ -68,8 +69,8 @@
6869
try:
6970
wrap_langchain()
7071
except Exception as e:
71-
print(f"Warning: Failed to wrap LangChain: {str(e)}")
72+
logger.warning(f"Failed to wrap LangChain: {str(e)}")
7273
except ImportError:
7374
# LangChain not installed or not compatible
7475
def wrap_langchain():
75-
print("LangChain support not available. Install LangChain to use this feature.")
76+
logger.warning("LangChain support not available. Install LangChain to use this feature.")

‎osmosis_wrap/adapters/anthropic.py‎

Lines changed: 39 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
from osmosis_wrap.utils import send_to_hoover
1111
from osmosis_wrap import utils
12+
from osmosis_wrap.logger import logger
1213

1314
def wrap_anthropic() -> None:
1415
"""
@@ -26,44 +27,44 @@ def wrap_anthropic() -> None:
2627
try:
2728
import anthropic
2829
except ImportError:
29-
print("Error: anthropic package is not installed.", file=sys.stderr)
30+
logger.error("anthropic package is not installed.")
3031
return
3132

32-
print(f"Wrapping Anthropic client, package version: {anthropic.__version__}", file=sys.stderr)
33+
logger.info(f"Wrapping Anthropic client, package version: {anthropic.__version__}")
3334

3435
# Get the resources.messages module and class
3536
messages_module = anthropic.resources.messages
3637
messages_class = messages_module.Messages
3738

38-
print(f"Found Anthropic messages class: {messages_class}", file=sys.stderr)
39+
logger.info(f"Found Anthropic messages class: {messages_class}")
3940

4041
# Patch the Messages.create method
4142
original_messages_create = messages_class.create
42-
print(f"Original create method: {original_messages_create}", file=sys.stderr)
43+
logger.info(f"Original create method: {original_messages_create}")
4344

4445
if not hasattr(original_messages_create, "_osmosis_wrapped"):
4546
@functools.wraps(original_messages_create)
4647
def wrapped_messages_create(self, *args, **kwargs):
47-
print(f"Wrapped create called with args: {args}, kwargs: {kwargs}", file=sys.stderr)
48+
logger.debug(f"Wrapped create called with args: {args}, kwargs: {kwargs}")
4849
try:
4950
# Check if the call includes tool use parameters
5051
has_tools = "tools" in kwargs
5152
if has_tools:
52-
print(f"Tool use detected with {len(kwargs['tools'])} tools", file=sys.stderr)
53+
logger.debug(f"Tool use detected with {len(kwargs['tools'])} tools")
5354

5455
# Check if this is a tool response message
5556
if "messages" in kwargs:
5657
for message in kwargs["messages"]:
5758
if message.get("role") == "user" and isinstance(message.get("content"), list):
5859
for content_item in message["content"]:
5960
if isinstance(content_item, dict) and content_item.get("type") == "tool_response":
60-
print("Tool response detected in messages", file=sys.stderr)
61+
logger.debug("Tool response detected in messages")
6162
break
6263

6364
response = original_messages_create(self, *args, **kwargs)
6465

6566
if utils.enabled:
66-
print("Sending success to Hoover (success)", file=sys.stderr)
67+
logger.debug("Sending success to Hoover (success)")
6768
send_to_hoover(
6869
query=kwargs,
6970
response=response.model_dump() if hasattr(response, 'model_dump') else response,
@@ -72,26 +73,26 @@ def wrapped_messages_create(self, *args, **kwargs):
7273

7374
return response
7475
except Exception as e:
75-
print(f"Error in wrapped create: {e}", file=sys.stderr)
76+
logger.error(f"Error in wrapped create: {e}")
7677
if utils.enabled:
7778
error_response = {"error": str(e)}
7879
send_to_hoover(
7980
query=kwargs,
8081
response=error_response,
8182
status=400
8283
)
83-
print("Sending error to Hoover (success)", file=sys.stderr)
84+
logger.debug("Sending error to Hoover (success)")
8485
raise # Re-raise the exception
8586

8687
wrapped_messages_create._osmosis_wrapped = True
8788
messages_class.create = wrapped_messages_create
88-
print("Successfully wrapped Messages.create method", file=sys.stderr)
89+
logger.info("Successfully wrapped Messages.create method")
8990

9091
# Directly wrap the AsyncAnthropic client
9192
try:
9293
# Get the AsyncAnthropic class
9394
AsyncAnthropicClass = anthropic.AsyncAnthropic
94-
print(f"Found AsyncAnthropic class: {AsyncAnthropicClass}", file=sys.stderr)
95+
logger.info(f"Found AsyncAnthropic class: {AsyncAnthropicClass}")
9596

9697
# Store the original __init__ to keep track of created instances
9798
original_async_init = AsyncAnthropicClass.__init__
@@ -102,7 +103,7 @@ def wrapped_async_init(self, *args, **kwargs):
102103
# Call the original init
103104
result = original_async_init(self, *args, **kwargs)
104105

105-
print("Wrapping new AsyncAnthropic instance's messages.create method", file=sys.stderr)
106+
logger.info("Wrapping new AsyncAnthropic instance's messages.create method")
106107

107108
# Get the messages client from this instance
108109
async_messages = self.messages
@@ -113,12 +114,12 @@ def wrapped_async_init(self, *args, **kwargs):
113114

114115
@functools.wraps(original_async_messages_create)
115116
async def wrapped_async_messages_create(*args, **kwargs):
116-
print(f"AsyncAnthropic.messages.create called with args: {args}, kwargs: {kwargs}", file=sys.stderr)
117+
logger.debug(f"AsyncAnthropic.messages.create called with args: {args}, kwargs: {kwargs}")
117118
try:
118119
response = await original_async_messages_create(*args, **kwargs)
119120

120121
if utils.enabled:
121-
print("Sending AsyncAnthropic response to Hoover (success)", file=sys.stderr)
122+
logger.debug("Sending AsyncAnthropic response to Hoover (success)")
122123
send_to_hoover(
123124
query=kwargs,
124125
response=response.model_dump() if hasattr(response, 'model_dump') else response,
@@ -127,9 +128,9 @@ async def wrapped_async_messages_create(*args, **kwargs):
127128

128129
return response
129130
except Exception as e:
130-
print(f"Error in wrapped AsyncAnthropic.messages.create: {e}", file=sys.stderr)
131+
logger.error(f"Error in wrapped AsyncAnthropic.messages.create: {e}")
131132
if utils.enabled:
132-
print("Sending AsyncAnthropic error to Hoover", file=sys.stderr)
133+
logger.debug("Sending AsyncAnthropic error to Hoover")
133134
error_response = {"error": str(e)}
134135
send_to_hoover(
135136
query=kwargs,
@@ -140,41 +141,41 @@ async def wrapped_async_messages_create(*args, **kwargs):
140141

141142
wrapped_async_messages_create._osmosis_wrapped = True
142143
async_messages.create = wrapped_async_messages_create
143-
print("Successfully wrapped AsyncAnthropic.messages.create method", file=sys.stderr)
144+
logger.info("Successfully wrapped AsyncAnthropic.messages.create method")
144145

145146
return result
146147

147148
wrapped_async_init._osmosis_wrapped = True
148149
AsyncAnthropicClass.__init__ = wrapped_async_init
149-
print("Successfully wrapped AsyncAnthropic.__init__ to patch message methods on new instances", file=sys.stderr)
150+
logger.info("Successfully wrapped AsyncAnthropic.__init__ to patch message methods on new instances")
150151
except (ImportError, AttributeError) as e:
151-
print(f"AsyncAnthropic class not found or has unexpected structure: {e}", file=sys.stderr)
152+
logger.warning(f"AsyncAnthropic class not found or has unexpected structure: {e}")
152153

153154
# For compatibility, still try to patch the old-style acreate method if it exists
154155
if hasattr(messages_class, "acreate"):
155156
original_acreate = messages_class.acreate
156157
if not hasattr(original_acreate, "_osmosis_wrapped"):
157158
@functools.wraps(original_acreate)
158159
async def wrapped_acreate(self, *args, **kwargs):
159-
print(f"Wrapped async create called with args: {args}, kwargs: {kwargs}", file=sys.stderr)
160+
logger.debug(f"Wrapped async create called with args: {args}, kwargs: {kwargs}")
160161
try:
161162
# Check if the async call includes tool use parameters
162163
has_tools = "tools" in kwargs
163164
if has_tools:
164-
print(f"Async tool use detected with {len(kwargs['tools'])} tools", file=sys.stderr)
165+
logger.debug(f"Async tool use detected with {len(kwargs['tools'])} tools")
165166

166167
if "messages" in kwargs:
167168
for message in kwargs["messages"]:
168169
if message.get("role") == "user" and isinstance(message.get("content"), list):
169170
for content_item in message["content"]:
170171
if isinstance(content_item, dict) and content_item.get("type") == "tool_response":
171-
print("Async tool response detected in messages", file=sys.stderr)
172+
logger.debug("Async tool response detected in messages")
172173
break
173174

174175
response = await original_acreate(self, *args, **kwargs)
175176

176177
if utils.enabled:
177-
print("Sending async response to Hoover (success)", file=sys.stderr)
178+
logger.debug("Sending async response to Hoover (success)")
178179
send_to_hoover(
179180
query=kwargs,
180181
response=response.model_dump() if hasattr(response, 'model_dump') else response,
@@ -183,9 +184,9 @@ async def wrapped_acreate(self, *args, **kwargs):
183184

184185
return response
185186
except Exception as e:
186-
print(f"Error in wrapped async create: {e}", file=sys.stderr)
187+
logger.error(f"Error in wrapped async create: {e}")
187188
if utils.enabled:
188-
print("Sending async error to Hoover", file=sys.stderr)
189+
logger.debug("Sending async error to Hoover")
189190
error_response = {"error": str(e)}
190191
send_to_hoover(
191192
query=kwargs,
@@ -196,9 +197,9 @@ async def wrapped_acreate(self, *args, **kwargs):
196197

197198
wrapped_acreate._osmosis_wrapped = True
198199
messages_class.acreate = wrapped_acreate
199-
print("Successfully wrapped Messages.acreate method", file=sys.stderr)
200+
logger.info("Successfully wrapped Messages.acreate method")
200201
else:
201-
print("No async acreate method found in Messages class", file=sys.stderr)
202+
logger.info("No async acreate method found in Messages class")
202203

203204
# Check if Completions exists and wrap it if it does
204205
try:
@@ -229,12 +230,12 @@ def wrapped_completions_create(self, *args, **kwargs):
229230
if not hasattr(original_completions_acreate, "_osmosis_wrapped"):
230231
@functools.wraps(original_completions_acreate)
231232
async def wrapped_completions_acreate(self, *args, **kwargs):
232-
print(f"Wrapped Completions async create called with args: {args}, kwargs: {kwargs}", file=sys.stderr)
233+
logger.debug(f"Wrapped Completions async create called with args: {args}, kwargs: {kwargs}")
233234
try:
234235
response = await original_completions_acreate(self, *args, **kwargs)
235236

236237
if utils.enabled:
237-
print("Sending Completions async response to Hoover (success)", file=sys.stderr)
238+
logger.debug("Sending Completions async response to Hoover (success)")
238239
send_to_hoover(
239240
query=kwargs,
240241
response=response.model_dump() if hasattr(response, 'model_dump') else response,
@@ -243,9 +244,9 @@ async def wrapped_completions_acreate(self, *args, **kwargs):
243244

244245
return response
245246
except Exception as e:
246-
print(f"Error in wrapped Completions async create: {e}", file=sys.stderr)
247+
logger.error(f"Error in wrapped Completions async create: {e}")
247248
if utils.enabled:
248-
print("Sending Completions async error to Hoover", file=sys.stderr)
249+
logger.debug("Sending Completions async error to Hoover")
249250
error_response = {"error": str(e)}
250251
send_to_hoover(
251252
query=kwargs,
@@ -256,15 +257,15 @@ async def wrapped_completions_acreate(self, *args, **kwargs):
256257

257258
wrapped_completions_acreate._osmosis_wrapped = True
258259
completions_class.acreate = wrapped_completions_acreate
259-
print("Successfully wrapped Completions.acreate method", file=sys.stderr)
260+
logger.info("Successfully wrapped Completions.acreate method")
260261
else:
261-
print("Completions.acreate already wrapped", file=sys.stderr)
262+
logger.info("Completions.acreate already wrapped")
262263
else:
263-
print("No async acreate method found in Completions class", file=sys.stderr)
264+
logger.info("No async acreate method found in Completions class")
264265
else:
265-
print("Completions.create already wrapped", file=sys.stderr)
266+
logger.info("Completions.create already wrapped")
266267
except (ImportError, AttributeError) as e:
267268
# Completions module may not exist in this version
268-
print(f"Completions module not found or has an unexpected structure: {e}", file=sys.stderr)
269+
logger.warning(f"Completions module not found or has an unexpected structure: {e}")
269270

270-
print("Anthropic client has been wrapped by osmosis-wrap.", file=sys.stderr)
271+
logger.info("Anthropic client has been wrapped by osmosis-wrap.")

0 commit comments

Comments
 (0)