klaude-code 1.6.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

klaude_code/llm/google/input.py

@@ -0,0 +1,215 @@
+# pyright: reportReturnType=false
+# pyright: reportArgumentType=false
+# pyright: reportUnknownMemberType=false
+# 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
+
+
+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
+
+    return types.Blob(data=decoded, mime_type=media_type)
+
+
+def _image_part_to_part(image: model.ImageURLPart) -> types.Part:
+    url = image.image_url.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:
+    parts: list[types.Part] = []
+    for text in group.text_parts:
+        parts.append(types.Part(text=text + "\n"))
+    for image in group.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]:
+    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:
+        merged_text = merge_reminder_text(
+            group.tool_result.output or "<system-reminder>Tool ran without output or errors</system-reminder>",
+            group.reminder_texts,
+        )
+        has_text = merged_text.strip() != ""
+
+        images = list(group.tool_result.images or []) + list(group.reminder_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}
+        )
+
+        function_response = types.FunctionResponse(
+            id=group.tool_result.call_id,
+            name=group.tool_result.tool_name or "",
+            response=response_payload,
+            parts=image_parts if (has_images and supports_multimodal_function_response) else None,
+        )
+        response_parts.append(types.Part(function_response=function_response))
+
+        if has_images and not supports_multimodal_function_response:
+            extra_image_contents.append(
+                types.Content(role="user", parts=[types.Part(text="Tool result image:"), *image_parts])
+            )
+
+    contents: list[types.Content] = []
+    if response_parts:
+        contents.append(types.Content(role="user", parts=response_parts))
+    contents.extend(extra_image_contents)
+    return contents
+
+
+def _assistant_group_to_content(group: AssistantGroup, model_name: str | None) -> types.Content | None:
+    parts: list[types.Part] = []
+
+    degraded_thinking_texts: list[str] = []
+    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:
+        parts.append(
+            types.Part(
+                text=pending_thought_text,
+                thought=True,
+                thought_signature=pending_thought_signature,
+            )
+        )
+
+    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],
+    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)
+                if content is not None:
+                    contents.append(content)
+
+    flush_tool_groups()
+    return contents
+
+
+def convert_tool_schema(tools: list[llm_param.ToolSchema] | None) -> list[types.Tool]:
+    if tools is None or len(tools) == 0:
+        return []
+    declarations = [
+        types.FunctionDeclaration(
+            name=tool.name,
+            description=tool.description,
+            parameters_json_schema=tool.parameters,
+        )
+        for tool in tools
+    ]
+    return [types.Tool(function_declarations=declarations)]

klaude_code/llm/registry.py

@@ -1,3 +1,4 @@
+import importlib
 from collections.abc import Callable
 from typing import TYPE_CHECKING, TypeVar
 
@@ -21,15 +22,19 @@ def _load_protocol(protocol: llm_param.LLMClientProtocol) -> None:
 
     # Import only the needed module to trigger @register decorator
     if protocol == llm_param.LLMClientProtocol.ANTHROPIC:
-        from . import anthropic as _
+        importlib.import_module("klaude_code.llm.anthropic")
+    elif protocol == llm_param.LLMClientProtocol.BEDROCK:
+        importlib.import_module("klaude_code.llm.bedrock")
     elif protocol == llm_param.LLMClientProtocol.CODEX:
-        from . import codex as _
+        importlib.import_module("klaude_code.llm.codex")
     elif protocol == llm_param.LLMClientProtocol.OPENAI:
-        from . import openai_compatible as _
+        importlib.import_module("klaude_code.llm.openai_compatible")
     elif protocol == llm_param.LLMClientProtocol.OPENROUTER:
-        from . import openrouter as _
+        importlib.import_module("klaude_code.llm.openrouter")
     elif protocol == llm_param.LLMClientProtocol.RESPONSES:
-        from . import responses as _  # noqa: F401
+        importlib.import_module("klaude_code.llm.responses")
+    elif protocol == llm_param.LLMClientProtocol.GOOGLE:
+        importlib.import_module("klaude_code.llm.google")
 
 
 def register(name: llm_param.LLMClientProtocol) -> Callable[[_T], _T]:

klaude_code/protocol/events.py

