Spaces:
Sleeping
Sleeping
| from typing import List | |
| from fastapi import FastAPI, HTTPException | |
| from fastapi.responses import HTMLResponse | |
| from pydantic import BaseModel | |
| from src.prompt_loader import PromptLoader | |
| from src.search_engine import PromptSearchEngine | |
| # Constants | |
| SEED = 42 | |
| DATA_SIZE = 100 | |
| # Initialize the prompt loader and search engine | |
| prompts = PromptLoader(seed=SEED).load_data(size=DATA_SIZE) | |
| engine = PromptSearchEngine(prompts) | |
| # Initialize FastAPI | |
| app = FastAPI() | |
| # Request and Response Models | |
| class QueryRequest(BaseModel): | |
| query: str | |
| n: int = 5 | |
| class SimilarPrompt(BaseModel): | |
| score: float | |
| prompt: str | |
| class QueryResponse(BaseModel): | |
| similar_prompts: List[SimilarPrompt] | |
| # API endpoint | |
| async def get_most_similar(query_request: QueryRequest): | |
| 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(): | |
| 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> | |
| <p>For more info, visit <a href="https://github.com/your-repository">GitHub</a>.</p> | |
| </div> | |
| </body> | |
| </html> | |
| """ | |
| ) | |