Update routes.py

routes.py

@@ -2,6 +2,8 @@ import os
 import shutil
 import tempfile
 from fastapi import APIRouter, HTTPException, UploadFile, File, Form
+from fastapi.encoders import jsonable_encoder
+from bson import ObjectId
 from models import InitializeBotResponse, DocumentPath, NewChatResponse, QueryRequest, QueryResponse
 from trainer_manager import get_trainer
 from config import CUSTOM_PROMPT
@@ -9,34 +11,60 @@ from prompt_templates import PromptTemplates
 
 router = APIRouter()
 
+def select_prompt(prompt_type: str):
+    """Selects a prompt template based on the provided prompt_type."""
+    if prompt_type == "university":
+        return PromptTemplates.get_university_chatbot_prompt()
+    elif prompt_type == "quiz_solving":
+        return PromptTemplates.get_quiz_solving_prompt()
+    elif prompt_type == "assignment_solving":
+        return PromptTemplates.get_assignment_solving_prompt()
+    elif prompt_type == "paper_solving":
+        return PromptTemplates.get_paper_solving_prompt()
+    elif prompt_type == "quiz_creation":
+        return PromptTemplates.get_quiz_creation_prompt()
+    elif prompt_type == "assignment_creation":
+        return PromptTemplates.get_assignment_creation_prompt()
+    elif prompt_type == "paper_creation":
+        return PromptTemplates.get_paper_creation_prompt()
+    elif prompt_type == "check_quiz":
+        return PromptTemplates.get_check_quiz_prompt()
+    elif prompt_type == "check_assignment":
+        return PromptTemplates.get_check_assignment_prompt()
+    elif prompt_type == "check_paper":
+        return PromptTemplates.get_check_paper_prompt()
+    else:
+        # Fallback to a default prompt if no matching type is found
+        return PromptTemplates.get_quiz_solving_prompt()
+
 @router.post("/initialize_bot", response_model=InitializeBotResponse)
-def initialize_bot(custom_prompt: str):
+def initialize_bot(prompt_type: str):
     """
     Initializes a new bot with a custom prompt template.
-    The custom_prompt parameter should be provided (e.g., via navigation) and is used as the bot's prompt template.
+    The prompt_type query parameter determines which prompt to use.
     """
     trainer = get_trainer()
     try:
         bot_id = trainer.initialize_bot_id()
-        trainer.set_custom_prompt_template(bot_id, custom_prompt)
+        prompt = select_prompt(prompt_type)
+        trainer.set_custom_prompt_template(bot_id, prompt)
         return InitializeBotResponse(bot_id=bot_id)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @router.post("/upload_document")
 async def upload_document(bot_id: str = Form(...), file: UploadFile = File(...)):
+    """
+    Saves the uploaded file temporarily and adds it to the bot's knowledge base.
+    """
     trainer = get_trainer()
     try:
-        # Save the uploaded file to a temporary location
         with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.filename)[1]) as tmp:
             contents = await file.read()
             tmp.write(contents)
             tmp_path = tmp.name
 
-        # Add the document from the temporary file path
         trainer.add_document_from_path(tmp_path, bot_id)
-
-        # Remove the temporary file
         os.remove(tmp_path)
         return {"message": "Document uploaded and added successfully."}
     except Exception as e:
@@ -44,6 +72,9 @@ async def upload_document(bot_id: str = Form(...), file: UploadFile = File(...))
 
 @router.post("/create_bot/{bot_id}")
 def create_bot(bot_id: str):
+    """
+    Creates the bot (builds its index) for the given bot_id.
+    """
     trainer = get_trainer()
     try:
         trainer.create_bot(bot_id)
@@ -53,6 +84,9 @@ def create_bot(bot_id: str):
 
 @router.post("/new_chat/{bot_id}", response_model=NewChatResponse)
 def new_chat(bot_id: str):
+    """
+    Creates a new chat session for the bot.
+    """
     trainer = get_trainer()
     try:
         chat_id = trainer.new_chat(bot_id)
@@ -62,6 +96,9 @@ def new_chat(bot_id: str):
 
 @router.post("/query", response_model=QueryResponse)
 def send_query(query_request: QueryRequest):
+    """
+    Processes a user's query and returns the bot's response along with any web sources.
+    """
     trainer = get_trainer()
     try:
         response, web_sources = trainer.get_response(
@@ -73,6 +110,9 @@ def send_query(query_request: QueryRequest):
 
 @router.get("/list_chats/{bot_id}")
 def list_chats(bot_id: str):
+    """
+    Returns a list of previous chat sessions for the specified bot.
+    """
     trainer = get_trainer()
     try:
         chats = trainer.list_chats(bot_id)
@@ -80,15 +120,15 @@ def list_chats(bot_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-from fastapi.encoders import jsonable_encoder
-from bson import ObjectId
-
 @router.get("/chat_history/{chat_id}")
 def chat_history(chat_id: str, bot_id: str):
+    """
+    Returns the chat history for a given chat session.
+    ObjectId instances are converted to strings.
+    """
     trainer = get_trainer()
     try:
         history = trainer.get_chat_by_id(chat_id=chat_id)
-        # Convert ObjectId instances to strings before returning the response.
         return jsonable_encoder(history, custom_encoder={ObjectId: str})
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))