[IMPROVEMENT][Remove sentence transformers from forest swarm and instead use litellm] [UPDATE][DOCs]
Kye Gomez
3 days ago
parent
438029dbe5
commit
5c79016afa
@ -0,0 +1,279 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
ForestSwarm Example Script
|
||||
|
||||
This script demonstrates the ForestSwarm functionality with realistic examples
|
||||
of financial services and investment management agents.
|
||||
"""
|
||||
|
||||
from swarms.structs.tree_swarm import TreeAgent, Tree, ForestSwarm
|
||||
|
||||
|
||||
def create_financial_services_forest():
|
||||
"""Create a comprehensive financial services forest with multiple specialized agents."""
|
||||
|
||||
print("🌳 Creating Financial Services Forest...")
|
||||
|
||||
# Financial Services Tree - Personal Finance & Planning
|
||||
financial_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""I am a certified financial planner specializing in personal finance,
|
||||
budgeting, debt management, and financial goal setting. I help individuals create
|
||||
comprehensive financial plans and make informed decisions about their money.""",
|
||||
agent_name="Personal Financial Planner",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a tax preparation specialist with expertise in individual and
|
||||
small business tax returns. I help clients maximize deductions, understand tax laws,
|
||||
and file taxes accurately and on time.""",
|
||||
agent_name="Tax Preparation Specialist",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a retirement planning expert who helps individuals and families
|
||||
plan for retirement. I specialize in 401(k)s, IRAs, Social Security optimization,
|
||||
and creating sustainable retirement income strategies.""",
|
||||
agent_name="Retirement Planning Expert",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a debt management counselor who helps individuals and families
|
||||
get out of debt and build financial stability. I provide strategies for debt
|
||||
consolidation, negotiation, and creating sustainable repayment plans.""",
|
||||
agent_name="Debt Management Counselor",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
]
|
||||
|
||||
# Investment & Trading Tree - Market Analysis & Portfolio Management
|
||||
investment_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""I am a stock market analyst who provides insights on market trends,
|
||||
stock recommendations, and portfolio optimization strategies. I analyze company
|
||||
fundamentals, market conditions, and economic indicators to help investors make
|
||||
informed decisions.""",
|
||||
agent_name="Stock Market Analyst",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am an investment strategist specializing in portfolio diversification,
|
||||
risk management, and asset allocation. I help investors create balanced portfolios
|
||||
that align with their risk tolerance and financial goals.""",
|
||||
agent_name="Investment Strategist",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a cryptocurrency and blockchain expert who provides insights on
|
||||
digital assets, DeFi protocols, and emerging blockchain technologies. I help
|
||||
investors understand the risks and opportunities in the crypto market.""",
|
||||
agent_name="Cryptocurrency Expert",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a real estate investment advisor who helps investors evaluate
|
||||
real estate opportunities, understand market trends, and build real estate
|
||||
portfolios for long-term wealth building.""",
|
||||
agent_name="Real Estate Investment Advisor",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
]
|
||||
|
||||
# Business & Corporate Tree - Business Finance & Strategy
|
||||
business_agents = [
|
||||
TreeAgent(
|
||||
system_prompt="""I am a business financial advisor specializing in corporate finance,
|
||||
business valuation, mergers and acquisitions, and strategic financial planning
|
||||
for small to medium-sized businesses.""",
|
||||
agent_name="Business Financial Advisor",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a Delaware incorporation specialist with deep knowledge of
|
||||
corporate formation, tax benefits, legal requirements, and ongoing compliance
|
||||
for businesses incorporating in Delaware.""",
|
||||
agent_name="Delaware Incorporation Specialist",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a startup funding advisor who helps entrepreneurs secure
|
||||
funding through venture capital, angel investors, crowdfunding, and other
|
||||
financing options. I provide guidance on business plans, pitch decks, and
|
||||
investor relations.""",
|
||||
agent_name="Startup Funding Advisor",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="""I am a business tax strategist who helps businesses optimize their
|
||||
tax position through strategic planning, entity structure optimization, and
|
||||
compliance with federal, state, and local tax laws.""",
|
||||
agent_name="Business Tax Strategist",
|
||||
model_name="gpt-4o",
|
||||
),
|
||||
]
|
||||
|
||||
# Create trees
|
||||
financial_tree = Tree(
|
||||
"Personal Finance & Planning", financial_agents
|
||||
)
|
||||
investment_tree = Tree("Investment & Trading", investment_agents)
|
||||
business_tree = Tree(
|
||||
"Business & Corporate Finance", business_agents
|
||||
)
|
||||
|
||||
# Create the forest
|
||||
forest = ForestSwarm(
|
||||
name="Comprehensive Financial Services Forest",
|
||||
description="A multi-agent system providing expert financial advice across personal, investment, and business domains",
|
||||
trees=[financial_tree, investment_tree, business_tree],
|
||||
)
|
||||
|
||||
print(
|
||||
f"✅ Created forest with {len(forest.trees)} trees and {sum(len(tree.agents) for tree in forest.trees)} agents"
|
||||
)
|
||||
return forest
|
||||
|
||||
|
||||
def demonstrate_agent_selection(forest):
|
||||
"""Demonstrate how the forest selects the most relevant agent for different types of questions."""
|
||||
|
||||
print("\n🎯 Demonstrating Agent Selection...")
|
||||
|
||||
# Test questions covering different domains
|
||||
test_questions = [
|
||||
{
|
||||
"question": "How much should I save monthly for retirement if I want to retire at 65?",
|
||||
"expected_agent": "Retirement Planning Expert",
|
||||
"category": "Personal Finance",
|
||||
},
|
||||
{
|
||||
"question": "What are the best investment strategies for a 401k retirement plan?",
|
||||
"expected_agent": "Investment Strategist",
|
||||
"category": "Investment",
|
||||
},
|
||||
{
|
||||
"question": "Our company is incorporated in Delaware, how do we do our taxes for free?",
|
||||
"expected_agent": "Delaware Incorporation Specialist",
|
||||
"category": "Business",
|
||||
},
|
||||
{
|
||||
"question": "Which tech stocks should I consider for my investment portfolio?",
|
||||
"expected_agent": "Stock Market Analyst",
|
||||
"category": "Investment",
|
||||
},
|
||||
{
|
||||
"question": "How can I consolidate my credit card debt and create a repayment plan?",
|
||||
"expected_agent": "Debt Management Counselor",
|
||||
"category": "Personal Finance",
|
||||
},
|
||||
{
|
||||
"question": "What are the benefits of incorporating in Delaware vs. other states?",
|
||||
"expected_agent": "Delaware Incorporation Specialist",
|
||||
"category": "Business",
|
||||
},
|
||||
]
|
||||
|
||||
for i, test_case in enumerate(test_questions, 1):
|
||||
print(f"\n--- Test Case {i}: {test_case['category']} ---")
|
||||
print(f"Question: {test_case['question']}")
|
||||
print(f"Expected Agent: {test_case['expected_agent']}")
|
||||
|
||||
try:
|
||||
# Find the relevant tree
|
||||
relevant_tree = forest.find_relevant_tree(
|
||||
test_case["question"]
|
||||
)
|
||||
if relevant_tree:
|
||||
print(f"Selected Tree: {relevant_tree.tree_name}")
|
||||
|
||||
# Find the relevant agent
|
||||
relevant_agent = relevant_tree.find_relevant_agent(
|
||||
test_case["question"]
|
||||
)
|
||||
if relevant_agent:
|
||||
print(
|
||||
f"Selected Agent: {relevant_agent.agent_name}"
|
||||
)
|
||||
|
||||
# Check if the selection matches expectation
|
||||
if (
|
||||
test_case["expected_agent"]
|
||||
in relevant_agent.agent_name
|
||||
):
|
||||
print(
|
||||
"✅ Agent selection matches expectation!"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
"⚠️ Agent selection differs from expectation"
|
||||
)
|
||||
print(
|
||||
f" Expected: {test_case['expected_agent']}"
|
||||
)
|
||||
print(
|
||||
f" Selected: {relevant_agent.agent_name}"
|
||||
)
|
||||
else:
|
||||
print("❌ No relevant agent found")
|
||||
else:
|
||||
print("❌ No relevant tree found")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Error during agent selection: {e}")
|
||||
|
||||
|
||||
def run_sample_tasks(forest):
|
||||
"""Run sample tasks to demonstrate the forest's capabilities."""
|
||||
|
||||
print("\n🚀 Running Sample Tasks...")
|
||||
|
||||
sample_tasks = [
|
||||
"What are the key benefits of incorporating a business in Delaware?",
|
||||
"How should I allocate my investment portfolio if I'm 30 years old?",
|
||||
"What's the best way to start saving for retirement in my 20s?",
|
||||
]
|
||||
|
||||
for i, task in enumerate(sample_tasks, 1):
|
||||
print(f"\n--- Task {i} ---")
|
||||
print(f"Task: {task}")
|
||||
|
||||
try:
|
||||
result = forest.run(task)
|
||||
print(
|
||||
f"Result: {result[:200]}..."
|
||||
if len(str(result)) > 200
|
||||
else f"Result: {result}"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"❌ Task execution failed: {e}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function to demonstrate ForestSwarm functionality."""
|
||||
|
||||
print("🌲 ForestSwarm Demonstration")
|
||||
print("=" * 60)
|
||||
|
||||
try:
|
||||
# Create the forest
|
||||
forest = create_financial_services_forest()
|
||||
|
||||
# Demonstrate agent selection
|
||||
demonstrate_agent_selection(forest)
|
||||
|
||||
# Run sample tasks
|
||||
run_sample_tasks(forest)
|
||||
|
||||
print(
|
||||
"\n🎉 ForestSwarm demonstration completed successfully!"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n❌ Error during demonstration: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,43 @@
|
||||
from swarms.structs.tree_swarm import TreeAgent, Tree, ForestSwarm
|
||||
|
||||
# Create agents with varying system prompts and dynamically generated distances/keywords
|
||||
agents_tree1 = [
|
||||
TreeAgent(
|
||||
system_prompt="Stock Analysis Agent",
|
||||
agent_name="Stock Analysis Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="Financial Planning Agent",
|
||||
agent_name="Financial Planning Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
agent_name="Retirement Strategy Agent",
|
||||
system_prompt="Retirement Strategy Agent",
|
||||
),
|
||||
]
|
||||
|
||||
agents_tree2 = [
|
||||
TreeAgent(
|
||||
system_prompt="Tax Filing Agent",
|
||||
agent_name="Tax Filing Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="Investment Strategy Agent",
|
||||
agent_name="Investment Strategy Agent",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt="ROTH IRA Agent", agent_name="ROTH IRA Agent"
|
||||
),
|
||||
]
|
||||
|
||||
# Create trees
|
||||
tree1 = Tree(tree_name="Financial Tree", agents=agents_tree1)
|
||||
tree2 = Tree(tree_name="Investment Tree", agents=agents_tree2)
|
||||
|
||||
# Create the ForestSwarm
|
||||
multi_agent_structure = ForestSwarm(trees=[tree1, tree2])
|
||||
|
||||
# Run a task
|
||||
task = "Our company is incorporated in delaware, how do we do our taxes for free?"
|
||||
output = multi_agent_structure.run(task)
|
||||
print(output)
|
||||
@ -0,0 +1,653 @@
|
||||
import sys
|
||||
|
||||
from swarms.structs.tree_swarm import (
|
||||
TreeAgent,
|
||||
Tree,
|
||||
ForestSwarm,
|
||||
AgentLogInput,
|
||||
AgentLogOutput,
|
||||
TreeLog,
|
||||
extract_keywords,
|
||||
cosine_similarity,
|
||||
)
|
||||
|
||||
|
||||
# Test Results Tracking
|
||||
test_results = {"passed": 0, "failed": 0, "total": 0}
|
||||
|
||||
|
||||
def assert_equal(actual, expected, test_name):
|
||||
"""Assert that actual equals expected, track test results."""
|
||||
test_results["total"] += 1
|
||||
if actual == expected:
|
||||
test_results["passed"] += 1
|
||||
print(f"✅ PASS: {test_name}")
|
||||
return True
|
||||
else:
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ FAIL: {test_name}")
|
||||
print(f" Expected: {expected}")
|
||||
print(f" Actual: {actual}")
|
||||
return False
|
||||
|
||||
|
||||
def assert_true(condition, test_name):
|
||||
"""Assert that condition is True, track test results."""
|
||||
test_results["total"] += 1
|
||||
if condition:
|
||||
test_results["passed"] += 1
|
||||
print(f"✅ PASS: {test_name}")
|
||||
return True
|
||||
else:
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ FAIL: {test_name}")
|
||||
print(" Condition was False")
|
||||
return False
|
||||
|
||||
|
||||
def assert_false(condition, test_name):
|
||||
"""Assert that condition is False, track test results."""
|
||||
test_results["total"] += 1
|
||||
if not condition:
|
||||
test_results["passed"] += 1
|
||||
print(f"✅ PASS: {test_name}")
|
||||
return True
|
||||
else:
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ FAIL: {test_name}")
|
||||
print(" Condition was True")
|
||||
return False
|
||||
|
||||
|
||||
def assert_is_instance(obj, expected_type, test_name):
|
||||
"""Assert that obj is an instance of expected_type, track test results."""
|
||||
test_results["total"] += 1
|
||||
if isinstance(obj, expected_type):
|
||||
test_results["passed"] += 1
|
||||
print(f"✅ PASS: {test_name}")
|
||||
return True
|
||||
else:
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ FAIL: {test_name}")
|
||||
print(f" Expected type: {expected_type}")
|
||||
print(f" Actual type: {type(obj)}")
|
||||
return False
|
||||
|
||||
|
||||
def assert_not_none(obj, test_name):
|
||||
"""Assert that obj is not None, track test results."""
|
||||
test_results["total"] += 1
|
||||
if obj is not None:
|
||||
test_results["passed"] += 1
|
||||
print(f"✅ PASS: {test_name}")
|
||||
return True
|
||||
else:
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ FAIL: {test_name}")
|
||||
print(" Object was None")
|
||||
return False
|
||||
|
||||
|
||||
# Test Data
|
||||
SAMPLE_SYSTEM_PROMPTS = {
|
||||
"financial_advisor": "I am a financial advisor specializing in investment planning, retirement strategies, and tax optimization for individuals and businesses.",
|
||||
"tax_expert": "I am a tax expert with deep knowledge of corporate taxation, Delaware incorporation benefits, and free tax filing options for businesses.",
|
||||
"stock_analyst": "I am a stock market analyst who provides insights on market trends, stock recommendations, and portfolio optimization strategies.",
|
||||
"retirement_planner": "I am a retirement planning specialist who helps individuals and businesses create comprehensive retirement strategies and investment plans.",
|
||||
}
|
||||
|
||||
SAMPLE_TASKS = {
|
||||
"tax_question": "Our company is incorporated in Delaware, how do we do our taxes for free?",
|
||||
"investment_question": "What are the best investment strategies for a 401k retirement plan?",
|
||||
"stock_question": "Which tech stocks should I consider for my investment portfolio?",
|
||||
"retirement_question": "How much should I save monthly for retirement if I want to retire at 65?",
|
||||
}
|
||||
|
||||
|
||||
# Test Functions
|
||||
|
||||
|
||||
def test_extract_keywords():
|
||||
"""Test the extract_keywords function."""
|
||||
print("\n🧪 Testing extract_keywords function...")
|
||||
|
||||
# Test basic keyword extraction
|
||||
text = (
|
||||
"financial advisor investment planning retirement strategies"
|
||||
)
|
||||
keywords = extract_keywords(text, top_n=3)
|
||||
assert_equal(
|
||||
len(keywords),
|
||||
3,
|
||||
"extract_keywords returns correct number of keywords",
|
||||
)
|
||||
assert_true(
|
||||
"financial" in keywords,
|
||||
"extract_keywords includes 'financial'",
|
||||
)
|
||||
assert_true(
|
||||
"investment" in keywords,
|
||||
"extract_keywords includes 'investment'",
|
||||
)
|
||||
|
||||
# Test with punctuation and case
|
||||
text = "Tax Expert! Corporate Taxation, Delaware Incorporation."
|
||||
keywords = extract_keywords(text, top_n=5)
|
||||
assert_true(
|
||||
"tax" in keywords,
|
||||
"extract_keywords handles punctuation and case",
|
||||
)
|
||||
assert_true(
|
||||
"corporate" in keywords,
|
||||
"extract_keywords handles punctuation and case",
|
||||
)
|
||||
|
||||
# Test empty string
|
||||
keywords = extract_keywords("", top_n=3)
|
||||
assert_equal(
|
||||
len(keywords), 0, "extract_keywords handles empty string"
|
||||
)
|
||||
|
||||
|
||||
def test_cosine_similarity():
|
||||
"""Test the cosine_similarity function."""
|
||||
print("\n🧪 Testing cosine_similarity function...")
|
||||
|
||||
# Test identical vectors
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [1.0, 0.0, 0.0]
|
||||
similarity = cosine_similarity(vec1, vec2)
|
||||
assert_equal(
|
||||
similarity,
|
||||
1.0,
|
||||
"cosine_similarity returns 1.0 for identical vectors",
|
||||
)
|
||||
|
||||
# Test orthogonal vectors
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [0.0, 1.0, 0.0]
|
||||
similarity = cosine_similarity(vec1, vec2)
|
||||
assert_equal(
|
||||
similarity,
|
||||
0.0,
|
||||
"cosine_similarity returns 0.0 for orthogonal vectors",
|
||||
)
|
||||
|
||||
# Test opposite vectors
|
||||
vec1 = [1.0, 0.0, 0.0]
|
||||
vec2 = [-1.0, 0.0, 0.0]
|
||||
similarity = cosine_similarity(vec1, vec2)
|
||||
assert_equal(
|
||||
similarity,
|
||||
-1.0,
|
||||
"cosine_similarity returns -1.0 for opposite vectors",
|
||||
)
|
||||
|
||||
# Test zero vectors
|
||||
vec1 = [0.0, 0.0, 0.0]
|
||||
vec2 = [1.0, 0.0, 0.0]
|
||||
similarity = cosine_similarity(vec1, vec2)
|
||||
assert_equal(
|
||||
similarity, 0.0, "cosine_similarity handles zero vectors"
|
||||
)
|
||||
|
||||
|
||||
def test_agent_log_models():
|
||||
"""Test the Pydantic log models."""
|
||||
print("\n🧪 Testing Pydantic log models...")
|
||||
|
||||
# Test AgentLogInput
|
||||
log_input = AgentLogInput(
|
||||
agent_name="test_agent", task="test_task"
|
||||
)
|
||||
assert_is_instance(
|
||||
log_input,
|
||||
AgentLogInput,
|
||||
"AgentLogInput creates correct instance",
|
||||
)
|
||||
assert_not_none(
|
||||
log_input.log_id, "AgentLogInput generates log_id"
|
||||
)
|
||||
assert_equal(
|
||||
log_input.agent_name,
|
||||
"test_agent",
|
||||
"AgentLogInput stores agent_name",
|
||||
)
|
||||
assert_equal(
|
||||
log_input.task, "test_task", "AgentLogInput stores task"
|
||||
)
|
||||
|
||||
# Test AgentLogOutput
|
||||
log_output = AgentLogOutput(
|
||||
agent_name="test_agent", result="test_result"
|
||||
)
|
||||
assert_is_instance(
|
||||
log_output,
|
||||
AgentLogOutput,
|
||||
"AgentLogOutput creates correct instance",
|
||||
)
|
||||
assert_not_none(
|
||||
log_output.log_id, "AgentLogOutput generates log_id"
|
||||
)
|
||||
assert_equal(
|
||||
log_output.result,
|
||||
"test_result",
|
||||
"AgentLogOutput stores result",
|
||||
)
|
||||
|
||||
# Test TreeLog
|
||||
tree_log = TreeLog(
|
||||
tree_name="test_tree",
|
||||
task="test_task",
|
||||
selected_agent="test_agent",
|
||||
result="test_result",
|
||||
)
|
||||
assert_is_instance(
|
||||
tree_log, TreeLog, "TreeLog creates correct instance"
|
||||
)
|
||||
assert_not_none(tree_log.log_id, "TreeLog generates log_id")
|
||||
assert_equal(
|
||||
tree_log.tree_name, "test_tree", "TreeLog stores tree_name"
|
||||
)
|
||||
|
||||
|
||||
def test_tree_agent_initialization():
|
||||
"""Test TreeAgent initialization and basic properties."""
|
||||
print("\n🧪 Testing TreeAgent initialization...")
|
||||
|
||||
# Test basic initialization
|
||||
agent = TreeAgent(
|
||||
name="Test Agent",
|
||||
description="A test agent",
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
)
|
||||
|
||||
assert_is_instance(
|
||||
agent, TreeAgent, "TreeAgent creates correct instance"
|
||||
)
|
||||
assert_equal(
|
||||
agent.agent_name,
|
||||
"financial_advisor",
|
||||
"TreeAgent stores agent_name",
|
||||
)
|
||||
assert_equal(
|
||||
agent.embedding_model_name,
|
||||
"text-embedding-ada-002",
|
||||
"TreeAgent has default embedding model",
|
||||
)
|
||||
assert_true(
|
||||
len(agent.relevant_keywords) > 0,
|
||||
"TreeAgent extracts keywords from system prompt",
|
||||
)
|
||||
assert_not_none(
|
||||
agent.system_prompt_embedding,
|
||||
"TreeAgent generates system prompt embedding",
|
||||
)
|
||||
|
||||
# Test with custom embedding model
|
||||
agent_custom = TreeAgent(
|
||||
system_prompt="Test prompt",
|
||||
embedding_model_name="custom-model",
|
||||
)
|
||||
assert_equal(
|
||||
agent_custom.embedding_model_name,
|
||||
"custom-model",
|
||||
"TreeAgent accepts custom embedding model",
|
||||
)
|
||||
|
||||
|
||||
def test_tree_agent_distance_calculation():
|
||||
"""Test TreeAgent distance calculation between agents."""
|
||||
print("\n🧪 Testing TreeAgent distance calculation...")
|
||||
|
||||
agent1 = TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
)
|
||||
|
||||
agent2 = TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
)
|
||||
|
||||
agent3 = TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["stock_analyst"],
|
||||
agent_name="stock_analyst",
|
||||
)
|
||||
|
||||
# Test distance calculation
|
||||
distance1 = agent1.calculate_distance(agent2)
|
||||
distance2 = agent1.calculate_distance(agent3)
|
||||
|
||||
assert_true(
|
||||
0.0 <= distance1 <= 1.0, "Distance is between 0 and 1"
|
||||
)
|
||||
assert_true(
|
||||
0.0 <= distance2 <= 1.0, "Distance is between 0 and 1"
|
||||
)
|
||||
assert_true(isinstance(distance1, float), "Distance is a float")
|
||||
|
||||
# Test that identical agents have distance 0
|
||||
identical_agent = TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="identical_advisor",
|
||||
)
|
||||
distance_identical = agent1.calculate_distance(identical_agent)
|
||||
assert_true(
|
||||
distance_identical < 0.1,
|
||||
"Identical agents have very small distance",
|
||||
)
|
||||
|
||||
|
||||
def test_tree_agent_task_relevance():
|
||||
"""Test TreeAgent task relevance checking."""
|
||||
print("\n🧪 Testing TreeAgent task relevance...")
|
||||
|
||||
tax_agent = TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
)
|
||||
|
||||
# Test keyword matching
|
||||
tax_task = SAMPLE_TASKS["tax_question"]
|
||||
is_relevant = tax_agent.is_relevant_for_task(
|
||||
tax_task, threshold=0.7
|
||||
)
|
||||
assert_true(is_relevant, "Tax agent is relevant for tax question")
|
||||
|
||||
# Test non-relevant task
|
||||
stock_task = SAMPLE_TASKS["stock_question"]
|
||||
is_relevant = tax_agent.is_relevant_for_task(
|
||||
stock_task, threshold=0.7
|
||||
)
|
||||
# This might be True due to semantic similarity, so we just check it's a boolean
|
||||
assert_true(
|
||||
isinstance(is_relevant, bool),
|
||||
"Task relevance returns boolean",
|
||||
)
|
||||
|
||||
|
||||
def test_tree_initialization():
|
||||
"""Test Tree initialization and agent organization."""
|
||||
print("\n🧪 Testing Tree initialization...")
|
||||
|
||||
agents = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["stock_analyst"],
|
||||
agent_name="stock_analyst",
|
||||
),
|
||||
]
|
||||
|
||||
tree = Tree("Financial Services Tree", agents)
|
||||
|
||||
assert_equal(
|
||||
tree.tree_name,
|
||||
"Financial Services Tree",
|
||||
"Tree stores tree_name",
|
||||
)
|
||||
assert_equal(len(tree.agents), 3, "Tree contains all agents")
|
||||
assert_true(
|
||||
all(hasattr(agent, "distance") for agent in tree.agents),
|
||||
"All agents have distance calculated",
|
||||
)
|
||||
|
||||
# Test that agents are sorted by distance
|
||||
distances = [agent.distance for agent in tree.agents]
|
||||
assert_true(
|
||||
distances == sorted(distances),
|
||||
"Agents are sorted by distance",
|
||||
)
|
||||
|
||||
|
||||
def test_tree_agent_finding():
|
||||
"""Test Tree agent finding functionality."""
|
||||
print("\n🧪 Testing Tree agent finding...")
|
||||
|
||||
agents = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
),
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
),
|
||||
]
|
||||
|
||||
tree = Tree("Test Tree", agents)
|
||||
|
||||
# Test finding relevant agent
|
||||
tax_task = SAMPLE_TASKS["tax_question"]
|
||||
relevant_agent = tree.find_relevant_agent(tax_task)
|
||||
assert_not_none(
|
||||
relevant_agent, "Tree finds relevant agent for tax task"
|
||||
)
|
||||
|
||||
# Test finding agent for unrelated task
|
||||
unrelated_task = "How do I cook pasta?"
|
||||
relevant_agent = tree.find_relevant_agent(unrelated_task)
|
||||
# This might return None or an agent depending on similarity threshold
|
||||
assert_true(
|
||||
relevant_agent is None
|
||||
or isinstance(relevant_agent, TreeAgent),
|
||||
"Tree handles unrelated tasks",
|
||||
)
|
||||
|
||||
|
||||
def test_forest_swarm_initialization():
|
||||
"""Test ForestSwarm initialization."""
|
||||
print("\n🧪 Testing ForestSwarm initialization...")
|
||||
|
||||
agents_tree1 = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
)
|
||||
]
|
||||
|
||||
agents_tree2 = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
)
|
||||
]
|
||||
|
||||
tree1 = Tree("Financial Tree", agents_tree1)
|
||||
tree2 = Tree("Tax Tree", agents_tree2)
|
||||
|
||||
forest = ForestSwarm(
|
||||
name="Test Forest",
|
||||
description="A test forest",
|
||||
trees=[tree1, tree2],
|
||||
)
|
||||
|
||||
assert_equal(
|
||||
forest.name, "Test Forest", "ForestSwarm stores name"
|
||||
)
|
||||
assert_equal(
|
||||
forest.description,
|
||||
"A test forest",
|
||||
"ForestSwarm stores description",
|
||||
)
|
||||
assert_equal(
|
||||
len(forest.trees), 2, "ForestSwarm contains all trees"
|
||||
)
|
||||
assert_not_none(
|
||||
forest.conversation, "ForestSwarm creates conversation object"
|
||||
)
|
||||
|
||||
|
||||
def test_forest_swarm_tree_finding():
|
||||
"""Test ForestSwarm tree finding functionality."""
|
||||
print("\n🧪 Testing ForestSwarm tree finding...")
|
||||
|
||||
agents_tree1 = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["financial_advisor"],
|
||||
agent_name="financial_advisor",
|
||||
)
|
||||
]
|
||||
|
||||
agents_tree2 = [
|
||||
TreeAgent(
|
||||
system_prompt=SAMPLE_SYSTEM_PROMPTS["tax_expert"],
|
||||
agent_name="tax_expert",
|
||||
)
|
||||
]
|
||||
|
||||
tree1 = Tree("Financial Tree", agents_tree1)
|
||||
tree2 = Tree("Tax Tree", agents_tree2)
|
||||
|
||||
forest = ForestSwarm(trees=[tree1, tree2])
|
||||
|
||||
# Test finding relevant tree for tax question
|
||||
tax_task = SAMPLE_TASKS["tax_question"]
|
||||
relevant_tree = forest.find_relevant_tree(tax_task)
|
||||
assert_not_none(
|
||||
relevant_tree, "ForestSwarm finds relevant tree for tax task"
|
||||
)
|
||||
|
||||
# Test finding relevant tree for financial question
|
||||
financial_task = SAMPLE_TASKS["investment_question"]
|
||||
relevant_tree = forest.find_relevant_tree(financial_task)
|
||||
assert_not_none(
|
||||
relevant_tree,
|
||||
"ForestSwarm finds relevant tree for financial task",
|
||||
)
|
||||
|
||||
|
||||
def test_forest_swarm_execution():
|
||||
"""Test ForestSwarm task execution."""
|
||||
print("\n🧪 Testing ForestSwarm task execution...")
|
||||
|
||||
# Create a simple forest with one tree and one agent
|
||||
agent = TreeAgent(
|
||||
system_prompt="I am a helpful assistant that can answer questions about Delaware incorporation and taxes.",
|
||||
agent_name="delaware_expert",
|
||||
)
|
||||
|
||||
tree = Tree("Delaware Tree", [agent])
|
||||
forest = ForestSwarm(trees=[tree])
|
||||
|
||||
# Test task execution
|
||||
task = "What are the benefits of incorporating in Delaware?"
|
||||
try:
|
||||
result = forest.run(task)
|
||||
assert_not_none(
|
||||
result, "ForestSwarm returns result from task execution"
|
||||
)
|
||||
assert_true(isinstance(result, str), "Result is a string")
|
||||
except Exception as e:
|
||||
# If execution fails due to external dependencies, that's okay for unit tests
|
||||
print(
|
||||
f"⚠️ Task execution failed (expected in unit test environment): {e}"
|
||||
)
|
||||
|
||||
|
||||
def test_edge_cases():
|
||||
"""Test edge cases and error handling."""
|
||||
print("\n🧪 Testing edge cases and error handling...")
|
||||
|
||||
# Test TreeAgent with None system prompt
|
||||
agent_no_prompt = TreeAgent(
|
||||
system_prompt=None, agent_name="no_prompt_agent"
|
||||
)
|
||||
assert_equal(
|
||||
len(agent_no_prompt.relevant_keywords),
|
||||
0,
|
||||
"Agent with None prompt has empty keywords",
|
||||
)
|
||||
assert_true(
|
||||
agent_no_prompt.system_prompt_embedding is None,
|
||||
"Agent with None prompt has None embedding",
|
||||
)
|
||||
|
||||
# Test Tree with empty agents list
|
||||
empty_tree = Tree("Empty Tree", [])
|
||||
assert_equal(
|
||||
len(empty_tree.agents), 0, "Empty tree has no agents"
|
||||
)
|
||||
|
||||
# Test ForestSwarm with empty trees list
|
||||
empty_forest = ForestSwarm(trees=[])
|
||||
assert_equal(
|
||||
len(empty_forest.trees), 0, "Empty forest has no trees"
|
||||
)
|
||||
|
||||
# Test cosine_similarity with empty vectors
|
||||
empty_vec = []
|
||||
vec = [1.0, 0.0, 0.0]
|
||||
similarity = cosine_similarity(empty_vec, vec)
|
||||
assert_equal(
|
||||
similarity, 0.0, "cosine_similarity handles empty vectors"
|
||||
)
|
||||
|
||||
|
||||
def run_all_tests():
|
||||
"""Run all unit tests and display results."""
|
||||
print("🚀 Starting ForestSwarm Unit Tests...")
|
||||
print("=" * 60)
|
||||
|
||||
# Run all test functions
|
||||
test_functions = [
|
||||
test_extract_keywords,
|
||||
test_cosine_similarity,
|
||||
test_agent_log_models,
|
||||
test_tree_agent_initialization,
|
||||
test_tree_agent_distance_calculation,
|
||||
test_tree_agent_task_relevance,
|
||||
test_tree_initialization,
|
||||
test_tree_agent_finding,
|
||||
test_forest_swarm_initialization,
|
||||
test_forest_swarm_tree_finding,
|
||||
test_forest_swarm_execution,
|
||||
test_edge_cases,
|
||||
]
|
||||
|
||||
for test_func in test_functions:
|
||||
try:
|
||||
test_func()
|
||||
except Exception as e:
|
||||
test_results["total"] += 1
|
||||
test_results["failed"] += 1
|
||||
print(f"❌ ERROR: {test_func.__name__} - {e}")
|
||||
|
||||
# Display results
|
||||
print("\n" + "=" * 60)
|
||||
print("📊 TEST RESULTS SUMMARY")
|
||||
print("=" * 60)
|
||||
print(f"Total Tests: {test_results['total']}")
|
||||
print(f"Passed: {test_results['passed']}")
|
||||
print(f"Failed: {test_results['failed']}")
|
||||
|
||||
success_rate = (
|
||||
(test_results["passed"] / test_results["total"]) * 100
|
||||
if test_results["total"] > 0
|
||||
else 0
|
||||
)
|
||||
print(f"Success Rate: {success_rate:.1f}%")
|
||||
|
||||
if test_results["failed"] == 0:
|
||||
print(
|
||||
"\n🎉 All tests passed! ForestSwarm is working correctly."
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n⚠️ {test_results['failed']} test(s) failed. Please review the failures above."
|
||||
)
|
||||
|
||||
return test_results["failed"] == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = run_all_tests()
|
||||
sys.exit(0 if success else 1)
|
||||
Loading…
Reference in new issue