Underground-Digital/Workflow-Engine

1from enum import Enum2from typing import Any, Literal, Optional, Union3 4from pydantic import BaseModel5 6 7class AgentToolEntity(BaseModel):8    """9    Agent Tool Entity.10    """11 12    provider_type: Literal["builtin", "api", "workflow"]13    provider_id: str14    tool_name: str15    tool_parameters: dict[str, Any] = {}16 17 18class AgentPromptEntity(BaseModel):19    """20    Agent Prompt Entity.21    """22 23    first_prompt: str24    next_iteration: str25 26 27class AgentScratchpadUnit(BaseModel):28    """29    Agent First Prompt Entity.30    """31 32    class Action(BaseModel):33        """34        Action Entity.35        """36 37        action_name: str38        action_input: Union[dict, str]39 40        def to_dict(self) -> dict:41            """42            Convert to dictionary.43            """44            return {45                "action": self.action_name,46                "action_input": self.action_input,47            }48 49    agent_response: Optional[str] = None50    thought: Optional[str] = None51    action_str: Optional[str] = None52    observation: Optional[str] = None53    action: Optional[Action] = None54 55    def is_final(self) -> bool:56        """57        Check if the scratchpad unit is final.58        """59        return self.action is None or (60            "final" in self.action.action_name.lower() and "answer" in self.action.action_name.lower()61        )62 63 64class AgentEntity(BaseModel):65    """66    Agent Entity.67    """68 69    class Strategy(Enum):70        """71        Agent Strategy.72        """73 74        CHAIN_OF_THOUGHT = "chain-of-thought"75        FUNCTION_CALLING = "function-calling"76 77    provider: str78    model: str79    strategy: Strategy80    prompt: Optional[AgentPromptEntity] = None81    tools: list[AgentToolEntity] = None82    max_iteration: int = 583