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

grasp_agents/openai/converters.py

@@ -1,31 +1,31 @@
-from collections.abc import AsyncIterator, Iterable
+from collections.abc import Iterable
 from typing import Any
 
 from pydantic import BaseModel
 
-from ..typing.completion import Completion, CompletionChunk
+from ..typing.completion import Completion, Usage
+from ..typing.completion_chunk import 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,
-    ChatCompletionMessage,
-    ChatCompletionSystemMessageParam,
-    ChatCompletionToolChoiceOptionParam,
-    ChatCompletionToolMessageParam,
-    ChatCompletionToolParam,
-    ChatCompletionUsage,
-    ChatCompletionUserMessageParam,
+    OpenAIAssistantMessageParam,
+    OpenAICompletion,
+    OpenAICompletionChunk,
+    OpenAICompletionUsage,
+    OpenAIContentPartParam,
+    OpenAIMessage,
+    OpenAISystemMessageParam,
+    OpenAIToolChoiceOptionParam,
+    OpenAIToolMessageParam,
+    OpenAIToolParam,
+    OpenAIUserMessageParam,
 )
+from .completion_chunk_converters import from_api_completion_chunk
 from .completion_converters import (
     from_api_completion,
-    from_api_completion_chunk,
-    from_api_completion_chunk_iterator,
+    from_api_completion_usage,
     to_api_completion,
 )
 from .content_converters import from_api_content, to_api_content
@@ -46,119 +46,95 @@ class OpenAIConverters(Converters):
     @staticmethod
     def to_system_message(
         system_message: SystemMessage, **kwargs: Any
-    ) -> ChatCompletionSystemMessageParam:
+    ) -> OpenAISystemMessageParam:
         return to_api_system_message(system_message, **kwargs)
 
     @staticmethod
     def from_system_message(
-        raw_message: ChatCompletionSystemMessageParam,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_message: OpenAISystemMessageParam, name: str | None = None, **kwargs: Any
     ) -> SystemMessage:
-        return from_api_system_message(raw_message, model_id=model_id, **kwargs)
+        return from_api_system_message(raw_message, name=name, **kwargs)
 
     @staticmethod
     def to_user_message(
         user_message: UserMessage, **kwargs: Any
-    ) -> ChatCompletionUserMessageParam:
+    ) -> OpenAIUserMessageParam:
         return to_api_user_message(user_message, **kwargs)
 
     @staticmethod
     def from_user_message(
-        raw_message: ChatCompletionUserMessageParam,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_message: OpenAIUserMessageParam, name: str | None = None, **kwargs: Any
     ) -> UserMessage:
-        return from_api_user_message(raw_message, model_id=model_id, **kwargs)
+        return from_api_user_message(raw_message, name=name, **kwargs)
 
     @staticmethod
     def to_assistant_message(
         assistant_message: AssistantMessage, **kwargs: Any
-    ) -> ChatCompletionAssistantMessageParam:
+    ) -> OpenAIAssistantMessageParam:
         return to_api_assistant_message(assistant_message, **kwargs)
 
+    @staticmethod
+    def from_completion_usage(raw_usage: OpenAICompletionUsage, **kwargs: Any) -> Usage:
+        return from_api_completion_usage(raw_usage, **kwargs)
+
     @staticmethod
     def from_assistant_message(
-        raw_message: ChatCompletionMessage,
-        raw_usage: ChatCompletionUsage,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_message: OpenAIMessage, name: str | None = None, **kwargs: Any
     ) -> AssistantMessage:
-        return from_api_assistant_message(
-            raw_message, raw_usage, model_id=model_id, **kwargs
-        )
+        return from_api_assistant_message(raw_message, name=name, **kwargs)
 
     @staticmethod
     def to_tool_message(
         tool_message: ToolMessage, **kwargs: Any
-    ) -> ChatCompletionToolMessageParam:
+    ) -> OpenAIToolMessageParam:
         return to_api_tool_message(tool_message, **kwargs)
 
     @staticmethod
     def from_tool_message(
-        raw_message: ChatCompletionToolMessageParam,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_message: OpenAIToolMessageParam, name: str | None = None, **kwargs: Any
     ) -> ToolMessage:
