xml-pipeline/xml_pipeline/message_bus/__init__.py
dullfig 8b11323a8b Add token budget enforcement and usage tracking
Token Budget System:
- ThreadBudgetRegistry tracks per-thread token usage with configurable limits
- BudgetExhaustedError raised when thread exceeds max_tokens_per_thread
- Integrates with LLMRouter to block LLM calls when budget exhausted
- Automatic cleanup when threads are pruned

Usage Tracking (for production billing):
- UsageTracker emits events after each LLM completion
- Subscribers receive UsageEvent with tokens, latency, estimated cost
- Cost estimation for common models (Grok, Claude, GPT, etc.)
- Aggregate stats by agent, model, and totals

Configuration:
- max_tokens_per_thread in organism.yaml (default 100k)
- LLMRouter.complete() accepts thread_id and metadata parameters

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-27 21:07:43 -08:00

121 lines
2.8 KiB
Python

"""
message_bus — Stream-based message pump for AgentServer v2.1
The message pump handles message flow through the organism:
- YAML config → bootstrap → pump → handlers → responses → loop
Key classes:
StreamPump Main pump class (queue-backed, aiostream-powered)
SystemPipeline Entry point for external messages (console, webhook)
ConfigLoader Load organism.yaml and resolve imports
Listener Runtime listener with handler and routing info
MessageState Message flowing through pipeline steps
Usage:
from xml_pipeline.message_bus import StreamPump, SystemPipeline, bootstrap
pump = await bootstrap("config/organism.yaml")
system = SystemPipeline(pump)
# Inject from console
thread_id = await system.inject_console("@greeter Dan", user="admin")
await pump.run()
"""
from xml_pipeline.message_bus.stream_pump import (
StreamPump,
ConfigLoader,
Listener,
ListenerConfig,
OrganismConfig,
bootstrap,
get_stream_pump,
set_stream_pump,
reset_stream_pump,
# Event hooks
PumpEvent,
MessageReceivedEvent,
MessageSentEvent,
AgentStateEvent,
ThreadEvent,
ReloadEvent,
)
from xml_pipeline.message_bus.message_state import (
MessageState,
HandlerMetadata,
)
from xml_pipeline.message_bus.system_pipeline import (
SystemPipeline,
ExternalMessage,
)
from xml_pipeline.message_bus.sequence_registry import (
SequenceState,
SequenceRegistry,
get_sequence_registry,
reset_sequence_registry,
)
from xml_pipeline.message_bus.buffer_registry import (
BufferState,
BufferItemResult,
BufferRegistry,
get_buffer_registry,
reset_buffer_registry,
)
from xml_pipeline.message_bus.budget_registry import (
ThreadBudget,
ThreadBudgetRegistry,
BudgetExhaustedError,
get_budget_registry,
configure_budget_registry,
reset_budget_registry,
)
__all__ = [
# Pump
"StreamPump",
"ConfigLoader",
"Listener",
"ListenerConfig",
"OrganismConfig",
"bootstrap",
"get_stream_pump",
"set_stream_pump",
"reset_stream_pump",
# Event hooks
"PumpEvent",
"MessageReceivedEvent",
"MessageSentEvent",
"AgentStateEvent",
"ThreadEvent",
"ReloadEvent",
# Message state
"MessageState",
"HandlerMetadata",
# System pipeline
"SystemPipeline",
"ExternalMessage",
# Sequence registry
"SequenceState",
"SequenceRegistry",
"get_sequence_registry",
"reset_sequence_registry",
# Buffer registry
"BufferState",
"BufferItemResult",
"BufferRegistry",
"get_buffer_registry",
"reset_buffer_registry",
# Budget registry
"ThreadBudget",
"ThreadBudgetRegistry",
"BudgetExhaustedError",
"get_budget_registry",
"configure_budget_registry",
"reset_budget_registry",
]