@@ -16,6 +16,7 @@ class EndEvent(BaseModel):
 class ErrorEvent(BaseModel):
     error_message: str
     can_retry: bool = False
+    session_id: str | None = None
 
 
 class TaskStartEvent(BaseModel):

klaude_code/protocol/llm_param.py

@@ -12,7 +12,9 @@ class LLMClientProtocol(Enum):
     RESPONSES = "responses"
     OPENROUTER = "openrouter"
     ANTHROPIC = "anthropic"
+    BEDROCK = "bedrock"
     CODEX = "codex"
+    GOOGLE = "google"
 
 
 class ToolSchema(BaseModel):
@@ -91,8 +93,15 @@ class LLMConfigProviderParameter(BaseModel):
     protocol: LLMClientProtocol
     base_url: str | None = None
     api_key: str | None = None
+    # Azure OpenAI
     is_azure: bool = False
     azure_api_version: str | None = None
+    # AWS Bedrock configuration
+    aws_access_key: str | None = None
+    aws_secret_key: str | None = None
+    aws_region: str | None = None
+    aws_session_token: str | None = None
+    aws_profile: str | None = None
 
 
 class LLMConfigModelParameter(BaseModel):

klaude_code/session/export.py

@@ -308,13 +308,17 @@ def _try_render_todo_args(arguments: str, tool_name: str) -> str | None:
         return None
 
 
-def _render_sub_agent_result(content: str) -> str:
+def _render_sub_agent_result(content: str, description: str | None = None) -> str:
     # Try to format as JSON for better readability
     try:
         parsed = json.loads(content)
         formatted = "```json\n" + json.dumps(parsed, ensure_ascii=False, indent=2) + "\n```"
     except (json.JSONDecodeError, TypeError):
         formatted = content
