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

grasp_agents/message_history.py

@@ -0,0 +1,140 @@
+import logging
+from collections.abc import Iterator, Sequence
+from copy import deepcopy
+
+from .typing.io import LLMPrompt
+from .typing.message import Message, Messages, SystemMessage
+
+logger = logging.getLogger(__name__)
+
+
+class MessageHistory:
+    def __init__(self, sys_prompt: LLMPrompt | None = None) -> None:
+        self._sys_prompt = sys_prompt
+        self._conversations: list[Messages]
+        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_target_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._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._conversations:
+            conversation.append(message)
+
+    def add_message_list(self, message_list: Sequence[Message]) -> None:
+        for message in message_list:
+            self.add_message(message)
+
+    def __len__(self) -> int:
+        return len(self._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._conversations)
+
+    def __iter__(self) -> Iterator[tuple[Message, ...]]:
+        for idx in range(len(self)):
+            yield tuple(conversation[idx] for conversation in self._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_target_batch_size(
+        self, target_batch_size: int
+    ) -> None:
+        assert self.batch_size == 1, "Batch size must be 1 to duplicate conversation"
+        self._conversations = [
+            deepcopy(self._conversations[0]) for _ in range(target_batch_size)
+        ]
+
+    @property
+    def conversations(self) -> list[Messages]:
+        return self._conversations
+
+    @property
+    def batch_size(self) -> int:
+        return len(self._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: Messages = []
+        if self._sys_prompt is not None:
+            conv.append(SystemMessage(content=self._sys_prompt))
+
+        self._conversations = [deepcopy(conv) for _ in range(batch_size)]
+
+    def erase(self) -> None:
+        self._conversations = [[]]

grasp_agents/openai/__init__.py

@@ -1,87 +1,105 @@
 # pyright: reportUnusedImport=false
 
 from openai._streaming import (
-    AsyncStream as ChatCompletionAsyncStream,  # type: ignore[import] # noqa: PLC2701
+    AsyncStream as OpenAIAsyncStream,  # type: ignore[import] # noqa: PLC2701
+)
+from openai.types import CompletionUsage as OpenAICompletionUsage
+from openai.types.chat.chat_completion import ChatCompletion as OpenAICompletion
+from openai.types.chat.chat_completion import (
+    ChoiceLogprobs as OpenAIChoiceLogprobs,
 )
-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,
+    ChatCompletionAssistantMessageParam as OpenAIAssistantMessageParam,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk as OpenAICompletionChunk,
+)
+from openai.types.chat.chat_completion_chunk import (
+    Choice as CompletionChunkChoice,
+)
+from openai.types.chat.chat_completion_chunk import (
+    Choice as OpenAIChunkChoice,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDelta as CompletionChunkChoiceDelta,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDelta as OpenAIChunkChoiceDelta,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDeltaToolCall as OpenAIChunkChoiceDeltaToolCall,
 )
-from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
 from openai.types.chat.chat_completion_content_part_image_param import (
-    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartImageParam as OpenAIContentPartImageParam,
 )
 from openai.types.chat.chat_completion_content_part_image_param import (
-    ImageURL as ChatCompletionImageURL,
+    ImageURL as OpenAIImageURL,
 )
 from openai.types.chat.chat_completion_content_part_param import (
-    ChatCompletionContentPartParam,
+    ChatCompletionContentPartParam as OpenAIContentPartParam,
 )
 from openai.types.chat.chat_completion_content_part_text_param import (
-    ChatCompletionContentPartTextParam,
+    ChatCompletionContentPartTextParam as OpenAIContentPartTextParam,
 )
 from openai.types.chat.chat_completion_developer_message_param import (
-    ChatCompletionDeveloperMessageParam,
+    ChatCompletionDeveloperMessageParam as OpenAIDeveloperMessageParam,
 )
 from openai.types.chat.chat_completion_function_message_param import (
-    ChatCompletionFunctionMessageParam,
+    ChatCompletionFunctionMessageParam as OpenAIFunctionMessageParam,
+)
+from openai.types.chat.chat_completion_message import (
+    ChatCompletionMessage as OpenAIMessage,
 )
-from openai.types.chat.chat_completion_message import ChatCompletionMessage
 from openai.types.chat.chat_completion_message_param import (
-    ChatCompletionMessageParam,
+    ChatCompletionMessageParam as OpenAIMessageParam,
 )
 from openai.types.chat.chat_completion_message_tool_call_param import (
-    ChatCompletionMessageToolCallParam,
+    ChatCompletionMessageToolCallParam as OpenAIToolCallParam,
 )
 from openai.types.chat.chat_completion_message_tool_call_param import (
-    Function as ChatCompletionToolCallFunction,
+    Function as OpenAIToolCallFunction,
 )
 from openai.types.chat.chat_completion_named_tool_choice_param import (
-    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionNamedToolChoiceParam as OpenAINamedToolChoiceParam,
 )
 from openai.types.chat.chat_completion_named_tool_choice_param import (
-    Function as ChatCompletionNamedToolChoiceFunction,
+    Function as OpenAINamedToolChoiceFunction,
 )
 from openai.types.chat.chat_completion_prediction_content_param import (
-    ChatCompletionPredictionContentParam,
+    ChatCompletionPredictionContentParam as OpenAIPredictionContentParam,
 )
 from openai.types.chat.chat_completion_stream_options_param import (
-    ChatCompletionStreamOptionsParam,
+    ChatCompletionStreamOptionsParam as OpenAIStreamOptionsParam,
 )
 from openai.types.chat.chat_completion_system_message_param import (
-    ChatCompletionSystemMessageParam,
+    ChatCompletionSystemMessageParam as OpenAISystemMessageParam,
 )
 from openai.types.chat.chat_completion_tool_choice_option_param import (
-    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolChoiceOptionParam as OpenAIToolChoiceOptionParam,
 )
 from openai.types.chat.chat_completion_tool_message_param import (
-    ChatCompletionToolMessageParam,
+    ChatCompletionToolMessageParam as OpenAIToolMessageParam,
 )
 from openai.types.chat.chat_completion_tool_param import (
-    ChatCompletionToolParam,
+    ChatCompletionToolParam as OpenAIToolParam,
 )
 from openai.types.chat.chat_completion_user_message_param import (
-    ChatCompletionUserMessageParam,
+    ChatCompletionUserMessageParam as OpenAIUserMessageParam,
 )
 from openai.types.chat.parsed_chat_completion import (
-    ParsedChatCompletion,
-    ParsedChatCompletionMessage,
-    ParsedChoice,
-)
-from openai.types.shared_params.function_definition import (
-    FunctionDefinition as ChatCompletionFunctionDefinition,
+    ParsedChatCompletion as OpenAIParsedCompletion,
 )
-from openai.types.shared_params.response_format_json_object import (
-    ResponseFormatJSONObject,
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChatCompletionMessage as OpenAIParsedMessage,
 )
-from openai.types.shared_params.response_format_json_schema import (
-    ResponseFormatJSONSchema,
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChoice as OpenAIParsedChoice,
 )
-from openai.types.shared_params.response_format_text import (
-    ResponseFormatText,
+from openai.types.shared_params.function_definition import (
+    FunctionDefinition as OpenAIFunctionDefinition,
 )
 
+# from openai.beta.types.chat.
 from .openai_llm import OpenAILLM, OpenAILLMSettings
 
 __all__ = ["OpenAILLM", "OpenAILLMSettings"]

grasp_agents/openai/completion_chunk_converters.py

@@ -0,0 +1,78 @@
+from ..typing.completion_chunk import (
+    CompletionChunk,
+    CompletionChunkChoice,
+    CompletionChunkChoiceDelta,
+    CompletionChunkDeltaToolCall,
+)
+from . import OpenAICompletionChunk  # type: ignore[import]
+from .completion_converters import from_api_completion_usage
+
+
+def from_api_completion_chunk(
+    api_completion_chunk: OpenAICompletionChunk, name: str | None = None
+) -> CompletionChunk:
+    if api_completion_chunk.choices is None:  # type: ignore
+        raise RuntimeError(
+            f"Completion chunk API error: "
+            f"{getattr(api_completion_chunk, 'error', None)}"
+        )
+
+    choices: list[CompletionChunkChoice] = []
+
+    for api_choice in api_completion_chunk.choices:
+        index = api_choice.index
+        finish_reason = api_choice.finish_reason
+
+        if api_choice.delta is None:  # type: ignore
+            raise RuntimeError(
+                "API returned None for delta content in completion chunk "
+                f"with finish_reason: {finish_reason}."
+            )
+        # if api_choice.delta.content is None:
+        #     raise RuntimeError(
+        #         "API returned None for delta content in completion chunk "
+        #         f"with finish_reason: {finish_reason}."
+        #     )
+
+        delta = CompletionChunkChoiceDelta(
+            content=api_choice.delta.content,
+            refusal=api_choice.delta.refusal,
+            role=api_choice.delta.role,
+            tool_calls=[
+                CompletionChunkDeltaToolCall(
+                    id=tool_call.id,
+                    index=tool_call.index,
+                    tool_name=tool_call.function.name,
+                    tool_arguments=tool_call.function.arguments,
+                )
+                for tool_call in (api_choice.delta.tool_calls or [])
+                if tool_call.function
+            ],
+        )
+
+        choice = CompletionChunkChoice(
+            index=index,
+            delta=delta,
+            finish_reason=finish_reason,
+            # OpenAI logprobs have identical, but separately defined, types for
+            # completion chunks and completions
+            logprobs=api_choice.logprobs,
+        )
+
+        choices.append(choice)
+
+    usage = (
+        from_api_completion_usage(api_completion_chunk.usage)
+        if api_completion_chunk.usage
+        else None
+    )
+
+    return CompletionChunk(
+        id=api_completion_chunk.id,
+        model=api_completion_chunk.model,
+        name=name,
+        created=api_completion_chunk.created,
+        system_fingerprint=api_completion_chunk.system_fingerprint,
+        choices=choices,
+        usage=usage,
+    )

grasp_agents/openai/completion_converters.py

@@ -1,18 +1,34 @@
-from collections.abc import AsyncIterator
-
-from ..typing.completion import Completion, CompletionChoice, CompletionChunk
-from . import (
-    ChatCompletion,
-    ChatCompletionAsyncStream,  # type: ignore[import]
-    ChatCompletionChunk,
-)
+from ..typing.completion import Completion, CompletionChoice, Usage
+from . import OpenAICompletion, OpenAICompletionUsage
 from .message_converters import from_api_assistant_message
 
 
+def from_api_completion_usage(api_usage: OpenAICompletionUsage) -> Usage:
+    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)
+
+    return Usage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        reasoning_tokens=reasoning_tokens,
+        cached_tokens=cached_tokens,
+    )
+
+
 def from_api_completion(
-    api_completion: ChatCompletion, model_id: str | None = None
+    api_completion: OpenAICompletion, name: str | None = None
 ) -> Completion:
     choices: list[CompletionChoice] = []
