feat(agent): add support for workflow agents with updated constraints and configuration

This commit is contained in:
Davidson Gomes 2025-05-06 18:37:13 -03:00
parent 64e483533d
commit 0fc47aaa57
5 changed files with 97 additions and 8 deletions

View File

@ -0,0 +1,49 @@
"""worflow_agent
Revision ID: ebac70616dab
Revises: c107446e38aa
Create Date: 2025-05-06 17:05:26.884902
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = "ebac70616dab"
down_revision: Union[str, None] = "c107446e38aa"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
# Remover a constraint antiga
op.drop_constraint("check_agent_type", "agents", type_="check")
# Adicionar a nova constraint que inclui o tipo 'workflow'
op.create_check_constraint(
"check_agent_type",
"agents",
"type IN ('llm', 'sequential', 'parallel', 'loop', 'a2a', 'workflow')",
)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
# Remover a constraint nova
op.drop_constraint("check_agent_type", "agents", type_="check")
# Restaurar a constraint anterior sem o tipo 'workflow'
op.create_check_constraint(
"check_agent_type",
"agents",
"type IN ('llm', 'sequential', 'parallel', 'loop', 'a2a')",
)
# ### end Alembic commands ###

View File

@ -81,7 +81,7 @@ class Agent(Base):
__table_args__ = (
CheckConstraint(
"type IN ('llm', 'sequential', 'parallel', 'loop', 'a2a')",
"type IN ('llm', 'sequential', 'parallel', 'loop', 'a2a', 'workflow')",
name="check_agent_type",
),
)

View File

@ -1,4 +1,4 @@
from typing import List, Optional, Dict, Union
from typing import List, Optional, Dict, Union, Any
from pydantic import BaseModel, Field
from uuid import UUID
import secrets
@ -44,6 +44,13 @@ class CustomMCPServerConfig(BaseModel):
from_attributes = True
class FlowNodes(BaseModel):
"""Configuration of workflow nodes"""
nodes: List[Any]
edges: List[Any]
class HTTPToolParameter(BaseModel):
"""Parameter of an HTTP tool"""
@ -133,6 +140,9 @@ class LLMConfig(BaseModel):
sub_agents: Optional[List[UUID]] = Field(
default=None, description="List of IDs of sub-agents"
)
workflow: Optional[FlowNodes] = Field(
default=None, description="Workflow configuration"
)
class Config:
from_attributes = True
@ -175,3 +185,20 @@ class LoopConfig(BaseModel):
class Config:
from_attributes = True
class WorkflowConfig(BaseModel):
"""Configuration for workflow agents"""
workflow: Dict[str, Any] = Field(
..., description="Workflow configuration with nodes and edges"
)
sub_agents: Optional[List[UUID]] = Field(
default_factory=list, description="List of IDs of sub-agents used in workflow"
)
api_key: Optional[str] = Field(
default_factory=generate_api_key, description="API key for the workflow agent"
)
class Config:
from_attributes = True

View File

@ -57,7 +57,7 @@ class AgentBase(BaseModel):
)
description: Optional[str] = Field(None, description="Agent description")
type: str = Field(
..., description="Agent type (llm, sequential, parallel, loop, a2a)"
..., description="Agent type (llm, sequential, parallel, loop, a2a, workflow)"
)
model: Optional[str] = Field(
None, description="Agent model (required only for llm type)"
@ -69,9 +69,7 @@ class AgentBase(BaseModel):
agent_card_url: Optional[str] = Field(
None, description="Agent card URL (required for a2a type)"
)
config: Optional[Union[LLMConfig, Dict[str, Any]]] = Field(
None, description="Agent configuration based on type"
)
config: Any = Field(None, description="Agent configuration based on type")
@validator("name")
def validate_name(cls, v, values):
@ -87,9 +85,9 @@ class AgentBase(BaseModel):
@validator("type")
def validate_type(cls, v):
if v not in ["llm", "sequential", "parallel", "loop", "a2a"]:
if v not in ["llm", "sequential", "parallel", "loop", "a2a", "workflow"]:
raise ValueError(
"Invalid agent type. Must be: llm, sequential, parallel, loop or a2a"
"Invalid agent type. Must be: llm, sequential, parallel, loop, a2a or workflow"
)
return v
@ -122,6 +120,10 @@ class AgentBase(BaseModel):
if "type" not in values:
return v
# Para agentes workflow, não fazemos nenhuma validação
if "type" in values and values["type"] == "workflow":
return v
if not v and values.get("type") != "a2a":
raise ValueError(
f"Configuration is required for {values.get('type')} agent type"
@ -147,6 +149,7 @@ class AgentBase(BaseModel):
raise ValueError(
f'Agent {values["type"]} must have at least one sub-agent'
)
return v

View File

@ -147,6 +147,16 @@ async def create_agent(db: Session, agent: AgentCreate) -> Agent:
detail=f"Failed to process agent card: {str(e)}",
)
# Para agentes workflow, não fazemos nenhuma validação específica
# apenas garantimos que config é um dicionário
elif agent.type == "workflow":
if not isinstance(agent.config, dict):
agent.config = {}
# Garantir a API key
if "api_key" not in agent.config or not agent.config["api_key"]:
agent.config["api_key"] = generate_api_key()
# Additional sub-agent validation (for non-llm and non-a2a types)
elif agent.type != "llm":
if not isinstance(agent.config, dict):