Improved Tool Calling Context

[EItanya]

Currently autogen studio supports Otel tracing for its agent runtime, however this support does not extend for any other network calls in the system. In order to perform distributed tracing it would be wonderful to be able to pass the trace context through to LLM calls as well as any tool calls.

In addition to passing tracing context, this system could be generic enough to allow for passing generic context variables through the call stack for authentication scenarios. For example, let's say you created an agent using agentchat which you wanted to perform function calls using the credentials of the user who initiated the task, it's currently impossible because tool calls only take in the function args, and a cancellation token.

I am by no means a Python expert, and I'm not sure how Context is usually handled, but I am suggesting modifying the tool interface to be something like:

class Context:
    @abstractmethod
    async def set(self, key: str, val: Any)
    @abstractmethod
    async def get(self, key: str) -> Any

    @abstractmethod
    async def run(self, args: ArgsT, cancellation_token: CancellationToken, ctx: Context | None = None) -> ReturnT: ...

    async def run_json(self, args: Mapping[str, Any], ctx: Context, cancellation_token: CancellationToken) -> Any:
        return_value = await self.run(self._args_type.model_validate(args), ctx, cancellation_token)
        return return_value

Again I want to make clear that the above example is probably not the most Pythonic code ever, so this is more meant to start the discussion rather than be a fixed solution.

I also want to point out that opentelemetry has a Context idea that we already use for tracing in the runtimes, so we could extend that concept for this use-case. I'm happy to do deeper research into how we can extend the existing Tracing work to be more generic, but I wanted to open the conversation first.

@victordibia and @ekzhu this is further work that I think could help productionize agentchat/studio. Happy to work more on it but wanted feedback first.

[ekzhu]

For OTEL, we can follow what we do for LLMCallEvent and emit a ToolCallEvent in BaseTool:

See:

autogen/python/packages/autogen-core/src/autogen_core/logging.py

Lines 9 to 10 in dd0781a

	class LLMCallEvent:
	def __init__(

And where it is emitted:

autogen/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py

Lines 559 to 560 in dd0781a

	logger.info(
	LLMCallEvent(

Created #5745

[ekzhu]

I see. At this point I am not sure. If it's something that can be done easily by wrapping the tools of adding it to the tools directly, we would like to see more examples of it first before adding to the framework.

[EItanya]

I outlined a number of simple examples above such as tracing and user authentication.

I don't think this is something that can be added without framework changes. Fundamentally in order to do this a developer would need the ability to have a variable or some other generic context passed all the way through the call chain. In practice this would most likely start with run_stream on either a team or an agent, and then be passed all the way through the call stack to either an LLM call, a tool call, or even potentially a DB call for Memory. If you think there's a way to accomplish this within the framework I'm all ears, but I think this sort of generic behavior may need to be added directly.

[ekzhu]

is it possible to use contextvar?

https://docs.python.org/3/library/contextvars.html#contextvars.Context

We use this as well for agent runtime:

autogen/python/packages/autogen-core/src/autogen_core/_agent_instantiation.py

Lines 9 to 10 in dd0781a

	class AgentInstantiationContext:
	"""A static class that provides context for agent instantiation.

[EItanya]

I think this will allow me to solve the problem in the short-term, but this will only work for SingleAgentRuntime, it will not work for any network based runtime. In that case it will be necessary to pass this metadata over the protocol.

I feel that I may not be explaining myself super well so I will try to come to next week's office hours

[ekzhu]

You explained well. I think I get the idea. Let's talk more during the office hour.