klaude-code 1.2.8__py3-none-any.whl → 1.2.10__py3-none-any.whl

klaude_code/llm/input_common.py

@@ -49,10 +49,6 @@ class AssistantGroup:
 
     text_content: str | None = None
     tool_calls: list[model.ToolCallItem] = field(default_factory=lambda: [])
-    reasoning_text: list[model.ReasoningTextItem] = field(default_factory=lambda: [])
-    reasoning_encrypted: list[model.ReasoningEncryptedItem] = field(default_factory=lambda: [])
-    # Preserve original ordering of reasoning items for providers that
-    # need to emit them as an ordered stream (e.g. OpenRouter).
     reasoning_items: list[model.ReasoningTextItem | model.ReasoningEncryptedItem] = field(default_factory=lambda: [])
 
 
@@ -184,10 +180,8 @@ def parse_message_groups(history: list[model.ConversationItem]) -> list[MessageG
                         case model.ToolCallItem():
                             group.tool_calls.append(item)
                         case model.ReasoningTextItem():
-                            group.reasoning_text.append(item)
                             group.reasoning_items.append(item)
                         case model.ReasoningEncryptedItem():
-                            group.reasoning_encrypted.append(item)
                             group.reasoning_items.append(item)
                         case _:
                             pass

klaude_code/llm/openai_compatible/client.py

@@ -1,15 +1,14 @@
 import json
 from collections.abc import AsyncGenerator
-from typing import Literal, override
+from typing import override
 
 import httpx
 import openai
 
-
 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.openai_compatible.stream_processor import StreamStateManager
 from klaude_code.llm.registry import register
 from klaude_code.llm.usage import MetadataTracker, convert_usage
 from klaude_code.protocol import llm_param, model
@@ -48,10 +47,10 @@ class OpenAICompatibleClient(LLMClientABC):
         messages = convert_history_to_input(param.input, param.system, param.model)
         tools = convert_tool_schema(param.tools)
 
-        metadata_tracker = MetadataTracker(cost_config=self._config.cost)
+        metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
 
         extra_body = {}
-        extra_headers = {"extra": json.dumps({"session_id": param.session_id})}
+        extra_headers = {"extra": json.dumps({"session_id": param.session_id}, sort_keys=True)}
 
         if param.thinking:
             extra_body["thinking"] = {
@@ -74,42 +73,7 @@ class OpenAICompatibleClient(LLMClientABC):
             extra_headers=extra_headers,
         )
 
-        stage: Literal["waiting", "reasoning", "assistant", "tool", "done"] = "waiting"
-        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
-
-        def flush_reasoning_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_reasoning
-            if not accumulated_reasoning:
-                return []
-            item = model.ReasoningTextItem(
-                content="".join(accumulated_reasoning),
-                response_id=response_id,
-                model=str(param.model),
-            )
-            accumulated_reasoning = []
-            return [item]
-
-        def flush_assistant_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_content
-            if len(accumulated_content) == 0:
-                return []
-            item = model.AssistantMessageItem(
-                content="".join(accumulated_content),
-                response_id=response_id,
-            )
-            accumulated_content = []
-            return [item]
-
-        def flush_tool_call_items() -> list[model.ToolCallItem]:
-            nonlocal accumulated_tool_calls
-            items: list[model.ToolCallItem] = accumulated_tool_calls.get()
-            if items:
-                accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
-            return items
+        state = StreamStateManager(param_model=str(param.model))
 
         try:
             async for event in await stream:
@@ -118,14 +82,13 @@ class OpenAICompatibleClient(LLMClientABC):
                     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 not state.response_id and event.id:
+                    state.set_response_id(event.id)
+                    yield model.StartItem(response_id=event.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_tracker.set_usage(convert_usage(event.usage, param.context_limit))
+                    metadata_tracker.set_usage(convert_usage(event.usage, param.context_limit, param.max_tokens))
                 if event.model:
                     metadata_tracker.set_model_name(event.model)
                 if provider := getattr(event, "provider", None):
@@ -141,6 +104,7 @@ class OpenAICompatibleClient(LLMClientABC):
                         convert_usage(
                             openai.types.CompletionUsage.model_validate(getattr(event.choices[0], "usage")),
                             param.context_limit,
+                            param.max_tokens,
                         )
                     )
 
@@ -152,60 +116,53 @@ class OpenAICompatibleClient(LLMClientABC):
                     reasoning_content = getattr(delta, "reasoning_content")
                 if reasoning_content:
                     metadata_tracker.record_token()
-                    stage = "reasoning"
-                    accumulated_reasoning.append(reasoning_content)
+                    state.stage = "reasoning"
+                    state.accumulated_reasoning.append(reasoning_content)
 
                 # Assistant
                 if delta.content and (
-                    stage == "assistant" or delta.content.strip()
+                    state.stage == "assistant" or delta.content.strip()
                 ):  # Process all content in assistant stage, filter empty content in reasoning stage
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "tool":
-                        for item in flush_tool_call_items():
+                    elif state.stage == "tool":
+                        for item in state.flush_tool_calls():
                             yield item
-                    stage = "assistant"
-                    accumulated_content.append(delta.content)
+                    state.stage = "assistant"
+                    state.accumulated_content.append(delta.content)
                     yield model.AssistantMessageDelta(
                         content=delta.content,
-                        response_id=response_id,
+                        response_id=state.response_id,
                     )
 
                 # Tool
                 if delta.tool_calls and len(delta.tool_calls) > 0:
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "assistant":
-                        for item in flush_assistant_items():
+                    elif state.stage == "assistant":
+                        for item in state.flush_assistant():
                             yield item
-                    stage = "tool"
+                    state.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)
+                        if tc.index not in state.emitted_tool_start_indices and tc.function and tc.function.name:
+                            state.emitted_tool_start_indices.add(tc.index)
                             yield model.ToolCallStartItem(
-                                response_id=response_id,
+                                response_id=state.response_id,
                                 call_id=tc.id or "",
                                 name=tc.function.name,
                             )
