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

klaude_code/llm/claude/client.py

@@ -0,0 +1,105 @@
+import json
+from collections.abc import AsyncGenerator
+from typing import override
+
+import anthropic
+import httpx
+from anthropic import APIError
+
+from klaude_code.auth.claude.exceptions import ClaudeNotLoggedInError
+from klaude_code.auth.claude.oauth import ClaudeOAuth
+from klaude_code.auth.claude.token_manager import ClaudeTokenManager
+from klaude_code.const import (
+    ANTHROPIC_BETA_FINE_GRAINED_TOOL_STREAMING,
+    ANTHROPIC_BETA_INTERLEAVED_THINKING,
+    ANTHROPIC_BETA_OAUTH,
+    LLM_HTTP_TIMEOUT_CONNECT,
+    LLM_HTTP_TIMEOUT_READ,
+    LLM_HTTP_TIMEOUT_TOTAL,
+)
+from klaude_code.llm.anthropic.client import build_payload, parse_anthropic_stream
+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.usage import MetadataTracker, error_stream_items
+from klaude_code.protocol import llm_param, message
+from klaude_code.trace import DebugType, log_debug
+
+_CLAUDE_OAUTH_REQUIRED_BETAS: tuple[str, ...] = (
+    ANTHROPIC_BETA_OAUTH,
+    ANTHROPIC_BETA_FINE_GRAINED_TOOL_STREAMING,
+)
+
+
+@register(llm_param.LLMClientProtocol.CLAUDE_OAUTH)
+class ClaudeClient(LLMClientABC):
+    """Claude OAuth client using Anthropic messages API with Bearer auth token."""
+
+    def __init__(self, config: llm_param.LLMConfigParameter):
+        super().__init__(config)
+
+        if config.base_url:
+            raise ValueError("CLAUDE protocol does not support custom base_url")
+
+        self._token_manager = ClaudeTokenManager()
+        self._oauth = ClaudeOAuth(self._token_manager)
+
+        if not self._token_manager.is_logged_in():
+            raise ClaudeNotLoggedInError("Claude authentication required. Run 'klaude login claude' first.")
+
+        self.client = self._create_client()
+
+    def _create_client(self) -> anthropic.AsyncAnthropic:
+        token = self._oauth.ensure_valid_token()
+        return anthropic.AsyncAnthropic(
+            auth_token=token,
+            timeout=httpx.Timeout(LLM_HTTP_TIMEOUT_TOTAL, connect=LLM_HTTP_TIMEOUT_CONNECT, read=LLM_HTTP_TIMEOUT_READ),
+        )
+
+    def _ensure_valid_token(self) -> None:
+        state = self._token_manager.get_state()
+        if state is None:
+            raise ClaudeNotLoggedInError("Not logged in to Claude. Run 'klaude login claude' first.")
+
+        if state.is_expired():
+            self._oauth.refresh()
+            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[message.LLMStreamItem]:
+        self._ensure_valid_token()
+        param = apply_config_defaults(param, self.get_llm_config())
+
+        metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
+
+        # Anthropic OAuth requires the oauth beta flag
+        extra_betas = list(_CLAUDE_OAUTH_REQUIRED_BETAS)
+        payload = build_payload(param, extra_betas=extra_betas)
+
+        # Keep the interleaved-thinking beta in sync with configured thinking.
+        if not (param.thinking and param.thinking.type == "enabled"):
+            payload["betas"] = [b for b in payload.get("betas", []) if b != ANTHROPIC_BETA_INTERLEAVED_THINKING]
+
+        log_debug(
+            json.dumps(payload, ensure_ascii=False, default=str),
+            style="yellow",
+            debug_type=DebugType.LLM_PAYLOAD,
+        )
+
+        stream = self.client.beta.messages.create(
+            **payload,
+            extra_headers={"extra": json.dumps({"session_id": param.session_id}, sort_keys=True)},
+        )
+
+        try:
+            async for item in parse_anthropic_stream(stream, param, metadata_tracker):
+                yield item
+        except (APIError, httpx.HTTPError) as e:
+            error_message = f"{e.__class__.__name__} {e!s}"
+            for item in error_stream_items(metadata_tracker, error=error_message):
+                yield item

klaude_code/llm/client.py

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
 from typing import ParamSpec, TypeVar, cast
 
-from klaude_code.protocol import llm_param, model
+from klaude_code.protocol import llm_param, message
 
 
 class LLMClientABC(ABC):
@@ -15,9 +15,10 @@ class LLMClientABC(ABC):
         pass
 
     @abstractmethod
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
+        if False:  # pragma: no cover
+            yield cast(message.LLMStreamItem, None)
         raise NotImplementedError
-        yield cast(model.ConversationItem, None)
 
     def get_llm_config(self) -> llm_param.LLMConfigParameter:
         return self._config