+    usage: Usage | None = None
+
     if api_completion.choices is None:  # type: ignore
         # Some providers return None for the choices when there is an error
         # TODO: add custom error types
@@ -20,36 +36,43 @@ def from_api_completion(
             f"Completion API error: {getattr(api_completion, 'error', None)}"
         )
     for api_choice in api_completion.choices:
-        # TODO: currently no way to assign individual message usages when len(choices) > 1
+        index = api_choice.index
         finish_reason = api_choice.finish_reason
+
         # Some providers return None for the message when finish_reason is other than "stop"
         if api_choice.message is None:  # type: ignore
             raise RuntimeError(
                 f"API returned None for message with finish_reason: {finish_reason}"
             )
-        message = from_api_assistant_message(
-            api_choice.message, api_completion.usage, model_id=model_id
-        )
-        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
+        # usage = from_api_completion_usage(api_usage) if api_usage else None
 
+        usage = (
+            from_api_completion_usage(api_completion.usage)
+            if api_completion.usage
+            else None
+        )
+        message = from_api_assistant_message(api_choice.message, name=name)
 
-def from_api_completion_chunk(
-    api_completion_chunk: ChatCompletionChunk, model_id: str | None = None
-) -> CompletionChunk:
-    delta = api_completion_chunk.choices[0].delta.content
+        choices.append(
+            CompletionChoice(
+                index=index,
+                message=message,
+                finish_reason=finish_reason,
+                logprobs=api_choice.logprobs,
+            )
+        )
 
