unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/language_model/schemas.py

@@ -1,22 +1,46 @@
 import json
 import math
 from enum import StrEnum
-from typing import Any, Optional, Self
+from typing import Any, Literal, Self, TypeVar
 from uuid import uuid4
 
 from humps import camelize
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.chat_completion_message_function_tool_call_param import (
+    ChatCompletionMessageFunctionToolCallParam,
+    Function,
+)
+from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+from openai.types.responses import (
+    EasyInputMessageParam,
+    FunctionToolParam,
+    ResponseCodeInterpreterToolCall,
+    ResponseFunctionToolCallParam,
+    ResponseOutputItem,
+    ResponseOutputMessage,
+)
+from openai.types.responses.response_input_param import FunctionCallOutput
+from openai.types.responses.response_output_text import AnnotationContainerFileCitation
+from openai.types.shared_params.function_definition import FunctionDefinition
 from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
     PrivateAttr,
     RootModel,
+    field_serializer,
     field_validator,
     model_serializer,
     model_validator,
 )
-from typing_extensions import deprecated
+from typing_extensions import deprecated, overload
 
+from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.utils import format_message
 
 # set config to convert camelCase to snake_case
@@ -27,6 +51,8 @@ model_config = ConfigDict(
 )
 
 
+# Equivalent to
+# from openai.types.chat.chat_completion_role import ChatCompletionRole
 class LanguageModelMessageRole(StrEnum):
     USER = "user"
     SYSTEM = "system"
@@ -34,33 +60,155 @@ class LanguageModelMessageRole(StrEnum):
     TOOL = "tool"
 
 
+# This is tailored to the unique backend
+class LanguageModelStreamResponseMessage(BaseModel):
+    model_config = model_config
+
+    id: str
+    previous_message_id: (
+        str | None
+    )  # Stream response can return a null previous_message_id if an assisstant message is manually added
+    role: LanguageModelMessageRole
+    text: str
+    original_text: str | None = None
+    references: list[ContentReference] = []
+
+    # TODO make sdk return role in lowercase
+    # Currently needed as sdk returns role in uppercase
+    @field_validator("role", mode="before")
+    def set_role(cls, value: str):
+        return value.lower()
+
+
 class LanguageModelFunction(BaseModel):
     model_config = model_config
 
-    id: str | None = None
+    id: str = Field(default_factory=lambda: uuid4().hex)
     name: str
-    arguments: Optional[dict[str, Any] | str] = None  # type: ignore
+    arguments: dict[str, Any] | None = None
 
     @field_validator("arguments", mode="before")
-    def set_arguments(cls, value):
+    def set_arguments(cls, value: Any) -> Any:
         if isinstance(value, str):
             return json.loads(value)
         return value
 
     @field_validator("id", mode="before")
-    def randomize_id(cls, value):
-        return uuid4().hex
+    def randomize_id(cls, value: Any) -> Any:
+        if value is None or value == "":
+            return uuid4().hex
+        return value
 
     @model_serializer()
     def serialize_model(self):
         seralization = {}
-        if self.id:
-            seralization["id"] = self.id
         seralization["name"] = self.name
         if self.arguments:
             seralization["arguments"] = json.dumps(self.arguments)
         return seralization
 
