tatlock/src/chat/service.py

"""
Chat completion service.
Currently returns mock responses with lorem ipsum.
TODO: Integrate with Ollama/PydanticAI in future.
"""
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,
)


# Mock lorem ipsum response
MOCK_RESPONSE = (
    "Lorem ipsum dolor sit amet, consectetur adipiscing elit. "
    "Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. "
    "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris."
)


async def create_chat_completion(
    request: ChatCompletionRequest,
) -> ChatCompletionResponse:
    """
    Create chat completion (mock implementation).
    
    Args:
        request: Chat completion request
        
    Returns:
        Mock chat completion response with lorem ipsum
    """
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created_at = int(time.time())
    
    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=MOCK_RESPONSE,
                ),
                finish_reason=constants.FINISH_REASON_STOP,
            )
        ],
        usage=ChatCompletionUsage(
            prompt_tokens=len(" ".join(m.content for m in request.messages).split()),
            completion_tokens=len(MOCK_RESPONSE.split()),
            total_tokens=len(" ".join(m.content for m in request.messages).split())
                        + len(MOCK_RESPONSE.split()),
        ),
    )


async def create_chat_completion_stream(
    request: ChatCompletionRequest,
) -> AsyncGenerator[ChatCompletionChunk, None]:
    """
    Create streaming chat completion (mock implementation).
    
    Args:
        request: Chat completion request with stream=True
        
    Yields:
        Mock chat completion chunks with lorem ipsum
    """
    completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
    created_at = int(time.time())
    
    # Split response into words for streaming simulation
    words = MOCK_RESPONSE.split()
    
    # 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 words
    for word in words:
        yield ChatCompletionChunk(
            id=completion_id,
            object=constants.CHAT_COMPLETION_CHUNK_OBJECT,
            created=created_at,
            model=request.model,
            choices=[
                ChatCompletionChunkChoice(
                    index=0,
                    delta=ChatCompletionChunkDelta(content=f"{word} "),
                    finish_reason=None,
                )
            ],
        )
    
    # 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,
            )
        ],
    )