unique_toolkit 1.15.0__py3-none-any.whl → 1.16.1__py3-none-any.whl

unique_toolkit/chat/responses_api.py

@@ -0,0 +1,461 @@
+import logging
+from typing import Any, NamedTuple, Sequence
+
+import unique_sdk
+from openai.types.responses import (
+    ResponseIncludable,
+    ResponseInputItemParam,
+    ResponseInputParam,
+    ResponseOutputItem,
+    ResponseTextConfigParam,
+    ToolParam,
+    response_create_params,
+)
+from openai.types.shared_params import Metadata, Reasoning
+from pydantic import BaseModel, TypeAdapter, ValidationError
+
+from unique_toolkit.agentic.tools.utils.execution.execution import (
+    failsafe,
+)
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.language_model.constants import (
+    DEFAULT_COMPLETE_TEMPERATURE,
+)
+from unique_toolkit.language_model.functions import (
+    SearchContext,
+    _clamp_temperature,
+    _to_search_context,
+)
+from unique_toolkit.language_model.infos import (
+    LanguageModelInfo,
+    LanguageModelName,
+)
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelMessage,
+    LanguageModelMessageOptions,
+    LanguageModelMessageRole,
+    LanguageModelMessages,
+    LanguageModelSystemMessage,
+    LanguageModelToolDescription,
+    LanguageModelToolMessage,
+    LanguageModelUserMessage,
+    ResponsesLanguageModelStreamResponse,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _convert_tools_to_openai(
+    tools: Sequence[LanguageModelToolDescription | ToolParam],
+) -> list[ToolParam]:
+    openai_tools = []
+    for tool in tools:
+        if isinstance(tool, LanguageModelToolDescription):
+            openai_tools.append(tool.to_openai(mode="responses"))
+        else:
+            openai_tools.append(tool)
+    return openai_tools
+
+
+def _convert_message_to_openai(
+    message: LanguageModelMessageOptions,
+) -> ResponseInputParam:
+    res = []
+    match message:
+        case LanguageModelAssistantMessage():
+            return message.to_openai(mode="responses")  # type: ignore
+        case (
+            LanguageModelUserMessage()
+            | LanguageModelSystemMessage()
+            | LanguageModelToolMessage()
+        ):
+            return [message.to_openai(mode="responses")]
+        case _:
+            return _convert_message_to_openai(_convert_to_specific_message(message))
+    return res
+
+
+def _convert_to_specific_message(
+    message: LanguageModelMessage,
+) -> "LanguageModelSystemMessage | LanguageModelUserMessage | LanguageModelAssistantMessage":
+    match message.role:
+        case LanguageModelMessageRole.SYSTEM:
+            return LanguageModelSystemMessage(content=message.content)
+        case LanguageModelMessageRole.USER:
+            return LanguageModelUserMessage(content=message.content)
+        case LanguageModelMessageRole.ASSISTANT:
+            return LanguageModelAssistantMessage(content=message.content)
+        case LanguageModelMessageRole.TOOL:
+            raise ValueError(
+                "Cannot convert message with role `tool`. Please use `LanguageModelToolMessage` instead."
+            )
+
+
+def _convert_messages_to_openai(
+    messages: Sequence[
+        ResponseInputItemParam | LanguageModelMessageOptions | ResponseOutputItem
+    ],
+) -> ResponseInputParam:
+    res = []
+    for message in messages:
+        if isinstance(message, LanguageModelMessageOptions):
+            res.extend(_convert_message_to_openai(message))
+        elif isinstance(
+            message, dict
+        ):  # Openai uses dicts for their input and BaseModel as output
+            res.append(message)
+        else:
+            assert isinstance(message, BaseModel)
+            res.append(message.model_dump(exclude_defaults=True))
+
+    return res
+
+
+class _ResponsesParams(NamedTuple):
+    temperature: float
+    model_name: str
+    search_context: SearchContext | None
+    messages: str | ResponseInputParam
+    tools: list[ToolParam] | None
+    reasoning: Reasoning | None
+    text: ResponseTextConfigParam | None
+
+
+def _prepare_responses_params_util(
+    model_name: LanguageModelName | str,
+    content_chunks: list[ContentChunk] | None,
+    temperature: float,
+    tools: Sequence[LanguageModelToolDescription | ToolParam] | None,
+    messages: str
+    | LanguageModelMessages
+    | Sequence[
+        ResponseInputItemParam | LanguageModelMessageOptions | ResponseOutputItem
+    ],
+    reasoning: Reasoning | None,
+    text: ResponseTextConfigParam | None,
+    other_options: dict | None = None,
+) -> _ResponsesParams:
+    search_context = (
+        _to_search_context(content_chunks) if content_chunks is not None else None
+    )
+
+    model = model_name.name if isinstance(model_name, LanguageModelName) else model_name
+
+    tools_res = _convert_tools_to_openai(tools) if tools is not None else None
+
+    if other_options is not None:
+        # Key word argument takes precedence
+        reasoning = reasoning or _attempt_extract_reasoning_from_options(other_options)
+        text = text or _attempt_extract_verbosity_from_options(other_options)
+
+    if isinstance(model_name, LanguageModelName):
+        model_info = LanguageModelInfo.from_name(model_name)
+
+        if model_info.temperature_bounds is not None and temperature is not None:
+            temperature = _clamp_temperature(temperature, model_info.temperature_bounds)
+
+        if (
+            reasoning is None
+            and model_info.default_options is not None
+            and "reasoning_effort" in model_info.default_options
+        ):
+            reasoning = Reasoning(effort=model_info.default_options["reasoning_effort"])
+
+        if (
+            reasoning is not None
+            and tools_res is not None
+            and any(tool["type"] == "code_interpreter" for tool in tools_res)
+            and "effort" in reasoning
+            and reasoning["effort"] == "minimal"
+        ):
+            logger.warning(
+                "Code interpreter cannot be used with `minimal` effort. Switching to `low`."
+            )
+            reasoning["effort"] = (
+                "low"  # Code interpreter cannot be used with minimal effort
+            )
+
+    messages_res = None
+    if isinstance(messages, LanguageModelMessages):
+        messages_res = _convert_messages_to_openai(messages.root)
+    elif isinstance(messages, list):
+        messages_res = _convert_messages_to_openai(messages)
+    else:
+        assert isinstance(messages, str)
+        messages_res = messages
+
+    return _ResponsesParams(
+        temperature, model, search_context, messages_res, tools_res, reasoning, text
+    )
+
+
+@failsafe(
+    failure_return_value=None,
+    exceptions=(ValidationError,),
+    log_exc_info=False,
+    logger=logger,
+)
+def _attempt_extract_reasoning_from_options(options: dict) -> Reasoning | None:
+    reasoning = None
+
+    # Responses API
+    if "reasoning" in options:
+        reasoning = options["reasoning"]
+
+    # Completions API
+    elif "reasoning_effort" in options:
+        reasoning = {"effort": options["reasoning_effort"]}
+    if "reasoningEffort" in options:
+        reasoning = {"effort": options["reasoningEffort"]}
+
+    if reasoning is not None:
+        return TypeAdapter(Reasoning).validate_python(reasoning)
+
+    return None
+
+
+@failsafe(
+    failure_return_value=None,
+    exceptions=(ValidationError,),
+    log_exc_info=False,
+    logger=logger,
+)
+def _attempt_extract_verbosity_from_options(
+    options: dict,
+) -> ResponseTextConfigParam | None:
+    reasoning = None
+
+    # Responses API
+    if "text" in options:
+        reasoning = options["text"]
+
+    # Completions API
+    elif "verbosity" in options:
+        reasoning = {"verbosity": options["verbosity"]}
+
+    if reasoning is not None:
+        return TypeAdapter(ResponseTextConfigParam).validate_python(reasoning)
+
+    return None
+
+
+def _prepare_responses_args(
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    params: _ResponsesParams,
+    debug_info: dict | None,
+    start_text: str | None,
+    include: list[ResponseIncludable] | None,
+    instructions: str | None,
+    max_output_tokens: int | None,
+    metadata: Metadata | None,
+    parallel_tool_calls: bool | None,
+    tool_choice: response_create_params.ToolChoice | None,
+    top_p: float | None,
+    other_options: dict | None = None,
+) -> dict[str, Any]:
+    options = {}
+
+    options["company_id"] = company_id
+    options["user_id"] = user_id
+
+    options["model"] = params.model_name
+
+    if params.search_context is not None:
+        options["searchContext"] = params.search_context
+
+    options["chatId"] = chat_id
+    options["assistantId"] = assistant_id
+    options["assistantMessageId"] = assistant_message_id
+    options["userMessageId"] = user_message_id
+
+    if debug_info is not None:
+        options["debugInfo"] = debug_info
+    if start_text is not None:
+        options["startText"] = start_text
+
+    options["input"] = params.messages
+
+    openai_options: unique_sdk.Integrated.CreateStreamResponsesOpenaiParams = {}
+
+    if params.temperature is not None:
+        openai_options["temperature"] = params.temperature
+
+    if params.reasoning is not None:
+        openai_options["reasoning"] = params.reasoning
+
+    if params.text is not None:
+        openai_options["text"] = params.text
+
+    if include is not None:
+        openai_options["include"] = include
+
+    if instructions is not None:
+        openai_options["instructions"] = instructions
+
+    if max_output_tokens is not None:
+        openai_options["max_output_tokens"] = max_output_tokens
+
+    if metadata is not None:
+        openai_options["metadata"] = metadata
+
+    if parallel_tool_calls is not None:
+        openai_options["parallel_tool_calls"] = parallel_tool_calls
+
+    if tool_choice is not None:
+        openai_options["tool_choice"] = tool_choice
+
+    if params.tools is not None:
+        openai_options["tools"] = params.tools
+
+    if top_p is not None:
+        openai_options["top_p"] = top_p
+
+    # allow any other openai.resources.responses.Response.create options
+    if other_options is not None:
+        openai_options.update(other_options)  # type: ignore
+
+    options["options"] = openai_options
+
+    return options
+
+
+def stream_responses_with_references(
+    *,
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    model_name: LanguageModelName | str,
+    messages: str
+    | LanguageModelMessages
+    | Sequence[
+        ResponseInputItemParam | LanguageModelMessageOptions | ResponseOutputItem
+    ],
+    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:
+    responses_params = _prepare_responses_params_util(
+        model_name=model_name,
+        content_chunks=content_chunks,
+        temperature=temperature,
+        tools=tools,
+        messages=messages,
+        reasoning=reasoning,
+        text=text,
+        other_options=other_options,
+    )
+
+    responses_args = _prepare_responses_args(
+        company_id=company_id,
+        user_id=user_id,
+        assistant_message_id=assistant_message_id,
+        user_message_id=user_message_id,
+        chat_id=chat_id,
+        assistant_id=assistant_id,
+        params=responses_params,
+        debug_info=debug_info,
+        start_text=start_text,
+        include=include,
+        instructions=instructions,
+        max_output_tokens=max_output_tokens,
+        metadata=metadata,
+        parallel_tool_calls=parallel_tool_calls,
+        tool_choice=tool_choice,
+        top_p=top_p,
+        other_options=other_options,
+    )
+
+    return ResponsesLanguageModelStreamResponse.model_validate(
+        unique_sdk.Integrated.responses_stream(
+            **responses_args,
+        )
+    )
+
+
+async def stream_responses_with_references_async(
+    *,
+    company_id: str,
+    user_id: str,
+    assistant_message_id: str,
+    user_message_id: str,
+    chat_id: str,
+    assistant_id: str,
+    model_name: LanguageModelName | str,
+    messages: str
+    | LanguageModelMessages
+    | Sequence[
+        ResponseInputItemParam | LanguageModelMessageOptions | ResponseOutputItem
+    ],
+    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:
+    responses_params = _prepare_responses_params_util(
+        model_name=model_name,
+        content_chunks=content_chunks,
+        temperature=temperature,
+        tools=tools,
+        messages=messages,
+        reasoning=reasoning,
+        text=text,
+        other_options=other_options,
+    )
+
+    responses_args = _prepare_responses_args(
+        company_id=company_id,
+        user_id=user_id,
+        assistant_message_id=assistant_message_id,
+        user_message_id=user_message_id,
+        chat_id=chat_id,
+        assistant_id=assistant_id,
+        params=responses_params,
+        debug_info=debug_info,
+        start_text=start_text,
+        include=include,
+        instructions=instructions,
+        max_output_tokens=max_output_tokens,
+        metadata=metadata,
+        parallel_tool_calls=parallel_tool_calls,
+        tool_choice=tool_choice,
+        top_p=top_p,
+        other_options=other_options,
+    )
+
+    return ResponsesLanguageModelStreamResponse.model_validate(
+        await unique_sdk.Integrated.responses_stream_async(
+            **responses_args,
+        )
+    )

unique_toolkit/language_model/functions.py

@@ -1,10 +1,11 @@
 import copy
 import logging
 from datetime import UTC, datetime
-from typing import Any, cast
+from typing import Any, Sequence, cast
 
 import humps
 import unique_sdk
+from openai.types.chat import ChatCompletionToolChoiceOptionParam
 from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 from pydantic import BaseModel
 
@@ -159,7 +160,7 @@ async def complete_async(
 
 def _add_tools_to_options(
     options: dict,
-    tools: list[LanguageModelTool | LanguageModelToolDescription] | None,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None,
 ) -> dict:
     if tools:
         options["tools"] = [
@@ -172,7 +173,12 @@ def _add_tools_to_options(
     return options
 
 
-def _to_search_context(chunks: list[ContentChunk]) -> dict | None:
+SearchContext = list[unique_sdk.Integrated.SearchResult]
+
+
+def _to_search_context(
+    chunks: list[ContentChunk],
+) -> SearchContext | None:
     if not chunks:
         return None
     return [
@@ -211,12 +217,12 @@ def _prepare_completion_params_util(
     messages: LanguageModelMessages,
     model_name: LanguageModelName | str,
     temperature: float,
-    tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
     structured_output_model: type[BaseModel] | None = None,
     structured_output_enforce_schema: bool = False,
-) -> tuple[dict, str, dict, dict | None]:
+) -> tuple[dict, str, dict, SearchContext | None]:
     """Prepare common parameters for completion requests.
 
     Returns
@@ -259,12 +265,12 @@ def _prepare_completion_params_util(
 def _prepare_openai_completion_params_util(
     model_name: LanguageModelName | str,
     temperature: float,
-    tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
     structured_output_model: type[BaseModel] | None = None,
     structured_output_enforce_schema: bool = False,
-) -> tuple[dict, str, dict | None]:
+) -> tuple[dict, str, SearchContext | None]:
     """Prepare common parameters for completion requests.
 
     Returns
@@ -330,18 +336,28 @@ def _prepare_all_completions_params_util(
     messages: LanguageModelMessages | list[ChatCompletionMessageParam],
     model_name: LanguageModelName | str,
     temperature: float,
-    tools: list[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
     other_options: dict | None = None,
     content_chunks: list[ContentChunk] | None = None,
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
     structured_output_model: type[BaseModel] | None = None,
     structured_output_enforce_schema: bool = False,
 ) -> tuple[
     dict,
     str,
     list[unique_sdk.Integrated.ChatCompletionRequestMessage],
-    dict | None,
+    SearchContext | None,
 ]:
     model_info = None
+
+    other_options = copy.deepcopy(other_options)
+
+    if tool_choice is not None:
+        if other_options is None:
+            other_options = {}
+        if "toolChoice" not in other_options:
+            other_options["toolChoice"] = tool_choice  # Backend expects CamelCase
+
     if isinstance(model_name, LanguageModelName):
         model_info = LanguageModelInfo.from_name(model_name)
         other_options = _prepare_other_options(