You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
swarms/playground/demos/chatbot/server/server_models.py

60 lines
1.5 KiB

""" Chatbot Server API Models """
from swarms.prompts.chat_prompt import Role
from strenum import StrEnum
from pydantic import BaseModel
from swarms.prompts import QA_PROMPT_TEMPLATE_STR as DefaultSystemPrompt
class AIModel(BaseModel):
""" Defines the model a user selected. """
id: str
name: str
maxLength: int
tokenLimit: int
class State(StrEnum):
""" State of RAGFile that's been uploaded. """
UNAVAILABLE = "UNAVAILABLE"
PROCESSING = "PROCESSING"
PROCESSED = "PROCESSED"
class RAGFile(BaseModel):
""" Defines a file uploaded by the users for RAG processing. """
filename: str
title: str
username: str
state: State = State.UNAVAILABLE
class RAGFiles(BaseModel):
""" Defines a list of RAGFile objects. """
files: list[RAGFile]
class Message(BaseModel):
""" Defines the type of a Message with a role and content. """
role: Role
content: str
class ChatRequest(BaseModel):
""" The model for a ChatRequest for theChatbot Chat POST endpoint"""
id: str
model: AIModel = AIModel(
id="NousResearch/Meta-Llama-3.1-8B-Instruct",
name="NousResearch/Meta-Llama-3.1-8B-Instruct",
maxLength=2048,
tokenLimit=2048,
)
messages: list[Message] = [
Message(role=Role.AI, content="Hello, how may I help you?"),
Message(role=Role.HUMAN, content="What is Swarms?"),
]
maxTokens: int = 2048
temperature: float = 0
prompt: str = DefaultSystemPrompt
file: RAGFile = RAGFile(filename="None", title="None", username="None")