feat(agents): AgentDependencies in AgentRunContext #780 (#781)

Author: maxpill

docs/how-to/agents/define_and_use_agents.md

@@ -84,6 +84,41 @@ In this scenario, the agent recognizes that the follow-up question "What about T
 AgentResult(content='The current temperature in Tokyo is 10°C.', ...)
 ```
 
+## Binding dependencies via AgentRunContext
+You can bind your external dependencies before the first access and safely use them in tools. After first attribute lookup, the dependencies container freezes to prevent mutation during a run.
+
+```python
+from dataclasses import dataclass
+from ragbits.agents import Agent, AgentRunContext
+from ragbits.core.llms.mock import MockLLM, MockLLMOptions
+
+@dataclass
+class Deps:
+    api_host: str
+
+def get_api_host(context: AgentRunContext | None) -> str:
+    """Return the API host taken from the bound dependencies in context."""
+    assert context is not None
+    return context.deps.api_host
+
+async def main() -> None:
+    llm = MockLLM(
+        default_options=MockLLMOptions(
+            response="Using dependencies from context.",
+            tool_calls=[{"name": "get_api_host", "arguments": "{}", "id": "example", "type": "function"}],
+        )
+    )
+    agent = Agent(llm=llm, prompt="Retrieve API host", tools=[get_api_host])
+
+    context = AgentRunContext()
+    context.deps.value = Deps(api_host="https://api.local")
+
+    result = await agent.run("What host are we using?", context=context)
+    print(result.tool_calls[0].result)
+```
+
+See the runnable example in `examples/agents/dependencies.py`.
+
 ## Streaming agent responses
 For use cases where you want to process partial outputs from the LLM as they arrive (e.g., in chat UIs), the [`Agent`][ragbits.agents.Agent] class supports streaming through the `run_streaming()` method.
 

examples/agents/dependencies.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+
+from ragbits.agents import Agent, AgentRunContext
+from ragbits.core.llms.mock import MockLLM, MockLLMOptions
+
+
+@dataclass
+class Deps:
+    """Simple dependency container used by the agent tools."""
+
+    api_host: str
+
+
+def get_api_host(context: AgentRunContext | None) -> str:
+    """Return the API host taken from the bound dependencies in context."""
+    if context is None:
+        raise ValueError("Missing AgentRunContext")
+    return context.deps.api_host
+
+
+async def main() -> None:
+    """Bind dependencies, run the agent, and access them inside a tool."""
+    llm = MockLLM(
+        default_options=MockLLMOptions(
+            response="Using dependencies from context.",
+            tool_calls=[
+                {
+                    "name": "get_api_host",
+                    "arguments": "{}",
+                    "id": "example",
+                    "type": "function",
+                }
+            ],
+        )
+    )
+
+    agent = Agent(llm=llm, prompt="Retrieve API host", tools=[get_api_host])
+
+    context = AgentRunContext()
+    context.deps.value = Deps(api_host="https://api.local")
+
+    result = await agent.run("What host are we using?", context=context)
+    print(result.tool_calls[0].result)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    asyncio.run(main())

packages/ragbits-agents/CHANGELOG.md

@@ -3,8 +3,9 @@
 ## Unreleased
 - Add tool_choice parameter to agent interface (#738)
 - Add PydanticAI agents support (#755)
-- Add unique ID to each agent instance for better tracking and identification
-- Add audit tracing for tool calls to improve observability and debugging
+- Add unique ID to each agent instance for better tracking and identification (#790)
+- Add audit tracing for tool calls to improve observability and debugging (#790)
+- Add AgentDependencies type (#781)
 
 ## 1.2.2 (2025-08-08)
 

packages/ragbits-agents/src/ragbits/agents/__init__.py

@@ -1,8 +1,17 @@
-from ragbits.agents._main import Agent, AgentOptions, AgentResult, AgentResultStreaming, AgentRunContext, ToolCallResult
+from ragbits.agents._main import (
+    Agent,
+    AgentDependencies,
+    AgentOptions,
+    AgentResult,
+    AgentResultStreaming,
+    AgentRunContext,
+    ToolCallResult,
+)
 from ragbits.agents.types import QuestionAnswerAgent, QuestionAnswerPromptInput, QuestionAnswerPromptOutput
 
 __all__ = [
     "Agent",
+    "AgentDependencies",
     "AgentOptions",
     "AgentResult",
     "AgentResultStreaming",

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -7,7 +7,7 @@
 from datetime import timedelta
 from inspect import iscoroutinefunction
 from types import ModuleType, SimpleNamespace
-from typing import ClassVar, Generic, cast, overload
+from typing import ClassVar, Generic, TypeVar, cast, overload
 
 from pydantic import BaseModel, Field
 from typing_extensions import Self
@@ -82,11 +82,70 @@ class AgentOptions(Options, Generic[LLMClientOptionsT]):
     or None, agent will run forever"""
 
 
-class AgentRunContext(BaseModel):
+DepsT = TypeVar("DepsT")
+
+
+class AgentDependencies(BaseModel, Generic[DepsT]):
     """
-    Context for the agent run.
+    Container for agent runtime dependencies.
+
+    Becomes immutable after first attribute access.
     """
 
+    model_config = {"arbitrary_types_allowed": True}
+
+    _frozen: bool
+    _value: DepsT | None
+
+    def __init__(self, value: DepsT | None = None) -> None:
+        super().__init__()
+        self._value = value
+        self._frozen = False
+
+    def __setattr__(self, name: str, value: object) -> None:
+        is_frozen = False
+        if name != "_frozen":
+            try:
+                is_frozen = object.__getattribute__(self, "_frozen")
+            except AttributeError:
+                is_frozen = False
+
+        if is_frozen and name not in {"_frozen"}:
+            raise RuntimeError("Dependencies are immutable after first access")
+
+        super().__setattr__(name, value)
+
+    @property
+    def value(self) -> DepsT | None:
+        return self._value
+
+    @value.setter
+    def value(self, value: DepsT) -> None:
+        if self._frozen:
+            raise RuntimeError("Dependencies are immutable after first access")
+        self._value = value
+
+    def _freeze(self) -> None:
+        if not self._frozen:
+            self._frozen = True
+
+    def __getattr__(self, name: str) -> object:
+        value = object.__getattribute__(self, "_value")
+        if value is None:
+            raise AttributeError(name)
+        self._freeze()
+        return getattr(value, name)
+
+    def __contains__(self, key: str) -> bool:
+        value = object.__getattribute__(self, "_value")
+        return hasattr(value, key) if value is not None else False
+
+
+class AgentRunContext(BaseModel, Generic[DepsT]):
+    """Context for the agent run."""
+
+    deps: AgentDependencies[DepsT] = Field(default_factory=lambda: AgentDependencies())
+    """Container for external dependencies."""
     usage: Usage = Field(default_factory=Usage)
     """The usage of the agent."""
 

packages/ragbits-agents/tests/unit/test_agent.py

@@ -325,7 +325,7 @@ async def test_agent_run_tools_with_context(
 
 @pytest.mark.parametrize("method", [_run, _run_streaming])
 async def test_agent_run_context_is_updated(llm_without_tool_call: MockLLM, method: Callable):
-    context = AgentRunContext()
+    context: AgentRunContext = AgentRunContext()
     agent: Agent = Agent(
         llm=llm_without_tool_call,
         prompt="NOT IMPORTANT",

packages/ragbits-core/src/ragbits/core/utils/function_schema.py

@@ -6,7 +6,7 @@
 import logging
 from collections.abc import Callable, Generator
 from types import UnionType
-from typing import Any, get_args, get_origin, get_type_hints
+from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints
 
 from griffe import Docstring, DocstringSectionKind
 from pydantic import BaseModel, Field, create_model
@@ -201,10 +201,20 @@ def _is_context_variable(ann: Any) -> bool:  # noqa: ANN401
     """
     Check if a type annotation is AgentRunContext.
     """
-    if hasattr(ann, "__name__") and ann.__name__ == "AgentRunContext":
+    if ann is None:
+        return False
+
+    origin = get_origin(ann)
+    if origin is not None:
+        if getattr(origin, "__name__", None) == "AgentRunContext":
+            return True
+        if origin in (Union, UnionType):
+            return any(_is_context_variable(arg) for arg in get_args(ann))
+        if origin is Annotated:
+            args = get_args(ann)
+            return bool(args) and _is_context_variable(args[0])
+
+    if getattr(ann, "__name__", None) == "AgentRunContext":
         return True
 
-    if get_origin(ann) is UnionType:
-        return any(_is_context_variable(arg) for arg in get_args(ann))
-
     return isinstance(ann, str) and "AgentRunContext" in ann