klaude-code 1.8.0__py3-none-any.whl → 2.0.0__py3-none-any.whl

klaude_code/llm/google/input.py

@@ -4,89 +4,80 @@
 # pyright: reportAttributeAccessIssue=false
 
 import json
-from base64 import b64decode
-from binascii import Error as BinasciiError
 from typing import Any
 
 from google.genai import types
 
-from klaude_code.llm.input_common import AssistantGroup, ToolGroup, UserGroup, merge_reminder_text, parse_message_groups
-from klaude_code.protocol import llm_param, model
+from klaude_code.const import EMPTY_TOOL_OUTPUT_MESSAGE
+from klaude_code.llm.image import parse_data_url
+from klaude_code.llm.input_common import (
+    DeveloperAttachment,
+    attach_developer_messages,
+    merge_reminder_text,
+    split_thinking_parts,
+)
+from klaude_code.protocol import llm_param, message
 
 
 def _data_url_to_blob(url: str) -> types.Blob:
-    header_and_media = url.split(",", 1)
-    if len(header_and_media) != 2:
-        raise ValueError("Invalid data URL for image: missing comma separator")
-    header, base64_data = header_and_media
-    if not header.startswith("data:"):
-        raise ValueError("Invalid data URL for image: missing data: prefix")
-    if ";base64" not in header:
-        raise ValueError("Invalid data URL for image: missing base64 marker")
-
-    media_type = header[5:].split(";", 1)[0]
-    base64_payload = base64_data.strip()
-    if base64_payload == "":
-        raise ValueError("Inline image data is empty")
-
-    try:
-        decoded = b64decode(base64_payload, validate=True)
-    except (BinasciiError, ValueError) as exc:
-        raise ValueError("Inline image data is not valid base64") from exc
-
+    media_type, _, decoded = parse_data_url(url)
     return types.Blob(data=decoded, mime_type=media_type)
 
 
-def _image_part_to_part(image: model.ImageURLPart) -> types.Part:
-    url = image.image_url.url
+def _image_part_to_part(image: message.ImageURLPart) -> types.Part:
+    url = image.url
     if url.startswith("data:"):
         return types.Part(inline_data=_data_url_to_blob(url))
     # Best-effort: Gemini supports file URIs, and may accept public HTTPS URLs.
     return types.Part(file_data=types.FileData(file_uri=url))
 
 
-def _user_group_to_content(group: UserGroup) -> types.Content:
+def _user_message_to_content(msg: message.UserMessage, attachment: DeveloperAttachment) -> types.Content:
     parts: list[types.Part] = []
-    for text in group.text_parts:
-        parts.append(types.Part(text=text + "\n"))
-    for image in group.images:
+    for part in msg.parts:
+        if isinstance(part, message.TextPart):
+            parts.append(types.Part(text=part.text))
+        elif isinstance(part, message.ImageURLPart):
+            parts.append(_image_part_to_part(part))
+    if attachment.text:
+        parts.append(types.Part(text=attachment.text))
+    for image in attachment.images:
         parts.append(_image_part_to_part(image))
     if not parts:
         parts.append(types.Part(text=""))
     return types.Content(role="user", parts=parts)
 
 
-def _tool_groups_to_content(groups: list[ToolGroup], model_name: str | None) -> list[types.Content]:
+def _tool_messages_to_contents(
+    msgs: list[tuple[message.ToolResultMessage, DeveloperAttachment]], model_name: str | None
+) -> list[types.Content]:
     supports_multimodal_function_response = bool(model_name and "gemini-3" in model_name.lower())
 
     response_parts: list[types.Part] = []
     extra_image_contents: list[types.Content] = []
 
-    for group in groups:
+    for msg, attachment in msgs:
         merged_text = merge_reminder_text(
-            group.tool_result.output or "<system-reminder>Tool ran without output or errors</system-reminder>",
-            group.reminder_texts,
+            msg.output_text or EMPTY_TOOL_OUTPUT_MESSAGE,
+            attachment.text,
         )
         has_text = merged_text.strip() != ""
 
