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>
121 lines
2.8 KiB
Python
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",
|
|
]
|