klaude-code 1.9.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

klaude_code/protocol/message.py

@@ -0,0 +1,250 @@
+"""
+Models for LLM API input and response items.
+
+History is persisted as HistoryEvent (messages + error/task metadata).
+Streaming-only items are emitted at runtime but never persisted.
+"""
+
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+from klaude_code.protocol.model import (
+    AtPatternParseResult,
+    CommandOutput,
+    StopReason,
+    TaskMetadata,
+    TaskMetadataItem,
+    ToolResultUIExtra,
+    ToolSideEffect,
+    ToolStatus,
+    Usage,
+)
+
+# Stream items
+
+
+class ToolCallStartItem(BaseModel):
+    """Transient streaming signal when LLM starts a tool call.
+
+    This is NOT persisted to conversation history. Used only for
+    real-time UI feedback (e.g., "Calling Bash …").
+    """
+
+    response_id: str | None = None
+    call_id: str
+    name: str
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class AssistantTextDelta(BaseModel):
+    response_id: str | None = None
+    content: str
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class AssistantImageDelta(BaseModel):
+    """Streaming signal indicating an image has been saved to disk."""
+
+    response_id: str | None = None
+    file_path: str
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class ThinkingTextDelta(BaseModel):
+    response_id: str | None = None
+    content: str
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class StreamErrorItem(BaseModel):
+    error: str
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+# Part types
+
+
+class TextPart(BaseModel):
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ImageURLPart(BaseModel):
+    type: Literal["image_url"] = "image_url"
+    url: str
+    id: str | None = None
+
+
+class ImageFilePart(BaseModel):
+    type: Literal["image_file"] = "image_file"
+    file_path: str
+    mime_type: str | None = None
+    byte_size: int | None = None
+    sha256: str | None = None
+
+
+class ThinkingTextPart(BaseModel):
+    type: Literal["thinking_text"] = "thinking_text"
+    id: str | None = None
+    text: str
+    model_id: str | None = None
+
+
+class ThinkingSignaturePart(BaseModel):
+    type: Literal["thinking_signature"] = "thinking_signature"
+    id: str | None = None
+    signature: str
+    model_id: str | None = None
+    format: str | None = None
+
+
+class ToolCallPart(BaseModel):
+    type: Literal["tool_call"] = "tool_call"
+    call_id: str
+    id: str | None = None  # Responses API: fc_xxx, different from call_id
+    tool_name: str
+    arguments_json: str
+
+
+Part = Annotated[
+    TextPart | ImageURLPart | ImageFilePart | ThinkingTextPart | ThinkingSignaturePart | ToolCallPart,
+    Field(discriminator="type"),
+]
+
+
+def _empty_parts() -> list[Part]:
+    return []
+
+
+# Message types
+
+
+class MessageBase(BaseModel):
+    id: str | None = None
+    created_at: datetime = Field(default_factory=datetime.now)
+    response_id: str | None = None
+
+
+class SystemMessage(MessageBase):
+    role: Literal["system"] = "system"
+    parts: list[TextPart]
+
+
+class DeveloperMessage(MessageBase):
+    role: Literal["developer"] = "developer"
+    parts: list[Part]
+
+    # Special fields for reminders UI
+    memory_paths: list[str] | None = None
+    memory_mentioned: dict[str, list[str]] | None = None  # memory_path -> list of @ patterns mentioned in it
+    external_file_changes: list[str] | None = None
+    todo_use: bool | None = None
+    at_files: list[AtPatternParseResult] | None = None
+    command_output: CommandOutput | None = None
+    user_image_count: int | None = None
+    skill_name: str | None = None  # Skill name activated via $skill syntax
+
+
+class UserMessage(MessageBase):
+    role: Literal["user"] = "user"
+    parts: list[Part]
+
+
+class AssistantMessage(MessageBase):
+    role: Literal["assistant"] = "assistant"
+    parts: list[Part]
+    usage: Usage | None = None
+    stop_reason: StopReason | None = None
+
+
+class ToolResultMessage(MessageBase):
+    role: Literal["tool"] = "tool"
+    call_id: str = ""
+    tool_name: str = ""
+    status: ToolStatus
+    output_text: str
+    parts: list[Part] = Field(default_factory=_empty_parts)
+    ui_extra: ToolResultUIExtra | None = None
+    side_effects: list[ToolSideEffect] | None = None
+    task_metadata: TaskMetadata | None = None  # Sub-agent task metadata for propagation to main agent
+
+    @field_validator("parts")
+    @classmethod
+    def _ensure_non_text_parts(cls, parts: list[Part]) -> list[Part]:
+        if any(isinstance(part, TextPart) for part in parts):
+            raise ValueError("ToolResultMessage.parts must not include text parts")
+        return parts
+
+
+Message = SystemMessage | DeveloperMessage | UserMessage | AssistantMessage | ToolResultMessage
+
+HistoryEvent = Message | StreamErrorItem | TaskMetadataItem
+
+StreamItem = AssistantTextDelta | AssistantImageDelta | ThinkingTextDelta | ToolCallStartItem
+
+LLMStreamItem = HistoryEvent | StreamItem
+
+
+# User input
+
+
+class UserInputPayload(BaseModel):
+    """Structured payload for user input containing text and optional images.
+
+    This is the unified data structure for user input across the entire
+    UI -> CLI -> Executor -> Agent -> Task chain.
+    """
+
+    text: str
+    images: list[ImageURLPart] | None = None
+
+
+# Helper functions
+
+
+def text_parts_from_str(text: str | None) -> list[Part]:
+    if not text:
+        return []
+    return [TextPart(text=text)]
+
+
+def parts_from_text_and_images(text: str | None, images: list[ImageURLPart] | None) -> list[Part]:
+    parts: list[Part] = []
+    if text:
+        parts.append(TextPart(text=text))
+    if images:
+        parts.extend(images)
+    return parts
+
+
+def join_text_parts(parts: Sequence[Part]) -> str:
+    return "".join(part.text for part in parts if isinstance(part, TextPart))
+
+
+def format_saved_images(images: Sequence[ImageFilePart], text: str = "") -> str:
+    """Format saved image paths with optional text prefix.
+
+    Args:
+        images: List of ImageFilePart with file paths to format.
+        text: Optional text content to prepend.
+
+    Returns:
+        Formatted string with image paths. Single image uses "Saved image at xxx",
+        multiple images use list format.
+    """
+    valid_paths = [img.file_path for img in images if img.file_path]
+    if not valid_paths:
+        return text
+
+    if len(valid_paths) == 1:
+        image_text = f"Saved image at {valid_paths[0]}"
+    else:
+        image_lines = "\n".join(f"- {path}" for path in valid_paths)
+        image_text = f"Saved images:\n{image_lines}"
+
+    if text.strip():
+        return f"{text}\n\n{image_text}"
+    return image_text

