From 371db12265982aa9006ca20ae326ca608acaf9e8 Mon Sep 17 00:00:00 2001 From: haikdc Date: Thu, 2 Apr 2026 10:27:39 -0700 Subject: [PATCH] [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 --- backend/apps/agents/manager/HaikFix/Agent.py | 64 +++++++------- .../HaikFix/run_agent_loop/run_agent_loop.py | 38 +-------- .../manager/HaikFix/shared_structs/Message.py | 53 ------------ .../HaikFix/shared_structs/Message/Message.py | 47 +++++++++++ .../shared_structs/Message/sub_types.py | 83 +++++++++++++++++++ .../HaikFix/shared_structs/MessageLog.py | 51 ++++++++++++ .../HaikFix/shared_structs/PromptChunks.py | 39 --------- 7 files changed, 217 insertions(+), 158 deletions(-) delete mode 100644 backend/apps/agents/manager/HaikFix/shared_structs/Message.py create mode 100644 backend/apps/agents/manager/HaikFix/shared_structs/Message/Message.py create mode 100644 backend/apps/agents/manager/HaikFix/shared_structs/Message/sub_types.py create mode 100644 backend/apps/agents/manager/HaikFix/shared_structs/MessageLog.py delete mode 100644 backend/apps/agents/manager/HaikFix/shared_structs/PromptChunks.py 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