-                    accumulated_tool_calls.add(delta.tool_calls)
+                    state.accumulated_tool_calls.add(delta.tool_calls)
         except (openai.OpenAIError, httpx.HTTPError) as e:
             yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
 
         # Finalize
-        for item in flush_reasoning_items():
+        for item in state.flush_all():
             yield item
 
-        for item in flush_assistant_items():
-            yield item
-
-        if stage == "tool":
-            for tool_call_item in flush_tool_call_items():
-                yield tool_call_item
-
-        metadata_tracker.set_response_id(response_id)
+        metadata_tracker.set_response_id(state.response_id)
         yield metadata_tracker.finalize()

klaude_code/llm/openai_compatible/input.py

@@ -10,7 +10,8 @@ from klaude_code.llm.input_common import AssistantGroup, ToolGroup, UserGroup, m
 from klaude_code.protocol import llm_param, model
 
 
-def _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
+def user_group_to_openai_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
+    """Convert a UserGroup to an OpenAI-compatible chat message."""
     parts: list[ChatCompletionContentPartParam] = []
     for text in group.text_parts:
         parts.append({"type": "text", "text": text + "\n"})
@@ -21,7 +22,8 @@ def _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
     return {"role": "user", "content": parts}
 
 
-def _tool_group_to_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
+def tool_group_to_openai_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
+    """Convert a ToolGroup to an OpenAI-compatible chat message."""
     merged_text = merge_reminder_text(
         group.tool_result.output or "<system-reminder>Tool ran without output or errors</system-reminder>",
         group.reminder_texts,
@@ -82,9 +84,9 @@ def convert_history_to_input(
     for group in parse_message_groups(history):
         match group:
             case UserGroup():
-                messages.append(_user_group_to_message(group))
+                messages.append(user_group_to_openai_message(group))
             case ToolGroup():
-                messages.append(_tool_group_to_message(group))
+                messages.append(tool_group_to_openai_message(group))
             case AssistantGroup():
                 messages.append(_assistant_group_to_message(group))
 

klaude_code/llm/openai_compatible/stream_processor.py

@@ -0,0 +1,82 @@
+"""Shared stream processing utilities for OpenAI-compatible clients.
+
+This module provides a reusable stream state manager that handles the common
+logic for accumulating and flushing reasoning, assistant content, and tool calls
+across different LLM providers (OpenAI-compatible, OpenRouter).
+"""
+
+from typing import Callable, Literal
+
+from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.protocol import model
+
+StreamStage = Literal["waiting", "reasoning", "assistant", "tool"]
+
+
+class StreamStateManager:
+    """Manages streaming state and provides flush operations for accumulated content.
+
+    This class encapsulates the common state management logic used by both
+    OpenAI-compatible and OpenRouter clients, reducing code duplication.
+    """
+
+    def __init__(
+        self,
+        param_model: str,
+        response_id: str | None = None,
+        reasoning_flusher: Callable[[], list[model.ConversationItem]] | None = None,
+    ):
+        self.param_model = param_model
+        self.response_id = response_id
+        self.stage: StreamStage = "waiting"
+        self.accumulated_reasoning: list[str] = []
+        self.accumulated_content: list[str] = []
+        self.accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
+        self.emitted_tool_start_indices: set[int] = set()
+        self._reasoning_flusher = reasoning_flusher
+
+    def set_response_id(self, response_id: str) -> None:
+        """Set the response ID once received from the stream."""
+        self.response_id = response_id
+        self.accumulated_tool_calls.response_id = response_id  # pyright: ignore[reportAttributeAccessIssue]
+
+    def flush_reasoning(self) -> list[model.ConversationItem]:
+        """Flush accumulated reasoning content and return items."""
+        if self._reasoning_flusher is not None:
+            return self._reasoning_flusher()
+        if not self.accumulated_reasoning:
+            return []
+        item = model.ReasoningTextItem(
+            content="".join(self.accumulated_reasoning),
+            response_id=self.response_id,
+            model=self.param_model,
+        )
+        self.accumulated_reasoning = []
+        return [item]
+
+    def flush_assistant(self) -> list[model.ConversationItem]:
+        """Flush accumulated assistant content and return items."""
+        if not self.accumulated_content:
+            return []
+        item = model.AssistantMessageItem(
+            content="".join(self.accumulated_content),
+            response_id=self.response_id,
+        )
+        self.accumulated_content = []
+        return [item]
+
+    def flush_tool_calls(self) -> list[model.ToolCallItem]:
+        """Flush accumulated tool calls and return items."""
+        items: list[model.ToolCallItem] = self.accumulated_tool_calls.get()
+        if items:
+            self.accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
+        return items
+
+    def flush_all(self) -> list[model.ConversationItem]:
+        """Flush all accumulated content in order: reasoning, assistant, tool calls."""
+        items: list[model.ConversationItem] = []
+        items.extend(self.flush_reasoning())
+        items.extend(self.flush_assistant())
+        if self.stage == "tool":
+            items.extend(self.flush_tool_calls())
+        return items

klaude_code/llm/openrouter/client.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncGenerator
-from typing import Literal, override
+from typing import override
 
 import httpx
 import openai
@@ -7,7 +7,7 @@ import openai
 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_tool_schema
-from klaude_code.llm.openai_compatible.tool_call_accumulator import BasicToolCallAccumulator, ToolCallAccumulatorABC
+from klaude_code.llm.openai_compatible.stream_processor import StreamStateManager
 from klaude_code.llm.openrouter.input import convert_history_to_input, is_claude_model
 from klaude_code.llm.openrouter.reasoning_handler import ReasoningDetail, ReasoningStreamHandler
 from klaude_code.llm.registry import register
@@ -38,7 +38,7 @@ class OpenRouterClient(LLMClientABC):
         messages = convert_history_to_input(param.input, param.system, param.model)
         tools = convert_tool_schema(param.tools)
 
-        metadata_tracker = MetadataTracker(cost_config=self._config.cost)
+        metadata_tracker = MetadataTracker(cost_config=self.get_llm_config().cost)
 
         extra_body: dict[str, object] = {
             "usage": {"include": True}  # To get the cache tokens at the end of the response
@@ -73,40 +73,18 @@ class OpenRouterClient(LLMClientABC):
             max_tokens=param.max_tokens,
             tools=tools,
             verbosity=param.verbosity,
-            extra_body=extra_body,  # pyright: ignore[reportUnknownArgumentType]
+            extra_body=extra_body,
             extra_headers=extra_headers,  # pyright: ignore[reportUnknownArgumentType]
         )
 
-        stage: Literal["waiting", "reasoning", "assistant", "tool", "done"] = "waiting"
-        response_id: str | None = None
-        accumulated_content: list[str] = []
-        accumulated_tool_calls: ToolCallAccumulatorABC = BasicToolCallAccumulator()
-        emitted_tool_start_indices: set[int] = set()
         reasoning_handler = ReasoningStreamHandler(
             param_model=str(param.model),
-            response_id=response_id,
+            response_id=None,
+        )
+        state = StreamStateManager(
+            param_model=str(param.model),
+            reasoning_flusher=reasoning_handler.flush,
         )
-
-        def flush_reasoning_items() -> list[model.ConversationItem]:
-            return reasoning_handler.flush()
-
-        def flush_assistant_items() -> list[model.ConversationItem]:
-            nonlocal accumulated_content
-            if len(accumulated_content) == 0:
-                return []
-            item = model.AssistantMessageItem(
-                content="".join(accumulated_content),
-                response_id=response_id,
-            )
-            accumulated_content = []
-            return [item]
-
-        def flush_tool_call_items() -> list[model.ToolCallItem]:
-            nonlocal accumulated_tool_calls
-            items: list[model.ToolCallItem] = accumulated_tool_calls.get()
-            if items:
-                accumulated_tool_calls.chunks_by_step = []  # pyright: ignore[reportAttributeAccessIssue]
-            return items
 
         try:
             async for event in await stream:
@@ -115,15 +93,14 @@ class OpenRouterClient(LLMClientABC):
                     style="blue",
                     debug_type=DebugType.LLM_STREAM,
                 )
-                if not response_id and event.id:
-                    response_id = event.id
-                    reasoning_handler.set_response_id(response_id)
-                    accumulated_tool_calls.response_id = response_id
-                    yield model.StartItem(response_id=response_id)
+                if not state.response_id and event.id:
+                    state.set_response_id(event.id)
+                    reasoning_handler.set_response_id(event.id)
+                    yield model.StartItem(response_id=event.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_tracker.set_usage(convert_usage(event.usage, param.context_limit))
+                    metadata_tracker.set_usage(convert_usage(event.usage, param.context_limit, param.max_tokens))
                 if event.model:
                     metadata_tracker.set_model_name(event.model)
                 if provider := getattr(event, "provider", None):
@@ -140,7 +117,7 @@ class OpenRouterClient(LLMClientABC):
                         try:
                             reasoning_detail = ReasoningDetail.model_validate(item)
                             metadata_tracker.record_token()
-                            stage = "reasoning"
+                            state.stage = "reasoning"
                             for conversation_item in reasoning_handler.on_detail(reasoning_detail):
                                 yield conversation_item
                         except Exception as e:
@@ -148,53 +125,46 @@ class OpenRouterClient(LLMClientABC):
 
                 # Assistant
                 if delta.content and (
-                    stage == "assistant" or delta.content.strip()
+                    state.stage == "assistant" or delta.content.strip()
                 ):  # Process all content in assistant stage, filter empty content in reasoning stage
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    stage = "assistant"
-                    accumulated_content.append(delta.content)
+                    state.stage = "assistant"
+                    state.accumulated_content.append(delta.content)
                     yield model.AssistantMessageDelta(
                         content=delta.content,
-                        response_id=response_id,
+                        response_id=state.response_id,
                     )
 
                 # Tool
                 if delta.tool_calls and len(delta.tool_calls) > 0:
                     metadata_tracker.record_token()
-                    if stage == "reasoning":
-                        for item in flush_reasoning_items():
+                    if state.stage == "reasoning":
+                        for item in state.flush_reasoning():
                             yield item
-                    elif stage == "assistant":
-                        for item in flush_assistant_items():
+                    elif state.stage == "assistant":
+                        for item in state.flush_assistant():
                             yield item
-                    stage = "tool"
+                    state.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)
+                        if tc.index not in state.emitted_tool_start_indices and tc.function and tc.function.name:
+                            state.emitted_tool_start_indices.add(tc.index)
                             yield model.ToolCallStartItem(
-                                response_id=response_id,
+                                response_id=state.response_id,
                                 call_id=tc.id or "",
                                 name=tc.function.name,
                             )
-                    accumulated_tool_calls.add(delta.tool_calls)
+                    state.accumulated_tool_calls.add(delta.tool_calls)
 
         except (openai.OpenAIError, httpx.HTTPError) as e:
             yield model.StreamErrorItem(error=f"{e.__class__.__name__} {str(e)}")
 
         # Finalize
-        for item in flush_reasoning_items():
-            yield item
-
-        for item in flush_assistant_items():
+        for item in state.flush_all():
             yield item
 
-        if stage == "tool":
-            for tool_call_item in flush_tool_call_items():
-                yield tool_call_item
-
-        metadata_tracker.set_response_id(response_id)
+        metadata_tracker.set_response_id(state.response_id)
         yield metadata_tracker.finalize()

klaude_code/llm/openrouter/input.py

@@ -7,9 +7,9 @@
 # pyright: reportGeneralTypeIssues=false
 
 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.llm.input_common import AssistantGroup, ToolGroup, UserGroup, parse_message_groups
+from klaude_code.llm.openai_compatible.input import tool_group_to_openai_message, user_group_to_openai_message
 from klaude_code.protocol import model
 
 
@@ -25,29 +25,6 @@ def is_gemini_model(model_name: str | None) -> bool:
     return model_name is not None and model_name.startswith("google/gemini")
 
 
-def _user_group_to_message(group: UserGroup) -> chat.ChatCompletionMessageParam:
-    parts: list[ChatCompletionContentPartParam] = []
-    for text in group.text_parts:
-        parts.append({"type": "text", "text": text + "\n"})
-    for image in group.images:
-        parts.append({"type": "image_url", "image_url": {"url": image.image_url.url}})
-    if not parts:
-        parts.append({"type": "text", "text": ""})
-    return {"role": "user", "content": parts}
-
-
-def _tool_group_to_message(group: ToolGroup) -> chat.ChatCompletionMessageParam:
-    merged_text = merge_reminder_text(
-        group.tool_result.output or "<system-reminder>Tool ran without output or errors</system-reminder>",
-        group.reminder_texts,
-    )
-    return {
-        "role": "tool",
-        "content": [{"type": "text", "text": merged_text}],
-        "tool_call_id": group.tool_result.call_id,
-    }
-
-
 def _assistant_group_to_message(group: AssistantGroup, model_name: str | None) -> chat.ChatCompletionMessageParam:
     assistant_message: dict[str, object] = {"role": "assistant"}
 
@@ -150,9 +127,9 @@ def convert_history_to_input(
     for group in parse_message_groups(history):
         match group:
             case UserGroup():
-                messages.append(_user_group_to_message(group))
+                messages.append(user_group_to_openai_message(group))
             case ToolGroup():
-                messages.append(_tool_group_to_message(group))
+                messages.append(tool_group_to_openai_message(group))
             case AssistantGroup():
                 messages.append(_assistant_group_to_message(group, model_name))
 

klaude_code/llm/registry.py

@@ -1,22 +1,48 @@
-from typing import Callable, TypeVar
+from typing import TYPE_CHECKING, Callable, TypeVar
 
-from klaude_code.llm.client import LLMClientABC
 from klaude_code.protocol import llm_param
 
-_REGISTRY: dict[llm_param.LLMClientProtocol, type[LLMClientABC]] = {}
+if TYPE_CHECKING:
+    from klaude_code.llm.client import LLMClientABC
 
-T = TypeVar("T", bound=LLMClientABC)
+_T = TypeVar("_T", bound=type["LLMClientABC"])
 
+# Track which protocols have been loaded
+_loaded_protocols: set[llm_param.LLMClientProtocol] = set()
+_REGISTRY: dict[llm_param.LLMClientProtocol, type["LLMClientABC"]] = {}
 
-def register(name: llm_param.LLMClientProtocol) -> Callable[[type[T]], type[T]]:
-    def _decorator(cls: type[T]) -> type[T]:
+
+def _load_protocol(protocol: llm_param.LLMClientProtocol) -> None:
+    """Load the module for a specific protocol on demand."""
+    if protocol in _loaded_protocols:
+        return
+    _loaded_protocols.add(protocol)
+
+    # Import only the needed module to trigger @register decorator
+    if protocol == llm_param.LLMClientProtocol.ANTHROPIC:
+        from . import anthropic as _  # noqa: F401
+    elif protocol == llm_param.LLMClientProtocol.CODEX:
+        from . import codex as _  # noqa: F401
+    elif protocol == llm_param.LLMClientProtocol.OPENAI:
+        from . import openai_compatible as _  # noqa: F401
+    elif protocol == llm_param.LLMClientProtocol.OPENROUTER:
+        from . import openrouter as _  # noqa: F401
+    elif protocol == llm_param.LLMClientProtocol.RESPONSES:
+        from . import responses as _  # noqa: F401
+
+
+def register(name: llm_param.LLMClientProtocol) -> Callable[[_T], _T]:
+    """Decorator to register an LLM client class for a protocol."""
+
+    def _decorator(cls: _T) -> _T:
         _REGISTRY[name] = cls
         return cls
 
     return _decorator
 
 
-def create_llm_client(config: llm_param.LLMConfigParameter) -> LLMClientABC:
+def create_llm_client(config: llm_param.LLMConfigParameter) -> "LLMClientABC":
+    _load_protocol(config.protocol)
     if config.protocol not in _REGISTRY:
         raise ValueError(f"Unknown LLMClient protocol: {config.protocol}")
     return _REGISTRY[config.protocol].create(config)