feat: Added new API method withLlmCustomAttributes to run a function in a LLM context

[MikeVaz]

Hi team. Here is our AIM's humble effort to implement adding LLM custom attributes

Description

During GA for the AIM product, custom attributes on the transaction level are prefixed with llm. were set up such that they would be added to all LLM events generated in that transaction from that point onwards, however, there has been demand for a more granular version of this.

This would be intended to support use cases such as A/B testing individual parts of flows (for example, the first call out to the LLM in a transaction might go to flow A, and the resulting second call in the transaction might go down flow B based on output from a tool).

Solution highlight

• We went with an option 4 based on the DACI implementing withLlmCustomAttributes() method
• In addition setLlmCustomAttributesCallback can be provided but removed from this PR for now.
• We originally tried to use an existing context manager on the agent instance to run a LLM context. this.agent._contextManager.runInContext.
• But the issue below made us move to using an isolated context manager on the transaction. transaction._llmContextManager = new AsyncLocalStorage();.
• Provided attributes will be normalized and later extracted from the context to be added to the custom event
• A new method withLlmCustomAttributes added to the APIs plus recordEvent method updated for openai, bedrock and langchain instrumentations.

How to Test

• We have tested it in our sample application. We have a dedicated branch.
• Changes are covered by unit and integration tests as well.

Related Issues

• One issue is revealed when Langchain versioned test is failing. We would like to consult with the team on that. I suspect something in the langchain flow overrides the context viathis.agent._contextManager.setContext({....}). A solutions could be to go with a dedicated isolated context manager transaction._llmContextManager = new AsyncLocalStorage();.

Problem

Consider we have a flow like below

└  Start Transaction
     └ withLlmCustomAttributes {llm.attribute: 1}
         │- Open AI call #1
         │   └ recordEvent with llm context {llm.attribute:1}
         └ switch context via agent._contextManager.runInContext({custom: 1})
                └ Open AI call #2
                     └ recordEvent when llm context lost     

Solution 1 (currently implemented)

We created an isolated context manager so it would never collide with agent's one.

Solution 2

Alternatively we can refactor agent's context manager to store parent context in a WeakMap

    setParentContext(child, parent) {
        this._parentContexMap = this._parentContextMap || new WeaMap()
        this._parentContextMap.set(child, parent)
    }
    runInContext(context, callback, cbThis, args) {
         const toInvoke = cbThis ? callback.bind(cbThis) : callback
         
         // Something like
         this.setParentContext(context, this.getContext())
         
         if (args) {
           return this._asyncLocalStorage.run(context, toInvoke, ...args)
         }
     
         return this._asyncLocalStorage.run(context, toInvoke)
       }

Now we can traverse related contexts to aggregate all llm context attributes attached along the way.

[CLAassistant]

All committers have signed the CLA.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
0 out of 3 committers have signed the CLA.

❌ mvazhenin
❌ MikeVaz
❌ RyanKadri

---

mvazhenin seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[bizob2828]

I don't think it has to be synchronously. The context manager should properly bind all functions run within the context

[MikeVaz]

I went with how they describe it in the nodejs docs.

Runs a function synchronously within a context and returns its return value. The store is not accessible outside of the callback function. The store is accessible to any asynchronous operations created within the callback.

[bizob2828]

The store is accessible to any asynchronous operations created within the callback.

[MikeVaz]

@bizob2828 @jsumners-nr I have updated my PR to address most of the feedback. I also got an idea why langchain was failing. Basically it would record events but only when virtual_llm = true and is_response = true events has the same context. The rest would somehow lose the context. Do you have ideas?

[jsumners-nr]

With the latest changes, I am not seeing any test failures. What does a test look like that generates the issue you are seeing?

[MikeVaz]

With the latest changes, I am not seeing any test failures. What does a test look like that generates the issue you are seeing?

It is not failing since I am filtering events

      const responses = events.filter((event) => {
          const [, chainEvent] = event
          return chainEvent.virtual_llm === true && chainEvent.is_response === true
        })

Removing filter would make test to fail

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.15%. Comparing base (0c2ee2f) to head (21e4b5f).
Report is 9 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2437      +/-   ##
==========================================
+ Coverage   96.20%   97.15%   +0.94%     
==========================================
  Files         288      289       +1     
  Lines       45305    45415     +110     
==========================================
+ Hits        43586    44123     +537     
+ Misses       1719     1292     -427     

Flag	Coverage Δ
integration-tests-cjs-18.x	73.80% <39.17%> (?)
integration-tests-cjs-20.x	73.82% <39.17%> (?)
integration-tests-cjs-22.x	73.85% <39.17%> (?)
integration-tests-esm-18.x	49.15% <39.17%> (?)
integration-tests-esm-20.x	49.15% <39.17%> (?)
integration-tests-esm-22.x	49.18% <39.17%> (?)
unit-tests-18.x	88.67% <82.69%> (+<0.01%)	⬆️
unit-tests-20.x	88.67% <82.69%> (+<0.01%)	⬆️
unit-tests-22.x	88.68% <82.69%> (+<0.01%)	⬆️
versioned-tests-18.x	78.84% <88.46%> (-0.06%)	⬇️
versioned-tests-20.x	78.88% <88.46%> (-0.03%)	⬇️
versioned-tests-22.x	78.88% <88.46%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[bizob2828]

feedback from the last round