klaude_code/llm/codex/client.py

@@ -12,13 +12,20 @@ from openai.types.responses.response_create_params import ResponseCreateParamsSt
 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.const import (
+    CODEX_BASE_URL,
+    CODEX_USER_AGENT,
+    LLM_HTTP_TIMEOUT_CONNECT,
+    LLM_HTTP_TIMEOUT_READ,
+    LLM_HTTP_TIMEOUT_TOTAL,
+)
 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.llm.usage import MetadataTracker, error_stream_items
+from klaude_code.protocol import llm_param, message
 from klaude_code.trace import DebugType, log_debug
 
 
@@ -57,17 +64,14 @@ def build_payload(param: llm_param.LLMCallParameter) -> ResponseCreateParamsStre
     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",
+    "User-Agent": CODEX_USER_AGENT,
     "OpenAI-Beta": "responses=experimental",
 }
 
 
-@register(llm_param.LLMClientProtocol.CODEX)
+@register(llm_param.LLMClientProtocol.CODEX_OAUTH)
 class CodexClient(LLMClientABC):
     """LLM client for Codex API using ChatGPT subscription."""
 
@@ -90,7 +94,7 @@ class CodexClient(LLMClientABC):
         return AsyncOpenAI(
             api_key=state.access_token,
             base_url=CODEX_BASE_URL,
-            timeout=httpx.Timeout(300.0, connect=15.0, read=285.0),
+            timeout=httpx.Timeout(LLM_HTTP_TIMEOUT_TOTAL, connect=LLM_HTTP_TIMEOUT_CONNECT, read=LLM_HTTP_TIMEOUT_READ),
             default_headers={
                 **CODEX_HEADERS,
                 "chatgpt-account-id": state.account_id,
@@ -114,7 +118,7 @@ class CodexClient(LLMClientABC):
         return cls(config)
 
     @override
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
         # Ensure token is valid before API call
         self._ensure_valid_token()
 
@@ -142,7 +146,9 @@ class CodexClient(LLMClientABC):
                 extra_headers=extra_headers,
             )
         except (openai.OpenAIError, httpx.HTTPError) as e:
-            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            error_message = f"{e.__class__.__name__} {e!s}"
+            for item in error_stream_items(metadata_tracker, error=error_message):
+                yield item
             return
 
         async for item in parse_responses_stream(stream, param, metadata_tracker):

klaude_code/llm/google/client.py

@@ -5,7 +5,7 @@
 
 import json
 from collections.abc import AsyncGenerator, AsyncIterator
-from typing import Any, cast, override
+from typing import Any, Literal, cast, override
 from uuid import uuid4
 
 import httpx
@@ -26,7 +26,7 @@ from klaude_code.llm.google.input import convert_history_to_contents, convert_to
 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.protocol import llm_param, message, model
 from klaude_code.trace import DebugType, log_debug
 
 
@@ -114,25 +114,74 @@ def _merge_partial_args(dst: dict[str, Any], partial_args: list[Any] | None) ->
         dst[key] = _partial_arg_value(partial)
 
 
+def _map_finish_reason(reason: str) -> model.StopReason | None:
+    normalized = reason.strip().lower()
+    mapping: dict[str, model.StopReason] = {
+        "stop": "stop",
+        "end_turn": "stop",
+        "max_tokens": "length",
+        "length": "length",
+        "tool_use": "tool_use",
+        "safety": "error",
+        "recitation": "error",
+        "other": "error",
+        "content_filter": "error",
+        "blocked": "error",
+        "blocklist": "error",
+        "cancelled": "aborted",
+        "canceled": "aborted",
+        "aborted": "aborted",
+    }
+    return mapping.get(normalized)
+
+
 async def parse_google_stream(
     stream: AsyncIterator[Any],
     param: llm_param.LLMCallParameter,
     metadata_tracker: MetadataTracker,
-) -> AsyncGenerator[model.ConversationItem]:
+) -> AsyncGenerator[message.LLMStreamItem]:
     response_id: str | None = None
-    started = False
+    stage: Literal["waiting", "thinking", "assistant", "tool"] = "waiting"
 
     accumulated_text: list[str] = []
     accumulated_thoughts: list[str] = []
     thought_signature: str | None = None
+    assistant_parts: list[message.Part] = []
 
     # 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()
+    completed_tool_items: set[str] = set()
 
     last_usage_metadata: UsageMetadata | None = None
