tatlock/src/chat/service.py

"""
Chat completion service.

Wrapper around Responses API that converts to Chat Completions format.
Embeds reasoning in <think> tags for Open WebUI compatibility.
"""
import asyncio
import time
import uuid
from typing import AsyncGenerator

from src.chat import constants
from src.chat.schemas import (
    ChatCompletionChunk,
    ChatCompletionChunkChoice,
    ChatCompletionChunkDelta,
    ChatCompletionChoice,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionUsage,
    ChatMessage,
)
from src.responses.schemas import ResponseRequest
from src.responses.service import create_response, create_response_with_steward


async def create_chat_completion(
    request: ChatCompletionRequest,
) -> ChatCompletionResponse:
    """
    Create chat completion by wrapping Responses API.

    Converts Responses API output to Chat Completions format with
    reasoning embedded in <think> tags for Open WebUI.

    Args:
        request: Chat completion request

    Returns:
        Chat completion response with reasoning as <think> tags
    """
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created_at = int(time.time())

    # Convert Chat request to Responses request
    input_messages = [
        {"role": msg.role, "content": msg.content}
        for msg in request.messages
    ]

    response_request = ResponseRequest(
        model=request.model,
        input=input_messages,
        reasoning={"effort": "medium", "summary": "auto"},  # Enable reasoning
        temperature=request.temperature or 1.0,
        max_output_tokens=request.max_tokens,
        stop=request.stop if isinstance(request.stop, list) else ([request.stop] if request.stop else None),
    )

    # Call Responses API (will use Steward for Tatlock)
    model_id = request.model
    if "." in model_id:
        model_id = model_id.split(".", 1)[1]

    use_steward = model_id.lower() == "tatlock"

    if use_steward:
        response = await create_response_with_steward(response_request)
    else:
        response = await create_response(response_request)

    # Convert Responses API output to Chat format
    content_parts = []

    for item in response.output:
        if item.type == "reasoning":
            reasoning_text = "\n".join(item.summary)
            content_parts.append(f"<think>\n{reasoning_text}\n</think>\n\n")
        elif item.type == "message":
            content_parts.append(item.content[0].text)

    content = "".join(content_parts)

    return ChatCompletionResponse(
        id=completion_id,
        object=constants.CHAT_COMPLETION_OBJECT,
        created=created_at,
        model=request.model,
        choices=[
            ChatCompletionChoice(
                index=0,
                message=ChatMessage(
                    role=constants.ROLE_ASSISTANT,
                    content=content,
                ),
                finish_reason=constants.FINISH_REASON_STOP,
            )
        ],
        usage=ChatCompletionUsage(
            prompt_tokens=response.usage.input_tokens,
            completion_tokens=response.usage.output_tokens,
            total_tokens=response.usage.total_tokens,
        ),
    )


async def create_chat_completion_stream(
    request: ChatCompletionRequest,
) -> AsyncGenerator[ChatCompletionChunk, None]:
    """
    Create streaming chat completion by wrapping Responses API.

    Streams reasoning in <think> tags followed by message content.

    Args:
        request: Chat completion request with stream=True

    Yields:
        Chat completion chunks with reasoning as <think> tags
    """
    from src.responses.streaming import StreamingCoordinator, StreamEventType

    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created_at = int(time.time())

    # Convert Chat request to Responses request
    input_messages = [
        {"role": msg.role, "content": msg.content}
        for msg in request.messages
    ]

    response_request = ResponseRequest(
        model=request.model,
        input=input_messages,
        reasoning={"effort": "medium", "summary": "auto"},
        temperature=request.temperature or 1.0,
        max_output_tokens=request.max_tokens,
        stop=request.stop if isinstance(request.stop, list) else ([request.stop] if request.stop else None),
        stream=True,
    )

    # Determine if we should use Steward
    model_id = request.model
    if "." in model_id:
        model_id = model_id.split(".", 1)[1]

    use_steward = model_id.lower() == "tatlock"

    # First chunk with role
    yield ChatCompletionChunk(
        id=completion_id,
        object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
        created=created_at,
        model=request.model,
        choices=[
            ChatCompletionChunkChoice(
                index=0,
                delta=ChatCompletionChunkDelta(role=constants.ROLE_ASSISTANT),
                finish_reason=None,
            )
        ],
    )

    # Stream from Responses API
    coordinator = StreamingCoordinator()
    in_reasoning = False

    if use_steward:
        stream_generator = coordinator.stream_response_with_steward(response_request)
    else:
        stream_generator = coordinator.stream_response(response_request)

    async for event in stream_generator:
        if event.event == StreamEventType.REASONING_SUMMARY_DELTA:
            # Start <think> block if needed
            if not in_reasoning:
                yield ChatCompletionChunk(
                    id=completion_id,
                    object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
                    created=created_at,
                    model=request.model,
                    choices=[
                        ChatCompletionChunkChoice(
                            index=0,
                            delta=ChatCompletionChunkDelta(content="<think>\n"),
                            finish_reason=None,
                        )
                    ],
                )
                in_reasoning = True

            # Stream reasoning delta
            yield ChatCompletionChunk(
                id=completion_id,
                object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
                created=created_at,
                model=request.model,
                choices=[
                    ChatCompletionChunkChoice(
                        index=0,
                        delta=ChatCompletionChunkDelta(content=event.delta),
                        finish_reason=None,
                    )
                ],
            )

        elif event.event == StreamEventType.REASONING_SUMMARY_DONE:
            # Close <think> block
            if in_reasoning:
                yield ChatCompletionChunk(
                    id=completion_id,
                    object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
                    created=created_at,
                    model=request.model,
                    choices=[
                        ChatCompletionChunkChoice(
                            index=0,
                            delta=ChatCompletionChunkDelta(content="</think>\n\n"),
                            finish_reason=None,
                        )
                    ],
                )
                in_reasoning = False

        elif event.event == StreamEventType.OUTPUT_TEXT_DELTA:
            # Stream message content
            yield ChatCompletionChunk(
                id=completion_id,
                object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
                created=created_at,
                model=request.model,
                choices=[
                    ChatCompletionChunkChoice(
                        index=0,
                        delta=ChatCompletionChunkDelta(content=event.delta),
                        finish_reason=None,
                    )
                ],
            )

        elif event.event == StreamEventType.RESPONSE_DONE:
            # Final chunk with finish_reason
            yield ChatCompletionChunk(
                id=completion_id,
                object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
                created=created_at,
                model=request.model,
                choices=[
                    ChatCompletionChunkChoice(
                        index=0,
                        delta=ChatCompletionChunkDelta(),
                        finish_reason=constants.FINISH_REASON_STOP,
                    )
                ],
            )