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

klaude_code/core/agent.py

@@ -5,34 +5,23 @@ from dataclasses import dataclass
 from typing import Protocol
 
 from klaude_code.core.prompt import get_system_prompt as load_system_prompt
-from klaude_code.core.reminders import (
-    Reminder,
-    get_main_agent_reminders,
-    get_sub_agent_reminders,
-    get_vanilla_reminders,
-)
-from klaude_code.core.sub_agent import get_sub_agent_profile
+from klaude_code.core.reminders import Reminder, load_agent_reminders
 from klaude_code.core.task import TaskExecutionContext, TaskExecutor
-from klaude_code.core.tool.tool_context import TodoContext
-from klaude_code.core.tool.tool_registry import (
-    get_main_agent_tools,
-    get_registry,
-    get_sub_agent_tools,
-    get_vanilla_tools,
-)
-from klaude_code.llm.client import LLMClientABC
-from klaude_code.protocol import events, llm_parameter, model, tools
+from klaude_code.core.tool import TodoContext, get_registry, load_agent_tools
+from klaude_code.llm import LLMClientABC
+from klaude_code.protocol import events, llm_param, model, tools
+from klaude_code.protocol.model import UserInputPayload
 from klaude_code.session import Session
 from klaude_code.trace import DebugType, log_debug
 
 
 @dataclass(frozen=True)
 class AgentProfile:
-    """Encapsulates the active LLM client plus prompt/tools/reminders."""
+    """Encapsulates the active LLM client plus prompts/tools/reminders."""
 
     llm_client: LLMClientABC
     system_prompt: str | None
-    tools: list[llm_parameter.ToolSchema]
+    tools: list[llm_param.ToolSchema]
     reminders: list[Reminder]
 
 
@@ -47,7 +36,7 @@ class ModelProfileProvider(Protocol):
 
 
 class DefaultModelProfileProvider(ModelProfileProvider):
-    """Default provider backed by global prompt/tool/reminder registries."""
+    """Default provider backed by global prompts/tool/reminder registries."""
 
     def build_profile(
         self,
@@ -55,23 +44,11 @@ class DefaultModelProfileProvider(ModelProfileProvider):
         sub_agent_type: tools.SubAgentType | None = None,
     ) -> AgentProfile:
         model_name = llm_client.model_name
-
-        if sub_agent_type is None:
-            prompt_key = "main"
-            tool_list = get_main_agent_tools(model_name)
-            reminders = get_main_agent_reminders(model_name)
-        else:
-            prompt_key = get_sub_agent_profile(sub_agent_type).name
-            tool_list = get_sub_agent_tools(model_name, sub_agent_type)
-            reminders = get_sub_agent_reminders(model_name)
-
-        system_prompt = load_system_prompt(model_name, prompt_key)
-
         return AgentProfile(
             llm_client=llm_client,
-            system_prompt=system_prompt,
-            tools=tool_list,
-            reminders=reminders,
+            system_prompt=load_system_prompt(model_name, sub_agent_type),
+            tools=load_agent_tools(model_name, sub_agent_type),
+            reminders=load_agent_reminders(model_name, sub_agent_type),
         )
 
 
@@ -83,11 +60,12 @@ class VanillaModelProfileProvider(ModelProfileProvider):
         llm_client: LLMClientABC,
         sub_agent_type: tools.SubAgentType | None = None,
     ) -> AgentProfile:
+        model_name = llm_client.model_name
         return AgentProfile(
             llm_client=llm_client,
             system_prompt=None,
-            tools=get_vanilla_tools(),
-            reminders=get_vanilla_reminders(),
+            tools=load_agent_tools(model_name, vanilla=True),
+            reminders=load_agent_reminders(model_name, vanilla=True),
         )
 
 
@@ -96,11 +74,8 @@ class Agent:
         self,
         session: Session,
         profile: AgentProfile,
-        *,
-        model_profile_provider: ModelProfileProvider | None = None,
     ):
         self.session: Session = session
-        self.model_profile_provider: ModelProfileProvider = model_profile_provider or DefaultModelProfileProvider()
         self.profile: AgentProfile | None = None
         # Active task executor, if any
         self._current_task: TaskExecutor | None = None
@@ -124,9 +99,13 @@ class Agent:
 
         # Record an interrupt marker in the session history
         self.session.append_history([model.InterruptItem()])
-        log_debug(f"Session {self.session.id} interrupted", style="yellow", debug_type=DebugType.EXECUTION)
+        log_debug(
+            f"Session {self.session.id} interrupted",
+            style="yellow",
+            debug_type=DebugType.EXECUTION,
+        )
 
-    async def run_task(self, user_input: str) -> AsyncGenerator[events.Event, None]:
+    async def run_task(self, user_input: UserInputPayload) -> AsyncGenerator[events.Event, None]:
         context = TaskExecutionContext(
             session_id=self.session.id,
             profile=self._require_profile(),
@@ -158,7 +137,9 @@ class Agent:
             return
 
         yield events.ReplayHistoryEvent(
-            events=list(self.session.get_history_item()), updated_at=self.session.updated_at, session_id=self.session.id
+            events=list(self.session.get_history_item()),
+            updated_at=self.session.updated_at,
+            session_id=self.session.id,
         )
 
     async def _process_reminder(self, reminder: Reminder) -> AsyncGenerator[events.DeveloperMessageEvent, None]:
@@ -172,14 +153,8 @@ class Agent:
         """Apply a fully constructed profile to the agent."""
 
         self.profile = profile
-        self.session.model_name = profile.llm_client.model_name
-
-    def build_model_profile(
-        self,
-        llm_client: LLMClientABC,
-        sub_agent_type: tools.SubAgentType | None = None,
-    ) -> AgentProfile:
-        return self.model_profile_provider.build_profile(llm_client, sub_agent_type)
+        if not self.session.model_name:
+            self.session.model_name = profile.llm_client.model_name
 
     def get_llm_client(self) -> LLMClientABC:
         return self._require_profile().llm_client