-        images = list(group.tool_result.images or []) + list(group.reminder_images)
+        images = [part for part in msg.parts if isinstance(part, message.ImageURLPart)] + attachment.images
         image_parts: list[types.Part] = []
         for image in images:
             try:
                 image_parts.append(_image_part_to_part(image))
             except ValueError:
-                # Skip invalid data URLs
                 continue
 
         has_images = len(image_parts) > 0
         response_value = merged_text if has_text else "(see attached image)" if has_images else ""
-        response_payload = (
-            {"error": response_value} if group.tool_result.status == "error" else {"output": response_value}
-        )
+        response_payload = {"error": response_value} if msg.status != "success" else {"output": response_value}
 
         function_response = types.FunctionResponse(
-            id=group.tool_result.call_id,
-            name=group.tool_result.tool_name or "",
+            id=msg.call_id,
+            name=msg.tool_name,
             response=response_payload,
             parts=image_parts if (has_images and supports_multimodal_function_response) else None,
         )
@@ -104,100 +95,95 @@ def _tool_groups_to_content(groups: list[ToolGroup], model_name: str | None) ->
     return contents
 
 
-def _assistant_group_to_content(group: AssistantGroup, model_name: str | None) -> types.Content | None:
+def _assistant_message_to_content(msg: message.AssistantMessage, model_name: str | None) -> types.Content | None:
     parts: list[types.Part] = []
-
-    degraded_thinking_texts: list[str] = []
+    native_thinking_parts, degraded_thinking_texts = split_thinking_parts(msg, model_name)
+    native_thinking_ids = {id(part) for part in native_thinking_parts}
     pending_thought_text: str | None = None
     pending_thought_signature: str | None = None
 
-    for item in group.reasoning_items:
-        match item:
-            case model.ReasoningTextItem():
-                if not item.content:
-                    continue
-                if model_name is not None and item.model is not None and item.model != model_name:
-                    degraded_thinking_texts.append(item.content)
-                else:
-                    pending_thought_text = item.content
-            case model.ReasoningEncryptedItem():
-                if not (
-                    model_name is not None
-                    and item.model == model_name
-                    and item.encrypted_content
-                    and (item.format or "").startswith("google")
-                    and pending_thought_text
-                ):
-                    continue
-                pending_thought_signature = item.encrypted_content
-                parts.append(
-                    types.Part(
-                        text=pending_thought_text,
-                        thought=True,
-                        thought_signature=pending_thought_signature,
-                    )
-                )
-                pending_thought_text = None
-                pending_thought_signature = None
-
-    if pending_thought_text:
+    def flush_thought() -> None:
+        nonlocal pending_thought_text, pending_thought_signature
+        if pending_thought_text is None and pending_thought_signature is None:
+            return
         parts.append(
             types.Part(
-                text=pending_thought_text,
+                text=pending_thought_text or "",
                 thought=True,
                 thought_signature=pending_thought_signature,
             )
         )
+        pending_thought_text = None
+        pending_thought_signature = None
+
+    for part in msg.parts:
+        if isinstance(part, message.ThinkingTextPart):
+            if id(part) not in native_thinking_ids:
+                continue
+            pending_thought_text = part.text
+            continue
+        if isinstance(part, message.ThinkingSignaturePart):
+            if id(part) not in native_thinking_ids:
+                continue
+            if part.signature and (part.format or "").startswith("google"):
+                pending_thought_signature = part.signature
+            continue
+
+        flush_thought()
+        if isinstance(part, message.TextPart):
+            parts.append(types.Part(text=part.text))
+        elif isinstance(part, message.ToolCallPart):
+            args: dict[str, Any]
+            if part.arguments_json:
+                try:
+                    args = json.loads(part.arguments_json)
+                except json.JSONDecodeError:
+                    args = {"_raw": part.arguments_json}
+            else:
+                args = {}
+            parts.append(types.Part(function_call=types.FunctionCall(id=part.call_id, name=part.tool_name, args=args)))
+
+    flush_thought()
 
     if degraded_thinking_texts:
         parts.insert(0, types.Part(text="<thinking>\n" + "\n".join(degraded_thinking_texts) + "\n</thinking>"))
 
