diff --git a/src/chat/__init__.py b/src/chat/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/chat/constants.py b/src/chat/constants.py new file mode 100644 index 0000000..8ad52a0 --- /dev/null +++ b/src/chat/constants.py @@ -0,0 +1,15 @@ +"""Chat completion constants.""" + +# OpenAI API object types +CHAT_COMPLETION_OBJECT = "chat.completion" +CHAT_COMPLETION_CHUNK_OBJECT = "chat.completion.chunk" + +# Finish reasons +FINISH_REASON_STOP = "stop" +FINISH_REASON_LENGTH = "length" +FINISH_REASON_ERROR = "error" + +# Roles +ROLE_SYSTEM = "system" +ROLE_USER = "user" +ROLE_ASSISTANT = "assistant" diff --git a/src/chat/dependencies.py b/src/chat/dependencies.py new file mode 100644 index 0000000..e69de29 diff --git a/src/chat/router.py b/src/chat/router.py new file mode 100644 index 0000000..fbcc851 --- /dev/null +++ b/src/chat/router.py @@ -0,0 +1,68 @@ +""" +Chat completion router. +OpenAI-compatible /v1/chat/completions endpoint. +""" +import json +import logging +from typing import AsyncGenerator + +from fastapi import APIRouter +from sse_starlette.sse import EventSourceResponse + +from src.chat import service +from src.chat.schemas import ( + ChatCompletionRequest, + ChatCompletionResponse, +) + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/chat", tags=["chat"]) + + +async def _stream_response( + request: ChatCompletionRequest, +) -> AsyncGenerator[dict, None]: + """ + Generate SSE stream for chat completion. + + EventSourceResponse adds "data: " prefix automatically. + We just yield the dict/string content. + """ + try: + async for chunk in service.create_chat_completion_stream(request): + # Yield dict - EventSourceResponse will format as SSE + yield {"data": chunk.model_dump_json()} + + # Send [DONE] message + yield {"data": "[DONE]"} + + except Exception as e: + logger.error(f"Error in streaming response: {e}") + error_data = {"error": {"message": str(e), "type": "internal_error"}} + yield {"data": json.dumps(error_data)} + + +@router.post("/completions", response_model=ChatCompletionResponse) +async def create_chat_completion( + request: ChatCompletionRequest, +) -> ChatCompletionResponse | EventSourceResponse: + """ + Create chat completion (OpenAI-compatible). + + Supports both regular and streaming responses. + Currently returns mock lorem ipsum responses. + + Args: + request: Chat completion request + + Returns: + Chat completion response or SSE stream + """ + logger.info(f"Chat completion request for model: {request.model}") + + if request.stream: + logger.info("Streaming response requested") + return EventSourceResponse(_stream_response(request)) + + return await service.create_chat_completion(request) diff --git a/src/chat/schemas.py b/src/chat/schemas.py new file mode 100644 index 0000000..aed6045 --- /dev/null +++ b/src/chat/schemas.py @@ -0,0 +1,73 @@ +""" +OpenAI-compatible chat completion schemas. +Following OpenAI API specification for compatibility. +""" +from typing import Literal + +from pydantic import Field + +from src.core.models import CustomBaseModel + + +class ChatMessage(CustomBaseModel): + """OpenAI-compatible chat message.""" + role: Literal["system", "user", "assistant"] + content: str + name: str | None = None + + +class ChatCompletionRequest(CustomBaseModel): + """OpenAI-compatible chat completion request.""" + model: str = Field(..., description="Model to use for completion") + messages: list[ChatMessage] = Field(..., description="List of messages") + temperature: float | None = Field(default=0.7, ge=0.0, le=2.0) + top_p: float | None = Field(default=1.0, ge=0.0, le=1.0) + max_tokens: int | None = Field(default=None, ge=1) + stream: bool = Field(default=False, description="Enable streaming") + stop: str | list[str] | None = None + + +class ChatCompletionChoice(CustomBaseModel): + """Choice in chat completion response.""" + index: int + message: ChatMessage + finish_reason: str | None + + +class ChatCompletionUsage(CustomBaseModel): + """Token usage information.""" + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(CustomBaseModel): + """OpenAI-compatible chat completion response.""" + id: str + object: str = "chat.completion" + created: int + model: str + choices: list[ChatCompletionChoice] + usage: ChatCompletionUsage | None = None + + +class ChatCompletionChunkDelta(CustomBaseModel): + """Delta in streaming chunk.""" + role: str | None = None + content: str | None = None + + +class ChatCompletionChunkChoice(CustomBaseModel): + """Choice in streaming chunk.""" + index: int + delta: ChatCompletionChunkDelta + finish_reason: str | None = None + + +class ChatCompletionChunk(CustomBaseModel): + """OpenAI-compatible streaming chunk.""" + id: str + object: str = "chat.completion.chunk" + created: int + model: str + choices: list[ChatCompletionChunkChoice] diff --git a/src/chat/service.py b/src/chat/service.py new file mode 100644 index 0000000..9cea384 --- /dev/null +++ b/src/chat/service.py @@ -0,0 +1,132 @@ +""" +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, + ) + ], + )