Spaces:
Sleeping
Sleeping
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() | |
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)) | |
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( | |
""" | |
<!DOCTYPE html> | |
<html lang="en"> | |
<head> | |
<meta charset="UTF-8"> | |
<meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
<title>Prompt Search Engine</title> | |
<style> | |
body { font-family: Arial, sans-serif; margin: 20px; } | |
h1 { color: #333; } | |
p { margin-bottom: 10px; } | |
code { background: #f4f4f4; padding: 2px 4px; border-radius: 4px; } | |
.container { max-width: 800px; margin: 0 auto; } | |
</style> | |
</head> | |
<body> | |
<div class="container"> | |
<h1>Prompt Search Engine API</h1> | |
<p>Use this API to find similar prompts based on a query.</p> | |
<h2>POST /most_similar</h2> | |
<p><strong>Request:</strong> <code>{"query": "string", "n": 1}</code></p> | |
<p><strong>Response:</strong> <code>{"similar_prompts": [{"score": 0.95, "prompt": "Example prompt 1"}]}</code></p> | |
</div> | |
</body> | |
</html> | |
""" | |
) | |