-        return from_api_tool_message(raw_message, model_id=model_id, **kwargs)
+        return from_api_tool_message(raw_message, name=name, **kwargs)
 
     @staticmethod
     def to_tool(
-        tool: BaseTool[BaseModel, Any, Any], **kwargs: Any
-    ) -> ChatCompletionToolParam:
-        return to_api_tool(tool, **kwargs)
+        tool: BaseTool[BaseModel, Any, Any], strict: bool | None = None, **kwargs: Any
+    ) -> OpenAIToolParam:
+        return to_api_tool(tool, strict=strict, **kwargs)
 
     @staticmethod
     def to_tool_choice(
         tool_choice: ToolChoice, **kwargs: Any
-    ) -> ChatCompletionToolChoiceOptionParam:
+    ) -> OpenAIToolChoiceOptionParam:
         return to_api_tool_choice(tool_choice, **kwargs)
 
     @staticmethod
-    def to_content(
-        content: Content, **kwargs: Any
-    ) -> Iterable[ChatCompletionContentPartParam]:
+    def to_content(content: Content, **kwargs: Any) -> Iterable[OpenAIContentPartParam]:
         return to_api_content(content, **kwargs)
 
     @staticmethod
     def from_content(
-        raw_content: str | Iterable[ChatCompletionContentPartParam],
-        **kwargs: Any,
+        raw_content: str | Iterable[OpenAIContentPartParam], **kwargs: Any
     ) -> Content:
         return from_api_content(raw_content, **kwargs)
 
     @staticmethod
-    def to_completion(completion: Completion, **kwargs: Any) -> ChatCompletion:
+    def to_completion(completion: Completion, **kwargs: Any) -> OpenAICompletion:
         return to_api_completion(completion, **kwargs)
 
     @staticmethod
     def from_completion(
-        raw_completion: ChatCompletion,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_completion: OpenAICompletion, name: str | None = None, **kwargs: Any
     ) -> Completion:
-        return from_api_completion(raw_completion, model_id=model_id, **kwargs)
+        return from_api_completion(raw_completion, name=name, **kwargs)
 
     @staticmethod
     def to_completion_chunk(
         chunk: CompletionChunk, **kwargs: Any
-    ) -> ChatCompletionChunk:
+    ) -> OpenAICompletionChunk:
         raise NotImplementedError
 
     @staticmethod
     def from_completion_chunk(
-        raw_chunk: ChatCompletionChunk,
-        model_id: str | None = None,
-        **kwargs: Any,
+        raw_chunk: OpenAICompletionChunk, name: 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
-        )
+        return from_api_completion_chunk(raw_chunk, name=name, **kwargs)

grasp_agents/openai/message_converters.py

@@ -1,76 +1,62 @@
 from typing import TypeAlias
 
+from ..typing.content import Content
 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,
+    OpenAIAssistantMessageParam,
+    OpenAIDeveloperMessageParam,
+    OpenAIFunctionMessageParam,
+    OpenAIMessage,
+    OpenAISystemMessageParam,
+    OpenAIToolCallFunction,
+    OpenAIToolCallParam,
+    OpenAIToolMessageParam,
+    OpenAIUserMessageParam,
 )
 from .content_converters import from_api_content, to_api_content
 
