Skip to content

Commit 901b579

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 fedcb07 commit 901b579

File tree

1 file changed

+49
-41
lines changed

1 file changed

+49
-41
lines changed

sentry_sdk/integrations/langchain.py

Lines changed: 49 additions & 41 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,57 @@ 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:
424-
return f(*args, **kwargs)
431+
return f(
432+
callback_manager_cls,
433+
inheritable_callbacks,
434+
local_callbacks,
435+
*args,
436+
**kwargs,
437+
)
425438

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

465473
return new_configure

0 commit comments

Comments
 (0)