34 lines
1.4 KiB
Python
34 lines
1.4 KiB
Python
import ray
|
|
from pydantic_ai import Agent
|
|
from pretor.core.workflow_manager.workflow import WorkStep
|
|
|
|
@ray.remote
|
|
class ControlNode:
|
|
def __init__(self, agent: Agent):
|
|
self.agent = agent
|
|
|
|
async def execute_step(self, step: WorkStep) -> WorkStep:
|
|
if step.action == "dispatch_model":
|
|
# The WorkStep schema from workflow manager may pass target info differently, assuming `input` here or simple `desc`
|
|
result = await self.dispatch_model({}, f"Execute task: {step.desc}")
|
|
step.output = result
|
|
elif step.action == "dispatch_tool":
|
|
# Simulating parsing of tool and args from `desc` or `input`
|
|
result = await self.dispatch_tool("simulated_tool", {"desc": step.desc})
|
|
step.output = str(result)
|
|
else:
|
|
result = await self.agent.run(f"Execute generic step: {step.model_dump()}")
|
|
step.output = result.data
|
|
|
|
step.status = "completed"
|
|
return step
|
|
|
|
async def dispatch_model(self, model_info: dict, prompt: str) -> str:
|
|
# In a real system, we'd select a smaller/specific model based on model_info
|
|
result = await self.agent.run(prompt)
|
|
return result.data
|
|
|
|
async def dispatch_tool(self, tool_name: str, tool_args: dict) -> dict:
|
|
# Simulated tool dispatch
|
|
return {"tool": tool_name, "status": "executed", "args": tool_args}
|