mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-13 21:27:41 +02:00
[Haik]: ok so new Agent class shld be done. In next push ima swap around the folder struct, then im gonna move onto other stuff bc the current agent manager shit seems like bs that can just be removed entirely
This commit is contained in:
haikdc
@@ -1,5 +1,6 @@
|
||||
# TODO: NON HAIK DEPS: ws_manager
|
||||
|
||||
from copy import deepcopy
|
||||
from uuid import uuid4
|
||||
import asyncio
|
||||
import os
|
||||
@@ -11,63 +12,45 @@ from typeguard import typechecked
|
||||
|
||||
from backend.apps.agents.manager.ws_manager import ws_manager
|
||||
from backend.apps.agents.manager.HaikFix.run_agent_loop.run_agent_loop import run_agent_loop
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.Message import Message
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.PromptChunks import ImageChunk
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.Message.Message import Message
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.ApprovalRequest import ApprovalRequest
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.MessageLog import MessageLog
|
||||
|
||||
os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")
|
||||
|
||||
class ContextPath(BaseModel):
|
||||
path: str
|
||||
type: Literal["file", "directory"]
|
||||
|
||||
class Skill(BaseModel):
|
||||
name: str
|
||||
content: str
|
||||
|
||||
|
||||
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"
|
||||
children: List["Agent"] = Field(default_factory=list)
|
||||
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,
|
||||
prompt: str,
|
||||
images: Optional[List[ImageChunk]] = None,
|
||||
) -> None:
|
||||
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
|
||||
|
||||
user_msg = Message(
|
||||
role="user",
|
||||
content=prompt,
|
||||
branch_id=self.branch_id,
|
||||
parent_id=self.parent_id,
|
||||
images=images,
|
||||
)
|
||||
|
||||
await ws_manager.emit_message(self.session_id, user_msg)
|
||||
await ws_manager.emit_message(self.session_id, msg)
|
||||
self.status = "running"
|
||||
await ws_manager.emit_status(self.session_id, "running", self)
|
||||
await ws_manager.emit_status(self.session_id, "running")
|
||||
self.messages.append(msg)
|
||||
|
||||
self.task = asyncio.create_task(run_agent_loop(
|
||||
prompt=prompt,
|
||||
images=images,
|
||||
msg=msg,
|
||||
options=self.config,
|
||||
branch_id=self.branch_id,
|
||||
parent_id=self.parent_id,
|
||||
@@ -75,7 +58,7 @@ class Agent(BaseModel):
|
||||
|
||||
@typechecked
|
||||
async def stop_agent(self):
|
||||
for child in self.children:
|
||||
for child in self.sub_agents:
|
||||
await child.stop_agent()
|
||||
|
||||
if self.task and not self.task.done():
|
||||
@@ -85,10 +68,27 @@ class Agent(BaseModel):
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# if self.session:
|
||||
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", self.session)
|
||||
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
|
||||
@@ -8,49 +8,19 @@ from claude_agent_sdk.types import StreamEvent
|
||||
from backend.apps.agents.manager.HaikFix.run_agent_loop.helpers.handle_stream_event import handle_stream_event
|
||||
from backend.apps.agents.manager.HaikFix.run_agent_loop.helpers.handle_assistant_message import handle_assistant_message
|
||||
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.PromptChunks import (
|
||||
ImageChunk, ImageChunkDict, TextChunk, TextChunkDict
|
||||
)
|
||||
# from backend.apps.agents.manager.HaikFix.shared_structs.Message.sub_types import ImageChunk, TextChunk
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.Message.Message import Message, PromptMsgDict
|
||||
from typing import List, Dict, Literal, Union, Optional
|
||||
|
||||
@typechecked
|
||||
def build_image_prompt_content(prompt: str, images: List[ImageChunk]) -> List[TextChunk | ImageChunk]:
|
||||
content: List[Union[ImageChunkDict, TextChunkDict]] = [TextChunk(text=prompt).to_dict()]
|
||||
for img in images:
|
||||
content.append(img.to_dict())
|
||||
return content
|
||||
|
||||
PromptMsgDict = Dict[
|
||||
Literal["type", "message"],
|
||||
Dict[
|
||||
Literal["role", "content"],
|
||||
List[
|
||||
Union[ImageChunkDict, TextChunkDict]
|
||||
]
|
||||
]
|
||||
]
|
||||
|
||||
@typechecked
|
||||
def build_prompt_msg(prompt: str, images: Optional[List[ImageChunk]]) -> PromptMsgDict:
|
||||
content = build_image_prompt_content(prompt, images)
|
||||
return {
|
||||
"type": "user",
|
||||
"message": {
|
||||
"role": "user",
|
||||
"content": content
|
||||
}
|
||||
}
|
||||
|
||||
@typechecked
|
||||
async def run_agent_loop(
|
||||
prompt: str,
|
||||
images: Optional[List[ImageChunk]] = None,
|
||||
msg: Message,
|
||||
options: ClaudeAgentOptions | None = None,
|
||||
branch_id: str | None = None,
|
||||
):
|
||||
"""Run the Claude Agent SDK query loop for a session."""
|
||||
|
||||
prompt_msg = build_prompt_msg(prompt=prompt, images=images)
|
||||
prompt_msg: PromptMsgDict = msg.to_prompt()
|
||||
|
||||
async def prompt_stream():
|
||||
yield prompt_msg
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
from typing import Optional, Literal, Union, List
|
||||
from pydantic import BaseModel, Field
|
||||
from datetime import datetime
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
########################################################
|
||||
# Message Content Types
|
||||
########################################################
|
||||
|
||||
class ToolCallContent(BaseModel):
|
||||
id: str
|
||||
tool: str
|
||||
input: dict
|
||||
|
||||
class ToolResultContent(BaseModel):
|
||||
text: str
|
||||
tool_name: Optional[str] = None
|
||||
elapsed_ms: Optional[float] = None
|
||||
sub_session_id: Optional[str] = None
|
||||
|
||||
MessageContent = Union[str, ToolCallContent, ToolResultContent]
|
||||
|
||||
|
||||
########################################################
|
||||
# Additional Message Types
|
||||
########################################################
|
||||
|
||||
class ContextPath(BaseModel):
|
||||
path: str
|
||||
type: Literal["file", "directory"]
|
||||
|
||||
class SkillMeta(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
# NOTE: content is omitted in backend to save space
|
||||
|
||||
class ImageMeta(BaseModel):
|
||||
data: str # base64-encoded
|
||||
media_type: str = "image/png"
|
||||
|
||||
class Message(BaseModel):
|
||||
id: str = Field(default_factory=lambda: uuid4().hex)
|
||||
role: Literal["user", "assistant", "tool_call", "tool_result", "system"]
|
||||
content: MessageContent
|
||||
timestamp: datetime = Field(default_factory=datetime.now)
|
||||
branch_id: str = "main"
|
||||
parent_id: Optional[str] = None
|
||||
context_paths: Optional[List[ContextPath]] = None
|
||||
attached_skills: Optional[List[SkillMeta]] = None
|
||||
forced_tools: Optional[List[str]] = None
|
||||
images: Optional[List[ImageMeta]] = None
|
||||
hidden: bool = False
|
||||
@@ -0,0 +1,47 @@
|
||||
from typing import Optional, Literal, Union, List, Dict
|
||||
from pydantic import BaseModel, Field
|
||||
from datetime import datetime
|
||||
from uuid import uuid4
|
||||
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.Message.sub_types import (
|
||||
MessageContent, ContextPath, SkillMeta, ImageChunkDict, TextChunkDict, TextChunk, ImageChunk
|
||||
)
|
||||
from typeguard import typechecked
|
||||
|
||||
PromptMsgDict = Dict[
|
||||
Literal["type", "message"],
|
||||
Dict[
|
||||
Literal["role", "content"],
|
||||
List[
|
||||
Union[ImageChunkDict, TextChunkDict]
|
||||
]
|
||||
]
|
||||
]
|
||||
|
||||
class Message(BaseModel):
|
||||
id: str = Field(default_factory=lambda: uuid4().hex)
|
||||
role: Literal["user", "assistant", "tool_call", "tool_result", "system"]
|
||||
content: MessageContent
|
||||
timestamp: datetime = Field(default_factory=datetime.now)
|
||||
branch_id: str = "main"
|
||||
parent_id: Optional[str] = None
|
||||
context_paths: Optional[List[ContextPath]] = None
|
||||
attached_skills: Optional[List[SkillMeta]] = None
|
||||
forced_tools: Optional[List[str]] = None
|
||||
images: Optional[List[ImageChunk]] = None
|
||||
hidden: bool = False
|
||||
|
||||
@typechecked
|
||||
def to_prompt(self) -> PromptMsgDict:
|
||||
assert isinstance(self.content, str), "Content must be a string"
|
||||
prompt: str = self.content
|
||||
content: List[Union[ImageChunkDict, TextChunkDict]] = [TextChunk(text=prompt).to_dict()]
|
||||
for img in self.images:
|
||||
content.append(img.to_dict())
|
||||
return {
|
||||
"type": "user",
|
||||
"message": {
|
||||
"role": "user",
|
||||
"content": content
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
from typing import Optional, Literal, Union, Dict
|
||||
from pydantic import BaseModel
|
||||
|
||||
from typeguard import typechecked
|
||||
|
||||
########################################################
|
||||
# Promp Chunk Types
|
||||
########################################################
|
||||
|
||||
TextChunkDict = Dict[Literal["type", "text"], str]
|
||||
class TextChunk(BaseModel):
|
||||
type: str = "text"
|
||||
text: str
|
||||
|
||||
@typechecked
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
|
||||
@typechecked
|
||||
def to_dict(self) -> TextChunkDict:
|
||||
return {
|
||||
"type": self.type,
|
||||
"text": self.text,
|
||||
}
|
||||
|
||||
|
||||
ImageChunkDict = Dict[Literal["type", "data", "media_type"], str]
|
||||
class ImageChunk(BaseModel):
|
||||
type: str = "base64"
|
||||
data: str
|
||||
media_type: str
|
||||
|
||||
@typechecked
|
||||
def __init__(self, data: str, media_type: str) -> None:
|
||||
self.data = data
|
||||
self.media_type = media_type
|
||||
|
||||
@typechecked
|
||||
def to_dict(self) -> ImageChunkDict:
|
||||
return {
|
||||
"type": self.type,
|
||||
"data": self.data,
|
||||
"media_type": self.media_type,
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
########################################################
|
||||
# Message Content Types
|
||||
########################################################
|
||||
|
||||
class ToolCallContent(BaseModel):
|
||||
id: str
|
||||
tool: str
|
||||
input: dict
|
||||
|
||||
class ToolResultContent(BaseModel):
|
||||
text: str
|
||||
tool_name: Optional[str] = None
|
||||
elapsed_ms: Optional[float] = None
|
||||
sub_session_id: Optional[str] = None
|
||||
|
||||
MessageContent = Union[str, ToolCallContent, ToolResultContent]
|
||||
# NOTE: is a string, tool call, or tool result
|
||||
|
||||
|
||||
########################################################
|
||||
# Additional Message Types
|
||||
########################################################
|
||||
|
||||
class ContextPath(BaseModel):
|
||||
path: str
|
||||
type: Literal["file", "directory"]
|
||||
|
||||
class SkillMeta(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
# NOTE: content is omitted in backend to save space
|
||||
|
||||
class ImageMeta(BaseModel):
|
||||
data: str # base64-encoded
|
||||
media_type: str = "image/png"
|
||||
@@ -0,0 +1,51 @@
|
||||
from typeguard import typechecked
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List
|
||||
|
||||
from backend.apps.agents.manager.HaikFix.shared_structs.Message import Message
|
||||
|
||||
|
||||
class MessageLog(BaseModel):
|
||||
"""Ordered log of conversation messages for an Agent.
|
||||
|
||||
Backed by a plain list — optimized for sequential append (the hot path)
|
||||
while still supporting ID-based lookup and slicing for branching.
|
||||
"""
|
||||
|
||||
messages: List[Message] = Field(default_factory=list)
|
||||
|
||||
@typechecked
|
||||
def append(self, msg: Message) -> None:
|
||||
"""Append a message to the end of the log. O(1) amortized."""
|
||||
self.messages.append(msg)
|
||||
|
||||
@typechecked
|
||||
def get(self, message_id: str) -> Message | None:
|
||||
"""Look up a single message by ID.
|
||||
|
||||
Returns None if no message with the given ID exists.
|
||||
O(n) scan — fine for infrequent lookups on conversation-sized data.
|
||||
"""
|
||||
return next((m for m in self.messages if m.id == message_id), None)
|
||||
|
||||
@typechecked
|
||||
def slice_to(self, message_id: str) -> List[Message]:
|
||||
"""Return all messages from the start up to and including the given ID.
|
||||
|
||||
Used by Agent.branch() to snapshot the conversation history at a
|
||||
specific fork point. Raises ValueError if the ID isn't found.
|
||||
"""
|
||||
for i, m in enumerate[Message](self.messages):
|
||||
if m.id == message_id:
|
||||
return self._messages[:i + 1]
|
||||
raise ValueError(f"Message {message_id} not found")
|
||||
|
||||
@typechecked
|
||||
def all(self) -> List[Message]:
|
||||
"""Return a shallow copy of the full message list."""
|
||||
return list[Message](self.messages)
|
||||
|
||||
@typechecked
|
||||
def __len__(self) -> int:
|
||||
"""Return the number of messages in the log."""
|
||||
return len(self.messages)
|
||||
@@ -1,39 +0,0 @@
|
||||
from typing import Dict, Literal
|
||||
from pydantic import BaseModel
|
||||
from typeguard import typechecked
|
||||
|
||||
TextChunkDict = Dict[Literal["type", "text"], str]
|
||||
class TextChunk(BaseModel):
|
||||
type: str = "text"
|
||||
text: str
|
||||
|
||||
@typechecked
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
|
||||
@typechecked
|
||||
def to_dict(self) -> TextChunkDict:
|
||||
return {
|
||||
"type": self.type,
|
||||
"text": self.text,
|
||||
}
|
||||
|
||||
|
||||
ImageChunkDict = Dict[Literal["type", "data", "media_type"], str]
|
||||
class ImageChunk(BaseModel):
|
||||
type: str = "base64"
|
||||
data: str
|
||||
media_type: str
|
||||
|
||||
@typechecked
|
||||
def __init__(self, data: str, media_type: str) -> None:
|
||||
self.data = data
|
||||
self.media_type = media_type
|
||||
|
||||
@typechecked
|
||||
def to_dict(self) -> ImageChunkDict:
|
||||
return {
|
||||
"type": self.type,
|
||||
"data": self.data,
|
||||
"media_type": self.media_type,
|
||||
}
|
||||
Reference in New Issue
Block a user