grasp_agents 0.2.11__py3-none-any.whl → 0.3.2__py3-none-any.whl

grasp_agents/typing/completion_chunk.py

@@ -0,0 +1,173 @@
+import time
+from collections import defaultdict
+from collections.abc import Sequence
+from uuid import uuid4
+
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceLogprobs as CompletionChunkChoiceLogprobs,
+)
+from openai.types.chat.chat_completion_token_logprob import (
+    ChatCompletionTokenLogprob as CompletionTokenLogprob,
+)
+from pydantic import BaseModel, Field
+
+from .completion import (
+    Completion,
+    CompletionChoice,
+    CompletionChoiceLogprobs,
+    FinishReason,
+    Usage,
+)
+from .message import AssistantMessage, ToolCall
+
+
+class CompletionChunkDeltaToolCall(BaseModel):
+    id: str | None
+    index: int
+    tool_name: str | None
+    tool_arguments: str | None
+
+
+class CompletionChunkChoiceDelta(BaseModel):
+    content: str | None = None
+    refusal: str | None
+    role: str | None
+    tool_calls: list[CompletionChunkDeltaToolCall] | None
+
+
+class CompletionChunkChoice(BaseModel):
+    delta: CompletionChunkChoiceDelta
+    finish_reason: FinishReason | None
+    index: int
+    logprobs: CompletionChunkChoiceLogprobs | None = None
+
+
+class CompletionChunk(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4())[:8])
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str
+    name: str | None = None
+    system_fingerprint: str | None = None
+    choices: list[CompletionChunkChoice]
+    usage: Usage | None = None
+
+
+def combine_completion_chunks(chunks: list[CompletionChunk]) -> Completion:
+    if not chunks:
+        raise ValueError("Cannot combine an empty list of completion chunks.")
+
+    model_list = {chunk.model for chunk in chunks}
+    if len(model_list) > 1:
+        raise ValueError("All chunks must have the same model.")
+    model = model_list.pop()
+
+    name_list = {chunk.name for chunk in chunks}
+    if len(name_list) > 1:
+        raise ValueError("All chunks must have the same name.")
+    name = name_list.pop()
+
+    system_fingerprints_list = {chunk.system_fingerprint for chunk in chunks}
+    if len(system_fingerprints_list) > 1:
+        raise ValueError("All chunks must have the same system fingerprint.")
+    system_fingerprint = system_fingerprints_list.pop()
+
+    created_list = [chunk.created for chunk in chunks]
+    created = max(created_list)
+
+    # Usage is found in the last completion chunk if requested
+    usage = chunks[-1].usage
+
+    logp_contents_per_choice: defaultdict[int, list[CompletionTokenLogprob]] = (
+        defaultdict(list)
+    )
+    logp_refusals_per_choice: defaultdict[int, list[CompletionTokenLogprob]] = (
+        defaultdict(list)
+    )
+    logprobs_per_choice: defaultdict[int, CompletionChoiceLogprobs | None] = (
+        defaultdict(lambda: None)
+    )
+
+    finish_reasons_per_choice: defaultdict[int, FinishReason | None] = defaultdict(
+        lambda: None
+    )
+
+    contents_per_choice: defaultdict[int, str] = defaultdict(lambda: "")
+    refusals_per_choice: defaultdict[int, str] = defaultdict(lambda: "")
+
+    tool_calls_per_choice: defaultdict[
+        int, Sequence[CompletionChunkDeltaToolCall] | None
+    ] = defaultdict(lambda: None)
+
+    messages_per_choice: dict[int, AssistantMessage] = {}
+
+    for chunk in chunks:
+        for choice in chunk.choices:
+            index = choice.index
+
+            # Concatenate content and refusal tokens for each choice
+            contents_per_choice[index] += choice.delta.content or ""
+            refusals_per_choice[index] += choice.delta.refusal or ""
+
+            # Concatenate logprobs for content and refusal tokens for each choice
+            if choice.logprobs is not None:
+                logp_contents_per_choice[index].extend(choice.logprobs.content or [])
+                logp_refusals_per_choice[index].extend(choice.logprobs.refusal or [])
+
+            # Take the last finish reason for each choice
+            finish_reasons_per_choice[index] = choice.finish_reason
+
+            # Tool calls should be in the last chunk for each choice
+            tool_calls_per_choice[index] = choice.delta.tool_calls
+
+    for index in finish_reasons_per_choice:
+        tool_calls: list[ToolCall] = []
+        if tool_calls_per_choice[index] is not None:
+            for _tool_call in tool_calls_per_choice[index]:  # type: ignore
+                if (
+                    _tool_call.id is None
+                    or _tool_call.tool_name is None
+                    or _tool_call.tool_arguments is None
+                ):
+                    raise ValueError(
+                        "Completion chunk tool calls must have id, tool_name, "
+                        "and tool_arguments set."
+                    )
+                tool_calls.append(
+                    ToolCall(
+                        id=_tool_call.id,
+                        tool_name=_tool_call.tool_name,
+                        tool_arguments=_tool_call.tool_arguments,
+                    )
+                )
+
+        messages_per_choice[index] = AssistantMessage(
+            name=name,
+            content=contents_per_choice[index] or "<empty>",
+            refusal=(refusals_per_choice[index] or None),
+            tool_calls=(tool_calls or None),
+        )
+
+        if logp_contents_per_choice[index] or logp_refusals_per_choice[index]:
+            logprobs_per_choice[index] = CompletionChoiceLogprobs(
+                content=logp_contents_per_choice[index],
+                refusal=logp_refusals_per_choice[index],
+            )
+
+    choices = [
+        CompletionChoice(
+            index=index,
+            message=message,
+            finish_reason=finish_reasons_per_choice[index],
+            logprobs=logprobs_per_choice[index],
+        )
+        for index, message in messages_per_choice.items()
+    ]
+
+    return Completion(
+        model=model,
+        name=name,
+        created=created,
+        system_fingerprint=system_fingerprint,
+        choices=choices,
+        usage=usage,
+    )