+    stop_reason: model.StopReason | None = None
+
+    def flush_thinking() -> None:
+        nonlocal thought_signature
+        if accumulated_thoughts:
+            assistant_parts.append(
+                message.ThinkingTextPart(
+                    text="".join(accumulated_thoughts),
+                    model_id=str(param.model),
+                )
+            )
+            accumulated_thoughts.clear()
+        if thought_signature:
+            assistant_parts.append(
+                message.ThinkingSignaturePart(
+                    signature=thought_signature,
+                    model_id=str(param.model),
+                    format="google_thought_signature",
+                )
+            )
+            thought_signature = None
+
+    def flush_text() -> None:
+        if not accumulated_text:
+            return
+        assistant_parts.append(message.TextPart(text="".join(accumulated_text)))
+        accumulated_text.clear()
 
     async for chunk in stream:
         log_debug(
@@ -143,33 +192,44 @@ async def parse_google_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
+        finish_reason = getattr(candidate0, "finish_reason", None) if candidate0 else None
+        if finish_reason is not None:
+            if isinstance(finish_reason, str):
+                reason_value = finish_reason
+            else:
+                reason_value = getattr(finish_reason, "name", None) or str(finish_reason)
+            stop_reason = _map_finish_reason(reason_value)
         content = getattr(candidate0, "content", None) if candidate0 else None
-        parts = getattr(content, "parts", None) if content else None
-        if not parts:
+        content_parts = getattr(content, "parts", None) if content else None
+        if not content_parts:
             continue
 
-        for part in parts:
+        for part in content_parts:
             if getattr(part, "text", None) is not None:
-                metadata_tracker.record_token()
                 text = part.text
+                if not text:
+                    continue
+                metadata_tracker.record_token()
                 if getattr(part, "thought", False) is True:
+                    if stage == "assistant":
+                        flush_text()
+                    stage = "thinking"
                     accumulated_thoughts.append(text)
                     if getattr(part, "thought_signature", None):
                         thought_signature = part.thought_signature
-                    yield model.ReasoningTextDelta(content=text, response_id=response_id)
+                    yield message.ThinkingTextDelta(content=text, response_id=response_id)
                 else:
+                    if stage == "thinking":
+                        flush_thinking()
+                    stage = "assistant"
                     accumulated_text.append(text)
-                    yield model.AssistantMessageDelta(content=text, response_id=response_id)
+                    yield message.AssistantTextDelta(content=text, response_id=response_id)
 
             function_call = getattr(part, "function_call", None)
             if function_call is None:
@@ -182,17 +242,23 @@ async def parse_google_stream(
 
             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)
+                yield message.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),
+                if stage == "thinking":
+                    flush_thinking()
+                if stage == "assistant":
+                    flush_text()
+                stage = "tool"
+                assistant_parts.append(
+                    message.ToolCallPart(
+                        call_id=call_id,
+                        tool_name=name,
+                        arguments_json=json.dumps(args_obj, ensure_ascii=False),
+                    )
                 )
+                completed_tool_items.add(call_id)
                 continue
 
             partial_args = getattr(function_call, "partial_args", None)
@@ -201,53 +267,49 @@ async def parse_google_stream(
                 _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),
+            if will_continue is False and call_id in partial_args_by_call and call_id not in completed_tool_items:
+                if stage == "thinking":
+                    flush_thinking()
+                if stage == "assistant":
+                    flush_text()
+                stage = "tool"
+                assistant_parts.append(
+                    message.ToolCallPart(
+                        call_id=call_id,
+                        tool_name=name,
+                        arguments_json=json.dumps(partial_args_by_call[call_id], ensure_ascii=False),
+                    )
                 )
+                completed_tool_items.add(call_id)
 
     # 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:
+        if call_id in completed_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",
+        assistant_parts.append(
+            message.ToolCallPart(
+                call_id=call_id,
+                tool_name=name,
+                arguments_json=json.dumps(args, ensure_ascii=False),
             )
+        )
 
-    if accumulated_text:
-        metadata_tracker.record_token()
-        yield model.AssistantMessageItem(content="".join(accumulated_text), response_id=response_id)
+    flush_thinking()
+    flush_text()
 
     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()
+    metadata = metadata_tracker.finalize()
+    yield message.AssistantMessage(
+        parts=assistant_parts,
+        response_id=response_id,
+        usage=metadata,
+        stop_reason=stop_reason,
+    )
 
 
 @register(llm_param.LLMClientProtocol.GOOGLE)
@@ -270,7 +332,7 @@ class GoogleClient(LLMClientABC):
         return cls(config)
 
     @override
-    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem]:
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[message.LLMStreamItem]:
         param = apply_config_defaults(param, self.get_llm_config())
         metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
 
@@ -297,13 +359,13 @@ class GoogleClient(LLMClientABC):
                 config=config,
             )
         except (APIError, ClientError, ServerError, httpx.HTTPError) as e:
-            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
-            yield metadata_tracker.finalize()
+            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield message.AssistantMessage(parts=[], response_id=None, usage=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()
+            yield message.StreamErrorItem(error=f"{e.__class__.__name__} {e!s}")
+            yield message.AssistantMessage(parts=[], response_id=None, usage=metadata_tracker.finalize())