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

klaude_code/llm/codex/client.py

@@ -0,0 +1,149 @@
+"""Codex LLM client using ChatGPT subscription via OAuth."""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import override
+
+import httpx
+import openai
+from openai import AsyncOpenAI
+from openai.types.responses.response_create_params import ResponseCreateParamsStreaming
+
+from klaude_code.auth.codex.exceptions import CodexNotLoggedInError
+from klaude_code.auth.codex.oauth import CodexOAuth
+from klaude_code.auth.codex.token_manager import CodexTokenManager
+from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.input_common import apply_config_defaults
+from klaude_code.llm.registry import register
+from klaude_code.llm.responses.client import parse_responses_stream
+from klaude_code.llm.responses.input import convert_history_to_input, convert_tool_schema
+from klaude_code.llm.usage import MetadataTracker
+from klaude_code.protocol import llm_param, model
+from klaude_code.trace import DebugType, log_debug
+
+
+def build_payload(param: llm_param.LLMCallParameter) -> ResponseCreateParamsStreaming:
+    """Build Codex API request parameters."""
+    inputs = convert_history_to_input(param.input, param.model)
+    tools = convert_tool_schema(param.tools)
+
+    session_id = param.session_id or ""
+
+    payload: ResponseCreateParamsStreaming = {
+        "model": str(param.model),
+        "tool_choice": "auto",
+        "parallel_tool_calls": True,
+        "include": [
+            "reasoning.encrypted_content",
+        ],
+        "store": False,
+        "stream": True,
+        "input": inputs,
+        "instructions": param.system,
+        "tools": tools,
+        "prompt_cache_key": session_id,
+        # max_output_token and temperature is not supported in Codex API
+    }
+
+    if param.thinking and param.thinking.reasoning_effort:
+        payload["reasoning"] = {
+            "effort": param.thinking.reasoning_effort,
+            "summary": param.thinking.reasoning_summary,
+        }
+
+    if param.verbosity:
+        payload["text"] = {"verbosity": param.verbosity}
+
+    return payload
+
+
+# Codex API configuration
+CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
+CODEX_HEADERS = {
+    "originator": "codex_cli_rs",
+    # Mocked Codex-style user agent string
+    "User-Agent": "codex_cli_rs/0.0.0-klaude",
+    "OpenAI-Beta": "responses=experimental",
+}
+
+
+@register(llm_param.LLMClientProtocol.CODEX)
+class CodexClient(LLMClientABC):
+    """LLM client for Codex API using ChatGPT subscription."""
+
+    def __init__(self, config: llm_param.LLMConfigParameter):
+        super().__init__(config)
+        self._token_manager = CodexTokenManager()
+        self._oauth = CodexOAuth(self._token_manager)
+
+        if not self._token_manager.is_logged_in():
+            raise CodexNotLoggedInError("Codex authentication required. Run 'klaude login codex' first.")
+
+        self.client = self._create_client()
+
+    def _create_client(self) -> AsyncOpenAI:
+        """Create OpenAI client with Codex configuration."""
+        state = self._token_manager.get_state()
+        if state is None:
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+
+        return AsyncOpenAI(
+            api_key=state.access_token,
+            base_url=CODEX_BASE_URL,
+            timeout=httpx.Timeout(300.0, connect=15.0, read=285.0),
+            default_headers={
+                **CODEX_HEADERS,
+                "chatgpt-account-id": state.account_id,
+            },
+        )
+
+    def _ensure_valid_token(self) -> None:
+        """Ensure token is valid, refresh if needed."""
+        state = self._token_manager.get_state()
+        if state is None:
+            raise CodexNotLoggedInError("Not logged in to Codex. Run 'klaude login codex' first.")
+
+        if state.is_expired():
+            self._oauth.refresh()
+            # Recreate client with new token
+            self.client = self._create_client()
+
+    @classmethod
+    @override
+    def create(cls, config: llm_param.LLMConfigParameter) -> "LLMClientABC":
+        return cls(config)
+
+    @override
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem]:
+        # Ensure token is valid before API call
+        self._ensure_valid_token()
+
+        param = apply_config_defaults(param, self.get_llm_config())
+
+        metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
+
+        payload = build_payload(param)
+
+        session_id = param.session_id or ""
+        extra_headers: dict[str, str] = {}
+        if session_id:
+            # Must send conversation_id/session_id headers to improve ChatGPT backend prompt cache hit rate.
+            extra_headers["conversation_id"] = session_id
+            extra_headers["session_id"] = session_id
+
+        log_debug(
+            json.dumps(payload, ensure_ascii=False, default=str),
+            style="yellow",
+            debug_type=DebugType.LLM_PAYLOAD,
+        )
+        try:
+            stream = await self.client.responses.create(
+                **payload,
+                extra_headers=extra_headers,
+            )
+        except (openai.OpenAIError, httpx.HTTPError) as e:
+            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            return
+
+        async for item in parse_responses_stream(stream, param, metadata_tracker):
+            yield item

