Aksh Parekh
1 day ago
committed by
GitHub
commit
161fee9368
@ -1,20 +1,18 @@
|
|||||||
import orjson
|
import json
|
||||||
from dotenv import load_dotenv
|
|
||||||
|
|
||||||
from swarms.structs.auto_swarm_builder import AutoSwarmBuilder
|
from swarms import AutoSwarmBuilder
|
||||||
|
|
||||||
load_dotenv()
|
|
||||||
|
|
||||||
swarm = AutoSwarmBuilder(
|
swarm = AutoSwarmBuilder(
|
||||||
name="My Swarm",
|
name="My Swarm",
|
||||||
description="My Swarm Description",
|
description="My Swarm Description",
|
||||||
verbose=True,
|
verbose=True,
|
||||||
max_loops=1,
|
max_loops=1,
|
||||||
return_agents=True,
|
execution_type="return-agents",
|
||||||
|
model_name="gpt-4.1",
|
||||||
)
|
)
|
||||||
|
|
||||||
result = swarm.run(
|
result = swarm.run(
|
||||||
task="Build a swarm to write a research paper on the topic of AI"
|
task="Build a swarm to write a research paper on the topic of AI"
|
||||||
)
|
)
|
||||||
|
|
||||||
print(orjson.dumps(result, option=orjson.OPT_INDENT_2).decode())
|
print(json.dumps(result, indent=2))
|
||||||
@ -0,0 +1,171 @@
|
|||||||
|
# Medical AOP Example
|
||||||
|
|
||||||
|
A real-world demonstration of the Agent Orchestration Protocol (AOP) using medical agents deployed as MCP tools.
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
This example showcases how to:
|
||||||
|
- Deploy multiple medical agents as MCP tools via AOP
|
||||||
|
- Use discovery tools for dynamic agent collaboration
|
||||||
|
- Execute real tool calls with structured schemas
|
||||||
|
- Integrate with keyless APIs for enhanced context
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```mermaid
|
||||||
|
graph LR
|
||||||
|
A[Medical Agents] --> B[AOP MCP Server<br/>Port 8000]
|
||||||
|
B --> C[Client<br/>Cursor/Python]
|
||||||
|
B --> D[Discovery Tools]
|
||||||
|
B --> E[Tool Execution]
|
||||||
|
|
||||||
|
subgraph "Medical Agents"
|
||||||
|
F[Chief Medical Officer]
|
||||||
|
G[Virologist]
|
||||||
|
H[Internist]
|
||||||
|
I[Medical Coder]
|
||||||
|
J[Diagnostic Synthesizer]
|
||||||
|
end
|
||||||
|
|
||||||
|
A --> F
|
||||||
|
A --> G
|
||||||
|
A --> H
|
||||||
|
A --> I
|
||||||
|
A --> J
|
||||||
|
```
|
||||||
|
|
||||||
|
### Medical Agents
|
||||||
|
- **Chief Medical Officer**: Coordination, diagnosis, triage
|
||||||
|
- **Virologist**: Viral disease analysis and ICD-10 coding
|
||||||
|
- **Internist**: Internal medicine evaluation and HCC tagging
|
||||||
|
- **Medical Coder**: ICD-10 code assignment and compliance
|
||||||
|
- **Diagnostic Synthesizer**: Final report synthesis with confidence levels
|
||||||
|
|
||||||
|
## Files
|
||||||
|
|
||||||
|
| File | Description |
|
||||||
|
|------|-------------|
|
||||||
|
| `medical_aop/server.py` | AOP server exposing medical agents as MCP tools |
|
||||||
|
| `medical_aop/client.py` | Discovery client with real tool execution |
|
||||||
|
| `README.md` | This documentation |
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
### 1. Start the AOP Server
|
||||||
|
```bash
|
||||||
|
python -m examples.aop_examples.medical_aop.server
|
||||||
|
```
|
||||||
|
|
||||||
|
### 2. Configure Cursor MCP Integration
|
||||||
|
|
||||||
|
Add to `~/.cursor/mcp.json`:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"mcpServers": {
|
||||||
|
"Medical AOP": {
|
||||||
|
"type": "http",
|
||||||
|
"url": "http://localhost:8000/mcp"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 3. Use in Cursor
|
||||||
|
|
||||||
|
Enable "Medical AOP" in Cursor's MCP settings, then:
|
||||||
|
|
||||||
|
#### Discover agents:
|
||||||
|
```
|
||||||
|
Call tool discover_agents with: {}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Execute medical coding:
|
||||||
|
```
|
||||||
|
Call tool Medical Coder with: {"task":"Patient: 45M, egfr 59 ml/min/1.73; non-African American. Provide ICD-10 suggestions and coding notes.","priority":"normal","include_images":false}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Review infection control:
|
||||||
|
```
|
||||||
|
Call tool Chief Medical Officer with: {"task":"Review current hospital infection control protocols in light of recent MRSA outbreak in ICU. Provide executive summary, policy adjustment recommendations, and estimated implementation costs.","priority":"high"}
|
||||||
|
```
|
||||||
|
|
||||||
|
### 4. Run Python Client
|
||||||
|
```bash
|
||||||
|
python -m examples.aop_examples.medical_aop.client
|
||||||
|
```
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
### Structured Schemas
|
||||||
|
- Custom input/output schemas with validation
|
||||||
|
- Priority levels (low/normal/high)
|
||||||
|
- Image processing support
|
||||||
|
- Confidence scoring
|
||||||
|
|
||||||
|
### Discovery Tools
|
||||||
|
| Tool | Description |
|
||||||
|
|------|-------------|
|
||||||
|
| `discover_agents` | List all available agents |
|
||||||
|
| `get_agent_details` | Detailed agent information |
|
||||||
|
| `search_agents` | Keyword-based agent search |
|
||||||
|
| `list_agents` | Simple agent name list |
|
||||||
|
|
||||||
|
### Real-world Integration
|
||||||
|
- Keyless API integration (disease.sh for epidemiology data)
|
||||||
|
- Structured medical coding workflows
|
||||||
|
- Executive-level policy recommendations
|
||||||
|
- Cost estimation and implementation timelines
|
||||||
|
|
||||||
|
## Response Format
|
||||||
|
|
||||||
|
All tools return consistent JSON:
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"result": "Agent response text",
|
||||||
|
"success": true,
|
||||||
|
"error": null,
|
||||||
|
"confidence": 0.95,
|
||||||
|
"codes": ["N18.3", "Z51.11"]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
|
||||||
|
### Server Settings
|
||||||
|
| Setting | Value |
|
||||||
|
|---------|-------|
|
||||||
|
| Port | 8000 |
|
||||||
|
| Transport | streamable-http |
|
||||||
|
| Timeouts | 40-50 seconds per agent |
|
||||||
|
| Logging | INFO level with traceback enabled |
|
||||||
|
|
||||||
|
### Agent Metadata
|
||||||
|
Each agent includes:
|
||||||
|
- Tags for categorization
|
||||||
|
- Capabilities for matching
|
||||||
|
- Role classification
|
||||||
|
- Model configuration
|
||||||
|
|
||||||
|
## Best Practices
|
||||||
|
|
||||||
|
1. **Use structured inputs**: Leverage the custom schemas for better results
|
||||||
|
2. **Chain agents**: Pass results between agents for comprehensive analysis
|
||||||
|
3. **Monitor timeouts**: Adjust based on task complexity
|
||||||
|
4. **Validate responses**: Check the `success` field in all responses
|
||||||
|
5. **Use discovery**: Query available agents before hardcoding tool names
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
| Issue | Solution |
|
||||||
|
|-------|----------|
|
||||||
|
| Connection refused | Ensure server is running on port 8000 |
|
||||||
|
| Tool not found | Use `discover_agents` to verify available tools |
|
||||||
|
| Timeout errors | Increase timeout values for complex tasks |
|
||||||
|
| Schema validation | Ensure input matches the defined JSON schema |
|
||||||
|
|
||||||
|
## References
|
||||||
|
|
||||||
|
- [AOP Reference](https://docs.swarms.world/en/latest/swarms/structs/aop/)
|
||||||
|
- [MCP Integration](https://docs.swarms.ai/examples/mcp-integration)
|
||||||
|
- [Protocol Overview](https://docs.swarms.world/en/latest/protocol/overview/)
|
||||||
@ -0,0 +1,164 @@
|
|||||||
|
# AOP Server Setup Example
|
||||||
|
|
||||||
|
This example demonstrates how to set up an Agent Orchestration Protocol (AOP) server with multiple specialized agents.
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
The AOP server allows you to deploy multiple agents that can be discovered and called by other agents or clients in the network. This example shows how to create a server with specialized agents for different tasks.
|
||||||
|
|
||||||
|
## Code Example
|
||||||
|
|
||||||
|
```python
|
||||||
|
from swarms import Agent
|
||||||
|
from swarms.structs.aop import (
|
||||||
|
AOP,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create specialized agents
|
||||||
|
research_agent = Agent(
|
||||||
|
agent_name="Research-Agent",
|
||||||
|
agent_description="Expert in research, data collection, and information gathering",
|
||||||
|
model_name="anthropic/claude-sonnet-4-5",
|
||||||
|
max_loops=1,
|
||||||
|
top_p=None,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
system_prompt="""You are a research specialist. Your role is to:
|
||||||
|
1. Gather comprehensive information on any given topic
|
||||||
|
2. Analyze data from multiple sources
|
||||||
|
3. Provide well-structured research findings
|
||||||
|
4. Cite sources and maintain accuracy
|
||||||
|
5. Present findings in a clear, organized manner
|
||||||
|
|
||||||
|
Always provide detailed, factual information with proper context.""",
|
||||||
|
)
|
||||||
|
|
||||||
|
analysis_agent = Agent(
|
||||||
|
agent_name="Analysis-Agent",
|
||||||
|
agent_description="Expert in data analysis, pattern recognition, and generating insights",
|
||||||
|
model_name="anthropic/claude-sonnet-4-5",
|
||||||
|
max_loops=1,
|
||||||
|
top_p=None,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
system_prompt="""You are an analysis specialist. Your role is to:
|
||||||
|
1. Analyze data and identify patterns
|
||||||
|
2. Generate actionable insights
|
||||||
|
3. Create visualizations and summaries
|
||||||
|
4. Provide statistical analysis
|
||||||
|
5. Make data-driven recommendations
|
||||||
|
|
||||||
|
Focus on extracting meaningful insights from information.""",
|
||||||
|
)
|
||||||
|
|
||||||
|
writing_agent = Agent(
|
||||||
|
agent_name="Writing-Agent",
|
||||||
|
agent_description="Expert in content creation, editing, and communication",
|
||||||
|
model_name="anthropic/claude-sonnet-4-5",
|
||||||
|
max_loops=1,
|
||||||
|
top_p=None,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
system_prompt="""You are a writing specialist. Your role is to:
|
||||||
|
1. Create engaging, well-structured content
|
||||||
|
2. Edit and improve existing text
|
||||||
|
3. Adapt tone and style for different audiences
|
||||||
|
4. Ensure clarity and coherence
|
||||||
|
5. Follow best practices in writing
|
||||||
|
|
||||||
|
Always produce high-quality, professional content.""",
|
||||||
|
)
|
||||||
|
|
||||||
|
code_agent = Agent(
|
||||||
|
agent_name="Code-Agent",
|
||||||
|
agent_description="Expert in programming, code review, and software development",
|
||||||
|
model_name="anthropic/claude-sonnet-4-5",
|
||||||
|
max_loops=1,
|
||||||
|
top_p=None,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
system_prompt="""You are a coding specialist. Your role is to:
|
||||||
|
1. Write clean, efficient code
|
||||||
|
2. Debug and fix issues
|
||||||
|
3. Review and optimize code
|
||||||
|
4. Explain programming concepts
|
||||||
|
5. Follow best practices and standards
|
||||||
|
|
||||||
|
Always provide working, well-documented code.""",
|
||||||
|
)
|
||||||
|
|
||||||
|
financial_agent = Agent(
|
||||||
|
agent_name="Financial-Agent",
|
||||||
|
agent_description="Expert in financial analysis, market research, and investment insights",
|
||||||
|
model_name="anthropic/claude-sonnet-4-5",
|
||||||
|
max_loops=1,
|
||||||
|
top_p=None,
|
||||||
|
dynamic_temperature_enabled=True,
|
||||||
|
system_prompt="""You are a financial specialist. Your role is to:
|
||||||
|
1. Analyze financial data and markets
|
||||||
|
2. Provide investment insights
|
||||||
|
3. Assess risk and opportunities
|
||||||
|
4. Create financial reports
|
||||||
|
5. Explain complex financial concepts
|
||||||
|
|
||||||
|
Always provide accurate, well-reasoned financial analysis.""",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Basic usage - individual agent addition
|
||||||
|
deployer = AOP("MyAgentServer", verbose=True, port=5932)
|
||||||
|
|
||||||
|
agents = [
|
||||||
|
research_agent,
|
||||||
|
analysis_agent,
|
||||||
|
writing_agent,
|
||||||
|
code_agent,
|
||||||
|
financial_agent,
|
||||||
|
]
|
||||||
|
|
||||||
|
deployer.add_agents_batch(agents)
|
||||||
|
|
||||||
|
deployer.run()
|
||||||
|
```
|
||||||
|
|
||||||
|
## Key Components
|
||||||
|
|
||||||
|
### 1. Agent Creation
|
||||||
|
|
||||||
|
Each agent is created with:
|
||||||
|
|
||||||
|
- **agent_name**: Unique identifier for the agent
|
||||||
|
- **agent_description**: Brief description of the agent's capabilities
|
||||||
|
- **model_name**: The language model to use
|
||||||
|
- **system_prompt**: Detailed instructions defining the agent's role and behavior
|
||||||
|
|
||||||
|
### 2. AOP Server Setup
|
||||||
|
|
||||||
|
- **Server Name**: "MyAgentServer" - identifies your server
|
||||||
|
- **Port**: 5932 - the port where the server will run
|
||||||
|
- **Verbose**: True - enables detailed logging
|
||||||
|
|
||||||
|
### 3. Agent Registration
|
||||||
|
|
||||||
|
- **add_agents_batch()**: Registers multiple agents at once
|
||||||
|
- Agents become available for discovery and remote calls
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
1. **Start the Server**: Run the script to start the AOP server
|
||||||
|
2. **Agent Discovery**: Other agents or clients can discover available agents
|
||||||
|
3. **Remote Calls**: Agents can be called remotely by their names
|
||||||
|
|
||||||
|
## Server Features
|
||||||
|
|
||||||
|
- **Agent Discovery**: Automatically registers agents for network discovery
|
||||||
|
- **Remote Execution**: Agents can be called from other network nodes
|
||||||
|
- **Load Balancing**: Distributes requests across available agents
|
||||||
|
- **Health Monitoring**: Tracks agent status and availability
|
||||||
|
|
||||||
|
## Configuration Options
|
||||||
|
|
||||||
|
- **Port**: Change the port number as needed
|
||||||
|
- **Verbose**: Set to False for reduced logging
|
||||||
|
- **Server Name**: Use a descriptive name for your server
|
||||||
|
|
||||||
|
## Next Steps
|
||||||
|
|
||||||
|
- See [AOP Cluster Example](aop_cluster_example.md) for multi-server setups
|
||||||
|
- Check [AOP Reference](../structs/aop.md) for advanced configuration options
|
||||||
|
- Explore agent communication patterns in the examples directory
|
||||||
@ -1,68 +0,0 @@
|
|||||||
import json
|
|
||||||
import asyncio
|
|
||||||
|
|
||||||
from swarms.structs.aop import AOPCluster
|
|
||||||
from swarms.tools.mcp_client_tools import execute_tool_call_simple
|
|
||||||
|
|
||||||
|
|
||||||
async def discover_agents_example():
|
|
||||||
"""Example of how to call the discover_agents tool."""
|
|
||||||
|
|
||||||
# Create AOP cluster connection
|
|
||||||
aop_cluster = AOPCluster(
|
|
||||||
urls=["http://localhost:5932/mcp"],
|
|
||||||
transport="streamable-http",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Check if discover_agents tool is available
|
|
||||||
discover_tool = aop_cluster.find_tool_by_server_name(
|
|
||||||
"discover_agents"
|
|
||||||
)
|
|
||||||
if discover_tool:
|
|
||||||
try:
|
|
||||||
# Create the tool call request
|
|
||||||
tool_call_request = {
|
|
||||||
"type": "function",
|
|
||||||
"function": {
|
|
||||||
"name": "discover_agents",
|
|
||||||
"arguments": json.dumps(
|
|
||||||
{}
|
|
||||||
), # No specific agent name = get all
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
# Execute the tool call
|
|
||||||
result = await execute_tool_call_simple(
|
|
||||||
response=tool_call_request,
|
|
||||||
server_path="http://localhost:5932/mcp",
|
|
||||||
output_type="dict",
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
print(json.dumps(result, indent=2))
|
|
||||||
|
|
||||||
# Parse the result
|
|
||||||
if isinstance(result, list) and len(result) > 0:
|
|
||||||
discovery_data = result[0]
|
|
||||||
if discovery_data.get("success"):
|
|
||||||
agents = discovery_data.get("agents", [])
|
|
||||||
return agents
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
|
|
||||||
except Exception:
|
|
||||||
return None
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
"""Main function to run the discovery example."""
|
|
||||||
# Run the async function
|
|
||||||
return asyncio.run(discover_agents_example())
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@ -0,0 +1,47 @@
|
|||||||
|
import json
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from swarms.structs.aop import AOPCluster
|
||||||
|
from swarms.tools.mcp_client_tools import execute_tool_call_simple
|
||||||
|
|
||||||
|
|
||||||
|
async def discover_agents_example():
|
||||||
|
"""
|
||||||
|
Discover all agents using the AOPCluster and print the result.
|
||||||
|
"""
|
||||||
|
aop_cluster = AOPCluster(
|
||||||
|
urls=["http://localhost:5932/mcp"],
|
||||||
|
transport="streamable-http",
|
||||||
|
)
|
||||||
|
tool = aop_cluster.find_tool_by_server_name("discover_agents")
|
||||||
|
if not tool:
|
||||||
|
print("discover_agents tool not found.")
|
||||||
|
return None
|
||||||
|
|
||||||
|
tool_call_request = {
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "discover_agents",
|
||||||
|
"arguments": "{}",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await execute_tool_call_simple(
|
||||||
|
response=tool_call_request,
|
||||||
|
server_path="http://localhost:5932/mcp",
|
||||||
|
output_type="dict",
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
print(json.dumps(result, indent=2))
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""
|
||||||
|
Run the discover_agents_example coroutine.
|
||||||
|
"""
|
||||||
|
asyncio.run(discover_agents_example())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,149 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Example demonstrating the AOP queue system for agent execution.
|
||||||
|
|
||||||
|
This example shows how to use the new queue-based execution system
|
||||||
|
in the AOP framework for improved performance and reliability.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import time
|
||||||
|
from swarms import Agent
|
||||||
|
from swarms.structs.aop import AOP
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Demonstrate AOP queue functionality."""
|
||||||
|
|
||||||
|
# Create some sample agents
|
||||||
|
agent1 = Agent(
|
||||||
|
agent_name="Research Agent",
|
||||||
|
agent_description="Specialized in research tasks",
|
||||||
|
model_name="gpt-4",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
agent2 = Agent(
|
||||||
|
agent_name="Writing Agent",
|
||||||
|
agent_description="Specialized in writing tasks",
|
||||||
|
model_name="gpt-4",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create AOP with queue enabled
|
||||||
|
aop = AOP(
|
||||||
|
server_name="Queue Demo Cluster",
|
||||||
|
description="A demonstration of queue-based agent execution",
|
||||||
|
queue_enabled=True,
|
||||||
|
max_workers_per_agent=2, # 2 workers per agent
|
||||||
|
max_queue_size_per_agent=100, # Max 100 tasks per queue
|
||||||
|
processing_timeout=60, # 60 second timeout
|
||||||
|
retry_delay=2.0, # 2 second delay between retries
|
||||||
|
verbose=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Add agents to the cluster
|
||||||
|
print("Adding agents to cluster...")
|
||||||
|
aop.add_agent(agent1, tool_name="researcher")
|
||||||
|
aop.add_agent(agent2, tool_name="writer")
|
||||||
|
|
||||||
|
# Get initial queue stats
|
||||||
|
print("\nInitial queue stats:")
|
||||||
|
stats = aop.get_queue_stats()
|
||||||
|
print(f"Stats: {stats}")
|
||||||
|
|
||||||
|
# Add some tasks to the queues
|
||||||
|
print("\nAdding tasks to queues...")
|
||||||
|
|
||||||
|
# Add high priority research task
|
||||||
|
research_task_id = aop.task_queues["researcher"].add_task(
|
||||||
|
task="Research the latest developments in quantum computing",
|
||||||
|
priority=10, # High priority
|
||||||
|
max_retries=2,
|
||||||
|
)
|
||||||
|
print(f"Added research task: {research_task_id}")
|
||||||
|
|
||||||
|
# Add medium priority writing task
|
||||||
|
writing_task_id = aop.task_queues["writer"].add_task(
|
||||||
|
task="Write a summary of AI trends in 2024",
|
||||||
|
priority=5, # Medium priority
|
||||||
|
max_retries=3,
|
||||||
|
)
|
||||||
|
print(f"Added writing task: {writing_task_id}")
|
||||||
|
|
||||||
|
# Add multiple low priority tasks
|
||||||
|
for i in range(3):
|
||||||
|
task_id = aop.task_queues["researcher"].add_task(
|
||||||
|
task=f"Research task {i+1}: Analyze market trends",
|
||||||
|
priority=1, # Low priority
|
||||||
|
max_retries=1,
|
||||||
|
)
|
||||||
|
print(f"Added research task {i+1}: {task_id}")
|
||||||
|
|
||||||
|
# Get updated queue stats
|
||||||
|
print("\nUpdated queue stats:")
|
||||||
|
stats = aop.get_queue_stats()
|
||||||
|
print(f"Stats: {stats}")
|
||||||
|
|
||||||
|
# Monitor task progress
|
||||||
|
print("\nMonitoring task progress...")
|
||||||
|
for _ in range(10): # Monitor for 10 iterations
|
||||||
|
time.sleep(1)
|
||||||
|
|
||||||
|
# Check research task status
|
||||||
|
research_status = aop.get_task_status(
|
||||||
|
"researcher", research_task_id
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f"Research task status: {research_status['task']['status'] if research_status['success'] else 'Error'}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check writing task status
|
||||||
|
writing_status = aop.get_task_status(
|
||||||
|
"writer", writing_task_id
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f"Writing task status: {writing_status['task']['status'] if writing_status['success'] else 'Error'}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Get current queue stats
|
||||||
|
current_stats = aop.get_queue_stats()
|
||||||
|
if current_stats["success"]:
|
||||||
|
for agent_name, agent_stats in current_stats[
|
||||||
|
"stats"
|
||||||
|
].items():
|
||||||
|
print(
|
||||||
|
f"{agent_name}: {agent_stats['pending_tasks']} pending, {agent_stats['processing_tasks']} processing, {agent_stats['completed_tasks']} completed"
|
||||||
|
)
|
||||||
|
|
||||||
|
print("---")
|
||||||
|
|
||||||
|
# Demonstrate queue management
|
||||||
|
print("\nDemonstrating queue management...")
|
||||||
|
|
||||||
|
# Pause the research agent queue
|
||||||
|
print("Pausing research agent queue...")
|
||||||
|
aop.pause_agent_queue("researcher")
|
||||||
|
|
||||||
|
# Get queue status
|
||||||
|
research_queue_status = aop.task_queues["researcher"].get_status()
|
||||||
|
print(f"Research queue status: {research_queue_status.value}")
|
||||||
|
|
||||||
|
# Resume the research agent queue
|
||||||
|
print("Resuming research agent queue...")
|
||||||
|
aop.resume_agent_queue("researcher")
|
||||||
|
|
||||||
|
# Clear all queues
|
||||||
|
print("Clearing all queues...")
|
||||||
|
cleared = aop.clear_all_queues()
|
||||||
|
print(f"Cleared tasks: {cleared}")
|
||||||
|
|
||||||
|
# Final stats
|
||||||
|
print("\nFinal queue stats:")
|
||||||
|
final_stats = aop.get_queue_stats()
|
||||||
|
print(f"Final stats: {final_stats}")
|
||||||
|
|
||||||
|
print("\nQueue demonstration completed!")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,88 @@
|
|||||||
|
import json
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from swarms.structs.aop import AOPCluster
|
||||||
|
from swarms.tools.mcp_client_tools import execute_tool_call_simple
|
||||||
|
from mcp import ClientSession
|
||||||
|
from mcp.client.streamable_http import streamablehttp_client
|
||||||
|
|
||||||
|
|
||||||
|
async def discover_agents_example():
|
||||||
|
"""
|
||||||
|
Discover all agents using the AOPCluster and print the result.
|
||||||
|
"""
|
||||||
|
aop_cluster = AOPCluster(
|
||||||
|
urls=["http://localhost:5932/mcp"],
|
||||||
|
transport="streamable-http",
|
||||||
|
)
|
||||||
|
tool = aop_cluster.find_tool_by_server_name("discover_agents")
|
||||||
|
if not tool:
|
||||||
|
print("discover_agents tool not found.")
|
||||||
|
return None
|
||||||
|
|
||||||
|
tool_call_request = {
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "discover_agents",
|
||||||
|
"arguments": "{}",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
result = await execute_tool_call_simple(
|
||||||
|
response=tool_call_request,
|
||||||
|
server_path="http://localhost:5932/mcp",
|
||||||
|
output_type="dict",
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
print(json.dumps(result, indent=2))
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
async def raw_mcp_discover_agents_example():
|
||||||
|
"""
|
||||||
|
Call the MCP server directly using the raw MCP client to execute the
|
||||||
|
built-in "discover_agents" tool and print the JSON result.
|
||||||
|
|
||||||
|
This demonstrates how to:
|
||||||
|
- Initialize an MCP client over streamable HTTP
|
||||||
|
- List available tools (optional)
|
||||||
|
- Call a specific tool by name with arguments
|
||||||
|
"""
|
||||||
|
url = "http://localhost:5932/mcp"
|
||||||
|
|
||||||
|
# Open a raw MCP client connection
|
||||||
|
async with streamablehttp_client(url, timeout=10) as ctx:
|
||||||
|
if len(ctx) == 2:
|
||||||
|
read, write = ctx
|
||||||
|
else:
|
||||||
|
read, write, *_ = ctx
|
||||||
|
|
||||||
|
async with ClientSession(read, write) as session:
|
||||||
|
# Initialize the MCP session and optionally inspect tools
|
||||||
|
await session.initialize()
|
||||||
|
|
||||||
|
# Optional: list tools (uncomment to print)
|
||||||
|
# tools = await session.list_tools()
|
||||||
|
# print(json.dumps(tools.model_dump(), indent=2))
|
||||||
|
|
||||||
|
# Call the built-in discovery tool with empty arguments
|
||||||
|
result = await session.call_tool(
|
||||||
|
name="discover_agents",
|
||||||
|
arguments={},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert to dict for pretty printing
|
||||||
|
print(json.dumps(result.model_dump(), indent=2))
|
||||||
|
return result.model_dump()
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""
|
||||||
|
Run the helper-based and raw MCP client discovery examples.
|
||||||
|
"""
|
||||||
|
asyncio.run(discover_agents_example())
|
||||||
|
asyncio.run(raw_mcp_discover_agents_example())
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,107 @@
|
|||||||
|
import json
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from mcp import ClientSession
|
||||||
|
from mcp.client.streamable_http import streamablehttp_client
|
||||||
|
|
||||||
|
|
||||||
|
async def call_agent_tool_raw(
|
||||||
|
url: str,
|
||||||
|
tool_name: str,
|
||||||
|
task: str,
|
||||||
|
img: str | None = None,
|
||||||
|
imgs: list[str] | None = None,
|
||||||
|
correct_answer: str | None = None,
|
||||||
|
) -> dict:
|
||||||
|
"""
|
||||||
|
Call a specific agent tool on an MCP server using the raw MCP client.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
url: MCP server URL (e.g., "http://localhost:5932/mcp").
|
||||||
|
tool_name: Name of the tool/agent to invoke.
|
||||||
|
task: Task prompt to execute.
|
||||||
|
img: Optional single image path/URL.
|
||||||
|
imgs: Optional list of image paths/URLs.
|
||||||
|
correct_answer: Optional expected answer for validation.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A dict containing the tool's JSON response.
|
||||||
|
"""
|
||||||
|
# Open a raw MCP client connection over streamable HTTP
|
||||||
|
async with streamablehttp_client(url, timeout=30) as ctx:
|
||||||
|
if len(ctx) == 2:
|
||||||
|
read, write = ctx
|
||||||
|
else:
|
||||||
|
read, write, *_ = ctx
|
||||||
|
|
||||||
|
async with ClientSession(read, write) as session:
|
||||||
|
# Initialize the MCP session
|
||||||
|
await session.initialize()
|
||||||
|
|
||||||
|
# Prepare arguments in the canonical AOP tool format
|
||||||
|
arguments: dict = {"task": task}
|
||||||
|
if img is not None:
|
||||||
|
arguments["img"] = img
|
||||||
|
if imgs is not None:
|
||||||
|
arguments["imgs"] = imgs
|
||||||
|
if correct_answer is not None:
|
||||||
|
arguments["correct_answer"] = correct_answer
|
||||||
|
|
||||||
|
# Invoke the tool by name
|
||||||
|
result = await session.call_tool(
|
||||||
|
name=tool_name, arguments=arguments
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert to dict for return/printing
|
||||||
|
return result.model_dump()
|
||||||
|
|
||||||
|
|
||||||
|
async def list_available_tools(url: str) -> dict:
|
||||||
|
"""
|
||||||
|
List tools from an MCP server using the raw client.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
url: MCP server URL (e.g., "http://localhost:5932/mcp").
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A dict representation of the tools listing.
|
||||||
|
"""
|
||||||
|
async with streamablehttp_client(url, timeout=30) as ctx:
|
||||||
|
if len(ctx) == 2:
|
||||||
|
read, write = ctx
|
||||||
|
else:
|
||||||
|
read, write, *_ = ctx
|
||||||
|
|
||||||
|
async with ClientSession(read, write) as session:
|
||||||
|
await session.initialize()
|
||||||
|
tools = await session.list_tools()
|
||||||
|
return tools.model_dump()
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
"""
|
||||||
|
Demonstration entrypoint: list tools, then call a specified tool with a task.
|
||||||
|
"""
|
||||||
|
url = "http://localhost:5932/mcp"
|
||||||
|
tool_name = "Research-Agent" # Change to your agent tool name
|
||||||
|
task = "Summarize the latest advances in agent orchestration protocols."
|
||||||
|
|
||||||
|
# List tools
|
||||||
|
tools_info = asyncio.run(list_available_tools(url))
|
||||||
|
print("Available tools:")
|
||||||
|
print(json.dumps(tools_info, indent=2))
|
||||||
|
|
||||||
|
# Call the tool
|
||||||
|
print(f"\nCalling tool '{tool_name}' with task...\n")
|
||||||
|
result = asyncio.run(
|
||||||
|
call_agent_tool_raw(
|
||||||
|
url=url,
|
||||||
|
tool_name=tool_name,
|
||||||
|
task=task,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(json.dumps(result, indent=2))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,113 @@
|
|||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
from typing import Dict
|
||||||
|
|
||||||
|
import requests
|
||||||
|
|
||||||
|
from swarms.structs.aop import AOPCluster
|
||||||
|
from swarms.tools.mcp_client_tools import execute_tool_call_simple
|
||||||
|
|
||||||
|
|
||||||
|
def _select_tools_by_keyword(tools: list, keyword: str) -> list:
|
||||||
|
"""
|
||||||
|
Return tools whose name or description contains the keyword
|
||||||
|
(case-insensitive).
|
||||||
|
"""
|
||||||
|
kw = keyword.lower()
|
||||||
|
selected = []
|
||||||
|
for t in tools:
|
||||||
|
name = t.get("function", {}).get("name", "")
|
||||||
|
desc = t.get("function", {}).get("description", "")
|
||||||
|
if kw in name.lower() or kw in desc.lower():
|
||||||
|
selected.append(t)
|
||||||
|
return selected
|
||||||
|
|
||||||
|
|
||||||
|
def _example_payload_from_schema(tools: list, tool_name: str) -> dict:
|
||||||
|
"""
|
||||||
|
Construct a minimal example payload for a given tool using its JSON schema.
|
||||||
|
Falls back to a generic 'task' if schema not present.
|
||||||
|
"""
|
||||||
|
for t in tools:
|
||||||
|
fn = t.get("function", {})
|
||||||
|
if fn.get("name") == tool_name:
|
||||||
|
schema = fn.get("parameters", {})
|
||||||
|
required = schema.get("required", [])
|
||||||
|
props = schema.get("properties", {})
|
||||||
|
payload = {}
|
||||||
|
for r in required:
|
||||||
|
if r in props:
|
||||||
|
if props[r].get("type") == "string":
|
||||||
|
payload[r] = (
|
||||||
|
"Example patient case: 45M, egfr 59 ml/min/1.73"
|
||||||
|
)
|
||||||
|
elif props[r].get("type") == "boolean":
|
||||||
|
payload[r] = False
|
||||||
|
else:
|
||||||
|
payload[r] = None
|
||||||
|
if not payload:
|
||||||
|
payload = {
|
||||||
|
"task": "Provide ICD-10 suggestions for the case above"
|
||||||
|
}
|
||||||
|
return payload
|
||||||
|
return {"task": "Provide ICD-10 suggestions for the case above"}
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
cluster = AOPCluster(
|
||||||
|
urls=["http://localhost:8000/mcp"],
|
||||||
|
transport="streamable-http",
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = cluster.get_tools(output_type="dict")
|
||||||
|
print(f"Tools: {len(tools)}")
|
||||||
|
|
||||||
|
coding_tools = _select_tools_by_keyword(tools, "coder")
|
||||||
|
names = [t.get("function", {}).get("name") for t in coding_tools]
|
||||||
|
print(f"Coding-related tools: {names}")
|
||||||
|
|
||||||
|
# Build a real payload for "Medical Coder" and execute the tool call
|
||||||
|
tool_name = "Medical Coder"
|
||||||
|
payload: Dict[str, object] = _example_payload_from_schema(tools, tool_name)
|
||||||
|
|
||||||
|
# Enrich with public keyless data (epidemiology context via disease.sh)
|
||||||
|
try:
|
||||||
|
epi = requests.get(
|
||||||
|
"https://disease.sh/v3/covid-19/countries/USA?strict=true",
|
||||||
|
timeout=5,
|
||||||
|
)
|
||||||
|
if epi.ok:
|
||||||
|
data = epi.json()
|
||||||
|
epi_summary = (
|
||||||
|
f"US COVID-19 context: cases={data.get('cases')}, "
|
||||||
|
f"todayCases={data.get('todayCases')}, deaths={data.get('deaths')}"
|
||||||
|
)
|
||||||
|
base_task = payload.get("task") or ""
|
||||||
|
payload["task"] = (
|
||||||
|
f"{base_task}\n\nEpidemiology context (no key API): {epi_summary}"
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
print("Calling tool:", tool_name)
|
||||||
|
request = {
|
||||||
|
"function": {
|
||||||
|
"name": tool_name,
|
||||||
|
"arguments": payload,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
result = asyncio.run(
|
||||||
|
execute_tool_call_simple(
|
||||||
|
response=request,
|
||||||
|
server_path="http://localhost:8000/mcp",
|
||||||
|
output_type="json",
|
||||||
|
transport="streamable-http",
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print("Response:")
|
||||||
|
print(result)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -0,0 +1,166 @@
|
|||||||
|
# Import medical agents defined in the demo module
|
||||||
|
from examples.demos.medical.medical_coder_agent import (chief_medical_officer,
|
||||||
|
internist,
|
||||||
|
medical_coder,
|
||||||
|
synthesizer,
|
||||||
|
virologist)
|
||||||
|
from swarms.structs.aop import AOP
|
||||||
|
|
||||||
|
|
||||||
|
def _enrich_agents_metadata() -> None:
|
||||||
|
"""
|
||||||
|
Add lightweight tags/capabilities/roles to imported agents for
|
||||||
|
better discovery results.
|
||||||
|
"""
|
||||||
|
chief_medical_officer.tags = [
|
||||||
|
"coordination",
|
||||||
|
"diagnosis",
|
||||||
|
"triage",
|
||||||
|
]
|
||||||
|
chief_medical_officer.capabilities = [
|
||||||
|
"case-intake",
|
||||||
|
"differential",
|
||||||
|
"planning",
|
||||||
|
]
|
||||||
|
chief_medical_officer.role = "coordinator"
|
||||||
|
|
||||||
|
virologist.tags = ["virology", "infectious-disease"]
|
||||||
|
virologist.capabilities = ["viral-analysis", "icd10-suggestion"]
|
||||||
|
virologist.role = "specialist"
|
||||||
|
|
||||||
|
internist.tags = ["internal-medicine", "evaluation"]
|
||||||
|
internist.capabilities = [
|
||||||
|
"system-review",
|
||||||
|
"hcc-tagging",
|
||||||
|
"risk-stratification",
|
||||||
|
]
|
||||||
|
internist.role = "specialist"
|
||||||
|
|
||||||
|
medical_coder.tags = ["coding", "icd10", "compliance"]
|
||||||
|
medical_coder.capabilities = [
|
||||||
|
"code-assignment",
|
||||||
|
"documentation-review",
|
||||||
|
]
|
||||||
|
medical_coder.role = "coder"
|
||||||
|
|
||||||
|
synthesizer.tags = ["synthesis", "reporting"]
|
||||||
|
synthesizer.capabilities = [
|
||||||
|
"evidence-reconciliation",
|
||||||
|
"final-report",
|
||||||
|
]
|
||||||
|
synthesizer.role = "synthesizer"
|
||||||
|
|
||||||
|
|
||||||
|
def _medical_input_schema() -> dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"task": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Patient case or instruction for the agent",
|
||||||
|
},
|
||||||
|
"priority": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["low", "normal", "high"],
|
||||||
|
"description": "Processing priority",
|
||||||
|
},
|
||||||
|
"include_images": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Whether to consider linked images if provided",
|
||||||
|
"default": False,
|
||||||
|
},
|
||||||
|
"img": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Optional image path/URL",
|
||||||
|
},
|
||||||
|
"imgs": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {"type": "string"},
|
||||||
|
"description": "Optional list of images",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["task"],
|
||||||
|
"additionalProperties": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _medical_output_schema() -> dict:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"result": {"type": "string"},
|
||||||
|
"success": {"type": "boolean"},
|
||||||
|
"error": {"type": "string"},
|
||||||
|
"confidence": {
|
||||||
|
"type": "number",
|
||||||
|
"minimum": 0,
|
||||||
|
"maximum": 1,
|
||||||
|
"description": "Optional confidence in the assessment",
|
||||||
|
},
|
||||||
|
"codes": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {"type": "string"},
|
||||||
|
"description": "Optional list of suggested ICD-10 codes",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["result", "success"],
|
||||||
|
"additionalProperties": True,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
"""
|
||||||
|
Start an AOP MCP server that exposes the medical agents as tools with
|
||||||
|
structured schemas and per-agent settings.
|
||||||
|
"""
|
||||||
|
_enrich_agents_metadata()
|
||||||
|
|
||||||
|
deployer = AOP(
|
||||||
|
server_name="Medical-AOP-Server",
|
||||||
|
port=8000,
|
||||||
|
verbose=False,
|
||||||
|
traceback_enabled=True,
|
||||||
|
log_level="INFO",
|
||||||
|
transport="streamable-http",
|
||||||
|
)
|
||||||
|
|
||||||
|
input_schema = _medical_input_schema()
|
||||||
|
output_schema = _medical_output_schema()
|
||||||
|
|
||||||
|
# Register each agent with a modest, role-appropriate timeout
|
||||||
|
deployer.add_agent(
|
||||||
|
chief_medical_officer,
|
||||||
|
timeout=45,
|
||||||
|
input_schema=input_schema,
|
||||||
|
output_schema=output_schema,
|
||||||
|
)
|
||||||
|
deployer.add_agent(
|
||||||
|
virologist,
|
||||||
|
timeout=40,
|
||||||
|
input_schema=input_schema,
|
||||||
|
output_schema=output_schema,
|
||||||
|
)
|
||||||
|
deployer.add_agent(
|
||||||
|
internist,
|
||||||
|
timeout=40,
|
||||||
|
input_schema=input_schema,
|
||||||
|
output_schema=output_schema,
|
||||||
|
)
|
||||||
|
deployer.add_agent(
|
||||||
|
medical_coder,
|
||||||
|
timeout=50,
|
||||||
|
input_schema=input_schema,
|
||||||
|
output_schema=output_schema,
|
||||||
|
)
|
||||||
|
deployer.add_agent(
|
||||||
|
synthesizer,
|
||||||
|
timeout=45,
|
||||||
|
input_schema=input_schema,
|
||||||
|
output_schema=output_schema,
|
||||||
|
)
|
||||||
|
|
||||||
|
deployer.run()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@ -1,122 +0,0 @@
|
|||||||
"""
|
|
||||||
Example demonstrating the use of uvloop for running multiple agents concurrently.
|
|
||||||
|
|
||||||
This example shows how to use the new uvloop-based functions:
|
|
||||||
- run_agents_concurrently_uvloop: For running multiple agents with the same task
|
|
||||||
- run_agents_with_tasks_uvloop: For running agents with different tasks
|
|
||||||
|
|
||||||
uvloop provides significant performance improvements over standard asyncio,
|
|
||||||
especially for I/O-bound operations and concurrent task execution.
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
from swarms.structs.multi_agent_exec import (
|
|
||||||
run_agents_concurrently_uvloop,
|
|
||||||
run_agents_with_tasks_uvloop,
|
|
||||||
)
|
|
||||||
from swarms.structs.agent import Agent
|
|
||||||
|
|
||||||
|
|
||||||
def create_example_agents(num_agents: int = 3):
|
|
||||||
"""Create example agents for demonstration."""
|
|
||||||
agents = []
|
|
||||||
for i in range(num_agents):
|
|
||||||
agent = Agent(
|
|
||||||
agent_name=f"Agent_{i+1}",
|
|
||||||
system_prompt=f"You are Agent {i+1}, a helpful AI assistant.",
|
|
||||||
model_name="gpt-4o-mini", # Using a lightweight model for examples
|
|
||||||
max_loops=1,
|
|
||||||
autosave=False,
|
|
||||||
verbose=False,
|
|
||||||
)
|
|
||||||
agents.append(agent)
|
|
||||||
return agents
|
|
||||||
|
|
||||||
|
|
||||||
def example_same_task():
|
|
||||||
"""Example: Running multiple agents with the same task using uvloop."""
|
|
||||||
print("=== Example 1: Same Task for All Agents (uvloop) ===")
|
|
||||||
|
|
||||||
agents = create_example_agents(3)
|
|
||||||
task = (
|
|
||||||
"Write a one-sentence summary about artificial intelligence."
|
|
||||||
)
|
|
||||||
|
|
||||||
print(f"Running {len(agents)} agents with the same task...")
|
|
||||||
print(f"Task: {task}")
|
|
||||||
|
|
||||||
try:
|
|
||||||
results = run_agents_concurrently_uvloop(agents, task)
|
|
||||||
|
|
||||||
print("\nResults:")
|
|
||||||
for i, result in enumerate(results, 1):
|
|
||||||
print(f"Agent {i}: {result}")
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Error: {e}")
|
|
||||||
|
|
||||||
|
|
||||||
def example_different_tasks():
|
|
||||||
"""Example: Running agents with different tasks using uvloop."""
|
|
||||||
print(
|
|
||||||
"\n=== Example 2: Different Tasks for Each Agent (uvloop) ==="
|
|
||||||
)
|
|
||||||
|
|
||||||
agents = create_example_agents(3)
|
|
||||||
tasks = [
|
|
||||||
"Explain what machine learning is in simple terms.",
|
|
||||||
"Describe the benefits of cloud computing.",
|
|
||||||
"What are the main challenges in natural language processing?",
|
|
||||||
]
|
|
||||||
|
|
||||||
print(f"Running {len(agents)} agents with different tasks...")
|
|
||||||
|
|
||||||
try:
|
|
||||||
results = run_agents_with_tasks_uvloop(agents, tasks)
|
|
||||||
|
|
||||||
print("\nResults:")
|
|
||||||
for i, (result, task) in enumerate(zip(results, tasks), 1):
|
|
||||||
print(f"Agent {i} (Task: {task[:50]}...):")
|
|
||||||
print(f" Response: {result}")
|
|
||||||
print()
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Error: {e}")
|
|
||||||
|
|
||||||
|
|
||||||
def performance_comparison():
|
|
||||||
"""Demonstrate the performance benefit of uvloop vs standard asyncio."""
|
|
||||||
print("\n=== Performance Comparison ===")
|
|
||||||
|
|
||||||
# Note: This is a conceptual example. In practice, you'd need to measure actual performance
|
|
||||||
print("uvloop vs Standard asyncio:")
|
|
||||||
print("• uvloop: Cython-based event loop, ~2-4x faster")
|
|
||||||
print("• Better for I/O-bound operations")
|
|
||||||
print("• Lower latency and higher throughput")
|
|
||||||
print("• Especially beneficial for concurrent agent execution")
|
|
||||||
print("• Automatic fallback to asyncio if uvloop unavailable")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
# Check if API key is available
|
|
||||||
if not os.getenv("OPENAI_API_KEY"):
|
|
||||||
print(
|
|
||||||
"Please set your OPENAI_API_KEY environment variable to run this example."
|
|
||||||
)
|
|
||||||
print("Example: export OPENAI_API_KEY='your-api-key-here'")
|
|
||||||
exit(1)
|
|
||||||
|
|
||||||
print("🚀 uvloop Multi-Agent Execution Examples")
|
|
||||||
print("=" * 50)
|
|
||||||
|
|
||||||
# Run examples
|
|
||||||
example_same_task()
|
|
||||||
example_different_tasks()
|
|
||||||
performance_comparison()
|
|
||||||
|
|
||||||
print("\n✅ Examples completed!")
|
|
||||||
print("\nTo use uvloop functions in your code:")
|
|
||||||
print(
|
|
||||||
"from swarms.structs.multi_agent_exec import run_agents_concurrently_uvloop"
|
|
||||||
)
|
|
||||||
print("results = run_agents_concurrently_uvloop(agents, task)")
|
|
||||||
@ -0,0 +1,30 @@
|
|||||||
|
from swarms.structs.agent import Agent
|
||||||
|
from swarms.structs.multi_agent_exec import (
|
||||||
|
run_agents_concurrently_uvloop,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def create_example_agents(num_agents: int = 3):
|
||||||
|
"""Create example agents for demonstration."""
|
||||||
|
agents = []
|
||||||
|
for i in range(num_agents):
|
||||||
|
agent = Agent(
|
||||||
|
agent_name=f"Agent_{i+1}",
|
||||||
|
system_prompt=f"You are Agent {i+1}, a helpful AI assistant.",
|
||||||
|
model_name="gpt-4o-mini", # Using a lightweight model for examples
|
||||||
|
max_loops=1,
|
||||||
|
autosave=False,
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
agents.append(agent)
|
||||||
|
return agents
|
||||||
|
|
||||||
|
|
||||||
|
agents = create_example_agents(3)
|
||||||
|
|
||||||
|
task = "Write a one-sentence summary about artificial intelligence."
|
||||||
|
|
||||||
|
|
||||||
|
results = run_agents_concurrently_uvloop(agents, task)
|
||||||
|
|
||||||
|
print(results)
|
||||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,140 @@
|
|||||||
|
"""
|
||||||
|
Custom docstring parser implementation to replace the docstring_parser package.
|
||||||
|
|
||||||
|
This module provides a simple docstring parser that extracts parameter information
|
||||||
|
and descriptions from Python docstrings in Google/NumPy style format.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import List, Optional, NamedTuple
|
||||||
|
|
||||||
|
|
||||||
|
class DocstringParam(NamedTuple):
|
||||||
|
"""Represents a parameter in a docstring."""
|
||||||
|
|
||||||
|
arg_name: str
|
||||||
|
description: str
|
||||||
|
|
||||||
|
|
||||||
|
class DocstringInfo(NamedTuple):
|
||||||
|
"""Represents parsed docstring information."""
|
||||||
|
|
||||||
|
short_description: Optional[str]
|
||||||
|
params: List[DocstringParam]
|
||||||
|
|
||||||
|
|
||||||
|
def parse(docstring: str) -> DocstringInfo:
|
||||||
|
"""
|
||||||
|
Parse a docstring and extract parameter information and description.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
docstring (str): The docstring to parse.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
DocstringInfo: Parsed docstring information containing short description and parameters.
|
||||||
|
"""
|
||||||
|
if not docstring or not docstring.strip():
|
||||||
|
return DocstringInfo(short_description=None, params=[])
|
||||||
|
|
||||||
|
# Clean up the docstring
|
||||||
|
lines = [line.strip() for line in docstring.strip().split("\n")]
|
||||||
|
|
||||||
|
# Extract short description (first non-empty line that's not a section header)
|
||||||
|
short_description = None
|
||||||
|
for line in lines:
|
||||||
|
if line and not line.startswith(
|
||||||
|
(
|
||||||
|
"Args:",
|
||||||
|
"Parameters:",
|
||||||
|
"Returns:",
|
||||||
|
"Yields:",
|
||||||
|
"Raises:",
|
||||||
|
"Note:",
|
||||||
|
"Example:",
|
||||||
|
"Examples:",
|
||||||
|
)
|
||||||
|
):
|
||||||
|
short_description = line
|
||||||
|
break
|
||||||
|
|
||||||
|
# Extract parameters
|
||||||
|
params = []
|
||||||
|
|
||||||
|
# Look for Args: or Parameters: section
|
||||||
|
in_args_section = False
|
||||||
|
current_param = None
|
||||||
|
|
||||||
|
for line in lines:
|
||||||
|
# Check if we're entering the Args/Parameters section
|
||||||
|
if line.lower().startswith(("args:", "parameters:")):
|
||||||
|
in_args_section = True
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Check if we're leaving the Args/Parameters section
|
||||||
|
if (
|
||||||
|
in_args_section
|
||||||
|
and line
|
||||||
|
and not line.startswith(" ")
|
||||||
|
and not line.startswith("\t")
|
||||||
|
):
|
||||||
|
# Check if this is a new section header
|
||||||
|
if line.lower().startswith(
|
||||||
|
(
|
||||||
|
"returns:",
|
||||||
|
"yields:",
|
||||||
|
"raises:",
|
||||||
|
"note:",
|
||||||
|
"example:",
|
||||||
|
"examples:",
|
||||||
|
"see also:",
|
||||||
|
"see_also:",
|
||||||
|
)
|
||||||
|
):
|
||||||
|
in_args_section = False
|
||||||
|
if current_param:
|
||||||
|
params.append(current_param)
|
||||||
|
current_param = None
|
||||||
|
continue
|
||||||
|
|
||||||
|
if in_args_section and line:
|
||||||
|
# Check if this line starts a new parameter (starts with parameter name)
|
||||||
|
# Pattern: param_name (type): description
|
||||||
|
param_match = re.match(
|
||||||
|
r"^(\w+)\s*(?:\([^)]*\))?\s*:\s*(.+)$", line
|
||||||
|
)
|
||||||
|
if param_match:
|
||||||
|
# Save previous parameter if exists
|
||||||
|
if current_param:
|
||||||
|
params.append(current_param)
|
||||||
|
|
||||||
|
param_name = param_match.group(1)
|
||||||
|
param_desc = param_match.group(2).strip()
|
||||||
|
current_param = DocstringParam(
|
||||||
|
arg_name=param_name, description=param_desc
|
||||||
|
)
|
||||||
|
elif current_param and (
|
||||||
|
line.startswith(" ") or line.startswith("\t")
|
||||||
|
):
|
||||||
|
# This is a continuation of the current parameter description
|
||||||
|
current_param = DocstringParam(
|
||||||
|
arg_name=current_param.arg_name,
|
||||||
|
description=current_param.description
|
||||||
|
+ " "
|
||||||
|
+ line.strip(),
|
||||||
|
)
|
||||||
|
elif not line.startswith(" ") and not line.startswith(
|
||||||
|
"\t"
|
||||||
|
):
|
||||||
|
# This might be a new section, stop processing args
|
||||||
|
in_args_section = False
|
||||||
|
if current_param:
|
||||||
|
params.append(current_param)
|
||||||
|
current_param = None
|
||||||
|
|
||||||
|
# Add the last parameter if it exists
|
||||||
|
if current_param:
|
||||||
|
params.append(current_param)
|
||||||
|
|
||||||
|
return DocstringInfo(
|
||||||
|
short_description=short_description, params=params
|
||||||
|
)
|
||||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
|
After Width: | Height: | Size: 175 KiB |
|
After Width: | Height: | Size: 178 KiB |
|
After Width: | Height: | Size: 130 KiB |
|
After Width: | Height: | Size: 75 KiB |
|
After Width: | Height: | Size: 66 KiB |
|
|
After Width: | Height: | Size: 15 KiB |
|
Before Width: | Height: | Size: 19 MiB After Width: | Height: | Size: 19 MiB |
|
Before Width: | Height: | Size: 492 KiB After Width: | Height: | Size: 492 KiB |
@ -0,0 +1,293 @@
|
|||||||
|
"""
|
||||||
|
Tests for bug #1115 fix in AutoSwarmBuilder.
|
||||||
|
|
||||||
|
This test module verifies the fix for AttributeError when creating agents
|
||||||
|
from AgentSpec Pydantic models in AutoSwarmBuilder.
|
||||||
|
|
||||||
|
Bug: https://github.com/kyegomez/swarms/issues/1115
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from swarms.structs.agent import Agent
|
||||||
|
from swarms.structs.auto_swarm_builder import (
|
||||||
|
AgentSpec,
|
||||||
|
AutoSwarmBuilder,
|
||||||
|
)
|
||||||
|
from swarms.structs.ma_utils import set_random_models_for_agents
|
||||||
|
|
||||||
|
|
||||||
|
class TestAutoSwarmBuilderFix:
|
||||||
|
"""Tests for bug #1115 fix in AutoSwarmBuilder."""
|
||||||
|
|
||||||
|
def test_create_agents_from_specs_with_dict(self):
|
||||||
|
"""Test that create_agents_from_specs handles dict input correctly."""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
# Create specs as a dictionary
|
||||||
|
specs = {
|
||||||
|
"agents": [
|
||||||
|
{
|
||||||
|
"agent_name": "test_agent_1",
|
||||||
|
"description": "Test agent 1 description",
|
||||||
|
"system_prompt": "You are a helpful assistant",
|
||||||
|
"model_name": "gpt-4o-mini",
|
||||||
|
"max_loops": 1,
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(specs)
|
||||||
|
|
||||||
|
# Verify agents were created correctly
|
||||||
|
assert len(agents) == 1
|
||||||
|
assert isinstance(agents[0], Agent)
|
||||||
|
assert agents[0].agent_name == "test_agent_1"
|
||||||
|
|
||||||
|
# Verify description was mapped to agent_description
|
||||||
|
assert hasattr(agents[0], "agent_description")
|
||||||
|
assert (
|
||||||
|
agents[0].agent_description == "Test agent 1 description"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_create_agents_from_specs_with_pydantic(self):
|
||||||
|
"""Test that create_agents_from_specs handles Pydantic model input correctly.
|
||||||
|
|
||||||
|
This is the main test for bug #1115 - it verifies that AgentSpec
|
||||||
|
Pydantic models can be unpacked correctly.
|
||||||
|
"""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
# Create specs as Pydantic AgentSpec objects
|
||||||
|
agent_spec = AgentSpec(
|
||||||
|
agent_name="test_agent_pydantic",
|
||||||
|
description="Pydantic test agent",
|
||||||
|
system_prompt="You are a helpful assistant",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
specs = {"agents": [agent_spec]}
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(specs)
|
||||||
|
|
||||||
|
# Verify agents were created correctly
|
||||||
|
assert len(agents) == 1
|
||||||
|
assert isinstance(agents[0], Agent)
|
||||||
|
assert agents[0].agent_name == "test_agent_pydantic"
|
||||||
|
|
||||||
|
# Verify description was mapped to agent_description
|
||||||
|
assert hasattr(agents[0], "agent_description")
|
||||||
|
assert agents[0].agent_description == "Pydantic test agent"
|
||||||
|
|
||||||
|
def test_parameter_name_mapping(self):
|
||||||
|
"""Test that 'description' field maps to 'agent_description' correctly."""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
# Test with dict that has 'description'
|
||||||
|
specs = {
|
||||||
|
"agents": [
|
||||||
|
{
|
||||||
|
"agent_name": "mapping_test",
|
||||||
|
"description": "This should map to agent_description",
|
||||||
|
"system_prompt": "You are helpful",
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(specs)
|
||||||
|
|
||||||
|
assert len(agents) == 1
|
||||||
|
agent = agents[0]
|
||||||
|
|
||||||
|
# Verify description was mapped
|
||||||
|
assert hasattr(agent, "agent_description")
|
||||||
|
assert (
|
||||||
|
agent.agent_description
|
||||||
|
== "This should map to agent_description"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_create_agents_from_specs_mixed_input(self):
|
||||||
|
"""Test that create_agents_from_specs handles mixed dict and Pydantic input."""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
# Mix of dict and Pydantic objects
|
||||||
|
dict_spec = {
|
||||||
|
"agent_name": "dict_agent",
|
||||||
|
"description": "Dict agent description",
|
||||||
|
"system_prompt": "You are helpful",
|
||||||
|
}
|
||||||
|
|
||||||
|
pydantic_spec = AgentSpec(
|
||||||
|
agent_name="pydantic_agent",
|
||||||
|
description="Pydantic agent description",
|
||||||
|
system_prompt="You are smart",
|
||||||
|
)
|
||||||
|
|
||||||
|
specs = {"agents": [dict_spec, pydantic_spec]}
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(specs)
|
||||||
|
|
||||||
|
# Verify both agents were created
|
||||||
|
assert len(agents) == 2
|
||||||
|
assert all(isinstance(agent, Agent) for agent in agents)
|
||||||
|
|
||||||
|
# Verify both have correct descriptions
|
||||||
|
dict_agent = next(
|
||||||
|
a for a in agents if a.agent_name == "dict_agent"
|
||||||
|
)
|
||||||
|
pydantic_agent = next(
|
||||||
|
a for a in agents if a.agent_name == "pydantic_agent"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert (
|
||||||
|
dict_agent.agent_description == "Dict agent description"
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
pydantic_agent.agent_description
|
||||||
|
== "Pydantic agent description"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_set_random_models_for_agents_with_valid_agents(
|
||||||
|
self,
|
||||||
|
):
|
||||||
|
"""Test set_random_models_for_agents with proper Agent objects."""
|
||||||
|
# Create proper Agent objects
|
||||||
|
agents = [
|
||||||
|
Agent(
|
||||||
|
agent_name="agent1",
|
||||||
|
system_prompt="You are agent 1",
|
||||||
|
max_loops=1,
|
||||||
|
),
|
||||||
|
Agent(
|
||||||
|
agent_name="agent2",
|
||||||
|
system_prompt="You are agent 2",
|
||||||
|
max_loops=1,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
# Set random models
|
||||||
|
model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
|
||||||
|
result = set_random_models_for_agents(
|
||||||
|
agents=agents, model_names=model_names
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify results
|
||||||
|
assert len(result) == 2
|
||||||
|
assert all(isinstance(agent, Agent) for agent in result)
|
||||||
|
assert all(hasattr(agent, "model_name") for agent in result)
|
||||||
|
assert all(
|
||||||
|
agent.model_name in model_names for agent in result
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_set_random_models_for_agents_with_single_agent(
|
||||||
|
self,
|
||||||
|
):
|
||||||
|
"""Test set_random_models_for_agents with a single agent."""
|
||||||
|
agent = Agent(
|
||||||
|
agent_name="single_agent",
|
||||||
|
system_prompt="You are helpful",
|
||||||
|
max_loops=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
model_names = ["gpt-4o-mini", "gpt-4o"]
|
||||||
|
result = set_random_models_for_agents(
|
||||||
|
agents=agent, model_names=model_names
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(result, Agent)
|
||||||
|
assert hasattr(result, "model_name")
|
||||||
|
assert result.model_name in model_names
|
||||||
|
|
||||||
|
def test_set_random_models_for_agents_with_none(self):
|
||||||
|
"""Test set_random_models_for_agents with None returns random model name."""
|
||||||
|
model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
|
||||||
|
result = set_random_models_for_agents(
|
||||||
|
agents=None, model_names=model_names
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(result, str)
|
||||||
|
assert result in model_names
|
||||||
|
|
||||||
|
@pytest.mark.skip(
|
||||||
|
reason="This test requires API key and makes LLM calls"
|
||||||
|
)
|
||||||
|
def test_auto_swarm_builder_return_agents_objects_integration(
|
||||||
|
self,
|
||||||
|
):
|
||||||
|
"""Integration test for AutoSwarmBuilder with execution_type='return-agents-objects'.
|
||||||
|
|
||||||
|
This test requires OPENAI_API_KEY and makes actual LLM calls.
|
||||||
|
Run manually with: pytest -k test_auto_swarm_builder_return_agents_objects_integration -v
|
||||||
|
"""
|
||||||
|
builder = AutoSwarmBuilder(
|
||||||
|
execution_type="return-agents-objects",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
max_loops=1,
|
||||||
|
verbose=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
agents = builder.run(
|
||||||
|
"Create a team of 2 data analysis agents with specific roles"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify agents were created
|
||||||
|
assert isinstance(agents, list)
|
||||||
|
assert len(agents) >= 1
|
||||||
|
assert all(isinstance(agent, Agent) for agent in agents)
|
||||||
|
assert all(hasattr(agent, "agent_name") for agent in agents)
|
||||||
|
assert all(
|
||||||
|
hasattr(agent, "agent_description") for agent in agents
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_agent_spec_to_agent_all_fields(self):
|
||||||
|
"""Test that all AgentSpec fields are properly passed to Agent."""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
agent_spec = AgentSpec(
|
||||||
|
agent_name="full_test_agent",
|
||||||
|
description="Full test description",
|
||||||
|
system_prompt="You are a comprehensive test agent",
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
auto_generate_prompt=False,
|
||||||
|
max_tokens=4096,
|
||||||
|
temperature=0.7,
|
||||||
|
role="worker",
|
||||||
|
max_loops=3,
|
||||||
|
goal="Test all parameters",
|
||||||
|
)
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(
|
||||||
|
{"agents": [agent_spec]}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert len(agents) == 1
|
||||||
|
agent = agents[0]
|
||||||
|
|
||||||
|
# Verify all fields were set
|
||||||
|
assert agent.agent_name == "full_test_agent"
|
||||||
|
assert agent.agent_description == "Full test description"
|
||||||
|
# Agent may modify system_prompt by adding additional instructions
|
||||||
|
assert (
|
||||||
|
"You are a comprehensive test agent"
|
||||||
|
in agent.system_prompt
|
||||||
|
)
|
||||||
|
assert agent.max_loops == 3
|
||||||
|
assert agent.max_tokens == 4096
|
||||||
|
assert agent.temperature == 0.7
|
||||||
|
|
||||||
|
def test_create_agents_from_specs_empty_list(self):
|
||||||
|
"""Test that create_agents_from_specs handles empty agent list."""
|
||||||
|
builder = AutoSwarmBuilder()
|
||||||
|
|
||||||
|
specs = {"agents": []}
|
||||||
|
|
||||||
|
agents = builder.create_agents_from_specs(specs)
|
||||||
|
|
||||||
|
assert isinstance(agents, list)
|
||||||
|
assert len(agents) == 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# Run tests with pytest
|
||||||
|
pytest.main([__file__, "-v", "--tb=short"])
|
||||||
@ -0,0 +1,150 @@
|
|||||||
|
from pydantic import BaseModel
|
||||||
|
from swarms.tools.pydantic_to_json import (
|
||||||
|
base_model_to_openai_function,
|
||||||
|
multi_base_model_to_openai_function,
|
||||||
|
)
|
||||||
|
from swarms.tools.base_tool import BaseTool
|
||||||
|
|
||||||
|
|
||||||
|
# Test Pydantic model
|
||||||
|
class TestModel(BaseModel):
|
||||||
|
"""A test model for validation."""
|
||||||
|
|
||||||
|
name: str
|
||||||
|
age: int
|
||||||
|
email: str = "test@example.com"
|
||||||
|
|
||||||
|
|
||||||
|
def test_base_model_to_openai_function():
|
||||||
|
"""Test that base_model_to_openai_function accepts output_str parameter."""
|
||||||
|
print(
|
||||||
|
"Testing base_model_to_openai_function with output_str=False..."
|
||||||
|
)
|
||||||
|
result_dict = base_model_to_openai_function(
|
||||||
|
TestModel, output_str=False
|
||||||
|
)
|
||||||
|
print(f"✓ Dict result type: {type(result_dict)}")
|
||||||
|
print(f"✓ Dict result keys: {list(result_dict.keys())}")
|
||||||
|
|
||||||
|
print(
|
||||||
|
"\nTesting base_model_to_openai_function with output_str=True..."
|
||||||
|
)
|
||||||
|
result_str = base_model_to_openai_function(
|
||||||
|
TestModel, output_str=True
|
||||||
|
)
|
||||||
|
print(f"✓ String result type: {type(result_str)}")
|
||||||
|
print(f"✓ String result preview: {result_str[:100]}...")
|
||||||
|
|
||||||
|
|
||||||
|
def test_multi_base_model_to_openai_function():
|
||||||
|
"""Test that multi_base_model_to_openai_function handles output_str correctly."""
|
||||||
|
print(
|
||||||
|
"\nTesting multi_base_model_to_openai_function with output_str=False..."
|
||||||
|
)
|
||||||
|
result_dict = multi_base_model_to_openai_function(
|
||||||
|
[TestModel], output_str=False
|
||||||
|
)
|
||||||
|
print(f"✓ Dict result type: {type(result_dict)}")
|
||||||
|
print(f"✓ Dict result keys: {list(result_dict.keys())}")
|
||||||
|
|
||||||
|
print(
|
||||||
|
"\nTesting multi_base_model_to_openai_function with output_str=True..."
|
||||||
|
)
|
||||||
|
result_str = multi_base_model_to_openai_function(
|
||||||
|
[TestModel], output_str=True
|
||||||
|
)
|
||||||
|
print(f"✓ String result type: {type(result_str)}")
|
||||||
|
print(f"✓ String result preview: {result_str[:100]}...")
|
||||||
|
|
||||||
|
|
||||||
|
def test_base_tool_methods():
|
||||||
|
"""Test that BaseTool methods handle output_str parameter correctly."""
|
||||||
|
print(
|
||||||
|
"\nTesting BaseTool.base_model_to_dict with output_str=False..."
|
||||||
|
)
|
||||||
|
tool = BaseTool()
|
||||||
|
result_dict = tool.base_model_to_dict(TestModel, output_str=False)
|
||||||
|
print(f"✓ Dict result type: {type(result_dict)}")
|
||||||
|
print(f"✓ Dict result keys: {list(result_dict.keys())}")
|
||||||
|
|
||||||
|
print(
|
||||||
|
"\nTesting BaseTool.base_model_to_dict with output_str=True..."
|
||||||
|
)
|
||||||
|
result_str = tool.base_model_to_dict(TestModel, output_str=True)
|
||||||
|
print(f"✓ String result type: {type(result_str)}")
|
||||||
|
print(f"✓ String result preview: {result_str[:100]}...")
|
||||||
|
|
||||||
|
print(
|
||||||
|
"\nTesting BaseTool.multi_base_models_to_dict with output_str=False..."
|
||||||
|
)
|
||||||
|
result_dict = tool.multi_base_models_to_dict(
|
||||||
|
[TestModel], output_str=False
|
||||||
|
)
|
||||||
|
print(f"✓ Dict result type: {type(result_dict)}")
|
||||||
|
print(f"✓ Dict result length: {len(result_dict)}")
|
||||||
|
|
||||||
|
print(
|
||||||
|
"\nTesting BaseTool.multi_base_models_to_dict with output_str=True..."
|
||||||
|
)
|
||||||
|
result_str = tool.multi_base_models_to_dict(
|
||||||
|
[TestModel], output_str=True
|
||||||
|
)
|
||||||
|
print(f"✓ String result type: {type(result_str)}")
|
||||||
|
print(f"✓ String result preview: {result_str[:100]}...")
|
||||||
|
|
||||||
|
|
||||||
|
def test_agent_integration():
|
||||||
|
"""Test that the Agent class can use the fixed methods without errors."""
|
||||||
|
print("\nTesting Agent integration...")
|
||||||
|
try:
|
||||||
|
from swarms import Agent
|
||||||
|
|
||||||
|
# Create a simple agent with a tool schema
|
||||||
|
agent = Agent(
|
||||||
|
model_name="gpt-4o-mini",
|
||||||
|
tool_schema=TestModel,
|
||||||
|
max_loops=1,
|
||||||
|
verbose=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# This should not raise an error anymore
|
||||||
|
agent.handle_tool_schema_ops()
|
||||||
|
print(
|
||||||
|
"✓ Agent.handle_tool_schema_ops() completed successfully"
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"✗ Agent integration failed: {e}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
print("=" * 60)
|
||||||
|
print("Testing output_str parameter fix")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
try:
|
||||||
|
test_base_model_to_openai_function()
|
||||||
|
test_multi_base_model_to_openai_function()
|
||||||
|
test_base_tool_methods()
|
||||||
|
|
||||||
|
if test_agent_integration():
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print(
|
||||||
|
"✅ All tests passed! The output_str parameter fix is working correctly."
|
||||||
|
)
|
||||||
|
print("=" * 60)
|
||||||
|
else:
|
||||||
|
print("\n" + "=" * 60)
|
||||||
|
print(
|
||||||
|
"❌ Some tests failed. Please check the implementation."
|
||||||
|
)
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\n❌ Test failed with error: {e}")
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
traceback.print_exc()
|
||||||
@ -0,0 +1,431 @@
|
|||||||
|
"""
|
||||||
|
Test suite for the custom docstring parser implementation.
|
||||||
|
|
||||||
|
This module contains comprehensive tests to ensure the docstring parser
|
||||||
|
works correctly with various docstring formats and edge cases.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from swarms.utils.docstring_parser import (
|
||||||
|
parse,
|
||||||
|
DocstringParam,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestDocstringParser:
|
||||||
|
"""Test cases for the docstring parser functionality."""
|
||||||
|
|
||||||
|
def test_empty_docstring(self):
|
||||||
|
"""Test parsing of empty docstring."""
|
||||||
|
result = parse("")
|
||||||
|
assert result.short_description is None
|
||||||
|
assert result.params == []
|
||||||
|
|
||||||
|
def test_none_docstring(self):
|
||||||
|
"""Test parsing of None docstring."""
|
||||||
|
result = parse(None)
|
||||||
|
assert result.short_description is None
|
||||||
|
assert result.params == []
|
||||||
|
|
||||||
|
def test_whitespace_only_docstring(self):
|
||||||
|
"""Test parsing of whitespace-only docstring."""
|
||||||
|
result = parse(" \n \t \n ")
|
||||||
|
assert result.short_description is None
|
||||||
|
assert result.params == []
|
||||||
|
|
||||||
|
def test_simple_docstring_no_args(self):
|
||||||
|
"""Test parsing of simple docstring without Args section."""
|
||||||
|
docstring = """
|
||||||
|
This is a simple function.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: A simple string
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description == "This is a simple function."
|
||||||
|
)
|
||||||
|
assert result.params == []
|
||||||
|
|
||||||
|
def test_docstring_with_args(self):
|
||||||
|
"""Test parsing of docstring with Args section."""
|
||||||
|
docstring = """
|
||||||
|
This is a test function.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
param2 (int): Second parameter
|
||||||
|
param3 (bool, optional): Third parameter with default
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Return value description
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description == "This is a test function."
|
||||||
|
assert len(result.params) == 3
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Second parameter"
|
||||||
|
)
|
||||||
|
assert result.params[2] == DocstringParam(
|
||||||
|
"param3", "Third parameter with default"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_parameters_section(self):
|
||||||
|
"""Test parsing of docstring with Parameters section."""
|
||||||
|
docstring = """
|
||||||
|
Another test function.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
name (str): The name parameter
|
||||||
|
age (int): The age parameter
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None: Nothing is returned
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description == "Another test function."
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"name", "The name parameter"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"age", "The age parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_multiline_param_description(self):
|
||||||
|
"""Test parsing of docstring with multiline parameter descriptions."""
|
||||||
|
docstring = """
|
||||||
|
Function with multiline descriptions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): This is a very long description
|
||||||
|
that spans multiple lines and should be
|
||||||
|
properly concatenated.
|
||||||
|
param2 (int): Short description
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Result
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with multiline descriptions."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 2
|
||||||
|
expected_desc = "This is a very long description that spans multiple lines and should be properly concatenated."
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", expected_desc
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Short description"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_without_type_annotations(self):
|
||||||
|
"""Test parsing of docstring without type annotations."""
|
||||||
|
docstring = """
|
||||||
|
Function without type annotations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1: First parameter without type
|
||||||
|
param2: Second parameter without type
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Result
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function without type annotations."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter without type"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Second parameter without type"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_pydantic_style_docstring(self):
|
||||||
|
"""Test parsing of Pydantic-style docstring."""
|
||||||
|
docstring = """
|
||||||
|
Convert a Pydantic model to a dictionary representation of functions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
pydantic_type (type[BaseModel]): The Pydantic model type to convert.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict[str, Any]: A dictionary representation of the functions.
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Convert a Pydantic model to a dictionary representation of functions."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"pydantic_type", "The Pydantic model type to convert."
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_various_sections(self):
|
||||||
|
"""Test parsing of docstring with multiple sections."""
|
||||||
|
docstring = """
|
||||||
|
Complex function with multiple sections.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_data (dict): Input data dictionary
|
||||||
|
validate (bool): Whether to validate input
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: Processed data
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If input is invalid
|
||||||
|
|
||||||
|
Note:
|
||||||
|
This is a note section
|
||||||
|
|
||||||
|
Example:
|
||||||
|
>>> result = complex_function({"key": "value"})
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Complex function with multiple sections."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"input_data", "Input data dictionary"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"validate", "Whether to validate input"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_see_also_section(self):
|
||||||
|
"""Test parsing of docstring with See Also section."""
|
||||||
|
docstring = """
|
||||||
|
Function with See Also section.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
|
||||||
|
See Also:
|
||||||
|
related_function: For related functionality
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with See Also section."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_see_also_underscore_section(self):
|
||||||
|
"""Test parsing of docstring with See_Also section (underscore variant)."""
|
||||||
|
docstring = """
|
||||||
|
Function with See_Also section.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
|
||||||
|
See_Also:
|
||||||
|
related_function: For related functionality
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with See_Also section."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_yields_section(self):
|
||||||
|
"""Test parsing of docstring with Yields section."""
|
||||||
|
docstring = """
|
||||||
|
Generator function.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
items (list): List of items to process
|
||||||
|
|
||||||
|
Yields:
|
||||||
|
str: Processed item
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description == "Generator function."
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"items", "List of items to process"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_raises_section(self):
|
||||||
|
"""Test parsing of docstring with Raises section."""
|
||||||
|
docstring = """
|
||||||
|
Function that can raise exceptions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
value (int): Value to process
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If value is negative
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function that can raise exceptions."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"value", "Value to process"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_examples_section(self):
|
||||||
|
"""Test parsing of docstring with Examples section."""
|
||||||
|
docstring = """
|
||||||
|
Function with examples.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
x (int): Input value
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
>>> result = example_function(5)
|
||||||
|
>>> print(result)
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description == "Function with examples."
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam("x", "Input value")
|
||||||
|
|
||||||
|
def test_docstring_with_note_section(self):
|
||||||
|
"""Test parsing of docstring with Note section."""
|
||||||
|
docstring = """
|
||||||
|
Function with a note.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (str): Input data
|
||||||
|
|
||||||
|
Note:
|
||||||
|
This function is deprecated
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description == "Function with a note."
|
||||||
|
assert len(result.params) == 1
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"data", "Input data"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_complex_type_annotations(self):
|
||||||
|
"""Test parsing of docstring with complex type annotations."""
|
||||||
|
docstring = """
|
||||||
|
Function with complex types.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (List[Dict[str, Any]]): Complex data structure
|
||||||
|
callback (Callable[[str], int]): Callback function
|
||||||
|
optional (Optional[str], optional): Optional parameter
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Union[str, None]: Result or None
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description == "Function with complex types."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 3
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"data", "Complex data structure"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"callback", "Callback function"
|
||||||
|
)
|
||||||
|
assert result.params[2] == DocstringParam(
|
||||||
|
"optional", "Optional parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_no_description(self):
|
||||||
|
"""Test parsing of docstring with no description, only Args."""
|
||||||
|
docstring = """
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
param2 (int): Second parameter
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert result.short_description is None
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Second parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_empty_args_section(self):
|
||||||
|
"""Test parsing of docstring with empty Args section."""
|
||||||
|
docstring = """
|
||||||
|
Function with empty Args section.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Result
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with empty Args section."
|
||||||
|
)
|
||||||
|
assert result.params == []
|
||||||
|
|
||||||
|
def test_docstring_with_mixed_indentation(self):
|
||||||
|
"""Test parsing of docstring with mixed indentation."""
|
||||||
|
docstring = """
|
||||||
|
Function with mixed indentation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
with continuation
|
||||||
|
param2 (int): Second parameter
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with mixed indentation."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter with continuation"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Second parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_docstring_with_tab_indentation(self):
|
||||||
|
"""Test parsing of docstring with tab indentation."""
|
||||||
|
docstring = """
|
||||||
|
Function with tab indentation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
param1 (str): First parameter
|
||||||
|
param2 (int): Second parameter
|
||||||
|
"""
|
||||||
|
result = parse(docstring)
|
||||||
|
assert (
|
||||||
|
result.short_description
|
||||||
|
== "Function with tab indentation."
|
||||||
|
)
|
||||||
|
assert len(result.params) == 2
|
||||||
|
assert result.params[0] == DocstringParam(
|
||||||
|
"param1", "First parameter"
|
||||||
|
)
|
||||||
|
assert result.params[1] == DocstringParam(
|
||||||
|
"param2", "Second parameter"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
pytest.main([__file__])
|
||||||
@ -1,120 +0,0 @@
|
|||||||
import pytest
|
|
||||||
|
|
||||||
from swarms.utils import print_class_parameters
|
|
||||||
|
|
||||||
|
|
||||||
class TestObject:
|
|
||||||
def __init__(self, value1, value2: int):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class TestObject2:
|
|
||||||
def __init__(self: "TestObject2", value1, value2: int = 5):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
def test_class_with_complex_parameters():
|
|
||||||
class ComplexArgs:
|
|
||||||
def __init__(self, value1: list, value2: dict = {}):
|
|
||||||
pass
|
|
||||||
|
|
||||||
output = {"value1": "<class 'list'>", "value2": "<class 'dict'>"}
|
|
||||||
assert (
|
|
||||||
print_class_parameters(ComplexArgs, api_format=True) == output
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def test_empty_class():
|
|
||||||
class Empty:
|
|
||||||
pass
|
|
||||||
|
|
||||||
with pytest.raises(Exception):
|
|
||||||
print_class_parameters(Empty)
|
|
||||||
|
|
||||||
|
|
||||||
def test_class_with_no_annotations():
|
|
||||||
class NoAnnotations:
|
|
||||||
def __init__(self, value1, value2):
|
|
||||||
pass
|
|
||||||
|
|
||||||
output = {
|
|
||||||
"value1": "<class 'inspect._empty'>",
|
|
||||||
"value2": "<class 'inspect._empty'>",
|
|
||||||
}
|
|
||||||
assert (
|
|
||||||
print_class_parameters(NoAnnotations, api_format=True)
|
|
||||||
== output
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def test_class_with_partial_annotations():
|
|
||||||
class PartialAnnotations:
|
|
||||||
def __init__(self, value1, value2: int):
|
|
||||||
pass
|
|
||||||
|
|
||||||
output = {
|
|
||||||
"value1": "<class 'inspect._empty'>",
|
|
||||||
"value2": "<class 'int'>",
|
|
||||||
}
|
|
||||||
assert (
|
|
||||||
print_class_parameters(PartialAnnotations, api_format=True)
|
|
||||||
== output
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"obj, expected",
|
|
||||||
[
|
|
||||||
(
|
|
||||||
TestObject,
|
|
||||||
{
|
|
||||||
"value1": "<class 'inspect._empty'>",
|
|
||||||
"value2": "<class 'int'>",
|
|
||||||
},
|
|
||||||
),
|
|
||||||
(
|
|
||||||
TestObject2,
|
|
||||||
{
|
|
||||||
"value1": "<class 'inspect._empty'>",
|
|
||||||
"value2": "<class 'int'>",
|
|
||||||
},
|
|
||||||
),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_parametrized_class_parameters(obj, expected):
|
|
||||||
assert print_class_parameters(obj, api_format=True) == expected
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"value",
|
|
||||||
[
|
|
||||||
int,
|
|
||||||
float,
|
|
||||||
str,
|
|
||||||
list,
|
|
||||||
set,
|
|
||||||
dict,
|
|
||||||
bool,
|
|
||||||
tuple,
|
|
||||||
complex,
|
|
||||||
bytes,
|
|
||||||
bytearray,
|
|
||||||
memoryview,
|
|
||||||
range,
|
|
||||||
frozenset,
|
|
||||||
slice,
|
|
||||||
object,
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_not_class_exception(value):
|
|
||||||
with pytest.raises(Exception):
|
|
||||||
print_class_parameters(value)
|
|
||||||
|
|
||||||
|
|
||||||
def test_api_format_flag():
|
|
||||||
assert print_class_parameters(TestObject2, api_format=True) == {
|
|
||||||
"value1": "<class 'inspect._empty'>",
|
|
||||||
"value2": "<class 'int'>",
|
|
||||||
}
|
|
||||||
print_class_parameters(TestObject)
|
|
||||||
# TODO: Capture printed output and assert correctness.
|
|
||||||
Loading…
Reference in new issue