mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 12:47:55 -10:00
252 lines
9.0 KiB
Python
252 lines
9.0 KiB
Python
from praisonaiagents import Agent, Task, AgentTeam
|
|
import random
|
|
from typing import List, Dict, Union
|
|
import json
|
|
from pydantic import BaseModel
|
|
|
|
# Add Pydantic models for data validation
|
|
class Person(BaseModel):
|
|
name: str
|
|
age: int
|
|
job: str
|
|
city: str
|
|
salary: int
|
|
|
|
class ProcessedPerson(Person):
|
|
salary_category: str
|
|
age_group: str
|
|
processing_status: str
|
|
|
|
class DataList(BaseModel):
|
|
items: List[Dict]
|
|
|
|
class ValidationResult(BaseModel):
|
|
validation_result: str
|
|
details: str = ""
|
|
|
|
def random_data_of_individuals() -> List[Dict]:
|
|
"""Generate random individual data"""
|
|
names = ["John", "Jane", "Mike", "Sarah", "David", "Emma"]
|
|
jobs = ["Engineer", "Doctor", "Teacher", "Artist", "Developer"]
|
|
cities = ["New York", "London", "Tokyo", "Paris", "Berlin"]
|
|
|
|
data = []
|
|
for _ in range(random.randint(3, 7)):
|
|
person = {
|
|
"name": random.choice(names),
|
|
"age": random.randint(25, 65),
|
|
"job": random.choice(jobs),
|
|
"city": random.choice(cities),
|
|
"salary": random.randint(30000, 120000)
|
|
}
|
|
data.append(person)
|
|
return data
|
|
|
|
def process_data_of_individuals(data: Union[List[Dict], Dict, str]) -> Dict:
|
|
"""Process individual data by adding categories and analysis"""
|
|
try:
|
|
print("\n[DEBUG] process_data_of_individuals input:", data)
|
|
print("[DEBUG] Current workflow state:")
|
|
print("- total_records:", workflow.get_state("total_records"))
|
|
print("- current_index:", workflow.get_state("current_index"))
|
|
print("- remaining:", workflow.get_state("remaining"))
|
|
|
|
# Get the items list from the collect_task result
|
|
collect_result = None
|
|
for task in workflow.tasks.values():
|
|
if task.name == "collect_data" and task.result:
|
|
try:
|
|
collect_data = json.loads(task.result.raw)
|
|
collect_result = collect_data.get("items", [])
|
|
print("[DEBUG] Found collect_data items:", len(collect_result))
|
|
except:
|
|
print("[DEBUG] Failed to parse collect_data result")
|
|
|
|
# Handle string input by trying to parse it as JSON
|
|
if isinstance(data, str):
|
|
if ":" in data and not data.strip().startswith("{"):
|
|
# Convert string format to dictionary
|
|
pairs = [pair.strip() for pair in data.split(",")]
|
|
data_dict = {}
|
|
for pair in pairs:
|
|
key, value = pair.split(":")
|
|
key = key.strip().lower()
|
|
value = value.strip()
|
|
if key == "age" or key == "salary":
|
|
value = int(value)
|
|
data_dict[key] = value
|
|
data = data_dict
|
|
else:
|
|
data = json.loads(data)
|
|
print("[DEBUG] Parsed data:", data)
|
|
|
|
# Handle single record
|
|
if isinstance(data, dict):
|
|
person = data
|
|
# Initialize total records if not set
|
|
total_records = workflow.get_state("total_records")
|
|
if total_records is None and collect_result:
|
|
total_records = len(collect_result)
|
|
workflow.set_state("total_records", total_records)
|
|
print(f"[DEBUG] Initialized total_records to {total_records}")
|
|
|
|
current_index = workflow.get_state("current_index", 0)
|
|
total_records = total_records or 1
|
|
remaining = total_records - (current_index + 1)
|
|
workflow.set_state("remaining", remaining)
|
|
print(f"[DEBUG] Processing record {current_index + 1}/{total_records}")
|
|
|
|
elif isinstance(data, list):
|
|
if len(data) == 0:
|
|
raise ValueError("Empty data list")
|
|
person = data[0]
|
|
workflow.set_state("total_records", len(data))
|
|
workflow.set_state("current_index", 0)
|
|
workflow.set_state("remaining", len(data) - 1)
|
|
print(f"[DEBUG] First record from list of {len(data)} items")
|
|
else:
|
|
raise ValueError("Input must be a dictionary or list with at least one record")
|
|
|
|
processed_person = person.copy()
|
|
|
|
# Add salary category
|
|
salary = person.get("salary", 0)
|
|
if salary < 50000:
|
|
processed_person["salary_category"] = "entry"
|
|
elif salary < 90000:
|
|
processed_person["salary_category"] = "mid"
|
|
else:
|
|
processed_person["salary_category"] = "senior"
|
|
|
|
# Add age group
|
|
age = person.get("age", 0)
|
|
if age < 35:
|
|
processed_person["age_group"] = "young"
|
|
elif age < 50:
|
|
processed_person["age_group"] = "mid"
|
|
else:
|
|
processed_person["age_group"] = "senior"
|
|
|
|
# Add processing status using workflow state
|
|
remaining = workflow.get_state("remaining", 0)
|
|
current_index = workflow.get_state("current_index", 0)
|
|
total_records = workflow.get_state("total_records", 1)
|
|
|
|
# Update current index for next iteration
|
|
workflow.set_state("current_index", current_index + 1)
|
|
|
|
print(f"[DEBUG] Status check - remaining: {remaining}, current_index: {current_index}, total_records: {total_records}")
|
|
|
|
if remaining <= 0 and current_index >= total_records - 1:
|
|
print("[DEBUG] Setting status to 'all records processed'")
|
|
processed_person["processing_status"] = "all records processed"
|
|
else:
|
|
print(f"[DEBUG] More records to process. Remaining: {remaining}")
|
|
processed_person["processing_status"] = f"more records to process ({remaining} remaining)"
|
|
|
|
print("[DEBUG] Final processed person:", processed_person)
|
|
return processed_person
|
|
|
|
except Exception as e:
|
|
print(f"[DEBUG] Error processing data: {str(e)}")
|
|
return {"error": str(e), "processing_status": "error occurred"}
|
|
|
|
# Create agents
|
|
data_agent = Agent(
|
|
name="DataCollector",
|
|
role="Data collection specialist",
|
|
goal="Collect and validate data about individuals",
|
|
backstory="Expert in gathering and validating demographic data",
|
|
tools=[random_data_of_individuals],
|
|
reflection=False
|
|
)
|
|
|
|
process_agent = Agent(
|
|
name="DataProcessor",
|
|
role="Data processor",
|
|
goal="Process and categorize individual data",
|
|
backstory="Expert in data analysis and categorization",
|
|
tools=[process_data_of_individuals],
|
|
reflection=False
|
|
)
|
|
|
|
# Modify tasks to use Pydantic models
|
|
collect_task = Task(
|
|
description="Collect random individual data using the random_data_of_individuals tool. Return as a JSON object with 'items' array.",
|
|
expected_output="List of individual records with basic information",
|
|
agent=data_agent,
|
|
name="collect_data",
|
|
tools=[random_data_of_individuals],
|
|
is_start=True,
|
|
next_tasks=["validate_data"],
|
|
output_json=DataList
|
|
)
|
|
|
|
validate_task = Task(
|
|
description="""Validate the collected data. Check if:
|
|
1. All required fields are present (name, age, job, city, salary)
|
|
2. Age is between 25 and 65
|
|
3. Salary is between 30000 and 120000
|
|
Return validation_result as 'valid' or 'invalid' with optional details.""",
|
|
expected_output="Validation result indicating if data is valid or invalid",
|
|
agent=data_agent,
|
|
name="validate_data",
|
|
task_type="decision",
|
|
condition={
|
|
"valid": ["process_data"],
|
|
"invalid": ["collect_data"]
|
|
},
|
|
output_json=ValidationResult
|
|
)
|
|
|
|
process_task = Task(
|
|
description="""Process one record at a time from the input data.
|
|
Current progress will be shown in Loop Status.
|
|
|
|
For the current record:
|
|
1. Use process_data_of_individuals tool to add categories
|
|
2. Return the processed record with remaining count
|
|
|
|
Current remaining: {remaining}
|
|
Current item: {current_item}
|
|
|
|
Process this record and indicate if more records remain.""",
|
|
expected_output="Processed record with categories and status",
|
|
agent=process_agent,
|
|
name="process_data",
|
|
tools=[process_data_of_individuals],
|
|
task_type="loop",
|
|
condition={
|
|
"more records to process": ["process_data"],
|
|
"all records processed": []
|
|
},
|
|
context=[collect_task],
|
|
output_json=ProcessedPerson
|
|
)
|
|
|
|
# Create Agents instance with workflow process
|
|
workflow = AgentTeam(
|
|
agents=[data_agent, process_agent],
|
|
tasks=[collect_task, validate_task, process_task],
|
|
process="workflow"
|
|
)
|
|
|
|
# Run the workflow
|
|
result = workflow.start()
|
|
|
|
# Print results
|
|
print("\nWorkflow Results:")
|
|
print("----------------")
|
|
for task_id, task in workflow.tasks.items():
|
|
print(f"\nTask: {task.name}")
|
|
print(f"Status: {task.status}")
|
|
if task.result:
|
|
print("Output:")
|
|
try:
|
|
# Try to format as pretty JSON
|
|
import json
|
|
output = json.loads(task.result.raw)
|
|
print(json.dumps(output, indent=2))
|
|
except:
|
|
# If not JSON, print raw output
|
|
print(task.result.raw[:500] + "..." if len(task.result.raw) > 500 else task.result.raw) |