mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-09 03:39:37 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1125c835ec | ||
|
|
889d04a843 | ||
|
|
fba811c477 | ||
|
|
df687d86c3 | ||
|
|
afa37b0782 |
No files matched your search
@@ -60,6 +60,7 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
|
||||
- [Planning & Workflows](#planning--workflows)
|
||||
- [Specialised Agents](#specialised-agents)
|
||||
- [MCP Protocol](#mcp-protocol)
|
||||
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
|
||||
- [Safety & Control](#safety--control)
|
||||
- [Developer Tools](#developer-tools)
|
||||
- [📊 Process Types & Patterns](#-process-types--patterns)
|
||||
@@ -188,6 +189,7 @@ npm install praisonai
|
||||
| 🧠 Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| 📄 Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
|
||||
| 💻 Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) |
|
||||
| ✏️ AI Code Editing Tools | [Example](#ai-code-editing) | [📖](https://docs.praison.ai/features/code) |
|
||||
| 📚 RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
|
||||
| 🤔 Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) |
|
||||
| 🔢 Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) |
|
||||
@@ -1628,6 +1630,48 @@ server.run() # stdio for Claude Desktop
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| Feature | Docs |
|
||||
@@ -2218,6 +2262,13 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
- 🔐 **MCP Security** - Origin validation, DNS rebinding prevention, secure sessions
|
||||
- 🔄 **MCP Resumability** - SSE stream recovery via Last-Event-ID
|
||||
|
||||
### A2A & A2UI Protocols
|
||||
|
||||
- 🔗 **A2A Protocol** - Agent-to-Agent communication for inter-agent collaboration
|
||||
- 🖼️ **A2UI Protocol** - Agent-to-User Interface for generating rich UIs from agents
|
||||
- 📋 **UI Templates** - ChatTemplate, ListTemplate, FormTemplate, DashboardTemplate
|
||||
- 🔧 **Surface Builder** - Fluent API for building declarative UIs
|
||||
|
||||
### Safety & Control
|
||||
|
||||
- 🤝 **Agent Handoffs** - Transfer context between specialised agents
|
||||
|
||||
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
|
||||
# Install Python packages (using latest versions)
|
||||
RUN pip install --no-cache-dir \
|
||||
praisonai_tools \
|
||||
"praisonai>=2.3.36" \
|
||||
"praisonai>=2.3.41" \
|
||||
"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>=2.3.36" \
|
||||
"praisonai>=2.3.41" \
|
||||
"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>=2.3.36" \
|
||||
"praisonai>=2.3.41" \
|
||||
"praisonai[ui]" \
|
||||
"praisonai[crewai]"
|
||||
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
PraisonAI A2A Server Example
|
||||
|
||||
Expose a PraisonAI Agent as an A2A (Agent2Agent) Server.
|
||||
This enables other AI agents to discover and communicate with your agent.
|
||||
|
||||
Run:
|
||||
uvicorn a2a-server:app --reload
|
||||
|
||||
Endpoints:
|
||||
GET /.well-known/agent.json - Agent Card for discovery
|
||||
GET /status - Server status
|
||||
"""
|
||||
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Search results for: {query}"
|
||||
|
||||
def calculate(expression: str) -> str:
|
||||
"""Calculate a mathematical expression."""
|
||||
try:
|
||||
return f"Result: {eval(expression)}"
|
||||
except Exception:
|
||||
return "Invalid expression"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics and answer questions",
|
||||
tools=[search_web, calculate]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(
|
||||
agent=agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0"
|
||||
)
|
||||
|
||||
# Create FastAPI app
|
||||
app = FastAPI(
|
||||
title="PraisonAI A2A Server",
|
||||
description="A2A-compatible agent server"
|
||||
)
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -9,6 +9,8 @@ Usage:
|
||||
export OPENAI_API_KEY=your_key
|
||||
python benchmarks/execution_benchmark.py
|
||||
python benchmarks/execution_benchmark.py --no-save
|
||||
python benchmarks/execution_benchmark.py --iterations 5
|
||||
python benchmarks/execution_benchmark.py --model gpt-4o --iterations 10
|
||||
"""
|
||||
|
||||
import time
|
||||
@@ -16,14 +18,10 @@ import os
|
||||
import argparse
|
||||
from importlib.metadata import version as get_version
|
||||
|
||||
# Check API key
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print("ERROR: Set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
MODEL = "gpt-4o-mini"
|
||||
ITERATIONS = 3
|
||||
PROMPT = "What is 2+2? Reply with just the number."
|
||||
# Defaults (can be overridden via CLI)
|
||||
DEFAULT_MODEL = "gpt-4o-mini"
|
||||
DEFAULT_ITERATIONS = 3
|
||||
DEFAULT_PROMPT = "What is 2+2? Reply with just the number."
|
||||
|
||||
|
||||
def get_package_versions():
|
||||
@@ -42,24 +40,24 @@ def get_package_versions():
|
||||
return versions
|
||||
|
||||
|
||||
def benchmark_praisonai():
|
||||
def benchmark_praisonai(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI agent execution using agent.start()"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI ---")
|
||||
print(f"Method: agent.start()")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=MODEL,
|
||||
llm=model,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(PROMPT)
|
||||
response = agent.start(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -70,24 +68,24 @@ def benchmark_praisonai():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_praisonai_litellm():
|
||||
def benchmark_praisonai_litellm(model, iterations, prompt):
|
||||
"""Benchmark PraisonAI with LiteLLM backend using agent.start()"""
|
||||
from praisonaiagents import Agent
|
||||
|
||||
print("\n--- PraisonAI (LiteLLM) ---")
|
||||
print(f"Method: agent.start()")
|
||||
print("Method: agent.start()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
instructions="You are a helpful assistant. Be very brief.",
|
||||
llm=f"openai/{MODEL}",
|
||||
llm=f"openai/{model}",
|
||||
verbose=False
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.start(PROMPT)
|
||||
response = agent.start(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -98,23 +96,23 @@ def benchmark_praisonai_litellm():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_agno():
|
||||
def benchmark_agno(model, iterations, prompt):
|
||||
"""Benchmark Agno agent execution using agent.run()"""
|
||||
from agno.agent import Agent
|
||||
from agno.models.openai import OpenAIChat
|
||||
|
||||
print("\n--- Agno ---")
|
||||
print(f"Method: agent.run()")
|
||||
print("Method: agent.run()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
model=OpenAIChat(id=MODEL),
|
||||
model=OpenAIChat(id=model),
|
||||
instructions=["You are a helpful assistant. Be very brief."],
|
||||
)
|
||||
|
||||
start = time.perf_counter()
|
||||
response = agent.run(PROMPT)
|
||||
response = agent.run(prompt)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
times.append(elapsed)
|
||||
@@ -127,25 +125,25 @@ def benchmark_agno():
|
||||
return avg
|
||||
|
||||
|
||||
def benchmark_crewai():
|
||||
def benchmark_crewai(model, iterations, prompt):
|
||||
"""Benchmark CrewAI agent execution using crew.kickoff()"""
|
||||
from crewai import Agent, Task, Crew
|
||||
|
||||
print("\n--- CrewAI ---")
|
||||
print(f"Method: crew.kickoff()")
|
||||
print("Method: crew.kickoff()")
|
||||
|
||||
times = []
|
||||
for i in range(ITERATIONS):
|
||||
for i in range(iterations):
|
||||
agent = Agent(
|
||||
role="Calculator",
|
||||
goal="Answer math questions",
|
||||
backstory="You are a helpful assistant.",
|
||||
llm=MODEL,
|
||||
llm=model,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
task = Task(
|
||||
description=PROMPT,
|
||||
description=prompt,
|
||||
expected_output="A number",
|
||||
agent=agent
|
||||
)
|
||||
@@ -168,7 +166,7 @@ def benchmark_crewai():
|
||||
return avg
|
||||
|
||||
|
||||
def save_results(results: dict):
|
||||
def save_results(results: dict, model: str, iterations: int, prompt: str):
|
||||
"""Save results to markdown file."""
|
||||
import os
|
||||
from datetime import datetime
|
||||
@@ -183,9 +181,9 @@ def save_results(results: dict):
|
||||
with open(filepath, 'w') as f:
|
||||
f.write('# PraisonAI Agents - Real Execution Benchmark\n\n')
|
||||
f.write(f'**Generated:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}\n')
|
||||
f.write(f'**Model:** {MODEL}\n')
|
||||
f.write(f'**Iterations:** {ITERATIONS}\n')
|
||||
f.write(f'**Prompt:** "{PROMPT}"\n\n')
|
||||
f.write(f'**Model:** {model}\n')
|
||||
f.write(f'**Iterations:** {iterations}\n')
|
||||
f.write(f'**Prompt:** "{prompt}"\n\n')
|
||||
f.write('## Results\n\n')
|
||||
f.write('| Framework | Method | Avg Time | Relative |\n')
|
||||
f.write('|-----------|--------|----------|----------|\n')
|
||||
@@ -224,14 +222,26 @@ def save_results(results: dict):
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='PraisonAI Agents - Real Execution Benchmark')
|
||||
parser.add_argument('--no-save', action='store_true', help='Do not save results to file')
|
||||
parser.add_argument('--model', '-m', type=str, default=DEFAULT_MODEL, help=f'Model to use (default: {DEFAULT_MODEL})')
|
||||
parser.add_argument('--iterations', '-i', type=int, default=DEFAULT_ITERATIONS, help=f'Number of iterations (default: {DEFAULT_ITERATIONS})')
|
||||
parser.add_argument('--prompt', '-p', type=str, default=DEFAULT_PROMPT, help=f'Prompt to use (default: "{DEFAULT_PROMPT}")')
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check API key
|
||||
if not os.environ.get('OPENAI_API_KEY'):
|
||||
print("ERROR: Set OPENAI_API_KEY environment variable")
|
||||
exit(1)
|
||||
|
||||
model = args.model
|
||||
iterations = args.iterations
|
||||
prompt = args.prompt
|
||||
|
||||
print('=' * 60)
|
||||
print('PraisonAI Agents - Real Execution Benchmark')
|
||||
print('=' * 60)
|
||||
print(f'\nModel: {MODEL}')
|
||||
print(f'Iterations: {ITERATIONS}')
|
||||
print(f'Prompt: "{PROMPT}"')
|
||||
print(f'\nModel: {model}')
|
||||
print(f'Iterations: {iterations}')
|
||||
print(f'Prompt: "{prompt}"')
|
||||
|
||||
versions = get_package_versions()
|
||||
print('\nPackage versions:')
|
||||
@@ -242,22 +252,22 @@ if __name__ == '__main__':
|
||||
|
||||
# Run benchmarks
|
||||
try:
|
||||
results['PraisonAI'] = benchmark_praisonai()
|
||||
results['PraisonAI'] = benchmark_praisonai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI error: {e}")
|
||||
|
||||
try:
|
||||
results['PraisonAI (LiteLLM)'] = benchmark_praisonai_litellm()
|
||||
results['PraisonAI (LiteLLM)'] = benchmark_praisonai_litellm(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"PraisonAI (LiteLLM) error: {e}")
|
||||
|
||||
try:
|
||||
results['Agno'] = benchmark_agno()
|
||||
results['Agno'] = benchmark_agno(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"Agno error: {e}")
|
||||
|
||||
try:
|
||||
results['CrewAI'] = benchmark_crewai()
|
||||
results['CrewAI'] = benchmark_crewai(model, iterations, prompt)
|
||||
except Exception as e:
|
||||
print(f"CrewAI error: {e}")
|
||||
|
||||
@@ -281,6 +291,6 @@ if __name__ == '__main__':
|
||||
|
||||
# Save results (unless --no-save)
|
||||
if results and not args.no_save:
|
||||
save_results(results)
|
||||
save_results(results, model, iterations, prompt)
|
||||
elif args.no_save:
|
||||
print('\nResults not saved (--no-save flag used)')
|
||||
@@ -305,3 +305,21 @@ __all__.extend([
|
||||
'LineRange'
|
||||
])
|
||||
|
||||
# AG-UI support (optional, lazy loaded)
|
||||
try:
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
__all__.append('AGUI')
|
||||
_agui_available = True
|
||||
except ImportError:
|
||||
_agui_available = False
|
||||
AGUI = None
|
||||
|
||||
# A2A support (optional, lazy loaded)
|
||||
try:
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
__all__.append('A2A')
|
||||
_a2a_available = True
|
||||
except ImportError:
|
||||
_a2a_available = False
|
||||
A2A = None
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
UI Integrations for PraisonAI Agents
|
||||
|
||||
This module provides UI protocol integrations for exposing PraisonAI Agents
|
||||
via various frontend protocols.
|
||||
|
||||
Available integrations:
|
||||
- agui: AG-UI (Agent-User Interface) protocol for CopilotKit and compatible frontends
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
|
||||
__all__ = ["AGUI"]
|
||||
@@ -0,0 +1,74 @@
|
||||
"""
|
||||
A2A Protocol Integration for PraisonAI Agents
|
||||
|
||||
This module provides A2A (Agent2Agent) protocol support,
|
||||
enabling PraisonAI Agents to be exposed as A2A Servers
|
||||
for agent-to-agent communication.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Agent Card available at: GET /.well-known/agent.json
|
||||
# A2A endpoint at: POST /a2a (JSON-RPC)
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.a2a.a2a import A2A
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
# Enums
|
||||
TaskState,
|
||||
Role,
|
||||
# Part types
|
||||
TextPart,
|
||||
FilePart,
|
||||
DataPart,
|
||||
# Core types
|
||||
Message,
|
||||
TaskStatus,
|
||||
Task,
|
||||
Artifact,
|
||||
# Agent Card types
|
||||
AgentSkill,
|
||||
AgentCapabilities,
|
||||
AgentCard,
|
||||
# Request/Response types
|
||||
MessageConfig,
|
||||
SendMessageRequest,
|
||||
# Streaming events
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Main class
|
||||
"A2A",
|
||||
# Enums
|
||||
"TaskState",
|
||||
"Role",
|
||||
# Part types
|
||||
"TextPart",
|
||||
"FilePart",
|
||||
"DataPart",
|
||||
# Core types
|
||||
"Message",
|
||||
"TaskStatus",
|
||||
"Task",
|
||||
"Artifact",
|
||||
# Agent Card types
|
||||
"AgentSkill",
|
||||
"AgentCapabilities",
|
||||
"AgentCard",
|
||||
# Request/Response types
|
||||
"MessageConfig",
|
||||
"SendMessageRequest",
|
||||
# Streaming events
|
||||
"TaskStatusUpdateEvent",
|
||||
"TaskArtifactUpdateEvent",
|
||||
]
|
||||
@@ -0,0 +1,154 @@
|
||||
"""
|
||||
A2A - Main A2A Interface Class
|
||||
|
||||
Exposes PraisonAI Agents via the A2A (Agent2Agent) protocol.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.a2a.types import AgentCard
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
from fastapi import APIRouter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class A2A:
|
||||
"""
|
||||
A2A Interface for PraisonAI Agents.
|
||||
|
||||
Exposes a PraisonAI Agent or PraisonAIAgents workflow via the A2A protocol,
|
||||
enabling agent-to-agent communication with other A2A-compatible systems.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Agent Card at: GET /.well-known/agent.json
|
||||
# A2A endpoint at: POST /a2a
|
||||
|
||||
Args:
|
||||
agent: Single PraisonAI Agent instance
|
||||
agents: PraisonAIAgents instance for multi-agent workflows
|
||||
name: Name for the A2A endpoint (defaults to agent name)
|
||||
description: Description of the agent
|
||||
url: URL where the A2A endpoint is hosted
|
||||
version: Version string for the agent
|
||||
prefix: URL prefix for the router
|
||||
tags: OpenAPI tags for the router
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Optional["Agent"] = None,
|
||||
agents: Optional["PraisonAIAgents"] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
url: str = "http://localhost:8000/a2a",
|
||||
version: str = "1.0.0",
|
||||
prefix: str = "",
|
||||
tags: Optional[List[str]] = None,
|
||||
):
|
||||
if agent is None and agents is None:
|
||||
raise ValueError("A2A requires an agent or agents instance")
|
||||
|
||||
self.agent = agent
|
||||
self.agents = agents
|
||||
self.url = url
|
||||
self.version = version
|
||||
self.prefix = prefix
|
||||
self.tags = tags or ["A2A"]
|
||||
|
||||
# Set name from agent if not provided
|
||||
if name:
|
||||
self.name = name
|
||||
elif agent and hasattr(agent, 'name'):
|
||||
self.name = agent.name
|
||||
elif agents and hasattr(agents, 'name'):
|
||||
self.name = agents.name
|
||||
else:
|
||||
self.name = "PraisonAI Agent"
|
||||
|
||||
# Set description
|
||||
if description:
|
||||
self.description = description
|
||||
elif agent and hasattr(agent, 'role'):
|
||||
self.description = agent.role
|
||||
else:
|
||||
self.description = "PraisonAI Agent via A2A"
|
||||
|
||||
# Initialize task store
|
||||
self.task_store = TaskStore()
|
||||
|
||||
# Router cache
|
||||
self._router: Optional["APIRouter"] = None
|
||||
self._agent_card: Optional[AgentCard] = None
|
||||
|
||||
def get_agent_card(self) -> AgentCard:
|
||||
"""
|
||||
Get the Agent Card for this A2A instance.
|
||||
|
||||
Returns:
|
||||
AgentCard object for A2A discovery
|
||||
"""
|
||||
if self._agent_card is None:
|
||||
self._agent_card = generate_agent_card(
|
||||
agent=self.agent,
|
||||
url=self.url,
|
||||
version=self.version,
|
||||
streaming=True,
|
||||
)
|
||||
return self._agent_card
|
||||
|
||||
def get_router(self) -> "APIRouter":
|
||||
"""
|
||||
Get the FastAPI router for this A2A instance.
|
||||
|
||||
Returns:
|
||||
FastAPI APIRouter with A2A endpoints
|
||||
"""
|
||||
if self._router is None:
|
||||
self._router = self._create_router()
|
||||
return self._router
|
||||
|
||||
def _create_router(self) -> "APIRouter":
|
||||
"""Create FastAPI router with A2A endpoints."""
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
router = APIRouter(prefix=self.prefix, tags=self.tags)
|
||||
|
||||
# Agent Card endpoint (well-known URI)
|
||||
@router.get("/.well-known/agent.json")
|
||||
async def get_agent_card():
|
||||
"""Return the Agent Card for discovery."""
|
||||
card = self.get_agent_card()
|
||||
return JSONResponse(content=card.model_dump(by_alias=True, exclude_none=True))
|
||||
|
||||
# Status endpoint
|
||||
@router.get("/status")
|
||||
async def get_status():
|
||||
"""Return server status."""
|
||||
return {
|
||||
"status": "ok",
|
||||
"name": self.name,
|
||||
"version": self.version,
|
||||
}
|
||||
|
||||
# TODO: Add JSON-RPC endpoint for message/send
|
||||
# TODO: Add streaming endpoint for message/stream
|
||||
# TODO: Add task management endpoints
|
||||
|
||||
return router
|
||||
@@ -0,0 +1,123 @@
|
||||
"""
|
||||
A2A Agent Card Generation
|
||||
|
||||
Generate Agent Cards from PraisonAI Agents for A2A discovery.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
AgentCard,
|
||||
AgentCapabilities,
|
||||
AgentSkill,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent
|
||||
|
||||
|
||||
def extract_skills_from_tools(tools: Optional[List]) -> List[AgentSkill]:
|
||||
"""
|
||||
Extract AgentSkill objects from a list of tools.
|
||||
|
||||
Args:
|
||||
tools: List of tool functions or objects
|
||||
|
||||
Returns:
|
||||
List of AgentSkill objects
|
||||
"""
|
||||
if not tools:
|
||||
return []
|
||||
|
||||
skills = []
|
||||
|
||||
for tool in tools:
|
||||
# Skip non-callable items
|
||||
if not callable(tool):
|
||||
continue
|
||||
|
||||
# Get function name
|
||||
name = getattr(tool, '__name__', None)
|
||||
if not name:
|
||||
continue
|
||||
|
||||
# Get description from docstring
|
||||
docstring = inspect.getdoc(tool)
|
||||
description = docstring.split('\n')[0] if docstring else f"Tool: {name}"
|
||||
|
||||
# Create skill
|
||||
skill = AgentSkill(
|
||||
id=name,
|
||||
name=name,
|
||||
description=description,
|
||||
tags=["tool"],
|
||||
)
|
||||
skills.append(skill)
|
||||
|
||||
return skills
|
||||
|
||||
|
||||
def generate_agent_card(
|
||||
agent: "Agent",
|
||||
url: str,
|
||||
version: str = "1.0.0",
|
||||
streaming: bool = False,
|
||||
push_notifications: bool = False,
|
||||
) -> AgentCard:
|
||||
"""
|
||||
Generate an A2A Agent Card from a PraisonAI Agent.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
url: URL where the A2A endpoint is hosted
|
||||
version: Version string for the agent
|
||||
streaming: Whether streaming is supported
|
||||
push_notifications: Whether push notifications are supported
|
||||
|
||||
Returns:
|
||||
AgentCard object for A2A discovery
|
||||
"""
|
||||
# Get agent name
|
||||
name = getattr(agent, 'name', None) or "PraisonAI Agent"
|
||||
|
||||
# Build description from role, goal, or instructions
|
||||
description_parts = []
|
||||
|
||||
role = getattr(agent, 'role', None)
|
||||
if role:
|
||||
description_parts.append(role)
|
||||
|
||||
goal = getattr(agent, 'goal', None)
|
||||
if goal:
|
||||
description_parts.append(goal)
|
||||
|
||||
instructions = getattr(agent, 'instructions', None)
|
||||
if instructions and not description_parts:
|
||||
# Use first line of instructions if no role/goal
|
||||
description_parts.append(instructions.split('\n')[0][:200])
|
||||
|
||||
description = ". ".join(description_parts) if description_parts else "PraisonAI Agent"
|
||||
|
||||
# Extract skills from tools
|
||||
tools = getattr(agent, 'tools', None)
|
||||
skills = extract_skills_from_tools(tools)
|
||||
|
||||
# Build capabilities
|
||||
capabilities = AgentCapabilities(
|
||||
streaming=streaming,
|
||||
push_notifications=push_notifications,
|
||||
)
|
||||
|
||||
# Create Agent Card
|
||||
card = AgentCard(
|
||||
name=name,
|
||||
url=url,
|
||||
version=version,
|
||||
description=description,
|
||||
capabilities=capabilities,
|
||||
skills=skills if skills else None,
|
||||
provider={"name": "PraisonAI"},
|
||||
)
|
||||
|
||||
return card
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
A2A Message Conversion
|
||||
|
||||
Convert between A2A protocol messages and PraisonAI internal format.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Message,
|
||||
Role,
|
||||
TextPart,
|
||||
Artifact,
|
||||
)
|
||||
|
||||
|
||||
def a2a_to_praisonai_messages(messages: List[Message]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert A2A messages to PraisonAI/OpenAI chat format.
|
||||
|
||||
Args:
|
||||
messages: List of A2A Message objects
|
||||
|
||||
Returns:
|
||||
List of dicts in OpenAI chat format {"role": ..., "content": ...}
|
||||
"""
|
||||
result = []
|
||||
|
||||
for msg in messages:
|
||||
# Map A2A role to OpenAI role
|
||||
role = "user" if msg.role == Role.USER else "assistant"
|
||||
|
||||
# Combine all text parts
|
||||
content_parts = []
|
||||
for part in msg.parts:
|
||||
if isinstance(part, TextPart):
|
||||
content_parts.append(part.text)
|
||||
elif hasattr(part, 'text'):
|
||||
content_parts.append(part.text)
|
||||
|
||||
content = " ".join(content_parts)
|
||||
|
||||
result.append({
|
||||
"role": role,
|
||||
"content": content,
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def praisonai_to_a2a_message(
|
||||
response: str,
|
||||
message_id: Optional[str] = None,
|
||||
context_id: Optional[str] = None,
|
||||
task_id: Optional[str] = None,
|
||||
) -> Message:
|
||||
"""
|
||||
Convert PraisonAI response to A2A Message.
|
||||
|
||||
Args:
|
||||
response: Text response from agent
|
||||
message_id: Optional message ID (auto-generated if not provided)
|
||||
context_id: Optional context ID
|
||||
task_id: Optional task ID
|
||||
|
||||
Returns:
|
||||
A2A Message object
|
||||
"""
|
||||
msg_id = message_id or f"msg-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
return Message(
|
||||
message_id=msg_id,
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text=response)],
|
||||
context_id=context_id,
|
||||
task_id=task_id,
|
||||
)
|
||||
|
||||
|
||||
def create_artifact(
|
||||
content: str,
|
||||
artifact_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
) -> Artifact:
|
||||
"""
|
||||
Create an A2A Artifact from response content.
|
||||
|
||||
Args:
|
||||
content: Text content for the artifact
|
||||
artifact_id: Optional artifact ID (auto-generated if not provided)
|
||||
name: Optional artifact name
|
||||
description: Optional artifact description
|
||||
|
||||
Returns:
|
||||
A2A Artifact object
|
||||
"""
|
||||
art_id = artifact_id or f"art-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
return Artifact(
|
||||
artifact_id=art_id,
|
||||
name=name,
|
||||
description=description,
|
||||
parts=[TextPart(text=content)],
|
||||
)
|
||||
|
||||
|
||||
def extract_user_input(messages: List[Message]) -> str:
|
||||
"""
|
||||
Extract the last user input from a list of messages.
|
||||
|
||||
Args:
|
||||
messages: List of A2A Message objects
|
||||
|
||||
Returns:
|
||||
Text content of the last user message, or empty string if none
|
||||
"""
|
||||
# Find the last user message
|
||||
for msg in reversed(messages):
|
||||
if msg.role == Role.USER:
|
||||
# Combine text parts
|
||||
content_parts = []
|
||||
for part in msg.parts:
|
||||
if isinstance(part, TextPart):
|
||||
content_parts.append(part.text)
|
||||
elif hasattr(part, 'text'):
|
||||
content_parts.append(part.text)
|
||||
return " ".join(content_parts)
|
||||
|
||||
return ""
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
A2A Streaming Events
|
||||
|
||||
SSE streaming support for A2A task updates.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import AsyncIterator, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Task,
|
||||
TaskState,
|
||||
TaskStatus,
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
Artifact,
|
||||
TextPart,
|
||||
)
|
||||
|
||||
|
||||
class A2AEventEncoder:
|
||||
"""Encode A2A events for SSE streaming."""
|
||||
|
||||
@staticmethod
|
||||
def encode_event(event_type: str, data: dict) -> str:
|
||||
"""
|
||||
Encode an event for SSE streaming.
|
||||
|
||||
Args:
|
||||
event_type: Type of event (e.g., "task.status", "task.artifact")
|
||||
data: Event data as dict
|
||||
|
||||
Returns:
|
||||
SSE-formatted string
|
||||
"""
|
||||
json_data = json.dumps(data, default=str)
|
||||
return f"event: {event_type}\ndata: {json_data}\n\n"
|
||||
|
||||
@staticmethod
|
||||
def encode_task_status(event: TaskStatusUpdateEvent) -> str:
|
||||
"""Encode TaskStatusUpdateEvent for SSE."""
|
||||
data = event.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task.status", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_task_artifact(event: TaskArtifactUpdateEvent) -> str:
|
||||
"""Encode TaskArtifactUpdateEvent for SSE."""
|
||||
data = event.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task.artifact", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_task(task: Task) -> str:
|
||||
"""Encode Task object for SSE."""
|
||||
data = task.model_dump(by_alias=True, exclude_none=True)
|
||||
return A2AEventEncoder.encode_event("task", data)
|
||||
|
||||
@staticmethod
|
||||
def encode_done() -> str:
|
||||
"""Encode stream completion event."""
|
||||
return "event: done\ndata: {}\n\n"
|
||||
|
||||
|
||||
def create_status_event(
|
||||
task_id: str,
|
||||
state: TaskState,
|
||||
context_id: Optional[str] = None,
|
||||
final: bool = False,
|
||||
) -> TaskStatusUpdateEvent:
|
||||
"""
|
||||
Create a TaskStatusUpdateEvent.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
state: New task state
|
||||
context_id: Optional context ID
|
||||
final: Whether this is the final event
|
||||
|
||||
Returns:
|
||||
TaskStatusUpdateEvent
|
||||
"""
|
||||
return TaskStatusUpdateEvent(
|
||||
task_id=task_id,
|
||||
context_id=context_id,
|
||||
status=TaskStatus(state=state),
|
||||
final=final,
|
||||
)
|
||||
|
||||
|
||||
def create_artifact_event(
|
||||
task_id: str,
|
||||
content: str,
|
||||
artifact_id: Optional[str] = None,
|
||||
context_id: Optional[str] = None,
|
||||
append: bool = False,
|
||||
last_chunk: bool = False,
|
||||
) -> TaskArtifactUpdateEvent:
|
||||
"""
|
||||
Create a TaskArtifactUpdateEvent.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
content: Artifact content
|
||||
artifact_id: Optional artifact ID
|
||||
context_id: Optional context ID
|
||||
append: Whether to append to existing artifact
|
||||
last_chunk: Whether this is the last chunk
|
||||
|
||||
Returns:
|
||||
TaskArtifactUpdateEvent
|
||||
"""
|
||||
import uuid
|
||||
art_id = artifact_id or f"art-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
artifact = Artifact(
|
||||
artifact_id=art_id,
|
||||
parts=[TextPart(text=content)],
|
||||
)
|
||||
|
||||
return TaskArtifactUpdateEvent(
|
||||
task_id=task_id,
|
||||
context_id=context_id,
|
||||
artifact=artifact,
|
||||
append=append,
|
||||
last_chunk=last_chunk,
|
||||
)
|
||||
|
||||
|
||||
async def stream_agent_response(
|
||||
agent,
|
||||
user_input: str,
|
||||
task_id: str,
|
||||
context_id: Optional[str] = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
Stream agent response as A2A events.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
user_input: User input text
|
||||
task_id: Task ID
|
||||
context_id: Optional context ID
|
||||
|
||||
Yields:
|
||||
SSE-formatted event strings
|
||||
"""
|
||||
encoder = A2AEventEncoder()
|
||||
|
||||
# Emit working status
|
||||
working_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.WORKING,
|
||||
context_id=context_id,
|
||||
)
|
||||
yield encoder.encode_task_status(working_event)
|
||||
|
||||
try:
|
||||
# Get response from agent
|
||||
response = agent.chat(user_input)
|
||||
|
||||
# Emit artifact with response
|
||||
artifact_event = create_artifact_event(
|
||||
task_id=task_id,
|
||||
content=str(response),
|
||||
context_id=context_id,
|
||||
last_chunk=True,
|
||||
)
|
||||
yield encoder.encode_task_artifact(artifact_event)
|
||||
|
||||
# Emit completed status
|
||||
completed_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.COMPLETED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(completed_event)
|
||||
|
||||
except Exception:
|
||||
# Emit failed status
|
||||
failed_event = create_status_event(
|
||||
task_id=task_id,
|
||||
state=TaskState.FAILED,
|
||||
context_id=context_id,
|
||||
final=True,
|
||||
)
|
||||
yield encoder.encode_task_status(failed_event)
|
||||
|
||||
# Emit done
|
||||
yield encoder.encode_done()
|
||||
@@ -0,0 +1,209 @@
|
||||
"""
|
||||
A2A Task Store
|
||||
|
||||
In-memory storage for A2A Task lifecycle management.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from praisonaiagents.ui.a2a.types import (
|
||||
Task,
|
||||
TaskState,
|
||||
TaskStatus,
|
||||
Message,
|
||||
Artifact,
|
||||
)
|
||||
|
||||
|
||||
class TaskStore:
|
||||
"""
|
||||
In-memory store for A2A Tasks.
|
||||
|
||||
Manages task lifecycle including creation, status updates,
|
||||
artifact storage, and message history.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty task store."""
|
||||
self._tasks: Dict[str, Task] = {}
|
||||
|
||||
def create_task(self, message: Message, context_id: Optional[str] = None) -> Task:
|
||||
"""
|
||||
Create a new task from an incoming message.
|
||||
|
||||
Args:
|
||||
message: The initial message that triggers the task
|
||||
context_id: Optional context ID (uses message's context_id if not provided)
|
||||
|
||||
Returns:
|
||||
Newly created Task in SUBMITTED state
|
||||
"""
|
||||
task_id = f"task-{uuid.uuid4().hex[:12]}"
|
||||
ctx_id = context_id or message.context_id
|
||||
|
||||
# Create initial status
|
||||
status = TaskStatus(
|
||||
state=TaskState.SUBMITTED,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
# Create task with initial message in history
|
||||
task = Task(
|
||||
id=task_id,
|
||||
context_id=ctx_id,
|
||||
status=status,
|
||||
history=[message],
|
||||
artifacts=None,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = task
|
||||
return task
|
||||
|
||||
def get_task(self, task_id: str) -> Optional[Task]:
|
||||
"""
|
||||
Retrieve a task by its ID.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to look up
|
||||
|
||||
Returns:
|
||||
Task if found, None otherwise
|
||||
"""
|
||||
return self._tasks.get(task_id)
|
||||
|
||||
def update_status(
|
||||
self,
|
||||
task_id: str,
|
||||
state: TaskState,
|
||||
message: Optional[Message] = None,
|
||||
) -> Optional[Task]:
|
||||
"""
|
||||
Update the status of a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
state: New task state
|
||||
message: Optional status message
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Create new status
|
||||
new_status = TaskStatus(
|
||||
state=state,
|
||||
message=message,
|
||||
timestamp=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
# Update task (create new instance for immutability)
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=new_status,
|
||||
artifacts=task.artifacts,
|
||||
history=task.history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def add_artifact(self, task_id: str, artifact: Artifact) -> Optional[Task]:
|
||||
"""
|
||||
Add an artifact to a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
artifact: Artifact to add
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Get existing artifacts or create new list
|
||||
artifacts = list(task.artifacts) if task.artifacts else []
|
||||
artifacts.append(artifact)
|
||||
|
||||
# Update task
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=task.status,
|
||||
artifacts=artifacts,
|
||||
history=task.history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def add_to_history(self, task_id: str, message: Message) -> Optional[Task]:
|
||||
"""
|
||||
Add a message to task history.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to update
|
||||
message: Message to add to history
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
task = self._tasks.get(task_id)
|
||||
if not task:
|
||||
return None
|
||||
|
||||
# Get existing history or create new list
|
||||
history = list(task.history) if task.history else []
|
||||
history.append(message)
|
||||
|
||||
# Update task
|
||||
updated = Task(
|
||||
id=task.id,
|
||||
context_id=task.context_id,
|
||||
status=task.status,
|
||||
artifacts=task.artifacts,
|
||||
history=history,
|
||||
metadata=task.metadata,
|
||||
)
|
||||
|
||||
self._tasks[task_id] = updated
|
||||
return updated
|
||||
|
||||
def list_tasks(self, context_id: Optional[str] = None) -> List[Task]:
|
||||
"""
|
||||
List all tasks, optionally filtered by context.
|
||||
|
||||
Args:
|
||||
context_id: Optional context ID to filter by
|
||||
|
||||
Returns:
|
||||
List of tasks
|
||||
"""
|
||||
if context_id:
|
||||
return [t for t in self._tasks.values() if t.context_id == context_id]
|
||||
return list(self._tasks.values())
|
||||
|
||||
def cancel_task(self, task_id: str) -> Optional[Task]:
|
||||
"""
|
||||
Cancel a task.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to cancel
|
||||
|
||||
Returns:
|
||||
Updated Task if found, None otherwise
|
||||
"""
|
||||
return self.update_status(task_id, TaskState.CANCELLED)
|
||||
|
||||
def clear(self):
|
||||
"""Clear all tasks from the store."""
|
||||
self._tasks.clear()
|
||||
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
A2A Protocol Types
|
||||
|
||||
Pydantic models for the Agent2Agent (A2A) protocol.
|
||||
Based on the A2A Protocol Specification v1.0.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Enums
|
||||
# =============================================================================
|
||||
|
||||
class TaskState(str, Enum):
|
||||
"""
|
||||
Defines the possible lifecycle states of a Task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.3.
|
||||
"""
|
||||
SUBMITTED = "submitted"
|
||||
WORKING = "working"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
INPUT_REQUIRED = "input_required"
|
||||
AUTH_REQUIRED = "auth_required"
|
||||
|
||||
|
||||
class Role(str, Enum):
|
||||
"""
|
||||
Defines the sender of a message in A2A protocol communication.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.5.
|
||||
"""
|
||||
USER = "user"
|
||||
AGENT = "agent"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Part Types
|
||||
# =============================================================================
|
||||
|
||||
class TextPart(BaseModel):
|
||||
"""
|
||||
A text content part.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.6.
|
||||
"""
|
||||
text: str
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class FilePart(BaseModel):
|
||||
"""
|
||||
A file content part.
|
||||
|
||||
Can contain either a URI reference or inline bytes.
|
||||
Based on A2A Protocol Specification Section 4.1.7.
|
||||
"""
|
||||
file_uri: Optional[str] = Field(default=None, alias="fileWithUri")
|
||||
file_bytes: Optional[bytes] = Field(default=None, alias="fileWithBytes")
|
||||
media_type: Optional[str] = Field(default=None, alias="mediaType")
|
||||
name: Optional[str] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class DataPart(BaseModel):
|
||||
"""
|
||||
A structured data part (JSON blob).
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.8.
|
||||
"""
|
||||
data: Dict[str, Any]
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# Union type for all Part types
|
||||
Part = Union[TextPart, FilePart, DataPart]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Core Types
|
||||
# =============================================================================
|
||||
|
||||
class Message(BaseModel):
|
||||
"""
|
||||
A communication turn between client and server.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
message_id: str = Field(alias="messageId")
|
||||
role: Role
|
||||
parts: List[Part]
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
task_id: Optional[str] = Field(default=None, alias="taskId")
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
extensions: Optional[Dict[str, str]] = None
|
||||
reference_task_ids: Optional[List[str]] = Field(default=None, alias="referenceTaskIds")
|
||||
|
||||
|
||||
class TaskStatus(BaseModel):
|
||||
"""
|
||||
Container for the status of a task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.2.
|
||||
"""
|
||||
state: TaskState
|
||||
message: Optional[Message] = None
|
||||
timestamp: Optional[str] = None
|
||||
|
||||
|
||||
class Artifact(BaseModel):
|
||||
"""
|
||||
An output generated by the agent as a result of a task.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.9.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
artifact_id: str = Field(alias="artifactId")
|
||||
name: Optional[str] = None
|
||||
description: Optional[str] = None
|
||||
parts: List[Part]
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
extensions: Optional[Dict[str, str]] = None
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
"""
|
||||
The core unit of action for A2A.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.1.1.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
status: TaskStatus
|
||||
artifacts: Optional[List[Artifact]] = None
|
||||
history: Optional[List[Message]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Agent Card Types
|
||||
# =============================================================================
|
||||
|
||||
class AgentSkill(BaseModel):
|
||||
"""
|
||||
Describes a skill/capability of an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
tags: Optional[List[str]] = None
|
||||
examples: Optional[List[str]] = None
|
||||
input_modes: Optional[List[str]] = Field(default=None, alias="inputModes")
|
||||
output_modes: Optional[List[str]] = Field(default=None, alias="outputModes")
|
||||
|
||||
|
||||
class AgentCapabilities(BaseModel):
|
||||
"""
|
||||
Describes the capabilities of an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.4.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
streaming: bool = False
|
||||
push_notifications: bool = Field(default=False, alias="pushNotifications")
|
||||
state_transition_history: bool = Field(default=False, alias="stateTransitionHistory")
|
||||
|
||||
|
||||
class AgentCard(BaseModel):
|
||||
"""
|
||||
JSON metadata document describing an agent.
|
||||
|
||||
Published by an A2A Server for discovery.
|
||||
Based on A2A Protocol Specification Section 8.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
name: str
|
||||
url: str
|
||||
version: str
|
||||
description: Optional[str] = None
|
||||
capabilities: AgentCapabilities
|
||||
skills: Optional[List[AgentSkill]] = None
|
||||
default_input_modes: Optional[List[str]] = Field(default=None, alias="defaultInputModes")
|
||||
default_output_modes: Optional[List[str]] = Field(default=None, alias="defaultOutputModes")
|
||||
provider: Optional[Dict[str, str]] = None
|
||||
documentation_url: Optional[str] = Field(default=None, alias="documentationUrl")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Request/Response Types
|
||||
# =============================================================================
|
||||
|
||||
class MessageConfig(BaseModel):
|
||||
"""
|
||||
Configuration for message sending.
|
||||
|
||||
Based on A2A Protocol Specification Section 3.2.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
accept_content_types: Optional[List[str]] = Field(default=None, alias="acceptContentTypes")
|
||||
blocking: Optional[bool] = None
|
||||
history_length: Optional[int] = Field(default=None, alias="historyLength")
|
||||
|
||||
|
||||
class SendMessageRequest(BaseModel):
|
||||
"""
|
||||
Request to send a message to an agent.
|
||||
|
||||
Based on A2A Protocol Specification Section 3.1.1.
|
||||
"""
|
||||
message: Message
|
||||
configuration: Optional[MessageConfig] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Streaming Event Types
|
||||
# =============================================================================
|
||||
|
||||
class TaskStatusUpdateEvent(BaseModel):
|
||||
"""
|
||||
Event sent by agent to notify client of task status change.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.2.1.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
task_id: str = Field(alias="taskId")
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
status: TaskStatus
|
||||
final: bool = False
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class TaskArtifactUpdateEvent(BaseModel):
|
||||
"""
|
||||
Event sent when an artifact is generated.
|
||||
|
||||
Based on A2A Protocol Specification Section 4.2.2.
|
||||
"""
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
task_id: str = Field(alias="taskId")
|
||||
context_id: Optional[str] = Field(default=None, alias="contextId")
|
||||
artifact: Artifact
|
||||
append: bool = False
|
||||
last_chunk: bool = Field(default=False, alias="lastChunk")
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
A2UI (Agent-to-User Interface) Protocol Integration for PraisonAI Agents
|
||||
|
||||
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,
|
||||
)
|
||||
"""
|
||||
|
||||
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,
|
||||
create_a2ui_part,
|
||||
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,
|
||||
)
|
||||
|
||||
__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_MIME_TYPE",
|
||||
"A2UI_EXTENSION_URI",
|
||||
"STANDARD_CATALOG_ID",
|
||||
"create_a2ui_part",
|
||||
"is_a2ui_part",
|
||||
"get_a2ui_agent_extension",
|
||||
"AgentExtension",
|
||||
# Surface builder
|
||||
"Surface",
|
||||
# Agent wrapper
|
||||
"A2UIAgent",
|
||||
# Templates
|
||||
"ChatTemplate",
|
||||
"ListTemplate",
|
||||
"FormTemplate",
|
||||
"DashboardTemplate",
|
||||
]
|
||||
@@ -0,0 +1,358 @@
|
||||
"""
|
||||
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})
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
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
|
||||
)
|
||||
@@ -0,0 +1,364 @@
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
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",
|
||||
]
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,523 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
AG-UI Protocol Integration for PraisonAI Agents
|
||||
|
||||
This module provides AG-UI (Agent-User Interface) protocol support,
|
||||
enabling PraisonAI Agents to be exposed via a standardized streaming API
|
||||
compatible with CopilotKit and other AG-UI frontends.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui.agui import AGUI
|
||||
|
||||
__all__ = ["AGUI"]
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
AGUI - Main AG-UI Interface Class
|
||||
|
||||
Exposes PraisonAI Agents via the AG-UI protocol.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from typing import AsyncIterator, List, Optional, TYPE_CHECKING
|
||||
|
||||
from praisonaiagents.ui.agui.types import (
|
||||
BaseEvent,
|
||||
RunAgentInput,
|
||||
)
|
||||
from praisonaiagents.ui.agui.conversion import (
|
||||
agui_messages_to_praisonai,
|
||||
validate_state,
|
||||
extract_user_input,
|
||||
)
|
||||
from praisonaiagents.ui.agui.streaming import (
|
||||
create_run_error_event,
|
||||
async_stream_agent_response,
|
||||
async_stream_agents_response,
|
||||
)
|
||||
from praisonaiagents.ui.agui.encoder import EventEncoder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from praisonaiagents import Agent, PraisonAIAgents
|
||||
from fastapi import APIRouter
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AGUI:
|
||||
"""
|
||||
AG-UI Interface for PraisonAI Agents.
|
||||
|
||||
Exposes a PraisonAI Agent or PraisonAIAgents workflow via the AG-UI protocol,
|
||||
enabling integration with CopilotKit and other AG-UI compatible frontends.
|
||||
|
||||
Usage:
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Assistant", role="Helper", goal="Help users")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
Args:
|
||||
agent: Single PraisonAI Agent instance
|
||||
agents: PraisonAIAgents instance for multi-agent workflows
|
||||
name: Name for the AG-UI endpoint (defaults to agent name)
|
||||
description: Description of the agent
|
||||
prefix: URL prefix for the router (e.g., "/api/v1")
|
||||
tags: OpenAPI tags for the router
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Optional["Agent"] = None,
|
||||
agents: Optional["PraisonAIAgents"] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
prefix: str = "",
|
||||
tags: Optional[List[str]] = None,
|
||||
):
|
||||
if agent is None and agents is None:
|
||||
raise ValueError("AGUI requires an agent or agents instance")
|
||||
|
||||
self.agent = agent
|
||||
self.agents = agents
|
||||
self.prefix = prefix
|
||||
self.tags = tags or ["AGUI"]
|
||||
|
||||
# Set name from agent if not provided
|
||||
if name:
|
||||
self.name = name
|
||||
elif agent and hasattr(agent, 'name'):
|
||||
self.name = agent.name
|
||||
elif agents and hasattr(agents, 'name'):
|
||||
self.name = agents.name
|
||||
else:
|
||||
self.name = "PraisonAI Agent"
|
||||
|
||||
# Set description
|
||||
if description:
|
||||
self.description = description
|
||||
elif agent and hasattr(agent, 'role'):
|
||||
self.description = agent.role
|
||||
else:
|
||||
self.description = "PraisonAI Agent via AG-UI"
|
||||
|
||||
self._router: Optional["APIRouter"] = None
|
||||
|
||||
def get_router(self) -> "APIRouter":
|
||||
"""
|
||||
Get the FastAPI router for this AGUI instance.
|
||||
|
||||
Returns:
|
||||
FastAPI APIRouter with AG-UI endpoints
|
||||
"""
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
if self._router is not None:
|
||||
return self._router
|
||||
|
||||
self._router = APIRouter(prefix=self.prefix, tags=self.tags)
|
||||
|
||||
# Attach routes
|
||||
self._attach_routes(self._router)
|
||||
|
||||
return self._router
|
||||
|
||||
def _attach_routes(self, router: "APIRouter") -> None:
|
||||
"""Attach AG-UI routes to the router."""
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
encoder = EventEncoder()
|
||||
|
||||
@router.post("/agui")
|
||||
async def run_agent_agui(run_input: RunAgentInput):
|
||||
"""Run the agent via AG-UI protocol."""
|
||||
async def event_generator():
|
||||
async for event in self._run_agent(run_input):
|
||||
yield encoder.encode(event)
|
||||
|
||||
return StreamingResponse(
|
||||
event_generator(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"Connection": "keep-alive",
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "POST, GET, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "*",
|
||||
},
|
||||
)
|
||||
|
||||
@router.get("/status")
|
||||
async def get_status():
|
||||
"""Get agent status."""
|
||||
return {"status": "available"}
|
||||
|
||||
async def _run_agent(self, run_input: RunAgentInput) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Run the agent with the given input.
|
||||
|
||||
Args:
|
||||
run_input: AG-UI run input
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
run_id = run_input.run_id or str(uuid.uuid4())
|
||||
thread_id = run_input.thread_id
|
||||
|
||||
try:
|
||||
# Convert messages
|
||||
messages = run_input.messages or []
|
||||
|
||||
# Extract user input
|
||||
user_input = extract_user_input(messages)
|
||||
|
||||
# Validate state
|
||||
session_state = validate_state(run_input.state, thread_id)
|
||||
|
||||
# Get user_id from forwarded_props if available
|
||||
user_id = None
|
||||
if run_input.forwarded_props and isinstance(run_input.forwarded_props, dict):
|
||||
user_id = run_input.forwarded_props.get("user_id")
|
||||
|
||||
# Run appropriate handler
|
||||
if self.agent:
|
||||
async for event in async_stream_agent_response(
|
||||
agent=self.agent,
|
||||
user_input=user_input,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
session_state=session_state,
|
||||
messages=agui_messages_to_praisonai(messages),
|
||||
):
|
||||
yield event
|
||||
|
||||
elif self.agents:
|
||||
async for event in async_stream_agents_response(
|
||||
agents=self.agents,
|
||||
user_input=user_input,
|
||||
thread_id=thread_id,
|
||||
run_id=run_id,
|
||||
session_state=session_state,
|
||||
):
|
||||
yield event
|
||||
|
||||
else:
|
||||
yield create_run_error_event("No agent or agents configured")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error running agent: {e}", exc_info=True)
|
||||
yield create_run_error_event(str(e))
|
||||
@@ -0,0 +1,211 @@
|
||||
"""
|
||||
AG-UI Message Conversion
|
||||
|
||||
Converts between AG-UI message format and PraisonAI message format.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from praisonaiagents.ui.agui.types import Message, ToolCall, FunctionCall
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def agui_messages_to_praisonai(messages: List[Message]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert AG-UI messages to PraisonAI format.
|
||||
|
||||
Args:
|
||||
messages: List of AG-UI Message objects
|
||||
|
||||
Returns:
|
||||
List of message dictionaries in PraisonAI format
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
# Track tool_call_ids that have results to avoid duplicates
|
||||
tool_call_ids_with_results: Set[str] = set()
|
||||
for msg in messages:
|
||||
if msg.role == "tool" and msg.tool_call_id:
|
||||
tool_call_ids_with_results.add(msg.tool_call_id)
|
||||
|
||||
result: List[Dict[str, Any]] = []
|
||||
seen_tool_call_ids: Set[str] = set()
|
||||
|
||||
for msg in messages:
|
||||
# Skip system messages - agent builds its own
|
||||
if msg.role == "system":
|
||||
continue
|
||||
|
||||
if msg.role == "tool":
|
||||
# Deduplicate tool results
|
||||
if msg.tool_call_id in seen_tool_call_ids:
|
||||
logger.debug(f"Skipping duplicate tool result: {msg.tool_call_id}")
|
||||
continue
|
||||
if msg.tool_call_id:
|
||||
seen_tool_call_ids.add(msg.tool_call_id)
|
||||
|
||||
result.append({
|
||||
"role": "tool",
|
||||
"content": msg.content or "",
|
||||
"tool_call_id": msg.tool_call_id,
|
||||
})
|
||||
|
||||
elif msg.role == "assistant":
|
||||
msg_dict: Dict[str, Any] = {
|
||||
"role": "assistant",
|
||||
"content": msg.content or "",
|
||||
}
|
||||
|
||||
if msg.tool_calls:
|
||||
# Include all tool_calls (filtering only when there are results to match)
|
||||
tool_calls_list = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
for tc in msg.tool_calls
|
||||
]
|
||||
if tool_calls_list:
|
||||
msg_dict["tool_calls"] = tool_calls_list
|
||||
|
||||
result.append(msg_dict)
|
||||
|
||||
elif msg.role == "user":
|
||||
result.append({
|
||||
"role": "user",
|
||||
"content": msg.content or "",
|
||||
})
|
||||
|
||||
else:
|
||||
logger.warning(f"Unknown message role: {msg.role}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def praisonai_messages_to_agui(messages: List[Dict[str, Any]]) -> List[Message]:
|
||||
"""
|
||||
Convert PraisonAI messages to AG-UI format.
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries in PraisonAI format
|
||||
|
||||
Returns:
|
||||
List of AG-UI Message objects
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
|
||||
result: List[Message] = []
|
||||
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "tool":
|
||||
result.append(Message(
|
||||
role="tool",
|
||||
content=content,
|
||||
tool_call_id=msg.get("tool_call_id"),
|
||||
))
|
||||
|
||||
elif role == "assistant":
|
||||
tool_calls = None
|
||||
if "tool_calls" in msg and msg["tool_calls"]:
|
||||
tool_calls = [
|
||||
ToolCall(
|
||||
id=tc.get("id", ""),
|
||||
function=FunctionCall(
|
||||
name=tc.get("function", {}).get("name", ""),
|
||||
arguments=tc.get("function", {}).get("arguments", "{}"),
|
||||
),
|
||||
)
|
||||
for tc in msg["tool_calls"]
|
||||
]
|
||||
|
||||
result.append(Message(
|
||||
role="assistant",
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
))
|
||||
|
||||
elif role == "user":
|
||||
result.append(Message(
|
||||
role="user",
|
||||
content=content,
|
||||
))
|
||||
|
||||
elif role == "system":
|
||||
result.append(Message(
|
||||
role="system",
|
||||
content=content,
|
||||
))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def validate_state(state: Any, thread_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Validate and convert state to a dictionary.
|
||||
|
||||
Args:
|
||||
state: The state to validate (can be dict, Pydantic model, or other)
|
||||
thread_id: Thread ID for logging
|
||||
|
||||
Returns:
|
||||
State as a dictionary, or None if invalid
|
||||
"""
|
||||
if state is None:
|
||||
return None
|
||||
|
||||
if isinstance(state, dict):
|
||||
return state
|
||||
|
||||
if isinstance(state, BaseModel):
|
||||
try:
|
||||
return state.model_dump()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try to_dict method
|
||||
if hasattr(state, "to_dict") and callable(getattr(state, "to_dict")):
|
||||
try:
|
||||
result = state.to_dict()
|
||||
if isinstance(result, dict):
|
||||
return result
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning(
|
||||
f"State must be a dict, got {type(state).__name__}. "
|
||||
f"State will be ignored. Thread: {thread_id}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def extract_user_input(messages: List[Message]) -> str:
|
||||
"""
|
||||
Extract the last user message content.
|
||||
|
||||
Args:
|
||||
messages: List of AG-UI messages
|
||||
|
||||
Returns:
|
||||
Content of the last user message, or empty string
|
||||
"""
|
||||
if not messages:
|
||||
return ""
|
||||
|
||||
for msg in reversed(messages):
|
||||
if msg.role == "user" and msg.content:
|
||||
return msg.content
|
||||
|
||||
return ""
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
AG-UI Event Encoder
|
||||
|
||||
Encodes AG-UI events to Server-Sent Events (SSE) format.
|
||||
"""
|
||||
|
||||
from praisonaiagents.ui.agui.types import BaseEvent
|
||||
|
||||
|
||||
class EventEncoder:
|
||||
"""Encodes AG-UI events to SSE format."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def get_content_type(self) -> str:
|
||||
"""Returns the content type for SSE."""
|
||||
return "text/event-stream"
|
||||
|
||||
def encode(self, event: BaseEvent) -> str:
|
||||
"""
|
||||
Encode an event to SSE format.
|
||||
|
||||
Args:
|
||||
event: The event to encode
|
||||
|
||||
Returns:
|
||||
SSE formatted string
|
||||
"""
|
||||
return f"data: {event.model_dump_json(by_alias=True, exclude_none=True)}\n\n"
|
||||
@@ -0,0 +1,372 @@
|
||||
"""
|
||||
AG-UI Event Streaming
|
||||
|
||||
Provides utilities for streaming PraisonAI agent responses as AG-UI events.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Set
|
||||
|
||||
from praisonaiagents.ui.agui.types import (
|
||||
BaseEvent,
|
||||
TextMessageStartEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
ToolCallStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
RunStartedEvent,
|
||||
RunFinishedEvent,
|
||||
RunErrorEvent,
|
||||
StepStartedEvent,
|
||||
StepFinishedEvent,
|
||||
StateSnapshotEvent,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventBuffer:
|
||||
"""Buffer to manage event ordering constraints."""
|
||||
|
||||
active_tool_call_ids: Set[str] = field(default_factory=set)
|
||||
ended_tool_call_ids: Set[str] = field(default_factory=set)
|
||||
current_text_message_id: str = ""
|
||||
next_text_message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
pending_tool_calls_parent_id: str = ""
|
||||
|
||||
def start_tool_call(self, tool_call_id: str) -> None:
|
||||
"""Start a new tool call."""
|
||||
self.active_tool_call_ids.add(tool_call_id)
|
||||
|
||||
def end_tool_call(self, tool_call_id: str) -> None:
|
||||
"""End a tool call."""
|
||||
self.active_tool_call_ids.discard(tool_call_id)
|
||||
self.ended_tool_call_ids.add(tool_call_id)
|
||||
|
||||
def start_text_message(self) -> str:
|
||||
"""Start a new text message and return its ID."""
|
||||
self.current_text_message_id = self.next_text_message_id
|
||||
self.next_text_message_id = str(uuid.uuid4())
|
||||
return self.current_text_message_id
|
||||
|
||||
def get_parent_message_id_for_tool_call(self) -> str:
|
||||
"""Get the message ID to use as parent for tool calls."""
|
||||
if self.pending_tool_calls_parent_id:
|
||||
return self.pending_tool_calls_parent_id
|
||||
return self.current_text_message_id
|
||||
|
||||
|
||||
# Event Creation Functions
|
||||
|
||||
def create_text_message_events(
|
||||
content: str,
|
||||
message_id: Optional[str] = None,
|
||||
role: str = "assistant"
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Create text message events for a complete message.
|
||||
|
||||
Args:
|
||||
content: The message content
|
||||
message_id: Optional message ID (generated if not provided)
|
||||
role: Message role (default: assistant)
|
||||
|
||||
Yields:
|
||||
TextMessageStartEvent, TextMessageContentEvent(s), TextMessageEndEvent
|
||||
"""
|
||||
msg_id = message_id or str(uuid.uuid4())
|
||||
|
||||
yield TextMessageStartEvent(message_id=msg_id, role=role)
|
||||
|
||||
if content:
|
||||
yield TextMessageContentEvent(message_id=msg_id, delta=content)
|
||||
|
||||
yield TextMessageEndEvent(message_id=msg_id)
|
||||
|
||||
|
||||
def stream_text_chunks(
|
||||
chunks: Iterator[str],
|
||||
message_id: Optional[str] = None,
|
||||
role: str = "assistant"
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Stream text chunks as AG-UI events.
|
||||
|
||||
Args:
|
||||
chunks: Iterator of text chunks
|
||||
message_id: Optional message ID
|
||||
role: Message role
|
||||
|
||||
Yields:
|
||||
AG-UI events for the text stream
|
||||
"""
|
||||
msg_id = message_id or str(uuid.uuid4())
|
||||
|
||||
yield TextMessageStartEvent(message_id=msg_id, role=role)
|
||||
|
||||
for chunk in chunks:
|
||||
if chunk:
|
||||
yield TextMessageContentEvent(message_id=msg_id, delta=chunk)
|
||||
|
||||
yield TextMessageEndEvent(message_id=msg_id)
|
||||
|
||||
|
||||
def create_tool_call_events(
|
||||
tool_call_id: str,
|
||||
tool_name: str,
|
||||
arguments: str,
|
||||
parent_message_id: Optional[str] = None
|
||||
) -> Iterator[BaseEvent]:
|
||||
"""
|
||||
Create tool call events.
|
||||
|
||||
Args:
|
||||
tool_call_id: Unique ID for the tool call
|
||||
tool_name: Name of the tool being called
|
||||
arguments: JSON string of arguments
|
||||
parent_message_id: Optional parent message ID
|
||||
|
||||
Yields:
|
||||
ToolCallStartEvent, ToolCallArgsEvent, ToolCallEndEvent
|
||||
"""
|
||||
yield ToolCallStartEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
tool_call_name=tool_name,
|
||||
parent_message_id=parent_message_id
|
||||
)
|
||||
|
||||
yield ToolCallArgsEvent(
|
||||
tool_call_id=tool_call_id,
|
||||
delta=arguments
|
||||
)
|
||||
|
||||
yield ToolCallEndEvent(tool_call_id=tool_call_id)
|
||||
|
||||
|
||||
def create_tool_result_event(
|
||||
tool_call_id: str,
|
||||
content: str,
|
||||
message_id: Optional[str] = None
|
||||
) -> ToolCallResultEvent:
|
||||
"""
|
||||
Create a tool result event.
|
||||
|
||||
Args:
|
||||
tool_call_id: ID of the tool call this is a result for
|
||||
content: Result content
|
||||
message_id: Optional message ID
|
||||
|
||||
Returns:
|
||||
ToolCallResultEvent
|
||||
"""
|
||||
return ToolCallResultEvent(
|
||||
message_id=message_id or str(uuid.uuid4()),
|
||||
tool_call_id=tool_call_id,
|
||||
content=content
|
||||
)
|
||||
|
||||
|
||||
# Run Lifecycle Events
|
||||
|
||||
def create_run_started_event(thread_id: str, run_id: str) -> RunStartedEvent:
|
||||
"""Create a run started event."""
|
||||
return RunStartedEvent(thread_id=thread_id, run_id=run_id)
|
||||
|
||||
|
||||
def create_run_finished_event(
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
result: Optional[Any] = None
|
||||
) -> RunFinishedEvent:
|
||||
"""Create a run finished event."""
|
||||
return RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
|
||||
|
||||
|
||||
def create_run_error_event(
|
||||
message: str,
|
||||
code: Optional[str] = None
|
||||
) -> RunErrorEvent:
|
||||
"""Create a run error event."""
|
||||
return RunErrorEvent(message=message, code=code)
|
||||
|
||||
|
||||
# Step Events
|
||||
|
||||
def create_step_started_event(step_name: str) -> StepStartedEvent:
|
||||
"""Create a step started event."""
|
||||
return StepStartedEvent(step_name=step_name)
|
||||
|
||||
|
||||
def create_step_finished_event(step_name: str) -> StepFinishedEvent:
|
||||
"""Create a step finished event."""
|
||||
return StepFinishedEvent(step_name=step_name)
|
||||
|
||||
|
||||
# State Events
|
||||
|
||||
def create_state_snapshot_event(state: Dict[str, Any]) -> StateSnapshotEvent:
|
||||
"""Create a state snapshot event."""
|
||||
return StateSnapshotEvent(snapshot=state)
|
||||
|
||||
|
||||
# Async Streaming
|
||||
|
||||
async def async_stream_response(
|
||||
response_stream: AsyncIterator[str],
|
||||
thread_id: str,
|
||||
run_id: str
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream an async response as AG-UI events.
|
||||
|
||||
Args:
|
||||
response_stream: Async iterator of response chunks
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
message_started = False
|
||||
|
||||
try:
|
||||
async for chunk in response_stream:
|
||||
if not message_started:
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
message_started = True
|
||||
|
||||
if chunk:
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=chunk)
|
||||
|
||||
if message_started:
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
if message_started:
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
yield create_run_error_event(str(e))
|
||||
|
||||
|
||||
async def async_stream_agent_response(
|
||||
agent,
|
||||
user_input: str,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
session_state: Optional[Dict[str, Any]] = None,
|
||||
messages: Optional[List[Dict[str, Any]]] = None
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream an agent's response as AG-UI events.
|
||||
|
||||
Args:
|
||||
agent: PraisonAI Agent instance
|
||||
user_input: User input string
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
session_state: Optional session state
|
||||
messages: Optional message history
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
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:
|
||||
raise ValueError("Agent must have 'chat' or 'achat' method")
|
||||
|
||||
# Emit text message events
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
|
||||
if response:
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=str(response))
|
||||
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
yield create_run_error_event(str(e))
|
||||
|
||||
|
||||
async def async_stream_agents_response(
|
||||
agents,
|
||||
user_input: str,
|
||||
thread_id: str,
|
||||
run_id: str,
|
||||
session_state: Optional[Dict[str, Any]] = None
|
||||
) -> AsyncIterator[BaseEvent]:
|
||||
"""
|
||||
Stream a PraisonAIAgents workflow response as AG-UI events.
|
||||
|
||||
Args:
|
||||
agents: PraisonAIAgents instance
|
||||
user_input: User input string
|
||||
thread_id: Thread ID
|
||||
run_id: Run ID
|
||||
session_state: Optional session state
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# Emit run started
|
||||
yield create_run_started_event(thread_id, run_id)
|
||||
|
||||
message_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Check if agents has async start method
|
||||
if hasattr(agents, 'astart'):
|
||||
result = await agents.astart(user_input)
|
||||
elif hasattr(agents, 'start'):
|
||||
# Run sync start in executor
|
||||
loop = asyncio.get_event_loop()
|
||||
result = await loop.run_in_executor(None, agents.start, user_input)
|
||||
else:
|
||||
raise ValueError("Agents must have 'start' or 'astart' method")
|
||||
|
||||
# Emit text message events with result
|
||||
yield TextMessageStartEvent(message_id=message_id, role="assistant")
|
||||
|
||||
if result:
|
||||
# Handle different result types
|
||||
if isinstance(result, dict):
|
||||
content = result.get("output", str(result))
|
||||
elif hasattr(result, "raw"):
|
||||
content = result.raw
|
||||
else:
|
||||
content = str(result)
|
||||
|
||||
yield TextMessageContentEvent(message_id=message_id, delta=content)
|
||||
|
||||
yield TextMessageEndEvent(message_id=message_id)
|
||||
|
||||
# Emit run finished
|
||||
yield create_run_finished_event(thread_id, run_id)
|
||||
|
||||
except Exception as e:
|
||||
yield create_run_error_event(str(e))
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
AG-UI Protocol Types
|
||||
|
||||
This module defines the event types and message types for the AG-UI protocol,
|
||||
following the ag-ui-protocol specification.
|
||||
"""
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class EventType(str, Enum):
|
||||
"""AG-UI Event Types."""
|
||||
TEXT_MESSAGE_START = "TEXT_MESSAGE_START"
|
||||
TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT"
|
||||
TEXT_MESSAGE_END = "TEXT_MESSAGE_END"
|
||||
TOOL_CALL_START = "TOOL_CALL_START"
|
||||
TOOL_CALL_ARGS = "TOOL_CALL_ARGS"
|
||||
TOOL_CALL_END = "TOOL_CALL_END"
|
||||
TOOL_CALL_RESULT = "TOOL_CALL_RESULT"
|
||||
STATE_SNAPSHOT = "STATE_SNAPSHOT"
|
||||
STATE_DELTA = "STATE_DELTA"
|
||||
RUN_STARTED = "RUN_STARTED"
|
||||
RUN_FINISHED = "RUN_FINISHED"
|
||||
RUN_ERROR = "RUN_ERROR"
|
||||
STEP_STARTED = "STEP_STARTED"
|
||||
STEP_FINISHED = "STEP_FINISHED"
|
||||
CUSTOM = "CUSTOM"
|
||||
|
||||
|
||||
class BaseEvent(BaseModel):
|
||||
"""Base event for all AG-UI events."""
|
||||
type: EventType
|
||||
timestamp: Optional[int] = None
|
||||
|
||||
model_config = {"use_enum_values": True}
|
||||
|
||||
|
||||
# Text Message Events
|
||||
class TextMessageStartEvent(BaseEvent):
|
||||
"""Event indicating the start of a text message."""
|
||||
type: Literal["TEXT_MESSAGE_START"] = "TEXT_MESSAGE_START"
|
||||
message_id: str
|
||||
role: str = "assistant"
|
||||
|
||||
|
||||
class TextMessageContentEvent(BaseEvent):
|
||||
"""Event containing text message content."""
|
||||
type: Literal["TEXT_MESSAGE_CONTENT"] = "TEXT_MESSAGE_CONTENT"
|
||||
message_id: str
|
||||
delta: str
|
||||
|
||||
|
||||
class TextMessageEndEvent(BaseEvent):
|
||||
"""Event indicating the end of a text message."""
|
||||
type: Literal["TEXT_MESSAGE_END"] = "TEXT_MESSAGE_END"
|
||||
message_id: str
|
||||
|
||||
|
||||
# Tool Call Events
|
||||
class ToolCallStartEvent(BaseEvent):
|
||||
"""Event indicating the start of a tool call."""
|
||||
type: Literal["TOOL_CALL_START"] = "TOOL_CALL_START"
|
||||
tool_call_id: str
|
||||
tool_call_name: str
|
||||
parent_message_id: Optional[str] = None
|
||||
|
||||
|
||||
class ToolCallArgsEvent(BaseEvent):
|
||||
"""Event containing tool call arguments."""
|
||||
type: Literal["TOOL_CALL_ARGS"] = "TOOL_CALL_ARGS"
|
||||
tool_call_id: str
|
||||
delta: str
|
||||
|
||||
|
||||
class ToolCallEndEvent(BaseEvent):
|
||||
"""Event indicating the end of a tool call."""
|
||||
type: Literal["TOOL_CALL_END"] = "TOOL_CALL_END"
|
||||
tool_call_id: str
|
||||
|
||||
|
||||
class ToolCallResultEvent(BaseEvent):
|
||||
"""Event containing tool call result."""
|
||||
type: Literal["TOOL_CALL_RESULT"] = "TOOL_CALL_RESULT"
|
||||
message_id: str
|
||||
tool_call_id: str
|
||||
content: str
|
||||
role: Optional[Literal["tool"]] = "tool"
|
||||
|
||||
|
||||
# Run Lifecycle Events
|
||||
class RunStartedEvent(BaseEvent):
|
||||
"""Event indicating a run has started."""
|
||||
type: Literal["RUN_STARTED"] = "RUN_STARTED"
|
||||
thread_id: str
|
||||
run_id: str
|
||||
parent_run_id: Optional[str] = None
|
||||
|
||||
|
||||
class RunFinishedEvent(BaseEvent):
|
||||
"""Event indicating a run has finished."""
|
||||
type: Literal["RUN_FINISHED"] = "RUN_FINISHED"
|
||||
thread_id: str
|
||||
run_id: str
|
||||
result: Optional[Any] = None
|
||||
|
||||
|
||||
class RunErrorEvent(BaseEvent):
|
||||
"""Event indicating a run error."""
|
||||
type: Literal["RUN_ERROR"] = "RUN_ERROR"
|
||||
message: str
|
||||
code: Optional[str] = None
|
||||
|
||||
|
||||
# Step Events
|
||||
class StepStartedEvent(BaseEvent):
|
||||
"""Event indicating a step has started."""
|
||||
type: Literal["STEP_STARTED"] = "STEP_STARTED"
|
||||
step_name: str
|
||||
|
||||
|
||||
class StepFinishedEvent(BaseEvent):
|
||||
"""Event indicating a step has finished."""
|
||||
type: Literal["STEP_FINISHED"] = "STEP_FINISHED"
|
||||
step_name: str
|
||||
|
||||
|
||||
# State Events
|
||||
class StateSnapshotEvent(BaseEvent):
|
||||
"""Event containing a state snapshot."""
|
||||
type: Literal["STATE_SNAPSHOT"] = "STATE_SNAPSHOT"
|
||||
snapshot: Dict[str, Any]
|
||||
|
||||
|
||||
# Message Types
|
||||
class FunctionCall(BaseModel):
|
||||
"""Function call in a tool call."""
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class ToolCall(BaseModel):
|
||||
"""Tool call in a message."""
|
||||
id: str
|
||||
function: FunctionCall
|
||||
type: str = "function"
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
"""AG-UI Message type."""
|
||||
role: Literal["user", "assistant", "system", "tool"]
|
||||
content: Optional[str] = None
|
||||
tool_calls: Optional[List[ToolCall]] = None
|
||||
tool_call_id: Optional[str] = None
|
||||
name: Optional[str] = None
|
||||
|
||||
|
||||
class RunAgentInput(BaseModel):
|
||||
"""Input for running an agent."""
|
||||
thread_id: str
|
||||
run_id: Optional[str] = None
|
||||
messages: Optional[List[Message]] = Field(default_factory=list)
|
||||
state: Optional[Dict[str, Any]] = None
|
||||
forwarded_props: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
# Union type for all events
|
||||
Event = Union[
|
||||
TextMessageStartEvent,
|
||||
TextMessageContentEvent,
|
||||
TextMessageEndEvent,
|
||||
ToolCallStartEvent,
|
||||
ToolCallArgsEvent,
|
||||
ToolCallEndEvent,
|
||||
ToolCallResultEvent,
|
||||
StateSnapshotEvent,
|
||||
RunStartedEvent,
|
||||
RunFinishedEvent,
|
||||
RunErrorEvent,
|
||||
StepStartedEvent,
|
||||
StepFinishedEvent,
|
||||
]
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.14"
|
||||
version = "0.1.17"
|
||||
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""A2A Protocol Tests"""
|
||||
@@ -0,0 +1,135 @@
|
||||
"""
|
||||
Tests for Main A2A Class
|
||||
|
||||
TDD: Write tests first, then implement A2A class.
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AClass:
|
||||
"""Tests for A2A main class."""
|
||||
|
||||
def test_a2a_creation_with_agent(self):
|
||||
"""Test creating A2A with single agent."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
assert a2a.agent is not None
|
||||
assert a2a.name == "Test"
|
||||
|
||||
def test_a2a_creation_with_custom_name(self):
|
||||
"""Test creating A2A with custom name."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, name="Custom Name")
|
||||
|
||||
assert a2a.name == "Custom Name"
|
||||
|
||||
def test_a2a_requires_agent(self):
|
||||
"""Test A2A requires agent or agents."""
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
A2A()
|
||||
|
||||
def test_a2a_get_agent_card(self):
|
||||
"""Test getting Agent Card from A2A."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Card Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
card = a2a.get_agent_card()
|
||||
|
||||
assert card.name == "Card Test"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
|
||||
def test_a2a_get_router(self):
|
||||
"""Test getting FastAPI router from A2A."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Router Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
router = a2a.get_router()
|
||||
|
||||
assert router is not None
|
||||
# Check router has routes
|
||||
assert len(router.routes) > 0
|
||||
|
||||
|
||||
class TestA2ARouter:
|
||||
"""Tests for A2A FastAPI router endpoints."""
|
||||
|
||||
def test_agent_card_endpoint(self):
|
||||
"""Test /.well-known/agent.json endpoint."""
|
||||
from fastapi.testclient import TestClient
|
||||
from fastapi import FastAPI
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Endpoint Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/.well-known/agent.json")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "Endpoint Test"
|
||||
|
||||
def test_status_endpoint(self):
|
||||
"""Test /status endpoint."""
|
||||
from fastapi.testclient import TestClient
|
||||
from fastapi import FastAPI
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
agent = Agent(name="Status Test", role="Tester", goal="Test")
|
||||
a2a = A2A(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "ok"
|
||||
|
||||
|
||||
class TestA2AWithTools:
|
||||
"""Tests for A2A with agent tools."""
|
||||
|
||||
def test_a2a_agent_card_includes_skills(self):
|
||||
"""Test Agent Card includes skills from tools."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a import A2A
|
||||
|
||||
def my_tool(x: str) -> str:
|
||||
"""A test tool."""
|
||||
return x
|
||||
|
||||
agent = Agent(
|
||||
name="Tool Agent",
|
||||
role="Helper",
|
||||
goal="Help",
|
||||
tools=[my_tool]
|
||||
)
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
card = a2a.get_agent_card()
|
||||
|
||||
assert card.skills is not None
|
||||
assert len(card.skills) >= 1
|
||||
@@ -0,0 +1,200 @@
|
||||
"""
|
||||
Tests for A2A Agent Card Generation
|
||||
|
||||
TDD: Write tests first, then implement agent_card module.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class TestAgentCardGeneration:
|
||||
"""Tests for generating Agent Card from PraisonAI Agent."""
|
||||
|
||||
def test_generate_agent_card_basic(self):
|
||||
"""Test generating Agent Card from basic Agent."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Test Agent",
|
||||
role="Test Role",
|
||||
goal="Test Goal"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.name == "Test Agent"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
assert card.version == "1.0.0"
|
||||
|
||||
def test_generate_agent_card_with_description(self):
|
||||
"""Test Agent Card includes description from role."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Research Agent",
|
||||
role="Research Analyst",
|
||||
goal="Research topics thoroughly"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.description is not None
|
||||
assert "Research" in card.description
|
||||
|
||||
def test_generate_agent_card_with_instructions(self):
|
||||
"""Test Agent Card from Agent with instructions."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
instructions="You are a helpful assistant that answers questions."
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.name is not None
|
||||
assert card.description is not None
|
||||
|
||||
def test_generate_agent_card_capabilities(self):
|
||||
"""Test Agent Card has correct capabilities."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
|
||||
card = generate_agent_card(
|
||||
agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
streaming=True
|
||||
)
|
||||
|
||||
assert card.capabilities.streaming is True
|
||||
|
||||
|
||||
class TestSkillsExtraction:
|
||||
"""Tests for extracting skills from Agent tools."""
|
||||
|
||||
def test_extract_skills_from_function(self):
|
||||
"""Test extracting skill from a function tool."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Search Agent",
|
||||
role="Searcher",
|
||||
goal="Search the web",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert card.skills is not None
|
||||
assert len(card.skills) >= 1
|
||||
|
||||
skill_names = [s.name for s in card.skills]
|
||||
assert "search_web" in skill_names
|
||||
|
||||
def test_extract_skills_with_description(self):
|
||||
"""Test skill description comes from function docstring."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def calculate(expression: str) -> str:
|
||||
"""Calculate a mathematical expression."""
|
||||
return str(eval(expression))
|
||||
|
||||
agent = Agent(
|
||||
name="Calculator",
|
||||
role="Math Helper",
|
||||
goal="Do math",
|
||||
tools=[calculate]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
calc_skill = next((s for s in card.skills if s.name == "calculate"), None)
|
||||
assert calc_skill is not None
|
||||
assert "mathematical" in calc_skill.description.lower()
|
||||
|
||||
def test_extract_skills_multiple_tools(self):
|
||||
"""Test extracting skills from multiple tools."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
def tool_one(x: str) -> str:
|
||||
"""First tool."""
|
||||
return x
|
||||
|
||||
def tool_two(y: str) -> str:
|
||||
"""Second tool."""
|
||||
return y
|
||||
|
||||
agent = Agent(
|
||||
name="Multi Tool Agent",
|
||||
role="Helper",
|
||||
goal="Help",
|
||||
tools=[tool_one, tool_two]
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
assert len(card.skills) >= 2
|
||||
|
||||
def test_no_skills_when_no_tools(self):
|
||||
"""Test Agent Card with no tools has empty skills."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(
|
||||
name="Simple Agent",
|
||||
role="Helper",
|
||||
goal="Help"
|
||||
)
|
||||
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
# Skills should be None or empty list
|
||||
assert card.skills is None or len(card.skills) == 0
|
||||
|
||||
|
||||
class TestAgentCardSerialization:
|
||||
"""Tests for Agent Card JSON serialization."""
|
||||
|
||||
def test_agent_card_to_json(self):
|
||||
"""Test Agent Card can be serialized to JSON."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
card = generate_agent_card(agent, url="http://localhost:8000/a2a")
|
||||
|
||||
json_str = card.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
|
||||
assert "name" in data
|
||||
assert "url" in data
|
||||
assert "capabilities" in data
|
||||
|
||||
def test_agent_card_camel_case_keys(self):
|
||||
"""Test Agent Card JSON uses camelCase keys."""
|
||||
from praisonaiagents import Agent
|
||||
from praisonaiagents.ui.a2a.agent_card import generate_agent_card
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
card = generate_agent_card(
|
||||
agent,
|
||||
url="http://localhost:8000/a2a",
|
||||
streaming=True
|
||||
)
|
||||
|
||||
# Use by_alias to get camelCase
|
||||
data = card.model_dump(by_alias=True)
|
||||
|
||||
# Check capabilities uses camelCase
|
||||
assert "pushNotifications" in data["capabilities"]
|
||||
@@ -0,0 +1,190 @@
|
||||
"""
|
||||
Tests for A2A Message Conversion
|
||||
|
||||
TDD: Write tests first, then implement conversion module.
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AToPraisonAI:
|
||||
"""Tests for converting A2A messages to PraisonAI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting A2A user message to PraisonAI format."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello, how are you?")]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[0]["content"] == "Hello, how are you?"
|
||||
|
||||
def test_convert_agent_message(self):
|
||||
"""Test converting A2A agent message to PraisonAI format."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="I'm doing well, thanks!")]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert result[0]["content"] == "I'm doing well, thanks!"
|
||||
|
||||
def test_convert_multiple_parts(self):
|
||||
"""Test converting message with multiple text parts."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
a2a_msg = Message(
|
||||
message_id="msg-3",
|
||||
role=Role.USER,
|
||||
parts=[
|
||||
TextPart(text="First part."),
|
||||
TextPart(text="Second part.")
|
||||
]
|
||||
)
|
||||
|
||||
result = a2a_to_praisonai_messages([a2a_msg])
|
||||
|
||||
assert "First part" in result[0]["content"]
|
||||
assert "Second part" in result[0]["content"]
|
||||
|
||||
def test_convert_conversation_history(self):
|
||||
"""Test converting full conversation history."""
|
||||
from praisonaiagents.ui.a2a.conversion import a2a_to_praisonai_messages
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.USER, parts=[TextPart(text="Hi")]),
|
||||
Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="Hello!")]),
|
||||
Message(message_id="m3", role=Role.USER, parts=[TextPart(text="How are you?")]),
|
||||
]
|
||||
|
||||
result = a2a_to_praisonai_messages(messages)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[2]["role"] == "user"
|
||||
|
||||
|
||||
class TestPraisonAIToA2A:
|
||||
"""Tests for converting PraisonAI responses to A2A format."""
|
||||
|
||||
def test_convert_string_response(self):
|
||||
"""Test converting string response to A2A Message."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "This is the agent's response."
|
||||
|
||||
result = praisonai_to_a2a_message(response)
|
||||
|
||||
assert result.role.value == "agent"
|
||||
assert len(result.parts) == 1
|
||||
assert result.parts[0].text == response
|
||||
|
||||
def test_convert_response_with_context(self):
|
||||
"""Test converting response with context_id."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "Response text"
|
||||
|
||||
result = praisonai_to_a2a_message(response, context_id="ctx-123")
|
||||
|
||||
assert result.context_id == "ctx-123"
|
||||
|
||||
def test_convert_response_with_task(self):
|
||||
"""Test converting response with task_id."""
|
||||
from praisonaiagents.ui.a2a.conversion import praisonai_to_a2a_message
|
||||
|
||||
response = "Response text"
|
||||
|
||||
result = praisonai_to_a2a_message(response, task_id="task-456")
|
||||
|
||||
assert result.task_id == "task-456"
|
||||
|
||||
|
||||
class TestArtifactConversion:
|
||||
"""Tests for converting responses to A2A Artifacts."""
|
||||
|
||||
def test_create_artifact_from_response(self):
|
||||
"""Test creating Artifact from response."""
|
||||
from praisonaiagents.ui.a2a.conversion import create_artifact
|
||||
|
||||
response = "Generated content here."
|
||||
|
||||
artifact = create_artifact(response)
|
||||
|
||||
assert artifact.artifact_id is not None
|
||||
assert len(artifact.parts) == 1
|
||||
assert artifact.parts[0].text == response
|
||||
|
||||
def test_create_artifact_with_name(self):
|
||||
"""Test creating Artifact with name."""
|
||||
from praisonaiagents.ui.a2a.conversion import create_artifact
|
||||
|
||||
response = "Report content"
|
||||
|
||||
artifact = create_artifact(response, name="Report", description="Generated report")
|
||||
|
||||
assert artifact.name == "Report"
|
||||
assert artifact.description == "Generated report"
|
||||
|
||||
|
||||
class TestExtractUserInput:
|
||||
"""Tests for extracting user input from A2A messages."""
|
||||
|
||||
def test_extract_from_single_message(self):
|
||||
"""Test extracting user input from single message."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
msg = Message(
|
||||
message_id="m1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="What is AI?")]
|
||||
)
|
||||
|
||||
result = extract_user_input([msg])
|
||||
|
||||
assert result == "What is AI?"
|
||||
|
||||
def test_extract_from_conversation(self):
|
||||
"""Test extracting last user input from conversation."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.USER, parts=[TextPart(text="First question")]),
|
||||
Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="Answer")]),
|
||||
Message(message_id="m3", role=Role.USER, parts=[TextPart(text="Follow up question")]),
|
||||
]
|
||||
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == "Follow up question"
|
||||
|
||||
def test_extract_empty_when_no_user_message(self):
|
||||
"""Test extracting returns empty when no user message."""
|
||||
from praisonaiagents.ui.a2a.conversion import extract_user_input
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
messages = [
|
||||
Message(message_id="m1", role=Role.AGENT, parts=[TextPart(text="Hello")]),
|
||||
]
|
||||
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == ""
|
||||
@@ -0,0 +1,116 @@
|
||||
"""
|
||||
Tests for A2A Streaming Events
|
||||
"""
|
||||
|
||||
|
||||
class TestA2AEventEncoder:
|
||||
"""Tests for A2AEventEncoder."""
|
||||
|
||||
def test_encode_event(self):
|
||||
"""Test encoding a generic event."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
|
||||
result = A2AEventEncoder.encode_event("test", {"key": "value"})
|
||||
|
||||
assert "event: test" in result
|
||||
assert "data:" in result
|
||||
assert "key" in result
|
||||
|
||||
def test_encode_task_status(self):
|
||||
"""Test encoding TaskStatusUpdateEvent."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
from praisonaiagents.ui.a2a.types import TaskStatusUpdateEvent, TaskStatus, TaskState
|
||||
|
||||
event = TaskStatusUpdateEvent(
|
||||
task_id="task-123",
|
||||
status=TaskStatus(state=TaskState.WORKING),
|
||||
final=False,
|
||||
)
|
||||
|
||||
result = A2AEventEncoder.encode_task_status(event)
|
||||
|
||||
assert "event: task.status" in result
|
||||
assert "task-123" in result
|
||||
|
||||
def test_encode_task_artifact(self):
|
||||
"""Test encoding TaskArtifactUpdateEvent."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
from praisonaiagents.ui.a2a.types import TaskArtifactUpdateEvent, Artifact, TextPart
|
||||
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
event = TaskArtifactUpdateEvent(
|
||||
task_id="task-456",
|
||||
artifact=artifact,
|
||||
last_chunk=True,
|
||||
)
|
||||
|
||||
result = A2AEventEncoder.encode_task_artifact(event)
|
||||
|
||||
assert "event: task.artifact" in result
|
||||
assert "art-1" in result
|
||||
|
||||
def test_encode_done(self):
|
||||
"""Test encoding done event."""
|
||||
from praisonaiagents.ui.a2a.streaming import A2AEventEncoder
|
||||
|
||||
result = A2AEventEncoder.encode_done()
|
||||
|
||||
assert "event: done" in result
|
||||
|
||||
|
||||
class TestCreateEvents:
|
||||
"""Tests for event creation helpers."""
|
||||
|
||||
def test_create_status_event(self):
|
||||
"""Test creating status event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_status_event
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
|
||||
event = create_status_event(
|
||||
task_id="task-1",
|
||||
state=TaskState.WORKING,
|
||||
)
|
||||
|
||||
assert event.task_id == "task-1"
|
||||
assert event.status.state == TaskState.WORKING
|
||||
|
||||
def test_create_status_event_final(self):
|
||||
"""Test creating final status event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_status_event
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
|
||||
event = create_status_event(
|
||||
task_id="task-2",
|
||||
state=TaskState.COMPLETED,
|
||||
final=True,
|
||||
)
|
||||
|
||||
assert event.final is True
|
||||
|
||||
def test_create_artifact_event(self):
|
||||
"""Test creating artifact event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_artifact_event
|
||||
|
||||
event = create_artifact_event(
|
||||
task_id="task-3",
|
||||
content="Generated content",
|
||||
)
|
||||
|
||||
assert event.task_id == "task-3"
|
||||
assert event.artifact is not None
|
||||
assert len(event.artifact.parts) == 1
|
||||
|
||||
def test_create_artifact_event_last_chunk(self):
|
||||
"""Test creating last chunk artifact event."""
|
||||
from praisonaiagents.ui.a2a.streaming import create_artifact_event
|
||||
|
||||
event = create_artifact_event(
|
||||
task_id="task-4",
|
||||
content="Final content",
|
||||
last_chunk=True,
|
||||
)
|
||||
|
||||
assert event.last_chunk is True
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
Tests for A2A Task Store
|
||||
|
||||
TDD: Write tests first, then implement task_store module.
|
||||
"""
|
||||
|
||||
|
||||
class TestTaskStore:
|
||||
"""Tests for TaskStore class."""
|
||||
|
||||
def test_create_task(self):
|
||||
"""Test creating a new task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
|
||||
assert task.id is not None
|
||||
assert task.status.state == TaskState.SUBMITTED
|
||||
|
||||
def test_create_task_with_context(self):
|
||||
"""Test creating task with context_id."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")],
|
||||
context_id="ctx-123"
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
|
||||
assert task.context_id == "ctx-123"
|
||||
|
||||
def test_get_task(self):
|
||||
"""Test retrieving a task by ID."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-3",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
created = store.create_task(msg)
|
||||
retrieved = store.get_task(created.id)
|
||||
|
||||
assert retrieved is not None
|
||||
assert retrieved.id == created.id
|
||||
|
||||
def test_get_task_not_found(self):
|
||||
"""Test getting non-existent task returns None."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
|
||||
store = TaskStore()
|
||||
result = store.get_task("non-existent-id")
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_update_status(self):
|
||||
"""Test updating task status."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-4",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
updated = store.update_status(task.id, TaskState.WORKING)
|
||||
|
||||
assert updated.status.state == TaskState.WORKING
|
||||
|
||||
def test_update_status_to_completed(self):
|
||||
"""Test updating task to completed state."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-5",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
store.update_status(task.id, TaskState.WORKING)
|
||||
updated = store.update_status(task.id, TaskState.COMPLETED)
|
||||
|
||||
assert updated.status.state == TaskState.COMPLETED
|
||||
|
||||
def test_add_artifact(self):
|
||||
"""Test adding artifact to task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, Artifact
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-6",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
|
||||
updated = store.add_artifact(task.id, artifact)
|
||||
|
||||
assert updated.artifacts is not None
|
||||
assert len(updated.artifacts) == 1
|
||||
|
||||
def test_add_message_to_history(self):
|
||||
"""Test adding message to task history."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
msg1 = Message(
|
||||
message_id="msg-7",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Question")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg1)
|
||||
|
||||
msg2 = Message(
|
||||
message_id="msg-8",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Answer")]
|
||||
)
|
||||
|
||||
updated = store.add_to_history(task.id, msg2)
|
||||
|
||||
assert updated.history is not None
|
||||
assert len(updated.history) >= 1
|
||||
|
||||
def test_list_tasks(self):
|
||||
"""Test listing all tasks."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
|
||||
for i in range(3):
|
||||
msg = Message(
|
||||
message_id=f"msg-{i}",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text=f"Message {i}")]
|
||||
)
|
||||
store.create_task(msg)
|
||||
|
||||
tasks = store.list_tasks()
|
||||
|
||||
assert len(tasks) == 3
|
||||
|
||||
def test_cancel_task(self):
|
||||
"""Test cancelling a task."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart, TaskState
|
||||
|
||||
store = TaskStore()
|
||||
msg = Message(
|
||||
message_id="msg-cancel",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
|
||||
task = store.create_task(msg)
|
||||
cancelled = store.cancel_task(task.id)
|
||||
|
||||
assert cancelled.status.state == TaskState.CANCELLED
|
||||
|
||||
|
||||
class TestTaskStoreContextManagement:
|
||||
"""Tests for context management in TaskStore."""
|
||||
|
||||
def test_tasks_grouped_by_context(self):
|
||||
"""Test getting tasks by context_id."""
|
||||
from praisonaiagents.ui.a2a.task_store import TaskStore
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
|
||||
store = TaskStore()
|
||||
|
||||
# Create tasks in context A
|
||||
for i in range(2):
|
||||
msg = Message(
|
||||
message_id=f"msg-a-{i}",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text=f"A{i}")],
|
||||
context_id="context-a"
|
||||
)
|
||||
store.create_task(msg)
|
||||
|
||||
# Create task in context B
|
||||
msg_b = Message(
|
||||
message_id="msg-b-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="B1")],
|
||||
context_id="context-b"
|
||||
)
|
||||
store.create_task(msg_b)
|
||||
|
||||
tasks_a = store.list_tasks(context_id="context-a")
|
||||
tasks_b = store.list_tasks(context_id="context-b")
|
||||
|
||||
assert len(tasks_a) == 2
|
||||
assert len(tasks_b) == 1
|
||||
@@ -0,0 +1,395 @@
|
||||
"""
|
||||
Tests for A2A Protocol Types
|
||||
|
||||
TDD: Write tests first, then implement types.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class TestTaskState:
|
||||
"""Tests for TaskState enum."""
|
||||
|
||||
def test_task_state_submitted(self):
|
||||
"""Test SUBMITTED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.SUBMITTED == "submitted"
|
||||
|
||||
def test_task_state_working(self):
|
||||
"""Test WORKING state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.WORKING == "working"
|
||||
|
||||
def test_task_state_completed(self):
|
||||
"""Test COMPLETED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.COMPLETED == "completed"
|
||||
|
||||
def test_task_state_failed(self):
|
||||
"""Test FAILED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.FAILED == "failed"
|
||||
|
||||
def test_task_state_cancelled(self):
|
||||
"""Test CANCELLED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.CANCELLED == "cancelled"
|
||||
|
||||
def test_task_state_input_required(self):
|
||||
"""Test INPUT_REQUIRED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.INPUT_REQUIRED == "input_required"
|
||||
|
||||
def test_task_state_auth_required(self):
|
||||
"""Test AUTH_REQUIRED state exists."""
|
||||
from praisonaiagents.ui.a2a.types import TaskState
|
||||
assert TaskState.AUTH_REQUIRED == "auth_required"
|
||||
|
||||
|
||||
class TestRole:
|
||||
"""Tests for Role enum."""
|
||||
|
||||
def test_role_user(self):
|
||||
"""Test USER role exists."""
|
||||
from praisonaiagents.ui.a2a.types import Role
|
||||
assert Role.USER == "user"
|
||||
|
||||
def test_role_agent(self):
|
||||
"""Test AGENT role exists."""
|
||||
from praisonaiagents.ui.a2a.types import Role
|
||||
assert Role.AGENT == "agent"
|
||||
|
||||
|
||||
class TestPart:
|
||||
"""Tests for Part models."""
|
||||
|
||||
def test_text_part_creation(self):
|
||||
"""Test TextPart can be created."""
|
||||
from praisonaiagents.ui.a2a.types import TextPart
|
||||
part = TextPart(text="Hello, world!")
|
||||
assert part.text == "Hello, world!"
|
||||
|
||||
def test_file_part_with_uri(self):
|
||||
"""Test FilePart with URI."""
|
||||
from praisonaiagents.ui.a2a.types import FilePart
|
||||
part = FilePart(
|
||||
file_uri="https://example.com/file.pdf",
|
||||
media_type="application/pdf",
|
||||
name="document.pdf"
|
||||
)
|
||||
assert part.file_uri == "https://example.com/file.pdf"
|
||||
assert part.media_type == "application/pdf"
|
||||
|
||||
def test_file_part_with_bytes(self):
|
||||
"""Test FilePart with bytes."""
|
||||
from praisonaiagents.ui.a2a.types import FilePart
|
||||
part = FilePart(
|
||||
file_bytes=b"file content",
|
||||
media_type="text/plain",
|
||||
name="test.txt"
|
||||
)
|
||||
assert part.file_bytes == b"file content"
|
||||
|
||||
def test_data_part_creation(self):
|
||||
"""Test DataPart can be created."""
|
||||
from praisonaiagents.ui.a2a.types import DataPart
|
||||
part = DataPart(data={"key": "value", "count": 42})
|
||||
assert part.data == {"key": "value", "count": 42}
|
||||
|
||||
|
||||
class TestMessage:
|
||||
"""Tests for Message model."""
|
||||
|
||||
def test_message_creation(self):
|
||||
"""Test Message can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-123",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
assert msg.message_id == "msg-123"
|
||||
assert msg.role == Role.USER
|
||||
assert len(msg.parts) == 1
|
||||
|
||||
def test_message_with_context(self):
|
||||
"""Test Message with context_id."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-456",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Response")],
|
||||
context_id="ctx-789"
|
||||
)
|
||||
assert msg.context_id == "ctx-789"
|
||||
|
||||
def test_message_with_task(self):
|
||||
"""Test Message with task_id."""
|
||||
from praisonaiagents.ui.a2a.types import Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-789",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Working on it")],
|
||||
task_id="task-123"
|
||||
)
|
||||
assert msg.task_id == "task-123"
|
||||
|
||||
|
||||
class TestTaskStatus:
|
||||
"""Tests for TaskStatus model."""
|
||||
|
||||
def test_task_status_creation(self):
|
||||
"""Test TaskStatus can be created."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState
|
||||
status = TaskStatus(state=TaskState.WORKING)
|
||||
assert status.state == TaskState.WORKING
|
||||
|
||||
def test_task_status_with_message(self):
|
||||
"""Test TaskStatus with message."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState, Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.AGENT,
|
||||
parts=[TextPart(text="Processing...")]
|
||||
)
|
||||
status = TaskStatus(state=TaskState.WORKING, message=msg)
|
||||
assert status.message is not None
|
||||
|
||||
def test_task_status_with_timestamp(self):
|
||||
"""Test TaskStatus with timestamp."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatus, TaskState
|
||||
ts = datetime.utcnow().isoformat()
|
||||
status = TaskStatus(state=TaskState.COMPLETED, timestamp=ts)
|
||||
assert status.timestamp == ts
|
||||
|
||||
|
||||
class TestTask:
|
||||
"""Tests for Task model."""
|
||||
|
||||
def test_task_creation(self):
|
||||
"""Test Task can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState
|
||||
task = Task(
|
||||
id="task-123",
|
||||
status=TaskStatus(state=TaskState.SUBMITTED)
|
||||
)
|
||||
assert task.id == "task-123"
|
||||
assert task.status.state == TaskState.SUBMITTED
|
||||
|
||||
def test_task_with_context(self):
|
||||
"""Test Task with context_id."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState
|
||||
task = Task(
|
||||
id="task-456",
|
||||
context_id="ctx-789",
|
||||
status=TaskStatus(state=TaskState.WORKING)
|
||||
)
|
||||
assert task.context_id == "ctx-789"
|
||||
|
||||
def test_task_with_artifacts(self):
|
||||
"""Test Task with artifacts."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState, Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Result")]
|
||||
)
|
||||
task = Task(
|
||||
id="task-789",
|
||||
status=TaskStatus(state=TaskState.COMPLETED),
|
||||
artifacts=[artifact]
|
||||
)
|
||||
assert len(task.artifacts) == 1
|
||||
|
||||
def test_task_with_history(self):
|
||||
"""Test Task with message history."""
|
||||
from praisonaiagents.ui.a2a.types import Task, TaskStatus, TaskState, Message, Role, TextPart
|
||||
msg1 = Message(message_id="m1", role=Role.USER, parts=[TextPart(text="Q")])
|
||||
msg2 = Message(message_id="m2", role=Role.AGENT, parts=[TextPart(text="A")])
|
||||
task = Task(
|
||||
id="task-abc",
|
||||
status=TaskStatus(state=TaskState.COMPLETED),
|
||||
history=[msg1, msg2]
|
||||
)
|
||||
assert len(task.history) == 2
|
||||
|
||||
|
||||
class TestArtifact:
|
||||
"""Tests for Artifact model."""
|
||||
|
||||
def test_artifact_creation(self):
|
||||
"""Test Artifact can be created."""
|
||||
from praisonaiagents.ui.a2a.types import Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-123",
|
||||
parts=[TextPart(text="Generated content")]
|
||||
)
|
||||
assert artifact.artifact_id == "art-123"
|
||||
|
||||
def test_artifact_with_name(self):
|
||||
"""Test Artifact with name and description."""
|
||||
from praisonaiagents.ui.a2a.types import Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-456",
|
||||
name="Report",
|
||||
description="Generated report",
|
||||
parts=[TextPart(text="Report content")]
|
||||
)
|
||||
assert artifact.name == "Report"
|
||||
assert artifact.description == "Generated report"
|
||||
|
||||
|
||||
class TestAgentSkill:
|
||||
"""Tests for AgentSkill model."""
|
||||
|
||||
def test_agent_skill_creation(self):
|
||||
"""Test AgentSkill can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="search",
|
||||
name="Web Search",
|
||||
description="Search the web for information"
|
||||
)
|
||||
assert skill.id == "search"
|
||||
assert skill.name == "Web Search"
|
||||
|
||||
def test_agent_skill_with_tags(self):
|
||||
"""Test AgentSkill with tags."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="calc",
|
||||
name="Calculator",
|
||||
description="Perform calculations",
|
||||
tags=["math", "utility"]
|
||||
)
|
||||
assert skill.tags == ["math", "utility"]
|
||||
|
||||
def test_agent_skill_with_examples(self):
|
||||
"""Test AgentSkill with examples."""
|
||||
from praisonaiagents.ui.a2a.types import AgentSkill
|
||||
skill = AgentSkill(
|
||||
id="translate",
|
||||
name="Translator",
|
||||
description="Translate text",
|
||||
examples=["Translate hello to Spanish", "What is bonjour in English?"]
|
||||
)
|
||||
assert len(skill.examples) == 2
|
||||
|
||||
|
||||
class TestAgentCapabilities:
|
||||
"""Tests for AgentCapabilities model."""
|
||||
|
||||
def test_capabilities_creation(self):
|
||||
"""Test AgentCapabilities can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCapabilities
|
||||
caps = AgentCapabilities(streaming=True)
|
||||
assert caps.streaming is True
|
||||
|
||||
def test_capabilities_defaults(self):
|
||||
"""Test AgentCapabilities defaults."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCapabilities
|
||||
caps = AgentCapabilities()
|
||||
assert caps.streaming is False
|
||||
assert caps.push_notifications is False
|
||||
|
||||
|
||||
class TestAgentCard:
|
||||
"""Tests for AgentCard model."""
|
||||
|
||||
def test_agent_card_creation(self):
|
||||
"""Test AgentCard can be created."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities
|
||||
card = AgentCard(
|
||||
name="Research Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities(streaming=True)
|
||||
)
|
||||
assert card.name == "Research Agent"
|
||||
assert card.url == "http://localhost:8000/a2a"
|
||||
|
||||
def test_agent_card_with_skills(self):
|
||||
"""Test AgentCard with skills."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities, AgentSkill
|
||||
skill = AgentSkill(id="search", name="Search", description="Search web")
|
||||
card = AgentCard(
|
||||
name="Search Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities(),
|
||||
skills=[skill]
|
||||
)
|
||||
assert len(card.skills) == 1
|
||||
|
||||
def test_agent_card_json_serialization(self):
|
||||
"""Test AgentCard can be serialized to JSON."""
|
||||
from praisonaiagents.ui.a2a.types import AgentCard, AgentCapabilities
|
||||
card = AgentCard(
|
||||
name="Test Agent",
|
||||
url="http://localhost:8000/a2a",
|
||||
version="1.0.0",
|
||||
capabilities=AgentCapabilities()
|
||||
)
|
||||
json_str = card.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
assert data["name"] == "Test Agent"
|
||||
|
||||
|
||||
class TestSendMessageRequest:
|
||||
"""Tests for SendMessageRequest model."""
|
||||
|
||||
def test_send_message_request_creation(self):
|
||||
"""Test SendMessageRequest can be created."""
|
||||
from praisonaiagents.ui.a2a.types import SendMessageRequest, Message, Role, TextPart
|
||||
msg = Message(
|
||||
message_id="msg-1",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Hello")]
|
||||
)
|
||||
req = SendMessageRequest(message=msg)
|
||||
assert req.message.message_id == "msg-1"
|
||||
|
||||
def test_send_message_request_with_config(self):
|
||||
"""Test SendMessageRequest with configuration."""
|
||||
from praisonaiagents.ui.a2a.types import SendMessageRequest, Message, Role, TextPart, MessageConfig
|
||||
msg = Message(
|
||||
message_id="msg-2",
|
||||
role=Role.USER,
|
||||
parts=[TextPart(text="Search for AI")]
|
||||
)
|
||||
config = MessageConfig(
|
||||
accept_content_types=["text/plain", "application/json"]
|
||||
)
|
||||
req = SendMessageRequest(message=msg, configuration=config)
|
||||
assert req.configuration is not None
|
||||
|
||||
|
||||
class TestStreamingEvents:
|
||||
"""Tests for streaming event models."""
|
||||
|
||||
def test_task_status_update_event(self):
|
||||
"""Test TaskStatusUpdateEvent creation."""
|
||||
from praisonaiagents.ui.a2a.types import TaskStatusUpdateEvent, TaskStatus, TaskState
|
||||
event = TaskStatusUpdateEvent(
|
||||
task_id="task-123",
|
||||
status=TaskStatus(state=TaskState.WORKING),
|
||||
final=False
|
||||
)
|
||||
assert event.task_id == "task-123"
|
||||
assert event.final is False
|
||||
|
||||
def test_task_artifact_update_event(self):
|
||||
"""Test TaskArtifactUpdateEvent creation."""
|
||||
from praisonaiagents.ui.a2a.types import TaskArtifactUpdateEvent, Artifact, TextPart
|
||||
artifact = Artifact(
|
||||
artifact_id="art-1",
|
||||
parts=[TextPart(text="Partial result")]
|
||||
)
|
||||
event = TaskArtifactUpdateEvent(
|
||||
task_id="task-456",
|
||||
artifact=artifact,
|
||||
append=True,
|
||||
last_chunk=False
|
||||
)
|
||||
assert event.append is True
|
||||
assert event.last_chunk is False
|
||||
@@ -0,0 +1 @@
|
||||
"""A2UI Tests Package."""
|
||||
@@ -0,0 +1,201 @@
|
||||
"""
|
||||
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.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)
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
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,228 @@
|
||||
"""
|
||||
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"
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,318 @@
|
||||
"""
|
||||
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"}
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for AG-UI integration."""
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
Test AGUI Class - TDD Tests for Main Interface Class
|
||||
|
||||
Phase 5: AGUI Class Tests
|
||||
- Test AGUI initialization with Agent
|
||||
- Test AGUI initialization with PraisonAIAgents
|
||||
- Test get_router returns FastAPI router
|
||||
- Test validation errors
|
||||
"""
|
||||
|
||||
|
||||
class TestAGUIInitialization:
|
||||
"""Test AGUI class initialization."""
|
||||
|
||||
def test_agui_init_with_agent(self):
|
||||
"""Test AGUI initialization with single Agent."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.agent == agent
|
||||
assert agui.agents is None
|
||||
|
||||
def test_agui_init_with_agents(self):
|
||||
"""Test AGUI initialization with PraisonAIAgents."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
task = Task(description="Test task", expected_output="Result", agent=agent)
|
||||
agents = PraisonAIAgents(agents=[agent], tasks=[task])
|
||||
|
||||
agui = AGUI(agents=agents)
|
||||
|
||||
assert agui.agents == agents
|
||||
assert agui.agent is None
|
||||
|
||||
def test_agui_requires_agent_or_agents(self):
|
||||
"""Test AGUI raises error if neither agent nor agents provided."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
import pytest
|
||||
|
||||
with pytest.raises(ValueError, match="requires an agent or agents"):
|
||||
AGUI()
|
||||
|
||||
def test_agui_custom_prefix(self):
|
||||
"""Test AGUI with custom prefix."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, prefix="/api/v1")
|
||||
|
||||
assert agui.prefix == "/api/v1"
|
||||
|
||||
def test_agui_custom_tags(self):
|
||||
"""Test AGUI with custom tags."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, tags=["Custom", "Tags"])
|
||||
|
||||
assert agui.tags == ["Custom", "Tags"]
|
||||
|
||||
def test_agui_default_tags(self):
|
||||
"""Test AGUI default tags."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.tags == ["AGUI"]
|
||||
|
||||
|
||||
class TestAGUIRouter:
|
||||
"""Test AGUI router generation."""
|
||||
|
||||
def test_get_router_returns_api_router(self):
|
||||
"""Test get_router returns FastAPI APIRouter."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import APIRouter
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
assert isinstance(router, APIRouter)
|
||||
|
||||
def test_router_has_agui_endpoint(self):
|
||||
"""Test router has /agui POST endpoint."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
# Check routes
|
||||
route_paths = [route.path for route in router.routes]
|
||||
assert "/agui" in route_paths
|
||||
|
||||
def test_router_has_status_endpoint(self):
|
||||
"""Test router has /status GET endpoint."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
# Check routes
|
||||
route_paths = [route.path for route in router.routes]
|
||||
assert "/status" in route_paths
|
||||
|
||||
def test_router_with_prefix(self):
|
||||
"""Test router respects prefix."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, prefix="/api/v1")
|
||||
|
||||
router = agui.get_router()
|
||||
|
||||
assert router.prefix == "/api/v1"
|
||||
|
||||
|
||||
class TestAGUIIntegration:
|
||||
"""Test AGUI integration with FastAPI."""
|
||||
|
||||
def test_include_router_in_fastapi(self):
|
||||
"""Test including AGUI router in FastAPI app."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
# Check app has routes
|
||||
route_paths = [route.path for route in app.routes]
|
||||
assert "/agui" in route_paths
|
||||
assert "/status" in route_paths
|
||||
|
||||
|
||||
class TestAGUIName:
|
||||
"""Test AGUI name and description."""
|
||||
|
||||
def test_agui_name(self):
|
||||
"""Test AGUI name property."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, name="My Agent")
|
||||
|
||||
assert agui.name == "My Agent"
|
||||
|
||||
def test_agui_description(self):
|
||||
"""Test AGUI description property."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent, description="A test agent")
|
||||
|
||||
assert agui.description == "A test agent"
|
||||
|
||||
def test_agui_default_name_from_agent(self):
|
||||
"""Test AGUI uses agent name by default."""
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="TestAgent", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
assert agui.name == "TestAgent"
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Test AG-UI Message Conversion - TDD Tests for Message Format Conversion
|
||||
|
||||
Phase 3: Message Conversion Tests
|
||||
- Test converting AG-UI messages to PraisonAI format
|
||||
- Test converting PraisonAI messages to AG-UI format
|
||||
- Test handling of tool calls
|
||||
- Test handling of tool results
|
||||
- Test edge cases (empty messages, missing fields)
|
||||
"""
|
||||
|
||||
class TestAGUIToPraisonAIConversion:
|
||||
"""Test converting AG-UI messages to PraisonAI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting a user message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="user", content="Hello")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[0]["content"] == "Hello"
|
||||
|
||||
def test_convert_assistant_message(self):
|
||||
"""Test converting an assistant message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="assistant", content="Hi there!")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert result[0]["content"] == "Hi there!"
|
||||
|
||||
def test_convert_tool_message(self):
|
||||
"""Test converting a tool result message."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [Message(role="tool", content="Search results", tool_call_id="tc-123")]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "tool"
|
||||
assert result[0]["content"] == "Search results"
|
||||
assert result[0]["tool_call_id"] == "tc-123"
|
||||
|
||||
def test_convert_assistant_with_tool_calls(self):
|
||||
"""Test converting assistant message with tool calls."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message, ToolCall, FunctionCall
|
||||
|
||||
agui_messages = [
|
||||
Message(
|
||||
role="assistant",
|
||||
content="",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="tc-123",
|
||||
function=FunctionCall(name="search", arguments='{"query": "test"}')
|
||||
)
|
||||
]
|
||||
)
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "assistant"
|
||||
assert "tool_calls" in result[0]
|
||||
assert len(result[0]["tool_calls"]) == 1
|
||||
assert result[0]["tool_calls"][0]["function"]["name"] == "search"
|
||||
|
||||
def test_convert_multiple_messages(self):
|
||||
"""Test converting multiple messages."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [
|
||||
Message(role="user", content="Hello"),
|
||||
Message(role="assistant", content="Hi!"),
|
||||
Message(role="user", content="How are you?"),
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0]["role"] == "user"
|
||||
assert result[1]["role"] == "assistant"
|
||||
assert result[2]["role"] == "user"
|
||||
|
||||
def test_convert_empty_messages(self):
|
||||
"""Test converting empty message list."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
|
||||
result = agui_messages_to_praisonai([])
|
||||
assert result == []
|
||||
|
||||
def test_skip_system_messages(self):
|
||||
"""Test that system messages are skipped (agent builds its own)."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
agui_messages = [
|
||||
Message(role="system", content="You are helpful"),
|
||||
Message(role="user", content="Hello"),
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
# System message should be skipped
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
|
||||
|
||||
class TestPraisonAIToAGUIConversion:
|
||||
"""Test converting PraisonAI messages to AG-UI format."""
|
||||
|
||||
def test_convert_user_message(self):
|
||||
"""Test converting a user message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [{"role": "user", "content": "Hello"}]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "user"
|
||||
assert result[0].content == "Hello"
|
||||
|
||||
def test_convert_assistant_message(self):
|
||||
"""Test converting an assistant message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [{"role": "assistant", "content": "Hi there!"}]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "assistant"
|
||||
assert result[0].content == "Hi there!"
|
||||
|
||||
def test_convert_tool_result(self):
|
||||
"""Test converting a tool result message."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [
|
||||
{"role": "tool", "content": "Search results", "tool_call_id": "tc-123"}
|
||||
]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "tool"
|
||||
assert result[0].content == "Search results"
|
||||
assert result[0].tool_call_id == "tc-123"
|
||||
|
||||
def test_convert_assistant_with_tool_calls(self):
|
||||
"""Test converting assistant message with tool calls."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
praison_messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "tc-123",
|
||||
"function": {
|
||||
"name": "search",
|
||||
"arguments": '{"query": "test"}'
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
result = praisonai_messages_to_agui(praison_messages)
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].role == "assistant"
|
||||
assert result[0].tool_calls is not None
|
||||
assert len(result[0].tool_calls) == 1
|
||||
assert result[0].tool_calls[0].function.name == "search"
|
||||
|
||||
def test_convert_empty_messages(self):
|
||||
"""Test converting empty message list."""
|
||||
from praisonaiagents.ui.agui.conversion import praisonai_messages_to_agui
|
||||
|
||||
result = praisonai_messages_to_agui([])
|
||||
assert result == []
|
||||
|
||||
|
||||
class TestStateConversion:
|
||||
"""Test state conversion utilities."""
|
||||
|
||||
def test_validate_dict_state(self):
|
||||
"""Test validating dict state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
state = {"key": "value"}
|
||||
result = validate_state(state, "thread-123")
|
||||
|
||||
assert result == {"key": "value"}
|
||||
|
||||
def test_validate_none_state(self):
|
||||
"""Test validating None state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
result = validate_state(None, "thread-123")
|
||||
assert result is None
|
||||
|
||||
def test_validate_pydantic_state(self):
|
||||
"""Test validating Pydantic model state."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
from pydantic import BaseModel
|
||||
|
||||
class TestState(BaseModel):
|
||||
key: str = "value"
|
||||
|
||||
state = TestState()
|
||||
result = validate_state(state, "thread-123")
|
||||
|
||||
assert result == {"key": "value"}
|
||||
|
||||
def test_validate_invalid_state_returns_none(self):
|
||||
"""Test that invalid state returns None."""
|
||||
from praisonaiagents.ui.agui.conversion import validate_state
|
||||
|
||||
# A string is not a valid state
|
||||
result = validate_state("invalid", "thread-123")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestToolCallDeduplication:
|
||||
"""Test tool call deduplication logic."""
|
||||
|
||||
def test_deduplicate_tool_results(self):
|
||||
"""Test that duplicate tool results are deduplicated."""
|
||||
from praisonaiagents.ui.agui.conversion import agui_messages_to_praisonai
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
# Same tool_call_id appearing twice
|
||||
agui_messages = [
|
||||
Message(role="tool", content="Result 1", tool_call_id="tc-123"),
|
||||
Message(role="tool", content="Result 2", tool_call_id="tc-123"), # Duplicate
|
||||
]
|
||||
result = agui_messages_to_praisonai(agui_messages)
|
||||
|
||||
# Should only have one result
|
||||
assert len(result) == 1
|
||||
assert result[0]["content"] == "Result 1" # First one wins
|
||||
|
||||
|
||||
class TestExtractUserInput:
|
||||
"""Test extracting user input from messages."""
|
||||
|
||||
def test_extract_last_user_message(self):
|
||||
"""Test extracting the last user message."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
messages = [
|
||||
Message(role="user", content="First message"),
|
||||
Message(role="assistant", content="Response"),
|
||||
Message(role="user", content="Last message"),
|
||||
]
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == "Last message"
|
||||
|
||||
def test_extract_from_empty_messages(self):
|
||||
"""Test extracting from empty messages."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
|
||||
result = extract_user_input([])
|
||||
assert result == ""
|
||||
|
||||
def test_extract_no_user_messages(self):
|
||||
"""Test extracting when no user messages exist."""
|
||||
from praisonaiagents.ui.agui.conversion import extract_user_input
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
|
||||
messages = [
|
||||
Message(role="assistant", content="Hello"),
|
||||
]
|
||||
result = extract_user_input(messages)
|
||||
|
||||
assert result == ""
|
||||
@@ -0,0 +1,311 @@
|
||||
"""
|
||||
Test AG-UI Integration - TDD Tests for End-to-End Integration
|
||||
|
||||
Phase 7: Integration Tests
|
||||
- Test full flow with single Agent
|
||||
- Test full flow with PraisonAIAgents
|
||||
- Test tool calling flow
|
||||
- Test error recovery
|
||||
- Test state persistence
|
||||
"""
|
||||
|
||||
class TestSingleAgentIntegration:
|
||||
"""Test integration with single Agent."""
|
||||
|
||||
def test_single_agent_chat_flow(self):
|
||||
"""Test complete chat flow with single agent."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
agent = Agent(
|
||||
name="Assistant",
|
||||
role="Helpful Assistant",
|
||||
goal="Help users with their questions",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Send a message
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Say hello"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse SSE events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
|
||||
# Should have run started event
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_STARTED" in event_types
|
||||
assert "RUN_FINISHED" in event_types or "RUN_ERROR" in event_types
|
||||
|
||||
|
||||
class TestMultiAgentIntegration:
|
||||
"""Test integration with PraisonAIAgents."""
|
||||
|
||||
def test_multi_agent_workflow_flow(self):
|
||||
"""Test complete workflow with multiple agents."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent, Task, PraisonAIAgents
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
# Create agents
|
||||
researcher = Agent(
|
||||
name="Researcher",
|
||||
role="Research Analyst",
|
||||
goal="Research topics thoroughly",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
writer = Agent(
|
||||
name="Writer",
|
||||
role="Content Writer",
|
||||
goal="Write engaging content",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
|
||||
# Create tasks
|
||||
research_task = Task(
|
||||
description="Research the topic: {topic}",
|
||||
expected_output="Research findings",
|
||||
agent=researcher
|
||||
)
|
||||
write_task = Task(
|
||||
description="Write about the research",
|
||||
expected_output="Written article",
|
||||
agent=writer
|
||||
)
|
||||
|
||||
# Create PraisonAIAgents
|
||||
agents = PraisonAIAgents(
|
||||
agents=[researcher, writer],
|
||||
tasks=[research_task, write_task]
|
||||
)
|
||||
|
||||
agui = AGUI(agents=agents)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Research AI"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestToolCallingIntegration:
|
||||
"""Test integration with tool calling."""
|
||||
|
||||
def test_agent_with_tools(self):
|
||||
"""Test agent with tools integration."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
def search_tool(query: str) -> str:
|
||||
"""Search for information."""
|
||||
return f"Search results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="SearchAgent",
|
||||
role="Search Assistant",
|
||||
goal="Help users search for information",
|
||||
tools=[search_tool],
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Search for Python tutorials"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
|
||||
# Check for tool call events (may or may not have them depending on LLM response)
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_STARTED" in event_types
|
||||
|
||||
|
||||
class TestStateManagement:
|
||||
"""Test state management in AG-UI."""
|
||||
|
||||
def test_state_passed_to_agent(self):
|
||||
"""Test that state is passed to agent."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(
|
||||
name="StatefulAgent",
|
||||
role="Stateful Assistant",
|
||||
goal="Remember context",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": {"user_name": "John", "preferences": {"language": "en"}}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestConversationHistory:
|
||||
"""Test conversation history handling."""
|
||||
|
||||
def test_multiple_messages_in_history(self):
|
||||
"""Test handling multiple messages in conversation history."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(
|
||||
name="ChatAgent",
|
||||
role="Chat Assistant",
|
||||
goal="Have conversations",
|
||||
llm="gpt-4o-mini"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [
|
||||
{"role": "user", "content": "My name is John"},
|
||||
{"role": "assistant", "content": "Hello John!"},
|
||||
{"role": "user", "content": "What is my name?"}
|
||||
]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestErrorRecovery:
|
||||
"""Test error recovery scenarios."""
|
||||
|
||||
def test_graceful_error_handling(self):
|
||||
"""Test graceful error handling."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
import json
|
||||
|
||||
# Agent with invalid LLM should handle error gracefully
|
||||
agent = Agent(
|
||||
name="ErrorAgent",
|
||||
role="Error Test",
|
||||
goal="Test error handling",
|
||||
llm="invalid-model-name"
|
||||
)
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "test-thread",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should not crash, should return error event
|
||||
assert response.status_code == 200
|
||||
|
||||
# Parse events
|
||||
events = []
|
||||
for line in response.text.split("\n"):
|
||||
if line.startswith("data: "):
|
||||
try:
|
||||
event_data = json.loads(line[6:])
|
||||
events.append(event_data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Should have error event
|
||||
event_types = [e.get("type") for e in events]
|
||||
assert "RUN_ERROR" in event_types or "RUN_FINISHED" in event_types
|
||||
|
||||
|
||||
class TestPerformance:
|
||||
"""Test performance - ensure no impact on existing functionality."""
|
||||
|
||||
def test_agui_import_does_not_slow_package(self):
|
||||
"""Test that importing AGUI doesn't slow down package import."""
|
||||
import time
|
||||
|
||||
start = time.time()
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
import_time = time.time() - start
|
||||
|
||||
# Import should be fast (< 1 second)
|
||||
assert import_time < 1.0
|
||||
|
||||
def test_agui_lazy_loading(self):
|
||||
"""Test that AGUI uses lazy loading where possible."""
|
||||
# This test ensures we don't import heavy dependencies until needed
|
||||
import sys
|
||||
|
||||
# Remove agui from cache if present
|
||||
modules_to_remove = [k for k in sys.modules.keys() if 'agui' in k]
|
||||
for mod in modules_to_remove:
|
||||
del sys.modules[mod]
|
||||
|
||||
# Import should not load FastAPI until get_router is called
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
# FastAPI should be loaded when get_router is called
|
||||
router = agui.get_router()
|
||||
assert router is not None
|
||||
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
Test AG-UI Router - TDD Tests for FastAPI HTTP Endpoints
|
||||
|
||||
Phase 6: FastAPI Router Tests
|
||||
- Test /agui POST endpoint
|
||||
- Test /status GET endpoint
|
||||
- Test request validation
|
||||
- Test streaming response
|
||||
- Test error handling
|
||||
"""
|
||||
|
||||
class TestStatusEndpoint:
|
||||
"""Test /status endpoint."""
|
||||
|
||||
def test_status_endpoint_returns_available(self):
|
||||
"""Test status endpoint returns available status."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
response = client.get("/status")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "available"}
|
||||
|
||||
|
||||
class TestAGUIEndpoint:
|
||||
"""Test /agui POST endpoint."""
|
||||
|
||||
def test_agui_endpoint_accepts_post(self):
|
||||
"""Test /agui endpoint accepts POST requests."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Minimal valid request
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return streaming response (200)
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_agui_endpoint_returns_streaming_response(self):
|
||||
"""Test /agui endpoint returns streaming response."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Check content type is event-stream
|
||||
assert "text/event-stream" in response.headers.get("content-type", "")
|
||||
|
||||
def test_agui_endpoint_requires_thread_id(self):
|
||||
"""Test /agui endpoint requires thread_id."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Missing thread_id
|
||||
response = client.post("/agui", json={
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return validation error
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
class TestAGUIEndpointWithState:
|
||||
"""Test /agui endpoint with state."""
|
||||
|
||||
def test_agui_endpoint_accepts_state(self):
|
||||
"""Test /agui endpoint accepts state parameter."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"state": {"custom_key": "custom_value"}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointWithRunId:
|
||||
"""Test /agui endpoint with run_id."""
|
||||
|
||||
def test_agui_endpoint_accepts_run_id(self):
|
||||
"""Test /agui endpoint accepts run_id parameter."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"run_id": "run-456",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
def test_agui_endpoint_generates_run_id_if_missing(self):
|
||||
"""Test /agui endpoint generates run_id if not provided."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should succeed even without run_id
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointCORS:
|
||||
"""Test /agui endpoint CORS headers."""
|
||||
|
||||
def test_agui_endpoint_cors_headers(self):
|
||||
"""Test /agui endpoint includes CORS headers."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Check CORS headers
|
||||
assert response.headers.get("access-control-allow-origin") == "*"
|
||||
|
||||
|
||||
class TestAGUIEndpointErrorHandling:
|
||||
"""Test /agui endpoint error handling."""
|
||||
|
||||
def test_agui_endpoint_handles_agent_error(self):
|
||||
"""Test /agui endpoint handles agent errors gracefully."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create agent that will fail
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# This should not crash the server
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}]
|
||||
})
|
||||
|
||||
# Should return 200 with error event in stream
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestAGUIEndpointWithForwardedProps:
|
||||
"""Test /agui endpoint with forwarded_props."""
|
||||
|
||||
def test_agui_endpoint_accepts_forwarded_props(self):
|
||||
"""Test /agui endpoint accepts forwarded_props."""
|
||||
from fastapi.testclient import TestClient
|
||||
from praisonaiagents.ui.agui import AGUI
|
||||
from praisonaiagents import Agent
|
||||
from fastapi import FastAPI
|
||||
|
||||
agent = Agent(name="Test", role="Tester", goal="Test things")
|
||||
agui = AGUI(agent=agent)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(agui.get_router())
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/agui", json={
|
||||
"thread_id": "thread-123",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"forwarded_props": {"user_id": "user-123"}
|
||||
})
|
||||
|
||||
assert response.status_code == 200
|
||||
@@ -0,0 +1,240 @@
|
||||
"""
|
||||
Test AG-UI Event Streaming - TDD Tests for Streaming PraisonAI Responses
|
||||
|
||||
Phase 4: Event Streaming Tests
|
||||
- Test streaming text content as AG-UI events
|
||||
- Test streaming tool calls as AG-UI events
|
||||
- Test run lifecycle events (start, finish, error)
|
||||
- Test step events
|
||||
- Test event ordering
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
|
||||
class TestTextMessageStreaming:
|
||||
"""Test streaming text messages."""
|
||||
|
||||
def test_stream_simple_text(self):
|
||||
"""Test streaming simple text content."""
|
||||
from praisonaiagents.ui.agui.streaming import create_text_message_events
|
||||
|
||||
events = list(create_text_message_events(
|
||||
content="Hello, world!",
|
||||
message_id="msg-123"
|
||||
))
|
||||
|
||||
# Should have start, content, end events
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[0].message_id == "msg-123"
|
||||
assert events[1].type == "TEXT_MESSAGE_CONTENT"
|
||||
assert events[1].delta == "Hello, world!"
|
||||
assert events[2].type == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_stream_chunked_text(self):
|
||||
"""Test streaming text in chunks."""
|
||||
from praisonaiagents.ui.agui.streaming import stream_text_chunks
|
||||
|
||||
chunks = ["Hello", ", ", "world", "!"]
|
||||
events = list(stream_text_chunks(chunks, message_id="msg-123"))
|
||||
|
||||
# Should have start, 4 content events, end
|
||||
assert len(events) == 6
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[1].delta == "Hello"
|
||||
assert events[2].delta == ", "
|
||||
assert events[3].delta == "world"
|
||||
assert events[4].delta == "!"
|
||||
assert events[5].type == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_stream_empty_text(self):
|
||||
"""Test streaming empty text."""
|
||||
from praisonaiagents.ui.agui.streaming import create_text_message_events
|
||||
|
||||
events = list(create_text_message_events(content="", message_id="msg-123"))
|
||||
|
||||
# Should still have start and end, but no content
|
||||
assert len(events) == 2
|
||||
assert events[0].type == "TEXT_MESSAGE_START"
|
||||
assert events[1].type == "TEXT_MESSAGE_END"
|
||||
|
||||
|
||||
class TestToolCallStreaming:
|
||||
"""Test streaming tool calls."""
|
||||
|
||||
def test_stream_tool_call(self):
|
||||
"""Test streaming a tool call."""
|
||||
from praisonaiagents.ui.agui.streaming import create_tool_call_events
|
||||
|
||||
events = list(create_tool_call_events(
|
||||
tool_call_id="tc-123",
|
||||
tool_name="search",
|
||||
arguments='{"query": "test"}',
|
||||
parent_message_id="msg-123"
|
||||
))
|
||||
|
||||
# Should have start, args, end events
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "TOOL_CALL_START"
|
||||
assert events[0].tool_call_id == "tc-123"
|
||||
assert events[0].tool_call_name == "search"
|
||||
assert events[1].type == "TOOL_CALL_ARGS"
|
||||
assert events[1].delta == '{"query": "test"}'
|
||||
assert events[2].type == "TOOL_CALL_END"
|
||||
|
||||
def test_stream_tool_result(self):
|
||||
"""Test streaming a tool result."""
|
||||
from praisonaiagents.ui.agui.streaming import create_tool_result_event
|
||||
|
||||
event = create_tool_result_event(
|
||||
tool_call_id="tc-123",
|
||||
content="Search results...",
|
||||
message_id="msg-456"
|
||||
)
|
||||
|
||||
assert event.type == "TOOL_CALL_RESULT"
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.content == "Search results..."
|
||||
|
||||
|
||||
class TestRunLifecycleEvents:
|
||||
"""Test run lifecycle events."""
|
||||
|
||||
def test_create_run_started_event(self):
|
||||
"""Test creating run started event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_started_event
|
||||
|
||||
event = create_run_started_event(thread_id="thread-123", run_id="run-456")
|
||||
|
||||
assert event.type == "RUN_STARTED"
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
|
||||
def test_create_run_finished_event(self):
|
||||
"""Test creating run finished event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_finished_event
|
||||
|
||||
event = create_run_finished_event(thread_id="thread-123", run_id="run-456")
|
||||
|
||||
assert event.type == "RUN_FINISHED"
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
|
||||
def test_create_run_error_event(self):
|
||||
"""Test creating run error event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_run_error_event
|
||||
|
||||
event = create_run_error_event(message="Something went wrong", code="ERR001")
|
||||
|
||||
assert event.type == "RUN_ERROR"
|
||||
assert event.message == "Something went wrong"
|
||||
assert event.code == "ERR001"
|
||||
|
||||
|
||||
class TestStepEvents:
|
||||
"""Test step events."""
|
||||
|
||||
def test_create_step_started_event(self):
|
||||
"""Test creating step started event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_step_started_event
|
||||
|
||||
event = create_step_started_event(step_name="research")
|
||||
|
||||
assert event.type == "STEP_STARTED"
|
||||
assert event.step_name == "research"
|
||||
|
||||
def test_create_step_finished_event(self):
|
||||
"""Test creating step finished event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_step_finished_event
|
||||
|
||||
event = create_step_finished_event(step_name="research")
|
||||
|
||||
assert event.type == "STEP_FINISHED"
|
||||
assert event.step_name == "research"
|
||||
|
||||
|
||||
class TestEventBuffer:
|
||||
"""Test EventBuffer for managing event ordering."""
|
||||
|
||||
def test_event_buffer_creation(self):
|
||||
"""Test EventBuffer creation."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
assert buffer.active_tool_call_ids == set()
|
||||
assert buffer.current_text_message_id == ""
|
||||
|
||||
def test_event_buffer_start_tool_call(self):
|
||||
"""Test starting a tool call in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
buffer.start_tool_call("tc-123")
|
||||
|
||||
assert "tc-123" in buffer.active_tool_call_ids
|
||||
|
||||
def test_event_buffer_end_tool_call(self):
|
||||
"""Test ending a tool call in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
buffer.start_tool_call("tc-123")
|
||||
buffer.end_tool_call("tc-123")
|
||||
|
||||
assert "tc-123" not in buffer.active_tool_call_ids
|
||||
assert "tc-123" in buffer.ended_tool_call_ids
|
||||
|
||||
def test_event_buffer_start_text_message(self):
|
||||
"""Test starting a text message in buffer."""
|
||||
from praisonaiagents.ui.agui.streaming import EventBuffer
|
||||
|
||||
buffer = EventBuffer()
|
||||
msg_id = buffer.start_text_message()
|
||||
|
||||
assert msg_id != ""
|
||||
assert buffer.current_text_message_id == msg_id
|
||||
|
||||
|
||||
class TestAsyncStreaming:
|
||||
"""Test async streaming functionality."""
|
||||
|
||||
def test_async_stream_response(self):
|
||||
"""Test async streaming of agent response."""
|
||||
async def run_test():
|
||||
from praisonaiagents.ui.agui.streaming import async_stream_response
|
||||
|
||||
# Mock response chunks
|
||||
async def mock_chunks():
|
||||
yield "Hello"
|
||||
yield ", "
|
||||
yield "world!"
|
||||
|
||||
events = []
|
||||
async for event in async_stream_response(
|
||||
response_stream=mock_chunks(),
|
||||
thread_id="thread-123",
|
||||
run_id="run-456"
|
||||
):
|
||||
events.append(event)
|
||||
|
||||
# Should have run_started, text events, run_finished
|
||||
assert len(events) > 0
|
||||
assert events[0].type == "RUN_STARTED"
|
||||
assert events[-1].type == "RUN_FINISHED"
|
||||
|
||||
asyncio.run(run_test())
|
||||
|
||||
|
||||
class TestStateSnapshotEvents:
|
||||
"""Test state snapshot events."""
|
||||
|
||||
def test_create_state_snapshot_event(self):
|
||||
"""Test creating state snapshot event."""
|
||||
from praisonaiagents.ui.agui.streaming import create_state_snapshot_event
|
||||
|
||||
state = {"task": "research", "progress": 50}
|
||||
event = create_state_snapshot_event(state)
|
||||
|
||||
assert event.type == "STATE_SNAPSHOT"
|
||||
assert event.snapshot == state
|
||||
@@ -0,0 +1,327 @@
|
||||
"""
|
||||
Test AG-UI Types - TDD Tests for Event and Message Types
|
||||
|
||||
Phase 2: Core Types Tests
|
||||
- Test EventType enum values
|
||||
- Test BaseEvent creation
|
||||
- Test all event type classes
|
||||
- Test message type classes
|
||||
- Test RunAgentInput validation
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
|
||||
class TestEventType:
|
||||
"""Test EventType enum."""
|
||||
|
||||
def test_event_type_text_message_start(self):
|
||||
"""Test TEXT_MESSAGE_START event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_START == "TEXT_MESSAGE_START"
|
||||
|
||||
def test_event_type_text_message_content(self):
|
||||
"""Test TEXT_MESSAGE_CONTENT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_CONTENT == "TEXT_MESSAGE_CONTENT"
|
||||
|
||||
def test_event_type_text_message_end(self):
|
||||
"""Test TEXT_MESSAGE_END event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TEXT_MESSAGE_END == "TEXT_MESSAGE_END"
|
||||
|
||||
def test_event_type_tool_call_start(self):
|
||||
"""Test TOOL_CALL_START event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_START == "TOOL_CALL_START"
|
||||
|
||||
def test_event_type_tool_call_args(self):
|
||||
"""Test TOOL_CALL_ARGS event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_ARGS == "TOOL_CALL_ARGS"
|
||||
|
||||
def test_event_type_tool_call_end(self):
|
||||
"""Test TOOL_CALL_END event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_END == "TOOL_CALL_END"
|
||||
|
||||
def test_event_type_tool_call_result(self):
|
||||
"""Test TOOL_CALL_RESULT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.TOOL_CALL_RESULT == "TOOL_CALL_RESULT"
|
||||
|
||||
def test_event_type_run_started(self):
|
||||
"""Test RUN_STARTED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_STARTED == "RUN_STARTED"
|
||||
|
||||
def test_event_type_run_finished(self):
|
||||
"""Test RUN_FINISHED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_FINISHED == "RUN_FINISHED"
|
||||
|
||||
def test_event_type_run_error(self):
|
||||
"""Test RUN_ERROR event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.RUN_ERROR == "RUN_ERROR"
|
||||
|
||||
def test_event_type_step_started(self):
|
||||
"""Test STEP_STARTED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STEP_STARTED == "STEP_STARTED"
|
||||
|
||||
def test_event_type_step_finished(self):
|
||||
"""Test STEP_FINISHED event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STEP_FINISHED == "STEP_FINISHED"
|
||||
|
||||
def test_event_type_state_snapshot(self):
|
||||
"""Test STATE_SNAPSHOT event type exists."""
|
||||
from praisonaiagents.ui.agui.types import EventType
|
||||
assert EventType.STATE_SNAPSHOT == "STATE_SNAPSHOT"
|
||||
|
||||
|
||||
class TestBaseEvent:
|
||||
"""Test BaseEvent class."""
|
||||
|
||||
def test_base_event_creation(self):
|
||||
"""Test BaseEvent can be created."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
assert event.type == EventType.RUN_STARTED
|
||||
|
||||
def test_base_event_timestamp_optional(self):
|
||||
"""Test BaseEvent timestamp is optional."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
assert event.timestamp is None
|
||||
|
||||
def test_base_event_to_dict(self):
|
||||
"""Test BaseEvent can be converted to dict."""
|
||||
from praisonaiagents.ui.agui.types import BaseEvent, EventType
|
||||
event = BaseEvent(type=EventType.RUN_STARTED)
|
||||
data = event.model_dump()
|
||||
assert data["type"] == "RUN_STARTED"
|
||||
|
||||
|
||||
class TestTextMessageEvents:
|
||||
"""Test text message event classes."""
|
||||
|
||||
def test_text_message_start_event(self):
|
||||
"""Test TextMessageStartEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.role == "assistant"
|
||||
assert event.type == "TEXT_MESSAGE_START"
|
||||
|
||||
def test_text_message_content_event(self):
|
||||
"""Test TextMessageContentEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageContentEvent
|
||||
event = TextMessageContentEvent(message_id="msg-123", delta="Hello")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.delta == "Hello"
|
||||
assert event.type == "TEXT_MESSAGE_CONTENT"
|
||||
|
||||
def test_text_message_end_event(self):
|
||||
"""Test TextMessageEndEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageEndEvent
|
||||
event = TextMessageEndEvent(message_id="msg-123")
|
||||
assert event.message_id == "msg-123"
|
||||
assert event.type == "TEXT_MESSAGE_END"
|
||||
|
||||
|
||||
class TestToolCallEvents:
|
||||
"""Test tool call event classes."""
|
||||
|
||||
def test_tool_call_start_event(self):
|
||||
"""Test ToolCallStartEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallStartEvent
|
||||
event = ToolCallStartEvent(
|
||||
tool_call_id="tc-123",
|
||||
tool_call_name="search",
|
||||
parent_message_id="msg-123"
|
||||
)
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.tool_call_name == "search"
|
||||
assert event.parent_message_id == "msg-123"
|
||||
assert event.type == "TOOL_CALL_START"
|
||||
|
||||
def test_tool_call_args_event(self):
|
||||
"""Test ToolCallArgsEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallArgsEvent
|
||||
event = ToolCallArgsEvent(tool_call_id="tc-123", delta='{"query": "test"}')
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.delta == '{"query": "test"}'
|
||||
assert event.type == "TOOL_CALL_ARGS"
|
||||
|
||||
def test_tool_call_end_event(self):
|
||||
"""Test ToolCallEndEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallEndEvent
|
||||
event = ToolCallEndEvent(tool_call_id="tc-123")
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.type == "TOOL_CALL_END"
|
||||
|
||||
def test_tool_call_result_event(self):
|
||||
"""Test ToolCallResultEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCallResultEvent
|
||||
event = ToolCallResultEvent(
|
||||
message_id="msg-456",
|
||||
tool_call_id="tc-123",
|
||||
content="Search results..."
|
||||
)
|
||||
assert event.message_id == "msg-456"
|
||||
assert event.tool_call_id == "tc-123"
|
||||
assert event.content == "Search results..."
|
||||
assert event.type == "TOOL_CALL_RESULT"
|
||||
|
||||
|
||||
class TestRunEvents:
|
||||
"""Test run lifecycle event classes."""
|
||||
|
||||
def test_run_started_event(self):
|
||||
"""Test RunStartedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunStartedEvent
|
||||
event = RunStartedEvent(thread_id="thread-123", run_id="run-456")
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
assert event.type == "RUN_STARTED"
|
||||
|
||||
def test_run_finished_event(self):
|
||||
"""Test RunFinishedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunFinishedEvent
|
||||
event = RunFinishedEvent(thread_id="thread-123", run_id="run-456")
|
||||
assert event.thread_id == "thread-123"
|
||||
assert event.run_id == "run-456"
|
||||
assert event.type == "RUN_FINISHED"
|
||||
|
||||
def test_run_error_event(self):
|
||||
"""Test RunErrorEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import RunErrorEvent
|
||||
event = RunErrorEvent(message="Something went wrong")
|
||||
assert event.message == "Something went wrong"
|
||||
assert event.type == "RUN_ERROR"
|
||||
|
||||
|
||||
class TestStepEvents:
|
||||
"""Test step event classes."""
|
||||
|
||||
def test_step_started_event(self):
|
||||
"""Test StepStartedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StepStartedEvent
|
||||
event = StepStartedEvent(step_name="research")
|
||||
assert event.step_name == "research"
|
||||
assert event.type == "STEP_STARTED"
|
||||
|
||||
def test_step_finished_event(self):
|
||||
"""Test StepFinishedEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StepFinishedEvent
|
||||
event = StepFinishedEvent(step_name="research")
|
||||
assert event.step_name == "research"
|
||||
assert event.type == "STEP_FINISHED"
|
||||
|
||||
|
||||
class TestStateEvents:
|
||||
"""Test state event classes."""
|
||||
|
||||
def test_state_snapshot_event(self):
|
||||
"""Test StateSnapshotEvent creation."""
|
||||
from praisonaiagents.ui.agui.types import StateSnapshotEvent
|
||||
event = StateSnapshotEvent(snapshot={"key": "value"})
|
||||
assert event.snapshot == {"key": "value"}
|
||||
assert event.type == "STATE_SNAPSHOT"
|
||||
|
||||
|
||||
class TestMessageTypes:
|
||||
"""Test message type classes."""
|
||||
|
||||
def test_user_message(self):
|
||||
"""Test UserMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="user", content="Hello")
|
||||
assert msg.role == "user"
|
||||
assert msg.content == "Hello"
|
||||
|
||||
def test_assistant_message(self):
|
||||
"""Test AssistantMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="assistant", content="Hi there!")
|
||||
assert msg.role == "assistant"
|
||||
assert msg.content == "Hi there!"
|
||||
|
||||
def test_tool_message(self):
|
||||
"""Test ToolMessage creation."""
|
||||
from praisonaiagents.ui.agui.types import Message
|
||||
msg = Message(role="tool", content="Result", tool_call_id="tc-123")
|
||||
assert msg.role == "tool"
|
||||
assert msg.content == "Result"
|
||||
assert msg.tool_call_id == "tc-123"
|
||||
|
||||
|
||||
class TestToolCall:
|
||||
"""Test ToolCall type."""
|
||||
|
||||
def test_tool_call_creation(self):
|
||||
"""Test ToolCall creation."""
|
||||
from praisonaiagents.ui.agui.types import ToolCall, FunctionCall
|
||||
tc = ToolCall(
|
||||
id="tc-123",
|
||||
function=FunctionCall(name="search", arguments='{"query": "test"}')
|
||||
)
|
||||
assert tc.id == "tc-123"
|
||||
assert tc.function.name == "search"
|
||||
assert tc.function.arguments == '{"query": "test"}'
|
||||
|
||||
|
||||
class TestRunAgentInput:
|
||||
"""Test RunAgentInput type."""
|
||||
|
||||
def test_run_agent_input_creation(self):
|
||||
"""Test RunAgentInput creation."""
|
||||
from praisonaiagents.ui.agui.types import RunAgentInput, Message
|
||||
input_data = RunAgentInput(
|
||||
thread_id="thread-123",
|
||||
run_id="run-456",
|
||||
messages=[Message(role="user", content="Hello")]
|
||||
)
|
||||
assert input_data.thread_id == "thread-123"
|
||||
assert input_data.run_id == "run-456"
|
||||
assert len(input_data.messages) == 1
|
||||
|
||||
def test_run_agent_input_optional_fields(self):
|
||||
"""Test RunAgentInput optional fields."""
|
||||
from praisonaiagents.ui.agui.types import RunAgentInput
|
||||
input_data = RunAgentInput(thread_id="thread-123")
|
||||
assert input_data.thread_id == "thread-123"
|
||||
assert input_data.run_id is None
|
||||
assert input_data.messages is None or len(input_data.messages) == 0
|
||||
assert input_data.state is None
|
||||
|
||||
|
||||
class TestEventSerialization:
|
||||
"""Test event serialization to JSON."""
|
||||
|
||||
def test_event_json_serialization(self):
|
||||
"""Test events can be serialized to JSON."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
json_str = event.model_dump_json()
|
||||
data = json.loads(json_str)
|
||||
assert data["type"] == "TEXT_MESSAGE_START"
|
||||
assert data["message_id"] == "msg-123"
|
||||
|
||||
def test_event_sse_format(self):
|
||||
"""Test events can be formatted as SSE."""
|
||||
from praisonaiagents.ui.agui.types import TextMessageStartEvent
|
||||
from praisonaiagents.ui.agui.encoder import EventEncoder
|
||||
|
||||
encoder = EventEncoder()
|
||||
event = TextMessageStartEvent(message_id="msg-123", role="assistant")
|
||||
sse = encoder.encode(event)
|
||||
|
||||
assert sse.startswith("data: ")
|
||||
assert sse.endswith("\n\n")
|
||||
# Parse the JSON part
|
||||
json_part = sse[6:-2] # Remove "data: " prefix and "\n\n" suffix
|
||||
data = json.loads(json_part)
|
||||
assert data["type"] == "TEXT_MESSAGE_START"
|
||||
Generated
+1
-1
@@ -2422,7 +2422,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.14"
|
||||
version = "0.1.17"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
||||
@@ -60,6 +60,7 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
|
||||
- [Planning & Workflows](#planning--workflows)
|
||||
- [Specialised Agents](#specialised-agents)
|
||||
- [MCP Protocol](#mcp-protocol)
|
||||
- [A2A & A2UI Protocols](#a2a--a2ui-protocols)
|
||||
- [Safety & Control](#safety--control)
|
||||
- [Developer Tools](#developer-tools)
|
||||
- [📊 Process Types & Patterns](#-process-types--patterns)
|
||||
@@ -188,6 +189,7 @@ npm install praisonai
|
||||
| 🧠 Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [📖](https://docs.praison.ai/concepts/memory) |
|
||||
| 📄 Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [📖](https://docs.praison.ai/features/chat-with-pdf) |
|
||||
| 💻 Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [📖](https://docs.praison.ai/features/codeagent) |
|
||||
| ✏️ AI Code Editing Tools | [Example](#ai-code-editing) | [📖](https://docs.praison.ai/features/code) |
|
||||
| 📚 RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [📖](https://docs.praison.ai/features/rag) |
|
||||
| 🤔 Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [📖](https://docs.praison.ai/features/async) |
|
||||
| 🔢 Math Agents | [Example](examples/python/agents/math-agent.py) | [📖](https://docs.praison.ai/features/mathagent) |
|
||||
@@ -1628,6 +1630,48 @@ server.run() # stdio for Claude Desktop
|
||||
| Security | Origin validation, DNS rebinding prevention |
|
||||
| WebSocket | Auto-reconnect with exponential backoff |
|
||||
|
||||
## A2A (Agent2Agent Protocol)
|
||||
|
||||
PraisonAI supports the [A2A Protocol](https://a2a-protocol.org) for agent-to-agent communication, enabling your agents to be discovered and collaborate with other AI agents.
|
||||
|
||||
### A2A Server (Expose Agent as A2A Server)
|
||||
```python
|
||||
from praisonaiagents import Agent, A2A
|
||||
from fastapi import FastAPI
|
||||
|
||||
# Create an agent with tools
|
||||
def search_web(query: str) -> str:
|
||||
"""Search the web for information."""
|
||||
return f"Results for: {query}"
|
||||
|
||||
agent = Agent(
|
||||
name="Research Assistant",
|
||||
role="Research Analyst",
|
||||
goal="Help users research topics",
|
||||
tools=[search_web]
|
||||
)
|
||||
|
||||
# Expose as A2A Server
|
||||
a2a = A2A(agent=agent, url="http://localhost:8000/a2a")
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(a2a.get_router())
|
||||
|
||||
# Run: uvicorn app:app --reload
|
||||
# Agent Card: GET /.well-known/agent.json
|
||||
# Status: GET /status
|
||||
```
|
||||
|
||||
### A2A Features
|
||||
| Feature | Description |
|
||||
|---------|-------------|
|
||||
| Agent Card | JSON metadata for agent discovery |
|
||||
| Skills Extraction | Auto-generate skills from tools |
|
||||
| Task Management | Stateful task lifecycle |
|
||||
| Streaming | SSE streaming for real-time updates |
|
||||
|
||||
> **Documentation**: [docs.praison.ai/a2a](https://docs.praison.ai/a2a) | **Examples**: [examples/python/a2a](https://github.com/MervinPraison/PraisonAI/tree/main/examples/python/a2a)
|
||||
|
||||
## CLI Features
|
||||
|
||||
| Feature | Docs |
|
||||
@@ -2218,6 +2262,13 @@ PraisonAI provides zero-dependency persistent memory for agents. For detailed ex
|
||||
- 🔐 **MCP Security** - Origin validation, DNS rebinding prevention, secure sessions
|
||||
- 🔄 **MCP Resumability** - SSE stream recovery via Last-Event-ID
|
||||
|
||||
### A2A & A2UI Protocols
|
||||
|
||||
- 🔗 **A2A Protocol** - Agent-to-Agent communication for inter-agent collaboration
|
||||
- 🖼️ **A2UI Protocol** - Agent-to-User Interface for generating rich UIs from agents
|
||||
- 📋 **UI Templates** - ChatTemplate, ListTemplate, FormTemplate, DashboardTemplate
|
||||
- 🔧 **Surface Builder** - Fluent API for building declarative UIs
|
||||
|
||||
### Safety & Control
|
||||
|
||||
- 🤝 **Agent Handoffs** - Transfer context between specialised agents
|
||||
|
||||
@@ -3,8 +3,8 @@ class Praisonai < Formula
|
||||
|
||||
desc "AI tools for various AI applications"
|
||||
homepage "https://github.com/MervinPraison/PraisonAI"
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.36.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.36.tar.gz | shasum -a 256`.split.first
|
||||
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.41.tar.gz"
|
||||
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.41.tar.gz | shasum -a 256`.split.first
|
||||
license "MIT"
|
||||
|
||||
depends_on "[email protected]"
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
PraisonAI Code - AI-powered code editing tools.
|
||||
|
||||
This module provides tools for AI agents to read, write, and modify code files
|
||||
with precision using search/replace diff strategies similar to Kilo Code.
|
||||
|
||||
Features:
|
||||
- read_file: Read file contents with optional line ranges
|
||||
- write_file: Create or overwrite files
|
||||
- list_files: List directory contents with filtering
|
||||
- apply_diff: Apply SEARCH/REPLACE diffs with fuzzy matching
|
||||
- search_replace: Multiple search/replace operations
|
||||
- execute_command: Run shell commands safely
|
||||
|
||||
Example:
|
||||
from praisonai.code import read_file, write_file, apply_diff
|
||||
|
||||
# Read a file
|
||||
content = read_file("path/to/file.py")
|
||||
|
||||
# Write a file
|
||||
write_file("path/to/new_file.py", "print('hello')")
|
||||
|
||||
# Apply a diff
|
||||
result = apply_diff("path/to/file.py", diff_content)
|
||||
"""
|
||||
|
||||
from .tools import (
|
||||
read_file,
|
||||
write_file,
|
||||
list_files,
|
||||
apply_diff,
|
||||
search_replace,
|
||||
execute_command,
|
||||
)
|
||||
|
||||
from .diff import (
|
||||
DiffResult,
|
||||
apply_search_replace_diff,
|
||||
parse_diff_blocks,
|
||||
)
|
||||
|
||||
from .utils import (
|
||||
FileAccessController,
|
||||
add_line_numbers,
|
||||
strip_line_numbers,
|
||||
)
|
||||
|
||||
from .agent_tools import (
|
||||
set_workspace,
|
||||
get_workspace,
|
||||
code_read_file,
|
||||
code_write_file,
|
||||
code_list_files,
|
||||
code_apply_diff,
|
||||
code_search_replace,
|
||||
code_execute_command,
|
||||
CODE_TOOLS,
|
||||
create_diff_block,
|
||||
create_multi_diff,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Main tools (low-level)
|
||||
"read_file",
|
||||
"write_file",
|
||||
"list_files",
|
||||
"apply_diff",
|
||||
"search_replace",
|
||||
"execute_command",
|
||||
# Agent tools (high-level, string returns)
|
||||
"set_workspace",
|
||||
"get_workspace",
|
||||
"code_read_file",
|
||||
"code_write_file",
|
||||
"code_list_files",
|
||||
"code_apply_diff",
|
||||
"code_search_replace",
|
||||
"code_execute_command",
|
||||
"CODE_TOOLS",
|
||||
"create_diff_block",
|
||||
"create_multi_diff",
|
||||
# Diff utilities
|
||||
"DiffResult",
|
||||
"apply_search_replace_diff",
|
||||
"parse_diff_blocks",
|
||||
# Utilities
|
||||
"FileAccessController",
|
||||
"add_line_numbers",
|
||||
"strip_line_numbers",
|
||||
]
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
Agent-compatible tool wrappers for PraisonAI Code.
|
||||
|
||||
These functions are designed to be used as tools with PraisonAI agents.
|
||||
They return string results suitable for agent consumption.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from .tools import (
|
||||
read_file as _read_file,
|
||||
write_file as _write_file,
|
||||
list_files as _list_files,
|
||||
apply_diff as _apply_diff,
|
||||
search_replace as _search_replace,
|
||||
execute_command as _execute_command,
|
||||
)
|
||||
from .tools.apply_diff import create_diff_block, create_multi_diff
|
||||
|
||||
|
||||
# Global workspace setting for tools
|
||||
_workspace_root: Optional[str] = None
|
||||
|
||||
|
||||
def set_workspace(workspace_path: str) -> None:
|
||||
"""
|
||||
Set the workspace root for all code tools.
|
||||
|
||||
Args:
|
||||
workspace_path: Absolute path to the workspace root
|
||||
"""
|
||||
global _workspace_root
|
||||
_workspace_root = os.path.abspath(workspace_path)
|
||||
|
||||
|
||||
def get_workspace() -> Optional[str]:
|
||||
"""Get the current workspace root."""
|
||||
return _workspace_root
|
||||
|
||||
|
||||
def code_read_file(
|
||||
path: str,
|
||||
start_line: Optional[int] = None,
|
||||
end_line: Optional[int] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Read the contents of a file.
|
||||
|
||||
Use this tool to read source code files. The content is returned with
|
||||
line numbers for easy reference when making edits.
|
||||
|
||||
Args:
|
||||
path: Path to the file (relative to workspace or absolute)
|
||||
start_line: First line to read (1-indexed, optional)
|
||||
end_line: Last line to read (1-indexed, optional)
|
||||
|
||||
Returns:
|
||||
File content with line numbers, or error message
|
||||
|
||||
Example:
|
||||
>>> content = code_read_file("src/main.py")
|
||||
>>> content = code_read_file("src/main.py", start_line=10, end_line=50)
|
||||
"""
|
||||
result = _read_file(
|
||||
path=path,
|
||||
start_line=start_line,
|
||||
end_line=end_line,
|
||||
add_line_nums=True,
|
||||
workspace=_workspace_root,
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
header = f"File: {path}"
|
||||
if start_line or end_line:
|
||||
header += f" (lines {result['start_line']}-{result['end_line']} of {result['total_lines']})"
|
||||
else:
|
||||
header += f" ({result['total_lines']} lines)"
|
||||
return f"{header}\n\n{result['content']}"
|
||||
else:
|
||||
return f"Error reading {path}: {result['error']}"
|
||||
|
||||
|
||||
def code_write_file(path: str, content: str) -> str:
|
||||
"""
|
||||
Write content to a file, creating it if it doesn't exist.
|
||||
|
||||
Use this tool to create new files or completely replace existing files.
|
||||
For partial modifications, use code_apply_diff instead.
|
||||
|
||||
Args:
|
||||
path: Path to the file (relative to workspace or absolute)
|
||||
content: Content to write to the file
|
||||
|
||||
Returns:
|
||||
Success message or error
|
||||
|
||||
Example:
|
||||
>>> result = code_write_file("src/new_module.py", "def hello():\\n print('Hello!')")
|
||||
"""
|
||||
result = _write_file(
|
||||
path=path,
|
||||
content=content,
|
||||
workspace=_workspace_root,
|
||||
create_directories=True,
|
||||
strip_code_fences=True,
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
action = "Created" if result['created'] else "Updated"
|
||||
return f"{action} file: {path} ({result['bytes_written']} bytes)"
|
||||
else:
|
||||
return f"Error writing {path}: {result['error']}"
|
||||
|
||||
|
||||
def code_list_files(
|
||||
path: str = ".",
|
||||
recursive: bool = False,
|
||||
extensions: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
List files in a directory.
|
||||
|
||||
Use this tool to explore the codebase structure.
|
||||
|
||||
Args:
|
||||
path: Directory path (relative to workspace or absolute)
|
||||
recursive: Whether to list recursively
|
||||
extensions: Comma-separated list of extensions to filter (e.g., "py,js,ts")
|
||||
|
||||
Returns:
|
||||
Formatted list of files and directories
|
||||
|
||||
Example:
|
||||
>>> files = code_list_files("src", recursive=True, extensions="py")
|
||||
"""
|
||||
ext_list = None
|
||||
if extensions:
|
||||
ext_list = [e.strip() for e in extensions.split(',')]
|
||||
|
||||
result = _list_files(
|
||||
path=path,
|
||||
recursive=recursive,
|
||||
workspace=_workspace_root,
|
||||
extensions=ext_list,
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
lines = [f"Contents of {path}:"]
|
||||
|
||||
# Add directories
|
||||
for d in result['directories']:
|
||||
lines.append(f" 📁 {d['path']}/")
|
||||
|
||||
# Add files
|
||||
for f in result['files']:
|
||||
size = f.get('size', 0)
|
||||
if size < 1024:
|
||||
size_str = f"{size}B"
|
||||
elif size < 1024 * 1024:
|
||||
size_str = f"{size / 1024:.1f}KB"
|
||||
else:
|
||||
size_str = f"{size / (1024 * 1024):.1f}MB"
|
||||
lines.append(f" 📄 {f['path']} ({size_str})")
|
||||
|
||||
if result.get('truncated'):
|
||||
lines.append(f"\n ... (showing {result['total_count']} items, more available)")
|
||||
|
||||
return '\n'.join(lines)
|
||||
else:
|
||||
return f"Error listing {path}: {result['error']}"
|
||||
|
||||
|
||||
def code_apply_diff(path: str, diff: str) -> str:
|
||||
"""
|
||||
Apply a SEARCH/REPLACE diff to modify a file.
|
||||
|
||||
Use this tool for precise, surgical edits to existing files.
|
||||
The SEARCH section must match existing content exactly.
|
||||
|
||||
Diff Format:
|
||||
<<<<<<< SEARCH
|
||||
:start_line:N
|
||||
-------
|
||||
[exact content to find]
|
||||
=======
|
||||
[new content to replace with]
|
||||
>>>>>>> REPLACE
|
||||
|
||||
Args:
|
||||
path: Path to the file to modify
|
||||
diff: The SEARCH/REPLACE diff content
|
||||
|
||||
Returns:
|
||||
Success message with number of changes, or error details
|
||||
|
||||
Example:
|
||||
>>> diff = '''<<<<<<< SEARCH
|
||||
... :start_line:10
|
||||
... -------
|
||||
... def old_function():
|
||||
... pass
|
||||
... =======
|
||||
... def new_function():
|
||||
... return True
|
||||
... >>>>>>> REPLACE'''
|
||||
>>> result = code_apply_diff("src/main.py", diff)
|
||||
"""
|
||||
result = _apply_diff(
|
||||
path=path,
|
||||
diff=diff,
|
||||
workspace=_workspace_root,
|
||||
fuzzy_threshold=0.95, # Allow slight fuzzy matching
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
msg = f"Successfully applied {result['applied_count']} change(s) to {path}"
|
||||
if result.get('failed_blocks'):
|
||||
msg += f"\nWarning: {len(result['failed_blocks'])} block(s) failed to apply"
|
||||
return msg
|
||||
else:
|
||||
error_msg = f"Error applying diff to {path}: {result['error']}"
|
||||
if result.get('failed_blocks'):
|
||||
for fb in result['failed_blocks'][:3]: # Show first 3 failures
|
||||
error_msg += f"\n- {fb.get('error', 'Unknown error')}"
|
||||
return error_msg
|
||||
|
||||
|
||||
def code_search_replace(
|
||||
path: str,
|
||||
search: str,
|
||||
replace: str,
|
||||
is_regex: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Perform a search and replace operation on a file.
|
||||
|
||||
Use this for simple text replacements. For complex multi-part edits,
|
||||
use code_apply_diff instead.
|
||||
|
||||
Args:
|
||||
path: Path to the file
|
||||
search: Text to search for (or regex pattern if is_regex=True)
|
||||
replace: Text to replace with
|
||||
is_regex: Whether to treat search as a regex pattern
|
||||
|
||||
Returns:
|
||||
Success message with replacement count, or error
|
||||
|
||||
Example:
|
||||
>>> result = code_search_replace("src/main.py", "old_name", "new_name")
|
||||
"""
|
||||
result = _search_replace(
|
||||
path=path,
|
||||
operations=[{
|
||||
'search': search,
|
||||
'replace': replace,
|
||||
'is_regex': is_regex,
|
||||
}],
|
||||
workspace=_workspace_root,
|
||||
)
|
||||
|
||||
if result['success']:
|
||||
return f"Replaced {result['total_replacements']} occurrence(s) in {path}"
|
||||
else:
|
||||
return f"Error: {result['error']}"
|
||||
|
||||
|
||||
def code_execute_command(
|
||||
command: str,
|
||||
cwd: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Execute a shell command.
|
||||
|
||||
Use this tool to run commands like tests, linters, or build tools.
|
||||
|
||||
Args:
|
||||
command: The command to execute
|
||||
cwd: Working directory (relative to workspace or absolute)
|
||||
|
||||
Returns:
|
||||
Command output (stdout and stderr)
|
||||
|
||||
Example:
|
||||
>>> result = code_execute_command("python -m pytest tests/")
|
||||
>>> result = code_execute_command("npm test", cwd="frontend")
|
||||
"""
|
||||
work_dir = cwd
|
||||
if work_dir and _workspace_root and not os.path.isabs(work_dir):
|
||||
work_dir = os.path.join(_workspace_root, work_dir)
|
||||
elif not work_dir and _workspace_root:
|
||||
work_dir = _workspace_root
|
||||
|
||||
result = _execute_command(
|
||||
command=command,
|
||||
cwd=work_dir,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
output_parts = []
|
||||
|
||||
if result['stdout']:
|
||||
output_parts.append(f"stdout:\n{result['stdout']}")
|
||||
|
||||
if result['stderr']:
|
||||
output_parts.append(f"stderr:\n{result['stderr']}")
|
||||
|
||||
if result['success']:
|
||||
status = f"Command completed successfully (exit code: {result['exit_code']})"
|
||||
else:
|
||||
if result.get('error'):
|
||||
status = f"Command failed: {result['error']}"
|
||||
else:
|
||||
status = f"Command failed (exit code: {result['exit_code']})"
|
||||
|
||||
output_parts.insert(0, status)
|
||||
|
||||
return '\n\n'.join(output_parts)
|
||||
|
||||
|
||||
# Export all agent tools
|
||||
CODE_TOOLS = [
|
||||
code_read_file,
|
||||
code_write_file,
|
||||
code_list_files,
|
||||
code_apply_diff,
|
||||
code_search_replace,
|
||||
code_execute_command,
|
||||
]
|
||||
|
||||
__all__ = [
|
||||
'set_workspace',
|
||||
'get_workspace',
|
||||
'code_read_file',
|
||||
'code_write_file',
|
||||
'code_list_files',
|
||||
'code_apply_diff',
|
||||
'code_search_replace',
|
||||
'code_execute_command',
|
||||
'CODE_TOOLS',
|
||||
'create_diff_block',
|
||||
'create_multi_diff',
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
Diff handling for PraisonAI Code module.
|
||||
|
||||
Provides SEARCH/REPLACE diff parsing and application with fuzzy matching.
|
||||
"""
|
||||
|
||||
from .diff_strategy import (
|
||||
DiffResult,
|
||||
DiffBlock,
|
||||
apply_search_replace_diff,
|
||||
parse_diff_blocks,
|
||||
validate_diff_format,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"DiffResult",
|
||||
"DiffBlock",
|
||||
"apply_search_replace_diff",
|
||||
"parse_diff_blocks",
|
||||
"validate_diff_format",
|
||||
]
|
||||
@@ -0,0 +1,432 @@
|
||||
"""
|
||||
Multi-Search-Replace Diff Strategy for PraisonAI Code.
|
||||
|
||||
This module implements a diff strategy that applies SEARCH/REPLACE blocks
|
||||
to file content with fuzzy matching support, similar to Kilo Code's approach.
|
||||
|
||||
Diff Format:
|
||||
<<<<<<< SEARCH
|
||||
:start_line:N
|
||||
-------
|
||||
[exact content to find]
|
||||
=======
|
||||
[new content to replace with]
|
||||
>>>>>>> REPLACE
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from ..utils.text_utils import (
|
||||
get_similarity,
|
||||
fuzzy_search,
|
||||
get_indentation,
|
||||
)
|
||||
from ..utils.file_utils import (
|
||||
strip_line_numbers,
|
||||
every_line_has_line_numbers,
|
||||
add_line_numbers,
|
||||
detect_line_ending,
|
||||
)
|
||||
|
||||
|
||||
# Default buffer lines for fuzzy search around start_line hint
|
||||
BUFFER_LINES = 40
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffBlock:
|
||||
"""
|
||||
Represents a single SEARCH/REPLACE block.
|
||||
|
||||
Attributes:
|
||||
search_content: The content to search for
|
||||
replace_content: The content to replace with
|
||||
start_line: Optional line number hint (1-indexed)
|
||||
end_line: Optional end line hint (1-indexed)
|
||||
"""
|
||||
search_content: str
|
||||
replace_content: str
|
||||
start_line: Optional[int] = None
|
||||
end_line: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffResult:
|
||||
"""
|
||||
Result of applying a diff operation.
|
||||
|
||||
Attributes:
|
||||
success: Whether the diff was applied successfully
|
||||
content: The resulting content (if successful)
|
||||
error: Error message (if failed)
|
||||
applied_count: Number of blocks successfully applied
|
||||
failed_blocks: List of failed block results
|
||||
"""
|
||||
success: bool
|
||||
content: Optional[str] = None
|
||||
error: Optional[str] = None
|
||||
applied_count: int = 0
|
||||
failed_blocks: List[dict] = field(default_factory=list)
|
||||
|
||||
|
||||
def _unescape_markers(content: str) -> str:
|
||||
"""
|
||||
Unescape escaped diff markers in content.
|
||||
|
||||
Args:
|
||||
content: Content with potentially escaped markers
|
||||
|
||||
Returns:
|
||||
Content with markers unescaped
|
||||
"""
|
||||
replacements = [
|
||||
(r'^\\\<\<\<\<\<\<\<', '<<<<<<<'),
|
||||
(r'^\\=======', '======='),
|
||||
(r'^\\\>\>\>\>\>\>\>', '>>>>>>>'),
|
||||
(r'^\\-------', '-------'),
|
||||
(r'^\\:end_line:', ':end_line:'),
|
||||
(r'^\\:start_line:', ':start_line:'),
|
||||
]
|
||||
|
||||
result = content
|
||||
for pattern, replacement in replacements:
|
||||
result = re.sub(pattern, replacement, result, flags=re.MULTILINE)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def validate_diff_format(diff_content: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
Validate that a diff string has correct marker sequencing.
|
||||
|
||||
Args:
|
||||
diff_content: The diff content to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
# State machine states
|
||||
STATE_START = 0
|
||||
STATE_AFTER_SEARCH = 1
|
||||
STATE_AFTER_SEPARATOR = 2
|
||||
|
||||
state = STATE_START
|
||||
line_num = 0
|
||||
|
||||
SEARCH_PATTERN = re.compile(r'^<<<<<<< SEARCH>?$')
|
||||
SEP = '======='
|
||||
REPLACE = '>>>>>>> REPLACE'
|
||||
|
||||
lines = diff_content.split('\n')
|
||||
|
||||
for line in lines:
|
||||
line_num += 1
|
||||
marker = line.strip()
|
||||
|
||||
# Check for line markers in REPLACE sections
|
||||
if state == STATE_AFTER_SEPARATOR:
|
||||
if marker.startswith(':start_line:') and not line.strip().startswith('\\:start_line:'):
|
||||
return False, f"ERROR: Line marker ':start_line:' found in REPLACE section at line {line_num}"
|
||||
if marker.startswith(':end_line:') and not line.strip().startswith('\\:end_line:'):
|
||||
return False, f"ERROR: Line marker ':end_line:' found in REPLACE section at line {line_num}"
|
||||
|
||||
if state == STATE_START:
|
||||
if marker == SEP:
|
||||
return False, f"ERROR: Unexpected '=======' at line {line_num}, expected '<<<<<<< SEARCH'"
|
||||
if marker == REPLACE:
|
||||
return False, f"ERROR: Unexpected '>>>>>>> REPLACE' at line {line_num}"
|
||||
if SEARCH_PATTERN.match(marker):
|
||||
state = STATE_AFTER_SEARCH
|
||||
|
||||
elif state == STATE_AFTER_SEARCH:
|
||||
if SEARCH_PATTERN.match(marker):
|
||||
return False, f"ERROR: Duplicate '<<<<<<< SEARCH' at line {line_num}"
|
||||
if marker == REPLACE:
|
||||
return False, f"ERROR: Missing '=======' before '>>>>>>> REPLACE' at line {line_num}"
|
||||
if marker == SEP:
|
||||
state = STATE_AFTER_SEPARATOR
|
||||
|
||||
elif state == STATE_AFTER_SEPARATOR:
|
||||
if SEARCH_PATTERN.match(marker):
|
||||
return False, f"ERROR: Unexpected '<<<<<<< SEARCH' at line {line_num}, expected '>>>>>>> REPLACE'"
|
||||
if marker == SEP:
|
||||
return False, f"ERROR: Duplicate '=======' at line {line_num}"
|
||||
if marker == REPLACE:
|
||||
state = STATE_START
|
||||
|
||||
if state != STATE_START:
|
||||
expected = "'======='" if state == STATE_AFTER_SEARCH else "'>>>>>>> REPLACE'"
|
||||
return False, f"ERROR: Unexpected end of diff, expected {expected}"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
def parse_diff_blocks(diff_content: str) -> Tuple[List[DiffBlock], Optional[str]]:
|
||||
"""
|
||||
Parse SEARCH/REPLACE blocks from diff content.
|
||||
|
||||
Args:
|
||||
diff_content: The diff content to parse
|
||||
|
||||
Returns:
|
||||
Tuple of (list of DiffBlock, error message if any)
|
||||
"""
|
||||
# Validate format first
|
||||
is_valid, error = validate_diff_format(diff_content)
|
||||
if not is_valid:
|
||||
return [], error
|
||||
|
||||
# Regex to match SEARCH/REPLACE blocks
|
||||
# Groups: 1=start_line line, 2=start_line number, 3=end_line line, 4=end_line number,
|
||||
# 5=separator line, 6=search content, 7=replace content
|
||||
pattern = re.compile(
|
||||
r'(?:^|\n)(?<!\\)<<<<<<< SEARCH>?\s*\n'
|
||||
r'((?::start_line:\s*(\d+)\s*\n))?'
|
||||
r'((?::end_line:\s*(\d+)\s*\n))?'
|
||||
r'((?<!\\)-------\s*\n)?'
|
||||
r'([\s\S]*?)(?:\n)?'
|
||||
r'(?:(?<=\n)(?<!\\)=======\s*\n)'
|
||||
r'([\s\S]*?)(?:\n)?'
|
||||
r'(?:(?<=\n)(?<!\\)>>>>>>> REPLACE)(?=\n|$)',
|
||||
re.MULTILINE
|
||||
)
|
||||
|
||||
matches = list(pattern.finditer(diff_content))
|
||||
|
||||
if not matches:
|
||||
return [], "Invalid diff format - no valid SEARCH/REPLACE blocks found"
|
||||
|
||||
blocks = []
|
||||
for match in matches:
|
||||
start_line = int(match.group(2)) if match.group(2) else None
|
||||
end_line = int(match.group(4)) if match.group(4) else None
|
||||
search_content = match.group(6) or ""
|
||||
replace_content = match.group(7) or ""
|
||||
|
||||
blocks.append(DiffBlock(
|
||||
search_content=search_content,
|
||||
replace_content=replace_content,
|
||||
start_line=start_line,
|
||||
end_line=end_line,
|
||||
))
|
||||
|
||||
# Sort by start_line (blocks without start_line go last)
|
||||
blocks.sort(key=lambda b: b.start_line if b.start_line else float('inf'))
|
||||
|
||||
return blocks, None
|
||||
|
||||
|
||||
def apply_search_replace_diff(
|
||||
original_content: str,
|
||||
diff_content: str,
|
||||
fuzzy_threshold: float = 1.0,
|
||||
buffer_lines: int = BUFFER_LINES,
|
||||
) -> DiffResult:
|
||||
"""
|
||||
Apply SEARCH/REPLACE diff blocks to original content.
|
||||
|
||||
Args:
|
||||
original_content: The original file content
|
||||
diff_content: The diff content with SEARCH/REPLACE blocks
|
||||
fuzzy_threshold: Similarity threshold (0.0-1.0, 1.0 = exact match)
|
||||
buffer_lines: Number of lines to search around start_line hint
|
||||
|
||||
Returns:
|
||||
DiffResult with success status and resulting content
|
||||
"""
|
||||
# Parse diff blocks
|
||||
blocks, parse_error = parse_diff_blocks(diff_content)
|
||||
if parse_error:
|
||||
return DiffResult(success=False, error=parse_error)
|
||||
|
||||
if not blocks:
|
||||
return DiffResult(
|
||||
success=False,
|
||||
error="No valid SEARCH/REPLACE blocks found in diff"
|
||||
)
|
||||
|
||||
# Detect line ending from original content
|
||||
line_ending = detect_line_ending(original_content)
|
||||
result_lines = original_content.split('\n')
|
||||
if original_content.endswith('\n'):
|
||||
# Handle trailing newline
|
||||
if result_lines and result_lines[-1] == '':
|
||||
result_lines = result_lines[:-1]
|
||||
|
||||
delta = 0 # Track line number changes from previous replacements
|
||||
applied_count = 0
|
||||
failed_blocks = []
|
||||
|
||||
for block in blocks:
|
||||
search_content = _unescape_markers(block.search_content)
|
||||
replace_content = _unescape_markers(block.replace_content)
|
||||
|
||||
# Handle line numbers in content
|
||||
has_all_line_numbers = (
|
||||
(every_line_has_line_numbers(search_content) and every_line_has_line_numbers(replace_content)) or
|
||||
(every_line_has_line_numbers(search_content) and not replace_content.strip())
|
||||
)
|
||||
|
||||
start_line = block.start_line
|
||||
if has_all_line_numbers and not start_line:
|
||||
# Extract start line from first line number in search content
|
||||
first_line = search_content.split('\n')[0]
|
||||
match = re.match(r'\s*(\d+)', first_line)
|
||||
if match:
|
||||
start_line = int(match.group(1))
|
||||
|
||||
if has_all_line_numbers:
|
||||
search_content = strip_line_numbers(search_content)
|
||||
replace_content = strip_line_numbers(replace_content)
|
||||
|
||||
# Validate search content
|
||||
if search_content == replace_content:
|
||||
failed_blocks.append({
|
||||
'error': 'Search and replace content are identical - no changes would be made',
|
||||
'search': search_content[:100],
|
||||
})
|
||||
continue
|
||||
|
||||
search_lines = search_content.split('\n') if search_content else []
|
||||
replace_lines = replace_content.split('\n') if replace_content else []
|
||||
|
||||
if not search_lines or (len(search_lines) == 1 and not search_lines[0]):
|
||||
failed_blocks.append({
|
||||
'error': 'Empty search content is not allowed',
|
||||
})
|
||||
continue
|
||||
|
||||
# Apply delta to start_line
|
||||
adjusted_start = start_line + delta if start_line else None
|
||||
|
||||
# Initialize search
|
||||
match_index = -1
|
||||
best_score = 0.0
|
||||
best_match_content = ""
|
||||
search_chunk = '\n'.join(search_lines)
|
||||
|
||||
# Determine search bounds
|
||||
search_start_index = 0
|
||||
search_end_index = len(result_lines)
|
||||
|
||||
if adjusted_start:
|
||||
# Try exact match at start_line first
|
||||
exact_start_index = adjusted_start - 1 # Convert to 0-indexed
|
||||
exact_end_index = exact_start_index + len(search_lines)
|
||||
|
||||
if 0 <= exact_start_index < len(result_lines):
|
||||
original_chunk = '\n'.join(result_lines[exact_start_index:exact_end_index])
|
||||
similarity = get_similarity(original_chunk, search_chunk)
|
||||
|
||||
if similarity >= fuzzy_threshold:
|
||||
match_index = exact_start_index
|
||||
best_score = similarity
|
||||
best_match_content = original_chunk
|
||||
else:
|
||||
# Set bounds for buffered search
|
||||
search_start_index = max(0, adjusted_start - buffer_lines - 1)
|
||||
search_end_index = min(len(result_lines), adjusted_start + len(search_lines) + buffer_lines)
|
||||
|
||||
# If no match found yet, try fuzzy search
|
||||
if match_index == -1:
|
||||
score, idx, content = fuzzy_search(
|
||||
result_lines, search_chunk, search_start_index, search_end_index
|
||||
)
|
||||
match_index = idx
|
||||
best_score = score
|
||||
best_match_content = content
|
||||
|
||||
# Check if match meets threshold
|
||||
if match_index == -1 or best_score < fuzzy_threshold:
|
||||
# Try aggressive line number stripping as fallback
|
||||
aggressive_search = strip_line_numbers(search_content, aggressive=True)
|
||||
aggressive_replace = strip_line_numbers(replace_content, aggressive=True)
|
||||
|
||||
if aggressive_search != search_content:
|
||||
aggressive_lines = aggressive_search.split('\n') if aggressive_search else []
|
||||
aggressive_chunk = '\n'.join(aggressive_lines)
|
||||
|
||||
score, idx, content = fuzzy_search(
|
||||
result_lines, aggressive_chunk, search_start_index, search_end_index
|
||||
)
|
||||
|
||||
if idx != -1 and score >= fuzzy_threshold:
|
||||
match_index = idx
|
||||
best_score = score
|
||||
best_match_content = content
|
||||
search_content = aggressive_search
|
||||
replace_content = aggressive_replace
|
||||
search_lines = aggressive_lines
|
||||
replace_lines = replace_content.split('\n') if replace_content else []
|
||||
|
||||
# Still no match - report error
|
||||
if match_index == -1 or best_score < fuzzy_threshold:
|
||||
line_info = f" at line {start_line}" if start_line else ""
|
||||
|
||||
# Build context for error message
|
||||
context_start = max(0, (adjusted_start or 1) - 1 - buffer_lines)
|
||||
context_end = min(len(result_lines), (adjusted_start or 1) + len(search_lines) + buffer_lines)
|
||||
original_context = '\n'.join(result_lines[context_start:context_end])
|
||||
|
||||
failed_blocks.append({
|
||||
'error': f"No sufficiently similar match found{line_info} ({int(best_score * 100)}% similar, needs {int(fuzzy_threshold * 100)}%)",
|
||||
'search_content': search_chunk[:200],
|
||||
'best_match': best_match_content[:200] if best_match_content else None,
|
||||
'similarity': best_score,
|
||||
'context': add_line_numbers(original_context, context_start + 1)[:500],
|
||||
})
|
||||
continue
|
||||
|
||||
# Apply the replacement with indentation preservation
|
||||
matched_lines = result_lines[match_index:match_index + len(search_lines)]
|
||||
|
||||
# Get indentation from original
|
||||
original_indents = [get_indentation(line) for line in matched_lines]
|
||||
search_indents = [get_indentation(line) for line in search_lines]
|
||||
|
||||
# Apply replacement with preserved indentation
|
||||
indented_replace_lines = []
|
||||
for i, line in enumerate(replace_lines):
|
||||
matched_indent = original_indents[0] if original_indents else ""
|
||||
current_indent = get_indentation(line)
|
||||
search_base_indent = search_indents[0] if search_indents else ""
|
||||
|
||||
# Calculate relative indentation
|
||||
search_base_level = len(search_base_indent)
|
||||
current_level = len(current_indent)
|
||||
relative_level = current_level - search_base_level
|
||||
|
||||
if relative_level < 0:
|
||||
final_indent = matched_indent[:max(0, len(matched_indent) + relative_level)]
|
||||
else:
|
||||
final_indent = matched_indent + current_indent[search_base_level:]
|
||||
|
||||
indented_replace_lines.append(final_indent + line.lstrip())
|
||||
|
||||
# Construct final content
|
||||
before_match = result_lines[:match_index]
|
||||
after_match = result_lines[match_index + len(search_lines):]
|
||||
result_lines = before_match + indented_replace_lines + after_match
|
||||
|
||||
# Update delta
|
||||
delta += len(replace_lines) - len(matched_lines)
|
||||
applied_count += 1
|
||||
|
||||
# Build final content
|
||||
final_content = line_ending.join(result_lines)
|
||||
|
||||
if applied_count == 0:
|
||||
return DiffResult(
|
||||
success=False,
|
||||
error="No blocks were successfully applied",
|
||||
failed_blocks=failed_blocks,
|
||||
)
|
||||
|
||||
return DiffResult(
|
||||
success=True,
|
||||
content=final_content,
|
||||
applied_count=applied_count,
|
||||
failed_blocks=failed_blocks,
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
Code editing tools for PraisonAI agents.
|
||||
|
||||
Provides tools for:
|
||||
- read_file: Read file contents with optional line ranges
|
||||
- write_file: Create or overwrite files
|
||||
- list_files: List directory contents
|
||||
- apply_diff: Apply SEARCH/REPLACE diffs
|
||||
- search_replace: Multiple search/replace operations
|
||||
- execute_command: Run shell commands
|
||||
"""
|
||||
|
||||
from .read_file import read_file
|
||||
from .write_file import write_file
|
||||
from .list_files import list_files
|
||||
from .apply_diff import apply_diff
|
||||
from .search_replace import search_replace
|
||||
from .execute_command import execute_command
|
||||
|
||||
__all__ = [
|
||||
"read_file",
|
||||
"write_file",
|
||||
"list_files",
|
||||
"apply_diff",
|
||||
"search_replace",
|
||||
"execute_command",
|
||||
]
|
||||
@@ -0,0 +1,221 @@
|
||||
"""
|
||||
Apply Diff Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to apply SEARCH/REPLACE diffs to files
|
||||
with fuzzy matching support.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from ..diff.diff_strategy import apply_search_replace_diff
|
||||
from ..utils.file_utils import (
|
||||
file_exists,
|
||||
is_path_within_directory,
|
||||
)
|
||||
|
||||
|
||||
def apply_diff(
|
||||
path: str,
|
||||
diff: str,
|
||||
workspace: Optional[str] = None,
|
||||
fuzzy_threshold: float = 1.0,
|
||||
buffer_lines: int = 40,
|
||||
backup: bool = False,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Apply a SEARCH/REPLACE diff to a file.
|
||||
|
||||
This tool applies precise, targeted modifications to an existing file
|
||||
by searching for specific sections of content and replacing them.
|
||||
|
||||
Diff Format:
|
||||
<<<<<<< SEARCH
|
||||
:start_line:N
|
||||
-------
|
||||
[exact content to find]
|
||||
=======
|
||||
[new content to replace with]
|
||||
>>>>>>> REPLACE
|
||||
|
||||
Args:
|
||||
path: Path to the file (absolute or relative to workspace)
|
||||
diff: The SEARCH/REPLACE diff content
|
||||
workspace: Workspace root directory (for relative paths)
|
||||
fuzzy_threshold: Similarity threshold (0.0-1.0, 1.0 = exact match)
|
||||
buffer_lines: Lines to search around start_line hint
|
||||
backup: Whether to create a backup before modifying
|
||||
encoding: File encoding (default: utf-8)
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- path: str
|
||||
- applied_count: int (number of blocks applied)
|
||||
- failed_blocks: list (details of failed blocks)
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> diff = '''
|
||||
... <<<<<<< SEARCH
|
||||
... :start_line:1
|
||||
... -------
|
||||
... def old_function():
|
||||
... pass
|
||||
... =======
|
||||
... def new_function():
|
||||
... return True
|
||||
... >>>>>>> REPLACE
|
||||
... '''
|
||||
>>> result = apply_diff("src/main.py", diff)
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Security check
|
||||
if workspace:
|
||||
if not is_path_within_directory(abs_path, workspace):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Path '{path}' is outside the workspace",
|
||||
'path': path,
|
||||
'applied_count': 0,
|
||||
}
|
||||
|
||||
# Check if file exists
|
||||
if not file_exists(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"File not found: {path}",
|
||||
'path': path,
|
||||
'applied_count': 0,
|
||||
}
|
||||
|
||||
try:
|
||||
# Read original content
|
||||
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
|
||||
original_content = f.read()
|
||||
|
||||
# Apply the diff
|
||||
result = apply_search_replace_diff(
|
||||
original_content=original_content,
|
||||
diff_content=diff,
|
||||
fuzzy_threshold=fuzzy_threshold,
|
||||
buffer_lines=buffer_lines,
|
||||
)
|
||||
|
||||
if not result.success:
|
||||
return {
|
||||
'success': False,
|
||||
'error': result.error or "Failed to apply diff",
|
||||
'path': path,
|
||||
'applied_count': result.applied_count,
|
||||
'failed_blocks': result.failed_blocks,
|
||||
}
|
||||
|
||||
# Create backup if requested
|
||||
backup_path = None
|
||||
if backup:
|
||||
import time
|
||||
timestamp = int(time.time())
|
||||
backup_path = f"{abs_path}.backup.{timestamp}"
|
||||
with open(backup_path, 'w', encoding=encoding) as f:
|
||||
f.write(original_content)
|
||||
|
||||
# Write the modified content
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(result.content)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'path': path,
|
||||
'absolute_path': abs_path,
|
||||
'applied_count': result.applied_count,
|
||||
'failed_blocks': result.failed_blocks if result.failed_blocks else None,
|
||||
'backup_path': backup_path,
|
||||
}
|
||||
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Permission denied: {path}",
|
||||
'path': path,
|
||||
'applied_count': 0,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error applying diff to {path}: {str(e)}",
|
||||
'path': path,
|
||||
'applied_count': 0,
|
||||
}
|
||||
|
||||
|
||||
def create_diff_block(
|
||||
search_content: str,
|
||||
replace_content: str,
|
||||
start_line: Optional[int] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Create a properly formatted SEARCH/REPLACE diff block.
|
||||
|
||||
Args:
|
||||
search_content: The content to search for
|
||||
replace_content: The content to replace with
|
||||
start_line: Optional line number hint
|
||||
|
||||
Returns:
|
||||
Formatted diff block string
|
||||
|
||||
Example:
|
||||
>>> diff = create_diff_block(
|
||||
... "def old():\\n pass",
|
||||
... "def new():\\n return True",
|
||||
... start_line=10
|
||||
... )
|
||||
"""
|
||||
lines = ["<<<<<<< SEARCH"]
|
||||
|
||||
if start_line:
|
||||
lines.append(f":start_line:{start_line}")
|
||||
|
||||
lines.append("-------")
|
||||
lines.append(search_content)
|
||||
lines.append("=======")
|
||||
lines.append(replace_content)
|
||||
lines.append(">>>>>>> REPLACE")
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
def create_multi_diff(blocks: list) -> str:
|
||||
"""
|
||||
Create a diff with multiple SEARCH/REPLACE blocks.
|
||||
|
||||
Args:
|
||||
blocks: List of dicts with 'search', 'replace', and optional 'start_line'
|
||||
|
||||
Returns:
|
||||
Combined diff string with all blocks
|
||||
|
||||
Example:
|
||||
>>> diff = create_multi_diff([
|
||||
... {'search': 'old1', 'replace': 'new1', 'start_line': 5},
|
||||
... {'search': 'old2', 'replace': 'new2', 'start_line': 20},
|
||||
... ])
|
||||
"""
|
||||
diff_blocks = []
|
||||
|
||||
for block in blocks:
|
||||
diff_block = create_diff_block(
|
||||
search_content=block.get('search', ''),
|
||||
replace_content=block.get('replace', ''),
|
||||
start_line=block.get('start_line'),
|
||||
)
|
||||
diff_blocks.append(diff_block)
|
||||
|
||||
return '\n\n'.join(diff_blocks)
|
||||
@@ -0,0 +1,275 @@
|
||||
"""
|
||||
Execute Command Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to execute shell commands safely.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import shlex
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
|
||||
# Commands that are generally safe to auto-run
|
||||
SAFE_COMMANDS = {
|
||||
'ls', 'dir', 'pwd', 'echo', 'cat', 'head', 'tail', 'wc',
|
||||
'grep', 'find', 'which', 'whereis', 'type', 'file',
|
||||
'date', 'whoami', 'hostname', 'uname',
|
||||
'python', 'python3', 'node', 'npm', 'npx', 'pip', 'pip3',
|
||||
'git', 'cargo', 'go', 'ruby', 'java', 'javac',
|
||||
}
|
||||
|
||||
# Commands that should never be auto-run
|
||||
DANGEROUS_COMMANDS = {
|
||||
'rm', 'rmdir', 'del', 'format', 'mkfs',
|
||||
'dd', 'shred', 'chmod', 'chown', 'chgrp',
|
||||
'kill', 'killall', 'pkill',
|
||||
'shutdown', 'reboot', 'halt', 'poweroff',
|
||||
'sudo', 'su', 'doas',
|
||||
'curl', 'wget', 'ssh', 'scp', 'rsync',
|
||||
'mv', 'cp', # Can be destructive
|
||||
}
|
||||
|
||||
|
||||
def execute_command(
|
||||
command: str,
|
||||
cwd: Optional[str] = None,
|
||||
workspace: Optional[str] = None,
|
||||
timeout: int = 120,
|
||||
capture_output: bool = True,
|
||||
shell: bool = True,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute a shell command and return the result.
|
||||
|
||||
This tool executes shell commands with safety checks and output capture.
|
||||
|
||||
Args:
|
||||
command: The command to execute
|
||||
cwd: Working directory for the command
|
||||
workspace: Workspace root (for security validation)
|
||||
timeout: Command timeout in seconds
|
||||
capture_output: Whether to capture stdout/stderr
|
||||
shell: Whether to run through shell
|
||||
env: Additional environment variables
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- exit_code: int
|
||||
- stdout: str
|
||||
- stderr: str
|
||||
- command: str
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> result = execute_command("python --version")
|
||||
>>> print(result['stdout'])
|
||||
"""
|
||||
if not command or not command.strip():
|
||||
return {
|
||||
'success': False,
|
||||
'error': "Empty command",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
|
||||
# Resolve working directory
|
||||
if cwd:
|
||||
if workspace and not os.path.isabs(cwd):
|
||||
work_dir = os.path.abspath(os.path.join(workspace, cwd))
|
||||
else:
|
||||
work_dir = os.path.abspath(cwd)
|
||||
elif workspace:
|
||||
work_dir = workspace
|
||||
else:
|
||||
work_dir = os.getcwd()
|
||||
|
||||
# Validate working directory exists
|
||||
if not os.path.isdir(work_dir):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Working directory not found: {cwd}",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
|
||||
# Prepare environment
|
||||
cmd_env = os.environ.copy()
|
||||
if env:
|
||||
cmd_env.update(env)
|
||||
|
||||
# Set PAGER to cat to avoid interactive pagers
|
||||
cmd_env['PAGER'] = 'cat'
|
||||
|
||||
try:
|
||||
# Execute command
|
||||
if shell:
|
||||
result = subprocess.run(
|
||||
command,
|
||||
shell=True,
|
||||
cwd=work_dir,
|
||||
capture_output=capture_output,
|
||||
timeout=timeout,
|
||||
env=cmd_env,
|
||||
text=True,
|
||||
)
|
||||
else:
|
||||
# Parse command for non-shell execution
|
||||
args = shlex.split(command)
|
||||
result = subprocess.run(
|
||||
args,
|
||||
cwd=work_dir,
|
||||
capture_output=capture_output,
|
||||
timeout=timeout,
|
||||
env=cmd_env,
|
||||
text=True,
|
||||
)
|
||||
|
||||
return {
|
||||
'success': result.returncode == 0,
|
||||
'exit_code': result.returncode,
|
||||
'stdout': result.stdout if capture_output else '',
|
||||
'stderr': result.stderr if capture_output else '',
|
||||
'command': command,
|
||||
'cwd': work_dir,
|
||||
}
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Command timed out after {timeout} seconds",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
except FileNotFoundError as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Command not found: {str(e)}",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': "Permission denied",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error executing command: {str(e)}",
|
||||
'command': command,
|
||||
'exit_code': -1,
|
||||
'stdout': '',
|
||||
'stderr': '',
|
||||
}
|
||||
|
||||
|
||||
def is_safe_command(command: str) -> bool:
|
||||
"""
|
||||
Check if a command is considered safe to auto-run.
|
||||
|
||||
Args:
|
||||
command: The command to check
|
||||
|
||||
Returns:
|
||||
True if the command is considered safe
|
||||
"""
|
||||
if not command:
|
||||
return False
|
||||
|
||||
# Get the base command (first word)
|
||||
parts = shlex.split(command)
|
||||
if not parts:
|
||||
return False
|
||||
|
||||
base_cmd = os.path.basename(parts[0]).lower()
|
||||
|
||||
# Check against dangerous commands
|
||||
if base_cmd in DANGEROUS_COMMANDS:
|
||||
return False
|
||||
|
||||
# Check against safe commands
|
||||
if base_cmd in SAFE_COMMANDS:
|
||||
return True
|
||||
|
||||
# Default to unsafe
|
||||
return False
|
||||
|
||||
|
||||
def get_command_info(command: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get information about a command without executing it.
|
||||
|
||||
Args:
|
||||
command: The command to analyze
|
||||
|
||||
Returns:
|
||||
Dictionary with command analysis
|
||||
"""
|
||||
if not command:
|
||||
return {
|
||||
'valid': False,
|
||||
'error': "Empty command",
|
||||
}
|
||||
|
||||
try:
|
||||
parts = shlex.split(command)
|
||||
base_cmd = parts[0] if parts else ''
|
||||
|
||||
return {
|
||||
'valid': True,
|
||||
'base_command': base_cmd,
|
||||
'arguments': parts[1:] if len(parts) > 1 else [],
|
||||
'is_safe': is_safe_command(command),
|
||||
'is_dangerous': os.path.basename(base_cmd).lower() in DANGEROUS_COMMANDS,
|
||||
}
|
||||
except ValueError as e:
|
||||
return {
|
||||
'valid': False,
|
||||
'error': f"Invalid command syntax: {str(e)}",
|
||||
}
|
||||
|
||||
|
||||
def run_python(
|
||||
code: str,
|
||||
cwd: Optional[str] = None,
|
||||
timeout: int = 60,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Execute Python code and return the result.
|
||||
|
||||
Args:
|
||||
code: Python code to execute
|
||||
cwd: Working directory
|
||||
timeout: Execution timeout
|
||||
|
||||
Returns:
|
||||
Dictionary with execution result
|
||||
"""
|
||||
# Create a temporary command to run the code
|
||||
import sys
|
||||
python_cmd = sys.executable
|
||||
|
||||
# Escape the code for command line
|
||||
escaped_code = code.replace('\\', '\\\\').replace('"', '\\"')
|
||||
command = f'{python_cmd} -c "{escaped_code}"'
|
||||
|
||||
return execute_command(
|
||||
command=command,
|
||||
cwd=cwd,
|
||||
timeout=timeout,
|
||||
)
|
||||
@@ -0,0 +1,274 @@
|
||||
"""
|
||||
List Files Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to list directory contents with filtering options.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from ..utils.ignore_utils import (
|
||||
should_ignore_path,
|
||||
load_gitignore_patterns,
|
||||
)
|
||||
|
||||
|
||||
def list_files(
|
||||
path: str,
|
||||
recursive: bool = False,
|
||||
workspace: Optional[str] = None,
|
||||
max_files: int = 200,
|
||||
include_hidden: bool = False,
|
||||
extensions: Optional[List[str]] = None,
|
||||
respect_gitignore: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
List files and directories in a given path.
|
||||
|
||||
This tool lists the contents of a directory, optionally recursively,
|
||||
with support for filtering and gitignore patterns.
|
||||
|
||||
Args:
|
||||
path: Path to the directory (absolute or relative to workspace)
|
||||
recursive: Whether to list recursively
|
||||
workspace: Workspace root directory (for relative paths)
|
||||
max_files: Maximum number of files to return
|
||||
include_hidden: Whether to include hidden files (starting with .)
|
||||
extensions: List of file extensions to include (e.g., ['py', 'js'])
|
||||
respect_gitignore: Whether to respect .gitignore patterns
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- files: List of file entries with path, type, size
|
||||
- directories: List of directory entries
|
||||
- total_count: int (total items found)
|
||||
- truncated: bool (True if max_files was reached)
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> result = list_files("src", recursive=True, extensions=['py'])
|
||||
>>> for f in result['files']:
|
||||
... print(f['path'])
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Check if directory exists
|
||||
if not os.path.isdir(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Directory not found: {path}",
|
||||
'files': [],
|
||||
'directories': [],
|
||||
'total_count': 0,
|
||||
}
|
||||
|
||||
# Load gitignore patterns
|
||||
ignore_patterns = []
|
||||
if respect_gitignore:
|
||||
# Load from workspace root if available
|
||||
if workspace:
|
||||
ignore_patterns = load_gitignore_patterns(workspace)
|
||||
# Also load from the target directory
|
||||
ignore_patterns.extend(load_gitignore_patterns(abs_path))
|
||||
|
||||
files = []
|
||||
directories = []
|
||||
total_count = 0
|
||||
truncated = False
|
||||
|
||||
try:
|
||||
if recursive:
|
||||
# Walk directory tree
|
||||
for root, dirs, filenames in os.walk(abs_path):
|
||||
# Filter hidden directories
|
||||
if not include_hidden:
|
||||
dirs[:] = [d for d in dirs if not d.startswith('.')]
|
||||
|
||||
# Filter ignored directories
|
||||
if respect_gitignore:
|
||||
base_dir = workspace or abs_path
|
||||
dirs[:] = [
|
||||
d for d in dirs
|
||||
if not should_ignore_path(
|
||||
os.path.join(root, d),
|
||||
ignore_patterns,
|
||||
base_dir,
|
||||
is_directory=True
|
||||
)
|
||||
]
|
||||
|
||||
# Process files
|
||||
for filename in filenames:
|
||||
if total_count >= max_files:
|
||||
truncated = True
|
||||
break
|
||||
|
||||
# Skip hidden files
|
||||
if not include_hidden and filename.startswith('.'):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(root, filename)
|
||||
rel_path = os.path.relpath(file_path, abs_path)
|
||||
|
||||
# Check gitignore
|
||||
if respect_gitignore:
|
||||
base_dir = workspace or abs_path
|
||||
if should_ignore_path(file_path, ignore_patterns, base_dir):
|
||||
continue
|
||||
|
||||
# Filter by extension
|
||||
if extensions:
|
||||
ext = os.path.splitext(filename)[1].lstrip('.')
|
||||
if ext.lower() not in [e.lower() for e in extensions]:
|
||||
continue
|
||||
|
||||
# Get file info
|
||||
try:
|
||||
stat = os.stat(file_path)
|
||||
files.append({
|
||||
'path': rel_path,
|
||||
'name': filename,
|
||||
'size': stat.st_size,
|
||||
'modified': stat.st_mtime,
|
||||
})
|
||||
total_count += 1
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
if truncated:
|
||||
break
|
||||
|
||||
# Add directories
|
||||
for dirname in dirs:
|
||||
if total_count >= max_files:
|
||||
truncated = True
|
||||
break
|
||||
|
||||
dir_path = os.path.join(root, dirname)
|
||||
rel_path = os.path.relpath(dir_path, abs_path)
|
||||
|
||||
directories.append({
|
||||
'path': rel_path,
|
||||
'name': dirname,
|
||||
})
|
||||
total_count += 1
|
||||
else:
|
||||
# List only top level
|
||||
for entry in os.listdir(abs_path):
|
||||
if total_count >= max_files:
|
||||
truncated = True
|
||||
break
|
||||
|
||||
# Skip hidden files
|
||||
if not include_hidden and entry.startswith('.'):
|
||||
continue
|
||||
|
||||
entry_path = os.path.join(abs_path, entry)
|
||||
|
||||
# Check gitignore
|
||||
if respect_gitignore:
|
||||
base_dir = workspace or abs_path
|
||||
is_dir = os.path.isdir(entry_path)
|
||||
if should_ignore_path(entry_path, ignore_patterns, base_dir, is_dir):
|
||||
continue
|
||||
|
||||
if os.path.isfile(entry_path):
|
||||
# Filter by extension
|
||||
if extensions:
|
||||
ext = os.path.splitext(entry)[1].lstrip('.')
|
||||
if ext.lower() not in [e.lower() for e in extensions]:
|
||||
continue
|
||||
|
||||
try:
|
||||
stat = os.stat(entry_path)
|
||||
files.append({
|
||||
'path': entry,
|
||||
'name': entry,
|
||||
'size': stat.st_size,
|
||||
'modified': stat.st_mtime,
|
||||
})
|
||||
total_count += 1
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
elif os.path.isdir(entry_path):
|
||||
directories.append({
|
||||
'path': entry,
|
||||
'name': entry,
|
||||
})
|
||||
total_count += 1
|
||||
|
||||
# Sort results
|
||||
files.sort(key=lambda x: x['path'])
|
||||
directories.sort(key=lambda x: x['path'])
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'files': files,
|
||||
'directories': directories,
|
||||
'total_count': total_count,
|
||||
'truncated': truncated,
|
||||
'path': path,
|
||||
}
|
||||
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Permission denied: {path}",
|
||||
'files': [],
|
||||
'directories': [],
|
||||
'total_count': 0,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error listing {path}: {str(e)}",
|
||||
'files': [],
|
||||
'directories': [],
|
||||
'total_count': 0,
|
||||
}
|
||||
|
||||
|
||||
def format_file_list(result: Dict[str, Any], show_size: bool = True) -> str:
|
||||
"""
|
||||
Format a list_files result as a readable string.
|
||||
|
||||
Args:
|
||||
result: Result from list_files()
|
||||
show_size: Whether to show file sizes
|
||||
|
||||
Returns:
|
||||
Formatted string representation
|
||||
"""
|
||||
if not result['success']:
|
||||
return f"Error: {result.get('error', 'Unknown error')}"
|
||||
|
||||
lines = []
|
||||
|
||||
# Add directories first
|
||||
for d in result['directories']:
|
||||
lines.append(f"📁 {d['path']}/")
|
||||
|
||||
# Add files
|
||||
for f in result['files']:
|
||||
if show_size:
|
||||
size = f.get('size', 0)
|
||||
if size < 1024:
|
||||
size_str = f"{size}B"
|
||||
elif size < 1024 * 1024:
|
||||
size_str = f"{size / 1024:.1f}KB"
|
||||
else:
|
||||
size_str = f"{size / (1024 * 1024):.1f}MB"
|
||||
lines.append(f"📄 {f['path']} ({size_str})")
|
||||
else:
|
||||
lines.append(f"📄 {f['path']}")
|
||||
|
||||
if result.get('truncated'):
|
||||
lines.append(f"\n... (truncated, showing {result['total_count']} of more items)")
|
||||
|
||||
return '\n'.join(lines)
|
||||
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
Read File Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to read file contents with optional line ranges
|
||||
and line number annotations.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from ..utils.file_utils import (
|
||||
add_line_numbers,
|
||||
is_binary_file,
|
||||
file_exists,
|
||||
is_path_within_directory,
|
||||
)
|
||||
|
||||
|
||||
def read_file(
|
||||
path: str,
|
||||
start_line: Optional[int] = None,
|
||||
end_line: Optional[int] = None,
|
||||
add_line_nums: bool = True,
|
||||
workspace: Optional[str] = None,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Read file contents with optional line range and line numbers.
|
||||
|
||||
This tool reads the contents of a file, optionally limiting to specific
|
||||
line ranges and adding line number annotations for AI context.
|
||||
|
||||
Args:
|
||||
path: Path to the file (absolute or relative to workspace)
|
||||
start_line: First line to read (1-indexed, inclusive)
|
||||
end_line: Last line to read (1-indexed, inclusive)
|
||||
add_line_nums: Whether to add line numbers to output
|
||||
workspace: Workspace root directory (for relative paths)
|
||||
encoding: File encoding (default: utf-8)
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- content: str (file content, possibly with line numbers)
|
||||
- total_lines: int (total lines in file)
|
||||
- start_line: int (actual start line read)
|
||||
- end_line: int (actual end line read)
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> result = read_file("src/main.py", start_line=1, end_line=50)
|
||||
>>> if result['success']:
|
||||
... print(result['content'])
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Security check - ensure path is within workspace if specified
|
||||
if workspace:
|
||||
if not is_path_within_directory(abs_path, workspace):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Path '{path}' is outside the workspace",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
|
||||
# Check if file exists
|
||||
if not file_exists(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"File not found: {path}",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
|
||||
# Check if binary
|
||||
if is_binary_file(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Cannot read binary file: {path}",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
|
||||
try:
|
||||
# Read file content
|
||||
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
total_lines = len(lines)
|
||||
|
||||
# Determine line range
|
||||
actual_start = start_line if start_line and start_line > 0 else 1
|
||||
actual_end = end_line if end_line and end_line > 0 else total_lines
|
||||
|
||||
# Clamp to valid range
|
||||
actual_start = max(1, min(actual_start, total_lines))
|
||||
actual_end = max(actual_start, min(actual_end, total_lines))
|
||||
|
||||
# Extract lines (convert to 0-indexed)
|
||||
selected_lines = lines[actual_start - 1:actual_end]
|
||||
content = ''.join(selected_lines)
|
||||
|
||||
# Remove trailing newline if present
|
||||
if content.endswith('\n'):
|
||||
content = content[:-1]
|
||||
|
||||
# Add line numbers if requested
|
||||
if add_line_nums:
|
||||
content = add_line_numbers(content, actual_start)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'content': content,
|
||||
'total_lines': total_lines,
|
||||
'start_line': actual_start,
|
||||
'end_line': actual_end,
|
||||
'path': path,
|
||||
}
|
||||
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Permission denied: {path}",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
except UnicodeDecodeError as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Encoding error reading {path}: {e}",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error reading {path}: {str(e)}",
|
||||
'content': None,
|
||||
'total_lines': 0,
|
||||
}
|
||||
|
||||
|
||||
def read_multiple_files(
|
||||
files: List[Dict[str, Any]],
|
||||
workspace: Optional[str] = None,
|
||||
add_line_nums: bool = True,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Read multiple files in a single operation.
|
||||
|
||||
Args:
|
||||
files: List of file specifications, each with:
|
||||
- path: str (required)
|
||||
- start_line: int (optional)
|
||||
- end_line: int (optional)
|
||||
workspace: Workspace root directory
|
||||
add_line_nums: Whether to add line numbers
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool (True if all files read successfully)
|
||||
- results: List of individual file results
|
||||
- errors: List of error messages
|
||||
|
||||
Example:
|
||||
>>> files = [
|
||||
... {'path': 'src/main.py', 'start_line': 1, 'end_line': 50},
|
||||
... {'path': 'src/utils.py'},
|
||||
... ]
|
||||
>>> result = read_multiple_files(files, workspace='/project')
|
||||
"""
|
||||
results = []
|
||||
errors = []
|
||||
all_success = True
|
||||
|
||||
for file_spec in files:
|
||||
path = file_spec.get('path')
|
||||
if not path:
|
||||
errors.append("Missing 'path' in file specification")
|
||||
all_success = False
|
||||
continue
|
||||
|
||||
result = read_file(
|
||||
path=path,
|
||||
start_line=file_spec.get('start_line'),
|
||||
end_line=file_spec.get('end_line'),
|
||||
add_line_nums=add_line_nums,
|
||||
workspace=workspace,
|
||||
)
|
||||
|
||||
results.append(result)
|
||||
|
||||
if not result['success']:
|
||||
all_success = False
|
||||
errors.append(result.get('error', f"Failed to read {path}"))
|
||||
|
||||
return {
|
||||
'success': all_success,
|
||||
'results': results,
|
||||
'errors': errors if errors else None,
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
Search and Replace Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to perform multiple search/replace operations
|
||||
on a file in a single operation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from ..utils.file_utils import (
|
||||
file_exists,
|
||||
is_path_within_directory,
|
||||
)
|
||||
|
||||
|
||||
def search_replace(
|
||||
path: str,
|
||||
operations: List[Dict[str, Any]],
|
||||
workspace: Optional[str] = None,
|
||||
backup: bool = False,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Perform multiple search/replace operations on a file.
|
||||
|
||||
This tool allows precise modifications to file content by performing
|
||||
multiple search and replace operations in sequence.
|
||||
|
||||
Args:
|
||||
path: Path to the file (absolute or relative to workspace)
|
||||
operations: List of operations, each with:
|
||||
- search: str (text to find)
|
||||
- replace: str (text to replace with)
|
||||
- is_regex: bool (optional, treat search as regex)
|
||||
- count: int (optional, max replacements, -1 for all)
|
||||
workspace: Workspace root directory
|
||||
backup: Whether to create a backup before modifying
|
||||
encoding: File encoding
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- path: str
|
||||
- operations_applied: int
|
||||
- total_replacements: int
|
||||
- failed_operations: list
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> operations = [
|
||||
... {'search': 'old_name', 'replace': 'new_name'},
|
||||
... {'search': r'def (\\w+)\\(', 'replace': r'def renamed_\\\\1(', 'is_regex': True},
|
||||
... ]
|
||||
>>> result = search_replace("src/main.py", operations)
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Security check
|
||||
if workspace:
|
||||
if not is_path_within_directory(abs_path, workspace):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Path '{path}' is outside the workspace",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
}
|
||||
|
||||
# Check if file exists
|
||||
if not file_exists(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"File not found: {path}",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
}
|
||||
|
||||
if not operations:
|
||||
return {
|
||||
'success': False,
|
||||
'error': "No operations provided",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
}
|
||||
|
||||
try:
|
||||
# Read original content
|
||||
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
|
||||
content = f.read()
|
||||
|
||||
original_content = content
|
||||
operations_applied = 0
|
||||
total_replacements = 0
|
||||
failed_operations = []
|
||||
|
||||
# Apply each operation
|
||||
for i, op in enumerate(operations):
|
||||
search = op.get('search')
|
||||
replace = op.get('replace', '')
|
||||
is_regex = op.get('is_regex', False)
|
||||
count = op.get('count', -1) # -1 means replace all
|
||||
|
||||
if not search:
|
||||
failed_operations.append({
|
||||
'index': i,
|
||||
'error': "Missing 'search' field",
|
||||
})
|
||||
continue
|
||||
|
||||
try:
|
||||
if is_regex:
|
||||
# Regex replacement
|
||||
if count == -1:
|
||||
new_content, num_subs = re.subn(search, replace, content)
|
||||
else:
|
||||
new_content, num_subs = re.subn(search, replace, content, count=count)
|
||||
else:
|
||||
# Literal replacement
|
||||
if count == -1:
|
||||
num_subs = content.count(search)
|
||||
new_content = content.replace(search, replace)
|
||||
else:
|
||||
num_subs = min(content.count(search), count)
|
||||
new_content = content
|
||||
for _ in range(count):
|
||||
if search in new_content:
|
||||
new_content = new_content.replace(search, replace, 1)
|
||||
else:
|
||||
break
|
||||
|
||||
if num_subs > 0:
|
||||
content = new_content
|
||||
operations_applied += 1
|
||||
total_replacements += num_subs
|
||||
else:
|
||||
failed_operations.append({
|
||||
'index': i,
|
||||
'search': search[:50],
|
||||
'error': "No matches found",
|
||||
})
|
||||
|
||||
except re.error as e:
|
||||
failed_operations.append({
|
||||
'index': i,
|
||||
'search': search[:50],
|
||||
'error': f"Invalid regex: {str(e)}",
|
||||
})
|
||||
|
||||
# Check if any changes were made
|
||||
if content == original_content:
|
||||
return {
|
||||
'success': False,
|
||||
'error': "No changes were made - no matches found",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
'failed_operations': failed_operations,
|
||||
}
|
||||
|
||||
# Create backup if requested
|
||||
backup_path = None
|
||||
if backup:
|
||||
import time
|
||||
timestamp = int(time.time())
|
||||
backup_path = f"{abs_path}.backup.{timestamp}"
|
||||
with open(backup_path, 'w', encoding=encoding) as f:
|
||||
f.write(original_content)
|
||||
|
||||
# Write the modified content
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(content)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'path': path,
|
||||
'absolute_path': abs_path,
|
||||
'operations_applied': operations_applied,
|
||||
'total_replacements': total_replacements,
|
||||
'failed_operations': failed_operations if failed_operations else None,
|
||||
'backup_path': backup_path,
|
||||
}
|
||||
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Permission denied: {path}",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error performing search/replace on {path}: {str(e)}",
|
||||
'path': path,
|
||||
'operations_applied': 0,
|
||||
'total_replacements': 0,
|
||||
}
|
||||
|
||||
|
||||
def simple_replace(
|
||||
path: str,
|
||||
search: str,
|
||||
replace: str,
|
||||
workspace: Optional[str] = None,
|
||||
is_regex: bool = False,
|
||||
count: int = -1,
|
||||
backup: bool = False,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Perform a single search/replace operation on a file.
|
||||
|
||||
Convenience wrapper around search_replace for single operations.
|
||||
|
||||
Args:
|
||||
path: Path to the file
|
||||
search: Text to find
|
||||
replace: Text to replace with
|
||||
workspace: Workspace root directory
|
||||
is_regex: Whether search is a regex pattern
|
||||
count: Max replacements (-1 for all)
|
||||
backup: Whether to create a backup
|
||||
encoding: File encoding
|
||||
|
||||
Returns:
|
||||
Dictionary with success status and details
|
||||
"""
|
||||
return search_replace(
|
||||
path=path,
|
||||
operations=[{
|
||||
'search': search,
|
||||
'replace': replace,
|
||||
'is_regex': is_regex,
|
||||
'count': count,
|
||||
}],
|
||||
workspace=workspace,
|
||||
backup=backup,
|
||||
encoding=encoding,
|
||||
)
|
||||
@@ -0,0 +1,217 @@
|
||||
"""
|
||||
Write File Tool for PraisonAI Code.
|
||||
|
||||
Provides functionality to create or overwrite files with content.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from ..utils.file_utils import (
|
||||
file_exists,
|
||||
create_directories_for_file,
|
||||
is_path_within_directory,
|
||||
detect_line_ending,
|
||||
normalize_line_endings,
|
||||
)
|
||||
from ..utils.text_utils import (
|
||||
unescape_html_entities,
|
||||
strip_markdown_code_fences,
|
||||
)
|
||||
|
||||
|
||||
def write_file(
|
||||
path: str,
|
||||
content: str,
|
||||
workspace: Optional[str] = None,
|
||||
create_directories: bool = True,
|
||||
backup: bool = False,
|
||||
strip_code_fences: bool = True,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Write content to a file, creating it if it doesn't exist.
|
||||
|
||||
This tool writes content to a file, optionally creating parent
|
||||
directories and handling common AI output artifacts like code fences.
|
||||
|
||||
Args:
|
||||
path: Path to the file (absolute or relative to workspace)
|
||||
content: Content to write to the file
|
||||
workspace: Workspace root directory (for relative paths)
|
||||
create_directories: Whether to create parent directories
|
||||
backup: Whether to create a backup of existing files
|
||||
strip_code_fences: Whether to strip markdown code fences
|
||||
encoding: File encoding (default: utf-8)
|
||||
|
||||
Returns:
|
||||
Dictionary with:
|
||||
- success: bool
|
||||
- path: str (the path written to)
|
||||
- created: bool (True if file was created, False if overwritten)
|
||||
- backup_path: str (if backup was created)
|
||||
- error: str (if success is False)
|
||||
|
||||
Example:
|
||||
>>> result = write_file("src/new_file.py", "print('hello')")
|
||||
>>> if result['success']:
|
||||
... print(f"Wrote to {result['path']}")
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Security check - ensure path is within workspace if specified
|
||||
if workspace:
|
||||
if not is_path_within_directory(abs_path, workspace):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Path '{path}' is outside the workspace",
|
||||
'path': path,
|
||||
}
|
||||
|
||||
# Process content
|
||||
processed_content = content
|
||||
|
||||
# Strip markdown code fences if requested
|
||||
if strip_code_fences:
|
||||
processed_content = strip_markdown_code_fences(processed_content)
|
||||
|
||||
# Unescape HTML entities (common in AI output)
|
||||
processed_content = unescape_html_entities(processed_content)
|
||||
|
||||
# Check if file exists
|
||||
file_existed = file_exists(abs_path)
|
||||
backup_path = None
|
||||
|
||||
try:
|
||||
# Create backup if requested and file exists
|
||||
if backup and file_existed:
|
||||
import time
|
||||
timestamp = int(time.time())
|
||||
backup_path = f"{abs_path}.backup.{timestamp}"
|
||||
with open(abs_path, 'r', encoding=encoding, errors='replace') as src:
|
||||
with open(backup_path, 'w', encoding=encoding) as dst:
|
||||
dst.write(src.read())
|
||||
|
||||
# Create parent directories if needed
|
||||
if create_directories:
|
||||
if not create_directories_for_file(abs_path):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Failed to create directories for {path}",
|
||||
'path': path,
|
||||
}
|
||||
|
||||
# Preserve line endings if file exists
|
||||
if file_existed:
|
||||
try:
|
||||
with open(abs_path, 'r', encoding=encoding, errors='replace') as f:
|
||||
original_content = f.read()
|
||||
original_line_ending = detect_line_ending(original_content)
|
||||
processed_content = normalize_line_endings(processed_content, original_line_ending)
|
||||
except Exception:
|
||||
pass # If we can't read, just use the content as-is
|
||||
|
||||
# Write the file
|
||||
with open(abs_path, 'w', encoding=encoding) as f:
|
||||
f.write(processed_content)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'path': path,
|
||||
'absolute_path': abs_path,
|
||||
'created': not file_existed,
|
||||
'backup_path': backup_path,
|
||||
'bytes_written': len(processed_content.encode(encoding)),
|
||||
}
|
||||
|
||||
except PermissionError:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Permission denied: {path}",
|
||||
'path': path,
|
||||
}
|
||||
except OSError as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"OS error writing {path}: {str(e)}",
|
||||
'path': path,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error writing {path}: {str(e)}",
|
||||
'path': path,
|
||||
}
|
||||
|
||||
|
||||
def append_to_file(
|
||||
path: str,
|
||||
content: str,
|
||||
workspace: Optional[str] = None,
|
||||
create_if_missing: bool = True,
|
||||
encoding: str = 'utf-8',
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Append content to an existing file.
|
||||
|
||||
Args:
|
||||
path: Path to the file
|
||||
content: Content to append
|
||||
workspace: Workspace root directory
|
||||
create_if_missing: Create file if it doesn't exist
|
||||
encoding: File encoding
|
||||
|
||||
Returns:
|
||||
Dictionary with success status and details
|
||||
"""
|
||||
# Resolve path
|
||||
if workspace and not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(workspace, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Security check
|
||||
if workspace:
|
||||
if not is_path_within_directory(abs_path, workspace):
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Path '{path}' is outside the workspace",
|
||||
'path': path,
|
||||
}
|
||||
|
||||
file_existed = file_exists(abs_path)
|
||||
|
||||
if not file_existed and not create_if_missing:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"File not found: {path}",
|
||||
'path': path,
|
||||
}
|
||||
|
||||
try:
|
||||
# Create directories if needed
|
||||
if not file_existed:
|
||||
create_directories_for_file(abs_path)
|
||||
|
||||
# Append to file
|
||||
with open(abs_path, 'a', encoding=encoding) as f:
|
||||
f.write(content)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'path': path,
|
||||
'absolute_path': abs_path,
|
||||
'created': not file_existed,
|
||||
'bytes_appended': len(content.encode(encoding)),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
return {
|
||||
'success': False,
|
||||
'error': f"Error appending to {path}: {str(e)}",
|
||||
'path': path,
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Utility functions for PraisonAI Code module.
|
||||
"""
|
||||
|
||||
from .file_utils import (
|
||||
add_line_numbers,
|
||||
strip_line_numbers,
|
||||
every_line_has_line_numbers,
|
||||
normalize_line_endings,
|
||||
get_file_extension,
|
||||
is_binary_file,
|
||||
create_directories_for_file,
|
||||
file_exists,
|
||||
)
|
||||
|
||||
from .text_utils import (
|
||||
normalize_string,
|
||||
unescape_html_entities,
|
||||
get_similarity,
|
||||
)
|
||||
|
||||
from .ignore_utils import (
|
||||
FileAccessController,
|
||||
load_gitignore_patterns,
|
||||
should_ignore_path,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# File utilities
|
||||
"add_line_numbers",
|
||||
"strip_line_numbers",
|
||||
"every_line_has_line_numbers",
|
||||
"normalize_line_endings",
|
||||
"get_file_extension",
|
||||
"is_binary_file",
|
||||
"create_directories_for_file",
|
||||
"file_exists",
|
||||
# Text utilities
|
||||
"normalize_string",
|
||||
"unescape_html_entities",
|
||||
"get_similarity",
|
||||
# Ignore utilities
|
||||
"FileAccessController",
|
||||
"load_gitignore_patterns",
|
||||
"should_ignore_path",
|
||||
]
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
File utility functions for PraisonAI Code module.
|
||||
|
||||
Provides utilities for file operations including:
|
||||
- Line number handling (add/strip)
|
||||
- File existence and type checking
|
||||
- Directory creation
|
||||
- Line ending normalization
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
def add_line_numbers(content: str, start_line: int = 1) -> str:
|
||||
"""
|
||||
Add line numbers to content in the format: " N | content"
|
||||
|
||||
Args:
|
||||
content: The text content to add line numbers to
|
||||
start_line: The starting line number (1-indexed)
|
||||
|
||||
Returns:
|
||||
Content with line numbers prefixed to each line
|
||||
|
||||
Example:
|
||||
>>> add_line_numbers("hello\\nworld", 1)
|
||||
' 1 | hello\\n 2 | world'
|
||||
"""
|
||||
if not content:
|
||||
return ""
|
||||
|
||||
lines = content.split('\n')
|
||||
max_line_num = start_line + len(lines) - 1
|
||||
width = len(str(max_line_num))
|
||||
|
||||
numbered_lines = []
|
||||
for i, line in enumerate(lines):
|
||||
line_num = start_line + i
|
||||
numbered_lines.append(f"{line_num:>{width}} | {line}")
|
||||
|
||||
return '\n'.join(numbered_lines)
|
||||
|
||||
|
||||
def strip_line_numbers(content: str, aggressive: bool = False) -> str:
|
||||
"""
|
||||
Strip line numbers from content.
|
||||
|
||||
Handles formats like:
|
||||
- " 1 | content"
|
||||
- "1| content"
|
||||
- " 1\tcontent"
|
||||
|
||||
Args:
|
||||
content: Content with line numbers
|
||||
aggressive: If True, also strip formats like "1:" or just leading numbers
|
||||
|
||||
Returns:
|
||||
Content with line numbers removed
|
||||
"""
|
||||
if not content:
|
||||
return ""
|
||||
|
||||
lines = content.split('\n')
|
||||
stripped_lines = []
|
||||
|
||||
# Pattern for standard format: optional spaces, digits, optional spaces, pipe/tab, content
|
||||
standard_pattern = re.compile(r'^\s*\d+\s*[|\t]\s?(.*)$')
|
||||
# Aggressive pattern: digits followed by colon or just leading digits with space
|
||||
aggressive_pattern = re.compile(r'^\s*\d+[:\s]\s?(.*)$')
|
||||
|
||||
for line in lines:
|
||||
match = standard_pattern.match(line)
|
||||
if match:
|
||||
stripped_lines.append(match.group(1))
|
||||
elif aggressive:
|
||||
match = aggressive_pattern.match(line)
|
||||
if match:
|
||||
stripped_lines.append(match.group(1))
|
||||
else:
|
||||
stripped_lines.append(line)
|
||||
else:
|
||||
stripped_lines.append(line)
|
||||
|
||||
return '\n'.join(stripped_lines)
|
||||
|
||||
|
||||
def every_line_has_line_numbers(content: str) -> bool:
|
||||
"""
|
||||
Check if every non-empty line in content has line numbers.
|
||||
|
||||
Args:
|
||||
content: The content to check
|
||||
|
||||
Returns:
|
||||
True if all non-empty lines have line numbers
|
||||
"""
|
||||
if not content or not content.strip():
|
||||
return False
|
||||
|
||||
lines = content.split('\n')
|
||||
pattern = re.compile(r'^\s*\d+\s*[|\t]')
|
||||
|
||||
for line in lines:
|
||||
if line.strip(): # Only check non-empty lines
|
||||
if not pattern.match(line):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def normalize_line_endings(content: str, target: str = '\n') -> str:
|
||||
"""
|
||||
Normalize line endings to a consistent format.
|
||||
|
||||
Args:
|
||||
content: The content to normalize
|
||||
target: The target line ending ('\\n' or '\\r\\n')
|
||||
|
||||
Returns:
|
||||
Content with normalized line endings
|
||||
"""
|
||||
# First normalize all to \n, then convert to target
|
||||
normalized = content.replace('\r\n', '\n').replace('\r', '\n')
|
||||
if target == '\r\n':
|
||||
normalized = normalized.replace('\n', '\r\n')
|
||||
return normalized
|
||||
|
||||
|
||||
def detect_line_ending(content: str) -> str:
|
||||
"""
|
||||
Detect the predominant line ending in content.
|
||||
|
||||
Args:
|
||||
content: The content to analyze
|
||||
|
||||
Returns:
|
||||
'\\r\\n' if Windows-style, '\\n' otherwise
|
||||
"""
|
||||
if '\r\n' in content:
|
||||
return '\r\n'
|
||||
return '\n'
|
||||
|
||||
|
||||
def get_file_extension(file_path: str) -> str:
|
||||
"""
|
||||
Get the file extension from a path.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
File extension without the dot, or empty string
|
||||
"""
|
||||
ext = os.path.splitext(file_path)[1]
|
||||
return ext[1:] if ext else ""
|
||||
|
||||
|
||||
def is_binary_file(file_path: str, sample_size: int = 8192) -> bool:
|
||||
"""
|
||||
Check if a file is binary by reading a sample.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
sample_size: Number of bytes to sample
|
||||
|
||||
Returns:
|
||||
True if the file appears to be binary
|
||||
"""
|
||||
try:
|
||||
with open(file_path, 'rb') as f:
|
||||
sample = f.read(sample_size)
|
||||
|
||||
# Check for null bytes (common in binary files)
|
||||
if b'\x00' in sample:
|
||||
return True
|
||||
|
||||
# Check the ratio of non-text characters
|
||||
text_chars = bytearray({7, 8, 9, 10, 12, 13, 27} | set(range(0x20, 0x100)) - {0x7f})
|
||||
non_text = sum(1 for byte in sample if byte not in text_chars)
|
||||
|
||||
# If more than 30% non-text characters, consider it binary
|
||||
return len(sample) > 0 and (non_text / len(sample)) > 0.30
|
||||
|
||||
except (IOError, OSError):
|
||||
return False
|
||||
|
||||
|
||||
def create_directories_for_file(file_path: str) -> bool:
|
||||
"""
|
||||
Create parent directories for a file path if they don't exist.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
True if directories were created or already exist
|
||||
"""
|
||||
try:
|
||||
parent_dir = os.path.dirname(file_path)
|
||||
if parent_dir:
|
||||
os.makedirs(parent_dir, exist_ok=True)
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def file_exists(file_path: str) -> bool:
|
||||
"""
|
||||
Check if a file exists.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
True if the file exists and is a file (not directory)
|
||||
"""
|
||||
return os.path.isfile(file_path)
|
||||
|
||||
|
||||
def get_relative_path(file_path: str, base_path: str) -> str:
|
||||
"""
|
||||
Get the relative path from base_path to file_path.
|
||||
|
||||
Args:
|
||||
file_path: The target file path
|
||||
base_path: The base directory path
|
||||
|
||||
Returns:
|
||||
Relative path string
|
||||
"""
|
||||
try:
|
||||
return os.path.relpath(file_path, base_path)
|
||||
except ValueError:
|
||||
# On Windows, relpath fails for paths on different drives
|
||||
return file_path
|
||||
|
||||
|
||||
def is_path_within_directory(file_path: str, directory: str) -> bool:
|
||||
"""
|
||||
Check if a file path is within a directory (prevents path traversal).
|
||||
|
||||
Args:
|
||||
file_path: The file path to check
|
||||
directory: The directory that should contain the file
|
||||
|
||||
Returns:
|
||||
True if file_path is within directory
|
||||
"""
|
||||
# Resolve to absolute paths
|
||||
abs_file = os.path.abspath(file_path)
|
||||
abs_dir = os.path.abspath(directory)
|
||||
|
||||
# Ensure directory ends with separator for proper prefix matching
|
||||
if not abs_dir.endswith(os.sep):
|
||||
abs_dir += os.sep
|
||||
|
||||
return abs_file.startswith(abs_dir) or abs_file == abs_dir.rstrip(os.sep)
|
||||
|
||||
|
||||
def read_file_lines(file_path: str, start_line: int = 1, end_line: Optional[int] = None) -> Tuple[str, int]:
|
||||
"""
|
||||
Read specific lines from a file.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
start_line: First line to read (1-indexed)
|
||||
end_line: Last line to read (inclusive), None for all remaining
|
||||
|
||||
Returns:
|
||||
Tuple of (content, total_lines)
|
||||
"""
|
||||
with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
total_lines = len(lines)
|
||||
|
||||
# Convert to 0-indexed
|
||||
start_idx = max(0, start_line - 1)
|
||||
end_idx = end_line if end_line else total_lines
|
||||
|
||||
selected_lines = lines[start_idx:end_idx]
|
||||
content = ''.join(selected_lines)
|
||||
|
||||
# Remove trailing newline if present
|
||||
if content.endswith('\n'):
|
||||
content = content[:-1]
|
||||
|
||||
return content, total_lines
|
||||
|
||||
|
||||
def count_file_lines(file_path: str) -> int:
|
||||
"""
|
||||
Count the number of lines in a file efficiently.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
Number of lines in the file
|
||||
"""
|
||||
count = 0
|
||||
with open(file_path, 'rb') as f:
|
||||
for _ in f:
|
||||
count += 1
|
||||
return count
|
||||
@@ -0,0 +1,308 @@
|
||||
"""
|
||||
File access control utilities for PraisonAI Code module.
|
||||
|
||||
Provides utilities for:
|
||||
- Loading and parsing .gitignore patterns
|
||||
- Checking if paths should be ignored
|
||||
- Controlling file access within workspaces
|
||||
"""
|
||||
|
||||
import os
|
||||
import fnmatch
|
||||
from typing import List, Optional, Set
|
||||
|
||||
|
||||
def load_gitignore_patterns(directory: str) -> List[str]:
|
||||
"""
|
||||
Load .gitignore patterns from a directory.
|
||||
|
||||
Args:
|
||||
directory: The directory to search for .gitignore
|
||||
|
||||
Returns:
|
||||
List of gitignore patterns
|
||||
"""
|
||||
patterns = []
|
||||
gitignore_path = os.path.join(directory, '.gitignore')
|
||||
|
||||
if os.path.isfile(gitignore_path):
|
||||
try:
|
||||
with open(gitignore_path, 'r', encoding='utf-8') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
# Skip empty lines and comments
|
||||
if line and not line.startswith('#'):
|
||||
patterns.append(line)
|
||||
except (IOError, OSError):
|
||||
pass
|
||||
|
||||
return patterns
|
||||
|
||||
|
||||
def _pattern_to_regex_parts(pattern: str) -> tuple:
|
||||
"""
|
||||
Convert a gitignore pattern to matching components.
|
||||
|
||||
Args:
|
||||
pattern: A gitignore pattern
|
||||
|
||||
Returns:
|
||||
Tuple of (is_negation, is_directory_only, cleaned_pattern)
|
||||
"""
|
||||
is_negation = pattern.startswith('!')
|
||||
if is_negation:
|
||||
pattern = pattern[1:]
|
||||
|
||||
is_directory_only = pattern.endswith('/')
|
||||
if is_directory_only:
|
||||
pattern = pattern[:-1]
|
||||
|
||||
return is_negation, is_directory_only, pattern
|
||||
|
||||
|
||||
def should_ignore_path(
|
||||
path: str,
|
||||
patterns: List[str],
|
||||
base_dir: str,
|
||||
is_directory: bool = False
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a path should be ignored based on gitignore patterns.
|
||||
|
||||
Args:
|
||||
path: The path to check (relative or absolute)
|
||||
patterns: List of gitignore patterns
|
||||
base_dir: The base directory for relative path calculation
|
||||
is_directory: Whether the path is a directory
|
||||
|
||||
Returns:
|
||||
True if the path should be ignored
|
||||
"""
|
||||
# Get relative path
|
||||
if os.path.isabs(path):
|
||||
try:
|
||||
rel_path = os.path.relpath(path, base_dir)
|
||||
except ValueError:
|
||||
rel_path = path
|
||||
else:
|
||||
rel_path = path
|
||||
|
||||
# Normalize path separators
|
||||
rel_path = rel_path.replace(os.sep, '/')
|
||||
|
||||
# Check each pattern
|
||||
ignored = False
|
||||
|
||||
for pattern in patterns:
|
||||
is_negation, is_directory_only, clean_pattern = _pattern_to_regex_parts(pattern)
|
||||
|
||||
# Skip directory-only patterns for files
|
||||
if is_directory_only and not is_directory:
|
||||
continue
|
||||
|
||||
# Check if pattern matches
|
||||
matches = False
|
||||
|
||||
# Handle patterns with /
|
||||
if '/' in clean_pattern:
|
||||
# Pattern with / matches from root
|
||||
if clean_pattern.startswith('/'):
|
||||
clean_pattern = clean_pattern[1:]
|
||||
matches = fnmatch.fnmatch(rel_path, clean_pattern)
|
||||
else:
|
||||
# Pattern without / matches any path component
|
||||
path_parts = rel_path.split('/')
|
||||
for part in path_parts:
|
||||
if fnmatch.fnmatch(part, clean_pattern):
|
||||
matches = True
|
||||
break
|
||||
# Also check full path for ** patterns
|
||||
if not matches and '**' in clean_pattern:
|
||||
matches = fnmatch.fnmatch(rel_path, clean_pattern)
|
||||
|
||||
if matches:
|
||||
ignored = not is_negation
|
||||
|
||||
return ignored
|
||||
|
||||
|
||||
class FileAccessController:
|
||||
"""
|
||||
Controls file access within a workspace.
|
||||
|
||||
Respects .gitignore patterns and can be configured with
|
||||
additional ignore/allow patterns.
|
||||
|
||||
Attributes:
|
||||
workspace_root: The root directory of the workspace
|
||||
ignore_patterns: Patterns for files to ignore
|
||||
protected_patterns: Patterns for files that are read-only
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workspace_root: str,
|
||||
load_gitignore: bool = True,
|
||||
additional_ignore_patterns: Optional[List[str]] = None,
|
||||
protected_patterns: Optional[List[str]] = None
|
||||
):
|
||||
"""
|
||||
Initialize the FileAccessController.
|
||||
|
||||
Args:
|
||||
workspace_root: The root directory of the workspace
|
||||
load_gitignore: Whether to load .gitignore patterns
|
||||
additional_ignore_patterns: Extra patterns to ignore
|
||||
protected_patterns: Patterns for read-only files
|
||||
"""
|
||||
self.workspace_root = os.path.abspath(workspace_root)
|
||||
self.ignore_patterns: List[str] = []
|
||||
self.protected_patterns: List[str] = protected_patterns or []
|
||||
self._ignored_cache: Set[str] = set()
|
||||
|
||||
# Load .gitignore patterns
|
||||
if load_gitignore:
|
||||
self.ignore_patterns.extend(load_gitignore_patterns(self.workspace_root))
|
||||
|
||||
# Add additional patterns
|
||||
if additional_ignore_patterns:
|
||||
self.ignore_patterns.extend(additional_ignore_patterns)
|
||||
|
||||
# Always ignore common patterns
|
||||
default_ignores = [
|
||||
'.git',
|
||||
'.git/**',
|
||||
'__pycache__',
|
||||
'__pycache__/**',
|
||||
'*.pyc',
|
||||
'.DS_Store',
|
||||
'node_modules',
|
||||
'node_modules/**',
|
||||
]
|
||||
for pattern in default_ignores:
|
||||
if pattern not in self.ignore_patterns:
|
||||
self.ignore_patterns.append(pattern)
|
||||
|
||||
def is_path_allowed(self, path: str) -> bool:
|
||||
"""
|
||||
Check if a path is allowed for access.
|
||||
|
||||
Args:
|
||||
path: The path to check
|
||||
|
||||
Returns:
|
||||
True if the path is allowed
|
||||
"""
|
||||
# Resolve to absolute path
|
||||
if not os.path.isabs(path):
|
||||
abs_path = os.path.abspath(os.path.join(self.workspace_root, path))
|
||||
else:
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Check if within workspace
|
||||
if not self._is_within_workspace(abs_path):
|
||||
return False
|
||||
|
||||
# Check if ignored
|
||||
is_dir = os.path.isdir(abs_path)
|
||||
if should_ignore_path(abs_path, self.ignore_patterns, self.workspace_root, is_dir):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def is_write_protected(self, path: str) -> bool:
|
||||
"""
|
||||
Check if a path is write-protected.
|
||||
|
||||
Args:
|
||||
path: The path to check
|
||||
|
||||
Returns:
|
||||
True if the path is write-protected
|
||||
"""
|
||||
if not self.protected_patterns:
|
||||
return False
|
||||
|
||||
# Get relative path
|
||||
if os.path.isabs(path):
|
||||
try:
|
||||
rel_path = os.path.relpath(path, self.workspace_root)
|
||||
except ValueError:
|
||||
rel_path = path
|
||||
else:
|
||||
rel_path = path
|
||||
|
||||
rel_path = rel_path.replace(os.sep, '/')
|
||||
|
||||
for pattern in self.protected_patterns:
|
||||
if fnmatch.fnmatch(rel_path, pattern):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def validate_access(self, path: str, write: bool = False) -> bool:
|
||||
"""
|
||||
Validate if access to a path is allowed.
|
||||
|
||||
Args:
|
||||
path: The path to validate
|
||||
write: Whether write access is needed
|
||||
|
||||
Returns:
|
||||
True if access is allowed
|
||||
"""
|
||||
if not self.is_path_allowed(path):
|
||||
return False
|
||||
|
||||
if write and self.is_write_protected(path):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def _is_within_workspace(self, abs_path: str) -> bool:
|
||||
"""
|
||||
Check if an absolute path is within the workspace.
|
||||
|
||||
Args:
|
||||
abs_path: The absolute path to check
|
||||
|
||||
Returns:
|
||||
True if within workspace
|
||||
"""
|
||||
# Normalize paths
|
||||
workspace = os.path.normpath(self.workspace_root)
|
||||
target = os.path.normpath(abs_path)
|
||||
|
||||
# Check if target starts with workspace path
|
||||
return target.startswith(workspace + os.sep) or target == workspace
|
||||
|
||||
def get_relative_path(self, path: str) -> str:
|
||||
"""
|
||||
Get the relative path from workspace root.
|
||||
|
||||
Args:
|
||||
path: The path to convert
|
||||
|
||||
Returns:
|
||||
Relative path string
|
||||
"""
|
||||
if os.path.isabs(path):
|
||||
try:
|
||||
return os.path.relpath(path, self.workspace_root)
|
||||
except ValueError:
|
||||
return path
|
||||
return path
|
||||
|
||||
def get_absolute_path(self, path: str) -> str:
|
||||
"""
|
||||
Get the absolute path for a relative path.
|
||||
|
||||
Args:
|
||||
path: The relative path
|
||||
|
||||
Returns:
|
||||
Absolute path string
|
||||
"""
|
||||
if os.path.isabs(path):
|
||||
return path
|
||||
return os.path.abspath(os.path.join(self.workspace_root, path))
|
||||
@@ -0,0 +1,276 @@
|
||||
"""
|
||||
Text utility functions for PraisonAI Code module.
|
||||
|
||||
Provides utilities for text processing including:
|
||||
- String normalization (smart quotes, unicode)
|
||||
- HTML entity unescaping
|
||||
- Similarity calculation using Levenshtein distance
|
||||
"""
|
||||
|
||||
import html
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
|
||||
def normalize_string(text: str) -> str:
|
||||
"""
|
||||
Normalize a string for comparison.
|
||||
|
||||
Handles:
|
||||
- Smart quotes to regular quotes
|
||||
- Unicode normalization
|
||||
- Consistent whitespace
|
||||
|
||||
Args:
|
||||
text: The text to normalize
|
||||
|
||||
Returns:
|
||||
Normalized text string
|
||||
"""
|
||||
if not text:
|
||||
return ""
|
||||
|
||||
# Unicode normalization (NFC form)
|
||||
normalized = unicodedata.normalize('NFC', text)
|
||||
|
||||
# Replace smart quotes with regular quotes
|
||||
quote_replacements = {
|
||||
'\u2018': "'", # Left single quote
|
||||
'\u2019': "'", # Right single quote
|
||||
'\u201C': '"', # Left double quote
|
||||
'\u201D': '"', # Right double quote
|
||||
'\u2013': '-', # En dash
|
||||
'\u2014': '-', # Em dash
|
||||
'\u2026': '...', # Ellipsis
|
||||
}
|
||||
|
||||
for smart, regular in quote_replacements.items():
|
||||
normalized = normalized.replace(smart, regular)
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def unescape_html_entities(text: str) -> str:
|
||||
"""
|
||||
Unescape HTML entities in text.
|
||||
|
||||
Args:
|
||||
text: Text with HTML entities
|
||||
|
||||
Returns:
|
||||
Text with entities converted to characters
|
||||
|
||||
Example:
|
||||
>>> unescape_html_entities("<div>")
|
||||
'<div>'
|
||||
"""
|
||||
if not text:
|
||||
return ""
|
||||
return html.unescape(text)
|
||||
|
||||
|
||||
def levenshtein_distance(s1: str, s2: str) -> int:
|
||||
"""
|
||||
Calculate the Levenshtein distance between two strings.
|
||||
|
||||
This is the minimum number of single-character edits (insertions,
|
||||
deletions, or substitutions) required to change one string into the other.
|
||||
|
||||
Args:
|
||||
s1: First string
|
||||
s2: Second string
|
||||
|
||||
Returns:
|
||||
The edit distance between the strings
|
||||
"""
|
||||
if len(s1) < len(s2):
|
||||
return levenshtein_distance(s2, s1)
|
||||
|
||||
if len(s2) == 0:
|
||||
return len(s1)
|
||||
|
||||
previous_row = range(len(s2) + 1)
|
||||
|
||||
for i, c1 in enumerate(s1):
|
||||
current_row = [i + 1]
|
||||
for j, c2 in enumerate(s2):
|
||||
# Cost is 0 if characters match, 1 otherwise
|
||||
insertions = previous_row[j + 1] + 1
|
||||
deletions = current_row[j] + 1
|
||||
substitutions = previous_row[j] + (c1 != c2)
|
||||
current_row.append(min(insertions, deletions, substitutions))
|
||||
previous_row = current_row
|
||||
|
||||
return previous_row[-1]
|
||||
|
||||
|
||||
def get_similarity(original: str, search: str) -> float:
|
||||
"""
|
||||
Calculate similarity ratio between two strings using Levenshtein distance.
|
||||
|
||||
Args:
|
||||
original: The original string
|
||||
search: The string to compare against
|
||||
|
||||
Returns:
|
||||
Similarity ratio from 0.0 (completely different) to 1.0 (identical)
|
||||
"""
|
||||
if not search:
|
||||
return 0.0
|
||||
|
||||
# Normalize both strings for comparison
|
||||
normalized_original = normalize_string(original)
|
||||
normalized_search = normalize_string(search)
|
||||
|
||||
if normalized_original == normalized_search:
|
||||
return 1.0
|
||||
|
||||
# Calculate Levenshtein distance
|
||||
dist = levenshtein_distance(normalized_original, normalized_search)
|
||||
|
||||
# Calculate similarity ratio
|
||||
max_length = max(len(normalized_original), len(normalized_search))
|
||||
if max_length == 0:
|
||||
return 1.0
|
||||
|
||||
return 1.0 - (dist / max_length)
|
||||
|
||||
|
||||
def fuzzy_search(
|
||||
lines: list,
|
||||
search_chunk: str,
|
||||
start_index: int,
|
||||
end_index: int
|
||||
) -> tuple:
|
||||
"""
|
||||
Perform a middle-out fuzzy search to find the best match for search_chunk.
|
||||
|
||||
Searches outward from the middle of the range to find the slice of lines
|
||||
that is most similar to the search chunk.
|
||||
|
||||
Args:
|
||||
lines: List of lines to search in
|
||||
search_chunk: The text to search for
|
||||
start_index: Start of search range (inclusive)
|
||||
end_index: End of search range (exclusive)
|
||||
|
||||
Returns:
|
||||
Tuple of (best_score, best_match_index, best_match_content)
|
||||
"""
|
||||
best_score = 0.0
|
||||
best_match_index = -1
|
||||
best_match_content = ""
|
||||
|
||||
search_lines = search_chunk.split('\n')
|
||||
search_len = len(search_lines)
|
||||
|
||||
if search_len == 0:
|
||||
return (0.0, -1, "")
|
||||
|
||||
# Middle-out search from the midpoint
|
||||
mid_point = (start_index + end_index) // 2
|
||||
left_index = mid_point
|
||||
right_index = mid_point + 1
|
||||
|
||||
while left_index >= start_index or right_index <= end_index - search_len:
|
||||
# Check left side
|
||||
if left_index >= start_index:
|
||||
original_chunk = '\n'.join(lines[left_index:left_index + search_len])
|
||||
similarity = get_similarity(original_chunk, search_chunk)
|
||||
if similarity > best_score:
|
||||
best_score = similarity
|
||||
best_match_index = left_index
|
||||
best_match_content = original_chunk
|
||||
left_index -= 1
|
||||
|
||||
# Check right side
|
||||
if right_index <= end_index - search_len:
|
||||
original_chunk = '\n'.join(lines[right_index:right_index + search_len])
|
||||
similarity = get_similarity(original_chunk, search_chunk)
|
||||
if similarity > best_score:
|
||||
best_score = similarity
|
||||
best_match_index = right_index
|
||||
best_match_content = original_chunk
|
||||
right_index += 1
|
||||
|
||||
return (best_score, best_match_index, best_match_content)
|
||||
|
||||
|
||||
def strip_markdown_code_fences(content: str) -> str:
|
||||
"""
|
||||
Strip markdown code fences from content.
|
||||
|
||||
Handles content that starts with ```language and ends with ```.
|
||||
|
||||
Args:
|
||||
content: Content potentially wrapped in code fences
|
||||
|
||||
Returns:
|
||||
Content with code fences removed
|
||||
"""
|
||||
if not content:
|
||||
return ""
|
||||
|
||||
lines = content.split('\n')
|
||||
|
||||
# Check if starts with code fence
|
||||
if lines and lines[0].strip().startswith('```'):
|
||||
lines = lines[1:]
|
||||
|
||||
# Check if ends with code fence
|
||||
if lines and lines[-1].strip() == '```':
|
||||
lines = lines[:-1]
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
def get_indentation(line: str) -> str:
|
||||
"""
|
||||
Extract the leading whitespace (indentation) from a line.
|
||||
|
||||
Args:
|
||||
line: The line to extract indentation from
|
||||
|
||||
Returns:
|
||||
The leading whitespace string (tabs and/or spaces)
|
||||
"""
|
||||
match = re.match(r'^[\t ]*', line)
|
||||
return match.group(0) if match else ""
|
||||
|
||||
|
||||
def preserve_indentation(original_lines: list, replacement_lines: list) -> list:
|
||||
"""
|
||||
Apply replacement lines while preserving the indentation from original.
|
||||
|
||||
Args:
|
||||
original_lines: The original lines being replaced
|
||||
replacement_lines: The new lines to insert
|
||||
|
||||
Returns:
|
||||
Replacement lines with adjusted indentation
|
||||
"""
|
||||
if not original_lines or not replacement_lines:
|
||||
return replacement_lines
|
||||
|
||||
# Get base indentation from first original line
|
||||
original_base_indent = get_indentation(original_lines[0])
|
||||
|
||||
# Get base indentation from first replacement line
|
||||
replacement_base_indent = get_indentation(replacement_lines[0])
|
||||
|
||||
result = []
|
||||
for line in replacement_lines:
|
||||
current_indent = get_indentation(line)
|
||||
content = line.lstrip()
|
||||
|
||||
# Calculate relative indentation
|
||||
if len(current_indent) >= len(replacement_base_indent):
|
||||
relative_indent = current_indent[len(replacement_base_indent):]
|
||||
new_indent = original_base_indent + relative_indent
|
||||
else:
|
||||
# Line has less indentation than base, adjust accordingly
|
||||
new_indent = original_base_indent[:max(0, len(original_base_indent) - (len(replacement_base_indent) - len(current_indent)))]
|
||||
|
||||
result.append(new_indent + content)
|
||||
|
||||
return result
|
||||
@@ -57,7 +57,7 @@ class CloudDeployer:
|
||||
file.write("FROM python:3.11-slim\n")
|
||||
file.write("WORKDIR /app\n")
|
||||
file.write("COPY . .\n")
|
||||
file.write("RUN pip install flask praisonai==2.3.36 gunicorn markdown\n")
|
||||
file.write("RUN pip install flask praisonai==2.3.41 gunicorn markdown\n")
|
||||
file.write("EXPOSE 8080\n")
|
||||
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "2.3.36"
|
||||
__version__ = "2.3.41"
|
||||
@@ -12,7 +12,7 @@ dependencies = [
|
||||
"rich>=13.7",
|
||||
"markdown>=3.5",
|
||||
"pyparsing>=3.0.0",
|
||||
"praisonaiagents>=0.1.14",
|
||||
"praisonaiagents>=0.1.17",
|
||||
"python-dotenv>=0.19.0",
|
||||
"instructor>=1.3.3",
|
||||
"PyYAML>=6.0",
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"$schema": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json",
|
||||
"name": "io.github.MervinPraison/praisonai",
|
||||
"title": "PraisonAI",
|
||||
"description": "AI Agents Framework with Self Reflection and MCP support",
|
||||
"repository": {
|
||||
"url": "https://github.com/MervinPraison/PraisonAI",
|
||||
"source": "github"
|
||||
},
|
||||
"version": "2.3.41",
|
||||
"packages": [
|
||||
{
|
||||
"registryType": "pypi",
|
||||
"identifier": "praisonai",
|
||||
"version": "2.3.41",
|
||||
"transport": {
|
||||
"type": "stdio"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
"""Tests for praisonai.code module."""
|
||||
@@ -0,0 +1,235 @@
|
||||
"""
|
||||
Integration tests for praisonai.code agent tools.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
class TestAgentTools:
|
||||
"""Tests for agent-compatible tool wrappers."""
|
||||
|
||||
def test_set_and_get_workspace(self):
|
||||
"""Test setting and getting workspace."""
|
||||
from praisonai.code import set_workspace, get_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
assert get_workspace() == temp_dir
|
||||
|
||||
def test_code_read_file(self):
|
||||
"""Test code_read_file agent tool."""
|
||||
from praisonai.code import code_read_file, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
# Create test file
|
||||
file_path = os.path.join(temp_dir, "test.py")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("def hello():\n print('hello')\n")
|
||||
|
||||
result = code_read_file("test.py")
|
||||
|
||||
assert "File: test.py" in result
|
||||
assert "def hello():" in result
|
||||
assert "1 |" in result # Line numbers
|
||||
|
||||
def test_code_write_file(self):
|
||||
"""Test code_write_file agent tool."""
|
||||
from praisonai.code import code_write_file, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
result = code_write_file("new_file.py", "print('hello')")
|
||||
|
||||
assert "Created" in result
|
||||
assert os.path.exists(os.path.join(temp_dir, "new_file.py"))
|
||||
|
||||
def test_code_list_files(self):
|
||||
"""Test code_list_files agent tool."""
|
||||
from praisonai.code import code_list_files, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
# Create some files
|
||||
open(os.path.join(temp_dir, "file1.py"), 'w').close()
|
||||
open(os.path.join(temp_dir, "file2.py"), 'w').close()
|
||||
os.makedirs(os.path.join(temp_dir, "subdir"))
|
||||
|
||||
result = code_list_files(".")
|
||||
|
||||
assert "file1.py" in result
|
||||
assert "file2.py" in result
|
||||
assert "subdir" in result
|
||||
|
||||
def test_code_apply_diff(self):
|
||||
"""Test code_apply_diff agent tool."""
|
||||
from praisonai.code import code_apply_diff, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
# Create test file
|
||||
file_path = os.path.join(temp_dir, "test.py")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("def old_function():\n return 'old'\n")
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def old_function():
|
||||
return 'old'
|
||||
=======
|
||||
def new_function():
|
||||
return 'new'
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = code_apply_diff("test.py", diff)
|
||||
|
||||
assert "Successfully applied" in result
|
||||
|
||||
# Verify file was modified
|
||||
with open(file_path, 'r') as f:
|
||||
content = f.read()
|
||||
assert "def new_function():" in content
|
||||
|
||||
def test_code_search_replace(self):
|
||||
"""Test code_search_replace agent tool."""
|
||||
from praisonai.code import code_search_replace, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
# Create test file
|
||||
file_path = os.path.join(temp_dir, "test.py")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("old_name = 'value'\nprint(old_name)\n")
|
||||
|
||||
result = code_search_replace("test.py", "old_name", "new_name")
|
||||
|
||||
assert "Replaced" in result
|
||||
assert "2" in result # 2 occurrences
|
||||
|
||||
def test_code_execute_command(self):
|
||||
"""Test code_execute_command agent tool."""
|
||||
from praisonai.code import code_execute_command, set_workspace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
result = code_execute_command("echo 'Hello from test'")
|
||||
|
||||
assert "completed successfully" in result
|
||||
assert "Hello" in result
|
||||
|
||||
def test_code_tools_list(self):
|
||||
"""Test CODE_TOOLS list contains all tools."""
|
||||
from praisonai.code import CODE_TOOLS
|
||||
|
||||
assert len(CODE_TOOLS) == 6
|
||||
|
||||
tool_names = [t.__name__ for t in CODE_TOOLS]
|
||||
assert "code_read_file" in tool_names
|
||||
assert "code_write_file" in tool_names
|
||||
assert "code_list_files" in tool_names
|
||||
assert "code_apply_diff" in tool_names
|
||||
assert "code_search_replace" in tool_names
|
||||
assert "code_execute_command" in tool_names
|
||||
|
||||
|
||||
class TestDiffHelpers:
|
||||
"""Tests for diff helper functions."""
|
||||
|
||||
def test_create_diff_block(self):
|
||||
"""Test creating a diff block."""
|
||||
from praisonai.code import create_diff_block
|
||||
|
||||
diff = create_diff_block(
|
||||
search_content="old content",
|
||||
replace_content="new content",
|
||||
start_line=10,
|
||||
)
|
||||
|
||||
assert "<<<<<<< SEARCH" in diff
|
||||
assert ":start_line:10" in diff
|
||||
assert "old content" in diff
|
||||
assert "=======" in diff
|
||||
assert "new content" in diff
|
||||
assert ">>>>>>> REPLACE" in diff
|
||||
|
||||
def test_create_multi_diff(self):
|
||||
"""Test creating multiple diff blocks."""
|
||||
from praisonai.code import create_multi_diff
|
||||
|
||||
blocks = [
|
||||
{'search': 'old1', 'replace': 'new1', 'start_line': 5},
|
||||
{'search': 'old2', 'replace': 'new2', 'start_line': 20},
|
||||
]
|
||||
|
||||
diff = create_multi_diff(blocks)
|
||||
|
||||
# Should contain two SEARCH/REPLACE blocks
|
||||
assert diff.count("<<<<<<< SEARCH") == 2
|
||||
assert diff.count(">>>>>>> REPLACE") == 2
|
||||
assert "old1" in diff
|
||||
assert "new1" in diff
|
||||
assert "old2" in diff
|
||||
assert "new2" in diff
|
||||
|
||||
|
||||
class TestEndToEndWorkflow:
|
||||
"""End-to-end workflow tests."""
|
||||
|
||||
def test_read_modify_write_workflow(self):
|
||||
"""Test a complete read-modify-write workflow."""
|
||||
from praisonai.code import (
|
||||
set_workspace, code_read_file, code_apply_diff, code_execute_command
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
set_workspace(temp_dir)
|
||||
|
||||
# Create a Python file
|
||||
file_path = os.path.join(temp_dir, "calculator.py")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write('''def add(a, b):
|
||||
return a + b
|
||||
|
||||
def subtract(a, b):
|
||||
return a - b
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(add(2, 3))
|
||||
''')
|
||||
|
||||
# Read the file
|
||||
content = code_read_file("calculator.py")
|
||||
assert "def add(a, b):" in content
|
||||
|
||||
# Apply a diff to add a multiply function
|
||||
diff = '''<<<<<<< SEARCH
|
||||
:start_line:4
|
||||
-------
|
||||
def subtract(a, b):
|
||||
return a - b
|
||||
=======
|
||||
def subtract(a, b):
|
||||
return a - b
|
||||
|
||||
def multiply(a, b):
|
||||
return a * b
|
||||
>>>>>>> REPLACE'''
|
||||
|
||||
result = code_apply_diff("calculator.py", diff)
|
||||
assert "Successfully applied" in result
|
||||
|
||||
# Verify the change
|
||||
content = code_read_file("calculator.py")
|
||||
assert "def multiply(a, b):" in content
|
||||
|
||||
# Run the file to verify it's valid Python
|
||||
result = code_execute_command("python calculator.py")
|
||||
assert "5" in result # add(2, 3) = 5
|
||||
@@ -0,0 +1,341 @@
|
||||
"""
|
||||
Tests for praisonai.code.diff.diff_strategy module.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
class TestValidateDiffFormat:
|
||||
"""Tests for validate_diff_format function."""
|
||||
|
||||
def test_valid_single_block(self):
|
||||
"""Test valid single SEARCH/REPLACE block."""
|
||||
from praisonai.code.diff.diff_strategy import validate_diff_format
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
old content
|
||||
=======
|
||||
new content
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
is_valid, error = validate_diff_format(diff)
|
||||
assert is_valid is True
|
||||
assert error is None
|
||||
|
||||
def test_valid_multiple_blocks(self):
|
||||
"""Test valid multiple SEARCH/REPLACE blocks."""
|
||||
from praisonai.code.diff.diff_strategy import validate_diff_format
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
old1
|
||||
=======
|
||||
new1
|
||||
>>>>>>> REPLACE
|
||||
|
||||
<<<<<<< SEARCH
|
||||
:start_line:10
|
||||
-------
|
||||
old2
|
||||
=======
|
||||
new2
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
is_valid, error = validate_diff_format(diff)
|
||||
assert is_valid is True
|
||||
|
||||
def test_missing_replace_marker(self):
|
||||
"""Test missing REPLACE marker."""
|
||||
from praisonai.code.diff.diff_strategy import validate_diff_format
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
old content
|
||||
=======
|
||||
new content"""
|
||||
|
||||
is_valid, error = validate_diff_format(diff)
|
||||
assert is_valid is False
|
||||
assert ">>>>>>> REPLACE" in error
|
||||
|
||||
def test_missing_separator(self):
|
||||
"""Test missing separator."""
|
||||
from praisonai.code.diff.diff_strategy import validate_diff_format
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
old content
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
is_valid, error = validate_diff_format(diff)
|
||||
assert is_valid is False
|
||||
|
||||
|
||||
class TestParseDiffBlocks:
|
||||
"""Tests for parse_diff_blocks function."""
|
||||
|
||||
def test_parse_single_block(self):
|
||||
"""Test parsing single block."""
|
||||
from praisonai.code.diff.diff_strategy import parse_diff_blocks
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:5
|
||||
-------
|
||||
def old():
|
||||
pass
|
||||
=======
|
||||
def new():
|
||||
return True
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
blocks, error = parse_diff_blocks(diff)
|
||||
|
||||
assert error is None
|
||||
assert len(blocks) == 1
|
||||
assert blocks[0].start_line == 5
|
||||
assert "def old():" in blocks[0].search_content
|
||||
assert "def new():" in blocks[0].replace_content
|
||||
|
||||
def test_parse_multiple_blocks(self):
|
||||
"""Test parsing multiple blocks."""
|
||||
from praisonai.code.diff.diff_strategy import parse_diff_blocks
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
first
|
||||
=======
|
||||
first_new
|
||||
>>>>>>> REPLACE
|
||||
|
||||
<<<<<<< SEARCH
|
||||
:start_line:10
|
||||
-------
|
||||
second
|
||||
=======
|
||||
second_new
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
blocks, error = parse_diff_blocks(diff)
|
||||
|
||||
assert error is None
|
||||
assert len(blocks) == 2
|
||||
# Should be sorted by start_line
|
||||
assert blocks[0].start_line == 1
|
||||
assert blocks[1].start_line == 10
|
||||
|
||||
def test_parse_block_without_start_line(self):
|
||||
"""Test parsing block without start_line hint."""
|
||||
from praisonai.code.diff.diff_strategy import parse_diff_blocks
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
-------
|
||||
old content
|
||||
=======
|
||||
new content
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
blocks, error = parse_diff_blocks(diff)
|
||||
|
||||
assert error is None
|
||||
assert len(blocks) == 1
|
||||
assert blocks[0].start_line is None
|
||||
|
||||
|
||||
class TestApplySearchReplaceDiff:
|
||||
"""Tests for apply_search_replace_diff function."""
|
||||
|
||||
def test_exact_match_replacement(self):
|
||||
"""Test exact match replacement."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """def hello():
|
||||
print("hello")
|
||||
|
||||
def world():
|
||||
print("world")"""
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def hello():
|
||||
print("hello")
|
||||
=======
|
||||
def hello():
|
||||
print("Hello, World!")
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_search_replace_diff(original, diff)
|
||||
|
||||
assert result.success is True
|
||||
assert result.applied_count == 1
|
||||
assert 'print("Hello, World!")' in result.content
|
||||
|
||||
def test_multiple_replacements(self):
|
||||
"""Test multiple replacements in one diff."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """line1
|
||||
line2
|
||||
line3
|
||||
line4
|
||||
line5"""
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
line1
|
||||
=======
|
||||
LINE1
|
||||
>>>>>>> REPLACE
|
||||
|
||||
<<<<<<< SEARCH
|
||||
:start_line:5
|
||||
-------
|
||||
line5
|
||||
=======
|
||||
LINE5
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_search_replace_diff(original, diff)
|
||||
|
||||
assert result.success is True
|
||||
assert result.applied_count == 2
|
||||
assert "LINE1" in result.content
|
||||
assert "LINE5" in result.content
|
||||
|
||||
def test_fuzzy_matching(self):
|
||||
"""Test fuzzy matching with threshold."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """def hello():
|
||||
print("hello")"""
|
||||
|
||||
# Slightly different content (extra space)
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def hello():
|
||||
print("hello")
|
||||
=======
|
||||
def hello():
|
||||
print("goodbye")
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
# With exact matching, this should work
|
||||
result = apply_search_replace_diff(original, diff, fuzzy_threshold=1.0)
|
||||
|
||||
assert result.success is True
|
||||
assert 'print("goodbye")' in result.content
|
||||
|
||||
def test_no_match_found(self):
|
||||
"""Test when no match is found."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """def hello():
|
||||
print("hello")"""
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def nonexistent():
|
||||
pass
|
||||
=======
|
||||
def new():
|
||||
pass
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_search_replace_diff(original, diff)
|
||||
|
||||
assert result.success is False
|
||||
assert len(result.failed_blocks) > 0
|
||||
|
||||
def test_indentation_preservation(self):
|
||||
"""Test that indentation is preserved."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """class MyClass:
|
||||
def method(self):
|
||||
pass"""
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:2
|
||||
-------
|
||||
def method(self):
|
||||
pass
|
||||
=======
|
||||
def method(self):
|
||||
return True
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_search_replace_diff(original, diff)
|
||||
|
||||
assert result.success is True
|
||||
# Check indentation is preserved
|
||||
lines = result.content.split('\n')
|
||||
assert lines[1].startswith(' ') # Method should be indented
|
||||
assert lines[2].startswith(' ') # Body should be double indented
|
||||
|
||||
def test_identical_search_replace_error(self):
|
||||
"""Test that identical search/replace is rejected."""
|
||||
from praisonai.code.diff.diff_strategy import apply_search_replace_diff
|
||||
|
||||
original = """def hello():
|
||||
pass"""
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def hello():
|
||||
pass
|
||||
=======
|
||||
def hello():
|
||||
pass
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_search_replace_diff(original, diff)
|
||||
|
||||
# Should fail because search and replace are identical
|
||||
assert result.success is False
|
||||
|
||||
|
||||
class TestDiffIntegration:
|
||||
"""Integration tests for diff operations."""
|
||||
|
||||
def test_apply_diff_to_file(self):
|
||||
"""Test applying diff to an actual file."""
|
||||
from praisonai.code.tools.apply_diff import apply_diff
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Create test file
|
||||
file_path = os.path.join(temp_dir, "test.py")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("""def old_function():
|
||||
return "old"
|
||||
|
||||
def another():
|
||||
pass""")
|
||||
|
||||
diff = """<<<<<<< SEARCH
|
||||
:start_line:1
|
||||
-------
|
||||
def old_function():
|
||||
return "old"
|
||||
=======
|
||||
def new_function():
|
||||
return "new"
|
||||
>>>>>>> REPLACE"""
|
||||
|
||||
result = apply_diff(file_path, diff)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['applied_count'] == 1
|
||||
|
||||
# Verify file was modified
|
||||
with open(file_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
assert "def new_function():" in content
|
||||
assert 'return "new"' in content
|
||||
@@ -0,0 +1,228 @@
|
||||
"""
|
||||
Tests for praisonai.code.utils.file_utils module.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
class TestAddLineNumbers:
|
||||
"""Tests for add_line_numbers function."""
|
||||
|
||||
def test_add_line_numbers_basic(self):
|
||||
"""Test basic line number addition."""
|
||||
from praisonai.code.utils.file_utils import add_line_numbers
|
||||
|
||||
content = "line1\nline2\nline3"
|
||||
result = add_line_numbers(content)
|
||||
|
||||
assert "1 | line1" in result
|
||||
assert "2 | line2" in result
|
||||
assert "3 | line3" in result
|
||||
|
||||
def test_add_line_numbers_custom_start(self):
|
||||
"""Test line numbers with custom start."""
|
||||
from praisonai.code.utils.file_utils import add_line_numbers
|
||||
|
||||
content = "line1\nline2"
|
||||
result = add_line_numbers(content, start_line=10)
|
||||
|
||||
assert "10 | line1" in result
|
||||
assert "11 | line2" in result
|
||||
|
||||
def test_add_line_numbers_empty(self):
|
||||
"""Test with empty content."""
|
||||
from praisonai.code.utils.file_utils import add_line_numbers
|
||||
|
||||
result = add_line_numbers("")
|
||||
assert result == ""
|
||||
|
||||
def test_add_line_numbers_single_line(self):
|
||||
"""Test with single line."""
|
||||
from praisonai.code.utils.file_utils import add_line_numbers
|
||||
|
||||
result = add_line_numbers("hello")
|
||||
assert "1 | hello" in result
|
||||
|
||||
|
||||
class TestStripLineNumbers:
|
||||
"""Tests for strip_line_numbers function."""
|
||||
|
||||
def test_strip_line_numbers_basic(self):
|
||||
"""Test basic line number stripping."""
|
||||
from praisonai.code.utils.file_utils import strip_line_numbers
|
||||
|
||||
content = " 1 | line1\n 2 | line2"
|
||||
result = strip_line_numbers(content)
|
||||
|
||||
assert result == "line1\nline2"
|
||||
|
||||
def test_strip_line_numbers_no_numbers(self):
|
||||
"""Test with content that has no line numbers."""
|
||||
from praisonai.code.utils.file_utils import strip_line_numbers
|
||||
|
||||
content = "line1\nline2"
|
||||
result = strip_line_numbers(content)
|
||||
|
||||
assert result == "line1\nline2"
|
||||
|
||||
def test_strip_line_numbers_tab_separator(self):
|
||||
"""Test with tab separator."""
|
||||
from praisonai.code.utils.file_utils import strip_line_numbers
|
||||
|
||||
content = "1\tline1\n2\tline2"
|
||||
result = strip_line_numbers(content)
|
||||
|
||||
assert result == "line1\nline2"
|
||||
|
||||
|
||||
class TestEveryLineHasLineNumbers:
|
||||
"""Tests for every_line_has_line_numbers function."""
|
||||
|
||||
def test_all_lines_have_numbers(self):
|
||||
"""Test when all lines have numbers."""
|
||||
from praisonai.code.utils.file_utils import every_line_has_line_numbers
|
||||
|
||||
content = " 1 | line1\n 2 | line2"
|
||||
assert every_line_has_line_numbers(content) is True
|
||||
|
||||
def test_no_lines_have_numbers(self):
|
||||
"""Test when no lines have numbers."""
|
||||
from praisonai.code.utils.file_utils import every_line_has_line_numbers
|
||||
|
||||
content = "line1\nline2"
|
||||
assert every_line_has_line_numbers(content) is False
|
||||
|
||||
def test_mixed_lines(self):
|
||||
"""Test when some lines have numbers."""
|
||||
from praisonai.code.utils.file_utils import every_line_has_line_numbers
|
||||
|
||||
content = " 1 | line1\nline2"
|
||||
assert every_line_has_line_numbers(content) is False
|
||||
|
||||
def test_empty_content(self):
|
||||
"""Test with empty content."""
|
||||
from praisonai.code.utils.file_utils import every_line_has_line_numbers
|
||||
|
||||
assert every_line_has_line_numbers("") is False
|
||||
|
||||
|
||||
class TestNormalizeLineEndings:
|
||||
"""Tests for normalize_line_endings function."""
|
||||
|
||||
def test_normalize_crlf_to_lf(self):
|
||||
"""Test converting CRLF to LF."""
|
||||
from praisonai.code.utils.file_utils import normalize_line_endings
|
||||
|
||||
content = "line1\r\nline2\r\n"
|
||||
result = normalize_line_endings(content, '\n')
|
||||
|
||||
assert result == "line1\nline2\n"
|
||||
assert '\r' not in result
|
||||
|
||||
def test_normalize_lf_to_crlf(self):
|
||||
"""Test converting LF to CRLF."""
|
||||
from praisonai.code.utils.file_utils import normalize_line_endings
|
||||
|
||||
content = "line1\nline2\n"
|
||||
result = normalize_line_endings(content, '\r\n')
|
||||
|
||||
assert result == "line1\r\nline2\r\n"
|
||||
|
||||
|
||||
class TestIsBinaryFile:
|
||||
"""Tests for is_binary_file function."""
|
||||
|
||||
def test_text_file(self):
|
||||
"""Test with a text file."""
|
||||
from praisonai.code.utils.file_utils import is_binary_file
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
||||
f.write("Hello, world!")
|
||||
temp_path = f.name
|
||||
|
||||
try:
|
||||
assert is_binary_file(temp_path) is False
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_binary_file(self):
|
||||
"""Test with a binary file."""
|
||||
from praisonai.code.utils.file_utils import is_binary_file
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode='wb', suffix='.bin', delete=False) as f:
|
||||
f.write(b'\x00\x01\x02\x03\x04\x05')
|
||||
temp_path = f.name
|
||||
|
||||
try:
|
||||
assert is_binary_file(temp_path) is True
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
|
||||
class TestFileExists:
|
||||
"""Tests for file_exists function."""
|
||||
|
||||
def test_existing_file(self):
|
||||
"""Test with existing file."""
|
||||
from praisonai.code.utils.file_utils import file_exists
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False) as f:
|
||||
temp_path = f.name
|
||||
|
||||
try:
|
||||
assert file_exists(temp_path) is True
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_nonexistent_file(self):
|
||||
"""Test with nonexistent file."""
|
||||
from praisonai.code.utils.file_utils import file_exists
|
||||
|
||||
assert file_exists("/nonexistent/path/file.txt") is False
|
||||
|
||||
def test_directory(self):
|
||||
"""Test with directory (should return False)."""
|
||||
from praisonai.code.utils.file_utils import file_exists
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
assert file_exists(temp_dir) is False
|
||||
|
||||
|
||||
class TestIsPathWithinDirectory:
|
||||
"""Tests for is_path_within_directory function."""
|
||||
|
||||
def test_path_within(self):
|
||||
"""Test path within directory."""
|
||||
from praisonai.code.utils.file_utils import is_path_within_directory
|
||||
|
||||
assert is_path_within_directory("/home/user/project/file.py", "/home/user/project") is True
|
||||
|
||||
def test_path_outside(self):
|
||||
"""Test path outside directory."""
|
||||
from praisonai.code.utils.file_utils import is_path_within_directory
|
||||
|
||||
assert is_path_within_directory("/home/other/file.py", "/home/user/project") is False
|
||||
|
||||
def test_path_traversal(self):
|
||||
"""Test path traversal attempt."""
|
||||
from praisonai.code.utils.file_utils import is_path_within_directory
|
||||
|
||||
# This should resolve to /home/other, which is outside /home/user/project
|
||||
assert is_path_within_directory("/home/user/project/../other/file.py", "/home/user/project") is False
|
||||
|
||||
|
||||
class TestCreateDirectoriesForFile:
|
||||
"""Tests for create_directories_for_file function."""
|
||||
|
||||
def test_create_directories(self):
|
||||
"""Test creating directories for a file path."""
|
||||
from praisonai.code.utils.file_utils import create_directories_for_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "subdir1", "subdir2", "file.txt")
|
||||
|
||||
result = create_directories_for_file(file_path)
|
||||
|
||||
assert result is True
|
||||
assert os.path.isdir(os.path.dirname(file_path))
|
||||
@@ -0,0 +1,379 @@
|
||||
"""
|
||||
Tests for praisonai.code.tools module.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
|
||||
class TestReadFile:
|
||||
"""Tests for read_file tool."""
|
||||
|
||||
def test_read_entire_file(self):
|
||||
"""Test reading entire file."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("line1\nline2\nline3")
|
||||
|
||||
result = read_file(file_path, add_line_nums=False)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['content'] == "line1\nline2\nline3"
|
||||
assert result['total_lines'] == 3
|
||||
|
||||
def test_read_with_line_numbers(self):
|
||||
"""Test reading with line numbers."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("line1\nline2\nline3")
|
||||
|
||||
result = read_file(file_path, add_line_nums=True)
|
||||
|
||||
assert result['success'] is True
|
||||
assert "1 | line1" in result['content']
|
||||
assert "2 | line2" in result['content']
|
||||
|
||||
def test_read_line_range(self):
|
||||
"""Test reading specific line range."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("line1\nline2\nline3\nline4\nline5")
|
||||
|
||||
result = read_file(file_path, start_line=2, end_line=4, add_line_nums=False)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['start_line'] == 2
|
||||
assert result['end_line'] == 4
|
||||
assert "line2" in result['content']
|
||||
assert "line4" in result['content']
|
||||
assert "line1" not in result['content']
|
||||
assert "line5" not in result['content']
|
||||
|
||||
def test_read_nonexistent_file(self):
|
||||
"""Test reading nonexistent file."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
result = read_file("/nonexistent/path/file.txt")
|
||||
|
||||
assert result['success'] is False
|
||||
assert "not found" in result['error'].lower()
|
||||
|
||||
def test_read_with_workspace(self):
|
||||
"""Test reading with workspace path."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "subdir", "test.txt")
|
||||
os.makedirs(os.path.dirname(file_path))
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("content")
|
||||
|
||||
result = read_file("subdir/test.txt", workspace=temp_dir, add_line_nums=False)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['content'] == "content"
|
||||
|
||||
def test_read_outside_workspace(self):
|
||||
"""Test reading file outside workspace is blocked."""
|
||||
from praisonai.code.tools.read_file import read_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
result = read_file("../../../etc/passwd", workspace=temp_dir)
|
||||
|
||||
assert result['success'] is False
|
||||
assert "outside" in result['error'].lower()
|
||||
|
||||
|
||||
class TestWriteFile:
|
||||
"""Tests for write_file tool."""
|
||||
|
||||
def test_write_new_file(self):
|
||||
"""Test writing a new file."""
|
||||
from praisonai.code.tools.write_file import write_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "new_file.txt")
|
||||
|
||||
result = write_file(file_path, "Hello, World!")
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['created'] is True
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
assert f.read() == "Hello, World!"
|
||||
|
||||
def test_overwrite_existing_file(self):
|
||||
"""Test overwriting an existing file."""
|
||||
from praisonai.code.tools.write_file import write_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "existing.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("old content")
|
||||
|
||||
result = write_file(file_path, "new content")
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['created'] is False
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
assert f.read() == "new content"
|
||||
|
||||
def test_create_directories(self):
|
||||
"""Test creating parent directories."""
|
||||
from praisonai.code.tools.write_file import write_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "a", "b", "c", "file.txt")
|
||||
|
||||
result = write_file(file_path, "content", create_directories=True)
|
||||
|
||||
assert result['success'] is True
|
||||
assert os.path.exists(file_path)
|
||||
|
||||
def test_strip_code_fences(self):
|
||||
"""Test stripping markdown code fences."""
|
||||
from praisonai.code.tools.write_file import write_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.py")
|
||||
|
||||
content = """```python
|
||||
def hello():
|
||||
pass
|
||||
```"""
|
||||
|
||||
result = write_file(file_path, content, strip_code_fences=True)
|
||||
|
||||
assert result['success'] is True
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
saved_content = f.read()
|
||||
|
||||
assert "```" not in saved_content
|
||||
assert "def hello():" in saved_content
|
||||
|
||||
def test_write_with_backup(self):
|
||||
"""Test writing with backup."""
|
||||
from praisonai.code.tools.write_file import write_file
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("original")
|
||||
|
||||
result = write_file(file_path, "modified", backup=True)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['backup_path'] is not None
|
||||
assert os.path.exists(result['backup_path'])
|
||||
|
||||
with open(result['backup_path'], 'r') as f:
|
||||
assert f.read() == "original"
|
||||
|
||||
|
||||
class TestListFiles:
|
||||
"""Tests for list_files tool."""
|
||||
|
||||
def test_list_files_basic(self):
|
||||
"""Test basic file listing."""
|
||||
from praisonai.code.tools.list_files import list_files
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Create some files
|
||||
open(os.path.join(temp_dir, "file1.txt"), 'w').close()
|
||||
open(os.path.join(temp_dir, "file2.py"), 'w').close()
|
||||
os.makedirs(os.path.join(temp_dir, "subdir"))
|
||||
|
||||
result = list_files(temp_dir)
|
||||
|
||||
assert result['success'] is True
|
||||
assert len(result['files']) == 2
|
||||
assert len(result['directories']) == 1
|
||||
|
||||
def test_list_files_recursive(self):
|
||||
"""Test recursive file listing."""
|
||||
from praisonai.code.tools.list_files import list_files
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Create nested structure
|
||||
os.makedirs(os.path.join(temp_dir, "a", "b"))
|
||||
open(os.path.join(temp_dir, "file1.txt"), 'w').close()
|
||||
open(os.path.join(temp_dir, "a", "file2.txt"), 'w').close()
|
||||
open(os.path.join(temp_dir, "a", "b", "file3.txt"), 'w').close()
|
||||
|
||||
result = list_files(temp_dir, recursive=True)
|
||||
|
||||
assert result['success'] is True
|
||||
assert len(result['files']) == 3
|
||||
|
||||
def test_list_files_with_extension_filter(self):
|
||||
"""Test filtering by extension."""
|
||||
from praisonai.code.tools.list_files import list_files
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
open(os.path.join(temp_dir, "file1.txt"), 'w').close()
|
||||
open(os.path.join(temp_dir, "file2.py"), 'w').close()
|
||||
open(os.path.join(temp_dir, "file3.py"), 'w').close()
|
||||
|
||||
result = list_files(temp_dir, extensions=['py'])
|
||||
|
||||
assert result['success'] is True
|
||||
assert len(result['files']) == 2
|
||||
assert all(f['name'].endswith('.py') for f in result['files'])
|
||||
|
||||
def test_list_files_exclude_hidden(self):
|
||||
"""Test excluding hidden files."""
|
||||
from praisonai.code.tools.list_files import list_files
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
open(os.path.join(temp_dir, "visible.txt"), 'w').close()
|
||||
open(os.path.join(temp_dir, ".hidden"), 'w').close()
|
||||
|
||||
result = list_files(temp_dir, include_hidden=False)
|
||||
|
||||
assert result['success'] is True
|
||||
assert len(result['files']) == 1
|
||||
assert result['files'][0]['name'] == "visible.txt"
|
||||
|
||||
def test_list_files_max_limit(self):
|
||||
"""Test max files limit."""
|
||||
from praisonai.code.tools.list_files import list_files
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Create many files
|
||||
for i in range(10):
|
||||
open(os.path.join(temp_dir, f"file{i}.txt"), 'w').close()
|
||||
|
||||
result = list_files(temp_dir, max_files=5)
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['total_count'] == 5
|
||||
assert result['truncated'] is True
|
||||
|
||||
|
||||
class TestSearchReplace:
|
||||
"""Tests for search_replace tool."""
|
||||
|
||||
def test_simple_replace(self):
|
||||
"""Test simple string replacement."""
|
||||
from praisonai.code.tools.search_replace import search_replace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("Hello, old world!")
|
||||
|
||||
result = search_replace(file_path, [
|
||||
{'search': 'old', 'replace': 'new'}
|
||||
])
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['total_replacements'] == 1
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
assert f.read() == "Hello, new world!"
|
||||
|
||||
def test_multiple_operations(self):
|
||||
"""Test multiple search/replace operations."""
|
||||
from praisonai.code.tools.search_replace import search_replace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("foo bar foo baz")
|
||||
|
||||
result = search_replace(file_path, [
|
||||
{'search': 'foo', 'replace': 'FOO'},
|
||||
{'search': 'bar', 'replace': 'BAR'},
|
||||
])
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['operations_applied'] == 2
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
assert "FOO" in content
|
||||
assert "BAR" in content
|
||||
|
||||
def test_regex_replace(self):
|
||||
"""Test regex replacement."""
|
||||
from praisonai.code.tools.search_replace import search_replace
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = os.path.join(temp_dir, "test.txt")
|
||||
with open(file_path, 'w') as f:
|
||||
f.write("def func1(): pass\ndef func2(): pass")
|
||||
|
||||
result = search_replace(file_path, [
|
||||
{'search': r'def (\w+)\(\)', 'replace': r'def renamed_\1()', 'is_regex': True}
|
||||
])
|
||||
|
||||
assert result['success'] is True
|
||||
|
||||
with open(file_path, 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
assert "renamed_func1" in content
|
||||
assert "renamed_func2" in content
|
||||
|
||||
|
||||
class TestExecuteCommand:
|
||||
"""Tests for execute_command tool."""
|
||||
|
||||
def test_simple_command(self):
|
||||
"""Test simple command execution."""
|
||||
from praisonai.code.tools.execute_command import execute_command
|
||||
|
||||
result = execute_command("echo 'Hello, World!'")
|
||||
|
||||
assert result['success'] is True
|
||||
assert result['exit_code'] == 0
|
||||
assert "Hello" in result['stdout']
|
||||
|
||||
def test_command_with_cwd(self):
|
||||
"""Test command with working directory."""
|
||||
from praisonai.code.tools.execute_command import execute_command
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
result = execute_command("pwd", cwd=temp_dir)
|
||||
|
||||
assert result['success'] is True
|
||||
assert temp_dir in result['stdout'] or os.path.basename(temp_dir) in result['stdout']
|
||||
|
||||
def test_command_failure(self):
|
||||
"""Test command that fails."""
|
||||
from praisonai.code.tools.execute_command import execute_command
|
||||
|
||||
result = execute_command("exit 1", shell=True)
|
||||
|
||||
assert result['success'] is False
|
||||
assert result['exit_code'] == 1
|
||||
|
||||
def test_command_timeout(self):
|
||||
"""Test command timeout."""
|
||||
from praisonai.code.tools.execute_command import execute_command
|
||||
|
||||
result = execute_command("sleep 10", timeout=1)
|
||||
|
||||
assert result['success'] is False
|
||||
assert "timed out" in result['error'].lower()
|
||||
|
||||
def test_is_safe_command(self):
|
||||
"""Test safe command detection."""
|
||||
from praisonai.code.tools.execute_command import is_safe_command
|
||||
|
||||
assert is_safe_command("ls -la") is True
|
||||
assert is_safe_command("echo hello") is True
|
||||
assert is_safe_command("rm -rf /") is False
|
||||
assert is_safe_command("sudo apt install") is False
|
||||
Generated
+4
-4
@@ -4197,7 +4197,7 @@ requires-dist = [
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.0.22" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.1.14" },
|
||||
{ name = "praisonaiagents", specifier = ">=0.1.17" },
|
||||
{ name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
|
||||
{ name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
|
||||
{ name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
|
||||
@@ -4252,7 +4252,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "praisonaiagents"
|
||||
version = "0.1.14"
|
||||
version = "0.1.17"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
@@ -4261,9 +4261,9 @@ dependencies = [
|
||||
{ name = "pydantic" },
|
||||
{ name = "rich" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/aa/19/3493e61f00afc8f2dc6beed5407dd5d18e36f70c96107aa0d8f23c136f26/praisonaiagents-0.1.14.tar.gz", hash = "sha256:3e86b548fa6a44f37b84fc9588d275ede34a303d2f5ac3d0c596f4e2464131b9", size = 621584, upload-time = "2025-12-18T14:39:37.194Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/3c/ed/fd26940dd2ee4506f2d6e178841b36d09f90f6a817ea052e0c1a9473f132/praisonaiagents-0.1.17.tar.gz", hash = "sha256:fa3bb9e77c6dbd16ed1fdb933d3669fac2208b1e58a2a5d5986fd630fef975c1", size = 647705, upload-time = "2025-12-18T21:52:30.244Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/f7/618f7278370091eb62c118331646010d139dd88a8477c8bea91e18ac274e/praisonaiagents-0.1.14-py3-none-any.whl", hash = "sha256:034c862bb293248c29f8bf6ce4a910cb17bc8b980dd5cc38d790a4cb16a08729", size = 548184, upload-time = "2025-12-18T14:39:34.989Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/94/40/261bbe4e0597442bd8d3aef9450c315063735970bb8f32f95cc4f0267b7b/praisonaiagents-0.1.17-py3-none-any.whl", hash = "sha256:61b98ede790e14fd7a21f4fc51729ba0a4b32ab1165f3ad9ee966757fa581c0f", size = 583244, upload-time = "2025-12-18T21:52:28.004Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
Reference in new issue
Block a user