+
+    if description:
+        formatted = f"# {description}\n\n{formatted}"
+
     encoded = _escape_html(formatted)
     return (
         f'<div class="sub-agent-result-container">'
@@ -628,7 +632,15 @@ def _format_tool_call(tool_call: model.ToolCallItem, result: model.ToolResultIte
 
         if result.output and not should_hide_text:
             if is_sub_agent_tool(tool_call.name):
-                items_to_render.append(_render_sub_agent_result(result.output))
+                description = None
+                try:
+                    args = json.loads(tool_call.arguments)
+                    if isinstance(args, dict):
+                        typed_args = cast(dict[str, Any], args)
+                        description = cast(str | None, typed_args.get("description"))
+                except (json.JSONDecodeError, TypeError):
+                    pass
+                items_to_render.append(_render_sub_agent_result(result.output, description))
             else:
                 items_to_render.append(_render_text_block(result.output))
 

klaude_code/session/session.py

@@ -62,6 +62,7 @@ class Session(BaseModel):
     need_todo_not_used_cooldown_counter: int = Field(exclude=True, default=0)
 
     _messages_count_cache: int | None = PrivateAttr(default=None)
+    _user_messages_cache: list[str] | None = PrivateAttr(default=None)
     _store: JsonlSessionStore = PrivateAttr(default_factory=get_default_store)
 
     @property
@@ -78,6 +79,20 @@ class Session(BaseModel):
     def _invalidate_messages_count_cache(self) -> None:
         self._messages_count_cache = None
 
+    @property
+    def user_messages(self) -> list[str]:
+        """All user message contents in this session.
+
+        This is used for session selection UI and search, and is also persisted
+        in meta.json to avoid scanning events.jsonl for every session.
+        """
+
+        if self._user_messages_cache is None:
+            self._user_messages_cache = [
+                it.content for it in self.conversation_history if isinstance(it, model.UserMessageItem) and it.content
+            ]
+        return self._user_messages_cache
+
     @staticmethod
     def _project_key() -> str:
         return _project_key_from_cwd()
@@ -178,6 +193,18 @@ class Session(BaseModel):
         self.conversation_history.extend(items)
         self._invalidate_messages_count_cache()
 
+        new_user_messages = [
+            it.content for it in items if isinstance(it, model.UserMessageItem) and it.content
+        ]
+        if new_user_messages:
+            if self._user_messages_cache is None:
+                # Build from full history once to ensure correctness when resuming older sessions.
+                self._user_messages_cache = [
+                    it.content for it in self.conversation_history if isinstance(it, model.UserMessageItem) and it.content
+                ]
+            else:
+                self._user_messages_cache.extend(new_user_messages)
+
         if self.created_at <= 0:
             self.created_at = time.time()
         self.updated_at = time.time()
@@ -188,6 +215,7 @@ class Session(BaseModel):
             sub_agent_state=self.sub_agent_state,
             file_tracker=self.file_tracker,
             todos=list(self.todos),
+            user_messages=self.user_messages,
             created_at=self.created_at,
             updated_at=self.updated_at,
             messages_count=self.messages_count,
@@ -218,7 +246,9 @@ class Session(BaseModel):
         forked.file_tracker = {k: v.model_copy(deep=True) for k, v in self.file_tracker.items()}
         forked.todos = [todo.model_copy(deep=True) for todo in self.todos]
 
-        history_to_copy = self.conversation_history[:until_index] if until_index is not None else self.conversation_history
+        history_to_copy = (
+            self.conversation_history[:until_index] if until_index is not None else self.conversation_history
+        )
         items = [it.model_copy(deep=True) for it in history_to_copy]
         if items:
             forked.append_history(items)
@@ -309,7 +339,7 @@ class Session(BaseModel):
                 case model.DeveloperMessageItem() as dm:
                     yield events.DeveloperMessageEvent(session_id=self.id, item=dm)
                 case model.StreamErrorItem() as se:
-                    yield events.ErrorEvent(error_message=se.error, can_retry=False)
+                    yield events.ErrorEvent(error_message=se.error, can_retry=False, session_id=self.id)
                 case _:
                     continue
             prev_item = it
@@ -376,6 +406,17 @@ class Session(BaseModel):
                 pass
             return messages
 
+        def _maybe_backfill_user_messages(*, meta_path: Path, meta: dict[str, Any], user_messages: list[str]) -> None:
+            if isinstance(meta.get("user_messages"), list):
+                return
+            meta["user_messages"] = user_messages
+            try:
+                tmp_path = meta_path.with_suffix(".json.tmp")
+                tmp_path.write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
+                tmp_path.replace(meta_path)
+            except OSError:
+                return
+
         items: list[Session.SessionMetaBrief] = []
         for meta_path in store.iter_meta_files():
             data = _read_json_dict(meta_path)
@@ -388,7 +429,15 @@ class Session(BaseModel):
             created = float(data.get("created_at", meta_path.stat().st_mtime))
             updated = float(data.get("updated_at", meta_path.stat().st_mtime))
             work_dir = str(data.get("work_dir", ""))
-            user_messages = _get_user_messages(sid)
+
+            user_messages_raw = data.get("user_messages")
+            if isinstance(user_messages_raw, list) and all(
+                isinstance(m, str) for m in cast(list[object], user_messages_raw)
+            ):
+                user_messages = cast(list[str], user_messages_raw)
+            else:
+                user_messages = _get_user_messages(sid)
+                _maybe_backfill_user_messages(meta_path=meta_path, meta=data, user_messages=user_messages)
             messages_count = int(data.get("messages_count", -1))
             model_name = data.get("model_name") if isinstance(data.get("model_name"), str) else None
 

klaude_code/session/store.py

@@ -193,6 +193,7 @@ def build_meta_snapshot(
     sub_agent_state: model.SubAgentState | None,
     file_tracker: dict[str, model.FileStatus],
     todos: list[model.TodoItem],
+    user_messages: list[str],
     created_at: float,
     updated_at: float,
     messages_count: int,
@@ -206,6 +207,8 @@ def build_meta_snapshot(
         "sub_agent_state": sub_agent_state.model_dump(mode="json") if sub_agent_state else None,
         "file_tracker": {path: status.model_dump(mode="json") for path, status in file_tracker.items()},
         "todos": [todo.model_dump(mode="json", exclude_defaults=True) for todo in todos],
+        # Cache user messages to avoid scanning events.jsonl during session listing.
+        "user_messages": list(user_messages),
         "created_at": created_at,
         "updated_at": updated_at,
         "messages_count": messages_count,