klaude_code/protocol/model.py

@@ -2,13 +2,15 @@ from datetime import datetime
 from enum import Enum
 from typing import Annotated, Any, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, computed_field
+from pydantic import BaseModel, Field, computed_field
 
-from klaude_code import const
+from klaude_code.const import DEFAULT_MAX_TOKENS
 from klaude_code.protocol.commands import CommandName
 from klaude_code.protocol.tools import SubAgentType
 
 RoleType = Literal["system", "developer", "user", "assistant", "tool"]
+StopReason = Literal["stop", "length", "tool_use", "error", "aborted"]
+ToolStatus = Literal["success", "error", "aborted"]
 TodoStatusType = Literal["pending", "in_progress", "completed"]
 
 
@@ -18,6 +20,7 @@ class Usage(BaseModel):
     cached_tokens: int = 0
     reasoning_tokens: int = 0
     output_tokens: int = 0
+    image_tokens: int = 0  # Image generation tokens
 
     # Context window tracking
     context_size: int | None = None  # Peak total_tokens seen (for context usage display)
@@ -31,7 +34,13 @@ class Usage(BaseModel):
     input_cost: float | None = None  # Cost for non-cached input tokens
     output_cost: float | None = None  # Cost for output tokens (including reasoning)
     cache_read_cost: float | None = None  # Cost for cached tokens
+    image_cost: float | None = None  # Cost for image generation tokens
     currency: str = "USD"  # Currency for cost display (USD or CNY)
+    response_id: str | None = None
+    model_name: str = ""
+    provider: str | None = None  # OpenRouter's provider name
+    task_duration_s: float | None = None
+    created_at: datetime = Field(default_factory=datetime.now)
 
     @computed_field
     @property
