Update routes.py

routes.py

@@ -6,15 +6,47 @@ from fastapi import APIRouter, HTTPException, UploadFile, File, Form
 from models import InitializeBotResponse, DocumentPath, NewChatResponse, QueryRequest, QueryResponse
 from trainer_manager import get_trainer
 from config import CUSTOM_PROMPT
+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():
+def initialize_bot(prompt_type: str = "quiz_solving"):
+    """
+    Initializes a new bot with a custom 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)
+        selected_prompt = select_prompt(prompt_type)
+        trainer.set_custom_prompt_template(bot_id, selected_prompt)
         return InitializeBotResponse(bot_id=bot_id)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -32,7 +64,7 @@ def add_document(document: DocumentPath):
 async def upload_document(bot_id: str = Form(...), file: UploadFile = File(...)):
     trainer = get_trainer()
     try:
-        # Save the uploaded file to a temporary directory
+        # 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)