-    return CompletionChunk(delta=delta, model_id=model_id)
+    return Completion(
+        id=api_completion.id,
+        created=api_completion.created,
+        system_fingerprint=api_completion.system_fingerprint,
+        usage=usage,
+        choices=choices,
+        model=api_completion.model,
+        name=name,
+    )
 
 
-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)
+def to_api_completion(completion: Completion) -> OpenAICompletion:
+    raise NotImplementedError

grasp_agents/openai/content_converters.py

@@ -8,10 +8,10 @@ from ..typing.content import (
     ImageData,
 )
 from . import (
-    ChatCompletionContentPartImageParam,
-    ChatCompletionContentPartParam,
-    ChatCompletionContentPartTextParam,
-    ChatCompletionImageURL,
+    OpenAIContentPartImageParam,
+    OpenAIContentPartParam,
+    OpenAIContentPartTextParam,
+    OpenAIImageURL,
 )
 
 BASE64_PREFIX = "data:image/jpeg;base64,"
@@ -26,7 +26,7 @@ def image_data_to_str(image_data: ImageData) -> str:
 
 
 def from_api_content(
-    api_content: str | Iterable[ChatCompletionContentPartParam],
+    api_content: str | Iterable[OpenAIContentPartParam],
 ) -> "Content":
     if isinstance(api_content, str):
         return Content(parts=[ContentPartText(data=api_content)])
