feat: format

Author: wgzesg-bd

arkitect/core/component/agent/base_agent.py

@@ -18,7 +18,7 @@
 from pydantic import BaseModel
 from volcenginesdkarkruntime import AsyncArk
 
-from arkitect.core.component.llm_event_stream.model import ContextInterruption, State
+from arkitect.core.component.llm_event_stream.model import State
 from arkitect.core.component.tool import MCPClient
 from arkitect.types.llm.model import ArkChatParameters
 from arkitect.types.responses.event import BaseEvent
@@ -33,17 +33,19 @@ class PreAgentCallHook(abc.ABC):
     async def pre_agent_call(
         self,
         state: State,
-    ) -> AsyncIterable[BaseEvent | ContextInterruption]:
-        pass
+    ) -> AsyncIterable[BaseEvent]:
+        return
+        yield
 
 
 class PostAgentCallHook(abc.ABC):
     @abc.abstractmethod
     async def post_agent_call(
         self,
         state: State,
-    ) -> AsyncIterable[BaseEvent | ContextInterruption]:
-        pass
+    ) -> AsyncIterable[BaseEvent]:
+        return
+        yield
 
 
 class BaseAgent(abc.ABC, BaseModel):
@@ -66,7 +68,8 @@ class BaseAgent(abc.ABC, BaseModel):
     # stream run step
     @abc.abstractmethod
     async def _astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
-        pass
+        return
+        yield
 
     async def astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
         if self.pre_agent_call_hook:

arkitect/core/component/agent/parallel_agent.py

@@ -12,19 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
-# Licensed under the 【火山方舟】原型应用软件自用许可协议
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#     https://www.volcengine.com/docs/82379/1433703
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import asyncio
-from typing import AsyncIterable
+from typing import Any, AsyncIterable
 
 from arkitect.core.component.agent.base_agent import BaseAgent
 from arkitect.core.component.llm_event_stream.model import State