grasp_agents/typing/converters.py

@@ -1,10 +1,10 @@
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator
 from typing import Any
 
 from pydantic import BaseModel
 
-from .completion import Completion, CompletionChunk
+from .completion import Completion, Usage
+from .completion_chunk import CompletionChunk
 from .content import Content
 from .message import AssistantMessage, Message, SystemMessage, ToolMessage, UserMessage
 from .tool import BaseTool, ToolChoice
@@ -38,9 +38,12 @@ class Converters(ABC):
 
     @staticmethod
     @abstractmethod
-    def from_assistant_message(
-        raw_message: Any, raw_usage: Any, **kwargs: Any
-    ) -> AssistantMessage:
+    def from_completion_usage(raw_usage: Any, **kwargs: Any) -> Usage:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def from_assistant_message(raw_message: Any, **kwargs: Any) -> AssistantMessage:
         pass
 
     @staticmethod
@@ -103,10 +106,3 @@ class Converters(ABC):
     @abstractmethod
     def from_completion_chunk(raw_chunk: Any, **kwargs: Any) -> CompletionChunk:
         pass
-
-    @staticmethod
-    @abstractmethod
-    def from_completion_chunk_iterator(
-        raw_chunk_iterator: AsyncIterator[Any], **kwargs: Any
-    ) -> AsyncIterator[CompletionChunk]:
-        pass

grasp_agents/typing/events.py

@@ -0,0 +1,86 @@
+import time
+from enum import StrEnum
+from typing import Any, Generic, Literal, TypeVar
+
+from pydantic import BaseModel, Field
+
+from ..packet import Packet
+from .completion import Completion
+from .completion_chunk import CompletionChunk
+from .message import AssistantMessage, SystemMessage, ToolCall, ToolMessage, UserMessage
+
+
+class EventSourceType(StrEnum):
+    LLM = "llm"
+    AGENT = "agent"
+    USER = "user"
+    TOOL = "tool"
+    PROCESSOR = "processor"
+
+
+class EventType(StrEnum):
+    SYS_MSG = "system_message"
+    USR_MSG = "user_message"
+    TOOL_MSG = "tool_message"
+    TOOL_CALL = "tool_call"
+    GEN_MSG = "gen_message"
+    COMP = "completion"
+    COMP_CHUNK = "completion_chunk"
+    PACKET = "packet"
+    PROC_OUT = "processor_output"
+
+
+_T = TypeVar("_T")
+
+
+class Event(BaseModel, Generic[_T], frozen=True):
+    type: EventType
+    source: EventSourceType
+    created: int = Field(default_factory=lambda: int(time.time()))
+    name: str | None = None
+    data: _T
+
+
+class CompletionEvent(Event[Completion], frozen=True):
+    type: Literal[EventType.COMP] = EventType.COMP
+    source: Literal[EventSourceType.LLM] = EventSourceType.LLM
+
+
+class CompletionChunkEvent(Event[CompletionChunk], frozen=True):
+    type: Literal[EventType.COMP_CHUNK] = EventType.COMP_CHUNK
+    source: Literal[EventSourceType.LLM] = EventSourceType.LLM
+
+
+class GenMessageEvent(Event[AssistantMessage], frozen=True):
+    type: Literal[EventType.GEN_MSG] = EventType.GEN_MSG
+    source: Literal[EventSourceType.LLM] = EventSourceType.LLM
+
+
+class ToolCallEvent(Event[ToolCall], frozen=True):
+    type: Literal[EventType.TOOL_CALL] = EventType.TOOL_CALL
+    source: Literal[EventSourceType.AGENT] = EventSourceType.AGENT
+
+
+class ToolMessageEvent(Event[ToolMessage], frozen=True):
+    type: Literal[EventType.TOOL_MSG] = EventType.TOOL_MSG
+    source: Literal[EventSourceType.TOOL] = EventSourceType.TOOL
+
+
+class UserMessageEvent(Event[UserMessage], frozen=True):
+    type: Literal[EventType.USR_MSG] = EventType.USR_MSG
+    source: Literal[EventSourceType.USER] = EventSourceType.USER
+
+
+class SystemMessageEvent(Event[SystemMessage], frozen=True):
+    type: Literal[EventType.SYS_MSG] = EventType.SYS_MSG
+    source: Literal[EventSourceType.AGENT] = EventSourceType.AGENT
+
+
+class PacketEvent(Event[Packet[Any]], frozen=True):
+    type: Literal[EventType.PACKET] = EventType.PACKET
+    source: Literal[EventSourceType.PROCESSOR] = EventSourceType.PROCESSOR
+
+
+class ProcOutputEvent(Event[Any], frozen=True):
+    type: Literal[EventType.PROC_OUT] = EventType.PROC_OUT
+    source: Literal[EventSourceType.PROCESSOR] = EventSourceType.PROCESSOR

grasp_agents/typing/io.py

@@ -1,25 +1,16 @@
-from collections.abc import Mapping
 from typing import TypeAlias, TypeVar
 
 from pydantic import BaseModel
 
-from .content import ImageData
-
-AgentID: TypeAlias = str
-
-
-class AgentState(BaseModel):
-    pass
+ProcName: TypeAlias = str
 
 
 class LLMPromptArgs(BaseModel):
     pass
 
 
-InT = TypeVar("InT", contravariant=True)  # noqa: PLC0105
-OutT = TypeVar("OutT", covariant=True)  # noqa: PLC0105
-StateT = TypeVar("StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
+InT_contra = TypeVar("InT_contra", contravariant=True)
+OutT_co = TypeVar("OutT_co", covariant=True)
+MemT_co = TypeVar("MemT_co", covariant=True)
 
 LLMPrompt: TypeAlias = str
-LLMFormattedSystemArgs: TypeAlias = Mapping[str, str | int | bool]
-LLMFormattedArgs: TypeAlias = Mapping[str, str | int | bool | ImageData]

grasp_agents/typing/message.py

@@ -4,7 +4,7 @@ from enum import StrEnum
 from typing import Annotated, Any, Literal, TypeAlias
 from uuid import uuid4
 
-from pydantic import BaseModel, Field, NonNegativeFloat, NonNegativeInt
+from pydantic import BaseModel, Field
 from pydantic.json import pydantic_encoder
 
 from .content import Content, ImageData
@@ -18,85 +18,48 @@ class Role(StrEnum):
     TOOL = "tool"
 
 
-class Usage(BaseModel):
-    input_tokens: NonNegativeInt = 0
-    output_tokens: NonNegativeInt = 0
-    reasoning_tokens: NonNegativeInt | None = None
-    cached_tokens: NonNegativeInt | None = None
-    cost: NonNegativeFloat | None = None
-
-    def __add__(self, add_usage: "Usage") -> "Usage":
-        input_tokens = self.input_tokens + add_usage.input_tokens
-        output_tokens = self.output_tokens + add_usage.output_tokens
-        if self.reasoning_tokens is not None or add_usage.reasoning_tokens is not None:
-            reasoning_tokens = (self.reasoning_tokens or 0) + (
-                add_usage.reasoning_tokens or 0
-            )
-        else:
-            reasoning_tokens = None
-
-        if self.cached_tokens is not None or add_usage.cached_tokens is not None:
-            cached_tokens = (self.cached_tokens or 0) + (add_usage.cached_tokens or 0)
-        else:
-            cached_tokens = None
-
-        cost = (
-            (self.cost or 0.0) + add_usage.cost
-            if (add_usage.cost is not None)
-            else None
-        )
-        return Usage(
-            input_tokens=input_tokens,
-            output_tokens=output_tokens,
-            reasoning_tokens=reasoning_tokens,
-            cached_tokens=cached_tokens,
-            cost=cost,
-        )
-
-
 class MessageBase(BaseModel):
-    message_id: Hashable = Field(default_factory=lambda: str(uuid4())[:8])
-    model_id: str | None = None
+    id: Hashable = Field(default_factory=lambda: str(uuid4())[:8])
+    name: str | None = None
 
 
 class AssistantMessage(MessageBase):
     role: Literal[Role.ASSISTANT] = Role.ASSISTANT
     content: str | None
-    usage: Usage | None = None
     tool_calls: Sequence[ToolCall] | None = None
     refusal: str | None = None
 
 
 class UserMessage(MessageBase):
     role: Literal[Role.USER] = Role.USER
-    content: Content
+    content: Content | str
 
     @classmethod
-    def from_text(cls, text: str, model_id: str | None = None) -> "UserMessage":
-        return cls(content=Content.from_text(text), model_id=model_id)
+    def from_text(cls, text: str, name: str | None = None) -> "UserMessage":
+        return cls(content=Content.from_text(text), name=name)
 
     @classmethod
     def from_formatted_prompt(
         cls,
         prompt_template: str,
         prompt_args: Mapping[str, str | int | bool | ImageData] | None = None,
-        model_id: str | None = None,
+        name: str | None = None,
     ) -> "UserMessage":
         content = Content.from_formatted_prompt(
             prompt_template=prompt_template, prompt_args=prompt_args
         )
 
-        return cls(content=content, model_id=model_id)
+        return cls(content=content, name=name)
 
     @classmethod
     def from_content_parts(
         cls,
         content_parts: Sequence[str | ImageData],
-        model_id: str | None = None,
+        name: str | None = None,
     ) -> "UserMessage":
         content = Content.from_content_parts(content_parts)
 
-        return cls(content=content, model_id=model_id)
+        return cls(content=content, name=name)
 
 
 class SystemMessage(MessageBase):
@@ -114,13 +77,12 @@ class ToolMessage(MessageBase):
         cls,
         tool_output: Any,
         tool_call: ToolCall,
-        model_id: str | None = None,
         indent: int = 2,
     ) -> "ToolMessage":
         return cls(
             content=json.dumps(tool_output, default=pydantic_encoder, indent=indent),
             tool_call_id=tool_call.id,
-            model_id=model_id,
+            name=tool_call.tool_name,
         )
 
 
@@ -129,4 +91,4 @@ Message = Annotated[
     Field(discriminator="role"),
 ]
 
-Conversation: TypeAlias = list[Message]
+Messages: TypeAlias = list[Message]

grasp_agents/typing/tool.py

@@ -1,23 +1,31 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Generic,
+    Literal,
+    TypeAlias,
+    TypeVar,
+)
 
 from pydantic import BaseModel, PrivateAttr, TypeAdapter
 
 from ..generics_utils import AutoInstanceAttributesMixin
 
 if TYPE_CHECKING:
-    from ..run_context import CtxT, RunContextWrapper
+    from ..run_context import CtxT, RunContext
 else:
     CtxT = TypeVar("CtxT")
 
-    class RunContextWrapper(Generic[CtxT]):
-        """Runtime placeholder so RunContextWrapper[CtxT] works"""
+    class RunContext(Generic[CtxT]):
+        """Runtime placeholder so RunContext[CtxT] works"""
 
 
-_ToolInT = TypeVar("_ToolInT", bound=BaseModel, contravariant=True)  # noqa: PLC0105
-_ToolOutT = TypeVar("_ToolOutT", covariant=True)  # noqa: PLC0105
+_InT_contra = TypeVar("_InT_contra", bound=BaseModel, contravariant=True)
+_OutT_co = TypeVar("_OutT_co", covariant=True)
 
 
 class ToolCall(BaseModel):
@@ -27,45 +35,63 @@ class ToolCall(BaseModel):
 
 
 class BaseTool(
-    AutoInstanceAttributesMixin, BaseModel, ABC, Generic[_ToolInT, _ToolOutT, CtxT]
+    AutoInstanceAttributesMixin,
+    BaseModel,
+    ABC,
+    Generic[_InT_contra, _OutT_co, CtxT],
 ):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
-        0: "_in_schema",
-        1: "_out_schema",
+        0: "_in_type",
+        1: "_out_type",
     }
 
     name: str
     description: str
 
-    _in_schema: type[_ToolInT] = PrivateAttr()
-    _out_schema: type[_ToolOutT] = PrivateAttr()
+    _in_type: type[_InT_contra] = PrivateAttr()
+    _out_type: type[_OutT_co] = PrivateAttr()
+
+    # _in_type_adapter: TypeAdapter[_InT_contra] = PrivateAttr()
+    # _out_type_adapter: TypeAdapter[_OutT_co] = PrivateAttr()
 
-    # Supported by OpenAI API
-    strict: bool | None = None
+    # def model_post_init(self, context: Any) -> None:
+    #     self._in_type_adapter = TypeAdapter(self._in_type)
+    #     self._out_type_adapter = TypeAdapter(self._out_type)
 
     @property
-    def in_schema(self) -> type[_ToolInT]:  # type: ignore[reportInvalidTypeVarUse]
+    def in_type(self) -> type[_InT_contra]:  # type: ignore[reportInvalidTypeVarUse]
         # Exposing the type of a contravariant variable only, should be type safe
-        return self._in_schema
+        return self._in_type
 
     @property
-    def out_schema(self) -> type[_ToolOutT]:
-        return self._out_schema
+    def out_type(self) -> type[_OutT_co]:
+        return self._out_type
+
+    # @property
+    # def in_type_adapter(self) -> TypeAdapter[_InT_contra]:
+    #     return self._in_type_adapter
+
+    # @property
+    # def out_type_adapter(self) -> TypeAdapter[_OutT_co]:
+    #     return self._out_type_adapter
 
     @abstractmethod
     async def run(
-        self, inp: _ToolInT, ctx: RunContextWrapper[CtxT] | None = None
-    ) -> _ToolOutT:
+        self, inp: _InT_contra, ctx: RunContext[CtxT] | None = None
+    ) -> _OutT_co:
         pass
 
     async def __call__(
-        self, ctx: RunContextWrapper[CtxT] | None = None, **kwargs: Any
-    ) -> _ToolOutT:
-        result = await self.run(self._in_schema(**kwargs), ctx=ctx)
+        self, ctx: RunContext[CtxT] | None = None, **kwargs: Any
+    ) -> _OutT_co:
+        input_args = TypeAdapter(self._in_type).validate_python(kwargs)
+        output = await self.run(input_args, ctx=ctx)
 
-        return TypeAdapter(self._out_schema).validate_python(result)
+        return TypeAdapter(self._out_type).validate_python(output)
 
 
-ToolChoice: TypeAlias = (
-    Literal["none", "auto", "required"] | BaseTool[BaseModel, Any, Any]
-)
+class NamedToolChoice(BaseModel):
+    name: str
+
+
+ToolChoice: TypeAlias = Literal["none", "auto", "required"] | NamedToolChoice

grasp_agents/usage_tracker.py

@@ -7,7 +7,7 @@ import yaml
 from pydantic import BaseModel, Field
 from termcolor import colored
 
-from .typing.message import AssistantMessage, Message, Usage
+from .typing.completion import Completion, Usage
 
 logger = logging.getLogger(__name__)
 
@@ -58,20 +58,20 @@ class UsageTracker(BaseModel):
         usage.cost = (input_cost + output_cost + reasoning_cost + cached_cost) / 1e6
 
     def update(
-        self, messages: Sequence[Message], model_name: str | None = None
+        self, completions: Sequence[Completion], model_name: str | None = None
     ) -> None:
         if model_name is not None and self.costs_dict is not None:
             model_costs_dict = self.costs_dict.get(model_name.split(":", 1)[-1])
         else:
             model_costs_dict = None
 
-        for message in messages:
-            if isinstance(message, AssistantMessage) and message.usage is not None:
+        for completion in completions:
+            if completion.usage is not None:
                 if model_costs_dict is not None:
                     self._add_cost_to_usage(
-                        usage=message.usage, model_costs_dict=model_costs_dict
+                        usage=completion.usage, model_costs_dict=model_costs_dict
                     )
-                self.total_usage += message.usage
+                self.total_usage += completion.usage
 
     def reset(self) -> None:
         self.total_usage = Usage()

grasp_agents/utils.py

@@ -3,7 +3,7 @@ import asyncio
 import json
 import re
 from collections.abc import Coroutine, Mapping
-from datetime import datetime
+from datetime import UTC, datetime
 from logging import getLogger
 from pathlib import Path
 from typing import Any, TypeVar, overload
@@ -126,7 +126,7 @@ def read_contents_from_file(
             return Path(file_path).read_bytes()
         return Path(file_path).read_text()
     except FileNotFoundError:
-        logger.error(f"File {file_path} not found.")
+        logger.exception(f"File {file_path} not found.")
         return ""
 
 
@@ -157,4 +157,4 @@ async def asyncio_gather_with_pbar(
 
 
 def get_timestamp() -> str:
-    return datetime.now().strftime("%Y%m%d_%H%M%S")
+    return datetime.now(UTC).strftime("%Y%m%d_%H%M%S")