diff --git a/backend/apps/agents/manager/HaikFix/Agent.py b/backend/apps/agents/manager/HaikFix/Agent.py index ce0d39e6..ac8ef645 100644 --- a/backend/apps/agents/manager/HaikFix/Agent.py +++ b/backend/apps/agents/manager/HaikFix/Agent.py @@ -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) \ No newline at end of file + 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 \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/run_agent_loop/run_agent_loop.py b/backend/apps/agents/manager/HaikFix/run_agent_loop/run_agent_loop.py index 6300615a..d579a7ff 100644 --- a/backend/apps/agents/manager/HaikFix/run_agent_loop/run_agent_loop.py +++ b/backend/apps/agents/manager/HaikFix/run_agent_loop/run_agent_loop.py @@ -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 diff --git a/backend/apps/agents/manager/HaikFix/shared_structs/Message.py b/backend/apps/agents/manager/HaikFix/shared_structs/Message.py deleted file mode 100644 index 8e07b5cd..00000000 --- a/backend/apps/agents/manager/HaikFix/shared_structs/Message.py +++ /dev/null @@ -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 diff --git a/backend/apps/agents/manager/HaikFix/shared_structs/Message/Message.py b/backend/apps/agents/manager/HaikFix/shared_structs/Message/Message.py new file mode 100644 index 00000000..bb0f9314 --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/shared_structs/Message/Message.py @@ -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 + } + } diff --git a/backend/apps/agents/manager/HaikFix/shared_structs/Message/sub_types.py b/backend/apps/agents/manager/HaikFix/shared_structs/Message/sub_types.py new file mode 100644 index 00000000..ede5e54c --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/shared_structs/Message/sub_types.py @@ -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" \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/shared_structs/MessageLog.py b/backend/apps/agents/manager/HaikFix/shared_structs/MessageLog.py new file mode 100644 index 00000000..572fc469 --- /dev/null +++ b/backend/apps/agents/manager/HaikFix/shared_structs/MessageLog.py @@ -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) \ No newline at end of file diff --git a/backend/apps/agents/manager/HaikFix/shared_structs/PromptChunks.py b/backend/apps/agents/manager/HaikFix/shared_structs/PromptChunks.py deleted file mode 100644 index 07a5ff26..00000000 --- a/backend/apps/agents/manager/HaikFix/shared_structs/PromptChunks.py +++ /dev/null @@ -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, - } \ No newline at end of file