"""Hook-first Bub framework runtime.""" from __future__ import annotations import contextlib from collections.abc import AsyncGenerator, AsyncIterator, Iterator from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any import pluggy import typer from dotenv import load_dotenv from loguru import logger from republic import AsyncTapeStore, RepublicError, TapeContext from republic.core.errors import ErrorKind from republic.tape import TapeStore from bub import configure from bub.envelope import content_of, field_of, unpack_batch from bub.hook_runtime import _SKIP_VALUE, HookRuntime from bub.hookspecs import BUB_HOOK_NAMESPACE, BubHookSpecs from bub.types import Envelope, MessageHandler, OutboundChannelRouter, TurnResult if TYPE_CHECKING: from bub.channels.base import Channel load_dotenv() DEFAULT_HOME = Path.home() / ".bub" DEFAULT_CONFIG_FILE = (DEFAULT_HOME / "config.yml").resolve() @dataclass(frozen=True) class PluginStatus: is_success: bool detail: str | None = None class BubFramework: """Minimal framework core. Everything grows from hook skills.""" def __init__(self, config_file: Path = DEFAULT_CONFIG_FILE) -> None: self.workspace = Path.cwd().resolve() self.config_file = config_file.resolve() self._plugin_manager = pluggy.PluginManager(BUB_HOOK_NAMESPACE) self._plugin_manager.add_hookspecs(BubHookSpecs) self._hook_runtime = HookRuntime(self._plugin_manager) self._plugin_status: dict[str, PluginStatus] = {} self._outbound_router: OutboundChannelRouter | None = None self._tape_store: TapeStore | AsyncTapeStore | None = None configure.load(self.config_file) def _load_builtin_hooks(self) -> None: from bub.builtin.hook_impl import BuiltinImpl impl = BuiltinImpl(self) try: self._plugin_manager.register(impl, name="builtin") except Exception as exc: self._plugin_status["builtin"] = PluginStatus(is_success=False, detail=str(exc)) else: self._plugin_status["builtin"] = PluginStatus(is_success=True) def load_hooks(self) -> None: import importlib.metadata pending_plugins: list[tuple[str, Any]] = [] self._load_builtin_hooks() for entry_point in importlib.metadata.entry_points(group="bub"): try: plugin = entry_point.load() except Exception as exc: logger.warning(f"Failed to load plugin '{entry_point.name}': {exc}") self._plugin_status[entry_point.name] = PluginStatus(is_success=False, detail=str(exc)) else: pending_plugins.append((entry_point.name, plugin)) for plugin_name, plugin in pending_plugins: try: if callable(plugin): # Support entry points that are classes plugin = plugin(self) self._plugin_manager.register(plugin, name=plugin_name) except Exception as exc: logger.warning(f"Failed to initialize plugin '{plugin_name}': {exc}") self._plugin_status[plugin_name] = PluginStatus(is_success=False, detail=str(exc)) else: self._plugin_status[plugin_name] = PluginStatus(is_success=True) def create_cli_app(self) -> typer.Typer: """Create CLI app by collecting commands from hooks. Can be used for custom CLI entry point.""" app = typer.Typer(name="bub", help="Batteries-included, hook-first AI framework", add_completion=False) @app.callback(invoke_without_command=True) def _main( ctx: typer.Context, workspace: str | None = typer.Option(None, "--workspace", "-w", help="Path to the workspace"), ) -> None: if workspace: self.workspace = Path(workspace).resolve() ctx.obj = self self._hook_runtime.call_many_sync("register_cli_commands", app=app) return app async def process_inbound(self, inbound: Envelope, stream_output: bool = False) -> TurnResult: """Run one inbound message through hooks and return turn result.""" try: session_id = await self._hook_runtime.call_first( "resolve_session", message=inbound ) or self._default_session_id(inbound) if isinstance(inbound, dict): inbound.setdefault("session_id", session_id) state = {"_runtime_workspace": str(self.workspace)} for hook_state in reversed( await self._hook_runtime.call_many("load_state", message=inbound, session_id=session_id) ): if isinstance(hook_state, dict): state.update(hook_state) prompt = await self._hook_runtime.call_first( "build_prompt", message=inbound, session_id=session_id, state=state ) if not prompt: prompt = content_of(inbound) model_output = "" try: model_output = await self._run_model(inbound, prompt, session_id, state, stream_output) finally: await self._hook_runtime.call_many( "save_state", session_id=session_id, state=state, message=inbound, model_output=model_output, ) outbounds = await self._collect_outbounds(inbound, session_id, state, model_output) for outbound in outbounds: await self._hook_runtime.call_many("dispatch_outbound", message=outbound) return TurnResult(session_id=session_id, prompt=prompt, model_output=model_output, outbounds=outbounds) except Exception as exc: logger.exception("Error processing inbound message") await self._hook_runtime.notify_error(stage="turn", error=exc, message=inbound) raise async def _run_model( self, inbound: Envelope, prompt: str | list[dict], session_id: str, state: dict[str, Any], stream_output: bool, ) -> str: if not stream_output: output = await self._hook_runtime.run_model(prompt=prompt, session_id=session_id, state=state) if output is None: await self._hook_runtime.notify_error( stage="run_model", error=RuntimeError("no model skill returned output"), message=inbound, ) return prompt if isinstance(prompt, str) else content_of(inbound) return output stream = await self._hook_runtime.run_model_stream(prompt=prompt, session_id=session_id, state=state) if stream is None: await self._hook_runtime.notify_error( stage="run_model", error=RuntimeError("no model skill returned output"), message=inbound, ) return prompt if isinstance(prompt, str) else content_of(inbound) else: parts: list[str] = [] if self._outbound_router is not None: stream = self._outbound_router.wrap_stream(inbound, stream) async for event in stream: if event.kind == "text": parts.append(str(event.data.get("delta", ""))) elif event.kind == "error": # Turn "kind" to enum type otherwise the RepublicError's __str__ won't work well data = { **event.data, "kind": ErrorKind(event.data.get("kind", "unknown")), } await self._hook_runtime.notify_error( stage="run_model", error=RepublicError(**data), message=inbound ) return "".join(parts) def hook_report(self) -> dict[str, list[str]]: """Return hook implementation summary for diagnostics.""" return self._hook_runtime.hook_report() def bind_outbound_router(self, router: OutboundChannelRouter | None) -> None: self._outbound_router = router async def dispatch_via_router(self, message: Envelope) -> bool: if self._outbound_router is None: return False return await self._outbound_router.dispatch_output(message) async def quit_via_router(self, session_id: str) -> None: if self._outbound_router is not None: await self._outbound_router.quit(session_id) @staticmethod def _default_session_id(message: Envelope) -> str: session_id = field_of(message, "session_id") if session_id is not None: return str(session_id) channel = str(field_of(message, "channel", "default")) chat_id = str(field_of(message, "chat_id", "default")) return f"{channel}:{chat_id}" async def _collect_outbounds( self, message: Envelope, session_id: str, state: dict[str, Any], model_output: str, ) -> list[Envelope]: batches = await self._hook_runtime.call_many( "render_outbound", message=message, session_id=session_id, state=state, model_output=model_output, ) outbounds: list[Envelope] = [] for batch in batches: outbounds.extend(unpack_batch(batch)) if outbounds: return outbounds fallback: dict[str, Any] = { "content": model_output, "session_id": session_id, } channel = field_of(message, "channel") chat_id = field_of(message, "chat_id") if channel is not None: fallback["channel"] = channel if chat_id is not None: fallback["chat_id"] = chat_id return [fallback] def get_channels(self, message_handler: MessageHandler) -> dict[str, Channel]: channels: dict[str, Channel] = {} for result in self._hook_runtime.call_many_sync("provide_channels", message_handler=message_handler): for channel in result: if channel.name not in channels: channels[channel.name] = channel return channels @contextlib.asynccontextmanager async def running(self) -> AsyncGenerator[contextlib.AsyncExitStack, None]: async with contextlib.AsyncExitStack() as stack: tape_store = self._hook_runtime.call_first_sync("provide_tape_store") # Allow plugins to return either TapeStore/AsyncTapeStore instances or context managers for them # This benefits plugins that need to initialize and clean up resources with the tape store. if isinstance(tape_store, AsyncIterator): tape_store = await stack.enter_async_context(contextlib.asynccontextmanager(lambda: tape_store)()) elif isinstance(tape_store, Iterator): tape_store = stack.enter_context(contextlib.contextmanager(lambda: tape_store)()) self._tape_store = tape_store try: yield stack finally: self._tape_store = None def get_tape_store(self) -> TapeStore | AsyncTapeStore | None: return self._tape_store def get_system_prompt(self, prompt: str | list[dict], state: dict[str, Any]) -> str: return "\n\n".join( result for result in reversed(self._hook_runtime.call_many_sync("system_prompt", prompt=prompt, state=state)) if result ) def build_tape_context(self) -> TapeContext: return self._hook_runtime.call_first_sync("build_tape_context") def collect_onboard_config(self) -> dict[str, Any]: current_config: dict[str, Any] = {} for impl in reversed(list(self._hook_runtime._iter_hookimpls("onboard_config"))): result = self._hook_runtime._invoke_impl_sync( hook_name="onboard_config", impl=impl, call_kwargs={"current_config": current_config}, kwargs={"current_config": current_config}, ) if result is _SKIP_VALUE: continue if result is None: continue if not isinstance(result, dict): raise TypeError("hook.onboard_config must return dict or None") configure.merge(current_config, result) return configure.validate(current_config)