Skip to content

Commit f2bd2a6

Browse files
ref(langchain): Greatly simplify _wrap_configure
Resolving #4443 requires some changes to this method, but the current `args`/`kwargs` business makes the method difficult to reason through. This PR simplifies the logic by listing out the parameters we need to access, so we don't need to access them through `args` and `kwargs`. We also cut down on the amount of branching and the amount of variables (`new_callbacks` vs `existing_callbacks`). Behavior does not change in this PR; we will introduce behavior changes in a subsequent PR
1 parent f71d223 commit f2bd2a6

File tree

1 file changed

+35
-40
lines changed

1 file changed

+35
-40
lines changed

sentry_sdk/integrations/langchain.py

Lines changed: 35 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from langchain_core.callbacks import (
2323
manager,
2424
BaseCallbackHandler,
25+
Callbacks,
2526
)
2627
from langchain_core.agents import AgentAction, AgentFinish
2728
except ImportError:
@@ -416,50 +417,44 @@ def _wrap_configure(f):
416417
# type: (Callable[..., Any]) -> Callable[..., Any]
417418

418419
@wraps(f)
419-
def new_configure(*args, **kwargs):
420-
# type: (Any, Any) -> Any
420+
def new_configure(
421+
callback_manager_cls, # type: type
422+
inheritable_callbacks=None, # type: Callbacks
423+
local_callbacks=None, # type: Callbacks
424+
*args, # type: Any
425+
**kwargs, # type: Any
426+
):
427+
# type: (...) -> Any
421428

422429
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
423430
if integration is None:
424431
return f(*args, **kwargs)
425432

426-
with capture_internal_exceptions():
427-
new_callbacks = [] # type: List[BaseCallbackHandler]
428-
if "local_callbacks" in kwargs:
429-
existing_callbacks = kwargs["local_callbacks"]
430-
kwargs["local_callbacks"] = new_callbacks
431-
elif len(args) > 2:
432-
existing_callbacks = args[2]
433-
args = (
434-
args[0],
435-
args[1],
436-
new_callbacks,
437-
) + args[3:]
438-
else:
439-
existing_callbacks = []
440-
441-
if existing_callbacks:
442-
if isinstance(existing_callbacks, list):
443-
for cb in existing_callbacks:
444-
new_callbacks.append(cb)
445-
elif isinstance(existing_callbacks, BaseCallbackHandler):
446-
new_callbacks.append(existing_callbacks)
447-
else:
448-
logger.debug("Unknown callback type: %s", existing_callbacks)
449-
450-
already_added = False
451-
for callback in new_callbacks:
452-
if isinstance(callback, SentryLangchainCallback):
453-
already_added = True
454-
455-
if not already_added:
456-
new_callbacks.append(
457-
SentryLangchainCallback(
458-
integration.max_spans,
459-
integration.include_prompts,
460-
integration.tiktoken_encoding_name,
461-
)
462-
)
463-
return f(*args, **kwargs)
433+
callbacks_list = local_callbacks or []
434+
435+
if isinstance(callbacks_list, BaseCallbackHandler):
436+
callbacks_list = [callbacks_list]
437+
else:
438+
logger.debug("Unknown callback type: %s", callbacks_list)
439+
callbacks_list = []
440+
441+
if not any(isinstance(cb, SentryLangchainCallback) for cb in callbacks_list):
442+
# Avoid mutating the existing callbacks list
443+
callbacks_list = [
444+
*callbacks_list,
445+
SentryLangchainCallback(
446+
integration.max_spans,
447+
integration.include_prompts,
448+
integration.tiktoken_encoding_name,
449+
),
450+
]
451+
452+
return f(
453+
callback_manager_cls,
454+
inheritable_callbacks,
455+
callbacks_list,
456+
*args,
457+
**kwargs,
458+
)
464459

465460
return new_configure

0 commit comments

Comments
 (0)