Update routes.py

routes.py

@@ -4,22 +4,18 @@ 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, CreateBotRequest, CreateBotResponse
-)
+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()
 trainer = get_trainer()
 
-
 @router.post("/initialize_bot", response_model=InitializeBotResponse)
 def initialize_bot():
     """
-    Initializes a new bot and returns the bot_id.
-    The prompt type is now passed during bot creation, not initialization.
+    Initializes a new bot and returns the generated bot_id.
     """
     try:
         bot_id = trainer.initialize_bot_id()
@@ -28,59 +24,68 @@ def initialize_bot():
         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)
-        os.remove(tmp_path)
 
+        # 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, prompt_type: str = Form(...)):
+    """
+    Creates (builds) the bot for the given bot_id using the specified prompt type.
+    This endpoint assigns the appropriate custom prompt template before finalizing the bot creation.
+    """
+    try:
+        # Determine the appropriate prompt template based on the provided 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()
+        
+        # Set the custom prompt template and create the bot.
+        trainer.set_custom_prompt_template(bot_id, prompt_template)
+        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):
     """
@@ -98,17 +103,42 @@ def send_query(query_request: QueryRequest):
     """
     Processes a query and returns the bot's response along with any web sources.
     """
+    prompt_type = query_request.prompt_type
     try:
+        # Assign the prompt template based on the provided 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()
+        
+        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):
     """
@@ -125,6 +155,7 @@ def list_chats(bot_id: str):
 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)