grasp_agents 0.5.6__py3-none-any.whl → 0.5.8__py3-none-any.whl

grasp_agents/processors/parallel_processor.py

@@ -1,16 +1,13 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator, Sequence
-from typing import Any, ClassVar, Generic,  cast
-
+from typing import Any, ClassVar, Generic, cast
 
 from ..errors import PacketRoutingError
 from ..memory import MemT
 from ..packet import Packet
 from ..run_context import CtxT, RunContext
-from ..typing.events import (
-    Event, ProcPacketOutputEvent, ProcPayloadOutputEvent
-)
+from ..typing.events import Event, ProcPacketOutputEvent, ProcPayloadOutputEvent
 from ..typing.io import InT, OutT
 from ..utils import stream_concurrent
 from .base_processor import BaseProcessor, with_retry, with_retry_stream
@@ -18,7 +15,9 @@ from .base_processor import BaseProcessor, with_retry, with_retry_stream
 logger = logging.getLogger(__name__)
 
 
-class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, CtxT]):
+class ParallelProcessor(
+    BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT, MemT, CtxT]
+):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
@@ -33,7 +32,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> OutT:
-        return cast(OutT, in_args)
+        return cast("OutT", in_args)
 
     async def _process_stream(
         self,
@@ -44,7 +43,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        output = cast(OutT, in_args)
+        output = cast("OutT", in_args)
         yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
 
     def _validate_parallel_recipients(
@@ -59,7 +58,7 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
                 message="Parallel runs must return the same recipients "
                 f"[proc_name={self.name}; call_id={call_id}]",
             )
-    
+
     @with_retry
     async def _run_single(
         self,
@@ -86,7 +85,6 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
 
         return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
 
-
     async def _run_parallel(
         self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT] | None = None
     ) -> Packet[OutT]:
@@ -125,8 +123,10 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         )
 
         if val_in_args and len(val_in_args) > 1:
