Skip to content

Commit e4dadc3

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 e4dadc3

File tree

1 file changed

+42
-40
lines changed

1 file changed

+42
-40
lines changed

sentry_sdk/integrations/langchain.py

Lines changed: 42 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,51 @@ 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+
elif not isinstance(callbacks_list, list):
438+
logger.debug("Unknown callback type: %s", callbacks_list)
439+
# Just proceed with original function call
440+
return f(
441+
callback_manager_cls,
442+
inheritable_callbacks,
443+
local_callbacks,
444+
*args,
445+
**kwargs,
446+
)
447+
448+
if not any(isinstance(cb, SentryLangchainCallback) for cb in callbacks_list):
449+
# Avoid mutating the existing callbacks list
450+
callbacks_list = [
451+
*callbacks_list,
452+
SentryLangchainCallback(
453+
integration.max_spans,
454+
integration.include_prompts,
455+
integration.tiktoken_encoding_name,
456+
),
457+
]
458+
459+
return f(
460+
callback_manager_cls,
461+
inheritable_callbacks,
462+
callbacks_list,
463+
*args,
464+
**kwargs,
465+
)
464466

465467
return new_configure

0 commit comments

Comments
 (0)