routes.py

# routes.py
import os
import shutil
import tempfile
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

router = APIRouter()

@router.post("/initialize_bot", response_model=InitializeBotResponse)
def initialize_bot():
    trainer = get_trainer()
    try:
        bot_id = trainer.initialize_bot_id()
        trainer.set_custom_prompt_template(bot_id, CUSTOM_PROMPT)
        return InitializeBotResponse(bot_id=bot_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/add_document")
def add_document(document: DocumentPath):
    trainer = get_trainer()
    try:
        trainer.add_document_from_path(document.data_path, document.bot_id)
        return {"message": "Document added successfully."}
    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(...)):
    trainer = get_trainer()
    try:
        # Save the uploaded file to a temporary directory
        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:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/create_bot/{bot_id}")
def create_bot(bot_id: str):
    trainer = get_trainer()
    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):
    trainer = get_trainer()
    try:
        chat_id = trainer.new_chat(bot_id)
        return NewChatResponse(chat_id=chat_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@router.post("/query", response_model=QueryResponse)
def send_query(query_request: QueryRequest):
    trainer = get_trainer()
    try:
        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):
    trainer = get_trainer()
    try:
        chats = trainer.list_chats(bot_id)
        return chats
    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):
    trainer = get_trainer()
    try:
        history = trainer.get_chat_by_id(chat_id=chat_id)
        return history
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))