feat: refine

Author: wgzesg-bd

arkitect/core/component/agent/base_agent.py

@@ -16,9 +16,11 @@
 from typing import Any, AsyncIterable, Callable, Union
 
 from pydantic import BaseModel
+from volcenginesdkarkruntime import AsyncArk
 
-from arkitect.core.component.llm_event_stream.model import ContextInterruption, NewState
+from arkitect.core.component.llm_event_stream.model import ContextInterruption, State
 from arkitect.core.component.tool import MCPClient
+from arkitect.types.llm.model import ArkChatParameters
 from arkitect.types.responses.event import BaseEvent
 
 """
@@ -30,7 +32,7 @@ class PreAgentCallHook(abc.ABC):
     @abc.abstractmethod
     async def pre_agent_call(
         self,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -39,7 +41,7 @@ class PostAgentCallHook(abc.ABC):
     @abc.abstractmethod
     async def post_agent_call(
         self,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -51,6 +53,8 @@ class BaseAgent(abc.ABC, BaseModel):
     tools: list[Union[MCPClient | Callable]] = []
     sub_agents: list["BaseAgent"] = []
     instruction: str | None = None
+    parameters: ArkChatParameters | None = None
+    client: AsyncArk | None = None
 
     pre_agent_call_hook: PreAgentCallHook | None = None
     post_agent_call_hook: PostAgentCallHook | None = None
@@ -61,12 +65,10 @@ class BaseAgent(abc.ABC, BaseModel):
 
     # stream run step
     @abc.abstractmethod
-    async def _astream(
-        self, state: NewState, **kwargs: Any
-    ) -> AsyncIterable[BaseEvent]:
+    async def _astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
         pass
 
-    async def astream(self, state: NewState, **kwargs: Any) -> AsyncIterable[BaseEvent]:
+    async def astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
 
         if self.pre_agent_call_hook:
             async for event in self.pre_agent_call_hook.pre_agent_call(state):
@@ -81,9 +83,7 @@ async def astream(self, state: NewState, **kwargs: Any) -> AsyncIterable[BaseEve
             async for event in self.post_agent_call_hook.post_agent_call(state):
                 yield event
 
-    async def __call__(
-        self, state: NewState, **kwargs: Any
-    ) -> AsyncIterable[BaseEvent]:
+    async def __call__(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
         async for event in self.astream(state, **kwargs):
             yield event
 

arkitect/core/component/agent/default_agent.py

@@ -18,7 +18,14 @@
 
 from arkitect.core.component.agent import BaseAgent
 from arkitect.core.component.llm_event_stream.llm_event_stream import LLMEventStream
-from arkitect.core.component.llm_event_stream.model import NewState
+
+from arkitect.core.component.llm_event_stream.hooks import (
+    PostLLMCallHook,
+    PostToolCallHook,
+    PreLLMCallHook,
+    PreToolCallHook,
+)
+from arkitect.core.component.llm_event_stream.model import State
 from arkitect.types.responses.event import BaseEvent
 
 """
@@ -31,17 +38,26 @@ class DefaultAgent(BaseAgent):
         "arbitrary_types_allowed": True,
     }
 
+    pre_tool_call_hook: PreToolCallHook | None = None
+    post_tool_call_hook: PostToolCallHook | None = None
+    pre_llm_call_hook: PreLLMCallHook | None = None
+    post_llm_call_hook: PostLLMCallHook | None = None
+
     # stream run step
