split main and run

main.py

@@ -0,0 +1,49 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
+from src.search_engine import PromptSearchEngine
+from src.prompt_loader import PromptLoader
+
+# 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
+@app.post("/most_similar", response_model=QueryResponse)
+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))

run.py

@@ -1,57 +1,4 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from typing import List
-from src.search_engine import PromptSearchEngine
-from src.prompt_loader import PromptLoader
-import os
+import uvicorn
 
-# 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
-@app.post("/most_similar", response_model=QueryResponse)
-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))
-
-
-# Run the server with: uvicorn main:app --reload
 if __name__ == "__main__":
-    import uvicorn
-
     uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)