Update routes.py

routes.py

@@ -4,86 +4,82 @@ 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 models import (
+    InitializeBotResponse, DocumentPath, NewChatResponse,
+    QueryRequest, QueryResponse, CreateBotRequest, CreateBotResponse
+)
 from trainer_manager import get_trainer
-from config import CUSTOM_PROMPT
 from prompt_templates import PromptTemplates
 
 router = APIRouter()
 trainer = get_trainer()
+
+
 @router.post("/initialize_bot", response_model=InitializeBotResponse)
-def initialize_bot(prompt_type: str):
+def initialize_bot():
     """
-    Initializes a new bot with a custom prompt template.
-    The custom_prompt parameter should be provided by the frontend (for example, via navigation).
+    Initializes a new bot and returns the bot_id.
+    The prompt type is now passed during bot creation, not initialization.
     """
-    
     try:
         bot_id = trainer.initialize_bot_id()
-        
-        # Assign the prompt template based on the prompt type.
-        if prompt_type == "university":
-            prompt_template = PromptTemplates.get_university_chatbot_prompt()
-        elif prompt_type == "quiz_solving":
-            prompt_template = PromptTemplates.get_quiz_solving_prompt()
-        elif prompt_type == "assignment_solving":
-            prompt_template = PromptTemplates.get_assignment_solving_prompt()
-        elif prompt_type == "paper_solving":
-            prompt_template = PromptTemplates.get_paper_solving_prompt()
-        elif prompt_type == "quiz_creation":
-            prompt_template = PromptTemplates.get_quiz_creation_prompt()
-        elif prompt_type == "assignment_creation":
-            prompt_template = PromptTemplates.get_assignment_creation_prompt()
-        elif prompt_type == "paper_creation":
-            prompt_template = PromptTemplates.get_paper_creation_prompt()
-        elif prompt_type == "check_quiz":
-            prompt_template = PromptTemplates.get_check_quiz_prompt()
-        elif prompt_type == "check_assignment":
-            prompt_template = PromptTemplates.get_check_assignment_prompt()
-        elif prompt_type == "check_paper":
-            prompt_template = PromptTemplates.get_check_paper_prompt()
-        else:
-            # Default prompt template if prompt_type is not recognized.
-            prompt_template = PromptTemplates.get_quiz_solving_prompt()
-        
-        # Pass the assigned prompt template to the bot.
-        trainer.set_custom_prompt_template(bot_id, prompt_template)
         return InitializeBotResponse(bot_id=bot_id)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