@@ -48,7 +48,7 @@ def from_api_content(
                     detail=detail,
                 )
             else:
-                image_data = ImageData.from_url(img_url=url, detail=detail)  # type: ignore
+                image_data = ImageData.from_url(img_url=url, detail=detail)
             content_part = ContentPartImage(data=image_data)
 
         content_parts.append(content_part)  # type: ignore
@@ -58,19 +58,18 @@ def from_api_content(
 
 def to_api_content(
     content: Content,
-) -> Iterable[ChatCompletionContentPartParam]:
-    api_content: list[ChatCompletionContentPartParam] = []
+) -> Iterable[OpenAIContentPartParam]:
+    api_content: list[OpenAIContentPartParam] = []
     for content_part in content.parts:
-        api_content_part: ChatCompletionContentPartParam
+        api_content_part: OpenAIContentPartParam
         if isinstance(content_part, ContentPartText):
-            api_content_part = ChatCompletionContentPartTextParam(
-                type="text",
-                text=content_part.data,
+            api_content_part = OpenAIContentPartTextParam(
+                type="text", text=content_part.data
             )
         else:
-            api_content_part = ChatCompletionContentPartImageParam(
+            api_content_part = OpenAIContentPartImageParam(
                 type="image_url",
-                image_url=ChatCompletionImageURL(
+                image_url=OpenAIImageURL(
                     url=image_data_to_str(content_part.data),
                     detail=content_part.data.detail,
                 ),