-    async def _astream(
-        self, state: NewState, **kwargs: Any
-    ) -> AsyncIterable[BaseEvent]:
+    async def _astream(self, state: State, **kwargs: Any) -> AsyncIterable[BaseEvent]:
         event_stream = LLMEventStream(
             model=self.model,
             agent_name=self.name,
             tools=self.tools,
             sub_agents=self.sub_agents,
             state=state,
             instruction=self.instruction,
+            pre_tool_call_hook=self.pre_tool_call_hook,
+            post_tool_call_hook=self.post_tool_call_hook,
+            pre_llm_call_hook=self.pre_llm_call_hook,
+            post_llm_call_hook=self.post_llm_call_hook,
+            parameters=self.parameters,
+            client=self.client,
         )
         await event_stream.init()
         resp_stream = await event_stream.completions.create(

arkitect/core/component/agent/parallel_agent.py

@@ -27,7 +27,7 @@
 from typing import AsyncIterable
 
 from arkitect.core.component.agent.base_agent import BaseAgent
-from arkitect.core.component.llm_event_stream.model import NewState
+from arkitect.core.component.llm_event_stream.model import State
 from arkitect.types.responses.event import BaseEvent
 
 """
@@ -77,7 +77,7 @@ async def _merge_agent_run(
 
 class ParallelAgent(BaseAgent):
     # stream run step
-    async def _astream(self, state: NewState, **kwargs) -> AsyncIterable[BaseEvent]:
+    async def _astream(self, state: State, **kwargs) -> 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/checkpoint.py

@@ -14,7 +14,7 @@
 import uuid
 from pydantic import BaseModel, ConfigDict, Field
 
-from arkitect.core.component.llm_event_stream.model import NewState
+from arkitect.core.component.llm_event_stream.model import State
 
 
 class Checkpoint(BaseModel):
@@ -39,7 +39,7 @@ class Checkpoint(BaseModel):
     """The name of the app."""
     user_id: str
     """The user id of the checkpoint."""
-    state: NewState = Field(default_factory=NewState)
+    state: State = Field(default_factory=State)
     """The state of the checkpoint."""
     last_update_time: float = 0.0
     """The last update time of the checkpoint."""

arkitect/core/component/checkpoint/in_memory_checkpoint_service.py

@@ -18,7 +18,7 @@
     BaseCheckpointService,
 )
 from arkitect.core.component.checkpoint.checkpoint import Checkpoint
-from arkitect.core.component.llm_event_stream.model import NewState
+from arkitect.core.component.llm_event_stream.model import State
 from arkitect.utils.common import Singleton
 
 
@@ -39,7 +39,7 @@ async def create_checkpoint(
                 id=checkpoint_id,
                 app_name=app_name,
                 user_id=user_id,
-                state=NewState(),
+                state=State(),
                 last_update_time=datetime.now().timestamp(),
                 create_time=datetime.now().timestamp(),
             )

arkitect/core/component/checkpoint/redis_checkpoint_service.py

@@ -19,7 +19,7 @@
     BaseCheckpointService,
 )
 from arkitect.core.component.checkpoint.checkpoint import Checkpoint
-from arkitect.core.component.llm_event_stream.model import NewState
+from arkitect.core.component.llm_event_stream.model import State
 from arkitect.utils.common import Singleton
 
 
@@ -48,7 +48,7 @@ async def create_checkpoint(
                 id=checkpoint_id,
                 app_name=app_name,
                 user_id=user_id,
-                state=NewState(),
+                state=State(),
                 last_update_time=datetime.now().timestamp(),
                 create_time=datetime.now().timestamp(),
             )

arkitect/core/component/llm_event_stream/chat_completion.py

@@ -23,15 +23,18 @@
 )
 
 from arkitect.core.component.tool.tool_pool import ToolPool
-from arkitect.types.llm.model import Message
+from arkitect.types.llm.model import Message, ArkChatParameters
 from arkitect.types.responses.event import BaseEvent, MessageEvent, StateUpdateEvent
 
-from .model import NewState
+from .model import State
 
 
 class _AsyncCompletions(AsyncCompletions):
-    def __init__(self, client: AsyncArk, state: NewState):
+    def __init__(
+        self, client: AsyncArk, state: State, parameters: ArkChatParameters | None
+    ):
         self._state = state
+        self.parameters = parameters
         super().__init__(client)
 
     async def create_event_stream(
@@ -41,11 +44,7 @@ async def create_event_stream(
         tool_pool: ToolPool | None = None,
         **kwargs: Dict[str, Any],
     ) -> AsyncIterable[BaseEvent]:
-        parameters = (
-            self._state.parameters.__dict__
-            if self._state.parameters is not None
-            else {}
-        )
+        parameters = self.parameters.__dict__ if self.parameters is not None else {}
         if tool_pool:
             tools = await tool_pool.list_tools()
             parameters["tools"] = [t.model_dump() for t in tools]
@@ -103,10 +102,16 @@ async def iterator() -> AsyncIterable[BaseEvent]:
 
 
 class _AsyncChat(AsyncChat):
-    def __init__(self, client: AsyncArk, state: NewState):
+    def __init__(
+        self,
+        client: AsyncArk,
+        state: State,
+        parameters: ArkChatParameters | None,
+    ):
         self._state = state
+        self.parameters = parameters
         super().__init__(client)
 
     @property
     def completions(self) -> _AsyncCompletions:
-        return _AsyncCompletions(self._client, self._state)
+        return _AsyncCompletions(self._client, self._state, self.parameters)

arkitect/core/component/llm_event_stream/context_completion.py

@@ -29,11 +29,11 @@
 from arkitect.core.component.tool.tool_pool import ToolPool
 from arkitect.types.responses.event import BaseEvent
 
-from .model import NewState
+from .model import State
 
 
 class _AsyncCompletions(AsyncCompletions):
-    def __init__(self, client: AsyncArk, state: NewState):
+    def __init__(self, client: AsyncArk, state: State):
         self._state = state
         super().__init__(client)
 
@@ -89,7 +89,7 @@ async def create_event_stream(
 
 
 class _AsyncContext(AsyncContext):
-    def __init__(self, client: AsyncArk, state: NewState):
+    def __init__(self, client: AsyncArk, state: State):
         self._state = state
         super().__init__(client)
 

arkitect/core/component/llm_event_stream/hooks.py

@@ -12,25 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import abc
-from typing import Any, Optional, Union
-
-from typing import Any, AsyncIterable, Optional
+from typing import Any, AsyncIterable, Optional, Union
 
 
 from arkitect.types.responses.event import (
     BaseEvent,
 )
 
-from .model import ContextInterruption, NewState
+from .model import ContextInterruption, State
 
-from .model import NewState
+from .model import State
 
 
 class HookInterruptException(Exception):
     def __init__(
         self,
         reason: str,
-        state: Optional[NewState] = None,
+        state: Optional[State] = None,
         details: Optional[Any] = None,
     ):
         self.reason = reason
@@ -44,7 +42,7 @@ async def pre_tool_call(
         self,
         name: str,
         arguments: str,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -57,7 +55,7 @@ async def post_tool_call(
         arguments: str,
         response: Any,
         exception: Optional[Exception],
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -66,7 +64,7 @@ class PreLLMCallHook(abc.ABC):
     @abc.abstractmethod
     async def pre_llm_call(
         self,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -75,7 +73,7 @@ class PostLLMCallHook(abc.ABC):
     @abc.abstractmethod
     async def post_llm_call(
         self,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         pass
 
@@ -93,7 +91,7 @@ async def pre_tool_call(
         self,
         name: str,
         arguments: str,
-        state: NewState,
+        state: State,
     ) -> AsyncIterable[BaseEvent | ContextInterruption]:
         if len(state.events) == 0:
             return