diff --git a/backend/apps/agents/HaikFix/Agent/run_agent_loop/helpers/handle_assistant_message.py b/backend/apps/agents/HaikFix/Agent/run_agent_loop/helpers/handle_assistant_message.py index a6148a95..fa260c81 100644 --- a/backend/apps/agents/HaikFix/Agent/run_agent_loop/helpers/handle_assistant_message.py +++ b/backend/apps/agents/HaikFix/Agent/run_agent_loop/helpers/handle_assistant_message.py @@ -1,5 +1,5 @@ from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import Message -from backend.apps.agents.HaikFix.Agent.shared_structs.Message.sub_types import ToolCallContent +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.agent_outputs import ToolCallContent from claude_agent_sdk.types import TextBlock, ToolUseBlock, AssistantMessage from backend.apps.agents.manager.ws_manager import ws_manager from typeguard import typechecked diff --git a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/Message.py b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/Message.py index 6146a4e4..aa0bc9b9 100644 --- a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/Message.py +++ b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/Message.py @@ -1,47 +1,78 @@ -from typing import Optional, Literal, Union, List, Dict +# Message.py + +from typing import List, Literal, Dict, ClassVar, Annotated, Union from pydantic import BaseModel, Field from datetime import datetime from uuid import uuid4 - -from backend.apps.agents.HaikFix.Agent.shared_structs.Message.sub_types import ( - MessageContent, ContextPath, SkillMeta, ImageChunkDict, TextChunkDict, TextChunk, ImageChunk -) from typeguard import typechecked +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.agent_inputs import ( + PromptBlock, TextPromptBlock, ImagePromptBlock, ImageSource, ContextPath, SkillMeta +) +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.agent_outputs import ( + ToolCallContent, ToolResultContent +) + +class Message(BaseModel): + id: str = Field(default_factory=lambda: uuid4().hex) + timestamp: datetime = Field(default_factory=datetime.now) + branch_id: str = "main" + hidden: bool = False + + PromptMsgDict = Dict[ Literal["type", "message"], Dict[ Literal["role", "content"], - List[ - Union[ImageChunkDict, TextChunkDict] - ] + List[PromptBlock] ] ] -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 +class UserMessage(Message): + role: Literal["user"] = "user" + content: str + images: List[str] = Field(default_factory=list) # base64 strings + image_media_types: List[str] = Field(default_factory=list) + context_paths: List[ContextPath] = Field(default_factory=list) + attached_skills: List[SkillMeta] = Field(default_factory=list) + forced_tools: List[str] = Field(default_factory=list) @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()) + blocks: List[PromptBlock] = [TextPromptBlock(type="text", text=self.content)] + for data, media_type in zip[tuple[str, str]](self.images, self.image_media_types): + blocks.append(ImagePromptBlock( + type="image", + source=ImageSource(type="base64", media_type=media_type, data=data), + )) return { - "type": "user", + "type": self.role, "message": { - "role": "user", - "content": content + "role": self.role, + "content": blocks } } + +class AssistantMessage(Message): + role: Literal["assistant"] = "assistant" + content: str + + +class ToolCallMessage(Message): + role: Literal["tool_call"] = "tool_call" + content: ToolCallContent + + +class ToolResultMessage(Message): + role: Literal["tool_result"] = "tool_result" + content: ToolResultContent + + +class SystemMessage(Message): + role: Literal["system"] = "system" + content: str + +AnyMessage = Annotated[ + Union[UserMessage, AssistantMessage, ToolCallMessage, ToolResultMessage, SystemMessage], + Field(discriminator="role"), +] \ No newline at end of file diff --git a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_inputs.py b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_inputs.py new file mode 100644 index 00000000..5ffb22d2 --- /dev/null +++ b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_inputs.py @@ -0,0 +1,33 @@ +# sub_types.py + +from typing import Literal, Union, TypedDict +from pydantic import BaseModel + + +# --- Prompt content blocks (sent to SDK via to_prompt) --- + +class TextPromptBlock(TypedDict): + type: Literal["text"] + text: str + +class ImageSource(TypedDict): + type: Literal["base64"] + media_type: str + data: str + +class ImagePromptBlock(TypedDict): + type: Literal["image"] + source: "ImageSource" + + +PromptBlock = Union[TextPromptBlock, ImagePromptBlock] + +# --- Attachments (user message only) --- + +class ContextPath(BaseModel): + path: str + type: Literal["file", "directory"] + +class SkillMeta(BaseModel): + id: str + name: str \ No newline at end of file diff --git a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_outputs.py b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_outputs.py new file mode 100644 index 00000000..519773fc --- /dev/null +++ b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/agent_outputs.py @@ -0,0 +1,26 @@ +from typing import Dict, Any, TypedDict, List +from typing_extensions import NotRequired +from pydantic import BaseModel + + +# --- Stored content types --- +class ToolCallContent(BaseModel): + id: str + tool: str + input: Dict[str, Any] + +class ToolResultContent(BaseModel): + tool_use_id: str + text: str + is_error: bool = False + + +# --- Tool response types --- + +class TextContent(TypedDict): + type: str + text: str + +class ToolResponse(TypedDict): + content: List[TextContent] + is_error: NotRequired[bool] \ No newline at end of file diff --git a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/sub_types.py b/backend/apps/agents/HaikFix/Agent/shared_structs/Message/sub_types.py deleted file mode 100644 index 8fc94610..00000000 --- a/backend/apps/agents/HaikFix/Agent/shared_structs/Message/sub_types.py +++ /dev/null @@ -1,79 +0,0 @@ -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 \ No newline at end of file diff --git a/backend/apps/agents/HaikFix/Agent/shared_structs/MessageLog.py b/backend/apps/agents/HaikFix/Agent/shared_structs/MessageLog.py index 572fc469..034db4f9 100644 --- a/backend/apps/agents/HaikFix/Agent/shared_structs/MessageLog.py +++ b/backend/apps/agents/HaikFix/Agent/shared_structs/MessageLog.py @@ -2,7 +2,7 @@ from typeguard import typechecked from pydantic import BaseModel, Field from typing import List -from backend.apps.agents.manager.HaikFix.shared_structs.Message import Message +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import AnyMessage class MessageLog(BaseModel): @@ -12,15 +12,15 @@ class MessageLog(BaseModel): while still supporting ID-based lookup and slicing for branching. """ - messages: List[Message] = Field(default_factory=list) + messages: List[AnyMessage] = Field(default_factory=list) @typechecked - def append(self, msg: Message) -> None: + def append(self, msg: AnyMessage) -> 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: + def get(self, message_id: str) -> AnyMessage | None: """Look up a single message by ID. Returns None if no message with the given ID exists. @@ -29,21 +29,21 @@ class MessageLog(BaseModel): return next((m for m in self.messages if m.id == message_id), None) @typechecked - def slice_to(self, message_id: str) -> List[Message]: + def slice_to(self, message_id: str) -> List[AnyMessage]: """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): + for i, m in enumerate[AnyMessage](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]: + def all(self) -> List[AnyMessage]: """Return a shallow copy of the full message list.""" - return list[Message](self.messages) + return list[AnyMessage](self.messages) @typechecked def __len__(self) -> int: diff --git a/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_create_agent_handler.py b/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_create_agent_handler.py index 9609c342..5f4de248 100644 --- a/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_create_agent_handler.py +++ b/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_create_agent_handler.py @@ -1,20 +1,21 @@ from typing import Dict, Any, TypedDict from typeguard import typechecked -from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import Message -from backend.apps.agents.HaikFix.Agent.Agent import Agent +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.agent_outputs import ToolResponse +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import ( + AnyMessage, UserMessage, AssistantMessage, +) +from backend.apps.agents.HaikFix.Agent.Agent import Agent class CreateAgentInput(TypedDict): task: str @typechecked def make_create_agent_handler(parent: Agent): - async def handler(args: CreateAgentInput) -> Dict[str, Any]: + + async def handler(args: CreateAgentInput) -> ToolResponse: task_message = args["task"] - if not task_message: - return {"content": [{"type": "text", "text": "Error: task is required"}], "is_error": True} - child = Agent( model=parent.model, mode=parent.mode, @@ -23,23 +24,22 @@ def make_create_agent_handler(parent: Agent): parent_id=parent.session_id, ) parent.sub_agents.append(child) - - msg = Message(role="user", content=task_message, branch_id=child.branch_id) + msg = UserMessage(content=task_message, branch_id=child.branch_id) await child.send_message(msg) - - # Wait for the child to finish if child.task: await child.task - - # Extract last assistant response last_response = "No response from sub-agent." - for m in reversed(child.messages.messages): - if m.role == "assistant" and isinstance(m.content, str): + for m in reversed[AnyMessage](child.messages.messages): + if isinstance(m, AssistantMessage): last_response = m.content break - - return {"content": [{"type": "text", "text": ( - f"**Sub-Agent Result** (session: {child.session_id})\n\n{last_response}" - )}]} + return { + "content": [ + { + "type": "text", + "text": (f"**Sub-Agent Result** (session: {child.session_id})\n\n{last_response}") + } + ] + } return handler \ No newline at end of file diff --git a/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_invoke_agent_handler.py b/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_invoke_agent_handler.py new file mode 100644 index 00000000..92a87072 --- /dev/null +++ b/backend/apps/agents/HaikFix/tools/BUILTIN_TOOLKIT/components/custom_handlers/make_invoke_agent_handler.py @@ -0,0 +1,56 @@ +# make_invoke_agent_handler.py + +from typing import Dict, TypedDict +from typeguard import typechecked +from uuid import uuid4 +import asyncio + +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.agent_outputs import ToolResponse +from backend.apps.agents.HaikFix.Agent.shared_structs.Message.Message import ( + AnyMessage, UserMessage, AssistantMessage, +) +from backend.apps.agents.HaikFix.Agent.Agent import Agent + + +class InvokeAgentInput(TypedDict): + session_id: str + message: str + + +@typechecked +def make_invoke_agent_handler(agent_registry: Dict[str, Agent]): + async def handler(args: InvokeAgentInput) -> ToolResponse: + session_id = args["session_id"] + message = args["message"] + + source = agent_registry.get(session_id) + if not source: + return {"content": [{"type": "text", "text": f"Error: session {session_id} not found"}], "is_error": True} + + if len(source.messages) == 0: + return {"content": [{"type": "text", "text": "Error: source agent has no messages"}], "is_error": True} + + fork = source.branch(source.messages.messages[-1].id) + agent_registry[fork.session_id] = fork + + msg = UserMessage(content=message, branch_id=fork.branch_id) + await fork.send_message(msg) + if fork.task: + await fork.task + + last_response = "No response from invoked agent." + for m in reversed[AnyMessage](fork.messages.messages): + if isinstance(m, AssistantMessage): + last_response = m.content + break + + return { + "content": [ + { + "type": "text", + "text": (f"**Invoked Agent Result** (forked session: {fork.session_id})\n\n{last_response}") + } + ] + } + + return handler \ No newline at end of file