104 lines
2.9 KiB
Python
104 lines
2.9 KiB
Python
import json
|
|
from enum import Enum
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
# Message interface between assistant providers and the assistant client
|
|
# Inspired by https://github.com/anthropics/claude-agent-sdk-python/blob/29c12cd80b256e88f321b2b8f1f5a88445077aa5/src/claude_agent_sdk/types.py
|
|
|
|
|
|
class TextBlock(BaseModel):
|
|
"""Text content block."""
|
|
|
|
text: str
|
|
|
|
|
|
class ThinkingBlock(BaseModel):
|
|
"""Thinking content block."""
|
|
|
|
thinking: str
|
|
signature: str
|
|
|
|
|
|
class ToolUseBlock(BaseModel):
|
|
"""Tool use content block."""
|
|
|
|
id: str
|
|
name: str
|
|
input: dict[str, Any]
|
|
|
|
|
|
class ToolResultBlock(BaseModel):
|
|
"""Tool result content block."""
|
|
|
|
tool_use_id: str
|
|
content: str | list[dict[str, Any]] | None = None
|
|
is_error: bool | None = None
|
|
|
|
|
|
ContentBlock = TextBlock | ThinkingBlock | ToolUseBlock | ToolResultBlock
|
|
|
|
|
|
class Message(BaseModel):
|
|
"""Structured message representation for assistant conversations.
|
|
|
|
Uses standard chat message format with role and content fields.
|
|
Can be extended in the future to support multi-modal content.
|
|
"""
|
|
|
|
role: Literal["user", "assistant", "system"] = Field(description="Role of the message sender")
|
|
content: str | list[ContentBlock] = Field(description="Content of the message")
|
|
|
|
|
|
class EventType(str, Enum):
|
|
MESSAGE = "message"
|
|
STREAM_EVENT = "stream_event"
|
|
DONE = "done"
|
|
ERROR = "error"
|
|
INTERRUPTED = "interrupted"
|
|
PERMISSION_REQUEST = "permission_request"
|
|
|
|
def __str__(self):
|
|
return self.value
|
|
|
|
|
|
class Event(BaseModel):
|
|
"""A common event format parsed from the raw assistant provider output."""
|
|
|
|
type: EventType
|
|
data: dict[str, Any]
|
|
|
|
def to_sse_event(self) -> str:
|
|
"""Convert the event to an SSE event string."""
|
|
return f"event: {self.type}\ndata: {json.dumps(self.data)}\n\n"
|
|
|
|
@classmethod
|
|
def from_error(cls, error: str) -> "Event":
|
|
return cls(type=EventType.ERROR, data={"error": error})
|
|
|
|
@classmethod
|
|
def from_message(cls, message: Message) -> "Event":
|
|
return cls(type=EventType.MESSAGE, data={"message": message.model_dump()})
|
|
|
|
@classmethod
|
|
def from_stream_event(cls, event: dict[str, Any]) -> "Event":
|
|
return cls(type=EventType.STREAM_EVENT, data={"event": event})
|
|
|
|
@classmethod
|
|
def from_result(cls, result: Any, session_id: str) -> "Event":
|
|
return cls(type=EventType.DONE, data={"result": result, "session_id": session_id})
|
|
|
|
@classmethod
|
|
def from_interrupted(cls) -> "Event":
|
|
return cls(type=EventType.INTERRUPTED, data={"message": "Assistant was interrupted"})
|
|
|
|
@classmethod
|
|
def from_permission_request(
|
|
cls, request_id: str, tool_name: str, tool_input: dict[str, Any]
|
|
) -> "Event":
|
|
return cls(
|
|
type=EventType.PERMISSION_REQUEST,
|
|
data={"request_id": request_id, "tool_name": tool_name, "tool_input": tool_input},
|
|
)
|