-            return await self._run_parallel(in_args=val_in_args, call_id=call_id, ctx=ctx)
-        
+            return await self._run_parallel(
+                in_args=val_in_args, call_id=call_id, ctx=ctx
+            )
+
         return await self._run_single(
             chat_inputs=chat_inputs,
             in_args=val_in_args[0] if val_in_args else None,
@@ -231,7 +231,9 @@ class ParallelProcessor(BaseProcessor[InT, OutT, MemT, CtxT], Generic[InT, OutT,
         )
 
         if val_in_args and len(val_in_args) > 1:
-            stream = self._run_parallel_stream(in_args=val_in_args, call_id=call_id, ctx=ctx)
+            stream = self._run_parallel_stream(
+                in_args=val_in_args, call_id=call_id, ctx=ctx
+            )
         else:
             stream = self._run_single_stream(
                 chat_inputs=chat_inputs,

grasp_agents/runner.py

@@ -6,7 +6,7 @@ from typing import Any, Generic
 from .errors import RunnerError
 from .packet import Packet, StartPacket
 from .packet_pool import END_PROC_NAME, PacketPool
-from .processors.processor import Processor
+from .processors.base_processor import BaseProcessor
 from .run_context import CtxT, RunContext
 from .typing.events import Event, ProcPacketOutputEvent, RunResultEvent
 from .typing.io import OutT
@@ -17,8 +17,8 @@ logger = logging.getLogger(__name__)
 class Runner(Generic[OutT, CtxT]):
     def __init__(
         self,
-        entry_proc: Processor[Any, Any, Any, CtxT],
-        procs: Sequence[Processor[Any, Any, Any, CtxT]],
+        entry_proc: BaseProcessor[Any, Any, Any, CtxT],
+        procs: Sequence[BaseProcessor[Any, Any, Any, CtxT]],
         ctx: RunContext[CtxT] | None = None,
     ) -> None:
         if entry_proc not in procs:
@@ -34,7 +34,6 @@ class Runner(Generic[OutT, CtxT]):
         self._entry_proc = entry_proc
         self._procs = procs
         self._ctx = ctx or RunContext[CtxT]()
-        self._packet_pool: PacketPool[CtxT] = PacketPool()
 
     @property
     def ctx(self) -> RunContext[CtxT]:
@@ -49,9 +48,10 @@ class Runner(Generic[OutT, CtxT]):
 
     async def _packet_handler(
         self,
-        proc: Processor[Any, Any, Any, CtxT],
-        pool: PacketPool[CtxT],
         packet: Packet[Any],
+        *,
+        proc: BaseProcessor[Any, Any, Any, CtxT],
+        pool: PacketPool,
         ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
@@ -72,9 +72,10 @@ class Runner(Generic[OutT, CtxT]):
 
     async def _packet_handler_stream(
         self,
-        proc: Processor[Any, Any, Any, CtxT],
-        pool: PacketPool[CtxT],
         packet: Packet[Any],
+        *,
+        proc: BaseProcessor[Any, Any, Any, CtxT],
+        pool: PacketPool,
         ctx: RunContext[CtxT],
         **run_kwargs: Any,
     ) -> None:
@@ -99,18 +100,18 @@ class Runner(Generic[OutT, CtxT]):
 
         await pool.post(out_packet)
 
-    async def run(
-        self,
-        chat_input: Any = "start",
-        **run_args: Any,
-    ) -> Packet[OutT]:
-        async with PacketPool[CtxT]() as pool:
+    async def run(self, chat_input: Any = "start", **run_args: Any) -> Packet[OutT]:
+        async with PacketPool() as pool:
             for proc in self._procs:
                 pool.register_packet_handler(
                     proc_name=proc.name,
-                    handler=partial(self._packet_handler, proc, pool),
-                    ctx=self._ctx,
-                    **run_args,
+                    handler=partial(
+                        self._packet_handler,
+                        proc=proc,
+                        pool=pool,
+                        ctx=self._ctx,
+                        **run_args,
+                    ),
                 )
             await pool.post(
                 StartPacket[Any](
@@ -120,17 +121,19 @@ class Runner(Generic[OutT, CtxT]):
             return await pool.final_result()
 
     async def run_stream(
-        self,
-        chat_input: Any = "start",
-        **run_args: Any,
+        self, chat_input: Any = "start", **run_args: Any
     ) -> AsyncIterator[Event[Any]]:
-        async with PacketPool[CtxT]() as pool:
+        async with PacketPool() as pool:
             for proc in self._procs:
                 pool.register_packet_handler(
                     proc_name=proc.name,
-                    handler=partial(self._packet_handler_stream, proc, pool),
-                    ctx=self._ctx,
-                    **run_args,
+                    handler=partial(
+                        self._packet_handler_stream,
+                        proc=proc,
+                        pool=pool,
+                        ctx=self._ctx,
+                        **run_args,
+                    ),
                 )
             await pool.post(
                 StartPacket[Any](

grasp_agents/typing/completion_chunk.py

@@ -15,7 +15,7 @@ from openai.types.chat.chat_completion_chunk import (
 from openai.types.chat.chat_completion_token_logprob import (
     ChatCompletionTokenLogprob as OpenAITokenLogprob,
 )
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError, field_validator
 
 from ..errors import CombineCompletionChunksError
 from .completion import Completion, CompletionChoice, FinishReason, Usage
@@ -38,13 +38,77 @@ class CompletionChunkDeltaToolCall(BaseModel):
 class CompletionChunkChoiceDelta(BaseModel):
     content: str | None = None
     refusal: str | None = None
-    role: Role | None
-    tool_calls: list[CompletionChunkDeltaToolCall] | None
+    role: Role | None = None
+    tool_calls: list[CompletionChunkDeltaToolCall] | None = None
     reasoning_content: str | None = None
     thinking_blocks: list[ThinkingBlock | RedactedThinkingBlock] | None = None
     annotations: list[LiteLLMAnnotation] | None = None
     provider_specific_fields: dict[str, Any] | None = None
 
+    @property
+    def thinking_delta(self) -> "CompletionChunkChoiceDelta | None":
+        return (
+            CompletionChunkChoiceDelta(
+                reasoning_content=self.reasoning_content,
+                thinking_blocks=self.thinking_blocks,
+                role=self.role,
+                provider_specific_fields=self.provider_specific_fields,
+            )
+            if self.reasoning_content or self.thinking_blocks
+            else None
+        )
+
+    @property
+    def tool_call_deltas(self) -> "list[CompletionChunkChoiceDelta] | None":
+        return (
+            [
+                CompletionChunkChoiceDelta(
+                    tool_calls=[tool_call],
+                    role=self.role,
+                    provider_specific_fields=self.provider_specific_fields,
+                )
+                for tool_call in self.tool_calls
+            ]
+            if self.tool_calls
+            else None
+        )
+
+    @property
+    def response_delta(self) -> "CompletionChunkChoiceDelta | None":
+        return (
+            CompletionChunkChoiceDelta(
+                content=self.content,
+                role=self.role,
+                provider_specific_fields=self.provider_specific_fields,
+            )
+            if self.content
+            else None
+        )
+
+    @property
+    def annotations_delta(self) -> "CompletionChunkChoiceDelta | None":
+        return (
+            CompletionChunkChoiceDelta(
+                annotations=self.annotations,
+                role=self.role,
+                provider_specific_fields=self.provider_specific_fields,
+            )
+            if self.annotations
+            else None
+        )
+
+    @property
+    def refusal_delta(self) -> "CompletionChunkChoiceDelta | None":
+        return (
+            CompletionChunkChoiceDelta(
+                refusal=self.refusal,
+                role=self.role,
+                provider_specific_fields=self.provider_specific_fields,
+            )
+            if self.refusal
+            else None
+        )
+
 
 class CompletionChunkChoice(BaseModel):
     delta: CompletionChunkChoiceDelta
@@ -66,6 +130,241 @@ class CompletionChunk(BaseModel):
     response_ms: float | None = None
     hidden_params: dict[str, Any] | None = None
 
+    def split_into_specialized(
+        self,
+    ) -> "list[CompletionChunk]":
+        if len(self.choices) != 1:
+            raise ValidationError(
+                "CompletionChunk must have exactly one choice for specialization."
+            )
+        delta = self.choices[0].delta
+
+        specialized_chunks: list[CompletionChunk] = []
+
+        thinking_delta = delta.thinking_delta
+        tool_call_deltas = delta.tool_call_deltas
+        response_delta = delta.response_delta
+        annotations_delta = delta.annotations_delta
+        refusal_delta = delta.refusal_delta
+
+        if thinking_delta is not None:
+            new_choice = self.choices[0].model_copy(update={"delta": thinking_delta})
+            new_chunk = self.model_copy(update={"choices": [new_choice]})
+            specialized_chunks.append(
+                ThinkingChunk.model_validate(new_chunk.model_dump())
+            )
+
+        if tool_call_deltas:
+            for delta_tool_call in tool_call_deltas:
+                new_choice = self.choices[0].model_copy(
+                    update={"delta": delta_tool_call}
+                )
+                new_chunk = self.model_copy(update={"choices": [new_choice]})
+                specialized_chunks.append(
+                    ToolCallChunk.model_validate(new_chunk.model_dump())
+                )
+
+        if response_delta is not None:
+            new_choice = self.choices[0].model_copy(update={"delta": response_delta})
+            new_chunk = self.model_copy(update={"choices": [new_choice]})
+            specialized_chunks.append(
+                ResponseChunk.model_validate(new_chunk.model_dump())
+            )
+
+        if annotations_delta is not None:
+            new_choice = self.choices[0].model_copy(update={"delta": annotations_delta})
+            new_chunk = self.model_copy(update={"choices": [new_choice]})
+            specialized_chunks.append(
+                AnnotationsChunk.model_validate(new_chunk.model_dump())
+            )
+
+        if refusal_delta is not None:
+            new_choice = self.choices[0].model_copy(update={"delta": refusal_delta})
+            new_chunk = self.model_copy(update={"choices": [new_choice]})
+            specialized_chunks.append(
+                RefusalChunk.model_validate(new_chunk.model_dump())
+            )
+
+        return specialized_chunks
+
+
+class ResponseChunk(CompletionChunk):
+    @field_validator("choices")
+    @classmethod
+    def validate_response_chunk(
+        cls, choices: list[CompletionChunkChoice]
+    ) -> list[CompletionChunkChoice]:
+        if len(choices) != 1:
+            raise ValidationError("ResponseChunk must have exactly one choice.")
+
+        delta = choices[0].delta
+
+        if not delta.content:
+            raise ValidationError("ResponseChunk must have content in deltas.")
+
+        if (
+            delta.reasoning_content is not None
+            or delta.thinking_blocks is not None
+            or delta.tool_calls is not None
+            or delta.refusal is not None
+            or delta.annotations is not None
+        ):
+            raise ValidationError(
+                "ResponseChunk should not have reasoning content, thinking blocks, "
+                "tool calls, refusal, or annotations in deltas."
+            )
+
+        return choices
+
+    @property
+    def response(self) -> str:
+        assert self.choices[0].delta.content
+        return self.choices[0].delta.content
+
+
+class ThinkingChunk(CompletionChunk):
+    @field_validator("choices")
+    @classmethod
+    def validate_thinking_chunk(
+        cls, choices: list[CompletionChunkChoice]
+    ) -> list[CompletionChunkChoice]:
+        if len(choices) != 1:
+            raise ValidationError("ThinkingChunk must have exactly one choice.")
+
+        delta = choices[0].delta
+
+        if not (delta.thinking_blocks or delta.reasoning_content):
+            raise ValidationError(
+                "ThinkingChunk must have reasoning content or "
+                "at least one thinking block."
+            )
+        if (
+            delta.content is not None
+            or delta.tool_calls is not None
+            or delta.refusal is not None
+            or delta.annotations is not None
+        ):
+            raise ValidationError(
+                "ThinkingChunk should not have content, tool calls, "
+                "refusal, or annotations in deltas."
+            )
+
+        return choices
+
+    @property
+    def thinking(self) -> str | list[ThinkingBlock | RedactedThinkingBlock]:
+        delta = self.choices[0].delta
+        if delta.reasoning_content:
+            return delta.reasoning_content
+        if delta.thinking_blocks:
+            return delta.thinking_blocks
+        raise ValueError("ThinkingChunk has no reasoning_content or thinking_blocks")
+
+
+class ToolCallChunk(CompletionChunk):
+    @field_validator("choices")
+    @classmethod
+    def validate_tool_call_chunk(
+        cls, choices: list[CompletionChunkChoice]
+    ) -> list[CompletionChunkChoice]:
+        if len(choices) != 1:
+            raise ValidationError("ToolCallChunk must have exactly one choice.")
+
+        delta = choices[0].delta
+
+        if not delta.tool_calls:
+            raise ValidationError("ToolCallChunk must have tool calls in deltas.")
+        if len(delta.tool_calls) != 1:
+            raise ValidationError(
+                "ToolCallChunk must have exactly one tool call in deltas."
+            )
+
+        if (
+            delta.reasoning_content is not None
+            or delta.thinking_blocks is not None
+            or delta.content is not None
+            or delta.refusal is not None
+            or delta.annotations is not None
+        ):
+            raise ValidationError(
+                "ToolCallChunk should not have reasoning content, thinking blocks, "
+                "content, refusal, or annotations in deltas."
+            )
+
+        return choices
+
+    @property
+    def tool_call(self) -> CompletionChunkDeltaToolCall:
+        assert self.choices[0].delta.tool_calls is not None
+        return self.choices[0].delta.tool_calls[0]
+
+
+class AnnotationsChunk(CompletionChunk):
+    @field_validator("choices")
+    @classmethod
+    def validate_annotations_chunk(
+        cls, choices: list[CompletionChunkChoice]
+    ) -> list[CompletionChunkChoice]:
+        if len(choices) != 1:
+            raise ValidationError("AnnotationsChunk must have exactly one choice.")
+
+        delta = choices[0].delta
+
+        if not delta.annotations:
+            raise ValidationError("AnnotationsChunk must have annotations in deltas.")
+
+        if (
+            delta.reasoning_content is not None
+            or delta.thinking_blocks is not None
+            or delta.content is not None
+            or delta.tool_calls is not None
+            or delta.refusal is not None
+        ):
+            raise ValidationError(
+                "AnnotationsChunk should not have reasoning content, thinking blocks, "
+                "content, tool calls, or refusal in deltas."
+            )
+
+        return choices
+
+    @property
+    def annotations(self) -> list[LiteLLMAnnotation]:
+        assert self.choices[0].delta.annotations is not None
+        return self.choices[0].delta.annotations
+
+
+class RefusalChunk(CompletionChunk):
+    @field_validator("choices")
+    @classmethod
+    def validate_refusal_chunk(
+        cls, choices: list[CompletionChunkChoice]
+    ) -> list[CompletionChunkChoice]:
+        if len(choices) != 1:
+            raise ValidationError("RefusalChunk must have exactly one choice.")
+
+        delta = choices[0].delta
+
+        if not delta.refusal:
+            raise ValidationError("RefusalChunk must have refusal in deltas.")
+
+        if (
+            delta.reasoning_content is not None
+            or delta.thinking_blocks is not None
+            or delta.content is not None
+            or delta.tool_calls is not None
+            or delta.annotations is not None
+        ):
+            raise ValidationError(
+                "RefusalChunk should not have reasoning content, thinking blocks, "
+                "content, tool calls, or annotations in deltas."
+            )
+
+        return choices
+
+    @property
+    def refusal(self) -> str | None:
+        return self.choices[0].delta.refusal
+
 
 def combine_completion_chunks(chunks: list[CompletionChunk]) -> Completion:
     if not chunks: