import os
from dotenv import load_dotenv
from fastapi import APIRouter, HTTPException
from fastapi.responses import HTMLResponse
from backend.models import QueryRequest, QueryResponse, SimilarPrompt
from src.prompt_loader import PromptLoader
from src.search_engine import PromptSearchEngine
# Load environment variables from .env file
load_dotenv()
# Constants
SEED = int(os.getenv("SEED", 42))
DATASET_SIZE = int(os.getenv("DATASET_SIZE", 1000))
# Initialize the prompt loader and search engine
prompts = PromptLoader(seed=SEED).load_data(size=DATASET_SIZE)
engine = PromptSearchEngine(prompts)
# Initialize the API router
router = APIRouter()
@router.post("/most_similar", response_model=QueryResponse)
async def get_most_similar(query_request: QueryRequest) -> QueryResponse:
"""
Endpoint to retrieve the most similar prompts based on a user query.
Args:
query_request (QueryRequest): The request payload containing the user query and the number of similar prompts to retrieve.
Returns:
QueryResponse: A response containing a list of similar prompts and their similarity scores.
Raises:
HTTPException: If an internal server error occurs while processing the request.
"""
try:
similar_prompts = engine.most_similar(
query=query_request.query, n=query_request.n
)
response = QueryResponse(
similar_prompts=[
SimilarPrompt(score=score, prompt=prompt)
for score, prompt in similar_prompts
]
)
return response
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@router.get("/", response_class=HTMLResponse)
async def home_page() -> HTMLResponse:
"""
Endpoint to serve a simple HTML page with information about the API.
Returns:
HTMLResponse: An HTML page providing an overview of the API and how to use it.
"""
return HTMLResponse(
"""
Prompt Search Engine
Prompt Search Engine API
Use this API to find similar prompts based on a query.