-OpenAIMessage: TypeAlias = (
-    ChatCompletionAssistantMessageParam
-    | ChatCompletionToolMessageParam
-    | ChatCompletionUserMessageParam
-    | ChatCompletionDeveloperMessageParam
-    | ChatCompletionSystemMessageParam
-    | ChatCompletionFunctionMessageParam
+OpenAIMessageType: TypeAlias = (
+    OpenAIAssistantMessageParam
+    | OpenAIToolMessageParam
+    | OpenAIUserMessageParam
+    | OpenAIDeveloperMessageParam
+    | OpenAISystemMessageParam
+    | OpenAIFunctionMessageParam
 )
 
 
 def from_api_user_message(
-    api_message: ChatCompletionUserMessageParam, model_id: str | None = None
+    api_message: OpenAIUserMessageParam, name: str | None = None
 ) -> UserMessage:
     content = from_api_content(api_message["content"])
+    name = api_message.get("name")
 
-    return UserMessage(content=content, model_id=model_id)
+    return UserMessage(content=content, name=name)
 
 
-def to_api_user_message(message: UserMessage) -> ChatCompletionUserMessageParam:
-    api_content = to_api_content(message.content)
+def to_api_user_message(message: UserMessage) -> OpenAIUserMessageParam:
+    api_content = (
+        to_api_content(message.content)
+        if isinstance(message.content, Content)
+        else message.content
+    )
+    api_name = message.name
+    api_message = OpenAIUserMessageParam(role="user", content=api_content)
+    if api_name is not None:
+        api_message["name"] = api_name
 
-    return ChatCompletionUserMessageParam(role="user", content=api_content)
+    return api_message
 
 
 def from_api_assistant_message(
-    api_message: ChatCompletionMessage,
-    api_usage: ChatCompletionUsage | None = None,
-    model_id: str | None = None,
+    api_message: OpenAIMessage, name: str | None = None
 ) -> AssistantMessage:
-    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 = [
@@ -84,23 +70,22 @@ def from_api_assistant_message(
 
     return AssistantMessage(
         content=api_message.content,
-        usage=usage,
         tool_calls=tool_calls,
         refusal=api_message.refusal,
-        model_id=model_id,
+        name=name,
     )
 
 
 def to_api_assistant_message(
     message: AssistantMessage,
-) -> ChatCompletionAssistantMessageParam:
+) -> OpenAIAssistantMessageParam:
     api_tool_calls = None
     if message.tool_calls is not None:
         api_tool_calls = [
-            ChatCompletionMessageToolCallParam(
+            OpenAIToolCallParam(
                 type="function",
                 id=tool_call.id,
-                function=ChatCompletionToolCallFunction(
+                function=OpenAIToolCallFunction(
                     name=tool_call.tool_name,
                     arguments=tool_call.tool_arguments,
                 ),
@@ -108,48 +93,49 @@ def to_api_assistant_message(
             for tool_call in message.tool_calls
         ]
 
-    api_message = ChatCompletionAssistantMessageParam(
-        role="assistant",
-        content=message.content,
-        tool_calls=api_tool_calls or [],
-        refusal=message.refusal,
-    )
+    api_message = OpenAIAssistantMessageParam(role="assistant", content=message.content)
+
+    if message.name is not None:
+        api_message["name"] = message.name
+    if api_tool_calls is not None:
+        api_message["tool_calls"] = api_tool_calls or []
+    if message.refusal is not None:
+        api_message["refusal"] = message.refusal
+
+    # TODO: hack
     if message.content is None:
         # Some API providers return None in the generated content without errors,
         # even though None in the input content is not accepted.
         api_message["content"] = "<empty>"
-    if api_tool_calls is None:
-        api_message.pop("tool_calls")
-    if message.refusal is None:
-        api_message.pop("refusal")
 
     return api_message
 
 
 def from_api_system_message(
-    api_message: ChatCompletionSystemMessageParam,
-    model_id: str | None = None,
+    api_message: OpenAISystemMessageParam, name: str | None = None
 ) -> SystemMessage:
-    return SystemMessage(content=api_message["content"], model_id=model_id)  # type: ignore
+    return SystemMessage(content=api_message["content"], name=name)  # type: ignore
 
 
-def to_api_system_message(
-    message: SystemMessage,
-) -> ChatCompletionSystemMessageParam:
-    return ChatCompletionSystemMessageParam(role="system", content=message.content)
+def to_api_system_message(message: SystemMessage) -> OpenAISystemMessageParam:
+    api_message = OpenAISystemMessageParam(role="system", content=message.content)
+    if message.name is not None:
+        api_message["name"] = message.name
+
+    return api_message
 
 
 def from_api_tool_message(
-    api_message: ChatCompletionToolMessageParam, model_id: str | None = None
+    api_message: OpenAIToolMessageParam, name: str | None = None
 ) -> ToolMessage:
     return ToolMessage(
         content=api_message["content"],  # type: ignore
         tool_call_id=api_message["tool_call_id"],
-        model_id=model_id,
+        name=name,
     )
 
 
-def to_api_tool_message(message: ToolMessage) -> ChatCompletionToolMessageParam:
-    return ChatCompletionToolMessageParam(
+def to_api_tool_message(message: ToolMessage) -> OpenAIToolMessageParam:
+    return OpenAIToolMessageParam(
         role="tool", content=message.content, tool_call_id=message.tool_call_id
     )

grasp_agents/openai/openai_llm.py

@@ -1,45 +1,45 @@
 import logging
-from collections.abc import Iterable, Mapping
+from collections.abc import AsyncIterator, Iterable, Mapping
 from copy import deepcopy
-from typing import Any, Literal
+from typing import Any, Literal, NamedTuple
 
 from openai import AsyncOpenAI
 from openai._types import NOT_GIVEN  # type: ignore[import]
+from openai.lib.streaming.chat import (
+    AsyncChatCompletionStreamManager as OpenAIAsyncChatCompletionStreamManager,
+)
+from openai.lib.streaming.chat import ChunkEvent as OpenAIChunkEvent
 from pydantic import BaseModel
 
 from ..cloud_llm import CloudLLM, CloudLLMSettings
 from ..http_client import AsyncHTTPClientParams
 from ..rate_limiting.rate_limiter_chunked import RateLimiterC
-from ..typing.message import AssistantMessage, Conversation
+from ..typing.message import AssistantMessage, Messages
 from ..typing.tool import BaseTool
 from . import (
-    ChatCompletion,
-    ChatCompletionAsyncStream,  # type: ignore[import]
-    ChatCompletionChunk,
-    ChatCompletionMessageParam,
-    ChatCompletionPredictionContentParam,
-    ChatCompletionStreamOptionsParam,
-    ChatCompletionToolChoiceOptionParam,
-    ChatCompletionToolParam,
-    ParsedChatCompletion,
-    # ResponseFormatJSONObject,
-    # ResponseFormatJSONSchema,
-    # ResponseFormatText,
+    OpenAICompletion,
+    OpenAICompletionChunk,
+    OpenAIMessageParam,
+    OpenAIParsedCompletion,
+    OpenAIPredictionContentParam,
+    OpenAIStreamOptionsParam,
+    OpenAIToolChoiceOptionParam,
+    OpenAIToolParam,
 )
 from .converters import OpenAIConverters
 
 logger = logging.getLogger(__name__)
 
 
+class ToolCallSettings(NamedTuple):
+    strict: bool | None = None
+
+
 class OpenAILLMSettings(CloudLLMSettings, total=False):
     reasoning_effort: Literal["low", "medium", "high"] | None
 
     parallel_tool_calls: bool
 
-    # response_format: (
-    #     ResponseFormatText | ResponseFormatJSONSchema | ResponseFormatJSONObject
-    # )
-
     modalities: list[Literal["text", "audio"]] | None
 
     frequency_penalty: float | None
@@ -48,16 +48,23 @@ class OpenAILLMSettings(CloudLLMSettings, total=False):
     stop: str | list[str] | None
     logprobs: bool | None
     top_logprobs: int | None
-    n: int | None
 
-    prediction: ChatCompletionPredictionContentParam | None
+    prediction: OpenAIPredictionContentParam | None
 
-    stream_options: ChatCompletionStreamOptionsParam | None
+    stream_options: OpenAIStreamOptionsParam | None
 
     metadata: dict[str, str] | None
     store: bool | None
     user: str
 
+    strict_tool_args: bool
+
+    # response_format: (
+    #     OpenAIResponseFormatText
+    #     | OpenAIResponseFormatJSONSchema
+    #     | OpenAIResponseFormatJSONObject
+    # )
+
 
 class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
     def __init__(
@@ -74,7 +81,7 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         ) = None,
         async_openai_client_params: dict[str, Any] | None = None,
         # Rate limiting
-        rate_limiter: (RateLimiterC[Conversation, AssistantMessage] | None) = None,
+        rate_limiter: (RateLimiterC[Messages, AssistantMessage] | None) = None,
         rate_limiter_rpm: float | None = None,
         rate_limiter_chunk_size: int = 1000,
         rate_limiter_max_concurrency: int = 300,
@@ -101,46 +108,55 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
             **kwargs,
         )
 
-        async_openai_client_params_ = deepcopy(async_openai_client_params or {})
+        self._tool_call_settings = {
+            "strict": self._llm_settings.pop("strict_tool_args", None)
+        }
+
+        _async_openai_client_params = deepcopy(async_openai_client_params or {})
         if self._async_http_client is not None:
-            async_openai_client_params_["http_client"] = self._async_http_client
+            _async_openai_client_params["http_client"] = self._async_http_client
 
         # TODO: context manager for async client
         self._client: AsyncOpenAI = AsyncOpenAI(
             base_url=self._base_url,
             api_key=self._api_key,
-            **async_openai_client_params_,
+            **_async_openai_client_params,
         )
 
     async def _get_completion(
         self,
-        api_messages: Iterable[ChatCompletionMessageParam],
-        api_tools: list[ChatCompletionToolParam] | None = None,
-        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        api_messages: Iterable[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        n_choices: int | None = None,
         **api_llm_settings: Any,
-    ) -> ChatCompletion:
+    ) -> OpenAICompletion:
         tools = api_tools or NOT_GIVEN
         tool_choice = api_tool_choice or NOT_GIVEN
+        n = n_choices or NOT_GIVEN
 
         return await self._client.chat.completions.create(
             model=self._api_model_name,
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
+            n=n,
             stream=False,
             **api_llm_settings,
         )
 
     async def _get_parsed_completion(
         self,
-        api_messages: Iterable[ChatCompletionMessageParam],
-        api_tools: list[ChatCompletionToolParam] | None = None,
-        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        api_messages: Iterable[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
         api_response_format: type | None = None,
+        n_choices: int | None = None,
         **api_llm_settings: Any,
-    ) -> ParsedChatCompletion[Any]:
+    ) -> OpenAIParsedCompletion[Any]:
         tools = api_tools or NOT_GIVEN
         tool_choice = api_tool_choice or NOT_GIVEN
+        n = n_choices or NOT_GIVEN
         response_format = api_response_format or NOT_GIVEN
 
         return await self._client.beta.chat.completions.parse(
@@ -148,27 +164,70 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
-            response_format=response_format,  # type: ignore[arg-type]
+            response_format=response_format,
+            n=n,
             **api_llm_settings,
         )
 
     async def _get_completion_stream(
         self,
-        api_messages: Iterable[ChatCompletionMessageParam],
-        api_tools: list[ChatCompletionToolParam] | None = None,
-        api_tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        api_messages: Iterable[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        n_choices: int | None = None,
         **api_llm_settings: Any,
-    ) -> ChatCompletionAsyncStream[ChatCompletionChunk]:
-        assert not api_tools, "Tool use is not supported in streaming mode"
-
+    ) -> AsyncIterator[OpenAICompletionChunk]:
         tools = api_tools or NOT_GIVEN
         tool_choice = api_tool_choice or NOT_GIVEN
+        n = n_choices or NOT_GIVEN
 
-        return await self._client.chat.completions.create(
+        stream_generator = await self._client.chat.completions.create(
             model=self._api_model_name,
             messages=api_messages,
             tools=tools,
             tool_choice=tool_choice,
             stream=True,
+            n=n,
             **api_llm_settings,
         )
+
+        async def iterate() -> AsyncIterator[OpenAICompletionChunk]:
+            async with stream_generator as stream:
+                async for completion_chunk in stream:
+                    yield completion_chunk
+
+        return iterate()
+
+    async def _get_parsed_completion_stream(
+        self,
+        api_messages: Iterable[OpenAIMessageParam],
+        api_tools: list[OpenAIToolParam] | None = None,
+        api_tool_choice: OpenAIToolChoiceOptionParam | None = None,
+        api_response_format: type | None = None,
+        n_choices: int | None = None,
+        **api_llm_settings: Any,
+    ) -> AsyncIterator[OpenAICompletionChunk]:
+        tools = api_tools or NOT_GIVEN
+        tool_choice = api_tool_choice or NOT_GIVEN
+        response_format = api_response_format or NOT_GIVEN
+        n = n_choices or NOT_GIVEN
+
+        stream_manager: OpenAIAsyncChatCompletionStreamManager[
+            OpenAICompletionChunk
+        ] = self._client.beta.chat.completions.stream(
+            model=self._api_model_name,
+            messages=api_messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            response_format=response_format,
+            n=n,
+            **api_llm_settings,
+        )
+
+        async def iterate() -> AsyncIterator[OpenAICompletionChunk]:
+            async with stream_manager as stream:
+                async for chunk_event in stream:
+                    if isinstance(chunk_event, OpenAIChunkEvent):
+                        yield chunk_event.chunk
+
+        return iterate()