-    if group.text_content:
-        parts.append(types.Part(text=group.text_content))
-
-    for tc in group.tool_calls:
-        args: dict[str, Any]
-        if tc.arguments:
-            try:
-                args = json.loads(tc.arguments)
-            except json.JSONDecodeError:
-                args = {"_raw": tc.arguments}
-        else:
-            args = {}
-        parts.append(types.Part(function_call=types.FunctionCall(id=tc.call_id, name=tc.name, args=args)))
-
     if not parts:
         return None
     return types.Content(role="model", parts=parts)
 
 
 def convert_history_to_contents(
-    history: list[model.ConversationItem],
+    history: list[message.Message],
     model_name: str | None,
 ) -> list[types.Content]:
     contents: list[types.Content] = []
-    pending_tool_groups: list[ToolGroup] = []
-
-    def flush_tool_groups() -> None:
-        nonlocal pending_tool_groups
-        if pending_tool_groups:
-            contents.extend(_tool_groups_to_content(pending_tool_groups, model_name=model_name))
-            pending_tool_groups = []
-
-    for group in parse_message_groups(history):
-        match group:
-            case UserGroup():
-                flush_tool_groups()
-                contents.append(_user_group_to_content(group))
-            case ToolGroup():
-                pending_tool_groups.append(group)
-            case AssistantGroup():
-                flush_tool_groups()
-                content = _assistant_group_to_content(group, model_name=model_name)
+    pending_tool_messages: list[tuple[message.ToolResultMessage, DeveloperAttachment]] = []
+
+    def flush_tool_messages() -> None:
+        nonlocal pending_tool_messages
+        if pending_tool_messages:
+            contents.extend(_tool_messages_to_contents(pending_tool_messages, model_name=model_name))
+            pending_tool_messages = []
+
+    for msg, attachment in attach_developer_messages(history):
+        match msg:
+            case message.ToolResultMessage():
+                pending_tool_messages.append((msg, attachment))
+            case message.UserMessage():
+                flush_tool_messages()
+                contents.append(_user_message_to_content(msg, attachment))
+            case message.AssistantMessage():
+                flush_tool_messages()
+                content = _assistant_message_to_content(msg, model_name=model_name)
                 if content is not None:
                     contents.append(content)
+            case message.SystemMessage():
+                continue
+            case _:
+                continue
 
-    flush_tool_groups()
+    flush_tool_messages()
     return contents
 
 

klaude_code/llm/image.py

