You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
2.4 KiB
91 lines
2.4 KiB
import time
|
|
from typing import List, Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class AgentSchema(BaseModel):
|
|
name: str = None
|
|
system_prompt: str = None
|
|
task: str = None
|
|
response: str = None
|
|
|
|
|
|
class JambaSwarmRequest(BaseModel):
|
|
task: str = (None,)
|
|
plan: str = None
|
|
agents: List[AgentSchema] = None
|
|
timestamp: int = int(time.time())
|
|
|
|
|
|
class JambaSwarmResponse(BaseModel):
|
|
task: str = (None,)
|
|
plan: str = None
|
|
agents: List[AgentSchema] = None
|
|
timestamp: int = int(time.time())
|
|
response: str = None
|
|
|
|
|
|
class AgentSchema(BaseModel):
|
|
name: Optional[str] = None
|
|
system_prompt: Optional[str] = None
|
|
task: Optional[str] = None
|
|
response: Optional[str] = None
|
|
|
|
|
|
class DirectorSettings(BaseModel):
|
|
name: str
|
|
strategy: str
|
|
objectives: List[str]
|
|
|
|
|
|
class BossSettings(BaseModel):
|
|
name: str
|
|
decision_making_strategy: str
|
|
recruitment_strategy: str
|
|
|
|
|
|
class TaskDistribution(BaseModel):
|
|
task: str
|
|
assigned_agents: List[str]
|
|
|
|
|
|
class JambaSwarmRequest(BaseModel):
|
|
task: Optional[str] = None
|
|
plan: Optional[str] = None
|
|
agents: Optional[List[AgentSchema]] = None
|
|
director_settings: DirectorSettings
|
|
boss_settings: BossSettings
|
|
task_distribution: Optional[List[TaskDistribution]] = None
|
|
timestamp: int = int(time.time())
|
|
|
|
|
|
class JambaSwarmResponse(BaseModel):
|
|
task: Optional[str] = None
|
|
plan: Optional[str] = None
|
|
agents: Optional[List[AgentSchema]] = None
|
|
response: Optional[str] = None
|
|
timestamp: int = int(time.time())
|
|
|
|
|
|
# Sample usage:
|
|
|
|
|
|
# try:
|
|
# request = JambaSwarmRequest(
|
|
# task="Research on AI",
|
|
# plan="Execute a comprehensive research plan",
|
|
# agents=[
|
|
# AgentSchema(name="Agent1", system_prompt="Analyze recent AI papers", task="AI research task"),
|
|
# AgentSchema(name="Agent2", system_prompt="Summarize AI research findings", task="Summarization task"),
|
|
# ],
|
|
# director_settings=DirectorSettings(name="Director1", strategy="Hierarchical", objectives=["Efficiency", "Accuracy"]),
|
|
# boss_settings=BossSettings(name="Boss1", decision_making_strategy="Collaborative", recruitment_strategy="Pre-selected"),
|
|
# task_distribution=[
|
|
# TaskDistribution(task="Research on AI", assigned_agents=["Agent1", "Agent2"])
|
|
# ]
|
|
# )
|
|
# print(request.json())
|
|
# except ValidationError as e:
|
|
# print(e.json())
|