Spaces:
Sleeping
Sleeping
| import os | |
| 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 | |
| from dotenv import load_dotenv | |
| # 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": 5}</code></p> | |
| <p><strong>Response:</strong> <code>{"similar_prompts": [{"score": 0.95, "prompt": "Example prompt 1"}]}</code></p> | |
| </div> | |
| </body> | |
| </html> | |
| """ | |
| ) | |