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

unique_toolkit/language_model/schemas.py

@@ -1,16 +1,31 @@
 import json
 import math
 from enum import StrEnum
-from typing import Any, Self
+from typing import Any, Literal, Self, TypeVar
 from uuid import uuid4
 
 from humps import camelize
-from openai.types.chat import ChatCompletionAssistantMessageParam
+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,
@@ -23,7 +38,7 @@ from pydantic import (
     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
@@ -68,22 +83,21 @@ class LanguageModelStreamResponseMessage(BaseModel):
 class LanguageModelFunction(BaseModel):
     model_config = model_config
 
-    id: str | None = None
+    id: str = Field(default_factory=lambda: uuid4().hex)
     name: str
-    arguments: dict[str, Any] | str | None = 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):
-        if not value:
+    def randomize_id(cls, value: Any) -> Any:
+        if value is None or value == "":
             return uuid4().hex
-        else:
-            return value
+        return value
 
     @model_serializer()
     def serialize_model(self):
@@ -93,7 +107,7 @@ class LanguageModelFunction(BaseModel):
             seralization["arguments"] = json.dumps(self.arguments)
         return seralization
 
-    def __eq__(self, other: Self) -> bool:
+    def __eq__(self, other: object) -> bool:
         """Compare two tool calls based on name and arguments."""
         if not isinstance(other, LanguageModelFunction):
             return False
@@ -109,21 +123,43 @@ class LanguageModelFunction(BaseModel):
 
         return True
 
-    def to_openai_param(self) -> ChatCompletionMessageFunctionToolCallParam:
+    @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
 
-        return ChatCompletionMessageFunctionToolCallParam(
-            type="function",
-            id=self.id or "unknown_id",
-            function=Function(name=self.name, arguments=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,
+            )
 
 
-# This is tailored to the unique backend
 class LanguageModelStreamResponse(BaseModel):
     model_config = model_config
 
@@ -148,6 +184,32 @@ class LanguageModelStreamResponse(BaseModel):
         )
 
 
+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
 
@@ -190,7 +252,6 @@ class LanguageModelMessage(BaseModel):
         return format_message(self.role.capitalize(), message=message, num_tabs=1)
 
 
-# Equivalent to
 class LanguageModelSystemMessage(LanguageModelMessage):
     role: LanguageModelMessageRole = LanguageModelMessageRole.SYSTEM
 
@@ -198,6 +259,26 @@ 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
@@ -210,6 +291,26 @@ 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
@@ -244,8 +345,8 @@ class LanguageModelAssistantMessage(LanguageModelMessage):
     def from_stream_response(cls, response: LanguageModelStreamResponse):
         tool_calls = [
             LanguageModelFunctionCall(
-                id=None,
-                type=None,
+                id=f.id,
+                type="function",
                 function=f,
             )
             for f in response.tool_calls or []
@@ -260,6 +361,50 @@ class LanguageModelAssistantMessage(LanguageModelMessage):
             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
@@ -281,20 +426,50 @@ 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:
@@ -487,9 +662,10 @@ class LanguageModelToolDescription(BaseModel):
         ...,
         description="Description of what the tool is doing the tool",
     )
-    parameters: type[BaseModel] = Field(
+    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
@@ -499,16 +675,46 @@ class LanguageModelToolDescription(BaseModel):
     )
 
     @field_serializer("parameters")
-    def serialize_parameters(self, parameters: type[BaseModel]):
-        return parameters.model_json_schema()
-
-    def to_openai(self) -> ChatCompletionToolParam:
-        return ChatCompletionToolParam(
-            function=FunctionDefinition(
+    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,
-                description=self.description,
-                parameters=self.parameters.model_json_schema(),
+                parameters=_parameters_as_json_schema(self.parameters),
                 strict=self.strict,
-            ),
-            type="function",
-        )
+                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()

unique_toolkit/protocols/support.py