klaude_code/llm/google/__init__.py

@@ -0,0 +1,3 @@
+from .client import GoogleClient
+
+__all__ = ["GoogleClient"]

klaude_code/llm/google/client.py

@@ -0,0 +1,309 @@
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportAttributeAccessIssue=false
+
+import json
+from collections.abc import AsyncGenerator, AsyncIterator
+from typing import Any, cast, override
+from uuid import uuid4
+
+import httpx
+from google.genai import Client
+from google.genai.errors import APIError, ClientError, ServerError
+from google.genai.types import (
+    FunctionCallingConfig,
+    FunctionCallingConfigMode,
+    GenerateContentConfig,
+    HttpOptions,
+    ThinkingConfig,
+    ToolConfig,
+    UsageMetadata,
+)
+
+from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.google.input import convert_history_to_contents, convert_tool_schema
+from klaude_code.llm.input_common import apply_config_defaults
+from klaude_code.llm.registry import register
+from klaude_code.llm.usage import MetadataTracker
+from klaude_code.protocol import llm_param, model
+from klaude_code.trace import DebugType, log_debug
+
+
+def _build_config(param: llm_param.LLMCallParameter) -> GenerateContentConfig:
+    tool_list = convert_tool_schema(param.tools)
+    tool_config: ToolConfig | None = None
+
+    if tool_list:
+        tool_config = ToolConfig(
+            function_calling_config=FunctionCallingConfig(
+                mode=FunctionCallingConfigMode.AUTO,
+                # Gemini streams tool args; keep this enabled to maximize fidelity.
+                stream_function_call_arguments=True,
+            )
+        )
+
+    thinking_config: ThinkingConfig | None = None
+    if param.thinking and param.thinking.type == "enabled":
+        thinking_config = ThinkingConfig(
+            include_thoughts=True,
+            thinking_budget=param.thinking.budget_tokens,
+        )
+
+    return GenerateContentConfig(
+        system_instruction=param.system,
+        temperature=param.temperature,
+        max_output_tokens=param.max_tokens,
+        tools=tool_list or None,
+        tool_config=tool_config,
+        thinking_config=thinking_config,
+    )
+
+
+def _usage_from_metadata(
+    usage: UsageMetadata | None,
+    *,
+    context_limit: int | None,
+    max_tokens: int | None,
+) -> model.Usage | None:
+    if usage is None:
+        return None
+
+    cached = usage.cached_content_token_count or 0
+    prompt = usage.prompt_token_count or 0
+    response = usage.response_token_count or 0
+    thoughts = usage.thoughts_token_count or 0
+
+    total = usage.total_token_count
+    if total is None:
+        total = prompt + cached + response + thoughts
+
+    return model.Usage(
+        input_tokens=prompt + cached,
+        cached_tokens=cached,
+        output_tokens=response + thoughts,
+        reasoning_tokens=thoughts,
+        context_size=total,
+        context_limit=context_limit,
+        max_tokens=max_tokens,
+    )
+
+
+def _partial_arg_value(partial: Any) -> Any:
+    if getattr(partial, "string_value", None) is not None:
+        return partial.string_value
+    if getattr(partial, "number_value", None) is not None:
+        return partial.number_value
+    if getattr(partial, "bool_value", None) is not None:
+        return partial.bool_value
+    if getattr(partial, "null_value", None) is not None:
+        return None
+    return None
+
+
+def _merge_partial_args(dst: dict[str, Any], partial_args: list[Any] | None) -> None:
+    if not partial_args:
+        return
+    for partial in partial_args:
+        json_path = getattr(partial, "json_path", None)
+        if not isinstance(json_path, str) or not json_path.startswith("$."):
+            continue
+        key = json_path[2:]
+        if not key or any(ch in key for ch in "[]"):
+            continue
+        dst[key] = _partial_arg_value(partial)
+
+
+async def parse_google_stream(
+    stream: AsyncIterator[Any],
+    param: llm_param.LLMCallParameter,
+    metadata_tracker: MetadataTracker,
+) -> AsyncGenerator[model.ConversationItem]:
+    response_id: str | None = None
+    started = False
+
+    accumulated_text: list[str] = []
+    accumulated_thoughts: list[str] = []
+    thought_signature: str | None = None
+
+    # Track tool calls where args arrive as partial updates.
+    partial_args_by_call: dict[str, dict[str, Any]] = {}
+    started_tool_calls: dict[str, str] = {}  # call_id -> name
+    started_tool_items: set[str] = set()
+    emitted_tool_items: set[str] = set()
+
+    last_usage_metadata: UsageMetadata | None = None
+
+    async for chunk in stream:
+        log_debug(
+            chunk.model_dump_json(exclude_none=True),
+            style="blue",
+            debug_type=DebugType.LLM_STREAM,
+        )
+
+        if response_id is None:
+            response_id = getattr(chunk, "response_id", None) or uuid4().hex
+        assert response_id is not None
+        if not started:
+            started = True
+            yield model.StartItem(response_id=response_id)
+
+        if getattr(chunk, "usage_metadata", None) is not None:
+            last_usage_metadata = chunk.usage_metadata
+
+        candidates = getattr(chunk, "candidates", None) or []
+        candidate0 = candidates[0] if candidates else None
+        content = getattr(candidate0, "content", None) if candidate0 else None
+        parts = getattr(content, "parts", None) if content else None
+        if not parts:
+            continue
+
+        for part in parts:
+            if getattr(part, "text", None) is not None:
+                metadata_tracker.record_token()
+                text = part.text
+                if getattr(part, "thought", False) is True:
+                    accumulated_thoughts.append(text)
+                    if getattr(part, "thought_signature", None):
+                        thought_signature = part.thought_signature
+                    yield model.ReasoningTextDelta(content=text, response_id=response_id)
+                else:
+                    accumulated_text.append(text)
+                    yield model.AssistantMessageDelta(content=text, response_id=response_id)
+
+            function_call = getattr(part, "function_call", None)
+            if function_call is None:
+                continue
+
+            metadata_tracker.record_token()
+            call_id = getattr(function_call, "id", None) or uuid4().hex
+            name = getattr(function_call, "name", None) or ""
+            started_tool_calls.setdefault(call_id, name)
+
+            if call_id not in started_tool_items:
+                started_tool_items.add(call_id)
+                yield model.ToolCallStartItem(response_id=response_id, call_id=call_id, name=name)
+
+            args_obj = getattr(function_call, "args", None)
+            if args_obj is not None:
+                emitted_tool_items.add(call_id)
+                yield model.ToolCallItem(
+                    response_id=response_id,
+                    call_id=call_id,
+                    name=name,
+                    arguments=json.dumps(args_obj, ensure_ascii=False),
+                )
+                continue
+
+            partial_args = getattr(function_call, "partial_args", None)
+            if partial_args is not None:
+                acc = partial_args_by_call.setdefault(call_id, {})
+                _merge_partial_args(acc, partial_args)
+
+            will_continue = getattr(function_call, "will_continue", None)
+            if will_continue is False and call_id in partial_args_by_call and call_id not in emitted_tool_items:
+                emitted_tool_items.add(call_id)
+                yield model.ToolCallItem(
+                    response_id=response_id,
+                    call_id=call_id,
+                    name=name,
+                    arguments=json.dumps(partial_args_by_call[call_id], ensure_ascii=False),
+                )
+
+    # Flush any pending tool calls that never produced args.
+    for call_id, name in started_tool_calls.items():
+        if call_id in emitted_tool_items:
+            continue
+        args = partial_args_by_call.get(call_id, {})
+        emitted_tool_items.add(call_id)
+        yield model.ToolCallItem(
+            response_id=response_id,
+            call_id=call_id,
+            name=name,
+            arguments=json.dumps(args, ensure_ascii=False),
+        )
+
+    if accumulated_thoughts:
+        metadata_tracker.record_token()
+        yield model.ReasoningTextItem(
+            content="".join(accumulated_thoughts),
+            response_id=response_id,
+            model=str(param.model),
+        )
+        if thought_signature:
+            yield model.ReasoningEncryptedItem(
+                encrypted_content=thought_signature,
+                response_id=response_id,
+                model=str(param.model),
+                format="google_thought_signature",
+            )
+
+    if accumulated_text:
+        metadata_tracker.record_token()
+        yield model.AssistantMessageItem(content="".join(accumulated_text), response_id=response_id)
+
+    usage = _usage_from_metadata(last_usage_metadata, context_limit=param.context_limit, max_tokens=param.max_tokens)
+    if usage is not None:
+        metadata_tracker.set_usage(usage)
+    metadata_tracker.set_model_name(str(param.model))
+    metadata_tracker.set_response_id(response_id)
+    yield metadata_tracker.finalize()
+
+
+@register(llm_param.LLMClientProtocol.GOOGLE)
+class GoogleClient(LLMClientABC):
+    def __init__(self, config: llm_param.LLMConfigParameter):
+        super().__init__(config)
+        http_options: HttpOptions | None = None
+        if config.base_url:
+            # If base_url already contains version path, don't append api_version.
+            http_options = HttpOptions(base_url=str(config.base_url), api_version="")
+
+        self.client = Client(
+            api_key=config.api_key,
+            http_options=http_options,
+        )
+
+    @classmethod
+    @override
+    def create(cls, config: llm_param.LLMConfigParameter) -> "LLMClientABC":
+        return cls(config)
+
+    @override
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem]:
+        param = apply_config_defaults(param, self.get_llm_config())
+        metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
+
+        contents = convert_history_to_contents(param.input, model_name=str(param.model))
+        config = _build_config(param)
+
+        log_debug(
+            json.dumps(
+                {
+                    "model": str(param.model),
+                    "contents": [c.model_dump(exclude_none=True) for c in contents],
+                    "config": config.model_dump(exclude_none=True),
+                },
+                ensure_ascii=False,
+            ),
+            style="yellow",
+            debug_type=DebugType.LLM_PAYLOAD,
+        )
+
+        try:
+            stream = await self.client.aio.models.generate_content_stream(
+                model=str(param.model),
+                contents=cast(Any, contents),
+                config=config,
+            )
+        except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
+            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield metadata_tracker.finalize()
+            return
+
+        try:
+            async for item in parse_google_stream(stream, param=param, metadata_tracker=metadata_tracker):
+                yield item
+        except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
+            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield metadata_tracker.finalize()

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)]