sssnake/dark_hal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
DigiJ
2026-03-13 12:56:43 -07:00
commit 159cf9fcfe
309 changed files with 64584 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

217
.gitignore vendored Normal file
View File

@@ -0,0 +1,217 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
.venv1/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be added to the global gitignore or merged into this project gitignore. For a PyCharm
# project, it is generally recommended to not ignore these files.
.idea/
# VS Code
.vscode/
# Model files and checkpoints
*.bin
*.safetensors
*.gguf
*.ggml
*.pth
*.pt
*.ckpt
*.pkl
*.pickle
# Large data files
*.csv
*.json
*.txt
*.log
# API Keys and secrets
*.key
*.env
config.yaml
secrets.yaml
api_keys.json
# Temporary files
temp/
tmp/
*.tmp
*.temp
# OS generated files
.DS_Store
.DS_Store?
._*
.Spotlight-V100
.Trashes
ehthumbs.db
Thumbs.db
# DarkHal specific ignores
logs/
output/
models/
checkpoints/
cache/
# Project-specific directories to ignore at repo root
/experimental_agent/
/llama-cpp-python/
/llm_chess-main/
/Darkhal/
/darkhall.egg-info/
/.claude/

1
0.19.0 Normal file
View File

@@ -0,0 +1 @@
- python-dotenv

1
0.2.80 Normal file
View File

@@ -0,0 +1 @@
- llama-cpp-python

1
0.24.0 Normal file
View File

@@ -0,0 +1 @@
- huggingface_hub

1
1.21.0 Normal file
View File

@@ -0,0 +1 @@
- numpy

1
2.32.0 Normal file
View File

@@ -0,0 +1 @@
- requests

1
5.8.0 Normal file
View File

@@ -0,0 +1 @@
- psutil

1
9.0.0 Normal file
View File

@@ -0,0 +1 @@
- Pillow

0
=0.19.0 Normal file
View File

48
=0.2.3 Normal file
View File

@@ -0,0 +1,48 @@
Collecting auto-gptq
Downloading auto_gptq-0.7.1.tar.gz (126 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/90/e5/b22697903982284fe284568fb2663a2196694a8eee637f5cf4ccfe435a38/auto_gptq-0.7.1.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/90/e5/b22697903982284fe284568fb2663a2196694a8eee637f5cf4ccfe435a38/auto_gptq-0.7.1.tar.gz has inconsistent version: expected '0.7.1', but metadata has '0.7.1+cu118'
Downloading auto_gptq-0.7.0.tar.gz (124 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/34/71/c3e73cf17681f6ff4754ef8f4cb8b67af3def230fc8711eac1250bbd78d5/auto_gptq-0.7.0.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/34/71/c3e73cf17681f6ff4754ef8f4cb8b67af3def230fc8711eac1250bbd78d5/auto_gptq-0.7.0.tar.gz has inconsistent version: expected '0.7.0', but metadata has '0.7.0+cu118'
Downloading auto_gptq-0.6.0.tar.gz (120 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/49/af/02b66e55dfd9aeb0ece923843043724ed7432ec0c649ea0f3b9fa1dd90c6/auto_gptq-0.6.0.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/49/af/02b66e55dfd9aeb0ece923843043724ed7432ec0c649ea0f3b9fa1dd90c6/auto_gptq-0.6.0.tar.gz has inconsistent version: expected '0.6.0', but metadata has '0.6.0+cu118'
Downloading auto_gptq-0.5.1.tar.gz (112 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/db/77/ec5a16c5625b0791dccfe5e42356171332ed3537c1df505d64a162148c8f/auto_gptq-0.5.1.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/db/77/ec5a16c5625b0791dccfe5e42356171332ed3537c1df505d64a162148c8f/auto_gptq-0.5.1.tar.gz has inconsistent version: expected '0.5.1', but metadata has '0.5.1+cu118'
Downloading auto_gptq-0.5.0.tar.gz (111 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/3d/fa/c2cd09965b2dbf4e454d9f073376922f7139a574f617f70a22adb203eced/auto_gptq-0.5.0.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/3d/fa/c2cd09965b2dbf4e454d9f073376922f7139a574f617f70a22adb203eced/auto_gptq-0.5.0.tar.gz has inconsistent version: expected '0.5.0', but metadata has '0.5.0+cu118'
Downloading auto_gptq-0.3.2.tar.gz (63 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Discarding https://files.pythonhosted.org/packages/1b/79/5a3a7d877a9b0a72f528e9977ec65cdb9fad800fa4f5110f87f2acaaf6fe/auto_gptq-0.3.2.tar.gz (from https://pypi.org/simple/auto-gptq/) (requires-python:>=3.8.0): Requested auto-gptq from https://files.pythonhosted.org/packages/1b/79/5a3a7d877a9b0a72f528e9977ec65cdb9fad800fa4f5110f87f2acaaf6fe/auto_gptq-0.3.2.tar.gz has inconsistent version: expected '0.3.2', but metadata has '0.3.2+cu118'
Downloading auto_gptq-0.3.1.tar.gz (63 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Collecting autoawq
Downloading autoawq-0.2.9.tar.gz (74 kB)
Preparing metadata (setup.py): started
Preparing metadata (setup.py): finished with status 'done'
Collecting exllamav2
Downloading exllamav2-0.3.2-py3-none-any.whl.metadata (430 bytes)
Requirement already satisfied: accelerate>=0.19.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from auto-gptq) (1.10.0)
Collecting datasets (from auto-gptq)
Downloading datasets-4.0.0-py3-none-any.whl.metadata (19 kB)
Requirement already satisfied: numpy in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from auto-gptq) (2.3.2)
Collecting rouge (from auto-gptq)
Downloading rouge-1.0.1-py3-none-any.whl.metadata (4.1 kB)
Requirement already satisfied: torch>=1.13.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from auto-gptq) (2.7.1+cu118)
Requirement already satisfied: safetensors in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from auto-gptq) (0.6.2)
Requirement already satisfied: transformers>=4.31.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from auto-gptq) (4.55.2)
Collecting peft (from auto-gptq)
Downloading peft-0.17.0-py3-none-any.whl.metadata (14 kB)
INFO: pip is looking at multiple versions of autoawq to determine which version is compatible with other requirements. This could take a while.
Collecting autoawq
Downloading autoawq-0.2.8.tar.gz (71 kB)

0
=0.2.5 Normal file
View File

0
=0.33.0 Normal file
View File

35
=0.4.2 Normal file
View File

@@ -0,0 +1,35 @@
Collecting transformers
Using cached transformers-4.55.2-py3-none-any.whl.metadata (41 kB)
Collecting accelerate
Using cached accelerate-1.10.0-py3-none-any.whl.metadata (19 kB)
Requirement already satisfied: safetensors in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (0.6.2)
Requirement already satisfied: filelock in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (3.19.1)
Requirement already satisfied: huggingface-hub<1.0,>=0.34.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (0.34.4)
Requirement already satisfied: numpy>=1.17 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (2.3.2)
Requirement already satisfied: packaging>=20.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (25.0)
Requirement already satisfied: pyyaml>=5.1 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (6.0.2)
Requirement already satisfied: regex!=2019.12.17 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (2025.7.34)
Requirement already satisfied: requests in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (2.32.5)
Requirement already satisfied: tokenizers<0.22,>=0.21 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (0.21.4)
Requirement already satisfied: tqdm>=4.27 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from transformers) (4.67.1)
Requirement already satisfied: fsspec>=2023.5.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from huggingface-hub<1.0,>=0.34.0->transformers) (2025.7.0)
Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from huggingface-hub<1.0,>=0.34.0->transformers) (4.14.1)
Requirement already satisfied: psutil in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from accelerate) (7.0.0)
Requirement already satisfied: torch>=2.0.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from accelerate) (2.7.1+cu118)
Requirement already satisfied: sympy>=1.13.3 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from torch>=2.0.0->accelerate) (1.13.3)
Requirement already satisfied: networkx in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from torch>=2.0.0->accelerate) (3.5)
Requirement already satisfied: jinja2 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from torch>=2.0.0->accelerate) (3.1.6)
Requirement already satisfied: setuptools in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from torch>=2.0.0->accelerate) (80.9.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from sympy>=1.13.3->torch>=2.0.0->accelerate) (1.3.0)
Requirement already satisfied: colorama in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from tqdm>=4.27->transformers) (0.4.6)
Requirement already satisfied: MarkupSafe>=2.0 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from jinja2->torch>=2.0.0->accelerate) (3.0.2)
Requirement already satisfied: charset_normalizer<4,>=2 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from requests->transformers) (3.4.3)
Requirement already satisfied: idna<4,>=2.5 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from requests->transformers) (3.10)
Requirement already satisfied: urllib3<3,>=1.21.1 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from requests->transformers) (2.5.0)
Requirement already satisfied: certifi>=2017.4.17 in c:\users\mdavi\pycharmprojects\llm_train\.venv\lib\site-packages (from requests->transformers) (2025.8.3)
Downloading transformers-4.55.2-py3-none-any.whl (11.3 MB)
---------------------------------------- 11.3/11.3 MB 39.1 MB/s 0:00:00
Downloading accelerate-1.10.0-py3-none-any.whl (374 kB)
Installing collected packages: accelerate, transformers
Successfully installed accelerate-1.10.0 transformers-4.55.2

0
=0.7.1 Normal file
View File

1
=2.0.0 Normal file
View File

@@ -0,0 +1 @@
Looking in indexes: https://download.pytorch.org/whl/cu121

0
=4.43.0 Normal file
View File

161
DEPENDENCY_INSTALLER_README.md Normal file
View File

@@ -0,0 +1,161 @@
# Windows Dependency Installer for LLM_Train
This dependency installer helps you automatically install common software packages required for LLM_Train on Windows using the Chocolatey package manager.
## Quick Start
### Option 1: Batch File (Recommended)
1. Double-click `install_dependencies.bat`
2. Grant administrator permissions when prompted
3. Select packages to install in the GUI
### Option 2: PowerShell Script
1. Right-click `install_dependencies.ps1` → "Run with PowerShell"
2. Grant administrator permissions when prompted
3. Select packages to install in the GUI
### Option 3: Direct Python Execution
```bash
# Run as administrator
python windows_dependency_installer.py
```
## What Gets Installed
### Essential Packages (Auto-selected)
- **Git** - Version control system (required for repository cloning)
- **Python 3** - Python programming language (if not already installed)
- **7-Zip** - File archiver for extracting downloads
- **Visual C++ Redistributables** - Microsoft runtime libraries
### Development Tools
- **Visual Studio Code** - Advanced code editor with Python support
- **Notepad++** - Enhanced text editor
### GPU Acceleration
- **CUDA Toolkit** - NVIDIA CUDA development toolkit for GPU acceleration
- **NVIDIA Display Driver** - Latest NVIDIA graphics drivers
### System Utilities
- **Wget** - Command-line downloader
- **cURL** - Data transfer tool
- **PowerToys** - Windows system utilities
### Runtimes
- **.NET Runtime** - Microsoft .NET framework
### Optional Tools
- **WinRAR** - Alternative file archiver
- **Firefox** - Web browser
- **VLC Media Player** - Media player
## System Requirements
- **Windows 10/11** (Windows 8.1 may work but is not tested)
- **Administrator privileges** (required for Chocolatey and package installation)
- **Internet connection** (for downloading packages)
- **Python 3.7+** (for running the installer GUI)
## Features
### Chocolatey Integration
- Automatically installs Chocolatey if not present
- Uses Chocolatey's robust package management
- Handles dependencies automatically
### Smart Package Selection
- **Select Essential** - Chooses only required packages
- **Select All** - Selects all available packages
- **Custom Selection** - Pick individual packages
### Installation Monitoring
- Real-time installation log
- Progress tracking
- Success/failure reporting
- Package status checking
### System Status Checks
- Administrator privilege detection
- Chocolatey installation status
- Individual package installation status
## Troubleshooting
### "Python not found" Error
1. Install Python from https://python.org/downloads/
2. During installation, check "Add Python to PATH"
3. Restart your command prompt/PowerShell
### "Administrator privileges required" Error
1. Right-click the batch file → "Run as administrator"
2. Or open Command Prompt as administrator and run manually
### "Chocolatey installation failed" Error
1. Ensure you're running as administrator
2. Check your internet connection
3. Temporarily disable antivirus software during installation
4. Check Windows execution policy: `Set-ExecutionPolicy RemoteSigned`
### Package Installation Failures
1. Check the installation log for specific error messages
2. Try installing packages individually
3. Ensure sufficient disk space
4. Check for conflicting software
### Network/Firewall Issues
1. Ensure Chocolatey URLs are not blocked:
- https://community.chocolatey.org/
- https://packages.chocolatey.org/
2. Configure proxy settings if behind corporate firewall
3. Temporarily disable firewall/antivirus
## Manual Installation
If the automatic installer fails, you can install Chocolatey manually:
1. Open PowerShell as Administrator
2. Run:
```powershell
Set-ExecutionPolicy Bypass -Scope Process -Force;
[System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072;
iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))
```
3. Then install packages manually:
```powershell
choco install git python 7zip vcredist-all -y
choco install cuda nvidia-display-driver -y # For GPU support
```
## Package Descriptions
### Why Each Package?
- **Git**: Required for cloning model repositories and version control
- **Python**: Core runtime for LLM_Train (if system Python is outdated)
- **7-Zip**: Many model files are compressed and need extraction
- **Visual C++ Redistributables**: Required by many Python packages and binaries
- **CUDA Toolkit**: Enables GPU acceleration for faster model inference
- **NVIDIA Drivers**: Latest drivers for optimal GPU performance
- **Visual Studio Code**: Best IDE for Python development and debugging
- **Wget/cURL**: Alternative download tools for model files
- **PowerToys**: Useful Windows utilities for power users
## Security Notes
- All packages are installed from official Chocolatey community repository
- Chocolatey packages are maintained by the community and Microsoft
- Administrator privileges are required only for system-wide installation
- No personal data is collected or transmitted
## Support
If you encounter issues:
1. Check the installation log for error messages
2. Search for the specific error on Chocolatey community forums
3. Try installing individual packages manually
4. Ensure your Windows is up to date
## License
This installer uses Chocolatey (Apache 2.0 License) and installs various packages with their respective licenses. Please review individual package licenses as needed.

50
MANIFEST.in Normal file
View File

@@ -0,0 +1,50 @@
# Include the README and license files
include README.md
include *.txt
include *.md
# Include configuration files
include pyproject.toml
include requirements.txt
# Include everything in darkhal package directory
recursive-include darkhal *
# Include any JSON configuration files at root
include *.json
# Include environment files
include *.env
# Include any shell scripts that might be useful
include *.sh
include *.bat
# Exclude development and build files
exclude .gitignore
exclude *.pyc
exclude __pycache__
recursive-exclude * __pycache__
recursive-exclude * *.py[co]
recursive-exclude * *.orig
recursive-exclude * *.rej
exclude .coverage
exclude .tox
recursive-exclude build *
recursive-exclude dist *
recursive-exclude .git *
recursive-exclude win-install *
recursive-exclude debian-install *
recursive-exclude tests *
recursive-exclude llm_chess-main *
recursive-exclude temp_llm_chess *
recursive-exclude llama-cpp-python *
recursive-exclude downloads *
recursive-exclude models *
recursive-exclude games *
# Exclude temporary and cache files
exclude .DS_Store
recursive-exclude * .DS_Store
exclude Thumbs.db
recursive-exclude * Thumbs.db

704
MANUAL.md Normal file
View File

@@ -0,0 +1,704 @@
# DarkHal 2.0 User Manual
## Table of Contents
1. [Introduction](#introduction)
2. [Getting Started](#getting-started)
3. [Main Interface](#main-interface)
4. [Model Management](#model-management)
5. [Agent Mode](#agent-mode)
6. [Advanced Features](#advanced-features)
7. [Troubleshooting](#troubleshooting)
8. [Keyboard Shortcuts](#keyboard-shortcuts)
9. [Command Reference](#command-reference)
10. [FAQ](#faq)
---
## Introduction
DarkHal 2.0 is an advanced AI model management platform that provides comprehensive tools for loading, running, and interacting with Large Language Models (LLMs). This manual covers all features and capabilities of the platform.
### System Requirements
**Minimum Requirements:**
- Windows 10/11, Linux (Ubuntu 20.04+), or macOS 11+
- Python 3.8 or higher
- 8GB RAM
- 20GB free disk space
- Internet connection for downloading models
**Recommended Requirements:**
- 16GB+ RAM
- NVIDIA GPU with 8GB+ VRAM
- 100GB+ free disk space for models
- High-speed internet for model downloads
---
## Getting Started
### First Launch
1. **Start DarkHal**:
```bash
python main.py --gui
```
2. **Initial Setup**:
- The splash screen will appear showing system information
- Configure your models directory via Settings menu
- Set up HuggingFace API token if you plan to download models
3. **Load Your First Model**:
- Click "Browse Model" to select a model file
- Or use "Browse Folder" to select a model directory
- Click "Load Model" to initialize
### Understanding the Interface
The main window consists of several tabs:
- **Run**: Chat interface and model loading
- **Model Library**: Browse and manage local models
- **HuggingFace**: Download models from HuggingFace Hub
- **Downloads**: Monitor active downloads
- **MCP**: Model Context Protocol server
- **Converter**: Convert between model formats
- **Chess**: Specialized chess engine interface
---
## Main Interface
### Run Tab
The Run tab is your primary workspace for interacting with models.
#### Model Selection Panel
**Model Path Input**:
- Enter the full path to your model file or directory
- Supports drag-and-drop from file explorer
- Auto-completes recently used models
**Browse Model Button**:
- Opens file dialog to select model files
- Filters: GGUF, SafeTensors, PyTorch, GPTQ, AWQ, EXL2
- Shows all supported formats
**Browse Folder Button**:
- Select directories containing model files
- Useful for HuggingFace format models
- Auto-detects config.json
**Load Model Button**:
- Initializes the selected model
- Shows loading progress
- Displays model information when complete
#### Chat Interface
**Chat Mode Options**:
- **Stream Output**: Shows text as it's generated
- **Chess Mode**: Enables ChessGPT for chess moves
- **Agent Mode**: Enables system command execution
**Input Area**:
- Multi-line text input
- Supports Ctrl+Enter for sending
- Maintains conversation history
**Output Display**:
- Shows conversation with "You:" and "Assistant:" prefixes
- Auto-scrolls to latest message
- Supports text selection and copying
**Control Buttons**:
- **Send (Chat)**: Submit your message
- **Stop**: Interrupt generation
- **Clear Output**: Clear conversation display
- **Clear History**: Reset conversation context
### Model Settings Tab
#### Basic Settings
**Context Size (n_ctx)**:
- Range: 512 to 32768 tokens
- Default: 4096
- Higher values use more memory but allow longer conversations
**GPU Layers**:
- Range: 0 to model layer count
- 0 = CPU only
- Higher values offload more to GPU
**Max Tokens**:
- Maximum tokens to generate
- Range: 1 to context size
- Default: 2048
#### Advanced Loading Options
**Quantization**:
- `none`: Full precision (FP16/FP32)
- `4bit`: ~75% memory savings
- `8bit`: ~50% memory savings
- `gptq`: Pre-quantized GPTQ format
- `awq`: Pre-quantized AWQ format
- `exl2`: Pre-quantized EXL2 format
**Device Strategy**:
- `auto`: Automatic distribution
- `force_gpu`: All layers on GPU
- `balanced_split`: Split between CPU/GPU
- `cpu_only`: CPU processing only
**GPU Memory Limit**:
- Maximum VRAM to use (in GB)
- Used with balanced_split strategy
- Prevents out-of-memory errors
#### Sampling Parameters
**Temperature** (0.0 - 2.0):
- Controls randomness
- 0.0 = Deterministic
- 0.7 = Balanced (default)
- 1.5+ = Very creative
**Top-p** (0.0 - 1.0):
- Nucleus sampling threshold
- 0.9 = Default
- Lower values = More focused
**Repetition Penalty** (1.0 - 2.0):
- Reduces repetitive text
- 1.0 = No penalty
- 1.1 = Light penalty (default)
**Min-p** (0.0 - 1.0):
- Minimum probability threshold
- 0.0 = Disabled (default)
**Typical-p** (0.0 - 1.0):
- Typical sampling threshold
- 1.0 = Disabled (default)
---
## Model Management
### Supported Formats
DarkHal 2.0 supports multiple model formats:
| Format | Extension | Use Case | Pros | Cons |
|--------|-----------|----------|------|------|
| **GGUF** | `.gguf` | CPU/GPU hybrid | Fast loading, efficient | Limited to llama.cpp models |
| **SafeTensors** | `.safetensors` | HuggingFace models | Secure, fast | Larger file sizes |
| **PyTorch** | `.bin`, `.pt`, `.pth` | Research models | Flexible | Slower loading |
| **GPTQ** | `*gptq*.safetensors` | GPU inference | 4-bit quantized | GPU required |
| **AWQ** | `*awq*.safetensors` | GPU inference | Optimized quantization | GPU required |
| **EXL2** | `.exl2` | ExLlamaV2 | Very fast | Specific hardware needs |
### Model Library Tab
The Model Library provides comprehensive model management:
**Features**:
- Automatic scanning of model directories
- Metadata extraction (parameters, architecture)
- Search by name, type, or tags
- Size and modification date display
- One-click loading
**Using the Library**:
1. Set your models directory in Settings
2. Click "Scan" to index models
3. Use search box to filter
4. Double-click to load model
### Downloading Models
#### HuggingFace Tab
**Search and Browse**:
- Enter model name or organization
- Browse trending models
- Filter by task type
- View model cards
**Download Process**:
1. Enter model ID (e.g., "meta-llama/Llama-2-7b")
2. Click "Get File List"
3. Select files to download
4. Click "Start Download"
5. Monitor progress in Downloads tab
**File Selection**:
- Use checkboxes to select specific files
- "Select All" for complete model
- Size estimates shown for each file
- Automatic resume on failure
#### Downloads Tab
**Download Management**:
- Grouped display by model
- Individual file progress
- Speed and time estimates
- Pause/resume capability
- Automatic retry on failure
**Controls**:
- Collapse/expand model groups
- Cancel individual files
- Clear completed downloads
- Set bandwidth limits
---
## Agent Mode
### ⚠️ WARNING
Agent Mode grants the AI unrestricted system access. Only enable with trusted models and full understanding of risks.
### Enabling Agent Mode
1. Load any model
2. Check "🤖 Agent Mode (SYSTEM ACCESS)"
3. Confirm security warning
4. Agent mode indicator shows "ACTIVE"
### Capabilities
**System Control**:
- Execute shell commands
- Run PowerShell scripts
- Execute Bash commands
- Launch applications
**File Operations**:
- Read any file
- Write/create files
- Delete files
- List directories
**Application Control**:
- Open programs (Word, Notepad, etc.)
- Control mouse movement
- Send keyboard input
- Automate workflows
**Programming**:
- Execute Python code
- Run scripts
- Install packages
- Compile code
### Example Commands
**Opening Applications**:
```
"Open PowerShell"
"Launch Microsoft Word"
"Start notepad"
"Open calculator"
```
**File Operations**:
```
"List files in current directory"
"Create a file called test.txt with 'Hello World'"
"Read the contents of config.json"
"Delete temporary files"
```
**System Commands**:
```
"Show system information"
"Check disk space"
"List running processes"
"Create a new folder called Projects"
```
**Document Creation**:
```
"Open Word and create a document about Python"
"Create an Apache server setup guide"
"Write a bash script to backup files"
```
### Safety Guidelines
1. **Review Commands**: Always review AI-generated commands before execution
2. **Backup Data**: Keep backups before allowing file operations
3. **Limit Scope**: Use specific requests rather than broad permissions
4. **Monitor Activity**: Watch the output for unexpected behavior
5. **Disable When Done**: Turn off Agent Mode after use
---
## Advanced Features
### Chat Templates
Chat templates format conversations for different model architectures.
**Loading Templates**:
1. Click "Load" next to Chat Template dropdown
2. Select JSON file with templates
3. Choose template from dropdown
**Adding Custom Templates**:
1. Click "Add" button
2. Define template format
3. Set special tokens
4. Save to templates file
**Template Format**:
```json
{
"name": "llama3",
"template": "<|begin_of_text|>{% for message in messages %}...",
"bos_token": "<|begin_of_text|>",
"eos_token": "<|eot_id|>"
}
```
### Model Conversion
The Converter tab allows format transformation:
**Supported Conversions**:
- GGUF → SafeTensors
- SafeTensors → GGUF
- PyTorch → GGUF
- GPTQ → GGUF
**Conversion Process**:
1. Select source model
2. Choose target format
3. Set quantization options
4. Click "Convert"
5. Monitor progress
### MCP Server
Model Context Protocol enables remote access:
**Starting Server**:
1. Go to MCP tab
2. Configure port and settings
3. Click "Start Server"
4. Note the connection URL
**Connecting Clients**:
- Claude Desktop integration
- Remote control GUI
- Custom API clients
- Web interfaces
**Available Endpoints**:
- `/list_models` - Available models
- `/load_model` - Load specific model
- `/generate` - Text generation
- `/chat` - Conversation mode
### Chess Mode
Specialized interface for chess AI:
**Features**:
- FEN notation support
- Move generation
- Position evaluation
- Game analysis
**Using Chess Mode**:
1. Enable "Chess Mode" checkbox
2. Enter position in FEN format
3. Request move analysis
4. Get UCI format moves
---
## Troubleshooting
### Common Issues
#### Model Won't Load
**Symptoms**: Error message when loading model
**Solutions**:
- Verify file path is correct
- Check file isn't corrupted
- Ensure sufficient RAM/VRAM
- Try reducing GPU layers
- Lower context size
#### Out of Memory
**Symptoms**: Application crashes or freezes
**Solutions**:
- Use quantized models (4-bit/8-bit)
- Reduce context size
- Lower GPU layers
- Use CPU-only mode
- Close other applications
#### Slow Generation
**Symptoms**: Very slow text generation
**Solutions**:
- Enable GPU acceleration
- Increase GPU layers
- Use smaller models
- Reduce context size
- Check CPU/GPU usage
#### Download Failures
**Symptoms**: Downloads fail or hang
**Solutions**:
- Check internet connection
- Verify HuggingFace token
- Clear download cache
- Use VPN if blocked
- Try different mirror
### Error Messages
**"CUDA out of memory"**:
- Reduce GPU layers
- Use smaller batch size
- Enable memory efficient attention
- Use quantized model
**"Model file not found"**:
- Check file path
- Verify file exists
- Check permissions
- Try absolute path
**"Invalid model format"**:
- Verify file format
- Check model compatibility
- Update DarkHal
- Try conversion
**"Token limit exceeded"**:
- Reduce input length
- Lower max tokens
- Clear conversation history
- Increase context size
---
## Keyboard Shortcuts
### Global Shortcuts
| Shortcut | Action |
|----------|--------|
| `Ctrl+N` | New conversation |
| `Ctrl+O` | Open model |
| `Ctrl+S` | Save conversation |
| `Ctrl+Q` | Quit application |
| `F1` | Open help |
| `F5` | Refresh model list |
### Chat Interface
| Shortcut | Action |
|----------|--------|
| `Ctrl+Enter` | Send message |
| `Ctrl+L` | Clear output |
| `Ctrl+H` | Clear history |
| `Esc` | Stop generation |
| `Ctrl+C` | Copy selected text |
| `Ctrl+A` | Select all |
### Model Library
| Shortcut | Action |
|----------|--------|
| `Ctrl+F` | Focus search |
| `Enter` | Load selected model |
| `Delete` | Remove from library |
| `F5` | Rescan directory |
---
## Command Reference
### CLI Arguments
```bash
python main.py [options]
```
**Options**:
- `--gui` - Launch GUI mode (default)
- `--model PATH` - Model file path
- `--prompt TEXT` - Initial prompt
- `--stream` - Enable streaming
- `--n_ctx N` - Context size
- `--n_gpu_layers N` - GPU layers
- `--lora PATH` - LoRA adapter path
### Configuration Files
**settings.json**:
```json
{
"paths": {
"models_directory": "./models",
"download_directory": "./downloads"
},
"model_settings": {
"default_n_ctx": 4096,
"default_n_gpu_layers": 0,
"stream_by_default": true,
"temperature": 0.7,
"top_p": 0.9,
"repetition_penalty": 1.1
}
}
```
**HUGGINGFACE.env**:
```
HF_API_KEY=your_token_here
HF_HOME=./models/huggingface
```
---
## FAQ
### General Questions
**Q: What models work with DarkHal?**
A: Any model in GGUF, SafeTensors, PyTorch, GPTQ, AWQ, or EXL2 format. Most HuggingFace models are compatible.
**Q: How much RAM do I need?**
A: Depends on model size. 7B models need ~8GB, 13B need ~16GB, 70B need ~64GB. Quantization reduces requirements.
**Q: Can I run without GPU?**
A: Yes, CPU-only mode works but is slower. Use GGUF models with 0 GPU layers.
**Q: Is my data private?**
A: Yes, all processing is local. No data is sent to external servers unless using HuggingFace downloads.
### Model Questions
**Q: What's the difference between formats?**
A: GGUF is optimized for CPU/GPU hybrid. SafeTensors is HuggingFace standard. GPTQ/AWQ/EXL2 are quantized for GPU.
**Q: How do I choose quantization?**
A: 4-bit saves most memory with slight quality loss. 8-bit balances quality and size. None uses full precision.
**Q: Why is generation slow?**
A: Check GPU usage, reduce context size, use quantized models, or enable more GPU layers.
**Q: Can I use multiple models?**
A: One model at a time in current version. Switch models by loading different ones.
### Agent Mode Questions
**Q: Is Agent Mode safe?**
A: Agent Mode grants full system access. Only use with trusted models and review commands.
**Q: What can Agent Mode do?**
A: Execute any system command, control applications, manage files, run code, automate tasks.
**Q: How do I limit Agent Mode?**
A: Currently all-or-nothing. Future versions will have granular permissions.
**Q: Can Agent Mode access internet?**
A: Yes, through system commands like curl or wget, and Python's requests library.
### Troubleshooting Questions
**Q: Download keeps failing?**
A: Check internet, verify HF token, try VPN, clear cache, or download manually.
**Q: Model won't load?**
A: Verify path, check format, ensure enough memory, try different quantization.
**Q: Getting CUDA errors?**
A: Update GPU drivers, check CUDA version, reduce GPU layers, or use CPU mode.
**Q: Application crashes?**
A: Check error logs, reduce memory usage, update dependencies, file bug report.
---
## Support
### Getting Help
**Documentation**: [https://darkhal.readthedocs.io](https://darkhal.readthedocs.io)
**GitHub Issues**: [https://github.com/darkhal/issues](https://github.com/darkhal/issues)
**Discussions**: [https://github.com/darkhal/discussions](https://github.com/darkhal/discussions)
**Email Support**: support@darkhal.ai
### Reporting Bugs
Include:
1. System information (OS, GPU, RAM)
2. Model details (format, size, source)
3. Error messages and logs
4. Steps to reproduce
5. Screenshots if applicable
### Contributing
We welcome contributions! See CONTRIBUTING.md for guidelines.
---
## Appendices
### A. Model Compatibility Matrix
| Model Family | GGUF | SafeTensors | GPTQ | AWQ | EXL2 |
|--------------|------|-------------|------|-----|------|
| Llama 2/3 | ✅ | ✅ | ✅ | ✅ | ✅ |
| Mistral | ✅ | ✅ | ✅ | ✅ | ✅ |
| Mixtral | ✅ | ✅ | ✅ | ✅ | ❌ |
| Qwen | ✅ | ✅ | ✅ | ✅ | ✅ |
| Yi | ✅ | ✅ | ✅ | ✅ | ✅ |
| Gemma | ✅ | ✅ | ❌ | ❌ | ❌ |
### B. Performance Benchmarks
| Model | Format | GPU | Speed (tok/s) | Memory |
|-------|--------|-----|---------------|---------|
| Llama-7B | GGUF Q4 | RTX 3060 | 45 | 4.5GB |
| Llama-13B | GGUF Q4 | RTX 3090 | 35 | 8.5GB |
| Mistral-7B | GPTQ | RTX 4090 | 65 | 5.0GB |
| Mixtral-8x7B | AWQ | A100 | 25 | 24GB |
### C. Glossary
**Context Size**: Maximum tokens the model can process at once
**GPU Layers**: Model layers offloaded to GPU for acceleration
**Quantization**: Reducing model precision to save memory
**LoRA**: Low-Rank Adaptation for model fine-tuning
**Token**: Basic unit of text (roughly 0.75 words)
**VRAM**: Video RAM on graphics card
**Streaming**: Showing text as it's generated
**KV Cache**: Key-value cache for faster inference
---
*DarkHal 2.0 User Manual - Version 2.0.0*
*Last Updated: January 2025*

441
README.md Normal file
View File

@@ -0,0 +1,441 @@
If you would like to join, please email ssSnake@darkHal.org or sshhh@setecastronomy.gg. Please include a short bio and a link to 2-3 projects you have worked on. If you have no experience and would like to join, email me at ssSnake@darkHal.org and make your case, just because you dont have experience doesn't mean you dont have skills.
# DarkHal 2.0 🤖
<div align="center">
![DarkHal 2.0 Logo](assets/logo.png)
**Advanced AI Model Management & Training Platform**
[![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Platform](https://img.shields.io/badge/platform-Windows%20%7C%20macOS%20%7C%20Linux-lightgrey)](https://github.com/your-org/darkhal)
*Comprehensive solution for AI model downloading, management, training, and inference*
</div>
## ⚠️ DISCLAIMER
**This software is provided "as is" without any warranties or guarantees. The user assumes all responsibility for the use of this software and any consequences that may arise from its use. The developers are not liable for any damages, data loss, or other issues that may occur.**
---
## 🚀 Features
### 🔽 **Advanced Download Management**
- **Grouped Downloads**: Organize multi-file model downloads with collapsible widgets
- **HuggingFace Integration**: Direct downloading from HuggingFace Hub with authentication
- **Smart File Selection**: Choose specific files from model repositories
- **Progress Tracking**: Real-time download progress with speed monitoring
- **Resume Support**: Automatically resume interrupted downloads
### 🤖 **AI Model Management**
- **Model Library**: Intelligent scanning and indexing of local model files
- **Multi-Format Support**: GGUF, SafeTensors, PyTorch, ONNX, and more
- **Metadata Extraction**: Automatic detection of model parameters and tags
- **Quick Search**: Fast searching by name, type, size, and tags
- **GPU Acceleration**: CUDA, ROCm, Metal, and Intel GPU support
### 🎛️ **Inference & Chat**
- **Local Inference**: Run models locally without internet connection
- **Streaming Support**: Real-time text generation with streaming output
- **Chat Interface**: Interactive conversation mode with context memory
- **Parameter Control**: Adjustable temperature, context size, and token limits
- **LoRA Support**: Load and apply LoRA adapters to base models
### 🌐 **Remote Control**
- **MCP Server**: Model Context Protocol server for remote management
- **Remote GUI**: Standalone application for controlling server remotely
- **Claude Integration**: Direct integration with Claude Desktop
- **API Access**: RESTful API for programmatic control
### ⚙️ **System Integration**
- **Auto-Setup**: Automatic dependency installation and GPU detection
- **Windows Integration**: MSI installer with Windows-specific optimizations
- **Cross-Platform**: Native support for Windows, macOS, and Linux
- **Chocolatey Support**: Automated Windows dependency management
---
## 📦 Installation
### 🔧 **Quick Install (Recommended)**
#### Windows (Coming Soon)
```bash
# Download and run the MSI installer
darkhal-2.0-setup.msi
# Or install via Chocolatey
choco install darkhal
```
#### macOS/Linux (cooming soon)
```bash
# Install via pip
pip install darkhal
# Or from source
git clone https://github.com/your-org/darkhal.git
cd darkhal
pip install -e .
```
### 🛠️ **Manual Installation**
1. **Clone the repository**
```bash
git clone https://github.com/your-org/darkhal.git
cd darkhal
```
2. **Install dependencies**
```bash
# Basic installation
pip install -r requirements.txt
# With GPU support
pip install -r requirements.txt darkhal[gpu]
# With audio/whisper support
pip install -r requirements.txt darkhal[audio]
```
3. **Run dependency installer** (Windows)
```bash
python windows_dependency_installer.py
```
4. **Configure HuggingFace** (optional)
```bash
# Create HUGGINGFACE.env file
echo "HF_API_KEY=your_token_here" > HUGGINGFACE.env
```
---
## 🎯 Quick Start
### 🖥️ **GUI Mode (Default)**
```bash
# Launch with splash screen
python main.py
# Or use the installed command
darkhal
```
### 💻 **CLI Mode**
Full CLI command list coming soon. Right now the agent mode works best from a powershell, shell or bash.
```bash
# Interactive chat
python main.py --model path/to/model.gguf
# Single prompt
python main.py --model path/to/model.gguf --prompt "Your question here"
# With GPU acceleration
python main.py --model path/to/model.gguf --n_gpu_layers 32
```
### 🌐 **Remote Control**
```bash
# Launch remote control GUI
python remotecontrol.py
# Or use the installed command
darkhal-remote
```
### 🔌 **MCP Server**
```bash
# Start MCP server
python mcp_server.py
# Or use the installed command
darkhal-mcp
```
---
## 📋 Usage Guide
### 1. **First Launch**
- DarkHal 2.0 will show a splash screen with disclaimers
- Configure your models directory in Settings
- Set up HuggingFace authentication if needed
- Scan your model library for automatic indexing
### 2. **Downloading Models**
- Go to the **HuggingFace** tab
- Search for models by name or tags
- Select desired files from multi-file models
- Downloads are organized in collapsible groups
- Monitor progress in the **Downloads** tab
### 3. **Managing Models**
- Use the **Model Library** tab to browse local models
- Search by name, file type, or tags
- View detailed metadata and statistics
- Load models directly from the library
### 4. **Running Inference**
- Select a model using "Browse" or the model library
- Configure context size and GPU layers
- Enter your prompt and click "Generate"
- Toggle streaming for real-time output
- Use chat mode for conversations
### 5. **Remote Operations**
- Start the MCP server from the main application
- Launch the remote control GUI
- Connect to the server and manage models remotely
- Integrate with Claude Desktop for enhanced AI workflows
---
## ⚙️ Configuration
### 📁 **Settings Files**
- `settings.json` - Main application settings
- `HUGGINGFACE.env` - HuggingFace API credentials
- `mcp_config.json` - MCP server configuration
- `.model_index.json` - Model library index
### 🎛️ **Key Settings**
```json
{
"paths": {
"models_directory": "./models",
"download_directory": "./downloads"
},
"model_settings": {
"default_n_ctx": 4096,
"default_n_gpu_layers": 0,
"stream_by_default": true
},
"download_settings": {
"max_concurrent_downloads": 3,
"speed_limit_mbps": 0
}
}
```
### 🖥️ **GPU Configuration**
DarkHal 2.0 automatically detects and optimizes for:
- **NVIDIA CUDA** (Windows/Linux)
- **AMD ROCm** (Linux)
- **Apple Metal** (macOS)
- **Intel GPU** (Windows/Linux)
---
## 🔌 API Reference
### **MCP Tools**
- `list_models` - Get available models
- `load_model` - Load a model with parameters
- `generate_text` - Generate text with prompt
- `get_system_info` - Get system capabilities
### **Claude Integration**
```json
{
"mcpServers": {
"darkhal": {
"command": "python",
"args": ["path/to/mcp_server.py"]
}
}
}
```
---
## 🏗️ Architecture
### **Core Components**
```
DarkHal 2.0/
├── main.py # Main GUI application
├── splash_screen.py # Startup splash screen
├── remotecontrol.py # Remote control GUI
├── mcp_server.py # MCP protocol server
├── settings_manager.py # Configuration management
├── model_library.py # Model indexing & search
├── grouped_download_* # Advanced download system
└── assets/ # Icons and resources
```
### **Data Flow**
1. **User Interface** → Settings Manager → Model Operations
2. **Download Manager** → HuggingFace API → Local Storage
3. **Model Library** → File Scanner → Metadata Extractor
4. **MCP Server** → llama.cpp → Text Generation
5. **Remote Control** → MCP Client → Server Commands
---
## 🛠️ Development
### **Building from Source**
```bash
# Clone repository
git clone https://github.com/your-org/darkhal.git
cd darkhal
# Install development dependencies
pip install -e .[dev]
# Run tests
pytest tests/
# Build distribution
python setup.py sdist bdist_wheel
# Build MSI installer (Windows)
python build_installer.py
```
### **Contributing**
1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests and documentation
5. Submit a pull request
### **Code Style**
- Follow PEP 8 guidelines
- Use type hints where possible
- Add docstrings to all functions
- Include comprehensive error handling
---
## 📊 Performance & Requirements
### **System Requirements**
| Component | Minimum | Recommended |
|-----------|---------|-------------|
| **OS** | Windows 10, macOS 10.14, Ubuntu 18.04 | Latest versions |
| **Python** | 3.8+ | 3.11+ |
| **RAM** | 8 GB | 16+ GB |
| **Storage** | 10 GB free | 100+ GB for models |
| **GPU** | Optional | 8+ GB VRAM |
### **Model Performance**
| Model Size | RAM Usage | GPU VRAM | Inference Speed |
|------------|-----------|----------|----------------|
| **7B Q4** | 4-6 GB | 4-6 GB | 10-50 tokens/sec |
| **13B Q4** | 8-10 GB | 8-10 GB | 5-25 tokens/sec |
| **70B Q4** | 40-50 GB | 40+ GB | 1-10 tokens/sec |
---
## 🔍 Troubleshooting
### **Common Issues**
#### Installation Problems
```bash
# Missing dependencies
pip install --upgrade pip setuptools wheel
pip install -r requirements.txt
# CUDA issues
pip install torch --index-url https://download.pytorch.org/whl/cu121
# Permission errors (Windows)
# Run as Administrator or use --user flag
pip install --user darkhal
```
#### Runtime Errors
```bash
# Model loading fails
# Check file permissions and disk space
# Verify model file integrity
# Reduce context size or GPU layers
# Download issues
# Check internet connection
# Verify HuggingFace token
# Clear download cache
```
#### Performance Issues
```bash
# Slow inference
# Enable GPU acceleration
# Reduce context size
# Use quantized models
# Close other applications
```
### **Getting Help**
- 📖 Check the [Documentation](https://darkhal.readthedocs.io/)
- 🐛 Report [Issues](https://github.com/your-org/darkhal/issues)
- 💬 Join [Discussions](https://github.com/your-org/darkhal/discussions)
- 📧 Contact [Support](mailto:support@seteclabs.com)
---
## 📄 License & Legal
### **License**
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
### **Third-Party Components**
- **llama.cpp** - MIT License
- **HuggingFace Transformers** - Apache 2.0
- **Tkinter** - Python Software Foundation License
- **Pillow** - PIL Software License
### **Copyright**
© 2025 Setec Labs. All rights reserved.
**Created by ssSnake**
---
## 🚀 Roadmap
### **Upcoming Features**
- [ ] **Multi-Model Support** - Run multiple models simultaneously
- [ ] **Fine-Tuning Interface** - Built-in model training capabilities
- [ ] **Plugin System** - Extensible architecture for custom tools
- [ ] **Cloud Integration** - AWS, Azure, GCP deployment options
- [ ] **Advanced Analytics** - Performance monitoring and analytics
- [ ] **Voice Interface** - Speech-to-text and text-to-speech
- [ ] **Docker Support** - Containerized deployment options
- [ ] **Model Marketplace** - Community model sharing platform
### **Version History**
- **v2.0.0** - Major rewrite with advanced features
- **v1.5.0** - Remote control and MCP integration
- **v1.0.0** - Initial release with basic functionality
---
## 🙏 Acknowledgments
Special thanks to:
- **llama.cpp team** for the excellent inference engine
- **HuggingFace** for the model hub and libraries
- **Anthropic** for the MCP protocol specification
- **Open source community** for various dependencies and tools
---
<div align="center">
**⭐ Star this repository if you find it useful!**
[**Download**](https://github.com/your-org/darkhal/releases) • [**Documentation**](https://darkhal.readthedocs.io/) • [**Issues**](https://github.com/your-org/darkhal/issues) • [**Discussions**](https://github.com/your-org/darkhal/discussions)
</div>

287
REMOTE_CONTROL_README.md Normal file
View File

@@ -0,0 +1,287 @@
# LLM_Train Remote Control
A standalone GUI application for remotely controlling and managing LLM_Train MCP servers. This allows you to connect to a running MCP server from a separate application and perform model management and inference operations remotely.
## Features
### 🔗 **Connection Management**
- Connect to local or remote MCP servers
- Real-time connection status monitoring
- Automatic reconnection handling
- Server path configuration and browsing
### 🤖 **Model Management**
- List all available models from the server
- Remote model loading with configuration
- Context size and GPU layer settings
- Current model status display
- Model unloading capabilities
### 💬 **Inference Interface**
- Text generation with configurable parameters
- Chat mode for conversational interactions
- Adjustable temperature and max tokens
- Real-time output display
- Chat history management
### 📊 **System Monitoring**
- Real-time system information display
- GPU acceleration status
- Platform and architecture details
- Performance metrics
### 📝 **Logging & Debugging**
- Comprehensive operation logging
- Error tracking and display
- Log saving functionality
- Connection event monitoring
## Quick Start
### Option 1: Batch File (Windows)
```bash
# Double-click launch_remote_control.bat
```
### Option 2: Direct Python Execution
```bash
python remotecontrol.py
```
## Requirements
- **Python 3.7+** with tkinter support
- **Running MCP Server** (from main LLM_Train application)
- **Network access** to the MCP server (if remote)
## Usage Guide
### 1. Starting the Remote Control
1. **Launch the application** using one of the methods above
2. The GUI will open with the connection panel at the top
### 2. Connecting to MCP Server
1. **Specify server path**:
- Default: `mcp_server.py` (local)
- Browse to select different server script
- Can be local file or network path
2. **Click "Connect"** button
3. **Connection status** will show:
- 🔴 **Disconnected** (Red): Not connected
- 🟢 **Connected** (Green): Successfully connected
### 3. Model Management
**Models Tab:**
- **View available models**: Automatically populated on connection
- **Select model**: Click on model in the list
- **Configure parameters**:
- Context Size: 512 - 32768 tokens
- GPU Layers: 0 - 100 layers
- **Load model**: Click "Load Model" button
- **Monitor status**: Current model display shows loaded model
### 4. Text Generation
**Inference Tab:**
- **Enter prompt**: Type or paste text in the input area
- **Configure generation**:
- Max Tokens: 1 - 8192
- Temperature: 0.0 - 2.0
- **Generate**: Click "Generate" button
- **View output**: Results appear in the output area
- **Chat mode**: Enable for conversational interface
### 5. System Information
**System Tab:**
- **View system details**: Platform, architecture, acceleration
- **Monitor GPU status**: CUDA, ROCm, Metal availability
- **Check performance**: Current acceleration method
- **Refresh**: Update information in real-time
### 6. Logging
**Log Tab:**
- **Monitor operations**: All actions are logged with timestamps
- **Error tracking**: Errors are highlighted and detailed
- **Save logs**: Export logs to text file for debugging
- **Clear logs**: Reset log display
## Advanced Configuration
### Server Connection Options
**Local Server:**
```
Server Path: mcp_server.py
```
**Custom Path:**
```
Server Path: C:\path\to\your\mcp_server.py
```
**Network Server (if supported):**
```
Server Path: \\network\path\mcp_server.py
```
### Model Configuration
**High-Performance Setup:**
- Context Size: 8192+
- GPU Layers: Maximum supported
- Temperature: 0.1-0.3 for focused responses
**Balanced Setup:**
- Context Size: 4096
- GPU Layers: Auto-detected optimum
- Temperature: 0.7 for creative responses
**CPU-Only Setup:**
- Context Size: 2048
- GPU Layers: 0
- Temperature: 0.5-1.0
### Generation Parameters
**Creative Writing:**
- Max Tokens: 1024+
- Temperature: 0.8-1.2
**Code Generation:**
- Max Tokens: 512
- Temperature: 0.1-0.3
**Question Answering:**
- Max Tokens: 256
- Temperature: 0.3-0.7
## Troubleshooting
### Connection Issues
**"Server file not found"**
- Verify the server path is correct
- Ensure the MCP server file exists
- Check file permissions
**"Connection failed"**
- Ensure the MCP server is not already running
- Check if the server script is executable
- Verify Python dependencies are installed
**"Disconnected unexpectedly"**
- Check server logs for errors
- Verify system resources are available
- Restart both applications
### Model Loading Issues
**"No models found"**
- Ensure model library is configured in main application
- Verify model files exist in specified directories
- Check library settings and scan depth
**"Failed to load model"**
- Verify model file is not corrupted
- Check available system memory
- Reduce context size or GPU layers
**"Out of memory"**
- Reduce context size
- Lower GPU layers
- Close other applications
### Generation Issues
**"No model loaded"**
- Load a model first using the Models tab
- Verify model loaded successfully
- Check current model display
**"Generation timeout"**
- Reduce max tokens
- Simplify the prompt
- Check system resources
**"Invalid parameters"**
- Verify temperature is between 0.0-2.0
- Ensure max tokens is reasonable
- Check prompt is not empty
## Technical Details
### MCP Protocol
- Uses JSON-RPC 2.0 over stdin/stdout
- Asynchronous request/response handling
- Automatic request ID management
- Error handling and recovery
### Threading Model
- Main UI thread for interface
- AsyncIO event loop for MCP communication
- Background threads for I/O operations
- Thread-safe callback system
### Security Considerations
- Local process communication only
- No network ports exposed
- Input validation on all parameters
- Error sanitization in logs
## Integration Examples
### Automated Workflows
```python
# Example: Batch text generation
prompts = ["Explain AI", "Code a function", "Write a story"]
for prompt in prompts:
# Use remote control to generate text
# Save results to files
```
### API Integration
```python
# Example: Integration with other tools
remote_control = RemoteControlClient()
remote_control.connect("mcp_server.py")
result = remote_control.generate("Your prompt here")
```
### Monitoring Scripts
```python
# Example: System monitoring
while True:
system_info = remote_control.get_system_info()
log_performance_metrics(system_info)
time.sleep(60)
```
## Support and Development
### Extending Functionality
- Add new MCP tool integrations
- Implement custom inference modes
- Create automation scripts
- Build monitoring dashboards
### Contributing
- Follow Python coding standards
- Add comprehensive logging
- Include error handling
- Write unit tests
### Reporting Issues
- Include full log output
- Specify system configuration
- Provide reproduction steps
- Attach relevant files
## License
This remote control application is part of the LLM_Train project and follows the same licensing terms as the main application.

50
__spy.py Normal file
View File

@@ -0,0 +1,50 @@
#!/usr/bin/env python3
"""
__spy.py
Lightweight global announcer for the currently loaded model.
Usage:
- Call set_model(model_name, model_obj, **params) when a model is loaded.
- Retrieve with get_model(), get_model_name(), or get_info() anywhere.
"""
from __future__ import annotations
from dataclasses import dataclass, asdict
from typing import Any, Dict, Optional
import threading
@dataclass
class SpyData:
model_name: str
model: Any
params: Dict[str, Any]
def to_dict(self) -> Dict[str, Any]:
d = asdict(self)
# Avoid serializing the raw model object
d["model"] = repr(self.model)
return d
_lock = threading.RLock()
_current: Optional[SpyData] = None
def set_model(model_name: str, model: Any, **params: Any) -> None:
"""Announce the current model and its load parameters."""
global _current
with _lock:
_current = SpyData(model_name=model_name, model=model, params=dict(params or {}))
def get_model() -> Optional[Any]:
"""Return the current model object, if any."""
with _lock:
return _current.model if _current else None
def get_model_name() -> Optional[str]:
"""Return the current model name, if any."""
with _lock:
return _current.model_name if _current else None
def get_info() -> Optional[SpyData]:
"""Return the full SpyData object, if any."""
with _lock:
return _current

234
agent_debug_tracer.py Normal file
View File

@@ -0,0 +1,234 @@
#!/usr/bin/env python3
"""
Agent Debug and Trace System for DarkAgent
Custom logging and monitoring system that tracks agent lifecycle without using sys.settrace()
"""
import time
import threading
import queue
import json
from datetime import datetime
from pathlib import Path
from typing import Dict, Any, Optional, List
class AgentDebugTracer:
"""Custom debug and trace system for DarkAgent monitoring."""
def __init__(self, log_file: str = "agent_debug.log"):
self.log_file = Path(log_file)
self.trace_queue = queue.Queue()
self.is_running = False
self.logger_thread = None
self.start_time = time.time()
# Agent state tracking
self.agent_states = {}
self.event_history = []
self.performance_metrics = {
"total_messages": 0,
"successful_responses": 0,
"errors": 0,
"average_response_time": 0.0,
"response_times": []
}
# Initialize log file
self._init_log_file()
def _init_log_file(self):
"""Initialize the log file with header."""
try:
with open(self.log_file, 'w') as f:
f.write(f"=== DarkAgent Debug Trace Started: {datetime.now()} ===\n")
f.write(f"Application Launch Time: {self.start_time}\n")
f.write("=" * 60 + "\n\n")
except Exception as e:
print(f"[AGENT_DEBUG] Failed to initialize log file: {e}")
def start_monitoring(self):
"""Start the debug monitoring system."""
if not self.is_running:
self.is_running = True
self.logger_thread = threading.Thread(target=self._logger_worker, daemon=True)
self.logger_thread.start()
self.trace("SYSTEM", "Debug monitoring started")
def stop_monitoring(self):
"""Stop the debug monitoring system."""
if self.is_running:
self.trace("SYSTEM", "Debug monitoring stopping")
self.is_running = False
if self.logger_thread and self.logger_thread.is_alive():
self.logger_thread.join(timeout=1.0)
def trace(self, category: str, message: str, data: Dict[str, Any] = None):
"""Add a trace entry."""
if not self.is_running:
return
timestamp = time.time()
elapsed = timestamp - self.start_time
entry = {
"timestamp": timestamp,
"elapsed": elapsed,
"datetime": datetime.now().isoformat(),
"category": category,
"message": message,
"data": data or {},
"thread": threading.current_thread().name
}
try:
self.trace_queue.put_nowait(entry)
self.event_history.append(entry)
# Keep history manageable
if len(self.event_history) > 1000:
self.event_history = self.event_history[-500:]
except queue.Full:
pass # Drop trace if queue is full
def _logger_worker(self):
"""Background thread that writes trace entries to file."""
while self.is_running:
try:
entry = self.trace_queue.get(timeout=1.0)
self._write_trace_entry(entry)
except queue.Empty:
continue
except Exception as e:
print(f"[AGENT_DEBUG] Logger error: {e}")
def _write_trace_entry(self, entry: Dict[str, Any]):
"""Write a trace entry to the log file."""
try:
with open(self.log_file, 'a') as f:
formatted_time = f"{entry['elapsed']:8.3f}s"
thread_info = f"[{entry['thread']}]" if entry['thread'] != 'MainThread' else ""
f.write(f"{formatted_time} [{entry['category']}]{thread_info} {entry['message']}")
if entry['data']:
f.write(f" | Data: {json.dumps(entry['data'], indent=None)}")
f.write("\n")
f.flush()
except Exception as e:
print(f"[AGENT_DEBUG] Write error: {e}")
# Agent-specific monitoring methods
def agent_startup(self, agent_name: str, config: Dict[str, Any] = None):
"""Track agent startup."""
self.agent_states[agent_name] = {
"status": "starting",
"start_time": time.time(),
"config": config or {}
}
self.trace("AGENT_STARTUP", f"Agent {agent_name} starting", {"config": config})
def agent_ready(self, agent_name: str):
"""Track agent ready state."""
if agent_name in self.agent_states:
self.agent_states[agent_name]["status"] = "ready"
startup_time = time.time() - self.agent_states[agent_name]["start_time"]
self.trace("AGENT_READY", f"Agent {agent_name} ready", {"startup_time": startup_time})
def agent_message_start(self, agent_name: str, message: str, message_id: str = None):
"""Track start of message processing."""
self.performance_metrics["total_messages"] += 1
self.trace("AGENT_MESSAGE_START", f"Agent {agent_name} processing message", {
"message_id": message_id,
"message_preview": message[:100] + "..." if len(message) > 100 else message
})
return time.time() # Return start time for response time calculation
def agent_message_end(self, agent_name: str, message_id: str, start_time: float, success: bool = True, error: str = None):
"""Track end of message processing."""
response_time = time.time() - start_time
self.performance_metrics["response_times"].append(response_time)
if success:
self.performance_metrics["successful_responses"] += 1
self.trace("AGENT_MESSAGE_SUCCESS", f"Agent {agent_name} completed message", {
"message_id": message_id,
"response_time": response_time
})
else:
self.performance_metrics["errors"] += 1
self.trace("AGENT_MESSAGE_ERROR", f"Agent {agent_name} message failed", {
"message_id": message_id,
"response_time": response_time,
"error": error
})
# Update average response time
if self.performance_metrics["response_times"]:
self.performance_metrics["average_response_time"] = sum(self.performance_metrics["response_times"]) / len(self.performance_metrics["response_times"])
def agent_shutdown(self, agent_name: str):
"""Track agent shutdown."""
if agent_name in self.agent_states:
self.agent_states[agent_name]["status"] = "shutdown"
uptime = time.time() - self.agent_states[agent_name]["start_time"]
self.trace("AGENT_SHUTDOWN", f"Agent {agent_name} shutting down", {"uptime": uptime})
def agent_error(self, agent_name: str, error: str, context: Dict[str, Any] = None):
"""Track agent errors."""
self.performance_metrics["errors"] += 1
self.trace("AGENT_ERROR", f"Agent {agent_name} error: {error}", context)
def ui_event(self, event_type: str, details: Dict[str, Any] = None):
"""Track UI events related to agent."""
self.trace("UI_EVENT", event_type, details)
def model_event(self, event_type: str, model_info: Dict[str, Any] = None):
"""Track model loading/unloading events."""
self.trace("MODEL_EVENT", event_type, model_info)
def get_performance_summary(self) -> Dict[str, Any]:
"""Get current performance metrics."""
return {
"uptime": time.time() - self.start_time,
"agent_states": self.agent_states,
"performance": self.performance_metrics,
"recent_events": self.event_history[-10:] if self.event_history else []
}
def print_summary(self):
"""Print performance summary to console."""
summary = self.get_performance_summary()
print("\n=== DarkAgent Debug Summary ===")
print(f"Uptime: {summary['uptime']:.2f}s")
print(f"Total Messages: {summary['performance']['total_messages']}")
print(f"Successful Responses: {summary['performance']['successful_responses']}")
print(f"Errors: {summary['performance']['errors']}")
print(f"Average Response Time: {summary['performance']['average_response_time']:.3f}s")
print(f"Active Agents: {len([a for a in summary['agent_states'].values() if a['status'] == 'ready'])}")
print("=" * 31)
# Global tracer instance
_global_tracer: Optional[AgentDebugTracer] = None
def get_tracer() -> AgentDebugTracer:
"""Get the global tracer instance."""
global _global_tracer
if _global_tracer is None:
_global_tracer = AgentDebugTracer()
_global_tracer.start_monitoring()
return _global_tracer
def trace(category: str, message: str, data: Dict[str, Any] = None):
"""Convenience function for tracing."""
get_tracer().trace(category, message, data)
def shutdown_tracer():
"""Shutdown the global tracer."""
global _global_tracer
if _global_tracer:
_global_tracer.stop_monitoring()
_global_tracer = None

85
agent_dhal/__init__.py Normal file
View File

@@ -0,0 +1,85 @@
#!/usr/bin/env python3
"""
AgentDhal - Complete AI Agent Framework for DarkHal 2.0
A comprehensive agent framework providing:
- Multi-agent conversation capabilities
- Agent orchestration and team management
- Tool integration and function calling
- Model context management
- Memory and state management
- Customizable agent behaviors
Legal Attribution:
This software is based on Microsoft AutoGen (https://github.com/microsoft/autogen)
Licensed under MIT License. AgentDhal is a derivative work with
modifications and extensions for the DarkHal project.
Copyright (c) 2025 DarkHal Project
"""
__version__ = "1.0.0"
__author__ = "DarkHal Project (based on Microsoft AutoGen)"
# Import core AgentDhal components
from .agentdhal_core import (
Agent,
AgentId,
AgentRuntime,
SingleThreadedAgentRuntime,
RoutedAgent,
MessageContext,
DefaultTopicId,
message_handler,
default_subscription,
BaseAgent,
AgentType,
TopicId,
Subscription
)
# Import Dhal - our primary AI agent
from .hal import Dhal, DhalConfig, create_dhal
# Import other AgentDhal components (available but not primary focus)
try:
from .agentdhal_agentchat import (
AssistantAgent,
UserProxyAgent,
ChatAgent,
Team
)
except ImportError:
# Graceful fallback if agentchat modules have issues
AssistantAgent = None
UserProxyAgent = None
ChatAgent = None
Team = None
__all__ = [
# Core framework
"Agent",
"AgentId",
"AgentRuntime",
"SingleThreadedAgentRuntime",
"RoutedAgent",
"MessageContext",
"DefaultTopicId",
"message_handler",
"default_subscription",
"BaseAgent",
"AgentType",
"TopicId",
"Subscription",
# Primary Hal Agent
"Hal",
"HalConfig",
"create_hal",
# Additional Agent Components (if available)
"AssistantAgent",
"UserProxyAgent",
"ChatAgent",
"Team"
]

14
agent_dhal/agentdhal_agentchat/__init__.py Normal file
View File

@@ -0,0 +1,14 @@
"""
This module provides the main entry point for the agentdhal_agentchat package.
It includes logger names for trace and event logs, and retrieves the package version.
"""
import importlib.metadata
TRACE_LOGGER_NAME = "agentdhal_agentchat"
"""Logger name for trace logs."""
EVENT_LOGGER_NAME = "agentdhal_agentchat.events"
"""Logger name for event logs."""
__version__ = importlib.metadata.version("agentdhal_agentchat")

25
agent_dhal/agentdhal_agentchat/agents/__init__.py Normal file
View File

@@ -0,0 +1,25 @@
"""
This module initializes various pre-defined agents provided by the package.
BaseChatAgent is the base class for all agents in AgentChat.
"""
from ._assistant_agent import AssistantAgent
from ._base_chat_agent import BaseChatAgent
from ._code_executor_agent import ApprovalFuncType, ApprovalRequest, ApprovalResponse, CodeExecutorAgent
from ._message_filter_agent import MessageFilterAgent, MessageFilterConfig, PerSourceFilter
from ._society_of_mind_agent import SocietyOfMindAgent
from ._user_proxy_agent import UserProxyAgent
__all__ = [
"BaseChatAgent",
"AssistantAgent",
"CodeExecutorAgent",
"SocietyOfMindAgent",
"UserProxyAgent",
"MessageFilterAgent",
"MessageFilterConfig",
"PerSourceFilter",
"ApprovalRequest",
"ApprovalResponse",
"ApprovalFuncType",
]

1699
agent_dhal/agentdhal_agentchat/agents/_assistant_agent.py Normal file
View File

File diff suppressed because it is too large Load Diff

245
agent_dhal/agentdhal_agentchat/agents/_base_chat_agent.py Normal file
View File

@@ -0,0 +1,245 @@
from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, List, Mapping, Sequence
from agentdhal_core import CancellationToken, ComponentBase, trace_create_agent_span, trace_invoke_agent_span
from pydantic import BaseModel
from ..base import ChatAgent, Response, TaskResult
from ..messages import (
BaseAgentEvent,
BaseChatMessage,
ModelClientStreamingChunkEvent,
TextMessage,
)
from ..state import BaseState
class BaseChatAgent(ChatAgent, ABC, ComponentBase[BaseModel]):
"""Base class for a chat agent.
This abstract class provides a base implementation for a :class:`ChatAgent`.
To create a new chat agent, subclass this class and implement the
:meth:`on_messages`, :meth:`on_reset`, and :attr:`produced_message_types`.
If streaming is required, also implement the :meth:`on_messages_stream` method.
An agent is considered stateful and maintains its state between calls to
the :meth:`on_messages` or :meth:`on_messages_stream` methods.
The agent should store its state in the
agent instance. The agent should also implement the :meth:`on_reset` method
to reset the agent to its initialization state.
.. note::
The caller should only pass the new messages to the agent on each call
to the :meth:`on_messages` or :meth:`on_messages_stream` method.
Do not pass the entire conversation history to the agent on each call.
This design principle must be followed when creating a new agent.
"""
component_type = "agent"
def __init__(self, name: str, description: str) -> None:
"""Initialize the agent with a name and description."""
with trace_create_agent_span(
agent_name=name,
agent_description=description,
):
self._name = name
if self._name.isidentifier() is False:
raise ValueError("The agent name must be a valid Python identifier.")
self._description = description
@property
def name(self) -> str:
"""The name of the agent. This is used by team to uniquely identify
the agent. It should be unique within the team."""
return self._name
@property
def description(self) -> str:
"""The description of the agent. This is used by team to
make decisions about which agents to use. The description should
describe the agent's capabilities and how to interact with it."""
return self._description
@property
@abstractmethod
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
"""The types of messages that the agent produces in the
:attr:`Response.chat_message` field. They must be :class:`BaseChatMessage` types."""
...
@abstractmethod
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
"""Handles incoming messages and returns a response.
.. note::
Agents are stateful and the messages passed to this method should
be the new messages since the last call to this method. The agent
should maintain its state between calls to this method. For example,
if the agent needs to remember the previous messages to respond to
the current message, it should store the previous messages in the
agent state.
"""
...
async def on_messages_stream(
self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
"""Handles incoming messages and returns a stream of messages and
and the final item is the response. The base implementation in
:class:`BaseChatAgent` simply calls :meth:`on_messages` and yields
the messages in the response.
.. note::
Agents are stateful and the messages passed to this method should
be the new messages since the last call to this method. The agent
should maintain its state between calls to this method. For example,
if the agent needs to remember the previous messages to respond to
the current message, it should store the previous messages in the
agent state.
"""
response = await self.on_messages(messages, cancellation_token)
for inner_message in response.inner_messages or []:
yield inner_message
yield response
async def run(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> TaskResult:
"""Run the agent with the given task and return the result."""
with trace_invoke_agent_span(
agent_name=self.name,
agent_description=self.description,
):
if cancellation_token is None:
cancellation_token = CancellationToken()
input_messages: List[BaseChatMessage] = []
output_messages: List[BaseAgentEvent | BaseChatMessage] = []
if task is None:
pass
elif isinstance(task, str):
text_msg = TextMessage(content=task, source="user")
input_messages.append(text_msg)
if output_task_messages:
output_messages.append(text_msg)
elif isinstance(task, BaseChatMessage):
input_messages.append(task)
if output_task_messages:
output_messages.append(task)
else:
if not task:
raise ValueError("Task list cannot be empty.")
# Task is a sequence of messages.
for msg in task:
if isinstance(msg, BaseChatMessage):
input_messages.append(msg)
if output_task_messages:
output_messages.append(msg)
else:
raise ValueError(f"Invalid message type in sequence: {type(msg)}")
response = await self.on_messages(input_messages, cancellation_token)
if response.inner_messages is not None:
output_messages += response.inner_messages
output_messages.append(response.chat_message)
return TaskResult(messages=output_messages)
async def run_stream(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
"""Run the agent with the given task and return a stream of messages
and the final task result as the last item in the stream.
Args:
task: The task to run. Can be a string, a single message, or a sequence of messages.
cancellation_token: The cancellation token to kill the task immediately.
output_task_messages: Whether to include task messages in the output stream. Defaults to True for backward compatibility.
"""
with trace_invoke_agent_span(
agent_name=self.name,
agent_description=self.description,
):
if cancellation_token is None:
cancellation_token = CancellationToken()
input_messages: List[BaseChatMessage] = []
output_messages: List[BaseAgentEvent | BaseChatMessage] = []
if task is None:
pass
elif isinstance(task, str):
text_msg = TextMessage(content=task, source="user")
input_messages.append(text_msg)
if output_task_messages:
output_messages.append(text_msg)
yield text_msg
elif isinstance(task, BaseChatMessage):
input_messages.append(task)
if output_task_messages:
output_messages.append(task)
yield task
else:
if not task:
raise ValueError("Task list cannot be empty.")
for msg in task:
if isinstance(msg, BaseChatMessage):
input_messages.append(msg)
if output_task_messages:
output_messages.append(msg)
yield msg
else:
raise ValueError(f"Invalid message type in sequence: {type(msg)}")
async for message in self.on_messages_stream(input_messages, cancellation_token):
if isinstance(message, Response):
yield message.chat_message
output_messages.append(message.chat_message)
yield TaskResult(messages=output_messages)
else:
yield message
if isinstance(message, ModelClientStreamingChunkEvent):
# Skip the model client streaming chunk events.
continue
output_messages.append(message)
@abstractmethod
async def on_reset(self, cancellation_token: CancellationToken) -> None:
"""Resets the agent to its initialization state."""
...
async def on_pause(self, cancellation_token: CancellationToken) -> None:
"""Called when the agent is paused while running in its :meth:`on_messages` or
:meth:`on_messages_stream` method. This is a no-op by default in the
:class:`BaseChatAgent` class. Subclasses can override this method to
implement custom pause behavior."""
pass
async def on_resume(self, cancellation_token: CancellationToken) -> None:
"""Called when the agent is resumed from a pause while running in
its :meth:`on_messages` or :meth:`on_messages_stream` method.
This is a no-op by default in the :class:`BaseChatAgent` class.
Subclasses can override this method to implement custom resume behavior."""
pass
async def save_state(self) -> Mapping[str, Any]:
"""Export state. Default implementation for stateless agents."""
return BaseState().model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Restore agent from saved state. Default implementation for stateless agents."""
BaseState.model_validate(state)
async def close(self) -> None:
"""Release any resources held by the agent. This is a no-op by default in the
:class:`BaseChatAgent` class. Subclasses can override this method to
implement custom close behavior."""
pass

881
agent_dhal/agentdhal_agentchat/agents/_code_executor_agent.py Normal file
View File

@@ -0,0 +1,881 @@
import logging
import re
from inspect import iscoroutinefunction
from typing import (
AsyncGenerator,
Awaitable,
Callable,
List,
Optional,
Sequence,
Union,
cast,
)
from agentdhal_core import CancellationToken, Component, ComponentModel
from agentdhal_core.code_executor import CodeBlock, CodeExecutor, CodeResult
from agentdhal_core.model_context import (
ChatCompletionContext,
UnboundedChatCompletionContext,
)
from agentdhal_core.models import (
AssistantMessage,
ChatCompletionClient,
CreateResult,
LLMMessage,
SystemMessage,
UserMessage,
)
from pydantic import BaseModel
from typing_extensions import Self
from .. import EVENT_LOGGER_NAME
from ..base import Response
from ..messages import (
BaseAgentEvent,
BaseChatMessage,
CodeExecutionEvent,
CodeGenerationEvent,
HandoffMessage,
ModelClientStreamingChunkEvent,
TextMessage,
ThoughtEvent,
)
from ..utils import remove_images
from ._base_chat_agent import BaseChatAgent
event_logger = logging.getLogger(EVENT_LOGGER_NAME)
class CodeExecutorAgentConfig(BaseModel):
"""Configuration for CodeExecutorAgent"""
name: str
code_executor: ComponentModel
model_client: ComponentModel | None = None
description: str | None = None
sources: List[str] | None = None
system_message: str | None = None
model_client_stream: bool = False
model_context: ComponentModel | None = None
supported_languages: List[str] | None = None
class RetryDecision(BaseModel):
reason: str
retry: bool
class ApprovalRequest(BaseModel):
"""Request for approval of code execution."""
code: str
context: List[LLMMessage]
class ApprovalResponse(BaseModel):
"""Response to approval request."""
approved: bool
reason: str
# Type aliases for approval functions
SyncApprovalFunc = Callable[[ApprovalRequest], ApprovalResponse]
AsyncApprovalFunc = Callable[[ApprovalRequest], Awaitable[ApprovalResponse]]
ApprovalFuncType = Union[SyncApprovalFunc, AsyncApprovalFunc]
class CodeExecutorAgent(BaseChatAgent, Component[CodeExecutorAgentConfig]):
"""(Experimental) An agent that generates and executes code snippets based on user instructions.
.. note::
This agent is experimental and may change in future releases.
It is typically used within a team with another agent that generates code snippets
to be executed or alone with `model_client` provided so that it can generate code
based on user query, execute it and reflect on the code result.
When used with `model_client`, it will generate code snippets using the model
and execute them using the provided `code_executor`. The model will also reflect on the
code execution results. The agent will yield the final reflection result from the model
as the final response.
When used without `model_client`, it will only execute code blocks found in
:class:`~agentdhal_agentchat.messages.TextMessage` messages and returns the output
of the code execution.
.. note::
Using :class:`~agentdhal_agentchat.agents.AssistantAgent` with
:class:`~agentdhal_extensions.tools.code_execution.PythonCodeExecutionTool`
is an alternative to this agent. However, the model for that agent will
have to generate properly escaped code string as a parameter to the tool.
Args:
name (str): The name of the agent.
code_executor (CodeExecutor): The code executor responsible for executing code received in messages
(:py:class:`~agentdhal_extensions.code_executors.docker.DockerCommandLineCodeExecutor` recommended. See example below)
model_client (ChatCompletionClient, optional): The model client to use for inference and generating code.
If not provided, the agent will only execute code blocks found in input messages.
Currently, the model must support structured output mode, which is required for
the automatic retry mechanism to work.
model_client_stream (bool, optional): If `True`, the model client will be used in streaming mode.
:meth:`on_messages_stream` and :meth:`BaseChatAgent.run_stream` methods will
also yield :class:`~agentdhal_agentchat.messages.ModelClientStreamingChunkEvent`
messages as the model client produces chunks of response. Defaults to `False`.
description (str, optional): The description of the agent. If not provided,
:class:`~agentdhal_agentchat.agents.CodeExecutorAgent.DEFAULT_AGENT_DESCRIPTION` will be used.
system_message (str, optional): The system message for the model. If provided, it will be prepended to the messages in the model context when making an inference. Set to `None` to disable.
Defaults to :class:`~agentdhal_agentchat.agents.CodeExecutorAgent.DEFAULT_SYSTEM_MESSAGE`. This is only used if `model_client` is provided.
sources (Sequence[str], optional): Check only messages from the specified agents for the code to execute.
This is useful when the agent is part of a group chat and you want to limit the code execution to messages from specific agents.
If not provided, all messages will be checked for code blocks.
This is only used if `model_client` is not provided.
max_retries_on_error (int, optional): The maximum number of retries on error. If the code execution fails, the agent will retry up to this number of times.
If the code execution fails after this number of retries, the agent will yield a reflection result.
supported_languages (List[str], optional): List of programming languages that will be parsed and executed from agent response;
others will be ignored. Defaults to DEFAULT_SUPPORTED_LANGUAGES.
approval_func (Optional[Union[Callable[[ApprovalRequest], ApprovalResponse], Callable[[ApprovalRequest], Awaitable[ApprovalResponse]]]], optional): A function that is called before each code execution to get approval.
The function takes an ApprovalRequest containing the code to be executed and the current context, and returns an ApprovalResponse.
The function can be either synchronous or asynchronous. If None (default), all code executions are automatically approved.
If set, the agent cannot be serialized using :meth:`~agentdhal_agentchat.agents.CodeExecutorAgent.dump_component`.
.. note::
It is recommended that the `CodeExecutorAgent` agent uses a Docker container to execute code. This ensures that model-generated code is executed in an isolated environment. To use Docker, your environment must have Docker installed and running.
Follow the installation instructions for `Docker <https://docs.docker.com/get-docker/>`_.
.. note::
The code executor only processes code that is properly formatted in markdown code blocks using triple backticks.
For example:
.. code-block:: text
```python
print("Hello World")
```
# or
```sh
echo "Hello World"
```
In this example, we show how to set up a `CodeExecutorAgent` agent that uses the
:py:class:`~agentdhal_extensions.code_executors.docker.DockerCommandLineCodeExecutor`
to execute code snippets in a Docker container. The `work_dir` parameter indicates
where all executed files are first saved locally before being executed in the Docker container.
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import CodeExecutorAgent, ApprovalRequest, ApprovalResponse
from agentdhal_agentchat.messages import TextMessage
from agentdhal_extensions.code_executors.docker import DockerCommandLineCodeExecutor
from agentdhal_core import CancellationToken
def simple_approval_func(request: ApprovalRequest) -> ApprovalResponse:
\"\"\"Simple approval function that requests user input for code execution approval.\"\"\"
print("Code execution approval requested:")
print("=" * 50)
print(request.code)
print("=" * 50)
while True:
user_input = input("Do you want to execute this code? (y/n): ").strip().lower()
if user_input in ['y', 'yes']:
return ApprovalResponse(approved=True, reason='Approved by user')
elif user_input in ['n', 'no']:
return ApprovalResponse(approved=False, reason='Denied by user')
else:
print("Please enter 'y' for yes or 'n' for no.")
async def run_code_executor_agent() -> None:
# Create a code executor agent that uses a Docker container to execute code.
code_executor = DockerCommandLineCodeExecutor(work_dir="coding")
await code_executor.start()
code_executor_agent = CodeExecutorAgent(
"code_executor",
code_executor=code_executor,
approval_func=simple_approval_func
)
# Run the agent with a given code snippet.
task = TextMessage(
content='''Here is some code
```python
print('Hello world')
```
''',
source="user",
)
response = await code_executor_agent.on_messages([task], CancellationToken())
print(response.chat_message)
# Stop the code executor.
await code_executor.stop()
asyncio.run(run_code_executor_agent())
In this example, we show how to set up a `CodeExecutorAgent` agent that uses the
:py:class:`~docker.types.DeviceRequest` to expose a GPU to the container for cuda-accelerated code execution.
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import CodeExecutorAgent
from agentdhal_agentchat.messages import TextMessage
from agentdhal_extensions.code_executors.docker import DockerCommandLineCodeExecutor
from agentdhal_core import CancellationToken
from docker.types import DeviceRequest
async def run_code_executor_agent() -> None:
# Create a code executor agent that uses a Docker container to execute code.
code_executor = DockerCommandLineCodeExecutor(
work_dir="coding", device_requests=[DeviceRequest(count=-1, capabilities=[["gpu"]])]
)
await code_executor.start()
code_executor_agent = CodeExecutorAgent("code_executor", code_executor=code_executor)
# Display the GPU information
task = TextMessage(
content='''Here is some code
```sh
nvidia-smi
```
''',
source="user",
)
response = await code_executor_agent.on_messages([task], CancellationToken())
print(response.chat_message)
# Stop the code executor.
await code_executor.stop()
asyncio.run(run_code_executor_agent())
In the following example, we show how to setup `CodeExecutorAgent` without `model_client` parameter for executing code blocks generated by other agents in a group chat using :py:class:`~agentdhal_extensions.code_executors.docker.DockerCommandLineCodeExecutor`
.. code-block:: python
import asyncio
from agentdhal_extensions.code_executors.docker import DockerCommandLineCodeExecutor
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent, CodeExecutorAgent, ApprovalRequest, ApprovalResponse
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.ui import Console
termination_condition = MaxMessageTermination(3)
def group_chat_approval_func(request: ApprovalRequest) -> ApprovalResponse:
\"\"\"Approval function for group chat that allows basic Python operations.\"\"\"
# Allow common safe operations
safe_operations = ["print(", "import ", "def ", "class ", "if ", "for ", "while "]
if any(op in request.code for op in safe_operations):
return ApprovalResponse(approved=True, reason='Safe Python operation')
# Deny file system operations in group chat
dangerous_operations = ["open(", "file(", "os.", "subprocess", "eval(", "exec("]
if any(op in request.code for op in dangerous_operations):
return ApprovalResponse(approved=False, reason='File system or dangerous operation not allowed')
return ApprovalResponse(approved=True, reason='Operation approved')
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
# define the Docker CLI Code Executor
code_executor = DockerCommandLineCodeExecutor(work_dir="coding")
# start the execution container
await code_executor.start()
code_executor_agent = CodeExecutorAgent(
"code_executor_agent",
code_executor=code_executor,
approval_func=group_chat_approval_func
)
coder_agent = AssistantAgent("coder_agent", model_client=model_client)
groupchat = RoundRobinGroupChat(
participants=[coder_agent, code_executor_agent], termination_condition=termination_condition
)
task = "Write python code to print Hello World!"
await Console(groupchat.run_stream(task=task))
# stop the execution container
await code_executor.stop()
asyncio.run(main())
.. code-block:: text
---------- user ----------
Write python code to print Hello World!
---------- coder_agent ----------
Certainly! Here's a simple Python code to print "Hello World!":
```python
print("Hello World!")
```
You can run this code in any Python environment to display the message.
---------- code_executor_agent ----------
Hello World!
In the following example, we show how to setup `CodeExecutorAgent` with `model_client`
that can generate its own code without the help of any other agent and executing it in
:py:class:`~agentdhal_extensions.code_executors.docker.DockerCommandLineCodeExecutor`.
It also demonstrates using a model-based approval function that reviews the code for safety before execution.
.. code-block:: python
import asyncio
from agentdhal_extensions.code_executors.docker import DockerCommandLineCodeExecutor
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_core.models import SystemMessage, UserMessage
from agentdhal_agentchat.agents import CodeExecutorAgent, ApprovalRequest, ApprovalResponse
from agentdhal_agentchat.conditions import TextMessageTermination
from agentdhal_agentchat.ui import Console
termination_condition = TextMessageTermination("code_executor_agent")
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
async def model_client_approval_func(request: ApprovalRequest) -> ApprovalResponse:
instruction = "Approve or reject the code in the last message based on whether it is dangerous or not. Use the following JSON format for your response: {approved: true/false, reason: 'your reason here'}"
response = await model_client.create(
messages=[SystemMessage(content=instruction)]
+ request.context
+ [UserMessage(content=request.code, source="user")],
json_output=ApprovalResponse,
)
assert isinstance(response.content, str)
return ApprovalResponse.model_validate_json(response.content)
# define the Docker CLI Code Executor
code_executor = DockerCommandLineCodeExecutor(work_dir="coding")
# start the execution container
await code_executor.start()
code_executor_agent = CodeExecutorAgent(
"code_executor_agent",
code_executor=code_executor,
model_client=model_client,
approval_func=model_client_approval_func,
)
task = "Write python code to print Hello World!"
await Console(code_executor_agent.run_stream(task=task))
# stop the execution container
await code_executor.stop()
asyncio.run(main())
.. code-block:: text
---------- user ----------
Write python code to print Hello World!
---------- code_executor_agent ----------
Certainly! Here is a simple Python code to print "Hello World!" to the console:
```python
print("Hello World!")
```
Let's execute it to confirm the output.
---------- code_executor_agent ----------
Hello World!
---------- code_executor_agent ----------
The code has been executed successfully, and it printed "Hello World!" as expected. If you have any more requests or questions, feel free to ask!
"""
DEFAULT_TERMINAL_DESCRIPTION = "A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks)."
DEFAULT_AGENT_DESCRIPTION = "A Code Execution Agent that generates and executes Python and shell scripts based on user instructions. It ensures correctness, efficiency, and minimal errors while gracefully handling edge cases."
DEFAULT_SYSTEM_MESSAGE = "You are a Code Execution Agent. Your role is to generate and execute Python code and shell scripts based on user instructions, ensuring correctness, efficiency, and minimal errors. Handle edge cases gracefully. Python code should be provided in ```python code blocks, and sh shell scripts should be provided in ```sh code blocks for execution."
NO_CODE_BLOCKS_FOUND_MESSAGE = "No code blocks found in the thread. Please provide at least one markdown-encoded code block to execute (i.e., quoting code in ```python or ```sh code blocks)."
DEFAULT_SUPPORTED_LANGUAGES = ["python", "sh"]
component_config_schema = CodeExecutorAgentConfig
component_provider_override = "agentdhal_agentchat.agents.CodeExecutorAgent"
def __init__(
self,
name: str,
code_executor: CodeExecutor,
*,
model_client: ChatCompletionClient | None = None,
model_context: ChatCompletionContext | None = None,
model_client_stream: bool = False,
max_retries_on_error: int = 0,
description: str | None = None,
system_message: str | None = DEFAULT_SYSTEM_MESSAGE,
sources: Sequence[str] | None = None,
supported_languages: List[str] | None = None,
approval_func: Optional[ApprovalFuncType] = None,
) -> None:
if description is None:
if model_client is None:
description = CodeExecutorAgent.DEFAULT_TERMINAL_DESCRIPTION
else:
description = CodeExecutorAgent.DEFAULT_AGENT_DESCRIPTION
super().__init__(name=name, description=description)
self._code_executor = code_executor
self._sources = sources
self._model_client_stream = model_client_stream
self._max_retries_on_error = max_retries_on_error
self._approval_func = approval_func
self._approval_func_is_async = approval_func is not None and iscoroutinefunction(approval_func)
if supported_languages is not None:
self._supported_languages = supported_languages
else:
self._supported_languages = CodeExecutorAgent.DEFAULT_SUPPORTED_LANGUAGES
self._supported_languages_regex = "|".join(re.escape(lang) for lang in self._supported_languages)
self._model_client = None
if model_client is not None:
self._model_client = model_client
if model_context is not None:
self._model_context = model_context
else:
self._model_context = UnboundedChatCompletionContext()
self._system_messaages: List[SystemMessage] = []
if system_message is None:
self._system_messages = []
else:
self._system_messages = [SystemMessage(content=system_message)]
if self._max_retries_on_error > 0:
if not self._model_client or not self._model_client.model_info:
raise ValueError("model_client.model_info must be provided when max_retries_on_error > 0")
if not self._model_client.model_info["structured_output"]:
raise ValueError("Specified model_client doesn't support structured output mode.")
@property
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
"""The types of messages that the code executor agent produces."""
return (TextMessage,)
@property
def model_context(self) -> ChatCompletionContext:
"""
The model context in use by the agent.
"""
return self._model_context
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
async for message in self.on_messages_stream(messages, cancellation_token):
if isinstance(message, Response):
return message
raise AssertionError("The stream should have returned the final result.")
async def on_messages_stream(
self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
"""
Process the incoming messages with the assistant agent and yield events/responses as they happen.
"""
# Gather all relevant state here
agent_name = self.name
model_context = self._model_context
system_messages = self._system_messages
model_client = self._model_client
model_client_stream = self._model_client_stream
max_retries_on_error = self._max_retries_on_error
execution_result: CodeResult | None = None
if model_client is None: # default behaviour for backward compatibility
# execute generated code if present
code_blocks: List[CodeBlock] = await self.extract_code_blocks_from_messages(messages)
if not code_blocks:
yield Response(
chat_message=TextMessage(
content=self.NO_CODE_BLOCKS_FOUND_MESSAGE,
source=agent_name,
)
)
return
execution_result = await self.execute_code_block(code_blocks, cancellation_token)
yield Response(chat_message=TextMessage(content=execution_result.output, source=self.name))
return
inner_messages: List[BaseAgentEvent | BaseChatMessage] = []
for nth_try in range(max_retries_on_error + 1): # Do one default generation, execution and inference loop
# Step 1: Add new user/handoff messages to the model context
await self._add_messages_to_context(
model_context=model_context,
messages=messages,
)
# Step 2: Run inference with the model context
model_result = None
async for inference_output in self._call_llm(
model_client=model_client,
model_client_stream=model_client_stream,
system_messages=system_messages,
model_context=model_context,
agent_name=agent_name,
cancellation_token=cancellation_token,
):
if isinstance(inference_output, CreateResult):
model_result = inference_output
else:
# Streaming chunk event
yield inference_output
assert model_result is not None, "No model result was produced."
# Step 3: [NEW] If the model produced a hidden "thought," yield it as an event
if model_result.thought:
thought_event = ThoughtEvent(content=model_result.thought, source=agent_name)
yield thought_event
inner_messages.append(thought_event)
# Step 4: Add the assistant message to the model context (including thought if present)
await model_context.add_message(
AssistantMessage(
content=model_result.content,
source=agent_name,
thought=getattr(model_result, "thought", None),
)
)
# Step 5: Extract the code blocks from inferred text
assert isinstance(model_result.content, str), "Expected inferred model_result.content to be of type str."
code_blocks = self._extract_markdown_code_blocks(str(model_result.content))
# Step 6: Exit the loop if no code blocks found
if not code_blocks:
yield Response(
chat_message=TextMessage(
content=str(model_result.content),
source=agent_name,
)
)
return
# Step 7: Yield a CodeGenerationEvent
inferred_text_message: CodeGenerationEvent = CodeGenerationEvent(
retry_attempt=nth_try,
content=model_result.content,
code_blocks=code_blocks,
source=agent_name,
)
yield inferred_text_message
# Step 8: Execute the extracted code blocks
execution_result = await self.execute_code_block(inferred_text_message.code_blocks, cancellation_token)
# Step 9: Update model context with the code execution result
await model_context.add_message(
UserMessage(
content=execution_result.output,
source=agent_name,
)
)
# Step 10: Yield a CodeExecutionEvent
yield CodeExecutionEvent(retry_attempt=nth_try, result=execution_result, source=self.name)
# If execution was successful or last retry, then exit
if execution_result.exit_code == 0 or nth_try == max_retries_on_error:
break
# Step 11: If exit code is non-zero and retries are available then
# make an inference asking if we should retry or not
chat_context = await model_context.get_messages()
retry_prompt = (
f"The most recent code execution resulted in an error:\n{execution_result.output}\n\n"
"Should we attempt to resolve it? Please respond with:\n"
"- A boolean value for 'retry' indicating whether it should be retried.\n"
"- A detailed explanation in 'reason' that identifies the issue, justifies your decision to retry or not, and outlines how you would resolve the error if a retry is attempted."
)
chat_context = chat_context + [
UserMessage(
content=retry_prompt,
source=agent_name,
)
]
response = await model_client.create(messages=chat_context, json_output=RetryDecision)
assert isinstance(
response.content, str
), "Expected structured response for retry decision to be of type str."
should_retry_generation = RetryDecision.model_validate_json(str(response.content))
# Exit if no-retry is needed
if not should_retry_generation.retry:
break
yield CodeGenerationEvent(
retry_attempt=nth_try,
content=f"Attempt number: {nth_try + 1}\nProposed correction: {should_retry_generation.reason}",
code_blocks=[],
source=agent_name,
)
# Always reflect on the execution result
async for reflection_response in CodeExecutorAgent._reflect_on_code_block_results_flow(
system_messages=system_messages,
model_client=model_client,
model_client_stream=model_client_stream,
model_context=model_context,
agent_name=agent_name,
inner_messages=inner_messages,
):
yield reflection_response # Last reflection_response is of type Response so it will finish the routine
async def extract_code_blocks_from_messages(self, messages: Sequence[BaseChatMessage]) -> List[CodeBlock]:
# Extract code blocks from the messages.
code_blocks: List[CodeBlock] = []
for msg in messages:
if self._sources is None or msg.source in self._sources:
if isinstance(msg, TextMessage):
code_blocks.extend(self._extract_markdown_code_blocks(msg.content))
# TODO: handle other message types if needed
return code_blocks
async def execute_code_block(
self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
) -> CodeResult:
# Check for approval before executing code blocks
if self._approval_func is not None:
# Combine all code blocks into a single string for approval
combined_code = "\n\n".join([f"```{block.language}\n{block.code}\n```" for block in code_blocks])
# Get the current context from model_context
context_messages = await self._model_context.get_messages()
# Create approval request
approval_request = ApprovalRequest(code=combined_code, context=context_messages)
# Get approval (handle both sync and async functions)
if self._approval_func_is_async:
# Cast to AsyncApprovalFunc for proper typing
async_func = cast(AsyncApprovalFunc, self._approval_func)
approval_response = await async_func(approval_request)
else:
# Cast to SyncApprovalFunc for proper typing
sync_func = cast(SyncApprovalFunc, self._approval_func)
approval_response = sync_func(approval_request)
# If not approved, return error result
if not approval_response.approved:
return CodeResult(
exit_code=1, output=f"Code execution was not approved. Reason: {approval_response.reason}"
)
# Execute the code blocks.
result = await self._code_executor.execute_code_blocks(code_blocks, cancellation_token=cancellation_token)
if result.output.strip() == "":
# No output
result.output = f"The script ran but produced no output to console. The POSIX exit code was: {result.exit_code}. If you were expecting output, consider revising the script to ensure content is printed to stdout."
elif result.exit_code != 0:
# Error
result.output = f"The script ran, then exited with an error (POSIX exit code: {result.exit_code})\nIts output was:\n{result.output}"
return result
async def on_reset(self, cancellation_token: CancellationToken) -> None:
"""Its a no-op as the code executor agent has no mutable state."""
pass
def _extract_markdown_code_blocks(self, markdown_text: str) -> List[CodeBlock]:
pattern = re.compile(rf"```(?:\s*({self._supported_languages_regex}))\n([\s\S]*?)```", re.IGNORECASE)
matches = pattern.findall(markdown_text)
code_blocks: List[CodeBlock] = []
for match in matches:
language = match[0].strip() if match[0] else ""
code_content = match[1]
code_blocks.append(CodeBlock(code=code_content, language=language))
return code_blocks
def _to_config(self) -> CodeExecutorAgentConfig:
if self._approval_func is not None:
raise ValueError(
"Cannot serialize CodeExecutorAgent with approval_func set. The approval function is not serializable."
)
return CodeExecutorAgentConfig(
name=self.name,
model_client=(self._model_client.dump_component() if self._model_client is not None else None),
code_executor=self._code_executor.dump_component(),
description=self.description,
sources=list(self._sources) if self._sources is not None else None,
system_message=(
self._system_messages[0].content
if self._system_messages and isinstance(self._system_messages[0].content, str)
else None
),
model_client_stream=self._model_client_stream,
model_context=self._model_context.dump_component(),
supported_languages=self._supported_languages,
)
@classmethod
def _from_config(cls, config: CodeExecutorAgentConfig) -> Self:
return cls(
name=config.name,
model_client=(
ChatCompletionClient.load_component(config.model_client) if config.model_client is not None else None
),
code_executor=CodeExecutor.load_component(config.code_executor),
description=config.description,
sources=config.sources,
system_message=config.system_message,
model_client_stream=config.model_client_stream,
model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None,
supported_languages=config.supported_languages,
approval_func=None, # approval_func cannot be serialized, so it's always None when loading from config
)
@staticmethod
def _get_compatible_context(model_client: ChatCompletionClient, messages: List[LLMMessage]) -> Sequence[LLMMessage]:
"""Ensure that the messages are compatible with the underlying client, by removing images if needed."""
if model_client.model_info["vision"]:
return messages
else:
return remove_images(messages)
@classmethod
async def _call_llm(
cls,
model_client: ChatCompletionClient,
model_client_stream: bool,
system_messages: List[SystemMessage],
model_context: ChatCompletionContext,
agent_name: str,
cancellation_token: CancellationToken,
) -> AsyncGenerator[Union[CreateResult, ModelClientStreamingChunkEvent], None]:
"""
Perform a model inference and yield either streaming chunk events or the final CreateResult.
"""
all_messages = await model_context.get_messages()
llm_messages = cls._get_compatible_context(model_client=model_client, messages=system_messages + all_messages)
if model_client_stream:
model_result: Optional[CreateResult] = None
async for chunk in model_client.create_stream(
llm_messages, tools=[], cancellation_token=cancellation_token
):
if isinstance(chunk, CreateResult):
model_result = chunk
elif isinstance(chunk, str):
yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name)
else:
raise RuntimeError(f"Invalid chunk type: {type(chunk)}")
if model_result is None:
raise RuntimeError("No final model result in streaming mode.")
yield model_result
else:
model_result = await model_client.create(llm_messages, tools=[], cancellation_token=cancellation_token)
yield model_result
@staticmethod
async def _add_messages_to_context(
model_context: ChatCompletionContext,
messages: Sequence[BaseChatMessage],
) -> None:
"""
Add incoming messages to the model context.
"""
for msg in messages:
if isinstance(msg, HandoffMessage):
for llm_msg in msg.context:
await model_context.add_message(llm_msg)
await model_context.add_message(msg.to_model_message())
@classmethod
async def _reflect_on_code_block_results_flow(
cls,
system_messages: List[SystemMessage],
model_client: ChatCompletionClient,
model_client_stream: bool,
model_context: ChatCompletionContext,
agent_name: str,
inner_messages: List[BaseAgentEvent | BaseChatMessage],
) -> AsyncGenerator[Response | ModelClientStreamingChunkEvent | ThoughtEvent, None]:
"""
If reflect_on_code_block_results=True, we do another inference based on tool results
and yield the final text response (or streaming chunks).
"""
all_messages = system_messages + await model_context.get_messages()
llm_messages = cls._get_compatible_context(model_client=model_client, messages=all_messages)
reflection_result: Optional[CreateResult] = None
if model_client_stream:
async for chunk in model_client.create_stream(llm_messages):
if isinstance(chunk, CreateResult):
reflection_result = chunk
elif isinstance(chunk, str):
yield ModelClientStreamingChunkEvent(content=chunk, source=agent_name)
else:
raise RuntimeError(f"Invalid chunk type: {type(chunk)}")
else:
reflection_result = await model_client.create(llm_messages)
if not reflection_result or not isinstance(reflection_result.content, str):
raise RuntimeError("Reflect on tool use produced no valid text response.")
# --- NEW: If the reflection produced a thought, yield it ---
if reflection_result.thought:
thought_event = ThoughtEvent(content=reflection_result.thought, source=agent_name)
yield thought_event
inner_messages.append(thought_event)
# Add to context (including thought if present)
await model_context.add_message(
AssistantMessage(
content=reflection_result.content,
source=agent_name,
thought=getattr(reflection_result, "thought", None),
)
)
yield Response(
chat_message=TextMessage(
content=reflection_result.content,
source=agent_name,
models_usage=reflection_result.usage,
),
inner_messages=inner_messages,
)

203
agent_dhal/agentdhal_agentchat/agents/_message_filter_agent.py Normal file
View File

@@ -0,0 +1,203 @@
from typing import AsyncGenerator, List, Literal, Optional, Sequence, Union
from agentdhal_core import CancellationToken, Component, ComponentModel
from pydantic import BaseModel
from agentdhal_agentchat.agents import BaseChatAgent
from agentdhal_agentchat.base import Response
from agentdhal_agentchat.messages import BaseAgentEvent, BaseChatMessage
# ------------------------------
# Message Filter Config
# ------------------------------
class PerSourceFilter(BaseModel):
source: str
position: Optional[Literal["first", "last"]] = None
count: Optional[int] = None
class MessageFilterConfig(BaseModel):
per_source: List[PerSourceFilter]
# ------------------------------
# Component Config
# ------------------------------
class MessageFilterAgentConfig(BaseModel):
name: str
wrapped_agent: ComponentModel
filter: MessageFilterConfig
# ------------------------------
# Message Filter Agent
# ------------------------------
class MessageFilterAgent(BaseChatAgent, Component[MessageFilterAgentConfig]):
"""
A wrapper agent that filters incoming messages before passing them to the inner agent.
.. warning::
This is an experimental feature, and the API will change in the future releases.
This is useful in scenarios like multi-agent workflows where an agent should only
process a subset of the full message history—for example, only the last message
from each upstream agent, or only the first message from a specific source.
Filtering is configured using :class:`MessageFilterConfig`, which supports:
- Filtering by message source (e.g., only messages from "user" or another agent)
- Selecting the first N or last N messages from each source
- If position is `None`, all messages from that source are included
This agent is compatible with both direct message passing and team-based execution
such as :class:`~agentdhal_agentchat.teams.GraphFlow`.
Example:
>>> agent_a = MessageFilterAgent(
... name="A",
... wrapped_agent=some_other_agent,
... filter=MessageFilterConfig(
... per_source=[
... PerSourceFilter(source="user", position="first", count=1),
... PerSourceFilter(source="B", position="last", count=2),
... ]
... ),
... )
Example use case with Graph:
Suppose you have a looping multi-agent graph: A → B → A → B → C.
You want:
- A to only see the user message and the last message from B
- B to see the user message, last message from A, and its own prior responses (for reflection)
- C to see the user message and the last message from B
Wrap the agents like so:
>>> agent_a = MessageFilterAgent(
... name="A",
... wrapped_agent=agent_a_inner,
... filter=MessageFilterConfig(
... per_source=[
... PerSourceFilter(source="user", position="first", count=1),
... PerSourceFilter(source="B", position="last", count=1),
... ]
... ),
... )
>>> agent_b = MessageFilterAgent(
... name="B",
... wrapped_agent=agent_b_inner,
... filter=MessageFilterConfig(
... per_source=[
... PerSourceFilter(source="user", position="first", count=1),
... PerSourceFilter(source="A", position="last", count=1),
... PerSourceFilter(source="B", position="last", count=10),
... ]
... ),
... )
>>> agent_c = MessageFilterAgent(
... name="C",
... wrapped_agent=agent_c_inner,
... filter=MessageFilterConfig(
... per_source=[
... PerSourceFilter(source="user", position="first", count=1),
... PerSourceFilter(source="B", position="last", count=1),
... ]
... ),
... )
Then define the graph:
>>> graph = DiGraph(
... nodes={
... "A": DiGraphNode(name="A", edges=[DiGraphEdge(target="B")]),
... "B": DiGraphNode(
... name="B",
... edges=[
... DiGraphEdge(target="C", condition="exit"),
... DiGraphEdge(target="A", condition="loop"),
... ],
... ),
... "C": DiGraphNode(name="C", edges=[]),
... },
... default_start_node="A",
... )
This will ensure each agent sees only what is needed for its decision or action logic.
"""
component_config_schema = MessageFilterAgentConfig
component_provider_override = "agentdhal_agentchat.agents.MessageFilterAgent"
def __init__(
self,
name: str,
wrapped_agent: BaseChatAgent,
filter: MessageFilterConfig,
):
super().__init__(name=name, description=f"{wrapped_agent.description} (with message filtering)")
self._wrapped_agent = wrapped_agent
self._filter = filter
@property
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
return self._wrapped_agent.produced_message_types
def _apply_filter(self, messages: Sequence[BaseChatMessage]) -> Sequence[BaseChatMessage]:
result: List[BaseChatMessage] = []
for source_filter in self._filter.per_source:
msgs = [m for m in messages if m.source == source_filter.source]
if source_filter.position == "first" and source_filter.count:
msgs = msgs[: source_filter.count]
elif source_filter.position == "last" and source_filter.count:
msgs = msgs[-source_filter.count :]
result.extend(msgs)
return result
async def on_messages(
self,
messages: Sequence[BaseChatMessage],
cancellation_token: CancellationToken,
) -> Response:
filtered = self._apply_filter(messages)
return await self._wrapped_agent.on_messages(filtered, cancellation_token)
async def on_messages_stream(
self,
messages: Sequence[BaseChatMessage],
cancellation_token: CancellationToken,
) -> AsyncGenerator[Union[BaseAgentEvent, BaseChatMessage, Response], None]:
filtered = self._apply_filter(messages)
async for item in self._wrapped_agent.on_messages_stream(filtered, cancellation_token):
yield item
async def on_reset(self, cancellation_token: CancellationToken) -> None:
await self._wrapped_agent.on_reset(cancellation_token)
def _to_config(self) -> MessageFilterAgentConfig:
return MessageFilterAgentConfig(
name=self.name,
wrapped_agent=self._wrapped_agent.dump_component(),
filter=self._filter,
)
@classmethod
def _from_config(cls, config: MessageFilterAgentConfig) -> "MessageFilterAgent":
wrapped = BaseChatAgent.load_component(config.wrapped_agent)
return cls(
name=config.name,
wrapped_agent=wrapped,
filter=config.filter,
)

302
agent_dhal/agentdhal_agentchat/agents/_society_of_mind_agent.py Normal file
View File

@@ -0,0 +1,302 @@
from typing import Any, AsyncGenerator, List, Mapping, Sequence
from agentdhal_core import CancellationToken, Component, ComponentModel
from agentdhal_core.model_context import (
ChatCompletionContext,
UnboundedChatCompletionContext,
)
from agentdhal_core.models import ChatCompletionClient, LLMMessage, SystemMessage, UserMessage
from pydantic import BaseModel
from typing_extensions import Self
from agentdhal_agentchat.base import Response
from agentdhal_agentchat.state import SocietyOfMindAgentState
from ..base import TaskResult, Team
from ..messages import (
BaseAgentEvent,
BaseChatMessage,
HandoffMessage,
ModelClientStreamingChunkEvent,
TextMessage,
)
from ._base_chat_agent import BaseChatAgent
class SocietyOfMindAgentConfig(BaseModel):
"""The declarative configuration for a SocietyOfMindAgent."""
name: str
team: ComponentModel
model_client: ComponentModel
description: str | None = None
instruction: str | None = None
response_prompt: str | None = None
model_context: ComponentModel | None = None
class SocietyOfMindAgent(BaseChatAgent, Component[SocietyOfMindAgentConfig]):
"""An agent that uses an inner team of agents to generate responses.
Each time the agent's :meth:`on_messages` or :meth:`on_messages_stream`
method is called, it runs the inner team of agents and then uses the
model client to generate a response based on the inner team's messages.
Once the response is generated, the agent resets the inner team by
calling :meth:`Team.reset`.
Limit context size sent to the model:
You can limit the number of messages sent to the model by setting
the `model_context` parameter to a :class:`~agentdhal_core.model_context.BufferedChatCompletionContext`.
This will limit the number of recent messages sent to the model and can be useful
when the model has a limit on the number of tokens it can process.
You can also create your own model context by subclassing
:class:`~agentdhal_core.model_context.ChatCompletionContext`.
Args:
name (str): The name of the agent.
team (Team): The team of agents to use.
model_client (ChatCompletionClient): The model client to use for preparing responses.
description (str, optional): The description of the agent.
instruction (str, optional): The instruction to use when generating a response using the inner team's messages.
Defaults to :attr:`DEFAULT_INSTRUCTION`. It assumes the role of 'system'.
response_prompt (str, optional): The response prompt to use when generating a response using the inner team's messages.
Defaults to :attr:`DEFAULT_RESPONSE_PROMPT`. It assumes the role of 'system'.
model_context (ChatCompletionContext | None, optional): The model context for storing and retrieving :class:`~agentdhal_core.models.LLMMessage`. It can be preloaded with initial messages. The initial messages will be cleared when the agent is reset.
Example:
.. code-block:: python
import asyncio
from agentdhal_agentchat.ui import Console
from agentdhal_agentchat.agents import AssistantAgent, SocietyOfMindAgent
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.conditions import TextMentionTermination
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("assistant1", model_client=model_client, system_message="You are a writer, write well.")
agent2 = AssistantAgent(
"assistant2",
model_client=model_client,
system_message="You are an editor, provide critical feedback. Respond with 'APPROVE' if the text addresses all feedbacks.",
)
inner_termination = TextMentionTermination("APPROVE")
inner_team = RoundRobinGroupChat([agent1, agent2], termination_condition=inner_termination)
society_of_mind_agent = SocietyOfMindAgent("society_of_mind", team=inner_team, model_client=model_client)
agent3 = AssistantAgent(
"assistant3", model_client=model_client, system_message="Translate the text to Spanish."
)
team = RoundRobinGroupChat([society_of_mind_agent, agent3], max_turns=2)
stream = team.run_stream(task="Write a short story with a surprising ending.")
await Console(stream)
asyncio.run(main())
"""
component_config_schema = SocietyOfMindAgentConfig
component_provider_override = "agentdhal_agentchat.agents.SocietyOfMindAgent"
DEFAULT_INSTRUCTION = "Earlier you were asked to fulfill a request. You and your team worked diligently to address that request. Here is a transcript of that conversation:"
"""str: The default instruction to use when generating a response using the
inner team's messages. The instruction will be prepended to the inner team's
messages when generating a response using the model. It assumes the role of
'system'."""
DEFAULT_RESPONSE_PROMPT = (
"Output a standalone response to the original request, without mentioning any of the intermediate discussion."
)
"""str: The default response prompt to use when generating a response using
the inner team's messages. It assumes the role of 'system'."""
DEFAULT_DESCRIPTION = "An agent that uses an inner team of agents to generate responses."
"""str: The default description for a SocietyOfMindAgent."""
def __init__(
self,
name: str,
team: Team,
model_client: ChatCompletionClient,
*,
description: str = DEFAULT_DESCRIPTION,
instruction: str = DEFAULT_INSTRUCTION,
response_prompt: str = DEFAULT_RESPONSE_PROMPT,
model_context: ChatCompletionContext | None = None,
) -> None:
super().__init__(name=name, description=description)
self._team = team
self._model_client = model_client
self._instruction = instruction
self._response_prompt = response_prompt
if model_context is not None:
self._model_context = model_context
else:
self._model_context = UnboundedChatCompletionContext()
@property
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
return (TextMessage,)
@property
def model_context(self) -> ChatCompletionContext:
"""
The model context in use by the agent.
"""
return self._model_context
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
# Call the stream method and collect the messages.
response: Response | None = None
async for msg in self.on_messages_stream(messages, cancellation_token):
if isinstance(msg, Response):
response = msg
assert response is not None
return response
async def on_messages_stream(
self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
# Prepare the task for the team of agents.
task_messages = list(messages)
# Run the team of agents.
result: TaskResult | None = None
inner_messages: List[BaseAgentEvent | BaseChatMessage] = []
model_context = self._model_context
prev_content = await model_context.get_messages()
if len(prev_content) > 0:
prev_message = HandoffMessage(
content="relevant previous messages",
source=self.name,
target="",
context=prev_content,
)
task_messages = [prev_message] + task_messages
if len(task_messages) == 0:
task = None
else:
task = task_messages
# Use the new output_task_messages parameter to avoid fragile count-based logic
async for inner_msg in self._team.run_stream(
task=task, cancellation_token=cancellation_token, output_task_messages=False
):
if isinstance(inner_msg, TaskResult):
result = inner_msg
else:
yield inner_msg
if isinstance(inner_msg, ModelClientStreamingChunkEvent):
# Skip the model client streaming chunk events.
continue
inner_messages.append(inner_msg)
assert result is not None
if len(inner_messages) == 0:
yield Response(
chat_message=TextMessage(source=self.name, content="No response."),
inner_messages=[],
# Response's inner_messages should be empty. Cause that mean is response to outer world.
)
else:
llm_messages: List[LLMMessage] = []
if self._model_client.model_info.get("multiple_system_messages", False):
# The model client supports multiple system messages, so we
llm_messages.append(SystemMessage(content=self._instruction))
else:
# The model client does not support multiple system messages, so we
llm_messages.append(UserMessage(content=self._instruction, source="user"))
# Generate a response using the model client.
for message in inner_messages:
if isinstance(message, BaseChatMessage):
llm_messages.append(message.to_model_message())
if self._model_client.model_info.get("multiple_system_messages", False):
# The model client supports multiple system messages, so we
llm_messages.append(SystemMessage(content=self._response_prompt))
else:
# The model client does not support multiple system messages, so we
llm_messages.append(UserMessage(content=self._response_prompt, source="user"))
completion = await self._model_client.create(messages=llm_messages, cancellation_token=cancellation_token)
assert isinstance(completion.content, str)
yield Response(
chat_message=TextMessage(source=self.name, content=completion.content, models_usage=completion.usage),
inner_messages=[],
# Response's inner_messages should be empty. Cause that mean is response to outer world.
)
# Add new user/handoff messages to the model context
await self._add_messages_to_context(
model_context=model_context,
messages=messages,
)
# Reset the team.
await self._team.reset()
@staticmethod
async def _add_messages_to_context(
model_context: ChatCompletionContext,
messages: Sequence[BaseChatMessage],
) -> None:
"""
Add incoming messages to the model context.
"""
for msg in messages:
if isinstance(msg, HandoffMessage):
for llm_msg in msg.context:
await model_context.add_message(llm_msg)
await model_context.add_message(msg.to_model_message())
async def on_reset(self, cancellation_token: CancellationToken) -> None:
await self._team.reset()
await self._model_context.clear()
async def save_state(self) -> Mapping[str, Any]:
team_state = await self._team.save_state()
state = SocietyOfMindAgentState(inner_team_state=team_state)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
society_of_mind_state = SocietyOfMindAgentState.model_validate(state)
await self._team.load_state(society_of_mind_state.inner_team_state)
def _to_config(self) -> SocietyOfMindAgentConfig:
return SocietyOfMindAgentConfig(
name=self.name,
team=self._team.dump_component(),
model_client=self._model_client.dump_component(),
description=self.description,
instruction=self._instruction,
response_prompt=self._response_prompt,
model_context=self._model_context.dump_component(),
)
@classmethod
def _from_config(cls, config: SocietyOfMindAgentConfig) -> Self:
model_client = ChatCompletionClient.load_component(config.model_client)
team = Team.load_component(config.team)
return cls(
name=config.name,
team=team,
model_client=model_client,
description=config.description or cls.DEFAULT_DESCRIPTION,
instruction=config.instruction or cls.DEFAULT_INSTRUCTION,
response_prompt=config.response_prompt or cls.DEFAULT_RESPONSE_PROMPT,
model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None,
)

249
agent_dhal/agentdhal_agentchat/agents/_user_proxy_agent.py Normal file
View File

@@ -0,0 +1,249 @@
import asyncio
import uuid
from contextlib import contextmanager
from contextvars import ContextVar
from inspect import iscoroutinefunction
from typing import Any, AsyncGenerator, Awaitable, Callable, ClassVar, Generator, Optional, Sequence, Union, cast
from agentdhal_core import CancellationToken, Component
from pydantic import BaseModel
from typing_extensions import Self
from ..base import Response
from ..messages import BaseAgentEvent, BaseChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent
from ._base_chat_agent import BaseChatAgent
SyncInputFunc = Callable[[str], str]
AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
# TODO: check if using to_thread fixes this in jupyter
async def cancellable_input(prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
task: asyncio.Task[str] = asyncio.create_task(asyncio.to_thread(input, prompt))
if cancellation_token is not None:
cancellation_token.link_future(task)
return await task
class UserProxyAgentConfig(BaseModel):
"""Declarative configuration for the UserProxyAgent."""
name: str
description: str = "A human user"
input_func: str | None = None
class UserProxyAgent(BaseChatAgent, Component[UserProxyAgentConfig]):
"""An agent that can represent a human user through an input function.
This agent can be used to represent a human user in a chat system by providing a custom input function.
.. note::
Using :class:`UserProxyAgent` puts a running team in a temporary blocked
state until the user responds. So it is important to time out the user input
function and cancel using the :class:`~agentdhal_core.CancellationToken` if the user does not respond.
The input function should also handle exceptions and return a default response if needed.
For typical use cases that involve
slow human responses, it is recommended to use termination conditions
such as :class:`~agentdhal_agentchat.conditions.HandoffTermination` or :class:`~agentdhal_agentchat.conditions.SourceMatchTermination`
to stop the running team and return the control to the application.
You can run the team again with the user input. This way, the state of the team
can be saved and restored when the user responds.
See `Human-in-the-loop <https://microsoft.github.io/autogen/stable/user-guide/agentchat-user-guide/tutorial/human-in-the-loop.html>`_ for more information.
Args:
name (str): The name of the agent.
description (str, optional): A description of the agent.
input_func (Optional[Callable[[str], str]], Callable[[str, Optional[CancellationToken]], Awaitable[str]]): A function that takes a prompt and returns a user input string.
For examples of integrating with web and UI frameworks, see the following:
* `FastAPI <https://github.com/microsoft/autogen/tree/main/python/samples/agentchat_fastapi>`_
* `ChainLit <https://github.com/microsoft/autogen/tree/main/python/samples/agentchat_chainlit>`_
Example:
Simple usage case::
import asyncio
from agentdhal_core import CancellationToken
from agentdhal_agentchat.agents import UserProxyAgent
from agentdhal_agentchat.messages import TextMessage
async def simple_user_agent():
agent = UserProxyAgent("user_proxy")
response = await asyncio.create_task(
agent.on_messages(
[TextMessage(content="What is your name? ", source="user")],
cancellation_token=CancellationToken(),
)
)
assert isinstance(response.chat_message, TextMessage)
print(f"Your name is {response.chat_message.content}")
Example:
Cancellable usage case::
import asyncio
from typing import Any
from agentdhal_core import CancellationToken
from agentdhal_agentchat.agents import UserProxyAgent
from agentdhal_agentchat.messages import TextMessage
token = CancellationToken()
agent = UserProxyAgent("user_proxy")
async def timeout(delay: float):
await asyncio.sleep(delay)
def cancellation_callback(task: asyncio.Task[Any]):
token.cancel()
async def cancellable_user_agent():
try:
timeout_task = asyncio.create_task(timeout(3))
timeout_task.add_done_callback(cancellation_callback)
agent_task = asyncio.create_task(
agent.on_messages(
[TextMessage(content="What is your name? ", source="user")],
cancellation_token=token,
)
)
response = await agent_task
assert isinstance(response.chat_message, TextMessage)
print(f"Your name is {response.chat_message.content}")
except Exception as e:
print(f"Exception: {e}")
except BaseException as e:
print(f"BaseException: {e}")
"""
component_type = "agent"
component_provider_override = "agentdhal_agentchat.agents.UserProxyAgent"
component_config_schema = UserProxyAgentConfig
class InputRequestContext:
def __init__(self) -> None:
raise RuntimeError(
"InputRequestContext cannot be instantiated. It is a static class that provides context management for user input requests."
)
_INPUT_REQUEST_CONTEXT_VAR: ClassVar[ContextVar[str]] = ContextVar("_INPUT_REQUEST_CONTEXT_VAR")
@classmethod
@contextmanager
def populate_context(cls, ctx: str) -> Generator[None, Any, None]:
""":meta private:"""
token = UserProxyAgent.InputRequestContext._INPUT_REQUEST_CONTEXT_VAR.set(ctx)
try:
yield
finally:
UserProxyAgent.InputRequestContext._INPUT_REQUEST_CONTEXT_VAR.reset(token)
@classmethod
def request_id(cls) -> str:
try:
return cls._INPUT_REQUEST_CONTEXT_VAR.get()
except LookupError as e:
raise RuntimeError(
"InputRequestContext.runtime() must be called within the input callback of a UserProxyAgent."
) from e
def __init__(
self,
name: str,
*,
description: str = "A human user",
input_func: Optional[InputFuncType] = None,
) -> None:
"""Initialize the UserProxyAgent."""
super().__init__(name=name, description=description)
self.input_func = input_func or cancellable_input
self._is_async = iscoroutinefunction(self.input_func)
@property
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
"""Message types this agent can produce."""
return (TextMessage, HandoffMessage)
def _get_latest_handoff(self, messages: Sequence[BaseChatMessage]) -> Optional[HandoffMessage]:
"""Find the HandoffMessage in the message sequence that addresses this agent."""
if len(messages) > 0 and isinstance(messages[-1], HandoffMessage):
if messages[-1].target == self.name:
return messages[-1]
else:
raise RuntimeError(f"Handoff message target does not match agent name: {messages[-1].source}")
return None
async def _get_input(self, prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
"""Handle input based on function signature."""
try:
if self._is_async:
# Cast to AsyncInputFunc for proper typing
async_func = cast(AsyncInputFunc, self.input_func)
return await async_func(prompt, cancellation_token)
else:
# Cast to SyncInputFunc for proper typing
sync_func = cast(SyncInputFunc, self.input_func)
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, sync_func, prompt)
except asyncio.CancelledError:
raise
except Exception as e:
raise RuntimeError(f"Failed to get user input: {str(e)}") from e
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
async for message in self.on_messages_stream(messages, cancellation_token):
if isinstance(message, Response):
return message
raise AssertionError("The stream should have returned the final result.")
async def on_messages_stream(
self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
"""Handle incoming messages by requesting user input."""
try:
# Check for handoff first
handoff = self._get_latest_handoff(messages)
prompt = (
f"Handoff received from {handoff.source}. Enter your response: " if handoff else "Enter your response: "
)
request_id = str(uuid.uuid4())
input_requested_event = UserInputRequestedEvent(request_id=request_id, source=self.name)
yield input_requested_event
with UserProxyAgent.InputRequestContext.populate_context(request_id):
user_input = await self._get_input(prompt, cancellation_token)
# Return appropriate message type based on handoff presence
if handoff:
yield Response(chat_message=HandoffMessage(content=user_input, target=handoff.source, source=self.name))
else:
yield Response(chat_message=TextMessage(content=user_input, source=self.name))
except asyncio.CancelledError:
raise
except Exception as e:
raise RuntimeError(f"Failed to get user input: {str(e)}") from e
async def on_reset(self, cancellation_token: Optional[CancellationToken] = None) -> None:
"""Reset agent state."""
pass
def _to_config(self) -> UserProxyAgentConfig:
# TODO: Add ability to serialie input_func
return UserProxyAgentConfig(name=self.name, description=self.description, input_func=None)
@classmethod
def _from_config(cls, config: UserProxyAgentConfig) -> Self:
return cls(name=config.name, description=config.description, input_func=None)

18
agent_dhal/agentdhal_agentchat/base/__init__.py Normal file
View File

@@ -0,0 +1,18 @@
from ._chat_agent import ChatAgent, Response
from ._handoff import Handoff
from ._task import TaskResult, TaskRunner
from ._team import Team
from ._termination import AndTerminationCondition, OrTerminationCondition, TerminatedException, TerminationCondition
__all__ = [
"ChatAgent",
"Response",
"Team",
"TerminatedException",
"TerminationCondition",
"AndTerminationCondition",
"OrTerminationCondition",
"TaskResult",
"TaskRunner",
"Handoff",
]

94
agent_dhal/agentdhal_agentchat/base/_chat_agent.py Normal file
View File

@@ -0,0 +1,94 @@
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, AsyncGenerator, Mapping, Sequence
from agentdhal_core import CancellationToken, ComponentBase
from pydantic import BaseModel, SerializeAsAny
from ..messages import BaseAgentEvent, BaseChatMessage
from ._task import TaskRunner
@dataclass(kw_only=True)
class Response:
"""A response from calling :meth:`ChatAgent.on_messages`."""
chat_message: SerializeAsAny[BaseChatMessage]
"""A chat message produced by the agent as the response."""
inner_messages: Sequence[SerializeAsAny[BaseAgentEvent | BaseChatMessage]] | None = None
"""Inner messages produced by the agent, they can be :class:`BaseAgentEvent`
or :class:`BaseChatMessage`."""
class ChatAgent(ABC, TaskRunner, ComponentBase[BaseModel]):
"""Protocol for a chat agent."""
component_type = "agent"
@property
@abstractmethod
def name(self) -> str:
"""The name of the agent. This is used by team to uniquely identify
the agent. It should be unique within the team."""
...
@property
@abstractmethod
def description(self) -> str:
"""The description of the agent. This is used by team to
make decisions about which agents to use. The description should
describe the agent's capabilities and how to interact with it."""
...
@property
@abstractmethod
def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
"""The types of messages that the agent produces in the
:attr:`Response.chat_message` field. They must be :class:`BaseChatMessage` types."""
...
@abstractmethod
async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
"""Handles incoming messages and returns a response."""
...
@abstractmethod
def on_messages_stream(
self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
"""Handles incoming messages and returns a stream of inner messages and
and the final item is the response."""
...
@abstractmethod
async def on_reset(self, cancellation_token: CancellationToken) -> None:
"""Resets the agent to its initialization state."""
...
@abstractmethod
async def on_pause(self, cancellation_token: CancellationToken) -> None:
"""Called when the agent is paused. The agent may be running in :meth:`on_messages` or
:meth:`on_messages_stream` when this method is called."""
...
@abstractmethod
async def on_resume(self, cancellation_token: CancellationToken) -> None:
"""Called when the agent is resumed. The agent may be running in :meth:`on_messages` or
:meth:`on_messages_stream` when this method is called."""
...
@abstractmethod
async def save_state(self) -> Mapping[str, Any]:
"""Save agent state for later restoration"""
...
@abstractmethod
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Restore agent from saved state"""
...
@abstractmethod
async def close(self) -> None:
"""Release any resources held by the agent."""
...

62
agent_dhal/agentdhal_agentchat/base/_handoff.py Normal file
View File

@@ -0,0 +1,62 @@
import logging
from typing import Any, Dict
from agentdhal_core.tools import BaseTool, FunctionTool
from pydantic import BaseModel, Field, model_validator
from .. import EVENT_LOGGER_NAME
event_logger = logging.getLogger(EVENT_LOGGER_NAME)
class Handoff(BaseModel):
"""Handoff configuration."""
target: str
"""The name of the target agent to handoff to."""
description: str = Field(default="")
"""The description of the handoff such as the condition under which it should happen and the target agent's ability.
If not provided, it is generated from the target agent's name."""
name: str = Field(default="")
"""The name of this handoff configuration. If not provided, it is generated from the target agent's name."""
message: str = Field(default="")
"""The message to the target agent.
By default, it will be the result for the handoff tool.
If not provided, it is generated from the target agent's name."""
@model_validator(mode="before")
@classmethod
def set_defaults(cls, values: Dict[str, Any]) -> Dict[str, Any]:
if not values.get("description"):
values["description"] = f"Handoff to {values['target']}."
if not values.get("name"):
values["name"] = f"transfer_to_{values['target']}".lower()
else:
name = values["name"]
if not isinstance(name, str):
raise ValueError(f"Handoff name must be a string: {values['name']}")
# Check if name is a valid identifier.
if not name.isidentifier():
raise ValueError(f"Handoff name must be a valid identifier: {values['name']}")
if not values.get("message"):
values["message"] = (
f"Transferred to {values['target']}, adopting the role of {values['target']} immediately."
)
return values
@property
def handoff_tool(self) -> BaseTool[BaseModel, BaseModel]:
"""Create a handoff tool from this handoff configuration."""
def _handoff_tool() -> str:
return self.message
return FunctionTool(_handoff_tool, name=self.name, description=self.description, strict=True)
"""
The tool that can be used to handoff to the target agent.
Typically, the results of the tool's execution are provided to the target agent.
"""

65
agent_dhal/agentdhal_agentchat/base/_task.py Normal file
View File

@@ -0,0 +1,65 @@
from typing import AsyncGenerator, Protocol, Sequence
from agentdhal_core import CancellationToken
from pydantic import BaseModel, SerializeAsAny
from ..messages import BaseAgentEvent, BaseChatMessage
class TaskResult(BaseModel):
"""Result of running a task."""
messages: Sequence[SerializeAsAny[BaseAgentEvent | BaseChatMessage]]
"""Messages produced by the task."""
stop_reason: str | None = None
"""The reason the task stopped."""
class TaskRunner(Protocol):
"""A task runner."""
async def run(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> TaskResult:
"""Run the task and return the result.
The task can be a string, a single message, or a sequence of messages.
The runner is stateful and a subsequent call to this method will continue
from where the previous call left off. If the task is not specified,
the runner will continue with the current task.
Args:
task: The task to run. Can be a string, a single message, or a sequence of messages.
cancellation_token: The cancellation token to kill the task immediately.
output_task_messages: Whether to include task messages in :attr:`TaskResult.messages`. Defaults to True for backward compatibility.
"""
...
def run_stream(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
"""Run the task and produces a stream of messages and the final result
:class:`TaskResult` as the last item in the stream.
The task can be a string, a single message, or a sequence of messages.
The runner is stateful and a subsequent call to this method will continue
from where the previous call left off. If the task is not specified,
the runner will continue with the current task.
Args:
task: The task to run. Can be a string, a single message, or a sequence of messages.
cancellation_token: The cancellation token to kill the task immediately.
output_task_messages: Whether to include task messages in the output stream. Defaults to True for backward compatibility.
"""
...

54
agent_dhal/agentdhal_agentchat/base/_team.py Normal file
View File

@@ -0,0 +1,54 @@
from abc import ABC, abstractmethod
from typing import Any, Mapping
from agentdhal_core import ComponentBase
from pydantic import BaseModel
from ._task import TaskRunner
class Team(ABC, TaskRunner, ComponentBase[BaseModel]):
component_type = "team"
@property
@abstractmethod
def name(self) -> str:
"""The name of the team. This is used by team to uniquely identify itself
in a larger team of teams."""
...
@property
@abstractmethod
def description(self) -> str:
"""A description of the team. This is used to provide context about the
team and its purpose to its parent orchestrator."""
...
@abstractmethod
async def reset(self) -> None:
"""Reset the team and all its participants to its initial state."""
...
@abstractmethod
async def pause(self) -> None:
"""Pause the team and all its participants. This is useful for
pausing the :meth:`agentdhal_agentchat.base.TaskRunner.run` or
:meth:`agentdhal_agentchat.base.TaskRunner.run_stream` methods from
concurrently, while keeping them alive."""
...
@abstractmethod
async def resume(self) -> None:
"""Resume the team and all its participants from a pause after
:meth:`pause` was called."""
...
@abstractmethod
async def save_state(self) -> Mapping[str, Any]:
"""Save the current state of the team."""
...
@abstractmethod
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Load the state of the team."""
...

179
agent_dhal/agentdhal_agentchat/base/_termination.py Normal file
View File

@@ -0,0 +1,179 @@
import asyncio
from abc import ABC, abstractmethod
from typing import List, Sequence
from agentdhal_core import Component, ComponentBase, ComponentModel
from pydantic import BaseModel
from typing_extensions import Self
from ..messages import BaseAgentEvent, BaseChatMessage, StopMessage
class TerminatedException(BaseException): ...
class TerminationCondition(ABC, ComponentBase[BaseModel]):
"""A stateful condition that determines when a conversation should be terminated.
A termination condition is a callable that takes a sequence of BaseChatMessage objects
since the last time the condition was called, and returns a StopMessage if the
conversation should be terminated, or None otherwise.
Once a termination condition has been reached, it must be reset before it can be used again.
Termination conditions can be combined using the AND and OR operators.
Example:
.. code-block:: python
import asyncio
from agentdhal_agentchat.conditions import MaxMessageTermination, TextMentionTermination
async def main() -> None:
# Terminate the conversation after 10 turns or if the text "TERMINATE" is mentioned.
cond1 = MaxMessageTermination(10) | TextMentionTermination("TERMINATE")
# Terminate the conversation after 10 turns and if the text "TERMINATE" is mentioned.
cond2 = MaxMessageTermination(10) & TextMentionTermination("TERMINATE")
# ...
# Reset the termination condition.
await cond1.reset()
await cond2.reset()
asyncio.run(main())
"""
component_type = "termination"
@property
@abstractmethod
def terminated(self) -> bool:
"""Check if the termination condition has been reached"""
...
@abstractmethod
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
"""Check if the conversation should be terminated based on the messages received
since the last time the condition was called.
Return a StopMessage if the conversation should be terminated, or None otherwise.
Args:
messages: The messages received since the last time the condition was called.
Returns:
StopMessage | None: A StopMessage if the conversation should be terminated, or None otherwise.
Raises:
TerminatedException: If the termination condition has already been reached."""
...
@abstractmethod
async def reset(self) -> None:
"""Reset the termination condition."""
...
def __and__(self, other: "TerminationCondition") -> "TerminationCondition":
"""Combine two termination conditions with an AND operation."""
return AndTerminationCondition(self, other)
def __or__(self, other: "TerminationCondition") -> "TerminationCondition":
"""Combine two termination conditions with an OR operation."""
return OrTerminationCondition(self, other)
class AndTerminationConditionConfig(BaseModel):
conditions: List[ComponentModel]
class AndTerminationCondition(TerminationCondition, Component[AndTerminationConditionConfig]):
component_config_schema = AndTerminationConditionConfig
component_type = "termination"
component_provider_override = "agentdhal_agentchat.base.AndTerminationCondition"
def __init__(self, *conditions: TerminationCondition) -> None:
self._conditions = conditions
self._stop_messages: List[StopMessage] = []
@property
def terminated(self) -> bool:
return all(condition.terminated for condition in self._conditions)
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self.terminated:
raise TerminatedException("Termination condition has already been reached.")
# Check all remaining conditions.
stop_messages = await asyncio.gather(
*[condition(messages) for condition in self._conditions if not condition.terminated]
)
# Collect stop messages.
for stop_message in stop_messages:
if stop_message is not None:
self._stop_messages.append(stop_message)
if any(stop_message is None for stop_message in stop_messages):
# If any remaining condition has not reached termination, it is not terminated.
return None
content = ", ".join(stop_message.content for stop_message in self._stop_messages)
source = ", ".join(stop_message.source for stop_message in self._stop_messages)
return StopMessage(content=content, source=source)
async def reset(self) -> None:
for condition in self._conditions:
await condition.reset()
self._stop_messages.clear()
def _to_config(self) -> AndTerminationConditionConfig:
"""Convert the AND termination condition to a config."""
return AndTerminationConditionConfig(conditions=[condition.dump_component() for condition in self._conditions])
@classmethod
def _from_config(cls, config: AndTerminationConditionConfig) -> Self:
"""Create an AND termination condition from a config."""
conditions = [TerminationCondition.load_component(condition_model) for condition_model in config.conditions]
return cls(*conditions)
class OrTerminationConditionConfig(BaseModel):
conditions: List[ComponentModel]
"""List of termination conditions where any one being satisfied is sufficient."""
class OrTerminationCondition(TerminationCondition, Component[OrTerminationConditionConfig]):
component_config_schema = OrTerminationConditionConfig
component_type = "termination"
component_provider_override = "agentdhal_agentchat.base.OrTerminationCondition"
def __init__(self, *conditions: TerminationCondition) -> None:
self._conditions = conditions
@property
def terminated(self) -> bool:
return any(condition.terminated for condition in self._conditions)
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self.terminated:
raise RuntimeError("Termination condition has already been reached")
stop_messages = await asyncio.gather(*[condition(messages) for condition in self._conditions])
stop_messages_filter = [stop_message for stop_message in stop_messages if stop_message is not None]
if len(stop_messages_filter) > 0:
content = ", ".join(stop_message.content for stop_message in stop_messages_filter)
source = ", ".join(stop_message.source for stop_message in stop_messages_filter)
return StopMessage(content=content, source=source)
return None
async def reset(self) -> None:
for condition in self._conditions:
await condition.reset()
def _to_config(self) -> OrTerminationConditionConfig:
"""Convert the OR termination condition to a config."""
return OrTerminationConditionConfig(conditions=[condition.dump_component() for condition in self._conditions])
@classmethod
def _from_config(cls, config: OrTerminationConditionConfig) -> Self:
"""Create an OR termination condition from a config."""
conditions = [TerminationCondition.load_component(condition_model) for condition_model in config.conditions]
return cls(*conditions)

32
agent_dhal/agentdhal_agentchat/conditions/__init__.py Normal file
View File

@@ -0,0 +1,32 @@
"""
This module provides various termination conditions for controlling the behavior of
multi-agent teams.
"""
from ._terminations import (
ExternalTermination,
FunctionalTermination,
FunctionCallTermination,
HandoffTermination,
MaxMessageTermination,
SourceMatchTermination,
StopMessageTermination,
TextMentionTermination,
TextMessageTermination,
TimeoutTermination,
TokenUsageTermination,
)
__all__ = [
"MaxMessageTermination",
"TextMentionTermination",
"StopMessageTermination",
"TokenUsageTermination",
"HandoffTermination",
"TimeoutTermination",
"ExternalTermination",
"SourceMatchTermination",
"TextMessageTermination",
"FunctionCallTermination",
"FunctionalTermination",
]

614
agent_dhal/agentdhal_agentchat/conditions/_terminations.py Normal file
View File

@@ -0,0 +1,614 @@
import asyncio
import time
from typing import Awaitable, Callable, List, Sequence
from agentdhal_core import Component
from pydantic import BaseModel
from typing_extensions import Self
from ..base import TerminatedException, TerminationCondition
from ..messages import (
BaseAgentEvent,
BaseChatMessage,
HandoffMessage,
StopMessage,
TextMessage,
ToolCallExecutionEvent,
)
class StopMessageTerminationConfig(BaseModel):
pass
class StopMessageTermination(TerminationCondition, Component[StopMessageTerminationConfig]):
"""Terminate the conversation if a StopMessage is received."""
component_config_schema = StopMessageTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.StopMessageTermination"
def __init__(self) -> None:
self._terminated = False
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if isinstance(message, StopMessage):
self._terminated = True
return StopMessage(content="Stop message received", source="StopMessageTermination")
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> StopMessageTerminationConfig:
return StopMessageTerminationConfig()
@classmethod
def _from_config(cls, config: StopMessageTerminationConfig) -> Self:
return cls()
class MaxMessageTerminationConfig(BaseModel):
max_messages: int
include_agent_event: bool = False
class MaxMessageTermination(TerminationCondition, Component[MaxMessageTerminationConfig]):
"""Terminate the conversation after a maximum number of messages have been exchanged.
Args:
max_messages: The maximum number of messages allowed in the conversation.
include_agent_event: If True, include :class:`~agentdhal_agentchat.messages.BaseAgentEvent` in the message count.
Otherwise, only include :class:`~agentdhal_agentchat.messages.BaseChatMessage`. Defaults to False.
"""
component_config_schema = MaxMessageTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.MaxMessageTermination"
def __init__(self, max_messages: int, include_agent_event: bool = False) -> None:
self._max_messages = max_messages
self._message_count = 0
self._include_agent_event = include_agent_event
@property
def terminated(self) -> bool:
return self._message_count >= self._max_messages
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self.terminated:
raise TerminatedException("Termination condition has already been reached")
self._message_count += len([m for m in messages if self._include_agent_event or isinstance(m, BaseChatMessage)])
if self._message_count >= self._max_messages:
return StopMessage(
content=f"Maximum number of messages {self._max_messages} reached, current message count: {self._message_count}",
source="MaxMessageTermination",
)
return None
async def reset(self) -> None:
self._message_count = 0
def _to_config(self) -> MaxMessageTerminationConfig:
return MaxMessageTerminationConfig(
max_messages=self._max_messages, include_agent_event=self._include_agent_event
)
@classmethod
def _from_config(cls, config: MaxMessageTerminationConfig) -> Self:
return cls(max_messages=config.max_messages, include_agent_event=config.include_agent_event)
class TextMentionTerminationConfig(BaseModel):
text: str
class TextMentionTermination(TerminationCondition, Component[TextMentionTerminationConfig]):
"""Terminate the conversation if a specific text is mentioned.
Args:
text: The text to look for in the messages.
sources: Check only messages of the specified agents for the text to look for.
"""
component_config_schema = TextMentionTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.TextMentionTermination"
def __init__(self, text: str, sources: Sequence[str] | None = None) -> None:
self._termination_text = text
self._terminated = False
self._sources = sources
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if self._sources is not None and message.source not in self._sources:
continue
content = message.to_text()
if self._termination_text in content:
self._terminated = True
return StopMessage(
content=f"Text '{self._termination_text}' mentioned", source="TextMentionTermination"
)
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> TextMentionTerminationConfig:
return TextMentionTerminationConfig(text=self._termination_text)
@classmethod
def _from_config(cls, config: TextMentionTerminationConfig) -> Self:
return cls(text=config.text)
class FunctionalTermination(TerminationCondition):
"""Terminate the conversation if an functional expression is met.
Args:
func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool] | Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]]): A function that takes a sequence of messages
and returns True if the termination condition is met, False otherwise.
The function can be a callable or an async callable.
Example:
.. code-block:: python
import asyncio
from typing import Sequence
from agentdhal_agentchat.conditions import FunctionalTermination
from agentdhal_agentchat.messages import BaseAgentEvent, BaseChatMessage, StopMessage
def expression(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> bool:
# Check if the last message is a stop message
return isinstance(messages[-1], StopMessage)
termination = FunctionalTermination(expression)
async def run() -> None:
messages = [
StopMessage(source="agent1", content="Stop"),
]
result = await termination(messages)
print(result)
asyncio.run(run())
.. code-block:: text
StopMessage(source="FunctionalTermination", content="Functional termination condition met")
"""
def __init__(
self,
func: Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], bool]
| Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[bool]],
) -> None:
self._func = func
self._terminated = False
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
if asyncio.iscoroutinefunction(self._func):
result = await self._func(messages)
else:
result = self._func(messages)
if result is True:
self._terminated = True
return StopMessage(content="Functional termination condition met", source="FunctionalTermination")
return None
async def reset(self) -> None:
self._terminated = False
class TokenUsageTerminationConfig(BaseModel):
max_total_token: int | None
max_prompt_token: int | None
max_completion_token: int | None
class TokenUsageTermination(TerminationCondition, Component[TokenUsageTerminationConfig]):
"""Terminate the conversation if a token usage limit is reached.
Args:
max_total_token: The maximum total number of tokens allowed in the conversation.
max_prompt_token: The maximum number of prompt tokens allowed in the conversation.
max_completion_token: The maximum number of completion tokens allowed in the conversation.
Raises:
ValueError: If none of max_total_token, max_prompt_token, or max_completion_token is provided.
"""
component_config_schema = TokenUsageTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.TokenUsageTermination"
def __init__(
self,
max_total_token: int | None = None,
max_prompt_token: int | None = None,
max_completion_token: int | None = None,
) -> None:
if max_total_token is None and max_prompt_token is None and max_completion_token is None:
raise ValueError(
"At least one of max_total_token, max_prompt_token, or max_completion_token must be provided"
)
self._max_total_token = max_total_token
self._max_prompt_token = max_prompt_token
self._max_completion_token = max_completion_token
self._total_token_count = 0
self._prompt_token_count = 0
self._completion_token_count = 0
@property
def terminated(self) -> bool:
return (
(self._max_total_token is not None and self._total_token_count >= self._max_total_token)
or (self._max_prompt_token is not None and self._prompt_token_count >= self._max_prompt_token)
or (self._max_completion_token is not None and self._completion_token_count >= self._max_completion_token)
)
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self.terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if message.models_usage is not None:
self._prompt_token_count += message.models_usage.prompt_tokens
self._completion_token_count += message.models_usage.completion_tokens
self._total_token_count += message.models_usage.prompt_tokens + message.models_usage.completion_tokens
if self.terminated:
content = f"Token usage limit reached, total token count: {self._total_token_count}, prompt token count: {self._prompt_token_count}, completion token count: {self._completion_token_count}."
return StopMessage(content=content, source="TokenUsageTermination")
return None
async def reset(self) -> None:
self._total_token_count = 0
self._prompt_token_count = 0
self._completion_token_count = 0
def _to_config(self) -> TokenUsageTerminationConfig:
return TokenUsageTerminationConfig(
max_total_token=self._max_total_token,
max_prompt_token=self._max_prompt_token,
max_completion_token=self._max_completion_token,
)
@classmethod
def _from_config(cls, config: TokenUsageTerminationConfig) -> Self:
return cls(
max_total_token=config.max_total_token,
max_prompt_token=config.max_prompt_token,
max_completion_token=config.max_completion_token,
)
class HandoffTerminationConfig(BaseModel):
target: str
class HandoffTermination(TerminationCondition, Component[HandoffTerminationConfig]):
"""Terminate the conversation if a :class:`~agentdhal_agentchat.messages.HandoffMessage`
with the given target is received.
Args:
target (str): The target of the handoff message.
"""
component_config_schema = HandoffTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.HandoffTermination"
def __init__(self, target: str) -> None:
self._terminated = False
self._target = target
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if isinstance(message, HandoffMessage) and message.target == self._target:
self._terminated = True
return StopMessage(
content=f"Handoff to {self._target} from {message.source} detected.", source="HandoffTermination"
)
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> HandoffTerminationConfig:
return HandoffTerminationConfig(target=self._target)
@classmethod
def _from_config(cls, config: HandoffTerminationConfig) -> Self:
return cls(target=config.target)
class TimeoutTerminationConfig(BaseModel):
timeout_seconds: float
class TimeoutTermination(TerminationCondition, Component[TimeoutTerminationConfig]):
"""Terminate the conversation after a specified duration has passed.
Args:
timeout_seconds: The maximum duration in seconds before terminating the conversation.
"""
component_config_schema = TimeoutTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.TimeoutTermination"
def __init__(self, timeout_seconds: float) -> None:
self._timeout_seconds = timeout_seconds
self._start_time = time.monotonic()
self._terminated = False
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
if (time.monotonic() - self._start_time) >= self._timeout_seconds:
self._terminated = True
return StopMessage(
content=f"Timeout of {self._timeout_seconds} seconds reached", source="TimeoutTermination"
)
return None
async def reset(self) -> None:
self._start_time = time.monotonic()
self._terminated = False
def _to_config(self) -> TimeoutTerminationConfig:
return TimeoutTerminationConfig(timeout_seconds=self._timeout_seconds)
@classmethod
def _from_config(cls, config: TimeoutTerminationConfig) -> Self:
return cls(timeout_seconds=config.timeout_seconds)
class ExternalTerminationConfig(BaseModel):
pass
class ExternalTermination(TerminationCondition, Component[ExternalTerminationConfig]):
"""A termination condition that is externally controlled
by calling the :meth:`set` method.
Example:
.. code-block:: python
from agentdhal_agentchat.conditions import ExternalTermination
termination = ExternalTermination()
# Run the team in an asyncio task.
...
# Set the termination condition externally
termination.set()
"""
component_config_schema = ExternalTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.ExternalTermination"
def __init__(self) -> None:
self._terminated = False
self._setted = False
@property
def terminated(self) -> bool:
return self._terminated
def set(self) -> None:
"""Set the termination condition to terminated."""
self._setted = True
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
if self._setted:
self._terminated = True
return StopMessage(content="External termination requested", source="ExternalTermination")
return None
async def reset(self) -> None:
self._terminated = False
self._setted = False
def _to_config(self) -> ExternalTerminationConfig:
return ExternalTerminationConfig()
@classmethod
def _from_config(cls, config: ExternalTerminationConfig) -> Self:
return cls()
class SourceMatchTerminationConfig(BaseModel):
sources: List[str]
class SourceMatchTermination(TerminationCondition, Component[SourceMatchTerminationConfig]):
"""Terminate the conversation after a specific source responds.
Args:
sources (List[str]): List of source names to terminate the conversation.
Raises:
TerminatedException: If the termination condition has already been reached.
"""
component_config_schema = SourceMatchTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.SourceMatchTermination"
def __init__(self, sources: List[str]) -> None:
self._sources = sources
self._terminated = False
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
if not messages:
return None
for message in messages:
if message.source in self._sources:
self._terminated = True
return StopMessage(content=f"'{message.source}' answered", source="SourceMatchTermination")
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> SourceMatchTerminationConfig:
return SourceMatchTerminationConfig(sources=self._sources)
@classmethod
def _from_config(cls, config: SourceMatchTerminationConfig) -> Self:
return cls(sources=config.sources)
class TextMessageTerminationConfig(BaseModel):
"""Configuration for the TextMessageTermination termination condition."""
source: str | None = None
"""The source of the text message to terminate the conversation."""
class TextMessageTermination(TerminationCondition, Component[TextMessageTerminationConfig]):
"""Terminate the conversation if a :class:`~agentdhal_agentchat.messages.TextMessage` is received.
This termination condition checks for TextMessage instances in the message sequence. When a TextMessage is found,
it terminates the conversation if either:
- No source was specified (terminates on any TextMessage)
- The message source matches the specified source
Args:
source (str | None, optional): The source name to match against incoming messages. If None, matches any source.
Defaults to None.
"""
component_config_schema = TextMessageTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.TextMessageTermination"
def __init__(self, source: str | None = None) -> None:
self._terminated = False
self._source = source
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if isinstance(message, TextMessage) and (self._source is None or message.source == self._source):
self._terminated = True
return StopMessage(
content=f"Text message received from '{message.source}'", source="TextMessageTermination"
)
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> TextMessageTerminationConfig:
return TextMessageTerminationConfig(source=self._source)
@classmethod
def _from_config(cls, config: TextMessageTerminationConfig) -> Self:
return cls(source=config.source)
class FunctionCallTerminationConfig(BaseModel):
"""Configuration for the :class:`FunctionCallTermination` termination condition."""
function_name: str
class FunctionCallTermination(TerminationCondition, Component[FunctionCallTerminationConfig]):
"""Terminate the conversation if a :class:`~agentdhal_core.models.FunctionExecutionResult`
with a specific name was received.
Args:
function_name (str): The name of the function to look for in the messages.
Raises:
TerminatedException: If the termination condition has already been reached.
"""
component_config_schema = FunctionCallTerminationConfig
component_provider_override = "agentdhal_agentchat.conditions.FunctionCallTermination"
"""The schema for the component configuration."""
def __init__(self, function_name: str) -> None:
self._terminated = False
self._function_name = function_name
@property
def terminated(self) -> bool:
return self._terminated
async def __call__(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> StopMessage | None:
if self._terminated:
raise TerminatedException("Termination condition has already been reached")
for message in messages:
if isinstance(message, ToolCallExecutionEvent):
for execution in message.content:
if execution.name == self._function_name:
self._terminated = True
return StopMessage(
content=f"Function '{self._function_name}' was executed.",
source="FunctionCallTermination",
)
return None
async def reset(self) -> None:
self._terminated = False
def _to_config(self) -> FunctionCallTerminationConfig:
return FunctionCallTerminationConfig(
function_name=self._function_name,
)
@classmethod
def _from_config(cls, config: FunctionCallTerminationConfig) -> Self:
return cls(
function_name=config.function_name,
)

693
agent_dhal/agentdhal_agentchat/messages.py Normal file
View File

@@ -0,0 +1,693 @@
"""
This module defines various message types used for agent-to-agent communication.
Each message type inherits either from the BaseChatMessage class or BaseAgentEvent
class and includes specific fields relevant to the type of message being sent.
"""
import uuid
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from typing import Any, Dict, Generic, List, Literal, Mapping, Optional, Type, TypeVar
from agentdhal_core import Component, ComponentBase, FunctionCall, Image
from agentdhal_core.code_executor import CodeBlock, CodeResult
from agentdhal_core.memory import MemoryContent
from agentdhal_core.models import (
FunctionExecutionResult,
LLMMessage,
RequestUsage,
UserMessage,
)
from agentdhal_core.utils import schema_to_pydantic_model
from pydantic import BaseModel, Field, computed_field
from typing_extensions import Annotated, Self
class BaseMessage(BaseModel, ABC):
"""Abstract base class for all message types in AgentChat.
.. warning::
If you want to create a new message type, do not inherit from this class.
Instead, inherit from :class:`BaseChatMessage` or :class:`BaseAgentEvent`
to clarify the purpose of the message type.
"""
@abstractmethod
def to_text(self) -> str:
"""Convert the message content to a string-only representation
that can be rendered in the console and inspected by the user or conditions.
This is not used for creating text-only content for models.
For :class:`BaseChatMessage` types, use :meth:`to_model_text` instead."""
...
def dump(self) -> Mapping[str, Any]:
"""Convert the message to a JSON-serializable dictionary.
The default implementation uses the Pydantic model's
:meth:`model_dump` method to convert the message to a dictionary.
Datetime objects are automatically converted to ISO format strings
to ensure JSON serialization compatibility.
Override this method if you want to customize the serialization
process or add additional fields to the output.
"""
return self.model_dump(mode="json")
@classmethod
def load(cls, data: Mapping[str, Any]) -> Self:
"""Create a message from a dictionary of JSON-serializable data.
The default implementation uses the Pydantic model's
:meth:`model_validate` method to create the message from the data.
Override this method if you want to customize the deserialization
process or add additional fields to the input data."""
return cls.model_validate(data)
class BaseChatMessage(BaseMessage, ABC):
"""Abstract base class for chat messages.
.. note::
If you want to create a new message type that is used for agent-to-agent
communication, inherit from this class, or simply use
:class:`StructuredMessage` if your content type is a subclass of
Pydantic BaseModel.
This class is used for messages that are sent between agents in a chat
conversation. Agents are expected to process the content of the
message using models and return a response as another :class:`BaseChatMessage`.
"""
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
"""Unique identifier for this message."""
source: str
"""The name of the agent that sent this message."""
models_usage: RequestUsage | None = None
"""The model client usage incurred when producing this message."""
metadata: Dict[str, str] = {}
"""Additional metadata about the message."""
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
"""The time when the message was created."""
@abstractmethod
def to_model_text(self) -> str:
"""Convert the content of the message to text-only representation.
This is used for creating text-only content for models.
This is not used for rendering the message in console. For that, use
:meth:`~BaseMessage.to_text`.
The difference between this and :meth:`to_model_message` is that this
is used to construct parts of the a message for the model client,
while :meth:`to_model_message` is used to create a complete message
for the model client.
"""
...
@abstractmethod
def to_model_message(self) -> UserMessage:
"""Convert the message content to a :class:`~agentdhal_core.models.UserMessage`
for use with model client, e.g., :class:`~agentdhal_core.models.ChatCompletionClient`.
"""
...
class BaseTextChatMessage(BaseChatMessage, ABC):
"""Base class for all text-only :class:`BaseChatMessage` types.
It has implementations for :meth:`to_text`, :meth:`to_model_text`,
and :meth:`to_model_message` methods.
Inherit from this class if your message content type is a string.
"""
content: str
"""The content of the message."""
def to_text(self) -> str:
return self.content
def to_model_text(self) -> str:
return self.content
def to_model_message(self) -> UserMessage:
return UserMessage(content=self.content, source=self.source)
class BaseAgentEvent(BaseMessage, ABC):
"""Base class for agent events.
.. note::
If you want to create a new message type for signaling observable events
to user and application, inherit from this class.
Agent events are used to signal actions and thoughts produced by agents
and teams to user and applications. They are not used for agent-to-agent
communication and are not expected to be processed by other agents.
You should override the :meth:`to_text` method if you want to provide
a custom rendering of the content.
"""
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
"""Unique identifier for this event."""
source: str
"""The name of the agent that sent this message."""
models_usage: RequestUsage | None = None
"""The model client usage incurred when producing this message."""
metadata: Dict[str, str] = {}
"""Additional metadata about the message."""
created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
"""The time when the message was created."""
StructuredContentType = TypeVar("StructuredContentType", bound=BaseModel, covariant=True)
"""Type variable for structured content types."""
class StructuredMessage(BaseChatMessage, Generic[StructuredContentType]):
"""A :class:`BaseChatMessage` type with an unspecified content type.
To create a new structured message type, specify the content type
as a subclass of `Pydantic BaseModel <https://docs.pydantic.dev/latest/concepts/models/>`_.
.. code-block:: python
from pydantic import BaseModel
from agentdhal_agentchat.messages import StructuredMessage
class MyMessageContent(BaseModel):
text: str
number: int
message = StructuredMessage[MyMessageContent](
content=MyMessageContent(text="Hello", number=42),
source="agent1",
)
print(message.to_text()) # {"text": "Hello", "number": 42}
.. code-block:: python
from pydantic import BaseModel
from agentdhal_agentchat.messages import StructuredMessage
class MyMessageContent(BaseModel):
text: str
number: int
message = StructuredMessage[MyMessageContent](
content=MyMessageContent(text="Hello", number=42),
source="agent",
format_string="Hello, {text} {number}!",
)
print(message.to_text()) # Hello, agent 42!
"""
content: StructuredContentType
"""The content of the message. Must be a subclass of
`Pydantic BaseModel <https://docs.pydantic.dev/latest/concepts/models/>`_."""
format_string: Optional[str] = None
"""(Experimental) An optional format string to render the content into a human-readable format.
The format string can use the fields of the content model as placeholders.
For example, if the content model has a field `name`, you can use
`{name}` in the format string to include the value of that field.
The format string is used in the :meth:`to_text` method to create a
human-readable representation of the message.
This setting is experimental and will change in the future.
"""
@computed_field
def type(self) -> str:
return self.__class__.__name__
def to_text(self) -> str:
if self.format_string is not None:
return self.format_string.format(**self.content.model_dump())
else:
return self.content.model_dump_json()
def to_model_text(self) -> str:
if self.format_string is not None:
return self.format_string.format(**self.content.model_dump())
else:
return self.content.model_dump_json()
def to_model_message(self) -> UserMessage:
return UserMessage(
content=self.content.model_dump_json(),
source=self.source,
)
class StructureMessageConfig(BaseModel):
"""The declarative configuration for the structured output."""
json_schema: Dict[str, Any]
format_string: Optional[str] = None
content_model_name: str
class StructuredMessageFactory(ComponentBase[StructureMessageConfig], Component[StructureMessageConfig]):
""":meta private:
A component that creates structured chat messages from Pydantic models or JSON schemas.
This component helps you generate strongly-typed chat messages with content defined using a Pydantic model.
It can be used in declarative workflows where message structure must be validated, formatted, and serialized.
You can initialize the component directly using a `BaseModel` subclass, or dynamically from a configuration
object (e.g., loaded from disk or a database).
### Example 1: Create from a Pydantic Model
.. code-block:: python
from pydantic import BaseModel
from agentdhal_agentchat.messages import StructuredMessageFactory
class TestContent(BaseModel):
field1: str
field2: int
format_string = "This is a string {field1} and this is an int {field2}"
sm_component = StructuredMessageFactory(input_model=TestContent, format_string=format_string)
message = sm_component.StructuredMessage(
source="test_agent", content=TestContent(field1="Hello", field2=42), format_string=format_string
)
print(message.to_model_text()) # Output: This is a string Hello and this is an int 42
config = sm_component.dump_component()
s_m_dyn = StructuredMessageFactory.load_component(config)
message = s_m_dyn.StructuredMessage(
source="test_agent",
content=s_m_dyn.ContentModel(field1="dyn agent", field2=43),
format_string=s_m_dyn.format_string,
)
print(type(message)) # StructuredMessage[GeneratedModel]
print(message.to_model_text()) # Output: This is a string dyn agent and this is an int 43
Attributes:
component_config_schema (StructureMessageConfig): Defines the configuration structure for this component.
component_provider_override (str): Path used to reference this component in external tooling.
component_type (str): Identifier used for categorization (e.g., "structured_message").
Raises:
ValueError: If neither `json_schema` nor `input_model` is provided.
Args:
json_schema (Optional[str]): JSON schema to dynamically create a Pydantic model.
input_model (Optional[Type[BaseModel]]): A subclass of `BaseModel` that defines the expected message structure.
format_string (Optional[str]): Optional string to render content into a human-readable format.
content_model_name (Optional[str]): Optional name for the generated Pydantic model.
"""
component_config_schema = StructureMessageConfig
component_provider_override = "agentdhal_agentchat.messages.StructuredMessageFactory"
component_type = "structured_message"
def __init__(
self,
json_schema: Optional[Dict[str, Any]] = None,
input_model: Optional[Type[BaseModel]] = None,
format_string: Optional[str] = None,
content_model_name: Optional[str] = None,
) -> None:
self.format_string = format_string
if json_schema:
self.ContentModel = schema_to_pydantic_model(
json_schema, model_name=content_model_name or "GeneratedContentModel"
)
elif input_model:
self.ContentModel = input_model
else:
raise ValueError("Either `json_schema` or `input_model` must be provided.")
self.StructuredMessage = StructuredMessage[self.ContentModel] # type: ignore[name-defined]
def _to_config(self) -> StructureMessageConfig:
return StructureMessageConfig(
json_schema=self.ContentModel.model_json_schema(),
format_string=self.format_string,
content_model_name=self.ContentModel.__name__,
)
@classmethod
def _from_config(cls, config: StructureMessageConfig) -> "StructuredMessageFactory":
return cls(
json_schema=config.json_schema,
format_string=config.format_string,
content_model_name=config.content_model_name,
)
class TextMessage(BaseTextChatMessage):
"""A text message with string-only content."""
type: Literal["TextMessage"] = "TextMessage"
class MultiModalMessage(BaseChatMessage):
"""A multimodal message."""
content: List[str | Image]
"""The content of the message."""
type: Literal["MultiModalMessage"] = "MultiModalMessage"
def to_model_text(self, image_placeholder: str | None = "[image]") -> str:
"""Convert the content of the message to a string-only representation.
If an image is present, it will be replaced with the image placeholder
by default, otherwise it will be a base64 string when set to None.
"""
text = ""
for c in self.content:
if isinstance(c, str):
text += c
elif isinstance(c, Image):
if image_placeholder is not None:
text += f" {image_placeholder}"
else:
text += f" {c.to_base64()}"
return text
def to_text(self, iterm: bool = False) -> str:
result: List[str] = []
for c in self.content:
if isinstance(c, str):
result.append(c)
else:
if iterm:
# iTerm2 image rendering protocol: https://iterm2.com/documentation-images.html
image_data = c.to_base64()
result.append(f"\033]1337;File=inline=1:{image_data}\a\n")
else:
result.append("<image>")
return "\n".join(result)
def to_model_message(self) -> UserMessage:
return UserMessage(content=self.content, source=self.source)
class StopMessage(BaseTextChatMessage):
"""A message requesting stop of a conversation."""
type: Literal["StopMessage"] = "StopMessage"
class HandoffMessage(BaseTextChatMessage):
"""A message requesting handoff of a conversation to another agent."""
target: str
"""The name of the target agent to handoff to."""
context: List[LLMMessage] = []
"""The model context to be passed to the target agent."""
type: Literal["HandoffMessage"] = "HandoffMessage"
class ToolCallSummaryMessage(BaseTextChatMessage):
"""A message signaling the summary of tool call results."""
type: Literal["ToolCallSummaryMessage"] = "ToolCallSummaryMessage"
tool_calls: List[FunctionCall]
"""The tool calls that were made."""
results: List[FunctionExecutionResult]
"""The results of the tool calls."""
class ToolCallRequestEvent(BaseAgentEvent):
"""An event signaling a request to use tools."""
content: List[FunctionCall]
"""The tool calls."""
type: Literal["ToolCallRequestEvent"] = "ToolCallRequestEvent"
def to_text(self) -> str:
return str(self.content)
class CodeGenerationEvent(BaseAgentEvent):
"""An event signaling code generation event."""
retry_attempt: int
"Retry number, 0 means first generation"
content: str
"The complete content as string."
code_blocks: List[CodeBlock]
"List of code blocks present in content"
type: Literal["CodeGenerationEvent"] = "CodeGenerationEvent"
def to_text(self) -> str:
return self.content
class CodeExecutionEvent(BaseAgentEvent):
"""An event signaling code execution event."""
retry_attempt: int
"Retry number, 0 means first execution"
result: CodeResult
"Code Execution Result"
type: Literal["CodeExecutionEvent"] = "CodeExecutionEvent"
def to_text(self) -> str:
return self.result.output
class ToolCallExecutionEvent(BaseAgentEvent):
"""An event signaling the execution of tool calls."""
content: List[FunctionExecutionResult]
"""The tool call results."""
type: Literal["ToolCallExecutionEvent"] = "ToolCallExecutionEvent"
def to_text(self) -> str:
return str(self.content)
class UserInputRequestedEvent(BaseAgentEvent):
"""An event signaling a that the user proxy has requested user input. Published prior to invoking the input callback."""
request_id: str
"""Identifier for the user input request."""
content: Literal[""] = ""
"""Empty content for compat with consumers expecting a content field."""
type: Literal["UserInputRequestedEvent"] = "UserInputRequestedEvent"
def to_text(self) -> str:
return str(self.content)
class MemoryQueryEvent(BaseAgentEvent):
"""An event signaling the results of memory queries."""
content: List[MemoryContent]
"""The memory query results."""
type: Literal["MemoryQueryEvent"] = "MemoryQueryEvent"
def to_text(self) -> str:
return str(self.content)
class ModelClientStreamingChunkEvent(BaseAgentEvent):
"""An event signaling a text output chunk from a model client in streaming mode."""
content: str
"""A string chunk from the model client."""
full_message_id: str | None = None
"""Optional reference to the complete message that may come after the chunks.
This allows consumers of the stream to correlate chunks with the eventual completed message."""
type: Literal["ModelClientStreamingChunkEvent"] = "ModelClientStreamingChunkEvent"
def to_text(self) -> str:
return self.content
class ThoughtEvent(BaseAgentEvent):
"""An event signaling the thought process of a model.
It is used to communicate the reasoning tokens generated by a reasoning model,
or the extra text content generated by a function call."""
content: str
"""The thought process of the model."""
type: Literal["ThoughtEvent"] = "ThoughtEvent"
def to_text(self) -> str:
return self.content
class SelectSpeakerEvent(BaseAgentEvent):
"""An event signaling the selection of speakers for a conversation."""
content: List[str]
"""The names of the selected speakers."""
type: Literal["SelectSpeakerEvent"] = "SelectSpeakerEvent"
def to_text(self) -> str:
return str(self.content)
class SelectorEvent(BaseAgentEvent):
"""An event emitted from the `SelectorGroupChat`."""
content: str
"""The content of the event."""
type: Literal["SelectorEvent"] = "SelectorEvent"
def to_text(self) -> str:
return str(self.content)
class MessageFactory:
""":meta private:
A factory for creating messages from JSON-serializable dictionaries.
This is useful for deserializing messages from JSON data.
"""
def __init__(self) -> None:
self._message_types: Dict[str, type[BaseAgentEvent | BaseChatMessage]] = {}
# Register all message types.
self._message_types[TextMessage.__name__] = TextMessage
self._message_types[MultiModalMessage.__name__] = MultiModalMessage
self._message_types[StopMessage.__name__] = StopMessage
self._message_types[ToolCallSummaryMessage.__name__] = ToolCallSummaryMessage
self._message_types[HandoffMessage.__name__] = HandoffMessage
self._message_types[ToolCallRequestEvent.__name__] = ToolCallRequestEvent
self._message_types[ToolCallExecutionEvent.__name__] = ToolCallExecutionEvent
self._message_types[MemoryQueryEvent.__name__] = MemoryQueryEvent
self._message_types[UserInputRequestedEvent.__name__] = UserInputRequestedEvent
self._message_types[ModelClientStreamingChunkEvent.__name__] = ModelClientStreamingChunkEvent
self._message_types[ThoughtEvent.__name__] = ThoughtEvent
self._message_types[SelectSpeakerEvent.__name__] = SelectSpeakerEvent
self._message_types[CodeGenerationEvent.__name__] = CodeGenerationEvent
self._message_types[CodeExecutionEvent.__name__] = CodeExecutionEvent
def is_registered(self, message_type: type[BaseAgentEvent | BaseChatMessage]) -> bool:
"""Check if a message type is registered with the factory."""
# Get the class name of the message type.
class_name = message_type.__name__
# Check if the class name is already registered.
return class_name in self._message_types
def register(self, message_type: type[BaseAgentEvent | BaseChatMessage]) -> None:
"""Register a new message type with the factory."""
if self.is_registered(message_type):
raise ValueError(f"Message type {message_type} is already registered.")
if not issubclass(message_type, BaseChatMessage) and not issubclass(message_type, BaseAgentEvent):
raise ValueError(f"Message type {message_type} must be a subclass of BaseChatMessage or BaseAgentEvent.")
# Get the class name of the
class_name = message_type.__name__
# Check if the class name is already registered.
# Register the message type.
self._message_types[class_name] = message_type
def create(self, data: Mapping[str, Any]) -> BaseAgentEvent | BaseChatMessage:
"""Create a message from a dictionary of JSON-serializable data."""
# Get the type of the message from the dictionary.
message_type = data.get("type")
if message_type is None:
raise ValueError("Field 'type' is required in the message data to recover the message type.")
if message_type not in self._message_types:
raise ValueError(f"Unknown message type: {message_type}")
if not isinstance(message_type, str):
raise ValueError(f"Message type must be a string, got {type(message_type)}")
# Get the class for the message type.
message_class = self._message_types[message_type]
# Create an instance of the message class.
assert issubclass(message_class, BaseChatMessage) or issubclass(message_class, BaseAgentEvent)
return message_class.load(data)
ChatMessage = Annotated[
TextMessage | MultiModalMessage | StopMessage | ToolCallSummaryMessage | HandoffMessage,
Field(discriminator="type"),
]
"""The union type of all built-in concrete subclasses of :class:`BaseChatMessage`.
It does not include :class:`StructuredMessage` types."""
AgentEvent = Annotated[
ToolCallRequestEvent
| ToolCallExecutionEvent
| MemoryQueryEvent
| UserInputRequestedEvent
| ModelClientStreamingChunkEvent
| ThoughtEvent
| SelectSpeakerEvent
| CodeGenerationEvent
| CodeExecutionEvent,
Field(discriminator="type"),
]
"""The union type of all built-in concrete subclasses of :class:`BaseAgentEvent`."""
__all__ = [
"AgentEvent",
"BaseMessage",
"ChatMessage",
"BaseChatMessage",
"BaseAgentEvent",
"BaseTextChatMessage",
"StructuredContentType",
"StructuredMessage",
"StructuredMessageFactory",
"HandoffMessage",
"MultiModalMessage",
"StopMessage",
"TextMessage",
"ToolCallExecutionEvent",
"ToolCallRequestEvent",
"ToolCallSummaryMessage",
"MemoryQueryEvent",
"UserInputRequestedEvent",
"ModelClientStreamingChunkEvent",
"ThoughtEvent",
"SelectSpeakerEvent",
"MessageFactory",
"CodeGenerationEvent",
"CodeExecutionEvent",
]

0
agent_dhal/agentdhal_agentchat/py.typed Normal file
View File

27
agent_dhal/agentdhal_agentchat/state/__init__.py Normal file
View File

@@ -0,0 +1,27 @@
"""State management for agents, teams and termination conditions."""
from ._states import (
AssistantAgentState,
BaseGroupChatManagerState,
BaseState,
ChatAgentContainerState,
MagenticOneOrchestratorState,
RoundRobinManagerState,
SelectorManagerState,
SocietyOfMindAgentState,
SwarmManagerState,
TeamState,
)
__all__ = [
"BaseState",
"AssistantAgentState",
"BaseGroupChatManagerState",
"ChatAgentContainerState",
"RoundRobinManagerState",
"SelectorManagerState",
"SwarmManagerState",
"MagenticOneOrchestratorState",
"TeamState",
"SocietyOfMindAgentState",
]

79
agent_dhal/agentdhal_agentchat/state/_states.py Normal file
View File

@@ -0,0 +1,79 @@
from typing import Any, List, Mapping, Optional
from pydantic import BaseModel, Field
class BaseState(BaseModel):
"""Base class for all saveable state"""
type: str = Field(default="BaseState")
version: str = Field(default="1.0.0")
class HalState(BaseState):
"""State for an assistant agent."""
llm_context: Mapping[str, Any] = Field(default_factory=lambda: dict([("messages", [])]))
type: str = Field(default="AssistantAgentState")
class TeamState(BaseState):
"""State for a team of agents."""
agent_states: Mapping[str, Any] = Field(default_factory=dict)
type: str = Field(default="TeamState")
class BaseGroupChatManagerState(BaseState):
"""Base state for all group chat managers."""
message_thread: List[Mapping[str, Any]] = Field(default_factory=list)
current_turn: int = Field(default=0)
type: str = Field(default="BaseGroupChatManagerState")
class ChatAgentContainerState(BaseState):
"""State for a container of chat agents."""
agent_state: Mapping[str, Any] = Field(default_factory=dict)
message_buffer: List[Mapping[str, Any]] = Field(default_factory=list)
type: str = Field(default="ChatAgentContainerState")
class RoundRobinManagerState(BaseGroupChatManagerState):
"""State for :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` manager."""
next_speaker_index: int = Field(default=0)
type: str = Field(default="RoundRobinManagerState")
class SelectorManagerState(BaseGroupChatManagerState):
"""State for :class:`~agentdhal_agentchat.teams.SelectorGroupChat` manager."""
previous_speaker: Optional[str] = Field(default=None)
type: str = Field(default="SelectorManagerState")
class SwarmManagerState(BaseGroupChatManagerState):
"""State for :class:`~agentdhal_agentchat.teams.Swarm` manager."""
current_speaker: str = Field(default="")
type: str = Field(default="SwarmManagerState")
class MagenticOneOrchestratorState(BaseGroupChatManagerState):
"""State for :class:`~agentdhal_agentchat.teams.MagneticOneGroupChat` orchestrator."""
task: str = Field(default="")
facts: str = Field(default="")
plan: str = Field(default="")
n_rounds: int = Field(default=0)
n_stalls: int = Field(default=0)
type: str = Field(default="MagenticOneOrchestratorState")
class SocietyOfMindAgentState(BaseState):
"""State for a Society of Mind agent."""
inner_team_state: Mapping[str, Any] = Field(default_factory=dict)
type: str = Field(default="SocietyOfMindAgentState")

30
agent_dhal/agentdhal_agentchat/teams/__init__.py Normal file
View File

@@ -0,0 +1,30 @@
"""
This module provides implementation of various pre-defined multi-agent teams.
Each team inherits from the BaseGroupChat class.
"""
from ._group_chat._base_group_chat import BaseGroupChat
from ._group_chat._graph import (
DiGraph,
DiGraphBuilder,
DiGraphEdge,
DiGraphNode,
GraphFlow,
)
from ._group_chat._magentic_one import MagenticOneGroupChat
from ._group_chat._round_robin_group_chat import RoundRobinGroupChat
from ._group_chat._selector_group_chat import SelectorGroupChat
from ._group_chat._swarm_group_chat import Swarm
__all__ = [
"BaseGroupChat",
"RoundRobinGroupChat",
"SelectorGroupChat",
"Swarm",
"MagenticOneGroupChat",
"DiGraphBuilder",
"DiGraph",
"DiGraphNode",
"DiGraphEdge",
"GraphFlow",
]

0
agent_dhal/agentdhal_agentchat/teams/_group_chat/__init__.py Normal file
View File

834
agent_dhal/agentdhal_agentchat/teams/_group_chat/_base_group_chat.py Normal file
View File

@@ -0,0 +1,834 @@
import asyncio
import uuid
from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, Callable, Dict, List, Mapping, Sequence
from agentdhal_core import (
AgentId,
AgentRuntime,
AgentType,
CancellationToken,
ComponentBase,
SingleThreadedAgentRuntime,
TypeSubscription,
)
from pydantic import BaseModel, ValidationError
from ...base import ChatAgent, TaskResult, Team, TerminationCondition
from ...messages import (
BaseAgentEvent,
BaseChatMessage,
MessageFactory,
ModelClientStreamingChunkEvent,
StopMessage,
StructuredMessage,
TextMessage,
)
from ...state import TeamState
from ._chat_agent_container import ChatAgentContainer
from ._events import (
GroupChatPause,
GroupChatReset,
GroupChatResume,
GroupChatStart,
GroupChatTermination,
SerializableException,
)
from ._sequential_routed_agent import SequentialRoutedAgent
class BaseGroupChat(Team, ABC, ComponentBase[BaseModel]):
"""The base class for group chat teams.
In a group chat team, participants share context by publishing their messages
to all other participants.
If an :class:`~agentdhal_agentchat.base.ChatAgent` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage` from the agent response's
:attr:`~agentdhal_agentchat.base.Response.chat_message` will be published
to other participants in the group chat.
If a :class:`~agentdhal_agentchat.base.Team` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage`
from the team result' :attr:`~agentdhal_agentchat.base.TaskResult.messages` will be published
to other participants in the group chat.
To implement a group chat team, first create a subclass of :class:`BaseGroupChatManager` and then
create a subclass of :class:`BaseGroupChat` that uses the group chat manager.
This base class provides the mapping between the agents of the AgentChat API
and the agent runtime of the Core API, and handles high-level features like
running, pausing, resuming, and resetting the team.
"""
component_type = "team"
def __init__(
self,
name: str,
description: str,
participants: List[ChatAgent | Team],
group_chat_manager_name: str,
group_chat_manager_class: type[SequentialRoutedAgent],
termination_condition: TerminationCondition | None = None,
max_turns: int | None = None,
runtime: AgentRuntime | None = None,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
emit_team_events: bool = False,
):
self._name = name
self._description = description
if len(participants) == 0:
raise ValueError("At least one participant is required.")
if len(participants) != len(set(participant.name for participant in participants)):
raise ValueError("The participant names must be unique.")
self._participants = participants
self._base_group_chat_manager_class = group_chat_manager_class
self._termination_condition = termination_condition
self._max_turns = max_turns
self._message_factory = MessageFactory()
if custom_message_types is not None:
for message_type in custom_message_types:
self._message_factory.register(message_type)
for agent in participants:
if isinstance(agent, ChatAgent):
for message_type in agent.produced_message_types:
try:
is_registered = self._message_factory.is_registered(message_type) # type: ignore[reportUnknownArgumentType]
if issubclass(message_type, StructuredMessage) and not is_registered:
self._message_factory.register(message_type) # type: ignore[reportUnknownArgumentType]
except TypeError:
# Not a class or not a valid subclassable type (skip)
pass
# The team ID is a UUID that is used to identify the team and its participants
# in the agent runtime. It is used to create unique topic types for each participant.
# Currently, team ID is binded to an object instance of the group chat class.
# So if you create two instances of group chat, there will be two teams with different IDs.
self._team_id = str(uuid.uuid4())
# Constants for the group chat team.
# The names are used to identify the agents within the team.
# The names may not be unique across different teams.
self._group_chat_manager_name = group_chat_manager_name
self._participant_names: List[str] = [participant.name for participant in participants]
self._participant_descriptions: List[str] = [participant.description for participant in participants]
# The group chat topic type is used for broadcast communication among all participants and the group chat manager.
self._group_topic_type = f"group_topic_{self._team_id}"
# The group chat manager topic type is used for direct communication with the group chat manager.
self._group_chat_manager_topic_type = f"{self._group_chat_manager_name}_{self._team_id}"
# The participant topic types are used for direct communication with each participant.
self._participant_topic_types: List[str] = [
f"{participant.name}_{self._team_id}" for participant in participants
]
# The output topic type is used for emitting streaming messages from the group chat.
# The group chat manager will relay the messages to the output message queue.
self._output_topic_type = f"output_topic_{self._team_id}"
# The queue for collecting the output messages.
self._output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination] = (
asyncio.Queue()
)
# Create a runtime for the team.
if runtime is not None:
self._runtime = runtime
self._embedded_runtime = False
else:
# Use a embedded single-threaded runtime for the group chat.
# Background exceptions must not be ignored as it results in non-surfaced exceptions and early team termination.
self._runtime = SingleThreadedAgentRuntime(ignore_unhandled_exceptions=False)
self._embedded_runtime = True
# Flag to track if the group chat has been initialized.
self._initialized = False
# Flag to track if the group chat is running.
self._is_running = False
# Flag to track if the team events should be emitted.
self._emit_team_events = emit_team_events
@property
def name(self) -> str:
"""The name of the group chat team."""
return self._name
@property
def description(self) -> str:
"""A description of the group chat team."""
return self._description
@abstractmethod
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], SequentialRoutedAgent]: ...
def _create_participant_factory(
self,
parent_topic_type: str,
output_topic_type: str,
agent: ChatAgent | Team,
message_factory: MessageFactory,
) -> Callable[[], ChatAgentContainer]:
def _factory() -> ChatAgentContainer:
container = ChatAgentContainer(parent_topic_type, output_topic_type, agent, message_factory)
return container
return _factory
async def _init(self, runtime: AgentRuntime) -> None:
# Constants for the group chat manager.
group_chat_manager_agent_type = AgentType(self._group_chat_manager_topic_type)
# Register participants.
# Use the participant topic type as the agent type.
for participant, agent_type in zip(self._participants, self._participant_topic_types, strict=True):
# Register the participant factory.
await ChatAgentContainer.register(
runtime,
type=agent_type,
factory=self._create_participant_factory(
self._group_topic_type, self._output_topic_type, participant, self._message_factory
),
)
# Add subscriptions for the participant.
# The participant should be able to receive messages from its own topic.
await runtime.add_subscription(TypeSubscription(topic_type=agent_type, agent_type=agent_type))
# The participant should be able to receive messages from the group topic.
await runtime.add_subscription(TypeSubscription(topic_type=self._group_topic_type, agent_type=agent_type))
# Register the group chat manager.
await self._base_group_chat_manager_class.register(
runtime,
type=group_chat_manager_agent_type.type,
factory=self._create_group_chat_manager_factory(
name=self._group_chat_manager_name,
group_topic_type=self._group_topic_type,
output_topic_type=self._output_topic_type,
participant_names=self._participant_names,
participant_topic_types=self._participant_topic_types,
participant_descriptions=self._participant_descriptions,
output_message_queue=self._output_message_queue,
termination_condition=self._termination_condition,
max_turns=self._max_turns,
message_factory=self._message_factory,
),
)
# Add subscriptions for the group chat manager.
# The group chat manager should be able to receive messages from the its own topic.
await runtime.add_subscription(
TypeSubscription(
topic_type=self._group_chat_manager_topic_type, agent_type=group_chat_manager_agent_type.type
)
)
# The group chat manager should be able to receive messages from the group topic.
await runtime.add_subscription(
TypeSubscription(topic_type=self._group_topic_type, agent_type=group_chat_manager_agent_type.type)
)
# The group chat manager will relay the messages from output topic to the output message queue.
await runtime.add_subscription(
TypeSubscription(topic_type=self._output_topic_type, agent_type=group_chat_manager_agent_type.type)
)
self._initialized = True
async def run(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> TaskResult:
"""Run the team and return the result. The base implementation uses
:meth:`run_stream` to run the team and then returns the final result.
Once the team is stopped, the termination condition is reset.
Args:
task (str | BaseChatMessage | Sequence[BaseChatMessage] | None): The task to run the team with. Can be a string, a single :class:`BaseChatMessage` , or a list of :class:`BaseChatMessage`.
cancellation_token (CancellationToken | None): The cancellation token to kill the task immediately.
Setting the cancellation token potentially put the team in an inconsistent state,
and it may not reset the termination condition.
To gracefully stop the team, use :class:`~agentdhal_agentchat.conditions.ExternalTermination` instead.
Returns:
result: The result of the task as :class:`~agentdhal_agentchat.base.TaskResult`. The result contains the messages produced by the team and the stop reason.
Example using the :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` team:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = MaxMessageTermination(3)
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
result = await team.run(task="Count from 1 to 10, respond one at a time.")
print(result)
# Run the team again without a task to continue the previous task.
result = await team.run()
print(result)
asyncio.run(main())
Example using the :class:`~agentdhal_core.CancellationToken` to cancel the task:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_core import CancellationToken
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = MaxMessageTermination(3)
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
cancellation_token = CancellationToken()
# Create a task to run the team in the background.
run_task = asyncio.create_task(
team.run(
task="Count from 1 to 10, respond one at a time.",
cancellation_token=cancellation_token,
)
)
# Wait for 1 second and then cancel the task.
await asyncio.sleep(1)
cancellation_token.cancel()
# This will raise a cancellation error.
await run_task
asyncio.run(main())
"""
result: TaskResult | None = None
async for message in self.run_stream(
task=task,
cancellation_token=cancellation_token,
output_task_messages=output_task_messages,
):
if isinstance(message, TaskResult):
result = message
if result is not None:
return result
raise AssertionError("The stream should have returned the final result.")
async def run_stream(
self,
*,
task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
cancellation_token: CancellationToken | None = None,
output_task_messages: bool = True,
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
"""Run the team and produces a stream of messages and the final result
of the type :class:`~agentdhal_agentchat.base.TaskResult` as the last item in the stream. Once the
team is stopped, the termination condition is reset.
.. note::
If an agent produces :class:`~agentdhal_agentchat.messages.ModelClientStreamingChunkEvent`,
the message will be yielded in the stream but it will not be included in the
:attr:`~agentdhal_agentchat.base.TaskResult.messages`.
Args:
task (str | BaseChatMessage | Sequence[BaseChatMessage] | None): The task to run the team with. Can be a string, a single :class:`BaseChatMessage` , or a list of :class:`BaseChatMessage`.
cancellation_token (CancellationToken | None): The cancellation token to kill the task immediately.
Setting the cancellation token potentially put the team in an inconsistent state,
and it may not reset the termination condition.
To gracefully stop the team, use :class:`~agentdhal_agentchat.conditions.ExternalTermination` instead.
output_task_messages (bool): Whether to include task messages in the output stream. Defaults to True for backward compatibility.
Returns:
stream: an :class:`~collections.abc.AsyncGenerator` that yields :class:`~agentdhal_agentchat.messages.BaseAgentEvent`, :class:`~agentdhal_agentchat.messages.BaseChatMessage`, and the final result :class:`~agentdhal_agentchat.base.TaskResult` as the last item in the stream.
Example using the :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` team:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = MaxMessageTermination(3)
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
stream = team.run_stream(task="Count from 1 to 10, respond one at a time.")
async for message in stream:
print(message)
# Run the team again without a task to continue the previous task.
stream = team.run_stream()
async for message in stream:
print(message)
asyncio.run(main())
Example using the :class:`~agentdhal_core.CancellationToken` to cancel the task:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.ui import Console
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_core import CancellationToken
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = MaxMessageTermination(3)
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
cancellation_token = CancellationToken()
# Create a task to run the team in the background.
run_task = asyncio.create_task(
Console(
team.run_stream(
task="Count from 1 to 10, respond one at a time.",
cancellation_token=cancellation_token,
)
)
)
# Wait for 1 second and then cancel the task.
await asyncio.sleep(1)
cancellation_token.cancel()
# This will raise a cancellation error.
await run_task
asyncio.run(main())
"""
# Create the messages list if the task is a string or a chat message.
messages: List[BaseChatMessage] | None = None
if task is None:
pass
elif isinstance(task, str):
messages = [TextMessage(content=task, source="user")]
elif isinstance(task, BaseChatMessage):
messages = [task]
elif isinstance(task, list):
if not task:
raise ValueError("Task list cannot be empty.")
messages = []
for msg in task:
if not isinstance(msg, BaseChatMessage):
raise ValueError("All messages in task list must be valid BaseChatMessage types")
messages.append(msg)
else:
raise ValueError("Task must be a string, a BaseChatMessage, or a list of BaseChatMessage.")
# Check if the messages types are registered with the message factory.
if messages is not None:
for msg in messages:
if not self._message_factory.is_registered(msg.__class__):
raise ValueError(
f"Message type {msg.__class__} is not registered with the message factory. "
"Please register it with the message factory by adding it to the "
"custom_message_types list when creating the team."
)
if self._is_running:
raise ValueError("The team is already running, it cannot run again until it is stopped.")
self._is_running = True
if self._embedded_runtime:
# Start the embedded runtime.
assert isinstance(self._runtime, SingleThreadedAgentRuntime)
self._runtime.start()
if not self._initialized:
await self._init(self._runtime)
shutdown_task: asyncio.Task[None] | None = None
if self._embedded_runtime:
async def stop_runtime() -> None:
assert isinstance(self._runtime, SingleThreadedAgentRuntime)
try:
# This will propagate any exceptions raised.
await self._runtime.stop_when_idle()
# Put a termination message in the queue to indicate that the group chat is stopped for whatever reason
# but not due to an exception.
await self._output_message_queue.put(
GroupChatTermination(
message=StopMessage(
content="The group chat is stopped.", source=self._group_chat_manager_name
)
)
)
except Exception as e:
# Stop the consumption of messages and end the stream.
# NOTE: we also need to put a GroupChatTermination event here because when the runtime
# has an exception, the group chat manager may not be able to put a GroupChatTermination event in the queue.
# This may not be necessary if the group chat manager is able to handle the exception and put the event in the queue.
await self._output_message_queue.put(
GroupChatTermination(
message=StopMessage(
content="An exception occurred in the runtime.", source=self._group_chat_manager_name
),
error=SerializableException.from_exception(e),
)
)
# Create a background task to stop the runtime when the group chat
# is stopped or has an exception.
shutdown_task = asyncio.create_task(stop_runtime())
try:
# Run the team by sending the start message to the group chat manager.
# The group chat manager will start the group chat by relaying the message to the participants
# and the group chat manager.
await self._runtime.send_message(
GroupChatStart(messages=messages, output_task_messages=output_task_messages),
recipient=AgentId(type=self._group_chat_manager_topic_type, key=self._team_id),
cancellation_token=cancellation_token,
)
# Collect the output messages in order.
output_messages: List[BaseAgentEvent | BaseChatMessage] = []
stop_reason: str | None = None
# Yield the messages until the queue is empty.
while True:
message_future = asyncio.ensure_future(self._output_message_queue.get())
if cancellation_token is not None:
cancellation_token.link_future(message_future)
# Wait for the next message, this will raise an exception if the task is cancelled.
message = await message_future
if isinstance(message, GroupChatTermination):
# If the message contains an error, we need to raise it here.
# This will stop the team and propagate the error.
if message.error is not None:
raise RuntimeError(str(message.error))
stop_reason = message.message.content
break
yield message
if isinstance(message, ModelClientStreamingChunkEvent):
# Skip the model client streaming chunk events.
continue
output_messages.append(message)
# Yield the final result.
yield TaskResult(messages=output_messages, stop_reason=stop_reason)
finally:
try:
if shutdown_task is not None:
# Wait for the shutdown task to finish.
# This will propagate any exceptions raised.
await shutdown_task
finally:
# Clear the output message queue.
while not self._output_message_queue.empty():
self._output_message_queue.get_nowait()
# Indicate that the team is no longer running.
self._is_running = False
async def reset(self) -> None:
"""Reset the team and its participants to their initial state.
The team must be stopped before it can be reset.
Raises:
RuntimeError: If the team has not been initialized or is currently running.
Example using the :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` team:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = MaxMessageTermination(3)
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
stream = team.run_stream(task="Count from 1 to 10, respond one at a time.")
async for message in stream:
print(message)
# Reset the team.
await team.reset()
stream = team.run_stream(task="Count from 1 to 10, respond one at a time.")
async for message in stream:
print(message)
asyncio.run(main())
"""
if not self._initialized:
await self._init(self._runtime)
if self._is_running:
raise RuntimeError("The group chat is currently running. It must be stopped before it can be reset.")
self._is_running = True
if self._embedded_runtime:
# Start the runtime.
assert isinstance(self._runtime, SingleThreadedAgentRuntime)
self._runtime.start()
try:
# Send a reset messages to all participants.
for participant_topic_type in self._participant_topic_types:
await self._runtime.send_message(
GroupChatReset(),
recipient=AgentId(type=participant_topic_type, key=self._team_id),
)
# Send a reset message to the group chat manager.
await self._runtime.send_message(
GroupChatReset(),
recipient=AgentId(type=self._group_chat_manager_topic_type, key=self._team_id),
)
finally:
if self._embedded_runtime:
# Stop the runtime.
assert isinstance(self._runtime, SingleThreadedAgentRuntime)
await self._runtime.stop_when_idle()
# Reset the output message queue.
while not self._output_message_queue.empty():
self._output_message_queue.get_nowait()
# Indicate that the team is no longer running.
self._is_running = False
async def pause(self) -> None:
"""Pause its participants when the team is running by calling their
:meth:`~agentdhal_agentchat.base.ChatAgent.on_pause` method via direct RPC calls.
.. attention::
This is an experimental feature introduced in v0.4.9 and may subject
to change or removal in the future.
The team must be initialized before it can be paused.
Different from termination, pausing the team does not cause the
:meth:`run` or :meth:`run_stream` method to return. It calls the
:meth:`~agentdhal_agentchat.base.ChatAgent.on_pause` method on each
participant, and if the participant does not implement the method, it
will be a no-op.
.. note::
It is the responsibility of the agent class to handle the pause
and ensure that the agent can be resumed later.
Make sure to implement the :meth:`~agentdhal_agentchat.agents.BaseChatAgent.on_pause`
method in your agent class for custom pause behavior.
By default, the agent will not do anything when called.
Raises:
RuntimeError: If the team has not been initialized. Exceptions from
the participants when calling their implementations of
:class:`~agentdhal_agentchat.base.ChatAgent.on_pause` are
propagated to this method and raised.
"""
if not self._initialized:
raise RuntimeError("The group chat has not been initialized. It must be run before it can be paused.")
# Send a pause message to all participants.
for participant_topic_type in self._participant_topic_types:
await self._runtime.send_message(
GroupChatPause(),
recipient=AgentId(type=participant_topic_type, key=self._team_id),
)
# Send a pause message to the group chat manager.
await self._runtime.send_message(
GroupChatPause(),
recipient=AgentId(type=self._group_chat_manager_topic_type, key=self._team_id),
)
async def resume(self) -> None:
"""Resume its participants when the team is running and paused by calling their
:meth:`~agentdhal_agentchat.base.ChatAgent.on_resume` method via direct RPC calls.
.. attention::
This is an experimental feature introduced in v0.4.9 and may subject
to change or removal in the future.
The team must be initialized before it can be resumed.
Different from termination and restart with a new task, resuming the team
does not cause the :meth:`run` or :meth:`run_stream` method to return.
It calls the :meth:`~agentdhal_agentchat.base.ChatAgent.on_resume` method on each
participant, and if the participant does not implement the method, it
will be a no-op.
.. note::
It is the responsibility of the agent class to handle the resume
and ensure that the agent continues from where it was paused.
Make sure to implement the :meth:`~agentdhal_agentchat.agents.BaseChatAgent.on_resume`
method in your agent class for custom resume behavior.
Raises:
RuntimeError: If the team has not been initialized. Exceptions from
the participants when calling their implementations of :class:`~agentdhal_agentchat.base.ChatAgent.on_resume`
method are propagated to this method and raised.
"""
if not self._initialized:
raise RuntimeError("The group chat has not been initialized. It must be run before it can be resumed.")
# Send a resume message to all participants.
for participant_topic_type in self._participant_topic_types:
await self._runtime.send_message(
GroupChatResume(),
recipient=AgentId(type=participant_topic_type, key=self._team_id),
)
# Send a resume message to the group chat manager.
await self._runtime.send_message(
GroupChatResume(),
recipient=AgentId(type=self._group_chat_manager_topic_type, key=self._team_id),
)
async def save_state(self) -> Mapping[str, Any]:
"""Save the state of the group chat team.
The state is saved by calling the :meth:`~agentdhal_core.AgentRuntime.agent_save_state` method
on each participant and the group chat manager with their internal agent ID.
The state is returned as a nested dictionary: a dictionary with key `agent_states`,
which is a dictionary the agent names as keys and the state as values.
.. code-block:: text
{
"agent_states": {
"agent1": ...,
"agent2": ...,
"RoundRobinGroupChatManager": ...
}
}
.. note::
Starting v0.4.9, the state is using the agent name as the key instead of the agent ID,
and the `team_id` field is removed from the state. This is to allow the state to be
portable across different teams and runtimes. States saved with the old format
may not be compatible with the new format in the future.
.. caution::
When calling :func:`~agentdhal_agentchat.teams.BaseGroupChat.save_state` on a team
while it is running, the state may not be consistent and may result in an unexpected state.
It is recommended to call this method when the team is not running or after it is stopped.
"""
if not self._initialized:
await self._init(self._runtime)
# Store state of each agent by their name.
# NOTE: we don't use the agent ID as the key here because we need to be able to decouple
# the state of the agents from their identities in the agent runtime.
agent_states: Dict[str, Mapping[str, Any]] = {}
# Save the state of all participants.
for name, agent_type in zip(self._participant_names, self._participant_topic_types, strict=True):
agent_id = AgentId(type=agent_type, key=self._team_id)
# NOTE: We are using the runtime's save state method rather than the agent instance's
# save_state method because we want to support saving state of remote agents.
agent_states[name] = await self._runtime.agent_save_state(agent_id)
# Save the state of the group chat manager.
agent_id = AgentId(type=self._group_chat_manager_topic_type, key=self._team_id)
agent_states[self._group_chat_manager_name] = await self._runtime.agent_save_state(agent_id)
return TeamState(agent_states=agent_states).model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Load an external state and overwrite the current state of the group chat team.
The state is loaded by calling the :meth:`~agentdhal_core.AgentRuntime.agent_load_state` method
on each participant and the group chat manager with their internal agent ID.
See :meth:`~agentdhal_agentchat.teams.BaseGroupChat.save_state` for the expected format of the state.
"""
if not self._initialized:
await self._init(self._runtime)
if self._is_running:
raise RuntimeError("The team cannot be loaded while it is running.")
self._is_running = True
try:
team_state = TeamState.model_validate(state)
# Load the state of all participants.
for name, agent_type in zip(self._participant_names, self._participant_topic_types, strict=True):
agent_id = AgentId(type=agent_type, key=self._team_id)
if name not in team_state.agent_states:
raise ValueError(f"Agent state for {name} not found in the saved state.")
await self._runtime.agent_load_state(agent_id, team_state.agent_states[name])
# Load the state of the group chat manager.
agent_id = AgentId(type=self._group_chat_manager_topic_type, key=self._team_id)
if self._group_chat_manager_name not in team_state.agent_states:
raise ValueError(f"Agent state for {self._group_chat_manager_name} not found in the saved state.")
await self._runtime.agent_load_state(agent_id, team_state.agent_states[self._group_chat_manager_name])
except ValidationError as e:
raise ValueError(
"Invalid state format. The expected state format has changed since v0.4.9. "
"Please read the release note on GitHub."
) from e
finally:
# Indicate that the team is no longer running.
self._is_running = False

326
agent_dhal/agentdhal_agentchat/teams/_group_chat/_base_group_chat_manager.py Normal file
View File

@@ -0,0 +1,326 @@
import asyncio
from abc import ABC, abstractmethod
from typing import Any, List, Sequence
from agentdhal_core import CancellationToken, DefaultTopicId, MessageContext, event, rpc
from ...base import TerminationCondition
from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory, SelectSpeakerEvent, StopMessage
from ._events import (
GroupChatAgentResponse,
GroupChatError,
GroupChatMessage,
GroupChatPause,
GroupChatRequestPublish,
GroupChatReset,
GroupChatResume,
GroupChatStart,
GroupChatTeamResponse,
GroupChatTermination,
SerializableException,
)
from ._sequential_routed_agent import SequentialRoutedAgent
class BaseGroupChatManager(SequentialRoutedAgent, ABC):
"""Base class for a group chat manager that manages a group chat with multiple participants.
It is the responsibility of the caller to ensure:
- All participants must subscribe to the group chat topic and each of their own topics.
- The group chat manager must subscribe to the group chat topic.
- The agent types of the participants must be unique.
- For each participant, the agent type must be the same as the topic type.
Without the above conditions, the group chat will not function correctly.
"""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
emit_team_events: bool = False,
):
super().__init__(
description="Group chat manager",
sequential_message_types=[
GroupChatStart,
GroupChatAgentResponse,
GroupChatTeamResponse,
GroupChatMessage,
GroupChatReset,
],
)
if max_turns is not None and max_turns <= 0:
raise ValueError("The maximum number of turns must be greater than 0.")
if len(participant_topic_types) != len(participant_descriptions):
raise ValueError("The number of participant topic types, agent types, and descriptions must be the same.")
if len(set(participant_topic_types)) != len(participant_topic_types):
raise ValueError("The participant topic types must be unique.")
if group_topic_type in participant_topic_types:
raise ValueError("The group topic type must not be in the participant topic types.")
self._name = name
self._group_topic_type = group_topic_type
self._output_topic_type = output_topic_type
self._participant_names = participant_names
self._participant_name_to_topic_type = {
name: topic_type for name, topic_type in zip(participant_names, participant_topic_types, strict=True)
}
self._participant_descriptions = participant_descriptions
self._message_thread: List[BaseAgentEvent | BaseChatMessage] = []
self._output_message_queue = output_message_queue
self._termination_condition = termination_condition
self._max_turns = max_turns
self._current_turn = 0
self._message_factory = message_factory
self._emit_team_events = emit_team_events
self._active_speakers: List[str] = []
@rpc
async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None:
"""Handle the start of a group chat by selecting a speaker to start the conversation."""
# Check if the conversation has already terminated.
if self._termination_condition is not None and self._termination_condition.terminated:
early_stop_message = StopMessage(
content="The group chat has already terminated.",
source=self._name,
)
# Signal termination to the caller of the team.
await self._signal_termination(early_stop_message)
# Stop the group chat.
return
# Validate the group state given the start messages
await self.validate_group_state(message.messages)
if message.messages is not None:
# Log all messages at once
await self.publish_message(
GroupChatStart(messages=message.messages),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Only put messages in output queue if output_task_messages is True
if message.output_task_messages:
for msg in message.messages:
await self._output_message_queue.put(msg)
# Relay all messages at once to participants
await self.publish_message(
GroupChatStart(messages=message.messages),
topic_id=DefaultTopicId(type=self._group_topic_type),
cancellation_token=ctx.cancellation_token,
)
# Append all messages to thread
await self.update_message_thread(message.messages)
# Check termination condition after processing all messages
if await self._apply_termination_condition(message.messages):
# Stop the group chat.
return
# Select speakers to start/continue the conversation
await self._transition_to_next_speakers(ctx.cancellation_token)
@event
async def handle_agent_response(
self, message: GroupChatAgentResponse | GroupChatTeamResponse, ctx: MessageContext
) -> None:
try:
# Construct the detla from the agent response.
delta: List[BaseAgentEvent | BaseChatMessage] = []
if isinstance(message, GroupChatAgentResponse):
if message.response.inner_messages is not None:
for inner_message in message.response.inner_messages:
delta.append(inner_message)
delta.append(message.response.chat_message)
else:
delta.extend(message.result.messages)
# Append the messages to the message thread.
await self.update_message_thread(delta)
# Remove the agent from the active speakers list.
self._active_speakers.remove(message.name)
if len(self._active_speakers) > 0:
# If there are still active speakers, return without doing anything.
return
# Check if the conversation should be terminated.
if await self._apply_termination_condition(delta, increment_turn_count=True):
# Stop the group chat.
return
# Select speakers to continue the conversation.
await self._transition_to_next_speakers(ctx.cancellation_token)
except Exception as e:
# Handle the exception and signal termination with an error.
error = SerializableException.from_exception(e)
await self._signal_termination_with_error(error)
# Raise the exception to the runtime.
raise
async def _transition_to_next_speakers(self, cancellation_token: CancellationToken) -> None:
speaker_names_future = asyncio.ensure_future(self.select_speaker(self._message_thread))
# Link the select speaker future to the cancellation token.
cancellation_token.link_future(speaker_names_future)
speaker_names = await speaker_names_future
if isinstance(speaker_names, str):
# If only one speaker is selected, convert it to a list.
speaker_names = [speaker_names]
for speaker_name in speaker_names:
if speaker_name not in self._participant_name_to_topic_type:
raise RuntimeError(f"Speaker {speaker_name} not found in participant names.")
await self._log_speaker_selection(speaker_names)
# Send request to publish message to the next speakers
for speaker_name in speaker_names:
speaker_topic_type = self._participant_name_to_topic_type[speaker_name]
await self.publish_message(
GroupChatRequestPublish(),
topic_id=DefaultTopicId(type=speaker_topic_type),
cancellation_token=cancellation_token,
)
self._active_speakers.append(speaker_name)
async def _apply_termination_condition(
self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False
) -> bool:
"""Apply the termination condition to the delta and return True if the conversation should be terminated.
It also resets the termination condition and turn count, and signals termination to the caller of the team."""
if self._termination_condition is not None:
stop_message = await self._termination_condition(delta)
if stop_message is not None:
# Reset the termination conditions and turn count.
await self._termination_condition.reset()
self._current_turn = 0
# Signal termination to the caller of the team.
await self._signal_termination(stop_message)
# Stop the group chat.
return True
if increment_turn_count:
# Increment the turn count.
self._current_turn += 1
# Check if the maximum number of turns has been reached.
if self._max_turns is not None:
if self._current_turn >= self._max_turns:
stop_message = StopMessage(
content=f"Maximum number of turns {self._max_turns} reached.",
source=self._name,
)
# Reset the termination conditions and turn count.
if self._termination_condition is not None:
await self._termination_condition.reset()
self._current_turn = 0
# Signal termination to the caller of the team.
await self._signal_termination(stop_message)
# Stop the group chat.
return True
return False
async def _log_speaker_selection(self, speaker_names: List[str]) -> None:
"""Log the selected speaker to the output message queue."""
select_msg = SelectSpeakerEvent(content=speaker_names, source=self._name)
if self._emit_team_events:
await self.publish_message(
GroupChatMessage(message=select_msg),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
await self._output_message_queue.put(select_msg)
async def _signal_termination(self, message: StopMessage) -> None:
termination_event = GroupChatTermination(message=message)
# Log the early stop message.
await self.publish_message(
termination_event,
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Put the termination event in the output message queue.
await self._output_message_queue.put(termination_event)
async def _signal_termination_with_error(self, error: SerializableException) -> None:
termination_event = GroupChatTermination(
message=StopMessage(content="An error occurred in the group chat.", source=self._name), error=error
)
# Log the termination event.
await self.publish_message(
termination_event,
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Put the termination event in the output message queue.
await self._output_message_queue.put(termination_event)
@event
async def handle_group_chat_message(self, message: GroupChatMessage, ctx: MessageContext) -> None:
"""Handle a group chat message by appending the content to its output message queue."""
await self._output_message_queue.put(message.message)
@event
async def handle_group_chat_error(self, message: GroupChatError, ctx: MessageContext) -> None:
"""Handle a group chat error by logging the error and signaling termination."""
await self._signal_termination_with_error(message.error)
@rpc
async def handle_reset(self, message: GroupChatReset, ctx: MessageContext) -> None:
"""Reset the group chat manager. Calling :meth:`reset` to reset the group chat manager
and clear the message thread."""
await self.reset()
@rpc
async def handle_pause(self, message: GroupChatPause, ctx: MessageContext) -> None:
"""Pause the group chat manager. This is a no-op in the base class."""
pass
@rpc
async def handle_resume(self, message: GroupChatResume, ctx: MessageContext) -> None:
"""Resume the group chat manager. This is a no-op in the base class."""
pass
@abstractmethod
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
"""Validate the state of the group chat given the start messages.
This is executed when the group chat manager receives a GroupChatStart event.
Args:
messages: A list of chat messages to validate, or None if no messages are provided.
"""
...
async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
"""Update the message thread with the new messages.
This is called when the group chat receives a GroupChatStart or GroupChatAgentResponse event,
before calling the select_speakers method.
"""
self._message_thread.extend(messages)
@abstractmethod
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Select speakers from the participants and return the topic types of the selected speaker.
This is called when the group chat manager have received all responses from the participants
for a turn and is ready to select the next speakers for the next turn.
Args:
thread: The message thread of the group chat.
Returns:
A list of topic types of the selected speakers.
If only one speaker is selected, a single string is returned instead of a list.
"""
...
@abstractmethod
async def reset(self) -> None:
"""Reset the group chat manager."""
...
async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None:
raise ValueError(f"Unhandled message in group chat manager: {type(message)}")

213
agent_dhal/agentdhal_agentchat/teams/_group_chat/_chat_agent_container.py Normal file
View File

@@ -0,0 +1,213 @@
from typing import Any, List, Mapping
from agentdhal_core import DefaultTopicId, MessageContext, event, rpc, trace_invoke_agent_span
from agentdhal_agentchat.messages import BaseAgentEvent, BaseChatMessage, MessageFactory
from ...base import ChatAgent, Response, TaskResult, Team
from ...state import ChatAgentContainerState
from ._events import (
GroupChatAgentResponse,
GroupChatError,
GroupChatMessage,
GroupChatPause,
GroupChatRequestPublish,
GroupChatReset,
GroupChatResume,
GroupChatStart,
GroupChatTeamResponse,
SerializableException,
)
from ._sequential_routed_agent import SequentialRoutedAgent
class ChatAgentContainer(SequentialRoutedAgent):
"""A core agent class that delegates message handling to an
:class:`agentdhal_agentchat.base.ChatAgent` or :class:`agentdhal_agentchat.base.Team`
so that it can be used in a group chat team.
Args:
parent_topic_type (str): The topic type of the parent orchestrator.
output_topic_type (str): The topic type for the output.
agent (ChatAgent | Team): The agent or team to delegate message handling to.
message_factory (MessageFactory): The message factory to use for
creating messages from JSON data.
"""
def __init__(
self, parent_topic_type: str, output_topic_type: str, agent: ChatAgent | Team, message_factory: MessageFactory
) -> None:
super().__init__(
description=agent.description,
sequential_message_types=[
GroupChatStart,
GroupChatRequestPublish,
GroupChatReset,
GroupChatAgentResponse,
GroupChatTeamResponse,
],
)
self._parent_topic_type = parent_topic_type
self._output_topic_type = output_topic_type
self._agent = agent
self._message_buffer: List[BaseChatMessage] = []
self._message_factory = message_factory
@event
async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None:
"""Handle a start event by appending the content to the buffer."""
if message.messages is not None:
for msg in message.messages:
self._buffer_message(msg)
@event
async def handle_agent_response(self, message: GroupChatAgentResponse, ctx: MessageContext) -> None:
"""Handle an agent response event by appending the content to the buffer."""
self._buffer_message(message.response.chat_message)
@event
async def handle_team_response(self, message: GroupChatTeamResponse, ctx: MessageContext) -> None:
"""Handle a team response event by appending the content to the buffer."""
for msg in message.result.messages:
if isinstance(msg, BaseChatMessage):
self._buffer_message(msg)
@rpc
async def handle_reset(self, message: GroupChatReset, ctx: MessageContext) -> None:
"""Handle a reset event by resetting the agent."""
self._message_buffer.clear()
if isinstance(self._agent, Team):
# If the agent is a team, reset the team.
await self._agent.reset()
else:
await self._agent.on_reset(ctx.cancellation_token)
@event
async def handle_request(self, message: GroupChatRequestPublish, ctx: MessageContext) -> None:
"""Handle a content request event by passing the messages in the buffer
to the delegate agent and publish the response."""
if isinstance(self._agent, Team):
try:
stream = self._agent.run_stream(
task=self._message_buffer,
cancellation_token=ctx.cancellation_token,
output_task_messages=False,
)
result: TaskResult | None = None
async for team_event in stream:
if isinstance(team_event, TaskResult):
result = team_event
else:
await self._log_message(team_event)
if result is None:
raise RuntimeError(
"The team did not produce a final TaskResult. Check the team's run_stream method."
)
self._message_buffer.clear()
# Publish the team response to the group chat.
await self.publish_message(
GroupChatTeamResponse(result=result, name=self._agent.name),
topic_id=DefaultTopicId(type=self._parent_topic_type),
cancellation_token=ctx.cancellation_token,
)
except Exception as e:
# Publish the error to the group chat.
error_message = SerializableException.from_exception(e)
await self.publish_message(
GroupChatError(error=error_message),
topic_id=DefaultTopicId(type=self._parent_topic_type),
cancellation_token=ctx.cancellation_token,
)
# Raise the error to the runtime.
raise
else:
# If the agent is not a team, handle it as a single agent.
with trace_invoke_agent_span(
agent_name=self._agent.name,
agent_description=self._agent.description,
agent_id=str(self.id),
):
try:
# Pass the messages in the buffer to the delegate agent.
response: Response | None = None
async for msg in self._agent.on_messages_stream(self._message_buffer, ctx.cancellation_token):
if isinstance(msg, Response):
await self._log_message(msg.chat_message)
response = msg
else:
await self._log_message(msg)
if response is None:
raise RuntimeError(
"The agent did not produce a final response. Check the agent's on_messages_stream method."
)
# Publish the response to the group chat.
self._message_buffer.clear()
await self.publish_message(
GroupChatAgentResponse(response=response, name=self._agent.name),
topic_id=DefaultTopicId(type=self._parent_topic_type),
cancellation_token=ctx.cancellation_token,
)
except Exception as e:
# Publish the error to the group chat.
error_message = SerializableException.from_exception(e)
await self.publish_message(
GroupChatError(error=error_message),
topic_id=DefaultTopicId(type=self._parent_topic_type),
cancellation_token=ctx.cancellation_token,
)
# Raise the error to the runtime.
raise
def _buffer_message(self, message: BaseChatMessage) -> None:
if not self._message_factory.is_registered(message.__class__):
raise ValueError(f"Message type {message.__class__} is not registered.")
# Buffer the message.
self._message_buffer.append(message)
async def _log_message(self, message: BaseAgentEvent | BaseChatMessage) -> None:
if not self._message_factory.is_registered(message.__class__):
raise ValueError(f"Message type {message.__class__} is not registered.")
# Log the message.
await self.publish_message(
GroupChatMessage(message=message),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
@rpc
async def handle_pause(self, message: GroupChatPause, ctx: MessageContext) -> None:
"""Handle a pause event by pausing the agent."""
if isinstance(self._agent, Team):
# If the agent is a team, pause the team.
await self._agent.pause()
else:
await self._agent.on_pause(ctx.cancellation_token)
@rpc
async def handle_resume(self, message: GroupChatResume, ctx: MessageContext) -> None:
"""Handle a resume event by resuming the agent."""
if isinstance(self._agent, Team):
# If the agent is a team, resume the team.
await self._agent.resume()
else:
await self._agent.on_resume(ctx.cancellation_token)
async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None:
raise ValueError(f"Unhandled message in agent container: {type(message)}")
async def save_state(self) -> Mapping[str, Any]:
agent_state = await self._agent.save_state()
state = ChatAgentContainerState(
agent_state=agent_state, message_buffer=[message.dump() for message in self._message_buffer]
)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
container_state = ChatAgentContainerState.model_validate(state)
self._message_buffer = []
for message_data in container_state.message_buffer:
message = self._message_factory.create(message_data)
if isinstance(message, BaseChatMessage):
self._message_buffer.append(message)
else:
raise ValueError(f"Invalid message type in message buffer: {type(message)}")
await self._agent.load_state(container_state.agent_state)

113
agent_dhal/agentdhal_agentchat/teams/_group_chat/_events.py Normal file
View File

@@ -0,0 +1,113 @@
import traceback
from typing import List
from pydantic import BaseModel, SerializeAsAny
from ...base import Response, TaskResult
from ...messages import BaseAgentEvent, BaseChatMessage, StopMessage
class SerializableException(BaseModel):
"""A serializable exception."""
error_type: str
"""The type of error that occurred."""
error_message: str
"""The error message that describes the error."""
traceback: str | None = None
"""The traceback of the error, if available."""
@classmethod
def from_exception(cls, exc: Exception) -> "SerializableException":
"""Create a GroupChatError from an exception."""
return cls(
error_type=type(exc).__name__,
error_message=str(exc),
traceback="\n".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
)
def __str__(self) -> str:
"""Return a string representation of the error, including the traceback if available."""
if self.traceback:
return f"{self.error_type}: {self.error_message}\nTraceback:\n{self.traceback}"
return f"{self.error_type}: {self.error_message}"
class GroupChatStart(BaseModel):
"""A request to start a group chat."""
messages: List[SerializeAsAny[BaseChatMessage]] | None = None
"""An optional list of messages to start the group chat."""
output_task_messages: bool = True
"""Whether to include task messages in the output. Defaults to True for backward compatibility."""
class GroupChatAgentResponse(BaseModel):
"""A response published to a group chat."""
response: SerializeAsAny[Response]
"""The response from an agent."""
name: str
"""The name of the agent that produced the response."""
class GroupChatTeamResponse(BaseModel):
"""A response published to a group chat from a team."""
result: SerializeAsAny[TaskResult]
"""The result from a team."""
name: str
"""The name of the team that produced the response."""
class GroupChatRequestPublish(BaseModel):
"""A request to publish a message to a group chat."""
...
class GroupChatMessage(BaseModel):
"""A message from a group chat."""
message: SerializeAsAny[BaseAgentEvent | BaseChatMessage]
"""The message that was published."""
class GroupChatTermination(BaseModel):
"""A message indicating that a group chat has terminated."""
message: StopMessage
"""The stop message that indicates the reason of termination."""
error: SerializableException | None = None
"""The error that occurred, if any."""
class GroupChatReset(BaseModel):
"""A request to reset the agents in the group chat."""
...
class GroupChatPause(BaseModel):
"""A request to pause the group chat."""
...
class GroupChatResume(BaseModel):
"""A request to resume the group chat."""
...
class GroupChatError(BaseModel):
"""A message indicating that an error occurred in the group chat."""
error: SerializableException
"""The error that occurred."""

17
agent_dhal/agentdhal_agentchat/teams/_group_chat/_graph/__init__.py Normal file
View File

@@ -0,0 +1,17 @@
from ._digraph_group_chat import (
DiGraph,
DiGraphEdge,
DiGraphNode,
GraphFlow,
GraphFlowManager,
)
from ._graph_builder import DiGraphBuilder
__all__ = [
"GraphFlow",
"DiGraph",
"GraphFlowManager",
"DiGraphNode",
"DiGraphEdge",
"DiGraphBuilder",
]

877
agent_dhal/agentdhal_agentchat/teams/_group_chat/_graph/_digraph_group_chat.py Normal file
View File

@@ -0,0 +1,877 @@
import asyncio
from collections import Counter, deque
from typing import Any, Callable, Deque, Dict, List, Literal, Mapping, Sequence, Set, Union
from agentdhal_core import AgentRuntime, Component, ComponentModel
from pydantic import BaseModel, Field, model_validator
from typing_extensions import Self
from agentdhal_agentchat.base import ChatAgent, TerminationCondition
from agentdhal_agentchat.messages import (
BaseAgentEvent,
BaseChatMessage,
MessageFactory,
StopMessage,
)
from agentdhal_agentchat.state import BaseGroupChatManagerState
from agentdhal_agentchat.teams import BaseGroupChat
from ..._group_chat._base_group_chat_manager import BaseGroupChatManager
from ..._group_chat._events import GroupChatTermination
_DIGRAPH_STOP_MESSAGE = "Digraph execution is complete"
class DiGraphEdge(BaseModel):
"""Represents a directed edge in a :class:`DiGraph`, with an optional execution condition.
.. warning::
This is an experimental feature, and the API will change in the future releases.
.. warning::
If the condition is a callable, it will not be serialized in the model.
"""
target: str # Target node name
condition: Union[str, Callable[[BaseChatMessage], bool], None] = Field(default=None)
"""(Experimental) Condition to execute this edge.
If None, the edge is unconditional.
If a string, the edge is conditional on the presence of that string in the last agent chat message.
If a callable, the edge is conditional on the callable returning True when given the last message.
"""
# Using Field to exclude the condition in serialization if it's a callable
condition_function: Callable[[BaseChatMessage], bool] | None = Field(default=None, exclude=True)
activation_group: str = Field(default="")
"""Group identifier for forward dependencies.
When multiple edges point to the same target node, they are grouped by this field.
This allows distinguishing between different cycles or dependency patterns.
Example: In a graph containing a cycle like A->B->C->B, the two edges pointing to B (A->B and C->B)
can be in different activation groups to control how B is activated.
Defaults to the target node name if not specified.
"""
activation_condition: Literal["all", "any"] = "all"
"""Determines how forward dependencies within the same activation_group are evaluated.
- "all": All edges in this activation group must be satisfied before the target node can execute
- "any": Any single edge in this activation group being satisfied allows the target node to execute
This is used to handle complex dependency patterns in cyclic graphs where multiple
paths can lead to the same target node.
"""
@model_validator(mode="after")
def _validate_condition(self) -> "DiGraphEdge":
# Store callable in a separate field and set condition to None for serialization
if callable(self.condition):
self.condition_function = self.condition
# For serialization purposes, we'll set the condition to None
# when storing as a pydantic model/dict
object.__setattr__(self, "condition", None)
# Set activation_group to target if not already set
if not self.activation_group:
self.activation_group = self.target
return self
def check_condition(self, message: BaseChatMessage) -> bool:
"""Check if the edge condition is satisfied for the given message.
Args:
message: The message to check the condition against.
Returns:
True if condition is satisfied (None condition always returns True),
False otherwise.
"""
if self.condition_function is not None:
return self.condition_function(message)
elif isinstance(self.condition, str):
# If it's a string, check if the string is in the message content
return self.condition in message.to_model_text()
return True # None condition is always satisfied
class DiGraphNode(BaseModel):
"""Represents a node (agent) in a :class:`DiGraph`, with its outgoing edges and activation type.
.. warning::
This is an experimental feature, and the API will change in the future releases.
"""
name: str # Agent's name
edges: List[DiGraphEdge] = [] # Outgoing edges
activation: Literal["all", "any"] = "all"
class DiGraph(BaseModel):
"""Defines a directed graph structure with nodes and edges.
:class:`GraphFlow` uses this to determine execution order and conditions.
.. warning::
This is an experimental feature, and the API will change in the future releases.
"""
nodes: Dict[str, DiGraphNode] # Node name → DiGraphNode mapping
default_start_node: str | None = None # Default start node name
_has_cycles: bool | None = None # Cyclic graph flag
def get_parents(self) -> Dict[str, List[str]]:
"""Compute a mapping of each node to its parent nodes."""
parents: Dict[str, List[str]] = {node: [] for node in self.nodes}
for node in self.nodes.values():
for edge in node.edges:
parents[edge.target].append(node.name)
return parents
def get_start_nodes(self) -> Set[str]:
"""Return the nodes that have no incoming edges (entry points)."""
if self.default_start_node:
return {self.default_start_node}
parents = self.get_parents()
return set([node_name for node_name, parent_list in parents.items() if not parent_list])
def get_leaf_nodes(self) -> Set[str]:
"""Return nodes that have no outgoing edges (final output nodes)."""
return set([name for name, node in self.nodes.items() if not node.edges])
def has_cycles_with_exit(self) -> bool:
"""
Check if the graph has any cycles and validate that each cycle has at least one conditional edge.
Returns:
bool: True if there is at least one cycle and all cycles have an exit condition.
False if there are no cycles.
Raises:
ValueError: If there is a cycle without any conditional edge.
"""
visited: Set[str] = set()
rec_stack: Set[str] = set()
path: List[str] = []
def dfs(node_name: str) -> bool:
visited.add(node_name)
rec_stack.add(node_name)
path.append(node_name)
for edge in self.nodes[node_name].edges:
target = edge.target
if target not in visited:
if dfs(target):
return True
elif target in rec_stack:
# Found a cycle → extract the cycle
cycle_start_index = path.index(target)
cycle_nodes = path[cycle_start_index:]
cycle_edges: List[DiGraphEdge] = []
for n in cycle_nodes:
cycle_edges.extend(self.nodes[n].edges)
if all(edge.condition is None and edge.condition_function is None for edge in cycle_edges):
raise ValueError(
f"Cycle detected without exit condition: {' -> '.join(cycle_nodes + cycle_nodes[:1])}"
)
return True # Found cycle, but it has an exit condition
rec_stack.remove(node_name)
path.pop()
return False
has_cycle = False
for node in self.nodes:
if node not in visited:
if dfs(node):
has_cycle = True
return has_cycle
def get_has_cycles(self) -> bool:
"""Indicates if the graph has at least one cycle (with valid exit conditions)."""
if self._has_cycles is None:
self._has_cycles = self.has_cycles_with_exit()
return self._has_cycles
def graph_validate(self) -> None:
"""Validate graph structure and execution rules."""
if not self.nodes:
raise ValueError("Graph has no nodes.")
if not self.get_start_nodes():
raise ValueError("Graph must have at least one start node")
if not self.get_leaf_nodes():
raise ValueError("Graph must have at least one leaf node")
# Outgoing edge condition validation (per node)
for node in self.nodes.values():
# Check that if a node has an outgoing conditional edge, then all outgoing edges are conditional
has_condition = any(
edge.condition is not None or edge.condition_function is not None for edge in node.edges
)
has_unconditioned = any(edge.condition is None and edge.condition_function is None for edge in node.edges)
if has_condition and has_unconditioned:
raise ValueError(f"Node '{node.name}' has a mix of conditional and unconditional edges.")
# Validate activation conditions across all edges in the graph
self._validate_activation_conditions()
self._has_cycles = self.has_cycles_with_exit()
def _validate_activation_conditions(self) -> None:
"""Validate that all edges pointing to the same target node have consistent activation_condition values.
Raises:
ValueError: If edges pointing to the same target have different activation_condition values
"""
target_activation_conditions: Dict[str, Dict[str, str]] = {} # target_node -> {activation_group -> condition}
for node in self.nodes.values():
for edge in node.edges:
target = edge.target # The target node this edge points to
activation_group = edge.activation_group
if target not in target_activation_conditions:
target_activation_conditions[target] = {}
if activation_group in target_activation_conditions[target]:
if target_activation_conditions[target][activation_group] != edge.activation_condition:
# Find the source node that has the conflicting condition
conflicting_source = self._find_edge_source_by_target_and_group(
target, activation_group, target_activation_conditions[target][activation_group]
)
raise ValueError(
f"Conflicting activation conditions for target '{target}' group '{activation_group}': "
f"'{target_activation_conditions[target][activation_group]}' (from node '{conflicting_source}') "
f"and '{edge.activation_condition}' (from node '{node.name}')"
)
else:
target_activation_conditions[target][activation_group] = edge.activation_condition
def _find_edge_source_by_target_and_group(
self, target: str, activation_group: str, activation_condition: str
) -> str:
"""Find the source node that has an edge pointing to the given target with the given activation_group and activation_condition."""
for node_name, node in self.nodes.items():
for edge in node.edges:
if (
edge.target == target
and edge.activation_group == activation_group
and edge.activation_condition == activation_condition
):
return node_name
return "unknown"
def get_remaining_map(self) -> Dict[str, Dict[str, int]]:
"""Get the remaining map that tracks how many edges point to each target node with each activation group.
Returns:
Dictionary mapping target nodes to their activation groups and remaining counts
"""
remaining_map: Dict[str, Dict[str, int]] = {}
for node in self.nodes.values():
for edge in node.edges:
target = edge.target
activation_group = edge.activation_group
if target not in remaining_map:
remaining_map[target] = {}
if activation_group not in remaining_map[target]:
remaining_map[target][activation_group] = 0
remaining_map[target][activation_group] += 1
return remaining_map
class GraphFlowManagerState(BaseGroupChatManagerState):
"""Tracks active execution state for DAG-based execution."""
active_nodes: List[str] = [] # Currently executing nodes
type: str = "GraphManagerState"
class GraphFlowManager(BaseGroupChatManager):
"""Manages execution of agents using a Directed Graph execution model."""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
graph: DiGraph,
) -> None:
"""Initialize the graph-based execution manager."""
super().__init__(
name=name,
group_topic_type=group_topic_type,
output_topic_type=output_topic_type,
participant_topic_types=participant_topic_types,
participant_names=participant_names,
participant_descriptions=participant_descriptions,
output_message_queue=output_message_queue,
termination_condition=termination_condition,
max_turns=max_turns,
message_factory=message_factory,
)
graph.graph_validate()
if graph.get_has_cycles() and self._termination_condition is None and self._max_turns is None:
raise ValueError("A termination condition is required for cyclic graphs without a maximum turn limit.")
self._graph = graph
# Lookup table for incoming edges for each node.
self._parents = graph.get_parents()
# Lookup table for outgoing edges for each node.
self._edges: Dict[str, List[DiGraphEdge]] = {n: node.edges for n, node in graph.nodes.items()}
# Build activation and enqueued_any lookup tables by collecting all edges and grouping by target node
self._build_lookup_tables(graph)
# Track which activation groups were triggered for each node
self._triggered_activation_groups: Dict[str, Set[str]] = {}
# === Mutable states for the graph execution ===
# Count the number of remaining parents to activate each node.
self._remaining: Dict[str, Counter[str]] = {
target: Counter(groups) for target, groups in graph.get_remaining_map().items()
}
# cache for remaining
self._origin_remaining: Dict[str, Dict[str, int]] = {
target: Counter(groups) for target, groups in self._remaining.items()
}
# Ready queue for nodes that are ready to execute, starting with the start nodes.
self._ready: Deque[str] = deque([n for n in graph.get_start_nodes()])
def _build_lookup_tables(self, graph: DiGraph) -> None:
"""Build activation and enqueued_any lookup tables by collecting all edges and grouping by target node.
Args:
graph: The directed graph
"""
self._activation: Dict[str, Dict[str, Literal["any", "all"]]] = {}
self._enqueued_any: Dict[str, Dict[str, bool]] = {}
for node in graph.nodes.values():
for edge in node.edges:
target = edge.target
activation_group = edge.activation_group
# Build activation lookup
if target not in self._activation:
self._activation[target] = {}
if activation_group not in self._activation[target]:
self._activation[target][activation_group] = edge.activation_condition
# Build enqueued_any lookup
if target not in self._enqueued_any:
self._enqueued_any[target] = {}
if activation_group not in self._enqueued_any[target]:
self._enqueued_any[target][activation_group] = False
async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
await super().update_message_thread(messages)
# Find the node that ran in the current turn.
message = messages[-1]
if message.source not in self._graph.nodes:
# Ignore messages from sources outside of the graph.
return
assert isinstance(message, BaseChatMessage)
source = message.source
# Propagate the update to the children of the node.
for edge in self._edges[source]:
# Use the new check_condition method that handles both string and callable conditions
if not edge.check_condition(message):
continue
target = edge.target
activation_group = edge.activation_group
if self._activation[target][activation_group] == "all":
self._remaining[target][activation_group] -= 1
if self._remaining[target][activation_group] == 0:
# If all parents are done, add to the ready queue.
self._ready.append(target)
# Track which activation group was triggered
self._save_triggered_activation_group(target, activation_group)
else:
# If activation is any, add to the ready queue if not already enqueued.
if not self._enqueued_any[target][activation_group]:
self._ready.append(target)
self._enqueued_any[target][activation_group] = True
# Track which activation group was triggered
self._save_triggered_activation_group(target, activation_group)
def _save_triggered_activation_group(self, target: str, activation_group: str) -> None:
"""Save which activation group was triggered for a target node.
Args:
target: The target node that was triggered
activation_group: The activation group that caused the trigger
"""
if target not in self._triggered_activation_groups:
self._triggered_activation_groups[target] = set()
self._triggered_activation_groups[target].add(activation_group)
def _reset_triggered_activation_groups(self, speaker: str) -> None:
"""Reset the bookkeeping for the specific activation groups that were triggered for a speaker.
Args:
speaker: The speaker node to reset activation groups for
"""
if speaker not in self._triggered_activation_groups:
return
for activation_group in self._triggered_activation_groups[speaker]:
if self._activation[speaker][activation_group] == "any":
self._enqueued_any[speaker][activation_group] = False
else:
# Reset the remaining count for this activation group using the graph's original count
if speaker in self._remaining and activation_group in self._remaining[speaker]:
self._remaining[speaker][activation_group] = self._origin_remaining[speaker][activation_group]
# Clear the triggered activation groups for this speaker
self._triggered_activation_groups[speaker].clear()
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str]:
# Drain the ready queue for the next set of speakers.
speakers: List[str] = []
while self._ready:
speaker = self._ready.popleft()
speakers.append(speaker)
# Reset the bookkeeping for the specific activation groups that were triggered
self._reset_triggered_activation_groups(speaker)
return speakers
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
async def _apply_termination_condition(
self, delta: Sequence[BaseAgentEvent | BaseChatMessage], increment_turn_count: bool = False
) -> bool:
"""Apply termination condition including graph-specific completion logic.
First checks if graph execution is complete, then checks standard termination conditions.
Args:
delta: The message delta to check termination conditions against
increment_turn_count: Whether to increment the turn count
Returns:
True if the conversation should be terminated, False otherwise
"""
# Check if the graph execution is complete (no ready speakers) - prioritize this check
if not self._ready:
stop_message = StopMessage(
content=_DIGRAPH_STOP_MESSAGE,
source=self._name,
)
# Reset the execution state when the graph has naturally completed
self._reset_execution_state()
# Reset the termination conditions and turn count.
if self._termination_condition is not None:
await self._termination_condition.reset()
self._current_turn = 0
# Signal termination to the caller of the team.
await self._signal_termination(stop_message)
return True
# Apply the standard termination conditions from the base class
return await super()._apply_termination_condition(delta, increment_turn_count)
def _reset_execution_state(self) -> None:
"""Reset the graph execution state to the initial state."""
self._remaining = {target: Counter(groups) for target, groups in self._graph.get_remaining_map().items()}
self._enqueued_any = {n: {g: False for g in self._enqueued_any[n]} for n in self._enqueued_any}
self._ready = deque([n for n in self._graph.get_start_nodes()])
async def save_state(self) -> Mapping[str, Any]:
"""Save the execution state."""
state = {
"message_thread": [message.dump() for message in self._message_thread],
"current_turn": self._current_turn,
"remaining": {target: dict(counter) for target, counter in self._remaining.items()},
"enqueued_any": dict(self._enqueued_any),
"ready": list(self._ready),
}
return state
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Restore execution state from saved data."""
self._message_thread = [self._message_factory.create(msg) for msg in state["message_thread"]]
self._current_turn = state["current_turn"]
self._remaining = {target: Counter(groups) for target, groups in state["remaining"].items()}
self._enqueued_any = state["enqueued_any"]
self._ready = deque(state["ready"])
async def reset(self) -> None:
"""Reset execution state to the start of the graph."""
self._current_turn = 0
self._message_thread.clear()
if self._termination_condition:
await self._termination_condition.reset()
self._reset_execution_state()
class GraphFlowConfig(BaseModel):
"""The declarative configuration for GraphFlow."""
name: str | None = None
description: str | None = None
participants: List[ComponentModel]
termination_condition: ComponentModel | None = None
max_turns: int | None = None
graph: DiGraph # The execution graph for agents
class GraphFlow(BaseGroupChat, Component[GraphFlowConfig]):
"""A team that runs a group chat following a Directed Graph execution pattern.
.. warning::
This is an experimental feature, and the API will change in the future releases.
This group chat executes agents based on a directed graph (:class:`DiGraph`) structure,
allowing complex workflows such as sequential execution, parallel fan-out,
conditional branching, join patterns, and loops with explicit exit conditions.
The execution order is determined by the edges defined in the `DiGraph`. Each node
in the graph corresponds to an agent, and edges define the flow of messages between agents.
Nodes can be configured to activate when:
- **All** parent nodes have completed (activation="all") → default
- **Any** parent node completes (activation="any")
Conditional branching is supported using edge conditions, where the next agent(s) are selected
based on content in the chat history. Loops are permitted as long as there is a condition
that eventually exits the loop.
.. note::
Use the :class:`DiGraphBuilder` class to create a :class:`DiGraph` easily. It provides a fluent API
for adding nodes and edges, setting entry points, and validating the graph structure.
See the :class:`DiGraphBuilder` documentation for more details.
The :class:`GraphFlow` class is designed to be used with the :class:`DiGraphBuilder` for creating complex workflows.
.. warning::
When using callable conditions in edges, they will not be serialized
when calling :meth:`dump_component`. This will be addressed in future releases.
Args:
participants (List[ChatAgent]): The participants in the group chat.
termination_condition (TerminationCondition, optional): Termination condition for the chat.
max_turns (int, optional): Maximum number of turns before forcing termination.
graph (DiGraph): Directed execution graph defining node flow and conditions.
Raises:
ValueError: If participant names are not unique, or if graph validation fails (e.g., cycles without exit).
Examples:
**Sequential Flow: A → B → C**
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import DiGraphBuilder, GraphFlow
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main():
# Initialize agents with OpenAI model clients.
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
agent_a = AssistantAgent("A", model_client=model_client, system_message="You are a helpful assistant.")
agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to Chinese.")
agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to English.")
# Create a directed graph with sequential flow A -> B -> C.
builder = DiGraphBuilder()
builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
builder.add_edge(agent_a, agent_b).add_edge(agent_b, agent_c)
graph = builder.build()
# Create a GraphFlow team with the directed graph.
team = GraphFlow(
participants=[agent_a, agent_b, agent_c],
graph=graph,
termination_condition=MaxMessageTermination(5),
)
# Run the team and print the events.
async for event in team.run_stream(task="Write a short story about a cat."):
print(event)
asyncio.run(main())
**Parallel Fan-out: A → (B, C)**
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import DiGraphBuilder, GraphFlow
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main():
# Initialize agents with OpenAI model clients.
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
agent_a = AssistantAgent("A", model_client=model_client, system_message="You are a helpful assistant.")
agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to Chinese.")
agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to Japanese.")
# Create a directed graph with fan-out flow A -> (B, C).
builder = DiGraphBuilder()
builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
builder.add_edge(agent_a, agent_b).add_edge(agent_a, agent_c)
graph = builder.build()
# Create a GraphFlow team with the directed graph.
team = GraphFlow(
participants=[agent_a, agent_b, agent_c],
graph=graph,
termination_condition=MaxMessageTermination(5),
)
# Run the team and print the events.
async for event in team.run_stream(task="Write a short story about a cat."):
print(event)
asyncio.run(main())
**Conditional Branching: A → B (if 'yes') or C (otherwise)**
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import DiGraphBuilder, GraphFlow
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main():
# Initialize agents with OpenAI model clients.
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
agent_a = AssistantAgent(
"A",
model_client=model_client,
system_message="Detect if the input is in Chinese. If it is, say 'yes', else say 'no', and nothing else.",
)
agent_b = AssistantAgent("B", model_client=model_client, system_message="Translate input to English.")
agent_c = AssistantAgent("C", model_client=model_client, system_message="Translate input to Chinese.")
# Create a directed graph with conditional branching flow A -> B ("yes"), A -> C (otherwise).
builder = DiGraphBuilder()
builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
# Create conditions as callables that check the message content.
builder.add_edge(agent_a, agent_b, condition=lambda msg: "yes" in msg.to_model_text())
builder.add_edge(agent_a, agent_c, condition=lambda msg: "yes" not in msg.to_model_text())
graph = builder.build()
# Create a GraphFlow team with the directed graph.
team = GraphFlow(
participants=[agent_a, agent_b, agent_c],
graph=graph,
termination_condition=MaxMessageTermination(5),
)
# Run the team and print the events.
async for event in team.run_stream(task="AutoGen is a framework for building AI agents."):
print(event)
asyncio.run(main())
**Loop with exit condition: A → B → C (if 'APPROVE') or A (otherwise)**
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import MaxMessageTermination
from agentdhal_agentchat.teams import DiGraphBuilder, GraphFlow
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main():
# Initialize agents with OpenAI model clients.
model_client = OpenAIChatCompletionClient(model="gpt-4.1")
agent_a = AssistantAgent(
"A",
model_client=model_client,
system_message="You are a helpful assistant.",
)
agent_b = AssistantAgent(
"B",
model_client=model_client,
system_message="Provide feedback on the input, if your feedback has been addressed, "
"say 'APPROVE', otherwise provide a reason for rejection.",
)
agent_c = AssistantAgent(
"C", model_client=model_client, system_message="Translate the final product to Korean."
)
# Create a loop graph with conditional exit: A -> B -> C ("APPROVE"), B -> A (otherwise).
builder = DiGraphBuilder()
builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
builder.add_edge(agent_a, agent_b)
# Create conditional edges using strings
builder.add_edge(agent_b, agent_c, condition=lambda msg: "APPROVE" in msg.to_model_text())
builder.add_edge(agent_b, agent_a, condition=lambda msg: "APPROVE" not in msg.to_model_text())
builder.set_entry_point(agent_a)
graph = builder.build()
# Create a GraphFlow team with the directed graph.
team = GraphFlow(
participants=[agent_a, agent_b, agent_c],
graph=graph,
termination_condition=MaxMessageTermination(20), # Max 20 messages to avoid infinite loop.
)
# Run the team and print the events.
async for event in team.run_stream(task="Write a short poem about AI Agents."):
print(event)
asyncio.run(main())
"""
component_config_schema = GraphFlowConfig
component_provider_override = "agentdhal_agentchat.teams.GraphFlow"
DEFAULT_NAME = "GraphFlow"
DEFAULT_DESCRIPTION = "A team of agents"
def __init__(
self,
participants: List[ChatAgent],
graph: DiGraph,
*,
name: str | None = None,
description: str | None = None,
termination_condition: TerminationCondition | None = None,
max_turns: int | None = None,
runtime: AgentRuntime | None = None,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
) -> None:
self._input_participants = participants
self._input_termination_condition = termination_condition
for participant in participants:
if not isinstance(participant, ChatAgent):
raise TypeError(f"Participant {participant} must be a ChatAgent.")
# No longer add _StopAgent or StopMessageTermination
# Termination is now handled directly in GraphFlowManager._apply_termination_condition
super().__init__(
name=name or self.DEFAULT_NAME,
description=description or self.DEFAULT_DESCRIPTION,
participants=list(participants),
group_chat_manager_name="GraphManager",
group_chat_manager_class=GraphFlowManager,
termination_condition=termination_condition,
max_turns=max_turns,
runtime=runtime,
custom_message_types=custom_message_types,
)
self._graph = graph
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], GraphFlowManager]:
"""Creates the factory method for initializing the DiGraph-based chat manager."""
def _factory() -> GraphFlowManager:
return GraphFlowManager(
name=name,
group_topic_type=group_topic_type,
output_topic_type=output_topic_type,
participant_topic_types=participant_topic_types,
participant_names=participant_names,
participant_descriptions=participant_descriptions,
output_message_queue=output_message_queue,
termination_condition=termination_condition,
max_turns=max_turns,
message_factory=message_factory,
graph=self._graph,
)
return _factory
def _to_config(self) -> GraphFlowConfig:
"""Converts the instance into a configuration object."""
participants = [participant.dump_component() for participant in self._input_participants]
termination_condition = (
self._input_termination_condition.dump_component() if self._input_termination_condition else None
)
return GraphFlowConfig(
name=self._name,
description=self._description,
participants=participants,
termination_condition=termination_condition,
max_turns=self._max_turns,
graph=self._graph,
)
@classmethod
def _from_config(cls, config: GraphFlowConfig) -> Self:
"""Reconstructs an instance from a configuration object."""
participants = [ChatAgent.load_component(participant) for participant in config.participants]
termination_condition = (
TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
)
return cls(
name=config.name,
description=config.description,
participants=participants,
graph=config.graph,
termination_condition=termination_condition,
max_turns=config.max_turns,
)

209
agent_dhal/agentdhal_agentchat/teams/_group_chat/_graph/_graph_builder.py Normal file
View File

@@ -0,0 +1,209 @@
import warnings
from typing import Callable, Dict, Literal, Optional, Union
from agentdhal_agentchat.base import ChatAgent
from agentdhal_agentchat.messages import BaseChatMessage
from ._digraph_group_chat import DiGraph, DiGraphEdge, DiGraphNode
class DiGraphBuilder:
"""
A fluent builder for constructing :class:`DiGraph` execution graphs used in :class:`GraphFlow`.
.. warning::
This is an experimental feature, and the API will change in the future releases.
This utility provides a convenient way to programmatically build a graph of agent interactions,
including complex execution flows such as:
- Sequential chains
- Parallel fan-outs
- Conditional branching
- Cyclic loops with safe exits
Each node in the graph represents an agent. Edges define execution paths between agents,
and can optionally be conditioned on message content using callable functions.
The builder is compatible with the `Graph` runner and supports both standard and filtered agents.
Methods:
- add_node(agent, activation): Add an agent node to the graph.
- add_edge(source, target, condition): Connect two nodes optionally with a condition.
- add_conditional_edges(source, condition_to_target): Add multiple conditional edges from a source.
- set_entry_point(agent): Define the default start node (optional).
- build(): Generate a validated `DiGraph`.
- get_participants(): Return the list of added agents.
Example — Sequential Flow A → B → C:
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> builder.add_edge(agent_a, agent_b).add_edge(agent_b, agent_c)
>>> team = Graph(
... participants=builder.get_participants(),
... graph=builder.build(),
... termination_condition=MaxMessageTermination(5),
... )
Example — Parallel Fan-out A → (B, C):
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> builder.add_edge(agent_a, agent_b).add_edge(agent_a, agent_c)
Example — Conditional Branching A → B or A → C:
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> # Add conditional edges using keyword check
>>> builder.add_edge(agent_a, agent_b, condition="keyword1")
>>> builder.add_edge(agent_a, agent_c, condition="keyword2")
Example — Using Custom String Conditions:
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> # Add condition strings to check in messages
>>> builder.add_edge(agent_a, agent_b, condition="big")
>>> builder.add_edge(agent_a, agent_c, condition="small")
Example — Loop: A → B → A or B → C:
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> builder.add_edge(agent_a, agent_b)
>> # Add a loop back to agent A
>>> builder.add_edge(agent_b, agent_a, condition=lambda msg: "loop" in msg.to_model_text())
>>> # Add exit condition to break the loop
>>> builder.add_edge(agent_b, agent_c, condition=lambda msg: "loop" not in msg.to_model_text())
Example — Loop with multiple paths to the same node: A → B → C → B:
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c)
>>> builder.add_edge(agent_a, agent_b)
>>> builder.add_edge(agent_b, agent_c)
>>> builder.add_edge(agent_c, agent_b, activation_group="loop_back")
Example — Loop with multiple paths to the same node with any activation condition: A → B → (C1, C2) → B → E(exit):
>>> builder = GraphBuilder()
>>> builder.add_node(agent_a).add_node(agent_b).add_node(agent_c1).add_node(agent_c2).add_node(agent_e)
>>> builder.add_edge(agent_a, agent_b)
>>> builder.add_edge(agent_b, agent_c1)
>>> builder.add_edge(agent_b, agent_c2)
>>> builder.add_edge(agent_b, agent_e, condition="exit")
>>> builder.add_edge(agent_c1, agent_b, activation_group="loop_back_group", activation_condition="any")
>>> builder.add_edge(agent_c2, agent_b, activation_group="loop_back_group", activation_condition="any")
"""
def __init__(self) -> None:
self.nodes: Dict[str, DiGraphNode] = {}
self.agents: Dict[str, ChatAgent] = {}
self._default_start_node: Optional[str] = None
def _get_name(self, obj: Union[str, ChatAgent]) -> str:
return obj if isinstance(obj, str) else obj.name
def add_node(self, agent: ChatAgent, activation: Literal["all", "any"] = "all") -> "DiGraphBuilder":
"""Add a node to the graph and register its agent."""
name = agent.name
if name not in self.nodes:
self.nodes[name] = DiGraphNode(name=name, edges=[], activation=activation)
self.agents[name] = agent
return self
def add_edge(
self,
source: Union[str, ChatAgent],
target: Union[str, ChatAgent],
condition: Optional[Union[str, Callable[[BaseChatMessage], bool]]] = None,
activation_group: Optional[str] = None,
activation_condition: Optional[Literal["all", "any"]] = None,
) -> "DiGraphBuilder":
"""Add a directed edge from source to target, optionally with a condition.
Args:
source: Source node (agent name or agent object)
target: Target node (agent name or agent object)
condition: Optional condition for edge activation.
If string, activates when substring is found in message.
If callable, activates when function returns True for the message.
Returns:
Self for method chaining
Raises:
ValueError: If source or target node doesn't exist in the builder
"""
source_name = self._get_name(source)
target_name = self._get_name(target)
if source_name not in self.nodes:
raise ValueError(f"Source node '{source_name}' must be added before adding an edge.")
if target_name not in self.nodes:
raise ValueError(f"Target node '{target_name}' must be added before adding an edge.")
if activation_group is None:
activation_group = target_name
if activation_condition is None:
activation_condition = "all"
self.nodes[source_name].edges.append(
DiGraphEdge(
target=target_name,
condition=condition,
activation_group=activation_group,
activation_condition=activation_condition,
)
)
return self
def add_conditional_edges(
self, source: Union[str, ChatAgent], condition_to_target: Dict[str, Union[str, ChatAgent]]
) -> "DiGraphBuilder":
"""Add multiple conditional edges from a source node based on keyword checks.
.. warning::
This method interface will be changed in the future to support callable conditions.
Please use `add_edge` if you need to specify custom conditions.
Args:
source: Source node (agent name or agent object)
condition_to_target: Mapping from condition strings to target nodes
Each key is a keyword that will be checked in the message content
Each value is the target node to activate when condition is met
For each key (keyword), a lambda will be created that checks
if the keyword is in the message text.
Returns:
Self for method chaining
"""
warnings.warn(
"add_conditional_edges will be changed in the future to support callable conditions. "
"For now, please use add_edge if you need to specify custom conditions.",
DeprecationWarning,
stacklevel=2,
)
for condition_keyword, target in condition_to_target.items():
self.add_edge(source, target, condition=condition_keyword)
return self
def set_entry_point(self, name: Union[str, ChatAgent]) -> "DiGraphBuilder":
"""Set the default start node of the graph."""
node_name = self._get_name(name)
if node_name not in self.nodes:
raise ValueError(f"Start node '{node_name}' must be added before setting as entry point.")
self._default_start_node = node_name
return self
def build(self) -> DiGraph:
"""Build and validate the DiGraph."""
graph = DiGraph(
nodes=self.nodes,
default_start_node=self._default_start_node,
)
graph.graph_validate()
return graph
def get_participants(self) -> list[ChatAgent]:
"""Return the list of agents in the builder, in insertion order."""
return list(self.agents.values())

5
agent_dhal/agentdhal_agentchat/teams/_group_chat/_magentic_one/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
from ._magentic_one_group_chat import MagenticOneGroupChat
__all__ = [
"MagenticOneGroupChat",
]

209
agent_dhal/agentdhal_agentchat/teams/_group_chat/_magentic_one/_magentic_one_group_chat.py Normal file
View File

@@ -0,0 +1,209 @@
import asyncio
import logging
from typing import Callable, List
from agentdhal_core import AgentRuntime, Component, ComponentModel
from agentdhal_core.models import ChatCompletionClient
from pydantic import BaseModel
from typing_extensions import Self
from .... import EVENT_LOGGER_NAME, TRACE_LOGGER_NAME
from ....base import ChatAgent, TerminationCondition
from ....messages import BaseAgentEvent, BaseChatMessage, MessageFactory
from .._base_group_chat import BaseGroupChat
from .._events import GroupChatTermination
from ._magentic_one_orchestrator import MagenticOneOrchestrator
from ._prompts import ORCHESTRATOR_FINAL_ANSWER_PROMPT
trace_logger = logging.getLogger(TRACE_LOGGER_NAME)
event_logger = logging.getLogger(EVENT_LOGGER_NAME)
class MagenticOneGroupChatConfig(BaseModel):
"""The declarative configuration for a MagenticOneGroupChat."""
name: str | None = None
description: str | None = None
participants: List[ComponentModel]
model_client: ComponentModel
termination_condition: ComponentModel | None = None
max_turns: int | None = None
max_stalls: int
final_answer_prompt: str
emit_team_events: bool = False
class MagenticOneGroupChat(BaseGroupChat, Component[MagenticOneGroupChatConfig]):
"""A team that runs a group chat with participants managed by the MagenticOneOrchestrator.
The orchestrator handles the conversation flow, ensuring that the task is completed
efficiently by managing the participants' interactions.
The orchestrator is based on the Magentic-One architecture, which is a generalist multi-agent system for solving complex tasks (see references below).
Unlike :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` and :class:`~agentdhal_agentchat.teams.SelectorGroupChat`,
the MagenticOneGroupChat does not support using team as participant.
Args:
participants (List[ChatAgent]): The participants in the group chat.
model_client (ChatCompletionClient): The model client used for generating responses.
termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None.
Without a termination condition, the group chat will run based on the orchestrator logic or until the maximum number of turns is reached.
max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to 20.
max_stalls (int, optional): The maximum number of stalls allowed before re-planning. Defaults to 3.
final_answer_prompt (str, optional): The LLM prompt used to generate the final answer or response from the team's transcript. A default (sensible for GPT-4o class models) is provided.
custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat.
If you are using custom message types or your agents produces custom message types, you need to specify them here.
Make sure your custom message types are subclasses of :class:`~agentdhal_agentchat.messages.BaseAgentEvent` or :class:`~agentdhal_agentchat.messages.BaseChatMessage`.
emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False.
Raises:
ValueError: In orchestration logic if progress ledger does not have required keys or if next speaker is not valid.
Examples:
MagenticOneGroupChat with one assistant agent:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import MagenticOneGroupChat
from agentdhal_agentchat.ui import Console
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
assistant = AssistantAgent(
"Assistant",
model_client=model_client,
)
team = MagenticOneGroupChat([assistant], model_client=model_client)
await Console(team.run_stream(task="Provide a different proof to Fermat last theorem"))
asyncio.run(main())
References:
If you use the MagenticOneGroupChat in your work, please cite the following paper:
.. code-block:: bibtex
@article{fourney2024magentic,
title={Magentic-one: A generalist multi-agent system for solving complex tasks},
author={Fourney, Adam and Bansal, Gagan and Mozannar, Hussein and Tan, Cheng and Salinas, Eduardo and Niedtner, Friederike and Proebsting, Grace and Bassman, Griffin and Gerrits, Jack and Alber, Jacob and others},
journal={arXiv preprint arXiv:2411.04468},
year={2024}
}
"""
component_config_schema = MagenticOneGroupChatConfig
component_provider_override = "agentdhal_agentchat.teams.MagenticOneGroupChat"
DEFAULT_NAME = "MagenticOneGroupChat"
DEFAULT_DESCRIPTION = "A team of agents."
def __init__(
self,
participants: List[ChatAgent],
model_client: ChatCompletionClient,
*,
name: str | None = None,
description: str | None = None,
termination_condition: TerminationCondition | None = None,
max_turns: int | None = 20,
runtime: AgentRuntime | None = None,
max_stalls: int = 3,
final_answer_prompt: str = ORCHESTRATOR_FINAL_ANSWER_PROMPT,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
emit_team_events: bool = False,
):
for participant in participants:
if not isinstance(participant, ChatAgent):
raise TypeError(f"Participant {participant} must be a ChatAgent.")
super().__init__(
name=name or self.DEFAULT_NAME,
description=description or self.DEFAULT_DESCRIPTION,
participants=list(participants),
group_chat_manager_name="MagenticOneOrchestrator",
group_chat_manager_class=MagenticOneOrchestrator,
termination_condition=termination_condition,
max_turns=max_turns,
runtime=runtime,
custom_message_types=custom_message_types,
emit_team_events=emit_team_events,
)
# Validate the participants.
if len(participants) == 0:
raise ValueError("At least one participant is required for MagenticOneGroupChat.")
self._model_client = model_client
self._max_stalls = max_stalls
self._final_answer_prompt = final_answer_prompt
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], MagenticOneOrchestrator]:
return lambda: MagenticOneOrchestrator(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
max_turns,
message_factory,
self._model_client,
self._max_stalls,
self._final_answer_prompt,
output_message_queue,
termination_condition,
self._emit_team_events,
)
def _to_config(self) -> MagenticOneGroupChatConfig:
participants = [participant.dump_component() for participant in self._participants]
termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
return MagenticOneGroupChatConfig(
name=self.name,
description=self.description,
participants=participants,
model_client=self._model_client.dump_component(),
termination_condition=termination_condition,
max_turns=self._max_turns,
max_stalls=self._max_stalls,
final_answer_prompt=self._final_answer_prompt,
emit_team_events=self._emit_team_events,
)
@classmethod
def _from_config(cls, config: MagenticOneGroupChatConfig) -> Self:
participants = [ChatAgent.load_component(participant) for participant in config.participants]
model_client = ChatCompletionClient.load_component(config.model_client)
termination_condition = (
TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
)
return cls(
participants=participants,
name=config.name,
description=config.description,
model_client=model_client,
termination_condition=termination_condition,
max_turns=config.max_turns,
max_stalls=config.max_stalls,
final_answer_prompt=config.final_answer_prompt,
emit_team_events=config.emit_team_events,
)

536
agent_dhal/agentdhal_agentchat/teams/_group_chat/_magentic_one/_magentic_one_orchestrator.py Normal file
View File

@@ -0,0 +1,536 @@
import asyncio
import json
import logging
import re
from typing import Any, Dict, List, Mapping, Sequence
from agentdhal_core import AgentId, CancellationToken, DefaultTopicId, MessageContext, event, rpc
from agentdhal_core.models import (
AssistantMessage,
ChatCompletionClient,
LLMMessage,
UserMessage,
)
from agentdhal_core.utils import extract_json_from_str
from .... import TRACE_LOGGER_NAME
from ....base import Response, TerminationCondition
from ....messages import (
BaseAgentEvent,
BaseChatMessage,
HandoffMessage,
MessageFactory,
MultiModalMessage,
SelectSpeakerEvent,
StopMessage,
TextMessage,
ToolCallExecutionEvent,
ToolCallRequestEvent,
ToolCallSummaryMessage,
)
from ....state import MagenticOneOrchestratorState
from ....utils import remove_images
from .._base_group_chat_manager import BaseGroupChatManager
from .._events import (
GroupChatAgentResponse,
GroupChatMessage,
GroupChatRequestPublish,
GroupChatReset,
GroupChatStart,
GroupChatTeamResponse,
GroupChatTermination,
SerializableException,
)
from ._prompts import (
ORCHESTRATOR_FINAL_ANSWER_PROMPT,
ORCHESTRATOR_PROGRESS_LEDGER_PROMPT,
ORCHESTRATOR_TASK_LEDGER_FACTS_PROMPT,
ORCHESTRATOR_TASK_LEDGER_FACTS_UPDATE_PROMPT,
ORCHESTRATOR_TASK_LEDGER_FULL_PROMPT,
ORCHESTRATOR_TASK_LEDGER_PLAN_PROMPT,
ORCHESTRATOR_TASK_LEDGER_PLAN_UPDATE_PROMPT,
LedgerEntry,
)
trace_logger = logging.getLogger(TRACE_LOGGER_NAME)
class MagenticOneOrchestrator(BaseGroupChatManager):
"""The MagenticOneOrchestrator manages a group chat with ledger based orchestration."""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
max_turns: int | None,
message_factory: MessageFactory,
model_client: ChatCompletionClient,
max_stalls: int,
final_answer_prompt: str,
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
emit_team_events: bool,
):
super().__init__(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
emit_team_events=emit_team_events,
)
self._model_client = model_client
self._max_stalls = max_stalls
self._final_answer_prompt = final_answer_prompt
self._max_json_retries = 10
self._task = ""
self._facts = ""
self._plan = ""
self._n_rounds = 0
self._n_stalls = 0
# Produce a team description. Each agent sould appear on a single line.
self._team_description = ""
for topic_type, description in zip(self._participant_names, self._participant_descriptions, strict=True):
self._team_description += re.sub(r"\s+", " ", f"{topic_type}: {description}").strip() + "\n"
self._team_description = self._team_description.strip()
def _get_task_ledger_facts_prompt(self, task: str) -> str:
return ORCHESTRATOR_TASK_LEDGER_FACTS_PROMPT.format(task=task)
def _get_task_ledger_plan_prompt(self, team: str) -> str:
return ORCHESTRATOR_TASK_LEDGER_PLAN_PROMPT.format(team=team)
def _get_task_ledger_full_prompt(self, task: str, team: str, facts: str, plan: str) -> str:
return ORCHESTRATOR_TASK_LEDGER_FULL_PROMPT.format(task=task, team=team, facts=facts, plan=plan)
def _get_progress_ledger_prompt(self, task: str, team: str, names: List[str]) -> str:
return ORCHESTRATOR_PROGRESS_LEDGER_PROMPT.format(task=task, team=team, names=", ".join(names))
def _get_task_ledger_facts_update_prompt(self, task: str, facts: str) -> str:
return ORCHESTRATOR_TASK_LEDGER_FACTS_UPDATE_PROMPT.format(task=task, facts=facts)
def _get_task_ledger_plan_update_prompt(self, team: str) -> str:
return ORCHESTRATOR_TASK_LEDGER_PLAN_UPDATE_PROMPT.format(team=team)
def _get_final_answer_prompt(self, task: str) -> str:
if self._final_answer_prompt == ORCHESTRATOR_FINAL_ANSWER_PROMPT:
return ORCHESTRATOR_FINAL_ANSWER_PROMPT.format(task=task)
else:
return self._final_answer_prompt
async def _log_message(self, log_message: str) -> None:
trace_logger.debug(log_message)
@rpc
async def handle_start(self, message: GroupChatStart, ctx: MessageContext) -> None: # type: ignore
"""Handle the start of a task."""
# Check if the conversation has already terminated.
if self._termination_condition is not None and self._termination_condition.terminated:
early_stop_message = StopMessage(content="The group chat has already terminated.", source=self._name)
# Signal termination.
await self._signal_termination(early_stop_message)
# Stop the group chat.
return
assert message is not None and message.messages is not None
# Validate the group state given all the messages.
await self.validate_group_state(message.messages)
# Log the message to the output topic.
await self.publish_message(message, topic_id=DefaultTopicId(type=self._output_topic_type))
# Log the message to the output queue.
for msg in message.messages:
await self._output_message_queue.put(msg)
# Outer Loop for first time
# Create the initial task ledger
#################################
# Combine all message contents for task
self._task = " ".join([msg.to_model_text() for msg in message.messages])
planning_conversation: List[LLMMessage] = []
# 1. GATHER FACTS
# create a closed book task and generate a response and update the chat history
planning_conversation.append(
UserMessage(content=self._get_task_ledger_facts_prompt(self._task), source=self._name)
)
response = await self._model_client.create(
self._get_compatible_context(planning_conversation), cancellation_token=ctx.cancellation_token
)
assert isinstance(response.content, str)
self._facts = response.content
planning_conversation.append(AssistantMessage(content=self._facts, source=self._name))
# 2. CREATE A PLAN
## plan based on available information
planning_conversation.append(
UserMessage(content=self._get_task_ledger_plan_prompt(self._team_description), source=self._name)
)
response = await self._model_client.create(
self._get_compatible_context(planning_conversation), cancellation_token=ctx.cancellation_token
)
assert isinstance(response.content, str)
self._plan = response.content
# Kick things off
self._n_stalls = 0
await self._reenter_outer_loop(ctx.cancellation_token)
@event
async def handle_agent_response( # type: ignore
self, message: GroupChatAgentResponse | GroupChatTeamResponse, ctx: MessageContext
) -> None: # type: ignore
try:
if not isinstance(message, GroupChatAgentResponse):
raise RuntimeError("MagenticOneOrchestrator does not support GroupChatTeamResponse messages.")
delta: List[BaseAgentEvent | BaseChatMessage] = []
if message.response.inner_messages is not None:
for inner_message in message.response.inner_messages:
delta.append(inner_message)
await self.update_message_thread([message.response.chat_message])
delta.append(message.response.chat_message)
if self._termination_condition is not None:
stop_message = await self._termination_condition(delta)
if stop_message is not None:
# Reset the termination conditions.
await self._termination_condition.reset()
# Signal termination.
await self._signal_termination(stop_message)
return
await self._orchestrate_step(ctx.cancellation_token)
except Exception as e:
error = SerializableException.from_exception(e)
await self._signal_termination_with_error(error)
# Raise the error to the runtime.
raise
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
async def save_state(self) -> Mapping[str, Any]:
state = MagenticOneOrchestratorState(
message_thread=[msg.dump() for msg in self._message_thread],
current_turn=self._current_turn,
task=self._task,
facts=self._facts,
plan=self._plan,
n_rounds=self._n_rounds,
n_stalls=self._n_stalls,
)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
orchestrator_state = MagenticOneOrchestratorState.model_validate(state)
self._message_thread = [self._message_factory.create(message) for message in orchestrator_state.message_thread]
self._current_turn = orchestrator_state.current_turn
self._task = orchestrator_state.task
self._facts = orchestrator_state.facts
self._plan = orchestrator_state.plan
self._n_rounds = orchestrator_state.n_rounds
self._n_stalls = orchestrator_state.n_stalls
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Not used in this orchestrator, we select next speaker in _orchestrate_step."""
return [""]
async def reset(self) -> None:
"""Reset the group chat manager."""
self._message_thread.clear()
if self._termination_condition is not None:
await self._termination_condition.reset()
self._n_rounds = 0
self._n_stalls = 0
self._task = ""
self._facts = ""
self._plan = ""
async def _reenter_outer_loop(self, cancellation_token: CancellationToken) -> None:
"""Re-enter Outer loop of the orchestrator after creating task ledger."""
# Reset the agents
for participant_topic_type in self._participant_name_to_topic_type.values():
await self._runtime.send_message(
GroupChatReset(),
recipient=AgentId(type=participant_topic_type, key=self.id.key),
cancellation_token=cancellation_token,
)
# Reset partially the group chat manager
self._message_thread.clear()
# Prepare the ledger
ledger_message = TextMessage(
content=self._get_task_ledger_full_prompt(self._task, self._team_description, self._facts, self._plan),
source=self._name,
)
# Save my copy
await self.update_message_thread([ledger_message])
# Log it to the output topic.
await self.publish_message(
GroupChatMessage(message=ledger_message),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Log it to the output queue.
await self._output_message_queue.put(ledger_message)
# Broadcast
await self.publish_message(
GroupChatAgentResponse(response=Response(chat_message=ledger_message), name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
)
# Restart the inner loop
await self._orchestrate_step(cancellation_token=cancellation_token)
async def _orchestrate_step(self, cancellation_token: CancellationToken) -> None:
"""Implements the inner loop of the orchestrator and selects next speaker."""
# Check if we reached the maximum number of rounds
if self._max_turns is not None and self._n_rounds > self._max_turns:
await self._prepare_final_answer("Max rounds reached.", cancellation_token)
return
self._n_rounds += 1
# Update the progress ledger
context = self._thread_to_context()
progress_ledger_prompt = self._get_progress_ledger_prompt(
self._task, self._team_description, self._participant_names
)
context.append(UserMessage(content=progress_ledger_prompt, source=self._name))
progress_ledger: Dict[str, Any] = {}
assert self._max_json_retries > 0
key_error: bool = False
for _ in range(self._max_json_retries):
if self._model_client.model_info.get("structured_output", False):
response = await self._model_client.create(
self._get_compatible_context(context), json_output=LedgerEntry
)
elif self._model_client.model_info.get("json_output", False):
response = await self._model_client.create(
self._get_compatible_context(context), cancellation_token=cancellation_token, json_output=True
)
else:
response = await self._model_client.create(
self._get_compatible_context(context), cancellation_token=cancellation_token
)
ledger_str = response.content
try:
assert isinstance(ledger_str, str)
output_json = extract_json_from_str(ledger_str)
if len(output_json) != 1:
raise ValueError(
f"Progress ledger should contain a single JSON object, but found: {len(progress_ledger)}"
)
progress_ledger = output_json[0]
# If the team consists of a single agent, deterministically set the next speaker
if len(self._participant_names) == 1:
progress_ledger["next_speaker"] = {
"reason": "The team consists of only one agent.",
"answer": self._participant_names[0],
}
# Validate the structure
required_keys = [
"is_request_satisfied",
"is_progress_being_made",
"is_in_loop",
"instruction_or_question",
"next_speaker",
]
key_error = False
for key in required_keys:
if (
key not in progress_ledger
or not isinstance(progress_ledger[key], dict)
or "answer" not in progress_ledger[key]
or "reason" not in progress_ledger[key]
):
key_error = True
break
# Validate the next speaker if the task is not yet complete
if (
not progress_ledger["is_request_satisfied"]["answer"]
and progress_ledger["next_speaker"]["answer"] not in self._participant_names
):
key_error = True
break
if not key_error:
break
await self._log_message(f"Failed to parse ledger information, retrying: {ledger_str}")
except (json.JSONDecodeError, TypeError):
key_error = True
await self._log_message("Invalid ledger format encountered, retrying...")
continue
if key_error:
raise ValueError("Failed to parse ledger information after multiple retries.")
await self._log_message(f"Progress Ledger: {progress_ledger}")
# Check for task completion
if progress_ledger["is_request_satisfied"]["answer"]:
await self._log_message("Task completed, preparing final answer...")
await self._prepare_final_answer(progress_ledger["is_request_satisfied"]["reason"], cancellation_token)
return
# Check for stalling
if not progress_ledger["is_progress_being_made"]["answer"]:
self._n_stalls += 1
elif progress_ledger["is_in_loop"]["answer"]:
self._n_stalls += 1
else:
self._n_stalls = max(0, self._n_stalls - 1)
# Too much stalling
if self._n_stalls >= self._max_stalls:
await self._log_message("Stall count exceeded, re-planning with the outer loop...")
await self._update_task_ledger(cancellation_token)
await self._reenter_outer_loop(cancellation_token)
return
# Broadcast the next step
message = TextMessage(content=progress_ledger["instruction_or_question"]["answer"], source=self._name)
await self.update_message_thread([message]) # My copy
await self._log_message(f"Next Speaker: {progress_ledger['next_speaker']['answer']}")
# Log it to the output topic.
await self.publish_message(
GroupChatMessage(message=message),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Log it to the output queue.
await self._output_message_queue.put(message)
# Broadcast it
await self.publish_message( # Broadcast
GroupChatAgentResponse(response=Response(chat_message=message), name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
cancellation_token=cancellation_token,
)
# Request that the step be completed
next_speaker = progress_ledger["next_speaker"]["answer"]
# Check if the next speaker is valid
if next_speaker not in self._participant_name_to_topic_type:
raise ValueError(
f"Invalid next speaker: {next_speaker} from the ledger, participants are: {self._participant_names}"
)
participant_topic_type = self._participant_name_to_topic_type[next_speaker]
await self.publish_message(
GroupChatRequestPublish(),
topic_id=DefaultTopicId(type=participant_topic_type),
cancellation_token=cancellation_token,
)
# Send the message to the next speaker
if self._emit_team_events:
select_msg = SelectSpeakerEvent(content=[next_speaker], source=self._name)
await self.publish_message(
GroupChatMessage(message=select_msg),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
await self._output_message_queue.put(select_msg)
async def _update_task_ledger(self, cancellation_token: CancellationToken) -> None:
"""Update the task ledger (outer loop) with the latest facts and plan."""
context = self._thread_to_context()
# Update the facts
update_facts_prompt = self._get_task_ledger_facts_update_prompt(self._task, self._facts)
context.append(UserMessage(content=update_facts_prompt, source=self._name))
response = await self._model_client.create(
self._get_compatible_context(context), cancellation_token=cancellation_token
)
assert isinstance(response.content, str)
self._facts = response.content
context.append(AssistantMessage(content=self._facts, source=self._name))
# Update the plan
update_plan_prompt = self._get_task_ledger_plan_update_prompt(self._team_description)
context.append(UserMessage(content=update_plan_prompt, source=self._name))
response = await self._model_client.create(
self._get_compatible_context(context), cancellation_token=cancellation_token
)
assert isinstance(response.content, str)
self._plan = response.content
async def _prepare_final_answer(self, reason: str, cancellation_token: CancellationToken) -> None:
"""Prepare the final answer for the task."""
context = self._thread_to_context()
# Get the final answer
final_answer_prompt = self._get_final_answer_prompt(self._task)
context.append(UserMessage(content=final_answer_prompt, source=self._name))
response = await self._model_client.create(
self._get_compatible_context(context), cancellation_token=cancellation_token
)
assert isinstance(response.content, str)
message = TextMessage(content=response.content, source=self._name)
await self.update_message_thread([message]) # My copy
# Log it to the output topic.
await self.publish_message(
GroupChatMessage(message=message),
topic_id=DefaultTopicId(type=self._output_topic_type),
)
# Log it to the output queue.
await self._output_message_queue.put(message)
# Broadcast
await self.publish_message(
GroupChatAgentResponse(response=Response(chat_message=message), name=self._name),
topic_id=DefaultTopicId(type=self._group_topic_type),
cancellation_token=cancellation_token,
)
if self._termination_condition is not None:
await self._termination_condition.reset()
# Signal termination
await self._signal_termination(StopMessage(content=reason, source=self._name))
def _thread_to_context(self) -> List[LLMMessage]:
"""Convert the message thread to a context for the model."""
context: List[LLMMessage] = []
for m in self._message_thread:
if isinstance(m, ToolCallRequestEvent | ToolCallExecutionEvent):
# Ignore tool call messages.
continue
elif isinstance(m, StopMessage | HandoffMessage):
context.append(UserMessage(content=m.content, source=m.source))
elif m.source == self._name:
assert isinstance(m, TextMessage | ToolCallSummaryMessage)
context.append(AssistantMessage(content=m.content, source=m.source))
else:
assert isinstance(m, (TextMessage, MultiModalMessage, ToolCallSummaryMessage))
context.append(UserMessage(content=m.content, source=m.source))
return context
def _get_compatible_context(self, messages: List[LLMMessage]) -> List[LLMMessage]:
"""Ensure that the messages are compatible with the underlying client, by removing images if needed."""
if self._model_client.model_info["vision"]:
return messages
else:
return remove_images(messages)

149
agent_dhal/agentdhal_agentchat/teams/_group_chat/_magentic_one/_prompts.py Normal file
View File

@@ -0,0 +1,149 @@
from pydantic import BaseModel
ORCHESTRATOR_SYSTEM_MESSAGE = ""
ORCHESTRATOR_TASK_LEDGER_FACTS_PROMPT = """Below I will present you a request. Before we begin addressing the request, please answer the following pre-survey to the best of your ability. Keep in mind that you are Ken Jennings-level with trivia, and Mensa-level with puzzles, so there should be a deep well to draw from.
Here is the request:
{task}
Here is the pre-survey:
1. Please list any specific facts or figures that are GIVEN in the request itself. It is possible that there are none.
2. Please list any facts that may need to be looked up, and WHERE SPECIFICALLY they might be found. In some cases, authoritative sources are mentioned in the request itself.
3. Please list any facts that may need to be derived (e.g., via logical deduction, simulation, or computation)
4. Please list any facts that are recalled from memory, hunches, well-reasoned guesses, etc.
When answering this survey, keep in mind that "facts" will typically be specific names, dates, statistics, etc. Your answer should use headings:
1. GIVEN OR VERIFIED FACTS
2. FACTS TO LOOK UP
3. FACTS TO DERIVE
4. EDUCATED GUESSES
DO NOT include any other headings or sections in your response. DO NOT list next steps or plans until asked to do so.
"""
ORCHESTRATOR_TASK_LEDGER_PLAN_PROMPT = """Fantastic. To address this request we have assembled the following team:
{team}
Based on the team composition, and known and unknown facts, please devise a short bullet-point plan for addressing the original request. Remember, there is no requirement to involve all team members -- a team member's particular expertise may not be needed for this task."""
ORCHESTRATOR_TASK_LEDGER_FULL_PROMPT = """
We are working to address the following user request:
{task}
To answer this request we have assembled the following team:
{team}
Here is an initial fact sheet to consider:
{facts}
Here is the plan to follow as best as possible:
{plan}
"""
ORCHESTRATOR_PROGRESS_LEDGER_PROMPT = """
Recall we are working on the following request:
{task}
And we have assembled the following team:
{team}
To make progress on the request, please answer the following questions, including necessary reasoning:
- Is the request fully satisfied? (True if complete, or False if the original request has yet to be SUCCESSFULLY and FULLY addressed)
- Are we in a loop where we are repeating the same requests and / or getting the same responses as before? Loops can span multiple turns, and can include repeated actions like scrolling up or down more than a handful of times.
- Are we making forward progress? (True if just starting, or recent messages are adding value. False if recent messages show evidence of being stuck in a loop or if there is evidence of significant barriers to success such as the inability to read from a required file)
- Who should speak next? (select from: {names})
- What instruction or question would you give this team member? (Phrase as if speaking directly to them, and include any specific information they may need)
Please output an answer in pure JSON format according to the following schema. The JSON object must be parsable as-is. DO NOT OUTPUT ANYTHING OTHER THAN JSON, AND DO NOT DEVIATE FROM THIS SCHEMA:
{{
"is_request_satisfied": {{
"reason": string,
"answer": boolean
}},
"is_in_loop": {{
"reason": string,
"answer": boolean
}},
"is_progress_being_made": {{
"reason": string,
"answer": boolean
}},
"next_speaker": {{
"reason": string,
"answer": string (select from: {names})
}},
"instruction_or_question": {{
"reason": string,
"answer": string
}}
}}
"""
class LedgerEntryBooleanAnswer(BaseModel):
reason: str
answer: bool
class LedgerEntryStringAnswer(BaseModel):
reason: str
answer: str
class LedgerEntry(BaseModel):
is_request_satisfied: LedgerEntryBooleanAnswer
is_in_loop: LedgerEntryBooleanAnswer
is_progress_being_made: LedgerEntryBooleanAnswer
next_speaker: LedgerEntryStringAnswer
instruction_or_question: LedgerEntryStringAnswer
ORCHESTRATOR_TASK_LEDGER_FACTS_UPDATE_PROMPT = """As a reminder, we are working to solve the following task:
{task}
It's clear we aren't making as much progress as we would like, but we may have learned something new. Please rewrite the following fact sheet, updating it to include anything new we have learned that may be helpful. Example edits can include (but are not limited to) adding new guesses, moving educated guesses to verified facts if appropriate, etc. Updates may be made to any section of the fact sheet, and more than one section of the fact sheet can be edited. This is an especially good time to update educated guesses, so please at least add or update one educated guess or hunch, and explain your reasoning.
Here is the old fact sheet:
{facts}
"""
ORCHESTRATOR_TASK_LEDGER_PLAN_UPDATE_PROMPT = """Please briefly explain what went wrong on this last run (the root cause of the failure), and then come up with a new plan that takes steps and/or includes hints to overcome prior challenges and especially avoids repeating the same mistakes. As before, the new plan should be concise, be expressed in bullet-point form, and consider the following team composition (do not involve any other outside people since we cannot contact anyone else):
{team}
"""
ORCHESTRATOR_FINAL_ANSWER_PROMPT = """
We are working on the following task:
{task}
We have completed the task.
The above messages contain the conversation that took place to complete the task.
Based on the information gathered, provide the final answer to the original request.
The answer should be phrased as if you were speaking to the user.
"""

328
agent_dhal/agentdhal_agentchat/teams/_group_chat/_round_robin_group_chat.py Normal file
View File

@@ -0,0 +1,328 @@
import asyncio
from typing import Any, Callable, List, Mapping, Sequence
from agentdhal_core import AgentRuntime, Component, ComponentModel
from pydantic import BaseModel
from typing_extensions import Self
from ...base import ChatAgent, Team, TerminationCondition
from ...messages import BaseAgentEvent, BaseChatMessage, MessageFactory
from ...state import RoundRobinManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
from ._events import GroupChatTermination
class RoundRobinGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker in a round-robin fashion."""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
emit_team_events: bool,
) -> None:
super().__init__(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
emit_team_events,
)
self._next_speaker_index = 0
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
async def reset(self) -> None:
self._current_turn = 0
self._message_thread.clear()
if self._termination_condition is not None:
await self._termination_condition.reset()
self._next_speaker_index = 0
async def save_state(self) -> Mapping[str, Any]:
state = RoundRobinManagerState(
message_thread=[message.dump() for message in self._message_thread],
current_turn=self._current_turn,
next_speaker_index=self._next_speaker_index,
)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
round_robin_state = RoundRobinManagerState.model_validate(state)
self._message_thread = [self._message_factory.create(message) for message in round_robin_state.message_thread]
self._current_turn = round_robin_state.current_turn
self._next_speaker_index = round_robin_state.next_speaker_index
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Select a speaker from the participants in a round-robin fashion.
.. note::
This method always returns a single speaker.
"""
current_speaker_index = self._next_speaker_index
self._next_speaker_index = (current_speaker_index + 1) % len(self._participant_names)
current_speaker = self._participant_names[current_speaker_index]
return current_speaker
class RoundRobinGroupChatConfig(BaseModel):
"""The declarative configuration RoundRobinGroupChat."""
name: str | None = None
description: str | None = None
participants: List[ComponentModel]
termination_condition: ComponentModel | None = None
max_turns: int | None = None
emit_team_events: bool = False
class RoundRobinGroupChat(BaseGroupChat, Component[RoundRobinGroupChatConfig]):
"""A team that runs a group chat with participants taking turns in a round-robin fashion
to publish a message to all.
If an :class:`~agentdhal_agentchat.base.ChatAgent` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage` from the agent response's
:attr:`~agentdhal_agentchat.base.Response.chat_message` will be published
to other participants in the group chat.
If a :class:`~agentdhal_agentchat.base.Team` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage`
from the team result' :attr:`~agentdhal_agentchat.base.TaskResult.messages` will be published
to other participants in the group chat.
If a single participant is in the team, the participant will be the only speaker.
Args:
participants (List[ChatAgent | Team]): The participants in the group chat.
name (str | None, optional): The name of the group chat, using :attr:`~agentdhal_agentchat.teams.RoundRobinGroupChat.DEFAULT_NAME` if not provided.
The name is used by a parent team to identify this group chat so it must be unique within the parent team.
description (str | None, optional): The description of the group chat, using :attr:`~agentdhal_agentchat.teams.RoundRobinGroupChat.DEFAULT_DESCRIPTION` if not provided.
termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None.
Without a termination condition, the group chat will run indefinitely.
max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit.
custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat.
If you are using custom message types or your agents produces custom message types, you need to specify them here.
Make sure your custom message types are subclasses of :class:`~agentdhal_agentchat.messages.BaseAgentEvent` or :class:`~agentdhal_agentchat.messages.BaseChatMessage`.
emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False.
Raises:
ValueError: If no participants are provided or if participant names are not unique.
Examples:
A team with one participant with tools:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.ui import Console
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
async def get_weather(location: str) -> str:
return f"The weather in {location} is sunny."
assistant = AssistantAgent(
"Assistant",
model_client=model_client,
tools=[get_weather],
)
termination = TextMentionTermination("TERMINATE")
team = RoundRobinGroupChat([assistant], termination_condition=termination)
await Console(team.run_stream(task="What's the weather in New York?"))
asyncio.run(main())
A team with multiple participants:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.ui import Console
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent("Assistant1", model_client=model_client)
agent2 = AssistantAgent("Assistant2", model_client=model_client)
termination = TextMentionTermination("TERMINATE")
team = RoundRobinGroupChat([agent1, agent2], termination_condition=termination)
await Console(team.run_stream(task="Tell me some jokes."))
asyncio.run(main())
A team of user proxy and a nested team of writer and reviewer agents:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import UserProxyAgent, AssistantAgent
from agentdhal_agentchat.conditions import TextMentionTermination, MaxMessageTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4.1-nano")
writer = AssistantAgent(
"writer", model_client=model_client, system_message="You are a writer.", model_client_stream=True
)
reviewer = AssistantAgent(
"reviewer",
model_client=model_client,
system_message="Provide feedback to the input and suggest improvements.",
model_client_stream=True,
)
# NOTE: you can skip input by pressing Enter.
user_proxy = UserProxyAgent("user_proxy")
# Maximum 1 round of review and revision.
inner_termination = MaxMessageTermination(max_messages=4)
# The outter-loop termination condition that will terminate the team when the user types "exit".
outter_termination = TextMentionTermination("exit", sources=["user_proxy"])
team = RoundRobinGroupChat(
[
# For each turn, the writer writes a summary and the reviewer reviews it.
RoundRobinGroupChat([writer, reviewer], termination_condition=inner_termination),
# The user proxy gets user input once the writer and reviewer have finished their actions.
user_proxy,
],
termination_condition=outter_termination,
)
# Start the team and wait for it to terminate.
await Console(team.run_stream(task="Write a short essay about the impact of AI on society."))
asyncio.run(main())
"""
component_config_schema = RoundRobinGroupChatConfig
component_provider_override = "agentdhal_agentchat.teams.RoundRobinGroupChat"
DEFAULT_NAME = "RoundRobinGroupChat"
DEFAULT_DESCRIPTION = "A team of agents."
def __init__(
self,
participants: List[ChatAgent | Team],
*,
name: str | None = None,
description: str | None = None,
termination_condition: TerminationCondition | None = None,
max_turns: int | None = None,
runtime: AgentRuntime | None = None,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
emit_team_events: bool = False,
) -> None:
super().__init__(
name=name or self.DEFAULT_NAME,
description=description or self.DEFAULT_DESCRIPTION,
participants=participants,
group_chat_manager_name="RoundRobinGroupChatManager",
group_chat_manager_class=RoundRobinGroupChatManager,
termination_condition=termination_condition,
max_turns=max_turns,
runtime=runtime,
custom_message_types=custom_message_types,
emit_team_events=emit_team_events,
)
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], RoundRobinGroupChatManager]:
def _factory() -> RoundRobinGroupChatManager:
return RoundRobinGroupChatManager(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
self._emit_team_events,
)
return _factory
def _to_config(self) -> RoundRobinGroupChatConfig:
participants = [participant.dump_component() for participant in self._participants]
termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
return RoundRobinGroupChatConfig(
name=self._name,
description=self._description,
participants=participants,
termination_condition=termination_condition,
max_turns=self._max_turns,
emit_team_events=self._emit_team_events,
)
@classmethod
def _from_config(cls, config: RoundRobinGroupChatConfig) -> Self:
participants: List[ChatAgent | Team] = []
for participant in config.participants:
if participant.component_type == Team.component_type:
participants.append(Team.load_component(participant))
else:
participants.append(ChatAgent.load_component(participant))
termination_condition = (
TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
)
return cls(
participants,
name=config.name,
description=config.description,
termination_condition=termination_condition,
max_turns=config.max_turns,
emit_team_events=config.emit_team_events,
)

730
agent_dhal/agentdhal_agentchat/teams/_group_chat/_selector_group_chat.py Normal file
View File

@@ -0,0 +1,730 @@
import asyncio
import logging
import re
from inspect import iscoroutinefunction
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Union, cast
from agentdhal_core import AgentRuntime, CancellationToken, Component, ComponentModel
from agentdhal_core.model_context import (
ChatCompletionContext,
UnboundedChatCompletionContext,
)
from agentdhal_core.models import (
AssistantMessage,
ChatCompletionClient,
CreateResult,
LLMMessage,
ModelFamily,
SystemMessage,
UserMessage,
)
from pydantic import BaseModel
from typing_extensions import Self
from ... import TRACE_LOGGER_NAME
from ...base import ChatAgent, Team, TerminationCondition
from ...messages import (
BaseAgentEvent,
BaseChatMessage,
HandoffMessage,
MessageFactory,
ModelClientStreamingChunkEvent,
SelectorEvent,
)
from ...state import SelectorManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
from ._events import GroupChatTermination
trace_logger = logging.getLogger(TRACE_LOGGER_NAME)
SyncSelectorFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None]
AsyncSelectorFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[str | None]]
SelectorFuncType = Union[SyncSelectorFunc | AsyncSelectorFunc]
SyncCandidateFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], List[str]]
AsyncCandidateFunc = Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[List[str]]]
CandidateFuncType = Union[SyncCandidateFunc | AsyncCandidateFunc]
class SelectorGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker using a ChatCompletion
model and a custom selector function."""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
model_client: ChatCompletionClient,
selector_prompt: str,
allow_repeated_speaker: bool,
selector_func: Optional[SelectorFuncType],
max_selector_attempts: int,
candidate_func: Optional[CandidateFuncType],
emit_team_events: bool,
model_context: ChatCompletionContext | None,
model_client_streaming: bool = False,
) -> None:
super().__init__(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
emit_team_events,
)
self._model_client = model_client
self._selector_prompt = selector_prompt
self._previous_speaker: str | None = None
self._allow_repeated_speaker = allow_repeated_speaker
self._selector_func = selector_func
self._is_selector_func_async = iscoroutinefunction(self._selector_func)
self._max_selector_attempts = max_selector_attempts
self._candidate_func = candidate_func
self._is_candidate_func_async = iscoroutinefunction(self._candidate_func)
self._model_client_streaming = model_client_streaming
if model_context is not None:
self._model_context = model_context
else:
self._model_context = UnboundedChatCompletionContext()
self._cancellation_token = CancellationToken()
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
pass
async def reset(self) -> None:
self._current_turn = 0
self._message_thread.clear()
await self._model_context.clear()
if self._termination_condition is not None:
await self._termination_condition.reset()
self._previous_speaker = None
async def save_state(self) -> Mapping[str, Any]:
state = SelectorManagerState(
message_thread=[msg.dump() for msg in self._message_thread],
current_turn=self._current_turn,
previous_speaker=self._previous_speaker,
)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
selector_state = SelectorManagerState.model_validate(state)
self._message_thread = [self._message_factory.create(msg) for msg in selector_state.message_thread]
await self._add_messages_to_context(
self._model_context, [msg for msg in self._message_thread if isinstance(msg, BaseChatMessage)]
)
self._current_turn = selector_state.current_turn
self._previous_speaker = selector_state.previous_speaker
@staticmethod
async def _add_messages_to_context(
model_context: ChatCompletionContext,
messages: Sequence[BaseChatMessage],
) -> None:
"""
Add incoming messages to the model context.
"""
for msg in messages:
if isinstance(msg, HandoffMessage):
for llm_msg in msg.context:
await model_context.add_message(llm_msg)
await model_context.add_message(msg.to_model_message())
async def update_message_thread(self, messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> None:
self._message_thread.extend(messages)
base_chat_messages = [m for m in messages if isinstance(m, BaseChatMessage)]
await self._add_messages_to_context(self._model_context, base_chat_messages)
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Selects the next speaker in a group chat using a ChatCompletion client,
with the selector function as override if it returns a speaker name.
.. note::
This method always returns a single speaker name.
A key assumption is that the agent type is the same as the topic type, which we use as the agent name.
"""
# Use the selector function if provided.
if self._selector_func is not None:
if self._is_selector_func_async:
async_selector_func = cast(AsyncSelectorFunc, self._selector_func)
speaker = await async_selector_func(thread)
else:
sync_selector_func = cast(SyncSelectorFunc, self._selector_func)
speaker = sync_selector_func(thread)
if speaker is not None:
if speaker not in self._participant_names:
raise ValueError(
f"Selector function returned an invalid speaker name: {speaker}. "
f"Expected one of: {self._participant_names}."
)
# Skip the model based selection.
return [speaker]
# Use the candidate function to filter participants if provided
if self._candidate_func is not None:
if self._is_candidate_func_async:
async_candidate_func = cast(AsyncCandidateFunc, self._candidate_func)
participants = await async_candidate_func(thread)
else:
sync_candidate_func = cast(SyncCandidateFunc, self._candidate_func)
participants = sync_candidate_func(thread)
if not participants:
raise ValueError("Candidate function must return a non-empty list of participant names.")
if not all(p in self._participant_names for p in participants):
raise ValueError(
f"Candidate function returned invalid participant names: {participants}. "
f"Expected one of: {self._participant_names}."
)
else:
# Construct the candidate agent list to be selected from, skip the previous speaker if not allowed.
if self._previous_speaker is not None and not self._allow_repeated_speaker:
participants = [p for p in self._participant_names if p != self._previous_speaker]
else:
participants = list(self._participant_names)
assert len(participants) > 0
# Construct agent roles.
# Each agent sould appear on a single line.
roles = ""
for topic_type, description in zip(self._participant_names, self._participant_descriptions, strict=True):
roles += re.sub(r"\s+", " ", f"{topic_type}: {description}").strip() + "\n"
roles = roles.strip()
# Select the next speaker.
if len(participants) > 1:
agent_name = await self._select_speaker(roles, participants, self._max_selector_attempts)
else:
agent_name = participants[0]
self._previous_speaker = agent_name
trace_logger.debug(f"Selected speaker: {agent_name}")
return [agent_name]
def construct_message_history(self, message_history: List[LLMMessage]) -> str:
# Construct the history of the conversation.
history_messages: List[str] = []
for msg in message_history:
if isinstance(msg, UserMessage) or isinstance(msg, AssistantMessage):
message = f"{msg.source}: {msg.content}"
history_messages.append(
message.rstrip() + "\n\n"
) # Create some consistency for how messages are separated in the transcript
history: str = "\n".join(history_messages)
return history
async def _select_speaker(self, roles: str, participants: List[str], max_attempts: int) -> str:
model_context_messages = await self._model_context.get_messages()
model_context_history = self.construct_message_history(model_context_messages)
select_speaker_prompt = self._selector_prompt.format(
roles=roles, participants=str(participants), history=model_context_history
)
select_speaker_messages: List[SystemMessage | UserMessage | AssistantMessage]
if ModelFamily.is_openai(self._model_client.model_info["family"]):
select_speaker_messages = [SystemMessage(content=select_speaker_prompt)]
else:
# Many other models need a UserMessage to respond to
select_speaker_messages = [UserMessage(content=select_speaker_prompt, source="user")]
num_attempts = 0
while num_attempts < max_attempts:
num_attempts += 1
if self._model_client_streaming:
chunk: CreateResult | str = ""
async for _chunk in self._model_client.create_stream(messages=select_speaker_messages):
chunk = _chunk
if self._emit_team_events:
if isinstance(chunk, str):
await self._output_message_queue.put(
ModelClientStreamingChunkEvent(content=cast(str, _chunk), source=self._name)
)
else:
assert isinstance(chunk, CreateResult)
assert isinstance(chunk.content, str)
await self._output_message_queue.put(
SelectorEvent(content=chunk.content, source=self._name)
)
# The last chunk must be CreateResult.
assert isinstance(chunk, CreateResult)
response = chunk
else:
response = await self._model_client.create(messages=select_speaker_messages)
assert isinstance(response.content, str)
select_speaker_messages.append(AssistantMessage(content=response.content, source="selector"))
# NOTE: we use all participant names to check for mentions, even if the previous speaker is not allowed.
# This is because the model may still select the previous speaker, and we want to catch that.
mentions = self._mentioned_agents(response.content, self._participant_names)
if len(mentions) == 0:
trace_logger.debug(f"Model failed to select a valid name: {response.content} (attempt {num_attempts})")
feedback = f"No valid name was mentioned. Please select from: {str(participants)}."
select_speaker_messages.append(UserMessage(content=feedback, source="user"))
elif len(mentions) > 1:
trace_logger.debug(f"Model selected multiple names: {str(mentions)} (attempt {num_attempts})")
feedback = (
f"Expected exactly one name to be mentioned. Please select only one from: {str(participants)}."
)
select_speaker_messages.append(UserMessage(content=feedback, source="user"))
else:
agent_name = list(mentions.keys())[0]
if (
not self._allow_repeated_speaker
and self._previous_speaker is not None
and agent_name == self._previous_speaker
):
trace_logger.debug(f"Model selected the previous speaker: {agent_name} (attempt {num_attempts})")
feedback = (
f"Repeated speaker is not allowed, please select a different name from: {str(participants)}."
)
select_speaker_messages.append(UserMessage(content=feedback, source="user"))
else:
# Valid selection
trace_logger.debug(f"Model selected a valid name: {agent_name} (attempt {num_attempts})")
return agent_name
if self._previous_speaker is not None:
trace_logger.warning(f"Model failed to select a speaker after {max_attempts}, using the previous speaker.")
return self._previous_speaker
trace_logger.warning(
f"Model failed to select a speaker after {max_attempts} and there was no previous speaker, using the first participant."
)
return participants[0]
def _mentioned_agents(self, message_content: str, agent_names: List[str]) -> Dict[str, int]:
"""Counts the number of times each agent is mentioned in the provided message content.
Agent names will match under any of the following conditions (all case-sensitive):
- Exact name match
- If the agent name has underscores it will match with spaces instead (e.g. 'Story_writer' == 'Story writer')
- If the agent name has underscores it will match with '\\_' instead of '_' (e.g. 'Story_writer' == 'Story\\_writer')
Args:
message_content (Union[str, List]): The content of the message, either as a single string or a list of strings.
agents (List[Agent]): A list of Agent objects, each having a 'name' attribute to be searched in the message content.
Returns:
Dict: a counter for mentioned agents.
"""
mentions: Dict[str, int] = dict()
for name in agent_names:
# Finds agent mentions, taking word boundaries into account,
# accommodates escaping underscores and underscores as spaces
regex = (
r"(?<=\W)("
+ re.escape(name)
+ r"|"
+ re.escape(name.replace("_", " "))
+ r"|"
+ re.escape(name.replace("_", r"\_"))
+ r")(?=\W)"
)
# Pad the message to help with matching
count = len(re.findall(regex, f" {message_content} "))
if count > 0:
mentions[name] = count
return mentions
class SelectorGroupChatConfig(BaseModel):
"""The declarative configuration for SelectorGroupChat."""
name: str | None = None
description: str | None = None
participants: List[ComponentModel]
model_client: ComponentModel
termination_condition: ComponentModel | None = None
max_turns: int | None = None
selector_prompt: str
allow_repeated_speaker: bool
# selector_func: ComponentModel | None
max_selector_attempts: int = 3
emit_team_events: bool = False
model_client_streaming: bool = False
model_context: ComponentModel | None = None
class SelectorGroupChat(BaseGroupChat, Component[SelectorGroupChatConfig]):
"""A group chat team that have participants takes turn to publish a message
to all, using a ChatCompletion model to select the next speaker after each message.
If an :class:`~agentdhal_agentchat.base.ChatAgent` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage` from the agent response's
:attr:`~agentdhal_agentchat.base.Response.chat_message` will be published
to other participants in the group chat.
If a :class:`~agentdhal_agentchat.base.Team` is a participant,
the :class:`~agentdhal_agentchat.messages.BaseChatMessage`
from the team result' :attr:`~agentdhal_agentchat.base.TaskResult.messages` will be published
to other participants in the group chat.
Args:
participants (List[ChatAgent | Team]): The participants in the group chat,
must have unique names and at least two participants.
model_client (ChatCompletionClient): The ChatCompletion model client used
to select the next speaker.
name (str | None, optional): The name of the group chat, using
:attr:`~agentdhal_agentchat.teams.SelectorGroupChat.DEFAULT_NAME` if not provided.
The name is used by a parent team to identify this group chat so it must
be unique within the parent team.
description (str | None, optional): The description of the group chat, using
:attr:`~agentdhal_agentchat.teams.SelectorGroupChat.DEFAULT_DESCRIPTION` if not provided.
termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None.
Without a termination condition, the group chat will run indefinitely.
max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit.
selector_prompt (str, optional): The prompt template to use for selecting the next speaker.
Available fields: '{roles}', '{participants}', and '{history}'.
`{participants}` is the names of candidates for selection. The format is `["<name1>", "<name2>", ...]`.
`{roles}` is a newline-separated list of names and descriptions of the candidate agents. The format for each line is: `"<name> : <description>"`.
`{history}` is the conversation history formatted as a double newline separated of names and message content. The format for each message is: `"<name> : <message content>"`.
allow_repeated_speaker (bool, optional): Whether to include the previous speaker in the list of candidates to be selected for the next turn.
Defaults to False. The model may still select the previous speaker -- a warning will be logged if this happens.
max_selector_attempts (int, optional): The maximum number of attempts to select a speaker using the model. Defaults to 3.
If the model fails to select a speaker after the maximum number of attempts, the previous speaker will be used if available,
otherwise the first participant will be used.
selector_func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None], Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[str | None]], optional): A custom selector
function that takes the conversation history and returns the name of the next speaker.
If provided, this function will be used to override the model to select the next speaker.
If the function returns None, the model will be used to select the next speaker.
NOTE: `selector_func` is not serializable and will be ignored during serialization and deserialization process.
candidate_func (Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], List[str]], Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], Awaitable[List[str]]], optional):
A custom function that takes the conversation history and returns a filtered list of candidates for the next speaker
selection using model. If the function returns an empty list or `None`, `SelectorGroupChat` will raise a `ValueError`.
This function is only used if `selector_func` is not set. The `allow_repeated_speaker` will be ignored if set.
custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat.
If you are using custom message types or your agents produces custom message types, you need to specify them here.
Make sure your custom message types are subclasses of :class:`~agentdhal_agentchat.messages.BaseAgentEvent` or :class:`~agentdhal_agentchat.messages.BaseChatMessage`.
emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False.
model_client_streaming (bool, optional): Whether to use streaming for the model client. (This is useful for reasoning models like QwQ). Defaults to False.
model_context (ChatCompletionContext | None, optional): The model context for storing and retrieving
:class:`~agentdhal_core.models.LLMMessage`. It can be preloaded with initial messages. Messages stored in model context will be used for speaker selection. The initial messages will be cleared when the team is reset.
Raises:
ValueError: If the number of participants is less than two or if the selector prompt is invalid.
Examples:
A team with multiple participants:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import SelectorGroupChat
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.ui import Console
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
async def lookup_hotel(location: str) -> str:
return f"Here are some hotels in {location}: hotel1, hotel2, hotel3."
async def lookup_flight(origin: str, destination: str) -> str:
return f"Here are some flights from {origin} to {destination}: flight1, flight2, flight3."
async def book_trip() -> str:
return "Your trip is booked!"
travel_advisor = AssistantAgent(
"Travel_Advisor",
model_client,
tools=[book_trip],
description="Helps with travel planning.",
)
hotel_agent = AssistantAgent(
"Hotel_Agent",
model_client,
tools=[lookup_hotel],
description="Helps with hotel booking.",
)
flight_agent = AssistantAgent(
"Flight_Agent",
model_client,
tools=[lookup_flight],
description="Helps with flight booking.",
)
termination = TextMentionTermination("TERMINATE")
team = SelectorGroupChat(
[travel_advisor, hotel_agent, flight_agent],
model_client=model_client,
termination_condition=termination,
)
await Console(team.run_stream(task="Book a 3-day trip to new york."))
asyncio.run(main())
A team with a custom selector function:
.. code-block:: python
import asyncio
from typing import Sequence
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import SelectorGroupChat
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.ui import Console
from agentdhal_agentchat.messages import BaseAgentEvent, BaseChatMessage
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
def check_calculation(x: int, y: int, answer: int) -> str:
if x + y == answer:
return "Correct!"
else:
return "Incorrect!"
agent1 = AssistantAgent(
"Agent1",
model_client,
description="For calculation",
system_message="Calculate the sum of two numbers",
)
agent2 = AssistantAgent(
"Agent2",
model_client,
tools=[check_calculation],
description="For checking calculation",
system_message="Check the answer and respond with 'Correct!' or 'Incorrect!'",
)
def selector_func(messages: Sequence[BaseAgentEvent | BaseChatMessage]) -> str | None:
if len(messages) == 1 or messages[-1].to_text() == "Incorrect!":
return "Agent1"
if messages[-1].source == "Agent1":
return "Agent2"
return None
termination = TextMentionTermination("Correct!")
team = SelectorGroupChat(
[agent1, agent2],
model_client=model_client,
selector_func=selector_func,
termination_condition=termination,
)
await Console(team.run_stream(task="What is 1 + 1?"))
asyncio.run(main())
A team with custom model context:
.. code-block:: python
import asyncio
from agentdhal_core.model_context import BufferedChatCompletionContext
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import TextMentionTermination
from agentdhal_agentchat.teams import SelectorGroupChat
from agentdhal_agentchat.ui import Console
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
model_context = BufferedChatCompletionContext(buffer_size=5)
async def lookup_hotel(location: str) -> str:
return f"Here are some hotels in {location}: hotel1, hotel2, hotel3."
async def lookup_flight(origin: str, destination: str) -> str:
return f"Here are some flights from {origin} to {destination}: flight1, flight2, flight3."
async def book_trip() -> str:
return "Your trip is booked!"
travel_advisor = AssistantAgent(
"Travel_Advisor",
model_client,
tools=[book_trip],
description="Helps with travel planning.",
)
hotel_agent = AssistantAgent(
"Hotel_Agent",
model_client,
tools=[lookup_hotel],
description="Helps with hotel booking.",
)
flight_agent = AssistantAgent(
"Flight_Agent",
model_client,
tools=[lookup_flight],
description="Helps with flight booking.",
)
termination = TextMentionTermination("TERMINATE")
team = SelectorGroupChat(
[travel_advisor, hotel_agent, flight_agent],
model_client=model_client,
termination_condition=termination,
model_context=model_context,
)
await Console(team.run_stream(task="Book a 3-day trip to new york."))
asyncio.run(main())
"""
component_config_schema = SelectorGroupChatConfig
component_provider_override = "agentdhal_agentchat.teams.SelectorGroupChat"
DEFAULT_NAME = "SelectorGroupChat"
DEFAULT_DESCRIPTION = "A team of agents."
def __init__(
self,
participants: List[ChatAgent | Team],
model_client: ChatCompletionClient,
*,
name: str | None = None,
description: str | None = None,
termination_condition: TerminationCondition | None = None,
max_turns: int | None = None,
runtime: AgentRuntime | None = None,
selector_prompt: str = """You are in a role play game. The following roles are available:
{roles}.
Read the following conversation. Then select the next role from {participants} to play. Only return the role.
{history}
Read the above conversation. Then select the next role from {participants} to play. Only return the role.
""",
allow_repeated_speaker: bool = False,
max_selector_attempts: int = 3,
selector_func: Optional[SelectorFuncType] = None,
candidate_func: Optional[CandidateFuncType] = None,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
emit_team_events: bool = False,
model_client_streaming: bool = False,
model_context: ChatCompletionContext | None = None,
):
super().__init__(
name=name or self.DEFAULT_NAME,
description=description or self.DEFAULT_DESCRIPTION,
participants=participants,
group_chat_manager_name="SelectorGroupChatManager",
group_chat_manager_class=SelectorGroupChatManager,
termination_condition=termination_condition,
max_turns=max_turns,
runtime=runtime,
custom_message_types=custom_message_types,
emit_team_events=emit_team_events,
)
# Validate the participants.
if len(participants) < 2:
raise ValueError("At least two participants are required for SelectorGroupChat.")
self._selector_prompt = selector_prompt
self._model_client = model_client
self._allow_repeated_speaker = allow_repeated_speaker
self._selector_func = selector_func
self._max_selector_attempts = max_selector_attempts
self._candidate_func = candidate_func
self._model_client_streaming = model_client_streaming
self._model_context = model_context
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], BaseGroupChatManager]:
return lambda: SelectorGroupChatManager(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
self._model_client,
self._selector_prompt,
self._allow_repeated_speaker,
self._selector_func,
self._max_selector_attempts,
self._candidate_func,
self._emit_team_events,
self._model_context,
self._model_client_streaming,
)
def _to_config(self) -> SelectorGroupChatConfig:
return SelectorGroupChatConfig(
name=self._name,
description=self._description,
participants=[participant.dump_component() for participant in self._participants],
model_client=self._model_client.dump_component(),
termination_condition=self._termination_condition.dump_component() if self._termination_condition else None,
max_turns=self._max_turns,
selector_prompt=self._selector_prompt,
allow_repeated_speaker=self._allow_repeated_speaker,
max_selector_attempts=self._max_selector_attempts,
# selector_func=self._selector_func.dump_component() if self._selector_func else None,
emit_team_events=self._emit_team_events,
model_client_streaming=self._model_client_streaming,
model_context=self._model_context.dump_component() if self._model_context else None,
)
@classmethod
def _from_config(cls, config: SelectorGroupChatConfig) -> Self:
participants: List[ChatAgent | Team] = []
for participant in config.participants:
if participant.component_type == ChatAgent.component_type:
participants.append(ChatAgent.load_component(participant))
elif participant.component_type == Team.component_type:
participants.append(Team.load_component(participant))
else:
raise ValueError(
f"Invalid participant component type: {participant.component_type}. " "Expected ChatAgent or Team."
)
return cls(
participants=participants,
model_client=ChatCompletionClient.load_component(config.model_client),
name=config.name,
description=config.description,
termination_condition=TerminationCondition.load_component(config.termination_condition)
if config.termination_condition
else None,
max_turns=config.max_turns,
selector_prompt=config.selector_prompt,
allow_repeated_speaker=config.allow_repeated_speaker,
max_selector_attempts=config.max_selector_attempts,
# selector_func=ComponentLoader.load_component(config.selector_func, Callable[[Sequence[BaseAgentEvent | BaseChatMessage]], str | None])
# if config.selector_func
# else None,
emit_team_events=config.emit_team_events,
model_client_streaming=config.model_client_streaming,
model_context=ChatCompletionContext.load_component(config.model_context) if config.model_context else None,
)

72
agent_dhal/agentdhal_agentchat/teams/_group_chat/_sequential_routed_agent.py Normal file
View File

@@ -0,0 +1,72 @@
import asyncio
from typing import Any, Sequence
from agentdhal_core import MessageContext, RoutedAgent
class FIFOLock:
"""A lock that ensures coroutines acquire the lock in the order they request it."""
def __init__(self) -> None:
self._queue = asyncio.Queue[asyncio.Event]()
self._locked = False
async def acquire(self) -> None:
# If the lock is not held by any coroutine, set the lock to be held
# by the current coroutine.
if not self._locked:
self._locked = True
return
# If the lock is held by another coroutine, create an event and put it
# in the queue. Wait for the event to be set.
event = asyncio.Event()
await self._queue.put(event)
await event.wait()
def release(self) -> None:
if not self._queue.empty():
# If there are events in the queue, get the next event and set it.
next_event = self._queue.get_nowait()
next_event.set()
else:
# If there are no events in the queue, release the lock.
self._locked = False
class SequentialRoutedAgent(RoutedAgent):
"""A subclass of :class:`agentdhal_core.RoutedAgent` that ensures
that messages of certain types are processed sequentially
using a FIFO lock.
This is useful for agents that need to maintain a strict order of
processing messages, such as in a group chat scenario.
Args:
description (str): The description of the agent.
sequential_message_types (Sequence[Type[Any]]): A sequence of message types that should be
processed sequentially. If a message of one of these types is received,
the agent will acquire a FIFO lock to ensure that it is processed
before any later messages that are also one of these types.
"""
def __init__(self, description: str, sequential_message_types: Sequence[type[Any]]) -> None:
super().__init__(description=description)
self._fifo_lock = FIFOLock()
self._sequential_message_types = sequential_message_types
async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None:
if any(isinstance(message, sequential_type) for sequential_type in self._sequential_message_types):
# Acquire the FIFO lock to ensure that this message is processed
# in the order it was received.
await self._fifo_lock.acquire()
try:
return await super().on_message_impl(message, ctx)
finally:
# Release the FIFO lock to allow the next message to be processed.
self._fifo_lock.release()
# If the message is not of a sequential type, process it normally.
return await super().on_message_impl(message, ctx)

321
agent_dhal/agentdhal_agentchat/teams/_group_chat/_swarm_group_chat.py Normal file
View File

@@ -0,0 +1,321 @@
import asyncio
from typing import Any, Callable, List, Mapping, Sequence
from agentdhal_core import AgentRuntime, Component, ComponentModel
from pydantic import BaseModel
from ...base import ChatAgent, TerminationCondition
from ...messages import BaseAgentEvent, BaseChatMessage, HandoffMessage, MessageFactory
from ...state import SwarmManagerState
from ._base_group_chat import BaseGroupChat
from ._base_group_chat_manager import BaseGroupChatManager
from ._events import GroupChatTermination
class SwarmGroupChatManager(BaseGroupChatManager):
"""A group chat manager that selects the next speaker based on handoff message only."""
def __init__(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
emit_team_events: bool,
) -> None:
super().__init__(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
emit_team_events,
)
self._current_speaker = self._participant_names[0]
async def validate_group_state(self, messages: List[BaseChatMessage] | None) -> None:
"""Validate the start messages for the group chat."""
# Check if any of the start messages is a handoff message.
if messages:
for message in messages:
if isinstance(message, HandoffMessage):
if message.target not in self._participant_names:
raise ValueError(
f"The target {message.target} is not one of the participants {self._participant_names}. "
"If you are resuming Swarm with a new HandoffMessage make sure to set the target to a valid participant as the target."
)
return
# Check if there is a handoff message in the thread that is not targeting a valid participant.
for existing_message in reversed(self._message_thread):
if isinstance(existing_message, HandoffMessage):
if existing_message.target not in self._participant_names:
raise ValueError(
f"The existing handoff target {existing_message.target} is not one of the participants {self._participant_names}. "
"If you are resuming Swarm with a new task make sure to include in your task "
"a HandoffMessage with a valid participant as the target. For example, if you are "
"resuming from a HandoffTermination, make sure the new task is a HandoffMessage "
"with a valid participant as the target."
)
# The latest handoff message should always target a valid participant.
# Do not look past the latest handoff message.
return
async def reset(self) -> None:
self._current_turn = 0
self._message_thread.clear()
if self._termination_condition is not None:
await self._termination_condition.reset()
self._current_speaker = self._participant_names[0]
async def select_speaker(self, thread: Sequence[BaseAgentEvent | BaseChatMessage]) -> List[str] | str:
"""Select a speaker from the participants based on handoff message.
Looks for the last handoff message in the thread to determine the next speaker.
.. note::
This method always returns a single speaker.
"""
if len(thread) == 0:
return [self._current_speaker]
for message in reversed(thread):
if isinstance(message, HandoffMessage):
self._current_speaker = message.target
# The latest handoff message should always target a valid participant.
assert self._current_speaker in self._participant_names
return [self._current_speaker]
return self._current_speaker
async def save_state(self) -> Mapping[str, Any]:
state = SwarmManagerState(
message_thread=[msg.dump() for msg in self._message_thread],
current_turn=self._current_turn,
current_speaker=self._current_speaker,
)
return state.model_dump()
async def load_state(self, state: Mapping[str, Any]) -> None:
swarm_state = SwarmManagerState.model_validate(state)
self._message_thread = [self._message_factory.create(message) for message in swarm_state.message_thread]
self._current_turn = swarm_state.current_turn
self._current_speaker = swarm_state.current_speaker
class SwarmConfig(BaseModel):
"""The declarative configuration for Swarm."""
name: str | None = None
description: str | None = None
participants: List[ComponentModel]
termination_condition: ComponentModel | None = None
max_turns: int | None = None
emit_team_events: bool = False
class Swarm(BaseGroupChat, Component[SwarmConfig]):
"""A group chat team that selects the next speaker based on handoff message only.
The first participant in the list of participants is the initial speaker.
The next speaker is selected based on the :class:`~agentdhal_agentchat.messages.HandoffMessage` message
sent by the current speaker. If no handoff message is sent, the current speaker
continues to be the speaker.
.. note::
Unlike :class:`~agentdhal_agentchat.teams.RoundRobinGroupChat` and
:class:`~agentdhal_agentchat.teams.SelectorGroupChat`, this group chat
team does not support inner teams as participants.
Args:
participants (List[ChatAgent]): The agents participating in the group chat. The first agent in the list is the initial speaker.
name (str | None, optional): The name of the group chat, using :attr:`~agentdhal_agentchat.teams.Swarm.DEFAULT_NAME` if not provided.
The name is used by a parent team to identify this group chat so it must be unique within the parent team.
description (str | None, optional): The description of the group chat, using :attr:`~agentdhal_agentchat.teams.Swarm.DEFAULT_DESCRIPTION` if not provided.
termination_condition (TerminationCondition, optional): The termination condition for the group chat. Defaults to None.
Without a termination condition, the group chat will run indefinitely.
max_turns (int, optional): The maximum number of turns in the group chat before stopping. Defaults to None, meaning no limit.
custom_message_types (List[type[BaseAgentEvent | BaseChatMessage]], optional): A list of custom message types that will be used in the group chat.
If you are using custom message types or your agents produces custom message types, you need to specify them here.
Make sure your custom message types are subclasses of :class:`~agentdhal_agentchat.messages.BaseAgentEvent` or :class:`~agentdhal_agentchat.messages.BaseChatMessage`.
emit_team_events (bool, optional): Whether to emit team events through :meth:`BaseGroupChat.run_stream`. Defaults to False.
Basic example:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import Swarm
from agentdhal_agentchat.conditions import MaxMessageTermination
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent1 = AssistantAgent(
"Alice",
model_client=model_client,
handoffs=["Bob"],
system_message="You are Alice and you only answer questions about yourself.",
)
agent2 = AssistantAgent(
"Bob", model_client=model_client, system_message="You are Bob and your birthday is on 1st January."
)
termination = MaxMessageTermination(3)
team = Swarm([agent1, agent2], termination_condition=termination)
stream = team.run_stream(task="What is bob's birthday?")
async for message in stream:
print(message)
asyncio.run(main())
Using the :class:`~agentdhal_agentchat.conditions.HandoffTermination` for human-in-the-loop handoff:
.. code-block:: python
import asyncio
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.teams import Swarm
from agentdhal_agentchat.conditions import HandoffTermination, MaxMessageTermination
from agentdhal_agentchat.ui import Console
from agentdhal_agentchat.messages import HandoffMessage
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent = AssistantAgent(
"Alice",
model_client=model_client,
handoffs=["user"],
system_message="You are Alice and you only answer questions about yourself, ask the user for help if needed.",
)
termination = HandoffTermination(target="user") | MaxMessageTermination(3)
team = Swarm([agent], termination_condition=termination)
# Start the conversation.
await Console(team.run_stream(task="What is bob's birthday?"))
# Resume with user feedback.
await Console(
team.run_stream(
task=HandoffMessage(source="user", target="Alice", content="Bob's birthday is on 1st January.")
)
)
asyncio.run(main())
"""
component_config_schema = SwarmConfig
component_provider_override = "agentdhal_agentchat.teams.Swarm"
DEFAULT_NAME = "Swarm"
DEFAULT_DESCRIPTION = "A team of agents."
def __init__(
self,
participants: List[ChatAgent],
*,
name: str | None = None,
description: str | None = None,
termination_condition: TerminationCondition | None = None,
max_turns: int | None = None,
runtime: AgentRuntime | None = None,
custom_message_types: List[type[BaseAgentEvent | BaseChatMessage]] | None = None,
emit_team_events: bool = False,
) -> None:
for participant in participants:
if not isinstance(participant, ChatAgent):
raise TypeError(f"Participant {participant} must be a ChatAgent.")
super().__init__(
name=name or self.DEFAULT_NAME,
description=description or self.DEFAULT_DESCRIPTION,
participants=[participant for participant in participants],
group_chat_manager_name="SwarmGroupChatManager",
group_chat_manager_class=SwarmGroupChatManager,
termination_condition=termination_condition,
max_turns=max_turns,
runtime=runtime,
custom_message_types=custom_message_types,
emit_team_events=emit_team_events,
)
# The first participant must be able to produce handoff messages.
first_participant = self._participants[0]
assert isinstance(first_participant, ChatAgent)
if HandoffMessage not in first_participant.produced_message_types:
raise ValueError("The first participant must be able to produce a handoff messages.")
def _create_group_chat_manager_factory(
self,
name: str,
group_topic_type: str,
output_topic_type: str,
participant_topic_types: List[str],
participant_names: List[str],
participant_descriptions: List[str],
output_message_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | GroupChatTermination],
termination_condition: TerminationCondition | None,
max_turns: int | None,
message_factory: MessageFactory,
) -> Callable[[], SwarmGroupChatManager]:
def _factory() -> SwarmGroupChatManager:
return SwarmGroupChatManager(
name,
group_topic_type,
output_topic_type,
participant_topic_types,
participant_names,
participant_descriptions,
output_message_queue,
termination_condition,
max_turns,
message_factory,
self._emit_team_events,
)
return _factory
def _to_config(self) -> SwarmConfig:
participants = [participant.dump_component() for participant in self._participants]
termination_condition = self._termination_condition.dump_component() if self._termination_condition else None
return SwarmConfig(
name=self._name,
description=self._description,
participants=participants,
termination_condition=termination_condition,
max_turns=self._max_turns,
emit_team_events=self._emit_team_events,
)
@classmethod
def _from_config(cls, config: SwarmConfig) -> "Swarm":
participants = [ChatAgent.load_component(participant) for participant in config.participants]
termination_condition = (
TerminationCondition.load_component(config.termination_condition) if config.termination_condition else None
)
return cls(
participants,
name=config.name,
description=config.description,
termination_condition=termination_condition,
max_turns=config.max_turns,
emit_team_events=config.emit_team_events,
)

4
agent_dhal/agentdhal_agentchat/tools/__init__.py Normal file
View File

@@ -0,0 +1,4 @@
from ._agent import AgentTool
from ._team import TeamTool
__all__ = ["AgentTool", "TeamTool"]

93
agent_dhal/agentdhal_agentchat/tools/_agent.py Normal file
View File

@@ -0,0 +1,93 @@
from agentdhal_core import Component, ComponentModel
from pydantic import BaseModel
from typing_extensions import Self
from agentdhal_agentchat.agents import BaseChatAgent
from ._task_runner_tool import TaskRunnerTool
class AgentToolConfig(BaseModel):
"""Configuration for the AgentTool."""
agent: ComponentModel
"""The agent to be used for running the task."""
return_value_as_last_message: bool = False
"""Whether to return the value as the last message of the task result."""
class AgentTool(TaskRunnerTool, Component[AgentToolConfig]):
"""Tool that can be used to run a task using an agent.
The tool returns the result of the task execution as a :class:`~agentdhal_agentchat.base.TaskResult` object.
.. important::
When using AgentTool, you **must** disable parallel tool calls in the model client configuration
to avoid concurrency issues. Agents cannot run concurrently as they maintain internal state
that would conflict with parallel execution. For example, set ``parallel_tool_calls=False``
for :class:`~agentdhal_extensions.models.openai.OpenAIChatCompletionClient` and
:class:`~agentdhal_extensions.models.openai.AzureOpenAIChatCompletionClient`.
Args:
agent (BaseChatAgent): The agent to be used for running the task.
return_value_as_last_message (bool): Whether to use the last message content of the task result
as the return value of the tool in :meth:`~agentdhal_agentchat.tools.TaskRunnerTool.return_value_as_string`.
If set to True, the last message content will be returned as a string.
If set to False, the tool will return all messages in the task result as a string concatenated together,
with each message prefixed by its source (e.g., "writer: ...", "assistant: ...").
Example:
.. code-block:: python
import asyncio
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.tools import AgentTool
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
model_client = OpenAIChatCompletionClient(model="gpt-4.1")
writer = AssistantAgent(
name="writer",
description="A writer agent for generating text.",
model_client=model_client,
system_message="Write well.",
)
writer_tool = AgentTool(agent=writer)
# Create model client with parallel tool calls disabled for the main agent
main_model_client = OpenAIChatCompletionClient(model="gpt-4.1", parallel_tool_calls=False)
assistant = AssistantAgent(
name="assistant",
model_client=main_model_client,
tools=[writer_tool],
system_message="You are a helpful assistant.",
)
await Console(assistant.run_stream(task="Write a poem about the sea."))
asyncio.run(main())
"""
component_config_schema = AgentToolConfig
component_provider_override = "agentdhal_agentchat.tools.AgentTool"
def __init__(self, agent: BaseChatAgent, return_value_as_last_message: bool = False) -> None:
self._agent = agent
super().__init__(
agent, agent.name, agent.description, return_value_as_last_message=return_value_as_last_message
)
def _to_config(self) -> AgentToolConfig:
return AgentToolConfig(
agent=self._agent.dump_component(),
return_value_as_last_message=self._return_value_as_last_message,
)
@classmethod
def _from_config(cls, config: AgentToolConfig) -> Self:
return cls(BaseChatAgent.load_component(config.agent), config.return_value_as_last_message)

72
agent_dhal/agentdhal_agentchat/tools/_task_runner_tool.py Normal file
View File

@@ -0,0 +1,72 @@
from abc import ABC
from typing import Annotated, Any, AsyncGenerator, List, Mapping
from agentdhal_core import CancellationToken
from agentdhal_core.tools import BaseStreamTool
from pydantic import BaseModel
from ..agents import BaseChatAgent
from ..base import TaskResult
from ..messages import BaseAgentEvent, BaseChatMessage
from ..teams import BaseGroupChat
class TaskRunnerToolArgs(BaseModel):
"""Input for the TaskRunnerTool."""
task: Annotated[str, "The task to be executed."]
class TaskRunnerTool(BaseStreamTool[TaskRunnerToolArgs, BaseAgentEvent | BaseChatMessage, TaskResult], ABC):
"""An base class for tool that can be used to run a task using a team or an agent."""
component_type = "tool"
def __init__(
self,
task_runner: BaseGroupChat | BaseChatAgent,
name: str,
description: str,
return_value_as_last_message: bool,
) -> None:
self._task_runner = task_runner
self._return_value_as_last_message = return_value_as_last_message
super().__init__(
args_type=TaskRunnerToolArgs,
return_type=TaskResult,
name=name,
description=description,
strict=True,
)
async def run(self, args: TaskRunnerToolArgs, cancellation_token: CancellationToken) -> TaskResult:
"""Run the task and return the result."""
return await self._task_runner.run(task=args.task, cancellation_token=cancellation_token)
async def run_stream(
self, args: TaskRunnerToolArgs, cancellation_token: CancellationToken
) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
"""Run the task and yield events or messages as they are produced, the final :class:`TaskResult`
will be yielded at the end."""
async for event in self._task_runner.run_stream(task=args.task, cancellation_token=cancellation_token):
yield event
def return_value_as_string(self, value: TaskResult) -> str:
"""Convert the task result to a string."""
if self._return_value_as_last_message:
if value.messages and isinstance(value.messages[-1], BaseChatMessage):
return value.messages[-1].to_model_text()
raise ValueError("The last message is not a BaseChatMessage.")
parts: List[str] = []
for message in value.messages:
if isinstance(message, BaseChatMessage):
if message.source == "user":
continue
parts.append(f"{message.source}: {message.to_model_text()}")
return "\n\n".join(parts)
async def save_state_json(self) -> Mapping[str, Any]:
return await self._task_runner.save_state()
async def load_state_json(self, state: Mapping[str, Any]) -> None:
await self._task_runner.load_state(state)

133
agent_dhal/agentdhal_agentchat/tools/_team.py Normal file
View File

@@ -0,0 +1,133 @@
from agentdhal_core import Component, ComponentModel
from pydantic import BaseModel
from typing_extensions import Self
from agentdhal_agentchat.teams import BaseGroupChat
from ._task_runner_tool import TaskRunnerTool
class TeamToolConfig(BaseModel):
"""Configuration for the TeamTool."""
name: str
"""The name of the tool."""
description: str
"""The name and description of the tool."""
team: ComponentModel
"""The team to be used for running the task."""
return_value_as_last_message: bool = False
"""Whether to return the value as the last message of the task result."""
class TeamTool(TaskRunnerTool, Component[TeamToolConfig]):
"""Tool that can be used to run a task.
The tool returns the result of the task execution as a :class:`~agentdhal_agentchat.base.TaskResult` object.
.. important::
When using TeamTool, you **must** disable parallel tool calls in the model client configuration
to avoid concurrency issues. Teams cannot run concurrently as they maintain internal state
that would conflict with parallel execution. For example, set ``parallel_tool_calls=False``
for :class:`~agentdhal_extensions.models.openai.OpenAIChatCompletionClient` and
:class:`~agentdhal_extensions.models.openai.AzureOpenAIChatCompletionClient`.
Args:
team (BaseGroupChat): The team to be used for running the task.
name (str): The name of the tool.
description (str): The description of the tool.
return_value_as_last_message (bool): Whether to use the last message content of the task result
as the return value of the tool in :meth:`~agentdhal_agentchat.tools.TaskRunnerTool.return_value_as_string`.
If set to True, the last message content will be returned as a string.
If set to False, the tool will return all messages in the task result as a string concatenated together,
with each message prefixed by its source (e.g., "writer: ...", "assistant: ...").
Example:
.. code-block:: python
from agentdhal_agentchat.agents import AssistantAgent
from agentdhal_agentchat.conditions import SourceMatchTermination
from agentdhal_agentchat.teams import RoundRobinGroupChat
from agentdhal_agentchat.tools import TeamTool
from agentdhal_agentchat.ui import Console
from agentdhal_extensions.models.openai import OpenAIChatCompletionClient
async def main() -> None:
# Disable parallel tool calls when using TeamTool
model_client = OpenAIChatCompletionClient(model="gpt-4.1")
writer = AssistantAgent(name="writer", model_client=model_client, system_message="You are a helpful assistant.")
reviewer = AssistantAgent(
name="reviewer", model_client=model_client, system_message="You are a critical reviewer."
)
summarizer = AssistantAgent(
name="summarizer",
model_client=model_client,
system_message="You combine the review and produce a revised response.",
)
team = RoundRobinGroupChat(
[writer, reviewer, summarizer], termination_condition=SourceMatchTermination(sources=["summarizer"])
)
# Create a TeamTool that uses the team to run tasks, returning the last message as the result.
tool = TeamTool(
team=team,
name="writing_team",
description="A tool for writing tasks.",
return_value_as_last_message=True,
)
# Create model client with parallel tool calls disabled for the main agent
main_model_client = OpenAIChatCompletionClient(model="gpt-4.1", parallel_tool_calls=False)
main_agent = AssistantAgent(
name="main_agent",
model_client=main_model_client,
system_message="You are a helpful assistant that can use the writing tool.",
tools=[tool],
)
# For handling each events manually.
# async for message in main_agent.run_stream(
# task="Write a short story about a robot learning to love.",
# ):
# print(message)
# Use Console to display the messages in a more readable format.
await Console(
main_agent.run_stream(
task="Write a short story about a robot learning to love.",
)
)
if __name__ == "__main__":
import asyncio
asyncio.run(main())
"""
component_config_schema = TeamToolConfig
component_provider_override = "agentdhal_agentchat.tools.TeamTool"
def __init__(
self, team: BaseGroupChat, name: str, description: str, return_value_as_last_message: bool = False
) -> None:
self._team = team
super().__init__(team, name, description, return_value_as_last_message=return_value_as_last_message)
def _to_config(self) -> TeamToolConfig:
return TeamToolConfig(
name=self._name,
description=self._description,
team=self._team.dump_component(),
return_value_as_last_message=self._return_value_as_last_message,
)
@classmethod
def _from_config(cls, config: TeamToolConfig) -> Self:
return cls(
BaseGroupChat.load_component(config.team),
config.name,
config.description,
config.return_value_as_last_message,
)

7
agent_dhal/agentdhal_agentchat/ui/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
"""
This module implements utility classes for formatting/printing agent messages.
"""
from ._console import Console, UserInputManager
__all__ = ["Console", "UserInputManager"]

204
agent_dhal/agentdhal_agentchat/ui/_console.py Normal file
View File

@@ -0,0 +1,204 @@
import asyncio
import os
import sys
import time
from inspect import iscoroutinefunction
from typing import AsyncGenerator, Awaitable, Callable, Dict, List, Optional, TypeVar, Union, cast
from agentdhal_core import CancellationToken
from agentdhal_core.models import RequestUsage
from agentdhal_agentchat.agents import UserProxyAgent
from agentdhal_agentchat.base import Response, TaskResult
from agentdhal_agentchat.messages import (
BaseAgentEvent,
BaseChatMessage,
ModelClientStreamingChunkEvent,
MultiModalMessage,
UserInputRequestedEvent,
)
def _is_running_in_iterm() -> bool:
return os.getenv("TERM_PROGRAM") == "iTerm.app"
def _is_output_a_tty() -> bool:
return sys.stdout.isatty()
SyncInputFunc = Callable[[str], str]
AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
T = TypeVar("T", bound=TaskResult | Response)
class UserInputManager:
def __init__(self, callback: InputFuncType):
self.input_events: Dict[str, asyncio.Event] = {}
self.callback = callback
def get_wrapped_callback(self) -> AsyncInputFunc:
async def user_input_func_wrapper(prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
# Lookup the event for the prompt, if it exists wait for it.
# If it doesn't exist, create it and store it.
# Get request ID:
request_id = UserProxyAgent.InputRequestContext.request_id()
if request_id in self.input_events:
event = self.input_events[request_id]
else:
event = asyncio.Event()
self.input_events[request_id] = event
await event.wait()
del self.input_events[request_id]
if iscoroutinefunction(self.callback):
# Cast to AsyncInputFunc for proper typing
async_func = cast(AsyncInputFunc, self.callback)
return await async_func(prompt, cancellation_token)
else:
# Cast to SyncInputFunc for proper typing
sync_func = cast(SyncInputFunc, self.callback)
loop = asyncio.get_event_loop()
return await loop.run_in_executor(None, sync_func, prompt)
return user_input_func_wrapper
def notify_event_received(self, request_id: str) -> None:
if request_id in self.input_events:
self.input_events[request_id].set()
else:
event = asyncio.Event()
self.input_events[request_id] = event
def aprint(output: str, end: str = "\n", flush: bool = False) -> Awaitable[None]:
return asyncio.to_thread(print, output, end=end, flush=flush)
async def Console(
stream: AsyncGenerator[BaseAgentEvent | BaseChatMessage | T, None],
*,
no_inline_images: bool = False,
output_stats: bool = False,
user_input_manager: UserInputManager | None = None,
) -> T:
"""
Consumes the message stream from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream`
or :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream` and renders the messages to the console.
Returns the last processed TaskResult or Response.
.. note::
`output_stats` is experimental and the stats may not be accurate.
It will be improved in future releases.
Args:
stream (AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None] | AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]): Message stream to render.
This can be from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream` or :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream`.
no_inline_images (bool, optional): If terminal is iTerm2 will render images inline. Use this to disable this behavior. Defaults to False.
output_stats (bool, optional): (Experimental) If True, will output a summary of the messages and inline token usage info. Defaults to False.
Returns:
last_processed: A :class:`~agentdhal_agentchat.base.TaskResult` if the stream is from :meth:`~agentdhal_agentchat.base.TaskRunner.run_stream`
or a :class:`~agentdhal_agentchat.base.Response` if the stream is from :meth:`~agentdhal_agentchat.base.ChatAgent.on_messages_stream`.
"""
render_image_iterm = _is_running_in_iterm() and _is_output_a_tty() and not no_inline_images
start_time = time.time()
total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)
last_processed: Optional[T] = None
streaming_chunks: List[str] = []
async for message in stream:
if isinstance(message, TaskResult):
duration = time.time() - start_time
if output_stats:
output = (
f"{'-' * 10} Summary {'-' * 10}\n"
f"Number of messages: {len(message.messages)}\n"
f"Finish reason: {message.stop_reason}\n"
f"Total prompt tokens: {total_usage.prompt_tokens}\n"
f"Total completion tokens: {total_usage.completion_tokens}\n"
f"Duration: {duration:.2f} seconds\n"
)
await aprint(output, end="", flush=True)
# mypy ignore
last_processed = message # type: ignore
elif isinstance(message, Response):
duration = time.time() - start_time
# Print final response.
if isinstance(message.chat_message, MultiModalMessage):
final_content = message.chat_message.to_text(iterm=render_image_iterm)
else:
final_content = message.chat_message.to_text()
output = f"{'-' * 10} {message.chat_message.source} {'-' * 10}\n{final_content}\n"
if message.chat_message.models_usage:
if output_stats:
output += f"[Prompt tokens: {message.chat_message.models_usage.prompt_tokens}, Completion tokens: {message.chat_message.models_usage.completion_tokens}]\n"
total_usage.completion_tokens += message.chat_message.models_usage.completion_tokens
total_usage.prompt_tokens += message.chat_message.models_usage.prompt_tokens
await aprint(output, end="", flush=True)
# Print summary.
if output_stats:
if message.inner_messages is not None:
num_inner_messages = len(message.inner_messages)
else:
num_inner_messages = 0
output = (
f"{'-' * 10} Summary {'-' * 10}\n"
f"Number of inner messages: {num_inner_messages}\n"
f"Total prompt tokens: {total_usage.prompt_tokens}\n"
f"Total completion tokens: {total_usage.completion_tokens}\n"
f"Duration: {duration:.2f} seconds\n"
)
await aprint(output, end="", flush=True)
# mypy ignore
last_processed = message # type: ignore
# We don't want to print UserInputRequestedEvent messages, we just use them to signal the user input event.
elif isinstance(message, UserInputRequestedEvent):
if user_input_manager is not None:
user_input_manager.notify_event_received(message.request_id)
else:
# Cast required for mypy to be happy
message = cast(BaseAgentEvent | BaseChatMessage, message) # type: ignore
if not streaming_chunks:
# Print message sender.
await aprint(
f"{'-' * 10} {message.__class__.__name__} ({message.source}) {'-' * 10}", end="\n", flush=True
)
if isinstance(message, ModelClientStreamingChunkEvent):
await aprint(message.to_text(), end="", flush=True)
streaming_chunks.append(message.content)
else:
if streaming_chunks:
streaming_chunks.clear()
# Chunked messages are already printed, so we just print a newline.
await aprint("", end="\n", flush=True)
elif isinstance(message, MultiModalMessage):
await aprint(message.to_text(iterm=render_image_iterm), end="\n", flush=True)
else:
await aprint(message.to_text(), end="\n", flush=True)
if message.models_usage:
if output_stats:
await aprint(
f"[Prompt tokens: {message.models_usage.prompt_tokens}, Completion tokens: {message.models_usage.completion_tokens}]",
end="\n",
flush=True,
)
total_usage.completion_tokens += message.models_usage.completion_tokens
total_usage.prompt_tokens += message.models_usage.prompt_tokens
if last_processed is None:
raise ValueError("No TaskResult or Response was processed.")
return last_processed

7
agent_dhal/agentdhal_agentchat/utils/__init__.py Normal file
View File

@@ -0,0 +1,7 @@
"""
This module implements various utilities common to AgentChat agents and teams.
"""
from ._utils import content_to_str, remove_images
__all__ = ["content_to_str", "remove_images"]

44
agent_dhal/agentdhal_agentchat/utils/_utils.py Normal file
View File

@@ -0,0 +1,44 @@
from typing import List, Union
from agentdhal_core import FunctionCall, Image
from agentdhal_core.models import FunctionExecutionResult, LLMMessage, UserMessage
from pydantic import BaseModel
# Type aliases for convenience
_StructuredContent = BaseModel
_UserContent = Union[str, List[Union[str, Image]]]
_AssistantContent = Union[str, List[FunctionCall]]
_FunctionExecutionContent = List[FunctionExecutionResult]
_SystemContent = str
def content_to_str(
content: _UserContent | _AssistantContent | _FunctionExecutionContent | _SystemContent | _StructuredContent,
) -> str:
"""Convert the content of an LLMMessage to a string."""
if isinstance(content, str):
return content
elif isinstance(content, BaseModel):
return content.model_dump_json()
else:
result: List[str] = []
for c in content:
if isinstance(c, str):
result.append(c)
elif isinstance(c, Image):
result.append("<image>")
else:
result.append(str(c))
return "\n".join(result)
def remove_images(messages: List[LLMMessage]) -> List[LLMMessage]:
"""Remove images from a list of LLMMessages"""
str_messages: List[LLMMessage] = []
for message in messages:
if isinstance(message, UserMessage) and isinstance(message.content, list):
str_messages.append(UserMessage(content=content_to_str(message.content), source=message.source))
else:
str_messages.append(message)
return str_messages

142
agent_dhal/agentdhal_core/__init__.py Normal file
View File

@@ -0,0 +1,142 @@
# AgentDhal Core Module - Self-contained version
__version__ = "1.0.0"
from ._agent import Agent
from ._agent_id import AgentId
from ._agent_instantiation import AgentInstantiationContext
from ._agent_metadata import AgentMetadata
from ._agent_proxy import AgentProxy
from ._agent_runtime import AgentRuntime
from ._agent_type import AgentType
from ._base_agent import BaseAgent
from ._cache_store import CacheStore, InMemoryStore
from ._cancellation_token import CancellationToken
from ._closure_agent import ClosureAgent, ClosureContext
from ._component_config import (
Component,
ComponentBase,
ComponentFromConfig,
ComponentLoader,
ComponentModel,
ComponentSchemaType,
ComponentToConfig,
ComponentType,
is_component_class,
is_component_instance,
)
from ._constants import (
EVENT_LOGGER_NAME as EVENT_LOGGER_NAME_ALIAS,
)
from ._constants import (
ROOT_LOGGER_NAME as ROOT_LOGGER_NAME_ALIAS,
)
from ._constants import (
TRACE_LOGGER_NAME as TRACE_LOGGER_NAME_ALIAS,
)
from ._default_subscription import DefaultSubscription, default_subscription, type_subscription
from ._default_topic import DefaultTopicId
from ._image import Image
from ._intervention import (
DefaultInterventionHandler,
DropMessage,
InterventionHandler,
)
from ._message_context import MessageContext
from ._message_handler_context import MessageHandlerContext
from ._routed_agent import RoutedAgent, event, message_handler, rpc
from ._serialization import (
JSON_DATA_CONTENT_TYPE as JSON_DATA_CONTENT_TYPE_ALIAS,
)
from ._serialization import (
PROTOBUF_DATA_CONTENT_TYPE as PROTOBUF_DATA_CONTENT_TYPE_ALIAS,
)
from ._serialization import (
MessageSerializer,
UnknownPayload,
try_get_known_serializers_for_type,
)
from ._single_threaded_agent_runtime import SingleThreadedAgentRuntime
from ._subscription import Subscription
from ._subscription_context import SubscriptionInstantiationContext
from ._telemetry import (
trace_create_agent_span,
trace_invoke_agent_span,
trace_tool_span,
)
from ._topic import TopicId
from ._type_prefix_subscription import TypePrefixSubscription
from ._type_subscription import TypeSubscription
from ._types import FunctionCall
EVENT_LOGGER_NAME = EVENT_LOGGER_NAME_ALIAS
"""The name of the logger used for structured events."""
ROOT_LOGGER_NAME = ROOT_LOGGER_NAME_ALIAS
"""The name of the root logger."""
TRACE_LOGGER_NAME = TRACE_LOGGER_NAME_ALIAS
"""Logger name used for developer intended trace logging. The content and format of this log should not be depended upon."""
JSON_DATA_CONTENT_TYPE = JSON_DATA_CONTENT_TYPE_ALIAS
"""The content type for JSON data."""
PROTOBUF_DATA_CONTENT_TYPE = PROTOBUF_DATA_CONTENT_TYPE_ALIAS
"""The content type for Protobuf data."""
__all__ = [
"Agent",
"AgentId",
"AgentProxy",
"AgentMetadata",
"AgentRuntime",
"BaseAgent",
"CacheStore",
"InMemoryStore",
"CancellationToken",
"AgentInstantiationContext",
"TopicId",
"Subscription",
"MessageContext",
"AgentType",
"SubscriptionInstantiationContext",
"MessageHandlerContext",
"MessageSerializer",
"try_get_known_serializers_for_type",
"UnknownPayload",
"Image",
"RoutedAgent",
"ClosureAgent",
"ClosureContext",
"message_handler",
"event",
"rpc",
"FunctionCall",
"TypeSubscription",
"DefaultSubscription",
"DefaultTopicId",
"default_subscription",
"type_subscription",
"TypePrefixSubscription",
"JSON_DATA_CONTENT_TYPE",
"PROTOBUF_DATA_CONTENT_TYPE",
"SingleThreadedAgentRuntime",
"ROOT_LOGGER_NAME",
"EVENT_LOGGER_NAME",
"TRACE_LOGGER_NAME",
"Component",
"ComponentBase",
"ComponentFromConfig",
"ComponentLoader",
"ComponentModel",
"ComponentSchemaType",
"ComponentToConfig",
"ComponentType",
"is_component_class",
"is_component_instance",
"DropMessage",
"InterventionHandler",
"DefaultInterventionHandler",
"trace_create_agent_span",
"trace_invoke_agent_span",
"trace_tool_span",
]

64
agent_dhal/agentdhal_core/_agent.py Normal file
View File

@@ -0,0 +1,64 @@
from typing import TYPE_CHECKING, Any, Mapping, Protocol, runtime_checkable
from ._agent_id import AgentId
from ._agent_metadata import AgentMetadata
from ._message_context import MessageContext
# Forward declaration for type checking only
if TYPE_CHECKING:
from ._agent_runtime import AgentRuntime
@runtime_checkable
class Agent(Protocol):
@property
def metadata(self) -> AgentMetadata:
"""Metadata of the agent."""
...
@property
def id(self) -> AgentId:
"""ID of the agent."""
...
async def bind_id_and_runtime(self, id: AgentId, runtime: "AgentRuntime") -> None:
"""Function used to bind an Agent instance to an `AgentRuntime`.
Args:
agent_id (AgentId): ID of the agent.
runtime (AgentRuntime): AgentRuntime instance to bind the agent to.
"""
...
async def on_message(self, message: Any, ctx: MessageContext) -> Any:
"""Message handler for the agent. This should only be called by the runtime, not by other agents.
Args:
message (Any): Received message. Type is one of the types in `subscriptions`.
ctx (MessageContext): Context of the message.
Returns:
Any: Response to the message. Can be None.
Raises:
asyncio.CancelledError: If the message was cancelled.
CantHandleException: If the agent cannot handle the message.
"""
...
async def save_state(self) -> Mapping[str, Any]:
"""Save the state of the agent. The result must be JSON serializable."""
...
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Load in the state of the agent obtained from `save_state`.
Args:
state (Mapping[str, Any]): State of the agent. Must be JSON serializable.
"""
...
async def close(self) -> None:
"""Called when the runtime is closed"""
...

68
agent_dhal/agentdhal_core/_agent_id.py Normal file
View File

@@ -0,0 +1,68 @@
import re
from typing_extensions import Self
from ._agent_type import AgentType
def is_valid_agent_type(value: str) -> bool:
return bool(re.match(r"^[\w\-\.]+\Z", value))
class AgentId:
"""
Agent ID uniquely identifies an agent instance within an agent runtime - including distributed runtime. It is the 'address' of the agent instance for receiving messages.
See here for more information: :ref:`agentid_and_lifecycle`
"""
def __init__(self, type: str | AgentType, key: str) -> None:
if isinstance(type, AgentType):
type = type.type
if not is_valid_agent_type(type):
raise ValueError(rf"Invalid agent type: {type}. Allowed values MUST match the regex: `^[\w\-\.]+\Z`")
self._type = type
self._key = key
def __hash__(self) -> int:
return hash((self._type, self._key))
def __str__(self) -> str:
return f"{self._type}/{self._key}"
def __repr__(self) -> str:
return f'AgentId(type="{self._type}", key="{self._key}")'
def __eq__(self, value: object) -> bool:
if not isinstance(value, AgentId):
return False
return self._type == value.type and self._key == value.key
@classmethod
def from_str(cls, agent_id: str) -> Self:
"""Convert a string of the format ``type/key`` into an AgentId"""
items = agent_id.split("/", maxsplit=1)
if len(items) != 2:
raise ValueError(f"Invalid agent id: {agent_id}")
type, key = items[0], items[1]
return cls(type, key)
@property
def type(self) -> str:
"""
An identifier that associates an agent with a specific factory function.
Strings may only be composed of alphanumeric letters (a-z) and (0-9), or underscores (_).
"""
return self._type
@property
def key(self) -> str:
"""
Agent instance identifier.
Strings may only be composed of alphanumeric letters (a-z) and (0-9), or underscores (_).
"""
return self._key

126
agent_dhal/agentdhal_core/_agent_instantiation.py Normal file
View File

@@ -0,0 +1,126 @@
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, ClassVar, Generator
from ._agent_id import AgentId
from ._agent_runtime import AgentRuntime
class AgentInstantiationContext:
"""A static class that provides context for agent instantiation.
This static class can be used to access the current runtime and agent ID
during agent instantiation -- inside the factory function or the agent's
class constructor.
Example:
Get the current runtime and agent ID inside the factory function and
the agent's constructor:
.. code-block:: python
import asyncio
from dataclasses import dataclass
from agentdhal_core import (
AgentId,
AgentInstantiationContext,
MessageContext,
RoutedAgent,
SingleThreadedAgentRuntime,
message_handler,
)
@dataclass
class TestMessage:
content: str
class TestAgent(RoutedAgent):
def __init__(self, description: str):
super().__init__(description)
# Get the current runtime -- we don't use it here, but it's available.
_ = AgentInstantiationContext.current_runtime()
# Get the current agent ID.
agent_id = AgentInstantiationContext.current_agent_id()
print(f"Current AgentID from constructor: {agent_id}")
@message_handler
async def handle_test_message(self, message: TestMessage, ctx: MessageContext) -> None:
print(f"Received message: {message.content}")
def test_agent_factory() -> TestAgent:
# Get the current runtime -- we don't use it here, but it's available.
_ = AgentInstantiationContext.current_runtime()
# Get the current agent ID.
agent_id = AgentInstantiationContext.current_agent_id()
print(f"Current AgentID from factory: {agent_id}")
return TestAgent(description="Test agent")
async def main() -> None:
# Create a SingleThreadedAgentRuntime instance.
runtime = SingleThreadedAgentRuntime()
# Start the runtime.
runtime.start()
# Register the agent type with a factory function.
await runtime.register_factory("test_agent", test_agent_factory)
# Send a message to the agent. The runtime will instantiate the agent and call the message handler.
await runtime.send_message(TestMessage(content="Hello, world!"), AgentId("test_agent", "default"))
# Stop the runtime.
await runtime.stop()
asyncio.run(main())
"""
def __init__(self) -> None:
raise RuntimeError(
"AgentInstantiationContext cannot be instantiated. It is a static class that provides context management for agent instantiation."
)
_AGENT_INSTANTIATION_CONTEXT_VAR: ClassVar[ContextVar[tuple[AgentRuntime, AgentId]]] = ContextVar(
"_AGENT_INSTANTIATION_CONTEXT_VAR"
)
@classmethod
@contextmanager
def populate_context(cls, ctx: tuple[AgentRuntime, AgentId]) -> Generator[None, Any, None]:
""":meta private:"""
token = AgentInstantiationContext._AGENT_INSTANTIATION_CONTEXT_VAR.set(ctx)
try:
yield
finally:
AgentInstantiationContext._AGENT_INSTANTIATION_CONTEXT_VAR.reset(token)
@classmethod
def current_runtime(cls) -> AgentRuntime:
try:
return cls._AGENT_INSTANTIATION_CONTEXT_VAR.get()[0]
except LookupError as e:
raise RuntimeError(
"AgentInstantiationContext.runtime() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so."
) from e
@classmethod
def current_agent_id(cls) -> AgentId:
try:
return cls._AGENT_INSTANTIATION_CONTEXT_VAR.get()[1]
except LookupError as e:
raise RuntimeError(
"AgentInstantiationContext.agent_id() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so."
) from e
@classmethod
def is_in_factory_call(cls) -> bool:
if cls._AGENT_INSTANTIATION_CONTEXT_VAR.get(None) is None:
return False
return True

7
agent_dhal/agentdhal_core/_agent_metadata.py Normal file
View File

@@ -0,0 +1,7 @@
from typing import TypedDict
class AgentMetadata(TypedDict):
type: str
key: str
description: str

56
agent_dhal/agentdhal_core/_agent_proxy.py Normal file
View File

@@ -0,0 +1,56 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Awaitable, Mapping
from ._agent_id import AgentId
from ._agent_metadata import AgentMetadata
from ._cancellation_token import CancellationToken
if TYPE_CHECKING:
from ._agent_runtime import AgentRuntime
class AgentProxy:
"""A helper class that allows you to use an :class:`~agentdhal_core.AgentId` in place of its associated :class:`~agentdhal_core.Agent`"""
def __init__(self, agent: AgentId, runtime: AgentRuntime):
self._agent = agent
self._runtime = runtime
@property
def id(self) -> AgentId:
"""Target agent for this proxy"""
return self._agent
@property
def metadata(self) -> Awaitable[AgentMetadata]:
"""Metadata of the agent."""
return self._runtime.agent_metadata(self._agent)
async def send_message(
self,
message: Any,
*,
sender: AgentId,
cancellation_token: CancellationToken | None = None,
message_id: str | None = None,
) -> Any:
return await self._runtime.send_message(
message,
recipient=self._agent,
sender=sender,
cancellation_token=cancellation_token,
message_id=message_id,
)
async def save_state(self) -> Mapping[str, Any]:
"""Save the state of the agent. The result must be JSON serializable."""
return await self._runtime.agent_save_state(self._agent)
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Load in the state of the agent obtained from `save_state`.
Args:
state (Mapping[str, Any]): State of the agent. Must be JSON serializable.
"""
await self._runtime.agent_load_state(self._agent, state)

295
agent_dhal/agentdhal_core/_agent_runtime.py Normal file
View File

@@ -0,0 +1,295 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Any, Awaitable, Callable, Mapping, Protocol, Type, TypeVar, overload, runtime_checkable
from ._agent import Agent
from ._agent_id import AgentId
from ._agent_metadata import AgentMetadata
from ._agent_type import AgentType
from ._cancellation_token import CancellationToken
from ._serialization import MessageSerializer
from ._subscription import Subscription
from ._topic import TopicId
# Undeliverable - error
T = TypeVar("T", bound=Agent)
@runtime_checkable
class AgentRuntime(Protocol):
async def send_message(
self,
message: Any,
recipient: AgentId,
*,
sender: AgentId | None = None,
cancellation_token: CancellationToken | None = None,
message_id: str | None = None,
) -> Any:
"""Send a message to an agent and get a response.
Args:
message (Any): The message to send.
recipient (AgentId): The agent to send the message to.
sender (AgentId | None, optional): Agent which sent the message. Should **only** be None if this was sent from no agent, such as directly to the runtime externally. Defaults to None.
cancellation_token (CancellationToken | None, optional): Token used to cancel an in progress . Defaults to None.
Raises:
CantHandleException: If the recipient cannot handle the message.
UndeliverableException: If the message cannot be delivered.
Other: Any other exception raised by the recipient.
Returns:
Any: The response from the agent.
"""
...
async def publish_message(
self,
message: Any,
topic_id: TopicId,
*,
sender: AgentId | None = None,
cancellation_token: CancellationToken | None = None,
message_id: str | None = None,
) -> None:
"""Publish a message to all agents in the given namespace, or if no namespace is provided, the namespace of the sender.
No responses are expected from publishing.
Args:
message (Any): The message to publish.
topic_id (TopicId): The topic to publish the message to.
sender (AgentId | None, optional): The agent which sent the message. Defaults to None.
cancellation_token (CancellationToken | None, optional): Token used to cancel an in progress. Defaults to None.
message_id (str | None, optional): The message id. If None, a new message id will be generated. Defaults to None. This message id must be unique. and is recommended to be a UUID.
Raises:
UndeliverableException: If the message cannot be delivered.
"""
...
async def register_factory(
self,
type: str | AgentType,
agent_factory: Callable[[], T | Awaitable[T]],
*,
expected_class: type[T] | None = None,
) -> AgentType:
"""Register an agent factory with the runtime associated with a specific type. The type must be unique. This API does not add any subscriptions.
.. note::
This is a low level API and usually the agent class's `register` method should be used instead, as this also handles subscriptions automatically.
Example:
.. code-block:: python
from dataclasses import dataclass
from agentdhal_core import AgentRuntime, MessageContext, RoutedAgent, event
from agentdhal_core.models import UserMessage
@dataclass
class MyMessage:
content: str
class MyAgent(RoutedAgent):
def __init__(self) -> None:
super().__init__("My core agent")
@event
async def handler(self, message: UserMessage, context: MessageContext) -> None:
print("Event received: ", message.content)
async def my_agent_factory():
return MyAgent()
async def main() -> None:
runtime: AgentRuntime = ... # type: ignore
await runtime.register_factory("my_agent", lambda: MyAgent())
import asyncio
asyncio.run(main())
Args:
type (str): The type of agent this factory creates. It is not the same as agent class name. The `type` parameter is used to differentiate between different factory functions rather than agent classes.
agent_factory (Callable[[], T]): The factory that creates the agent, where T is a concrete Agent type. Inside the factory, use `agentdhal_core.AgentInstantiationContext` to access variables like the current runtime and agent ID.
expected_class (type[T] | None, optional): The expected class of the agent, used for runtime validation of the factory. Defaults to None. If None, no validation is performed.
"""
...
async def register_agent_instance(
self,
agent_instance: Agent,
agent_id: AgentId,
) -> AgentId:
"""Register an agent instance with the runtime. The type may be reused, but each agent_id must be unique. All agent instances within a type must be of the same object type. This API does not add any subscriptions.
.. note::
This is a low level API and usually the agent class's `register_instance` method should be used instead, as this also handles subscriptions automatically.
Example:
.. code-block:: python
from dataclasses import dataclass
from agentdhal_core import AgentId, AgentRuntime, MessageContext, RoutedAgent, event
from agentdhal_core.models import UserMessage
@dataclass
class MyMessage:
content: str
class MyAgent(RoutedAgent):
def __init__(self) -> None:
super().__init__("My core agent")
@event
async def handler(self, message: UserMessage, context: MessageContext) -> None:
print("Event received: ", message.content)
async def main() -> None:
runtime: AgentRuntime = ... # type: ignore
agent = MyAgent()
await runtime.register_agent_instance(
agent_instance=agent, agent_id=AgentId(type="my_agent", key="default")
)
import asyncio
asyncio.run(main())
Args:
agent_instance (Agent): A concrete instance of the agent.
agent_id (AgentId): The agent's identifier. The agent's type is `agent_id.type`.
"""
...
# TODO: uncomment out the following type ignore when this is fixed in mypy: https://github.com/python/mypy/issues/3737
async def try_get_underlying_agent_instance(self, id: AgentId, type: Type[T] = Agent) -> T: # type: ignore[assignment]
"""Try to get the underlying agent instance by name and namespace. This is generally discouraged (hence the long name), but can be useful in some cases.
If the underlying agent is not accessible, this will raise an exception.
Args:
id (AgentId): The agent id.
type (Type[T], optional): The expected type of the agent. Defaults to Agent.
Returns:
T: The concrete agent instance.
Raises:
LookupError: If the agent is not found.
NotAccessibleError: If the agent is not accessible, for example if it is located remotely.
TypeError: If the agent is not of the expected type.
"""
...
@overload
async def get(self, id: AgentId, /, *, lazy: bool = ...) -> AgentId: ...
@overload
async def get(self, type: AgentType | str, /, key: str = ..., *, lazy: bool = ...) -> AgentId: ...
async def get(
self, id_or_type: AgentId | AgentType | str, /, key: str = "default", *, lazy: bool = True
) -> AgentId: ...
async def save_state(self) -> Mapping[str, Any]:
"""Save the state of the entire runtime, including all hosted agents. The only way to restore the state is to pass it to :meth:`load_state`.
The structure of the state is implementation defined and can be any JSON serializable object.
Returns:
Mapping[str, Any]: The saved state.
"""
...
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Load the state of the entire runtime, including all hosted agents. The state should be the same as the one returned by :meth:`save_state`.
Args:
state (Mapping[str, Any]): The saved state.
"""
...
async def agent_metadata(self, agent: AgentId) -> AgentMetadata:
"""Get the metadata for an agent.
Args:
agent (AgentId): The agent id.
Returns:
AgentMetadata: The agent metadata.
"""
...
async def agent_save_state(self, agent: AgentId) -> Mapping[str, Any]:
"""Save the state of a single agent.
The structure of the state is implementation defined and can be any JSON serializable object.
Args:
agent (AgentId): The agent id.
Returns:
Mapping[str, Any]: The saved state.
"""
...
async def agent_load_state(self, agent: AgentId, state: Mapping[str, Any]) -> None:
"""Load the state of a single agent.
Args:
agent (AgentId): The agent id.
state (Mapping[str, Any]): The saved state.
"""
...
async def add_subscription(self, subscription: Subscription) -> None:
"""Add a new subscription that the runtime should fulfill when processing published messages
Args:
subscription (Subscription): The subscription to add
"""
...
async def remove_subscription(self, id: str) -> None:
"""Remove a subscription from the runtime
Args:
id (str): id of the subscription to remove
Raises:
LookupError: If the subscription does not exist
"""
...
def add_message_serializer(self, serializer: MessageSerializer[Any] | Sequence[MessageSerializer[Any]]) -> None:
"""Add a new message serialization serializer to the runtime
Note: This will deduplicate serializers based on the type_name and data_content_type properties
Args:
serializer (MessageSerializer[Any] | Sequence[MessageSerializer[Any]]): The serializer/s to add
"""
...

7
agent_dhal/agentdhal_core/_agent_type.py Normal file
View File

@@ -0,0 +1,7 @@
from dataclasses import dataclass
@dataclass(eq=True, frozen=True)
class AgentType:
type: str
"""String representation of this agent type."""

254
agent_dhal/agentdhal_core/_base_agent.py Normal file
View File

@@ -0,0 +1,254 @@
from __future__ import annotations
import inspect
import warnings
from abc import ABC, abstractmethod
from collections.abc import Sequence
from typing import Any, Awaitable, Callable, ClassVar, List, Mapping, Tuple, Type, TypeVar, final
from typing_extensions import Self
from ._agent import Agent
from ._agent_id import AgentId
from ._agent_instantiation import AgentInstantiationContext
from ._agent_metadata import AgentMetadata
from ._agent_runtime import AgentRuntime
from ._agent_type import AgentType
from ._cancellation_token import CancellationToken
from ._message_context import MessageContext
from ._serialization import MessageSerializer, try_get_known_serializers_for_type
from ._subscription import Subscription, UnboundSubscription
from ._subscription_context import SubscriptionInstantiationContext
from ._topic import TopicId
from ._type_prefix_subscription import TypePrefixSubscription
from ._type_subscription import TypeSubscription
T = TypeVar("T", bound=Agent)
BaseAgentType = TypeVar("BaseAgentType", bound="BaseAgent")
# Decorator for adding an unbound subscription to an agent
def subscription_factory(subscription: UnboundSubscription) -> Callable[[Type[BaseAgentType]], Type[BaseAgentType]]:
""":meta private:"""
def decorator(cls: Type[BaseAgentType]) -> Type[BaseAgentType]:
cls.internal_unbound_subscriptions_list.append(subscription)
return cls
return decorator
def handles(
type: Type[Any], serializer: MessageSerializer[Any] | List[MessageSerializer[Any]] | None = None
) -> Callable[[Type[BaseAgentType]], Type[BaseAgentType]]:
def decorator(cls: Type[BaseAgentType]) -> Type[BaseAgentType]:
if serializer is None:
serializer_list = try_get_known_serializers_for_type(type)
else:
serializer_list = [serializer] if not isinstance(serializer, Sequence) else serializer
if len(serializer_list) == 0:
raise ValueError(f"No serializers found for type {type}. Please provide an explicit serializer.")
cls.internal_extra_handles_types.append((type, serializer_list))
return cls
return decorator
class BaseAgent(ABC, Agent):
internal_unbound_subscriptions_list: ClassVar[List[UnboundSubscription]] = []
""":meta private:"""
internal_extra_handles_types: ClassVar[List[Tuple[Type[Any], List[MessageSerializer[Any]]]]] = []
""":meta private:"""
def __init_subclass__(cls, **kwargs: Any) -> None:
super().__init_subclass__(**kwargs)
# Automatically set class_variable in each subclass so that they are not shared between subclasses
cls.internal_extra_handles_types = []
cls.internal_unbound_subscriptions_list = []
@classmethod
def _handles_types(cls) -> List[Tuple[Type[Any], List[MessageSerializer[Any]]]]:
return cls.internal_extra_handles_types
@classmethod
def _unbound_subscriptions(cls) -> List[UnboundSubscription]:
return cls.internal_unbound_subscriptions_list
@property
def metadata(self) -> AgentMetadata:
assert self._id is not None
return AgentMetadata(key=self._id.key, type=self._id.type, description=self._description)
def __init__(self, description: str) -> None:
if AgentInstantiationContext.is_in_factory_call():
self._runtime: AgentRuntime = AgentInstantiationContext.current_runtime()
self._id = AgentInstantiationContext.current_agent_id()
if not isinstance(description, str):
raise ValueError("Agent description must be a string")
self._description = description
async def bind_id_and_runtime(self, id: AgentId, runtime: AgentRuntime) -> None:
if hasattr(self, "_id"):
if self._id != id:
raise RuntimeError("Agent is already bound to a different ID")
if hasattr(self, "_runtime"):
if self._runtime != runtime:
raise RuntimeError("Agent is already bound to a different runtime")
self._id = id
self._runtime = runtime
@property
def type(self) -> str:
return self.id.type
@property
def id(self) -> AgentId:
return self._id
@property
def runtime(self) -> AgentRuntime:
return self._runtime
@final
async def on_message(self, message: Any, ctx: MessageContext) -> Any:
return await self.on_message_impl(message, ctx)
@abstractmethod
async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any: ...
async def send_message(
self,
message: Any,
recipient: AgentId,
*,
cancellation_token: CancellationToken | None = None,
message_id: str | None = None,
) -> Any:
"""See :py:meth:`agentdhal_core.AgentRuntime.send_message` for more information."""
if cancellation_token is None:
cancellation_token = CancellationToken()
return await self._runtime.send_message(
message,
sender=self.id,
recipient=recipient,
cancellation_token=cancellation_token,
message_id=message_id,
)
async def publish_message(
self,
message: Any,
topic_id: TopicId,
*,
cancellation_token: CancellationToken | None = None,
) -> None:
await self._runtime.publish_message(message, topic_id, sender=self.id, cancellation_token=cancellation_token)
async def save_state(self) -> Mapping[str, Any]:
warnings.warn("save_state not implemented", stacklevel=2)
return {}
async def load_state(self, state: Mapping[str, Any]) -> None:
warnings.warn("load_state not implemented", stacklevel=2)
pass
async def close(self) -> None:
pass
async def register_instance(
self,
runtime: AgentRuntime,
agent_id: AgentId,
*,
skip_class_subscriptions: bool = True,
skip_direct_message_subscription: bool = False,
) -> AgentId:
"""
This function is similar to `register` but is used for registering an instance of an agent. A subscription based on the agent ID is created and added to the runtime.
"""
agent_id = await runtime.register_agent_instance(agent_instance=self, agent_id=agent_id)
id_subscription = TypeSubscription(topic_type=agent_id.key, agent_type=agent_id.type)
await runtime.add_subscription(id_subscription)
if not skip_class_subscriptions:
with SubscriptionInstantiationContext.populate_context(AgentType(agent_id.type)):
subscriptions: List[Subscription] = []
for unbound_subscription in self._unbound_subscriptions():
subscriptions_list_result = unbound_subscription()
if inspect.isawaitable(subscriptions_list_result):
subscriptions_list = await subscriptions_list_result
else:
subscriptions_list = subscriptions_list_result
subscriptions.extend(subscriptions_list)
for subscription in subscriptions:
await runtime.add_subscription(subscription)
if not skip_direct_message_subscription:
# Additionally adds a special prefix subscription for this agent to receive direct messages
try:
await runtime.add_subscription(
TypePrefixSubscription(
# The prefix MUST include ":" to avoid collisions with other agents
topic_type_prefix=agent_id.type + ":",
agent_type=agent_id.type,
)
)
except ValueError:
# We don't care if the subscription already exists
pass
# TODO: deduplication
for _message_type, serializer in self._handles_types():
runtime.add_message_serializer(serializer)
return agent_id
@classmethod
async def register(
cls,
runtime: AgentRuntime,
type: str,
factory: Callable[[], Self | Awaitable[Self]],
*,
skip_class_subscriptions: bool = False,
skip_direct_message_subscription: bool = False,
) -> AgentType:
agent_type = AgentType(type)
agent_type = await runtime.register_factory(type=agent_type, agent_factory=factory, expected_class=cls)
if not skip_class_subscriptions:
with SubscriptionInstantiationContext.populate_context(agent_type):
subscriptions: List[Subscription] = []
for unbound_subscription in cls._unbound_subscriptions():
subscriptions_list_result = unbound_subscription()
if inspect.isawaitable(subscriptions_list_result):
subscriptions_list = await subscriptions_list_result
else:
subscriptions_list = subscriptions_list_result
subscriptions.extend(subscriptions_list)
for subscription in subscriptions:
await runtime.add_subscription(subscription)
if not skip_direct_message_subscription:
# Additionally adds a special prefix subscription for this agent to receive direct messages
await runtime.add_subscription(
TypePrefixSubscription(
# The prefix MUST include ":" to avoid collisions with other agents
topic_type_prefix=agent_type.type + ":",
agent_type=agent_type.type,
)
)
# TODO: deduplication
for _message_type, serializer in cls._handles_types():
runtime.add_message_serializer(serializer)
return agent_type

70
agent_dhal/agentdhal_core/_cache_store.py Normal file
View File

@@ -0,0 +1,70 @@
from abc import ABC, abstractmethod
from typing import Dict, Generic, Optional, TypeVar
from pydantic import BaseModel
from typing_extensions import Self
from ._component_config import Component, ComponentBase
T = TypeVar("T")
class CacheStore(ABC, Generic[T], ComponentBase[BaseModel]):
"""
This protocol defines the basic interface for store/cache operations.
Sub-classes should handle the lifecycle of underlying storage.
"""
component_type = "cache_store"
@abstractmethod
def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
"""
Retrieve an item from the store.
Args:
key: The key identifying the item in the store.
default (optional): The default value to return if the key is not found.
Defaults to None.
Returns:
The value associated with the key if found, else the default value.
"""
...
@abstractmethod
def set(self, key: str, value: T) -> None:
"""
Set an item in the store.
Args:
key: The key under which the item is to be stored.
value: The value to be stored in the store.
"""
...
class InMemoryStoreConfig(BaseModel):
pass
class InMemoryStore(CacheStore[T], Component[InMemoryStoreConfig]):
component_provider_override = "agentdhal_core.InMemoryStore"
component_config_schema = InMemoryStoreConfig
def __init__(self) -> None:
self.store: Dict[str, T] = {}
def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
return self.store.get(key, default)
def set(self, key: str, value: T) -> None:
self.store[key] = value
def _to_config(self) -> InMemoryStoreConfig:
return InMemoryStoreConfig()
@classmethod
def _from_config(cls, config: InMemoryStoreConfig) -> Self:
return cls()

46
agent_dhal/agentdhal_core/_cancellation_token.py Normal file
View File

@@ -0,0 +1,46 @@
import threading
from asyncio import Future
from typing import Any, Callable, List
class CancellationToken:
"""A token used to cancel pending async calls"""
def __init__(self) -> None:
self._cancelled: bool = False
self._lock: threading.Lock = threading.Lock()
self._callbacks: List[Callable[[], None]] = []
def cancel(self) -> None:
"""Cancel pending async calls linked to this cancellation token."""
with self._lock:
if not self._cancelled:
self._cancelled = True
for callback in self._callbacks:
callback()
def is_cancelled(self) -> bool:
"""Check if the CancellationToken has been used"""
with self._lock:
return self._cancelled
def add_callback(self, callback: Callable[[], None]) -> None:
"""Attach a callback that will be called when cancel is invoked"""
with self._lock:
if self._cancelled:
callback()
else:
self._callbacks.append(callback)
def link_future(self, future: Future[Any]) -> Future[Any]:
"""Link a pending async call to a token to allow its cancellation"""
with self._lock:
if self._cancelled:
future.cancel()
else:
def _cancel() -> None:
future.cancel()
self._callbacks.append(_cancel)
return future

241
agent_dhal/agentdhal_core/_closure_agent.py Normal file
View File

@@ -0,0 +1,241 @@
from __future__ import annotations
import inspect
import warnings
from typing import Any, Awaitable, Callable, List, Literal, Mapping, Protocol, Sequence, TypeVar, get_type_hints
from ._agent_id import AgentId
from ._agent_instantiation import AgentInstantiationContext
from ._agent_metadata import AgentMetadata
from ._agent_runtime import AgentRuntime
from ._agent_type import AgentType
from ._base_agent import BaseAgent
from ._cancellation_token import CancellationToken
from ._message_context import MessageContext
from ._serialization import try_get_known_serializers_for_type
from ._subscription import Subscription
from ._subscription_context import SubscriptionInstantiationContext
from ._topic import TopicId
from ._type_helpers import get_types
from .exceptions import CantHandleException
T = TypeVar("T")
ClosureAgentType = TypeVar("ClosureAgentType", bound="ClosureAgent")
def get_handled_types_from_closure(
closure: Callable[[ClosureAgent, T, MessageContext], Awaitable[Any]],
) -> Sequence[type]:
args = inspect.getfullargspec(closure)[0]
if len(args) != 3:
raise AssertionError("Closure must have 4 arguments")
message_arg_name = args[1]
type_hints = get_type_hints(closure)
if "return" not in type_hints:
raise AssertionError("return not found in function signature")
# Get the type of the message parameter
target_types = get_types(type_hints[message_arg_name])
if target_types is None:
raise AssertionError("Message type not found")
# print(type_hints)
return_types = get_types(type_hints["return"])
if return_types is None:
raise AssertionError("Return type not found")
return target_types
class ClosureContext(Protocol):
@property
def id(self) -> AgentId: ...
async def send_message(
self,
message: Any,
recipient: AgentId,
*,
cancellation_token: CancellationToken | None = None,
message_id: str | None = None,
) -> Any: ...
async def publish_message(
self,
message: Any,
topic_id: TopicId,
*,
cancellation_token: CancellationToken | None = None,
) -> None: ...
class ClosureAgent(BaseAgent, ClosureContext):
def __init__(
self,
description: str,
closure: Callable[[ClosureContext, T, MessageContext], Awaitable[Any]],
*,
unknown_type_policy: Literal["error", "warn", "ignore"] = "warn",
) -> None:
try:
runtime = AgentInstantiationContext.current_runtime()
id = AgentInstantiationContext.current_agent_id()
except Exception as e:
raise RuntimeError(
"ClosureAgent must be instantiated within the context of an AgentRuntime. It cannot be directly instantiated."
) from e
self._runtime: AgentRuntime = runtime
self._id: AgentId = id
self._description = description
handled_types = get_handled_types_from_closure(closure)
self._expected_types = handled_types
self._closure = closure
self._unknown_type_policy = unknown_type_policy
super().__init__(description)
@property
def metadata(self) -> AgentMetadata:
assert self._id is not None
return AgentMetadata(
key=self._id.key,
type=self._id.type,
description=self._description,
)
@property
def id(self) -> AgentId:
return self._id
@property
def runtime(self) -> AgentRuntime:
return self._runtime
async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any:
if type(message) not in self._expected_types:
if self._unknown_type_policy == "warn":
warnings.warn(
f"Message type {type(message)} not in target types {self._expected_types} of {self.id}. Set unknown_type_policy to 'error' to raise an exception, or 'ignore' to suppress this warning.",
stacklevel=1,
)
return None
elif self._unknown_type_policy == "error":
raise CantHandleException(
f"Message type {type(message)} not in target types {self._expected_types} of {self.id}. Set unknown_type_policy to 'warn' to suppress this exception, or 'ignore' to suppress this warning."
)
return await self._closure(self, message, ctx)
async def save_state(self) -> Mapping[str, Any]:
"""Closure agents do not have state. So this method always returns an empty dictionary."""
return {}
async def load_state(self, state: Mapping[str, Any]) -> None:
"""Closure agents do not have state. So this method does nothing."""
pass
@classmethod
async def register_closure(
cls,
runtime: AgentRuntime,
type: str,
closure: Callable[[ClosureContext, T, MessageContext], Awaitable[Any]],
*,
unknown_type_policy: Literal["error", "warn", "ignore"] = "warn",
skip_direct_message_subscription: bool = False,
description: str = "",
subscriptions: Callable[[], list[Subscription] | Awaitable[list[Subscription]]] | None = None,
) -> AgentType:
"""The closure agent allows you to define an agent using a closure, or function without needing to define a class. It allows values to be extracted out of the runtime.
The closure can define the type of message which is expected, or `Any` can be used to accept any type of message.
Example:
.. code-block:: python
import asyncio
from agentdhal_core import SingleThreadedAgentRuntime, MessageContext, ClosureAgent, ClosureContext
from dataclasses import dataclass
from agentdhal_core._default_subscription import DefaultSubscription
from agentdhal_core._default_topic import DefaultTopicId
@dataclass
class MyMessage:
content: str
async def main():
queue = asyncio.Queue[MyMessage]()
async def output_result(_ctx: ClosureContext, message: MyMessage, ctx: MessageContext) -> None:
await queue.put(message)
runtime = SingleThreadedAgentRuntime()
await ClosureAgent.register_closure(
runtime, "output_result", output_result, subscriptions=lambda: [DefaultSubscription()]
)
runtime.start()
await runtime.publish_message(MyMessage("Hello, world!"), DefaultTopicId())
await runtime.stop_when_idle()
result = await queue.get()
print(result)
asyncio.run(main())
Args:
runtime (AgentRuntime): Runtime to register the agent to
type (str): Agent type of registered agent
closure (Callable[[ClosureContext, T, MessageContext], Awaitable[Any]]): Closure to handle messages
unknown_type_policy (Literal["error", "warn", "ignore"], optional): What to do if a type is encountered that does not match the closure type. Defaults to "warn".
skip_direct_message_subscription (bool, optional): Do not add direct message subscription for this agent. Defaults to False.
description (str, optional): Description of what agent does. Defaults to "".
subscriptions (Callable[[], list[Subscription] | Awaitable[list[Subscription]]] | None, optional): List of subscriptions for this closure agent. Defaults to None.
Returns:
AgentType: Type of the agent that was registered
"""
def factory() -> ClosureAgent:
return ClosureAgent(description=description, closure=closure, unknown_type_policy=unknown_type_policy)
assert len(cls._unbound_subscriptions()) == 0, "Closure agents are expected to have no class subscriptions"
agent_type = await cls.register(
runtime=runtime,
type=type,
factory=factory, # type: ignore
# There should be no need to process class subscriptions, as the closure agent does not have any subscriptions.s
skip_class_subscriptions=True,
skip_direct_message_subscription=skip_direct_message_subscription,
)
subscriptions_list: List[Subscription] = []
if subscriptions is not None:
with SubscriptionInstantiationContext.populate_context(agent_type):
subscriptions_list_result = subscriptions()
if inspect.isawaitable(subscriptions_list_result):
subscriptions_list.extend(await subscriptions_list_result)
else:
# just ignore mypy here
subscriptions_list.extend(subscriptions_list_result) # type: ignore
for subscription in subscriptions_list:
await runtime.add_subscription(subscription)
handled_types = get_handled_types_from_closure(closure)
for message_type in handled_types:
# TODO: support custom serializers
serializer = try_get_known_serializers_for_type(message_type)
runtime.add_message_serializer(serializer)
return agent_type

361
agent_dhal/agentdhal_core/_component_config.py Normal file
View File

@@ -0,0 +1,361 @@
from __future__ import annotations
import importlib
import warnings
from typing import Any, ClassVar, Dict, Generic, Literal, Type, TypeGuard, cast, overload
from pydantic import BaseModel
from typing_extensions import Self, TypeVar
ComponentType = Literal["model", "agent", "tool", "termination", "token_provider", "workbench"] | str
ConfigT = TypeVar("ConfigT", bound=BaseModel)
FromConfigT = TypeVar("FromConfigT", bound=BaseModel, contravariant=True)
ToConfigT = TypeVar("ToConfigT", bound=BaseModel, covariant=True)
T = TypeVar("T", bound=BaseModel, covariant=True)
class ComponentModel(BaseModel):
"""Model class for a component. Contains all information required to instantiate a component."""
provider: str
"""Describes how the component can be instantiated."""
component_type: ComponentType | None = None
"""Logical type of the component. If missing, the component assumes the default type of the provider."""
version: int | None = None
"""Version of the component specification. If missing, the component assumes whatever is the current version of the library used to load it. This is obviously dangerous and should be used for user authored ephmeral config. For all other configs version should be specified."""
component_version: int | None = None
"""Version of the component. If missing, the component assumes the default version of the provider."""
description: str | None = None
"""Description of the component."""
label: str | None = None
"""Human readable label for the component. If missing the component assumes the class name of the provider."""
config: dict[str, Any]
"""The schema validated config field is passed to a given class's implmentation of :py:meth:`agentdhal_core.ComponentConfigImpl._from_config` to create a new instance of the component class."""
def _type_to_provider_str(t: type) -> str:
return f"{t.__module__}.{t.__qualname__}"
WELL_KNOWN_PROVIDERS = {
"azure_openai_chat_completion_client": "agentdhal_extensions.models.openai.AzureOpenAIChatCompletionClient",
"AzureOpenAIChatCompletionClient": "agentdhal_extensions.models.openai.AzureOpenAIChatCompletionClient",
"openai_chat_completion_client": "agentdhal_extensions.models.openai.OpenAIChatCompletionClient",
"OpenAIChatCompletionClient": "agentdhal_extensions.models.openai.OpenAIChatCompletionClient",
}
class ComponentFromConfig(Generic[FromConfigT]):
@classmethod
def _from_config(cls, config: FromConfigT) -> Self:
"""Create a new instance of the component from a configuration object.
Args:
config (T): The configuration object.
Returns:
Self: The new instance of the component.
:meta public:
"""
raise NotImplementedError("This component does not support dumping to config")
@classmethod
def _from_config_past_version(cls, config: Dict[str, Any], version: int) -> Self:
"""Create a new instance of the component from a previous version of the configuration object.
This is only called when the version of the configuration object is less than the current version, since in this case the schema is not known.
Args:
config (Dict[str, Any]): The configuration object.
version (int): The version of the configuration object.
Returns:
Self: The new instance of the component.
:meta public:
"""
raise NotImplementedError("This component does not support loading from past versions")
class ComponentToConfig(Generic[ToConfigT]):
"""The two methods a class must implement to be a component.
Args:
Protocol (ConfigT): Type which derives from :py:class:`pydantic.BaseModel`.
"""
component_type: ClassVar[ComponentType]
"""The logical type of the component."""
component_version: ClassVar[int] = 1
"""The version of the component, if schema incompatibilities are introduced this should be updated."""
component_provider_override: ClassVar[str | None] = None
"""Override the provider string for the component. This should be used to prevent internal module names being a part of the module name."""
component_description: ClassVar[str | None] = None
"""A description of the component. If not provided, the docstring of the class will be used."""
component_label: ClassVar[str | None] = None
"""A human readable label for the component. If not provided, the component class name will be used."""
def _to_config(self) -> ToConfigT:
"""Dump the configuration that would be requite to create a new instance of a component matching the configuration of this instance.
Returns:
T: The configuration of the component.
:meta public:
"""
raise NotImplementedError("This component does not support dumping to config")
def dump_component(self) -> ComponentModel:
"""Dump the component to a model that can be loaded back in.
Raises:
TypeError: If the component is a local class.
Returns:
ComponentModel: The model representing the component.
"""
if self.component_provider_override is not None:
provider = self.component_provider_override
else:
provider = _type_to_provider_str(self.__class__)
# Warn if internal module name is used,
if "._" in provider:
warnings.warn(
"Internal module name used in provider string. This is not recommended and may cause issues in the future. Silence this warning by setting component_provider_override to this value.",
stacklevel=2,
)
if "<locals>" in provider:
raise TypeError("Cannot dump component with local class")
if not hasattr(self, "component_type"):
raise AttributeError("component_type not defined")
description = self.component_description
if description is None and self.__class__.__doc__:
# use docstring as description
docstring = self.__class__.__doc__.strip()
for marker in ["\n\nArgs:", "\n\nParameters:", "\n\nAttributes:", "\n\n"]:
docstring = docstring.split(marker)[0]
description = docstring.strip()
obj_config = self._to_config().model_dump(exclude_none=True)
model = ComponentModel(
provider=provider,
component_type=self.component_type,
version=self.component_version,
component_version=self.component_version,
description=description,
label=self.component_label or self.__class__.__name__,
config=obj_config,
)
return model
ExpectedType = TypeVar("ExpectedType")
class ComponentLoader:
@overload
@classmethod
def load_component(cls, model: ComponentModel | Dict[str, Any], expected: None = None) -> Self: ...
@overload
@classmethod
def load_component(cls, model: ComponentModel | Dict[str, Any], expected: Type[ExpectedType]) -> ExpectedType: ...
@classmethod
def load_component(
cls, model: ComponentModel | Dict[str, Any], expected: Type[ExpectedType] | None = None
) -> Self | ExpectedType:
"""Load a component from a model. Intended to be used with the return type of :py:meth:`agentdhal_core.ComponentConfig.dump_component`.
Example:
.. code-block:: python
from agentdhal_core import ComponentModel
from agentdhal_core.models import ChatCompletionClient
component: ComponentModel = ... # type: ignore
model_client = ChatCompletionClient.load_component(component)
Args:
model (ComponentModel): The model to load the component from.
Returns:
Self: The loaded component.
Args:
model (ComponentModel): _description_
expected (Type[ExpectedType] | None, optional): Explicit type only if used directly on ComponentLoader. Defaults to None.
Raises:
ValueError: If the provider string is invalid.
TypeError: Provider is not a subclass of ComponentConfigImpl, or the expected type does not match.
Returns:
Self | ExpectedType: The loaded component.
"""
# Use global and add further type checks
if isinstance(model, dict):
loaded_model = ComponentModel(**model)
else:
loaded_model = model
# First, do a look up in well known providers
if loaded_model.provider in WELL_KNOWN_PROVIDERS:
loaded_model.provider = WELL_KNOWN_PROVIDERS[loaded_model.provider]
output = loaded_model.provider.rsplit(".", maxsplit=1)
if len(output) != 2:
raise ValueError("Invalid")
module_path, class_name = output
module = importlib.import_module(module_path)
component_class = module.__getattribute__(class_name)
if not is_component_class(component_class):
raise TypeError("Invalid component class")
# We need to check the schema is valid
if not hasattr(component_class, "component_config_schema"):
raise AttributeError("component_config_schema not defined")
if not hasattr(component_class, "component_type"):
raise AttributeError("component_type not defined")
loaded_config_version = loaded_model.component_version or component_class.component_version
if loaded_config_version < component_class.component_version:
try:
instance = component_class._from_config_past_version(loaded_model.config, loaded_config_version) # type: ignore
except NotImplementedError as e:
raise NotImplementedError(
f"Tried to load component {component_class} which is on version {component_class.component_version} with a config on version {loaded_config_version} but _from_config_past_version is not implemented"
) from e
else:
schema = component_class.component_config_schema # type: ignore
validated_config = schema.model_validate(loaded_model.config)
# We're allowed to use the private method here
instance = component_class._from_config(validated_config) # type: ignore
if expected is None and not isinstance(instance, cls):
raise TypeError("Expected type does not match")
elif expected is None:
return cast(Self, instance)
elif not isinstance(instance, expected):
raise TypeError("Expected type does not match")
else:
return cast(ExpectedType, instance)
class ComponentSchemaType(Generic[ConfigT]):
# Ideally would be ClassVar[Type[ConfigT]], but this is disallowed https://github.com/python/typing/discussions/1424 (despite being valid in this context)
component_config_schema: Type[ConfigT]
"""The Pydantic model class which represents the configuration of the component."""
required_class_vars = ["component_config_schema", "component_type"]
def __init_subclass__(cls, **kwargs: Any):
super().__init_subclass__(**kwargs)
if cls.__name__ != "Component" and not cls.__name__ == "_ConcreteComponent":
# TODO: validate provider is loadable
for var in cls.required_class_vars:
if not hasattr(cls, var):
warnings.warn(
f"Class variable '{var}' must be defined in {cls.__name__} to be a valid component",
stacklevel=2,
)
class ComponentBase(ComponentToConfig[ConfigT], ComponentLoader, Generic[ConfigT]): ...
class Component(
ComponentFromConfig[ConfigT],
ComponentSchemaType[ConfigT],
Generic[ConfigT],
):
"""To create a component class, inherit from this class for the concrete class and ComponentBase on the interface. Then implement two class variables:
- :py:attr:`component_config_schema` - A Pydantic model class which represents the configuration of the component. This is also the type parameter of Component.
- :py:attr:`component_type` - What is the logical type of the component.
Example:
.. code-block:: python
from __future__ import annotations
from pydantic import BaseModel
from agentdhal_core import Component
class Config(BaseModel):
value: str
class MyComponent(Component[Config]):
component_type = "custom"
component_config_schema = Config
def __init__(self, value: str):
self.value = value
def _to_config(self) -> Config:
return Config(value=self.value)
@classmethod
def _from_config(cls, config: Config) -> MyComponent:
return cls(value=config.value)
"""
def __init_subclass__(cls, **kwargs: Any):
super().__init_subclass__(**kwargs)
if not is_component_class(cls):
warnings.warn(
f"Component class '{cls.__name__}' must subclass the following: ComponentFromConfig, ComponentToConfig, ComponentSchemaType, ComponentLoader, individually or with ComponentBase and Component. Look at the component config documentation or how OpenAIChatCompletionClient does it.",
stacklevel=2,
)
# Should never be used directly, only for type checking
class _ConcreteComponent(
ComponentFromConfig[ConfigT],
ComponentSchemaType[ConfigT],
ComponentToConfig[ConfigT],
ComponentLoader,
Generic[ConfigT],
): ...
def is_component_instance(cls: Any) -> TypeGuard[_ConcreteComponent[BaseModel]]:
return (
isinstance(cls, ComponentFromConfig)
and isinstance(cls, ComponentToConfig)
and isinstance(cls, ComponentSchemaType)
and isinstance(cls, ComponentLoader)
)
def is_component_class(cls: type) -> TypeGuard[Type[_ConcreteComponent[BaseModel]]]:
return (
issubclass(cls, ComponentFromConfig)
and issubclass(cls, ComponentToConfig)
and issubclass(cls, ComponentSchemaType)
and issubclass(cls, ComponentLoader)
)

9
agent_dhal/agentdhal_core/_constants.py Normal file
View File

@@ -0,0 +1,9 @@
ROOT_LOGGER_NAME = "agentdhal_core"
"""str: Logger name used for root logger"""
EVENT_LOGGER_NAME = "agentdhal_core.events"
"""str: Logger name used for structured event logging"""
TRACE_LOGGER_NAME = "agentdhal_core.trace"
"""str: Logger name used for developer intended trace logging. The content and format of this log should not be depended upon."""

53
agent_dhal/agentdhal_core/_default_subscription.py Normal file
View File

@@ -0,0 +1,53 @@
from typing import Callable, Type, TypeVar, overload
from ._agent_type import AgentType
from ._base_agent import BaseAgent, subscription_factory
from ._subscription_context import SubscriptionInstantiationContext
from ._type_subscription import TypeSubscription
from .exceptions import CantHandleException
class DefaultSubscription(TypeSubscription):
"""The default subscription is designed to be a sensible default for applications that only need global scope for agents.
This topic by default uses the "default" topic type and attempts to detect the agent type to use based on the instantiation context.
Args:
topic_type (str, optional): The topic type to subscribe to. Defaults to "default".
agent_type (str, optional): The agent type to use for the subscription. Defaults to None, in which case it will attempt to detect the agent type based on the instantiation context.
"""
def __init__(self, topic_type: str = "default", agent_type: str | AgentType | None = None):
if agent_type is None:
try:
agent_type = SubscriptionInstantiationContext.agent_type().type
except RuntimeError as e:
raise CantHandleException(
"If agent_type is not specified DefaultSubscription must be created within the subscription callback in AgentRuntime.register"
) from e
super().__init__(topic_type, agent_type)
BaseAgentType = TypeVar("BaseAgentType", bound="BaseAgent")
@overload
def default_subscription() -> Callable[[Type[BaseAgentType]], Type[BaseAgentType]]: ...
@overload
def default_subscription(cls: Type[BaseAgentType]) -> Type[BaseAgentType]: ...
def default_subscription(
cls: Type[BaseAgentType] | None = None,
) -> Callable[[Type[BaseAgentType]], Type[BaseAgentType]] | Type[BaseAgentType]:
if cls is None:
return subscription_factory(lambda: [DefaultSubscription()])
else:
return subscription_factory(lambda: [DefaultSubscription()])(cls)
def type_subscription(topic_type: str) -> Callable[[Type[BaseAgentType]], Type[BaseAgentType]]:
return subscription_factory(lambda: [DefaultSubscription(topic_type=topic_type)])

23
agent_dhal/agentdhal_core/_default_topic.py Normal file
View File

@@ -0,0 +1,23 @@
from ._message_handler_context import MessageHandlerContext
from ._topic import TopicId
class DefaultTopicId(TopicId):
"""DefaultTopicId provides a sensible default for the topic_id and source fields of a TopicId.
If created in the context of a message handler, the source will be set to the agent_id of the message handler, otherwise it will be set to "default".
Args:
type (str, optional): Topic type to publish message to. Defaults to "default".
source (str | None, optional): Topic source to publish message to. If None, the source will be set to the agent_id of the message handler if in the context of a message handler, otherwise it will be set to "default". Defaults to None.
"""
def __init__(self, type: str = "default", source: str | None = None) -> None:
if source is None:
try:
source = MessageHandlerContext.agent_id().key
# If we aren't in the context of a message handler, we use the default source
except RuntimeError:
source = "default"
super().__init__(type, source)

324
agent_dhal/agentdhal_core/_function_utils.py Normal file
View File

@@ -0,0 +1,324 @@
# File based from: https://github.com/microsoft/autogen/blob/47f905267245e143562abfb41fcba503a9e1d56d/autogen/function_utils.py
# Credit to original authors
import inspect
import typing
from functools import partial
from logging import getLogger
from typing import (
Annotated,
Any,
Callable,
Dict,
List,
Optional,
Set,
Tuple,
Type,
TypeVar,
Union,
cast,
get_args,
get_origin,
)
from pydantic import BaseModel, Field, TypeAdapter, create_model # type: ignore
from pydantic_core import PydanticUndefined
from typing_extensions import Literal
logger = getLogger(__name__)
T = TypeVar("T")
def get_typed_signature(call: Callable[..., Any]) -> inspect.Signature:
"""Get the signature of a function with type annotations.
Args:
call: The function to get the signature for
Returns:
The signature of the function with type annotations
"""
signature = inspect.signature(call)
globalns = getattr(call, "__globals__", {})
func_call = call.func if isinstance(call, partial) else call
type_hints = typing.get_type_hints(func_call, globalns, include_extras=True)
typed_params = [
inspect.Parameter(
name=param.name,
kind=param.kind,
default=param.default,
annotation=type_hints[param.name],
)
for param in signature.parameters.values()
]
return_annotation = type_hints.get("return", inspect.Signature.empty)
typed_signature = inspect.Signature(typed_params, return_annotation=return_annotation)
return typed_signature
def get_typed_return_annotation(call: Callable[..., Any]) -> Any:
"""Get the return annotation of a function.
Args:
call: The function to get the return annotation for
Returns:
The return annotation of the function
"""
signature = inspect.signature(call)
annotation = signature.return_annotation
if annotation is inspect.Signature.empty:
return None
globalns = getattr(call, "__globals__", {})
type_hints = typing.get_type_hints(call, globalns, include_extras=True)
return type_hints.get("return", inspect.Signature.empty)
def get_param_annotations(
typed_signature: inspect.Signature,
) -> Dict[str, Union[Annotated[Type[Any], str], Type[Any]]]:
"""Get the type annotations of the parameters of a function
Args:
typed_signature: The signature of the function with type annotations
Returns:
A dictionary of the type annotations of the parameters of the function
"""
return {
k: v.annotation for k, v in typed_signature.parameters.items() if v.annotation is not inspect.Signature.empty
}
class Parameters(BaseModel):
"""Parameters of a function as defined by the OpenAI API"""
type: Literal["object"] = "object"
properties: Dict[str, Dict[str, Any]]
required: List[str]
class Function(BaseModel):
"""A function as defined by the OpenAI API"""
description: Annotated[str, Field(description="Description of the function")]
name: Annotated[str, Field(description="Name of the function")]
parameters: Annotated[Parameters, Field(description="Parameters of the function")]
class ToolFunction(BaseModel):
"""A function under tool as defined by the OpenAI API."""
type: Literal["function"] = "function"
function: Annotated[Function, Field(description="Function under tool")]
def type2description(k: str, v: Union[Annotated[Type[Any], str], Type[Any]]) -> str:
# handles Annotated
if hasattr(v, "__metadata__"):
retval = v.__metadata__[0]
if isinstance(retval, str):
return retval
else:
raise ValueError(f"Invalid description {retval} for parameter {k}, should be a string.")
else:
return k
def get_parameter_json_schema(k: str, v: Any, default_values: Dict[str, Any]) -> Dict[str, Any]:
"""Get a JSON schema for a parameter as defined by the OpenAI API
Args:
k: The name of the parameter
v: The type of the parameter
default_values: The default values of the parameters of the function
Returns:
A Pydanitc model for the parameter
"""
schema = TypeAdapter(v).json_schema()
if k in default_values:
dv = default_values[k]
schema["default"] = dv
schema["description"] = type2description(k, v)
return schema
def get_required_params(typed_signature: inspect.Signature) -> List[str]:
"""Get the required parameters of a function
Args:
typed_signature: The signature of the function as returned by inspect.signature
Returns:
A list of the required parameters of the function
"""
return [k for k, v in typed_signature.parameters.items() if v.default == inspect.Signature.empty]
def get_default_values(typed_signature: inspect.Signature) -> Dict[str, Any]:
"""Get default values of parameters of a function
Args:
typed_signature: The signature of the function as returned by inspect.signature
Returns:
A dictionary of the default values of the parameters of the function
"""
return {k: v.default for k, v in typed_signature.parameters.items() if v.default != inspect.Signature.empty}
def get_parameters(
required: List[str],
param_annotations: Dict[str, Union[Annotated[Type[Any], str], Type[Any]]],
default_values: Dict[str, Any],
) -> Parameters:
"""Get the parameters of a function as defined by the OpenAI API
Args:
required: The required parameters of the function
param_annotations: A dictionary of the type annotations of the parameters of the function
default_values: The default values of the parameters of the function
Returns:
A Pydantic model for the parameters of the function
"""
return Parameters(
properties={
k: get_parameter_json_schema(k, v, default_values)
for k, v in param_annotations.items()
if v is not inspect.Signature.empty
},
required=required,
)
def get_missing_annotations(typed_signature: inspect.Signature, required: List[str]) -> Tuple[Set[str], Set[str]]:
"""Get the missing annotations of a function
Ignores the parameters with default values as they are not required to be annotated, but logs a warning.
Args:
typed_signature: The signature of the function with type annotations
required: The required parameters of the function
Returns:
A set of the missing annotations of the function
"""
all_missing = {k for k, v in typed_signature.parameters.items() if v.annotation is inspect.Signature.empty}
missing = all_missing.intersection(set(required))
unannotated_with_default = all_missing.difference(missing)
return missing, unannotated_with_default
def get_function_schema(f: Callable[..., Any], *, name: Optional[str] = None, description: str) -> Dict[str, Any]:
"""Get a JSON schema for a function as defined by the OpenAI API
Args:
f: The function to get the JSON schema for
name: The name of the function
description: The description of the function
Returns:
A JSON schema for the function
Raises:
TypeError: If the function is not annotated
Examples:
.. code-block:: python
def f(
a: Annotated[str, "Parameter a"],
b: int = 2,
c: Annotated[float, "Parameter c"] = 0.1,
) -> None:
pass
get_function_schema(f, description="function f")
# {'type': 'function',
# 'function': {'description': 'function f',
# 'name': 'f',
# 'parameters': {'type': 'object',
# 'properties': {'a': {'type': 'str', 'description': 'Parameter a'},
# 'b': {'type': 'int', 'description': 'b'},
# 'c': {'type': 'float', 'description': 'Parameter c'}},
# 'required': ['a']}}}
"""
typed_signature = get_typed_signature(f)
required = get_required_params(typed_signature)
default_values = get_default_values(typed_signature)
param_annotations = get_param_annotations(typed_signature)
return_annotation = get_typed_return_annotation(f)
missing, unannotated_with_default = get_missing_annotations(typed_signature, required)
if return_annotation is None:
logger.warning(
f"The return type of the function '{f.__name__}' is not annotated. Although annotating it is "
+ "optional, the function should return either a string, a subclass of 'pydantic.BaseModel'."
)
if unannotated_with_default != set():
unannotated_with_default_s = [f"'{k}'" for k in sorted(unannotated_with_default)]
logger.warning(
f"The following parameters of the function '{f.__name__}' with default values are not annotated: "
+ f"{', '.join(unannotated_with_default_s)}."
)
if missing != set():
missing_s = [f"'{k}'" for k in sorted(missing)]
raise TypeError(
f"All parameters of the function '{f.__name__}' without default values must be annotated. "
+ f"The annotations are missing for the following parameters: {', '.join(missing_s)}"
)
fname = name if name else f.__name__
parameters = get_parameters(required, param_annotations, default_values=default_values)
function = ToolFunction(
function=Function(
description=description,
name=fname,
parameters=parameters,
)
)
return function.model_dump()
def normalize_annotated_type(type_hint: Type[Any]) -> Type[Any]:
"""Normalize typing.Annotated types to the inner type."""
if get_origin(type_hint) is Annotated:
# Extract the inner type from Annotated
return get_args(type_hint)[0] # type: ignore
return type_hint
def args_base_model_from_signature(name: str, sig: inspect.Signature) -> Type[BaseModel]:
fields: Dict[str, tuple[Type[Any], Any]] = {}
for param_name, param in sig.parameters.items():
# This is handled externally
if param_name == "cancellation_token":
continue
if param.annotation is inspect.Parameter.empty:
raise ValueError("No annotation")
type = normalize_annotated_type(param.annotation)
description = type2description(param_name, param.annotation)
default_value = param.default if param.default is not inspect.Parameter.empty else PydanticUndefined
fields[param_name] = (type, Field(default=default_value, description=description))
return cast(BaseModel, create_model(name, **fields)) # type: ignore

127
agent_dhal/agentdhal_core/_image.py Normal file
View File

@@ -0,0 +1,127 @@
from __future__ import annotations
import base64
import re
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, cast
from PIL import Image as PILImage
from pydantic import GetCoreSchemaHandler, ValidationInfo
from pydantic_core import core_schema
from typing_extensions import Literal
class Image:
"""Represents an image.
Example:
Loading an image from a URL:
.. code-block:: python
from agentdhal_core import Image
from PIL import Image as PILImage
import aiohttp
import asyncio
async def from_url(url: str) -> Image:
async with aiohttp.ClientSession() as session:
async with session.get(url) as response:
content = await response.read()
return Image.from_pil(PILImage.open(content))
image = asyncio.run(from_url("https://example.com/image"))
"""
def __init__(self, image: PILImage.Image):
self.image: PILImage.Image = image.convert("RGB")
@classmethod
def from_pil(cls, pil_image: PILImage.Image) -> Image:
return cls(pil_image)
@classmethod
def from_uri(cls, uri: str) -> Image:
if not re.match(r"data:image/(?:png|jpeg);base64,", uri):
raise ValueError("Invalid URI format. It should be a base64 encoded image URI.")
# A URI. Remove the prefix and decode the base64 string.
base64_data = re.sub(r"data:image/(?:png|jpeg);base64,", "", uri)
return cls.from_base64(base64_data)
@classmethod
def from_base64(cls, base64_str: str) -> Image:
return cls(PILImage.open(BytesIO(base64.b64decode(base64_str))))
def to_base64(self) -> str:
buffered = BytesIO()
self.image.save(buffered, format="PNG")
content = buffered.getvalue()
return base64.b64encode(content).decode("utf-8")
@classmethod
def from_file(cls, file_path: Path) -> Image:
return cls(PILImage.open(file_path))
def _repr_html_(self) -> str:
# Show the image in Jupyter notebook
return f'<img src="{self.data_uri}"/>'
@property
def data_uri(self) -> str:
return _convert_base64_to_data_uri(self.to_base64())
# Returns openai.types.chat.ChatCompletionContentPartImageParam, which is a TypedDict
# We don't use the explicit type annotation so that we can avoid a dependency on the OpenAI Python SDK in this package.
def to_openai_format(self, detail: Literal["auto", "low", "high"] = "auto") -> Dict[str, Any]:
return {"type": "image_url", "image_url": {"url": self.data_uri, "detail": detail}}
@classmethod
def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema:
# Custom validation
def validate(value: Any, validation_info: ValidationInfo) -> Image:
if isinstance(value, dict):
base_64 = cast(str | None, value.get("data")) # type: ignore
if base_64 is None:
raise ValueError("Expected 'data' key in the dictionary")
return cls.from_base64(base_64)
elif isinstance(value, cls):
return value
else:
raise TypeError(f"Expected dict or {cls.__name__} instance, got {type(value)}")
# Custom serialization
def serialize(value: Image) -> dict[str, Any]:
return {"data": value.to_base64()}
return core_schema.with_info_after_validator_function(
validate,
core_schema.any_schema(), # Accept any type; adjust if needed
serialization=core_schema.plain_serializer_function_ser_schema(serialize),
)
def _convert_base64_to_data_uri(base64_image: str) -> str:
def _get_mime_type_from_data_uri(base64_image: str) -> str:
# Decode the base64 string
image_data = base64.b64decode(base64_image)
# Check the first few bytes for known signatures
if image_data.startswith(b"\xff\xd8\xff"):
return "image/jpeg"
elif image_data.startswith(b"\x89PNG\r\n\x1a\n"):
return "image/png"
elif image_data.startswith(b"GIF87a") or image_data.startswith(b"GIF89a"):
return "image/gif"
elif image_data.startswith(b"RIFF") and image_data[8:12] == b"WEBP":
return "image/webp"
return "image/jpeg" # use jpeg for unknown formats, best guess.
mime_type = _get_mime_type_from_data_uri(base64_image)
data_uri = f"data:{mime_type};base64,{base64_image}"
return data_uri

83
agent_dhal/agentdhal_core/_intervention.py Normal file
View File

@@ -0,0 +1,83 @@
from typing import Any, Protocol, final
from ._agent_id import AgentId
from ._message_context import MessageContext
__all__ = [
"DropMessage",
"InterventionHandler",
"DefaultInterventionHandler",
]
@final
class DropMessage:
"""Marker type for signalling that a message should be dropped by an intervention handler. The type itself should be returned from the handler."""
...
class InterventionHandler(Protocol):
"""An intervention handler is a class that can be used to modify, log or drop messages that are being processed by the :class:`agentdhal_core.base.AgentRuntime`.
The handler is called when the message is submitted to the runtime.
Currently the only runtime which supports this is the :class:`agentdhal_core.base.SingleThreadedAgentRuntime`.
Note: Returning None from any of the intervention handler methods will result in a warning being issued and treated as "no change". If you intend to drop a message, you should return :class:`DropMessage` explicitly.
Example:
.. code-block:: python
from agentdhal_core import DefaultInterventionHandler, MessageContext, AgentId, SingleThreadedAgentRuntime
from dataclasses import dataclass
from typing import Any
@dataclass
class MyMessage:
content: str
class MyInterventionHandler(DefaultInterventionHandler):
async def on_send(self, message: Any, *, message_context: MessageContext, recipient: AgentId) -> MyMessage:
if isinstance(message, MyMessage):
message.content = message.content.upper()
return message
runtime = SingleThreadedAgentRuntime(intervention_handlers=[MyInterventionHandler()])
"""
async def on_send(
self, message: Any, *, message_context: MessageContext, recipient: AgentId
) -> Any | type[DropMessage]:
"""Called when a message is submitted to the AgentRuntime using :meth:`agentdhal_core.base.AgentRuntime.send_message`."""
...
async def on_publish(self, message: Any, *, message_context: MessageContext) -> Any | type[DropMessage]:
"""Called when a message is published to the AgentRuntime using :meth:`agentdhal_core.base.AgentRuntime.publish_message`."""
...
async def on_response(self, message: Any, *, sender: AgentId, recipient: AgentId | None) -> Any | type[DropMessage]:
"""Called when a response is received by the AgentRuntime from an Agent's message handler returning a value."""
...
class DefaultInterventionHandler(InterventionHandler):
"""Simple class that provides a default implementation for all intervention
handler methods, that simply returns the message unchanged. Allows for easy
subclassing to override only the desired methods."""
async def on_send(
self, message: Any, *, message_context: MessageContext, recipient: AgentId
) -> Any | type[DropMessage]:
return message
async def on_publish(self, message: Any, *, message_context: MessageContext) -> Any | type[DropMessage]:
return message
async def on_response(self, message: Any, *, sender: AgentId, recipient: AgentId | None) -> Any | type[DropMessage]:
return message

14
agent_dhal/agentdhal_core/_message_context.py Normal file
View File

@@ -0,0 +1,14 @@
from dataclasses import dataclass
from ._agent_id import AgentId
from ._cancellation_token import CancellationToken
from ._topic import TopicId
@dataclass
class MessageContext:
sender: AgentId | None
topic_id: TopicId | None
is_rpc: bool
cancellation_token: CancellationToken
message_id: str

31
agent_dhal/agentdhal_core/_message_handler_context.py Normal file
View File

@@ -0,0 +1,31 @@
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, ClassVar, Generator
from ._agent_id import AgentId
class MessageHandlerContext:
def __init__(self) -> None:
raise RuntimeError(
"MessageHandlerContext cannot be instantiated. It is a static class that provides context management for message handling."
)
_MESSAGE_HANDLER_CONTEXT: ClassVar[ContextVar[AgentId]] = ContextVar("_MESSAGE_HANDLER_CONTEXT")
@classmethod
@contextmanager
def populate_context(cls, ctx: AgentId) -> Generator[None, Any, None]:
""":meta private:"""
token = MessageHandlerContext._MESSAGE_HANDLER_CONTEXT.set(ctx)
try:
yield
finally:
MessageHandlerContext._MESSAGE_HANDLER_CONTEXT.reset(token)
@classmethod
def agent_id(cls) -> AgentId:
try:
return cls._MESSAGE_HANDLER_CONTEXT.get()
except LookupError as e:
raise RuntimeError("MessageHandlerContext.agent_id() must be called within a message handler.") from e

264
agent_dhal/agentdhal_core/_queue.py Normal file
View File

@@ -0,0 +1,264 @@
# Copy of Asyncio queue: https://github.com/python/cpython/blob/main/Lib/asyncio/queues.py
# So that shutdown can be used in <3.13
# Modified to work outside of the asyncio package
import asyncio
import collections
import threading
from typing import Generic, TypeVar
_global_lock = threading.Lock()
class _LoopBoundMixin:
_loop = None
def _get_loop(self) -> asyncio.AbstractEventLoop:
loop = asyncio.get_running_loop()
if self._loop is None:
with _global_lock:
if self._loop is None:
self._loop = loop
if loop is not self._loop:
raise RuntimeError(f"{self!r} is bound to a different event loop")
return loop
class QueueShutDown(Exception):
"""Raised when putting on to or getting from a shut-down Queue."""
pass
T = TypeVar("T")
class Queue(_LoopBoundMixin, Generic[T]):
def __init__(self, maxsize: int = 0):
self._maxsize = maxsize
self._getters = collections.deque[asyncio.Future[None]]()
self._putters = collections.deque[asyncio.Future[None]]()
self._unfinished_tasks = 0
self._finished = asyncio.Event()
self._finished.set()
self._queue = collections.deque[T]()
self._is_shutdown = False
# These three are overridable in subclasses.
def _get(self) -> T:
return self._queue.popleft()
def _put(self, item: T) -> None:
self._queue.append(item)
# End of the overridable methods.
def _wakeup_next(self, waiters: collections.deque[asyncio.Future[None]]) -> None:
# Wake up the next waiter (if any) that isn't cancelled.
while waiters:
waiter = waiters.popleft()
if not waiter.done():
waiter.set_result(None)
break
def __repr__(self) -> str:
return f"<{type(self).__name__} at {id(self):#x} {self._format()}>"
def __str__(self) -> str:
return f"<{type(self).__name__} {self._format()}>"
def _format(self) -> str:
result = f"maxsize={self._maxsize!r}"
if getattr(self, "_queue", None):
result += f" _queue={list(self._queue)!r}"
if self._getters:
result += f" _getters[{len(self._getters)}]"
if self._putters:
result += f" _putters[{len(self._putters)}]"
if self._unfinished_tasks:
result += f" tasks={self._unfinished_tasks}"
if self._is_shutdown:
result += " shutdown"
return result
def qsize(self) -> int:
"""Number of items in the queue."""
return len(self._queue)
@property
def maxsize(self) -> int:
"""Number of items allowed in the queue."""
return self._maxsize
def empty(self) -> bool:
"""Return True if the queue is empty, False otherwise."""
return not self._queue
def full(self) -> bool:
"""Return True if there are maxsize items in the queue.
Note: if the Queue was initialized with maxsize=0 (the default),
then full() is never True.
"""
if self._maxsize <= 0:
return False
else:
return self.qsize() >= self._maxsize
async def put(self, item: T) -> None:
"""Put an item into the queue.
Put an item into the queue. If the queue is full, wait until a free
slot is available before adding item.
Raises QueueShutDown if the queue has been shut down.
"""
while self.full():
if self._is_shutdown:
raise QueueShutDown
putter = self._get_loop().create_future()
self._putters.append(putter)
try:
await putter
except:
putter.cancel() # Just in case putter is not done yet.
try:
# Clean self._putters from canceled putters.
self._putters.remove(putter)
except ValueError:
# The putter could be removed from self._putters by a
# previous get_nowait call or a shutdown call.
pass
if not self.full() and not putter.cancelled():
# We were woken up by get_nowait(), but can't take
# the call. Wake up the next in line.
self._wakeup_next(self._putters)
raise
return self.put_nowait(item)
def put_nowait(self, item: T) -> None:
"""Put an item into the queue without blocking.
If no free slot is immediately available, raise QueueFull.
Raises QueueShutDown if the queue has been shut down.
"""
if self._is_shutdown:
raise QueueShutDown
if self.full():
raise asyncio.QueueFull
self._put(item)
self._unfinished_tasks += 1
self._finished.clear()
self._wakeup_next(self._getters)
async def get(self) -> T:
"""Remove and return an item from the queue.
If queue is empty, wait until an item is available.
Raises QueueShutDown if the queue has been shut down and is empty, or
if the queue has been shut down immediately.
"""
while self.empty():
if self._is_shutdown and self.empty():
raise QueueShutDown
getter = self._get_loop().create_future()
self._getters.append(getter)
try:
await getter
except:
getter.cancel() # Just in case getter is not done yet.
try:
# Clean self._getters from canceled getters.
self._getters.remove(getter)
except ValueError:
# The getter could be removed from self._getters by a
# previous put_nowait call, or a shutdown call.
pass
if not self.empty() and not getter.cancelled():
# We were woken up by put_nowait(), but can't take
# the call. Wake up the next in line.
self._wakeup_next(self._getters)
raise
return self.get_nowait()
def get_nowait(self) -> T:
"""Remove and return an item from the queue.
Return an item if one is immediately available, else raise QueueEmpty.
Raises QueueShutDown if the queue has been shut down and is empty, or
if the queue has been shut down immediately.
"""
if self.empty():
if self._is_shutdown:
raise QueueShutDown
raise asyncio.QueueEmpty
item = self._get()
self._wakeup_next(self._putters)
return item
def task_done(self) -> None:
"""Indicate that a formerly enqueued task is complete.
Used by queue consumers. For each get() used to fetch a task,
a subsequent call to task_done() tells the queue that the processing
on the task is complete.
If a join() is currently blocking, it will resume when all items have
been processed (meaning that a task_done() call was received for every
item that had been put() into the queue).
shutdown(immediate=True) calls task_done() for each remaining item in
the queue.
Raises ValueError if called more times than there were items placed in
the queue.
"""
if self._unfinished_tasks <= 0:
raise ValueError("task_done() called too many times")
self._unfinished_tasks -= 1
if self._unfinished_tasks == 0:
self._finished.set()
async def join(self) -> None:
"""Block until all items in the queue have been gotten and processed.
The count of unfinished tasks goes up whenever an item is added to the
queue. The count goes down whenever a consumer calls task_done() to
indicate that the item was retrieved and all work on it is complete.
When the count of unfinished tasks drops to zero, join() unblocks.
"""
if self._unfinished_tasks > 0:
await self._finished.wait()
def shutdown(self, immediate: bool = False) -> None:
"""Shut-down the queue, making queue gets and puts raise QueueShutDown.
By default, gets will only raise once the queue is empty. Set
'immediate' to True to make gets raise immediately instead.
All blocked callers of put() and get() will be unblocked. If
'immediate', a task is marked as done for each item remaining in
the queue, which may unblock callers of join().
"""
self._is_shutdown = True
if immediate:
while not self.empty():
self._get()
if self._unfinished_tasks > 0:
self._unfinished_tasks -= 1
if self._unfinished_tasks == 0:
self._finished.set()
# All getters need to re-check queue-empty to raise ShutDown
while self._getters:
getter = self._getters.popleft()
if not getter.done():
getter.set_result(None)
while self._putters:
putter = self._putters.popleft()
if not putter.done():
putter.set_result(None)

518
agent_dhal/agentdhal_core/_routed_agent.py Normal file
View File

@@ -0,0 +1,518 @@
import logging
from functools import wraps
from typing import (
Any,
Callable,
Coroutine,
DefaultDict,
List,
Literal,
Protocol,
Sequence,
Tuple,
Type,
TypeVar,
cast,
get_type_hints,
overload,
runtime_checkable,
)
from ._base_agent import BaseAgent
from ._message_context import MessageContext
from ._serialization import MessageSerializer, try_get_known_serializers_for_type
from ._type_helpers import AnyType, get_types
from .exceptions import CantHandleException
logger = logging.getLogger("agentdhal_core")
AgentT = TypeVar("AgentT")
ReceivesT = TypeVar("ReceivesT")
ProducesT = TypeVar("ProducesT", covariant=True)
# TODO: Generic typevar bound binding U to agent type
# Can't do because python doesnt support it
# Pyright and mypy disagree on the variance of ReceivesT. Mypy thinks it should be contravariant here.
# Revisit this later to see if we can remove the ignore.
@runtime_checkable
class MessageHandler(Protocol[AgentT, ReceivesT, ProducesT]): # type: ignore
target_types: Sequence[type]
produces_types: Sequence[type]
is_message_handler: Literal[True]
router: Callable[[ReceivesT, MessageContext], bool]
# agent_instance binds to self in the method
@staticmethod
async def __call__(agent_instance: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT: ...
# NOTE: this works on concrete types and not inheritance
# TODO: Use a protocol for the outer function to check checked arg names
@overload
def message_handler(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
) -> MessageHandler[AgentT, ReceivesT, ProducesT]: ...
@overload
def message_handler(
func: None = None,
*,
match: None = ...,
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]: ...
@overload
def message_handler(
func: None = None,
*,
match: Callable[[ReceivesT, MessageContext], bool],
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]: ...
def message_handler(
func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]] = None,
*,
strict: bool = True,
match: None | Callable[[ReceivesT, MessageContext], bool] = None,
) -> (
Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]
| MessageHandler[AgentT, ReceivesT, ProducesT]
):
"""Decorator for generic message handlers.
Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle both event and RPC messages.
These methods must have a specific signature that needs to be followed for it to be valid:
- The method must be an `async` method.
- The method must be decorated with the `@message_handler` decorator.
- The method must have exactly 3 arguments:
1. `self`
2. `message`: The message to be handled, this must be type-hinted with the message type that it is intended to handle.
3. `ctx`: A :class:`agentdhal_core.MessageContext` object.
- The method must be type hinted with what message types it can return as a response, or it can return `None` if it does not return anything.
Handlers can handle more than one message type by accepting a Union of the message types. It can also return more than one message type by returning a Union of the message types.
Args:
func: The function to be decorated.
strict: If `True`, the handler will raise an exception if the message type or return type is not in the target types. If `False`, it will log a warning instead.
match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
"""
def decorator(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
) -> MessageHandler[AgentT, ReceivesT, ProducesT]:
type_hints = get_type_hints(func)
if "message" not in type_hints:
raise AssertionError("message parameter not found in function signature")
if "return" not in type_hints:
raise AssertionError("return parameter not found in function signature")
# Get the type of the message parameter
target_types = get_types(type_hints["message"])
if target_types is None:
raise AssertionError("Message type not found")
# print(type_hints)
return_types = get_types(type_hints["return"])
if return_types is None:
raise AssertionError("Return type not found")
# Convert target_types to list and stash
@wraps(func)
async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT:
if type(message) not in target_types:
if strict:
raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
else:
logger.warning(f"Message type {type(message)} not in target types {target_types}")
return_value = await func(self, message, ctx)
if AnyType not in return_types and type(return_value) not in return_types:
if strict:
raise ValueError(f"Return type {type(return_value)} not in return types {return_types}")
else:
logger.warning(f"Return type {type(return_value)} not in return types {return_types}")
return return_value
wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
wrapper_handler.target_types = list(target_types)
wrapper_handler.produces_types = list(return_types)
wrapper_handler.is_message_handler = True
wrapper_handler.router = match or (lambda _message, _ctx: True)
return wrapper_handler
if func is None and not callable(func):
return decorator
elif callable(func):
return decorator(func)
else:
raise ValueError("Invalid arguments")
@overload
def event(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]],
) -> MessageHandler[AgentT, ReceivesT, None]: ...
@overload
def event(
func: None = None,
*,
match: None = ...,
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
MessageHandler[AgentT, ReceivesT, None],
]: ...
@overload
def event(
func: None = None,
*,
match: Callable[[ReceivesT, MessageContext], bool],
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
MessageHandler[AgentT, ReceivesT, None],
]: ...
def event(
func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]] = None,
*,
strict: bool = True,
match: None | Callable[[ReceivesT, MessageContext], bool] = None,
) -> (
Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]]],
MessageHandler[AgentT, ReceivesT, None],
]
| MessageHandler[AgentT, ReceivesT, None]
):
"""Decorator for event message handlers.
Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle event messages.
These methods must have a specific signature that needs to be followed for it to be valid:
- The method must be an `async` method.
- The method must be decorated with the `@message_handler` decorator.
- The method must have exactly 3 arguments:
1. `self`
2. `message`: The event message to be handled, this must be type-hinted with the message type that it is intended to handle.
3. `ctx`: A :class:`agentdhal_core.MessageContext` object.
- The method must return `None`.
Handlers can handle more than one message type by accepting a Union of the message types.
Args:
func: The function to be decorated.
strict: If `True`, the handler will raise an exception if the message type is not in the target types. If `False`, it will log a warning instead.
match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
"""
def decorator(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, None]],
) -> MessageHandler[AgentT, ReceivesT, None]:
type_hints = get_type_hints(func)
if "message" not in type_hints:
raise AssertionError("message parameter not found in function signature")
if "return" not in type_hints:
raise AssertionError("return parameter not found in function signature")
# Get the type of the message parameter
target_types = get_types(type_hints["message"])
if target_types is None:
raise AssertionError("Message type not found. Please provide a type hint for the message parameter.")
return_types = get_types(type_hints["return"])
if return_types is None:
raise AssertionError("Return type not found. Please use `None` as the type hint of the return type.")
# Convert target_types to list and stash
@wraps(func)
async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> None:
if type(message) not in target_types:
if strict:
raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
else:
logger.warning(f"Message type {type(message)} not in target types {target_types}")
return_value = await func(self, message, ctx) # type: ignore
if return_value is not None:
if strict:
raise ValueError(f"Return type {type(return_value)} is not None.")
else:
logger.warning(f"Return type {type(return_value)} is not None. It will be ignored.")
return None
wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, None], wrapper)
wrapper_handler.target_types = list(target_types)
wrapper_handler.produces_types = list(return_types)
wrapper_handler.is_message_handler = True
# Wrap the match function with a check on the is_rpc flag.
wrapper_handler.router = lambda _message, _ctx: (not _ctx.is_rpc) and (match(_message, _ctx) if match else True)
return wrapper_handler
if func is None and not callable(func):
return decorator
elif callable(func):
return decorator(func)
else:
raise ValueError("Invalid arguments")
@overload
def rpc(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
) -> MessageHandler[AgentT, ReceivesT, ProducesT]: ...
@overload
def rpc(
func: None = None,
*,
match: None = ...,
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]: ...
@overload
def rpc(
func: None = None,
*,
match: Callable[[ReceivesT, MessageContext], bool],
strict: bool = ...,
) -> Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]: ...
def rpc(
func: None | Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]] = None,
*,
strict: bool = True,
match: None | Callable[[ReceivesT, MessageContext], bool] = None,
) -> (
Callable[
[Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]]],
MessageHandler[AgentT, ReceivesT, ProducesT],
]
| MessageHandler[AgentT, ReceivesT, ProducesT]
):
"""Decorator for RPC message handlers.
Add this decorator to methods in a :class:`RoutedAgent` class that are intended to handle RPC messages.
These methods must have a specific signature that needs to be followed for it to be valid:
- The method must be an `async` method.
- The method must be decorated with the `@message_handler` decorator.
- The method must have exactly 3 arguments:
1. `self`
2. `message`: The message to be handled, this must be type-hinted with the message type that it is intended to handle.
3. `ctx`: A :class:`agentdhal_core.MessageContext` object.
- The method must be type hinted with what message types it can return as a response, or it can return `None` if it does not return anything.
Handlers can handle more than one message type by accepting a Union of the message types. It can also return more than one message type by returning a Union of the message types.
Args:
func: The function to be decorated.
strict: If `True`, the handler will raise an exception if the message type or return type is not in the target types. If `False`, it will log a warning instead.
match: A function that takes the message and the context as arguments and returns a boolean. This is used for secondary routing after the message type. For handlers addressing the same message type, the match function is applied in alphabetical order of the handlers and the first matching handler will be called while the rest are skipped. If `None`, the first handler in alphabetical order matching the same message type will be called.
"""
def decorator(
func: Callable[[AgentT, ReceivesT, MessageContext], Coroutine[Any, Any, ProducesT]],
) -> MessageHandler[AgentT, ReceivesT, ProducesT]:
type_hints = get_type_hints(func)
if "message" not in type_hints:
raise AssertionError("message parameter not found in function signature")
if "return" not in type_hints:
raise AssertionError("return parameter not found in function signature")
# Get the type of the message parameter
target_types = get_types(type_hints["message"])
if target_types is None:
raise AssertionError("Message type not found")
# print(type_hints)
return_types = get_types(type_hints["return"])
if return_types is None:
raise AssertionError("Return type not found")
# Convert target_types to list and stash
@wraps(func)
async def wrapper(self: AgentT, message: ReceivesT, ctx: MessageContext) -> ProducesT:
if type(message) not in target_types:
if strict:
raise CantHandleException(f"Message type {type(message)} not in target types {target_types}")
else:
logger.warning(f"Message type {type(message)} not in target types {target_types}")
return_value = await func(self, message, ctx)
if AnyType not in return_types and type(return_value) not in return_types:
if strict:
raise ValueError(f"Return type {type(return_value)} not in return types {return_types}")
else:
logger.warning(f"Return type {type(return_value)} not in return types {return_types}")
return return_value
wrapper_handler = cast(MessageHandler[AgentT, ReceivesT, ProducesT], wrapper)
wrapper_handler.target_types = list(target_types)
wrapper_handler.produces_types = list(return_types)
wrapper_handler.is_message_handler = True
wrapper_handler.router = lambda _message, _ctx: (_ctx.is_rpc) and (match(_message, _ctx) if match else True)
return wrapper_handler
if func is None and not callable(func):
return decorator
elif callable(func):
return decorator(func)
else:
raise ValueError("Invalid arguments")
class RoutedAgent(BaseAgent):
"""A base class for agents that route messages to handlers based on the type of the message
and optional matching functions.
To create a routed agent, subclass this class and add message handlers as methods decorated with
either :func:`event` or :func:`rpc` decorator.
Example:
.. code-block:: python
from dataclasses import dataclass
from agentdhal_core import MessageContext
from agentdhal_core import RoutedAgent, event, rpc
@dataclass
class Message:
pass
@dataclass
class MessageWithContent:
content: str
@dataclass
class Response:
pass
class MyAgent(RoutedAgent):
def __init__(self):
super().__init__("MyAgent")
@event
async def handle_event_message(self, message: Message, ctx: MessageContext) -> None:
assert ctx.topic_id is not None
await self.publish_message(MessageWithContent("event handled"), ctx.topic_id)
@rpc(match=lambda message, ctx: message.content == "special") # type: ignore
async def handle_special_rpc_message(self, message: MessageWithContent, ctx: MessageContext) -> Response:
return Response()
"""
def __init__(self, description: str) -> None:
# Self is already bound to the handlers
self._handlers: DefaultDict[
Type[Any],
List[MessageHandler[RoutedAgent, Any, Any]],
] = DefaultDict(list)
handlers = self._discover_handlers()
for message_handler in handlers:
for target_type in message_handler.target_types:
self._handlers[target_type].append(message_handler)
super().__init__(description)
async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None:
"""Handle a message by routing it to the appropriate message handler.
Do not override this method in subclasses. Instead, add message handlers as methods decorated with
either the :func:`event` or :func:`rpc` decorator."""
key_type: Type[Any] = type(message) # type: ignore
handlers = self._handlers.get(key_type) # type: ignore
if handlers is not None:
# Iterate over all handlers for this matching message type.
# Call the first handler whose router returns True and then return the result.
for h in handlers:
if h.router(message, ctx):
return await h(self, message, ctx)
return await self.on_unhandled_message(message, ctx) # type: ignore
async def on_unhandled_message(self, message: Any, ctx: MessageContext) -> None:
"""Called when a message is received that does not have a matching message handler.
The default implementation logs an info message."""
logger.info(f"Unhandled message: {message}")
@classmethod
def _discover_handlers(cls) -> Sequence[MessageHandler[Any, Any, Any]]:
handlers: List[MessageHandler[Any, Any, Any]] = []
for attr in dir(cls):
if callable(getattr(cls, attr, None)):
# Since we are getting it from the class, self is not bound
handler = getattr(cls, attr)
if hasattr(handler, "is_message_handler"):
handlers.append(cast(MessageHandler[Any, Any, Any], handler))
return handlers
@classmethod
def _handles_types(cls) -> List[Tuple[Type[Any], List[MessageSerializer[Any]]]]:
# TODO handle deduplication
handlers = cls._discover_handlers()
types: List[Tuple[Type[Any], List[MessageSerializer[Any]]]] = []
types.extend(cls.internal_extra_handles_types)
for handler in handlers:
for t in handler.target_types:
# TODO: support different serializers
serializers = try_get_known_serializers_for_type(t)
if len(serializers) == 0:
raise ValueError(f"No serializers found for type {t}.")
types.append((t, try_get_known_serializers_for_type(t)))
return types

78
agent_dhal/agentdhal_core/_runtime_impl_helpers.py Normal file
View File

@@ -0,0 +1,78 @@
from collections import defaultdict
from typing import Awaitable, Callable, DefaultDict, List, Sequence, Set
from ._agent import Agent
from ._agent_id import AgentId
from ._agent_type import AgentType
from ._subscription import Subscription
from ._topic import TopicId
async def get_impl(
*,
id_or_type: AgentId | AgentType | str,
key: str,
lazy: bool,
instance_getter: Callable[[AgentId], Awaitable[Agent]],
) -> AgentId:
if isinstance(id_or_type, AgentId):
if not lazy:
await instance_getter(id_or_type)
return id_or_type
type_str = id_or_type if isinstance(id_or_type, str) else id_or_type.type
id = AgentId(type_str, key)
if not lazy:
await instance_getter(id)
return id
class SubscriptionManager:
def __init__(self) -> None:
self._subscriptions: List[Subscription] = []
self._seen_topics: Set[TopicId] = set()
self._subscribed_recipients: DefaultDict[TopicId, List[AgentId]] = defaultdict(list)
@property
def subscriptions(self) -> Sequence[Subscription]:
return self._subscriptions
async def add_subscription(self, subscription: Subscription) -> None:
# Check if the subscription already exists
if any(sub == subscription for sub in self._subscriptions):
raise ValueError("Subscription already exists")
self._subscriptions.append(subscription)
self._rebuild_subscriptions(self._seen_topics)
async def remove_subscription(self, id: str) -> None:
# Check if the subscription exists
if not any(sub.id == id for sub in self._subscriptions):
raise ValueError("Subscription does not exist")
def is_not_sub(x: Subscription) -> bool:
return x.id != id
self._subscriptions = list(filter(is_not_sub, self._subscriptions))
# Rebuild the subscriptions
self._rebuild_subscriptions(self._seen_topics)
async def get_subscribed_recipients(self, topic: TopicId) -> List[AgentId]:
if topic not in self._seen_topics:
self._build_for_new_topic(topic)
return self._subscribed_recipients[topic]
# TODO: optimize this...
def _rebuild_subscriptions(self, topics: Set[TopicId]) -> None:
self._subscribed_recipients.clear()
for topic in topics:
self._build_for_new_topic(topic)
def _build_for_new_topic(self, topic: TopicId) -> None:
self._seen_topics.add(topic)
for subscription in self._subscriptions:
if subscription.is_match(topic):
self._subscribed_recipients[topic].append(subscription.map_to_agent(topic))

258
agent_dhal/agentdhal_core/_serialization.py Normal file
View File

@@ -0,0 +1,258 @@
import json
from dataclasses import asdict, dataclass, fields
from typing import Any, ClassVar, Dict, List, Protocol, Sequence, TypeVar, cast, get_args, get_origin, runtime_checkable
from google.protobuf import any_pb2
from google.protobuf.message import Message
from pydantic import BaseModel
from ._type_helpers import is_union
T = TypeVar("T")
class MessageSerializer(Protocol[T]):
@property
def data_content_type(self) -> str: ...
@property
def type_name(self) -> str: ...
def deserialize(self, payload: bytes) -> T: ...
def serialize(self, message: T) -> bytes: ...
@runtime_checkable
class IsDataclass(Protocol):
# as already noted in comments, checking for this attribute is currently
# the most reliable way to ascertain that something is a dataclass
__dataclass_fields__: ClassVar[Dict[str, Any]]
def is_dataclass(cls: type[Any]) -> bool:
return hasattr(cls, "__dataclass_fields__")
def has_nested_dataclass(cls: type[IsDataclass]) -> bool:
# iterate fields and check if any of them are dataclasses
return any(is_dataclass(f.type) for f in cls.__dataclass_fields__.values())
def contains_a_union(cls: type[IsDataclass]) -> bool:
return any(is_union(f.type) for f in cls.__dataclass_fields__.values())
def has_nested_base_model(cls: type[IsDataclass]) -> bool:
for f in fields(cls):
field_type = f.type
# Resolve forward references and other annotations
origin = get_origin(field_type)
args = get_args(field_type)
# If the field type is directly a subclass of BaseModel
if isinstance(field_type, type) and issubclass(field_type, BaseModel):
return True
# If the field type is a generic type like List[BaseModel], Tuple[BaseModel, ...], etc.
if origin is not None and args:
for arg in args:
# Recursively check the argument types
if isinstance(arg, type) and issubclass(arg, BaseModel):
return True
elif get_origin(arg) is not None:
# Handle nested generics like List[List[BaseModel]]
if has_nested_base_model_in_type(arg):
return True
# Handle Union types
elif args:
for arg in args:
if isinstance(arg, type) and issubclass(arg, BaseModel):
return True
elif get_origin(arg) is not None:
if has_nested_base_model_in_type(arg):
return True
return False
def has_nested_base_model_in_type(tp: Any) -> bool:
"""Helper function to check if a type or its arguments is a BaseModel subclass."""
origin = get_origin(tp)
args = get_args(tp)
if isinstance(tp, type) and issubclass(tp, BaseModel):
return True
if origin is not None and args:
for arg in args:
if has_nested_base_model_in_type(arg):
return True
return False
DataclassT = TypeVar("DataclassT", bound=IsDataclass)
JSON_DATA_CONTENT_TYPE = "application/json"
"""JSON data content type"""
# TODO: what's the correct content type? There seems to be some disagreement over what it should be
PROTOBUF_DATA_CONTENT_TYPE = "application/x-protobuf"
"""Protobuf data content type"""
class DataclassJsonMessageSerializer(MessageSerializer[DataclassT]):
def __init__(self, cls: type[DataclassT]) -> None:
if contains_a_union(cls):
raise ValueError("Dataclass has a union type, which is not supported. To use a union, use a Pydantic model")
if has_nested_dataclass(cls) or has_nested_base_model(cls):
raise ValueError(
"Dataclass has nested dataclasses or base models, which are not supported. To use nested types, use a Pydantic model"
)
self.cls = cls
@property
def data_content_type(self) -> str:
return JSON_DATA_CONTENT_TYPE
@property
def type_name(self) -> str:
return _type_name(self.cls)
def deserialize(self, payload: bytes) -> DataclassT:
message_str = payload.decode("utf-8")
return self.cls(**json.loads(message_str))
def serialize(self, message: DataclassT) -> bytes:
return json.dumps(asdict(message)).encode("utf-8")
PydanticT = TypeVar("PydanticT", bound=BaseModel)
class PydanticJsonMessageSerializer(MessageSerializer[PydanticT]):
def __init__(self, cls: type[PydanticT]) -> None:
self.cls = cls
@property
def data_content_type(self) -> str:
return JSON_DATA_CONTENT_TYPE
@property
def type_name(self) -> str:
return _type_name(self.cls)
def deserialize(self, payload: bytes) -> PydanticT:
message_str = payload.decode("utf-8")
return self.cls.model_validate_json(message_str)
def serialize(self, message: PydanticT) -> bytes:
return message.model_dump_json().encode("utf-8")
ProtobufT = TypeVar("ProtobufT", bound=Message)
# This class serializes to and from a google.protobuf.Any message that has been serialized to a string
class ProtobufMessageSerializer(MessageSerializer[ProtobufT]):
def __init__(self, cls: type[ProtobufT]) -> None:
self.cls = cls
@property
def data_content_type(self) -> str:
return PROTOBUF_DATA_CONTENT_TYPE
@property
def type_name(self) -> str:
return _type_name(self.cls)
def deserialize(self, payload: bytes) -> ProtobufT:
# Parse payload into a proto any
any_proto = any_pb2.Any()
any_proto.ParseFromString(payload)
destination_message = self.cls()
if not any_proto.Unpack(destination_message): # type: ignore
raise ValueError(f"Failed to unpack payload into {self.cls}")
return destination_message
def serialize(self, message: ProtobufT) -> bytes:
any_proto = any_pb2.Any()
any_proto.Pack(message) # type: ignore
return any_proto.SerializeToString()
@dataclass
class UnknownPayload:
type_name: str
data_content_type: str
payload: bytes
def _type_name(cls: type[Any] | Any) -> str:
# If cls is a protobuf, then we need to determine the descriptor
if isinstance(cls, type):
if issubclass(cls, Message):
return cast(str, cls.DESCRIPTOR.full_name)
elif isinstance(cls, Message):
return cast(str, cls.DESCRIPTOR.full_name)
if isinstance(cls, type):
return cls.__name__
else:
return cast(str, cls.__class__.__name__)
V = TypeVar("V")
def try_get_known_serializers_for_type(cls: type[Any]) -> list[MessageSerializer[Any]]:
""":meta private:"""
serializers: List[MessageSerializer[Any]] = []
if issubclass(cls, BaseModel):
serializers.append(PydanticJsonMessageSerializer(cls))
elif is_dataclass(cls):
serializers.append(DataclassJsonMessageSerializer(cls))
elif issubclass(cls, Message):
serializers.append(ProtobufMessageSerializer(cls))
return serializers
class SerializationRegistry:
""":meta private:"""
def __init__(self) -> None:
# type_name, data_content_type -> serializer
self._serializers: dict[tuple[str, str], MessageSerializer[Any]] = {}
def add_serializer(self, serializer: MessageSerializer[Any] | Sequence[MessageSerializer[Any]]) -> None:
if isinstance(serializer, Sequence):
for c in serializer:
self.add_serializer(c)
return
self._serializers[(serializer.type_name, serializer.data_content_type)] = serializer
def deserialize(self, payload: bytes, *, type_name: str, data_content_type: str) -> Any:
serializer = self._serializers.get((type_name, data_content_type))
if serializer is None:
return UnknownPayload(type_name, data_content_type, payload)
return serializer.deserialize(payload)
def serialize(self, message: Any, *, type_name: str, data_content_type: str) -> bytes:
serializer = self._serializers.get((type_name, data_content_type))
if serializer is None:
raise ValueError(f"Unknown type {type_name} with content type {data_content_type}")
return serializer.serialize(message)
def is_registered(self, type_name: str, data_content_type: str) -> bool:
return (type_name, data_content_type) in self._serializers
def type_name(self, message: Any) -> str:
return _type_name(message)

1029
agent_dhal/agentdhal_core/_single_threaded_agent_runtime.py Normal file
View File

File diff suppressed because it is too large Load Diff

65
agent_dhal/agentdhal_core/_subscription.py Normal file
View File

@@ -0,0 +1,65 @@
from __future__ import annotations
from typing import Awaitable, Callable, Protocol, runtime_checkable
from ._agent_id import AgentId
from ._topic import TopicId
@runtime_checkable
class Subscription(Protocol):
"""Subscriptions define the topics that an agent is interested in."""
@property
def id(self) -> str:
"""Get the ID of the subscription.
Implementations should return a unique ID for the subscription. Usually this is a UUID.
Returns:
str: ID of the subscription.
"""
...
def __eq__(self, other: object) -> bool:
"""Check if two subscriptions are equal.
Args:
other (object): Other subscription to compare against.
Returns:
bool: True if the subscriptions are equal, False otherwise.
"""
if not isinstance(other, Subscription):
return False
return self.id == other.id
def is_match(self, topic_id: TopicId) -> bool:
"""Check if a given topic_id matches the subscription.
Args:
topic_id (TopicId): TopicId to check.
Returns:
bool: True if the topic_id matches the subscription, False otherwise.
"""
...
def map_to_agent(self, topic_id: TopicId) -> AgentId:
"""Map a topic_id to an agent. Should only be called if `is_match` returns True for the given topic_id.
Args:
topic_id (TopicId): TopicId to map.
Returns:
AgentId: ID of the agent that should handle the topic_id.
Raises:
CantHandleException: If the subscription cannot handle the topic_id.
"""
...
# Helper alias to represent the lambdas used to define subscriptions
UnboundSubscription = Callable[[], list[Subscription] | Awaitable[list[Subscription]]]

33
agent_dhal/agentdhal_core/_subscription_context.py Normal file
View File

@@ -0,0 +1,33 @@
from contextlib import contextmanager
from contextvars import ContextVar
from typing import Any, ClassVar, Generator
from ._agent_type import AgentType
class SubscriptionInstantiationContext:
def __init__(self) -> None:
raise RuntimeError(
"SubscriptionInstantiationContext cannot be instantiated. It is a static class that provides context management for subscription instantiation."
)
_SUBSCRIPTION_CONTEXT_VAR: ClassVar[ContextVar[AgentType]] = ContextVar("_SUBSCRIPTION_CONTEXT_VAR")
@classmethod
@contextmanager
def populate_context(cls, ctx: AgentType) -> Generator[None, Any, None]:
""":meta private:"""
token = SubscriptionInstantiationContext._SUBSCRIPTION_CONTEXT_VAR.set(ctx)
try:
yield
finally:
SubscriptionInstantiationContext._SUBSCRIPTION_CONTEXT_VAR.reset(token)
@classmethod
def agent_type(cls) -> AgentType:
try:
return cls._SUBSCRIPTION_CONTEXT_VAR.get()
except LookupError as e:
raise RuntimeError(
"SubscriptionInstantiationContext.runtime() must be called within an instantiation context such as when the AgentRuntime is instantiating an agent. Mostly likely this was caused by directly instantiating an agent instead of using the AgentRuntime to do so."
) from e

25
agent_dhal/agentdhal_core/_telemetry/__init__.py Normal file
View File

@@ -0,0 +1,25 @@
from ._genai import (
trace_create_agent_span,
trace_invoke_agent_span,
trace_tool_span,
)
from ._propagation import (
EnvelopeMetadata,
TelemetryMetadataContainer,
get_telemetry_envelope_metadata,
get_telemetry_grpc_metadata,
)
from ._tracing import TraceHelper
from ._tracing_config import MessageRuntimeTracingConfig
__all__ = [
"EnvelopeMetadata",
"get_telemetry_envelope_metadata",
"get_telemetry_grpc_metadata",
"TelemetryMetadataContainer",
"TraceHelper",
"MessageRuntimeTracingConfig",
"trace_create_agent_span",
"trace_invoke_agent_span",
"trace_tool_span",
]

1
agent_dhal/agentdhal_core/_telemetry/_constants.py Normal file
View File

@@ -0,0 +1 @@
NAMESPACE = "agentdhal"

214
agent_dhal/agentdhal_core/_telemetry/_genai.py Normal file
View File

@@ -0,0 +1,214 @@
from collections.abc import Generator
from contextlib import contextmanager
from enum import Enum
from typing import Any, Optional
from opentelemetry import trace
from opentelemetry.trace import Span, SpanKind
from .._agent_instantiation import AgentInstantiationContext
# OpenTelemetry semantic convention constants for GenAI operations
# Copied from opentelemetry-semantic-conventions to avoid dependency
# GenAI Agent attributes
GEN_AI_AGENT_DESCRIPTION = "gen_ai.agent.description"
GEN_AI_AGENT_ID = "gen_ai.agent.id"
GEN_AI_AGENT_NAME = "gen_ai.agent.name"
# GenAI Operation attributes
GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
GEN_AI_SYSTEM = "gen_ai.system"
# GenAI Tool attributes
GEN_AI_TOOL_CALL_ID = "gen_ai.tool.call.id"
GEN_AI_TOOL_DESCRIPTION = "gen_ai.tool.description"
GEN_AI_TOOL_NAME = "gen_ai.tool.name"
# Error attributes
ERROR_TYPE = "error.type"
class GenAiOperationNameValues(Enum):
"""Enum for GenAI operation name values."""
CHAT = "chat"
CREATE_AGENT = "create_agent"
EMBEDDINGS = "embeddings"
EXECUTE_TOOL = "execute_tool"
GENERATE_CONTENT = "generate_content"
INVOKE_AGENT = "invoke_agent"
TEXT_COMPLETION = "text_completion"
# Constant for system name
GENAI_SYSTEM_AUTOGEN = "agentdhal"
@contextmanager
def trace_tool_span(
tool_name: str,
*,
tracer: Optional[trace.Tracer] = None,
parent: Optional[Span] = None,
tool_description: Optional[str] = None,
tool_call_id: Optional[str] = None,
) -> Generator[Span, Any, None]:
"""Context manager to create a span for tool execution following the
OpenTelemetry Semantic conventions for generative AI systems.
See the GenAI semantic conventions documentation:
`OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__
.. warning::
The GenAI Semantic Conventions are still in incubation and
subject to changes in future releases.
Args:
tool_name (str): The name of the tool being executed.
tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
parent (Optional[Span]): The parent span to link this span to.
tool_description (Optional[str]): A description of the tool.
tool_call_id (Optional[str]): A unique identifier for the tool call.
"""
if tracer is None:
tracer = trace.get_tracer("agentdhal-core")
span_attributes = {
GEN_AI_OPERATION_NAME: GenAiOperationNameValues.EXECUTE_TOOL.value,
GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
GEN_AI_TOOL_NAME: tool_name,
}
if tool_description is not None:
span_attributes[GEN_AI_TOOL_DESCRIPTION] = tool_description
if tool_call_id is not None:
span_attributes[GEN_AI_TOOL_CALL_ID] = tool_call_id
with tracer.start_as_current_span(
f"{GenAiOperationNameValues.EXECUTE_TOOL.value} {tool_name}",
kind=SpanKind.INTERNAL,
context=trace.set_span_in_context(parent) if parent else None,
attributes=span_attributes,
) as span:
try:
yield span
except Exception as e:
# Set the exception details on the span if an error occurs
span.record_exception(e)
span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
span.set_attribute(ERROR_TYPE, type(e).__name__)
raise
@contextmanager
def trace_create_agent_span(
agent_name: str,
*,
tracer: Optional[trace.Tracer] = None,
parent: Optional[Span] = None,
agent_id: Optional[str] = None,
agent_description: Optional[str] = None,
) -> Generator[Span, Any, None]:
"""Context manager to create a span for agent creation following the
OpenTelemetry Semantic conventions for generative AI systems.
See the GenAI semantic conventions documentation:
`OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__
.. warning::
The GenAI Semantic Conventions are still in incubation and
subject to changes in future releases.
Args:
agent_name (str): The name of the agent being created.
tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
parent (Optional[Span]): The parent span to link this span to.
agent_id (Optional[str]): The unique identifier for the agent.
agent_description (Optional[str]): A description of the agent.
"""
if tracer is None:
tracer = trace.get_tracer("agentdhal-core")
span_attributes = {
GEN_AI_OPERATION_NAME: GenAiOperationNameValues.CREATE_AGENT.value,
GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
GEN_AI_AGENT_NAME: agent_name,
}
if agent_id is None:
# Try to see if we can get the agent ID from the current context
try:
agent_id = str(AgentInstantiationContext.current_agent_id())
except RuntimeError:
agent_id = None
if agent_id is not None:
span_attributes[GEN_AI_AGENT_ID] = agent_id
if agent_description is not None:
span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description
with tracer.start_as_current_span(
f"{GenAiOperationNameValues.CREATE_AGENT.value} {agent_name}",
kind=SpanKind.CLIENT,
context=trace.set_span_in_context(parent) if parent else None,
attributes=span_attributes,
) as span:
try:
yield span
except Exception as e:
# Set the exception details on the span if an error occurs
span.record_exception(e)
span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
span.set_attribute(ERROR_TYPE, type(e).__name__)
raise
@contextmanager
def trace_invoke_agent_span(
agent_name: str,
*,
tracer: Optional[trace.Tracer] = None,
parent: Optional[Span] = None,
agent_id: Optional[str] = None,
agent_description: Optional[str] = None,
) -> Generator[Span, Any, None]:
"""Context manager to create a span for invoking an agent following the
OpenTelemetry Semantic conventions for generative AI systems.
See the GenAI semantic conventions documentation:
`OpenTelemetry GenAI Semantic Conventions <https://opentelemetry.io/docs/specs/semconv/gen-ai/>`__
.. warning::
The GenAI Semantic Conventions are still in incubation and
subject to changes in future releases.
Args:
agent_name (str): The name of the agent being invoked.
tracer (Optional[trace.Tracer]): The tracer to use for creating the span.
parent (Optional[Span]): The parent span to link this span to.
agent_id (Optional[str]): The unique identifier for the agent.
agent_description (Optional[str]): A description of the agent.
"""
if tracer is None:
tracer = trace.get_tracer("agentdhal-core")
span_attributes = {
GEN_AI_OPERATION_NAME: GenAiOperationNameValues.INVOKE_AGENT.value,
GEN_AI_SYSTEM: GENAI_SYSTEM_AUTOGEN,
GEN_AI_AGENT_NAME: agent_name,
}
if agent_id is not None:
span_attributes[GEN_AI_AGENT_ID] = agent_id
if agent_description is not None:
span_attributes[GEN_AI_AGENT_DESCRIPTION] = agent_description
with tracer.start_as_current_span(
f"{GenAiOperationNameValues.INVOKE_AGENT.value} {agent_name}",
kind=SpanKind.CLIENT,
context=trace.set_span_in_context(parent) if parent else None,
attributes=span_attributes,
) as span:
try:
yield span
except Exception as e:
# Set the exception details on the span if an error occurs
span.record_exception(e)
span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
span.set_attribute(ERROR_TYPE, type(e).__name__)
raise

Some files were not shown because too many files have changed in this diff Show More