@@ -50,7 +39,7 @@ async def _merge_agent_run(
         Event: The next event from the merged generator.
     """
     tasks = [
-        asyncio.create_task(events_for_one_agent.__anext__())
+        asyncio.create_task(events_for_one_agent.__anext__())  # type: ignore
         for events_for_one_agent in agent_runs
     ]
     pending_tasks = set(tasks)
@@ -66,7 +55,7 @@ async def _merge_agent_run(
                 # Find the generator that produced this event and move it on.
                 for i, original_task in enumerate(tasks):
                     if task == original_task:
-                        new_task = asyncio.create_task(agent_runs[i].__anext__())
+                        new_task = asyncio.create_task(agent_runs[i].__anext__())  # type: ignore
                         tasks[i] = new_task
                         pending_tasks.add(new_task)
                         break  # stop iterating once found
@@ -77,7 +66,7 @@ async def _merge_agent_run(
 
 class ParallelAgent(BaseAgent):
     # stream run step
-    async def _astream(self, state: State, **kwargs) -> AsyncIterable[BaseEvent]:
+    async def _astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
         agent_runs = [agent(state) for agent in self.sub_agents]
         async for event in _merge_agent_run(agent_runs):
             yield event

arkitect/core/component/checkpoint/base_checkpoint_service.py

@@ -13,27 +13,31 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
+from typing import Any
 
 from arkitect.core.component.checkpoint.checkpoint import Checkpoint
 
 
 class BaseCheckpointService(ABC):
     @abstractmethod
-    def create_checkpoint(
+    async def create_checkpoint(
         self,
         app_name: str,
         checkpoint_id: str,
         user_id: str,
         checkpoint: Checkpoint | None = None,
+        **kwargs: Any,
     ) -> Checkpoint:
         pass
 
     @abstractmethod
-    async def get_checkpoint(self, app_name: str, checkpoint_id: str) -> Checkpoint:
+    async def get_checkpoint(
+        self, app_name: str, checkpoint_id: str
+    ) -> Checkpoint | None:
         pass
 
     @abstractmethod
-    async def list_checkpoints(self, app_name: str) -> list[Checkpoint]:
+    async def list_checkpoints(self, app_name: str, **kwargs: Any) -> list[Checkpoint]:
         pass
 
     @abstractmethod

arkitect/core/component/checkpoint/in_memory_checkpoint_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from datetime import datetime
+from typing import Any
 
 from arkitect.core.component.checkpoint.base_checkpoint_service import (
     BaseCheckpointService,
@@ -23,7 +24,7 @@
 
 
 class InMemoryCheckpointService(BaseCheckpointService):
-    def __init__(self):
+    def __init__(self) -> None:
         # A map from app name to a map from user ID to a map from session ID to session.
         self.checkpoints: dict[str, dict[str, Checkpoint]] = {}
 
@@ -33,6 +34,7 @@ async def create_checkpoint(
         checkpoint_id: str,
         user_id: str,
         checkpoint: Checkpoint | None = None,
+        **kwargs: Any,
     ) -> Checkpoint:
         checkpoint = (
             Checkpoint(
@@ -53,10 +55,12 @@ async def create_checkpoint(
 
         return checkpoint
 
-    async def get_checkpoint(self, app_name: str, checkpoint_id: str) -> Checkpoint:
+    async def get_checkpoint(
+        self, app_name: str, checkpoint_id: str
+    ) -> Checkpoint | None:
         return self.checkpoints.get(app_name, {}).get(checkpoint_id, None)
 
-    async def list_checkpoints(self, app_name: str) -> list[Checkpoint]:
+    async def list_checkpoints(self, app_name: str, **kwargs: Any) -> list[Checkpoint]:
         return list(self.checkpoints.get(app_name, {}).values())
 
     async def update_checkpoint(

arkitect/core/component/checkpoint/redis_checkpoint_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from datetime import datetime
+from typing import Any
 
 from arkitect.core.client.redis import RedisClient
 from arkitect.core.component.checkpoint.base_checkpoint_service import (
@@ -42,6 +43,7 @@ async def create_checkpoint(
         checkpoint_id: str,
         user_id: str,
         checkpoint: Checkpoint | None = None,
+        **kwargs: Any,
     ) -> Checkpoint:
         checkpoint = (
             Checkpoint(
@@ -60,13 +62,19 @@ async def create_checkpoint(
             await self.redis_client.set(key, checkpoint.model_dump_json())
         return checkpoint
 
-    async def get_checkpoint(self, app_name: str, checkpoint_id: str) -> Checkpoint:
+    async def get_checkpoint(
+        self, app_name: str, checkpoint_id: str
+    ) -> Checkpoint | None:
         value = await self.redis_client.get(make_key(app_name, checkpoint_id))
         if value is None:
             return None
         return Checkpoint.model_validate_json(value)
 
-    async def list_checkpoints(self, app_name: str) -> list[Checkpoint]:
+    async def list_checkpoints(
+        self,
+        app_name: str,
+        **kwargs: Any,
+    ) -> list[Checkpoint]:
         keys, values = await self.redis_client.get_with_prefix(make_key(app_name, "*"))
         return [Checkpoint.model_validate_json(value) for value in values]
 

arkitect/core/component/context/context.py

@@ -95,11 +95,11 @@ async def handle_tool_call(self) -> bool:
                     tool_exception = e
 
                 self._ctx.state.messages.append(
-                    {
+                    {  # type: ignore
                         "role": "tool",
                         "tool_call_id": tool_call_param.get("id", ""),
                         "content": tool_resp if tool_resp else str(tool_exception),
-                    }
+                    }  # type: ignore
                 )
                 if self._ctx.post_tool_call_hook:
                     # post tool call hooks
@@ -327,7 +327,7 @@ async def tool_call_events() -> AsyncIterable[ToolChunk]:
                     tool_response=resp,
                 )
                 self._ctx.state.messages.append(
-                    {
+                    {  # type: ignore
                         "role": "tool",
                         "tool_call_id": tool_call.get("id", ""),
                         "content": resp,