parent
4e481a0cab
commit
35f7e890b6
@ -1,430 +0,0 @@
|
||||
"""
|
||||
Schemas for Talk Structurally, Act Hierarchically Framework
|
||||
|
||||
This module defines the Pydantic schemas used in the TalkHierarchical framework
|
||||
for structured communication and hierarchical evaluation.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Union
|
||||
from pydantic import BaseModel, Field
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class CommunicationType(str, Enum):
|
||||
"""Types of communication in the structured protocol"""
|
||||
MESSAGE = "message" # M_ij: Specific task instructions
|
||||
BACKGROUND = "background" # B_ij: Context and problem background
|
||||
INTERMEDIATE_OUTPUT = "intermediate_output" # I_ij: Intermediate results
|
||||
|
||||
|
||||
class AgentRole(str, Enum):
|
||||
"""Roles for agents in the hierarchical system"""
|
||||
SUPERVISOR = "supervisor"
|
||||
GENERATOR = "generator"
|
||||
EVALUATOR = "evaluator"
|
||||
REFINER = "refiner"
|
||||
COORDINATOR = "coordinator"
|
||||
|
||||
|
||||
class StructuredMessageSchema(BaseModel):
|
||||
"""Schema for structured communication messages"""
|
||||
message: str = Field(
|
||||
description="Specific task instructions (M_ij)",
|
||||
min_length=1
|
||||
)
|
||||
background: str = Field(
|
||||
description="Context and problem background (B_ij)",
|
||||
default=""
|
||||
)
|
||||
intermediate_output: str = Field(
|
||||
description="Intermediate results (I_ij)",
|
||||
default=""
|
||||
)
|
||||
sender: str = Field(
|
||||
description="Name of the sending agent",
|
||||
min_length=1
|
||||
)
|
||||
recipient: str = Field(
|
||||
description="Name of the receiving agent",
|
||||
min_length=1
|
||||
)
|
||||
timestamp: Optional[str] = Field(
|
||||
description="Timestamp of the message",
|
||||
default=None
|
||||
)
|
||||
communication_type: CommunicationType = Field(
|
||||
description="Type of communication",
|
||||
default=CommunicationType.MESSAGE
|
||||
)
|
||||
|
||||
|
||||
class HierarchicalOrderSchema(BaseModel):
|
||||
"""Schema for hierarchical task orders"""
|
||||
agent_name: str = Field(
|
||||
description="Name of the agent to receive the task",
|
||||
min_length=1
|
||||
)
|
||||
task: str = Field(
|
||||
description="Specific task description",
|
||||
min_length=1
|
||||
)
|
||||
communication_type: CommunicationType = Field(
|
||||
description="Type of communication to use",
|
||||
default=CommunicationType.MESSAGE
|
||||
)
|
||||
background_context: str = Field(
|
||||
description="Background context for the task",
|
||||
default=""
|
||||
)
|
||||
intermediate_output: str = Field(
|
||||
description="Intermediate output to pass along",
|
||||
default=""
|
||||
)
|
||||
priority: int = Field(
|
||||
description="Task priority (1-10, higher is more important)",
|
||||
default=5,
|
||||
ge=1,
|
||||
le=10
|
||||
)
|
||||
|
||||
|
||||
class EvaluationCriterionSchema(BaseModel):
|
||||
"""Schema for evaluation criteria"""
|
||||
name: str = Field(
|
||||
description="Name of the evaluation criterion",
|
||||
min_length=1
|
||||
)
|
||||
description: str = Field(
|
||||
description="Description of what this criterion evaluates",
|
||||
min_length=1
|
||||
)
|
||||
weight: float = Field(
|
||||
description="Weight of this criterion in overall evaluation (0-1)",
|
||||
default=1.0,
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
scale: str = Field(
|
||||
description="Scale for evaluation (e.g., '0-10', 'A-F', 'Poor-Excellent')",
|
||||
default="0-10"
|
||||
)
|
||||
|
||||
|
||||
class EvaluationResultSchema(BaseModel):
|
||||
"""Schema for evaluation results"""
|
||||
evaluator_name: str = Field(
|
||||
description="Name of the evaluator",
|
||||
min_length=1
|
||||
)
|
||||
criterion: str = Field(
|
||||
description="Evaluation criterion",
|
||||
min_length=1
|
||||
)
|
||||
score: float = Field(
|
||||
description="Evaluation score",
|
||||
ge=0.0
|
||||
)
|
||||
feedback: str = Field(
|
||||
description="Detailed feedback",
|
||||
min_length=1
|
||||
)
|
||||
confidence: float = Field(
|
||||
description="Confidence in evaluation (0-1)",
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Reasoning behind the evaluation",
|
||||
default=""
|
||||
)
|
||||
suggestions: List[str] = Field(
|
||||
description="Suggestions for improvement",
|
||||
default=[]
|
||||
)
|
||||
|
||||
|
||||
class EvaluationSummarySchema(BaseModel):
|
||||
"""Schema for evaluation summaries"""
|
||||
overall_score: float = Field(
|
||||
description="Overall evaluation score",
|
||||
ge=0.0
|
||||
)
|
||||
criterion_scores: Dict[str, float] = Field(
|
||||
description="Scores for each criterion",
|
||||
default={}
|
||||
)
|
||||
strengths: List[str] = Field(
|
||||
description="Identified strengths",
|
||||
default=[]
|
||||
)
|
||||
weaknesses: List[str] = Field(
|
||||
description="Identified weaknesses",
|
||||
default=[]
|
||||
)
|
||||
recommendations: List[str] = Field(
|
||||
description="Recommendations for improvement",
|
||||
default=[]
|
||||
)
|
||||
confidence: float = Field(
|
||||
description="Overall confidence in evaluation",
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
|
||||
|
||||
class AgentConfigSchema(BaseModel):
|
||||
"""Schema for agent configuration"""
|
||||
name: str = Field(
|
||||
description="Name of the agent",
|
||||
min_length=1
|
||||
)
|
||||
role: AgentRole = Field(
|
||||
description="Role of the agent in the system"
|
||||
)
|
||||
model_name: str = Field(
|
||||
description="Model to use for this agent",
|
||||
default="gpt-4o-mini"
|
||||
)
|
||||
system_prompt: str = Field(
|
||||
description="System prompt for the agent",
|
||||
min_length=1
|
||||
)
|
||||
tools: List[str] = Field(
|
||||
description="Tools available to this agent",
|
||||
default=[]
|
||||
)
|
||||
capabilities: List[str] = Field(
|
||||
description="Capabilities of this agent",
|
||||
default=[]
|
||||
)
|
||||
evaluation_criteria: List[str] = Field(
|
||||
description="Evaluation criteria this agent can assess",
|
||||
default=[]
|
||||
)
|
||||
|
||||
|
||||
class TalkHierarchicalConfigSchema(BaseModel):
|
||||
"""Schema for TalkHierarchical swarm configuration"""
|
||||
name: str = Field(
|
||||
description="Name of the swarm",
|
||||
default="TalkHierarchicalSwarm"
|
||||
)
|
||||
description: str = Field(
|
||||
description="Description of the swarm",
|
||||
default="Talk Structurally, Act Hierarchically Framework"
|
||||
)
|
||||
max_loops: int = Field(
|
||||
description="Maximum number of refinement loops",
|
||||
default=3,
|
||||
ge=1
|
||||
)
|
||||
enable_structured_communication: bool = Field(
|
||||
description="Enable structured communication protocol",
|
||||
default=True
|
||||
)
|
||||
enable_hierarchical_evaluation: bool = Field(
|
||||
description="Enable hierarchical evaluation system",
|
||||
default=True
|
||||
)
|
||||
shared_memory: bool = Field(
|
||||
description="Enable shared memory between agents",
|
||||
default=True
|
||||
)
|
||||
evaluation_criteria: List[EvaluationCriterionSchema] = Field(
|
||||
description="Evaluation criteria for the system",
|
||||
default=[]
|
||||
)
|
||||
agents: List[AgentConfigSchema] = Field(
|
||||
description="Configuration for all agents",
|
||||
default=[]
|
||||
)
|
||||
quality_threshold: float = Field(
|
||||
description="Quality threshold for stopping refinement",
|
||||
default=8.0,
|
||||
ge=0.0,
|
||||
le=10.0
|
||||
)
|
||||
|
||||
|
||||
class TalkHierarchicalStateSchema(BaseModel):
|
||||
"""Schema for TalkHierarchical swarm state"""
|
||||
conversation_history: List[StructuredMessageSchema] = Field(
|
||||
description="History of structured messages",
|
||||
default=[]
|
||||
)
|
||||
intermediate_outputs: Dict[str, str] = Field(
|
||||
description="Intermediate outputs from agents",
|
||||
default={}
|
||||
)
|
||||
evaluation_results: List[EvaluationResultSchema] = Field(
|
||||
description="Results from evaluations",
|
||||
default=[]
|
||||
)
|
||||
current_loop: int = Field(
|
||||
description="Current refinement loop",
|
||||
default=0
|
||||
)
|
||||
task: str = Field(
|
||||
description="Current task being processed",
|
||||
default=""
|
||||
)
|
||||
final_result: Optional[str] = Field(
|
||||
description="Final result of the workflow",
|
||||
default=None
|
||||
)
|
||||
|
||||
|
||||
class TalkHierarchicalResponseSchema(BaseModel):
|
||||
"""Schema for TalkHierarchical swarm response"""
|
||||
final_result: str = Field(
|
||||
description="Final result of the workflow"
|
||||
)
|
||||
total_loops: int = Field(
|
||||
description="Total number of loops executed",
|
||||
ge=1
|
||||
)
|
||||
conversation_history: List[StructuredMessageSchema] = Field(
|
||||
description="Complete conversation history"
|
||||
)
|
||||
evaluation_results: List[EvaluationResultSchema] = Field(
|
||||
description="All evaluation results"
|
||||
)
|
||||
intermediate_outputs: Dict[str, str] = Field(
|
||||
description="All intermediate outputs"
|
||||
)
|
||||
evaluation_summary: Optional[EvaluationSummarySchema] = Field(
|
||||
description="Summary of evaluations",
|
||||
default=None
|
||||
)
|
||||
performance_metrics: Dict[str, Union[float, int, str]] = Field(
|
||||
description="Performance metrics",
|
||||
default={}
|
||||
)
|
||||
|
||||
|
||||
class GraphNodeSchema(BaseModel):
|
||||
"""Schema for graph nodes in agent orchestration"""
|
||||
node_id: str = Field(
|
||||
description="Unique identifier for the node",
|
||||
min_length=1
|
||||
)
|
||||
agent_name: str = Field(
|
||||
description="Name of the agent at this node",
|
||||
min_length=1
|
||||
)
|
||||
role: AgentRole = Field(
|
||||
description="Role of the agent"
|
||||
)
|
||||
capabilities: List[str] = Field(
|
||||
description="Capabilities of this node",
|
||||
default=[]
|
||||
)
|
||||
connections: List[str] = Field(
|
||||
description="IDs of connected nodes",
|
||||
default=[]
|
||||
)
|
||||
conditions: Dict[str, str] = Field(
|
||||
description="Conditions for routing to this node",
|
||||
default={}
|
||||
)
|
||||
|
||||
|
||||
class GraphSchema(BaseModel):
|
||||
"""Schema for agent orchestration graph"""
|
||||
nodes: List[GraphNodeSchema] = Field(
|
||||
description="Nodes in the graph",
|
||||
default=[]
|
||||
)
|
||||
start_node: str = Field(
|
||||
description="ID of the start node",
|
||||
min_length=1
|
||||
)
|
||||
end_nodes: List[str] = Field(
|
||||
description="IDs of end nodes",
|
||||
default=[]
|
||||
)
|
||||
routing_rules: Dict[str, str] = Field(
|
||||
description="Rules for routing between nodes",
|
||||
default={}
|
||||
)
|
||||
max_depth: int = Field(
|
||||
description="Maximum depth of the graph",
|
||||
default=10,
|
||||
ge=1
|
||||
)
|
||||
|
||||
|
||||
# Response schemas for different agent types
|
||||
class GeneratorResponseSchema(BaseModel):
|
||||
"""Schema for generator agent responses"""
|
||||
content: str = Field(
|
||||
description="Generated content",
|
||||
min_length=1
|
||||
)
|
||||
intermediate_output: str = Field(
|
||||
description="Intermediate output for next agent",
|
||||
default=""
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Reasoning behind the generation",
|
||||
default=""
|
||||
)
|
||||
confidence: float = Field(
|
||||
description="Confidence in the generated content",
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
|
||||
|
||||
class EvaluatorResponseSchema(BaseModel):
|
||||
"""Schema for evaluator agent responses"""
|
||||
criterion: str = Field(
|
||||
description="Evaluation criterion",
|
||||
min_length=1
|
||||
)
|
||||
score: float = Field(
|
||||
description="Evaluation score",
|
||||
ge=0.0
|
||||
)
|
||||
feedback: str = Field(
|
||||
description="Detailed feedback",
|
||||
min_length=1
|
||||
)
|
||||
confidence: float = Field(
|
||||
description="Confidence in evaluation",
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Reasoning behind the evaluation",
|
||||
default=""
|
||||
)
|
||||
suggestions: List[str] = Field(
|
||||
description="Suggestions for improvement",
|
||||
default=[]
|
||||
)
|
||||
|
||||
|
||||
class RefinerResponseSchema(BaseModel):
|
||||
"""Schema for refiner agent responses"""
|
||||
refined_content: str = Field(
|
||||
description="Refined content",
|
||||
min_length=1
|
||||
)
|
||||
changes_made: List[str] = Field(
|
||||
description="List of changes made",
|
||||
default=[]
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description="Reasoning behind refinements",
|
||||
default=""
|
||||
)
|
||||
confidence: float = Field(
|
||||
description="Confidence in refinements",
|
||||
ge=0.0,
|
||||
le=1.0
|
||||
)
|
||||
feedback_addressed: List[str] = Field(
|
||||
description="Feedback points addressed",
|
||||
default=[]
|
||||
)
|
Loading…
Reference in new issue