+@router.post("/create_bot", response_model=CreateBotResponse)
+def create_bot(data: CreateBotRequest):
+    """
+    Finalizes bot creation by assigning a prompt template based on the prompt type.
+    """
+    try:
+        # Assign the appropriate prompt template based on prompt_type
+        prompt_type = data.prompt_type.lower()
+        prompt_template_mapping = {
+            "university": PromptTemplates.get_university_chatbot_prompt,
+            "quiz_solving": PromptTemplates.get_quiz_solving_prompt,
+            "assignment_solving": PromptTemplates.get_assignment_solving_prompt,
+            "paper_solving": PromptTemplates.get_paper_solving_prompt,
+            "quiz_creation": PromptTemplates.get_quiz_creation_prompt,
+            "assignment_creation": PromptTemplates.get_assignment_creation_prompt,
+            "paper_creation": PromptTemplates.get_paper_creation_prompt,
+            "check_quiz": PromptTemplates.get_check_quiz_prompt,
+            "check_assignment": PromptTemplates.get_check_assignment_prompt,
+            "check_paper": PromptTemplates.get_check_paper_prompt
+        }
+
+        prompt_template = prompt_template_mapping.get(prompt_type, PromptTemplates.get_quiz_solving_prompt)()
+
+        # Apply the selected prompt template to the bot
+        trainer.set_custom_prompt_template(data.bot_id, prompt_template)
+
+        # Complete bot creation
+        trainer.create_bot(data.bot_id)
+
+        return CreateBotResponse(bot_id=data.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.
     """
     try:
-        # Save the 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 using 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:
         raise HTTPException(status_code=500, detail=str(e))
 
-@router.post("/create_bot/{bot_id}")
-def create_bot(bot_id: str):
-    """
-    Creates (builds) the bot (e.g., builds its index) for the given bot_id.
-    """
-    try:
-        trainer.create_bot(bot_id)
-        return {"message": f"Bot {bot_id} created successfully."}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
 
 @router.post("/new_chat/{bot_id}", response_model=NewChatResponse)
 def new_chat(bot_id: str):
@@ -96,47 +92,40 @@ def new_chat(bot_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @router.post("/query", response_model=QueryResponse)
 def send_query(query_request: QueryRequest):
     """
     Processes a query and returns the bot's response along with any web sources.
     """
-    # trainer = get_trainer()
-    # trainer.load_bot(query_request.bot_id)
-    prompt_type = query_request.prompt_type
     try:
-        # Assign the prompt template based on the prompt type.
-        if prompt_type == "university":
-            prompt_template = PromptTemplates.get_university_chatbot_prompt()
-        elif prompt_type == "quiz_solving":
-            prompt_template = PromptTemplates.get_quiz_solving_prompt()
-        elif prompt_type == "assignment_solving":
-            prompt_template = PromptTemplates.get_assignment_solving_prompt()
-        elif prompt_type == "paper_solving":
-            prompt_template = PromptTemplates.get_paper_solving_prompt()
-        elif prompt_type == "quiz_creation":
-            prompt_template = PromptTemplates.get_quiz_creation_prompt()
-        elif prompt_type == "assignment_creation":
-            prompt_template = PromptTemplates.get_assignment_creation_prompt()
-        elif prompt_type == "paper_creation":
-            prompt_template = PromptTemplates.get_paper_creation_prompt()
-        elif prompt_type == "check_quiz":
-            prompt_template = PromptTemplates.get_check_quiz_prompt()
-        elif prompt_type == "check_assignment":
-            prompt_template = PromptTemplates.get_check_assignment_prompt()
-        elif prompt_type == "check_paper":
-            prompt_template = PromptTemplates.get_check_paper_prompt()
-        else:
-            # Default prompt template if prompt_type is not recognized.
-            prompt_template = PromptTemplates.get_quiz_solving_prompt()
+        prompt_type = query_request.prompt_type.lower()
+        prompt_template_mapping = {
+            "university": PromptTemplates.get_university_chatbot_prompt,
+            "quiz_solving": PromptTemplates.get_quiz_solving_prompt,
+            "assignment_solving": PromptTemplates.get_assignment_solving_prompt,
+            "paper_solving": PromptTemplates.get_paper_solving_prompt,
+            "quiz_creation": PromptTemplates.get_quiz_creation_prompt,
+            "assignment_creation": PromptTemplates.get_assignment_creation_prompt,
+            "paper_creation": PromptTemplates.get_paper_creation_prompt,
+            "check_quiz": PromptTemplates.get_check_quiz_prompt,
+            "check_assignment": PromptTemplates.get_check_assignment_prompt,
+            "check_paper": PromptTemplates.get_check_paper_prompt
+        }
+
+        prompt_template = prompt_template_mapping.get(prompt_type, PromptTemplates.get_quiz_solving_prompt)()
+
         trainer.set_custom_prompt_template(query_request.bot_id, prompt_template)
+
         response, web_sources = trainer.get_response(
             query_request.query, query_request.bot_id, query_request.chat_id
         )
+
         return QueryResponse(response=response, web_sources=web_sources)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @router.get("/list_chats/{bot_id}")
 def list_chats(bot_id: str):
     """
@@ -148,11 +137,11 @@ def list_chats(bot_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+
 @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 in the history are converted to strings.
     """
     try:
         history = trainer.get_chat_by_id(chat_id=chat_id)