@@ -42,8 +51,8 @@ class Usage(BaseModel):
     @computed_field
     @property
     def total_cost(self) -> float | None:
-        """Total cost computed from input + output + cache_read costs."""
-        costs = [self.input_cost, self.output_cost, self.cache_read_cost]
+        """Total cost computed from input + output + cache_read + image costs."""
+        costs = [self.input_cost, self.output_cost, self.cache_read_cost, self.image_cost]
         non_none = [c for c in costs if c is not None]
         return sum(non_none) if non_none else None
 
@@ -55,18 +64,94 @@ class Usage(BaseModel):
             return None
         if self.context_size is None:
             return None
-        effective_limit = self.context_limit - (self.max_tokens or const.DEFAULT_MAX_TOKENS)
+        effective_limit = self.context_limit - (self.max_tokens or DEFAULT_MAX_TOKENS)
         if effective_limit <= 0:
             return None
         return (self.context_size / effective_limit) * 100
 
 
-class TodoItem(BaseModel):
-    model_config = ConfigDict(populate_by_name=True)
+class TaskMetadata(BaseModel):
+    """Base metadata for a task execution (used by both main and sub-agents)."""
+
+    usage: Usage | None = None
+    model_name: str = ""
+    provider: str | None = None
+    task_duration_s: float | None = None
+    turn_count: int = 0
+
+    @staticmethod
+    def merge_usage(dst: Usage, src: Usage) -> None:
+        """Merge src usage into dst usage (in-place).
+
+        Accumulates token counts and cost components. Does not handle
+        special fields like throughput_tps, first_token_latency_ms,
+        context_size, or context_limit - those require custom logic.
+        """
+        dst.input_tokens += src.input_tokens
+        dst.cached_tokens += src.cached_tokens
+        dst.reasoning_tokens += src.reasoning_tokens
+        dst.output_tokens += src.output_tokens
+        dst.image_tokens += src.image_tokens
+
+        if src.input_cost is not None:
+            dst.input_cost = (dst.input_cost or 0.0) + src.input_cost
+        if src.output_cost is not None:
+            dst.output_cost = (dst.output_cost or 0.0) + src.output_cost
+        if src.cache_read_cost is not None:
+            dst.cache_read_cost = (dst.cache_read_cost or 0.0) + src.cache_read_cost
+        if src.image_cost is not None:
+            dst.image_cost = (dst.image_cost or 0.0) + src.image_cost
+
+    @staticmethod
+    def aggregate_by_model(metadata_list: list["TaskMetadata"]) -> list["TaskMetadata"]:
+        """Aggregate multiple TaskMetadata by (model_name, provider).
+
+        Returns a list sorted by total_cost descending.
+
+        Note: total_tokens and total_cost are now computed fields,
+        so we only accumulate the primary state fields here.
+        """
+        aggregated: dict[tuple[str, str | None], TaskMetadata] = {}
+
+        for meta in metadata_list:
+            if not meta.usage:
+                continue
+
+            key = (meta.model_name, meta.provider)
+            usage = meta.usage
+
+            if key not in aggregated:
+                aggregated[key] = TaskMetadata(
+                    model_name=meta.model_name,
+                    provider=meta.provider,
+                    usage=Usage(currency=usage.currency),
+                )
+
+            agg = aggregated[key]
+            if agg.usage is None:
+                continue
+
+            TaskMetadata.merge_usage(agg.usage, usage)
+
+        # Sort by total_cost descending
+        return sorted(
+            aggregated.values(),
+            key=lambda m: m.usage.total_cost if m.usage and m.usage.total_cost else 0.0,
+            reverse=True,
+        )
+
 
+class TaskMetadataItem(BaseModel):
+    """Aggregated metadata for a complete task, stored in conversation history."""
+
+    main_agent: TaskMetadata = Field(default_factory=TaskMetadata)  # Main agent metadata
+    sub_agent_task_metadata: list[TaskMetadata] = Field(default_factory=lambda: list[TaskMetadata]())
+    created_at: datetime = Field(default_factory=datetime.now)
+
+
+class TodoItem(BaseModel):
     content: str
     status: TodoStatusType
-    active_form: str = Field(default="", alias="activeForm")
 
 
 class FileStatus(BaseModel):
