Access observation id in langchain runnable

[snikch]

I've been really struggling today to get something which I think should be obvious. Inside a langgraph node (inside any langchain runnable I believe), how do I get the current observation id?

Using langfuse_context.get_current_observation_id() returns the root trace id.

I've worked around this by passing in a "run_id" to the langchain invoke call as both the run id and a configurable value, then inside the node I can get the langfuse handler from the langchain handlers and access the run by run id. This is quite hacky and it would be nice to access the observation id through the normal context method.

def get_langchain_observation_id(config: RunnableConfig) -> Optional[str]:
    """
    Get the langchain observation id via the config.
    For this to work, the RunnableConfig needs to have been configured with a run id via with_run_id.
    This allows Langchain to use the run id via the direct attachment, then allows us to get the run
    id via the configurable field. We can't use the run id on the config root at this point as the
    run id may not available at this point.

    Example:

    def extract_observation_id(messages: Any, config: RunnableConfig):
        # Add the observation to some parent scoped object
        state["observation_id"] = get_langchain_observation_id(config)
        return messages
    res = (RunnableLambda(extract_observation_id) | agent).invoke(state, with_run_id(config)))
    """
    callbacks = config.get("callbacks", None)
    if callbacks is None:
        return None
    if not isinstance(callbacks, CallbackManager):
        return None
    h = callbacks.handlers[0]
    if not isinstance(h, CallbackHandler):
        return None
    if h.runs is None or len(h.runs) == 0:
        return None
    run_id = config.get("configurable", {}).get("run_id", None)
    if run_id is None:
        return None
    if run_id not in h.runs:
        return None
    return h.runs[run_id].id

def with_run_id(config: RunnableConfig):
    """Wrap the configurable with a run id to enable get_langchain_observation_id to work"""
    config["run_id"] = uuid.uuid4()
    config.get("configurable", {})["run_id"] = config["run_id"]
    return config

[GMartin-dev]

I'm in a similar situation... I have my own override class for the langfuse callback with some modifications to get the observation id per generation. But is not pretty... basically I avoid clearing the "runs" dict and I collect the latest generation id.
Lucky me I already work with all runnables wrapped in a "preprocessing" Runnable Lambda, working on async you have to do some archaic validation to avoid race conditions.
THis is a complex problem to solve specially working async.

[snikch]

@GMartin-dev I've updated my question with the code I'm using. Thoughts on it compared to your approach?

[GMartin-dev]

@GMartin-dev I've updated my question with the code I'm using. Thoughts on it compared to your approach?

Yeap we are doing something similar. My approach is slightly diff since I added the observation getter in the Langfuse callback class and working async I have to wait until the actual LLM ends the generation.
So i'm keeping a generations tupple with id and the raw generation, and at the end of generation I ask for the macthing id.
I have to do this because same handlers manages a lot of concurrent and async generations that go into the a single trace.

This is a very simplified idea of the code, I do a few more customizations around fallback tokens calculations and other stuff.

from langfuse.callback import CallbackHandler

class CustomCallbackHandler(CallbackHandler):

    generations: list = []

    def get_generation_id(self, prediction: Any):
        observation_id = None
        for gen in self.generations:
            if gen[1] == prediction:
                observation_id = gen[0]
        return observation_id

  def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
    ...
    self.generations.append((self.runs[run_id].id, extracted_response))
    

[snikch]

Ah nice, so you instantiate your custom handler at the start of the chain rather than using langfuse_context.get_current_langchain_handler? I think that's a badly named method because it actually instantiates one if it doesn't find one.

This may also solve my other issue, which is that I actually want a generation up a number of levels, so I could store a map of run id to parent run ids.

[david1542]

Is there any news on that? I'm also interested in retrieving the observation id rather than the root trace id.

[david1542]

I managed to workaround this issue by also overriding the langfuse callback handler. Wrote a comment about it here.