diff --git a/beaker_kernel/lib/context.py b/beaker_kernel/lib/context.py index 1d736fb7..d1263d52 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,19 @@ 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: + 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 diff --git a/beaker_kernel/lib/integrations/adhoc.py b/beaker_kernel/lib/integrations/adhoc.py index d45ac055..73869055 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,58 @@ 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 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: """