@@ -193,7 +278,6 @@ class AtPatternParseResult(BaseModel):
     result: str
     tool_args: str
     operation: Literal["Read", "List"]
-    images: list["ImageURLPart"] | None = None
     mentioned_in: str | None = None  # Parent file that referenced this file
 
 
@@ -207,283 +291,9 @@ class SubAgentState(BaseModel):
     sub_agent_type: SubAgentType
     sub_agent_desc: str
     sub_agent_prompt: str
+    resume: str | None = None
     output_schema: dict[str, Any] | None = None
-
-
-"""
-Models for LLM API input and response items.
-
-A typical sequence of response items is:
-- [StartItem]
-- [ReasoningTextItem | ReasoningEncryptedItem]
-- [AssistantMessageDelta] × n
-- [AssistantMessageItem]
-- [ToolCallItem] × n
-- [ResponseMetadataItem]
-- Done
-
-A conversation history input contains:
-- [UserMessageItem]
-- [ReasoningTextItem | ReasoningEncryptedItem]
-- [AssistantMessageItem]
-- [ToolCallItem]
-- [ToolResultItem]
-- [InterruptItem]
-- [DeveloperMessageItem]
-
-When adding a new item, please also modify the following:
-- session/codec.py (ConversationItem registry derived from ConversationItem union)
-"""
-
-
-class StartItem(BaseModel):
-    response_id: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class InterruptItem(BaseModel):
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class SystemMessageItem(BaseModel):
-    id: str | None = None
-    role: RoleType = "system"
-    content: str | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class DeveloperMessageItem(BaseModel):
-    id: str | None = None
-    role: RoleType = "developer"
-    content: str | None = None  # For LLM input
-    images: list["ImageURLPart"] | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-    # Special fields for reminders UI
-    memory_paths: list[str] | None = None
-    memory_mentioned: dict[str, list[str]] | None = None  # memory_path -> list of @ patterns mentioned in it
-    external_file_changes: list[str] | None = None
-    todo_use: bool | None = None
-    at_files: list[AtPatternParseResult] | None = None
-    command_output: CommandOutput | None = None
-    user_image_count: int | None = None
-    skill_name: str | None = None  # Skill name activated via $skill syntax
-
-
-class ImageURLPart(BaseModel):
-    class ImageURL(BaseModel):
-        url: str
-        id: str | None = None
-
-    image_url: ImageURL
-
-
-class UserInputPayload(BaseModel):
-    """Structured payload for user input containing text and optional images.
-
-    This is the unified data structure for user input across the entire
-    UI -> CLI -> Executor -> Agent -> Task chain.
-    """
-
-    text: str
-    images: list["ImageURLPart"] | None = None
-
-
-class UserMessageItem(BaseModel):
-    id: str | None = None
-    role: RoleType = "user"
-    content: str | None = None
-    images: list[ImageURLPart] | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class AssistantMessageItem(BaseModel):
-    id: str | None = None
-    role: RoleType = "assistant"
-    content: str | None = None
-    response_id: str | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ReasoningTextItem(BaseModel):
-    id: str | None = None
-    response_id: str | None = None
-    content: str
-    model: str | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ReasoningEncryptedItem(BaseModel):
-    id: str | None = None
-    response_id: str | None = None
-    encrypted_content: str  # OpenAI encrypted content or Anthropic thinking signature
-    format: str | None = None
-    model: str | None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ToolCallStartItem(BaseModel):
-    """Transient streaming signal when LLM starts a tool call.
-
-    This is NOT persisted to conversation history. Used only for
-    real-time UI feedback (e.g., "Calling Bash ...").
-    """
-
-    response_id: str | None = None
-    call_id: str
-    name: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ToolCallItem(BaseModel):
-    id: str | None = None
-    response_id: str | None = None
-    call_id: str
-    name: str
-    arguments: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ToolResultItem(BaseModel):
-    call_id: str = ""  # This field will auto set by tool registry's run_tool
-    output: str | None = None
-    status: Literal["success", "error"]
-    tool_name: str | None = None  # This field will auto set by tool registry's run_tool
-    ui_extra: ToolResultUIExtra | None = None  # Extra data for UI display, e.g. diff render
-    images: list[ImageURLPart] | None = None
-    side_effects: list[ToolSideEffect] | None = None
-    task_metadata: "TaskMetadata | None" = None  # Sub-agent task metadata for propagation to main agent
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class AssistantMessageDelta(BaseModel):
-    response_id: str | None = None
-    content: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ReasoningTextDelta(BaseModel):
-    response_id: str | None = None
-    content: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class StreamErrorItem(BaseModel):
-    error: str
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class ResponseMetadataItem(BaseModel):
-    """Metadata for a single LLM response (turn-level)."""
-
-    response_id: str | None = None
-    usage: Usage | None = None
-    model_name: str = ""
-    provider: str | None = None  # OpenRouter's provider name
-    task_duration_s: float | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-class TaskMetadata(BaseModel):
-    """Base metadata for a task execution (used by both main and sub-agents)."""
-
-    usage: Usage | None = None
-    model_name: str = ""
-    provider: str | None = None
-    task_duration_s: float | None = None
-    turn_count: int = 0
-
-    @staticmethod
-    def merge_usage(dst: Usage, src: Usage) -> None:
-        """Merge src usage into dst usage (in-place).
-
-        Accumulates token counts and cost components. Does not handle
-        special fields like throughput_tps, first_token_latency_ms,
-        context_size, or context_limit - those require custom logic.
-        """
-        dst.input_tokens += src.input_tokens
-        dst.cached_tokens += src.cached_tokens
-        dst.reasoning_tokens += src.reasoning_tokens
-        dst.output_tokens += src.output_tokens
-
-        if src.input_cost is not None:
-            dst.input_cost = (dst.input_cost or 0.0) + src.input_cost
-        if src.output_cost is not None:
-            dst.output_cost = (dst.output_cost or 0.0) + src.output_cost
-        if src.cache_read_cost is not None:
-            dst.cache_read_cost = (dst.cache_read_cost or 0.0) + src.cache_read_cost
-
-    @staticmethod
-    def aggregate_by_model(metadata_list: list["TaskMetadata"]) -> list["TaskMetadata"]:
-        """Aggregate multiple TaskMetadata by (model_name, provider).
-
-        Returns a list sorted by total_cost descending.
-
-        Note: total_tokens and total_cost are now computed fields,
-        so we only accumulate the primary state fields here.
-        """
-        aggregated: dict[tuple[str, str | None], TaskMetadata] = {}
-
-        for meta in metadata_list:
-            if not meta.usage:
-                continue
-
-            key = (meta.model_name, meta.provider)
-            usage = meta.usage
-
-            if key not in aggregated:
-                aggregated[key] = TaskMetadata(
-                    model_name=meta.model_name,
-                    provider=meta.provider,
-                    usage=Usage(currency=usage.currency),
-                )
-
-            agg = aggregated[key]
-            if agg.usage is None:
-                continue
-
-            TaskMetadata.merge_usage(agg.usage, usage)
-
-        # Sort by total_cost descending
-        return sorted(
-            aggregated.values(),
-            key=lambda m: m.usage.total_cost if m.usage and m.usage.total_cost else 0.0,
-            reverse=True,
-        )
-
-
-class TaskMetadataItem(BaseModel):
-    """Aggregated metadata for a complete task, stored in conversation history."""
-
-    main_agent: TaskMetadata = Field(default_factory=TaskMetadata)  # Main agent metadata
-    sub_agent_task_metadata: list[TaskMetadata] = Field(default_factory=lambda: list[TaskMetadata]())
-    created_at: datetime = Field(default_factory=datetime.now)
-
-
-MessageItem = (
-    UserMessageItem
-    | AssistantMessageItem
-    | SystemMessageItem
-    | DeveloperMessageItem
-    | ReasoningTextItem
-    | ReasoningEncryptedItem
-    | ToolCallItem
-    | ToolResultItem
-)
-
-
-StreamItem = AssistantMessageDelta | ReasoningTextDelta
-
-ConversationItem = (
-    StartItem
-    | InterruptItem
-    | StreamErrorItem
-    | StreamItem
-    | MessageItem
-    | ResponseMetadataItem
-    | TaskMetadataItem
-    | ToolCallStartItem
-)
+    generation: dict[str, Any] | None = None
 
 
 def todo_list_str(todos: list[TodoItem]) -> str: