From bd7ee28d016154d3f7f9bd21e45628a79db7247f Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 10:26:26 -0500 Subject: [PATCH 1/6] create new tool for dynamically loading/deloading integration into agent context --- beaker_kernel/lib/agent.py | 6 ++++ beaker_kernel/lib/context.py | 10 ++++++ beaker_kernel/lib/integrations/adhoc.py | 42 +++++++++++++++++++++++-- 3 files changed, 56 insertions(+), 2 deletions(-) diff --git a/beaker_kernel/lib/agent.py b/beaker_kernel/lib/agent.py index ce953dad..e5b8c7eb 100644 --- a/beaker_kernel/lib/agent.py +++ b/beaker_kernel/lib/agent.py @@ -58,6 +58,12 @@ def __init__( for tool in self.tools.values(): set_tool_execution_context(tool) + async def post_loop(self, skip_summarization: bool = False): + # Clear active integrations so the next query starts with just summaries + if hasattr(self.context, 'active_integrations'): + self.context.active_integrations.clear() + await super().post_loop(skip_summarization=skip_summarization) + async def react_async(self, query: str, react_context: dict = None) -> str: return await super().react_async(query, react_context) diff --git a/beaker_kernel/lib/context.py b/beaker_kernel/lib/context.py index 1d736fb7..21925fdd 100644 --- a/beaker_kernel/lib/context.py +++ b/beaker_kernel/lib/context.py @@ -75,6 +75,7 @@ def __init__(self, beaker_kernel: "BeakerKernel", agent_cls: "BeakerAgent", conf integrations: list[BaseIntegrationProvider] = None): self.intercepts = [] self.integrations = integrations if integrations is not None else [] + self.active_integrations: set[str] = set() self.jinja_env = None self.templates = {} self.workflows = {} @@ -248,6 +249,15 @@ async def auto_context(self): integration.prompt for integration in self.integrations ] parts.append("---".join(integration_prompts)) + + # Include full documentation for integrations loaded via load_integration_docs + if self.active_integrations: + for provider in self.integrations: + if hasattr(provider, 'get_rendered_docs'): + for slug in list(self.active_integrations): + docs = provider.get_rendered_docs(slug) + if docs: + parts.append(f"## Loaded Integration Documentation: {slug}\n\n{docs}") content = "\n\n".join(parts) return content diff --git a/beaker_kernel/lib/integrations/adhoc.py b/beaker_kernel/lib/integrations/adhoc.py index d45ac055..71e6e0c8 100644 --- a/beaker_kernel/lib/integrations/adhoc.py +++ b/beaker_kernel/lib/integrations/adhoc.py @@ -369,6 +369,11 @@ def build_adhoc(self): substitutions = {} # handling None cases in failed renders keeps them editable but not usable by the agent rendered_apis = [spec.render(self, substitutions) for spec in self.specifications] + # Cache rendered documentation for direct context injection via load_integration_docs + self._rendered_docs = { + api.slug: api.documentation + for api in rendered_apis if api is not None + } self.adhoc_api = AdhocApi( apis=[api for api in rendered_apis if api is not None], **self.adhoc_config_options @@ -378,6 +383,10 @@ def refresh_adhoc_specs(self): # TODO: future way to not fully reinitialize to make it less slow. self.build_adhoc() + def get_rendered_docs(self, slug: str) -> Optional[str]: + """Returns the full rendered documentation for an integration by slug.""" + return self._rendered_docs.get(slug) + @property def prompt(self): agent_details = {spec.slug: spec.description for spec in self.specifications} @@ -386,9 +395,10 @@ def prompt(self): self.prompt_instructions if self.prompt_instructions else "", "" f"{self.display_name}:", - "You have access to the following integrations to use with the `draft_integration_code` and `consult_integration_docs` tools,", + "You have access to the following integrations to use with the `draft_integration_code`, `consult_integration_docs`, and `load_integration_docs` tools,", "as well as their descriptions for when and why you should use the given integration, delimited in three backticks.", - "Only use these integrations with the `draft_integration_code` and `consult_integration_docs` tools.", + "Use `load_integration_docs` to load an integration's full documentation directly into your context when you need to work extensively with it.", + "Use `draft_integration_code` and `consult_integration_docs` for quick, targeted requests.", "", delimiter ] @@ -569,6 +579,34 @@ async def add_example(self, integration: str, query: str, code: str, notes: str) return "Add resource tool failed." return f"Example has been added to {integration}." + @tool + async def load_integration_docs(self, integration: str, agent: AgentRef, loop: LoopControllerRef, react_context: ReactContextRef) -> str: + """ +Loads the full documentation for an integration directly into your context window. Use this when you need to +work extensively with an integration — for example, when writing or debugging code that interacts with the +integration's API. Once loaded, the documentation stays in your context for the remainder of the current query, +so you can reference it directly without additional tool calls. + +Prefer this tool when you anticipate multiple interactions with an integration. For quick one-off questions, +`consult_integration_docs` or `draft_integration_code` may be more efficient. + +Args: + integration (str): The name/slug of the integration to load documentation for. + +Returns: + str: The full documentation for the integration. + """ + docs = self.get_rendered_docs(integration) + if docs is None: + available = [spec.slug for spec in self.specifications] + return f"Integration '{integration}' not found. Available integrations: {', '.join(available)}" + + # Mark integration as active on the context so auto_context includes it on subsequent turns + if hasattr(agent, 'context') and hasattr(agent.context, 'active_integrations'): + agent.context.active_integrations.add(integration) + + return f"Documentation for '{integration}' has been loaded into your context:\n\n{docs}" + @tool async def draft_integration_code(self, integration: str, goal: str, agent: AgentRef, loop: LoopControllerRef, react_context: ReactContextRef) -> str: """ From d8f5cab3077332e51232f8a57a42fed4faae0558 Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 10:39:05 -0500 Subject: [PATCH 2/6] fix typing bug --- beaker_kernel/lib/integrations/adhoc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/beaker_kernel/lib/integrations/adhoc.py b/beaker_kernel/lib/integrations/adhoc.py index 71e6e0c8..3c31e398 100644 --- a/beaker_kernel/lib/integrations/adhoc.py +++ b/beaker_kernel/lib/integrations/adhoc.py @@ -371,7 +371,7 @@ def build_adhoc(self): rendered_apis = [spec.render(self, substitutions) for spec in self.specifications] # Cache rendered documentation for direct context injection via load_integration_docs self._rendered_docs = { - api.slug: api.documentation + api["slug"]: api["documentation"] for api in rendered_apis if api is not None } self.adhoc_api = AdhocApi( From 004306036faf9fbc9bb34438d38c7595d9d6d703 Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 11:57:15 -0500 Subject: [PATCH 3/6] debug auto_context for integrations --- beaker_kernel/lib/context.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/beaker_kernel/lib/context.py b/beaker_kernel/lib/context.py index 21925fdd..d1263d52 100644 --- a/beaker_kernel/lib/context.py +++ b/beaker_kernel/lib/context.py @@ -252,12 +252,16 @@ async def auto_context(self): # Include full documentation for integrations loaded via load_integration_docs if self.active_integrations: + logger.info(f"Active integrations for auto_context: {self.active_integrations}") for provider in self.integrations: if hasattr(provider, 'get_rendered_docs'): for slug in list(self.active_integrations): docs = provider.get_rendered_docs(slug) if docs: + logger.info(f"Injecting docs for '{slug}' into auto_context ({len(docs)} chars)") parts.append(f"## Loaded Integration Documentation: {slug}\n\n{docs}") + else: + logger.warning(f"No rendered docs found for active integration '{slug}'") content = "\n\n".join(parts) return content From 7412ef25ca5643815b972f0805b8029623186578 Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 12:18:03 -0500 Subject: [PATCH 4/6] remove post loop for testing --- beaker_kernel/lib/agent.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/beaker_kernel/lib/agent.py b/beaker_kernel/lib/agent.py index e5b8c7eb..ce953dad 100644 --- a/beaker_kernel/lib/agent.py +++ b/beaker_kernel/lib/agent.py @@ -58,12 +58,6 @@ def __init__( for tool in self.tools.values(): set_tool_execution_context(tool) - async def post_loop(self, skip_summarization: bool = False): - # Clear active integrations so the next query starts with just summaries - if hasattr(self.context, 'active_integrations'): - self.context.active_integrations.clear() - await super().post_loop(skip_summarization=skip_summarization) - async def react_async(self, query: str, react_context: dict = None) -> str: return await super().react_async(query, react_context) From 36bd29f3ad26d32a5be52f874d501cac8435538a Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 12:48:50 -0500 Subject: [PATCH 5/6] add in post loop to clear sys prompt of integration --- beaker_kernel/lib/agent.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/beaker_kernel/lib/agent.py b/beaker_kernel/lib/agent.py index ce953dad..94a3fddb 100644 --- a/beaker_kernel/lib/agent.py +++ b/beaker_kernel/lib/agent.py @@ -59,7 +59,12 @@ def __init__( set_tool_execution_context(tool) async def react_async(self, query: str, react_context: dict = None) -> str: - return await super().react_async(query, react_context) + try: + return await super().react_async(query, react_context) + finally: + # Clear loaded integration docs after each query to avoid context bloat + if hasattr(self.context, 'active_integrations'): + self.context.active_integrations.clear() async def execute(self, *args, **kwargs) -> str: return await super().execute(*args, **kwargs) From ddecfb00cde44907574b413fe97c101169ce6dbd Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 1 Apr 2026 13:15:48 -0500 Subject: [PATCH 6/6] try adding tool for unloading integration from context [managed by agent, not automated] --- beaker_kernel/lib/agent.py | 7 +------ beaker_kernel/lib/integrations/adhoc.py | 24 ++++++++++++++++++++++++ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/beaker_kernel/lib/agent.py b/beaker_kernel/lib/agent.py index 94a3fddb..ce953dad 100644 --- a/beaker_kernel/lib/agent.py +++ b/beaker_kernel/lib/agent.py @@ -59,12 +59,7 @@ def __init__( set_tool_execution_context(tool) async def react_async(self, query: str, react_context: dict = None) -> str: - try: - return await super().react_async(query, react_context) - finally: - # Clear loaded integration docs after each query to avoid context bloat - if hasattr(self.context, 'active_integrations'): - self.context.active_integrations.clear() + return await super().react_async(query, react_context) async def execute(self, *args, **kwargs) -> str: return await super().execute(*args, **kwargs) diff --git a/beaker_kernel/lib/integrations/adhoc.py b/beaker_kernel/lib/integrations/adhoc.py index 3c31e398..73869055 100644 --- a/beaker_kernel/lib/integrations/adhoc.py +++ b/beaker_kernel/lib/integrations/adhoc.py @@ -607,6 +607,30 @@ async def load_integration_docs(self, integration: str, agent: AgentRef, loop: L return f"Documentation for '{integration}' has been loaded into your context:\n\n{docs}" + @tool + async def unload_integration_docs(self, integration: str, agent: AgentRef, loop: LoopControllerRef, react_context: ReactContextRef) -> str: + """ +Unloads a previously loaded integration's documentation from your context window. Use this after you have +finished working with an integration to free up context space. + +You should call this tool when: +- You have completed the user's task involving the integration +- You no longer need to reference the integration's documentation + +Args: + integration (str): The name/slug of the integration to unload. + +Returns: + str: Confirmation that the documentation was unloaded. + """ + if hasattr(agent, 'context') and hasattr(agent.context, 'active_integrations'): + if integration in agent.context.active_integrations: + agent.context.active_integrations.discard(integration) + return f"Documentation for '{integration}' has been unloaded from your context." + else: + return f"Integration '{integration}' is not currently loaded." + return f"Integration '{integration}' is not currently loaded." + @tool async def draft_integration_code(self, integration: str, goal: str, agent: AgentRef, loop: LoopControllerRef, react_context: ReactContextRef) -> str: """