@@ -1,4 +1,18 @@
-from typing import Any, Awaitable, Protocol
+from typing import Awaitable, Protocol, Sequence
+
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+)
+from openai.types.responses import (
+    ResponseIncludable,
+    ResponseInputItemParam,
+    ResponseOutputItem,
+    ResponseTextConfigParam,
+    ToolParam,
+    response_create_params,
+)
+from openai.types.shared_params import Metadata, Reasoning
 
 from unique_toolkit.content import ContentChunk
 from unique_toolkit.language_model import (
@@ -13,6 +27,10 @@ from unique_toolkit.language_model.constants import (
     DEFAULT_COMPLETE_TEMPERATURE,
     DEFAULT_COMPLETE_TIMEOUT,
 )
+from unique_toolkit.language_model.schemas import (
+    LanguageModelMessageOptions,
+    ResponsesLanguageModelStreamResponse,
+)
 
 # As soon as we have multiple, remember
 # https://pypi.org/project/typing-protocol-intersection/
@@ -35,7 +53,7 @@ class SupportsComplete(Protocol):
         model_name: LanguageModelName | str,
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
-        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     ) -> Awaitable[LanguageModelResponse]: ...
 
 
@@ -45,19 +63,83 @@ class SupportCompleteWithReferences(Protocol):
         messages: LanguageModelMessages,
         model_name: LanguageModelName | str,
         content_chunks: list[ContentChunk] | None = None,
-        debug_info: dict[str, Any] = {},
+        debug_info: dict | None = None,
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
-        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        start_text: str | None = None,
+        tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        other_options: dict | None = None,
     ) -> LanguageModelStreamResponse: ...
 
-    def complete_with_references_async(
+    async def complete_with_references_async(
         self,
-        messages: LanguageModelMessages,
+        messages: LanguageModelMessages | list[ChatCompletionMessageParam],
         model_name: LanguageModelName | str,
         content_chunks: list[ContentChunk] | None = None,
-        debug_info: dict[str, Any] = {},
+        debug_info: dict | None = None,
         temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
         timeout: int = DEFAULT_COMPLETE_TIMEOUT,
-        tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
-    ) -> Awaitable[LanguageModelStreamResponse]: ...
+        tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
+        tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
+        start_text: str | None = None,
+        other_options: dict | None = None,
+    ) -> LanguageModelStreamResponse: ...
+
+
+class ResponsesSupportCompleteWithReferences(Protocol):
+    def complete_with_references(
+        self,
+        *,
+        model_name: LanguageModelName | str,
+        messages: str
+        | LanguageModelMessages
+        | Sequence[
+            ResponseInputItemParam
+            | LanguageModelMessageOptions
+            | ResponseOutputItem  # History is automatically convertible
+        ],
+        content_chunks: list[ContentChunk] | None = None,
+        tools: Sequence[LanguageModelToolDescription | ToolParam] | None = None,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        debug_info: dict | None = None,
+        start_text: str | None = None,
+        include: list[ResponseIncludable] | None = None,
+        instructions: str | None = None,
+        max_output_tokens: int | None = None,
+        metadata: Metadata | None = None,
+        parallel_tool_calls: bool | None = None,
+        text: ResponseTextConfigParam | None = None,
+        tool_choice: response_create_params.ToolChoice | None = None,
+        top_p: float | None = None,
+        reasoning: Reasoning | None = None,
+        other_options: dict | None = None,
+    ) -> ResponsesLanguageModelStreamResponse: ...
+
+    async def complete_with_references_async(
+        self,
+        *,
+        model_name: LanguageModelName | str,
+        messages: str
+        | LanguageModelMessages
+        | Sequence[
+            ResponseInputItemParam
+            | LanguageModelMessageOptions
+            | ResponseOutputItem  # History is automatically convertible
+        ],
+        content_chunks: list[ContentChunk] | None = None,
+        tools: Sequence[LanguageModelToolDescription | ToolParam] | None = None,
+        temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+        debug_info: dict | None = None,
+        start_text: str | None = None,
+        include: list[ResponseIncludable] | None = None,
+        instructions: str | None = None,
+        max_output_tokens: int | None = None,
+        metadata: Metadata | None = None,
+        parallel_tool_calls: bool | None = None,
+        text: ResponseTextConfigParam | None = None,
+        tool_choice: response_create_params.ToolChoice | None = None,
+        top_p: float | None = None,
+        reasoning: Reasoning | None = None,
+        other_options: dict | None = None,
+    ) -> ResponsesLanguageModelStreamResponse: ...

unique_toolkit/services/__init__.py

@@ -0,0 +1,7 @@
+from unique_toolkit.services.chat_service import ChatService
+from unique_toolkit.services.knowledge_base import KnowledgeBaseService
+
+__all__ = [
+    "ChatService",
+    "KnowledgeBaseService",
+]