service/dto.py

from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
from llama_index.core.llms import MessageRole

class MetadataRequest(BaseModel):
    title: str
    category : str
    author: str
    year: int
    publisher: str
    
class DeleteById(BaseModel):
    id : str

class UserPromptRequest(BaseModel): 
    session_id : str
    prompt : str
    streaming : bool
    
class BotResponse(BaseModel):
    role : str = "assistant"
    content : str
    raw_references : List
    references : List
    metadata : List
    scores : List

class BotResponseStreaming(BaseModel):
    role : str = "assistant"
    content : Optional[str] = None
    completed_content : Optional[str] = None
    reference : Optional[str] = None
    metadata : Optional[Dict] = None
    score : Optional[float] = None
    
class TestStreaming(BaseModel):
    role : str = "assistant"
    content : str
    
class ChatMessage(BaseModel):
    """Chat message."""

    role: MessageRole = MessageRole.ASSISTANT
    content: Optional[Any] = ""
    metadata: List

    def __str__(self) -> str:
        return f"{self.role.value}: {self.content}"