mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-14 05:37:40 +02:00
97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
# TODO: NON HAIK DEPS: ws_manager
|
|
|
|
from copy import deepcopy
|
|
from uuid import uuid4
|
|
import asyncio
|
|
import os
|
|
|
|
from claude_agent_sdk import ClaudeAgentOptions
|
|
from pydantic import BaseModel, Field, InstanceOf
|
|
from typing import List, Literal, Optional
|
|
from typeguard import typechecked
|
|
|
|
from backend.apps.agents.manager.ws_manager import ws_manager
|
|
from backend.apps.agents.HaikFix.Agent.run_agent_loop.run_agent_loop import run_agent_loop
|
|
from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import Message
|
|
from backend.apps.agents.HaikFix.Agent.shared_structs.ApprovalRequest import ApprovalRequest
|
|
from backend.apps.agents.HaikFix.Agent.shared_structs.MessageLog import MessageLog
|
|
|
|
os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")
|
|
|
|
|
|
# NOTE and TODO: we shld remove the ws streaming from this class bc it conflicts with the browser agent
|
|
|
|
class Agent(BaseModel):
|
|
model: str
|
|
mode: str
|
|
status: Literal["running", "waiting_approval", "completed", "error", "stopped"]
|
|
pending_approvals: List[ApprovalRequest] = Field(default_factory=list)
|
|
|
|
messages: MessageLog = Field(default_factory=MessageLog)
|
|
|
|
session_id: str = Field(default_factory=lambda: uuid4().hex)
|
|
config: ClaudeAgentOptions
|
|
|
|
branch_id: str = "main"
|
|
sub_agents: List["Agent"] = Field(default_factory=list)
|
|
sub_branches: List["Agent"] = Field(default_factory=list)
|
|
parent_id: Optional[str] = None
|
|
|
|
task: Optional[asyncio.Task] = None
|
|
lock: InstanceOf[asyncio.Lock] = Field(default_factory=asyncio.Lock)
|
|
|
|
@typechecked
|
|
async def send_message(self, msg: Message) -> None:
|
|
async with self.lock:
|
|
if self.task is not None and not self.task.done():
|
|
print("[Agent.send_message] Agent is already running")
|
|
return
|
|
|
|
await ws_manager.emit_message(self.session_id, msg)
|
|
self.status = "running"
|
|
await ws_manager.emit_status(self.session_id, "running")
|
|
self.messages.append(msg)
|
|
|
|
self.task = asyncio.create_task(run_agent_loop(
|
|
msg=msg,
|
|
options=self.config,
|
|
branch_id=self.branch_id,
|
|
parent_id=self.parent_id,
|
|
))
|
|
|
|
@typechecked
|
|
async def stop_agent(self):
|
|
for child in self.sub_agents:
|
|
await child.stop_agent()
|
|
|
|
if self.task and not self.task.done():
|
|
self.task.cancel()
|
|
try:
|
|
await self.task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
for req in list[ApprovalRequest](self.pending_approvals):
|
|
ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Agent stopped"})
|
|
self.pending_approvals = []
|
|
|
|
self.status = "stopped"
|
|
await ws_manager.emit_status(self.session_id, "stopped")
|
|
|
|
@typechecked
|
|
def branch(self, at_message_id: str) -> "Agent":
|
|
branch_id = uuid4().hex
|
|
branched_messages = deepcopy(self.messages.slice_to(at_message_id))
|
|
child = self.model_copy(deep=True, update={
|
|
"session_id": uuid4().hex,
|
|
"branch_id": branch_id,
|
|
"parent_id": self.session_id,
|
|
"status": "completed",
|
|
"messages": MessageLog(messages=branched_messages),
|
|
"sub_agents": [],
|
|
"pending_approvals": [],
|
|
"task": None,
|
|
"lock": asyncio.Lock(),
|
|
})
|
|
self.sub_branches.append(child)
|
|
return child |