feat(langchain): use callbacks

[tibor-reiss]

Some work on #541 (which will then allow to fix #1101 in a more maintainable way)

• I have added tests that cover my changes.
• If adding a new instrumentation or changing an existing one, I've added screenshots from some observability platform showing the change.
• PR name follows conventional commits format: feat(instrumentation): ... or fix(instrumentation): ....
• (If applicable) I have updated the documentation accordingly.

ToDo:

• Add async callback -> according to the docs, sync callbacks work

[tibor-reiss]

There is quite some duplication, maybe there is a better way to do this?

[nirga]

@tibor-reiss I think a better way will actually be something like #902. I think the contributor stopped working on it there so potentially we can pick this up (I was planning to do so sometime next week)

[tibor-reiss]

Hi @nirga, I have checked out the linked PR, and based on that (+research into callbacks) I managed to make it (in sync case) work for LLMChain, SequentialChain, see attached screenshot.

However, RunnableSequence - which seems to be the way forward in langchain (e.g. even LLMChain will be deprecated) - is not that simple: https://python.langchain.com/v0.2/docs/how_to/callbacks_attach/. I could not yet figure out a way how to hook into the with_config method...

Note: for now, below seems to work, but it's explicit:
tracer = TracerWrapper().get_tracer() runnable = runnable.with_config(callbacks=[SyncSpanCallbackHandler(tracer, "bla", "blue"), ])

I don't know how fast the API changes on langchain side, so not sure about the whole migration to callbacks via a decorator. Since there is the simple way of adding callbacks to e.g. LLMChain and RunnableSequence, it might be an option to go that route, e.g. SyncTraceloopCallback, and then llm_chain = LLMChain(..., callbacks=[..., SyncTracelooopCallback]).

Let me know your thoughts, please!

Ps.: we can open a new PR for the callback part first, and then rework this later...

[nirga]

Thanks @tibor-reiss this looks like a great step in the right direction! We should do something similar to our llamaindex instrumentation as well (we have an issue for that). I think this is still more robust than the way we do it today. I've been in touch with the folks from Langchain and hopefully we'll found a way for the new Runnable syntax.
I'll dig into their codebase as well to try to come up with a solution but I think your proposal here works well! (at least for now).

Mind fixing the lint / tests?

[nirga]

I'm guessing you want to keep the assertions here?

[tibor-reiss]

Yes, ofc, cleaned up the PR now.

[nirga]

Potentially what we can do (maybe outside the scope for this PR) - wrap a call to Runnable.invoke() and then create a RunnableBinding that wraps the original Runnable and invoke it

[tibor-reiss]

Potentially what we can do (maybe outside the scope for this PR) - wrap a call to Runnable.invoke() and then create a RunnableBinding that wraps the original Runnable and invoke it

Sounds like a good idea. I agree that smaller PRs seems advantageous here because of how fast langchain changes, and because the way forward is not so clear.

I have now edited the scope/description of the PR. As you said, it's probably better to first implement callbacks as far as possible, and do later RunnableSequence.

I'll try to add the async callback asap - but feel free to add it yourself :)

For this PR, what might still make sense is to add BaseChatModel. It's not trivial due to 2 reasons:

1. custom_chat_wrapper.py has quite some additional span attributes (shouldn't be a big issue)
2. There is no on_chat_model_end (yet). So one might have to call self.span_end() already in on_chat_model_start.

[nirga]

@tibor-reiss so let's do it gradually. Begin with this PR, and think about BaseChatModel later. The issue with no on_chat_model_end is that we won't have the right timing for the span, so might be better maybe to keep it as it is for now

I see you marked this as WIP, let me know when this is ready and I'll do a full review :)

[tibor-reiss]

@nirga please check
Note: according to the docs, sync callbacks work also in the async case.

[nirga]

Looks good overall @tibor-reiss! Great work :)
(I'm guessing we'll need to do it for other types of chains as well, right?)

[nirga]

Can we also test other parts of the callback instrumentation? For example, see the inputs and outputs are set on the spans, or give names to the LLMChains and see that the span name is set accordingly?

[tibor-reiss]

Added more checks.
Re name: the span name is set in the callback handler, so it can't be changed (there might be occasions where it can be, e.g. with_config)

[midhun1998]

@tibor-reiss I think a better way will actually be something like #902. I think the contributor stopped working on it there so potentially we can pick this up (I was planning to do so sometime next week)

Hi @nirga ,

I was away due to some personal reasons. I have some changes parked in my local and wanted to review them, but it looks like there is some progress in this PR and I'm not sure if I should continue with my PR 😞 . Let me know if you want me to close that one.

[nirga]

Maybe we can combine your work here as well @midhun1998 ?

[midhun1998]

Maybe we can combine your work here as well @midhun1998 ?

Works for me. I can add in my changes here or get my PR finished up quickly.

[tibor-reiss]

Looks good overall @tibor-reiss! Great work :) (I'm guessing we'll need to do it for other types of chains as well, right?)

Yes, chains can be added. I added now StuffDocumentsChain (with a test). Maybe there is a better way to add all children of the base class since Chain is already added and here is the callbacks attribute defined...

@midhun1998 feel free to edit

[nirga]

@tibor-reiss @midhun1998 I'm going to merge this so we can continue this work against main branch - I think it will be easier. Great work @tibor-reiss :) Thanks so much!