@@ -0,0 +1,123 @@
+"""Image processing utilities for LLM responses.
+
+This module provides reusable image handling primitives that can be shared
+across different LLM providers and protocols (OpenAI, Anthropic, etc.).
+"""
+
+from __future__ import annotations
+
+import hashlib
+import mimetypes
+import time
+from base64 import b64decode, b64encode
+from binascii import Error as BinasciiError
+from pathlib import Path
+
+from klaude_code.const import (
+    IMAGE_OUTPUT_MAX_BYTES,
+    TOOL_OUTPUT_TRUNCATION_DIR,
+    ProjectPaths,
+    project_key_from_cwd,
+)
+from klaude_code.protocol import message
+
+IMAGE_EXT_BY_MIME: dict[str, str] = {
+    "image/png": ".png",
+    "image/jpeg": ".jpg",
+    "image/jpg": ".jpg",
+    "image/webp": ".webp",
+    "image/gif": ".gif",
+}
+
+
+def parse_data_url(url: str) -> tuple[str, str, bytes]:
+    """Parse a base64 data URL and return (mime_type, base64_payload, decoded_bytes)."""
+
+    header_and_media = url.split(",", 1)
+    if len(header_and_media) != 2:
+        raise ValueError("Invalid data URL for image: missing comma separator")
+    header, base64_data = header_and_media
+    if not header.startswith("data:"):
+        raise ValueError("Invalid data URL for image: missing data: prefix")
+    if ";base64" not in header:
+        raise ValueError("Invalid data URL for image: missing base64 marker")
+
+    mime_type = header[5:].split(";", 1)[0]
+    base64_payload = base64_data.strip()
+    if base64_payload == "":
+        raise ValueError("Inline image data is empty")
+
+    try:
+        decoded = b64decode(base64_payload, validate=True)
+    except (BinasciiError, ValueError) as exc:
+        raise ValueError("Inline image data is not valid base64") from exc
+
+    return mime_type, base64_payload, decoded
+
+
+def parse_data_url_image(url: str) -> tuple[str, bytes]:
+    """Parse a base64 data URL and return (mime_type, decoded_bytes)."""
+
+    mime_type, _, decoded = parse_data_url(url)
+    return mime_type, decoded
+
+
+def get_assistant_image_output_dir(session_id: str | None) -> Path:
+    """Get the output directory for assistant-generated images."""
+    if session_id:
+        paths = ProjectPaths(project_key=project_key_from_cwd())
+        return paths.images_dir(session_id)
+    return Path(TOOL_OUTPUT_TRUNCATION_DIR) / "images"
+
+
+def save_assistant_image(
+    *, data_url: str, session_id: str | None, response_id: str | None, image_index: int
+) -> message.ImageFilePart:
+    """Decode a data URL image and save it to the session image artifacts directory."""
+
+    mime_type, decoded = parse_data_url_image(data_url)
+
+    if len(decoded) > IMAGE_OUTPUT_MAX_BYTES:
+        decoded_mb = len(decoded) / (1024 * 1024)
+        limit_mb = IMAGE_OUTPUT_MAX_BYTES / (1024 * 1024)
+        raise ValueError(f"Image output size ({decoded_mb:.2f}MB) exceeds limit ({limit_mb:.2f}MB)")
+
+    output_dir = get_assistant_image_output_dir(session_id)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    ext = IMAGE_EXT_BY_MIME.get(mime_type, ".bin")
+    response_part = (response_id or "unknown").replace("/", "_")
+    ts = time.time_ns()
+    file_path = output_dir / f"img-{response_part}-{image_index}-{ts}{ext}"
+    file_path.write_bytes(decoded)
+
+    return message.ImageFilePart(
+        file_path=str(file_path),
+        mime_type=mime_type,
+        byte_size=len(decoded),
+        sha256=hashlib.sha256(decoded).hexdigest(),
+    )
+
+
+def assistant_image_to_data_url(image: message.ImageFilePart) -> str:
+    """Load an assistant image from disk and encode it as a base64 data URL.
+
+    This is primarily used for multi-turn image editing, where providers require
+    sending the previous assistant message (including images) back to the model.
+    """
+
+    file_path = Path(image.file_path)
+    decoded = file_path.read_bytes()
+
+    if len(decoded) > IMAGE_OUTPUT_MAX_BYTES:
+        decoded_mb = len(decoded) / (1024 * 1024)
+        limit_mb = IMAGE_OUTPUT_MAX_BYTES / (1024 * 1024)
+        raise ValueError(f"Assistant image size ({decoded_mb:.2f}MB) exceeds limit ({limit_mb:.2f}MB)")
+
+    mime_type = image.mime_type
+    if not mime_type:
+        guessed, _ = mimetypes.guess_type(str(file_path))
+        mime_type = guessed or "application/octet-stream"
+
+    encoded = b64encode(decoded).decode("ascii")
+    return f"data:{mime_type};base64,{encoded}"