mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-11 03:31:27 -10:00
Compare commits
108
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7e8ca38749 | ||
|
|
5c8d1149fb | ||
|
|
d5f1114aaf | ||
|
|
d364fa5020 | ||
|
|
c1db6317c8 | ||
|
|
f88727cb75 | ||
|
|
7ab0479e82 | ||
|
|
0e0241d15f | ||
|
|
5815dc61bf | ||
|
|
e720c0bc3d | ||
|
|
cd448fa710 | ||
|
|
c35dadebdb | ||
|
|
f9ae9e9aa4 | ||
|
|
9d9a9a7ad8 | ||
|
|
44d9d8c4c0 | ||
|
|
bf2656cee9 | ||
|
|
d55cd989ea | ||
|
|
0305214695 | ||
|
|
3f26bf166c | ||
|
|
ab19a13cc4 | ||
|
|
28d487b66d | ||
|
|
a06239fb15 | ||
|
|
dabe2cf671 | ||
|
|
619787c9da | ||
|
|
2fa6a6efd0 | ||
|
|
8acf77c531 | ||
|
|
5ff79d39c9 | ||
|
|
6902147f13 | ||
|
|
321b6b28a9 | ||
|
|
1c5b58aa13 | ||
|
|
71834d39d7 | ||
|
|
6971d8ecc9 | ||
|
|
8348c2b7ab | ||
|
|
9cfc636a92 | ||
|
|
b43e93c502 | ||
|
|
b09450d9dc | ||
|
|
710d607f6e | ||
|
|
a7910c2eb0 | ||
|
|
c98e1a1972 | ||
|
|
aee9552dd0 | ||
|
|
6507b8e7ef | ||
|
|
50652531e7 | ||
|
|
a949579759 | ||
|
|
79702ba1ff | ||
|
|
e9b75bfef9 | ||
|
|
982b7dfe4f | ||
|
|
f372821b9e | ||
|
|
445003042e | ||
|
|
00b40e8cc4 | ||
|
|
6dca959eda | ||
|
|
5ba8782bf7 | ||
|
|
a2402063f8 | ||
|
|
302bbcbec2 | ||
|
|
ac309e99ea | ||
|
|
727ee626c8 | ||
|
|
1dc61d3946 | ||
|
|
50a0399f1c | ||
|
|
40bfd66797 | ||
|
|
0da4a90c80 | ||
|
|
dc75c51fea | ||
|
|
b95a7f6ded | ||
|
|
92953fdbe6 | ||
|
|
b3cd7ecfd0 | ||
|
|
01ea0f5a05 | ||
|
|
eddcae8d31 | ||
|
|
71f891003e | ||
|
|
13750a4ccd | ||
|
|
c9e1d1a202 | ||
|
|
5515290329 | ||
|
|
cb1cefd94b | ||
|
|
dbf8b0d6e1 | ||
|
|
a90666a19f | ||
|
|
89df701c2a | ||
|
|
480848c26d | ||
|
|
a064a286f3 | ||
|
|
c83aa5a35d | ||
|
|
9543799f95 | ||
|
|
7719a47892 | ||
|
|
7873afd6c9 | ||
|
|
8d05e93c21 | ||
|
|
8c2ef4102b | ||
|
|
5796583019 | ||
|
|
11c7e370e1 | ||
|
|
92e4540e81 | ||
|
|
69e0e91623 | ||
|
|
613309f39b | ||
|
|
6dbb645ba6 | ||
|
|
a4428fd0d0 | ||
|
|
348904557c | ||
|
|
9a5c1bce9e | ||
|
|
fee8e581f4 | ||
|
|
0d43cd98f2 | ||
|
|
fdcbdc7640 | ||
|
|
d6a1a26685 | ||
|
|
1d78204b0b | ||
|
|
31967f6a65 | ||
|
|
bba5d43cc1 | ||
|
|
c2dfa859c3 | ||
|
|
bb8a56e33c | ||
|
|
456c4d5853 | ||
|
|
efc01ab3d8 | ||
|
|
88c86db05d | ||
|
|
12d1811f0c | ||
|
|
6012d36dfd | ||
|
|
0f17585315 | ||
|
|
5900d8bba6 | ||
|
|
75e9085eb4 | ||
|
|
bfba25efaf |
No files matched your search
@@ -37,6 +37,14 @@ WHATSAPP_PHONE_NUMBER_ID=... # WhatsApp phone number ID
|
||||
# === Optional: Database ===
|
||||
DATABASE_URL=sqlite:///~/.praison/database.sqlite # Database connection string
|
||||
|
||||
# === Optional: Tool Security ===
|
||||
# ALLOWED_TOOLS=search,send_message # Whitelist specific tools to prevent name collisions
|
||||
# Useful in multi-environment setups with overlapping tool names
|
||||
# Leave unset to use all tools (with collision warnings)
|
||||
# Do NOT set to empty string: that's treated as configuration error
|
||||
# In CI, unknown tool names in the whitelist cause strict startup failure
|
||||
# HERMES_ONLY_TOOLS=search,send_message # (Deprecated) Backward compatibility alias for ALLOWED_TOOLS
|
||||
|
||||
# === Optional: Development ===
|
||||
LOGLEVEL=INFO # Logging level (DEBUG, INFO, WARNING, ERROR)
|
||||
CHAINLIT_AUTH_SECRET=... # Chainlit authentication secret
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
name: Cross-Repo Integration Tests
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
paths:
|
||||
- 'src/praisonai/praisonai/integration/**'
|
||||
- 'src/praisonai/tests/integration/test_aiui_*'
|
||||
- 'src/praisonai/pyproject.toml'
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
paths:
|
||||
- 'src/praisonai/praisonai/integration/**'
|
||||
- 'src/praisonai/tests/integration/test_aiui_*'
|
||||
- 'src/praisonai/pyproject.toml'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
cross-repo-integration:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- name: Checkout PraisonAI
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
path: PraisonAI
|
||||
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
|
||||
|
||||
- name: Checkout PraisonAIUI
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: MervinPraison/PraisonAIUI
|
||||
path: PraisonAIUI
|
||||
ref: main
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
# Install PraisonAI packages in development mode
|
||||
cd PraisonAI/src/praisonai-agents
|
||||
pip install -e .
|
||||
cd ../praisonai
|
||||
pip install -e ".[ui,dev]"
|
||||
# Install PraisonAIUI in development mode
|
||||
cd ../../../PraisonAIUI
|
||||
pip install -e .
|
||||
|
||||
- name: Verify versions
|
||||
run: |
|
||||
python -c "import praisonaiagents; print(f'praisonaiagents: {praisonaiagents.__version__}')"
|
||||
python -c "import praisonai; print(f'praisonai: {praisonai.__version__}')"
|
||||
python -c "import praisonaiui; print(f'praisonaiui: {praisonaiui.__version__}')"
|
||||
|
||||
- name: Run PraisonAI integration tests
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
python -m pytest tests/integration/test_aiui_* -v --timeout=300
|
||||
|
||||
- name: Run PraisonAIUI integration tests
|
||||
run: |
|
||||
cd PraisonAIUI
|
||||
if [ -f tests/integration/test_agentic_roundtrip.py ]; then
|
||||
python -m pytest tests/integration/test_agentic_roundtrip.py -v --timeout=300
|
||||
fi
|
||||
if [ -f tests/test_feature_sdk_backends.py ]; then
|
||||
python -m pytest tests/test_feature_sdk_backends.py -v --timeout=300
|
||||
fi
|
||||
|
||||
- name: Test Pattern C CLI
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
# Test that ui-gateway command exists
|
||||
python -m praisonai serve --help | grep ui-gateway
|
||||
|
||||
- name: Test public API exports
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
python -c "from praisonai import run_integrated_gateway; print('✓ run_integrated_gateway imported')"
|
||||
python -c "from praisonai import configure_host; print('✓ configure_host imported')"
|
||||
|
||||
- name: Integration smoke test
|
||||
env:
|
||||
PRAISONAI_TEST_MODE: "1"
|
||||
PYTHONPATH: ${{ github.workspace }}/PraisonAI/src/praisonai:${{ github.workspace }}/PraisonAI/src/praisonai-agents:${{ github.workspace }}/PraisonAIUI
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
python -c "
|
||||
import sys
|
||||
sys.path.insert(0, '.')
|
||||
from praisonai.integration.host_app import configure_host
|
||||
from praisonai import run_integrated_gateway
|
||||
print('✓ Integration imports successful')
|
||||
|
||||
# Test configure_host with new parameters
|
||||
try:
|
||||
configure_host(
|
||||
style='dashboard',
|
||||
context_paths=['AGENTS.md'],
|
||||
title='Test App'
|
||||
)
|
||||
print('✓ configure_host with new parameters works')
|
||||
except Exception as e:
|
||||
print(f'× configure_host failed: {e}')
|
||||
sys.exit(1)
|
||||
"
|
||||
|
||||
# Optional job with API key for agentic tests (if secret is available)
|
||||
agentic-integration:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
if: ${{ vars.RUN_AGENTIC_TESTS == 'true' }}
|
||||
needs: cross-repo-integration
|
||||
|
||||
steps:
|
||||
- name: Checkout PraisonAI
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
path: PraisonAI
|
||||
ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
|
||||
|
||||
- name: Checkout PraisonAIUI
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: MervinPraison/PraisonAIUI
|
||||
path: PraisonAIUI
|
||||
ref: main
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
cd PraisonAI/src/praisonai-agents
|
||||
pip install -e .
|
||||
cd ../praisonai
|
||||
pip install -e ".[ui,dev]"
|
||||
cd ../../../PraisonAIUI
|
||||
pip install -e .
|
||||
|
||||
- name: Run agentic integration test
|
||||
env:
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
run: |
|
||||
cd PraisonAI/src/praisonai
|
||||
if [ -n "$OPENAI_API_KEY" ]; then
|
||||
python -m pytest tests/integration/test_aiui_host_agentic.py -v --timeout=600
|
||||
else
|
||||
echo "OPENAI_API_KEY not available, skipping agentic tests"
|
||||
fi
|
||||
@@ -26,8 +26,8 @@ jobs:
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
persist-credentials: false
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Set up Python 3.11
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
@@ -176,6 +176,18 @@ jobs:
|
||||
run: praisonai queue list
|
||||
continue-on-error: true
|
||||
|
||||
# ==================== PACKAGING TESTS ====================
|
||||
- name: Test packaging/daemon entry points
|
||||
run: |
|
||||
echo "Testing both entry point methods..."
|
||||
python -m praisonai --version
|
||||
praisonai --version
|
||||
|
||||
- name: Test doctor packaging command
|
||||
run: |
|
||||
echo "Testing doctor packaging diagnostics..."
|
||||
praisonai doctor packaging --json
|
||||
|
||||
# ==================== REAL API TESTS ====================
|
||||
- name: Real API test - Agent run
|
||||
if: env.OPENAI_API_KEY != ''
|
||||
|
||||
@@ -382,7 +382,7 @@ praisonai flow
|
||||
|
||||
Open **http://localhost:7861** — use the **Agent** and **Agent Team** components to create sequential or parallel workflows. Connect Chat Input → Agent Team → Chat Output for instant multi-agent pipelines.
|
||||
|
||||
> 📖 [Full Flow docs](https://docs.praison.ai/docs/concepts/flow) — visual agent building, component reference, and deployment
|
||||
> 📖 [Full Flow docs](https://docs.praison.ai/docs/concepts/agentflow) — visual agent building, component reference, and deployment
|
||||
|
||||
### 8. PraisonAI UI 🤖 (Clean Chat)
|
||||
|
||||
|
||||
@@ -40,6 +40,7 @@ Methods:
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">chat_with_context</a>(message: str, context: 'ContextPack', **kwargs) -> str</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">checkpoints</a>() -> Optional[bool]</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">checkpoints</a>(value: Optional[bool]) -> None</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">clone_for_channel</a>() -> 'Agent'</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">close</a>() -> None</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">console</a>() -> Optional[Any]</code>
|
||||
* <code title="class Agent">Agent.<a href="./src/praisonai-agents/praisonaiagents/agent/agent.py">context_manager</a>() -> Optional[Any]</code>
|
||||
@@ -171,6 +172,7 @@ Methods:
|
||||
* <code title="class BaseTool">BaseTool.<a href="./src/praisonai-agents/praisonaiagents/tools/base.py">safe_run</a>(**kwargs) -> ToolResult</code>
|
||||
* <code title="class BaseTool">BaseTool.<a href="./src/praisonai-agents/praisonaiagents/tools/base.py">validate</a>() -> bool</code>
|
||||
* <code title="class BaseTool">BaseTool.<a href="./src/praisonai-agents/praisonaiagents/tools/base.py">validate_class</a>() -> bool</code>
|
||||
* <code title="class BaseTool">BaseTool.<a href="./src/praisonai-agents/praisonaiagents/tools/base.py">validate_schema_roundtrip</a>() -> bool</code>
|
||||
* <code title="class FunctionTool">FunctionTool.<a href="./src/praisonai-agents/praisonaiagents/tools/decorator.py">__call__</a>(*args, **kwargs) -> Any</code>
|
||||
* <code title="class FunctionTool">FunctionTool.<a href="./src/praisonai-agents/praisonaiagents/tools/decorator.py">check_availability</a>() -> tuple[bool, str]</code>
|
||||
* <code title="class FunctionTool">FunctionTool.<a href="./src/praisonai-agents/praisonaiagents/tools/decorator.py">injected_params</a>() -> Dict[str, Any]</code>
|
||||
@@ -183,6 +185,7 @@ Methods:
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">list_available_tools</a>(context: Optional[Dict[str, Any]] = None) -> List[Union[BaseTool, Callable]]</code>
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">list_base_tools</a>() -> List[BaseTool]</code>
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">list_tools</a>() -> List[str]</code>
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">list_tools_with_allowed_filter</a>(context: Optional[Dict[str, Any]] = None) -> List[str]</code>
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">register</a>(tool: Union[BaseTool, Callable], name: Optional[str] = None, overwrite: bool = False) -> None</code>
|
||||
* <code title="class ToolRegistry">ToolRegistry.<a href="./src/praisonai-agents/praisonaiagents/tools/registry.py">unregister</a>(name: str) -> bool</code>
|
||||
* <code title="class Tools">Tools.<a href="./src/praisonai-agents/praisonaiagents/tools/tools.py">internet_search</a>(*args, **kwargs)</code>
|
||||
@@ -476,8 +479,10 @@ Methods:
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">edit_skill</a>(name: str, content: str) -> dict</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_allowed_tools</a>(name: str) -> List[str]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_available_skills</a>() -> List[SkillMetadata]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_available_skills_by_state</a>(state: SkillState) -> List[LoadedSkill]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_instructions</a>(name: str) -> Optional[str]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_skill</a>(name: str) -> Optional[LoadedSkill]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_skills_diagnostics</a>() -> Dict[str, ValidationResult]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">get_user_invocable_skills</a>() -> List[LoadedSkill]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">invoke</a>(name: str, raw_args: str = '', session_id: Optional[str] = None, shell_exec: bool = False) -> Optional[str]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">load_resources</a>(name: str) -> bool</code>
|
||||
@@ -486,6 +491,7 @@ Methods:
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">skill_names</a>() -> List[str]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">skills</a>() -> List[LoadedSkill]</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">to_prompt</a>() -> str</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">validate_skill_capabilities</a>(skill_name: str, force_refresh: bool = False) -> ValidationResult</code>
|
||||
* <code title="class SkillManager">SkillManager.<a href="./src/praisonai-agents/praisonaiagents/skills/manager.py">write_skill_file</a>(name: str, file_path: str, file_content: str) -> dict</code>
|
||||
* <code title="class SkillMetadata">SkillMetadata.<a href="./src/praisonai-agents/praisonaiagents/skills/models.py">from_properties</a>(props: SkillProperties) -> 'SkillMetadata'</code>
|
||||
* <code title="class SkillProperties">SkillProperties.<a href="./src/praisonai-agents/praisonaiagents/skills/models.py">to_dict</a>() -> dict</code>
|
||||
@@ -601,6 +607,13 @@ Methods:
|
||||
* <code title="class ContextManager">ContextManager.<a href="./src/praisonai-agents/praisonaiagents/context/manager.py">set_tool_budget</a>(tool_name: str, max_tokens: int, protected: bool = False) -> None</code>
|
||||
* <code title="class ContextManager">ContextManager.<a href="./src/praisonai-agents/praisonaiagents/context/manager.py">truncate_tool_output</a>(tool_name: str, output: str) -> str</code>
|
||||
|
||||
# UI
|
||||
|
||||
Types:
|
||||
```python
|
||||
from praisonaiagents import A2UI
|
||||
```
|
||||
|
||||
# Config
|
||||
|
||||
Types:
|
||||
@@ -755,7 +768,7 @@ Methods:
|
||||
|
||||
Types:
|
||||
```python
|
||||
from praisonai import Agent, AgentApp, AgentOS, AnthropicManagedAgent, CloudProvider, DB, Deploy, DeployConfig, DeployType, HostedAgent, HostedAgentConfig, LocalAgent, LocalAgentConfig, LocalManagedAgent, LocalManagedConfig, ManagedAgent, ManagedConfig, PraisonAI, __version__
|
||||
from praisonai import Agent, AgentApp, AgentOS, AnthropicManagedAgent, CloudProvider, DB, Deploy, DeployConfig, DeployType, HostedAgent, HostedAgentConfig, LocalAgent, LocalAgentConfig, LocalManagedAgent, LocalManagedConfig, ManagedAgent, ManagedConfig, PraisonAI, __version__, arun, run, run_integrated_gateway
|
||||
```
|
||||
|
||||
# CLI
|
||||
@@ -857,8 +870,10 @@ Methods:
|
||||
* <code title="cli">praisonai doctor env <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor mcp <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor network <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor packaging <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor performance <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor selftest <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor skills <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor tools <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai doctor troubleshoot <a href="./src/praisonai/praisonai/cli/commands/doctor.py">--help</a></code>
|
||||
* <code title="cli">praisonai endpoints list <a href="./src/praisonai/praisonai/cli/commands/endpoints.py">--help</a></code>
|
||||
@@ -890,6 +905,7 @@ Methods:
|
||||
* <code title="cli">praisonai gateway send <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway start <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway status <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway stop <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai gateway uninstall <a href="./src/praisonai/praisonai/cli/commands/gateway.py">--help</a></code>
|
||||
* <code title="cli">praisonai github triage <a href="./src/praisonai/praisonai/cli/commands/github.py">--help</a></code>
|
||||
* <code title="cli">praisonai hooks add <a href="./src/praisonai/praisonai/cli/commands/hooks.py">--help</a></code>
|
||||
@@ -1052,6 +1068,7 @@ Methods:
|
||||
* <code title="cli">praisonai serve status <a href="./src/praisonai/praisonai/cli/commands/serve.py">--help</a></code>
|
||||
* <code title="cli">praisonai serve stop <a href="./src/praisonai/praisonai/cli/commands/serve.py">--help</a></code>
|
||||
* <code title="cli">praisonai serve ui <a href="./src/praisonai/praisonai/cli/commands/serve.py">--help</a></code>
|
||||
* <code title="cli">praisonai serve ui-gateway <a href="./src/praisonai/praisonai/cli/commands/serve.py">--help</a></code>
|
||||
* <code title="cli">praisonai serve unified <a href="./src/praisonai/praisonai/cli/commands/serve.py">--help</a></code>
|
||||
* <code title="cli">praisonai session delete <a href="./src/praisonai/praisonai/cli/commands/session.py">--help</a></code>
|
||||
* <code title="cli">praisonai session export <a href="./src/praisonai/praisonai/cli/commands/session.py">--help</a></code>
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.41" \
|
||||
"praisonai>=4.6.48" \
|
||||
"praisonai[chat]" \
|
||||
"embedchain[github,youtube]"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.41" \
|
||||
"praisonai>=4.6.48" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[chat]" \
|
||||
"praisonai[realtime]" \
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=4.6.41" \
|
||||
"praisonai>=4.6.48" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Minimal A2UI example — agent + send_a2ui_messages tool.
|
||||
|
||||
Install:
|
||||
pip install praisonaiagents[a2ui]
|
||||
|
||||
Run:
|
||||
python send_surface.py
|
||||
|
||||
For live rendering, use PraisonAIUI with AG-UI (/agui) — see integrate-a2ui-frontend docs.
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.tools.a2ui_tools import send_a2ui_messages
|
||||
|
||||
agent = Agent(
|
||||
name="ui-assistant",
|
||||
instructions=(
|
||||
"When asked for UI, call send_a2ui_messages with A2UI v0.9 JSON "
|
||||
"(createSurface on surface id 'main')."
|
||||
),
|
||||
tools=[send_a2ui_messages],
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
result = agent.start("Show a simple welcome card on surface main")
|
||||
print(result)
|
||||
@@ -0,0 +1,222 @@
|
||||
# Multi-Channel Bot Deployment Guide
|
||||
|
||||
This guide explains how to set up multiple bots on the same platform (e.g., multiple Telegram bots) for different roles using PraisonAI, similar to Hermes-style AI workforces.
|
||||
|
||||
## Overview
|
||||
|
||||
PraisonAI supports running multiple bots on the same platform, each routing to different agents with specialized knowledge and capabilities. This is useful for:
|
||||
|
||||
- **Role-based bots**: `@company_cfo_bot`, `@company_ops_bot`, `@company_content_bot`
|
||||
- **Department-specific agents**: Each with different knowledge brains and tool access
|
||||
- **Workflow separation**: Different bots handle different types of requests
|
||||
|
||||
## Quick Start with Onboard Wizard
|
||||
|
||||
The easiest way to set up multiple channels is using the onboard wizard:
|
||||
|
||||
```bash
|
||||
praisonai onboard
|
||||
```
|
||||
|
||||
The wizard now supports:
|
||||
1. **Multi-channel setup**: Add multiple bots for the same platform
|
||||
2. **Role-based naming**: Automatically generates `TELEGRAM_<ROLE>_BOT_TOKEN` environment variables
|
||||
3. **Agent routing**: Each channel routes to a different agent
|
||||
|
||||
### Onboard Wizard Flow
|
||||
|
||||
1. **Choose platform**: Start with your first platform (e.g., telegram)
|
||||
2. **Add more channels**: Wizard asks "Add another bot channel?"
|
||||
3. **Specify role**: Enter role name (e.g., "cfo", "ops", "content")
|
||||
4. **Configure tokens**: Each channel gets its own environment variable
|
||||
5. **Test connections**: Wizard tests each bot separately
|
||||
|
||||
## Environment Variable Convention
|
||||
|
||||
For multiple bots on the same platform, use this naming convention:
|
||||
|
||||
| Channel | Environment Variable | Bot Purpose |
|
||||
|---------|---------------------|-------------|
|
||||
| `telegram_cfo` | `TELEGRAM_CFO_BOT_TOKEN` | CFO-related tasks |
|
||||
| `telegram_ops` | `TELEGRAM_OPS_BOT_TOKEN` | Operations tasks |
|
||||
| `telegram_content` | `TELEGRAM_CONTENT_BOT_TOKEN` | Content creation |
|
||||
| `telegram_support` | `TELEGRAM_SUPPORT_BOT_TOKEN` | Customer support |
|
||||
|
||||
### Setting Up Tokens
|
||||
|
||||
1. **Create bots in @BotFather**:
|
||||
```text
|
||||
/newbot
|
||||
Bot name: Company CFO Bot
|
||||
Username: @company_cfo_bot
|
||||
-> Copy token
|
||||
|
||||
/newbot
|
||||
Bot name: Company Ops Bot
|
||||
Username: @company_ops_bot
|
||||
-> Copy token
|
||||
```
|
||||
|
||||
2. **Set environment variables**:
|
||||
```bash
|
||||
export TELEGRAM_CFO_BOT_TOKEN="123456789:ABC..."
|
||||
export TELEGRAM_OPS_BOT_TOKEN="987654321:DEF..."
|
||||
export TELEGRAM_CONTENT_BOT_TOKEN="456789123:GHI..."
|
||||
```
|
||||
|
||||
3. **Or use `~/.praisonai/.env`**:
|
||||
```env
|
||||
TELEGRAM_CFO_BOT_TOKEN=123456789:ABC...
|
||||
TELEGRAM_OPS_BOT_TOKEN=987654321:DEF...
|
||||
TELEGRAM_CONTENT_BOT_TOKEN=456789123:GHI...
|
||||
TELEGRAM_ALLOWED_USERS=123456789,987654321
|
||||
```
|
||||
|
||||
## Example Configuration
|
||||
|
||||
See `multi-telegram-hermes-workforce.yaml` for a complete example.
|
||||
|
||||
### Key Features
|
||||
|
||||
1. **Multiple Channels**:
|
||||
```yaml
|
||||
channels:
|
||||
telegram_cfo:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_CFO_BOT_TOKEN}
|
||||
routes:
|
||||
default: cfo
|
||||
|
||||
telegram_ops:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_OPS_BOT_TOKEN}
|
||||
routes:
|
||||
default: ops
|
||||
```
|
||||
|
||||
2. **Specialized Agents**:
|
||||
```yaml
|
||||
agents:
|
||||
cfo:
|
||||
instructions: "You are a CFO agent. Help with financial analysis..."
|
||||
model: gpt-4o-mini
|
||||
|
||||
ops:
|
||||
instructions: "You are an ops agent. Help with system monitoring..."
|
||||
model: gpt-4o-mini
|
||||
```
|
||||
|
||||
## Validation with Doctor
|
||||
|
||||
Use the doctor command to validate your multi-channel setup:
|
||||
|
||||
```bash
|
||||
praisonai doctor
|
||||
```
|
||||
|
||||
The doctor checks for:
|
||||
- ✅ **Unique tokens**: Each channel uses a different bot token
|
||||
- ✅ **Naming convention**: Follows `PLATFORM_ROLE_BOT_TOKEN` pattern
|
||||
- ✅ **Missing tokens**: All referenced environment variables are set
|
||||
- ❌ **Duplicate tokens**: Same token used by multiple channels (fails)
|
||||
|
||||
## Deployment
|
||||
|
||||
### Start Gateway
|
||||
|
||||
```bash
|
||||
# Start the multi-bot gateway
|
||||
praisonai gateway start --config multi-telegram-hermes-workforce.yaml
|
||||
|
||||
# Or install as daemon
|
||||
praisonai onboard # Wizard installs daemon automatically
|
||||
praisonai gateway status
|
||||
```
|
||||
|
||||
### Verify Setup
|
||||
|
||||
1. **Check health**: `curl http://127.0.0.1:8765/health`
|
||||
2. **View dashboard**: `praisonai claw` → http://127.0.0.1:8082
|
||||
3. **Test bots**: Send messages to each bot to verify routing
|
||||
|
||||
## Common Patterns
|
||||
|
||||
### Basic Multi-Role Setup
|
||||
```yaml
|
||||
# Single platform, multiple roles
|
||||
channels:
|
||||
telegram_support:
|
||||
token: ${TELEGRAM_SUPPORT_BOT_TOKEN}
|
||||
routes:
|
||||
default: support_agent
|
||||
|
||||
telegram_sales:
|
||||
token: ${TELEGRAM_SALES_BOT_TOKEN}
|
||||
routes:
|
||||
default: sales_agent
|
||||
```
|
||||
|
||||
### Cross-Platform Workforce
|
||||
```yaml
|
||||
# Multiple platforms, multiple roles
|
||||
channels:
|
||||
telegram_cfo:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_CFO_BOT_TOKEN}
|
||||
routes:
|
||||
default: cfo
|
||||
|
||||
discord_ops:
|
||||
platform: discord
|
||||
token: ${DISCORD_OPS_BOT_TOKEN}
|
||||
routes:
|
||||
default: ops
|
||||
|
||||
slack_content:
|
||||
platform: slack
|
||||
token: ${SLACK_CONTENT_BOT_TOKEN}
|
||||
routes:
|
||||
default: content
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **"Token is used by multiple channels"**
|
||||
- Each bot needs its own unique token from @BotFather
|
||||
- Check environment variables for duplicates
|
||||
|
||||
2. **"Channel skipped - no token"**
|
||||
- Set the environment variable: `export TELEGRAM_CFO_BOT_TOKEN="..."`
|
||||
- Or update `~/.praisonai/.env`
|
||||
|
||||
3. **"Bot not responding"**
|
||||
- Verify token is correct: `praisonai doctor`
|
||||
- Check bot is not already running elsewhere
|
||||
- Ensure allowlist includes your user ID
|
||||
|
||||
### Debugging Commands
|
||||
|
||||
```bash
|
||||
# Check configuration
|
||||
praisonai doctor
|
||||
|
||||
# View logs
|
||||
praisonai gateway logs
|
||||
|
||||
# Test connection
|
||||
curl -H "Authorization: Bearer $GATEWAY_AUTH_TOKEN" \
|
||||
http://127.0.0.1:8765/info
|
||||
```
|
||||
|
||||
## Migration from Single-Bot Setup
|
||||
|
||||
If you currently have a single bot setup, you can add more channels:
|
||||
|
||||
1. **Run onboard wizard**: `praisonai onboard`
|
||||
2. **Add new channels**: Wizard will preserve existing setup
|
||||
3. **Update configuration**: New channels added to existing `bot.yaml`
|
||||
4. **Restart gateway**: `praisonai gateway restart`
|
||||
|
||||
The wizard preserves backward compatibility with single-bot configurations.
|
||||
@@ -0,0 +1,19 @@
|
||||
# A2UI canvas agent — YAML example
|
||||
#
|
||||
# Usage:
|
||||
# pip install praisonaiagents[a2ui]
|
||||
# praisonai agent run --file a2ui-canvas.yaml --prompt "Show a welcome card"
|
||||
#
|
||||
# For live A2UI rendering, use PraisonAIUI (AG-UI / surfaces):
|
||||
# pip install praisonaiui
|
||||
# See examples/python/a2ui/send_surface.py and PraisonAIUI example 29.
|
||||
|
||||
agent:
|
||||
name: "ui-assistant"
|
||||
instructions: |
|
||||
You are a helpful assistant. When the user asks for UI elements
|
||||
(buttons, cards, forms), call send_a2ui_messages with valid A2UI v0.9 JSON.
|
||||
Use surface id "main" unless told otherwise.
|
||||
llm: gpt-4o-mini
|
||||
tools:
|
||||
- send_a2ui_messages
|
||||
@@ -0,0 +1,66 @@
|
||||
# Multi-Telegram Channel Hermes Workforce Example
|
||||
# Generated by `praisonai onboard` with multi-channel support
|
||||
# Start with: praisonai gateway start --config multi-telegram-hermes-workforce.yaml
|
||||
|
||||
gateway:
|
||||
host: "127.0.0.1"
|
||||
port: 8765
|
||||
|
||||
# Legacy single-agent block for backward compatibility
|
||||
agent:
|
||||
name: "assistant"
|
||||
instructions: "You are a helpful AI assistant. You can search the web, manage schedules, and remember past conversations. Use your tools proactively."
|
||||
model: "gpt-4o-mini"
|
||||
|
||||
# Multi-agent workforce with different specialized roles
|
||||
agents:
|
||||
cfo:
|
||||
instructions: "You are a CFO agent. You are a helpful AI assistant. Use your tools proactively to help with CFO-related tasks like financial analysis, budget planning, and cost optimization."
|
||||
model: gpt-4o-mini
|
||||
memory: false
|
||||
|
||||
ops:
|
||||
instructions: "You are an ops agent. You are a helpful AI assistant. Use your tools proactively to help with ops-related tasks like system monitoring, deployment, and infrastructure management."
|
||||
model: gpt-4o-mini
|
||||
memory: false
|
||||
|
||||
content:
|
||||
instructions: "You are a content agent. You are a helpful AI assistant. Use your tools proactively to help with content-related tasks like writing, editing, and content strategy."
|
||||
model: gpt-4o-mini
|
||||
memory: false
|
||||
|
||||
# Multiple Telegram channels - each with unique tokens and routing
|
||||
channels:
|
||||
telegram_cfo:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_CFO_BOT_TOKEN}
|
||||
group_policy: mention_only
|
||||
allowed_users: ${TELEGRAM_ALLOWED_USERS}
|
||||
# auto_approve_tools: true
|
||||
# Chat bots cannot show CLI approval prompts; safe tools auto-approved by default.
|
||||
routes:
|
||||
default: cfo
|
||||
|
||||
telegram_ops:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_OPS_BOT_TOKEN}
|
||||
group_policy: mention_only
|
||||
allowed_users: ${TELEGRAM_ALLOWED_USERS}
|
||||
# auto_approve_tools: true
|
||||
# Chat bots cannot show CLI approval prompts; safe tools auto-approved by default.
|
||||
routes:
|
||||
default: ops
|
||||
|
||||
telegram_content:
|
||||
platform: telegram
|
||||
token: ${TELEGRAM_CONTENT_BOT_TOKEN}
|
||||
group_policy: mention_only
|
||||
allowed_users: ${TELEGRAM_ALLOWED_USERS}
|
||||
# auto_approve_tools: true
|
||||
# Chat bots cannot show CLI approval prompts; safe tools auto-approved by default.
|
||||
routes:
|
||||
default: content
|
||||
|
||||
# Legacy routing block
|
||||
routing:
|
||||
default: assistant
|
||||
@@ -343,6 +343,7 @@ _LAZY_IMPORTS = {
|
||||
# UI (optional)
|
||||
'AGUI': ('praisonaiagents.ui.agui', 'AGUI'),
|
||||
'A2A': ('praisonaiagents.ui.a2a', 'A2A'),
|
||||
'A2UI': ('praisonaiagents.ui.a2ui', 'A2UI'),
|
||||
|
||||
# Feature configs
|
||||
'MemoryConfig': ('praisonaiagents.config.feature_configs', 'MemoryConfig'),
|
||||
@@ -814,6 +815,6 @@ def __dir__():
|
||||
# - Workflow, Task, Route, Parallel, Loop, Repeat, Pipeline, etc.
|
||||
# - MCP, FlowDisplay, track_workflow, FastContext, etc.
|
||||
# - Plan, PlanStep, TodoList, PlanningAgent, ApprovalCallback, etc.
|
||||
# - RAG, RAGConfig, RAGResult, AGUI, A2A, etc.
|
||||
# - RAG, RAGConfig, RAGResult, AGUI, A2A, A2UI, etc.
|
||||
# - All telemetry, display, and utility functions
|
||||
# ============================================================================
|
||||
@@ -1952,6 +1952,129 @@ Your Goal: {self.goal}
|
||||
# Sandbox configuration - initialize SandboxMixin
|
||||
super().__init__(sandbox=sandbox)
|
||||
|
||||
def __deepcopy__(self, memo: dict) -> "Agent":
|
||||
"""Custom deepcopy that creates fresh threading primitives.
|
||||
|
||||
threading.RLock (self.__cache_lock) and threading.Lock (self._cost_lock)
|
||||
cannot be deep-copied on CPython < 3.13. This hook deep-copies every
|
||||
other attribute normally and replaces the locks with new instances so
|
||||
that copy.deepcopy(agent) works in any Python version.
|
||||
"""
|
||||
import copy
|
||||
cls = self.__class__
|
||||
result = cls.__new__(cls)
|
||||
memo[id(self)] = result
|
||||
for k, v in self.__dict__.items():
|
||||
if k in ("_Agent__cache_lock",):
|
||||
object.__setattr__(result, k, threading.RLock())
|
||||
elif k == "_cost_lock":
|
||||
object.__setattr__(result, k, threading.Lock())
|
||||
else:
|
||||
object.__setattr__(result, k, copy.deepcopy(v, memo))
|
||||
return result
|
||||
|
||||
def clone_for_channel(self) -> "Agent":
|
||||
"""Return a fully independent copy of this agent for a gateway channel.
|
||||
|
||||
This method safely clones an agent for use in multi-channel scenarios
|
||||
by avoiding deepcopy issues with tools/handoffs and creating fresh
|
||||
instances of mutable objects like interrupt_controller.
|
||||
|
||||
Returns:
|
||||
Agent: A new Agent instance with copied configuration but fresh
|
||||
locks, interrupt controller, and shallow-copied tools.
|
||||
"""
|
||||
import copy
|
||||
from ..tools.base import BaseTool
|
||||
|
||||
# Build kwargs for new Agent instance, copying core attributes
|
||||
clone_kwargs = {
|
||||
# Core identity
|
||||
'name': self.name,
|
||||
'role': self.role,
|
||||
'goal': self.goal,
|
||||
'backstory': self.backstory,
|
||||
'instructions': self.instructions,
|
||||
|
||||
# LLM configuration
|
||||
'llm': self.llm,
|
||||
'base_url': getattr(self, 'base_url', None),
|
||||
'api_key': getattr(self, 'api_key', None),
|
||||
'auth': getattr(self, 'auth', None),
|
||||
|
||||
# Shallow copy tools to avoid deepcopy issues with nested objects
|
||||
'tools': list(self.tools) if self.tools else None,
|
||||
|
||||
# Skip handoffs entirely - they shouldn't be shared across channels
|
||||
# and can contain nested Agent instances that cause RLock issues
|
||||
'handoffs': None,
|
||||
|
||||
# Feature configurations - check for actual stored config objects
|
||||
'memory': getattr(self, '_memory_config', None),
|
||||
'knowledge': getattr(self, '_knowledge_config', None),
|
||||
'planning': getattr(self, '_planning_config', None),
|
||||
'reflection': getattr(self, '_reflection_config', None),
|
||||
'guardrails': getattr(self, '_guardrails_config', None),
|
||||
'web': getattr(self, '_web_config', None),
|
||||
'context': getattr(self, '_context_config', None),
|
||||
'autonomy': getattr(self, '_autonomy_config', None),
|
||||
'output': getattr(self, '_output_config', None),
|
||||
'execution': getattr(self, '_execution_config', None),
|
||||
'templates': getattr(self, '_template_config', None),
|
||||
'caching': getattr(self, '_caching_config', None),
|
||||
'hooks': getattr(self, '_hooks_config', None),
|
||||
'skills': getattr(self, '_skills_config', None),
|
||||
'approval': getattr(self, '_approval_config', None),
|
||||
'learn': getattr(self, '_learn_config', None),
|
||||
|
||||
# Tool configuration
|
||||
'tool_timeout': getattr(self, '_tool_timeout', None),
|
||||
'parallel_tool_calls': getattr(self, 'parallel_tool_calls', False),
|
||||
|
||||
# CLI backend
|
||||
'cli_backend': getattr(self, '_cli_backend', None),
|
||||
|
||||
# Create fresh interrupt controller to avoid shared state
|
||||
'interrupt_controller': None, # Let new instance create its own
|
||||
|
||||
# Sandbox config
|
||||
'sandbox': getattr(self, '_sandbox_config', None),
|
||||
}
|
||||
|
||||
# Handle deprecated parameters for backward compatibility
|
||||
import warnings
|
||||
|
||||
# Check for deprecated standalone attributes and emit warnings
|
||||
deprecated_attrs = [
|
||||
('allow_code_execution', 'execution'),
|
||||
('code_execution_mode', 'execution'),
|
||||
('auto_save', 'memory'),
|
||||
('rate_limiter', 'execution'),
|
||||
('allow_delegation', 'handoffs'),
|
||||
('verification_hooks', 'autonomy')
|
||||
]
|
||||
|
||||
for old_attr, new_param in deprecated_attrs:
|
||||
if hasattr(self, old_attr):
|
||||
value = getattr(self, old_attr)
|
||||
if value is not None and value != False: # Skip None/False defaults
|
||||
warnings.warn(
|
||||
f"Deprecated attribute '{old_attr}' found in agent clone. "
|
||||
f"Use '{new_param}=' parameter instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2
|
||||
)
|
||||
# For basic compatibility, include in clone_kwargs if not already set
|
||||
if old_attr == 'allow_code_execution' and clone_kwargs['execution'] is None:
|
||||
from ..config.execution import ExecutionConfig
|
||||
clone_kwargs['execution'] = ExecutionConfig(code_execution=value)
|
||||
elif old_attr == 'auto_save' and clone_kwargs['memory'] is None:
|
||||
from ..config.memory import MemoryConfig
|
||||
clone_kwargs['memory'] = MemoryConfig(auto_save=value)
|
||||
|
||||
# Create new Agent instance
|
||||
return self.__class__(**{k: v for k, v in clone_kwargs.items() if v is not None})
|
||||
|
||||
@property
|
||||
def _telemetry(self):
|
||||
"""Lazy-loaded telemetry instance for performance."""
|
||||
@@ -4616,60 +4739,57 @@ Answer:"""
|
||||
error=f"Agent guardrail validation error: {str(e)}"
|
||||
)
|
||||
|
||||
def _apply_guardrail_with_retry(self, response_text, prompt, temperature=1.0, tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Apply guardrail validation with retry logic.
|
||||
|
||||
Args:
|
||||
response_text: The response to validate
|
||||
prompt: Original prompt for regeneration if needed
|
||||
temperature: Temperature for regeneration
|
||||
tools: Tools for regeneration
|
||||
|
||||
Returns:
|
||||
str: The validated response text or None if validation fails after retries
|
||||
"""
|
||||
def _validate_with_guardrail(self, response_text):
|
||||
"""Validate response with guardrail. Returns (success, result, error)."""
|
||||
if not self._guardrail_fn:
|
||||
return response_text
|
||||
return True, response_text, None
|
||||
|
||||
from ..main import TaskOutput
|
||||
|
||||
task_output = TaskOutput(
|
||||
description="Agent response output",
|
||||
raw=response_text,
|
||||
agent=self.name
|
||||
)
|
||||
|
||||
guardrail_result = self._process_guardrail(task_output)
|
||||
|
||||
if guardrail_result.success:
|
||||
# Return the potentially modified result
|
||||
if guardrail_result.result and hasattr(guardrail_result.result, 'raw'):
|
||||
return True, guardrail_result.result.raw, None
|
||||
elif guardrail_result.result:
|
||||
return True, str(guardrail_result.result), None
|
||||
else:
|
||||
return True, response_text, None
|
||||
else:
|
||||
return False, None, guardrail_result.error
|
||||
|
||||
def _apply_guardrail_with_retry(self, response_text, prompt, temperature=1.0, tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Apply guardrail validation with retry logic (sync version)."""
|
||||
retry_count = 0
|
||||
current_response = response_text
|
||||
|
||||
while retry_count <= self.max_guardrail_retries:
|
||||
# Create TaskOutput object
|
||||
task_output = TaskOutput(
|
||||
description="Agent response output",
|
||||
raw=current_response,
|
||||
agent=self.name
|
||||
)
|
||||
success, result, error = self._validate_with_guardrail(current_response)
|
||||
|
||||
# Process guardrail
|
||||
guardrail_result = self._process_guardrail(task_output)
|
||||
|
||||
if guardrail_result.success:
|
||||
if success:
|
||||
logging.info(f"Agent {self.name}: Guardrail validation passed")
|
||||
# Return the potentially modified result
|
||||
if guardrail_result.result and hasattr(guardrail_result.result, 'raw'):
|
||||
return guardrail_result.result.raw
|
||||
elif guardrail_result.result:
|
||||
return str(guardrail_result.result)
|
||||
else:
|
||||
return current_response
|
||||
return result
|
||||
|
||||
# Guardrail failed
|
||||
if retry_count >= self.max_guardrail_retries:
|
||||
raise Exception(
|
||||
f"Agent {self.name} response failed guardrail validation after {self.max_guardrail_retries} retries. "
|
||||
f"Last error: {guardrail_result.error}"
|
||||
f"Last error: {error}"
|
||||
)
|
||||
|
||||
retry_count += 1
|
||||
logging.warning(f"Agent {self.name}: Guardrail validation failed (retry {retry_count}/{self.max_guardrail_retries}): {guardrail_result.error}")
|
||||
logging.warning(f"Agent {self.name}: Guardrail validation failed (retry {retry_count}/{self.max_guardrail_retries}): {error}")
|
||||
|
||||
# Regenerate response for retry
|
||||
try:
|
||||
retry_prompt = f"{prompt}\n\nNote: Previous response failed validation due to: {guardrail_result.error}. Please provide an improved response."
|
||||
retry_prompt = f"{prompt}\n\nNote: Previous response failed validation due to: {error}. Please provide an improved response."
|
||||
response = self._chat_completion([{"role": "user", "content": retry_prompt}], temperature, tools, task_name=task_name, task_description=task_description, task_id=task_id)
|
||||
if response and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
@@ -4678,10 +4798,46 @@ Answer:"""
|
||||
raise Exception("Failed to generate retry response")
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Error during guardrail retry: {e}")
|
||||
# If we can't regenerate, fail the guardrail
|
||||
raise Exception(f"Agent {self.name} guardrail retry failed: {e}")
|
||||
|
||||
return current_response
|
||||
|
||||
async def _aapply_guardrail_with_retry(self, response_text, prompt, temperature=1.0, tools=None, task_name=None, task_description=None, task_id=None):
|
||||
"""Apply guardrail validation with retry logic (async version)."""
|
||||
retry_count = 0
|
||||
current_response = response_text
|
||||
|
||||
while retry_count <= self.max_guardrail_retries:
|
||||
success, result, error = self._validate_with_guardrail(current_response)
|
||||
|
||||
if success:
|
||||
logging.info(f"Agent {self.name}: Guardrail validation passed")
|
||||
return result
|
||||
|
||||
# Guardrail failed
|
||||
if retry_count >= self.max_guardrail_retries:
|
||||
raise Exception(
|
||||
f"Agent {self.name} guardrail retry failed: {e}"
|
||||
f"Agent {self.name} response failed guardrail validation after {self.max_guardrail_retries} retries. "
|
||||
f"Last error: {error}"
|
||||
)
|
||||
|
||||
retry_count += 1
|
||||
logging.warning(f"Agent {self.name}: Guardrail validation failed (retry {retry_count}/{self.max_guardrail_retries}): {error}")
|
||||
|
||||
# Regenerate response for retry (async version)
|
||||
try:
|
||||
retry_prompt = f"{prompt}\n\nNote: Previous response failed validation due to: {error}. Please provide an improved response."
|
||||
response = await self._execute_unified_achat_completion([{"role": "user", "content": retry_prompt}], temperature, tools, task_name=task_name, task_description=task_description, task_id=task_id)
|
||||
if response and hasattr(response, 'choices') and response.choices:
|
||||
content = response.choices[0].message.content
|
||||
current_response = content.strip() if content else ""
|
||||
elif isinstance(response, str):
|
||||
current_response = response.strip()
|
||||
else:
|
||||
raise Exception("Failed to generate retry response")
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Error during guardrail retry: {e}")
|
||||
raise Exception(f"Agent {self.name} guardrail retry failed: {e}")
|
||||
|
||||
return current_response
|
||||
|
||||
|
||||
@@ -753,11 +753,11 @@ Your Goal: {self.goal}"""
|
||||
|
||||
# Create LLMError with contextual metadata
|
||||
error = LLMError(
|
||||
str(e),
|
||||
model_name=model_name,
|
||||
agent_id=self.name,
|
||||
str(e),
|
||||
model_name=model_name,
|
||||
agent_id=self.name,
|
||||
is_retryable=is_retryable,
|
||||
session_id=session_id
|
||||
context={"session_id": session_id},
|
||||
)
|
||||
|
||||
# Call error hook if available for error interception
|
||||
@@ -902,7 +902,7 @@ Your Goal: {self.goal}"""
|
||||
max_tokens=getattr(self, 'max_tokens', None),
|
||||
stream=stream,
|
||||
response_format=response_format,
|
||||
execute_tool_fn=getattr(self, 'execute_tool', None),
|
||||
execute_tool_fn=getattr(self, 'execute_tool_async', None),
|
||||
console=self.console if (self.verbose or stream) else None,
|
||||
display_fn=self._display_generating if self.verbose else None,
|
||||
stream_callback=stream_callback,
|
||||
@@ -935,6 +935,151 @@ Your Goal: {self.goal}"""
|
||||
logging.error(f"Unified async chat completion failed: {e}")
|
||||
raise
|
||||
|
||||
async def _finalize_unified_achat_response(
|
||||
self,
|
||||
response,
|
||||
original_prompt,
|
||||
temperature,
|
||||
tools,
|
||||
output_json,
|
||||
output_pydantic,
|
||||
reasoning_steps,
|
||||
stream,
|
||||
messages,
|
||||
chat_history_length,
|
||||
start_time,
|
||||
task_name=None,
|
||||
task_description=None,
|
||||
task_id=None,
|
||||
):
|
||||
"""Post-process unified achat completion (parity with sync _chat_impl)."""
|
||||
if not response:
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
if hasattr(response, "choices") and response.choices:
|
||||
msg = response.choices[0].message
|
||||
content = getattr(msg, "content", None)
|
||||
response_text = content.strip() if content else ""
|
||||
if (
|
||||
reasoning_steps
|
||||
and getattr(msg, "reasoning_content", None)
|
||||
):
|
||||
response_text = msg.reasoning_content.strip()
|
||||
else:
|
||||
extracted = self._extract_llm_response_content(response)
|
||||
response_text = (extracted or "").strip() if isinstance(extracted, str) else str(extracted or "")
|
||||
|
||||
if output_json or output_pydantic:
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(
|
||||
response_text, original_prompt, temperature, tools,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
self._execute_callback_and_display(
|
||||
original_prompt, validated_response, time.time() - start_time,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
return await self._atrigger_after_agent_hook(
|
||||
original_prompt, validated_response, start_time,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
if not self.self_reflect:
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(
|
||||
response_text, original_prompt, temperature, tools,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
self._execute_callback_and_display(
|
||||
original_prompt, validated_response, time.time() - start_time,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
return await self._atrigger_after_agent_hook(
|
||||
original_prompt, validated_response, start_time,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# Self-reflection (legacy OpenAI path)
|
||||
reflection_count = 0
|
||||
while True:
|
||||
reflection_prompt = f"""
|
||||
Reflect on your previous response: '{response_text}'.
|
||||
{self.reflect_prompt if self.reflect_prompt else "Identify any flaws, improvements, or actions."}
|
||||
Provide a "satisfactory" status ('yes' or 'no').
|
||||
Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"""
|
||||
reflection_messages = messages + [
|
||||
{"role": "assistant", "content": response_text},
|
||||
{"role": "user", "content": reflection_prompt},
|
||||
]
|
||||
try:
|
||||
if self._openai_client is None:
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
return await self._atrigger_after_agent_hook(
|
||||
original_prompt, response_text, start_time,
|
||||
)
|
||||
reflection_response = await self._openai_client.async_client.beta.chat.completions.parse(
|
||||
model=self.reflect_llm if self.reflect_llm else self.llm,
|
||||
messages=reflection_messages,
|
||||
temperature=temperature,
|
||||
response_format=_get_display_functions()['ReflectionOutput'],
|
||||
)
|
||||
reflection_output = reflection_response.choices[0].message.parsed
|
||||
if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
|
||||
break
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
break
|
||||
regenerate_messages = reflection_messages + [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (
|
||||
f"Self Reflection: {reflection_output.reflection} "
|
||||
f"Satisfactory?: {reflection_output.satisfactory}"
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": "Now regenerate your response using the reflection you made"},
|
||||
]
|
||||
new_response = await self._openai_client.async_client.chat.completions.create(
|
||||
model=self.llm,
|
||||
messages=regenerate_messages,
|
||||
temperature=temperature,
|
||||
)
|
||||
response_text = new_response.choices[0].message.content
|
||||
reflection_count += 1
|
||||
except Exception as e:
|
||||
logging.error("Reflection parsing failed.", exc_info=True)
|
||||
reflection_count += 1
|
||||
if reflection_count >= self.max_reflect:
|
||||
break
|
||||
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(
|
||||
response_text, original_prompt, temperature, tools,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
self._execute_callback_and_display(
|
||||
original_prompt, validated_response, time.time() - start_time,
|
||||
task_name, task_description, task_id,
|
||||
)
|
||||
return await self._atrigger_after_agent_hook(
|
||||
original_prompt, validated_response, start_time,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for OpenAI client: {e}")
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
def _execute_callback_and_display(self, prompt: str, response: str, generation_time: float, task_name=None, task_description=None, task_id=None):
|
||||
"""Helper method to execute callbacks and display interaction.
|
||||
|
||||
@@ -1934,7 +2079,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
# Apply guardrail validation for custom LLM response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id)
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(normalized_content, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(prompt, validated_response, start_time)
|
||||
@@ -2042,8 +2187,216 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
task_id=task_id,
|
||||
response_format=response_format
|
||||
)
|
||||
# Continue to handle structured outputs and other processing
|
||||
# Don't return immediately - fall through to existing logic
|
||||
|
||||
# Process response - mirror sync _chat_impl behavior (lines ~1544-1602)
|
||||
if not response:
|
||||
# Rollback chat history on response failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# Extract response content using the same method as sync
|
||||
response_text = self._extract_llm_response_content(response)
|
||||
if isinstance(response_text, str):
|
||||
response_text = response_text.strip()
|
||||
|
||||
# Handle output_json or output_pydantic if specified
|
||||
if output_json or output_pydantic:
|
||||
# Add to chat history and return raw response
|
||||
# User message already added before LLM call via _build_messages
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
# Apply guardrail validation even for JSON output
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for JSON output: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# For regular responses (no self-reflection)
|
||||
if not self.self_reflect:
|
||||
# User message already added before LLM call via _build_messages
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB (non-reflect path)
|
||||
self._persist_message("assistant", response_text)
|
||||
if self.verbose:
|
||||
logging.debug(f"Agent {self.name} final response: {response_text}")
|
||||
# Return only reasoning content if reasoning_steps is True
|
||||
if reasoning_steps and hasattr(response, 'choices') and response.choices and hasattr(response.choices[0].message, 'reasoning_content') and response.choices[0].message.reasoning_content:
|
||||
# Apply guardrail to reasoning content
|
||||
try:
|
||||
validated_reasoning = await self._aapply_guardrail_with_retry(response.choices[0].message.reasoning_content, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_reasoning, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_reasoning, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed for reasoning content: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
else:
|
||||
# Apply guardrail to regular response content
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# If self-reflection is enabled, implement reflection logic
|
||||
if self.self_reflect:
|
||||
# Implement async self-reflection similar to sync path
|
||||
reflection_prompt = f"""
|
||||
Reflect on your previous response: '{response_text}'.
|
||||
{self.reflect_prompt if self.reflect_prompt else "Identify any flaws, improvements, or actions."}
|
||||
Provide a "satisfactory" status ('yes' or 'no').
|
||||
Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
"""
|
||||
|
||||
# Add reflection prompt to messages
|
||||
reflection_messages = messages + [
|
||||
{"role": "assistant", "content": response_text},
|
||||
{"role": "user", "content": reflection_prompt}
|
||||
]
|
||||
|
||||
reflection_count = 0
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Check if OpenAI client is available for self-reflection
|
||||
if self._openai_client is None:
|
||||
# For custom LLMs, self-reflection with structured output is not supported
|
||||
if self.verbose:
|
||||
_get_display_functions()['display_self_reflection'](f"Agent {self.name}: Self-reflection with structured output is not supported for custom LLM providers. Skipping reflection.", console=self.console)
|
||||
# Return the original response without reflection
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
reflection_response = await self._openai_client.async_client.beta.chat.completions.parse(
|
||||
model=self.reflect_llm if self.reflect_llm else self.llm,
|
||||
messages=reflection_messages,
|
||||
temperature=temperature,
|
||||
response_format=_get_display_functions()['ReflectionOutput']
|
||||
)
|
||||
|
||||
reflection_output = reflection_response.choices[0].message.parsed
|
||||
|
||||
if self.verbose:
|
||||
_get_display_functions()['display_self_reflection'](f"Agent {self.name} self reflection (using {self.reflect_llm if self.reflect_llm else self.llm}): reflection='{reflection_output.reflection}' satisfactory='{reflection_output.satisfactory}'", console=self.console)
|
||||
|
||||
# Only consider satisfactory after minimum reflections
|
||||
if reflection_output.satisfactory == "yes" and reflection_count >= self.min_reflect - 1:
|
||||
if self.verbose:
|
||||
_get_display_functions()['display_self_reflection']("Agent marked the response as satisfactory after meeting minimum reflections", console=self.console)
|
||||
# Add to chat history and return
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# Check if we've hit max reflections
|
||||
if reflection_count >= self.max_reflect - 1:
|
||||
if self.verbose:
|
||||
_get_display_functions()['display_self_reflection']("Maximum reflection count reached, returning current response", console=self.console)
|
||||
# Add to chat history and return
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after max reflections: {e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
|
||||
# Regenerate response based on reflection
|
||||
regenerate_messages = reflection_messages + [
|
||||
{"role": "assistant", "content": f"Self Reflection: {reflection_output.reflection} Satisfactory?: {reflection_output.satisfactory}"},
|
||||
{"role": "user", "content": "Now regenerate your response using the reflection you made"}
|
||||
]
|
||||
|
||||
new_response = await self._execute_unified_achat_completion(
|
||||
messages=regenerate_messages,
|
||||
temperature=temperature,
|
||||
tools=formatted_tools,
|
||||
stream=stream,
|
||||
reasoning_steps=reasoning_steps,
|
||||
task_name=task_name,
|
||||
task_description=task_description,
|
||||
task_id=task_id,
|
||||
response_format=response_format
|
||||
)
|
||||
|
||||
if new_response:
|
||||
new_response_text = self._extract_llm_response_content(new_response)
|
||||
if isinstance(new_response_text, str):
|
||||
response_text = new_response_text.strip()
|
||||
# Update reflection_messages to include the new response for next iteration
|
||||
reflection_messages = regenerate_messages + [
|
||||
{"role": "assistant", "content": response_text}
|
||||
]
|
||||
|
||||
reflection_count += 1
|
||||
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
_get_display_functions()['display_error'](f"Error in parsing self-reflection json {e}. Retrying", console=self.console)
|
||||
logging.error("Reflection parsing failed.", exc_info=True)
|
||||
reflection_count += 1
|
||||
if reflection_count >= self.max_reflect:
|
||||
# Return original response after max reflection attempts
|
||||
self._append_to_chat_history({"role": "assistant", "content": response_text})
|
||||
# Persist assistant message to DB
|
||||
self._persist_message("assistant", response_text)
|
||||
try:
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
except Exception as guard_e:
|
||||
logging.error(f"Agent {self.name}: Guardrail validation failed after reflection error: {guard_e}")
|
||||
# Rollback chat history on guardrail failure
|
||||
self._truncate_chat_history(chat_history_length)
|
||||
return None
|
||||
continue
|
||||
|
||||
# This should never be reached due to the returns above
|
||||
# But adding as safety fallback
|
||||
logging.warning(f"Agent {self.name}: Unexpected code path reached in unified dispatch")
|
||||
return None
|
||||
else:
|
||||
# LEGACY: Check if OpenAI client is available
|
||||
if self._openai_client is None:
|
||||
@@ -2180,7 +2533,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
# Apply guardrail validation for OpenAI client response
|
||||
try:
|
||||
validated_response = self._apply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
validated_response = await self._aapply_guardrail_with_retry(response_text, original_prompt, temperature, tools, task_name, task_description, task_id)
|
||||
# Execute callback after validation
|
||||
self._execute_callback_and_display(original_prompt, validated_response, time.time() - start_time, task_name, task_description, task_id)
|
||||
return await self._atrigger_after_agent_hook(original_prompt, validated_response, start_time)
|
||||
|
||||
@@ -403,9 +403,33 @@ class MemoryMixin:
|
||||
self._session_store.add_user_message(self._session_id, content)
|
||||
elif role == "assistant":
|
||||
self._session_store.add_assistant_message(self._session_id, content)
|
||||
self._persist_session_stats()
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to persist message to session store: {e}")
|
||||
|
||||
def _persist_session_stats(self):
|
||||
"""Flush agent cost/token stats into session JSON metadata."""
|
||||
store = getattr(self, "_session_store", None)
|
||||
session_id = getattr(self, "_session_id", None)
|
||||
if store is None or session_id is None:
|
||||
return
|
||||
if not hasattr(store, "update_session_metadata"):
|
||||
return
|
||||
try:
|
||||
summary = getattr(self, "cost_summary", {}) or {}
|
||||
llm = getattr(self, "llm", None)
|
||||
model = str(llm) if llm is not None else None
|
||||
store.update_session_metadata(
|
||||
session_id,
|
||||
model=model,
|
||||
total_tokens=summary.get("tokens_in", 0) + summary.get("tokens_out", 0),
|
||||
cost=summary.get("cost", 0),
|
||||
source=getattr(self, "_session_source", "chat"),
|
||||
agent_id=getattr(self, "agent_id", None) or getattr(self, "_agent_id", None),
|
||||
)
|
||||
except Exception as e:
|
||||
logging.debug(f"Session stats persist skipped: {e}")
|
||||
|
||||
def session_id(self) -> Optional[str]:
|
||||
"""Get the current session ID."""
|
||||
return self._session_id
|
||||
|
||||
@@ -26,6 +26,21 @@ if TYPE_CHECKING:
|
||||
class ToolExecutionMixin:
|
||||
"""Mixin providing toolexecution methods for the Agent class."""
|
||||
|
||||
def _get_existing_stream_emitter(self):
|
||||
"""Return an already-initialized stream emitter without creating one."""
|
||||
emitter = getattr(self, "_stream_emitter", None)
|
||||
if emitter is not None:
|
||||
return emitter
|
||||
|
||||
# Support name-mangled private attributes across class renames/inheritance.
|
||||
for cls in type(self).mro():
|
||||
mangled = f"_{cls.__name__}__stream_emitter"
|
||||
if hasattr(self, mangled):
|
||||
emitter = getattr(self, mangled, None)
|
||||
if emitter is not None:
|
||||
return emitter
|
||||
return None
|
||||
|
||||
def _resolve_tool_names(self, tool_names):
|
||||
"""Resolve tool names to actual tool instances from registry.
|
||||
|
||||
@@ -160,8 +175,9 @@ class ToolExecutionMixin:
|
||||
|
||||
# Emit TOOL_CALL_START to stream_emitter (for AIUI/AG-UI consumers)
|
||||
# Zero overhead when no callbacks registered
|
||||
if hasattr(self, '_Agent__stream_emitter') and getattr(self, "_Agent__stream_emitter", None) is not None and getattr(self, "_Agent__stream_emitter", None).has_callbacks:
|
||||
getattr(self, "_Agent__stream_emitter", None).emit(StreamEvent(
|
||||
_stream_emitter = self._get_existing_stream_emitter()
|
||||
if _stream_emitter is not None and _stream_emitter.has_callbacks:
|
||||
_stream_emitter.emit(StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_START,
|
||||
timestamp=_tool_start_perf,
|
||||
tool_call={
|
||||
@@ -275,10 +291,10 @@ class ToolExecutionMixin:
|
||||
|
||||
# Emit TOOL_CALL_RESULT to stream_emitter (for AIUI/AG-UI consumers)
|
||||
# Zero overhead when no callbacks registered
|
||||
if hasattr(self, '_Agent__stream_emitter') and getattr(self, "_Agent__stream_emitter", None) is not None and getattr(self, "_Agent__stream_emitter", None).has_callbacks:
|
||||
if _stream_emitter is not None and _stream_emitter.has_callbacks:
|
||||
# Truncate result for stream event (keep it reasonable for UI display)
|
||||
result_summary = str(result)[:500] if result else None
|
||||
getattr(self, "_Agent__stream_emitter", None).emit(StreamEvent(
|
||||
_stream_emitter.emit(StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_RESULT,
|
||||
timestamp=_time.perf_counter(),
|
||||
tool_call={
|
||||
@@ -290,6 +306,18 @@ class ToolExecutionMixin:
|
||||
agent_id=self.name,
|
||||
metadata={"duration_ms": _duration_ms},
|
||||
))
|
||||
_stream_emitter.emit(StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_END,
|
||||
timestamp=_time.perf_counter(),
|
||||
tool_call={
|
||||
"name": function_name,
|
||||
"arguments": arguments,
|
||||
"result": result_summary,
|
||||
"id": tool_call_id,
|
||||
},
|
||||
agent_id=self.name,
|
||||
metadata={"duration_ms": _duration_ms},
|
||||
))
|
||||
|
||||
# Trigger AFTER_TOOL hook
|
||||
from ..hooks import HookEvent, AfterToolInput
|
||||
@@ -926,4 +954,3 @@ class ToolExecutionMixin:
|
||||
def pending_approval_count(self) -> int:
|
||||
"""Number of approval requests still waiting."""
|
||||
return len(self._pending_approvals)
|
||||
|
||||
@@ -15,9 +15,9 @@ from enum import Enum
|
||||
|
||||
# Import token tracking
|
||||
try:
|
||||
from ..telemetry.token_collector import _token_collector
|
||||
from ..telemetry.token_collector import get_token_collector
|
||||
except ImportError:
|
||||
_token_collector = None
|
||||
get_token_collector = None
|
||||
|
||||
# Import async utility for hot-path usage
|
||||
try:
|
||||
@@ -1609,25 +1609,25 @@ class AgentTeam:
|
||||
original_on_task_start = self.on_task_start
|
||||
original_on_task_complete = self.on_task_complete
|
||||
|
||||
def composed_on_task_start(task):
|
||||
def composed_on_task_start(task, task_id):
|
||||
try:
|
||||
verbose_task_start_callback(task)
|
||||
verbose_task_start_callback(task, task_id)
|
||||
except Exception as e:
|
||||
logging.debug(f"Error in verbose task start callback: {e}")
|
||||
if original_on_task_start:
|
||||
try:
|
||||
original_on_task_start(task)
|
||||
original_on_task_start(task, task_id)
|
||||
except Exception as e:
|
||||
logging.debug(f"Error in original task start callback: {e}")
|
||||
|
||||
def composed_on_task_complete(task):
|
||||
def composed_on_task_complete(task, task_output):
|
||||
try:
|
||||
verbose_task_complete_callback(task)
|
||||
verbose_task_complete_callback(task, task_output)
|
||||
except Exception as e:
|
||||
logging.debug(f"Error in verbose task complete callback: {e}")
|
||||
if original_on_task_complete:
|
||||
try:
|
||||
original_on_task_complete(task)
|
||||
original_on_task_complete(task, task_output)
|
||||
except Exception as e:
|
||||
logging.debug(f"Error in original task complete callback: {e}")
|
||||
|
||||
@@ -1826,18 +1826,19 @@ class AgentTeam:
|
||||
|
||||
def get_token_usage_summary(self) -> Dict[str, Any]:
|
||||
"""Get a summary of token usage across all agents and tasks."""
|
||||
if not _token_collector:
|
||||
if not get_token_collector:
|
||||
return {"error": "Token tracking not available"}
|
||||
|
||||
return _token_collector.get_session_summary()
|
||||
return get_token_collector().get_session_summary()
|
||||
|
||||
def get_detailed_token_report(self) -> Dict[str, Any]:
|
||||
"""Get a detailed token usage report."""
|
||||
if not _token_collector:
|
||||
if not get_token_collector:
|
||||
return {"error": "Token tracking not available"}
|
||||
|
||||
summary = _token_collector.get_session_summary()
|
||||
recent = _token_collector.get_recent_interactions(limit=20)
|
||||
collector = get_token_collector()
|
||||
summary = collector.get_session_summary()
|
||||
recent = collector.get_recent_interactions(limit=20)
|
||||
|
||||
# Calculate cost estimates (example rates)
|
||||
cost_per_1k_input = 0.0005 # $0.0005 per 1K input tokens
|
||||
@@ -1861,11 +1862,11 @@ class AgentTeam:
|
||||
|
||||
def display_token_usage(self):
|
||||
"""Display token usage in a formatted table."""
|
||||
if not _token_collector:
|
||||
if not get_token_collector:
|
||||
print("Token tracking not available")
|
||||
return
|
||||
|
||||
summary = _token_collector.get_session_summary()
|
||||
summary = get_token_collector().get_session_summary()
|
||||
|
||||
print("\n" + "="*50)
|
||||
print("TOKEN USAGE SUMMARY")
|
||||
|
||||
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
Tool whitelist filter module with diagnostics.
|
||||
|
||||
This module provides a canonical implementation of the ALLOWED_TOOLS filter
|
||||
that prevents tool name collisions in multi-environment agent systems by
|
||||
whitelisting only specified tools.
|
||||
|
||||
The filter is designed to solve the "tool shadowing" problem where multiple
|
||||
modules register tools with overlapping or ambiguous names, causing agents
|
||||
to invoke the wrong implementation.
|
||||
|
||||
Environment Variables:
|
||||
ALLOWED_TOOLS: Primary variable (comma-separated tool names)
|
||||
HERMES_ONLY_TOOLS: Backward compatibility alias for ALLOWED_TOOLS
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
# Create filter from environment variable
|
||||
filter = AllowedToolsFilter()
|
||||
|
||||
# Filter tools
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
filtered_tools = filter.filter_tools(available_tools)
|
||||
|
||||
# Get diagnostics
|
||||
filter.log_diagnostics()
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Set, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AllowedToolsFilter:
|
||||
"""
|
||||
Canonical tool whitelist filter with diagnostics and backward compatibility.
|
||||
|
||||
This filter implements the semantics for tool filtering:
|
||||
- ALLOWED_TOOLS unset (or HERMES_ONLY_TOOLS): all tools visible (with warning about collisions)
|
||||
- ALLOWED_TOOLS empty string: error - must specify tools or unset
|
||||
- ALLOWED_TOOLS with values: only whitelisted tools visible
|
||||
- Unknown tools in list: warn and strip in dev, strict fail in CI
|
||||
|
||||
Backward Compatibility:
|
||||
- Supports both ALLOWED_TOOLS and HERMES_ONLY_TOOLS
|
||||
- ALLOWED_TOOLS takes precedence if both are set
|
||||
"""
|
||||
|
||||
def __init__(self, env_var_name: str = "ALLOWED_TOOLS"):
|
||||
"""
|
||||
Initialize the filter with backward compatibility support.
|
||||
|
||||
Args:
|
||||
env_var_name: Primary environment variable name (default: "ALLOWED_TOOLS")
|
||||
"""
|
||||
# Support both new (ALLOWED_TOOLS) and legacy (HERMES_ONLY_TOOLS) naming
|
||||
self.primary_var = env_var_name
|
||||
self.legacy_var = "HERMES_ONLY_TOOLS"
|
||||
|
||||
# ALLOWED_TOOLS takes precedence over HERMES_ONLY_TOOLS for backward compatibility
|
||||
self.env_value = os.environ.get(self.primary_var) or os.environ.get(self.legacy_var)
|
||||
self.env_var_name = self.primary_var if os.environ.get(self.primary_var) else self.legacy_var
|
||||
|
||||
self.is_ci = os.environ.get("CI", "").lower() in ("true", "1", "yes")
|
||||
self._whitelist: Optional[Set[str]] = self._parse_whitelist()
|
||||
self._diagnostics: Dict[str, Any] = {}
|
||||
|
||||
def _parse_whitelist(self) -> Optional[Set[str]]:
|
||||
"""
|
||||
Parse the environment variable value into a whitelist set.
|
||||
|
||||
Returns:
|
||||
None if unset, Set of tool names if set, raises ValueError for empty string
|
||||
"""
|
||||
if self.env_value is None:
|
||||
return None
|
||||
|
||||
# Empty string is an error - must be explicit
|
||||
if self.env_value.strip() == "":
|
||||
raise ValueError(
|
||||
f"{self.env_var_name} cannot be empty. Either unset it for all tools "
|
||||
"or provide a comma-separated list of tool names."
|
||||
)
|
||||
|
||||
# Parse comma-separated values, strip whitespace
|
||||
tools = {name.strip() for name in self.env_value.split(",") if name.strip()}
|
||||
|
||||
if not tools:
|
||||
raise ValueError(
|
||||
f"{self.env_var_name} contains no valid tool names. "
|
||||
"Provide comma-separated tool names or unset the variable."
|
||||
)
|
||||
|
||||
return tools
|
||||
|
||||
def filter_tools(self, available_tools: Union[Set[str], List[str], Dict[str, Any]]) -> Set[str]:
|
||||
"""
|
||||
Filter tools based on ALLOWED_TOOLS whitelist.
|
||||
|
||||
Args:
|
||||
available_tools: Available tools as set, list, or dict with tool names as keys
|
||||
|
||||
Returns:
|
||||
Filtered set of tool names
|
||||
|
||||
Raises:
|
||||
ValueError: If unknown tools in CI mode or configuration errors
|
||||
"""
|
||||
# Convert input to set of tool names
|
||||
if isinstance(available_tools, dict):
|
||||
tool_names = set(available_tools.keys())
|
||||
else:
|
||||
tool_names = set(available_tools)
|
||||
|
||||
# Store for diagnostics
|
||||
self._diagnostics["registered_before_filter"] = sorted(tool_names)
|
||||
|
||||
# If no whitelist, return all tools (with warning)
|
||||
if self._whitelist is None:
|
||||
logger.warning(
|
||||
"Tool whitelist is unset. All %d tools are visible. "
|
||||
"Consider using ALLOWED_TOOLS to prevent tool name collisions.",
|
||||
len(tool_names)
|
||||
)
|
||||
self._diagnostics["registered_after_filter"] = sorted(tool_names)
|
||||
self._diagnostics["dropped_tools"] = []
|
||||
self._diagnostics["unknown_tools"] = []
|
||||
return tool_names
|
||||
|
||||
# Find intersection and unknown tools
|
||||
available_whitelisted = tool_names.intersection(self._whitelist)
|
||||
unknown_tools = self._whitelist.difference(tool_names)
|
||||
dropped_tools = tool_names.difference(self._whitelist)
|
||||
|
||||
# Store for diagnostics
|
||||
self._diagnostics["registered_after_filter"] = sorted(available_whitelisted)
|
||||
self._diagnostics["dropped_tools"] = sorted(dropped_tools)
|
||||
self._diagnostics["unknown_tools"] = sorted(unknown_tools)
|
||||
|
||||
# Handle unknown tools based on environment
|
||||
if unknown_tools:
|
||||
unknown_list = sorted(unknown_tools)
|
||||
if self.is_ci:
|
||||
raise ValueError(
|
||||
f"Unknown tools in {self.env_var_name}: {unknown_list}. "
|
||||
"All specified tools must be available in CI mode."
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Unknown tools in %s will be ignored: %s",
|
||||
self.env_var_name,
|
||||
unknown_list
|
||||
)
|
||||
|
||||
if dropped_tools:
|
||||
logger.info(
|
||||
"Dropped %d tools not in %s whitelist: %s",
|
||||
len(dropped_tools),
|
||||
self.env_var_name,
|
||||
sorted(dropped_tools)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"%s filtered tools: %d available, %d after filter",
|
||||
self.env_var_name,
|
||||
len(tool_names),
|
||||
len(available_whitelisted)
|
||||
)
|
||||
|
||||
return available_whitelisted
|
||||
|
||||
def log_diagnostics(self) -> None:
|
||||
"""Log startup diagnostics section as specified in the issue."""
|
||||
if not self._diagnostics:
|
||||
logger.warning("No diagnostics available. Call filter_tools() first.")
|
||||
return
|
||||
|
||||
logger.info("=" * 50)
|
||||
logger.info("ALLOWED_TOOLS FILTER DIAGNOSTICS")
|
||||
logger.info("=" * 50)
|
||||
logger.info("%s=%s", self.env_var_name, self.env_value or "<unset>")
|
||||
logger.info("RegisteredBeforeFilter=%s", self._diagnostics.get("registered_before_filter", []))
|
||||
logger.info("RegisteredAfterFilter=%s", self._diagnostics.get("registered_after_filter", []))
|
||||
logger.info("DroppedTools=%s", self._diagnostics.get("dropped_tools", []))
|
||||
logger.info("UnknownTools=%s", self._diagnostics.get("unknown_tools", []))
|
||||
logger.info("IsCI=%s", self.is_ci)
|
||||
logger.info("=" * 50)
|
||||
|
||||
def get_diagnostics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get diagnostics data as a dictionary.
|
||||
|
||||
Returns:
|
||||
Dictionary with diagnostics information
|
||||
"""
|
||||
return {
|
||||
"env_var_name": self.env_var_name,
|
||||
"env_value": self.env_value,
|
||||
"is_ci": self.is_ci,
|
||||
"whitelist": list(self._whitelist) if self._whitelist else None,
|
||||
**self._diagnostics
|
||||
}
|
||||
|
||||
def is_enabled(self) -> bool:
|
||||
"""
|
||||
Check if filtering is enabled.
|
||||
|
||||
Returns:
|
||||
True if tool whitelist is set and filtering is active
|
||||
"""
|
||||
return self._whitelist is not None
|
||||
|
||||
def get_whitelist(self) -> Optional[Set[str]]:
|
||||
"""
|
||||
Get the current whitelist.
|
||||
|
||||
Returns:
|
||||
Set of whitelisted tool names, or None if not set
|
||||
"""
|
||||
return self._whitelist.copy() if self._whitelist else None
|
||||
|
||||
|
||||
def filter_tools_with_allowed_tools(
|
||||
available_tools: Union[Set[str], List[str], Dict[str, Any]],
|
||||
env_var_name: str = "ALLOWED_TOOLS",
|
||||
log_diagnostics: bool = True
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Convenience function to filter tools with ALLOWED_TOOLS.
|
||||
|
||||
Args:
|
||||
available_tools: Available tools to filter
|
||||
env_var_name: Environment variable name (default: "ALLOWED_TOOLS")
|
||||
log_diagnostics: Whether to log diagnostics (default: True)
|
||||
|
||||
Returns:
|
||||
Filtered set of tool names
|
||||
"""
|
||||
filter_instance = AllowedToolsFilter(env_var_name)
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
if log_diagnostics:
|
||||
filter_instance.log_diagnostics()
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
# Backward compatibility aliases
|
||||
filter_tools_with_hermes = filter_tools_with_allowed_tools
|
||||
hermes_filter = filter_tools_with_allowed_tools
|
||||
apply_hermes_filter = filter_tools_with_allowed_tools
|
||||
|
||||
# Legacy class alias for backward compatibility
|
||||
HermesToolFilter = AllowedToolsFilter
|
||||
@@ -42,6 +42,7 @@ class MemoryBackend(str, Enum):
|
||||
FILE = "file"
|
||||
SQLITE = "sqlite"
|
||||
REDIS = "redis"
|
||||
VALKEY = "valkey"
|
||||
POSTGRES = "postgres"
|
||||
MEM0 = "mem0"
|
||||
MONGODB = "mongodb"
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Hooks query protocol for UI/API consumers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class HooksQueryProtocol(Protocol):
|
||||
"""Read path for registered hooks."""
|
||||
|
||||
def list_hooks_for_api(self) -> List[Dict]:
|
||||
...
|
||||
@@ -439,3 +439,13 @@ def has_hook(event: Union[str, HookEvent]) -> bool:
|
||||
)
|
||||
|
||||
return get_default_registry().has_hooks(event)
|
||||
|
||||
|
||||
def list_hooks_for_api() -> List[Dict]:
|
||||
"""Flat hook list suitable for REST/UI consumers."""
|
||||
registry = get_default_registry()
|
||||
flat: List[Dict] = []
|
||||
for event, hooks in registry.list_hooks().items():
|
||||
for h in hooks:
|
||||
flat.append({"event": event, "source": "sdk", **h})
|
||||
return flat
|
||||
@@ -68,6 +68,14 @@ class HookEvent(str, Enum):
|
||||
SCHEDULE_REMOVE = "schedule_remove"
|
||||
SCHEDULE_TRIGGER = "schedule_trigger"
|
||||
|
||||
# Kanban task lifecycle (wrapper dispatcher + tools emit these)
|
||||
KANBAN_TASK_CREATED = "kanban_task_created"
|
||||
KANBAN_TASK_CLAIMED = "kanban_task_claimed"
|
||||
KANBAN_TASK_MOVED = "kanban_task_moved"
|
||||
KANBAN_TASK_DONE = "kanban_task_done"
|
||||
KANBAN_TASK_BLOCKED = "kanban_task_blocked"
|
||||
KANBAN_TASK_FAILED = "kanban_task_failed"
|
||||
|
||||
# Claude Code parity events
|
||||
USER_PROMPT_SUBMIT = "user_prompt_submit" # When user submits a prompt
|
||||
NOTIFICATION = "notification" # When notification is sent
|
||||
@@ -101,6 +109,34 @@ class HookInput:
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class KanbanHookInput(HookInput):
|
||||
"""Hook input for kanban task lifecycle events.
|
||||
|
||||
Used by wrapper dispatcher and tools when emitting kanban events.
|
||||
Observability adapters can subscribe to these via the hook registry.
|
||||
"""
|
||||
task_id: str = ""
|
||||
board: str = "default"
|
||||
status: str = ""
|
||||
assignee: Optional[str] = field(default=None)
|
||||
from_status: Optional[str] = field(default=None)
|
||||
to_status: Optional[str] = field(default=None)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
result = super().to_dict()
|
||||
result.update({
|
||||
"task_id": self.task_id,
|
||||
"board": self.board,
|
||||
"status": self.status,
|
||||
"assignee": self.assignee,
|
||||
"from_status": self.from_status,
|
||||
"to_status": self.to_status,
|
||||
})
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookOutput:
|
||||
"""Base hook output - common fields for all events."""
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
Kanban protocols and types for PraisonAI Agents.
|
||||
|
||||
This module provides the protocol contracts for kanban functionality,
|
||||
allowing the wrapper (praisonai) and PraisonAIUI to share a stable
|
||||
interface without coupling the core to SQLite or other heavy implementations.
|
||||
"""
|
||||
|
||||
from praisonaiagents.kanban.protocols import (
|
||||
KanbanStoreProtocol,
|
||||
KanbanTaskProtocol,
|
||||
KanbanCommentingProtocol,
|
||||
KanbanLinkingProtocol,
|
||||
VALID_KANBAN_STATUSES,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"KanbanStoreProtocol",
|
||||
"KanbanTaskProtocol",
|
||||
"KanbanCommentingProtocol",
|
||||
"KanbanLinkingProtocol",
|
||||
"VALID_KANBAN_STATUSES",
|
||||
]
|
||||
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
Kanban protocols for PraisonAI Agents.
|
||||
|
||||
Defines the protocol contracts that the wrapper (praisonai) and PraisonAIUI
|
||||
can use to implement kanban functionality with a stable interface.
|
||||
|
||||
This follows AGENTS.md §3.2 Protocol-First Design:
|
||||
- Protocols define WHAT (interface contract)
|
||||
- Implementations define HOW (concrete behavior)
|
||||
- Core SDK has protocols only, wrapper has heavy implementations
|
||||
"""
|
||||
|
||||
from typing import Protocol, TypedDict, runtime_checkable
|
||||
|
||||
|
||||
# Valid kanban statuses matching PraisonAIUI columns + archived
|
||||
VALID_KANBAN_STATUSES: frozenset[str] = frozenset([
|
||||
"triage",
|
||||
"todo",
|
||||
"ready",
|
||||
"running",
|
||||
"blocked",
|
||||
"review",
|
||||
"done",
|
||||
"archived"
|
||||
])
|
||||
|
||||
|
||||
class KanbanTaskProtocol(TypedDict, total=False):
|
||||
"""Typed dict shape for kanban task fields.
|
||||
|
||||
Defines the expected structure of task objects returned by
|
||||
KanbanStoreProtocol implementations.
|
||||
"""
|
||||
id: str
|
||||
title: str
|
||||
body: str
|
||||
status: str
|
||||
assignee: str | None
|
||||
priority: str | None
|
||||
tenant: str | None
|
||||
board: str
|
||||
created_at: float
|
||||
updated_at: float
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class KanbanStoreProtocol(Protocol):
|
||||
"""Protocol contract for kanban store implementations.
|
||||
|
||||
This protocol defines the core interface that PraisonAIUI expects
|
||||
for injected kanban stores. Wrapper implementations must
|
||||
implement all methods to be compatible.
|
||||
|
||||
Duck typing contract matches InjectedKanbanStore in PraisonAIUI.
|
||||
"""
|
||||
|
||||
def get_board(
|
||||
self,
|
||||
*,
|
||||
board: str = "default",
|
||||
tenant: str | None = None,
|
||||
include_archived: bool = False,
|
||||
) -> dict:
|
||||
"""Get board data with tasks grouped by status.
|
||||
|
||||
Args:
|
||||
board: Board name (default: "default")
|
||||
tenant: Optional tenant filter
|
||||
include_archived: Whether to include archived tasks
|
||||
|
||||
Returns:
|
||||
Board data with tasks organized by status columns
|
||||
"""
|
||||
...
|
||||
|
||||
def get_task(self, task_id: str) -> dict | None:
|
||||
"""Get a single task by ID.
|
||||
|
||||
Args:
|
||||
task_id: Task identifier
|
||||
|
||||
Returns:
|
||||
Task data or None if not found
|
||||
"""
|
||||
...
|
||||
|
||||
def create_task(self, data: dict) -> dict:
|
||||
"""Create a new task.
|
||||
|
||||
Args:
|
||||
data: Task creation data (title, body, status, etc.)
|
||||
|
||||
Returns:
|
||||
Created task data with assigned ID
|
||||
"""
|
||||
...
|
||||
|
||||
def update_task(self, task_id: str, data: dict) -> dict | None:
|
||||
"""Update an existing task.
|
||||
|
||||
Args:
|
||||
task_id: Task identifier
|
||||
data: Fields to update
|
||||
|
||||
Returns:
|
||||
Updated task data or None if not found
|
||||
"""
|
||||
...
|
||||
|
||||
def move_task(self, task_id: str, status: str) -> dict | None:
|
||||
"""Move a task to a different status.
|
||||
|
||||
Args:
|
||||
task_id: Task identifier
|
||||
status: New status (must be in VALID_KANBAN_STATUSES)
|
||||
|
||||
Returns:
|
||||
Updated task data or None if not found
|
||||
"""
|
||||
...
|
||||
|
||||
def bulk_update(self, task_ids: list[str], status: str) -> dict:
|
||||
"""Update multiple tasks to the same status.
|
||||
|
||||
Args:
|
||||
task_ids: List of task identifiers
|
||||
status: New status for all tasks
|
||||
|
||||
Returns:
|
||||
Results summary with success/failure counts
|
||||
"""
|
||||
...
|
||||
|
||||
def delete_task(self, task_id: str) -> bool:
|
||||
"""Delete a task.
|
||||
|
||||
Args:
|
||||
task_id: Task identifier
|
||||
|
||||
Returns:
|
||||
True if deleted, False if not found
|
||||
"""
|
||||
...
|
||||
|
||||
def list_events(self, since: float = 0.0, board: str = "default") -> list[dict]:
|
||||
"""List kanban events since a timestamp.
|
||||
|
||||
Args:
|
||||
since: Unix timestamp to filter events from
|
||||
board: Board name to filter events
|
||||
|
||||
Returns:
|
||||
List of event data
|
||||
"""
|
||||
...
|
||||
|
||||
def health(self) -> dict:
|
||||
"""Check store health status.
|
||||
|
||||
Returns:
|
||||
Health status information
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class KanbanCommentingProtocol(Protocol):
|
||||
"""Extension protocol for kanban task commenting functionality.
|
||||
|
||||
This protocol is implemented separately from KanbanStoreProtocol to allow
|
||||
stores to optionally support commenting without breaking isinstance checks
|
||||
on the core protocol.
|
||||
"""
|
||||
|
||||
def add_comment(self, task_id: str, text: str, author: str | None = None) -> dict | None:
|
||||
"""Add a comment to a task.
|
||||
|
||||
Args:
|
||||
task_id: Task identifier
|
||||
text: Comment text
|
||||
author: Optional comment author
|
||||
|
||||
Returns:
|
||||
Comment data or None if task not found
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class KanbanLinkingProtocol(Protocol):
|
||||
"""Extension protocol for kanban task linking functionality.
|
||||
|
||||
This protocol is implemented separately from KanbanStoreProtocol to allow
|
||||
stores to optionally support task relationships without breaking isinstance
|
||||
checks on the core protocol.
|
||||
"""
|
||||
|
||||
def link_tasks(self, parent_id: str, child_id: str) -> bool:
|
||||
"""Link two tasks in parent-child relationship.
|
||||
|
||||
Args:
|
||||
parent_id: Parent task identifier
|
||||
child_id: Child task identifier
|
||||
|
||||
Returns:
|
||||
True if linked successfully, False otherwise
|
||||
"""
|
||||
...
|
||||
|
||||
def unlink_tasks(self, parent_id: str, child_id: str) -> bool:
|
||||
"""Unlink parent-child task relationship.
|
||||
|
||||
Args:
|
||||
parent_id: Parent task identifier
|
||||
child_id: Child task identifier
|
||||
|
||||
Returns:
|
||||
True if unlinked successfully, False otherwise
|
||||
"""
|
||||
...
|
||||
@@ -174,7 +174,14 @@ class ChromaKnowledgeAdapter:
|
||||
|
||||
# Add where filter if any conditions specified
|
||||
if where_filter:
|
||||
query_kwargs["where"] = where_filter
|
||||
# ChromaDB validation expects exactly one condition in the where clause
|
||||
# When multiple conditions exist, use $and operator to combine them
|
||||
if len(where_filter) == 1:
|
||||
query_kwargs["where"] = where_filter
|
||||
else:
|
||||
# Convert multiple conditions to ChromaDB's $and format
|
||||
and_conditions = [{k: v} for k, v in where_filter.items()]
|
||||
query_kwargs["where"] = {"$and": and_conditions}
|
||||
|
||||
response = self.collection.query(**query_kwargs)
|
||||
|
||||
@@ -197,6 +204,10 @@ class ChromaKnowledgeAdapter:
|
||||
return SearchResult(results=items)
|
||||
|
||||
except Exception as e:
|
||||
# Log ChromaDB query errors for better debugging
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.warning(f"ChromaDB search failed: {e}")
|
||||
return SearchResult(results=[])
|
||||
|
||||
def add(self, content: Any, *, user_id: Optional[str] = None, agent_id: Optional[str] = None,
|
||||
|
||||
@@ -73,10 +73,10 @@ def _get_live():
|
||||
|
||||
# Import token tracking
|
||||
try:
|
||||
from ..telemetry.token_collector import TokenMetrics, _token_collector
|
||||
from ..telemetry.token_collector import TokenMetrics, get_token_collector
|
||||
except ImportError:
|
||||
TokenMetrics = None
|
||||
_token_collector = None
|
||||
get_token_collector = None
|
||||
|
||||
# Logging is already configured in _logging.py via __init__.py
|
||||
|
||||
@@ -4519,7 +4519,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
|
||||
def _track_token_usage(self, response: Dict[str, Any], model: str) -> Optional[TokenMetrics]:
|
||||
"""Extract and track token usage from LLM response."""
|
||||
if not TokenMetrics or not _token_collector:
|
||||
if not TokenMetrics or not get_token_collector:
|
||||
return None
|
||||
|
||||
# Note: metrics check moved to call sites for performance
|
||||
@@ -4551,7 +4551,7 @@ Output MUST be JSON with 'reflection' and 'satisfactory'.
|
||||
# Track in global collector
|
||||
# Extract provider from model string (e.g., 'openai' from 'openai/gpt-4o-mini')
|
||||
provider = model.split('/')[0] if '/' in model else 'litellm'
|
||||
_token_collector.track_tokens(
|
||||
get_token_collector().track_tokens(
|
||||
model=model,
|
||||
agent=self.current_agent_name,
|
||||
metrics=metrics,
|
||||
|
||||
@@ -6,7 +6,13 @@ No heavy imports - only stdlib.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..knowledge.models import SearchResultItem
|
||||
|
||||
|
||||
ResultItem = Union[Dict[str, Any], "SearchResultItem"]
|
||||
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
@@ -19,6 +25,39 @@ def _estimate_tokens(text: str) -> int:
|
||||
return len(text) // 4 + 1
|
||||
|
||||
|
||||
def _extract_value(item: ResultItem, key: str, default: Any = None) -> Any:
|
||||
"""
|
||||
Extract value from either dict or object (e.g., SearchResultItem).
|
||||
|
||||
Args:
|
||||
item: Dictionary or object with attributes
|
||||
key: Key/attribute name to extract
|
||||
default: Default value if key not found
|
||||
|
||||
Returns:
|
||||
Extracted value or default
|
||||
"""
|
||||
if isinstance(item, dict):
|
||||
return item.get(key, default)
|
||||
else:
|
||||
return getattr(item, key, default)
|
||||
|
||||
|
||||
def _extract_metadata_value(
|
||||
item: ResultItem,
|
||||
metadata: Dict[str, Any],
|
||||
key: str,
|
||||
default: Any = "",
|
||||
) -> Any:
|
||||
"""
|
||||
Extract metadata value with fallback to top-level item attribute.
|
||||
"""
|
||||
value = metadata.get(key)
|
||||
if value is None:
|
||||
return _extract_value(item, key, default)
|
||||
return value
|
||||
|
||||
|
||||
def _chunk_hash(text: str, source: Optional[str] = None) -> str:
|
||||
"""
|
||||
Generate stable hash for chunk deduplication.
|
||||
@@ -35,31 +74,38 @@ def _chunk_hash(text: str, source: Optional[str] = None) -> str:
|
||||
|
||||
|
||||
def deduplicate_chunks(
|
||||
results: List[Dict[str, Any]],
|
||||
results: List[ResultItem],
|
||||
similarity_threshold: float = 0.9,
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> List[ResultItem]:
|
||||
"""
|
||||
Deduplicate chunks by content hash.
|
||||
|
||||
Handles both dict format and SearchResultItem objects.
|
||||
|
||||
Args:
|
||||
results: List of search results with 'text' or 'memory' key
|
||||
results: List of search results (dicts or SearchResultItem objects)
|
||||
similarity_threshold: Not used currently (hash-based dedup)
|
||||
|
||||
Returns:
|
||||
Deduplicated list of results
|
||||
"""
|
||||
seen_hashes: Set[str] = set()
|
||||
unique_results: List[Dict[str, Any]] = []
|
||||
unique_results: List[ResultItem] = []
|
||||
|
||||
for result in results:
|
||||
# Skip None results
|
||||
if result is None:
|
||||
continue
|
||||
# Handle different result formats
|
||||
text = result.get("text") or result.get("memory", "")
|
||||
|
||||
# Handle different result formats (dict or SearchResultItem)
|
||||
text = _extract_value(result, "text") or _extract_value(result, "memory", "")
|
||||
|
||||
# CRITICAL: Handle metadata=None from mem0 - ensure always dict
|
||||
metadata = result.get("metadata") or {}
|
||||
source = metadata.get("source", "")
|
||||
metadata = _extract_value(result, "metadata") or {}
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
|
||||
source = _extract_metadata_value(result, metadata, "source", "")
|
||||
|
||||
chunk_id = _chunk_hash(text, source)
|
||||
|
||||
@@ -106,17 +152,19 @@ def truncate_context(
|
||||
|
||||
|
||||
def build_context(
|
||||
results: List[Dict[str, Any]],
|
||||
results: List[ResultItem],
|
||||
max_tokens: int = 4000,
|
||||
deduplicate: bool = True,
|
||||
separator: str = "\n\n---\n\n",
|
||||
include_source: bool = True,
|
||||
) -> Tuple[str, List[Dict[str, Any]]]:
|
||||
) -> Tuple[str, List[ResultItem]]:
|
||||
"""
|
||||
Build context string from retrieval results.
|
||||
|
||||
Handles both dict format and SearchResultItem objects.
|
||||
|
||||
Args:
|
||||
results: List of search results
|
||||
results: List of search results (dicts or SearchResultItem objects)
|
||||
max_tokens: Maximum tokens for context
|
||||
deduplicate: Whether to deduplicate chunks
|
||||
separator: Separator between chunks
|
||||
@@ -133,7 +181,7 @@ def build_context(
|
||||
results = deduplicate_chunks(results)
|
||||
|
||||
context_parts: List[str] = []
|
||||
used_results: List[Dict[str, Any]] = []
|
||||
used_results: List[ResultItem] = []
|
||||
current_tokens = 0
|
||||
separator_tokens = _estimate_tokens(separator)
|
||||
|
||||
@@ -141,17 +189,21 @@ def build_context(
|
||||
# Skip None results
|
||||
if result is None:
|
||||
continue
|
||||
# Handle different result formats
|
||||
text = result.get("text") or result.get("memory", "")
|
||||
|
||||
# Handle different result formats (dict or SearchResultItem)
|
||||
text = _extract_value(result, "text") or _extract_value(result, "memory", "")
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Build chunk text with optional source
|
||||
if include_source:
|
||||
# CRITICAL: Handle metadata=None from mem0 - ensure always dict
|
||||
metadata = result.get("metadata") or {}
|
||||
source = metadata.get("source", "")
|
||||
filename = metadata.get("filename", "")
|
||||
metadata = _extract_value(result, "metadata") or {}
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
|
||||
source = _extract_metadata_value(result, metadata, "source", "")
|
||||
filename = _extract_metadata_value(result, metadata, "filename", "")
|
||||
source_label = filename or source or f"Source {i + 1}"
|
||||
chunk_text = f"[{source_label}]\n{text}"
|
||||
else:
|
||||
@@ -194,11 +246,14 @@ class DefaultContextBuilder:
|
||||
|
||||
def build(
|
||||
self,
|
||||
results: List[Dict[str, Any]],
|
||||
results: List[ResultItem],
|
||||
max_tokens: int = 4000,
|
||||
deduplicate: bool = True,
|
||||
) -> str:
|
||||
"""Build context string from results."""
|
||||
"""Build context string from results.
|
||||
|
||||
Handles both dict format and SearchResultItem objects.
|
||||
"""
|
||||
context, _ = build_context(
|
||||
results=results,
|
||||
max_tokens=max_tokens,
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
# Scheduler — runner vs loop
|
||||
|
||||
## ScheduleRunner (stateless)
|
||||
|
||||
`ScheduleRunner` is **stateless**. Call `get_due_jobs()` whenever you want to check for jobs that should fire. It does not run background threads.
|
||||
|
||||
Use when:
|
||||
- CLI one-shot tick (`praisonai schedule tick`)
|
||||
- Host integration that polls on its own timer
|
||||
- Tests
|
||||
|
||||
## ScheduleLoop (optional daemon)
|
||||
|
||||
`ScheduleLoop` is an **optional background daemon**. You must call `.start()` explicitly; it polls the runner and invokes `on_trigger` callbacks.
|
||||
|
||||
Use when:
|
||||
- Long-running host process (Pattern B/C via `setup_bridges()`)
|
||||
- Standalone scheduler service
|
||||
|
||||
```python
|
||||
from praisonaiagents.scheduler import FileScheduleStore, ScheduleRunner, ScheduleLoop
|
||||
|
||||
store = FileScheduleStore()
|
||||
runner = ScheduleRunner(store)
|
||||
loop = ScheduleLoop(runner)
|
||||
loop.start() # optional daemon
|
||||
```
|
||||
|
||||
The wrapper host calls `ensure_schedule_runner()` which starts `ScheduleLoop` when the scheduler optional dependency is available.
|
||||
@@ -32,7 +32,22 @@ When you provide a `session_id` to an Agent:
|
||||
|
||||
### Default Storage Location
|
||||
|
||||
Sessions are stored in: `~/.praison/sessions/{session_id}.json`
|
||||
Sessions are stored in: `~/.praisonai/sessions/{session_id}.json`
|
||||
|
||||
### Session metadata contract (Issue #48)
|
||||
|
||||
List endpoints and UI dashboards may expose these optional root-level fields:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `model` | string | LLM model used in the session |
|
||||
| `total_tokens` | int | Cumulative input+output tokens |
|
||||
| `cost` | float | Estimated USD cost |
|
||||
| `agent_id` | string | Gateway or registry agent id |
|
||||
| `source` | string | Origin: `chat`, `gateway`, `cli`, `api` |
|
||||
| `agent_name` | string | Human-readable agent name |
|
||||
|
||||
Fields are persisted in session JSON `metadata` and mirrored at root on save. Populated automatically after each assistant turn when using the default session store.
|
||||
|
||||
### Session File Format
|
||||
|
||||
@@ -45,7 +60,13 @@ Sessions are stored in: `~/.praison/sessions/{session_id}.json`
|
||||
],
|
||||
"created_at": "2026-01-02T04:00:00+00:00",
|
||||
"updated_at": "2026-01-02T04:01:00+00:00",
|
||||
"agent_name": "Assistant"
|
||||
"agent_name": "Assistant",
|
||||
"agent_id": "agent-abc",
|
||||
"source": "chat",
|
||||
"model": "gpt-4o-mini",
|
||||
"total_tokens": 128,
|
||||
"cost": 0.0004,
|
||||
"metadata": {}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
@@ -144,6 +144,35 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._extended_cache: Dict[str, ExtendedSessionData] = {}
|
||||
|
||||
def _load_session_from_disk(self, session_id: str, filepath: str) -> ExtendedSessionData:
|
||||
"""Load extended session JSON from disk (caller must hold FileLock)."""
|
||||
if os.path.exists(filepath):
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return ExtendedSessionData.from_dict(data)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
pass
|
||||
return ExtendedSessionData(session_id=session_id)
|
||||
|
||||
def _modify_session_locked(
|
||||
self,
|
||||
session_id: str,
|
||||
mutator,
|
||||
*,
|
||||
error_label: str = "modify session",
|
||||
) -> bool:
|
||||
"""Locked read-modify-write preserving extended session fields."""
|
||||
result = super()._modify_session_locked(
|
||||
session_id, mutator, error_label=error_label
|
||||
)
|
||||
if result:
|
||||
with self._lock:
|
||||
cached = self._cache.get(session_id)
|
||||
if isinstance(cached, ExtendedSessionData):
|
||||
self._extended_cache[session_id] = cached
|
||||
return result
|
||||
|
||||
def add_message(
|
||||
self,
|
||||
@@ -157,25 +186,22 @@ class HierarchicalSessionStore(DefaultSessionStore):
|
||||
|
||||
Overrides parent to preserve extended session data.
|
||||
"""
|
||||
from .store import SessionMessage
|
||||
|
||||
# Load extended session (force reload to get latest)
|
||||
session = self._load_extended_session(session_id, force_reload=True)
|
||||
|
||||
|
||||
message = SessionMessage(
|
||||
role=role,
|
||||
content=content,
|
||||
timestamp=time.time(),
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
session.messages.append(message)
|
||||
|
||||
# Trim messages if over limit
|
||||
if len(session.messages) > self.max_messages:
|
||||
session.messages = session.messages[-self.max_messages:]
|
||||
|
||||
return self._save_extended_session(session)
|
||||
|
||||
def _apply(session: SessionData) -> None:
|
||||
session.messages.append(message)
|
||||
if len(session.messages) > self.max_messages:
|
||||
session.messages = session.messages[-self.max_messages :]
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="add message to session"
|
||||
)
|
||||
|
||||
def _load_extended_session(self, session_id: str, force_reload: bool = False) -> ExtendedSessionData:
|
||||
"""Load extended session from disk."""
|
||||
|
||||
@@ -129,3 +129,16 @@ class SessionStoreProtocol(Protocol):
|
||||
True if the session exists.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class CheckpointQueryProtocol(Protocol):
|
||||
"""Read path for session checkpoints / rollback snapshots."""
|
||||
|
||||
def list_checkpoints(self, session_id: str) -> List[Dict[str, Any]]:
|
||||
"""Return checkpoint metadata for a session."""
|
||||
...
|
||||
|
||||
def get_checkpoint(self, session_id: str, checkpoint_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Return a single checkpoint payload."""
|
||||
...
|
||||
@@ -22,7 +22,7 @@ else:
|
||||
_HAS_FCNTL = False
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from ..paths import get_sessions_dir
|
||||
|
||||
@@ -81,7 +81,7 @@ class SessionData:
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
data = {
|
||||
"session_id": self.session_id,
|
||||
"messages": [m.to_dict() for m in self.messages],
|
||||
"created_at": self.created_at,
|
||||
@@ -92,6 +92,10 @@ class SessionData:
|
||||
"gateway_session_id": self.gateway_session_id,
|
||||
"agent_id": self.agent_id,
|
||||
}
|
||||
for key in ("model", "llm", "total_tokens", "token_count", "cost", "source"):
|
||||
if key in self.metadata:
|
||||
data[key] = self.metadata[key]
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "SessionData":
|
||||
@@ -137,7 +141,8 @@ class FileLock:
|
||||
self._lock_path = filepath + ".lock"
|
||||
|
||||
def __enter__(self):
|
||||
self.acquire()
|
||||
if not self.acquire():
|
||||
raise IOError(f"Failed to acquire file lock for {self.filepath} after {self.timeout}s")
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
@@ -286,6 +291,72 @@ class DefaultSessionStore:
|
||||
self._cache[session_id] = session
|
||||
|
||||
return session
|
||||
|
||||
def _load_session_from_disk(self, session_id: str, filepath: str) -> SessionData:
|
||||
"""Load session JSON from disk (caller must hold FileLock)."""
|
||||
if os.path.exists(filepath):
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return SessionData.from_dict(data)
|
||||
except (json.JSONDecodeError, IOError):
|
||||
pass
|
||||
return SessionData(session_id=session_id)
|
||||
|
||||
def _read_session_fresh(self, session_id: str) -> SessionData:
|
||||
"""Reload session from disk and refresh the in-process cache."""
|
||||
filepath = self._get_session_path(session_id)
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
session = self._load_session_from_disk(session_id, filepath)
|
||||
with self._lock:
|
||||
self._cache[session_id] = session
|
||||
return session
|
||||
|
||||
def _modify_session_locked(
|
||||
self,
|
||||
session_id: str,
|
||||
mutator: Callable[[SessionData], None],
|
||||
*,
|
||||
error_label: str = "modify session",
|
||||
) -> bool:
|
||||
"""Apply mutator after reloading from disk under FileLock."""
|
||||
filepath = self._get_session_path(session_id)
|
||||
|
||||
with FileLock(filepath, self.lock_timeout):
|
||||
session = self._load_session_from_disk(session_id, filepath)
|
||||
mutator(session)
|
||||
session.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
if len(session.messages) > self.max_messages:
|
||||
session.messages = session.messages[-self.max_messages:]
|
||||
|
||||
try:
|
||||
dir_path = os.path.dirname(filepath) or "."
|
||||
os.makedirs(dir_path, exist_ok=True)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
encoding="utf-8",
|
||||
dir=dir_path,
|
||||
delete=False,
|
||||
suffix=".tmp",
|
||||
) as f:
|
||||
json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
|
||||
temp_path = f.name
|
||||
|
||||
os.replace(temp_path, filepath)
|
||||
|
||||
with self._lock:
|
||||
self._cache[session_id] = session
|
||||
|
||||
return True
|
||||
except (IOError, OSError) as e:
|
||||
logger.error(f"Failed to {error_label} {session_id}: {e}")
|
||||
try:
|
||||
if "temp_path" in locals():
|
||||
os.remove(temp_path)
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
return False
|
||||
|
||||
def _save_session(self, session: SessionData) -> bool:
|
||||
"""Save session to disk with atomic write."""
|
||||
@@ -435,13 +506,13 @@ class DefaultSessionStore:
|
||||
Returns:
|
||||
List of {"role": "user/assistant", "content": "..."} dicts.
|
||||
"""
|
||||
session = self._load_session(session_id)
|
||||
session = self._read_session_fresh(session_id)
|
||||
limit = max_messages or self.max_messages
|
||||
return session.get_chat_history(limit)
|
||||
|
||||
def get_session(self, session_id: str) -> SessionData:
|
||||
"""Get full session data."""
|
||||
return self._load_session(session_id)
|
||||
return self._read_session_fresh(session_id)
|
||||
|
||||
def set_agent_info(
|
||||
self,
|
||||
@@ -450,27 +521,42 @@ class DefaultSessionStore:
|
||||
user_id: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Set agent info for a session."""
|
||||
session = self._load_session(session_id)
|
||||
|
||||
with self._lock:
|
||||
|
||||
def _apply(session: SessionData) -> None:
|
||||
if agent_name:
|
||||
session.agent_name = agent_name
|
||||
if user_id:
|
||||
session.user_id = user_id
|
||||
self._cache[session_id] = session
|
||||
|
||||
return self._save_session(session)
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="set agent info for session"
|
||||
)
|
||||
|
||||
def clear_session(self, session_id: str) -> bool:
|
||||
"""Clear all messages from a session."""
|
||||
session = self._load_session(session_id)
|
||||
|
||||
with self._lock:
|
||||
session.messages.clear()
|
||||
self._cache[session_id] = session
|
||||
|
||||
return self._save_session(session)
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id,
|
||||
lambda session: session.messages.clear(),
|
||||
error_label="clear session",
|
||||
)
|
||||
|
||||
def update_session_metadata(self, session_id: str, **fields: Any) -> bool:
|
||||
"""Merge run stats / metadata fields into a persisted session."""
|
||||
if not fields:
|
||||
return True
|
||||
|
||||
def _apply(session: SessionData) -> None:
|
||||
for key, value in fields.items():
|
||||
if value is None:
|
||||
continue
|
||||
session.metadata[key] = value
|
||||
if key in ("agent_id", "agent_name", "user_id"):
|
||||
setattr(session, key, value)
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="update session metadata"
|
||||
)
|
||||
|
||||
def delete_session(self, session_id: str) -> bool:
|
||||
"""Delete a session completely."""
|
||||
filepath = self._get_session_path(session_id)
|
||||
@@ -500,10 +586,16 @@ class DefaultSessionStore:
|
||||
data = json.load(f)
|
||||
sessions.append({
|
||||
"session_id": data.get("session_id", filename[:-5]),
|
||||
"id": data.get("session_id", filename[:-5]),
|
||||
"agent_name": data.get("agent_name"),
|
||||
"agent_id": data.get("agent_id") or (data.get("metadata") or {}).get("agent_id"),
|
||||
"source": data.get("source") or (data.get("metadata") or {}).get("source"),
|
||||
"created_at": data.get("created_at"),
|
||||
"updated_at": data.get("updated_at"),
|
||||
"message_count": len(data.get("messages", [])),
|
||||
"model": data.get("model") or data.get("llm") or (data.get("metadata") or {}).get("model"),
|
||||
"total_tokens": data.get("total_tokens") or data.get("token_count") or (data.get("metadata") or {}).get("total_tokens"),
|
||||
"cost": data.get("cost") or (data.get("metadata") or {}).get("cost"),
|
||||
})
|
||||
except (json.JSONDecodeError, IOError):
|
||||
continue
|
||||
@@ -563,7 +655,7 @@ class DefaultSessionStore:
|
||||
List of SessionData objects for the specified agent
|
||||
"""
|
||||
session_ids = self.list_sessions_by_agent(agent_name, limit)
|
||||
return [self._load_session(sid) for sid in session_ids]
|
||||
return [self._read_session_fresh(sid) for sid in session_ids]
|
||||
|
||||
def get_agent_chat_history(
|
||||
self,
|
||||
@@ -617,16 +709,15 @@ class DefaultSessionStore:
|
||||
Returns:
|
||||
True if saved successfully
|
||||
"""
|
||||
session = self._load_session(session_id)
|
||||
|
||||
with self._lock:
|
||||
def _apply(session: SessionData) -> None:
|
||||
if gateway_session_id:
|
||||
session.gateway_session_id = gateway_session_id
|
||||
if agent_id:
|
||||
session.agent_id = agent_id
|
||||
self._cache[session_id] = session
|
||||
|
||||
return self._save_session(session)
|
||||
|
||||
return self._modify_session_locked(
|
||||
session_id, _apply, error_label="set gateway info for session"
|
||||
)
|
||||
|
||||
def get_by_gateway_session(self, gateway_session_id: str) -> Optional[SessionData]:
|
||||
"""Get session data linked to a gateway session.
|
||||
|
||||
@@ -19,7 +19,9 @@ Usage:
|
||||
__all__ = [
|
||||
# Models
|
||||
"SkillProperties",
|
||||
"SkillMetadata",
|
||||
"SkillMetadata",
|
||||
"SkillRequirements",
|
||||
"SkillState",
|
||||
"ParseError",
|
||||
"ValidationError",
|
||||
# Parser
|
||||
@@ -54,13 +56,17 @@ __all__ = [
|
||||
"SkillActivatedEvent",
|
||||
# Budget
|
||||
"SkillPromptBudget",
|
||||
# Capability validation
|
||||
"CapabilityValidator",
|
||||
"EnforcementLevel",
|
||||
"ValidationResult",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy load module components to avoid import overhead."""
|
||||
if name in ("SkillProperties", "SkillMetadata", "ParseError", "ValidationError"):
|
||||
from .models import SkillProperties, SkillMetadata, ParseError, ValidationError
|
||||
if name in ("SkillProperties", "SkillMetadata", "SkillRequirements", "SkillState", "ParseError", "ValidationError"):
|
||||
from .models import SkillProperties, SkillMetadata, SkillRequirements, SkillState, ParseError, ValidationError
|
||||
return locals()[name]
|
||||
|
||||
if name in ("parse_frontmatter", "find_skill_md", "read_properties"):
|
||||
@@ -130,5 +136,9 @@ def __getattr__(name: str):
|
||||
return None
|
||||
|
||||
return load_skill
|
||||
|
||||
if name in ("CapabilityValidator", "EnforcementLevel", "ValidationResult"):
|
||||
from .capability_validator import CapabilityValidator, EnforcementLevel, ValidationResult
|
||||
return locals()[name]
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Capability validation for Agent Skills.
|
||||
|
||||
This module implements the capability gate enforcement system for skills,
|
||||
checking that required tools, servers, and environment variables are available
|
||||
before allowing skill activation.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
from enum import Enum
|
||||
|
||||
from .models import SkillProperties, SkillRequirements, SkillState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EnforcementLevel(Enum):
|
||||
"""Enforcement level for capability gates."""
|
||||
DISABLED = "disabled" # No enforcement (existing behavior)
|
||||
TELEMETRY = "telemetry" # Log only, no blocking
|
||||
WARN = "warn" # Warning but allow activation
|
||||
STRICT = "strict" # Hard failure, block activation
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValidationResult:
|
||||
"""Result of capability validation for a skill."""
|
||||
skill_name: str
|
||||
state: SkillState
|
||||
satisfied_tools: List[str]
|
||||
missing_tools: List[str]
|
||||
satisfied_servers: List[str]
|
||||
missing_servers: List[str]
|
||||
satisfied_env_vars: List[str]
|
||||
missing_env_vars: List[str]
|
||||
warnings: List[str]
|
||||
errors: List[str]
|
||||
|
||||
@property
|
||||
def is_fully_satisfied(self) -> bool:
|
||||
"""Check if all requirements are satisfied."""
|
||||
return not any([
|
||||
self.missing_tools,
|
||||
self.missing_servers,
|
||||
self.missing_env_vars,
|
||||
self.errors
|
||||
])
|
||||
|
||||
@property
|
||||
def has_critical_missing(self) -> bool:
|
||||
"""Check if critical requirements are missing (tools, servers)."""
|
||||
return bool(self.missing_tools or self.missing_servers)
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"skill_name": self.skill_name,
|
||||
"state": self.state.value,
|
||||
"satisfied_tools": self.satisfied_tools,
|
||||
"missing_tools": self.missing_tools,
|
||||
"satisfied_servers": self.satisfied_servers,
|
||||
"missing_servers": self.missing_servers,
|
||||
"satisfied_env_vars": self.satisfied_env_vars,
|
||||
"missing_env_vars": self.missing_env_vars,
|
||||
"warnings": self.warnings,
|
||||
"errors": self.errors,
|
||||
"is_fully_satisfied": self.is_fully_satisfied,
|
||||
"has_critical_missing": self.has_critical_missing
|
||||
}
|
||||
|
||||
|
||||
class CapabilityValidator:
|
||||
"""Validates skill requirements against available capabilities.
|
||||
|
||||
This is the core capability gate enforcement engine. It checks skill
|
||||
requirements against the runtime context (available tools, servers, etc.)
|
||||
and determines whether a skill should be activated based on policy.
|
||||
"""
|
||||
|
||||
def __init__(self, enforcement_level: EnforcementLevel = EnforcementLevel.WARN):
|
||||
"""Initialize the validator.
|
||||
|
||||
Args:
|
||||
enforcement_level: How strictly to enforce requirements
|
||||
"""
|
||||
self.enforcement_level = enforcement_level
|
||||
self._tool_cache: Optional[Set[str]] = None
|
||||
self._server_cache: Optional[Set[str]] = None
|
||||
|
||||
def validate_skill(
|
||||
self,
|
||||
skill: SkillProperties,
|
||||
available_tools: Optional[Set[str]] = None,
|
||||
available_servers: Optional[Set[str]] = None
|
||||
) -> ValidationResult:
|
||||
"""Validate a skill's requirements against available capabilities.
|
||||
|
||||
Args:
|
||||
skill: Skill to validate
|
||||
available_tools: Set of available tool names (auto-detected if None)
|
||||
available_servers: Set of available server names (auto-detected if None)
|
||||
|
||||
Returns:
|
||||
ValidationResult with validation status and details
|
||||
"""
|
||||
if not skill.requirements or skill.requirements.is_empty():
|
||||
# No requirements specified - skill is always available
|
||||
return ValidationResult(
|
||||
skill_name=skill.name,
|
||||
state=SkillState.ACTIVE,
|
||||
satisfied_tools=[],
|
||||
missing_tools=[],
|
||||
satisfied_servers=[],
|
||||
missing_servers=[],
|
||||
satisfied_env_vars=[],
|
||||
missing_env_vars=[],
|
||||
warnings=[],
|
||||
errors=[]
|
||||
)
|
||||
|
||||
# Get available capabilities
|
||||
if available_tools is None:
|
||||
available_tools = self._get_available_tools()
|
||||
if available_servers is None:
|
||||
available_servers = self._get_available_servers()
|
||||
|
||||
requirements = skill.requirements
|
||||
|
||||
# Validate tools
|
||||
satisfied_tools = [tool for tool in requirements.tools if tool in available_tools]
|
||||
missing_tools = [tool for tool in requirements.tools if tool not in available_tools]
|
||||
|
||||
# Validate servers
|
||||
satisfied_servers = [srv for srv in requirements.servers if srv in available_servers]
|
||||
missing_servers = [srv for srv in requirements.servers if srv not in available_servers]
|
||||
|
||||
# Validate environment variables
|
||||
satisfied_env_vars = [var for var in requirements.env_vars if os.getenv(var)]
|
||||
missing_env_vars = [var for var in requirements.env_vars if not os.getenv(var)]
|
||||
|
||||
# Determine warnings and errors
|
||||
warnings = []
|
||||
errors = []
|
||||
|
||||
if missing_tools:
|
||||
msg = f"Missing required tools: {', '.join(missing_tools)}"
|
||||
if self.enforcement_level in (EnforcementLevel.STRICT,):
|
||||
errors.append(msg)
|
||||
else:
|
||||
warnings.append(msg)
|
||||
|
||||
if missing_servers:
|
||||
msg = f"Missing required servers: {', '.join(missing_servers)}"
|
||||
if self.enforcement_level in (EnforcementLevel.STRICT,):
|
||||
errors.append(msg)
|
||||
else:
|
||||
warnings.append(msg)
|
||||
|
||||
if missing_env_vars:
|
||||
msg = f"Missing required environment variables: {', '.join(missing_env_vars)}"
|
||||
if self.enforcement_level == EnforcementLevel.STRICT:
|
||||
errors.append(msg)
|
||||
else:
|
||||
warnings.append(msg)
|
||||
|
||||
# Determine skill state
|
||||
if errors:
|
||||
state = SkillState.UNAVAILABLE
|
||||
elif warnings:
|
||||
state = SkillState.DEGRADED
|
||||
else:
|
||||
state = SkillState.ACTIVE
|
||||
|
||||
result = ValidationResult(
|
||||
skill_name=skill.name,
|
||||
state=state,
|
||||
satisfied_tools=satisfied_tools,
|
||||
missing_tools=missing_tools,
|
||||
satisfied_servers=satisfied_servers,
|
||||
missing_servers=missing_servers,
|
||||
satisfied_env_vars=satisfied_env_vars,
|
||||
missing_env_vars=missing_env_vars,
|
||||
warnings=warnings,
|
||||
errors=errors
|
||||
)
|
||||
|
||||
# Log based on enforcement level
|
||||
self._log_validation_result(result)
|
||||
|
||||
return result
|
||||
|
||||
def _get_available_tools(self) -> Set[str]:
|
||||
"""Get set of available tool names from registry."""
|
||||
if self._tool_cache is None:
|
||||
try:
|
||||
from ..tools.registry import get_registry
|
||||
registry = get_registry()
|
||||
self._tool_cache = set(registry.list_tools())
|
||||
except ImportError:
|
||||
logger.debug("Tool registry not available")
|
||||
self._tool_cache = set()
|
||||
return self._tool_cache
|
||||
|
||||
def _get_available_servers(self) -> Set[str]:
|
||||
"""Get set of available MCP server names."""
|
||||
if self._server_cache is None:
|
||||
# TODO: Implement MCP server discovery
|
||||
# For now, return empty set - this can be extended later
|
||||
self._server_cache = set()
|
||||
return self._server_cache
|
||||
|
||||
def _log_validation_result(self, result: ValidationResult) -> None:
|
||||
"""Log validation result based on enforcement level."""
|
||||
if self.enforcement_level == EnforcementLevel.DISABLED:
|
||||
return
|
||||
|
||||
skill_name = result.skill_name
|
||||
|
||||
if result.state == SkillState.UNAVAILABLE:
|
||||
if self.enforcement_level == EnforcementLevel.STRICT:
|
||||
logger.error(f"Skill '{skill_name}' blocked: {'; '.join(result.errors)}")
|
||||
else:
|
||||
logger.warning(f"Skill '{skill_name}' has critical missing dependencies: {'; '.join(result.errors)}")
|
||||
|
||||
elif result.state == SkillState.DEGRADED:
|
||||
logger.warning(f"Skill '{skill_name}' partially available: {'; '.join(result.warnings)}")
|
||||
|
||||
elif result.state == SkillState.ACTIVE:
|
||||
logger.debug(f"Skill '{skill_name}' requirements satisfied")
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""Clear cached capability information."""
|
||||
self._tool_cache = None
|
||||
self._server_cache = None
|
||||
@@ -108,7 +108,7 @@ class SkillLoader:
|
||||
return False
|
||||
|
||||
try:
|
||||
content = skill_md.read_text()
|
||||
content = skill_md.read_text(encoding="utf-8")
|
||||
_, body = parse_frontmatter(content)
|
||||
skill.instructions = body
|
||||
return True
|
||||
@@ -136,7 +136,7 @@ class SkillLoader:
|
||||
for script_file in scripts_dir.iterdir():
|
||||
if script_file.is_file():
|
||||
try:
|
||||
scripts[script_file.name] = script_file.read_text()
|
||||
scripts[script_file.name] = script_file.read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
continue
|
||||
except PermissionError:
|
||||
@@ -166,7 +166,7 @@ class SkillLoader:
|
||||
for ref_file in refs_dir.iterdir():
|
||||
if ref_file.is_file():
|
||||
try:
|
||||
refs[ref_file.name] = ref_file.read_text()
|
||||
refs[ref_file.name] = ref_file.read_text(encoding="utf-8")
|
||||
except Exception:
|
||||
continue
|
||||
except PermissionError:
|
||||
|
||||
@@ -5,12 +5,13 @@ from typing import List, Optional, Dict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from .models import SkillMetadata
|
||||
from .models import SkillMetadata, SkillState
|
||||
from .discovery import discover_skills, get_default_skill_dirs
|
||||
from .loader import SkillLoader, LoadedSkill
|
||||
from .prompt import generate_skills_xml
|
||||
from .substitution import render_skill_body
|
||||
from .shell_render import render_shell_blocks
|
||||
from .capability_validator import CapabilityValidator, EnforcementLevel, ValidationResult
|
||||
|
||||
|
||||
class SkillManager:
|
||||
@@ -34,11 +35,21 @@ class SkillManager:
|
||||
manager.activate(skill)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the SkillManager."""
|
||||
def __init__(self, enforcement_level: Optional[EnforcementLevel] = None):
|
||||
"""Initialize the SkillManager.
|
||||
|
||||
Args:
|
||||
enforcement_level: Capability enforcement level (defaults to environment-based)
|
||||
"""
|
||||
self._skills: Dict[str, LoadedSkill] = {}
|
||||
self._loader = SkillLoader()
|
||||
self._discovered = False
|
||||
self._validation_cache: Dict[str, ValidationResult] = {}
|
||||
|
||||
# Initialize capability validator
|
||||
if enforcement_level is None:
|
||||
enforcement_level = self._get_default_enforcement_level()
|
||||
self._validator = CapabilityValidator(enforcement_level)
|
||||
|
||||
@property
|
||||
def skills(self) -> List[LoadedSkill]:
|
||||
@@ -126,15 +137,29 @@ class SkillManager:
|
||||
def get_available_skills(self) -> List[SkillMetadata]:
|
||||
"""Get metadata for all available skills.
|
||||
|
||||
This is used for system prompt injection. Skills with
|
||||
``disable-model-invocation: true`` are omitted so the LLM never
|
||||
sees them and cannot auto-trigger them.
|
||||
This is used for system prompt injection. Skills are filtered by:
|
||||
1. disable-model-invocation: true (excluded)
|
||||
2. Capability enforcement (if strict mode, exclude UNAVAILABLE skills)
|
||||
"""
|
||||
return [
|
||||
skill.metadata
|
||||
for skill in self._skills.values()
|
||||
if not getattr(skill.properties, "disable_model_invocation", False)
|
||||
]
|
||||
available = []
|
||||
for skill in self._skills.values():
|
||||
# Skip if explicitly disabled for model invocation
|
||||
if getattr(skill.properties, "disable_model_invocation", False):
|
||||
continue
|
||||
|
||||
# Check capability requirements if enforcement is strict
|
||||
if self._validator.enforcement_level == EnforcementLevel.STRICT:
|
||||
try:
|
||||
validation = self.validate_skill_capabilities(skill.properties.name)
|
||||
if validation.state == SkillState.UNAVAILABLE:
|
||||
continue # Skip unavailable skills in strict mode
|
||||
except Exception as e:
|
||||
logger.warning(f"Skipping skill '{skill.properties.name}' due to validation error: {e}")
|
||||
continue
|
||||
|
||||
available.append(skill.metadata)
|
||||
|
||||
return available
|
||||
|
||||
def get_user_invocable_skills(self) -> List[LoadedSkill]:
|
||||
"""Return skills the *user* may invoke via slash-commands.
|
||||
@@ -566,9 +591,90 @@ author: agent
|
||||
logger.debug("Failed to clean up temp file %s", temp_path, exc_info=True)
|
||||
raise
|
||||
|
||||
def validate_skill_capabilities(self, skill_name: str, force_refresh: bool = False) -> ValidationResult:
|
||||
"""Validate a skill's capability requirements.
|
||||
|
||||
Args:
|
||||
skill_name: Name of the skill to validate
|
||||
force_refresh: If True, bypass cache and re-validate
|
||||
|
||||
Returns:
|
||||
ValidationResult with capability status
|
||||
|
||||
Raises:
|
||||
ValueError: If skill not found
|
||||
"""
|
||||
skill = self.get_skill(skill_name)
|
||||
if skill is None:
|
||||
raise ValueError(f"Skill '{skill_name}' not found")
|
||||
|
||||
# Check cache first
|
||||
if not force_refresh and skill_name in self._validation_cache:
|
||||
return self._validation_cache[skill_name]
|
||||
|
||||
# Perform validation
|
||||
result = self._validator.validate_skill(skill.properties)
|
||||
self._validation_cache[skill_name] = result
|
||||
return result
|
||||
|
||||
def get_available_skills_by_state(self, state: SkillState) -> List[LoadedSkill]:
|
||||
"""Get skills filtered by their capability validation state.
|
||||
|
||||
Args:
|
||||
state: Desired skill state
|
||||
|
||||
Returns:
|
||||
List of skills in the specified state
|
||||
"""
|
||||
matching = []
|
||||
for skill in self._skills.values():
|
||||
try:
|
||||
result = self.validate_skill_capabilities(skill.properties.name)
|
||||
if result.state == state:
|
||||
matching.append(skill)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to validate skill '{skill.properties.name}': {e}")
|
||||
continue
|
||||
return matching
|
||||
|
||||
def get_skills_diagnostics(self) -> Dict[str, ValidationResult]:
|
||||
"""Get capability diagnostics for all skills.
|
||||
|
||||
Returns:
|
||||
Dict mapping skill names to their validation results
|
||||
"""
|
||||
diagnostics = {}
|
||||
for skill_name in self.skill_names:
|
||||
try:
|
||||
diagnostics[skill_name] = self.validate_skill_capabilities(skill_name)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to validate skill '{skill_name}': {e}")
|
||||
return diagnostics
|
||||
|
||||
def _get_default_enforcement_level(self) -> EnforcementLevel:
|
||||
"""Get default enforcement level from environment."""
|
||||
import os
|
||||
level_str = os.getenv('SKILL_CAPABILITY_ENFORCEMENT', 'warn').lower()
|
||||
|
||||
level_map = {
|
||||
'disabled': EnforcementLevel.DISABLED,
|
||||
'off': EnforcementLevel.DISABLED,
|
||||
'telemetry': EnforcementLevel.TELEMETRY,
|
||||
'log': EnforcementLevel.TELEMETRY,
|
||||
'warn': EnforcementLevel.WARN,
|
||||
'warning': EnforcementLevel.WARN,
|
||||
'strict': EnforcementLevel.STRICT,
|
||||
'hard': EnforcementLevel.STRICT,
|
||||
'fail': EnforcementLevel.STRICT,
|
||||
}
|
||||
|
||||
return level_map.get(level_str, EnforcementLevel.WARN)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all loaded skills."""
|
||||
self._skills.clear()
|
||||
self._validation_cache.clear()
|
||||
self._validator.clear_cache()
|
||||
self._discovered = False
|
||||
|
||||
def __len__(self) -> int:
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""Data models for Agent Skills."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
from typing import Optional, List, Union
|
||||
from pathlib import Path
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ParseError(Exception):
|
||||
@@ -15,6 +16,105 @@ class ValidationError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class SkillState(Enum):
|
||||
"""Skill activation state based on requirement validation."""
|
||||
ACTIVE = "active" # All requirements satisfied
|
||||
DEGRADED = "degraded" # Some requirements missing (soft warn)
|
||||
UNAVAILABLE = "unavailable" # Critical requirements missing (hard fail)
|
||||
UNKNOWN = "unknown" # Requirements not yet validated
|
||||
|
||||
|
||||
@dataclass
|
||||
class SkillRequirements:
|
||||
"""Capability requirements for a skill.
|
||||
|
||||
This represents the parsed and normalized requirements from skill frontmatter.
|
||||
Supports both Hermes/OpenClaw conventions and PraisonAI-specific extensions.
|
||||
|
||||
Attributes:
|
||||
servers: Required MCP servers or service endpoints
|
||||
tools: Required tool names (from tool registry)
|
||||
env_vars: Required environment variables (use sparingly)
|
||||
openclaw_hints: OpenClaw-specific metadata (passthrough)
|
||||
"""
|
||||
servers: List[str] = field(default_factory=list)
|
||||
tools: List[str] = field(default_factory=list)
|
||||
env_vars: List[str] = field(default_factory=list)
|
||||
openclaw_hints: dict = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_frontmatter(cls, metadata: dict) -> "SkillRequirements":
|
||||
"""Parse requirements from skill frontmatter.
|
||||
|
||||
Normalizes various naming conventions:
|
||||
- requires_servers, requires_tools (Hermes/OpenClaw style)
|
||||
- allowed-tools (existing PraisonAI convention)
|
||||
- openclaw (OpenClaw-specific hints)
|
||||
|
||||
Args:
|
||||
metadata: Raw frontmatter dict
|
||||
|
||||
Returns:
|
||||
SkillRequirements instance
|
||||
"""
|
||||
servers = []
|
||||
tools = []
|
||||
env_vars = []
|
||||
openclaw_hints = {}
|
||||
|
||||
# Parse servers
|
||||
if "requires_servers" in metadata:
|
||||
servers = cls._normalize_list(metadata["requires_servers"])
|
||||
elif "requires-servers" in metadata:
|
||||
servers = cls._normalize_list(metadata["requires-servers"])
|
||||
|
||||
# Parse tools (multiple naming conventions)
|
||||
if "requires_tools" in metadata:
|
||||
tools.extend(cls._normalize_list(metadata["requires_tools"]))
|
||||
elif "requires-tools" in metadata:
|
||||
tools.extend(cls._normalize_list(metadata["requires-tools"]))
|
||||
|
||||
# Support existing allowed-tools convention for backward compatibility
|
||||
if "allowed-tools" in metadata:
|
||||
tools.extend(cls._normalize_list(metadata["allowed-tools"]))
|
||||
|
||||
# Parse environment variables
|
||||
if "requires_env" in metadata:
|
||||
env_vars = cls._normalize_list(metadata["requires_env"])
|
||||
elif "requires-env" in metadata:
|
||||
env_vars = cls._normalize_list(metadata["requires-env"])
|
||||
|
||||
# OpenClaw hints (passthrough)
|
||||
if "openclaw" in metadata:
|
||||
openclaw_hints = metadata["openclaw"] if isinstance(metadata["openclaw"], dict) else {}
|
||||
|
||||
return cls(
|
||||
servers=servers,
|
||||
tools=tools,
|
||||
env_vars=env_vars,
|
||||
openclaw_hints=openclaw_hints
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_list(value: Union[str, List[str]]) -> List[str]:
|
||||
"""Normalize string or list to list of strings."""
|
||||
if isinstance(value, str):
|
||||
# Support both comma-separated and space-separated strings
|
||||
return [item.strip() for item in value.replace(',', ' ').split() if item.strip()]
|
||||
elif isinstance(value, (list, tuple)):
|
||||
return [str(item).strip() for item in value if str(item).strip()]
|
||||
else:
|
||||
return []
|
||||
|
||||
def is_empty(self) -> bool:
|
||||
"""Check if no requirements are specified."""
|
||||
return not any([self.servers, self.tools, self.env_vars, self.openclaw_hints])
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
"""Truth value: True if any requirements are specified."""
|
||||
return not self.is_empty()
|
||||
|
||||
|
||||
@dataclass
|
||||
class SkillProperties:
|
||||
"""Properties parsed from a skill's SKILL.md frontmatter.
|
||||
@@ -49,6 +149,8 @@ class SkillProperties:
|
||||
hooks: Optional[dict] = None
|
||||
paths: Optional[object] = None # list[str] glob patterns
|
||||
shell: Optional[str] = None # "bash" | "powershell"
|
||||
# Capability requirements (new)
|
||||
requirements: Optional[SkillRequirements] = None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary, excluding None values and empty metadata."""
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from .models import ParseError, ValidationError, SkillProperties
|
||||
from .models import ParseError, ValidationError, SkillProperties, SkillRequirements
|
||||
|
||||
|
||||
def find_skill_md(skill_dir: Path) -> Optional[Path]:
|
||||
@@ -100,7 +100,7 @@ def read_properties(skill_dir: Path) -> SkillProperties:
|
||||
if skill_md is None:
|
||||
raise ParseError(f"SKILL.md not found in {skill_dir}")
|
||||
|
||||
content = skill_md.read_text()
|
||||
content = skill_md.read_text(encoding="utf-8")
|
||||
metadata, _ = parse_frontmatter(content)
|
||||
|
||||
if "name" not in metadata:
|
||||
@@ -143,4 +143,6 @@ def read_properties(skill_dir: Path) -> SkillProperties:
|
||||
hooks=metadata.get("hooks") if isinstance(metadata.get("hooks"), dict) else None,
|
||||
paths=metadata.get("paths"),
|
||||
shell=metadata.get("shell"),
|
||||
# Parse capability requirements
|
||||
requirements=SkillRequirements.from_frontmatter(metadata),
|
||||
)
|
||||
@@ -7,7 +7,7 @@ plugin bundle, enterprise policy) without touching the core.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Protocol, Iterable, Optional, runtime_checkable
|
||||
from typing import Protocol, Iterable, Optional, runtime_checkable, List, Dict, Any
|
||||
|
||||
from .models import SkillProperties
|
||||
|
||||
@@ -45,7 +45,8 @@ class SkillMutatorProtocol(Protocol):
|
||||
"""Agent-managed skill CRUD operations.
|
||||
|
||||
Allows agents to create, edit, and manage their own skills at runtime.
|
||||
Implementations should provide safe-by-default behavior (e.g., propose mode).
|
||||
Implementations should provide safe-by-default behaviour: use ``propose=True``
|
||||
(default) to stage mutations for human approval before writing to disk.
|
||||
"""
|
||||
|
||||
def create(self, name: str, content: str, category: Optional[str] = None,
|
||||
@@ -163,3 +164,30 @@ class SkillMutatorProtocol(Protocol):
|
||||
Status message describing the action taken
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SkillsCatalogProtocol(Protocol):
|
||||
"""Metadata listing for skills — UI catalog pages."""
|
||||
|
||||
def list_skills(self) -> List[Dict[str, Any]]:
|
||||
"""Return skill metadata dicts (name, description, location)."""
|
||||
...
|
||||
|
||||
|
||||
def list_skills_for_api() -> List[Dict[str, Any]]:
|
||||
"""Default catalog adapter using SkillManager when available."""
|
||||
try:
|
||||
from .manager import SkillManager
|
||||
|
||||
manager = SkillManager()
|
||||
return [
|
||||
{
|
||||
"name": s.name,
|
||||
"description": s.description,
|
||||
"location": getattr(s, "location", ""),
|
||||
}
|
||||
for s in manager.get_available_skills()
|
||||
]
|
||||
except Exception:
|
||||
return []
|
||||
@@ -191,9 +191,11 @@ def validate(skill_dir: Path, strict: bool = False) -> List[str]:
|
||||
return ["Missing required file: SKILL.md"]
|
||||
|
||||
try:
|
||||
content = skill_md.read_text()
|
||||
content = skill_md.read_text(encoding="utf-8")
|
||||
metadata, _ = parse_frontmatter(content)
|
||||
except ParseError as e:
|
||||
return [str(e)]
|
||||
except UnicodeDecodeError as e:
|
||||
return [f"SKILL.md is not valid UTF-8: {e}"]
|
||||
|
||||
return validate_metadata(metadata, skill_dir, strict=strict)
|
||||
@@ -346,13 +346,15 @@ class _LazyStorageModule:
|
||||
|
||||
_BACKENDS = {
|
||||
"RedisBackend": "praisonai.storage",
|
||||
"MongoDBBackend": "praisonai.storage",
|
||||
"MongoDBBackend": "praisonai.storage",
|
||||
"PostgreSQLBackend": "praisonai.storage",
|
||||
"DynamoDBBackend": "praisonai.storage",
|
||||
"ValkeyBackend": "praisonai.storage",
|
||||
"ValkeySearchBackend": "praisonai.storage",
|
||||
# Legacy names for compatibility
|
||||
"RedisStorageAdapter": "praisonai.storage",
|
||||
"MongoDBStorageAdapter": "praisonai.storage",
|
||||
"PostgreSQLStorageAdapter": "praisonai.storage",
|
||||
"PostgreSQLStorageAdapter": "praisonai.storage",
|
||||
"DynamoDBStorageAdapter": "praisonai.storage",
|
||||
}
|
||||
|
||||
@@ -390,7 +392,7 @@ def get_backend(
|
||||
Factory function to get a storage backend.
|
||||
|
||||
Args:
|
||||
backend_type: Type of backend ("file", "sqlite", "redis", "mongodb", "postgresql", "dynamodb")
|
||||
backend_type: Type of backend ("file", "sqlite", "redis", "mongodb", "postgresql", "dynamodb", "valkey")
|
||||
**kwargs: Backend-specific arguments
|
||||
|
||||
Returns:
|
||||
@@ -431,9 +433,13 @@ def get_backend(
|
||||
# Lazy load from wrapper
|
||||
backend_class = getattr(_heavy_backends, "DynamoDBBackend")
|
||||
return backend_class(**kwargs)
|
||||
elif backend_type == "valkey":
|
||||
# Lazy load from wrapper
|
||||
backend_class = getattr(_heavy_backends, "ValkeyBackend")
|
||||
return backend_class(**kwargs)
|
||||
else:
|
||||
raise ValueError(f"Unknown backend type: {backend_type}. "
|
||||
f"Supported: file, sqlite, redis, mongodb, postgresql, dynamodb")
|
||||
f"Supported: file, sqlite, redis, mongodb, postgresql, dynamodb, valkey")
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""
|
||||
@@ -442,7 +448,8 @@ def __getattr__(name: str):
|
||||
This allows `from praisonaiagents.storage.backends import RedisBackend` to work
|
||||
while maintaining lazy loading.
|
||||
"""
|
||||
if name in {"RedisBackend", "MongoDBBackend", "PostgreSQLBackend", "DynamoDBBackend"}:
|
||||
if name in {"RedisBackend", "MongoDBBackend", "PostgreSQLBackend", "DynamoDBBackend",
|
||||
"ValkeyBackend", "ValkeySearchBackend"}:
|
||||
return getattr(_heavy_backends, name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
@@ -100,3 +100,107 @@ class InMemoryTokenUsageSink:
|
||||
def clear(self) -> None:
|
||||
"""Clear all records."""
|
||||
self.records.clear()
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class UsageQueryProtocol(Protocol):
|
||||
"""Read path for token usage — UI dashboards and analytics."""
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
"""Aggregate totals and breakdowns by model/agent."""
|
||||
...
|
||||
|
||||
def list_recent(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
"""Recent usage records, newest last."""
|
||||
...
|
||||
|
||||
def get_by_task(self, task_id: str) -> List[Dict[str, Any]]:
|
||||
...
|
||||
|
||||
def get_by_agent(self, agent_name: str) -> List[Dict[str, Any]]:
|
||||
...
|
||||
|
||||
|
||||
class InMemoryUsageQuery:
|
||||
"""Query adapter backed by InMemoryTokenUsageSink records."""
|
||||
|
||||
def __init__(self, sink: InMemoryTokenUsageSink):
|
||||
self._sink = sink
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
records = self._sink.records
|
||||
total_in = 0
|
||||
total_out = 0
|
||||
total_all = 0
|
||||
by_model: Dict[str, Dict[str, int]] = {}
|
||||
by_agent: Dict[str, Dict[str, int]] = {}
|
||||
|
||||
# Single-pass optimization: calculate everything in one loop
|
||||
for r in records:
|
||||
r_in = r.get("input_tokens", 0)
|
||||
r_out = r.get("output_tokens", 0)
|
||||
r_total = r.get("total_tokens", r_in + r_out)
|
||||
|
||||
total_in += r_in
|
||||
total_out += r_out
|
||||
total_all += r_total
|
||||
|
||||
model = r.get("model", "unknown")
|
||||
agent = r.get("agent_name", "unknown")
|
||||
|
||||
# Update by_model and by_agent efficiently
|
||||
for bucket, key in ((by_model, model), (by_agent, agent)):
|
||||
stats = bucket.setdefault(key, {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0})
|
||||
stats["input_tokens"] += r_in
|
||||
stats["output_tokens"] += r_out
|
||||
stats["total_tokens"] += r_total
|
||||
return {
|
||||
"total_requests": len(records),
|
||||
"total_input_tokens": total_in,
|
||||
"total_output_tokens": total_out,
|
||||
"total_tokens": total_all, # Use consistent calculation
|
||||
"by_model": by_model,
|
||||
"by_agent": by_agent,
|
||||
}
|
||||
|
||||
def list_recent(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
return self._sink.records[-limit:]
|
||||
|
||||
def get_by_task(self, task_id: str) -> List[Dict[str, Any]]:
|
||||
return self._sink.get_by_task(task_id)
|
||||
|
||||
def get_by_agent(self, agent_name: str) -> List[Dict[str, Any]]:
|
||||
return self._sink.get_by_agent(agent_name)
|
||||
|
||||
|
||||
class TokenCollectorUsageQuery:
|
||||
"""Query adapter reading from the global TokenCollector + optional sink."""
|
||||
|
||||
def __init__(self, collector: Any):
|
||||
self._collector = collector
|
||||
|
||||
def get_summary(self) -> Dict[str, Any]:
|
||||
summary = self._collector.get_session_summary()
|
||||
return {
|
||||
"total_requests": summary.get("total_interactions", 0),
|
||||
"total_input_tokens": summary.get("total_metrics", {}).get("input_tokens", 0),
|
||||
"total_output_tokens": summary.get("total_metrics", {}).get("output_tokens", 0),
|
||||
"total_tokens": summary.get("total_tokens", 0),
|
||||
"by_model": summary.get("by_model", {}),
|
||||
"by_agent": summary.get("by_agent", {}),
|
||||
}
|
||||
|
||||
def list_recent(self, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
return self._collector.get_recent_interactions(limit)
|
||||
|
||||
def get_by_task(self, task_id: str) -> List[Dict[str, Any]]:
|
||||
sink = getattr(self._collector, "_sink", None)
|
||||
if sink and hasattr(sink, "get_by_task"):
|
||||
return sink.get_by_task(task_id)
|
||||
return []
|
||||
|
||||
def get_by_agent(self, agent_name: str) -> List[Dict[str, Any]]:
|
||||
sink = getattr(self._collector, "_sink", None)
|
||||
if sink and hasattr(sink, "get_by_agent"):
|
||||
return sink.get_by_agent(agent_name)
|
||||
return []
|
||||
@@ -191,4 +191,9 @@ class TokenCollector:
|
||||
|
||||
|
||||
# Global token collector instance
|
||||
_token_collector = TokenCollector()
|
||||
_token_collector = TokenCollector()
|
||||
|
||||
|
||||
def get_token_collector() -> TokenCollector:
|
||||
"""Return the global TokenCollector singleton."""
|
||||
return _token_collector
|
||||
@@ -7,7 +7,7 @@ import logging
|
||||
from praisonaiagents._logging import get_logger
|
||||
|
||||
# Import dependencies
|
||||
from .token_collector import _token_collector, TokenMetrics
|
||||
from .token_collector import get_token_collector, TokenMetrics
|
||||
from .telemetry import get_telemetry
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -56,11 +56,12 @@ class TokenTelemetryBridge:
|
||||
|
||||
def export_token_metrics(self) -> Dict[str, Any]:
|
||||
"""Export token metrics for telemetry reporting."""
|
||||
if not _token_collector:
|
||||
collector = get_token_collector()
|
||||
if not collector:
|
||||
return {}
|
||||
|
||||
try:
|
||||
metrics = _token_collector.export_metrics()
|
||||
metrics = collector.export_metrics()
|
||||
|
||||
# Prepare telemetry-friendly format
|
||||
return {
|
||||
@@ -77,8 +78,7 @@ class TokenTelemetryBridge:
|
||||
|
||||
def reset_token_metrics(self):
|
||||
"""Reset token metrics collection."""
|
||||
if _token_collector:
|
||||
_token_collector.reset()
|
||||
get_token_collector().reset()
|
||||
|
||||
# Global telemetry bridge instance
|
||||
_token_telemetry_bridge = TokenTelemetryBridge()
|
||||
|
||||
@@ -10,7 +10,7 @@ _tools_lazy_cache = {}
|
||||
# Export core tool items for organized imports (lightweight)
|
||||
from .base import BaseTool, ToolResult, ToolValidationError, validate_tool
|
||||
from .decorator import tool, FunctionTool
|
||||
from .registry import get_registry, register_tool, get_tool, add_tool, has_tool, remove_tool, list_tools, list_available_tools, ToolRegistry
|
||||
from .registry import get_registry, register_tool, get_tool, add_tool, has_tool, remove_tool, list_tools, list_available_tools, list_tools_with_allowed_filter, list_tools_with_hermes_filter, ToolRegistry
|
||||
from .tools import Tools
|
||||
|
||||
# Export Injected type directly for easy access
|
||||
@@ -363,7 +363,7 @@ __all__ = list(TOOL_MAPPINGS.keys()) + [
|
||||
'BaseTool', 'ToolResult', 'ToolValidationError', 'validate_tool',
|
||||
'tool', 'FunctionTool',
|
||||
'get_registry', 'register_tool', 'get_tool', 'add_tool', 'has_tool', 'remove_tool',
|
||||
'list_tools', 'list_available_tools', 'ToolRegistry',
|
||||
'list_tools', 'list_available_tools', 'list_tools_with_allowed_filter', 'list_tools_with_hermes_filter', 'ToolRegistry',
|
||||
'Tools',
|
||||
# Validation and retry protocols
|
||||
'ValidationResult', 'ToolValidatorProtocol', 'AsyncToolValidatorProtocol', 'PassthroughValidator',
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
A2UI tools — optional declarative UI via Google's a2ui-agent-sdk.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.tools.a2ui_tools import send_a2ui_messages
|
||||
|
||||
agent = Agent(name="assistant", tools=[send_a2ui_messages])
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Union
|
||||
|
||||
from praisonaiagents.tools.decorator import tool
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
|
||||
def _unwrap_messages(messages: Union[List[Dict[str, Any]], Dict[str, Any], str]) -> List[Dict[str, Any]]:
|
||||
"""Minimal unwrap for common LLM arg shapes (not full UI normalisation)."""
|
||||
if isinstance(messages, str):
|
||||
try:
|
||||
messages = json.loads(messages)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ValueError(f"messages JSON string is not valid JSON: {exc}") from exc
|
||||
|
||||
if isinstance(messages, dict):
|
||||
if "messages" in messages:
|
||||
messages = messages["messages"]
|
||||
else:
|
||||
messages = [messages]
|
||||
|
||||
if not isinstance(messages, list):
|
||||
raise ValueError("messages must be a list of A2UI message dicts or a JSON string")
|
||||
|
||||
return messages
|
||||
|
||||
|
||||
@tool
|
||||
def send_a2ui_messages(messages: Union[List[Dict[str, Any]], Dict[str, Any], str]) -> Dict[str, Any]:
|
||||
"""
|
||||
Send validated A2UI JSON messages to a connected A2UI renderer client.
|
||||
|
||||
Pass a list of A2UI message dicts (createSurface, updateComponents, etc.)
|
||||
or a JSON string. Requires: pip install praisonaiagents[a2ui]
|
||||
|
||||
Args:
|
||||
messages: A2UI message list or JSON string from the model.
|
||||
|
||||
Returns:
|
||||
Dict with mime_type and serialised messages for A2A/AG-UI clients.
|
||||
"""
|
||||
from praisonaiagents.ui.a2ui.adapter import create_a2ui_part
|
||||
|
||||
messages = _unwrap_messages(messages)
|
||||
|
||||
part = create_a2ui_part({"messages": messages})
|
||||
data = getattr(getattr(part, "root", part), "data", messages)
|
||||
metadata = getattr(getattr(part, "root", part), "metadata", {}) or {}
|
||||
|
||||
return {
|
||||
"mime_type": metadata.get("mimeType", A2UI_MIME_TYPE),
|
||||
"messages": messages,
|
||||
"a2ui_part": data,
|
||||
}
|
||||
@@ -15,7 +15,7 @@ Usage:
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional, Type, get_type_hints
|
||||
from typing import Any, Dict, List, Optional, Type, get_type_hints
|
||||
import inspect
|
||||
import logging
|
||||
|
||||
@@ -234,18 +234,81 @@ class BaseTool(ABC):
|
||||
if getattr(self.run, '__isabstractmethod__', False):
|
||||
errors.append("Tool must implement the 'run()' method")
|
||||
|
||||
# Check parameters schema is valid
|
||||
# Check parameters schema is valid and OpenAI-compatible
|
||||
if self.parameters:
|
||||
if not isinstance(self.parameters, dict):
|
||||
errors.append("'parameters' must be a dictionary")
|
||||
elif "type" not in self.parameters:
|
||||
errors.append("'parameters' must have a 'type' field")
|
||||
elif "properties" not in self.parameters:
|
||||
errors.append("'parameters' must have a 'properties' field for OpenAI compatibility")
|
||||
|
||||
# Validate schema structure for OpenAI compatibility
|
||||
try:
|
||||
schema = self.get_schema()
|
||||
if not isinstance(schema, dict):
|
||||
errors.append("get_schema() must return a dictionary")
|
||||
elif schema.get("type") != "function":
|
||||
errors.append("get_schema() must return schema with type='function'")
|
||||
elif "function" not in schema:
|
||||
errors.append("get_schema() must have 'function' key")
|
||||
else:
|
||||
func = schema["function"]
|
||||
if not isinstance(func.get("parameters"), dict):
|
||||
errors.append("Schema function.parameters must be a dictionary")
|
||||
elif "properties" not in func["parameters"]:
|
||||
errors.append("Schema function.parameters must have 'properties' for OpenAI compatibility")
|
||||
except Exception as e:
|
||||
errors.append(f"Schema generation failed: {e}")
|
||||
|
||||
if errors:
|
||||
raise ToolValidationError(f"Tool '{self.name}' validation failed: {'; '.join(errors)}")
|
||||
|
||||
return True
|
||||
|
||||
def validate_schema_roundtrip(self) -> bool:
|
||||
"""Validate that schema can round-trip through OpenAI-style serialization.
|
||||
|
||||
This ensures the tool schema is compatible with LLM providers that expect
|
||||
OpenAI function calling format and that parameters can be properly parsed.
|
||||
|
||||
Raises:
|
||||
ToolValidationError: If round-trip fails
|
||||
|
||||
Returns:
|
||||
True if round-trip validation passes
|
||||
"""
|
||||
try:
|
||||
import json
|
||||
|
||||
# Get the schema
|
||||
schema = self.get_schema()
|
||||
|
||||
# Serialize to JSON (what gets sent to LLM)
|
||||
serialized = json.dumps(schema)
|
||||
|
||||
# Deserialize back
|
||||
deserialized = json.loads(serialized)
|
||||
|
||||
# Validate structure matches expectations
|
||||
if deserialized != schema:
|
||||
raise ToolValidationError(f"Schema serialization round-trip failed for tool '{self.name}'")
|
||||
|
||||
# Validate OpenAI format requirements
|
||||
func_schema = deserialized.get("function", {})
|
||||
parameters = func_schema.get("parameters", {})
|
||||
|
||||
if not isinstance(parameters.get("properties"), dict):
|
||||
raise ToolValidationError(f"Tool '{self.name}' parameters.properties must be a dict for OpenAI compatibility")
|
||||
|
||||
if "required" in parameters and not isinstance(parameters["required"], list):
|
||||
raise ToolValidationError(f"Tool '{self.name}' parameters.required must be a list")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
raise ToolValidationError(f"Schema round-trip validation failed for tool '{self.name}': {e}")
|
||||
|
||||
@classmethod
|
||||
def validate_class(cls) -> bool:
|
||||
"""Validate a tool class before instantiation.
|
||||
@@ -282,7 +345,9 @@ def validate_tool(tool: Any) -> bool:
|
||||
ToolValidationError: If validation fails
|
||||
"""
|
||||
if isinstance(tool, BaseTool):
|
||||
return tool.validate()
|
||||
tool.validate()
|
||||
tool.validate_schema_roundtrip()
|
||||
return True
|
||||
|
||||
if callable(tool):
|
||||
# For plain functions, check they have a name
|
||||
@@ -294,5 +359,67 @@ def validate_tool(tool: Any) -> bool:
|
||||
raise ToolValidationError(f"Invalid tool type: {type(tool)}")
|
||||
|
||||
|
||||
def validate_tool_schema_consistency(tools: List[Any]) -> bool:
|
||||
"""Validate a list of tools for schema consistency with OpenAI format.
|
||||
|
||||
This function ensures all tools in a list can be properly serialized
|
||||
and have consistent schema structures for use with LLM providers.
|
||||
|
||||
Args:
|
||||
tools: List of tool objects to validate
|
||||
|
||||
Returns:
|
||||
True if all tools are valid and consistent
|
||||
|
||||
Raises:
|
||||
ToolValidationError: If validation fails
|
||||
"""
|
||||
if not tools:
|
||||
return True
|
||||
|
||||
import json
|
||||
from .decorator import get_tool_schema
|
||||
schemas = []
|
||||
|
||||
for i, tool in enumerate(tools):
|
||||
try:
|
||||
# Validate individual tool
|
||||
validate_tool(tool)
|
||||
|
||||
# Get schema from different tool types
|
||||
if isinstance(tool, BaseTool):
|
||||
schema = tool.get_schema()
|
||||
elif hasattr(tool, 'get_schema') and callable(getattr(tool, 'get_schema')):
|
||||
schema = tool.get_schema()
|
||||
elif callable(tool):
|
||||
schema = get_tool_schema(tool)
|
||||
else:
|
||||
raise ToolValidationError(f"Cannot extract schema from tool at index {i}: {type(tool)}")
|
||||
|
||||
if not schema:
|
||||
raise ToolValidationError(f"Tool at index {i} returned empty schema")
|
||||
|
||||
# Validate JSON serialization
|
||||
try:
|
||||
json.dumps(schema)
|
||||
except (TypeError, ValueError) as e:
|
||||
raise ToolValidationError(f"Tool schema at index {i} is not JSON serializable: {e}")
|
||||
|
||||
schemas.append(schema)
|
||||
|
||||
except Exception as e:
|
||||
raise ToolValidationError(f"Tool validation failed at index {i}: {e}")
|
||||
|
||||
# Check for duplicate tool names
|
||||
names = set()
|
||||
for schema in schemas:
|
||||
name = schema.get("function", {}).get("name")
|
||||
if name in names:
|
||||
raise ToolValidationError(f"Duplicate tool name '{name}' found in tool list")
|
||||
names.add(name)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
# For backward compatibility - tools can also just be functions
|
||||
# The @tool decorator (in decorator.py) wraps functions into BaseTool instances
|
||||
@@ -214,6 +214,13 @@ def tool(
|
||||
availability=availability
|
||||
)
|
||||
|
||||
# Validate the tool at creation time for early error detection
|
||||
try:
|
||||
tool_instance.validate()
|
||||
tool_instance.validate_schema_roundtrip()
|
||||
except Exception as e:
|
||||
logging.warning(f"Tool validation warning for {tool_instance.name}: {e}")
|
||||
|
||||
# Register with global registry if available
|
||||
try:
|
||||
from .registry import get_registry
|
||||
|
||||
@@ -184,6 +184,41 @@ class ToolRegistry:
|
||||
|
||||
return available
|
||||
|
||||
def list_tools_with_allowed_filter(self, context: Optional[Dict[str, Any]] = None) -> List[str]:
|
||||
"""List tool names filtered by ALLOWED_TOOLS environment variable.
|
||||
|
||||
This applies the canonical ALLOWED_TOOLS filter to prevent tool
|
||||
name collisions in multi-environment agent systems.
|
||||
|
||||
Args:
|
||||
context: Optional context for availability checks
|
||||
|
||||
Returns:
|
||||
List of filtered tool names
|
||||
"""
|
||||
try:
|
||||
# Lazy import to avoid circular dependencies
|
||||
from ..allowed_tools_filter import filter_tools_with_allowed_tools
|
||||
|
||||
# Get all available tool names
|
||||
all_tools = self.list_tools()
|
||||
|
||||
# Apply ALLOWED_TOOLS filter with diagnostics
|
||||
filtered_tools = filter_tools_with_allowed_tools(all_tools, log_diagnostics=True)
|
||||
|
||||
return sorted(filtered_tools)
|
||||
|
||||
except ImportError:
|
||||
logging.warning("ALLOWED_TOOLS filter not available, returning all tools")
|
||||
return self.list_tools()
|
||||
except ValueError:
|
||||
# Preserve strict ALLOWED_TOOLS semantics (e.g., CI unknown tools, empty value)
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.error("Error applying ALLOWED_TOOLS filter: %s", e)
|
||||
# Fallback to all tools on error
|
||||
return self.list_tools()
|
||||
|
||||
def get_all(self) -> Dict[str, Union[BaseTool, Callable]]:
|
||||
"""Get all registered tools as a dict. Thread-safe."""
|
||||
with self._lock:
|
||||
@@ -381,3 +416,24 @@ def list_available_tools(context: Optional[Dict[str, Any]] = None) -> List[Union
|
||||
List of available tools
|
||||
"""
|
||||
return get_registry().list_available_tools(context)
|
||||
|
||||
|
||||
def list_tools_with_allowed_filter(context: Optional[Dict[str, Any]] = None) -> List[str]:
|
||||
"""List tool names filtered by ALLOWED_TOOLS from the global registry.
|
||||
|
||||
This is the canonical entry point for applying ALLOWED_TOOLS filtering
|
||||
across the PraisonAI ecosystem.
|
||||
|
||||
Args:
|
||||
context: Optional context for availability checks
|
||||
|
||||
Returns:
|
||||
List of filtered tool names
|
||||
"""
|
||||
return get_registry().list_tools_with_allowed_filter(context)
|
||||
|
||||
|
||||
# Backward compatibility alias
|
||||
def list_tools_with_hermes_filter(context: Optional[Dict[str, Any]] = None) -> List[str]:
|
||||
"""Backward compatibility alias for list_tools_with_allowed_filter."""
|
||||
return list_tools_with_allowed_filter(context)
|
||||
@@ -253,7 +253,7 @@ class A2A:
|
||||
from praisonaiagents.ui.a2a.conversion import (
|
||||
extract_user_input,
|
||||
praisonai_to_a2a_message,
|
||||
create_artifact,
|
||||
create_response_artifact,
|
||||
)
|
||||
|
||||
# Validate params
|
||||
@@ -286,7 +286,8 @@ class A2A:
|
||||
else:
|
||||
raise RuntimeError("No agent or agents configured")
|
||||
response_msg = praisonai_to_a2a_message(str(response), task_id=task_id)
|
||||
artifact = create_artifact(str(response))
|
||||
history = getattr(self.agent, "chat_history", None) if self.agent else None
|
||||
artifact = create_response_artifact(str(response), chat_history=history)
|
||||
self.task_store.add_to_history(task_id, response_msg)
|
||||
self.task_store.add_artifact(task_id, artifact)
|
||||
self.task_store.update_status(task_id, TaskState.COMPLETED)
|
||||
@@ -322,7 +323,8 @@ class A2A:
|
||||
response_msg = praisonai_to_a2a_message(
|
||||
str(response), task_id=task.id
|
||||
)
|
||||
artifact = create_artifact(str(response))
|
||||
history = getattr(self.agent, "chat_history", None) if self.agent else None
|
||||
artifact = create_response_artifact(str(response), chat_history=history)
|
||||
|
||||
# Update task store
|
||||
self.task_store.add_to_history(task.id, response_msg)
|
||||
|
||||
@@ -16,6 +16,7 @@ from praisonaiagents.ui.a2a.types import (
|
||||
DataPart,
|
||||
Artifact,
|
||||
)
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
|
||||
def a2a_to_praisonai_messages(messages: List[Message]) -> List[Dict[str, Any]]:
|
||||
@@ -107,6 +108,75 @@ def praisonai_to_a2a_message(
|
||||
)
|
||||
|
||||
|
||||
def find_a2ui_tool_result(
|
||||
chat_history: Optional[List[Dict[str, Any]]],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Return the last A2UI tool result dict from agent chat history, if any."""
|
||||
if not chat_history:
|
||||
return None
|
||||
for msg in reversed(chat_history):
|
||||
if msg.get("role") != "tool":
|
||||
continue
|
||||
content = msg.get("content")
|
||||
if content is None:
|
||||
continue
|
||||
data: Any = content if isinstance(content, dict) else None
|
||||
if data is None:
|
||||
try:
|
||||
data = json.loads(content)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
continue
|
||||
if isinstance(data, dict) and data.get("mime_type") == A2UI_MIME_TYPE:
|
||||
return data
|
||||
return None
|
||||
|
||||
|
||||
def create_a2ui_artifact(
|
||||
result: Dict[str, Any],
|
||||
artifact_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> Artifact:
|
||||
"""Create an A2A Artifact with a DataPart for an A2UI tool result."""
|
||||
art_id = artifact_id or f"art-{uuid.uuid4().hex[:12]}"
|
||||
messages = result.get("messages") or []
|
||||
a2ui_data = result.get("a2ui_part")
|
||||
if not isinstance(a2ui_data, dict):
|
||||
a2ui_data = {"messages": messages}
|
||||
elif "messages" not in a2ui_data:
|
||||
a2ui_data = {**a2ui_data, "messages": messages}
|
||||
|
||||
return Artifact(
|
||||
artifact_id=art_id,
|
||||
name=name,
|
||||
description=description,
|
||||
parts=[
|
||||
DataPart(
|
||||
data=a2ui_data,
|
||||
metadata={"mimeType": result.get("mime_type", A2UI_MIME_TYPE)},
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def create_response_artifact(
|
||||
content: str,
|
||||
chat_history: Optional[List[Dict[str, Any]]] = None,
|
||||
artifact_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> Artifact:
|
||||
"""Build artifact from response text, preferring A2UI DataPart when present in history."""
|
||||
a2ui = find_a2ui_tool_result(chat_history)
|
||||
if a2ui:
|
||||
return create_a2ui_artifact(
|
||||
a2ui, artifact_id=artifact_id, name=name, description=description
|
||||
)
|
||||
return create_artifact(
|
||||
content, artifact_id=artifact_id, name=name, description=description
|
||||
)
|
||||
|
||||
|
||||
def create_artifact(
|
||||
content: str,
|
||||
artifact_id: Optional[str] = None,
|
||||
|
||||
@@ -1,126 +1,44 @@
|
||||
"""
|
||||
A2UI (Agent-to-User Interface) Protocol Integration for PraisonAI Agents
|
||||
A2UI integration via Google's a2ui-agent-sdk (optional).
|
||||
|
||||
This module provides A2UI protocol support, enabling PraisonAI Agents
|
||||
to generate rich, interactive UIs declaratively via JSON.
|
||||
|
||||
A2UI is Google's open standard that allows agents to "speak UI" by sending
|
||||
declarative JSON describing UI components, which clients then render natively.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents.ui.a2ui import (
|
||||
# Message types
|
||||
CreateSurfaceMessage,
|
||||
UpdateComponentsMessage,
|
||||
UpdateDataModelMessage,
|
||||
DeleteSurfaceMessage,
|
||||
# Component types
|
||||
TextComponent,
|
||||
ImageComponent,
|
||||
ButtonComponent,
|
||||
CardComponent,
|
||||
RowComponent,
|
||||
ColumnComponent,
|
||||
# Extension helpers
|
||||
create_a2ui_part,
|
||||
is_a2ui_part,
|
||||
)
|
||||
Declarative agent UI uses the official SDK — not reimplemented in PraisonAI core.
|
||||
Install: pip install praisonaiagents[a2ui]
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.a2ui.types import (
|
||||
# Data binding
|
||||
PathBinding,
|
||||
resolve_string_or_path,
|
||||
# Action types
|
||||
ActionContext,
|
||||
Action,
|
||||
# Children types
|
||||
ChildrenTemplate,
|
||||
# Component types
|
||||
TextComponent,
|
||||
ImageComponent,
|
||||
IconComponent,
|
||||
ButtonComponent,
|
||||
CardComponent,
|
||||
RowComponent,
|
||||
ColumnComponent,
|
||||
ListComponent,
|
||||
TextFieldComponent,
|
||||
CheckBoxComponent,
|
||||
SliderComponent,
|
||||
DateTimeInputComponent,
|
||||
# Message types
|
||||
CreateSurfaceMessage,
|
||||
UpdateComponentsMessage,
|
||||
UpdateDataModelMessage,
|
||||
DeleteSurfaceMessage,
|
||||
# Data part
|
||||
A2UIDataPart,
|
||||
)
|
||||
|
||||
from praisonaiagents.ui.a2ui.extension import (
|
||||
A2UI_MIME_TYPE,
|
||||
A2UI_EXTENSION_URI,
|
||||
STANDARD_CATALOG_ID,
|
||||
from praisonaiagents.ui.a2ui.adapter import (
|
||||
create_a2ui_part,
|
||||
generate_a2ui_system_prompt,
|
||||
get_schema_manager,
|
||||
is_a2ui_part,
|
||||
get_a2ui_agent_extension,
|
||||
AgentExtension,
|
||||
)
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
from praisonaiagents.ui.a2ui.templates import (
|
||||
ChatTemplate,
|
||||
ListTemplate,
|
||||
FormTemplate,
|
||||
DashboardTemplate,
|
||||
parse_a2ui_response,
|
||||
)
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE, A2UIToolResultProtocol
|
||||
|
||||
__all__ = [
|
||||
# Data binding
|
||||
"PathBinding",
|
||||
"resolve_string_or_path",
|
||||
# Action types
|
||||
"ActionContext",
|
||||
"Action",
|
||||
# Children types
|
||||
"ChildrenTemplate",
|
||||
# Component types
|
||||
"TextComponent",
|
||||
"ImageComponent",
|
||||
"IconComponent",
|
||||
"ButtonComponent",
|
||||
"CardComponent",
|
||||
"RowComponent",
|
||||
"ColumnComponent",
|
||||
"ListComponent",
|
||||
"TextFieldComponent",
|
||||
"CheckBoxComponent",
|
||||
"SliderComponent",
|
||||
"DateTimeInputComponent",
|
||||
# Message types
|
||||
"CreateSurfaceMessage",
|
||||
"UpdateComponentsMessage",
|
||||
"UpdateDataModelMessage",
|
||||
"DeleteSurfaceMessage",
|
||||
# Data part
|
||||
"A2UIDataPart",
|
||||
# Extension
|
||||
"A2UI",
|
||||
"A2UI_MIME_TYPE",
|
||||
"A2UI_EXTENSION_URI",
|
||||
"STANDARD_CATALOG_ID",
|
||||
"A2UIToolResultProtocol",
|
||||
"create_a2ui_part",
|
||||
"generate_a2ui_system_prompt",
|
||||
"get_schema_manager",
|
||||
"is_a2ui_part",
|
||||
"get_a2ui_agent_extension",
|
||||
"AgentExtension",
|
||||
# Surface builder
|
||||
"Surface",
|
||||
# Agent wrapper
|
||||
"A2UIAgent",
|
||||
# Templates
|
||||
"ChatTemplate",
|
||||
"ListTemplate",
|
||||
"FormTemplate",
|
||||
"DashboardTemplate",
|
||||
"parse_a2ui_response",
|
||||
]
|
||||
|
||||
|
||||
class A2UI:
|
||||
"""
|
||||
Optional A2UI facade (lazy — imports a2ui-agent-sdk on first use).
|
||||
|
||||
Example:
|
||||
from praisonaiagents.ui import A2UI
|
||||
|
||||
part = A2UI.create_part({"createSurface": {...}})
|
||||
prompt = A2UI.system_prompt("You are a helpful assistant.")
|
||||
"""
|
||||
|
||||
create_part = staticmethod(create_a2ui_part)
|
||||
is_part = staticmethod(is_a2ui_part)
|
||||
parse_response = staticmethod(parse_a2ui_response)
|
||||
schema_manager = staticmethod(get_schema_manager)
|
||||
system_prompt = staticmethod(generate_a2ui_system_prompt)
|
||||
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
Thin adapter over Google's a2ui-agent-sdk (optional dependency).
|
||||
|
||||
Install: pip install praisonaiagents[a2ui]
|
||||
Spec: https://github.com/google/A2UI
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
_A2UI_INSTALL_HINT = (
|
||||
"A2UI support requires the optional dependency. "
|
||||
"Install with: pip install praisonaiagents[a2ui]"
|
||||
)
|
||||
|
||||
|
||||
def _import(name: str):
|
||||
try:
|
||||
import importlib
|
||||
|
||||
return importlib.import_module(name)
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
f"{_A2UI_INSTALL_HINT} "
|
||||
f"(failed to import '{name}'; "
|
||||
"if already installed, ensure a2a-sdk>=0.3.0,<1.0.0 and not a2a-sdk 1.x)"
|
||||
) from exc
|
||||
|
||||
|
||||
def create_a2ui_part(a2ui_data: Dict[str, Any]) -> Any:
|
||||
"""Wrap A2UI payload for A2A transport (MIME: application/json+a2ui)."""
|
||||
parts = _import("a2ui.a2a.parts")
|
||||
return parts.create_a2ui_part(a2ui_data)
|
||||
|
||||
|
||||
def is_a2ui_part(part: Any) -> bool:
|
||||
"""Return True if an A2A part contains A2UI data."""
|
||||
parts = _import("a2ui.a2a.parts")
|
||||
return parts.is_a2ui_part(part)
|
||||
|
||||
|
||||
def parse_a2ui_response(text: str) -> List[Any]:
|
||||
"""Parse and split LLM text into A2A parts (text + validated A2UI JSON)."""
|
||||
parser = _import("a2ui.parser.parser")
|
||||
return parser.parse_response(text)
|
||||
|
||||
|
||||
def get_schema_manager(
|
||||
version: str = "0.9",
|
||||
catalogs: Optional[List[Any]] = None,
|
||||
accepts_inline_catalogs: bool = False,
|
||||
) -> Any:
|
||||
"""Create an A2uiSchemaManager for LLM system prompts and validation."""
|
||||
constants = _import("a2ui.schema.constants")
|
||||
manager = _import("a2ui.schema.manager")
|
||||
basic = _import("a2ui.basic_catalog.provider")
|
||||
|
||||
version_const = constants.VERSION_0_9 if version == "0.9" else version
|
||||
catalog_configs = catalogs
|
||||
if catalog_configs is None:
|
||||
catalog_configs = [basic.BasicCatalog.get_config(version=version_const)]
|
||||
|
||||
return manager.A2uiSchemaManager(
|
||||
version=version_const,
|
||||
catalogs=catalog_configs,
|
||||
accepts_inline_catalogs=accepts_inline_catalogs,
|
||||
)
|
||||
|
||||
|
||||
def generate_a2ui_system_prompt(
|
||||
role_description: str,
|
||||
workflow_description: str = "",
|
||||
ui_description: str = "",
|
||||
*,
|
||||
version: str = "0.9",
|
||||
include_schema: bool = True,
|
||||
include_examples: bool = True,
|
||||
) -> str:
|
||||
"""Build an LLM system prompt with embedded A2UI schema and examples."""
|
||||
schema_manager = get_schema_manager(version=version)
|
||||
return schema_manager.generate_system_prompt(
|
||||
role_description=role_description,
|
||||
workflow_description=workflow_description,
|
||||
ui_description=ui_description,
|
||||
include_schema=include_schema,
|
||||
include_examples=include_examples,
|
||||
)
|
||||
@@ -1,358 +0,0 @@
|
||||
"""
|
||||
A2UI Agent Integration
|
||||
|
||||
Provides A2UIAgent wrapper for generating A2UI responses from PraisonAI agents.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import (
|
||||
PathBinding,
|
||||
A2UIDataPart,
|
||||
)
|
||||
from praisonaiagents.ui.a2ui.extension import (
|
||||
create_a2ui_part,
|
||||
STANDARD_CATALOG_ID,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
class A2UIAgent:
|
||||
"""
|
||||
A2UI wrapper for PraisonAI agents.
|
||||
|
||||
Provides methods to generate A2UI responses from agent interactions.
|
||||
|
||||
Example:
|
||||
>>> from praisonaiagents import Agent
|
||||
>>> from praisonaiagents.ui.a2ui import A2UIAgent
|
||||
>>>
|
||||
>>> agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
>>> a2ui_agent = A2UIAgent(agent=agent)
|
||||
>>>
|
||||
>>> # Render a simple text response
|
||||
>>> messages = a2ui_agent.render_text("Hello, World!")
|
||||
>>>
|
||||
>>> # Render a list of items
|
||||
>>> messages = a2ui_agent.render_list(
|
||||
... title="Results",
|
||||
... items=[{"name": "Item 1"}, {"name": "Item 2"}]
|
||||
... )
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: "Agent",
|
||||
surface_id: Optional[str] = None,
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
):
|
||||
"""
|
||||
Initialize A2UIAgent wrapper.
|
||||
|
||||
Args:
|
||||
agent: The PraisonAI Agent to wrap
|
||||
surface_id: Unique ID for the UI surface (auto-generated if not provided)
|
||||
catalog_id: The component catalog to use
|
||||
"""
|
||||
self.agent = agent
|
||||
self.surface_id = surface_id or f"surface-{uuid.uuid4().hex[:8]}"
|
||||
self.catalog_id = catalog_id
|
||||
self._surface: Optional[Surface] = None
|
||||
|
||||
def get_surface(self) -> Surface:
|
||||
"""
|
||||
Get the underlying Surface object.
|
||||
|
||||
Creates a new Surface if one doesn't exist.
|
||||
|
||||
Returns:
|
||||
The Surface object
|
||||
"""
|
||||
if self._surface is None:
|
||||
self._surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
return self._surface
|
||||
|
||||
def _reset_surface(self) -> Surface:
|
||||
"""Reset and return a new Surface."""
|
||||
self._surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
return self._surface
|
||||
|
||||
# =========================================================================
|
||||
# Rendering methods
|
||||
# =========================================================================
|
||||
|
||||
def render_text(
|
||||
self,
|
||||
text: str,
|
||||
title: Optional[str] = None,
|
||||
usage_hint: str = "body",
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Render a simple text response as A2UI.
|
||||
|
||||
Args:
|
||||
text: The text content to display
|
||||
title: Optional title heading
|
||||
usage_hint: Text style hint (h1, h2, body, caption)
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = self._reset_surface()
|
||||
|
||||
children = []
|
||||
|
||||
if title:
|
||||
surface.text(id="title", text=title, usage_hint="h1")
|
||||
children.append("title")
|
||||
|
||||
surface.text(id="content", text=text, usage_hint=usage_hint)
|
||||
children.append("content")
|
||||
|
||||
# Add root column
|
||||
surface.column(id="root", children=children)
|
||||
|
||||
return surface.to_messages()
|
||||
|
||||
def render_list(
|
||||
self,
|
||||
title: str,
|
||||
items: List[Dict[str, Any]],
|
||||
item_title_key: str = "title",
|
||||
item_description_key: str = "description",
|
||||
item_image_key: Optional[str] = None,
|
||||
item_action: Optional[str] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Render a list of items as A2UI cards.
|
||||
|
||||
Args:
|
||||
title: List title
|
||||
items: List of item dictionaries
|
||||
item_title_key: Key for item title in each dict
|
||||
item_description_key: Key for item description
|
||||
item_image_key: Optional key for item image URL
|
||||
item_action: Optional action name for item buttons
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = self._reset_surface()
|
||||
|
||||
# Title
|
||||
surface.text(id="list-title", text=title, usage_hint="h1")
|
||||
|
||||
# Create card for each item
|
||||
card_ids = []
|
||||
for i, item in enumerate(items):
|
||||
card_id = f"card-{i}"
|
||||
content_id = f"card-content-{i}"
|
||||
title_id = f"item-title-{i}"
|
||||
desc_id = f"item-desc-{i}"
|
||||
|
||||
# Item title
|
||||
surface.text(
|
||||
id=title_id,
|
||||
text=item.get(item_title_key, f"Item {i+1}"),
|
||||
usage_hint="h3"
|
||||
)
|
||||
|
||||
# Item description
|
||||
surface.text(
|
||||
id=desc_id,
|
||||
text=item.get(item_description_key, ""),
|
||||
usage_hint="body"
|
||||
)
|
||||
|
||||
content_children = [title_id, desc_id]
|
||||
|
||||
# Optional action button
|
||||
if item_action:
|
||||
btn_id = f"btn-{i}"
|
||||
btn_text_id = f"btn-text-{i}"
|
||||
surface.text(id=btn_text_id, text="Select")
|
||||
surface.button(
|
||||
id=btn_id,
|
||||
child=btn_text_id,
|
||||
action_name=item_action,
|
||||
action_context=[{"key": "index", "value": i}],
|
||||
primary=True
|
||||
)
|
||||
content_children.append(btn_id)
|
||||
|
||||
# Card content column
|
||||
surface.column(id=content_id, children=content_children)
|
||||
|
||||
# Card
|
||||
surface.card(id=card_id, child=content_id)
|
||||
card_ids.append(card_id)
|
||||
|
||||
# List column
|
||||
surface.column(id="item-list", children=card_ids)
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=["list-title", "item-list"])
|
||||
|
||||
# Set data
|
||||
surface.set_data("items", items)
|
||||
|
||||
return surface.to_messages()
|
||||
|
||||
def render_form(
|
||||
self,
|
||||
title: str,
|
||||
fields: List[Dict[str, str]],
|
||||
submit_action: str = "submit",
|
||||
submit_label: str = "Submit",
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Render a form as A2UI.
|
||||
|
||||
Args:
|
||||
title: Form title
|
||||
fields: List of field definitions with id, label, type
|
||||
submit_action: Action name for submit button
|
||||
submit_label: Label for submit button
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = self._reset_surface()
|
||||
|
||||
# Title
|
||||
surface.text(id="form-title", text=title, usage_hint="h1")
|
||||
|
||||
# Fields
|
||||
field_ids = []
|
||||
form_data = {}
|
||||
|
||||
for field in fields:
|
||||
field_id = field.get("id", f"field-{len(field_ids)}")
|
||||
label = field.get("label", field_id)
|
||||
field_type = field.get("type", "text")
|
||||
|
||||
surface.text_field(
|
||||
id=field_id,
|
||||
label=label,
|
||||
text=PathBinding(path=f"/{field_id}"),
|
||||
field_type=field_type,
|
||||
)
|
||||
field_ids.append(field_id)
|
||||
form_data[field_id] = field.get("default", "")
|
||||
|
||||
# Submit button
|
||||
surface.text(id="submit-text", text=submit_label)
|
||||
|
||||
# Build action context from fields
|
||||
action_context = [
|
||||
{"key": f["id"], "value": {"path": f"/{f['id']}"}}
|
||||
for f in fields
|
||||
]
|
||||
|
||||
surface.button(
|
||||
id="submit-btn",
|
||||
child="submit-text",
|
||||
action_name=submit_action,
|
||||
action_context=action_context,
|
||||
primary=True,
|
||||
)
|
||||
field_ids.append("submit-btn")
|
||||
|
||||
# Form column
|
||||
surface.column(id="form-fields", children=field_ids)
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=["form-title", "form-fields"])
|
||||
|
||||
# Set initial form data
|
||||
for key, value in form_data.items():
|
||||
surface.set_data(key, value)
|
||||
|
||||
return surface.to_messages()
|
||||
|
||||
# =========================================================================
|
||||
# Agent interaction methods
|
||||
# =========================================================================
|
||||
|
||||
async def chat_async(self, message: str) -> str:
|
||||
"""
|
||||
Send a message to the agent and get a response.
|
||||
|
||||
Args:
|
||||
message: The message to send
|
||||
|
||||
Returns:
|
||||
The agent's response text
|
||||
"""
|
||||
# Use the agent's chat method
|
||||
if hasattr(self.agent, 'chat'):
|
||||
response = self.agent.chat(message)
|
||||
return response
|
||||
elif hasattr(self.agent, 'run'):
|
||||
response = self.agent.run(message)
|
||||
return response
|
||||
else:
|
||||
raise NotImplementedError("Agent does not have chat or run method")
|
||||
|
||||
def chat(self, message: str) -> str:
|
||||
"""
|
||||
Send a message to the agent and get a response (sync).
|
||||
|
||||
Args:
|
||||
message: The message to send
|
||||
|
||||
Returns:
|
||||
The agent's response text
|
||||
"""
|
||||
if hasattr(self.agent, 'chat'):
|
||||
return self.agent.chat(message)
|
||||
elif hasattr(self.agent, 'run'):
|
||||
return self.agent.run(message)
|
||||
else:
|
||||
raise NotImplementedError("Agent does not have chat or run method")
|
||||
|
||||
# =========================================================================
|
||||
# Export methods
|
||||
# =========================================================================
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get the current surface as A2UI messages.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
return self.get_surface().to_messages()
|
||||
|
||||
def to_json(self, indent: Optional[int] = None) -> str:
|
||||
"""
|
||||
Get the current surface as JSON string.
|
||||
|
||||
Args:
|
||||
indent: JSON indentation level
|
||||
|
||||
Returns:
|
||||
JSON string of A2UI messages
|
||||
"""
|
||||
return self.get_surface().to_json(indent=indent)
|
||||
|
||||
def to_a2a_part(self) -> A2UIDataPart:
|
||||
"""
|
||||
Get the current surface as an A2A DataPart.
|
||||
|
||||
Returns:
|
||||
A2UIDataPart with A2UI MIME type
|
||||
"""
|
||||
messages = self.to_messages()
|
||||
# For A2A, we send messages as a list
|
||||
return create_a2ui_part({"messages": messages})
|
||||
@@ -1,115 +0,0 @@
|
||||
"""
|
||||
A2UI Extension for A2A Protocol
|
||||
|
||||
Provides helpers for integrating A2UI with the A2A (Agent-to-Agent) protocol.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from praisonaiagents.ui.a2ui.types import A2UIDataPart
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Constants
|
||||
# =============================================================================
|
||||
|
||||
A2UI_EXTENSION_URI = "https://a2ui.org/a2a-extension/a2ui/v0.9"
|
||||
A2UI_MIME_TYPE = "application/json+a2ui"
|
||||
STANDARD_CATALOG_ID = "https://raw.githubusercontent.com/google/A2UI/refs/heads/main/specification/0.9/json/standard_catalog_definition.json"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Agent Extension Type
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class AgentExtension:
|
||||
"""
|
||||
A2A Agent Extension configuration.
|
||||
|
||||
Describes an extension that an agent supports.
|
||||
"""
|
||||
uri: str
|
||||
description: str = ""
|
||||
params: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper Functions
|
||||
# =============================================================================
|
||||
|
||||
def create_a2ui_part(a2ui_data: Dict[str, Any]) -> A2UIDataPart:
|
||||
"""
|
||||
Create an A2UI data part for A2A integration.
|
||||
|
||||
Wraps A2UI data in an A2UIDataPart with the appropriate MIME type.
|
||||
|
||||
Args:
|
||||
a2ui_data: The A2UI data dictionary (e.g., createSurface, updateComponents)
|
||||
|
||||
Returns:
|
||||
An A2UIDataPart with the A2UI MIME type in metadata
|
||||
|
||||
Example:
|
||||
>>> part = create_a2ui_part({"createSurface": {"surfaceId": "main", "catalogId": "..."}})
|
||||
>>> part.metadata["mimeType"]
|
||||
'application/json+a2ui'
|
||||
"""
|
||||
return A2UIDataPart(
|
||||
data=a2ui_data,
|
||||
metadata={"mimeType": A2UI_MIME_TYPE}
|
||||
)
|
||||
|
||||
|
||||
def is_a2ui_part(part: A2UIDataPart) -> bool:
|
||||
"""
|
||||
Check if a data part contains A2UI data.
|
||||
|
||||
Args:
|
||||
part: The data part to check
|
||||
|
||||
Returns:
|
||||
True if the part contains A2UI data (has A2UI MIME type), False otherwise
|
||||
|
||||
Example:
|
||||
>>> part = create_a2ui_part({"createSurface": {...}})
|
||||
>>> is_a2ui_part(part)
|
||||
True
|
||||
"""
|
||||
if not isinstance(part, A2UIDataPart):
|
||||
return False
|
||||
if part.metadata is None:
|
||||
return False
|
||||
return part.metadata.get("mimeType") == A2UI_MIME_TYPE
|
||||
|
||||
|
||||
def get_a2ui_agent_extension(
|
||||
accepts_inline_custom_catalog: bool = False,
|
||||
) -> AgentExtension:
|
||||
"""
|
||||
Create the A2UI AgentExtension configuration.
|
||||
|
||||
Returns an AgentExtension that can be added to an A2A Agent Card
|
||||
to indicate A2UI support.
|
||||
|
||||
Args:
|
||||
accepts_inline_custom_catalog: Whether the agent accepts inline custom catalogs
|
||||
|
||||
Returns:
|
||||
The configured A2UI AgentExtension
|
||||
|
||||
Example:
|
||||
>>> ext = get_a2ui_agent_extension()
|
||||
>>> ext.uri
|
||||
'https://a2ui.org/a2a-extension/a2ui/v0.9'
|
||||
"""
|
||||
params: Optional[Dict[str, Any]] = None
|
||||
if accepts_inline_custom_catalog:
|
||||
params = {"acceptsInlineCustomCatalog": True}
|
||||
|
||||
return AgentExtension(
|
||||
uri=A2UI_EXTENSION_URI,
|
||||
description="Provides agent-driven UI using the A2UI JSON format.",
|
||||
params=params
|
||||
)
|
||||
@@ -1,364 +0,0 @@
|
||||
"""
|
||||
A2UI Surface Builder
|
||||
|
||||
Provides a fluent API for building A2UI surfaces with components and data.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
from praisonaiagents.ui.a2ui.types import (
|
||||
Component,
|
||||
TextComponent,
|
||||
ImageComponent,
|
||||
ButtonComponent,
|
||||
CardComponent,
|
||||
RowComponent,
|
||||
ColumnComponent,
|
||||
ListComponent,
|
||||
TextFieldComponent,
|
||||
Action,
|
||||
ActionContext,
|
||||
PathBinding,
|
||||
CreateSurfaceMessage,
|
||||
UpdateComponentsMessage,
|
||||
UpdateDataModelMessage,
|
||||
StringOrPath,
|
||||
ChildrenType,
|
||||
)
|
||||
|
||||
|
||||
class Surface:
|
||||
"""
|
||||
A2UI Surface builder for creating rich agent-generated UIs.
|
||||
|
||||
Provides a fluent API for adding components and data, then
|
||||
generating the A2UI JSON messages.
|
||||
|
||||
Example:
|
||||
>>> surface = Surface(surface_id="main")
|
||||
>>> surface.text(id="title", text="Hello World", usage_hint="h1")
|
||||
>>> surface.button(id="submit", child="submit-text", action_name="submit")
|
||||
>>> surface.set_data("title", "Hello World")
|
||||
>>> messages = surface.to_messages()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
surface_id: str,
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
):
|
||||
"""
|
||||
Initialize a new Surface.
|
||||
|
||||
Args:
|
||||
surface_id: Unique identifier for this surface
|
||||
catalog_id: The component catalog to use (defaults to standard catalog)
|
||||
"""
|
||||
self.surface_id = surface_id
|
||||
self.catalog_id = catalog_id
|
||||
self.components: List[Component] = []
|
||||
self.data: Dict[str, Any] = {}
|
||||
|
||||
def add(self, component: Component) -> "Surface":
|
||||
"""
|
||||
Add a component to the surface.
|
||||
|
||||
Args:
|
||||
component: The component to add
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
self.components.append(component)
|
||||
return self
|
||||
|
||||
def set_data(self, key: str, value: Any) -> "Surface":
|
||||
"""
|
||||
Set a value in the data model.
|
||||
|
||||
Args:
|
||||
key: The key in the data model
|
||||
value: The value to set
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
self.data[key] = value
|
||||
return self
|
||||
|
||||
# =========================================================================
|
||||
# Helper methods for common components
|
||||
# =========================================================================
|
||||
|
||||
def text(
|
||||
self,
|
||||
id: str,
|
||||
text: StringOrPath,
|
||||
usage_hint: Optional[str] = None,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a Text component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
text: Text content or path binding
|
||||
usage_hint: Style hint (h1, h2, h3, body, caption)
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(TextComponent(
|
||||
id=id,
|
||||
text=text,
|
||||
usage_hint=usage_hint,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def image(
|
||||
self,
|
||||
id: str,
|
||||
url: StringOrPath,
|
||||
fit: Optional[str] = None,
|
||||
usage_hint: Optional[str] = None,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add an Image component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
url: Image URL or path binding
|
||||
fit: Image fit (contain, cover, fill, none, scale-down)
|
||||
usage_hint: Size hint (icon, avatar, smallFeature, etc.)
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(ImageComponent(
|
||||
id=id,
|
||||
url=url,
|
||||
fit=fit,
|
||||
usage_hint=usage_hint,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def button(
|
||||
self,
|
||||
id: str,
|
||||
child: str,
|
||||
action_name: Optional[str] = None,
|
||||
action_context: Optional[List[Dict[str, Any]]] = None,
|
||||
primary: bool = False,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a Button component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
child: Child component ID (usually a Text component)
|
||||
action_name: Name of the action to trigger
|
||||
action_context: List of {key, value} dicts for action context
|
||||
primary: Whether this is a primary button
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
action = None
|
||||
if action_name:
|
||||
context = []
|
||||
if action_context:
|
||||
for ctx in action_context:
|
||||
value = ctx.get("value")
|
||||
if isinstance(value, dict) and "path" in value:
|
||||
value = PathBinding(path=value["path"])
|
||||
context.append(ActionContext(key=ctx["key"], value=value))
|
||||
action = Action(name=action_name, context=context)
|
||||
|
||||
return self.add(ButtonComponent(
|
||||
id=id,
|
||||
child=child,
|
||||
action=action,
|
||||
primary=primary,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def card(
|
||||
self,
|
||||
id: str,
|
||||
child: str,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a Card component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
child: Child component ID
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(CardComponent(
|
||||
id=id,
|
||||
child=child,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def row(
|
||||
self,
|
||||
id: str,
|
||||
children: ChildrenType,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a Row (horizontal layout) component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
children: List of child IDs or ChildrenTemplate
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(RowComponent(
|
||||
id=id,
|
||||
children=children,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def column(
|
||||
self,
|
||||
id: str,
|
||||
children: ChildrenType,
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a Column (vertical layout) component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
children: List of child IDs or ChildrenTemplate
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(ColumnComponent(
|
||||
id=id,
|
||||
children=children,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def list(
|
||||
self,
|
||||
id: str,
|
||||
children: ChildrenType,
|
||||
direction: str = "vertical",
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a List component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
children: List of child IDs or ChildrenTemplate for dynamic lists
|
||||
direction: "vertical" or "horizontal"
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(ListComponent(
|
||||
id=id,
|
||||
children=children,
|
||||
direction=direction,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
def text_field(
|
||||
self,
|
||||
id: str,
|
||||
label: StringOrPath,
|
||||
text: StringOrPath,
|
||||
field_type: str = "text",
|
||||
weight: Optional[float] = None,
|
||||
) -> "Surface":
|
||||
"""
|
||||
Add a TextField input component.
|
||||
|
||||
Args:
|
||||
id: Component ID
|
||||
label: Field label
|
||||
text: Text value or path binding
|
||||
field_type: Input type (text, number, email, etc.)
|
||||
weight: Flex weight
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
return self.add(TextFieldComponent(
|
||||
id=id,
|
||||
label=label,
|
||||
text=text,
|
||||
field_type=field_type,
|
||||
weight=weight,
|
||||
))
|
||||
|
||||
# =========================================================================
|
||||
# Message generation
|
||||
# =========================================================================
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate A2UI messages for this surface.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
messages: List[Dict[str, Any]] = []
|
||||
|
||||
# 1. CreateSurface message
|
||||
create_msg = CreateSurfaceMessage(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
messages.append(create_msg.to_dict())
|
||||
|
||||
# 2. UpdateComponents message (if we have components)
|
||||
if self.components:
|
||||
update_msg = UpdateComponentsMessage(
|
||||
surface_id=self.surface_id,
|
||||
components=self.components,
|
||||
)
|
||||
messages.append(update_msg.to_dict())
|
||||
|
||||
# 3. UpdateDataModel message (if we have data)
|
||||
if self.data:
|
||||
data_msg = UpdateDataModelMessage(
|
||||
surface_id=self.surface_id,
|
||||
value=self.data,
|
||||
)
|
||||
messages.append(data_msg.to_dict())
|
||||
|
||||
return messages
|
||||
|
||||
def to_json(self, indent: Optional[int] = None) -> str:
|
||||
"""
|
||||
Generate A2UI messages as a JSON string.
|
||||
|
||||
Args:
|
||||
indent: JSON indentation level (None for compact)
|
||||
|
||||
Returns:
|
||||
JSON string of A2UI messages
|
||||
"""
|
||||
return json.dumps(self.to_messages(), indent=indent)
|
||||
@@ -1,17 +0,0 @@
|
||||
"""
|
||||
A2UI Templates
|
||||
|
||||
Pre-built UI templates for common use cases.
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.a2ui.templates.chat import ChatTemplate
|
||||
from praisonaiagents.ui.a2ui.templates.list import ListTemplate
|
||||
from praisonaiagents.ui.a2ui.templates.form import FormTemplate
|
||||
from praisonaiagents.ui.a2ui.templates.dashboard import DashboardTemplate
|
||||
|
||||
__all__ = [
|
||||
"ChatTemplate",
|
||||
"ListTemplate",
|
||||
"FormTemplate",
|
||||
"DashboardTemplate",
|
||||
]
|
||||
@@ -1,122 +0,0 @@
|
||||
"""
|
||||
Chat Template for A2UI
|
||||
|
||||
Provides a conversational chat UI template.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
|
||||
|
||||
class ChatTemplate:
|
||||
"""
|
||||
Chat UI template for conversational interfaces.
|
||||
|
||||
Creates a chat-style UI with user and agent messages.
|
||||
|
||||
Example:
|
||||
>>> template = ChatTemplate(surface_id="chat")
|
||||
>>> template.add_user_message("Hello!")
|
||||
>>> template.add_agent_message("Hi there! How can I help?")
|
||||
>>> messages = template.to_messages()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
surface_id: str = "chat",
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
title: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize ChatTemplate.
|
||||
|
||||
Args:
|
||||
surface_id: Unique ID for the surface
|
||||
catalog_id: Component catalog to use
|
||||
title: Optional title for the chat
|
||||
"""
|
||||
self.surface_id = surface_id
|
||||
self.catalog_id = catalog_id
|
||||
self.title = title
|
||||
self.messages: List[Dict[str, Any]] = []
|
||||
|
||||
def add_user_message(self, content: str) -> "ChatTemplate":
|
||||
"""
|
||||
Add a user message.
|
||||
|
||||
Args:
|
||||
content: Message content
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.messages.append({
|
||||
"role": "user",
|
||||
"content": content,
|
||||
})
|
||||
return self
|
||||
|
||||
def add_agent_message(self, content: str) -> "ChatTemplate":
|
||||
"""
|
||||
Add an agent message.
|
||||
|
||||
Args:
|
||||
content: Message content
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.messages.append({
|
||||
"role": "agent",
|
||||
"content": content,
|
||||
})
|
||||
return self
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate A2UI messages for the chat.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
|
||||
children = []
|
||||
|
||||
# Title if provided
|
||||
if self.title:
|
||||
surface.text(id="chat-title", text=self.title, usage_hint="h1")
|
||||
children.append("chat-title")
|
||||
|
||||
# Message components
|
||||
msg_ids = []
|
||||
for i, msg in enumerate(self.messages):
|
||||
msg_id = f"msg-{i}"
|
||||
role = msg["role"]
|
||||
content = msg["content"]
|
||||
|
||||
# Style based on role
|
||||
if role == "user":
|
||||
surface.text(id=msg_id, text=f"You: {content}", usage_hint="body")
|
||||
else:
|
||||
surface.text(id=msg_id, text=f"Agent: {content}", usage_hint="body")
|
||||
|
||||
msg_ids.append(msg_id)
|
||||
|
||||
# Messages column
|
||||
if msg_ids:
|
||||
surface.column(id="messages", children=msg_ids)
|
||||
children.append("messages")
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=children)
|
||||
|
||||
# Set data
|
||||
surface.set_data("messages", self.messages)
|
||||
|
||||
return surface.to_messages()
|
||||
@@ -1,128 +0,0 @@
|
||||
"""
|
||||
Dashboard Template for A2UI
|
||||
|
||||
Provides a dashboard UI template with multiple panels.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
|
||||
|
||||
class DashboardTemplate:
|
||||
"""
|
||||
Dashboard UI template with multiple panels.
|
||||
|
||||
Creates a dashboard layout with cards for each panel.
|
||||
|
||||
Example:
|
||||
>>> template = DashboardTemplate(surface_id="dashboard", title="My Dashboard")
|
||||
>>> template.add_panel(id="stats", title="Statistics", content="100 items")
|
||||
>>> template.add_panel(id="chart", title="Chart", content="[Chart here]")
|
||||
>>> messages = template.to_messages()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
surface_id: str = "dashboard",
|
||||
title: str = "Dashboard",
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
columns: int = 2,
|
||||
):
|
||||
"""
|
||||
Initialize DashboardTemplate.
|
||||
|
||||
Args:
|
||||
surface_id: Unique ID for the surface
|
||||
title: Dashboard title
|
||||
catalog_id: Component catalog to use
|
||||
columns: Number of columns in the grid
|
||||
"""
|
||||
self.surface_id = surface_id
|
||||
self.title = title
|
||||
self.catalog_id = catalog_id
|
||||
self.columns = columns
|
||||
self.panels: List[Dict[str, Any]] = []
|
||||
|
||||
def add_panel(
|
||||
self,
|
||||
id: str,
|
||||
title: str,
|
||||
content: str,
|
||||
weight: Optional[float] = None,
|
||||
) -> "DashboardTemplate":
|
||||
"""
|
||||
Add a panel to the dashboard.
|
||||
|
||||
Args:
|
||||
id: Panel ID
|
||||
title: Panel title
|
||||
content: Panel content text
|
||||
weight: Optional flex weight
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.panels.append({
|
||||
"id": id,
|
||||
"title": title,
|
||||
"content": content,
|
||||
"weight": weight,
|
||||
})
|
||||
return self
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate A2UI messages for the dashboard.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
|
||||
# Title
|
||||
surface.text(id="dashboard-title", text=self.title, usage_hint="h1")
|
||||
|
||||
# Create cards for each panel
|
||||
card_ids = []
|
||||
for panel in self.panels:
|
||||
panel_id = panel["id"]
|
||||
content_id = f"{panel_id}-content"
|
||||
title_id = f"{panel_id}-title"
|
||||
text_id = f"{panel_id}-text"
|
||||
|
||||
# Panel title
|
||||
surface.text(id=title_id, text=panel["title"], usage_hint="h3")
|
||||
|
||||
# Panel content
|
||||
surface.text(id=text_id, text=panel["content"], usage_hint="body")
|
||||
|
||||
# Panel content column
|
||||
surface.column(id=content_id, children=[title_id, text_id])
|
||||
|
||||
# Panel card
|
||||
surface.card(id=panel_id, child=content_id, weight=panel.get("weight"))
|
||||
card_ids.append(panel_id)
|
||||
|
||||
# Arrange panels in rows based on columns setting
|
||||
row_ids = []
|
||||
for i in range(0, len(card_ids), self.columns):
|
||||
row_cards = card_ids[i:i + self.columns]
|
||||
row_id = f"row-{i // self.columns}"
|
||||
surface.row(id=row_id, children=row_cards)
|
||||
row_ids.append(row_id)
|
||||
|
||||
# Panels container
|
||||
surface.column(id="panels", children=row_ids)
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=["dashboard-title", "panels"])
|
||||
|
||||
# Set data
|
||||
surface.set_data("panels", self.panels)
|
||||
|
||||
return surface.to_messages()
|
||||
@@ -1,205 +0,0 @@
|
||||
"""
|
||||
Form Template for A2UI
|
||||
|
||||
Provides a form UI template for data collection.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import PathBinding
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
|
||||
|
||||
class FormTemplate:
|
||||
"""
|
||||
Form UI template for data collection.
|
||||
|
||||
Creates a form with various input fields and a submit button.
|
||||
|
||||
Example:
|
||||
>>> template = FormTemplate(surface_id="form", title="Contact")
|
||||
>>> template.add_text_field(id="name", label="Name")
|
||||
>>> template.add_email_field(id="email", label="Email")
|
||||
>>> template.set_submit_action("submit_form", "Send")
|
||||
>>> messages = template.to_messages()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
surface_id: str = "form",
|
||||
title: str = "Form",
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
):
|
||||
"""
|
||||
Initialize FormTemplate.
|
||||
|
||||
Args:
|
||||
surface_id: Unique ID for the surface
|
||||
title: Form title
|
||||
catalog_id: Component catalog to use
|
||||
"""
|
||||
self.surface_id = surface_id
|
||||
self.title = title
|
||||
self.catalog_id = catalog_id
|
||||
self.fields: List[Dict[str, Any]] = []
|
||||
self.submit_action = "submit"
|
||||
self.submit_label = "Submit"
|
||||
|
||||
def add_text_field(
|
||||
self,
|
||||
id: str,
|
||||
label: str,
|
||||
default: str = "",
|
||||
required: bool = False,
|
||||
) -> "FormTemplate":
|
||||
"""
|
||||
Add a text input field.
|
||||
|
||||
Args:
|
||||
id: Field ID
|
||||
label: Field label
|
||||
default: Default value
|
||||
required: Whether field is required
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.fields.append({
|
||||
"id": id,
|
||||
"label": label,
|
||||
"type": "text",
|
||||
"default": default,
|
||||
"required": required,
|
||||
})
|
||||
return self
|
||||
|
||||
def add_number_field(
|
||||
self,
|
||||
id: str,
|
||||
label: str,
|
||||
default: float = 0,
|
||||
required: bool = False,
|
||||
) -> "FormTemplate":
|
||||
"""
|
||||
Add a number input field.
|
||||
|
||||
Args:
|
||||
id: Field ID
|
||||
label: Field label
|
||||
default: Default value
|
||||
required: Whether field is required
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.fields.append({
|
||||
"id": id,
|
||||
"label": label,
|
||||
"type": "number",
|
||||
"default": default,
|
||||
"required": required,
|
||||
})
|
||||
return self
|
||||
|
||||
def add_email_field(
|
||||
self,
|
||||
id: str,
|
||||
label: str,
|
||||
default: str = "",
|
||||
required: bool = False,
|
||||
) -> "FormTemplate":
|
||||
"""
|
||||
Add an email input field.
|
||||
|
||||
Args:
|
||||
id: Field ID
|
||||
label: Field label
|
||||
default: Default value
|
||||
required: Whether field is required
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.fields.append({
|
||||
"id": id,
|
||||
"label": label,
|
||||
"type": "email",
|
||||
"default": default,
|
||||
"required": required,
|
||||
})
|
||||
return self
|
||||
|
||||
def set_submit_action(
|
||||
self,
|
||||
action: str,
|
||||
label: str = "Submit",
|
||||
) -> "FormTemplate":
|
||||
"""
|
||||
Set the submit button action.
|
||||
|
||||
Args:
|
||||
action: Action name
|
||||
label: Button label
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
self.submit_action = action
|
||||
self.submit_label = label
|
||||
return self
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate A2UI messages for the form.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
|
||||
# Title
|
||||
surface.text(id="form-title", text=self.title, usage_hint="h1")
|
||||
|
||||
# Fields
|
||||
field_ids = []
|
||||
for field in self.fields:
|
||||
field_id = field["id"]
|
||||
surface.text_field(
|
||||
id=field_id,
|
||||
label=field["label"],
|
||||
text=PathBinding(path=f"/{field_id}"),
|
||||
field_type=field["type"],
|
||||
)
|
||||
field_ids.append(field_id)
|
||||
|
||||
# Set default value
|
||||
surface.set_data(field_id, field.get("default", ""))
|
||||
|
||||
# Submit button
|
||||
surface.text(id="submit-text", text=self.submit_label)
|
||||
|
||||
action_context = [
|
||||
{"key": f["id"], "value": {"path": f"/{f['id']}"}}
|
||||
for f in self.fields
|
||||
]
|
||||
|
||||
surface.button(
|
||||
id="submit-btn",
|
||||
child="submit-text",
|
||||
action_name=self.submit_action,
|
||||
action_context=action_context,
|
||||
primary=True,
|
||||
)
|
||||
field_ids.append("submit-btn")
|
||||
|
||||
# Form fields column
|
||||
surface.column(id="fields", children=field_ids)
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=["form-title", "fields"])
|
||||
|
||||
return surface.to_messages()
|
||||
@@ -1,147 +0,0 @@
|
||||
"""
|
||||
List Template for A2UI
|
||||
|
||||
Provides a list UI template for displaying items.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
|
||||
|
||||
class ListTemplate:
|
||||
"""
|
||||
List UI template for displaying items.
|
||||
|
||||
Creates a list of cards with title, description, and optional image.
|
||||
|
||||
Example:
|
||||
>>> template = ListTemplate(surface_id="list", title="Results")
|
||||
>>> template.add_item(title="Item 1", description="First item")
|
||||
>>> template.add_item(title="Item 2", description="Second item")
|
||||
>>> messages = template.to_messages()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
surface_id: str = "list",
|
||||
title: str = "Items",
|
||||
catalog_id: str = STANDARD_CATALOG_ID,
|
||||
):
|
||||
"""
|
||||
Initialize ListTemplate.
|
||||
|
||||
Args:
|
||||
surface_id: Unique ID for the surface
|
||||
title: List title
|
||||
catalog_id: Component catalog to use
|
||||
"""
|
||||
self.surface_id = surface_id
|
||||
self.title = title
|
||||
self.catalog_id = catalog_id
|
||||
self.items: List[Dict[str, Any]] = []
|
||||
|
||||
def add_item(
|
||||
self,
|
||||
title: str,
|
||||
description: str = "",
|
||||
image_url: Optional[str] = None,
|
||||
action: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
) -> "ListTemplate":
|
||||
"""
|
||||
Add an item to the list.
|
||||
|
||||
Args:
|
||||
title: Item title
|
||||
description: Item description
|
||||
image_url: Optional image URL
|
||||
action: Optional action name for item button
|
||||
metadata: Optional additional metadata
|
||||
|
||||
Returns:
|
||||
self for chaining
|
||||
"""
|
||||
item = {
|
||||
"title": title,
|
||||
"description": description,
|
||||
}
|
||||
if image_url:
|
||||
item["image_url"] = image_url
|
||||
if action:
|
||||
item["action"] = action
|
||||
if metadata:
|
||||
item.update(metadata)
|
||||
|
||||
self.items.append(item)
|
||||
return self
|
||||
|
||||
def to_messages(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Generate A2UI messages for the list.
|
||||
|
||||
Returns:
|
||||
List of A2UI message dictionaries
|
||||
"""
|
||||
surface = Surface(
|
||||
surface_id=self.surface_id,
|
||||
catalog_id=self.catalog_id,
|
||||
)
|
||||
|
||||
# Title
|
||||
surface.text(id="list-title", text=self.title, usage_hint="h1")
|
||||
|
||||
# Create cards for each item
|
||||
card_ids = []
|
||||
for i, item in enumerate(self.items):
|
||||
card_id = f"card-{i}"
|
||||
content_id = f"content-{i}"
|
||||
title_id = f"title-{i}"
|
||||
desc_id = f"desc-{i}"
|
||||
|
||||
# Item title
|
||||
surface.text(id=title_id, text=item["title"], usage_hint="h3")
|
||||
|
||||
# Item description
|
||||
surface.text(id=desc_id, text=item.get("description", ""), usage_hint="body")
|
||||
|
||||
content_children = [title_id, desc_id]
|
||||
|
||||
# Optional image
|
||||
if item.get("image_url"):
|
||||
img_id = f"img-{i}"
|
||||
surface.image(id=img_id, url=item["image_url"], usage_hint="smallFeature")
|
||||
content_children.insert(0, img_id)
|
||||
|
||||
# Optional action button
|
||||
if item.get("action"):
|
||||
btn_id = f"btn-{i}"
|
||||
btn_text_id = f"btn-text-{i}"
|
||||
surface.text(id=btn_text_id, text="Select")
|
||||
surface.button(
|
||||
id=btn_id,
|
||||
child=btn_text_id,
|
||||
action_name=item["action"],
|
||||
action_context=[{"key": "index", "value": i}],
|
||||
primary=True,
|
||||
)
|
||||
content_children.append(btn_id)
|
||||
|
||||
# Card content
|
||||
surface.column(id=content_id, children=content_children)
|
||||
|
||||
# Card
|
||||
surface.card(id=card_id, child=content_id)
|
||||
card_ids.append(card_id)
|
||||
|
||||
# Items column
|
||||
surface.column(id="items", children=card_ids)
|
||||
|
||||
# Root column
|
||||
surface.column(id="root", children=["list-title", "items"])
|
||||
|
||||
# Set data
|
||||
surface.set_data("items", self.items)
|
||||
|
||||
return surface.to_messages()
|
||||
@@ -1,523 +0,0 @@
|
||||
"""
|
||||
A2UI Protocol Types
|
||||
|
||||
Pydantic-style classes for the A2UI (Agent-to-User Interface) protocol.
|
||||
Based on the A2UI Specification v0.9.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from enum import Enum
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Data Binding Types
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class PathBinding:
|
||||
"""
|
||||
A reference to a value in the data model.
|
||||
|
||||
Used for data binding in A2UI components.
|
||||
"""
|
||||
path: str
|
||||
|
||||
def to_dict(self) -> Dict[str, str]:
|
||||
return {"path": self.path}
|
||||
|
||||
|
||||
# Type alias for values that can be literal or path-bound
|
||||
StringOrPath = Union[str, PathBinding]
|
||||
NumberOrPath = Union[int, float, PathBinding]
|
||||
BooleanOrPath = Union[bool, PathBinding]
|
||||
|
||||
|
||||
def resolve_string_or_path(value: StringOrPath) -> Union[str, Dict[str, str]]:
|
||||
"""Resolve a StringOrPath to its serialized form."""
|
||||
if isinstance(value, PathBinding):
|
||||
return value.to_dict()
|
||||
return value
|
||||
|
||||
|
||||
def resolve_number_or_path(value: NumberOrPath) -> Union[int, float, Dict[str, str]]:
|
||||
"""Resolve a NumberOrPath to its serialized form."""
|
||||
if isinstance(value, PathBinding):
|
||||
return value.to_dict()
|
||||
return value
|
||||
|
||||
|
||||
def resolve_boolean_or_path(value: BooleanOrPath) -> Union[bool, Dict[str, str]]:
|
||||
"""Resolve a BooleanOrPath to its serialized form."""
|
||||
if isinstance(value, PathBinding):
|
||||
return value.to_dict()
|
||||
return value
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Action Types
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class ActionContext:
|
||||
"""
|
||||
A key-value pair for action context.
|
||||
|
||||
Used to pass data from UI to action handlers.
|
||||
"""
|
||||
key: str
|
||||
value: Union[str, int, float, bool, PathBinding]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
if isinstance(self.value, PathBinding):
|
||||
return {"key": self.key, "value": self.value.to_dict()}
|
||||
return {"key": self.key, "value": self.value}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Action:
|
||||
"""
|
||||
An action triggered by a UI component (e.g., button click).
|
||||
"""
|
||||
name: str
|
||||
context: List[ActionContext] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"context": [c.to_dict() for c in self.context]
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Children Types
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class ChildrenTemplate:
|
||||
"""
|
||||
A template for generating dynamic children from a data model list.
|
||||
"""
|
||||
component_id: str
|
||||
path: str
|
||||
|
||||
def to_dict(self) -> Dict[str, str]:
|
||||
return {
|
||||
"componentId": self.component_id,
|
||||
"path": self.path
|
||||
}
|
||||
|
||||
|
||||
# Type alias for children - either explicit list or template
|
||||
ChildrenType = Union[List[str], ChildrenTemplate]
|
||||
|
||||
|
||||
def resolve_children(children: ChildrenType) -> Union[List[str], Dict[str, str]]:
|
||||
"""Resolve children to serialized form."""
|
||||
if isinstance(children, ChildrenTemplate):
|
||||
return children.to_dict()
|
||||
return children
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Component Types
|
||||
# =============================================================================
|
||||
|
||||
class UsageHint(str, Enum):
|
||||
"""Text usage hints for styling."""
|
||||
H1 = "h1"
|
||||
H2 = "h2"
|
||||
H3 = "h3"
|
||||
H4 = "h4"
|
||||
H5 = "h5"
|
||||
CAPTION = "caption"
|
||||
BODY = "body"
|
||||
|
||||
|
||||
class ImageFit(str, Enum):
|
||||
"""Image fit options."""
|
||||
CONTAIN = "contain"
|
||||
COVER = "cover"
|
||||
FILL = "fill"
|
||||
NONE = "none"
|
||||
SCALE_DOWN = "scale-down"
|
||||
|
||||
|
||||
class ImageUsageHint(str, Enum):
|
||||
"""Image usage hints."""
|
||||
ICON = "icon"
|
||||
AVATAR = "avatar"
|
||||
SMALL_FEATURE = "smallFeature"
|
||||
MEDIUM_FEATURE = "mediumFeature"
|
||||
LARGE_FEATURE = "largeFeature"
|
||||
HEADER = "header"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextComponent:
|
||||
"""Text display component."""
|
||||
id: str
|
||||
text: StringOrPath
|
||||
usage_hint: Optional[str] = None
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Text",
|
||||
"text": resolve_string_or_path(self.text)
|
||||
}
|
||||
if self.usage_hint:
|
||||
result["usageHint"] = self.usage_hint
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageComponent:
|
||||
"""Image display component."""
|
||||
id: str
|
||||
url: StringOrPath
|
||||
fit: Optional[str] = None
|
||||
usage_hint: Optional[str] = None
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Image",
|
||||
"url": resolve_string_or_path(self.url)
|
||||
}
|
||||
if self.fit:
|
||||
result["fit"] = self.fit
|
||||
if self.usage_hint:
|
||||
result["usageHint"] = self.usage_hint
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class IconComponent:
|
||||
"""Icon display component."""
|
||||
id: str
|
||||
name: StringOrPath
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Icon",
|
||||
"name": resolve_string_or_path(self.name)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ButtonComponent:
|
||||
"""Button component with action."""
|
||||
id: str
|
||||
child: str
|
||||
action: Optional[Action] = None
|
||||
primary: bool = False
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Button",
|
||||
"child": self.child
|
||||
}
|
||||
if self.primary:
|
||||
result["primary"] = True
|
||||
if self.action:
|
||||
result["action"] = self.action.to_dict()
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class CardComponent:
|
||||
"""Card container component."""
|
||||
id: str
|
||||
child: str
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Card",
|
||||
"child": self.child
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class RowComponent:
|
||||
"""Horizontal layout component."""
|
||||
id: str
|
||||
children: ChildrenType
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Row",
|
||||
"children": resolve_children(self.children)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ColumnComponent:
|
||||
"""Vertical layout component."""
|
||||
id: str
|
||||
children: ChildrenType
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Column",
|
||||
"children": resolve_children(self.children)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ListComponent:
|
||||
"""List component for repeating items."""
|
||||
id: str
|
||||
children: ChildrenType
|
||||
direction: str = "vertical"
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "List",
|
||||
"direction": self.direction,
|
||||
"children": resolve_children(self.children)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextFieldComponent:
|
||||
"""Text input field component."""
|
||||
id: str
|
||||
label: StringOrPath
|
||||
text: StringOrPath
|
||||
field_type: str = "text"
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "TextField",
|
||||
"label": resolve_string_or_path(self.label),
|
||||
"text": resolve_string_or_path(self.text),
|
||||
"type": self.field_type
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckBoxComponent:
|
||||
"""Checkbox input component."""
|
||||
id: str
|
||||
label: StringOrPath
|
||||
checked: BooleanOrPath = False
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "CheckBox",
|
||||
"label": resolve_string_or_path(self.label),
|
||||
"checked": resolve_boolean_or_path(self.checked)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class SliderComponent:
|
||||
"""Slider input component."""
|
||||
id: str
|
||||
min_value: NumberOrPath
|
||||
max_value: NumberOrPath
|
||||
value: NumberOrPath
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "Slider",
|
||||
"min": resolve_number_or_path(self.min_value),
|
||||
"max": resolve_number_or_path(self.max_value),
|
||||
"value": resolve_number_or_path(self.value)
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class DateTimeInputComponent:
|
||||
"""Date/time input component."""
|
||||
id: str
|
||||
label: StringOrPath
|
||||
value: StringOrPath
|
||||
enable_date: bool = True
|
||||
enable_time: bool = False
|
||||
weight: Optional[float] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"id": self.id,
|
||||
"component": "DateTimeInput",
|
||||
"label": resolve_string_or_path(self.label),
|
||||
"value": resolve_string_or_path(self.value),
|
||||
"enableDate": self.enable_date,
|
||||
"enableTime": self.enable_time
|
||||
}
|
||||
if self.weight is not None:
|
||||
result["weight"] = self.weight
|
||||
return result
|
||||
|
||||
|
||||
# Type alias for any component
|
||||
Component = Union[
|
||||
TextComponent,
|
||||
ImageComponent,
|
||||
IconComponent,
|
||||
ButtonComponent,
|
||||
CardComponent,
|
||||
RowComponent,
|
||||
ColumnComponent,
|
||||
ListComponent,
|
||||
TextFieldComponent,
|
||||
CheckBoxComponent,
|
||||
SliderComponent,
|
||||
DateTimeInputComponent,
|
||||
]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Message Types
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class CreateSurfaceMessage:
|
||||
"""
|
||||
Message to create a new UI surface.
|
||||
|
||||
Signals the client to create a new surface and begin rendering it.
|
||||
"""
|
||||
surface_id: str
|
||||
catalog_id: str
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"createSurface": {
|
||||
"surfaceId": self.surface_id,
|
||||
"catalogId": self.catalog_id
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class UpdateComponentsMessage:
|
||||
"""
|
||||
Message to update surface components.
|
||||
|
||||
Updates a surface with a new set of components.
|
||||
"""
|
||||
surface_id: str
|
||||
components: List[Component]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"updateComponents": {
|
||||
"surfaceId": self.surface_id,
|
||||
"components": [c.to_dict() for c in self.components]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class UpdateDataModelMessage:
|
||||
"""
|
||||
Message to update the data model.
|
||||
|
||||
Updates the data model for an existing surface.
|
||||
"""
|
||||
surface_id: str
|
||||
value: Any
|
||||
path: str = "/"
|
||||
op: str = "replace"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {
|
||||
"updateDataModel": {
|
||||
"surfaceId": self.surface_id,
|
||||
"path": self.path,
|
||||
"op": self.op,
|
||||
"value": self.value
|
||||
}
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class DeleteSurfaceMessage:
|
||||
"""
|
||||
Message to delete a surface.
|
||||
|
||||
Signals the client to delete the surface.
|
||||
"""
|
||||
surface_id: str
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"deleteSurface": {
|
||||
"surfaceId": self.surface_id
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Data Part Type
|
||||
# =============================================================================
|
||||
|
||||
@dataclass
|
||||
class A2UIDataPart:
|
||||
"""
|
||||
A2UI data part for A2A integration.
|
||||
|
||||
Wraps A2UI data with the appropriate MIME type.
|
||||
"""
|
||||
data: Dict[str, Any]
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
result: Dict[str, Any] = {"data": self.data}
|
||||
if self.metadata:
|
||||
result["metadata"] = self.metadata
|
||||
return result
|
||||
@@ -4,12 +4,15 @@ AG-UI Event Streaming
|
||||
Provides utilities for streaming PraisonAI agent responses as AG-UI events.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Set
|
||||
|
||||
from praisonaiagents.ui.agui.types import (
|
||||
BaseEvent,
|
||||
CustomEvent,
|
||||
TextMessageStartEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
@@ -24,6 +27,7 @@ from praisonaiagents.ui.agui.types import (
|
||||
StepFinishedEvent,
|
||||
StateSnapshotEvent,
|
||||
)
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -256,6 +260,75 @@ async def async_stream_response(
|
||||
yield create_run_error_event(str(e))
|
||||
|
||||
|
||||
def stream_event_to_agui_events(
|
||||
event: Any,
|
||||
message_id: str,
|
||||
buffer: EventBuffer,
|
||||
) -> List[BaseEvent]:
|
||||
"""Convert a PraisonAI StreamEvent into zero or more AG-UI events."""
|
||||
from praisonaiagents.streaming.events import StreamEventType
|
||||
|
||||
if event.type == StreamEventType.DELTA_TEXT and event.content:
|
||||
if not buffer.current_text_message_id:
|
||||
buffer.start_text_message()
|
||||
return [
|
||||
TextMessageStartEvent(message_id=buffer.current_text_message_id, role="assistant"),
|
||||
TextMessageContentEvent(
|
||||
message_id=buffer.current_text_message_id, delta=event.content
|
||||
),
|
||||
]
|
||||
return [
|
||||
TextMessageContentEvent(
|
||||
message_id=buffer.current_text_message_id, delta=event.content
|
||||
)
|
||||
]
|
||||
|
||||
if event.type == StreamEventType.TOOL_CALL_START and event.tool_call:
|
||||
tool_call_id = event.tool_call.get("id") or str(uuid.uuid4())
|
||||
tool_name = event.tool_call.get("name", "tool")
|
||||
arguments = event.tool_call.get("arguments", {})
|
||||
args_str = json.dumps(arguments) if isinstance(arguments, dict) else str(arguments)
|
||||
parent_id = buffer.get_parent_message_id_for_tool_call() or message_id
|
||||
buffer.start_tool_call(tool_call_id)
|
||||
return list(create_tool_call_events(tool_call_id, tool_name, args_str, parent_id))
|
||||
|
||||
if event.type == StreamEventType.TOOL_CALL_RESULT and event.tool_call:
|
||||
tool_call_id = event.tool_call.get("id") or str(uuid.uuid4())
|
||||
raw_result = event.tool_call.get("result")
|
||||
content = event.content if event.content is not None else str(raw_result or "")
|
||||
buffer.end_tool_call(tool_call_id)
|
||||
events: List[BaseEvent] = [
|
||||
create_tool_result_event(tool_call_id, str(content), message_id)
|
||||
]
|
||||
if isinstance(raw_result, dict) and raw_result.get("mime_type") == A2UI_MIME_TYPE:
|
||||
events.append(
|
||||
CustomEvent(
|
||||
name="a2ui",
|
||||
value={
|
||||
"mime_type": A2UI_MIME_TYPE,
|
||||
"messages": raw_result.get("messages", []),
|
||||
"surface_id": _infer_a2ui_surface_id(raw_result),
|
||||
},
|
||||
)
|
||||
)
|
||||
return events
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def _infer_a2ui_surface_id(result: Dict[str, Any]) -> str:
|
||||
"""Infer surface id from send_a2ui_messages result (inline, no UI-layer deps)."""
|
||||
for msg in result.get("messages") or []:
|
||||
if isinstance(msg, dict):
|
||||
create = msg.get("createSurface")
|
||||
if isinstance(create, dict):
|
||||
sid = create.get("surfaceId") or create.get("id")
|
||||
if sid:
|
||||
return str(sid)
|
||||
sid = result.get("surface_id") or result.get("surfaceId")
|
||||
return str(sid) if sid else "main"
|
||||
|
||||
|
||||
async def async_stream_agent_response(
|
||||
agent,
|
||||
user_input: str,
|
||||
@@ -278,37 +351,63 @@ async def async_stream_agent_response(
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
buffer = EventBuffer()
|
||||
message_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Check if agent has async chat method
|
||||
if hasattr(agent, 'achat'):
|
||||
response = await agent.achat(user_input)
|
||||
elif hasattr(agent, 'chat'):
|
||||
# Run sync chat in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
response = await loop.run_in_executor(None, agent.chat, user_input)
|
||||
else:
|
||||
text_started = False
|
||||
text_ended = False
|
||||
original_emitter = getattr(agent, "stream_emitter", None)
|
||||
|
||||
def on_stream_event(event: Any) -> None:
|
||||
loop.call_soon_threadsafe(queue.put_nowait, ("stream", event))
|
||||
|
||||
if original_emitter is not None:
|
||||
original_emitter.add_callback(on_stream_event)
|
||||
|
||||
async def run_chat() -> Any:
|
||||
try:
|
||||
if hasattr(agent, "achat"):
|
||||
return await agent.achat(user_input, stream=True)
|
||||
if hasattr(agent, "chat"):
|
||||
return await loop.run_in_executor(
|
||||
None, lambda: agent.chat(user_input, stream=True)
|
||||
)
|
||||
raise ValueError("Agent must have 'chat' or 'achat' method")
|
||||
|
||||
# Emit text message events
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
|
||||
if response:
|
||||
finally:
|
||||
loop.call_soon_threadsafe(queue.put_nowait, ("done", None))
|
||||
|
||||
chat_task = asyncio.create_task(run_chat())
|
||||
|
||||
try:
|
||||
while True:
|
||||
kind, payload = await queue.get()
|
||||
if kind == "done":
|
||||
break
|
||||
for agui_event in stream_event_to_agui_events(payload, message_id, buffer):
|
||||
if isinstance(agui_event, TextMessageStartEvent):
|
||||
text_started = True
|
||||
yield agui_event
|
||||
|
||||
response = await chat_task
|
||||
|
||||
if text_started and buffer.current_text_message_id and not text_ended:
|
||||
yield TextMessageEndEvent(message_id=buffer.current_text_message_id)
|
||||
text_ended = True
|
||||
elif not text_started and response:
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=str(response))
|
||||
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
|
||||
except Exception as e:
|
||||
yield create_run_error_event(str(e))
|
||||
finally:
|
||||
if original_emitter is not None:
|
||||
original_emitter.remove_callback(on_stream_event)
|
||||
|
||||
|
||||
async def async_stream_agents_response(
|
||||
@@ -331,8 +430,6 @@ async def async_stream_agents_response(
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
@@ -344,7 +441,7 @@ async def async_stream_agents_response(
|
||||
result = await agents.astart(user_input)
|
||||
elif hasattr(agents, 'start'):
|
||||
# Run sync start in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
loop = asyncio.get_running_loop()
|
||||
result = await loop.run_in_executor(None, agents.start, user_input)
|
||||
else:
|
||||
raise ValueError("Agents must have 'start' or 'astart' method")
|
||||
|
||||
@@ -133,6 +133,13 @@ class StateSnapshotEvent(BaseEvent):
|
||||
snapshot: Dict[str, Any]
|
||||
|
||||
|
||||
class CustomEvent(BaseEvent):
|
||||
"""Custom extension event (e.g. structured A2UI payloads)."""
|
||||
type: Literal["CUSTOM"] = "CUSTOM"
|
||||
name: str
|
||||
value: Any
|
||||
|
||||
|
||||
# Message Types
|
||||
class FunctionCall(BaseModel):
|
||||
"""Function call in a tool call."""
|
||||
@@ -175,6 +182,7 @@ Event = Union[
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
StateSnapshotEvent,
|
||||
CustomEvent,
|
||||
RunStartedEvent,
|
||||
RunFinishedEvent,
|
||||
RunErrorEvent,
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
UI protocol types — lightweight contracts for agent-to-UI integration.
|
||||
|
||||
Zero runtime cost: TypedDict and constants only. No optional dependencies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, TypedDict, Union
|
||||
|
||||
# Standard MIME type for A2UI payloads (Google A2UI / A2A transport).
|
||||
A2UI_MIME_TYPE = "application/json+a2ui"
|
||||
|
||||
|
||||
class A2UIToolResultProtocol(TypedDict, total=False):
|
||||
"""
|
||||
Output shape from ``send_a2ui_messages`` — the integrator contract for any UI.
|
||||
|
||||
Example:
|
||||
if result.get("mime_type") == A2UI_MIME_TYPE:
|
||||
render(result["messages"])
|
||||
"""
|
||||
|
||||
mime_type: str
|
||||
messages: List[Dict[str, Any]]
|
||||
a2ui_part: Union[Dict[str, Any], List[Any]]
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "1.6.41"
|
||||
version = "1.6.48"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
@@ -106,6 +106,12 @@ sandbox = [
|
||||
"e2b-code-interpreter>=1.0.0",
|
||||
]
|
||||
|
||||
# Declarative agent UI (Google A2UI SDK — lazy loaded)
|
||||
a2ui = [
|
||||
"a2ui-agent-sdk>=0.2.1",
|
||||
"a2a-sdk>=0.3.0,<1.0.0",
|
||||
]
|
||||
|
||||
# Docker sandbox dependencies
|
||||
sandbox-docker = [
|
||||
"docker>=7.0.0",
|
||||
@@ -126,7 +132,8 @@ all = [
|
||||
"praisonaiagents[search]",
|
||||
"praisonaiagents[crawl]",
|
||||
"praisonaiagents[autonomy]",
|
||||
"praisonaiagents[os]"
|
||||
"praisonaiagents[os]",
|
||||
"praisonaiagents[a2ui]",
|
||||
]
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
|
||||
@@ -233,6 +233,47 @@ class TestA2AJsonRpc:
|
||||
assert result["status"]["state"] == "completed"
|
||||
assert result["artifacts"] is not None
|
||||
assert len(result["artifacts"]) > 0
|
||||
|
||||
def test_post_a2a_message_send_a2ui_artifact(self):
|
||||
"""message/send with A2UI tool result in chat_history → DataPart artifact."""
|
||||
import json
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
app, a2a = _make_app()
|
||||
a2ui_payload = {
|
||||
"mime_type": A2UI_MIME_TYPE,
|
||||
"messages": [{"createSurface": {"surfaceId": "main", "catalogId": "basic"}}],
|
||||
"a2ui_part": {"messages": [{"createSurface": {"surfaceId": "main", "catalogId": "basic"}}]},
|
||||
}
|
||||
a2a.agent.chat = MagicMock(return_value="Surface ready.")
|
||||
a2a.agent.chat_history = [
|
||||
{"role": "tool", "content": json.dumps(a2ui_payload)},
|
||||
]
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.post("/a2a", json={
|
||||
"jsonrpc": "2.0",
|
||||
"method": "message/send",
|
||||
"id": "a2ui-1",
|
||||
"params": {
|
||||
"message": {
|
||||
"messageId": "msg-a2ui",
|
||||
"role": "user",
|
||||
"parts": [{"text": "Show UI"}],
|
||||
}
|
||||
},
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()["result"]
|
||||
artifact = result["artifacts"][0]
|
||||
parts = artifact["parts"]
|
||||
data_parts = [p for p in parts if "data" in p]
|
||||
assert data_parts, "Expected DataPart in artifact"
|
||||
assert data_parts[0]["metadata"]["mimeType"] == A2UI_MIME_TYPE
|
||||
|
||||
def test_post_a2a_message_send_creates_task_in_store(self):
|
||||
"""message/send should create a task in TaskStore."""
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Tests for A2UI artifact conversion in A2A transport."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
|
||||
SAMPLE_A2UI_RESULT = {
|
||||
"mime_type": A2UI_MIME_TYPE,
|
||||
"messages": [{"createSurface": {"surfaceId": "main", "catalogId": "basic"}}],
|
||||
"a2ui_part": {"messages": [{"createSurface": {"surfaceId": "main", "catalogId": "basic"}}]},
|
||||
}
|
||||
|
||||
|
||||
class TestFindA2uiToolResult:
|
||||
def test_finds_last_a2ui_tool_message(self):
|
||||
from praisonaiagents.ui.a2a.conversion import find_a2ui_tool_result
|
||||
|
||||
history = [
|
||||
{"role": "user", "content": "show ui"},
|
||||
{"role": "tool", "content": json.dumps({"mime_type": "text/plain", "data": "x"})},
|
||||
{"role": "tool", "content": json.dumps(SAMPLE_A2UI_RESULT)},
|
||||
]
|
||||
result = find_a2ui_tool_result(history)
|
||||
assert result is not None
|
||||
assert result["mime_type"] == A2UI_MIME_TYPE
|
||||
|
||||
def test_returns_none_without_a2ui(self):
|
||||
from praisonaiagents.ui.a2a.conversion import find_a2ui_tool_result
|
||||
|
||||
assert find_a2ui_tool_result([]) is None
|
||||
assert find_a2ui_tool_result([{"role": "tool", "content": "plain text"}]) is None
|
||||
|
||||
|
||||
class TestCreateA2uiArtifact:
|
||||
def test_creates_datapart_with_mime(self):
|
||||
from praisonaiagents.ui.a2a.conversion import create_a2ui_artifact
|
||||
from praisonaiagents.ui.a2a.types import DataPart
|
||||
|
||||
artifact = create_a2ui_artifact(SAMPLE_A2UI_RESULT, name="ui-surface")
|
||||
assert artifact.name == "ui-surface"
|
||||
assert len(artifact.parts) == 1
|
||||
part = artifact.parts[0]
|
||||
assert isinstance(part, DataPart)
|
||||
assert part.metadata.get("mimeType") == A2UI_MIME_TYPE
|
||||
assert "messages" in part.data
|
||||
|
||||
|
||||
class TestCreateResponseArtifact:
|
||||
def test_prefers_a2ui_from_history(self):
|
||||
from praisonaiagents.ui.a2a.conversion import create_response_artifact
|
||||
from praisonaiagents.ui.a2a.types import DataPart
|
||||
|
||||
history = [{"role": "tool", "content": json.dumps(SAMPLE_A2UI_RESULT)}]
|
||||
artifact = create_response_artifact("Done.", chat_history=history)
|
||||
assert isinstance(artifact.parts[0], DataPart)
|
||||
|
||||
def test_falls_back_to_text(self):
|
||||
from praisonaiagents.ui.a2a.conversion import create_response_artifact
|
||||
from praisonaiagents.ui.a2a.types import TextPart
|
||||
|
||||
artifact = create_response_artifact("Hello", chat_history=[])
|
||||
assert isinstance(artifact.parts[0], TextPart)
|
||||
assert artifact.parts[0].text == "Hello"
|
||||
@@ -1 +0,0 @@
|
||||
"""A2UI Tests Package."""
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Tests for send_a2ui_messages minimal arg unwrap."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.tools.a2ui_tools import _unwrap_messages
|
||||
|
||||
try:
|
||||
import a2ui # noqa: F401
|
||||
|
||||
HAS_A2UI_SDK = True
|
||||
except ImportError:
|
||||
HAS_A2UI_SDK = False
|
||||
|
||||
SAMPLE_MSG = {"createSurface": {"surfaceId": "main", "catalogId": "basic"}}
|
||||
|
||||
|
||||
class TestUnwrapMessages:
|
||||
def test_list_input(self):
|
||||
assert _unwrap_messages([SAMPLE_MSG]) == [SAMPLE_MSG]
|
||||
|
||||
def test_json_string_list(self):
|
||||
assert _unwrap_messages(json.dumps([SAMPLE_MSG])) == [SAMPLE_MSG]
|
||||
|
||||
def test_wrapped_messages_dict(self):
|
||||
assert _unwrap_messages({"messages": [SAMPLE_MSG]}) == [SAMPLE_MSG]
|
||||
|
||||
def test_json_string_wrapped_object(self):
|
||||
assert _unwrap_messages(json.dumps({"messages": [SAMPLE_MSG]})) == [SAMPLE_MSG]
|
||||
|
||||
def test_single_message_dict(self):
|
||||
assert _unwrap_messages(SAMPLE_MSG) == [SAMPLE_MSG]
|
||||
|
||||
def test_invalid_json_raises(self):
|
||||
with pytest.raises(ValueError, match="not valid JSON"):
|
||||
_unwrap_messages("{not json")
|
||||
|
||||
def test_invalid_shape_raises(self):
|
||||
with pytest.raises(ValueError, match="must be"):
|
||||
_unwrap_messages(123)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_A2UI_SDK, reason="a2ui-agent-sdk not installed")
|
||||
class TestSendA2uiMessagesIntegration:
|
||||
def test_full_tool_call(self):
|
||||
from praisonaiagents.tools.a2ui_tools import send_a2ui_messages
|
||||
|
||||
result = send_a2ui_messages(messages=[SAMPLE_MSG])
|
||||
assert result["mime_type"] == "application/json+a2ui"
|
||||
assert len(result["messages"]) == 1
|
||||
@@ -0,0 +1,60 @@
|
||||
"""Tests for A2UI adapter (optional a2ui-agent-sdk dependency)."""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class TestA2UIAdapterImport:
|
||||
def test_a2ui_class_export(self):
|
||||
from praisonaiagents.ui.a2ui import A2UI
|
||||
|
||||
assert A2UI is not None
|
||||
assert hasattr(A2UI, "create_part")
|
||||
|
||||
def test_ui_lazy_a2ui_export(self):
|
||||
from praisonaiagents.ui import A2UI
|
||||
|
||||
assert A2UI is not None
|
||||
|
||||
def test_missing_optional_dep_raises_clear_error(self):
|
||||
pytest.importorskip("a2ui", reason="a2ui-agent-sdk not installed")
|
||||
# If installed, adapter imports should work
|
||||
from praisonaiagents.ui.a2ui.adapter import create_a2ui_part
|
||||
|
||||
part = create_a2ui_part({"createSurface": {"surfaceId": "main", "catalogId": "test"}})
|
||||
assert part is not None
|
||||
|
||||
|
||||
class TestA2UITool:
|
||||
def test_send_a2ui_messages_requires_sdk(self):
|
||||
from praisonaiagents.tools.a2ui_tools import send_a2ui_messages
|
||||
|
||||
if pytest.importorskip("a2ui", reason="skip without sdk"):
|
||||
result = send_a2ui_messages(
|
||||
messages=[{"createSurface": {"surfaceId": "x", "catalogId": "c"}}]
|
||||
)
|
||||
assert result["mime_type"] == "application/json+a2ui"
|
||||
assert len(result["messages"]) == 1
|
||||
|
||||
|
||||
class TestA2UIAdapterSmoke:
|
||||
def test_is_a2ui_part_with_sdk(self):
|
||||
pytest.importorskip("a2ui", reason="a2ui-agent-sdk not installed")
|
||||
from praisonaiagents.ui.a2ui.adapter import create_a2ui_part, is_a2ui_part
|
||||
|
||||
part = create_a2ui_part({"messages": [{"createSurface": {"surfaceId": "m"}}]})
|
||||
assert is_a2ui_part(part) is True
|
||||
|
||||
def test_missing_sdk_import_error_message(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
from praisonaiagents.ui.a2ui import adapter
|
||||
|
||||
def _raise(_name):
|
||||
raise ImportError(
|
||||
"A2UI support requires the optional dependency. "
|
||||
"Install with: pip install praisonaiagents[a2ui]"
|
||||
)
|
||||
|
||||
with patch.object(adapter, "_import", side_effect=_raise):
|
||||
with pytest.raises(ImportError, match="praisonaiagents\\[a2ui\\]"):
|
||||
adapter.create_a2ui_part({"messages": []})
|
||||
@@ -1,202 +0,0 @@
|
||||
"""
|
||||
Tests for A2UI Agent Integration
|
||||
|
||||
TDD: Write tests first, then implement the A2UIAgent class.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
|
||||
|
||||
# Skip tests if no API key
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY"),
|
||||
reason="OPENAI_API_KEY not set"
|
||||
)
|
||||
|
||||
|
||||
class TestA2UIAgent:
|
||||
"""Test A2UIAgent wrapper class."""
|
||||
|
||||
def test_a2ui_agent_creation(self):
|
||||
"""Test A2UIAgent creation with an agent."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
role="Test assistant",
|
||||
goal="Help with testing"
|
||||
)
|
||||
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
assert a2ui_agent.agent is agent
|
||||
assert a2ui_agent.surface_id is not None
|
||||
|
||||
def test_a2ui_agent_with_custom_surface_id(self):
|
||||
"""Test A2UIAgent with custom surface ID."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent, surface_id="custom-surface")
|
||||
|
||||
assert a2ui_agent.surface_id == "custom-surface"
|
||||
|
||||
def test_render_text_response(self):
|
||||
"""Test rendering a simple text response as A2UI."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
# Render a simple text response
|
||||
messages = a2ui_agent.render_text("Hello, World!")
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2 # createSurface + updateComponents
|
||||
|
||||
# Check createSurface
|
||||
assert "createSurface" in messages[0]
|
||||
|
||||
# Check updateComponents has a text component
|
||||
update_msg = next((m for m in messages if "updateComponents" in m), None)
|
||||
assert update_msg is not None
|
||||
components = update_msg["updateComponents"]["components"]
|
||||
assert any(c.get("component") == "Text" for c in components)
|
||||
|
||||
def test_render_list_response(self):
|
||||
"""Test rendering a list of items as A2UI cards."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
items = [
|
||||
{"title": "Item 1", "description": "First item"},
|
||||
{"title": "Item 2", "description": "Second item"},
|
||||
]
|
||||
|
||||
messages = a2ui_agent.render_list(
|
||||
title="My Items",
|
||||
items=items,
|
||||
item_title_key="title",
|
||||
item_description_key="description"
|
||||
)
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2
|
||||
|
||||
def test_render_form(self):
|
||||
"""Test rendering a form as A2UI."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
fields = [
|
||||
{"id": "name", "label": "Name", "type": "text"},
|
||||
{"id": "email", "label": "Email", "type": "text"},
|
||||
]
|
||||
|
||||
messages = a2ui_agent.render_form(
|
||||
title="Contact Form",
|
||||
fields=fields,
|
||||
submit_action="submit_form"
|
||||
)
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2
|
||||
|
||||
# Check for TextField components
|
||||
update_msg = next((m for m in messages if "updateComponents" in m), None)
|
||||
assert update_msg is not None
|
||||
components = update_msg["updateComponents"]["components"]
|
||||
text_fields = [c for c in components if c.get("component") == "TextField"]
|
||||
assert len(text_fields) >= 2
|
||||
|
||||
def test_get_surface(self):
|
||||
"""Test getting the underlying Surface object."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
surface = a2ui_agent.get_surface()
|
||||
|
||||
assert isinstance(surface, Surface)
|
||||
assert surface.surface_id == a2ui_agent.surface_id
|
||||
|
||||
|
||||
class TestA2UIAgentWithLLM:
|
||||
"""Test A2UIAgent with actual LLM calls."""
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_with_ui(self):
|
||||
"""Test chat that returns A2UI response."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
|
||||
agent = Agent(
|
||||
name="UIAgent",
|
||||
role="UI Generator",
|
||||
goal="Generate helpful UI responses",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
# This should work with the agent's chat
|
||||
response = await a2ui_agent.chat_async("Say hello")
|
||||
|
||||
assert response is not None
|
||||
# Response should contain text
|
||||
assert isinstance(response, str) or isinstance(response, dict)
|
||||
|
||||
|
||||
class TestA2UIAgentIntegration:
|
||||
"""Integration tests for A2UIAgent with A2A."""
|
||||
|
||||
def test_get_a2a_data_part(self):
|
||||
"""Test getting A2UI as A2A DataPart."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
from praisonaiagents.ui.a2ui.extension import is_a2ui_part
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
# Render something
|
||||
a2ui_agent.render_text("Hello")
|
||||
|
||||
# Get as A2A DataPart
|
||||
part = a2ui_agent.to_a2a_part()
|
||||
|
||||
assert part is not None
|
||||
assert is_a2ui_part(part)
|
||||
|
||||
def test_to_json(self):
|
||||
"""Test getting A2UI as JSON string."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui.agent import A2UIAgent
|
||||
import json
|
||||
|
||||
agent = Agent(name="TestAgent", role="Test", goal="Test")
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
# Render something
|
||||
a2ui_agent.render_text("Hello")
|
||||
|
||||
# Get as JSON
|
||||
json_str = a2ui_agent.to_json()
|
||||
|
||||
# Should be valid JSON
|
||||
parsed = json.loads(json_str)
|
||||
assert isinstance(parsed, list)
|
||||
@@ -1,115 +0,0 @@
|
||||
"""
|
||||
Tests for A2UI Extension Module
|
||||
|
||||
TDD: Write tests first, then implement the extension.
|
||||
"""
|
||||
|
||||
|
||||
class TestA2UIExtension:
|
||||
"""Test A2UI extension helpers for A2A integration."""
|
||||
|
||||
def test_a2ui_mime_type(self):
|
||||
"""Test A2UI MIME type constant."""
|
||||
from praisonaiagents.ui.a2ui.extension import A2UI_MIME_TYPE
|
||||
|
||||
assert A2UI_MIME_TYPE == "application/json+a2ui"
|
||||
|
||||
def test_a2ui_extension_uri(self):
|
||||
"""Test A2UI extension URI constant."""
|
||||
from praisonaiagents.ui.a2ui.extension import A2UI_EXTENSION_URI
|
||||
|
||||
assert "a2ui" in A2UI_EXTENSION_URI.lower()
|
||||
|
||||
def test_create_a2ui_part(self):
|
||||
"""Test create_a2ui_part wraps A2UI data in A2A DataPart."""
|
||||
from praisonaiagents.ui.a2ui.extension import create_a2ui_part, A2UI_MIME_TYPE
|
||||
|
||||
a2ui_data = {
|
||||
"createSurface": {
|
||||
"surfaceId": "main",
|
||||
"catalogId": "https://a2ui.dev/standard"
|
||||
}
|
||||
}
|
||||
|
||||
part = create_a2ui_part(a2ui_data)
|
||||
|
||||
# Should return a Part with DataPart root
|
||||
assert part is not None
|
||||
assert hasattr(part, 'data')
|
||||
assert part.data == a2ui_data
|
||||
assert part.metadata is not None
|
||||
assert part.metadata.get("mimeType") == A2UI_MIME_TYPE
|
||||
|
||||
def test_is_a2ui_part_true(self):
|
||||
"""Test is_a2ui_part returns True for A2UI parts."""
|
||||
from praisonaiagents.ui.a2ui.extension import create_a2ui_part, is_a2ui_part
|
||||
|
||||
a2ui_data = {"createSurface": {"surfaceId": "main", "catalogId": "test"}}
|
||||
part = create_a2ui_part(a2ui_data)
|
||||
|
||||
assert is_a2ui_part(part) is True
|
||||
|
||||
def test_is_a2ui_part_false(self):
|
||||
"""Test is_a2ui_part returns False for non-A2UI parts."""
|
||||
from praisonaiagents.ui.a2ui.extension import is_a2ui_part
|
||||
from praisonaiagents.ui.a2ui.types import A2UIDataPart
|
||||
|
||||
# Create a regular data part without A2UI mime type
|
||||
part = A2UIDataPart(data={"foo": "bar"}, metadata={"mimeType": "application/json"})
|
||||
|
||||
assert is_a2ui_part(part) is False
|
||||
|
||||
def test_get_a2ui_agent_extension(self):
|
||||
"""Test get_a2ui_agent_extension returns AgentExtension config."""
|
||||
from praisonaiagents.ui.a2ui.extension import get_a2ui_agent_extension
|
||||
|
||||
ext = get_a2ui_agent_extension()
|
||||
|
||||
assert ext is not None
|
||||
assert hasattr(ext, 'uri')
|
||||
assert "a2ui" in ext.uri.lower()
|
||||
assert hasattr(ext, 'description')
|
||||
|
||||
def test_get_a2ui_agent_extension_with_inline_catalog(self):
|
||||
"""Test get_a2ui_agent_extension with inline catalog support."""
|
||||
from praisonaiagents.ui.a2ui.extension import get_a2ui_agent_extension
|
||||
|
||||
ext = get_a2ui_agent_extension(accepts_inline_custom_catalog=True)
|
||||
|
||||
assert ext.params is not None
|
||||
assert ext.params.get("acceptsInlineCustomCatalog") is True
|
||||
|
||||
def test_standard_catalog_id(self):
|
||||
"""Test standard catalog ID constant."""
|
||||
from praisonaiagents.ui.a2ui.extension import STANDARD_CATALOG_ID
|
||||
|
||||
assert "standard_catalog" in STANDARD_CATALOG_ID.lower() or "a2ui" in STANDARD_CATALOG_ID.lower()
|
||||
|
||||
|
||||
class TestA2UIDataPart:
|
||||
"""Test A2UI DataPart type."""
|
||||
|
||||
def test_a2ui_data_part_creation(self):
|
||||
"""Test A2UIDataPart creation."""
|
||||
from praisonaiagents.ui.a2ui.types import A2UIDataPart
|
||||
|
||||
part = A2UIDataPart(
|
||||
data={"test": "value"},
|
||||
metadata={"mimeType": "application/json+a2ui"}
|
||||
)
|
||||
|
||||
assert part.data == {"test": "value"}
|
||||
assert part.metadata["mimeType"] == "application/json+a2ui"
|
||||
|
||||
def test_a2ui_data_part_to_dict(self):
|
||||
"""Test A2UIDataPart serialization."""
|
||||
from praisonaiagents.ui.a2ui.types import A2UIDataPart
|
||||
|
||||
part = A2UIDataPart(
|
||||
data={"createSurface": {"surfaceId": "main"}},
|
||||
metadata={"mimeType": "application/json+a2ui"}
|
||||
)
|
||||
|
||||
data = part.to_dict()
|
||||
assert "data" in data
|
||||
assert "metadata" in data
|
||||
@@ -1,225 +0,0 @@
|
||||
"""
|
||||
End-to-End Integration Tests for A2UI
|
||||
|
||||
Tests the complete A2UI workflow from agent to JSON output.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import pytest
|
||||
|
||||
|
||||
class TestA2UIEndToEnd:
|
||||
"""End-to-end integration tests."""
|
||||
|
||||
def test_complete_surface_workflow(self):
|
||||
"""Test complete workflow: create surface, add components, generate JSON."""
|
||||
from praisonaiagents.ui.a2ui import (
|
||||
Surface,
|
||||
PathBinding,
|
||||
)
|
||||
|
||||
# Create surface
|
||||
surface = Surface(surface_id="test-surface")
|
||||
|
||||
# Add components using fluent API
|
||||
surface.text(id="title", text="Welcome", usage_hint="h1")
|
||||
surface.text(id="subtitle", text=PathBinding(path="/subtitle"), usage_hint="h2")
|
||||
surface.button(
|
||||
id="action-btn",
|
||||
child="btn-text",
|
||||
action_name="do_action",
|
||||
action_context=[{"key": "id", "value": "123"}],
|
||||
primary=True
|
||||
)
|
||||
surface.text(id="btn-text", text="Click Me")
|
||||
surface.column(id="root", children=["title", "subtitle", "action-btn"])
|
||||
|
||||
# Set data
|
||||
surface.set_data("subtitle", "Hello World")
|
||||
|
||||
# Generate messages
|
||||
messages = surface.to_messages()
|
||||
|
||||
# Verify structure
|
||||
assert len(messages) == 3 # createSurface, updateComponents, updateDataModel
|
||||
|
||||
# Verify createSurface
|
||||
assert "createSurface" in messages[0]
|
||||
assert messages[0]["createSurface"]["surfaceId"] == "test-surface"
|
||||
|
||||
# Verify updateComponents
|
||||
assert "updateComponents" in messages[1]
|
||||
components = messages[1]["updateComponents"]["components"]
|
||||
assert len(components) == 5 # title, subtitle, btn, btn-text, root
|
||||
|
||||
# Verify updateDataModel
|
||||
assert "updateDataModel" in messages[2]
|
||||
assert messages[2]["updateDataModel"]["value"]["subtitle"] == "Hello World"
|
||||
|
||||
# Verify JSON output
|
||||
json_str = surface.to_json()
|
||||
parsed = json.loads(json_str)
|
||||
assert len(parsed) == 3
|
||||
|
||||
def test_a2ui_agent_render_workflow(self):
|
||||
"""Test A2UIAgent rendering workflow."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui import A2UIAgent
|
||||
|
||||
agent = Agent(
|
||||
name="TestAgent",
|
||||
role="Test assistant",
|
||||
goal="Help with testing"
|
||||
)
|
||||
|
||||
a2ui_agent = A2UIAgent(agent=agent, surface_id="agent-surface")
|
||||
|
||||
# Render text
|
||||
messages = a2ui_agent.render_text(
|
||||
text="This is a test response",
|
||||
title="Test Title"
|
||||
)
|
||||
|
||||
# Verify structure
|
||||
assert len(messages) >= 2
|
||||
assert "createSurface" in messages[0]
|
||||
|
||||
# Verify JSON export
|
||||
json_str = a2ui_agent.to_json()
|
||||
parsed = json.loads(json_str)
|
||||
assert isinstance(parsed, list)
|
||||
|
||||
def test_template_workflow(self):
|
||||
"""Test template-based workflow."""
|
||||
from praisonaiagents.ui.a2ui import (
|
||||
ChatTemplate,
|
||||
ListTemplate,
|
||||
FormTemplate,
|
||||
DashboardTemplate,
|
||||
)
|
||||
|
||||
# Chat template
|
||||
chat = ChatTemplate(surface_id="chat")
|
||||
chat.add_user_message("Hello!")
|
||||
chat.add_agent_message("Hi there!")
|
||||
chat_messages = chat.to_messages()
|
||||
assert len(chat_messages) >= 2
|
||||
|
||||
# List template
|
||||
list_template = ListTemplate(surface_id="list", title="Items")
|
||||
list_template.add_item(title="Item 1", description="First")
|
||||
list_template.add_item(title="Item 2", description="Second")
|
||||
list_messages = list_template.to_messages()
|
||||
assert len(list_messages) >= 2
|
||||
|
||||
# Form template
|
||||
form = FormTemplate(surface_id="form", title="Contact")
|
||||
form.add_text_field(id="name", label="Name")
|
||||
form.add_email_field(id="email", label="Email")
|
||||
form.set_submit_action("submit", "Send")
|
||||
form_messages = form.to_messages()
|
||||
assert len(form_messages) >= 2
|
||||
|
||||
# Dashboard template
|
||||
dashboard = DashboardTemplate(surface_id="dashboard", title="Dashboard")
|
||||
dashboard.add_panel(id="stats", title="Stats", content="100 items")
|
||||
dashboard.add_panel(id="chart", title="Chart", content="[Chart]")
|
||||
dashboard_messages = dashboard.to_messages()
|
||||
assert len(dashboard_messages) >= 2
|
||||
|
||||
def test_a2a_integration(self):
|
||||
"""Test A2A protocol integration."""
|
||||
from praisonaiagents.ui.a2ui import (
|
||||
Surface,
|
||||
create_a2ui_part,
|
||||
is_a2ui_part,
|
||||
A2UI_MIME_TYPE,
|
||||
)
|
||||
|
||||
# Create surface
|
||||
surface = Surface(surface_id="a2a-test")
|
||||
surface.text(id="msg", text="Hello from A2A")
|
||||
surface.column(id="root", children=["msg"])
|
||||
|
||||
# Get messages
|
||||
messages = surface.to_messages()
|
||||
|
||||
# Wrap in A2A DataPart
|
||||
part = create_a2ui_part({"messages": messages})
|
||||
|
||||
# Verify A2UI part
|
||||
assert is_a2ui_part(part)
|
||||
assert part.metadata["mimeType"] == A2UI_MIME_TYPE
|
||||
assert "messages" in part.data
|
||||
assert len(part.data["messages"]) >= 2
|
||||
|
||||
def test_data_binding_workflow(self):
|
||||
"""Test data binding with path references."""
|
||||
from praisonaiagents.ui.a2ui import (
|
||||
Surface,
|
||||
PathBinding,
|
||||
ChildrenTemplate,
|
||||
)
|
||||
|
||||
surface = Surface(surface_id="data-binding-test")
|
||||
|
||||
# Add components with data binding
|
||||
surface.text(id="title", text=PathBinding(path="/title"), usage_hint="h1")
|
||||
surface.text(id="item-template", text=PathBinding(path="/name"))
|
||||
|
||||
# Add list with template
|
||||
surface.list(
|
||||
id="items",
|
||||
children=ChildrenTemplate(component_id="item-template", path="/items"),
|
||||
direction="vertical"
|
||||
)
|
||||
|
||||
surface.column(id="root", children=["title", "items"])
|
||||
|
||||
# Set data
|
||||
surface.set_data("title", "My List")
|
||||
surface.set_data("items", [
|
||||
{"name": "Item 1"},
|
||||
{"name": "Item 2"},
|
||||
{"name": "Item 3"},
|
||||
])
|
||||
|
||||
messages = surface.to_messages()
|
||||
|
||||
# Verify data model
|
||||
data_msg = next((m for m in messages if "updateDataModel" in m), None)
|
||||
assert data_msg is not None
|
||||
assert data_msg["updateDataModel"]["value"]["title"] == "My List"
|
||||
assert len(data_msg["updateDataModel"]["value"]["items"]) == 3
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get("OPENAI_API_KEY"),
|
||||
reason="OPENAI_API_KEY not set"
|
||||
)
|
||||
class TestA2UIWithLLM:
|
||||
"""Tests that require LLM API calls."""
|
||||
|
||||
def test_agent_chat_and_render(self):
|
||||
"""Test agent chat followed by UI rendering."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2ui import A2UIAgent
|
||||
|
||||
agent = Agent(
|
||||
name="UIAgent",
|
||||
role="UI Generator",
|
||||
goal="Generate helpful responses",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
a2ui_agent = A2UIAgent(agent=agent)
|
||||
|
||||
# Chat with agent
|
||||
response = a2ui_agent.chat("What is 2+2?")
|
||||
|
||||
# Render response as UI
|
||||
messages = a2ui_agent.render_text(response, title="Answer")
|
||||
|
||||
assert len(messages) >= 2
|
||||
assert "createSurface" in messages[0]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Tests for UI protocol types (zero-dep contracts)."""
|
||||
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE, A2UIToolResultProtocol
|
||||
|
||||
|
||||
class TestA2UIProtocols:
|
||||
def test_mime_type_constant(self):
|
||||
assert A2UI_MIME_TYPE == "application/json+a2ui"
|
||||
|
||||
def test_tool_result_protocol_shape(self):
|
||||
result: A2UIToolResultProtocol = {
|
||||
"mime_type": A2UI_MIME_TYPE,
|
||||
"messages": [{"createSurface": {"surfaceId": "main"}}],
|
||||
"a2ui_part": {"messages": []},
|
||||
}
|
||||
assert result["mime_type"] == A2UI_MIME_TYPE
|
||||
|
||||
def test_exported_from_a2ui_package(self):
|
||||
from praisonaiagents.ui.a2ui import A2UI_MIME_TYPE, A2UIToolResultProtocol
|
||||
|
||||
assert A2UI_MIME_TYPE == "application/json+a2ui"
|
||||
assert A2UIToolResultProtocol is not None
|
||||
|
||||
def test_top_level_lazy_a2ui_export(self):
|
||||
from praisonaiagents import A2UI
|
||||
|
||||
assert A2UI is not None
|
||||
@@ -1,228 +0,0 @@
|
||||
"""
|
||||
Tests for A2UI Surface Builder
|
||||
|
||||
TDD: Write tests first, then implement the Surface class.
|
||||
"""
|
||||
|
||||
|
||||
class TestSurface:
|
||||
"""Test Surface builder class."""
|
||||
|
||||
def test_surface_creation(self):
|
||||
"""Test Surface creation with ID and catalog."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
|
||||
assert surface.surface_id == "main"
|
||||
assert surface.catalog_id is not None # Should have default
|
||||
|
||||
def test_surface_with_custom_catalog(self):
|
||||
"""Test Surface with custom catalog ID."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(
|
||||
surface_id="main",
|
||||
catalog_id="https://custom.catalog/v1"
|
||||
)
|
||||
|
||||
assert surface.catalog_id == "https://custom.catalog/v1"
|
||||
|
||||
def test_add_component(self):
|
||||
"""Test adding a component to the surface."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
|
||||
|
||||
surface.add(text)
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "title"
|
||||
|
||||
def test_add_multiple_components(self):
|
||||
"""Test adding multiple components."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent, ButtonComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.add(TextComponent(id="title", text="Hello"))
|
||||
surface.add(ButtonComponent(id="btn", child="btn-text"))
|
||||
|
||||
assert len(surface.components) == 2
|
||||
|
||||
def test_set_data(self):
|
||||
"""Test setting data model values."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.set_data("title", "Hello World")
|
||||
surface.set_data("count", 42)
|
||||
|
||||
assert surface.data["title"] == "Hello World"
|
||||
assert surface.data["count"] == 42
|
||||
|
||||
def test_set_nested_data(self):
|
||||
"""Test setting nested data model values."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.set_data("user", {"name": "John", "age": 30})
|
||||
|
||||
assert surface.data["user"]["name"] == "John"
|
||||
|
||||
def test_to_messages_creates_surface(self):
|
||||
"""Test to_messages generates CreateSurface message."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.add(TextComponent(id="title", text="Hello"))
|
||||
|
||||
messages = surface.to_messages()
|
||||
|
||||
assert len(messages) >= 1
|
||||
assert "createSurface" in messages[0]
|
||||
assert messages[0]["createSurface"]["surfaceId"] == "main"
|
||||
|
||||
def test_to_messages_includes_components(self):
|
||||
"""Test to_messages includes UpdateComponents message."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.add(TextComponent(id="title", text="Hello"))
|
||||
|
||||
messages = surface.to_messages()
|
||||
|
||||
# Should have createSurface and updateComponents
|
||||
assert len(messages) >= 2
|
||||
update_msg = next((m for m in messages if "updateComponents" in m), None)
|
||||
assert update_msg is not None
|
||||
assert len(update_msg["updateComponents"]["components"]) == 1
|
||||
|
||||
def test_to_messages_includes_data_model(self):
|
||||
"""Test to_messages includes UpdateDataModel message when data is set."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.add(TextComponent(id="title", text="Hello"))
|
||||
surface.set_data("title", "Hello World")
|
||||
|
||||
messages = surface.to_messages()
|
||||
|
||||
# Should have createSurface, updateComponents, and updateDataModel
|
||||
assert len(messages) == 3
|
||||
data_msg = next((m for m in messages if "updateDataModel" in m), None)
|
||||
assert data_msg is not None
|
||||
assert data_msg["updateDataModel"]["value"]["title"] == "Hello World"
|
||||
|
||||
def test_to_json(self):
|
||||
"""Test to_json returns JSON string."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
import json
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.add(TextComponent(id="title", text="Hello"))
|
||||
|
||||
json_str = surface.to_json()
|
||||
|
||||
# Should be valid JSON
|
||||
parsed = json.loads(json_str)
|
||||
assert isinstance(parsed, list)
|
||||
assert len(parsed) >= 2
|
||||
|
||||
|
||||
class TestSurfaceChaining:
|
||||
"""Test Surface method chaining."""
|
||||
|
||||
def test_add_returns_self(self):
|
||||
"""Test add() returns self for chaining."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
result = surface.add(TextComponent(id="title", text="Hello"))
|
||||
|
||||
assert result is surface
|
||||
|
||||
def test_set_data_returns_self(self):
|
||||
"""Test set_data() returns self for chaining."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
result = surface.set_data("key", "value")
|
||||
|
||||
assert result is surface
|
||||
|
||||
def test_fluent_api(self):
|
||||
"""Test fluent API for building surfaces."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent, ButtonComponent
|
||||
|
||||
surface = (
|
||||
Surface(surface_id="main")
|
||||
.add(TextComponent(id="title", text="Hello"))
|
||||
.add(ButtonComponent(id="btn", child="btn-text"))
|
||||
.set_data("title", "Hello World")
|
||||
)
|
||||
|
||||
assert len(surface.components) == 2
|
||||
assert surface.data["title"] == "Hello World"
|
||||
|
||||
|
||||
class TestSurfaceHelpers:
|
||||
"""Test Surface helper methods."""
|
||||
|
||||
def test_text_helper(self):
|
||||
"""Test text() helper method."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.text(id="title", text="Hello", usage_hint="h1")
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "title"
|
||||
|
||||
def test_button_helper(self):
|
||||
"""Test button() helper method."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.button(id="submit", child="submit-text", action_name="submit")
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "submit"
|
||||
|
||||
def test_card_helper(self):
|
||||
"""Test card() helper method."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.card(id="card-1", child="card-content")
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "card-1"
|
||||
|
||||
def test_row_helper(self):
|
||||
"""Test row() helper method."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.row(id="row-1", children=["col-1", "col-2"])
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "row-1"
|
||||
|
||||
def test_column_helper(self):
|
||||
"""Test column() helper method."""
|
||||
from praisonaiagents.ui.a2ui.surface import Surface
|
||||
|
||||
surface = Surface(surface_id="main")
|
||||
surface.column(id="col-1", children=["item-1", "item-2"])
|
||||
|
||||
assert len(surface.components) == 1
|
||||
assert surface.components[0].id == "col-1"
|
||||
@@ -1,210 +0,0 @@
|
||||
"""
|
||||
Tests for A2UI Templates
|
||||
|
||||
TDD: Write tests first, then implement the templates.
|
||||
"""
|
||||
|
||||
|
||||
class TestChatTemplate:
|
||||
"""Test ChatTemplate for conversational UIs."""
|
||||
|
||||
def test_chat_template_creation(self):
|
||||
"""Test ChatTemplate creation."""
|
||||
from praisonaiagents.ui.a2ui.templates import ChatTemplate
|
||||
|
||||
template = ChatTemplate(surface_id="chat")
|
||||
|
||||
assert template.surface_id == "chat"
|
||||
|
||||
def test_add_user_message(self):
|
||||
"""Test adding a user message."""
|
||||
from praisonaiagents.ui.a2ui.templates import ChatTemplate
|
||||
|
||||
template = ChatTemplate(surface_id="chat")
|
||||
template.add_user_message("Hello!")
|
||||
|
||||
assert len(template.messages) == 1
|
||||
assert template.messages[0]["role"] == "user"
|
||||
assert template.messages[0]["content"] == "Hello!"
|
||||
|
||||
def test_add_agent_message(self):
|
||||
"""Test adding an agent message."""
|
||||
from praisonaiagents.ui.a2ui.templates import ChatTemplate
|
||||
|
||||
template = ChatTemplate(surface_id="chat")
|
||||
template.add_agent_message("Hi there!")
|
||||
|
||||
assert len(template.messages) == 1
|
||||
assert template.messages[0]["role"] == "agent"
|
||||
|
||||
def test_to_messages(self):
|
||||
"""Test generating A2UI messages."""
|
||||
from praisonaiagents.ui.a2ui.templates import ChatTemplate
|
||||
|
||||
template = ChatTemplate(surface_id="chat")
|
||||
template.add_user_message("Hello!")
|
||||
template.add_agent_message("Hi there!")
|
||||
|
||||
messages = template.to_messages()
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2 # createSurface + updateComponents
|
||||
|
||||
|
||||
class TestListTemplate:
|
||||
"""Test ListTemplate for displaying lists of items."""
|
||||
|
||||
def test_list_template_creation(self):
|
||||
"""Test ListTemplate creation."""
|
||||
from praisonaiagents.ui.a2ui.templates import ListTemplate
|
||||
|
||||
template = ListTemplate(
|
||||
surface_id="list",
|
||||
title="My Items"
|
||||
)
|
||||
|
||||
assert template.surface_id == "list"
|
||||
assert template.title == "My Items"
|
||||
|
||||
def test_add_item(self):
|
||||
"""Test adding an item to the list."""
|
||||
from praisonaiagents.ui.a2ui.templates import ListTemplate
|
||||
|
||||
template = ListTemplate(surface_id="list", title="Items")
|
||||
template.add_item(title="Item 1", description="First item")
|
||||
|
||||
assert len(template.items) == 1
|
||||
assert template.items[0]["title"] == "Item 1"
|
||||
|
||||
def test_add_item_with_image(self):
|
||||
"""Test adding an item with an image."""
|
||||
from praisonaiagents.ui.a2ui.templates import ListTemplate
|
||||
|
||||
template = ListTemplate(surface_id="list", title="Items")
|
||||
template.add_item(
|
||||
title="Item 1",
|
||||
description="First item",
|
||||
image_url="https://example.com/image.jpg"
|
||||
)
|
||||
|
||||
assert template.items[0]["image_url"] == "https://example.com/image.jpg"
|
||||
|
||||
def test_to_messages(self):
|
||||
"""Test generating A2UI messages."""
|
||||
from praisonaiagents.ui.a2ui.templates import ListTemplate
|
||||
|
||||
template = ListTemplate(surface_id="list", title="Items")
|
||||
template.add_item(title="Item 1", description="First")
|
||||
template.add_item(title="Item 2", description="Second")
|
||||
|
||||
messages = template.to_messages()
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2
|
||||
|
||||
|
||||
class TestFormTemplate:
|
||||
"""Test FormTemplate for data collection forms."""
|
||||
|
||||
def test_form_template_creation(self):
|
||||
"""Test FormTemplate creation."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(
|
||||
surface_id="form",
|
||||
title="Contact Form"
|
||||
)
|
||||
|
||||
assert template.surface_id == "form"
|
||||
assert template.title == "Contact Form"
|
||||
|
||||
def test_add_text_field(self):
|
||||
"""Test adding a text field."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(surface_id="form", title="Form")
|
||||
template.add_text_field(id="name", label="Name")
|
||||
|
||||
assert len(template.fields) == 1
|
||||
assert template.fields[0]["id"] == "name"
|
||||
assert template.fields[0]["type"] == "text"
|
||||
|
||||
def test_add_number_field(self):
|
||||
"""Test adding a number field."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(surface_id="form", title="Form")
|
||||
template.add_number_field(id="age", label="Age")
|
||||
|
||||
assert template.fields[0]["type"] == "number"
|
||||
|
||||
def test_add_email_field(self):
|
||||
"""Test adding an email field."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(surface_id="form", title="Form")
|
||||
template.add_email_field(id="email", label="Email")
|
||||
|
||||
assert template.fields[0]["type"] == "email"
|
||||
|
||||
def test_set_submit_action(self):
|
||||
"""Test setting submit action."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(surface_id="form", title="Form")
|
||||
template.set_submit_action("submit_form", "Send")
|
||||
|
||||
assert template.submit_action == "submit_form"
|
||||
assert template.submit_label == "Send"
|
||||
|
||||
def test_to_messages(self):
|
||||
"""Test generating A2UI messages."""
|
||||
from praisonaiagents.ui.a2ui.templates import FormTemplate
|
||||
|
||||
template = FormTemplate(surface_id="form", title="Form")
|
||||
template.add_text_field(id="name", label="Name")
|
||||
template.add_email_field(id="email", label="Email")
|
||||
|
||||
messages = template.to_messages()
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2
|
||||
|
||||
|
||||
class TestDashboardTemplate:
|
||||
"""Test DashboardTemplate for multi-panel dashboards."""
|
||||
|
||||
def test_dashboard_template_creation(self):
|
||||
"""Test DashboardTemplate creation."""
|
||||
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
|
||||
|
||||
template = DashboardTemplate(
|
||||
surface_id="dashboard",
|
||||
title="My Dashboard"
|
||||
)
|
||||
|
||||
assert template.surface_id == "dashboard"
|
||||
assert template.title == "My Dashboard"
|
||||
|
||||
def test_add_panel(self):
|
||||
"""Test adding a panel."""
|
||||
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
|
||||
|
||||
template = DashboardTemplate(surface_id="dashboard", title="Dashboard")
|
||||
template.add_panel(id="panel-1", title="Panel 1", content="Some content")
|
||||
|
||||
assert len(template.panels) == 1
|
||||
assert template.panels[0]["id"] == "panel-1"
|
||||
|
||||
def test_to_messages(self):
|
||||
"""Test generating A2UI messages."""
|
||||
from praisonaiagents.ui.a2ui.templates import DashboardTemplate
|
||||
|
||||
template = DashboardTemplate(surface_id="dashboard", title="Dashboard")
|
||||
template.add_panel(id="panel-1", title="Panel 1", content="Content 1")
|
||||
template.add_panel(id="panel-2", title="Panel 2", content="Content 2")
|
||||
|
||||
messages = template.to_messages()
|
||||
|
||||
assert isinstance(messages, list)
|
||||
assert len(messages) >= 2
|
||||
@@ -1,318 +0,0 @@
|
||||
"""
|
||||
Tests for A2UI Types Module
|
||||
|
||||
TDD: Write tests first, then implement the types.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
class TestA2UIMessageTypes:
|
||||
"""Test A2UI message types: CreateSurface, UpdateComponents, UpdateDataModel, DeleteSurface."""
|
||||
|
||||
def test_create_surface_message(self):
|
||||
"""Test CreateSurface message creation."""
|
||||
from praisonaiagents.ui.a2ui.types import CreateSurfaceMessage
|
||||
|
||||
msg = CreateSurfaceMessage(
|
||||
surface_id="main",
|
||||
catalog_id="https://a2ui.dev/standard"
|
||||
)
|
||||
|
||||
assert msg.surface_id == "main"
|
||||
assert msg.catalog_id == "https://a2ui.dev/standard"
|
||||
|
||||
# Test serialization
|
||||
data = msg.to_dict()
|
||||
assert "createSurface" in data
|
||||
assert data["createSurface"]["surfaceId"] == "main"
|
||||
assert data["createSurface"]["catalogId"] == "https://a2ui.dev/standard"
|
||||
|
||||
def test_update_components_message(self):
|
||||
"""Test UpdateComponents message creation."""
|
||||
from praisonaiagents.ui.a2ui.types import UpdateComponentsMessage, TextComponent
|
||||
|
||||
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
|
||||
msg = UpdateComponentsMessage(
|
||||
surface_id="main",
|
||||
components=[text]
|
||||
)
|
||||
|
||||
assert msg.surface_id == "main"
|
||||
assert len(msg.components) == 1
|
||||
|
||||
# Test serialization
|
||||
data = msg.to_dict()
|
||||
assert "updateComponents" in data
|
||||
assert data["updateComponents"]["surfaceId"] == "main"
|
||||
assert len(data["updateComponents"]["components"]) == 1
|
||||
|
||||
def test_update_data_model_message(self):
|
||||
"""Test UpdateDataModel message creation."""
|
||||
from praisonaiagents.ui.a2ui.types import UpdateDataModelMessage
|
||||
|
||||
msg = UpdateDataModelMessage(
|
||||
surface_id="main",
|
||||
path="/",
|
||||
op="replace",
|
||||
value={"title": "Hello", "items": []}
|
||||
)
|
||||
|
||||
assert msg.surface_id == "main"
|
||||
assert msg.path == "/"
|
||||
assert msg.op == "replace"
|
||||
assert msg.value == {"title": "Hello", "items": []}
|
||||
|
||||
# Test serialization
|
||||
data = msg.to_dict()
|
||||
assert "updateDataModel" in data
|
||||
assert data["updateDataModel"]["surfaceId"] == "main"
|
||||
|
||||
def test_delete_surface_message(self):
|
||||
"""Test DeleteSurface message creation."""
|
||||
from praisonaiagents.ui.a2ui.types import DeleteSurfaceMessage
|
||||
|
||||
msg = DeleteSurfaceMessage(surface_id="main")
|
||||
|
||||
assert msg.surface_id == "main"
|
||||
|
||||
# Test serialization
|
||||
data = msg.to_dict()
|
||||
assert "deleteSurface" in data
|
||||
assert data["deleteSurface"]["surfaceId"] == "main"
|
||||
|
||||
|
||||
class TestA2UIComponentTypes:
|
||||
"""Test A2UI component types."""
|
||||
|
||||
def test_text_component(self):
|
||||
"""Test Text component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent
|
||||
|
||||
text = TextComponent(id="title", text="Hello World", usage_hint="h1")
|
||||
|
||||
assert text.id == "title"
|
||||
assert text.text == "Hello World"
|
||||
assert text.usage_hint == "h1"
|
||||
|
||||
# Test serialization
|
||||
data = text.to_dict()
|
||||
assert data["id"] == "title"
|
||||
assert data["component"] == "Text"
|
||||
assert data["text"] == "Hello World"
|
||||
assert data["usageHint"] == "h1"
|
||||
|
||||
def test_text_component_with_path(self):
|
||||
"""Test Text component with data binding path."""
|
||||
from praisonaiagents.ui.a2ui.types import TextComponent, PathBinding
|
||||
|
||||
text = TextComponent(id="title", text=PathBinding(path="/title"))
|
||||
|
||||
data = text.to_dict()
|
||||
assert data["text"] == {"path": "/title"}
|
||||
|
||||
def test_image_component(self):
|
||||
"""Test Image component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import ImageComponent
|
||||
|
||||
img = ImageComponent(
|
||||
id="hero",
|
||||
url="https://example.com/image.jpg",
|
||||
fit="cover",
|
||||
usage_hint="header"
|
||||
)
|
||||
|
||||
assert img.id == "hero"
|
||||
assert img.url == "https://example.com/image.jpg"
|
||||
assert img.fit == "cover"
|
||||
|
||||
data = img.to_dict()
|
||||
assert data["component"] == "Image"
|
||||
assert data["url"] == "https://example.com/image.jpg"
|
||||
|
||||
def test_button_component(self):
|
||||
"""Test Button component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import ButtonComponent, Action, ActionContext
|
||||
|
||||
action = Action(
|
||||
name="submit",
|
||||
context=[
|
||||
ActionContext(key="value", value="test")
|
||||
]
|
||||
)
|
||||
btn = ButtonComponent(
|
||||
id="submit-btn",
|
||||
child="submit-text",
|
||||
primary=True,
|
||||
action=action
|
||||
)
|
||||
|
||||
assert btn.id == "submit-btn"
|
||||
assert btn.child == "submit-text"
|
||||
assert btn.primary is True
|
||||
|
||||
data = btn.to_dict()
|
||||
assert data["component"] == "Button"
|
||||
assert data["action"]["name"] == "submit"
|
||||
|
||||
def test_card_component(self):
|
||||
"""Test Card component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import CardComponent
|
||||
|
||||
card = CardComponent(id="card-1", child="card-content")
|
||||
|
||||
assert card.id == "card-1"
|
||||
assert card.child == "card-content"
|
||||
|
||||
data = card.to_dict()
|
||||
assert data["component"] == "Card"
|
||||
assert data["child"] == "card-content"
|
||||
|
||||
def test_row_component(self):
|
||||
"""Test Row component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import RowComponent
|
||||
|
||||
row = RowComponent(id="row-1", children=["col-1", "col-2"])
|
||||
|
||||
assert row.id == "row-1"
|
||||
assert row.children == ["col-1", "col-2"]
|
||||
|
||||
data = row.to_dict()
|
||||
assert data["component"] == "Row"
|
||||
assert data["children"] == ["col-1", "col-2"]
|
||||
|
||||
def test_column_component(self):
|
||||
"""Test Column component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import ColumnComponent
|
||||
|
||||
col = ColumnComponent(id="col-1", children=["item-1", "item-2"])
|
||||
|
||||
assert col.id == "col-1"
|
||||
assert col.children == ["item-1", "item-2"]
|
||||
|
||||
data = col.to_dict()
|
||||
assert data["component"] == "Column"
|
||||
|
||||
def test_list_component(self):
|
||||
"""Test List component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import ListComponent, ChildrenTemplate
|
||||
|
||||
template = ChildrenTemplate(component_id="item-template", path="/items")
|
||||
lst = ListComponent(id="item-list", direction="vertical", children=template)
|
||||
|
||||
assert lst.id == "item-list"
|
||||
assert lst.direction == "vertical"
|
||||
|
||||
data = lst.to_dict()
|
||||
assert data["component"] == "List"
|
||||
assert data["children"]["componentId"] == "item-template"
|
||||
assert data["children"]["path"] == "/items"
|
||||
|
||||
def test_text_field_component(self):
|
||||
"""Test TextField component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import TextFieldComponent, PathBinding
|
||||
|
||||
field = TextFieldComponent(
|
||||
id="name-field",
|
||||
label="Name",
|
||||
text=PathBinding(path="/name"),
|
||||
field_type="text"
|
||||
)
|
||||
|
||||
assert field.id == "name-field"
|
||||
assert field.label == "Name"
|
||||
|
||||
data = field.to_dict()
|
||||
assert data["component"] == "TextField"
|
||||
assert data["label"] == "Name"
|
||||
assert data["text"] == {"path": "/name"}
|
||||
|
||||
def test_checkbox_component(self):
|
||||
"""Test CheckBox component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import CheckBoxComponent, PathBinding
|
||||
|
||||
cb = CheckBoxComponent(
|
||||
id="agree-cb",
|
||||
label="I agree",
|
||||
checked=PathBinding(path="/agreed")
|
||||
)
|
||||
|
||||
data = cb.to_dict()
|
||||
assert data["component"] == "CheckBox"
|
||||
assert data["label"] == "I agree"
|
||||
|
||||
def test_slider_component(self):
|
||||
"""Test Slider component creation."""
|
||||
from praisonaiagents.ui.a2ui.types import SliderComponent
|
||||
|
||||
slider = SliderComponent(
|
||||
id="volume",
|
||||
min_value=0,
|
||||
max_value=100,
|
||||
value=50
|
||||
)
|
||||
|
||||
data = slider.to_dict()
|
||||
assert data["component"] == "Slider"
|
||||
assert data["min"] == 0
|
||||
assert data["max"] == 100
|
||||
|
||||
|
||||
class TestDataBindingTypes:
|
||||
"""Test data binding types."""
|
||||
|
||||
def test_path_binding(self):
|
||||
"""Test PathBinding for data model references."""
|
||||
from praisonaiagents.ui.a2ui.types import PathBinding
|
||||
|
||||
path = PathBinding(path="/user/name")
|
||||
|
||||
assert path.path == "/user/name"
|
||||
assert path.to_dict() == {"path": "/user/name"}
|
||||
|
||||
def test_string_or_path_literal(self):
|
||||
"""Test StringOrPath with literal string."""
|
||||
from praisonaiagents.ui.a2ui.types import resolve_string_or_path
|
||||
|
||||
result = resolve_string_or_path("Hello")
|
||||
assert result == "Hello"
|
||||
|
||||
def test_string_or_path_binding(self):
|
||||
"""Test StringOrPath with path binding."""
|
||||
from praisonaiagents.ui.a2ui.types import resolve_string_or_path, PathBinding
|
||||
|
||||
result = resolve_string_or_path(PathBinding(path="/title"))
|
||||
assert result == {"path": "/title"}
|
||||
|
||||
|
||||
class TestActionTypes:
|
||||
"""Test action types for button events."""
|
||||
|
||||
def test_action_creation(self):
|
||||
"""Test Action creation."""
|
||||
from praisonaiagents.ui.a2ui.types import Action, ActionContext
|
||||
|
||||
action = Action(
|
||||
name="book_restaurant",
|
||||
context=[
|
||||
ActionContext(key="restaurantName", value="The Fancy Place"),
|
||||
ActionContext(key="address", value="123 Main St")
|
||||
]
|
||||
)
|
||||
|
||||
assert action.name == "book_restaurant"
|
||||
assert len(action.context) == 2
|
||||
|
||||
data = action.to_dict()
|
||||
assert data["name"] == "book_restaurant"
|
||||
assert len(data["context"]) == 2
|
||||
assert data["context"][0]["key"] == "restaurantName"
|
||||
|
||||
def test_action_context_with_path(self):
|
||||
"""Test ActionContext with path binding."""
|
||||
from praisonaiagents.ui.a2ui.types import ActionContext, PathBinding
|
||||
|
||||
ctx = ActionContext(key="name", value=PathBinding(path="/restaurant/name"))
|
||||
|
||||
data = ctx.to_dict()
|
||||
assert data["key"] == "name"
|
||||
assert data["value"] == {"path": "/restaurant/name"}
|
||||
@@ -98,6 +98,61 @@ class TestToolCallStreaming:
|
||||
assert event.content == "Search results..."
|
||||
|
||||
|
||||
class TestStreamEventBridge:
|
||||
"""Test PraisonAI stream_emitter → AG-UI event conversion."""
|
||||
|
||||
def test_delta_text_to_agui(self):
|
||||
from praisonaiagents.streaming.events import StreamEvent, StreamEventType
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer, stream_event_to_agui_events
|
||||
|
||||
buffer = EventBuffer()
|
||||
event = StreamEvent(type=StreamEventType.DELTA_TEXT, content="Hello")
|
||||
agui_events = stream_event_to_agui_events(event, "msg-1", buffer)
|
||||
|
||||
assert len(agui_events) == 2
|
||||
assert agui_events[0].type == "TEXT_MESSAGE_START"
|
||||
assert agui_events[1].type == "TEXT_MESSAGE_CONTENT"
|
||||
assert agui_events[1].delta == "Hello"
|
||||
|
||||
def test_tool_call_start_to_agui(self):
|
||||
from praisonaiagents.streaming.events import StreamEvent, StreamEventType
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer, stream_event_to_agui_events
|
||||
|
||||
buffer = EventBuffer()
|
||||
event = StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_START,
|
||||
tool_call={"id": "tc-1", "name": "search", "arguments": {"q": "test"}},
|
||||
)
|
||||
agui_events = stream_event_to_agui_events(event, "msg-1", buffer)
|
||||
|
||||
assert len(agui_events) == 3
|
||||
assert agui_events[0].type == "TOOL_CALL_START"
|
||||
assert agui_events[0].tool_call_name == "search"
|
||||
|
||||
def test_tool_call_result_a2ui_emits_custom_event(self):
|
||||
from praisonaiagents.streaming.events import StreamEvent, StreamEventType
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer, stream_event_to_agui_events
|
||||
from praisonaiagents.ui.protocols import A2UI_MIME_TYPE
|
||||
|
||||
buffer = EventBuffer()
|
||||
a2ui_result = {
|
||||
"mime_type": A2UI_MIME_TYPE,
|
||||
"messages": [{"createSurface": {"surfaceId": "panel-1"}}],
|
||||
}
|
||||
event = StreamEvent(
|
||||
type=StreamEventType.TOOL_CALL_RESULT,
|
||||
tool_call={"id": "tc-a2ui", "name": "send_a2ui_messages", "result": a2ui_result},
|
||||
)
|
||||
agui_events = stream_event_to_agui_events(event, "msg-1", buffer)
|
||||
|
||||
assert len(agui_events) == 2
|
||||
assert agui_events[0].type == "TOOL_CALL_RESULT"
|
||||
assert agui_events[1].type == "CUSTOM"
|
||||
assert agui_events[1].name == "a2ui"
|
||||
assert agui_events[1].value["surface_id"] == "panel-1"
|
||||
assert agui_events[1].value["messages"] == a2ui_result["messages"]
|
||||
|
||||
|
||||
class TestRunLifecycleEvents:
|
||||
"""Test run lifecycle events."""
|
||||
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Integration tests for the ALLOWED_TOOLS filter module."""
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
from praisonai.cli.features.agents import MultiAgentHandler
|
||||
|
||||
|
||||
class TestAllowedToolsFilterRegistryIntegration:
|
||||
"""Integration tests for ALLOWED_TOOLS filter with the tool registry."""
|
||||
|
||||
def test_registry_allowed_filter_unset(self):
|
||||
"""Test registry filter behavior when ALLOWED_TOOLS is unset."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
registry = get_registry()
|
||||
|
||||
# Mock some tools in the registry
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email']):
|
||||
filtered_tools = registry.list_tools_with_allowed_filter()
|
||||
|
||||
# Should return all tools when unset
|
||||
assert set(filtered_tools) == {'search_web', 'send_email'}
|
||||
|
||||
def test_registry_allowed_filter_with_whitelist(self):
|
||||
"""Test registry filter with ALLOWED_TOOLS whitelist."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
with patch.dict(os.environ, {
|
||||
"ALLOWED_TOOLS": "search_web,send_email,unknown_tool",
|
||||
"CI": "false",
|
||||
}):
|
||||
registry = get_registry()
|
||||
|
||||
# Mock some tools in the registry
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email', 'extract_pdf']):
|
||||
filtered_tools = registry.list_tools_with_allowed_filter()
|
||||
|
||||
# Should only include whitelisted tools that are available
|
||||
assert set(filtered_tools) == {'search_web', 'send_email'}
|
||||
|
||||
def test_registry_allowed_filter_empty_result(self):
|
||||
"""Test registry filter when no tools match whitelist."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
with patch.dict(os.environ, {
|
||||
"ALLOWED_TOOLS": "nonexistent_tool1,nonexistent_tool2",
|
||||
"CI": "false",
|
||||
}):
|
||||
registry = get_registry()
|
||||
|
||||
# Mock some tools in the registry
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email']):
|
||||
filtered_tools = registry.list_tools_with_allowed_filter()
|
||||
|
||||
# Should return empty list when no matches
|
||||
assert filtered_tools == []
|
||||
|
||||
def test_registry_allowed_filter_import_error_fallback(self):
|
||||
"""Test registry filter fallback when import fails."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
registry = get_registry()
|
||||
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email']):
|
||||
# Mock import error for the filter
|
||||
with patch('praisonaiagents.tools.registry.filter_tools_with_allowed_tools', side_effect=ImportError("Filter not available")):
|
||||
filtered_tools = registry.list_tools_with_allowed_filter()
|
||||
|
||||
# Should fallback to all tools
|
||||
assert set(filtered_tools) == {'search_web', 'send_email'}
|
||||
|
||||
def test_registry_allowed_filter_value_error_propagation(self):
|
||||
"""Test that ValueError from filter is properly propagated."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ""}): # Empty string should cause ValueError
|
||||
registry = get_registry()
|
||||
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email']):
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
registry.list_tools_with_allowed_filter()
|
||||
|
||||
def test_registry_allowed_filter_general_error_fallback(self):
|
||||
"""Test registry filter fallback on general errors."""
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
registry = get_registry()
|
||||
|
||||
with patch.object(registry, 'list_tools', return_value=['search_web', 'send_email']):
|
||||
# Mock a generic error during filtering
|
||||
with patch('praisonaiagents.tools.registry.filter_tools_with_allowed_tools', side_effect=RuntimeError("Some error")):
|
||||
filtered_tools = registry.list_tools_with_allowed_filter()
|
||||
|
||||
# Should fallback to all tools on general errors
|
||||
assert set(filtered_tools) == {'search_web', 'send_email'}
|
||||
|
||||
|
||||
class TestAllowedToolsFilterGlobalRegistryPath:
|
||||
"""Test ALLOWED_TOOLS filter through the global registry function."""
|
||||
|
||||
def test_global_list_tools_with_allowed_filter(self):
|
||||
"""Test the global list_tools_with_allowed_filter function."""
|
||||
from praisonaiagents.tools import list_tools_with_allowed_filter
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search_web,send_email"}):
|
||||
registry = get_registry()
|
||||
|
||||
with patch.object(registry, 'list_tools_with_allowed_filter', return_value=['search_web', 'send_email']) as mock_filter:
|
||||
result = list_tools_with_allowed_filter()
|
||||
|
||||
# Should call the registry method
|
||||
mock_filter.assert_called_once_with(None)
|
||||
assert result == ['search_web', 'send_email']
|
||||
|
||||
def test_global_backward_compatibility_alias(self):
|
||||
"""Test the backward compatibility alias for hermes filter."""
|
||||
from praisonaiagents.tools import list_tools_with_hermes_filter, list_tools_with_allowed_filter
|
||||
from praisonaiagents.tools.registry import get_registry
|
||||
|
||||
registry = get_registry()
|
||||
|
||||
with patch.object(registry, 'list_tools_with_allowed_filter', return_value=['search_web']) as mock_filter:
|
||||
# Both should call the same underlying function
|
||||
result1 = list_tools_with_allowed_filter()
|
||||
result2 = list_tools_with_hermes_filter()
|
||||
|
||||
assert result1 == result2 == ['search_web']
|
||||
assert mock_filter.call_count == 2
|
||||
|
||||
|
||||
class TestAllowedToolsFilterCLIIntegration:
|
||||
"""Integration tests for ALLOWED_TOOLS filter with CLI agent handler."""
|
||||
|
||||
def test_cli_agents_load_tools_no_filter(self):
|
||||
"""Test CLI agent tool loading without ALLOWED_TOOLS filter."""
|
||||
handler = MultiAgentHandler(verbose=False)
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
# Mock the tools import to avoid actual import dependencies
|
||||
with patch('praisonai.cli.features.agents.internet_search'), \
|
||||
patch('praisonai.cli.features.agents.read_file'):
|
||||
|
||||
tools = handler._load_tools(["internet_search", "read_file"])
|
||||
|
||||
# Should load all requested tools when no filter
|
||||
assert len(tools) == 2
|
||||
|
||||
def test_cli_agents_load_tools_with_filter(self):
|
||||
"""Test CLI agent tool loading with ALLOWED_TOOLS filter."""
|
||||
handler = MultiAgentHandler(verbose=False)
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "internet_search"}):
|
||||
# Mock the tools and filter
|
||||
mock_internet_search = MagicMock()
|
||||
mock_read_file = MagicMock()
|
||||
|
||||
with patch('praisonai.cli.features.agents.internet_search', mock_internet_search), \
|
||||
patch('praisonai.cli.features.agents.read_file', mock_read_file):
|
||||
|
||||
tools = handler._load_tools(["internet_search", "read_file"])
|
||||
|
||||
# Should only load tools allowed by filter
|
||||
assert len(tools) == 1
|
||||
assert mock_internet_search in tools
|
||||
assert mock_read_file not in tools
|
||||
|
||||
def test_cli_agents_load_tools_filter_error_handling(self):
|
||||
"""Test CLI tool loading with ALLOWED_TOOLS filter error handling."""
|
||||
handler = MultiAgentHandler(verbose=False)
|
||||
requested = ["internet_search", "read_file"]
|
||||
|
||||
with patch('os.environ.get', return_value=None):
|
||||
with patch('praisonai.cli.features.agents.internet_search'), \
|
||||
patch('praisonai.cli.features.agents.read_file'):
|
||||
baseline_tools = handler._load_tools(requested)
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ""}):
|
||||
with pytest.raises(ValueError):
|
||||
handler._load_tools(requested)
|
||||
|
||||
def test_cli_agents_load_tools_import_error_fallback(self):
|
||||
"""Test CLI tool loading fallback when ALLOWED_TOOLS filter import fails."""
|
||||
handler = MultiAgentHandler(verbose=False)
|
||||
|
||||
# Mock ImportError for the filter module
|
||||
with patch('praisonai.cli.features.agents.filter_tools_with_allowed_tools', side_effect=ImportError("Filter module not available")):
|
||||
with patch('praisonai.cli.features.agents.internet_search'), \
|
||||
patch('praisonai.cli.features.agents.read_file'):
|
||||
|
||||
tools = handler._load_tools(["internet_search", "read_file"])
|
||||
|
||||
# Should load all tools on import error
|
||||
assert len(tools) == 2
|
||||
|
||||
def test_cli_agents_load_tools_filter_value_error_propagation(self):
|
||||
"""Test that ValueError from filter is properly propagated in CLI."""
|
||||
handler = MultiAgentHandler(verbose=False)
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ""}): # Empty string causes ValueError
|
||||
with pytest.raises(ValueError, match="cannot be empty"):
|
||||
handler._load_tools(["internet_search"])
|
||||
|
||||
|
||||
class TestAllowedToolsFilterSpecSemantics:
|
||||
"""Test ALLOWED_TOOLS filter specification semantics."""
|
||||
|
||||
def test_spec_semantics_unset_all_tools_visible(self):
|
||||
"""Test spec: unset ALLOWED_TOOLS means all tools visible with warning."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
import logging
|
||||
with patch.object(logging.getLogger(), 'warning') as mock_warn:
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should return all tools
|
||||
assert filtered == available_tools
|
||||
|
||||
# Should warn about unset whitelist
|
||||
mock_warn.assert_called_once()
|
||||
assert "Tool whitelist is unset" in mock_warn.call_args[0][0]
|
||||
|
||||
def test_spec_semantics_empty_string_error(self):
|
||||
"""Test spec: empty ALLOWED_TOOLS string is an error."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ""}):
|
||||
with pytest.raises(ValueError, match="ALLOWED_TOOLS cannot be empty"):
|
||||
AllowedToolsFilter()
|
||||
|
||||
def test_spec_semantics_whitelist_only_specified_visible(self):
|
||||
"""Test spec: with ALLOWED_TOOLS values, only whitelisted tools visible."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf", "analyze_code"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should only return whitelisted tools
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
def test_spec_semantics_unknown_tools_dev_warn_strip(self):
|
||||
"""Test spec: unknown tools in dev mode are warned and stripped."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,unknown_tool", "CI": "false"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message"}
|
||||
|
||||
import logging
|
||||
with patch.object(logging.getLogger(), 'warning') as mock_warn:
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should return only known tools
|
||||
assert filtered == {"search"}
|
||||
|
||||
# Should warn about unknown tools
|
||||
mock_warn.assert_called()
|
||||
assert "unknown_tool" in str(mock_warn.call_args[0][2])
|
||||
|
||||
def test_spec_semantics_unknown_tools_ci_strict_fail(self):
|
||||
"""Test spec: unknown tools in CI mode cause strict failure."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,unknown_tool", "CI": "true"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message"}
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown tools in ALLOWED_TOOLS"):
|
||||
filter_instance.filter_tools(available_tools)
|
||||
|
||||
def test_spec_semantics_order_independence(self):
|
||||
"""Test spec: whitelist order doesn't matter."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"}):
|
||||
filter1 = AllowedToolsFilter()
|
||||
result1 = filter1.filter_tools(available_tools)
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "send_message,search"}):
|
||||
filter2 = AllowedToolsFilter()
|
||||
result2 = filter2.filter_tools(available_tools)
|
||||
|
||||
# Results should be the same regardless of order
|
||||
assert result1 == result2 == {"search", "send_message"}
|
||||
|
||||
def test_spec_semantics_duplicate_isolation(self):
|
||||
"""Test spec: duplicate tool names in whitelist are handled properly."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,search,send_message,search"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Duplicates should be de-duplicated
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
# Whitelist should be a set (no duplicates)
|
||||
whitelist = filter_instance.get_whitelist()
|
||||
assert whitelist == {"search", "send_message"}
|
||||
|
||||
def test_spec_backward_compatibility_hermes_precedence(self):
|
||||
"""Test spec: ALLOWED_TOOLS takes precedence over HERMES_ONLY_TOOLS."""
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter
|
||||
|
||||
with patch.dict(os.environ, {
|
||||
"ALLOWED_TOOLS": "search,send_message",
|
||||
"HERMES_ONLY_TOOLS": "extract_pdf"
|
||||
}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should use ALLOWED_TOOLS, not HERMES_ONLY_TOOLS
|
||||
assert filtered == {"search", "send_message"}
|
||||
assert filter_instance.env_var_name == "ALLOWED_TOOLS"
|
||||
@@ -349,6 +349,161 @@ class TestDbSessionAdapter:
|
||||
assert meta["agent_id"] == "abc"
|
||||
assert meta["total_input_tokens"] == 42
|
||||
|
||||
def test_clear_session_purges_db_history(self):
|
||||
"""clear_session must not leave stale messages reloadable from the DB."""
|
||||
from praisonai.integrations.db_session_adapter import DbSessionAdapter
|
||||
from praisonaiagents.db.protocol import DbMessage
|
||||
|
||||
persisted: list = []
|
||||
|
||||
mock_db = MagicMock()
|
||||
|
||||
def _on_start(agent_name, session_id):
|
||||
return list(persisted)
|
||||
|
||||
def _on_user(sid, content, **kw):
|
||||
persisted.append(DbMessage(role="user", content=content))
|
||||
|
||||
def _on_assistant(sid, content, **kw):
|
||||
persisted.append(DbMessage(role="assistant", content=content))
|
||||
|
||||
mock_db.on_agent_start.side_effect = _on_start
|
||||
mock_db.on_user_message.side_effect = _on_user
|
||||
mock_db.on_agent_message.side_effect = _on_assistant
|
||||
|
||||
conv = MagicMock()
|
||||
conv.delete_messages.side_effect = lambda sid, ids: persisted.clear()
|
||||
mock_db.conversation = conv
|
||||
|
||||
adapter = DbSessionAdapter(mock_db)
|
||||
sid = "bot-trim-session"
|
||||
adapter.add_message(sid, "user", "old")
|
||||
adapter.add_message(sid, "assistant", "old-reply")
|
||||
assert len(persisted) == 2
|
||||
|
||||
adapter.clear_session(sid)
|
||||
conv.delete_messages.assert_called_once_with(sid, None)
|
||||
assert persisted == []
|
||||
assert adapter.get_chat_history(sid) == []
|
||||
|
||||
adapter.add_message(sid, "user", "new-only")
|
||||
assert len(adapter.get_chat_history(sid)) == 1
|
||||
assert len(persisted) == 1
|
||||
|
||||
adapter2 = DbSessionAdapter(mock_db)
|
||||
assert adapter2.get_chat_history(sid) == [{"role": "user", "content": "new-only"}]
|
||||
|
||||
def test_metadata_survives_adapter_restart(self):
|
||||
"""set_metadata must persist to the conversation store for resume after restart."""
|
||||
from praisonai.integrations.db_session_adapter import DbSessionAdapter
|
||||
from praisonai.persistence.conversation.base import ConversationSession
|
||||
|
||||
sessions: dict = {}
|
||||
|
||||
def _get_session(sid):
|
||||
return sessions.get(sid)
|
||||
|
||||
def _create_session(session):
|
||||
sessions[session.session_id] = session
|
||||
|
||||
def _update_session(session):
|
||||
sessions[session.session_id] = session
|
||||
|
||||
conv = MagicMock()
|
||||
conv.get_session.side_effect = _get_session
|
||||
conv.create_session.side_effect = _create_session
|
||||
conv.update_session.side_effect = _update_session
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.on_agent_start.return_value = []
|
||||
mock_db.conversation = None # Prevent auto-creation of conversation attr
|
||||
mock_db._conversation_store = conv
|
||||
|
||||
sid = "meta-restart-session"
|
||||
adapter1 = DbSessionAdapter(mock_db)
|
||||
adapter1.set_metadata(
|
||||
sid,
|
||||
{
|
||||
"agent_id": "agent-xyz",
|
||||
"total_input_tokens": 99,
|
||||
"session_history": [{"turn": 1}],
|
||||
},
|
||||
)
|
||||
|
||||
adapter2 = DbSessionAdapter(mock_db)
|
||||
meta = adapter2.get_metadata(sid)
|
||||
assert meta["agent_id"] == "agent-xyz"
|
||||
assert meta["total_input_tokens"] == 99
|
||||
assert meta["session_history"] == [{"turn": 1}]
|
||||
|
||||
def test_managed_agent_metadata_persistence_e2e(self):
|
||||
"""Integration test: ManagedAgent metadata survives restart via real adapter."""
|
||||
import tempfile
|
||||
import os
|
||||
from praisonaiagents.managed import LocalManagedAgent
|
||||
from praisonai.integrations.db_session_adapter import DbSessionAdapter
|
||||
|
||||
# Create a real SQLite database for this test
|
||||
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
|
||||
db_path = tmp.name
|
||||
|
||||
try:
|
||||
# Mock PraisonDB that uses file-based conversation store
|
||||
mock_db = MagicMock()
|
||||
mock_db.on_agent_start.return_value = []
|
||||
|
||||
# Real conversation store backed by file
|
||||
from praisonai.persistence.conversation.base import ConversationSession
|
||||
from praisonai.persistence.conversation.file_store import FileConversationStore
|
||||
|
||||
conv_store = FileConversationStore(store_path=db_path + "_conv.json")
|
||||
mock_db.conversation = None
|
||||
mock_db._conversation_store = conv_store
|
||||
|
||||
# Create and run agent with metadata
|
||||
adapter = DbSessionAdapter(mock_db)
|
||||
agent = LocalManagedAgent(
|
||||
name="test-agent",
|
||||
instructions="You are a helpful assistant.",
|
||||
session_store=adapter,
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Start agent and collect session metadata
|
||||
session_id = "integration-test-session"
|
||||
adapter._ensure_session(session_id, agent.name)
|
||||
|
||||
# Set some metadata that should persist
|
||||
test_metadata = {
|
||||
"agent_id": "test-agent-123",
|
||||
"total_input_tokens": 150,
|
||||
"session_history": [{"turn": 1, "action": "started"}],
|
||||
"compute_instance": "test-instance-1"
|
||||
}
|
||||
adapter.set_metadata(session_id, test_metadata)
|
||||
|
||||
# Simulate process restart with new adapter and agent
|
||||
adapter2 = DbSessionAdapter(mock_db)
|
||||
agent2 = LocalManagedAgent(
|
||||
name="test-agent",
|
||||
instructions="You are a helpful assistant.",
|
||||
session_store=adapter2,
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Metadata should be restored from persistence
|
||||
restored_meta = adapter2.get_metadata(session_id)
|
||||
assert restored_meta["agent_id"] == "test-agent-123"
|
||||
assert restored_meta["total_input_tokens"] == 150
|
||||
assert restored_meta["session_history"] == [{"turn": 1, "action": "started"}]
|
||||
assert restored_meta["compute_instance"] == "test-instance-1"
|
||||
|
||||
finally:
|
||||
# Clean up temp files
|
||||
for path in [db_path, db_path + "_conv.json"]:
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 7. ManagedAgent factory wiring
|
||||
|
||||
@@ -6,6 +6,7 @@ Databases tested (all via local Docker):
|
||||
2. PostgreSQL — ConversationStore + ManagedAgent roundtrip
|
||||
3. MySQL — ConversationStore + ManagedAgent roundtrip
|
||||
4. Redis — StateStore + metadata persistence
|
||||
4a. Valkey — StateStore + metadata persistence
|
||||
5. MongoDB — StateStore + metadata persistence
|
||||
6. ClickHouse — Raw connectivity + data write/read
|
||||
7. JSON file — DefaultSessionStore + ManagedAgent roundtrip
|
||||
@@ -27,7 +28,6 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -562,6 +562,206 @@ class TestRedisStateStore:
|
||||
r.close()
|
||||
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 4a. Valkey — StateStore operations + metadata persistence
|
||||
# ===========================================================================
|
||||
class TestValkeyStateStore:
|
||||
"""Valkey StateStore: real operations against a live Valkey instance."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_valkey(self):
|
||||
"""Skip if Valkey not available."""
|
||||
try:
|
||||
from praisonai.persistence._valkey_client import create_valkey_client
|
||||
client = create_valkey_client(host="localhost", port=6379)
|
||||
client.ping()
|
||||
client.close()
|
||||
except Exception:
|
||||
pytest.skip("Valkey not available on localhost:6379")
|
||||
|
||||
def test_valkey_state_operations(self):
|
||||
"""StateStore get/set/delete/exists/hash with real Valkey."""
|
||||
from praisonai.persistence.state.valkey import ValkeyStateStore
|
||||
prefix = f"test_{uuid.uuid4().hex[:6]}:"
|
||||
store = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
# Basic set/get
|
||||
store.set("agent_id", "agent_valkey_001")
|
||||
assert store.get("agent_id") == "agent_valkey_001"
|
||||
|
||||
# JSON roundtrip
|
||||
state = {
|
||||
"agent_id": "agent_valkey_001",
|
||||
"agent_version": 3,
|
||||
"total_input_tokens": 500,
|
||||
"total_output_tokens": 200,
|
||||
"compute_instance_id": "docker_valkey_test",
|
||||
"session_history": [{"id": "s1", "status": "idle"}],
|
||||
}
|
||||
store.set_json("managed_state", state)
|
||||
recovered = store.get_json("managed_state")
|
||||
assert recovered["agent_id"] == "agent_valkey_001"
|
||||
assert recovered["total_input_tokens"] == 500
|
||||
assert recovered["compute_instance_id"] == "docker_valkey_test"
|
||||
assert len(recovered["session_history"]) == 1
|
||||
|
||||
# Hash operations
|
||||
store.hset("agent_meta", "version", "5")
|
||||
store.hset("agent_meta", "env_id", "env_abc")
|
||||
assert str(store.hget("agent_meta", "version")) == "5"
|
||||
all_meta = store.hgetall("agent_meta")
|
||||
assert str(all_meta["env_id"]) == "env_abc"
|
||||
|
||||
# Exists and delete
|
||||
assert store.exists("agent_id")
|
||||
store.delete("agent_id")
|
||||
assert not store.exists("agent_id")
|
||||
|
||||
finally:
|
||||
for k in store.keys():
|
||||
store.delete(k)
|
||||
store.close()
|
||||
|
||||
def test_valkey_managed_state_persist_restore(self):
|
||||
"""Simulate managed agent state persist/restore across store instances."""
|
||||
from praisonai.persistence.state.valkey import ValkeyStateStore
|
||||
prefix = f"mgd_{uuid.uuid4().hex[:6]}:"
|
||||
store = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
session_id = f"valkey_session_{uuid.uuid4().hex[:8]}"
|
||||
state = {
|
||||
"agent_id": "agent_valkey_mgd",
|
||||
"agent_version": 2,
|
||||
"environment_id": "env_valkey_001",
|
||||
"total_input_tokens": 1000,
|
||||
"total_output_tokens": 500,
|
||||
"compute_instance_id": "modal_valkey_xyz",
|
||||
"session_history": [
|
||||
{"id": session_id, "status": "idle", "title": "Valkey test"},
|
||||
],
|
||||
}
|
||||
store.set_json(f"managed:{session_id}", state)
|
||||
|
||||
# Simulate restart — new store instance, same prefix
|
||||
store2 = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
recovered = store2.get_json(f"managed:{session_id}")
|
||||
assert recovered is not None
|
||||
assert recovered["agent_id"] == "agent_valkey_mgd"
|
||||
assert recovered["total_input_tokens"] == 1000
|
||||
assert recovered["compute_instance_id"] == "modal_valkey_xyz"
|
||||
assert len(recovered["session_history"]) == 1
|
||||
|
||||
store.close()
|
||||
store2.close()
|
||||
finally:
|
||||
store3 = ValkeyStateStore(host="localhost", port=6379, prefix=prefix)
|
||||
for k in store3.keys():
|
||||
store3.delete(k)
|
||||
store3.close()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 4b. Valkey Vector Search — ValkeyVectorKnowledgeStore against a live instance
|
||||
# ===========================================================================
|
||||
class TestValkeyVectorSearch:
|
||||
"""Valkey vector search: real operations against a live ValkeySearch instance."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def check_valkey_search(self):
|
||||
"""Skip if Valkey is unreachable or ValkeySearch (FT.*) is not loaded."""
|
||||
try:
|
||||
from glide_sync import ft
|
||||
except ImportError:
|
||||
pytest.skip("valkey-glide-sync not installed")
|
||||
|
||||
try:
|
||||
from praisonai.persistence._valkey_client import create_valkey_client
|
||||
except ImportError as e:
|
||||
pytest.skip(f"praisonai not importable: {e}")
|
||||
|
||||
client = None
|
||||
try:
|
||||
client = create_valkey_client(host="localhost", port=6379)
|
||||
client.ping()
|
||||
except Exception as e:
|
||||
if client is not None:
|
||||
client.close()
|
||||
pytest.skip(f"Valkey not available on localhost:6379: {e}")
|
||||
|
||||
try:
|
||||
# Probe FT.* support; raises if module not loaded
|
||||
ft.list(client)
|
||||
except Exception as e:
|
||||
pytest.skip(f"ValkeySearch (FT.*) module not loaded on localhost:6379: {e}")
|
||||
finally:
|
||||
if client is not None:
|
||||
client.close()
|
||||
|
||||
def test_valkey_vector_search_roundtrip(self):
|
||||
"""Full create → insert → search → delete_collection roundtrip."""
|
||||
from praisonai.persistence.knowledge.valkey_vector import ValkeyVectorKnowledgeStore
|
||||
from praisonai.persistence.knowledge.base import KnowledgeDocument
|
||||
|
||||
prefix = f"praisontest_{uuid.uuid4().hex[:6]}_"
|
||||
collection = f"vec_{uuid.uuid4().hex[:6]}"
|
||||
dimension = 4
|
||||
store = ValkeyVectorKnowledgeStore(host="localhost", port=6379, prefix=prefix)
|
||||
|
||||
try:
|
||||
# Create index
|
||||
store.create_collection(collection, dimension=dimension, distance="cosine")
|
||||
assert store.collection_exists(collection)
|
||||
|
||||
# Insert documents with simple, distinguishable embeddings
|
||||
docs = [
|
||||
KnowledgeDocument(
|
||||
id="doc_a",
|
||||
content="alpha document",
|
||||
embedding=[1.0, 0.0, 0.0, 0.0],
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc_b",
|
||||
content="bravo document",
|
||||
embedding=[0.0, 1.0, 0.0, 0.0],
|
||||
),
|
||||
KnowledgeDocument(
|
||||
id="doc_c",
|
||||
content="charlie document",
|
||||
embedding=[0.0, 0.0, 1.0, 0.0],
|
||||
),
|
||||
]
|
||||
inserted = store.insert(collection, docs)
|
||||
assert set(inserted) == {"doc_a", "doc_b", "doc_c"}
|
||||
|
||||
# Vector search: query closest to doc_a
|
||||
results = store.search(collection, query_embedding=[1.0, 0.0, 0.0, 0.0], limit=2)
|
||||
assert len(results) >= 1, "Expected at least one search result"
|
||||
assert results[0].id == "doc_a", f"Expected doc_a as nearest neighbor, got {results[0].id}"
|
||||
assert results[0].content == "alpha document"
|
||||
|
||||
# Verify the bare ID is returned (no prefix leakage)
|
||||
for r in results:
|
||||
assert ":" not in r.id or r.id in {"doc_a", "doc_b", "doc_c"}
|
||||
|
||||
# get() roundtrip
|
||||
got = store.get(collection, ["doc_b"])
|
||||
assert len(got) == 1
|
||||
assert got[0].id == "doc_b"
|
||||
assert got[0].content == "bravo document"
|
||||
|
||||
# Delete one doc
|
||||
deleted = store.delete(collection, ids=["doc_c"])
|
||||
assert deleted == 1
|
||||
finally:
|
||||
try:
|
||||
store.delete_collection(collection)
|
||||
finally:
|
||||
store.close()
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 5. MongoDB — StateStore operations + metadata persistence
|
||||
# ===========================================================================
|
||||
@@ -673,7 +873,6 @@ class TestClickHouseConnectivity:
|
||||
def test_clickhouse_write_read(self):
|
||||
"""Write and read data from ClickHouse."""
|
||||
import clickhouse_connect
|
||||
|
||||
table = f"praison_test_{uuid.uuid4().hex[:8]}"
|
||||
c = clickhouse_connect.get_client(
|
||||
host="localhost", port=8123, username="clickhouse", password="clickhouse",
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
"""Test tool schema parity and OpenAI compatibility.
|
||||
|
||||
This module tests the fixes for issue #1696 to ensure BaseTool ↔ OpenAI-compatible
|
||||
tool schema parity in chat/agent dispatch.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import json
|
||||
from typing import Any, Dict, List
|
||||
from praisonaiagents.tools.base import BaseTool, ToolValidationError, validate_tool_schema_consistency
|
||||
from praisonaiagents.tools.decorator import tool, FunctionTool
|
||||
|
||||
|
||||
class MockTool(BaseTool):
|
||||
"""A simple mock tool for testing."""
|
||||
|
||||
name = "mock_tool"
|
||||
description = "A mock tool for testing"
|
||||
|
||||
def run(self, query: str, count: int = 1) -> str:
|
||||
return f"Mock result for {query} (count: {count})"
|
||||
|
||||
|
||||
class TestBasicToolSchemaValidation:
|
||||
"""Test basic tool schema validation and OpenAI compatibility."""
|
||||
|
||||
def test_base_tool_schema_structure(self):
|
||||
"""Test that BaseTool generates correct OpenAI schema structure."""
|
||||
tool = MockTool()
|
||||
schema = tool.get_schema()
|
||||
|
||||
# Validate top-level structure
|
||||
assert isinstance(schema, dict)
|
||||
assert schema["type"] == "function"
|
||||
assert "function" in schema
|
||||
|
||||
# Validate function structure
|
||||
func = schema["function"]
|
||||
assert func["name"] == "mock_tool"
|
||||
assert func["description"] == "A mock tool for testing"
|
||||
assert "parameters" in func
|
||||
|
||||
# Validate parameters structure
|
||||
params = func["parameters"]
|
||||
assert params["type"] == "object"
|
||||
assert "properties" in params
|
||||
assert "required" in params
|
||||
|
||||
# Validate specific properties
|
||||
props = params["properties"]
|
||||
assert "query" in props
|
||||
assert "count" in props
|
||||
assert props["query"]["type"] == "string"
|
||||
assert props["count"]["type"] == "integer"
|
||||
|
||||
# Validate required fields
|
||||
assert "query" in params["required"]
|
||||
assert "count" not in params["required"] # Has default value
|
||||
|
||||
def test_tool_schema_json_serialization(self):
|
||||
"""Test that tool schemas can be JSON serialized (round-trip test)."""
|
||||
tool = MockTool()
|
||||
schema = tool.get_schema()
|
||||
|
||||
# Should serialize without errors
|
||||
serialized = json.dumps(schema)
|
||||
assert isinstance(serialized, str)
|
||||
|
||||
# Should deserialize back to same structure
|
||||
deserialized = json.loads(serialized)
|
||||
assert deserialized == schema
|
||||
|
||||
def test_tool_validation_passes(self):
|
||||
"""Test that valid tools pass validation."""
|
||||
tool = MockTool()
|
||||
|
||||
# Basic validation should pass
|
||||
assert tool.validate() == True
|
||||
|
||||
# Round-trip validation should pass
|
||||
assert tool.validate_schema_roundtrip() == True
|
||||
|
||||
def test_function_tool_schema_parity(self):
|
||||
"""Test that @tool decorated functions have same schema structure."""
|
||||
@tool
|
||||
def test_function(query: str, count: int = 1) -> str:
|
||||
"""A test function for schema validation."""
|
||||
return f"Result for {query} (count: {count})"
|
||||
|
||||
schema = test_function.get_schema()
|
||||
|
||||
# Should have same structure as BaseTool
|
||||
assert schema["type"] == "function"
|
||||
func = schema["function"]
|
||||
assert func["name"] == "test_function"
|
||||
assert "parameters" in func
|
||||
|
||||
params = func["parameters"]
|
||||
assert params["type"] == "object"
|
||||
assert "properties" in params
|
||||
assert "required" in params
|
||||
|
||||
|
||||
class TestToolSchemaValidationErrors:
|
||||
"""Test that validation catches schema errors."""
|
||||
|
||||
def test_invalid_tool_missing_properties(self):
|
||||
"""Test validation fails for tools with invalid parameters."""
|
||||
|
||||
class BadTool(BaseTool):
|
||||
name = "bad_tool"
|
||||
description = "Bad tool"
|
||||
parameters = {"type": "object"} # Missing properties
|
||||
|
||||
def run(self, query: str) -> str:
|
||||
return query
|
||||
|
||||
tool = BadTool()
|
||||
with pytest.raises(ToolValidationError, match="must have a 'properties' field"):
|
||||
tool.validate()
|
||||
|
||||
def test_invalid_tool_missing_type(self):
|
||||
"""Test validation fails for tools with missing type."""
|
||||
|
||||
class BadTool(BaseTool):
|
||||
name = "bad_tool"
|
||||
description = "Bad tool"
|
||||
parameters = {"properties": {}} # Missing type
|
||||
|
||||
def run(self, query: str) -> str:
|
||||
return query
|
||||
|
||||
tool = BadTool()
|
||||
with pytest.raises(ToolValidationError, match="must have a 'type' field"):
|
||||
tool.validate()
|
||||
|
||||
|
||||
class TestToolListValidation:
|
||||
"""Test validation of tool lists for consistency."""
|
||||
|
||||
def test_valid_tool_list(self):
|
||||
"""Test that valid tool lists pass consistency checks."""
|
||||
@tool
|
||||
def tool_one(x: str) -> str:
|
||||
"""First tool."""
|
||||
return x
|
||||
|
||||
@tool
|
||||
def tool_two(y: int) -> int:
|
||||
"""Second tool."""
|
||||
return y
|
||||
|
||||
tool_three = MockTool()
|
||||
|
||||
tools = [tool_one, tool_two, tool_three]
|
||||
|
||||
# Should pass validation
|
||||
assert validate_tool_schema_consistency(tools) == True
|
||||
|
||||
def test_duplicate_tool_names_error(self):
|
||||
"""Test that duplicate tool names are caught."""
|
||||
@tool(name="duplicate_name")
|
||||
def tool_one(x: str) -> str:
|
||||
return x
|
||||
|
||||
@tool(name="duplicate_name")
|
||||
def tool_two(y: int) -> int:
|
||||
return y
|
||||
|
||||
tools = [tool_one, tool_two]
|
||||
|
||||
with pytest.raises(ToolValidationError, match="Duplicate tool name"):
|
||||
validate_tool_schema_consistency(tools)
|
||||
|
||||
def test_empty_tool_list(self):
|
||||
"""Test that empty tool lists are valid."""
|
||||
assert validate_tool_schema_consistency([]) == True
|
||||
|
||||
|
||||
class TestRealAgenticScenario:
|
||||
"""Test real agentic scenarios to ensure end-to-end compatibility."""
|
||||
|
||||
def test_tool_schema_agent_integration(self):
|
||||
"""Test that tools work with agent in real scenario."""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
@tool
|
||||
def echo_tool(message: str) -> str:
|
||||
"""Echo the input message."""
|
||||
return f"Echo: {message}"
|
||||
|
||||
# Create agent with tool - should not raise validation errors
|
||||
agent = Agent(
|
||||
name="test_agent",
|
||||
instructions="You are a helpful assistant",
|
||||
tools=[echo_tool]
|
||||
)
|
||||
|
||||
# Schema should be valid for LLM use
|
||||
schema = echo_tool.get_schema()
|
||||
assert schema["type"] == "function"
|
||||
|
||||
# Should be JSON serializable
|
||||
json.dumps(schema)
|
||||
|
||||
# Agent should format tools properly
|
||||
formatted_tools = agent._format_tools_for_completion(agent.tools)
|
||||
assert len(formatted_tools) == 1
|
||||
assert formatted_tools[0]["type"] == "function"
|
||||
|
||||
print("✅ Real agentic test passed: Tool schema integration works end-to-end")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Unit tests for unified achat dispatch (Issue #1703)."""
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def _fake_completion(content: str = "Mock async reply"):
|
||||
return SimpleNamespace(
|
||||
choices=[SimpleNamespace(message=SimpleNamespace(content=content, reasoning_content=None))]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_achat_returns_after_single_dispatch():
|
||||
"""achat must return after one unified dispatch, not loop forever."""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
agent._using_custom_llm = False
|
||||
agent.self_reflect = False
|
||||
|
||||
mock_unified = AsyncMock(return_value=_fake_completion())
|
||||
|
||||
with patch.object(agent, "_execute_unified_achat_completion", mock_unified), patch.object(
|
||||
agent, "_apply_guardrail_with_retry", return_value="Mock async reply"
|
||||
), patch.object(
|
||||
agent, "_atrigger_after_agent_hook", new_callable=AsyncMock, return_value="Mock async reply"
|
||||
), patch.object(
|
||||
agent, "_execute_callback_and_display"
|
||||
), patch.object(
|
||||
agent, "_build_messages", return_value=([{"role": "user", "content": "hi"}], "hi")
|
||||
), patch.object(
|
||||
agent, "_persist_message"
|
||||
), patch.object(
|
||||
agent, "_format_tools_for_completion", return_value=[]
|
||||
):
|
||||
result = await asyncio.wait_for(agent.achat("hi", stream=False), timeout=5)
|
||||
|
||||
assert result == "Mock async reply"
|
||||
mock_unified.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_achat_completion_uses_execute_tool_async():
|
||||
"""Async unified dispatch must wire execute_tool_async, not sync execute_tool."""
|
||||
agent = Agent(name="test", instructions="You are helpful", llm="gpt-4o-mini")
|
||||
agent._using_custom_llm = False
|
||||
|
||||
mock_dispatcher = MagicMock()
|
||||
mock_dispatcher.achat_completion = AsyncMock(return_value=_fake_completion())
|
||||
agent._unified_dispatcher = mock_dispatcher
|
||||
|
||||
with patch.object(agent, "_apply_guardrail_with_retry", return_value="ok"), patch.object(
|
||||
agent, "_atrigger_after_agent_hook", new_callable=AsyncMock, return_value="ok"
|
||||
), patch.object(agent, "_execute_callback_and_display"), patch.object(
|
||||
agent, "_build_messages", return_value=([{"role": "user", "content": "hi"}], "hi")
|
||||
), patch.object(agent, "_persist_message"), patch.object(
|
||||
agent, "_format_tools_for_completion", return_value=[]
|
||||
):
|
||||
await agent._achat_impl(
|
||||
prompt="hi",
|
||||
temperature=1.0,
|
||||
tools=None,
|
||||
output_json=None,
|
||||
output_pydantic=None,
|
||||
reasoning_steps=False,
|
||||
stream=False,
|
||||
task_name=None,
|
||||
task_description=None,
|
||||
task_id=None,
|
||||
config=None,
|
||||
force_retrieval=False,
|
||||
skip_retrieval=True,
|
||||
attachments=None,
|
||||
_trace_emitter=MagicMock(),
|
||||
)
|
||||
|
||||
call_kwargs = mock_dispatcher.achat_completion.call_args[1]
|
||||
tool_fn = call_kwargs.get("execute_tool_fn")
|
||||
assert tool_fn is not None
|
||||
assert getattr(tool_fn, "__name__", "") == "execute_tool_async"
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Tests for Agent.__deepcopy__ and Agent.clone_for_channel() (issue #1746).
|
||||
|
||||
Validates:
|
||||
- copy.deepcopy(agent) works without threading.RLock errors
|
||||
- clone_for_channel() returns an independent instance
|
||||
- Cloned agent has its own lock objects (not shared)
|
||||
- Key configuration attributes are preserved in the clone
|
||||
- Multiple clones are isolated from each other
|
||||
"""
|
||||
import copy
|
||||
import threading
|
||||
import pytest
|
||||
|
||||
|
||||
class TestAgentDeepCopy:
|
||||
"""Test Agent.__deepcopy__ and clone_for_channel()."""
|
||||
|
||||
def _make_agent(self, **kwargs):
|
||||
from praisonaiagents.agent.agent import Agent
|
||||
return Agent(name="TestAgent", instructions="test", **kwargs)
|
||||
|
||||
def test_deepcopy_does_not_raise(self):
|
||||
"""copy.deepcopy(agent) must not raise TypeError for RLock."""
|
||||
agent = self._make_agent()
|
||||
cloned = copy.deepcopy(agent)
|
||||
assert cloned is not agent
|
||||
|
||||
def test_clone_for_channel_returns_new_instance(self):
|
||||
agent = self._make_agent()
|
||||
cloned = agent.clone_for_channel()
|
||||
assert cloned is not agent
|
||||
|
||||
def test_clone_has_independent_cache_lock(self):
|
||||
"""__cache_lock must be a new RLock, not shared with original."""
|
||||
agent = self._make_agent()
|
||||
cloned = agent.clone_for_channel()
|
||||
assert cloned._Agent__cache_lock is not agent._Agent__cache_lock
|
||||
assert isinstance(cloned._Agent__cache_lock, type(threading.RLock()))
|
||||
|
||||
def test_clone_has_independent_cost_lock(self):
|
||||
"""_cost_lock must be a new Lock, not shared with original."""
|
||||
agent = self._make_agent()
|
||||
cloned = agent.clone_for_channel()
|
||||
assert cloned._cost_lock is not agent._cost_lock
|
||||
|
||||
def test_clone_preserves_name(self):
|
||||
agent = self._make_agent(name="MyAgent")
|
||||
assert agent.clone_for_channel().name == "MyAgent"
|
||||
|
||||
def test_clone_preserves_instructions(self):
|
||||
agent = self._make_agent(instructions="Do the thing")
|
||||
assert agent.clone_for_channel().instructions == "Do the thing"
|
||||
|
||||
def test_clone_preserves_role_goal_backstory(self):
|
||||
agent = self._make_agent(
|
||||
role="Analyst",
|
||||
goal="Analyse",
|
||||
backstory="Expert analyst",
|
||||
)
|
||||
cloned = agent.clone_for_channel()
|
||||
assert cloned.role == "Analyst"
|
||||
assert cloned.goal == "Analyse"
|
||||
assert cloned.backstory == "Expert analyst"
|
||||
|
||||
def test_clone_preserves_llm(self):
|
||||
agent = self._make_agent(llm="gpt-4o-mini")
|
||||
assert agent.clone_for_channel().llm == "gpt-4o-mini"
|
||||
|
||||
def test_multiple_clones_are_isolated(self):
|
||||
"""Regression: creating two clones (simulating 2nd+ gateway channel) must work."""
|
||||
agent = self._make_agent()
|
||||
clone1 = agent.clone_for_channel()
|
||||
clone2 = agent.clone_for_channel()
|
||||
assert clone1 is not clone2
|
||||
assert clone1._Agent__cache_lock is not clone2._Agent__cache_lock
|
||||
|
||||
def test_deepcopy_multiple_times(self):
|
||||
"""Successive deepcopies must all succeed (no cumulative state corruption)."""
|
||||
agent = self._make_agent()
|
||||
for _ in range(3):
|
||||
cloned = copy.deepcopy(agent)
|
||||
assert cloned is not agent
|
||||
@@ -0,0 +1,45 @@
|
||||
from types import SimpleNamespace
|
||||
import sys
|
||||
|
||||
|
||||
def test_start_verbose_composed_callbacks_forward_all_arguments(monkeypatch):
|
||||
from praisonaiagents import Agent, AgentTeam
|
||||
|
||||
on_start_calls = []
|
||||
on_complete_calls = []
|
||||
debug_messages = []
|
||||
|
||||
def on_task_start(task, task_id):
|
||||
on_start_calls.append((task, task_id))
|
||||
|
||||
def on_task_complete(task, task_output):
|
||||
on_complete_calls.append((task, task_output))
|
||||
|
||||
team = AgentTeam(
|
||||
agents=[
|
||||
Agent(name="Researcher", instructions="Research"),
|
||||
Agent(name="Writer", instructions="Write"),
|
||||
],
|
||||
hooks={
|
||||
"on_task_start": on_task_start,
|
||||
"on_task_complete": on_task_complete,
|
||||
},
|
||||
)
|
||||
|
||||
def fake_run_all_tasks():
|
||||
task_id, task = next(iter(team.tasks.items()))
|
||||
task_output = SimpleNamespace(raw="done")
|
||||
team.on_task_start(task, task_id)
|
||||
team.on_task_complete(task, task_output)
|
||||
|
||||
monkeypatch.setattr(team, "run_all_tasks", fake_run_all_tasks)
|
||||
monkeypatch.setattr(sys.stdout, "isatty", lambda: True)
|
||||
monkeypatch.setattr("praisonaiagents.agents.agents.logging.debug", lambda msg: debug_messages.append(str(msg)))
|
||||
|
||||
team.start(output="verbose")
|
||||
|
||||
assert len(on_start_calls) == 1
|
||||
assert len(on_complete_calls) == 1
|
||||
assert on_start_calls[0][1] == next(iter(team.tasks.keys()))
|
||||
assert on_complete_calls[0][1].raw == "done"
|
||||
assert all("Error in verbose task" not in message for message in debug_messages)
|
||||
@@ -48,6 +48,12 @@ class TestResolveMemory:
|
||||
result = resolve_memory("redis")
|
||||
assert isinstance(result, MemoryConfig)
|
||||
assert result.backend == MemoryBackend.REDIS
|
||||
|
||||
def test_string_valkey_backend(self):
|
||||
"""'valkey' string returns MemoryConfig with VALKEY backend."""
|
||||
result = resolve_memory("valkey")
|
||||
assert isinstance(result, MemoryConfig)
|
||||
assert result.backend == MemoryBackend.VALKEY
|
||||
|
||||
def test_string_custom_backend(self):
|
||||
"""Custom string backend is preserved."""
|
||||
|
||||
@@ -5,6 +5,8 @@ Tests Mem0Adapter lazy loading and normalization.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import importlib
|
||||
from types import SimpleNamespace
|
||||
import pytest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -174,3 +176,54 @@ class TestProtocolCompliance:
|
||||
|
||||
for method in required_methods:
|
||||
assert hasattr(Mem0Adapter, method), f"Missing method: {method}"
|
||||
|
||||
|
||||
class TestChromaKnowledgeAdapterWhereFilters:
|
||||
"""Tests for ChromaKnowledgeAdapter where-filter formatting."""
|
||||
|
||||
def _make_adapter(self):
|
||||
from praisonaiagents.knowledge.adapters.factories import ChromaKnowledgeAdapter
|
||||
|
||||
adapter = ChromaKnowledgeAdapter.__new__(ChromaKnowledgeAdapter)
|
||||
adapter.collection = MagicMock()
|
||||
adapter.collection.query.return_value = {
|
||||
"ids": [[]],
|
||||
"documents": [[]],
|
||||
"metadatas": [[]],
|
||||
"distances": [[]],
|
||||
}
|
||||
return adapter
|
||||
|
||||
def test_search_wraps_multiple_where_filters_with_and(self, monkeypatch):
|
||||
embedding_module = importlib.import_module("praisonaiagents.embedding")
|
||||
|
||||
adapter = self._make_adapter()
|
||||
monkeypatch.setattr(
|
||||
embedding_module,
|
||||
"embedding",
|
||||
lambda *args, **kwargs: SimpleNamespace(embeddings=[[0.1, 0.2]]),
|
||||
raising=False,
|
||||
)
|
||||
|
||||
adapter.search("q", user_id="u1", agent_id="a1", run_id="r1")
|
||||
|
||||
where = adapter.collection.query.call_args.kwargs["where"]
|
||||
assert where == {
|
||||
"$and": [{"user_id": "u1"}, {"agent_id": "a1"}, {"run_id": "r1"}]
|
||||
}
|
||||
|
||||
def test_search_keeps_single_where_filter_unwrapped(self, monkeypatch):
|
||||
embedding_module = importlib.import_module("praisonaiagents.embedding")
|
||||
|
||||
adapter = self._make_adapter()
|
||||
monkeypatch.setattr(
|
||||
embedding_module,
|
||||
"embedding",
|
||||
lambda *args, **kwargs: SimpleNamespace(embeddings=[[0.1, 0.2]]),
|
||||
raising=False,
|
||||
)
|
||||
|
||||
adapter.search("q", user_id="u1")
|
||||
|
||||
where = adapter.collection.query.call_args.kwargs["where"]
|
||||
assert where == {"user_id": "u1"}
|
||||
@@ -223,3 +223,188 @@ class TestDefaultContextBuilder:
|
||||
|
||||
context = builder.build(results)
|
||||
assert "\n\n" in context
|
||||
|
||||
|
||||
class TestExtractValue:
|
||||
"""Tests for the _extract_value helper function."""
|
||||
|
||||
def test_extract_value_from_dict(self):
|
||||
"""Test extracting value from dictionary."""
|
||||
from praisonaiagents.rag.context import _extract_value
|
||||
|
||||
item = {"text": "test content", "metadata": {"source": "test.pdf"}}
|
||||
|
||||
assert _extract_value(item, "text") == "test content"
|
||||
assert _extract_value(item, "metadata")["source"] == "test.pdf"
|
||||
assert _extract_value(item, "nonexistent", "default") == "default"
|
||||
|
||||
def test_extract_value_from_searchresultitem(self):
|
||||
"""Test extracting value from SearchResultItem object."""
|
||||
from praisonaiagents.rag.context import _extract_value
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
item = SearchResultItem(
|
||||
text="test content",
|
||||
source="test.pdf",
|
||||
filename="test.pdf",
|
||||
metadata={"extra": "data"}
|
||||
)
|
||||
|
||||
assert _extract_value(item, "text") == "test content"
|
||||
assert _extract_value(item, "source") == "test.pdf"
|
||||
assert _extract_value(item, "filename") == "test.pdf"
|
||||
assert _extract_value(item, "metadata")["extra"] == "data"
|
||||
assert _extract_value(item, "nonexistent", "default") == "default"
|
||||
|
||||
|
||||
class TestExtractMetadataValue:
|
||||
"""Tests for the _extract_metadata_value helper function."""
|
||||
|
||||
def test_extract_metadata_value_from_metadata(self):
|
||||
"""Test extracting value from metadata dict."""
|
||||
from praisonaiagents.rag.context import _extract_metadata_value
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
item = SearchResultItem(text="content", source="fallback.pdf")
|
||||
metadata = {"source": "metadata.pdf"}
|
||||
|
||||
# Should prefer metadata over top-level
|
||||
result = _extract_metadata_value(item, metadata, "source", "")
|
||||
assert result == "metadata.pdf"
|
||||
|
||||
def test_extract_metadata_value_fallback_to_toplevel(self):
|
||||
"""Test fallback to top-level when metadata doesn't have key."""
|
||||
from praisonaiagents.rag.context import _extract_metadata_value
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
item = SearchResultItem(text="content", source="toplevel.pdf")
|
||||
metadata = {} # Empty metadata
|
||||
|
||||
# Should fall back to top-level source
|
||||
result = _extract_metadata_value(item, metadata, "source", "")
|
||||
assert result == "toplevel.pdf"
|
||||
|
||||
def test_extract_metadata_value_with_default(self):
|
||||
"""Test fallback to default when neither metadata nor top-level have key."""
|
||||
from praisonaiagents.rag.context import _extract_metadata_value
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
item = SearchResultItem(text="content") # No source
|
||||
metadata = {} # No source in metadata
|
||||
|
||||
# Should use default
|
||||
result = _extract_metadata_value(item, metadata, "source", "default.pdf")
|
||||
assert result == "default.pdf"
|
||||
|
||||
|
||||
class TestSearchResultItemSupport:
|
||||
"""Tests for SearchResultItem object support in context functions."""
|
||||
|
||||
def test_deduplicate_chunks_with_searchresultitem(self):
|
||||
"""Test deduplication with SearchResultItem objects."""
|
||||
from praisonaiagents.rag.context import deduplicate_chunks
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
results = [
|
||||
SearchResultItem(text="Same content", source="a.pdf"),
|
||||
SearchResultItem(text="Same content", source="a.pdf"), # Duplicate
|
||||
SearchResultItem(text="Different content", source="b.pdf"),
|
||||
]
|
||||
|
||||
deduped = deduplicate_chunks(results)
|
||||
assert len(deduped) == 2 # One duplicate should be removed
|
||||
|
||||
def test_deduplicate_chunks_mixed_formats(self):
|
||||
"""Test deduplication with mixed dict and SearchResultItem formats."""
|
||||
from praisonaiagents.rag.context import deduplicate_chunks
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
results = [
|
||||
{"text": "Content A", "metadata": {"source": "a.pdf"}},
|
||||
SearchResultItem(text="Content A", source="a.pdf"), # Should be deduped
|
||||
SearchResultItem(text="Content B", source="b.pdf"),
|
||||
]
|
||||
|
||||
deduped = deduplicate_chunks(results)
|
||||
assert len(deduped) == 2 # One duplicate should be removed
|
||||
|
||||
def test_deduplicate_with_top_level_source(self):
|
||||
"""Test deduplication considers top-level source from SearchResultItem."""
|
||||
from praisonaiagents.rag.context import deduplicate_chunks
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
results = [
|
||||
SearchResultItem(text="Same text", source="source1.pdf"),
|
||||
SearchResultItem(text="Same text", source="source2.pdf"),
|
||||
]
|
||||
|
||||
deduped = deduplicate_chunks(results)
|
||||
assert len(deduped) == 2 # Different sources should not be deduped
|
||||
|
||||
def test_build_context_with_searchresultitem(self):
|
||||
"""Test building context with SearchResultItem objects."""
|
||||
from praisonaiagents.rag.context import build_context
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
results = [
|
||||
SearchResultItem(
|
||||
text="First content",
|
||||
source="source1.pdf",
|
||||
filename="file1.pdf"
|
||||
),
|
||||
SearchResultItem(
|
||||
text="Second content",
|
||||
source="source2.pdf",
|
||||
filename="file2.pdf"
|
||||
),
|
||||
]
|
||||
|
||||
context, used = build_context(results)
|
||||
|
||||
assert "First content" in context
|
||||
assert "Second content" in context
|
||||
assert "file1.pdf" in context # Should use filename
|
||||
assert "file2.pdf" in context
|
||||
assert len(used) == 2
|
||||
|
||||
def test_build_context_source_fallback(self):
|
||||
"""Test that build_context falls back to top-level source when metadata is empty."""
|
||||
from praisonaiagents.rag.context import build_context
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
# SearchResultItem with source but empty metadata
|
||||
results = [
|
||||
SearchResultItem(
|
||||
text="Content with source",
|
||||
source="fallback.pdf",
|
||||
metadata={} # Empty metadata, should fall back to top-level source
|
||||
)
|
||||
]
|
||||
|
||||
context, used = build_context(results, include_source=True)
|
||||
|
||||
assert "Content with source" in context
|
||||
assert "fallback.pdf" in context # Should use top-level source
|
||||
assert len(used) == 1
|
||||
|
||||
def test_build_context_mixed_formats_with_sources(self):
|
||||
"""Test building context with mixed dict/SearchResultItem formats."""
|
||||
from praisonaiagents.rag.context import build_context
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
results = [
|
||||
{"text": "Dict content", "metadata": {"filename": "dict.pdf"}},
|
||||
SearchResultItem(
|
||||
text="Object content",
|
||||
filename="object.pdf",
|
||||
metadata={}
|
||||
),
|
||||
]
|
||||
|
||||
context, used = build_context(results, include_source=True)
|
||||
|
||||
assert "Dict content" in context
|
||||
assert "Object content" in context
|
||||
assert "dict.pdf" in context
|
||||
assert "object.pdf" in context # Should use top-level filename
|
||||
assert len(used) == 2
|
||||
@@ -8,6 +8,7 @@ from praisonaiagents.rag.context import (
|
||||
deduplicate_chunks,
|
||||
build_context,
|
||||
)
|
||||
from praisonaiagents.knowledge.models import SearchResultItem
|
||||
|
||||
|
||||
class TestDeduplicateChunksNormalization:
|
||||
@@ -122,3 +123,25 @@ class TestEdgeCases:
|
||||
context, used = build_context(results)
|
||||
assert len(used) == 1
|
||||
assert "valid content" in context
|
||||
|
||||
|
||||
class TestSearchResultItemCompatibility:
|
||||
"""Tests for SearchResultItem object compatibility."""
|
||||
|
||||
def test_deduplicate_uses_source_attribute_fallback(self):
|
||||
"""Different source attributes should produce unique chunks."""
|
||||
results = [
|
||||
SearchResultItem(text="same content", source="a.pdf"),
|
||||
SearchResultItem(text="same content", source="b.pdf"),
|
||||
]
|
||||
deduped = deduplicate_chunks(results)
|
||||
assert len(deduped) == 2
|
||||
|
||||
def test_build_context_uses_filename_attribute_fallback(self):
|
||||
"""Source label should use object filename/source attributes."""
|
||||
results = [
|
||||
SearchResultItem(text="chunk", filename="doc.md"),
|
||||
]
|
||||
context, used = build_context(results, include_source=True)
|
||||
assert len(used) == 1
|
||||
assert "doc.md" in context
|
||||
@@ -103,6 +103,35 @@ class TestHierarchicalSessionStore:
|
||||
"""Clean up test fixtures."""
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_add_message_preserves_concurrent_writes(self):
|
||||
"""Concurrent add_message must not lose messages (same as DefaultSessionStore)."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = HierarchicalSessionStore(session_dir=tmpdir)
|
||||
reader = HierarchicalSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
history = writer.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
assert history[1]["content"] == "second"
|
||||
|
||||
def test_update_session_metadata_preserves_extended_fields(self):
|
||||
"""Metadata updates must not strip parent_id, snapshots, etc."""
|
||||
session_id = self.store.create_session(title="Parent")
|
||||
self.store.add_message(session_id, "user", "hello")
|
||||
parent_id = self.store.create_session(title="Child", parent_id=session_id)
|
||||
snapshot_id = self.store.create_snapshot(session_id, label="checkpoint")
|
||||
|
||||
assert self.store.update_session_metadata(session_id, model="gpt-4o-mini")
|
||||
|
||||
session = self.store.get_extended_session(session_id)
|
||||
assert session.metadata.get("model") == "gpt-4o-mini"
|
||||
assert session.title == "Parent"
|
||||
assert parent_id in session.children_ids
|
||||
assert any(s.id == snapshot_id for s in session.snapshots)
|
||||
|
||||
def test_create_session(self):
|
||||
"""Test creating a session."""
|
||||
session_id = self.store.create_session(title="Test Session")
|
||||
|
||||
@@ -282,6 +282,101 @@ class TestDefaultSessionStore:
|
||||
session = temp_store.get_session("session-1")
|
||||
assert session.agent_name == "TestBot"
|
||||
assert session.user_id == "user-123"
|
||||
|
||||
def test_update_session_metadata_preserves_messages(self, temp_store):
|
||||
"""Metadata updates must not drop messages added by another store instance."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = DefaultSessionStore(session_dir=tmpdir)
|
||||
reader = DefaultSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
assert reader.update_session_metadata(
|
||||
"session-1", model="gpt-4o-mini", total_tokens=42
|
||||
)
|
||||
|
||||
writer.invalidate_cache("session-1")
|
||||
history = writer.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
assert history[1]["content"] == "second"
|
||||
|
||||
session = writer.get_session("session-1")
|
||||
assert session.metadata.get("model") == "gpt-4o-mini"
|
||||
assert session.metadata.get("total_tokens") == 42
|
||||
|
||||
def test_set_agent_info_preserves_messages(self, temp_store):
|
||||
"""Agent info updates must not drop messages added by another store instance."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = DefaultSessionStore(session_dir=tmpdir)
|
||||
reader = DefaultSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
assert reader.set_agent_info("session-1", agent_name="TestBot", user_id="u-1")
|
||||
|
||||
writer.invalidate_cache("session-1")
|
||||
history = writer.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
assert history[1]["content"] == "second"
|
||||
|
||||
session = writer.get_session("session-1")
|
||||
assert session.agent_name == "TestBot"
|
||||
assert session.user_id == "u-1"
|
||||
|
||||
def test_get_chat_history_sees_writes_from_other_store(self, temp_store):
|
||||
"""Reads must reload from disk, not a stale in-memory cache."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = DefaultSessionStore(session_dir=tmpdir)
|
||||
reader = DefaultSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
history = reader.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
assert history[1]["content"] == "second"
|
||||
|
||||
def test_clear_session_preserves_new_messages(self, temp_store):
|
||||
"""Clear must reload from disk so concurrent adds are not lost."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = DefaultSessionStore(session_dir=tmpdir)
|
||||
reader = DefaultSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
assert reader.clear_session("session-1")
|
||||
|
||||
writer.invalidate_cache("session-1")
|
||||
history = writer.get_chat_history("session-1")
|
||||
assert len(history) == 0
|
||||
|
||||
def test_set_gateway_info_preserves_messages(self, temp_store):
|
||||
"""Gateway info updates must not drop messages added by another store instance."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
writer = DefaultSessionStore(session_dir=tmpdir)
|
||||
reader = DefaultSessionStore(session_dir=tmpdir)
|
||||
|
||||
writer.add_user_message("session-1", "first")
|
||||
reader._load_session("session-1")
|
||||
writer.add_user_message("session-1", "second")
|
||||
|
||||
assert reader.set_gateway_info("session-1", gateway_session_id="gw-123", agent_id="agent-456")
|
||||
|
||||
writer.invalidate_cache("session-1")
|
||||
history = writer.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
assert history[1]["content"] == "second"
|
||||
|
||||
session = writer.get_session("session-1")
|
||||
assert session.gateway_session_id == "gw-123"
|
||||
assert session.agent_id == "agent-456"
|
||||
|
||||
def test_concurrent_writes(self, temp_store):
|
||||
"""Test concurrent writes to same session."""
|
||||
@@ -313,7 +408,7 @@ class TestDefaultSessionStore:
|
||||
assert history == []
|
||||
|
||||
def test_invalidate_cache(self, temp_store):
|
||||
"""Test cache invalidation."""
|
||||
"""Test that reads always see latest disk state."""
|
||||
temp_store.add_user_message("session-1", "Hello")
|
||||
|
||||
# Modify file directly
|
||||
@@ -324,11 +419,11 @@ class TestDefaultSessionStore:
|
||||
with open(filepath, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
# Without invalidation, cache returns old data
|
||||
# Reads always reload from disk (no stale cache)
|
||||
history = temp_store.get_chat_history("session-1")
|
||||
assert len(history) == 1 # Cached
|
||||
assert len(history) == 2 # Fresh from disk
|
||||
|
||||
# After invalidation, new data is loaded
|
||||
# invalidate_cache still works but is now essentially a no-op for reads
|
||||
temp_store.invalidate_cache("session-1")
|
||||
history = temp_store.get_chat_history("session-1")
|
||||
assert len(history) == 2
|
||||
|
||||
@@ -0,0 +1,311 @@
|
||||
"""Tests for skill capability validation."""
|
||||
|
||||
import pytest
|
||||
import os
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from praisonaiagents.skills.models import SkillProperties, SkillRequirements, SkillState
|
||||
from praisonaiagents.skills.capability_validator import (
|
||||
CapabilityValidator,
|
||||
EnforcementLevel,
|
||||
ValidationResult
|
||||
)
|
||||
|
||||
|
||||
class TestSkillRequirements:
|
||||
"""Test SkillRequirements parsing and normalization."""
|
||||
|
||||
def test_from_frontmatter_empty(self):
|
||||
"""Test parsing empty frontmatter."""
|
||||
requirements = SkillRequirements.from_frontmatter({})
|
||||
assert requirements.is_empty()
|
||||
assert not requirements # Test __bool__
|
||||
|
||||
def test_from_frontmatter_tools_list(self):
|
||||
"""Test parsing tools from list."""
|
||||
metadata = {"requires_tools": ["web_search", "file_write"]}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.tools == ["web_search", "file_write"]
|
||||
assert requirements.servers == []
|
||||
assert not requirements.is_empty()
|
||||
|
||||
def test_from_frontmatter_tools_string(self):
|
||||
"""Test parsing tools from string."""
|
||||
metadata = {"requires-tools": "web_search file_write"}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.tools == ["web_search", "file_write"]
|
||||
|
||||
def test_from_frontmatter_backward_compatibility(self):
|
||||
"""Test backward compatibility with allowed-tools."""
|
||||
metadata = {"allowed-tools": ["existing_tool"]}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.tools == ["existing_tool"]
|
||||
|
||||
def test_from_frontmatter_servers(self):
|
||||
"""Test parsing server requirements."""
|
||||
metadata = {"requires_servers": ["mcp:filesystem", "http:internal"]}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.servers == ["mcp:filesystem", "http:internal"]
|
||||
|
||||
def test_from_frontmatter_env_vars(self):
|
||||
"""Test parsing environment variable requirements."""
|
||||
metadata = {"requires_env": ["API_KEY", "SECRET_TOKEN"]}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.env_vars == ["API_KEY", "SECRET_TOKEN"]
|
||||
|
||||
def test_from_frontmatter_openclaw_hints(self):
|
||||
"""Test parsing OpenClaw hints."""
|
||||
metadata = {"openclaw": {"version": "1.0", "compatibility": "hermes"}}
|
||||
requirements = SkillRequirements.from_frontmatter(metadata)
|
||||
assert requirements.openclaw_hints == {"version": "1.0", "compatibility": "hermes"}
|
||||
|
||||
def test_normalize_list(self):
|
||||
"""Test list normalization utility."""
|
||||
assert SkillRequirements._normalize_list("a b c") == ["a", "b", "c"]
|
||||
assert SkillRequirements._normalize_list("a,b,c") == ["a", "b", "c"]
|
||||
assert SkillRequirements._normalize_list("a, b c") == ["a", "b", "c"]
|
||||
assert SkillRequirements._normalize_list(["a", "b", "c"]) == ["a", "b", "c"]
|
||||
assert SkillRequirements._normalize_list("") == []
|
||||
assert SkillRequirements._normalize_list([]) == []
|
||||
assert SkillRequirements._normalize_list(123) == []
|
||||
|
||||
|
||||
class TestCapabilityValidator:
|
||||
"""Test capability validation logic."""
|
||||
|
||||
def test_enforcement_levels(self):
|
||||
"""Test different enforcement levels."""
|
||||
assert EnforcementLevel.DISABLED.value == "disabled"
|
||||
assert EnforcementLevel.WARN.value == "warn"
|
||||
assert EnforcementLevel.STRICT.value == "strict"
|
||||
|
||||
def test_validator_initialization(self):
|
||||
"""Test validator initialization."""
|
||||
validator = CapabilityValidator()
|
||||
assert validator.enforcement_level == EnforcementLevel.WARN
|
||||
|
||||
validator = CapabilityValidator(EnforcementLevel.STRICT)
|
||||
assert validator.enforcement_level == EnforcementLevel.STRICT
|
||||
|
||||
def test_validate_skill_no_requirements(self):
|
||||
"""Test validation of skill with no requirements."""
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=None
|
||||
)
|
||||
|
||||
validator = CapabilityValidator()
|
||||
result = validator.validate_skill(skill)
|
||||
|
||||
assert result.skill_name == "test-skill"
|
||||
assert result.state == SkillState.ACTIVE
|
||||
assert result.is_fully_satisfied
|
||||
assert not result.has_critical_missing
|
||||
|
||||
def test_validate_skill_empty_requirements(self):
|
||||
"""Test validation of skill with empty requirements."""
|
||||
requirements = SkillRequirements()
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
validator = CapabilityValidator()
|
||||
result = validator.validate_skill(skill)
|
||||
|
||||
assert result.state == SkillState.ACTIVE
|
||||
assert result.is_fully_satisfied
|
||||
|
||||
@patch.dict(os.environ, {"TEST_VAR": "value"})
|
||||
def test_validate_skill_satisfied_requirements(self):
|
||||
"""Test validation with all requirements satisfied."""
|
||||
requirements = SkillRequirements(
|
||||
tools=["available_tool"],
|
||||
servers=["available_server"],
|
||||
env_vars=["TEST_VAR"]
|
||||
)
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
available_tools = {"available_tool"}
|
||||
available_servers = {"available_server"}
|
||||
|
||||
validator = CapabilityValidator()
|
||||
result = validator.validate_skill(
|
||||
skill,
|
||||
available_tools=available_tools,
|
||||
available_servers=available_servers
|
||||
)
|
||||
|
||||
assert result.state == SkillState.ACTIVE
|
||||
assert result.satisfied_tools == ["available_tool"]
|
||||
assert result.satisfied_servers == ["available_server"]
|
||||
assert result.satisfied_env_vars == ["TEST_VAR"]
|
||||
assert result.is_fully_satisfied
|
||||
|
||||
def test_validate_skill_missing_tools_warn_mode(self):
|
||||
"""Test validation with missing tools in warn mode."""
|
||||
requirements = SkillRequirements(tools=["missing_tool"])
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
validator = CapabilityValidator(EnforcementLevel.WARN)
|
||||
result = validator.validate_skill(
|
||||
skill,
|
||||
available_tools=set(),
|
||||
available_servers=set()
|
||||
)
|
||||
|
||||
assert result.state == SkillState.DEGRADED
|
||||
assert result.missing_tools == ["missing_tool"]
|
||||
assert not result.is_fully_satisfied
|
||||
assert result.has_critical_missing
|
||||
assert len(result.warnings) == 1
|
||||
assert "missing_tool" in result.warnings[0]
|
||||
|
||||
def test_validate_skill_missing_tools_strict_mode(self):
|
||||
"""Test validation with missing tools in strict mode."""
|
||||
requirements = SkillRequirements(tools=["missing_tool"])
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
validator = CapabilityValidator(EnforcementLevel.STRICT)
|
||||
result = validator.validate_skill(
|
||||
skill,
|
||||
available_tools=set(),
|
||||
available_servers=set()
|
||||
)
|
||||
|
||||
assert result.state == SkillState.UNAVAILABLE
|
||||
assert result.missing_tools == ["missing_tool"]
|
||||
assert not result.is_fully_satisfied
|
||||
assert result.has_critical_missing
|
||||
assert len(result.errors) == 1
|
||||
assert "missing_tool" in result.errors[0]
|
||||
|
||||
def test_validate_skill_missing_env_vars_warn(self):
|
||||
"""Test validation with missing environment variables in warn mode."""
|
||||
requirements = SkillRequirements(env_vars=["MISSING_VAR"])
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
validator = CapabilityValidator(EnforcementLevel.WARN)
|
||||
result = validator.validate_skill(skill)
|
||||
|
||||
assert result.state == SkillState.DEGRADED
|
||||
assert result.missing_env_vars == ["MISSING_VAR"]
|
||||
assert len(result.warnings) == 1
|
||||
assert len(result.errors) == 0
|
||||
|
||||
def test_validate_skill_missing_env_vars_strict(self):
|
||||
"""Test validation with missing environment variables in strict mode."""
|
||||
requirements = SkillRequirements(env_vars=["MISSING_VAR"])
|
||||
skill = SkillProperties(
|
||||
name="test-skill",
|
||||
description="Test skill",
|
||||
requirements=requirements
|
||||
)
|
||||
|
||||
validator = CapabilityValidator(EnforcementLevel.STRICT)
|
||||
result = validator.validate_skill(skill)
|
||||
|
||||
assert result.state == SkillState.UNAVAILABLE
|
||||
assert result.missing_env_vars == ["MISSING_VAR"]
|
||||
assert len(result.warnings) == 0
|
||||
assert len(result.errors) == 1
|
||||
|
||||
def test_validation_result_to_dict(self):
|
||||
"""Test ValidationResult serialization."""
|
||||
result = ValidationResult(
|
||||
skill_name="test-skill",
|
||||
state=SkillState.DEGRADED,
|
||||
satisfied_tools=["tool1"],
|
||||
missing_tools=["tool2"],
|
||||
satisfied_servers=[],
|
||||
missing_servers=[],
|
||||
satisfied_env_vars=[],
|
||||
missing_env_vars=["VAR1"],
|
||||
warnings=["Missing tool2"],
|
||||
errors=[]
|
||||
)
|
||||
|
||||
data = result.to_dict()
|
||||
assert data["skill_name"] == "test-skill"
|
||||
assert data["state"] == "degraded"
|
||||
assert data["satisfied_tools"] == ["tool1"]
|
||||
assert data["missing_tools"] == ["tool2"]
|
||||
assert data["warnings"] == ["Missing tool2"]
|
||||
assert data["is_fully_satisfied"] == False
|
||||
assert data["has_critical_missing"] == True
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for the full capability validation system."""
|
||||
|
||||
def test_skill_parsing_with_requirements(self):
|
||||
"""Test that skills parse requirements correctly."""
|
||||
from praisonaiagents.skills.parser import read_properties
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
with TemporaryDirectory() as temp_dir:
|
||||
skill_dir = Path(temp_dir) / "test-skill"
|
||||
skill_dir.mkdir()
|
||||
|
||||
skill_md = skill_dir / "SKILL.md"
|
||||
skill_md.write_text('''---
|
||||
name: test-skill
|
||||
description: A test skill
|
||||
requires_tools:
|
||||
- web_search
|
||||
- file_write
|
||||
requires_servers:
|
||||
- mcp:filesystem
|
||||
requires_env:
|
||||
- API_KEY
|
||||
openclaw:
|
||||
version: "1.0"
|
||||
---
|
||||
|
||||
# Test Skill
|
||||
|
||||
This is a test skill.
|
||||
''')
|
||||
|
||||
properties = read_properties(skill_dir)
|
||||
assert properties.requirements is not None
|
||||
assert properties.requirements.tools == ["web_search", "file_write"]
|
||||
assert properties.requirements.servers == ["mcp:filesystem"]
|
||||
assert properties.requirements.env_vars == ["API_KEY"]
|
||||
assert properties.requirements.openclaw_hints == {"version": "1.0"}
|
||||
|
||||
def test_skill_manager_enforcement(self):
|
||||
"""Test that SkillManager respects enforcement levels."""
|
||||
from praisonaiagents.skills.manager import SkillManager
|
||||
|
||||
# Test that manager initializes with default enforcement
|
||||
manager = SkillManager()
|
||||
assert manager._validator.enforcement_level in [
|
||||
EnforcementLevel.WARN,
|
||||
EnforcementLevel.DISABLED,
|
||||
EnforcementLevel.STRICT,
|
||||
EnforcementLevel.TELEMETRY
|
||||
]
|
||||
|
||||
# Test explicit enforcement level
|
||||
strict_manager = SkillManager(EnforcementLevel.STRICT)
|
||||
assert strict_manager._validator.enforcement_level == EnforcementLevel.STRICT
|
||||
@@ -241,3 +241,89 @@ Body.
|
||||
loader.activate(skill)
|
||||
|
||||
assert skill.is_activated is True
|
||||
|
||||
def test_load_utf8_content(self):
|
||||
"""Test loading UTF-8 encoded skill content.
|
||||
|
||||
Regression test for Windows encoding issues (Issue #1695).
|
||||
Tests UTF-8 handling in SKILL.md body, scripts, and references.
|
||||
"""
|
||||
from praisonaiagents.skills.loader import SkillLoader
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
skill_dir = Path(tmpdir) / "utf8-loader-test"
|
||||
skill_dir.mkdir()
|
||||
|
||||
# Create SKILL.md with UTF-8 characters
|
||||
skill_content = """---
|
||||
name: utf8-loader-test
|
||||
description: Loader test with UTF-8 — em dashes — and more
|
||||
---
|
||||
|
||||
# UTF-8 Instructions
|
||||
|
||||
This skill handles UTF-8 content:
|
||||
- Smart quotes: "hello" and 'world'
|
||||
- Em dash: —
|
||||
- Arrows: → ← ↑ ↓
|
||||
- Accented: café, naïve, résumé
|
||||
- Emoji: 🔥 📝 ✨
|
||||
"""
|
||||
|
||||
with open(skill_dir / "SKILL.md", 'w', encoding='utf-8') as f:
|
||||
f.write(skill_content)
|
||||
|
||||
# Create scripts directory with UTF-8 script
|
||||
scripts_dir = skill_dir / "scripts"
|
||||
scripts_dir.mkdir()
|
||||
|
||||
script_content = """#!/bin/bash
|
||||
# Script with UTF-8: café processing → résumé generation
|
||||
echo "Processing naïve input with émphasis"
|
||||
"""
|
||||
|
||||
with open(scripts_dir / "process.sh", 'w', encoding='utf-8') as f:
|
||||
f.write(script_content)
|
||||
|
||||
# Create references directory with UTF-8 content
|
||||
refs_dir = skill_dir / "references"
|
||||
refs_dir.mkdir()
|
||||
|
||||
ref_content = """Reference with UTF-8:
|
||||
Café automation guide — step by step
|
||||
Arrow symbols: → ← ↑ ↓
|
||||
Special chars: "smart quotes" and émphasis
|
||||
"""
|
||||
|
||||
with open(refs_dir / "readme.txt", 'w', encoding='utf-8') as f:
|
||||
f.write(ref_content)
|
||||
|
||||
# Test loading and activation
|
||||
loader = SkillLoader()
|
||||
skill = loader.load_metadata(str(skill_dir))
|
||||
|
||||
assert skill is not None
|
||||
assert "—" in skill.properties.description # UTF-8 in metadata
|
||||
|
||||
# Test activation (loads SKILL.md body)
|
||||
success = loader.activate(skill)
|
||||
assert success is True
|
||||
assert "—" in skill.instructions # UTF-8 in instructions
|
||||
assert "café" in skill.instructions
|
||||
assert "🔥" in skill.instructions # Emoji
|
||||
|
||||
# Test script loading
|
||||
scripts = loader.load_scripts(skill)
|
||||
assert "process.sh" in scripts
|
||||
script_text = scripts["process.sh"]
|
||||
assert "café" in script_text
|
||||
assert "résumé" in script_text
|
||||
assert "naïve" in script_text
|
||||
|
||||
# Test references loading
|
||||
refs = loader.load_references(skill)
|
||||
assert "readme.txt" in refs
|
||||
ref_text = refs["readme.txt"]
|
||||
assert "Café" in ref_text
|
||||
assert "—" in ref_text
|
||||
assert "émphasis" in ref_text
|
||||
@@ -286,3 +286,54 @@ Body.
|
||||
assert False, "Should have raised ValidationError"
|
||||
except ValidationError as e:
|
||||
assert "description" in str(e).lower()
|
||||
|
||||
def test_read_properties_utf8_encoding(self):
|
||||
"""Test reading UTF-8 encoded SKILL.md files.
|
||||
|
||||
Regression test for Windows encoding issues (Issue #1695).
|
||||
Verifies that skills with non-ASCII characters load correctly.
|
||||
"""
|
||||
from praisonaiagents.skills.parser import read_properties
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
skill_dir = Path(tmpdir) / "utf8-skill"
|
||||
skill_dir.mkdir()
|
||||
|
||||
# Create SKILL.md with various UTF-8 characters
|
||||
utf8_content = """---
|
||||
name: utf8-test
|
||||
description: UTF-8 encoding test — quality — verification
|
||||
license: MIT
|
||||
compatibility: Supports café files → résumé processing
|
||||
metadata:
|
||||
author: "José María"
|
||||
notes: "Handles naïve input with émphasis"
|
||||
---
|
||||
|
||||
# UTF-8 Test Skill
|
||||
|
||||
This skill demonstrates UTF-8 character handling:
|
||||
- Em dash: —
|
||||
- Smart quotes: "hello" and 'world'
|
||||
- Arrow: →
|
||||
- Accented chars: café, naïve, résumé
|
||||
- Emoji: 🔥 📝 ✨
|
||||
|
||||
Works with internationalization!
|
||||
"""
|
||||
|
||||
# Explicitly write as UTF-8 to simulate real-world scenario
|
||||
with open(skill_dir / "SKILL.md", 'w', encoding='utf-8') as f:
|
||||
f.write(utf8_content)
|
||||
|
||||
# Test that parsing works without UnicodeDecodeError
|
||||
props = read_properties(skill_dir)
|
||||
|
||||
# Verify UTF-8 characters are preserved correctly
|
||||
assert props.name == "utf8-test"
|
||||
assert "—" in props.description # em dash
|
||||
assert "café" in props.compatibility # accented character
|
||||
assert "résumé" in props.compatibility # accented character
|
||||
assert "José María" in props.metadata["author"] # accented characters
|
||||
assert "naïve" in props.metadata["notes"] # accented character
|
||||
assert "émphasis" in props.metadata["notes"] # accented character
|
||||
@@ -295,3 +295,58 @@ description: A test skill for validation
|
||||
|
||||
assert len(errors) > 0
|
||||
assert any("directory" in e.lower() for e in errors)
|
||||
|
||||
def test_validate_utf8_skill(self):
|
||||
"""Test validation of UTF-8 encoded skill.
|
||||
|
||||
Regression test for Windows encoding issues (Issue #1695).
|
||||
Ensures validator can handle UTF-8 characters without errors.
|
||||
"""
|
||||
from praisonaiagents.skills.validator import validate
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
skill_dir = Path(tmpdir) / "utf8-validation-test"
|
||||
skill_dir.mkdir()
|
||||
|
||||
# Create SKILL.md with UTF-8 content
|
||||
utf8_content = """---
|
||||
name: utf8-validation-test
|
||||
description: Validation test with UTF-8 — em dashes — and accents
|
||||
license: MIT
|
||||
compatibility: Works with café files → résumé output
|
||||
---
|
||||
|
||||
# UTF-8 Validation Test
|
||||
|
||||
Testing validation with various UTF-8 characters.
|
||||
"""
|
||||
|
||||
# Write as UTF-8 explicitly
|
||||
with open(skill_dir / "SKILL.md", 'w', encoding='utf-8') as f:
|
||||
f.write(utf8_content)
|
||||
|
||||
# Validate should not produce errors for UTF-8 content
|
||||
errors = validate(skill_dir)
|
||||
|
||||
assert errors == [], f"Validation should pass for UTF-8 content, got errors: {errors}"
|
||||
|
||||
def test_validate_non_utf8_skill_returns_error(self):
|
||||
"""Test validation returns a structured error for non-UTF-8 SKILL.md."""
|
||||
from praisonaiagents.skills.validator import validate
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
skill_dir = Path(tmpdir) / "latin1-validation-test"
|
||||
skill_dir.mkdir()
|
||||
|
||||
latin1_content = """---
|
||||
name: latin1-validation-test
|
||||
description: Café skill
|
||||
---
|
||||
"""
|
||||
with open(skill_dir / "SKILL.md", "wb") as f:
|
||||
f.write(latin1_content.encode("cp1252"))
|
||||
|
||||
errors = validate(skill_dir)
|
||||
|
||||
assert len(errors) > 0
|
||||
assert any("utf-8" in e.lower() for e in errors)
|
||||
@@ -225,3 +225,19 @@ class TestToolCallIdPropagation:
|
||||
)
|
||||
|
||||
assert event.tool_call["id"] == "call_abc123xyz"
|
||||
|
||||
|
||||
class TestToolExecutionMixinStreamEmitterLookup:
|
||||
def test_mixin_resolves_name_mangled_stream_emitter_without_agent_class_name(self):
|
||||
from praisonaiagents.agent.tool_execution import ToolExecutionMixin
|
||||
from praisonaiagents.streaming.events import StreamEventEmitter
|
||||
|
||||
class CustomAgent(ToolExecutionMixin):
|
||||
def __init__(self):
|
||||
self.__stream_emitter = StreamEventEmitter()
|
||||
|
||||
agent = CustomAgent()
|
||||
emitter = agent._get_existing_stream_emitter()
|
||||
|
||||
assert emitter is not None
|
||||
assert isinstance(emitter, StreamEventEmitter)
|
||||
@@ -134,3 +134,37 @@ class TestTokenCollectorSinkIntegration:
|
||||
)
|
||||
assert sink.records[0]["task_id"] == "task-abc"
|
||||
assert sink.records[0]["agent_name"] == "researcher"
|
||||
|
||||
|
||||
class TestInMemoryUsageQuery:
|
||||
"""Test query adapter backed by in-memory sink records."""
|
||||
|
||||
def test_get_summary_uses_total_tokens_when_present(self):
|
||||
from praisonaiagents.telemetry.protocols import InMemoryTokenUsageSink, InMemoryUsageQuery
|
||||
|
||||
sink = InMemoryTokenUsageSink()
|
||||
sink.records.extend([
|
||||
{
|
||||
"task_id": "t1",
|
||||
"agent_name": "a1",
|
||||
"model": "m1",
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"total_tokens": 50,
|
||||
},
|
||||
{
|
||||
"task_id": "t2",
|
||||
"agent_name": "a1",
|
||||
"model": "m2",
|
||||
"input_tokens": 2,
|
||||
"output_tokens": 3,
|
||||
},
|
||||
])
|
||||
|
||||
summary = InMemoryUsageQuery(sink).get_summary()
|
||||
|
||||
assert summary["total_input_tokens"] == 12
|
||||
assert summary["total_output_tokens"] == 8
|
||||
assert summary["total_tokens"] == 55
|
||||
assert summary["by_model"]["m1"]["total_tokens"] == 50
|
||||
assert summary["by_model"]["m2"]["total_tokens"] == 5
|
||||
@@ -0,0 +1,300 @@
|
||||
"""Unit tests for the ALLOWED_TOOLS filter module."""
|
||||
import os
|
||||
import logging
|
||||
from unittest.mock import patch, MagicMock
|
||||
import pytest
|
||||
|
||||
from praisonaiagents.allowed_tools_filter import AllowedToolsFilter, filter_tools_with_allowed_tools
|
||||
|
||||
|
||||
class TestAllowedToolsFilter:
|
||||
"""Test cases for AllowedToolsFilter class."""
|
||||
|
||||
def test_filter_initialization_unset(self):
|
||||
"""Test filter initialization with no environment variable set."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
assert filter_instance.env_value is None
|
||||
assert not filter_instance.is_enabled()
|
||||
assert filter_instance.get_whitelist() is None
|
||||
|
||||
def test_filter_initialization_primary_var(self):
|
||||
"""Test filter initialization with ALLOWED_TOOLS set."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
assert filter_instance.env_value == "search,send_message"
|
||||
assert filter_instance.env_var_name == "ALLOWED_TOOLS"
|
||||
assert filter_instance.is_enabled()
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message"}
|
||||
|
||||
def test_filter_initialization_legacy_var(self):
|
||||
"""Test filter initialization with HERMES_ONLY_TOOLS set (backward compatibility)."""
|
||||
with patch.dict(os.environ, {"HERMES_ONLY_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
assert filter_instance.env_value == "search,send_message"
|
||||
assert filter_instance.env_var_name == "HERMES_ONLY_TOOLS"
|
||||
assert filter_instance.is_enabled()
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message"}
|
||||
|
||||
def test_filter_initialization_primary_precedence(self):
|
||||
"""Test that ALLOWED_TOOLS takes precedence over HERMES_ONLY_TOOLS."""
|
||||
with patch.dict(os.environ, {
|
||||
"ALLOWED_TOOLS": "search,send_message",
|
||||
"HERMES_ONLY_TOOLS": "old_tool"
|
||||
}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
assert filter_instance.env_value == "search,send_message"
|
||||
assert filter_instance.env_var_name == "ALLOWED_TOOLS"
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message"}
|
||||
|
||||
def test_filter_initialization_empty_string_error(self):
|
||||
"""Test that empty string raises ValueError."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ""}):
|
||||
with pytest.raises(ValueError, match="ALLOWED_TOOLS cannot be empty"):
|
||||
AllowedToolsFilter()
|
||||
|
||||
def test_filter_initialization_whitespace_only_error(self):
|
||||
"""Test that whitespace-only string raises ValueError."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": " "}):
|
||||
with pytest.raises(ValueError, match="ALLOWED_TOOLS cannot be empty"):
|
||||
AllowedToolsFilter()
|
||||
|
||||
def test_filter_initialization_empty_after_split_error(self):
|
||||
"""Test that comma-only string raises ValueError."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": ",,,"}):
|
||||
with pytest.raises(ValueError, match="contains no valid tool names"):
|
||||
AllowedToolsFilter()
|
||||
|
||||
def test_parse_whitelist_with_spaces(self):
|
||||
"""Test parsing comma-separated values with spaces."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": " search , send_message , extract_pdf "}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message", "extract_pdf"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,unknown_tool,send_message", "CI": "false"})
|
||||
def test_filter_tools_with_unknown_tools_dev_mode(self):
|
||||
"""Test filtering tools with unknown tools in development mode."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
with patch.object(logging.getLogger(), 'warning') as mock_warn:
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should return intersection of available and whitelisted tools
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
# Should warn about unknown tools
|
||||
mock_warn.assert_called()
|
||||
call_args = mock_warn.call_args[0]
|
||||
assert "Unknown tools" in call_args[0]
|
||||
assert ["unknown_tool"] == call_args[2]
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,unknown_tool,send_message", "CI": "true"})
|
||||
def test_filter_tools_with_unknown_tools_ci_mode(self):
|
||||
"""Test filtering tools with unknown tools in CI mode."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown tools in ALLOWED_TOOLS"):
|
||||
filter_instance.filter_tools(available_tools)
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_filter_tools_all_available(self):
|
||||
"""Test filtering when all whitelisted tools are available."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_filter_tools_partial_match(self):
|
||||
"""Test filtering when only some whitelisted tools are available."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
assert filtered == {"search"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_filter_tools_dict_input(self):
|
||||
"""Test filtering with dictionary input."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search": lambda: None, "send_message": lambda: None, "extract_pdf": lambda: None}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_filter_tools_list_input(self):
|
||||
"""Test filtering with list input."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = ["search", "send_message", "extract_pdf"]
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message,unknown_in_whitelist"})
|
||||
def test_diagnostics_data(self):
|
||||
"""Test diagnostics data collection."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
diagnostics = filter_instance.get_diagnostics()
|
||||
|
||||
assert diagnostics["env_var_name"] == "ALLOWED_TOOLS"
|
||||
assert diagnostics["env_value"] == "search,send_message,unknown_in_whitelist"
|
||||
assert diagnostics["whitelist"] == ["search", "send_message", "unknown_in_whitelist"]
|
||||
assert set(diagnostics["registered_before_filter"]) == {"search", "send_message", "extract_pdf"}
|
||||
assert set(diagnostics["registered_after_filter"]) == {"search", "send_message"}
|
||||
assert set(diagnostics["dropped_tools"]) == {"extract_pdf"}
|
||||
assert set(diagnostics["unknown_tools"]) == {"unknown_in_whitelist"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_log_diagnostics(self, caplog):
|
||||
"""Test diagnostics logging."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
filter_instance.filter_tools(available_tools)
|
||||
filter_instance.log_diagnostics()
|
||||
|
||||
log_text = caplog.text
|
||||
assert "ALLOWED_TOOLS FILTER DIAGNOSTICS" in log_text
|
||||
assert "ALLOWED_TOOLS=search,send_message" in log_text
|
||||
assert "RegisteredBeforeFilter" in log_text
|
||||
assert "RegisteredAfterFilter" in log_text
|
||||
assert "DroppedTools" in log_text
|
||||
assert "UnknownTools" in log_text
|
||||
|
||||
def test_log_diagnostics_before_filtering(self, caplog):
|
||||
"""Test logging diagnostics before calling filter_tools."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
filter_instance.log_diagnostics()
|
||||
|
||||
assert "No diagnostics available" in caplog.text
|
||||
|
||||
def test_filter_tools_unset_env_var(self, caplog):
|
||||
"""Test filtering with unset environment variable."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
|
||||
# Should return all tools
|
||||
assert filtered == available_tools
|
||||
assert "Tool whitelist is unset" in caplog.text
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "tool1", "CI": "false"})
|
||||
def test_no_matching_tools(self):
|
||||
"""Test filtering when no whitelisted tools match available tools."""
|
||||
filter_instance = AllowedToolsFilter()
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_instance.filter_tools(available_tools)
|
||||
assert filtered == set() # Empty set when no matches
|
||||
|
||||
def test_custom_env_var_name(self):
|
||||
"""Test using custom environment variable name."""
|
||||
with patch.dict(os.environ, {"CUSTOM_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter("CUSTOM_TOOLS")
|
||||
assert filter_instance.env_var_name == "CUSTOM_TOOLS"
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message"}
|
||||
|
||||
|
||||
class TestConvenienceFunction:
|
||||
"""Test cases for the convenience function."""
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"})
|
||||
def test_filter_tools_with_allowed_tools(self):
|
||||
"""Test the convenience function."""
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
filtered = filter_tools_with_allowed_tools(available_tools, log_diagnostics=False)
|
||||
assert filtered == {"search", "send_message"}
|
||||
|
||||
@patch.dict(os.environ, {"CUSTOM_TOOLS": "search"})
|
||||
def test_filter_tools_with_custom_env_var(self):
|
||||
"""Test the convenience function with custom env var."""
|
||||
available_tools = {"search", "send_message"}
|
||||
|
||||
filtered = filter_tools_with_allowed_tools(available_tools, env_var_name="CUSTOM_TOOLS", log_diagnostics=False)
|
||||
assert filtered == {"search"}
|
||||
|
||||
@patch.dict(os.environ, {"ALLOWED_TOOLS": "search"})
|
||||
def test_filter_tools_with_diagnostics(self, caplog):
|
||||
"""Test the convenience function with diagnostics enabled."""
|
||||
available_tools = {"search", "send_message"}
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
filtered = filter_tools_with_allowed_tools(available_tools, log_diagnostics=True)
|
||||
|
||||
assert filtered == {"search"}
|
||||
assert "ALLOWED_TOOLS FILTER DIAGNOSTICS" in caplog.text
|
||||
|
||||
|
||||
class TestBackwardCompatibility:
|
||||
"""Test backward compatibility aliases."""
|
||||
|
||||
@patch.dict(os.environ, {"HERMES_ONLY_TOOLS": "search,send_message"})
|
||||
def test_legacy_class_alias(self):
|
||||
"""Test HermesToolFilter alias works."""
|
||||
from praisonaiagents.allowed_tools_filter import HermesToolFilter
|
||||
|
||||
filter_instance = HermesToolFilter()
|
||||
assert isinstance(filter_instance, AllowedToolsFilter)
|
||||
assert filter_instance.env_var_name == "HERMES_ONLY_TOOLS"
|
||||
assert filter_instance.get_whitelist() == {"search", "send_message"}
|
||||
|
||||
@patch.dict(os.environ, {"HERMES_ONLY_TOOLS": "search,send_message"})
|
||||
def test_legacy_function_aliases(self):
|
||||
"""Test legacy function aliases work."""
|
||||
from praisonaiagents.allowed_tools_filter import (
|
||||
filter_tools_with_hermes, hermes_filter, apply_hermes_filter
|
||||
)
|
||||
|
||||
available_tools = {"search", "send_message", "extract_pdf"}
|
||||
|
||||
# All aliases should work and return the same result
|
||||
result1 = filter_tools_with_hermes(available_tools, log_diagnostics=False)
|
||||
result2 = hermes_filter(available_tools, log_diagnostics=False)
|
||||
result3 = apply_hermes_filter(available_tools, log_diagnostics=False)
|
||||
|
||||
expected = {"search", "send_message"}
|
||||
assert result1 == expected
|
||||
assert result2 == expected
|
||||
assert result3 == expected
|
||||
|
||||
def test_empty_tools_edge_case(self):
|
||||
"""Test filtering with empty available tools."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
|
||||
filtered = filter_instance.filter_tools(set())
|
||||
assert filtered == set()
|
||||
|
||||
# Also test with empty list and dict
|
||||
assert filter_instance.filter_tools([]) == set()
|
||||
assert filter_instance.filter_tools({}) == set()
|
||||
|
||||
def test_whitelist_immutability(self):
|
||||
"""Test that get_whitelist returns a copy, not the original."""
|
||||
with patch.dict(os.environ, {"ALLOWED_TOOLS": "search,send_message"}):
|
||||
filter_instance = AllowedToolsFilter()
|
||||
|
||||
whitelist1 = filter_instance.get_whitelist()
|
||||
whitelist2 = filter_instance.get_whitelist()
|
||||
|
||||
# Should be equal but not the same object
|
||||
assert whitelist1 == whitelist2
|
||||
assert whitelist1 is not whitelist2
|
||||
|
||||
# Modifying one shouldn't affect the other
|
||||
whitelist1.add("new_tool")
|
||||
assert whitelist1 != whitelist2
|
||||
Loaded 100 of 215 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user