mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-11 12:17:45 +02:00
[Haik]: Message and output classes refactored. Handlers for creating and invoking agents done. Now onto their actual tools, then the browser agent
This commit is contained in:
haikdc
+1
-1
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
]
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
+18
-18
@@ -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
|
||||
+56
@@ -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
|
||||
Reference in New Issue
Block a user