klaude-code 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

klaude_code/llm/openai_compatible/client.py

@@ -1,49 +1,24 @@
 import json
-import time
 from collections.abc import AsyncGenerator
-from typing import Callable, Literal, ParamSpec, TypeVar, override
+from typing import Literal, override
 
 import httpx
 import openai
 from openai import APIError, RateLimitError
 
-from klaude_code.llm.client import LLMClientABC
+from klaude_code.llm.client import LLMClientABC, call_with_logged_payload
+from klaude_code.llm.input_common import apply_config_defaults
 from klaude_code.llm.openai_compatible.input import convert_history_to_input, convert_tool_schema
 from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
 from klaude_code.llm.registry import register
-from klaude_code.protocol import model
-from klaude_code.protocol.llm_parameter import (
-    LLMCallParameter,
-    LLMClientProtocol,
-    LLMConfigParameter,
-    apply_config_defaults,
-)
-from klaude_code.protocol.model import StreamErrorItem
+from klaude_code.llm.usage import MetadataTracker, convert_usage
+from klaude_code.protocol import llm_param, model
 from klaude_code.trace import DebugType, log_debug
 
-P = ParamSpec("P")
-R = TypeVar("R")
 
-
-def call_with_logged_payload(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
-    """Call an SDK function while logging the JSON payload.
-
-    The function reuses the original callable's type signature via ParamSpec
-    so static type checkers can validate arguments at the call site.
-    """
-
-    payload = {k: v for k, v in kwargs.items() if v is not None}
-    log_debug(
-        json.dumps(payload, ensure_ascii=False, default=str),
-        style="yellow",
-        debug_type=DebugType.LLM_PAYLOAD,
-    )
-    return func(*args, **kwargs)
-
-
-@register(LLMClientProtocol.OPENAI)
+@register(llm_param.LLMClientProtocol.OPENAI)
 class OpenAICompatibleClient(LLMClientABC):
-    def __init__(self, config: LLMConfigParameter):
+    def __init__(self, config: llm_param.LLMConfigParameter):
         super().__init__(config)
         if config.is_azure:
             if not config.base_url:
@@ -64,18 +39,16 @@ class OpenAICompatibleClient(LLMClientABC):
 
     @classmethod
     @override
-    def create(cls, config: LLMConfigParameter) -> "LLMClientABC":
+    def create(cls, config: llm_param.LLMConfigParameter) -> "LLMClientABC":
         return cls(config)
 
     @override
-    async def call(self, param: LLMCallParameter) -> AsyncGenerator[model.ConversationItem, None]:
+    async def call(self, param: llm_param.LLMCallParameter) -> AsyncGenerator[model.ConversationItem, None]:
         param = apply_config_defaults(param, self.get_llm_config())
         messages = convert_history_to_input(param.input, param.system, param.model)
         tools = convert_tool_schema(param.tools)
 
-        request_start_time = time.time()
-        first_token_time: float | None = None
-        last_token_time: float | None = None
+        metadata_tracker = MetadataTracker(cost_config=self._config.cost)
 
         extra_body = {}
         extra_headers = {"extra": json.dumps({"session_id": param.session_id})}
@@ -105,8 +78,8 @@ class OpenAICompatibleClient(LLMClientABC):
         accumulated_reasoning: list[str] = []
         accumulated_content: list[str] = []
         accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
+        emitted_tool_start_indices: set[int] = set()
         response_id: str | None = None
-        metadata_item = model.ResponseMetadataItem()
 
         def flush_reasoning_items() -> list[model.ConversationItem]:
             nonlocal accumulated_reasoning
@@ -140,17 +113,23 @@ class OpenAICompatibleClient(LLMClientABC):
 
         try:
             async for event in await stream:
-                log_debug(event.model_dump_json(exclude_none=True), style="blue", debug_type=DebugType.LLM_STREAM)
+                log_debug(
+                    event.model_dump_json(exclude_none=True),
+                    style="blue",
+                    debug_type=DebugType.LLM_STREAM,
+                )
                 if not response_id and event.id:
                     response_id = event.id
                     accumulated_tool_calls.response_id = response_id
                     yield model.StartItem(response_id=response_id)
-                if event.usage is not None and event.usage.completion_tokens is not None:  # pyright: ignore[reportUnnecessaryComparison] gcp gemini will return None usage field
-                    metadata_item.usage = convert_usage(event.usage, param.context_limit)
+                if (
+                    event.usage is not None and event.usage.completion_tokens is not None  # pyright: ignore[reportUnnecessaryComparison] gcp gemini will return None usage field
+                ):
+                    metadata_tracker.set_usage(convert_usage(event.usage, param.context_limit))
                 if event.model:
-                    metadata_item.model_name = event.model
+                    metadata_tracker.set_model_name(event.model)
                 if provider := getattr(event, "provider", None):
-                    metadata_item.provider = str(provider)
+                    metadata_tracker.set_provider(str(provider))
 
                 if len(event.choices) == 0:
                     continue
@@ -158,9 +137,11 @@ class OpenAICompatibleClient(LLMClientABC):
 
                 # Support Kimi K2's usage field in choice
                 if hasattr(event.choices[0], "usage") and getattr(event.choices[0], "usage"):
-                    metadata_item.usage = convert_usage(
-                        openai.types.CompletionUsage.model_validate(getattr(event.choices[0], "usage")),
-                        param.context_limit,
+                    metadata_tracker.set_usage(
+                        convert_usage(
+                            openai.types.CompletionUsage.model_validate(getattr(event.choices[0], "usage")),
+                            param.context_limit,
+                        )
                     )
 
                 # Reasoning
@@ -170,9 +151,7 @@ class OpenAICompatibleClient(LLMClientABC):
                 if hasattr(delta, "reasoning_content") and getattr(delta, "reasoning_content"):
                     reasoning_content = getattr(delta, "reasoning_content")
                 if reasoning_content:
-                    if first_token_time is None:
-                        first_token_time = time.time()
-                    last_token_time = time.time()
+                    metadata_tracker.record_token()
                     stage = "reasoning"
                     accumulated_reasoning.append(reasoning_content)
 
@@ -180,9 +159,7 @@ class OpenAICompatibleClient(LLMClientABC):
                 if delta.content and (
                     stage == "assistant" or delta.content.strip()
                 ):  # Process all content in assistant stage, filter empty content in reasoning stage
-                    if first_token_time is None:
-                        first_token_time = time.time()
-                    last_token_time = time.time()
+                    metadata_tracker.record_token()
                     if stage == "reasoning":
                         for item in flush_reasoning_items():
                             yield item
@@ -198,9 +175,7 @@ class OpenAICompatibleClient(LLMClientABC):
 
                 # Tool
                 if delta.tool_calls and len(delta.tool_calls) > 0:
-                    if first_token_time is None:
-                        first_token_time = time.time()
-                    last_token_time = time.time()
+                    metadata_tracker.record_token()
                     if stage == "reasoning":
                         for item in flush_reasoning_items():
                             yield item
@@ -208,9 +183,18 @@ class OpenAICompatibleClient(LLMClientABC):
                         for item in flush_assistant_items():
                             yield item
                     stage = "tool"
+                    # Emit ToolCallStartItem for new tool calls
+                    for tc in delta.tool_calls:
+                        if tc.index not in emitted_tool_start_indices and tc.function and tc.function.name:
+                            emitted_tool_start_indices.add(tc.index)
+                            yield model.ToolCallStartItem(
+                                response_id=response_id,
+                                call_id=tc.id or "",
+                                name=tc.function.name,
+                            )
                     accumulated_tool_calls.add(delta.tool_calls)
         except (RateLimitError, APIError) as e:
-            yield StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
+            yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
 
         # Finalize
         for item in flush_reasoning_items():
@@ -223,31 +207,5 @@ class OpenAICompatibleClient(LLMClientABC):
             for tool_call_item in flush_tool_call_items():
                 yield tool_call_item
 
-        metadata_item.response_id = response_id
-
-        # Calculate performance metrics if we have timing data
-        if metadata_item.usage and first_token_time is not None:
-            metadata_item.usage.first_token_latency_ms = (first_token_time - request_start_time) * 1000
-
-            if last_token_time is not None and metadata_item.usage.output_tokens > 0:
-                time_duration = last_token_time - first_token_time
-                if time_duration >= 0.15:
-                    metadata_item.usage.throughput_tps = metadata_item.usage.output_tokens / time_duration
-
-        yield metadata_item
-
-
-def convert_usage(usage: openai.types.CompletionUsage, context_limit: int | None = None) -> model.Usage:
-    total_tokens = usage.total_tokens
-    context_usage_percent = (total_tokens / context_limit) * 100 if context_limit else None
-    return model.Usage(
-        input_tokens=usage.prompt_tokens,
-        cached_tokens=(usage.prompt_tokens_details.cached_tokens if usage.prompt_tokens_details else 0) or 0,
-        reasoning_tokens=(usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else 0)
-        or 0,
-        output_tokens=usage.completion_tokens,
-        total_tokens=total_tokens,
-        context_usage_percent=context_usage_percent,
-        throughput_tps=None,
-        first_token_latency_ms=None,
-    )
+        metadata_tracker.set_response_id(response_id)
+        yield metadata_tracker.finalize()

klaude_code/llm/openai_compatible/input.py

@@ -6,15 +6,8 @@
 from openai.types import chat
 from openai.types.chat import ChatCompletionContentPartParam
 
-from klaude_code.llm.input_common import (
-    AssistantGroup,
-    ToolGroup,
-    UserGroup,
-    merge_reminder_text,
-    parse_message_groups,
-)
-from klaude_code.protocol.llm_parameter import ToolSchema
-from klaude_code.protocol.model import ConversationItem, ImageURLPart
+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 _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
@@ -39,7 +32,9 @@ def _tool_group_to_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
     }
 
 
-def _assistant_group_to_message(group: AssistantGroup) -> chat.ChatCompletionMessageParam:
+def _assistant_group_to_message(
+    group: AssistantGroup,
+) -> chat.ChatCompletionMessageParam:
     assistant_message: dict[str, object] = {"role": "assistant"}
 
     if group.text_content:
@@ -61,13 +56,15 @@ def _assistant_group_to_message(group: AssistantGroup) -> chat.ChatCompletionMes
     return assistant_message
 
 
-def build_user_content_parts(images: list[ImageURLPart]) -> list[ChatCompletionContentPartParam]:
+def build_user_content_parts(
+    images: list[model.ImageURLPart],
+) -> list[ChatCompletionContentPartParam]:
     """Build content parts for images only. Used by OpenRouter."""
     return [{"type": "image_url", "image_url": {"url": image.image_url.url}} for image in images]
 
 
 def convert_history_to_input(
-    history: list[ConversationItem],
+    history: list[model.ConversationItem],
     system: str | None = None,
     model_name: str | None = None,
 ) -> list[chat.ChatCompletionMessageParam]:
@@ -79,9 +76,7 @@ def convert_history_to_input(
         system: System message.
         model_name: Model name. Not used in OpenAI-compatible, kept for API consistency.
     """
-    messages: list[chat.ChatCompletionMessageParam] = (
-        [{"role": "system", "content": system}] if system else []
-    )
+    messages: list[chat.ChatCompletionMessageParam] = [{"role": "system", "content": system}] if system else []
 
     for group in parse_message_groups(history):
         match group:
@@ -96,7 +91,7 @@ def convert_history_to_input(
 
 
 def convert_tool_schema(
-    tools: list[ToolSchema] | None,
+    tools: list[llm_param.ToolSchema] | None,
 ) -> list[chat.ChatCompletionToolParam]:
     if tools is None:
         return []

klaude_code/llm/openai_compatible/tool_call_accumulator.py

@@ -8,7 +8,7 @@ from klaude_code.protocol import model
 
 class ToolCallAccumulatorABC(ABC):
     @abstractmethod
-    def add(self, chunks: list[ChoiceDeltaToolCall]):
+    def add(self, chunks: list[ChoiceDeltaToolCall]) -> None:
         pass
 
     @abstractmethod
@@ -50,7 +50,7 @@ class BasicToolCallAccumulator(ToolCallAccumulatorABC, BaseModel):
     chunks_by_step: list[list[ChoiceDeltaToolCall]] = Field(default_factory=list)  # pyright: ignore[reportUnknownVariableType]
     response_id: str | None = None
 
-    def add(self, chunks: list[ChoiceDeltaToolCall]):
+    def add(self, chunks: list[ChoiceDeltaToolCall]) -> None:
         self.chunks_by_step.append(chunks)
 
     def get(self) -> list[model.ToolCallItem]: