grasp_agents 0.1.5__py3-none-any.whl

grasp_agents/memory.py

@@ -0,0 +1,150 @@
+import logging
+from collections.abc import Iterator, Sequence
+from copy import deepcopy
+
+from .typing.io import LLMPrompt
+from .typing.message import Conversation, Message, SystemMessage
+
+logger = logging.getLogger(__name__)
+
+
+class MessageHistory:
+    def __init__(self, sys_prompt: LLMPrompt | None = None) -> None:
+        self._sys_prompt = sys_prompt
+        self._batched_conversations: list[Conversation]
+        self.reset()
+
+    @property
+    def sys_prompt(self) -> LLMPrompt | None:
+        return self._sys_prompt
+
+    def add_message_batch(self, message_batch: Sequence[Message]) -> None:
+        """
+        Adds a batch of messages to the current batched conversations.
+        This method verifies that the size of the input message batch matches
+        the expected batch size (self.batch_size).
+        If there is a mismatch, the method adjusts by duplicating either
+        the message or the conversation as necessary:
+
+        - If the message batch contains exactly one message and
+            self.batch_size > 1, the single message is duplicated to match
+            the batch size.
+        - If the message batch contains multiple messages but
+            self.batch_size == 1, the entire conversation is duplicated to
+            accommodate each message in the batch.
+        - If the message batch size does not match self.batch_size and none of
+            the above adjustments apply, a ValueError is raised.
+
+        Afterwards, each message in the batch is appended to its corresponding
+        conversation in the batched conversations.
+
+        Args:
+            message_batch: A sequence of Message objects
+                representing the batch of messages to be added. Must align with
+                or be adjusted to match the current batch size.
+
+        Raises:
+            ValueError: If the message batch size does not match the current
+                batch size and cannot be automatically adjusted.
+
+        """
+        message_batch_size = len(message_batch)
+
+        if message_batch_size == 1 and self.batch_size > 1:
+            logger.info(
+                "Message batch size is 1, current batch size is "
+                f"{self.batch_size}: duplicating the message to match the "
+                "current batch size"
+            )
+            message_batch = self._duplicate_message_to_current_batch_size(message_batch)
+            message_batch_size = self.batch_size
+        elif message_batch_size > 1 and self.batch_size == 1:
+            logger.info(
+                f"Message batch size is {len(message_batch)}, current batch "
+                "size is 1: duplicating the conversation to match the message "
+                "batch size"
+            )
+            self._duplicate_conversation_to_message_batch_size(message_batch_size)
+        elif message_batch_size != self.batch_size:
+            raise ValueError(
+                f"Message batch size {message_batch_size} does not match "
+                f"current batch size {self.batch_size}"
+            )
+
+        for batch_id in range(message_batch_size):
+            self._batched_conversations[batch_id].append(message_batch[batch_id])
+
+    def add_message_batches(self, message_batches: Sequence[Sequence[Message]]) -> None:
+        for message_batch in message_batches:
+            self.add_message_batch(message_batch)
+
+    def add_message(self, message: Message) -> None:
+        for conversation in self._batched_conversations:
+            conversation.append(message)
+
+    def add_messages(self, messages: Sequence[Message]) -> None:
+        for message in messages:
+            self.add_message(message)
+
+    def __len__(self) -> int:
+        return len(self._batched_conversations[0])
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(len={len(self)}; bs={self.batch_size})"
+
+    def __getitem__(self, idx: int) -> tuple[Message, ...]:
+        return tuple(conversation[idx] for conversation in self._batched_conversations)
+
+    def __iter__(self) -> Iterator[tuple[Message, ...]]:
+        for idx in range(len(self)):
+            yield tuple(
+                conversation[idx] for conversation in self._batched_conversations
+            )
+
+    def _duplicate_message_to_current_batch_size(
+        self, message_batch: Sequence[Message]
+    ) -> Sequence[Message]:
+        assert len(message_batch) == 1, (
+            "Message batch size must be 1 to duplicate to current batch size"
+        )
+
+        return [deepcopy(message_batch[0]) for _ in range(self.batch_size)]
+
+    def _duplicate_conversation_to_message_batch_size(
+        self, target_batch_size: int
+    ) -> None:
+        assert self.batch_size == 1, "Batch size must be 1 to duplicate conversation"
+        self._batched_conversations = [
+            deepcopy(self._batched_conversations[0]) for _ in range(target_batch_size)
+        ]
+
+    @property
+    def batched_conversations(self) -> list[Conversation]:
+        return self._batched_conversations
+
+    @property
+    def batch_size(self) -> int:
+        return len(self._batched_conversations)
+
+    def reset(
+        self, sys_prompt: LLMPrompt | None = None, *, batch_size: int = 1
+    ) -> None:
+        if sys_prompt is not None:
+            self._sys_prompt = sys_prompt
+
+        conv: Conversation
+        if self._sys_prompt is not None:
+            conv = [SystemMessage(content=self._sys_prompt)]
+        else:
+            conv = []
+
+        self._batched_conversations = [deepcopy(conv) for _ in range(batch_size)]
+
+    def erase(self) -> None:
+        self._batched_conversations = [[]]
+
+    # def get_batch(self, batch_id: int) -> list[Message]:
+    #     return self._batched_conversations[batch_id]
+
+    # def iterate_conversations(self) -> Iterator[list[Message]]:
+    #     return iter(self._batched_conversations)

grasp_agents/openai/__init__.py

@@ -0,0 +1,83 @@
+# pyright: reportUnusedImport=false
+
+from openai._streaming import (
+    AsyncStream as ChatCompletionAsyncStream,  # type: ignore[import] # noqa: PLC2701
+)
+from openai.types import CompletionUsage as ChatCompletionUsage
+from openai.types.chat.chat_completion import ChatCompletion
+from openai.types.chat.chat_completion_assistant_message_param import (
+    ChatCompletionAssistantMessageParam,
+)
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from openai.types.chat.chat_completion_content_part_image_param import (
+    ChatCompletionContentPartImageParam,
+)
+from openai.types.chat.chat_completion_content_part_image_param import (
+    ImageURL as ChatCompletionImageURL,
+)
+from openai.types.chat.chat_completion_content_part_param import (
+    ChatCompletionContentPartParam,
+)
+from openai.types.chat.chat_completion_content_part_text_param import (
+    ChatCompletionContentPartTextParam,
+)
+from openai.types.chat.chat_completion_developer_message_param import (
+    ChatCompletionDeveloperMessageParam,
+)
+from openai.types.chat.chat_completion_function_message_param import (
+    ChatCompletionFunctionMessageParam,
+)
+from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from openai.types.chat.chat_completion_message_param import (
+    ChatCompletionMessageParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    ChatCompletionMessageToolCallParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    Function as ChatCompletionToolCallFunction,
+)
+from openai.types.chat.chat_completion_named_tool_choice_param import (
+    ChatCompletionNamedToolChoiceParam,
+)
+from openai.types.chat.chat_completion_named_tool_choice_param import (
+    Function as ChatCompletionNamedToolChoiceFunction,
+)
+from openai.types.chat.chat_completion_prediction_content_param import (
+    ChatCompletionPredictionContentParam,
+)
+from openai.types.chat.chat_completion_stream_options_param import (
+    ChatCompletionStreamOptionsParam,
+)
+from openai.types.chat.chat_completion_system_message_param import (
+    ChatCompletionSystemMessageParam,
+)
+from openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam,
+)
+from openai.types.chat.chat_completion_tool_message_param import (
+    ChatCompletionToolMessageParam,
+)
+from openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam,
+)
+from openai.types.chat.chat_completion_user_message_param import (
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChatCompletion,
+    ParsedChatCompletionMessage,
+    ParsedChoice,
+)
+from openai.types.shared_params.function_definition import (
+    FunctionDefinition as ChatCompletionFunctionDefinition,
+)
+from openai.types.shared_params.response_format_json_object import (
+    ResponseFormatJSONObject,
+)
+from openai.types.shared_params.response_format_json_schema import (
+    ResponseFormatJSONSchema,
+)
+from openai.types.shared_params.response_format_text import (
+    ResponseFormatText,
+)

grasp_agents/openai/completion_converters.py

@@ -0,0 +1,49 @@
+from collections.abc import AsyncIterator
+
+from ..typing.completion import Completion, CompletionChoice, CompletionChunk
+from . import (
+    ChatCompletion,
+    ChatCompletionAsyncStream,  # type: ignore[import]
+    ChatCompletionChunk,
+)
+from .message_converters import from_api_assistant_message
+
+
+def from_api_completion(
+    api_completion: ChatCompletion, model_id: str | None = None
+) -> Completion:
+    choices: list[CompletionChoice] = []
+    # TODO: add custom error type
+    if api_completion.choices is None:  # type: ignore
+        raise ValueError(
+            f"Completion API error: {getattr(api_completion, 'error', None)}"
+        )
+    for api_choice in api_completion.choices:
+        # TODO: no way to assign individual message usages when len(choices) > 1
+        message = from_api_assistant_message(
+            api_choice.message, api_completion.usage, model_id=model_id
+        )
+        finish_reason = api_choice.finish_reason
+        choices.append(CompletionChoice(message=message, finish_reason=finish_reason))
+
+    return Completion(choices=choices, model_id=model_id)
+
+
+def to_api_completion(completion: Completion) -> ChatCompletion:
+    raise NotImplementedError
+
+
+def from_api_completion_chunk(
+    api_completion_chunk: ChatCompletionChunk, model_id: str | None = None
+) -> CompletionChunk:
+    delta = api_completion_chunk.choices[0].delta.content
+
+    return CompletionChunk(delta=delta, model_id=model_id)
+
+
+async def from_api_completion_chunk_iterator(
+    api_completion_chunk_iterator: ChatCompletionAsyncStream[ChatCompletionChunk],
+    model_id: str | None = None,
+) -> AsyncIterator[CompletionChunk]:
+    async for api_chunk in api_completion_chunk_iterator:
+        yield from_api_completion_chunk(api_chunk, model_id=model_id)

grasp_agents/openai/content_converters.py

@@ -0,0 +1,80 @@
+from collections.abc import Iterable
+
+from ..typing.content import (
+    Content,
+    ContentPart,
+    ContentPartImage,
+    ContentPartText,
+    ImageData,
+)
+from . import (
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartParam,
+    ChatCompletionContentPartTextParam,
+    ChatCompletionImageURL,
+)
+
+BASE64_PREFIX = "data:image/jpeg;base64,"
+
+
+def image_data_to_str(image_data: ImageData) -> str:
+    if image_data.type == "url":
+        return str(image_data.url)
+    if image_data.type == "base64":
+        return f"{BASE64_PREFIX}{image_data.base64}"
+    raise ValueError(f"Unsupported image data type: {image_data.type}")
+
+
+def from_api_content(
+    api_content: str | Iterable[ChatCompletionContentPartParam],
+) -> "Content":
+    if isinstance(api_content, str):
+        return Content(parts=[ContentPartText(data=api_content)])
+
+    content_parts: list[ContentPart] = []
+    for api_content_part in api_content:
+        content_part: ContentPart
+
+        if api_content_part["type"] == "text":
+            text_data = api_content_part["text"]
+            content_part = ContentPartText(data=text_data)
+
+        elif api_content_part["type"] == "image_url":
+            url = api_content_part["image_url"]["url"]
+            detail = api_content_part["image_url"].get("detail")
+            if url.startswith(BASE64_PREFIX):
+                image_data = ImageData.from_base64(
+                    base64_encoding=url.removeprefix(BASE64_PREFIX),
+                    detail=detail,
+                )
+            else:
+                image_data = ImageData.from_url(img_url=url, detail=detail)  # type: ignore
+            content_part = ContentPartImage(data=image_data)
+
+        content_parts.append(content_part)  # type: ignore
+
+    return Content(parts=content_parts)
+
+
+def to_api_content(
+    content: Content,
+) -> Iterable[ChatCompletionContentPartParam]:
+    api_content: list[ChatCompletionContentPartParam] = []
+    for content_part in content.parts:
+        api_content_part: ChatCompletionContentPartParam
+        if isinstance(content_part, ContentPartText):
+            api_content_part = ChatCompletionContentPartTextParam(
+                type="text",
+                text=content_part.data,
+            )
+        else:
+            api_content_part = ChatCompletionContentPartImageParam(
+                type="image_url",
+                image_url=ChatCompletionImageURL(
+                    url=image_data_to_str(content_part.data),
+                    detail=content_part.data.detail,
+                ),
+            )
+        api_content.append(api_content_part)
+
+    return api_content

grasp_agents/openai/converters.py

@@ -0,0 +1,170 @@
+from collections.abc import AsyncIterator, Iterable
+from typing import Any
+
+from pydantic import BaseModel
+
+from ..typing.completion import Completion, CompletionChunk
+from ..typing.content import Content
+from ..typing.converters import Converters
+from ..typing.message import (
+    AssistantMessage,
+    SystemMessage,
+    ToolMessage,
+    UserMessage,
+)
+from ..typing.tool import BaseTool, ToolChoice
+from . import (
+    ChatCompletion,
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionAsyncStream,  # type: ignore[import]
+    ChatCompletionChunk,
+    ChatCompletionContentPartParam,
+    # ChatCompletionDeveloperMessageParam,
+    ChatCompletionMessage,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionToolParam,
+    ChatCompletionUsage,
+    ChatCompletionUserMessageParam,
+)
+from .completion_converters import (
+    from_api_completion,
+    from_api_completion_chunk,
+    from_api_completion_chunk_iterator,
+    to_api_completion,
+)
+from .content_converters import from_api_content, to_api_content
+from .message_converters import (
+    from_api_assistant_message,
+    from_api_system_message,
+    from_api_tool_message,
+    from_api_user_message,
+    to_api_assistant_message,
+    to_api_system_message,
+    to_api_tool_message,
+    to_api_user_message,
+)
+from .tool_converters import to_api_tool, to_api_tool_choice
+
+
+class OpenAIConverters(Converters):
+    @staticmethod
+    def to_system_message(
+        system_message: SystemMessage, **kwargs: Any
+    ) -> ChatCompletionSystemMessageParam:
+        return to_api_system_message(system_message, **kwargs)
+
+    @staticmethod
+    def from_system_message(
+        raw_message: ChatCompletionSystemMessageParam,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> SystemMessage:
+        return from_api_system_message(raw_message, model_id=model_id, **kwargs)
+
+    @staticmethod
+    def to_user_message(
+        user_message: UserMessage, **kwargs: Any
+    ) -> ChatCompletionUserMessageParam:
+        return to_api_user_message(user_message, **kwargs)
+
+    @staticmethod
+    def from_user_message(
+        raw_message: ChatCompletionUserMessageParam,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> UserMessage:
+        return from_api_user_message(raw_message, model_id=model_id, **kwargs)
+
+    @staticmethod
+    def to_assistant_message(
+        assistant_message: AssistantMessage, **kwargs: Any
+    ) -> ChatCompletionAssistantMessageParam:
+        return to_api_assistant_message(assistant_message, **kwargs)
+
+    @staticmethod
+    def from_assistant_message(
+        raw_message: ChatCompletionMessage,
+        raw_usage: ChatCompletionUsage,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> AssistantMessage:
+        return from_api_assistant_message(
+            raw_message, raw_usage, model_id=model_id, **kwargs
+        )
+
+    @staticmethod
+    def to_tool_message(
+        tool_message: ToolMessage, **kwargs: Any
+    ) -> ChatCompletionToolMessageParam:
+        return to_api_tool_message(tool_message, **kwargs)
+
+    @staticmethod
+    def from_tool_message(
+        raw_message: ChatCompletionToolMessageParam,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> ToolMessage:
+        return from_api_tool_message(raw_message, model_id=model_id, **kwargs)
+
+    @staticmethod
+    def to_tool(
+        tool: BaseTool[BaseModel, BaseModel, Any], **kwargs: Any
+    ) -> ChatCompletionToolParam:
+        return to_api_tool(tool, **kwargs)
+
+    @staticmethod
+    def to_tool_choice(
+        tool_choice: ToolChoice, **kwargs: Any
+    ) -> ChatCompletionToolChoiceOptionParam:
+        return to_api_tool_choice(tool_choice, **kwargs)
+
+    @staticmethod
+    def to_content(
+        content: Content, **kwargs: Any
+    ) -> Iterable[ChatCompletionContentPartParam]:
+        return to_api_content(content, **kwargs)
+
+    @staticmethod
+    def from_content(
+        raw_content: str | Iterable[ChatCompletionContentPartParam],
+        **kwargs: Any,
+    ) -> Content:
+        return from_api_content(raw_content, **kwargs)
+
+    @staticmethod
+    def to_completion(completion: Completion, **kwargs: Any) -> ChatCompletion:
+        return to_api_completion(completion, **kwargs)
+
+    @staticmethod
+    def from_completion(
+        raw_completion: ChatCompletion,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> Completion:
+        return from_api_completion(raw_completion, model_id=model_id, **kwargs)
+
+    @staticmethod
+    def to_completion_chunk(
+        chunk: CompletionChunk, **kwargs: Any
+    ) -> ChatCompletionChunk:
+        raise NotImplementedError
+
+    @staticmethod
+    def from_completion_chunk(
+        raw_chunk: ChatCompletionChunk,
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> CompletionChunk:
+        return from_api_completion_chunk(raw_chunk, model_id=model_id, **kwargs)
+
+    @staticmethod
+    def from_completion_chunk_iterator(  # type: ignore[override]
+        raw_chunk_iterator: ChatCompletionAsyncStream[ChatCompletionChunk],
+        model_id: str | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[CompletionChunk]:
+        return from_api_completion_chunk_iterator(
+            raw_chunk_iterator, model_id=model_id, **kwargs
+        )

grasp_agents/openai/message_converters.py

@@ -0,0 +1,155 @@
+from typing import TypeAlias
+
+from ..typing.message import (
+    AssistantMessage,
+    SystemMessage,
+    ToolMessage,
+    Usage,
+    UserMessage,
+)
+from ..typing.tool import ToolCall
+from . import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionDeveloperMessageParam,
+    ChatCompletionFunctionMessageParam,
+    ChatCompletionMessage,
+    ChatCompletionMessageToolCallParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolCallFunction,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUsage,
+    ChatCompletionUserMessageParam,
+)
+from .content_converters import from_api_content, to_api_content
+
+OpenAIMessage: TypeAlias = (
+    ChatCompletionAssistantMessageParam
+    | ChatCompletionToolMessageParam
+    | ChatCompletionUserMessageParam
+    | ChatCompletionDeveloperMessageParam
+    | ChatCompletionSystemMessageParam
+    | ChatCompletionFunctionMessageParam
+)
+
+
+def from_api_user_message(
+    api_message: ChatCompletionUserMessageParam, model_id: str | None = None
+) -> UserMessage:
+    content = from_api_content(api_message["content"])
+
+    return UserMessage(content=content, model_id=model_id)
+
+
+def to_api_user_message(message: UserMessage) -> ChatCompletionUserMessageParam:
+    api_content = to_api_content(message.content)
+
+    return ChatCompletionUserMessageParam(role="user", content=api_content)
+
+
+def from_api_assistant_message(
+    api_message: ChatCompletionMessage,
+    api_usage: ChatCompletionUsage | None = None,
+    model_id: str | None = None,
+) -> AssistantMessage:
+    content = api_message.content or ""
+    assert isinstance(content, str), (
+        "Only string content is currently supported in assistant messages"
+    )
+
+    usage = None
+    if api_usage is not None:
+        reasoning_tokens = None
+        cached_tokens = None
+
+        if api_usage.completion_tokens_details is not None:
+            reasoning_tokens = api_usage.completion_tokens_details.reasoning_tokens
+        if api_usage.prompt_tokens_details is not None:
+            cached_tokens = api_usage.prompt_tokens_details.cached_tokens
+
+        input_tokens = api_usage.prompt_tokens - (cached_tokens or 0)
+        output_tokens = api_usage.completion_tokens - (reasoning_tokens or 0)
+
+        usage = Usage(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            reasoning_tokens=reasoning_tokens,
+            cached_tokens=cached_tokens,
+        )
+
+    tool_calls = None
+    if api_message.tool_calls is not None:
+        tool_calls = [
+            ToolCall(
+                id=tool_call.id,
+                tool_name=tool_call.function.name,
+                tool_arguments=tool_call.function.arguments,
+            )
+            for tool_call in api_message.tool_calls
+        ]
+
+    return AssistantMessage(
+        content=content,
+        usage=usage,
+        tool_calls=tool_calls,
+        refusal=api_message.refusal,
+        model_id=model_id,
+    )
+
+
+def to_api_assistant_message(
+    message: AssistantMessage,
+) -> ChatCompletionAssistantMessageParam:
+    api_tool_calls = None
+    if message.tool_calls is not None:
+        api_tool_calls = [
+            ChatCompletionMessageToolCallParam(
+                type="function",
+                id=tool_call.id,
+                function=ChatCompletionToolCallFunction(
+                    name=tool_call.tool_name,
+                    arguments=tool_call.tool_arguments,
+                ),
+            )
+            for tool_call in message.tool_calls
+        ]
+
+    return ChatCompletionAssistantMessageParam(
+        role="assistant",
+        content=message.content,
+        tool_calls=api_tool_calls,  # type: ignore
+        refusal=message.refusal,
+    )
+
+
+def from_api_system_message(
+    api_message: ChatCompletionSystemMessageParam,
+    model_id: str | None = None,
+) -> SystemMessage:
+    return SystemMessage(content=api_message["content"], model_id=model_id)  # type: ignore
+
+
+def to_api_system_message(
+    message: SystemMessage,
+) -> ChatCompletionSystemMessageParam:
+    return ChatCompletionSystemMessageParam(role="system", content=message.content)
+    # return ChatCompletionSystemMessageParam(
+    #     role="system", content=message.content
+    # )
+
+
+def from_api_tool_message(
+    api_message: ChatCompletionToolMessageParam, model_id: str | None = None
+) -> ToolMessage:
+    return ToolMessage(
+        content=api_message["content"],  # type: ignore
+        tool_call_id=api_message["tool_call_id"],
+        model_id=model_id,
+    )
+
+
+def to_api_tool_message(message: ToolMessage) -> ChatCompletionToolMessageParam:
+    return ChatCompletionToolMessageParam(
+        role="tool",
+        content=message.content,
+        tool_call_id=message.tool_call_id,
+    )