admin/build-your-own-openclaw
1
0
Fork 0
mirror of https://github.com/czl9707/build-your-own-openclaw.git synced 2026-08-14 00:47:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

10-websocket

Browse Source
This commit is contained in:
zane
2026-03-10 22:35:02 -04:00
parent faa69751c3
commit 2b69d4727a
57 changed files with 4249 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
09-channels/src/mybot/channel/discord_channel.py
+2 -11
View File
@@ -41,12 +41,7 @@ class DiscordChannel(Channel[DiscordEventSource]):
platform_name = "discord"
def __init__(self, config: DiscordConfig):
"""
Initialize DiscordChannel.
Args:
config: Discord configuration
"""
"""Initialize DiscordChannel."""
self.config = config
self.client: discord.Client | None = None
self._running_task: asyncio.Task | None = None
@@ -54,11 +49,7 @@ class DiscordChannel(Channel[DiscordEventSource]):
async def run(
self, on_message: Callable[[str, DiscordEventSource], Awaitable[None]]
) -> None:
"""Run the Discord channel. Blocks until stop() is called.
Raises:
RuntimeError: If run() is called when already running.
"""
"""Run the Discord channel. Blocks until stop() is called."""
if self._running_task is not None:
raise RuntimeError("DiscordChannel already running")
09-channels/src/mybot/channel/telegram_channel.py
+2 -11
View File
@@ -42,12 +42,7 @@ class TelegramChannel(Channel[TelegramEventSource]):
platform_name = "telegram"
def __init__(self, config: TelegramConfig):
"""
Initialize TelegramChannel.
Args:
config: Telegram configuration
"""
"""Initialize TelegramChannel."""
self.config = config
self.application: Application | None = None
self._running_task: asyncio.Task | None = None
@@ -62,11 +57,7 @@ class TelegramChannel(Channel[TelegramEventSource]):
async def run(
self, on_message: Callable[[str, TelegramEventSource], Awaitable[None]]
) -> None:
"""Run the Telegram channel. Blocks until stop() is called.
Raises:
RuntimeError: If run() is called when already running.
"""
"""Run the Telegram channel. Blocks until stop() is called."""
if self.application is not None:
raise RuntimeError("TelegramChannel already running")
09-channels/src/mybot/core/context_guard.py
+6 -55
View File
@@ -30,14 +30,7 @@ class ContextGuard:
max_tool_result_chars: int = MAX_TOOL_RESULT_CHARS
def estimate_tokens(self, state: "SessionState") -> int:
"""Estimate token count for session state.
Args:
state: Session state to estimate
Returns:
Estimated token count
"""
"""Estimate token count for session state."""
if not state.messages:
return 0
return token_counter(
@@ -48,19 +41,7 @@ class ContextGuard:
self,
state: "SessionState",
) -> "SessionState":
"""Check token count, compact and roll session if needed.
Applies truncation to large tool results when over threshold,
then falls back to full compaction if still needed.
Args:
state: Current session state
Returns:
SessionState to use (same state if under threshold,
same state with truncated content if truncation was sufficient,
new rolled state if compaction was needed)
"""
"""Check token count, compact and roll session if needed."""
token_count = self.estimate_tokens(state)
if token_count < self.token_threshold:
@@ -80,14 +61,7 @@ class ContextGuard:
return min(compress_count, len(state.messages) - keep_count)
def _truncate_large_tool_results(self, messages: list[Message]) -> list[Message]:
"""Truncate oversized tool results to reduce context size.
Args:
messages: List of messages to process
Returns:
List of messages with large tool results truncated
"""
"""Truncate oversized tool results to reduce context size."""
result: list[Message] = []
for msg in messages:
if msg.get("role") == "tool":
@@ -112,14 +86,7 @@ class ContextGuard:
return result
def _serialize_messages_for_summary(self, messages: list[Message]) -> str:
"""Serialize messages to plain text for summarization.
Args:
messages: List of messages to serialize
Returns:
Plain text representation
"""
"""Serialize messages to plain text for summarization."""
lines = []
for msg in messages:
role = msg.get("role", "unknown")
@@ -143,15 +110,7 @@ class ContextGuard:
self,
state: "SessionState",
) -> "SessionState":
"""Compact history, roll to new session, return new messages.
Args:
state: Current session state
messages: Current full message list (with system prompt)
Returns:
Tuple of (compacted messages, new SessionState)
"""
"""Compact history, roll to new session, return new messages."""
new_session = state.agent.new_session(state.source)
self._clear_source_session_cache(str(state.source))
@@ -165,15 +124,7 @@ class ContextGuard:
self,
state: "SessionState",
) -> list[Message]:
"""Generate summary of older messages using agent's LLM.
Args:
state: Current session state
messages: History message list (without system prompt)
Returns:
Compacted message list with summary + recent messages
"""
"""Generate summary of older messages using agent's LLM."""
compress_count = self._compress_message_count(state)
old_messages = state.messages[:compress_count]
09-channels/src/mybot/utils/logging.py
+1 -7
View File
@@ -7,13 +7,7 @@ from mybot.utils.config import Config
from logging.handlers import RotatingFileHandler
def setup_logging(config: Config, console_output: bool = False) -> None:
"""
Set up logging for pickle-bot.
Args:
config: Application configuration
console_output: Whether to output logs to console (default: False)
"""
"""Set up logging for pickle-bot."""
format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
formatter = logging.Formatter(format_str)
10-websocket/README.md
+307
View File
@@ -0,0 +1,307 @@
# Step 09: Channels - Multi-Platform Support
Extend the agent to support multiple messaging platforms (CLI, Telegram, Discord) through a unified channel abstraction.
## What We Will Build
```
┌────────────────────────────────────────────────────────────────────┐
│ Server │
│ │
│ ┌─────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │EventBus │ │AgentWorker │ │DeliveryWorker│ │
│ │ │ │ │ │ │ │
│ └─────────────┘ └──────────────┘ └──────────────┘ │
│ ▲ ▲ │ │
│ │ │ │ │
│ │ ┌─────┴─────┐ │ │
│ │ │ Agent │ │ │
│ │ │ Session │ │ │
│ │ └───────────┘ │ │
│ │ │ │
│ ┌──────┴─────────────────────────────────────┴──────┐ │
│ │ ChannelWorker │ │
│ │ │ │
│ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │
│ │ │CLI │ │Telegram │ │Discord │ │ │
│ │ │Channel │ │Channel │ │Channel │ │ │
│ │ └─────────┘ └─────────┘ └─────────┘ │ │
│ └──────────────────────────────────────────────────┘ │
│ │
│ ┌──────────────────────────────────────────────────────┐ │
│ │ RoutingTable │ │
│ │ - Maps sources to agents │ │
│ │ - Manages session affinity │ │
│ └──────────────────────────────────────────────────────┘ │
└────────────────────────────────────────────────────────────────────┘
```
**Key Components:**
- **EventSource** - Abstract base for platform-specific event sources (CLI, Telegram, Discord)
- **Channel** - Abstract base for messaging platforms with run/reply/stop interface
- **ChannelWorker** - Manages multiple channels and publishes InboundEvents
- **DeliveryWorker** - Subscribes to OutboundEvents and delivers via appropriate channel
- **RoutingTable** - Maps sources to sessions and agents
- **Server** - Orchestrates all workers (EventBus, AgentWorker, ChannelWorker, DeliveryWorker)
## Key Changes
### 1. EventSource and Platform Sources ([src/mybot/core/events.py](src/mybot/core/events.py))
```python
class EventSource(ABC):
"""Abstract base for all event sources."""
_registry: ClassVar[dict[str, type["EventSource"]]] = {}
_namespace: ClassVar[str] = ""
@property
def platform_name(self) -> str | None:
if not self.is_platform:
return None
return self._namespace.split("-", 1)[1]
@classmethod
def from_string(cls, s: str) -> "EventSource":
"""Parse string to EventSource using namespace registry."""
namespace = s.split(":")[0]
source_cls = cls._registry.get(namespace)
return source_cls.from_string(s)
@dataclass
class CliEventSource(EventSource):
"""Source for CLI-originated events."""
_namespace = "platform-cli"
def __str__(self) -> str:
return "platform-cli:cli-user"
```
### 2. Channel Base Class ([src/mybot/channel/base.py](src/mybot/channel/base.py))
```python
class Channel(ABC, Generic[T]):
"""Abstract base for messaging platforms."""
@property
@abstractmethod
def platform_name(self) -> str:
pass
@abstractmethod
async def run(self, on_message: Callable[[str, T], Awaitable[None]]) -> None:
"""Run the channel. Blocks until stop() is called."""
pass
@abstractmethod
async def reply(self, content: str, source: T) -> None:
"""Reply to incoming message."""
pass
@abstractmethod
async def stop(self) -> None:
"""Stop listening and cleanup resources."""
pass
```
### 3. ChannelWorker ([src/mybot/server/channel_worker.py](src/mybot/server/channel_worker.py))
```python
class ChannelWorker(Worker):
"""Ingests messages from platforms, publishes INBOUND events."""
async def run(self) -> None:
"""Start all channels and process incoming messages."""
channel_tasks = [
channel.run(self._create_callback(channel.platform_name))
for channel in self.channels
]
await asyncio.gather(*channel_tasks)
def _create_callback(self, platform: str):
async def callback(message: str, source: EventSource) -> None:
session_id = self.context.routing_table.get_or_create_session_id(source)
event = InboundEvent(
session_id=session_id,
source=source,
content=message,
)
await self.context.eventbus.publish(event)
return callback
```
### 4. DeliveryWorker ([src/mybot/server/delivery_worker.py](src/mybot/server/delivery_worker.py))
```python
class DeliveryWorker(SubscriberWorker):
"""Delivers outbound messages to platforms."""
async def handle_event(self, event: OutboundEvent) -> None:
"""Handle an outbound message event."""
session_info = self._get_session_source(event.session_id)
source = self._get_delivery_source(session_info)
if source and source.platform_name:
channel = self._get_channel(source.platform_name)
if channel:
await channel.reply(event.content, source)
self.context.eventbus.ack(event)
```
### 5. RoutingTable ([src/mybot/core/routing.py](src/mybot/core/routing.py))
```python
@dataclass
class RoutingTable:
"""Routes sources to agents using regex bindings."""
def get_or_create_session_id(self, source: EventSource) -> str:
"""Get existing or create new session_id for source."""
source_str = str(source)
# Check for existing session (affinity)
source_session = self._context.config.sources.get(source_str)
if source_session:
return source_session.session_id
# Create new session
agent_id = self.resolve(source_str)
agent_def = self._context.agent_loader.load(agent_id)
agent = Agent(agent_def, self._context)
session = agent.new_session(source)
# Cache session
self._context.config.set_runtime(
f"sources.{source_str}", SourceSessionConfig(session_id=session.session_id)
)
return session.session_id
```
### 5. Server Updates ([src/mybot/server/server.py](src/mybot/server/server.py))
```python
class Server:
"""Orchestrates workers with queue-based communication."""
def _setup_workers(self) -> None:
# Create WebSocketWorker first and attach to context
ws_worker = WebSocketWorker(self.context)
self.context.websocket_worker = ws_worker
self.workers = [
self.context.eventbus,
AgentWorker(self.context),
DeliveryWorker(self.context),
ws_worker, # WebSocketWorker added to workers
]
if self.context.config.channels.enabled:
self.workers.append(ChannelWorker(self.context))
async def run(self) -> None:
self._setup_workers()
self._start_workers()
# Start API server if configured
if self.context.config.server:
self._api_task = asyncio.create_task(self._run_api())
await self._monitor_workers()
async def _run_api(self) -> None:
"""Run the WebSocket API server."""
app = create_app(self.context)
config = uvicorn.Config(
app,
host=self.context.config.server.host,
port=self.context.config.server.port,
)
server = uvicorn.Server(config)
await server.serve()
```
async def run(self) -> None:
self._setup_workers()
self._start_workers()
await self._monitor_workers()
```
## How to Run
**Start server with WebSocket:**
```bash
cd 10-websocket
uv run my-bot server
```
The server will start on http://0.0.0.0:8000 (configurable in config.user.yaml):
- **Web UI**: http://localhost:8000/
- **WebSocket**: ws://localhost:8000/ws
**Chat via WebSocket client:**
```javascript
const ws = new WebSocket('ws://localhost:8000/ws');
// Send message
ws.send(JSON.stringify({
source: "user-123",
content: "Hello, agent!"
}));
// Receive events
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
console.log(data.type, data.content);
};
```
**Or use the built-in web UI** by opening http://localhost:8000/ in your browser.
## Architecture Notes
**Event Flow:**
1. User sends message via platform (CLI, Telegram, Discord)
2. Channel receives message and creates EventSource
3. ChannelWorker creates/picks session via RoutingTable
4. ChannelWorker publishes InboundEvent to EventBus
5. AgentWorker processes event and generates response
6. AgentWorker publishes OutboundEvent to EventBus
7. DeliveryWorker receives OutboundEvent
8. DeliveryWorker looks up session's source and sends via appropriate channel
**Session Affinity:**
- Each EventSource (e.g., "platform-telegram:123:456") maps to one session
- First message creates session, subsequent messages reuse it
- Session ID cached in config.runtime.yaml
- Enables persistent conversations across restarts
**Multi-Agent Routing:**
- RoutingTable matches sources to agents via regex patterns
- Enables different agents for different platforms/users
- Falls back to default_agent if no match
## What's Next
Step 10 will add **WebSocket UI** - real-time web interface for interacting with agents.
```python
class DeliveryWorker(SubscriberWorker):
"""Delivers outbound messages to platforms."""
async def handle_event(self, event: OutboundEvent) -> None:
"""Handle an outbound message event."""
session_info = self._get_session_source(event.session_id)
source = self._get_delivery_source(session_info)
if source and source.platform_name:
channel = self._get_channel(source.platform_name)
if channel:
await channel.reply(event.content, source)
self.context.eventbus.ack(event)
```
10-websocket/pyproject.toml
+31
View File
@@ -0,0 +1,31 @@
[project]
name = "my-bot"
version = "0.1.0"
description = "Step 10: WebSocket UI - Real-Time Interface"
requires-python = ">=3.11"
dependencies = [
"litellm>=1.0.0",
"typer>=0.9.0",
"rich>=13.0.0",
"pydantic>=2.0.0",
"pyyaml>=6.0",
"httpx>=0.27.0",
"crawl4ai>=0.3.0",
"watchdog>=3.0.0",
"python-telegram-bot>=20.0",
"discord.py>=2.0",
"fastapi>=0.104.0",
"uvicorn[standard]>=0.24.0",
"websockets>=12.0",
]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["src/mybot"]
[project.scripts]
my-bot = "mybot.cli.main:app"
10-websocket/src/mybot/channel/__init__.py
+7
View File
@@ -0,0 +1,7 @@
"""Channel implementations for different platforms."""
from mybot.channel.base import Channel
from mybot.channel.telegram_channel import TelegramChannel
from mybot.channel.discord_channel import DiscordChannel
__all__ = ["Channel", "TelegramChannel", "DiscordChannel"]
10-websocket/src/mybot/channel/base.py
+57
View File
@@ -0,0 +1,57 @@
"""Abstract base class for channel implementations."""
from abc import ABC, abstractmethod
from typing import Callable, Awaitable, Generic, TypeVar, Any
from mybot.core.events import EventSource
from mybot.utils.config import Config
T = TypeVar("T", bound=EventSource)
class Channel(ABC, Generic[T]):
"""Abstract base for messaging platforms with EventSource-based context."""
@property
@abstractmethod
def platform_name(self) -> str:
"""Platform identifier."""
pass
@abstractmethod
async def run(self, on_message: Callable[[str, T], Awaitable[None]]) -> None:
"""Run the channel. Blocks until stop() is called."""
pass
@abstractmethod
def is_allowed(self, source: T) -> bool:
"""Check if sender is whitelisted."""
pass
@abstractmethod
async def reply(self, content: str, source: T) -> None:
"""Reply to incoming message."""
pass
@abstractmethod
async def stop(self) -> None:
"""Stop listening and cleanup resources."""
pass
@staticmethod
def from_config(config: Config) -> list["Channel[Any]"]:
"""Create channel instances from configuration."""
# Inline imports to avoid circular dependency
from mybot.channel.telegram_channel import TelegramChannel
from mybot.channel.discord_channel import DiscordChannel
channels: list["Channel[Any]"] = []
channel_config = config.channels
if channel_config.telegram and channel_config.telegram.enabled:
channels.append(TelegramChannel(channel_config.telegram))
if channel_config.discord and channel_config.discord.enabled:
channels.append(DiscordChannel(channel_config.discord))
return channels
10-websocket/src/mybot/channel/discord_channel.py
+150
View File
@@ -0,0 +1,150 @@
"""Discord channel implementation."""
import asyncio
from dataclasses import dataclass
import logging
from typing import Callable, Awaitable
import discord
from mybot.core.events import EventSource
from mybot.channel.base import Channel
from mybot.utils.config import DiscordConfig
logger = logging.getLogger(__name__)
@dataclass
class DiscordEventSource(EventSource):
"""Source for Discord-originated events."""
_namespace = "platform-discord"
user_id: str
channel_id: str
def __str__(self) -> str:
return f"platform-discord:{self.user_id}:{self.channel_id}"
@classmethod
def from_string(cls, s: str) -> "DiscordEventSource":
_, user_id, channel_id = s.split(":")
return cls(user_id=user_id, channel_id=channel_id)
@property
def platform_name(self) -> str:
return "discord"
class DiscordChannel(Channel[DiscordEventSource]):
"""Discord platform implementation using discord.py."""
platform_name = "discord"
def __init__(self, config: DiscordConfig):
"""Initialize DiscordChannel."""
self.config = config
self.client: discord.Client | None = None
self._running_task: asyncio.Task | None = None
async def run(
self, on_message: Callable[[str, DiscordEventSource], Awaitable[None]]
) -> None:
"""Run the Discord channel. Blocks until stop() is called."""
if self._running_task is not None:
raise RuntimeError("DiscordChannel already running")
logger.info(f"Channel enabled with platform: {self.platform_name}")
# Configure intents
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
self.client = discord.Client(intents=intents)
@self.client.event
async def _on_discord_message(message: discord.Message) -> None:
"""Handle incoming Discord message."""
# Ignore bot's own messages
if self.client and message.author == self.client.user:
return
# Check channel restriction (optional)
if (
self.config.channel_id
and str(message.channel.id) != self.config.channel_id
):
return
# Only handle text messages
if not message.content:
return
# Extract user_id (the person) and channel_id (the channel)
user_id = str(message.author.id)
channel_id = str(message.channel.id)
content = message.content
logger.info(
f"Received Discord message from user {user_id} in channel {channel_id}"
)
source = DiscordEventSource(user_id=user_id, channel_id=channel_id)
try:
await on_message(content, source)
except Exception as e:
logger.error(f"Error in message callback: {e}")
# Start the bot and store the task
self._running_task = asyncio.create_task(
self.client.start(self.config.bot_token)
)
logger.info("DiscordChannel started")
await self._running_task
def is_allowed(self, source: DiscordEventSource) -> bool:
"""Check if sender is whitelisted."""
if not self.config.allowed_user_ids:
return True
return source.user_id in self.config.allowed_user_ids
async def reply(self, content: str, source: DiscordEventSource) -> None:
"""Reply to incoming message in the same channel."""
if not self.client:
raise RuntimeError("DiscordChannel not started")
try:
channel = self.client.get_channel(int(source.channel_id))
if not channel:
raise ValueError(f"Channel {source.channel_id} not found")
# Type ignore: discord.py returns a union, but we know text channels have send()
await channel.send(content) # type: ignore[union-attr]
logger.debug(f"Sent Discord reply to {source.channel_id}")
except Exception as e:
logger.error(f"Failed to send Discord reply: {e}")
raise
async def stop(self) -> None:
"""Stop Discord bot and cleanup."""
# Idempotent: skip if not running
if self.client is None:
logger.debug("DiscordChannel not running, skipping stop")
return
await self.client.close()
# Wait for running task to complete
if self._running_task and not self._running_task.done():
try:
await asyncio.wait_for(self._running_task, timeout=2.0)
except asyncio.TimeoutError:
logger.warning("Running task did not complete in time")
except Exception:
pass # Task may have already failed
self.client = None
self._running_task = None
logger.info("DiscordChannel stopped")
10-websocket/src/mybot/channel/telegram_channel.py
+161
View File
@@ -0,0 +1,161 @@
"""Telegram channel implementation."""
import asyncio
from dataclasses import dataclass
import logging
from typing import Callable, Awaitable
from telegram import Update
from telegram.ext import Application, MessageHandler, filters, ContextTypes
from mybot.core.events import EventSource
from mybot.channel.base import Channel
from mybot.utils.config import TelegramConfig
logger = logging.getLogger(__name__)
@dataclass
class TelegramEventSource(EventSource):
"""Source for Telegram-originated events."""
_namespace = "platform-telegram"
user_id: str
chat_id: str
def __str__(self) -> str:
return f"platform-telegram:{self.user_id}:{self.chat_id}"
@classmethod
def from_string(cls, s: str) -> "TelegramEventSource":
_, user_id, chat_id = s.split(":")
return cls(user_id=user_id, chat_id=chat_id)
@property
def platform_name(self) -> str:
return "telegram"
class TelegramChannel(Channel[TelegramEventSource]):
"""Telegram platform implementation using python-telegram-bot."""
platform_name = "telegram"
def __init__(self, config: TelegramConfig):
"""Initialize TelegramChannel."""
self.config = config
self.application: Application | None = None
self._running_task: asyncio.Task | None = None
self._stop_event: asyncio.Event | None = None
def is_allowed(self, source: TelegramEventSource) -> bool:
"""Check if sender is whitelisted."""
if not self.config.allowed_user_ids:
return True
return source.user_id in self.config.allowed_user_ids
async def run(
self, on_message: Callable[[str, TelegramEventSource], Awaitable[None]]
) -> None:
"""Run the Telegram channel. Blocks until stop() is called."""
if self.application is not None:
raise RuntimeError("TelegramChannel already running")
logger.info(f"Channel enabled with platform: {self.platform_name}")
self.application = Application.builder().token(self.config.bot_token).build()
self._stop_event = asyncio.Event()
async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
"""Handle incoming Telegram message."""
if (
update.message
and update.message.text
and update.effective_chat
and update.message.from_user
):
# Extract user_id (the person) and chat_id (the conversation)
user_id = str(update.message.from_user.id)
chat_id = str(update.effective_chat.id)
message = update.message.text
logger.info(
f"Received Telegram message from user {user_id} in chat {chat_id}"
)
source = TelegramEventSource(user_id=user_id, chat_id=chat_id)
try:
await on_message(message, source)
except Exception as e:
logger.error(f"Error in message callback: {e}")
handler = MessageHandler(filters.TEXT, handle_message)
self.application.add_handler(handler)
# Start the bot
await self.application.initialize()
await self.application.start()
if self.application.updater:
await self.application.updater.start_polling()
logger.info("TelegramChannel started")
# Create the running task that monitors for stop
async def run_until_stopped():
"""Run until stop() is called or updater stops unexpectedly."""
while self.application and self.application.updater:
if self.application.updater.running:
if self._stop_event and self._stop_event.is_set():
return # Graceful stop
await asyncio.sleep(1)
else:
if self._stop_event and not self._stop_event.is_set():
raise RuntimeError("Telegram updater stopped unexpectedly")
return
self._running_task = asyncio.create_task(run_until_stopped())
await self._running_task
async def reply(self, content: str, source: TelegramEventSource) -> None:
"""Reply to incoming message."""
if not self.application:
raise RuntimeError("TelegramChannel not started")
try:
await self.application.bot.send_message(
chat_id=int(source.chat_id), text=content
)
logger.debug(f"Sent Telegram reply to {source.chat_id}")
except Exception as e:
logger.error(f"Failed to send Telegram reply: {e}")
raise
async def stop(self) -> None:
"""Stop Telegram bot and cleanup."""
# Idempotent: skip if not running
if self.application is None:
logger.debug("TelegramChannel not running, skipping stop")
return
# Signal the running task to stop
if self._stop_event:
self._stop_event.set()
if self.application.updater and self.application.updater.running:
await self.application.updater.stop()
await self.application.stop()
await self.application.shutdown()
# Wait for running task to complete
if self._running_task and not self._running_task.done():
try:
await asyncio.wait_for(self._running_task, timeout=2.0)
except asyncio.TimeoutError:
logger.warning("Running task did not complete in time")
except Exception:
pass # Task may have already failed
self.application = None
self._running_task = None
self._stop_event = None
logger.info("TelegramChannel stopped")
10-websocket/src/mybot/cli/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""CLI interface for my-bot."""
from mybot.cli.main import app
__all__ = ["app"]
10-websocket/src/mybot/cli/chat.py
+125
View File
@@ -0,0 +1,125 @@
"""Chat CLI command for interactive sessions with event-driven architecture."""
import asyncio
import typer
from rich.console import Console
from rich.panel import Panel
from rich.prompt import Prompt
from rich.text import Text
from mybot.core.agent import Agent
from mybot.core.context import SharedContext
from mybot.core.events import (
OutboundEvent,
InboundEvent,
CliEventSource,
)
from mybot.server import (
AgentWorker,
Worker,
)
from mybot.utils.config import Config, ConfigReloader
from mybot.utils.logging import setup_logging
class ChatLoop:
"""Interactive chat session using event-driven architecture."""
def __init__(self, config: Config, agent_id: str | None = None):
self.config = config
self.console = Console()
self.context = SharedContext(config=config, channels=[])
self.config_reloader = ConfigReloader(config)
self.workers: list[Worker] = [
self.context.eventbus,
AgentWorker(self.context),
]
self.response_queue: asyncio.Queue[OutboundEvent] = asyncio.Queue()
self.context.eventbus.subscribe(OutboundEvent, self.handle_outbound_event)
agent_id = agent_id or config.default_agent
self.agent_def = self.context.agent_loader.load(agent_id)
async def handle_outbound_event(self, event: OutboundEvent) -> None:
"""Handle outbound events by adding to response queue."""
await self.response_queue.put(event)
self.context.eventbus.ack(event)
def get_user_input(self) -> str:
"""Get user input with styled prompt."""
prompt_text = Text("You", style="cyan")
user_input = Prompt.ask(prompt_text, console=self.console)
return user_input.strip()
def display_agent_response(self, content: str) -> None:
"""Display agent response with styled prefix."""
prefix = Text(f"{self.agent_def.id}: ", style="green")
self.console.print(prefix, end="")
self.console.print(content)
async def run(self) -> None:
"""Run the interactive chat loop."""
self.console.print(
Panel(
Text("Welcome to my-bot!", style="bold cyan"),
title="Chat",
border_style="cyan",
)
)
self.console.print("Type '/help' for commands, 'quit' or 'exit' to end.\n")
self.config_reloader.start()
for worker in self.workers:
worker.start()
session_id = (
Agent(self.agent_def, self.context).new_session(CliEventSource()).session_id
)
try:
while True:
user_input = await asyncio.to_thread(self.get_user_input)
if user_input.lower() in ("quit", "exit", "q"):
self.console.print("\n[bold yellow]Goodbye![/bold yellow]")
break
if not user_input:
continue
event = InboundEvent(
session_id=session_id,
source=CliEventSource(),
content=user_input,
)
await self.context.eventbus.publish(event)
try:
response = await asyncio.wait_for(
self.response_queue.get(), timeout=60.0
)
self.display_agent_response(response.content)
except asyncio.TimeoutError:
self.console.print("[red]Agent response timed out[/red]")
self.console.print()
except (KeyboardInterrupt, EOFError):
self.console.print("\n[bold yellow]Goodbye![/bold yellow]")
finally:
for worker in self.workers:
await worker.stop()
self.config_reloader.stop()
def chat_command(ctx: typer.Context, agent_id: str | None = None) -> None:
"""Start interactive chat session."""
config = ctx.obj.get("config")
setup_logging(config, console_output=False)
chat_loop = ChatLoop(config, agent_id=agent_id)
asyncio.run(chat_loop.run())
10-websocket/src/mybot/cli/main.py
+80
View File
@@ -0,0 +1,80 @@
"""CLI interface for my-bot using Typer."""
from pathlib import Path
from typing import Annotated
import typer
from rich.console import Console
from mybot.cli.chat import chat_command
from mybot.cli.server import server_command
from mybot.utils.config import Config
app = typer.Typer(
name="my-bot",
help="my-bot: Personal AI Assistant",
no_args_is_help=True,
add_completion=True,
)
console = Console()
def workspace_callback(ctx: typer.Context, workspace: str) -> Path:
"""Store workspace path in context for later use."""
ctx.ensure_object(dict)
ctx.obj["workspace"] = Path(workspace)
return Path(workspace)
@app.callback()
def main(
ctx: typer.Context,
workspace: str = typer.Option(
"../default_workspace",
"--workspace",
"-w",
help="Path to workspace directory",
callback=workspace_callback,
),
) -> None:
"""Configuration is loaded from workspace/config.user.yaml by default."""
workspace_path = ctx.obj["workspace"]
config_file = workspace_path / "config.user.yaml"
if not config_file.exists():
console.print(f"[yellow]No configuration found at {config_file}[/yellow]")
raise typer.Exit(1)
try:
cfg = Config.load(workspace_path)
ctx.obj["config"] = cfg
except Exception as e:
console.print(f"[red]Error loading config: {e}[/red]")
raise typer.Exit(1)
@app.command("chat")
def chat(
ctx: typer.Context,
agent: Annotated[
str | None,
typer.Option(
"--agent",
"-a",
help="Agent ID to use (overrides default_agent from config)",
),
] = None,
) -> None:
"""Start interactive chat session."""
chat_command(ctx, agent_id=agent)
@app.command("server")
def server(ctx: typer.Context) -> None:
"""Start the 24/7 server for cron and messagebus execution."""
server_command(ctx)
if __name__ == "__main__":
app()
10-websocket/src/mybot/cli/server.py
+25
View File
@@ -0,0 +1,25 @@
"""Server CLI command for worker-based architecture."""
import asyncio
import typer
from mybot.core.context import SharedContext
from mybot.server.server import Server
from mybot.utils.logging import setup_logging
def server_command(ctx: typer.Context) -> None:
"""Start the 24/7 server for cron and messagebus execution."""
config = ctx.obj.get("config")
setup_logging(config, console_output=True)
typer.echo("Starting mybot server...")
typer.echo("Press Ctrl+C to stop")
try:
context = SharedContext(config)
asyncio.run(Server(context).run())
except KeyboardInterrupt:
typer.echo("\nServer stopped")
10-websocket/src/mybot/core/__init__.py
+20
View File
@@ -0,0 +1,20 @@
"""Core agent functionality."""
from .agent import Agent, AgentSession
from .agent_loader import (
AgentLoader,
AgentDef,
)
from .context import SharedContext
from .history import HistoryMessage, HistorySession, HistoryStore
__all__ = [
"Agent",
"AgentSession",
"AgentDef",
"AgentLoader",
"SharedContext",
"HistoryStore",
"HistoryMessage",
"HistorySession",
]
10-websocket/src/mybot/core/agent.py
+232
View File
@@ -0,0 +1,232 @@
import uuid
import json
import asyncio
from dataclasses import dataclass, field
from datetime import datetime
from typing import TYPE_CHECKING
from mybot.core.context_guard import ContextGuard
from mybot.core.session_state import SessionState
from mybot.core.events import EventSource, CliEventSource
from mybot.provider.llm import LLMProvider
from mybot.tools.registry import ToolRegistry
from mybot.tools.skill_tool import create_skill_tool
from mybot.tools.websearch_tool import create_websearch_tool
from mybot.tools.webread_tool import create_webread_tool
from litellm.types.completion import (
ChatCompletionMessageParam as Message,
ChatCompletionMessageToolCallParam,
)
if TYPE_CHECKING:
from mybot.core.context import SharedContext
from mybot.core.agent_loader import AgentDef
from mybot.provider.llm import LLMToolCall
class Agent:
"""A configured agent that creates and manages conversation sessions."""
def __init__(self, agent_def: "AgentDef", context: "SharedContext") -> None:
self.agent_def = agent_def
self.context = context
self.llm = LLMProvider.from_config(agent_def.llm)
def _build_tools(self) -> ToolRegistry:
"""Build a ToolRegistry with tools appropriate for the session."""
registry = ToolRegistry.with_builtins()
# Register skill tool if allowed
if self.agent_def.allow_skills:
skill_tool = create_skill_tool(self.context.skill_loader)
if skill_tool:
registry.register(skill_tool)
websearch_tool = create_websearch_tool(self.context)
if websearch_tool:
registry.register(websearch_tool)
webread_tool = create_webread_tool(self.context)
if webread_tool:
registry.register(webread_tool)
return registry
def _get_token_threshold(self) -> int:
"""Get token threshold based on model's context window."""
# Default to 80% of 200k context
return 160000
def new_session(
self,
source: EventSource | None = None,
session_id: str | None = None,
) -> "AgentSession":
"""Create a new conversation session."""
source = source or CliEventSource()
session_id = session_id or str(uuid.uuid4())
tools = self._build_tools()
# Create context guard for this session
context_guard = ContextGuard(
shared_context=self.context,
token_threshold=self._get_token_threshold(),
)
state = SessionState(
session_id=session_id,
agent=self,
messages=[],
source=source,
shared_context=self.context,
)
session = AgentSession(
agent=self,
state=state,
context_guard=context_guard,
tools=tools,
)
self.context.history_store.create_session(
self.agent_def.id, session_id, source
)
return session
def resume_session(self, session_id: str) -> "AgentSession":
"""Load an existing conversation session."""
session_query = [
session
for session in self.context.history_store.list_sessions()
if session.id == session_id
]
if not session_query:
raise ValueError(f"Session not found: {session_id}")
session_info = session_query[0]
source = session_info.get_source()
# Get all messages (no max_history limit)
history_messages = self.context.history_store.get_messages(session_id)
# Convert HistoryMessage to litellm Message format
messages: list[Message] = [msg.to_message() for msg in history_messages]
# Build tools for resumed session
tools = self._build_tools()
# Create context guard
context_guard = ContextGuard(
shared_context=self.context,
token_threshold=self._get_token_threshold(),
)
# Create SessionState with loaded messages
state = SessionState(
session_id=session_info.id,
agent=self,
messages=messages,
source=source,
shared_context=self.context,
)
return AgentSession(
agent=self,
state=state,
context_guard=context_guard,
tools=tools,
)
@dataclass
class AgentSession:
"""Chat orchestrator - operates on swappable SessionState."""
agent: Agent
state: SessionState
context_guard: ContextGuard
tools: ToolRegistry
started_at: datetime = field(default_factory=datetime.now)
@property
def session_id(self) -> str:
"""Delegate to state."""
return self.state.session_id
@property
def shared_context(self) -> "SharedContext":
"""Delegate to state."""
return self.state.shared_context
async def chat(self, message: str) -> str:
"""Send a message to the LLM and get a response."""
user_msg: Message = {"role": "user", "content": message}
self.state.add_message(user_msg)
tool_schemas = self.tools.get_tool_schemas()
while True:
messages = self.state.build_messages()
self.state = await self.context_guard.check_and_compact(self.state)
content, tool_calls = await self.agent.llm.chat(messages, tool_schemas)
tool_call_dicts: list[ChatCompletionMessageToolCallParam] = [
{
"id": tc.id,
"type": "function",
"function": {"name": tc.name, "arguments": tc.arguments},
}
for tc in tool_calls
]
assistant_msg: Message = {
"role": "assistant",
"content": content,
"tool_calls": tool_call_dicts,
}
self.state.add_message(assistant_msg)
if not tool_calls:
break
await self._handle_tool_calls(tool_calls)
continue
return content
async def _handle_tool_calls(
self,
tool_calls: list["LLMToolCall"],
) -> None:
"""Handle tool calls from the LLM response."""
tool_call_results = await asyncio.gather(
*[self._execute_tool_call(tool_call) for tool_call in tool_calls]
)
for tool_call, result in zip(tool_calls, tool_call_results):
tool_msg: Message = {
"role": "tool",
"content": result,
"tool_call_id": tool_call.id,
}
self.state.add_message(tool_msg)
async def _execute_tool_call(
self,
tool_call: "LLMToolCall",
) -> str:
"""Execute a single tool call."""
# Extract key arguments
try:
args = json.loads(tool_call.arguments)
except json.JSONDecodeError:
args = {}
try:
result = await self.tools.execute_tool(tool_call.name, session=self, **args)
except Exception as e:
result = f"Error executing tool: {e}"
return result
10-websocket/src/mybot/core/agent_loader.py
+77
View File
@@ -0,0 +1,77 @@
"""Agent definition loader."""
from typing import Any
from pydantic import BaseModel, ValidationError
from mybot.utils.config import Config, LLMConfig
from mybot.utils.def_loader import (
DefNotFoundError,
InvalidDefError,
parse_definition,
)
class AgentDef(BaseModel):
"""Loaded agent definition with merged settings."""
id: str
name: str
description: str = ""
agent_md: str
llm: LLMConfig
allow_skills: bool = False
class AgentLoader:
"""Loads agent definitions from AGENT.md files."""
@staticmethod
def from_config(config: Config) -> "AgentLoader":
return AgentLoader(config)
def __init__(self, config: Config):
"""Initialize AgentLoader."""
self.config = config
def load(self, agent_id: str) -> AgentDef:
"""Load agent by ID."""
agent_file = self.config.agents_path / agent_id / "AGENT.md"
if not agent_file.exists():
raise DefNotFoundError("agent", agent_id)
try:
content = agent_file.read_text()
agent_def = parse_definition(content, agent_id, self._parse_agent_def)
except InvalidDefError:
raise
except Exception as e:
raise InvalidDefError("agent", agent_id, str(e))
return agent_def
def _parse_agent_def(
self, def_id: str, frontmatter: dict[str, Any], body: str
) -> AgentDef:
"""Parse agent definition from frontmatter (callback for parse_definition)."""
llm_overrides = frontmatter.get("llm")
merged_llm = self._merge_llm_config(llm_overrides)
try:
return AgentDef(
id=def_id,
name=frontmatter["name"], # type: ignore[misc]
description=frontmatter.get("description", ""),
agent_md=body.strip(),
llm=merged_llm,
allow_skills=frontmatter.get("allow_skills", False),
)
except ValidationError as e:
raise InvalidDefError("agent", def_id, str(e))
def _merge_llm_config(self, agent_llm: dict[str, Any] | None) -> LLMConfig:
"""Deep merge agent's llm config with global defaults."""
base = self.config.llm.model_dump()
if agent_llm:
base = {**base, **agent_llm}
return LLMConfig(**base)
10-websocket/src/mybot/core/commands/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Slash commands module."""
from mybot.core.commands.base import Command
from mybot.core.commands.registry import CommandRegistry
__all__ = ["Command", "CommandRegistry"]
10-websocket/src/mybot/core/commands/base.py
+20
View File
@@ -0,0 +1,20 @@
"""Base classes for slash commands."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
class Command(ABC):
"""Base class for slash commands."""
name: str
aliases: list[str] = []
description: str = ""
@abstractmethod
async def execute(self, args: str, session: "AgentSession") -> str:
"""Execute the command and return response string."""
pass
10-websocket/src/mybot/core/commands/handlers.py
+126
View File
@@ -0,0 +1,126 @@
"""Built-in slash command handlers."""
from typing import TYPE_CHECKING
from mybot.core.commands.base import Command
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
class SessionCommand(Command):
"""Show current session details."""
name = "session"
description = "Show current session details"
async def execute(self, args: str, session: "AgentSession") -> str:
info = session.agent.history_store.history_store.get_session_info(session.session_id)
# Handle case where session not found in index
created_str = info.created_at if info else "Unknown"
lines = [
f"**Session ID:** `{session.session_id}`",
f"**Agent:** {session.agent.agent_def.name} (`{session.agent.agent_def.id}`)",
f"**Created:** {created_str}",
f"**Messages:** {len(session.state.messages)}",
]
return "\n".join(lines)
class HelpCommand(Command):
"""Show available commands."""
name = "help"
aliases = ["?"]
description = "Show available commands"
async def execute(self, args: str, session: "AgentSession") -> str:
lines = ["**Available Commands:**"]
for cmd in session.command_registry.list_commands():
names = [f"/{cmd.name}"] + [f"/{a}" for a in cmd.aliases]
lines.append(f"{', '.join(names)} - {cmd.description}")
return "\n".join(lines)
class CompactCommand(Command):
"""Trigger manual context compaction."""
name = "compact"
description = "Compact conversation context manually"
async def execute(self, args: str, session: "AgentSession") -> str:
# Force compaction regardless of threshold
await session.context_guard._compact_messages(session.state)
msg_count = len(session.state.messages)
return f"✓ Context compacted. {msg_count} messages retained."
class ContextCommand(Command):
"""Show session context information."""
name = "context"
description = "Show session context information"
async def execute(self, args: str, session: "AgentSession") -> str:
token_count = session.context_guard.estimate_tokens(session.state)
threshold = session.context_guard.token_threshold
usage_pct = (token_count / threshold) * 100 if threshold > 0 else 0
lines = [
f"**Messages:** {len(session.state.messages)}",
f"**Tokens:** {token_count:,} ({usage_pct:.1f}% of {threshold:,} threshold)",
]
return "\n".join(lines)
class ClearCommand(Command):
"""Clear conversation and start fresh."""
name = "clear"
description = "Clear conversation and start fresh"
async def execute(self, args: str, session: "AgentSession") -> str:
source_str = str(session.source)
if source_str in self.shared_context.config.sources:
del self.shared_context.config.sources[source_str]
self.shared_context.config.set_runtime(
"sources", self.shared_context.config.sources
)
return "✓ Conversation cleared. Next message starts fresh."
class SkillsCommand(Command):
"""List all skills or show skill details."""
name = "skills"
description = "List all skills or show skill details"
async def execute(self, args: str, session: "AgentSession") -> str:
if not args:
skills = session.agent.skill_loader.discover_skills()
if not skills:
return "No skills configured."
lines = ["**Skills:**"]
for skill in skills:
lines.append(f"- `{skill.id}`: {skill.description}")
return "\n".join(lines)
# Show specific skill details
skill_id = args.strip()
try:
skill = session.agent.skill_loader.load_skill(skill_id)
except FileNotFoundError:
return f"✗ Skill `{skill_id}` not found."
lines = [
f"**Skill:** `{skill.id}`",
f"**Name:** {skill.name}",
f"**Description:** {skill.description}",
f"\n---\n\n**SKILL.md:**\n```\n{skill.content}\n```",
]
return "\n".join(lines)
10-websocket/src/mybot/core/commands/registry.py
+76
View File
@@ -0,0 +1,76 @@
"""Command registry for managing slash commands."""
from typing import TYPE_CHECKING
from mybot.core.commands.base import Command
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
class CommandRegistry:
"""Registry for slash commands."""
def __init__(self) -> None:
self._commands: dict[str, Command] = {}
def register(self, cmd: Command) -> None:
"""Register a command and its aliases."""
self._commands[cmd.name] = cmd
for alias in cmd.aliases:
self._commands[alias] = cmd
def list_commands(self) -> list[Command]:
"""Return list of unique commands (deduplicated by name)."""
seen = set()
commands = []
for cmd in self._commands.values():
if cmd.name not in seen:
seen.add(cmd.name)
commands.append(cmd)
return commands
def resolve(self, input: str) -> tuple[Command, str] | None:
"""Parse input and return (command, args) if it matches."""
if not input.startswith("/"):
return None
parts = input[1:].split(None, 1)
if not parts:
return None
cmd_name = parts[0].lower()
args = parts[1] if len(parts) > 1 else ""
cmd = self._commands.get(cmd_name)
if cmd:
return (cmd, args)
return None
async def dispatch(self, input: str, session: "AgentSession") -> str | None:
"""Parse and execute a slash command."""
resolved = self.resolve(input)
if not resolved:
return None
cmd, args = resolved
return await cmd.execute(args, session)
@classmethod
def with_builtins(cls) -> "CommandRegistry":
"""Create registry with built-in commands registered."""
from mybot.core.commands.handlers import (
HelpCommand,
SkillsCommand,
SessionCommand,
CompactCommand,
ContextCommand,
)
registry = cls()
registry.register(HelpCommand())
registry.register(SkillsCommand())
registry.register(SessionCommand())
registry.register(CompactCommand())
registry.register(ContextCommand())
return registry
10-websocket/src/mybot/core/context.py
+42
View File
@@ -0,0 +1,42 @@
from typing import Any, TYPE_CHECKING
from mybot.core.agent_loader import AgentLoader
from mybot.core.commands.registry import CommandRegistry
from mybot.core.history import HistoryStore
from mybot.core.skill_loader import SkillLoader
from mybot.core.eventbus import EventBus
from mybot.channel.base import Channel
from mybot.utils.config import Config
if TYPE_CHECKING:
from mybot.server.websocket_worker import WebSocketWorker
class SharedContext:
"""Global shared state for the application."""
config: Config
history_store: HistoryStore
agent_loader: AgentLoader
skill_loader: SkillLoader
command_registry: CommandRegistry
channels: list[Channel[Any]]
eventbus: EventBus
websocket_worker: "WebSocketWorker | None"
def __init__(
self, config: Config, channels: list[Channel[Any]] | None = None
) -> None:
self.config = config
self.history_store = HistoryStore.from_config(config)
self.agent_loader = AgentLoader.from_config(config)
self.skill_loader = SkillLoader.from_config(config)
self.command_registry = CommandRegistry.with_builtins()
if channels is not None:
self.channels = channels
else:
self.channels = Channel.from_config(config)
self.eventbus = EventBus(self)
self.websocket_worker = None
10-websocket/src/mybot/core/context_guard.py
+163
View File
@@ -0,0 +1,163 @@
"""Context guard for proactive context window management."""
from dataclasses import dataclass
from typing import TYPE_CHECKING, cast
from litellm import token_counter
from litellm.types.completion import (
ChatCompletionMessageParam as Message,
ChatCompletionAssistantMessageParam,
ChatCompletionToolMessageParam,
)
from mybot.core.session_state import SessionState
if TYPE_CHECKING:
from mybot.core.context import SharedContext
from mybot.core.session_state import SessionState
# Default max size for tool result content before truncation
MAX_TOOL_RESULT_CHARS = 10000
@dataclass
class ContextGuard:
"""Manages context window size with proactive compaction."""
shared_context: "SharedContext"
token_threshold: int = 160000 # 80% of 200k context
max_tool_result_chars: int = MAX_TOOL_RESULT_CHARS
def estimate_tokens(self, state: "SessionState") -> int:
"""Estimate token count for session state."""
if not state.messages:
return 0
return token_counter(
model=state.agent.agent_def.llm.model, messages=state.build_messages()
)
async def check_and_compact(
self,
state: "SessionState",
) -> "SessionState":
"""Check token count, compact and roll session if needed."""
token_count = self.estimate_tokens(state)
if token_count < self.token_threshold:
return state
state.messages = self._truncate_large_tool_results(state.messages)
token_count = self.estimate_tokens(state)
if token_count < self.token_threshold:
return state
return await self.compact_and_roll(state)
def _compress_message_count(self, state: "SessionState") -> int:
keep_count = max(4, int(len(state.messages) * 0.2))
compress_count = max(2, int(len(state.messages) * 0.5))
return min(compress_count, len(state.messages) - keep_count)
def _truncate_large_tool_results(self, messages: list[Message]) -> list[Message]:
"""Truncate oversized tool results to reduce context size."""
result: list[Message] = []
for msg in messages:
if msg.get("role") == "tool":
content = msg.get("content", "")
if (
isinstance(content, str)
and len(content) > self.max_tool_result_chars
):
original_size = len(content)
truncated = content[: self.max_tool_result_chars]
truncated_content = (
f"{truncated}\n\n"
f"[Truncated - original size: {original_size} chars]"
)
msg = cast(
ChatCompletionToolMessageParam,
{**msg, "content": truncated_content},
)
result.append(msg)
return result
def _serialize_messages_for_summary(self, messages: list[Message]) -> str:
"""Serialize messages to plain text for summarization."""
lines = []
for msg in messages:
role = msg.get("role", "unknown")
content = msg.get("content", "")
# Handle tool calls in assistant messages
if role == "assistant" and msg.get("tool_calls"):
tool_names = [
tc.get("function", {}).get("name", "unknown")
for tc in (cast(ChatCompletionAssistantMessageParam, msg)).get(
"tool_calls", []
)
]
lines.append(
f"ASSISTANT: [used tools: {', '.join(tool_names)}] {content}"
)
else:
lines.append(f"{role.upper()}: {content}")
return "\n".join(lines)
async def compact_and_roll(
self,
state: "SessionState",
) -> "SessionState":
"""Compact history, roll to new session, return new messages."""
new_session = state.agent.new_session(state.source)
self._clear_source_session_cache(str(state.source))
compacted_history = await self._build_compacted_messages(state)
for message in compacted_history:
new_session.state.add_message(message)
return new_session.state
async def _build_compacted_messages(
self,
state: "SessionState",
) -> list[Message]:
"""Generate summary of older messages using agent's LLM."""
compress_count = self._compress_message_count(state)
old_messages = state.messages[:compress_count]
old_text = self._serialize_messages_for_summary(old_messages)
summary_prompt = f"""Summarize the conversation so far. Keep it factual and concise. Focus on key decisions, facts, and user preferences discovered:
{old_text}"""
response, _ = await state.agent.llm.chat(
[{"role": "user", "content": summary_prompt}],
[], # No tools needed
)
messages: list[Message] = []
messages.append(
{
"role": "user",
"content": f"[Previous conversation summary]\n{response}",
}
)
messages.append(
{
"role": "assistant",
"content": "Understood, I have the context.",
}
)
messages.extend(state.messages[compress_count:])
return messages
def _clear_source_session_cache(self, source_str: str) -> None:
if source_str in self.shared_context.config.sources:
del self.shared_context.config.sources[source_str]
self.shared_context.config.set_runtime(
"sources", self.shared_context.config.sources
)
10-websocket/src/mybot/core/eventbus.py
+145
View File
@@ -0,0 +1,145 @@
"""Central event bus for pub/sub event distribution."""
import asyncio
import json
import logging
import os
from collections import defaultdict
from typing import Awaitable, Callable, TypeVar, TYPE_CHECKING
from mybot.server.worker import Worker
from .events import (
Event,
OutboundEvent,
deserialize_event,
)
if TYPE_CHECKING:
from mybot.core.context import SharedContext
logger = logging.getLogger(__name__)
E = TypeVar("E", bound=Event)
Handler = Callable[[Event], Awaitable[None]]
class EventBus(Worker):
"""Central event bus with subscription support and async dispatch."""
def __init__(self, context: "SharedContext"):
super().__init__(context)
self.context = context
self._subscribers: dict[type[Event], list[Handler]] = defaultdict(list)
self._queue: asyncio.Queue[Event] = asyncio.Queue()
self.pending_dir = context.config.event_path / "pending"
self.pending_dir.mkdir(parents=True, exist_ok=True)
def subscribe(
self, event_class: type[E], handler: Callable[[E], Awaitable[None]]
) -> None:
"""Subscribe a handler to an event class."""
self._subscribers[event_class].append(handler)
logger.debug(f"Subscribed handler to {event_class.__name__} events")
def unsubscribe(self, handler: Handler) -> None:
"""Remove a handler from all subscriptions."""
for event_class in self._subscribers:
if handler in self._subscribers[event_class]:
self._subscribers[event_class].remove(handler)
logger.debug(f"Unsubscribed handler from {event_class.__name__} events")
async def publish(self, event: Event) -> None:
"""Publish an event to the internal queue (non-blocking)."""
await self._queue.put(event)
logger.debug(f"Queued {event.__class__.__name__} event from {event.source}")
async def run(self) -> None:
"""Process events from queue, starting with recovery."""
logger.info("EventBus started")
# Run recovery first
await self._recover()
try:
while True:
event = await self._queue.get()
try:
await self._dispatch(event)
except Exception as e:
logger.error(f"Error dispatching event: {e}")
finally:
self._queue.task_done()
except asyncio.CancelledError:
logger.info("EventBus stopping...")
raise
async def _dispatch(self, event: Event) -> None:
"""Persist if OUTBOUND, then notify subscribers."""
await self._persist_outbound(event)
await self._notify_subscribers(event)
logger.debug(f"Dispatched {event.__class__.__name__} event from {event.source}")
async def _notify_subscribers(self, event: Event) -> None:
"""Notify all subscribers of an event (waits for all handlers to complete)."""
handlers = self._subscribers.get(type(event), [])
if not handlers:
return
tasks = [handler(event) for handler in handlers]
results = await asyncio.gather(*tasks, return_exceptions=True)
for result in results:
if isinstance(result, Exception):
logger.error(f"Error in event handler: {result}")
async def _persist_outbound(self, event: Event) -> None:
"""Persist event to disk (only OUTBOUND events)."""
if not isinstance(event, OutboundEvent):
return
filename = f"{event.timestamp}_{event.session_id}.json"
final_path = self.pending_dir / filename
tmp_path = self.pending_dir / f".tmp.{os.getpid()}.{filename}"
data = json.dumps(event.to_dict(), ensure_ascii=False)
# Atomic write: tmp + fsync + rename
with open(tmp_path, "w", encoding="utf-8") as f:
f.write(data)
f.flush()
os.fsync(f.fileno())
os.replace(str(tmp_path), str(final_path))
logger.debug(f"Persisted event to {final_path}")
async def _recover(self) -> int:
"""Recover pending events from previous crash. Returns count recovered."""
pending_files = list(self.pending_dir.glob("*.json"))
if not pending_files:
return 0
logger.info(f"Recovering {len(pending_files)} pending events")
count = 0
for file_path in pending_files:
try:
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
# Use deserialize_event to handle typed events
event = deserialize_event(data)
await self._notify_subscribers(event)
count += 1
logger.debug(f"Recovered event from {file_path.name}")
except Exception as e:
logger.error(f"Failed to recover {file_path}: {e}")
logger.info(f"Recovered {count} events")
return count
def ack(self, event: Event) -> None:
"""Acknowledge successful delivery, delete persisted event."""
filename = f"{event.timestamp}_{event.session_id}.json"
final_path = self.pending_dir / filename
if final_path.exists():
final_path.unlink()
logger.debug(f"Acked and deleted {filename}")
10-websocket/src/mybot/core/events.py
+178
View File
@@ -0,0 +1,178 @@
"""Event types and data classes for the event bus."""
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, ClassVar
class EventSource(ABC):
"""Abstract base for all event sources."""
_registry: ClassVar[dict[str, type["EventSource"]]] = {}
_namespace: ClassVar[str] = ""
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
if cls._namespace:
cls._registry[cls._namespace] = cls
@property
def is_platform(self) -> bool:
return self._namespace.startswith("platform-")
@property
def is_agent(self) -> bool:
return self._namespace == "agent"
@property
def is_cron(self) -> bool:
return self._namespace == "cron"
@property
def platform_name(self) -> str | None:
if not self.is_platform:
return None
return self._namespace.split("-", 1)[1]
@classmethod
def from_string(cls, s: str) -> "EventSource":
"""Parse string to EventSource using namespace registry."""
namespace = s.split(":")[0]
source_cls = cls._registry.get(namespace)
if not source_cls:
raise ValueError(f"Unknown source namespace: {namespace}")
return source_cls.from_string(s)
@abstractmethod
def __str__(self) -> str: ...
@dataclass
class AgentEventSource(EventSource):
"""Source for agent-generated events."""
_namespace = "agent"
agent_id: str
def __str__(self) -> str:
return f"agent:{self.agent_id}"
@classmethod
def from_string(cls, s: str) -> "AgentEventSource":
_, agent_id = s.split(":", 1)
return cls(agent_id=agent_id)
@dataclass
class CliEventSource(EventSource):
"""Source for CLI-originated events."""
_namespace = "platform-cli"
def __str__(self) -> str:
return "platform-cli:cli-user"
@classmethod
def from_string(cls, s: str) -> "CliEventSource":
return cls()
@property
def platform_name(self) -> str:
return "cli"
@dataclass
class WebSocketEventSource(EventSource):
"""Event from WebSocket client."""
_namespace = "platform-ws"
user_id: str
@classmethod
def from_string(cls, s: str) -> "WebSocketEventSource":
"""Parse source string into WebSocketEventSource."""
parts = s.split(":", 1)
if len(parts) != 2 or parts[0] != cls._namespace or not parts[1]:
raise ValueError(f"Invalid WebSocketEventSource: {s}")
return cls(user_id=parts[1])
def __str__(self) -> str:
"""Convert to source string format."""
return f"{self._namespace}:{self.user_id}"
@property
def is_platform(self) -> bool:
"""WebSocket sources are platform sources."""
return True
@dataclass
class Event:
"""Base class for all typed events."""
session_id: str
source: EventSource # Changed from str to typed EventSource
content: str
timestamp: float = field(default_factory=time.time)
def to_dict(self) -> dict[str, Any]:
"""Serialize event to dictionary, including type."""
result: dict[str, Any] = {"type": self.__class__.__name__}
for field_name in self.__dataclass_fields__:
value = getattr(self, field_name)
if field_name == "source":
result[field_name] = str(value)
else:
result[field_name] = value
return result
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "Event":
"""Deserialize event from dictionary."""
kwargs = {}
for k, v in data.items():
if k == "type":
continue
if k == "source":
kwargs[k] = EventSource.from_string(v)
elif k in cls.__dataclass_fields__:
kwargs[k] = v
return cls(**kwargs)
@dataclass
class InboundEvent(Event):
"""Event for external work entering the system (platforms, cron, retry)."""
retry_count: int = 0
@dataclass
class OutboundEvent(Event):
"""Event for agent responses to deliver to platforms."""
error: str | None = None
# Registry mapping event class names to event classes
_EVENT_CLASSES: dict[str, type[Event]] = {
"InboundEvent": InboundEvent,
"OutboundEvent": OutboundEvent,
}
def serialize_event(event: Event) -> dict[str, Any]:
"""Serialize any event type to dict."""
return event.to_dict()
def deserialize_event(data: dict[str, Any]) -> Event:
"""Deserialize dict to appropriate event type."""
event_type: str = data.get("type", "")
event_class = _EVENT_CLASSES.get(event_type)
if event_class is None:
raise ValueError(f"Unknown event type: {event_type}")
return event_class.from_dict(data)
10-websocket/src/mybot/core/history.py
+230
View File
@@ -0,0 +1,230 @@
"""JSONL file-based conversation history backend."""
from datetime import datetime
from pathlib import Path
from typing import Any, Literal, TYPE_CHECKING
from pydantic import BaseModel, Field, field_validator
from mybot.core.events import EventSource
from litellm.types.completion import ChatCompletionMessageParam as Message
if TYPE_CHECKING:
from mybot.utils.config import Config
def _now_iso() -> str:
"""Return current datetime as ISO format string."""
return datetime.now().isoformat()
class HistorySession(BaseModel):
"""Session metadata - stored in index.jsonl."""
id: str
agent_id: str
source: str # Serialized EventSource (e.g., "platform-telegram:123:456")
title: str | None = None
message_count: int = 0
created_at: str
updated_at: str
@field_validator("source", mode="before")
@classmethod
def parse_source(cls, v: Any) -> str:
if hasattr(v, "__str__"):
return str(v)
return v
def get_source(self) -> EventSource:
"""Get the session's EventSource."""
return EventSource.from_string(self.source)
class HistoryMessage(BaseModel):
"""Single message - stored in session.jsonl."""
timestamp: str = Field(default_factory=_now_iso)
role: Literal["user", "assistant", "system", "tool"]
content: str
tool_calls: list[dict[str, Any]] | None = None
tool_call_id: str | None = None
@classmethod
def from_message(cls, message: Message) -> "HistoryMessage":
"""Create HistoryMessage from litellm Message format."""
tool_calls = None
if message.get("tool_calls"):
tool_calls = [
{
"id": tc.get("id"),
"type": tc.get("type", "function"),
"function": tc.get("function", {}),
}
for tc in message["tool_calls"]
]
tool_call_id = message.get("tool_call_id")
return cls(
role=message["role"],
content=str(message.get("content", "")),
tool_calls=tool_calls,
tool_call_id=tool_call_id,
)
def to_message(self) -> Message:
"""Convert HistoryMessage to litellm Message format."""
base: dict[str, Any] = {
"role": self.role,
"content": self.content,
}
if self.role == "assistant" and self.tool_calls:
return {
"role": "assistant",
"content": self.content,
"tool_calls": self.tool_calls,
}
if self.role == "tool" and self.tool_call_id:
base["tool_call_id"] = self.tool_call_id
return base
return base
class HistoryStore:
"""JSONL file-based history storage."""
@staticmethod
def from_config(config: "Config") -> "HistoryStore":
return HistoryStore(config.history_path)
def __init__(self, base_path: Path):
self.base_path = Path(base_path)
self.sessions_path = self.base_path / "sessions"
self.index_path = self.base_path / "index.jsonl"
self.base_path.mkdir(parents=True, exist_ok=True)
self.sessions_path.mkdir(parents=True, exist_ok=True)
def _session_path(self, session_id: str) -> Path:
"""Get the file path for a session."""
return self.sessions_path / f"{session_id}.jsonl"
def _read_index(self) -> list[HistorySession]:
"""Read all session entries from index.jsonl."""
if not self.index_path.exists():
return []
sessions = []
with open(self.index_path) as f:
for line in f:
line = line.strip()
if line:
try:
sessions.append(HistorySession.model_validate_json(line))
except Exception:
continue
return sessions
def _write_index(self, sessions: list[HistorySession]) -> None:
"""Write all session entries to index.jsonl."""
with open(self.index_path, "w") as f:
for session in sessions:
f.write(session.model_dump_json() + "\n")
def _find_session_index(
self, sessions: list[HistorySession], session_id: str
) -> int:
"""Find the index of a session in the list."""
for i, s in enumerate(sessions):
if s.id == session_id:
return i
return -1
def create_session(
self, agent_id: str, session_id: str, source: "EventSource",
) -> dict[str, Any]:
"""Create a new conversation session."""
now = _now_iso()
session = HistorySession(
id=session_id,
agent_id=agent_id,
source=source,
title=None,
message_count=0,
created_at=now,
updated_at=now,
)
# Append to index
with open(self.index_path, "a") as f:
f.write(session.model_dump_json() + "\n")
# Create session file
self._session_path(session_id).touch()
return session.model_dump()
def save_message(self, session_id: str, message: HistoryMessage) -> None:
"""Save a message to history."""
sessions = self._read_index()
idx = self._find_session_index(sessions, session_id)
if idx < 0:
raise ValueError(f"Session not found: {session_id}")
session = sessions[idx]
# Append message to session file
session_file = self._session_path(session_id)
with open(session_file, "a") as f:
f.write(message.model_dump_json() + "\n")
# Update index
session.message_count += 1
session.updated_at = _now_iso()
# Auto-generate title from first user message
if session.title is None and message.role == "user":
title = message.content[:50]
if len(message.content) > 50:
title += "..."
session.title = title
sessions.sort(key=lambda s: s.updated_at, reverse=True)
self._write_index(sessions)
def list_sessions(self) -> list[HistorySession]:
"""List all sessions, most recently updated first."""
sessions = self._read_index()
sessions.sort(key=lambda s: s.updated_at, reverse=True)
return sessions
def get_messages(self, session_id: str) -> list[HistoryMessage]:
"""Get all messages for a session."""
session_file = self._session_path(session_id)
if not session_file.exists():
return []
messages: list[HistoryMessage] = []
with open(session_file) as f:
for line in f:
line = line.strip()
if line:
try:
messages.append(HistoryMessage.model_validate_json(line))
except Exception:
continue
return messages
def get_session_info(self, session_id: str) -> HistorySession | None:
"""Get session metadata without loading messages."""
sessions = self._read_index()
for session in sessions:
if session.id == session_id:
return session
return None
10-websocket/src/mybot/core/session_state.py
+37
View File
@@ -0,0 +1,37 @@
"""Session state container with persistence helpers."""
from dataclasses import dataclass
from typing import TYPE_CHECKING
from litellm.types.completion import ChatCompletionMessageParam as Message
from mybot.core.history import HistoryMessage
if TYPE_CHECKING:
from mybot.core.agent import Agent
from mybot.core.context import SharedContext
from mybot.core.events import EventSource
@dataclass
class SessionState:
"""Pure conversation state + persistence."""
session_id: str
agent: "Agent"
messages: list[Message]
source: "EventSource"
shared_context: "SharedContext"
def add_message(self, message: Message) -> None:
"""Add message to in-memory list + persist."""
self.messages.append(message)
history_msg = HistoryMessage.from_message(message)
self.shared_context.history_store.save_message(self.session_id, history_msg)
def build_messages(self) -> list[Message]:
"""Build messages list with system prompt."""
system_prompt = self.agent.agent_def.agent_md
messages: list[Message] = [{"role": "system", "content": system_prompt}]
messages.extend(self.messages)
return messages
10-websocket/src/mybot/core/skill_loader.py
+69
View File
@@ -0,0 +1,69 @@
"""Skill loader for discovering and loading skills."""
import logging
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, ConfigDict, ValidationError
from mybot.utils.def_loader import DefNotFoundError, discover_definitions
if TYPE_CHECKING:
from mybot.utils.config import Config
logger = logging.getLogger(__name__)
class SkillDef(BaseModel):
"""Loaded skill definition."""
model_config = ConfigDict(extra="forbid")
id: str
name: str
description: str
content: str
class SkillLoader:
"""Load and manage skill definitions from filesystem."""
@staticmethod
def from_config(config: "Config") -> "SkillLoader":
"""Create SkillLoader from config."""
return SkillLoader(config)
def __init__(self, config: "Config"):
self.config = config
def discover_skills(self) -> list[SkillDef]:
"""Scan skills directory and return list of valid SkillDef."""
return discover_definitions(
self.config.skills_path, "SKILL.md", self._parse_skill_def
)
def _parse_skill_def(
self, def_id: str, frontmatter: dict[str, Any], body: str
) -> SkillDef | None:
"""Parse skill definition from frontmatter (callback for discover_definitions)."""
try:
return SkillDef(
id=def_id,
name=frontmatter["name"], # type: ignore[misc]
description=frontmatter["description"], # type: ignore[misc]
content=body.strip(),
)
except ValidationError as e:
logger.warning(f"Invalid skill '{def_id}': {e}")
return None
except KeyError as e:
logger.warning(f"Missing required field in skill '{def_id}': {e}")
return None
def load_skill(self, skill_id: str) -> SkillDef:
"""Load full skill definition by ID."""
skills = self.discover_skills()
for skill in skills:
if skill.id == skill_id:
return skill
raise DefNotFoundError("skill", skill_id)
10-websocket/src/mybot/provider/__init__.py
View File
10-websocket/src/mybot/provider/llm/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""LLM provider abstraction."""
from .base import LLMProvider, LLMToolCall
__all__ = ["LLMProvider", "LLMToolCall"]
10-websocket/src/mybot/provider/llm/base.py
+85
View File
@@ -0,0 +1,85 @@
"""Base LLM provider abstraction."""
from dataclasses import dataclass
from typing import Any, Optional, cast
from litellm import acompletion, Choices
from litellm.types.completion import ChatCompletionMessageParam as Message
from mybot.utils.config import LLMConfig
@dataclass
class LLMToolCall:
"""A tool/function call from the LLM."""
id: str
name: str
arguments: str # JSON string
class LLMProvider:
"""LLM provider using litellm for multi-provider support."""
def __init__(
self,
model: str,
api_key: str,
api_base: Optional[str] = None,
temperature: float = 0.7,
max_tokens: int = 2048,
**kwargs: Any,
):
"""Initialize LLM provider."""
self.model = model
self.api_key = api_key
self.api_base = api_base
self.temperature = temperature
self.max_tokens = max_tokens
self._settings = kwargs
@classmethod
def from_config(cls, config: LLMConfig) -> "LLMProvider":
"""Create provider from LLMConfig."""
return cls(
model=config.model,
api_key=config.api_key,
api_base=config.api_base,
temperature=config.temperature,
max_tokens=config.max_tokens,
)
async def chat(
self,
messages: list[Message],
tools: Optional[list[dict[str, Any]]] = None,
**kwargs: Any,
) -> tuple[str, list[LLMToolCall]]:
"""Default implementation using litellm. Subclasses can override."""
request_kwargs: dict[str, Any] = {
"model": self.model,
"messages": messages,
"api_key": self.api_key,
}
if self.api_base:
request_kwargs["api_base"] = self.api_base
if tools:
request_kwargs["tools"] = tools
request_kwargs.update(kwargs)
response = await acompletion(**request_kwargs)
message = cast(Choices, response.choices[0]).message
return (
message.content or "",
[
LLMToolCall(
id=tc["id"],
name=tc["function"]["name"],
arguments=tc["function"]["arguments"],
)
for tc in (message.tool_calls or [])
],
)
10-websocket/src/mybot/provider/web_read/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Web read provider module."""
from .base import ReadResult, WebReadProvider
__all__ = ["ReadResult", "WebReadProvider"]
10-websocket/src/mybot/provider/web_read/base.py
+41
View File
@@ -0,0 +1,41 @@
"""Base class for web read providers."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from mybot.utils.config import Config
class ReadResult(BaseModel):
"""Normalized result from reading a web page."""
url: str
title: str
content: str # Markdown content
error: str | None = None
class WebReadProvider(ABC):
"""Abstract base class for web page reading providers."""
@abstractmethod
async def read(self, url: str) -> ReadResult:
"""Read a web page and return normalized content."""
pass
@staticmethod
def from_config(config: "Config") -> "WebReadProvider":
"""Factory method to create provider from config."""
if config.webread is None:
raise ValueError("Webread not configured")
match config.webread.provider:
case "crawl4ai":
from .crawl4ai import Crawl4AIProvider
return Crawl4AIProvider()
case _:
raise ValueError(f"Unknown webread provider: {config.webread.provider}")
10-websocket/src/mybot/provider/web_read/crawl4ai.py
+36
View File
@@ -0,0 +1,36 @@
"""Crawl4AI provider for web page reading."""
from crawl4ai import AsyncWebCrawler
from .base import WebReadProvider, ReadResult
class Crawl4AIProvider(WebReadProvider):
"""Web read provider using Crawl4AI."""
def __init__(self):
"""Initialize Crawl4AI provider."""
pass
async def read(self, url: str) -> ReadResult:
"""Read a web page using Crawl4AI."""
try:
async with AsyncWebCrawler(verbose=False) as crawler:
result = await crawler.arun(url=url)
if not result.success:
raise Exception(result.error_message or "Failed to crawl page")
return ReadResult(
url=url,
title=(result.metadata.get("title", "") if result.metadata else ""),
content=result.markdown or "",
error=None,
)
except Exception as e:
return ReadResult(
url=url,
title="",
content="",
error=str(e),
)
10-websocket/src/mybot/provider/web_search/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Web search provider module."""
from .base import SearchResult, WebSearchProvider
__all__ = ["SearchResult", "WebSearchProvider"]
10-websocket/src/mybot/provider/web_search/base.py
+42
View File
@@ -0,0 +1,42 @@
"""Base class for web search providers."""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from mybot.utils.config import Config
class SearchResult(BaseModel):
"""Normalized search result from any provider."""
title: str
url: str
snippet: str
class WebSearchProvider(ABC):
"""Abstract base class for web search providers."""
@abstractmethod
async def search(self, query: str) -> list[SearchResult]:
"""Search the web and return normalized results."""
pass
@staticmethod
def from_config(config: "Config") -> "WebSearchProvider":
"""Factory method to create provider from config."""
if config.websearch is None:
raise ValueError("Websearch not configured")
match config.websearch.provider:
case "brave":
from .brave import BraveSearchProvider
return BraveSearchProvider(config)
case _:
raise ValueError(
f"Unknown websearch provider: {config.websearch.provider}"
)
10-websocket/src/mybot/provider/web_search/brave.py
+49
View File
@@ -0,0 +1,49 @@
"""Brave Search API provider."""
from typing import TYPE_CHECKING
import httpx
from .base import WebSearchProvider, SearchResult
if TYPE_CHECKING:
from mybot.utils.config import Config
class BraveSearchProvider(WebSearchProvider):
"""Web search provider using Brave Search API."""
BASE_URL = "https://api.search.brave.com/res/v1/web/search"
def __init__(self, config: "Config"):
"""Initialize Brave Search provider."""
self.api_key = config.websearch.api_key
async def search(self, query: str) -> list[SearchResult]:
"""Search the web using Brave Search API."""
async with httpx.AsyncClient() as client:
response = await client.get(
self.BASE_URL,
headers={
"Accept": "application/json",
"X-Subscription-Token": self.api_key,
},
params={
"q": query,
"count": 10,
},
timeout=30.0,
)
response.raise_for_status()
data = response.json()
results = []
for item in data.get("web", {}).get("results", []):
results.append(
SearchResult(
title=item.get("title", ""),
url=item.get("url", ""),
snippet=item.get("description", ""),
)
)
return results
10-websocket/src/mybot/server/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""Worker-based server architecture."""
from .worker import Worker, SubscriberWorker
from .delivery_worker import DeliveryWorker
from .agent_worker import AgentWorker
from .channel_worker import ChannelWorker
__all__ = [
"Worker",
"SubscriberWorker",
"DeliveryWorker",
"AgentWorker",
"ChannelWorker",
]
10-websocket/src/mybot/server/agent_worker.py
+113
View File
@@ -0,0 +1,113 @@
"""Agent worker for executing agent jobs."""
import asyncio
import logging
from dataclasses import replace
from .worker import SubscriberWorker
from mybot.core.agent import Agent
from mybot.core.events import (
AgentEventSource,
InboundEvent,
OutboundEvent,
)
from mybot.utils.def_loader import DefNotFoundError
# Maximum number of retry attempts for failed sessions
MAX_RETRIES = 3
logger = logging.getLogger(__name__)
class AgentWorker(SubscriberWorker):
"""Dispatches events to session executors."""
def __init__(self, context):
super().__init__(context)
# Auto-subscribe to events
self.context.eventbus.subscribe(InboundEvent, self.dispatch_event)
self.logger.info("AgentWorker subscribed to InboundEvent events")
async def dispatch_event(self, event: InboundEvent) -> None:
"""Create executor task for typed event."""
# Get agent_id from session (single source of truth)
session_info = self.context.history_store.get_session_info(event.session_id)
if not session_info:
logger.error(f"Session not found: {event.session_id}")
return
agent_id = session_info.agent_id
try:
agent_def = self.context.agent_loader.load(agent_id)
except DefNotFoundError as e:
logger.error(f"Agent not found: {agent_id}: {e}")
await self._emit_response(event, "", agent_def.id, str(e))
asyncio.create_task(self.exec_session(event, agent_def))
async def exec_session(self, event: InboundEvent, agent_def) -> None:
session_id = event.session_id
try:
agent = Agent(agent_def, self.context)
if session_id:
try:
session = agent.resume_session(session_id)
except ValueError:
logger.warning(f"Session {session_id} not found, creating new")
session = agent.new_session(session_id=session_id)
else:
session = agent.new_session()
session_id = session.session_id
# Check for slash command FIRST
if event.content.startswith("/"):
result = await self.context.command_registry.dispatch(
event.content, session
)
if result:
# Emit response and skip agent chat
await self._emit_response(event, result, agent_def.id)
logger.info(f"Command completed: {session_id}")
return
response = await session.chat(event.content)
logger.info(f"Session completed: {session_id}")
await self._emit_response(event, response, agent_def.id)
except Exception as e:
logger.error(f"Session failed: {e}")
if event.retry_count < MAX_RETRIES:
# Use dataclasses.replace() for retry logic
retry_event = replace(
event,
retry_count=event.retry_count + 1,
content=".", # Minimal message for retry
)
await self.context.eventbus.publish(retry_event)
else:
await self._emit_response(event, "", agent_def.id, str(e))
async def _emit_response(
self,
event: InboundEvent,
content: str,
agent_id: str,
error: str | None = None,
) -> None:
"""Emit response event with content."""
result_event = OutboundEvent(
session_id=event.session_id,
source=AgentEventSource(agent_id),
content=content,
error=str(error) if error else None,
)
await self.context.eventbus.publish(result_event)
10-websocket/src/mybot/server/app.py
+41
View File
@@ -0,0 +1,41 @@
"""FastAPI application with WebSocket support."""
from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from mybot.core.context import SharedContext
def create_app(context: SharedContext) -> FastAPI:
"""Create and configure the FastAPI application."""
app = FastAPI(
title="MyBot WebSocket Server",
description="WebSocket server for real-time agent communication",
version="0.1.0",
)
app.state.context = context
# Enable CORS for web clients
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# WebSocket endpoint
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
"""WebSocket endpoint for real-time event streaming and chat."""
await websocket.accept()
# Check if WebSocket worker is available
if context.websocket_worker is None:
await websocket.close(code=1013, reason="WebSocket not available")
return
# Hand off to worker
await context.websocket_worker.handle_connection(websocket)
return app
10-websocket/src/mybot/server/channel_worker.py
+97
View File
@@ -0,0 +1,97 @@
"""Channel worker for ingesting platform messages."""
import asyncio
import time
from typing import TYPE_CHECKING
from mybot.core.agent import Agent
from .worker import Worker
from mybot.core.events import EventSource, InboundEvent
from mybot.utils.config import SourceSessionConfig
if TYPE_CHECKING:
from mybot.core.context import SharedContext
class ChannelWorker(Worker):
"""Ingests messages from platforms, publishes INBOUND events to Channel."""
def __init__(self, context: "SharedContext"):
super().__init__(context)
self.channels = context.channels
self.channel_map = {channel.platform_name: channel for channel in self.channels}
async def run(self) -> None:
"""Start all channels and process incoming messages."""
self.logger.info(f"ChannelWorker started with {len(self.channels)} channel(es)")
channel_tasks = [
channel.run(self._create_callback(channel.platform_name))
for channel in self.channels
]
try:
await asyncio.gather(*channel_tasks)
except asyncio.CancelledError:
await asyncio.gather(*[channel.stop() for channel in self.channels])
raise
def _create_callback(self, platform: str):
"""Create callback for a specific platform."""
async def callback(message: str, source: EventSource) -> None:
try:
channel = self.channel_map[platform]
if not channel.is_allowed(source):
self.logger.debug(
f"Ignored non-whitelisted message from {platform}"
)
return
# Set default delivery source only on first non-CLI platform message
if source.is_platform and source.platform_name != "cli":
if not self.context.config.default_delivery_source:
source_str_value = str(source)
self.context.config.set_runtime(
"default_delivery_source", source_str_value
)
# Update in-memory value immediately for other workers
self.context.config.default_delivery_source = source_str_value
session_id = self._get_or_create_session_id(source)
# Publish INBOUND event with typed source
event = InboundEvent(
session_id=session_id,
source=source,
content=message,
timestamp=time.time(),
)
await self.context.eventbus.publish(event)
self.logger.debug(f"Published INBOUND event from {source}")
except Exception as e:
self.logger.error(f"Error processing message from {platform}: {e}")
return callback
def _get_or_create_session_id(self, source: EventSource) -> str:
"""Get or create session ID for a given source."""
source_str = str(source)
source_session = self.context.config.sources.get(source_str)
if source_session:
return source_session.session_id
agent_def = self.context.agent_loader.load(self.context.config.default_agent)
agent = Agent(agent_def, self.context)
session = agent.new_session(source)
# Cache the session
self.context.config.set_runtime(
f"sources.{source_str}", SourceSessionConfig(session_id=session.session_id)
)
return session.session_id
10-websocket/src/mybot/server/delivery_worker.py
+194
View File
@@ -0,0 +1,194 @@
"""Worker that delivers outbound messages to platforms."""
import asyncio
import logging
import random
from functools import lru_cache
from typing import TYPE_CHECKING, Any
from mybot.core.events import EventSource, OutboundEvent
from mybot.core.history import HistorySession
from .worker import SubscriberWorker
if TYPE_CHECKING:
from mybot.core.context import SharedContext
from mybot.channel.base import Channel
logger = logging.getLogger(__name__)
# Retry configuration
BACKOFF_MS = [5000, 25000, 120000, 600000] # 5s, 25s, 2min, 10min
MAX_RETRIES = 5
def compute_backoff_ms(retry_count: int) -> int:
"""Compute backoff time with jitter."""
if retry_count <= 0:
return 0
# Cap at last backoff value
idx = min(retry_count - 1, len(BACKOFF_MS) - 1)
base = BACKOFF_MS[idx]
# Add +/- 20% jitter
jitter = random.randint(-base // 5, base // 5)
return max(0, base + jitter)
# Platform message size limits
PLATFORM_LIMITS: dict[str, float] = {
"telegram": 4096,
"discord": 2000,
"cli": float("inf"), # no limit
}
def chunk_message(content: str, limit: int) -> list[str]:
"""Split message at paragraph boundaries, respecting limit."""
if len(content) <= limit:
return [content]
chunks = []
paragraphs = content.split("\n\n")
current = ""
for para in paragraphs:
# Try to add to current chunk
if current:
potential = current + "\n\n" + para
else:
potential = para
if len(potential) <= limit:
current = potential
else:
if current:
chunks.append(current)
# Handle paragraph that exceeds limit
if len(para) > limit:
# Hard split
for i in range(0, len(para), limit):
chunks.append(para[i : i + limit])
current = ""
else:
current = para
if current:
chunks.append(current)
return chunks
class DeliveryWorker(SubscriberWorker):
"""Worker that delivers outbound messages to platforms."""
def __init__(self, context: "SharedContext"):
super().__init__(context)
self.context.eventbus.subscribe(OutboundEvent, self.handle_event)
self.logger.info("DeliveryWorker subscribed to OUTBOUND events")
async def _deliver_with_retry(
self, chunks: list[str], source: "EventSource", channel: "Channel[Any]"
) -> bool:
"""Deliver all chunks with retry logic. Returns True on success."""
for attempt in range(1, MAX_RETRIES + 1):
try:
for chunk in chunks:
await channel.reply(chunk, source)
return True
except Exception as e:
if attempt < MAX_RETRIES:
backoff_ms = compute_backoff_ms(attempt)
self.logger.warning(
f"Delivery failed (attempt {attempt}/{MAX_RETRIES}), "
f"retrying in {backoff_ms}ms: {e}"
)
await asyncio.sleep(backoff_ms / 1000)
else:
self.logger.error(
f"Delivery failed after {MAX_RETRIES} attempts: {e}"
)
return False
return False
@lru_cache(maxsize=10)
def _get_session_source(self, session_id: str) -> HistorySession | None:
"""Get session info from HistoryStore (cached)."""
for session in self.context.history_store.list_sessions():
if session.id == session_id:
return session
return None
def _get_delivery_source(
self, session_info: HistorySession
) -> "EventSource | None":
"""Get the delivery source for a session."""
source = session_info.get_source()
# If source already has a platform, use it
if source.platform_name:
return source
# Try default delivery source for agent events
default_source_str = self.context.config.default_delivery_source
if default_source_str:
try:
source = EventSource.from_string(default_source_str)
if not source.platform_name:
self.logger.error(
f"default_delivery_source '{default_source_str}' is not a platform source"
)
return None
return source
except ValueError as e:
self.logger.error(f"Invalid default_delivery_source: {e}")
return None
else:
self.logger.warning(
f"No platform for session {session_info.id} and no default_delivery_source configured"
)
return None
async def handle_event(self, event: OutboundEvent) -> None:
"""Handle an outbound message event."""
try:
session_info = self._get_session_source(event.session_id)
if not session_info or not session_info.source:
self.logger.warning(
f"No source for session {event.session_id}, skipping delivery"
)
return
source = self._get_delivery_source(session_info)
if not source or not source.platform_name:
# No valid delivery source - don't ack, let event be retried
return
limit = PLATFORM_LIMITS.get(source.platform_name, float("inf"))
chunks = chunk_message(
event.content,
int(limit) if limit != float("inf") else len(event.content),
)
channel = self._get_channel(source.platform_name)
if channel:
success = await self._deliver_with_retry(chunks, source, channel)
if not success:
self.logger.error(f"Dropped message for session {event.session_id}")
self.context.eventbus.ack(event)
self.logger.info(
f"Delivered message to {source.platform_name} for session {event.session_id}"
)
except Exception as e:
self.logger.error(f"Failed to deliver message: {e}")
def _get_channel(self, platform: str) -> "Channel[Any] | None":
"""Get the message channel for a platform."""
for channel in self.context.channels:
if channel.platform_name == platform:
return channel
return None
10-websocket/src/mybot/server/server.py
+121
View File
@@ -0,0 +1,121 @@
"""Server orchestrator for worker-based architecture."""
import asyncio
import logging
from typing import TYPE_CHECKING
import uvicorn
from .worker import Worker
from .agent_worker import AgentWorker
from .delivery_worker import DeliveryWorker
from .channel_worker import ChannelWorker
from .websocket_worker import WebSocketWorker
from .app import create_app
from mybot.utils.config import ConfigReloader
if TYPE_CHECKING:
from mybot.core.context import SharedContext
logger = logging.getLogger(__name__)
class Server:
"""Orchestrates workers with queue-based communication."""
def __init__(self, context: "SharedContext"):
self.context = context
self.workers: list[Worker] = []
self._api_task: asyncio.Task | None = None
self.config_reloader: ConfigReloader = ConfigReloader(self.context.config)
async def run(self) -> None:
"""Start all workers and monitor for crashes."""
self._setup_workers()
self._start_workers()
# Start API server if configured
if self.context.config.api:
self._api_task = asyncio.create_task(self._run_api())
try:
await self._monitor_workers()
except asyncio.CancelledError:
logger.info("Server shutting down...")
await self._stop_all()
raise
def _setup_workers(self) -> None:
"""Create all workers."""
self.config_reloader.start()
# Create WebSocketWorker first and attach to context
ws_worker = WebSocketWorker(self.context)
self.context.websocket_worker = ws_worker
self.workers = [
self.context.eventbus, # EventBus (active worker)
AgentWorker(self.context), # SubscriberWorker
DeliveryWorker(self.context), # SubscriberWorker
ws_worker, # WebSocketWorker (SubscriberWorker)
]
if self.context.config.channels.enabled:
channels = self.context.channels
if channels:
self.workers.append(ChannelWorker(self.context))
logger.info(f"Channel enabled with {len(channels)} channel(es)")
else:
logger.warning("Channel enabled but no channels configured")
logger.info(f"Server setup complete with {len(self.workers)} core workers")
def _start_workers(self) -> None:
"""Start all workers as tasks."""
for worker in self.workers:
worker.start()
logger.info(f"Started {worker.__class__.__name__}")
async def _monitor_workers(self) -> None:
"""Monitor worker tasks, restart on crash."""
while True:
for worker in self.workers:
if worker.has_crashed():
exc = worker.get_exception()
if exc is None:
logger.warning(
f"{worker.__class__.__name__} exited unexpectedly"
)
else:
logger.error(f"{worker.__class__.__name__} crashed: {exc}")
worker.start()
logger.info(f"Restarted {worker.__class__.__name__}")
await asyncio.sleep(5)
async def _stop_all(self) -> None:
"""Stop all workers gracefully."""
for worker in self.workers:
await worker.stop()
# Stop config reloader
if self.config_reloader is not None:
self.config_reloader.stop()
async def _run_api(self) -> None:
"""Run the WebSocket API server."""
if not self.context.config.api:
return
app = create_app(self.context)
config = uvicorn.Config(
app,
host=self.context.config.api.host,
port=self.context.config.api.port,
)
server = uvicorn.Server(config)
logger.info(
f"WebSocket server started on {self.context.config.api.host}:{self.context.config.api.port}"
)
await server.serve()
10-websocket/src/mybot/server/websocket_worker.py
+149
View File
@@ -0,0 +1,149 @@
"""WebSocket worker for broadcasting events to connected clients."""
import logging
import time
import dataclasses
from typing import TYPE_CHECKING, Set
from mybot.core.agent import Agent
from fastapi import WebSocket
from fastapi.websockets import WebSocketDisconnect
from pydantic import ValidationError, BaseModel, Field
from .worker import SubscriberWorker
from mybot.core.events import Event, EventSource, InboundEvent, OutboundEvent, WebSocketEventSource
from mybot.utils.config import SourceSessionConfig
if TYPE_CHECKING:
from mybot.core.context import SharedContext
logger = logging.getLogger(__name__)
class WebSocketMessage(BaseModel):
"""Incoming WebSocket message from client."""
source: str = Field(..., min_length=1, description="Client identifier")
content: str = Field(..., min_length=1, description="Message content")
agent_id: str | None = Field(
None, description="Target agent ID (optional - uses routing if not specified)"
)
class WebSocketWorker(SubscriberWorker):
"""Manages WebSocket connections and event broadcasting."""
def __init__(self, context: "SharedContext"):
super().__init__(context)
self.clients: Set[WebSocket] = set()
# Auto-subscribe to event classes
for event_class in [InboundEvent, OutboundEvent]:
self.context.eventbus.subscribe(event_class, self.handle_event)
self.logger.info("WebSocketWorker subscribed to event types")
async def handle_connection(self, ws: WebSocket) -> None:
"""Handle a single WebSocket connection lifecycle."""
self.clients.add(ws)
self.logger.info(
f"WebSocket client connected. Total clients: {len(self.clients)}"
)
try:
await self._run_client_loop(ws)
finally:
self.clients.discard(ws)
self.logger.info(
f"WebSocket client disconnected. Total clients: {len(self.clients)}"
)
async def _run_client_loop(self, ws: WebSocket) -> None:
"""Run message receiving loop for a single client."""
while True:
try:
data = await ws.receive_json()
msg = WebSocketMessage(**data)
event = self._normalize_message(msg)
await self.context.eventbus.publish(event)
self.logger.debug(f"Emitted InboundEvent from WebSocket: {msg.source}")
except WebSocketDisconnect:
self.logger.info("Client disconnected normally")
break
except ValidationError as e:
await ws.send_json(
{"type": "error", "message": f"Validation error: {e}"}
)
self.logger.warning(f"Validation error from client: {e}")
except Exception as e:
self.logger.error(f"Unexpected error in client loop: {e}")
break
def _normalize_message(self, msg: "WebSocketMessage") -> InboundEvent:
"""Normalize WebSocketMessage to InboundEvent."""
source = WebSocketEventSource(user_id=msg.source)
agent_id = msg.agent_id
if agent_id is None:
agent_id = self.context.routing_table.resolve(str(source))
session_id = self._get_or_create_session_id(source)
return InboundEvent(
session_id=session_id,
source=source,
content=msg.content,
timestamp=time.time(),
)
async def handle_event(self, event: Event) -> None:
"""Handle EventBus event by broadcasting to WebSocket clients."""
if not self.clients:
return
# Serialize event to dict with type information
event_dict = {
"type": event.__class__.__name__,
}
event_dict.update(dataclasses.asdict(event))
# Convert EventSource to string for JSON serialization
if "source" in event_dict and hasattr(event.source, "__str__"):
event_dict["source"] = str(event.source)
# Broadcast to all clients
self.logger.debug(
f"Broadcasting {event.__class__.__name__} to {len(self.clients)} clients"
)
for client in list(self.clients):
try:
await client.send_json(event_dict)
except Exception as e:
self.logger.error(f"Failed to send to client: {e}")
self.clients.discard(client)
self.context.eventbus.ack(event)
def _get_or_create_session_id(self, source: "EventSource") -> str:
"""Get or create session ID for a given source."""
source_str = str(source)
source_session = self.context.config.sources.get(source_str)
if source_session:
return source_session.session_id
agent_def = self.context.agent_loader.load(self.context.config.default_agent)
agent = Agent(agent_def, self.context)
session = agent.new_session(source)
# Cache the session
self.context.config.set_runtime(
f"sources.{source_str}", SourceSessionConfig(session_id=session.session_id)
)
return session.session_id
10-websocket/src/mybot/server/worker.py
+60
View File
@@ -0,0 +1,60 @@
"""Base worker lifecycle management."""
import asyncio
import logging
from abc import ABC, abstractmethod
class Worker(ABC):
"""Base class for all workers with lifecycle management."""
def __init__(self, context):
self.context = context
self.logger = logging.getLogger(f"mybot.server.{self.__class__.__name__}")
self._task: asyncio.Task | None = None
@abstractmethod
async def run(self) -> None:
"""Main worker loop. Runs until cancelled."""
pass
def start(self) -> asyncio.Task:
"""Start the worker as an asyncio Task."""
self._task = asyncio.create_task(self.run())
return self._task
def is_running(self) -> bool:
"""Check if worker is actively running."""
return self._task is not None and not self._task.done()
def has_crashed(self) -> bool:
"""Check if worker crashed (done but not cancelled)."""
return (
self._task is not None and self._task.done() and not self._task.cancelled()
)
def get_exception(self) -> BaseException | None:
"""Get the exception if worker crashed, None otherwise."""
if self.has_crashed() and self._task is not None:
return self._task.exception()
return None
async def stop(self) -> None:
"""Gracefully stop the worker."""
if self._task:
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
class SubscriberWorker(Worker):
"""Worker that only subscribes to events, no active loop."""
async def run(self) -> None:
"""Wait for cancellation - actual work happens in event handlers."""
try:
await asyncio.Future()
except asyncio.CancelledError:
pass
10-websocket/src/mybot/tools/__init__.py
+7
View File
@@ -0,0 +1,7 @@
"""Tools module for agent capabilities."""
from mybot.tools.base import BaseTool, tool
from mybot.tools.builtin_tools import bash, edit_file, read_file, write_file
from mybot.tools.registry import ToolRegistry
__all__ = ["BaseTool", "tool", "ToolRegistry", "read_file", "write_file", "edit_file", "bash"]
10-websocket/src/mybot/tools/base.py
+63
View File
@@ -0,0 +1,63 @@
"""Base tool interface and decorator."""
import asyncio
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
class BaseTool(ABC):
"""Abstract base class for all tools."""
name: str
description: str
parameters: dict[str, Any] # JSON Schema for function calling
@abstractmethod
async def execute(self, session: "AgentSession", **kwargs: Any) -> str:
"""Execute the tool."""
def get_tool_schema(self) -> dict[str, Any]:
"""Get the tool/function schema for LiteLLM."""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters,
},
}
def tool(name: str, description: str, parameters: dict[str, Any]) -> Callable:
"""Decorator to register a function as a tool."""
def decorator(func: Callable) -> "FunctionTool":
return FunctionTool(name, description, parameters, func)
return decorator
class FunctionTool(BaseTool):
"""A tool created from a function using the @tool decorator."""
def __init__(
self,
name: str,
description: str,
parameters: dict[str, Any],
func: Callable,
):
self.name = name
self.description = description
self.parameters = parameters
self._func = func
async def execute(self, session: "AgentSession", **kwargs: Any) -> str:
"""Execute the underlying function."""
result = self._func(session=session, **kwargs)
if asyncio.iscoroutine(result):
result = await result
return str(result)
10-websocket/src/mybot/tools/builtin_tools.py
+133
View File
@@ -0,0 +1,133 @@
"""Built-in tools for agent capabilities."""
import asyncio
from pathlib import Path
from typing import TYPE_CHECKING
from mybot.tools.base import tool
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
# Filesystem tools
@tool(
name="read",
description="Read the contents of a text file",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Path to the file to read"},
},
"required": ["path"],
},
)
async def read_file(path: str, session: "AgentSession") -> str:
"""Read and return the contents of a file at the given path."""
try:
return Path(path).read_text()
except FileNotFoundError:
return f"Error: File not found: {path}"
except PermissionError:
return f"Error: Permission denied reading: {path}"
except IsADirectoryError:
return f"Error: Path is a directory, not a file: {path}"
except Exception as e:
return f"Error reading file: {e}"
@tool(
name="write",
description="Write content to a file",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Path to the file to write"},
"content": {
"type": "string",
"description": "Content to write to the file",
},
},
"required": ["path", "content"],
},
)
async def write_file(path: str, content: str, session: "AgentSession") -> str:
"""Write content to a file at the given path."""
try:
Path(path).write_text(content)
return f"Successfully wrote to: {path}"
except PermissionError:
return f"Error: Permission denied writing to: {path}"
except IsADirectoryError:
return f"Error: Path is a directory, not a file: {path}"
except Exception as e:
return f"Error writing file: {e}"
@tool(
name="edit",
description="Edit a file by replacing a string with new content",
parameters={
"type": "object",
"properties": {
"path": {"type": "string", "description": "Path to the file to edit"},
"old_text": {"type": "string", "description": "The text to replace"},
"new_text": {
"type": "string",
"description": "The new text to replace with",
},
},
"required": ["path", "old_text", "new_text"],
},
)
async def edit_file(
path: str, old_text: str, new_text: str, session: "AgentSession"
) -> str:
"""Edit a file by replacing old_text with new_text."""
try:
content = Path(path).read_text()
if old_text not in content:
return f"Error: '{old_text}' not found in {path}"
new_content = content.replace(old_text, new_text)
Path(path).write_text(new_content)
return f"Successfully edited {path}"
except FileNotFoundError:
return f"Error: File not found: {path}"
except PermissionError:
return f"Error: Permission denied editing: {path}"
except Exception as e:
return f"Error editing file: {e}"
# Shell tool
@tool(
name="bash",
description="Execute a bash shell command",
parameters={
"type": "object",
"properties": {
"command": {"type": "string", "description": "The bash command to execute"},
},
"required": ["command"],
},
)
async def bash(command: str, session: "AgentSession") -> str:
"""Execute a bash command and return the output."""
try:
process = await asyncio.create_subprocess_shell(
command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
stdout, stderr = await process.communicate()
output = stdout.decode() if stdout else ""
error = stderr.decode() if stderr else ""
if output and error:
return f"{output}\n{error}"
return output or error or "Command completed with no output"
except Exception as e:
return f"Error executing command: {e}"
10-websocket/src/mybot/tools/registry.py
+56
View File
@@ -0,0 +1,56 @@
"""Tool registry for managing available tools."""
from typing import TYPE_CHECKING, Any
from mybot.tools.base import BaseTool
from mybot.tools.builtin_tools import bash, edit_file, read_file, write_file
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
class ToolRegistry:
"""Registry for all available tools."""
def __init__(self) -> None:
"""Initialize an empty tool registry."""
self._tools: dict[str, BaseTool] = {}
def register(self, tool: BaseTool) -> None:
"""Register a tool."""
self._tools[tool.name] = tool
def get(self, name: str) -> BaseTool | None:
"""Get a tool by name."""
return self._tools.get(name)
def list_all(self) -> list[BaseTool]:
"""List all registered tools."""
return list(self._tools.values())
def get_tool_schemas(self) -> list[dict[str, Any]]:
"""Get tool schemas for all registered tools."""
return [tool.get_tool_schema() for tool in self._tools.values()]
async def execute_tool(
self, name: str, session: "AgentSession", **kwargs: Any
) -> str:
"""Execute a tool by name."""
tool = self.get(name)
if tool is None:
raise ValueError(f"Tool not found: {name}")
return await tool.execute(session=session, **kwargs)
@classmethod
def with_builtins(cls) -> "ToolRegistry":
"""Create a ToolRegistry with builtin tools already registered."""
registry = cls()
registry.register(read_file)
registry.register(write_file)
registry.register(edit_file)
registry.register(bash)
return registry
10-websocket/src/mybot/tools/skill_tool.py
+51
View File
@@ -0,0 +1,51 @@
"""Skill tool factory for creating dynamic skill tool."""
from typing import TYPE_CHECKING
from mybot.tools.base import tool
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
from mybot.core.skill_loader import SkillLoader
def create_skill_tool(skill_loader: "SkillLoader"):
"""Factory function to create skill tool with dynamic schema."""
skill_metadata = skill_loader.discover_skills()
if not skill_metadata:
return None
# Build XML description of available skills
skills_xml = "<skills>\n"
for meta in skill_metadata:
skills_xml += f' <skill name="{meta.name}">{meta.description}</skill>\n'
skills_xml += "</skills>"
# Build enum of skill IDs
skill_enum = [meta.id for meta in skill_metadata]
@tool(
name="skill",
description=f"Load and invoke a specialized skill. {skills_xml}",
parameters={
"type": "object",
"properties": {
"skill_name": {
"type": "string",
"enum": skill_enum,
"description": "The name of the skill to load",
}
},
"required": ["skill_name"],
},
)
async def skill_tool(skill_name: str, session: "AgentSession") -> str:
"""Load and return skill content."""
try:
skill_def = skill_loader.load_skill(skill_name)
return skill_def.content
except Exception:
return f"Error: Skill '{skill_name}' not found. It may have been removed or is unavailable."
return skill_tool
10-websocket/src/mybot/tools/webread_tool.py
+47
View File
@@ -0,0 +1,47 @@
"""Webread tool factory."""
from typing import TYPE_CHECKING
from mybot.tools.base import BaseTool, tool
from mybot.provider.web_read import WebReadProvider
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
from mybot.core.context import SharedContext
def create_webread_tool(context: "SharedContext") -> BaseTool | None:
"""Factory to create webread tool with injected context."""
if not context.config.webread:
return None
provider = WebReadProvider.from_config(context.config)
@tool(
name="webread",
description=(
"Read and extract content from a web page. "
"Returns the page content as markdown."
),
parameters={
"type": "object",
"properties": {
"url": {
"type": "string",
"description": "The URL to read",
}
},
"required": ["url"],
},
)
async def webread(url: str, session: "AgentSession") -> str:
"""Read a web page and return markdown content."""
result = await provider.read(url)
if result.error:
return f"Error reading {url}: {result.error}"
return f"**{result.title}**\n\n{result.content}"
return webread
10-websocket/src/mybot/tools/websearch_tool.py
+50
View File
@@ -0,0 +1,50 @@
"""Websearch tool factory."""
from typing import TYPE_CHECKING
from mybot.tools.base import BaseTool, tool
from mybot.provider.web_search import WebSearchProvider
if TYPE_CHECKING:
from mybot.core.agent import AgentSession
from mybot.core.context import SharedContext
def create_websearch_tool(context: "SharedContext") -> BaseTool | None:
"""Factory to create websearch tool with injected context."""
if not context.config.websearch:
return None
provider = WebSearchProvider.from_config(context.config)
@tool(
name="websearch",
description=(
"Search the web for information. "
"Returns a list of results with titles, URLs, and snippets."
),
parameters={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The search query",
}
},
"required": ["query"],
},
)
async def websearch(query: str, session: "AgentSession") -> str:
"""Search the web and return formatted results."""
results = await provider.search(query)
if not results:
return "No results found."
output = []
for i, r in enumerate(results, 1):
output.append(f"{i}. **{r.title}**\n {r.url}\n {r.snippet}")
return "\n\n".join(output)
return websearch
10-websocket/src/mybot/utils/__init__.py
+17
View File
@@ -0,0 +1,17 @@
"""Utilities package."""
from mybot.utils.def_loader import (
DefNotFoundError,
InvalidDefError,
discover_definitions,
parse_definition,
)
from mybot.utils.logging import setup_logging
__all__ = [
"DefNotFoundError",
"InvalidDefError",
"discover_definitions",
"parse_definition",
"setup_logging",
]
10-websocket/src/mybot/utils/config.py
+241
View File
@@ -0,0 +1,241 @@
"""Configuration management with hot reload support."""
import logging
from pathlib import Path
from typing import Any, Literal
import yaml
from pydantic import BaseModel, Field, field_validator, model_validator
from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer
class LLMConfig(BaseModel):
"""LLM provider configuration."""
provider: str
model: str
api_key: str
api_base: str | None = None
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
max_tokens: int = Field(default=2048, gt=0)
@field_validator("api_base")
@classmethod
def api_base_must_be_url(cls, v: str | None) -> str | None:
if v is not None and not v.startswith(("http://", "https://")):
raise ValueError("api_base must be a valid URL")
return v
class TelegramConfig(BaseModel):
"""Telegram platform configuration."""
enabled: bool = True
bot_token: str
allowed_user_ids: list[str] = Field(default_factory=list)
class DiscordConfig(BaseModel):
"""Discord platform configuration."""
enabled: bool = True
bot_token: str
channel_id: str | None = None
allowed_user_ids: list[str] = Field(default_factory=list)
class BraveWebSearchConfig(BaseModel):
"""Configuration for web search provider."""
provider: Literal["brave"] = "brave"
api_key: str
class Crawl4AIWebReadConfig(BaseModel):
"""Configuration for web read provider."""
provider: Literal["crawl4ai"] = "crawl4ai"
class SourceSessionConfig(BaseModel):
"""Session affinity configuration for a source."""
session_id: str
class ChannelConfig(BaseModel):
"""Channel configuration."""
enabled: bool = False
telegram: TelegramConfig | None = None
discord: DiscordConfig | None = None
class ApiConfig(BaseModel):
"""HTTP API configuration."""
host: str = "127.0.0.1"
port: int = Field(default=8000, gt=0, lt=65536)
class Config(BaseModel):
"""Main configuration with hot reload support."""
workspace: Path
llm: LLMConfig
default_agent: str
agents_path: Path = Field(default=Path("agents"))
skills_path: Path = Field(default=Path("skills"))
logging_path: Path = Field(default=Path(".logs"))
history_path: Path = Field(default=Path(".history"))
event_path: Path = Field(default=Path(".event"))
websearch: BraveWebSearchConfig | None = None
webread: Crawl4AIWebReadConfig | None = None
channels: ChannelConfig = Field(default_factory=ChannelConfig)
api: ApiConfig = Field(default_factory=ApiConfig)
sources: dict[str, SourceSessionConfig] = Field(default_factory=dict)
default_delivery_source: str | None = None
@model_validator(mode="after")
def resolve_paths(self) -> "Config":
"""Resolve relative paths to absolute using workspace."""
for field_name in (
"agents_path",
"skills_path",
"logging_path",
"history_path",
"event_path",
):
path = getattr(self, field_name)
if not path.is_absolute():
setattr(self, field_name, self.workspace / path)
return self
@classmethod
def load(cls, workspace_dir: Path) -> "Config":
"""Load configuration from workspace directory."""
config_data = cls._load_merged_configs(workspace_dir)
config_data["workspace"] = workspace_dir
return cls.model_validate(config_data)
@classmethod
def _load_merged_configs(cls, workspace_dir: Path) -> dict[str, Any]:
"""Load and merge user and runtime config files."""
config_data: dict[str, Any] = {}
user_config = workspace_dir / "config.user.yaml"
runtime_config = workspace_dir / "config.runtime.yaml"
if user_config.exists():
with open(user_config) as f:
config_data = cls._deep_merge(config_data, yaml.safe_load(f) or {})
if runtime_config.exists():
with open(runtime_config) as f:
config_data = cls._deep_merge(config_data, yaml.safe_load(f) or {})
return config_data
@staticmethod
def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
"""Deep merge override dict into base dict."""
result = base.copy()
for key, value in override.items():
if (
key in result
and isinstance(result[key], dict)
and isinstance(value, dict)
):
result[key] = Config._deep_merge(result[key], value)
else:
result[key] = value
return result
def _set_nested(self, obj: dict, key: str, value: Any) -> None:
"""Set a nested value in a dict using dot notation."""
keys = key.split(".")
for k in keys[:-1]:
if k not in obj or not isinstance(obj[k], dict):
obj[k] = {}
obj = obj[k]
obj[keys[-1]] = value
def _set_config_value(self, config_path: Path, key: str, value: Any) -> None:
"""Update a config value in a YAML file."""
# Load existing or start fresh
if config_path.exists():
with open(config_path) as f:
data = yaml.safe_load(f) or {}
else:
data = {}
if isinstance(value, BaseModel):
value = value.model_dump()
# Update the key (supports nested via dot notation)
self._set_nested(data, key, value)
# Write back
with open(config_path, "w") as f:
yaml.dump(data, f)
def set_user(self, key: str, value: Any) -> None:
"""Update a config value in config.user.yaml."""
self._set_config_value(self.workspace / "config.user.yaml", key, value)
def set_runtime(self, key: str, value: Any) -> None:
"""Update a runtime value in config.runtime.yaml."""
self._set_config_value(self.workspace / "config.runtime.yaml", key, value)
def reload(self) -> bool:
"""Re-read config.user.yaml and merge with runtime."""
try:
config_data = self._load_merged_configs(self.workspace)
config_data["workspace"] = self.workspace
# Create new instance and copy values
new_config = Config.model_validate(config_data)
# Update all fields from new config
for field_name in Config.model_fields:
setattr(self, field_name, getattr(new_config, field_name))
return True
except Exception as e:
logging.debug("Config reload failed: %s", e)
return False
class ConfigHandler(FileSystemEventHandler):
"""Handles config file modification events."""
def __init__(self, config: Config):
self._config = config
def on_modified(self, event):
"""Reload config when config.user.yaml changes."""
if not event.is_directory and event.src_path.endswith("config.user.yaml"):
self._config.reload()
class ConfigReloader:
"""Manages watchdog observer for config hot reload."""
def __init__(self, config: Config):
self._config = config
self._observer = Observer()
def start(self) -> None:
"""Start watching config file for changes."""
handler = ConfigHandler(self._config)
self._observer.schedule(handler, str(self._config.workspace), recursive=False)
self._observer.start()
def stop(self) -> None:
"""Stop watching."""
self._observer.stop()
self._observer.join()
del self._observer
10-websocket/src/mybot/utils/def_loader.py
+105
View File
@@ -0,0 +1,105 @@
"""Shared utilities for loading definition files (agents, skills, crons)."""
import logging
from pathlib import Path
from typing import Any, Callable, TypeVar
import yaml
T = TypeVar("T")
logger = logging.getLogger(__name__)
class DefNotFoundError(Exception):
"""Definition folder or file doesn't exist."""
def __init__(self, kind: str, def_id: str):
super().__init__(f"{kind.capitalize()} not found: {def_id}")
self.kind = kind
self.def_id = def_id
class InvalidDefError(Exception):
"""Definition file is malformed."""
def __init__(self, kind: str, def_id: str, reason: str):
super().__init__(f"Invalid {kind} '{def_id}': {reason}")
self.kind = kind
self.def_id = def_id
self.reason = reason
def parse_definition(
content: str,
def_id: str,
parse_fn: Callable[[str, dict[str, Any], str], T],
) -> T:
"""Parse YAML frontmatter + markdown body with type conversion."""
# Find frontmatter delimiters
if not content.startswith("---\n"):
body = content
return parse_fn(def_id, {}, body)
end_delimiter = content.find("\n---\n", 4)
if end_delimiter == -1:
body = content
return parse_fn(def_id, {}, body)
frontmatter_text = content[4:end_delimiter]
body = content[end_delimiter + 5 :]
raw_dict = yaml.safe_load(frontmatter_text) or {}
return parse_fn(def_id, raw_dict, body)
def discover_definitions(
path: Path,
filename: str,
parse_fn: Callable[[str, dict[str, Any], str], T | None],
) -> list[T]:
"""Scan directory for definition files."""
if not path.exists():
logger.warning(f"Definitions directory not found: {path}")
return []
results = []
for def_dir in path.iterdir():
if not def_dir.is_dir():
continue
def_file = def_dir / filename
if not def_file.exists():
logger.warning(f"No {filename} found in {def_dir.name}")
continue
try:
content = def_file.read_text()
result = parse_definition(content, def_dir.name, parse_fn)
if result is not None:
results.append(result)
except Exception as e:
logger.warning(f"Failed to parse {def_dir.name}: {e}")
continue
return results
def write_definition(
def_id: str,
frontmatter: dict[str, Any],
body: str,
base_path: Path,
filename: str,
) -> Path:
"""Write a definition file with YAML frontmatter and markdown body."""
def_dir = base_path / def_id
def_dir.mkdir(parents=True, exist_ok=True)
# Build file content with YAML frontmatter
yaml_content = yaml.dump(frontmatter, default_flow_style=False, sort_keys=False)
content = f"---\n{yaml_content}---\n\n{body.strip()}\n"
def_file = def_dir / filename
def_file.write_text(content)
return def_file
10-websocket/src/mybot/utils/logging.py
+34
View File
@@ -0,0 +1,34 @@
"""Logging configuration for mybot."""
import logging
import sys
from mybot.utils.config import Config
from logging.handlers import RotatingFileHandler
def setup_logging(config: Config, console_output: bool = False) -> None:
"""Set up logging for pickle-bot."""
format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
formatter = logging.Formatter(format_str)
# Console format is simpler (no timestamp)
console_format = "%(levelname)s - %(name)s - %(message)s"
console_formatter = logging.Formatter(console_format)
root_logger = logging.getLogger("mybot")
root_logger.setLevel(logging.DEBUG)
config.logging_path.mkdir(parents=True, exist_ok=True)
file_handler = RotatingFileHandler(
config.logging_path / "mybot.log", maxBytes=256 * 1024 * 128, backupCount=3
)
file_handler.setFormatter(formatter)
file_handler.setLevel(logging.DEBUG)
root_logger.addHandler(file_handler)
# Optionally log to console (for server mode)
if console_output:
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setFormatter(console_formatter)
console_handler.setLevel(logging.INFO)
root_logger.addHandler(console_handler)
GAP.md
+8 -4
View File
@@ -25,12 +25,16 @@ This document tracks features that exist in picklebot but will **never** be adde
## HTTP API Endpoints
**Picklebot has:**
- HTTP API endpoints for skills/agents/crons
- Full REST API with endpoints for skills/agents/crons/sessions/memories
- Complete API server with FastAPI routers
**Tutorial has:** None of these
**Tutorial has:**
- WebSocket-only FastAPI server (no REST API endpoints)
- Just the `/ws` endpoint for real-time communication
**Why it's fine:**
- Tutorial is for learning agent patterns, not production deployment
- Users can add these later if needed
- Tutorial focuses on real-time WebSocket communication
- REST API endpoints add complexity without teaching core concepts
- Users can add REST endpoints later if needed
---