+    def __eq__(self, other: object) -> bool:
+        """Compare two tool calls based on name and arguments."""
+        if not isinstance(other, LanguageModelFunction):
+            return False
+
+        if self.id != other.id:
+            return False
+
+        if self.name != other.name:
+            return False
+
+        if self.arguments != other.arguments:
+            return False
+
+        return True
+
+    @overload
+    def to_openai_param(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionMessageFunctionToolCallParam: ...
+
+    @overload
+    def to_openai_param(
+        self, mode: Literal["responses"]
+    ) -> ResponseFunctionToolCallParam: ...
+
+    def to_openai_param(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> ChatCompletionMessageFunctionToolCallParam | ResponseFunctionToolCallParam:
+        arguments = ""
+        if isinstance(self.arguments, dict):
+            arguments = json.dumps(self.arguments)
+        elif isinstance(self.arguments, str):
+            arguments = self.arguments
+
+        if mode == "completions":
+            return ChatCompletionMessageFunctionToolCallParam(
+                type="function",
+                id=self.id or "unknown_id",
+                function=Function(name=self.name, arguments=arguments),
+            )
+        elif mode == "responses":
+            if self.id is None:
+                raise ValueError("Missing tool call id")
+
+            return ResponseFunctionToolCallParam(
+                type="function_call",
+                call_id=self.id,
+                name=self.name,
+                arguments=arguments,
+            )
+
+
+class LanguageModelStreamResponse(BaseModel):
+    model_config = model_config
+
+    message: LanguageModelStreamResponseMessage
+    tool_calls: list[LanguageModelFunction] | None = None
+
+    def is_empty(self) -> bool:
+        """
+        Check if the stream response is empty.
+        An empty stream response has no text and no tool calls.
+        """
+        return not self.message.original_text and not self.tool_calls
+
+    def to_openai_param(self) -> ChatCompletionAssistantMessageParam:
+        return ChatCompletionAssistantMessageParam(
+            role="assistant",
+            audio=None,
+            content=self.message.text,
+            function_call=None,
+            refusal=None,
+            tool_calls=[t.to_openai_param() for t in self.tool_calls or []],
+        )
+
+
+OutputItemType = TypeVar("OutputItemType", bound=ResponseOutputItem)
+
+
+class ResponsesLanguageModelStreamResponse(LanguageModelStreamResponse):
+    output: list[ResponseOutputItem]
+
+    def filter_output(self, type: type[OutputItemType]) -> list[OutputItemType]:
+        return [item for item in self.output if isinstance(item, type)]
+
+    @property
+    def code_interpreter_calls(self) -> list[ResponseCodeInterpreterToolCall]:
+        return self.filter_output(ResponseCodeInterpreterToolCall)
+
+    @property
+    def container_files(self) -> list[AnnotationContainerFileCitation]:
+        container_files = []
+        messages = self.filter_output(ResponseOutputMessage)
+        for message in messages:
+            for content in message.content:
+                if content.type == "output_text":
+                    for annotation in content.annotations:
+                        if annotation.type == "container_file_citation":
+                            container_files.append(annotation)
+        return container_files
+
 
 class LanguageModelFunctionCall(BaseModel):
     model_config = model_config
@@ -69,10 +217,12 @@ class LanguageModelFunctionCall(BaseModel):
     type: str | None = None
     function: LanguageModelFunction
 
+    # TODO: Circular reference of types
+    @deprecated("Use LanguageModelAssistantMessage.from_functions instead.")
     @staticmethod
     def create_assistant_message_from_tool_calls(
         tool_calls: list[LanguageModelFunction],
-    ):
+    ) -> "LanguageModelAssistantMessage":
         assistant_message = LanguageModelAssistantMessage(
             content="",
             tool_calls=[
@@ -93,8 +243,7 @@ class LanguageModelMessage(BaseModel):
     content: str | list[dict] | None = None
 
     def __str__(self):
-        if not self.content:
-            message = ""
+        message = ""
         if isinstance(self.content, str):
             message = self.content
         elif isinstance(self.content, list):
@@ -110,6 +259,30 @@ class LanguageModelSystemMessage(LanguageModelMessage):
     def set_role(cls, value):
         return LanguageModelMessageRole.SYSTEM
 
+    @overload
+    def to_openai(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionSystemMessageParam: ...
+
+    @overload
+    def to_openai(self, mode: Literal["responses"]) -> EasyInputMessageParam: ...
+
+    def to_openai(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> ChatCompletionSystemMessageParam | EasyInputMessageParam:
+        content = self.content or ""
+        if not isinstance(content, str):
+            raise ValueError("Content must be a string")
+
+        if mode == "completions":
+            return ChatCompletionSystemMessageParam(role="system", content=content)
+        elif mode == "responses":
+            return EasyInputMessageParam(role="user", content=content)
+
+
+# Equivalent to
+# from openai.types.chat.chat_completion_user_message_param import ChatCompletionUserMessageParam
+
 
 class LanguageModelUserMessage(LanguageModelMessage):
     role: LanguageModelMessageRole = LanguageModelMessageRole.USER
@@ -118,7 +291,29 @@ class LanguageModelUserMessage(LanguageModelMessage):
     def set_role(cls, value):
         return LanguageModelMessageRole.USER
 
+    @overload
+    def to_openai(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionUserMessageParam: ...
+
+    @overload
+    def to_openai(self, mode: Literal["responses"]) -> EasyInputMessageParam: ...
+
+    def to_openai(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> ChatCompletionUserMessageParam | EasyInputMessageParam:
+        content = self.content or ""
+        if not isinstance(content, str):
+            raise ValueError("Content must be a string")
+
+        if mode == "completions":
+            return ChatCompletionUserMessageParam(role="user", content=content)
+        elif mode == "responses":
+            return EasyInputMessageParam(role="user", content=content)
+
 
+# Equivalent to
+# from openai.types.chat.chat_completion_assistant_message_param import ChatCompletionAssistantMessageParam
 class LanguageModelAssistantMessage(LanguageModelMessage):
     role: LanguageModelMessageRole = LanguageModelMessageRole.ASSISTANT
     parsed: dict | None = None
@@ -129,6 +324,91 @@ class LanguageModelAssistantMessage(LanguageModelMessage):
     def set_role(cls, value):
         return LanguageModelMessageRole.ASSISTANT
 
+    @classmethod
+    def from_functions(
+        cls,
+        tool_calls: list[LanguageModelFunction],
+    ):
+        return cls(
+            content="",
+            tool_calls=[
+                LanguageModelFunctionCall(
+                    id=tool_call.id,
+                    type="function",
+                    function=tool_call,
+                )
+                for tool_call in tool_calls
+            ],
+        )
+
+    @classmethod
+    def from_stream_response(cls, response: LanguageModelStreamResponse):
+        tool_calls = [
+            LanguageModelFunctionCall(
+                id=f.id,
+                type="function",
+                function=f,
+            )
+            for f in response.tool_calls or []
+        ]
+
+        tool_calls = tool_calls if len(tool_calls) > 0 else None
+
+        return cls(
+            content=response.message.text,
+            parsed=None,
+            refusal=None,
+            tool_calls=tool_calls,
+        )
+
+    @overload
+    def to_openai(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionAssistantMessageParam: ...
+
+    @overload
+    def to_openai(
+        self, mode: Literal["responses"]
+    ) -> list[EasyInputMessageParam | ResponseFunctionToolCallParam]: ...
+
+    def to_openai(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> (
+        ChatCompletionAssistantMessageParam
+        | list[EasyInputMessageParam | ResponseFunctionToolCallParam]
+    ):
+        content = self.content or ""
+        if not isinstance(content, str):
+            raise ValueError("Content must be a string")
+
+        if mode == "completions":
+            return ChatCompletionAssistantMessageParam(
+                role="assistant",
+                content=content,
+                tool_calls=[
+                    t.function.to_openai_param() for t in self.tool_calls or []
+                ],
+            )
+        elif mode == "responses":
+            """
+            Responses API does not support assistant messages with tool calls
+            """
+            res = []
+            if content != "":
+                res.append(EasyInputMessageParam(role="assistant", content=content))
+            if self.tool_calls:
+                res.extend(
+                    [
+                        t.function.to_openai_param(mode="responses")
+                        for t in self.tool_calls
+                    ]
+                )
+            return res
+
+
+# Equivalent to
+# from openai.types.chat.chat_completion_tool_message_param import ChatCompletionToolMessageParam
+
 
 class LanguageModelToolMessage(LanguageModelMessage):
     role: LanguageModelMessageRole = LanguageModelMessageRole.TOOL
@@ -146,15 +426,84 @@ class LanguageModelToolMessage(LanguageModelMessage):
     def set_role(cls, value):
         return LanguageModelMessageRole.TOOL
 
+    @overload
+    def to_openai(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionToolMessageParam: ...
+
+    @overload
+    def to_openai(self, mode: Literal["responses"]) -> FunctionCallOutput: ...
+
+    def to_openai(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> ChatCompletionToolMessageParam | FunctionCallOutput:
+        content = self.content or ""
+        if not isinstance(content, str):
+            raise ValueError("Content must be a string")
+
+        if mode == "completions":
+            return ChatCompletionToolMessageParam(
+                role="tool",
+                content=content,
+                tool_call_id=self.tool_call_id,
+            )
+        elif mode == "responses":
+            return FunctionCallOutput(
+                call_id=self.tool_call_id,
+                output=content,
+                type="function_call_output",
+            )
+
+
+# Equivalent implementation for list of
+# from openai.types.chat.chat_completion_tool_message_param import ChatCompletionToolMessageParam
+# with the addition of the builder
+
+LanguageModelMessageOptions = (
+    LanguageModelMessage
+    | LanguageModelToolMessage
+    | LanguageModelAssistantMessage
+    | LanguageModelSystemMessage
+    | LanguageModelUserMessage
+)
+
 
 class LanguageModelMessages(RootModel):
-    root: list[
-        LanguageModelMessage
-        | LanguageModelToolMessage
-        | LanguageModelAssistantMessage
-        | LanguageModelSystemMessage
-        | LanguageModelUserMessage
-    ]
+    root: list[LanguageModelMessageOptions]
+
+    @classmethod
+    def load_messages_to_root(cls, data: list[dict] | dict) -> Self:
+        """Convert list of dictionaries to appropriate message objects based on role."""
+        # Handle case where data is already wrapped in root
+        if isinstance(data, dict) and "root" in data:
+            messages_list = data["root"]
+        elif isinstance(data, list):
+            messages_list = data
+        else:
+            raise ValueError("Invalid data type")
+
+        # Convert the messages list
+        converted_messages = []
+        for item in messages_list:
+            if isinstance(item, dict):
+                role = item.get("role", "").lower()
+
+                # Map dictionary to appropriate message class based on role
+                if role == "system":
+                    converted_messages.append(LanguageModelSystemMessage(**item))
+                elif role == "user":
+                    converted_messages.append(LanguageModelUserMessage(**item))
+                elif role == "assistant":
+                    converted_messages.append(LanguageModelAssistantMessage(**item))
+                elif role == "tool":
+                    converted_messages.append(LanguageModelToolMessage(**item))
+                else:
+                    # Fallback to base LanguageModelMessage
+                    converted_messages.append(LanguageModelMessage(**item))
+            else:
+                # If it's already a message object, keep it as is
+                converted_messages.append(item)
+        return cls(root=converted_messages)
 
     def __str__(self):
         return "\n\n".join([str(message) for message in self.root])
@@ -174,6 +523,11 @@ class LanguageModelMessages(RootModel):
         return builder
 
 
+# This seems similar to
+# from openai.types.completion_choice import CompletionChoice
+# but is missing multiple attributes and uses message instead of text
+
+
 class LanguageModelCompletionChoice(BaseModel):
     model_config = model_config
 
@@ -182,38 +536,26 @@ class LanguageModelCompletionChoice(BaseModel):
     finish_reason: str
 
 
+# This seems similar to
+# from openai.types.completion import Completion
+# but is missing multiple attributes
 class LanguageModelResponse(BaseModel):
     model_config = model_config
 
     choices: list[LanguageModelCompletionChoice]
 
+    @classmethod
+    def from_stream_response(cls, response: LanguageModelStreamResponse):
+        choice = LanguageModelCompletionChoice(
+            index=0,
+            message=LanguageModelAssistantMessage.from_stream_response(response),
+            finish_reason="",
+        )
 
-class LanguageModelStreamResponseMessage(BaseModel):
-    model_config = model_config
-
-    id: str
-    previous_message_id: (
-        str | None
-    )  # Stream response can return a null previous_message_id if an assisstant message is manually added
-    role: LanguageModelMessageRole
-    text: str
-    original_text: str | None = None
-    references: list[dict[str, list | dict | str | int | float | bool]] = []  # type: ignore
-
-    # TODO make sdk return role in lowercase
-    # Currently needed as sdk returns role in uppercase
-    @field_validator("role", mode="before")
-    def set_role(cls, value: str):
-        return value.lower()
-
-
-class LanguageModelStreamResponse(BaseModel):
-    model_config = model_config
-
-    message: LanguageModelStreamResponseMessage
-    tool_calls: Optional[list[LanguageModelFunction]] = None
+        return cls(choices=[choice])
 
 
+# This is tailored for unique and only used in language model info
 class LanguageModelTokenLimits(BaseModel):
     token_limit_input: int
     token_limit_output: int
@@ -255,29 +597,46 @@ class LanguageModelTokenLimits(BaseModel):
 
                 data["token_limit_input"] = math.floor(fraction_input * token_limit)
                 data["token_limit_output"] = math.floor(
-                    (1 - fraction_input) * token_limit
+                    (1 - fraction_input) * token_limit,
                 )
                 data["_fraction_adaptpable"] = True
             return data
 
         raise ValueError(
-            'Either "token_limit_input" and "token_limit_output" must be provided together, or "token_limit" must be provided.'
+            'Either "token_limit_input" and "token_limit_output" must be provided together, or "token_limit" must be provided.',
         )
 
 
+# This is more restrictive than what openai allows
+
+
+@deprecated(
+    "Deprecated as `LanguageModelTool` is deprecated in favor of `LanguageModelToolDescription`",
+)
 class LanguageModelToolParameterProperty(BaseModel):
     type: str
     description: str
-    enum: Optional[list[Any]] = None
-    items: Optional[Self] = None
+    enum: list[Any] | None = None
+    items: Self | None = None
 
 
+# Looks most like
+# from openai.types.shared.function_parameters import FunctionParameters
+@deprecated(
+    "Deprecated as `LanguageModelTool` is deprecated in favor of `LanguageModelToolDescription`",
+)
 class LanguageModelToolParameters(BaseModel):
     type: str = "object"
     properties: dict[str, LanguageModelToolParameterProperty]
     required: list[str]
 
 
+# Looks most like
+# from openai.types.shared_params.function_definition import FunctionDefinition
+# but returns parameter is not known
+@deprecated(
+    "Deprecated as `LanguageModelTool` use `LanguageModelToolDescription` instead",
+)
 class LanguageModelTool(BaseModel):
     name: str = Field(
         ...,
@@ -286,8 +645,76 @@ class LanguageModelTool(BaseModel):
     )
     description: str
     parameters: (
-        LanguageModelToolParameters | dict
+        LanguageModelToolParameters | dict[str, Any]
     )  # dict represents json schema dumped from pydantic
     returns: LanguageModelToolParameterProperty | LanguageModelToolParameters | None = (
         None
     )
+
+
+class LanguageModelToolDescription(BaseModel):
+    name: str = Field(
+        ...,
+        pattern=r"^[a-zA-Z1-9_-]+$",
+        description="Name must adhere to the pattern ^[a-zA-Z1-9_-]+$",
+    )
+    description: str = Field(
+        ...,
+        description="Description of what the tool is doing the tool",
+    )
+    parameters: type[BaseModel] | dict[str, Any] = Field(
+        ...,
+        description="Pydantic model for the tool parameters",
+        union_mode="left_to_right",
+    )
+
+    # TODO: This should be default `True` but if this is the case the parameter_model needs to include additional properties
+    strict: bool = Field(
+        default=False,
+        description="Setting strict to true will ensure function calls reliably adhere to the function schema, instead of being best effort. If set to True the `parameter_model` set `model_config = {'extra':'forbid'}` must be set for on all BaseModels.",
+    )
+
+    @field_serializer("parameters")
+    def serialize_parameters(
+        self, parameters: type[BaseModel] | dict[str, Any]
+    ) -> dict[str, Any]:
+        return _parameters_as_json_schema(parameters)
+
+    @overload
+    def to_openai(
+        self, mode: Literal["completions"] = "completions"
+    ) -> ChatCompletionToolParam: ...
+
+    @overload
+    def to_openai(self, mode: Literal["responses"]) -> FunctionToolParam: ...
+
+    def to_openai(
+        self, mode: Literal["completions", "responses"] = "completions"
+    ) -> ChatCompletionToolParam | FunctionToolParam:
+        if mode == "completions":
+            return ChatCompletionToolParam(
+                function=FunctionDefinition(
+                    name=self.name,
+                    description=self.description,
+                    parameters=_parameters_as_json_schema(self.parameters),
+                    strict=self.strict,
+                ),
+                type="function",
+            )
+        elif mode == "responses":
+            return FunctionToolParam(
+                type="function",
+                name=self.name,
+                parameters=_parameters_as_json_schema(self.parameters),
+                strict=self.strict,
+                description=self.description,
+            )
+
+
+def _parameters_as_json_schema(
+    parameters: type[BaseModel] | dict[str, Any],
+) -> dict[str, Any]:
+    if isinstance(parameters, dict):
+        return parameters
+
+    return parameters.model_json_schema()