unique_toolkit 1.15.0__py3-none-any.whl → 1.16.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
@@ -109,21 +124,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 +185,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 +253,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 +260,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 +292,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 +346,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 +362,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 +427,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:
@@ -502,13 +678,32 @@ class LanguageModelToolDescription(BaseModel):
     def serialize_parameters(self, parameters: type[BaseModel]):
         return parameters.model_json_schema()
 
-    def to_openai(self) -> ChatCompletionToolParam:
-        return ChatCompletionToolParam(
-            function=FunctionDefinition(
+    @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=self.parameters.model_json_schema(),
+                    strict=self.strict,
+                ),
+                type="function",
+            )
+        elif mode == "responses":
+            return FunctionToolParam(
+                type="function",
                 name=self.name,
-                description=self.description,
                 parameters=self.parameters.model_json_schema(),
                 strict=self.strict,
-            ),
-            type="function",
-        )
+                description=self.description,
+            )

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",
+]