refactor everything

.env

@@ -0,0 +1,3 @@
+API_URL="https://lazarr19-prompt-engine.hf.space"
+SEED=42
+DATASET_SIZE=1000

README.md

@@ -43,6 +43,10 @@ Before you start, ensure you have the following tools installed:
    poetry shell
    ```
 
+### Environment Variables
+
+Environment variables are set in the .env file in the root of the repo.
+
 ### Backend Details
 
 The backend API provides the following endpoint:
@@ -76,6 +80,8 @@ In case you want to set up your own HuggingFace Space, you must create a HF toke
 git remote set-url origin https://USERNAME:[email protected]/spaces/USERNAME/REPO_NAME.git
 ```
 
+Also, pay attention to the HF Space parameters at the beginning of this README.md.
+
 ### User Interface
 
 The frontend UI is simple and includes:
@@ -86,7 +92,7 @@ The frontend UI is simple and includes:
 To start only the UI service you can run:
 
 ```bash
-poe frontend --api_url http://localhost:8000
+poe frontend
 ```
 
-The default api points to HF Space of this repository.
+Or you can do it with `https://lazarr19-prompt-engine.hf.space` if you prefer to hit the HF Space endpoint.

backend/main.py

@@ -0,0 +1,9 @@
+from fastapi import FastAPI
+
+from backend.routes import router
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Include routes from the routes module
+app.include_router(router)

backend/models.py

@@ -0,0 +1,29 @@
+from typing import List
+
+from pydantic import BaseModel
+
+
+class QueryRequest(BaseModel):
+    """
+    Represents the request model for querying similar prompts.
+    """
+
+    query: str
+    n: int = 5
+
+
+class SimilarPrompt(BaseModel):
+    """
+    Represents a single similar prompt with its similarity score.
+    """
+
+    score: float
+    prompt: str
+
+
+class QueryResponse(BaseModel):
+    """
+    Represents the response model containing a list of similar prompts.
+    """
+
+    similar_prompts: List[SimilarPrompt]

main.py → backend/routes.py

@@ -1,42 +1,38 @@
-from typing import List
+import os
 
-from fastapi import FastAPI, HTTPException
+from fastapi import APIRouter, HTTPException
 from fastapi.responses import HTMLResponse
-from pydantic import BaseModel
 
+from backend.models import QueryRequest, QueryResponse, SimilarPrompt
 from src.prompt_loader import PromptLoader
 from src.search_engine import PromptSearchEngine
 
 # Constants
-SEED = 42
-DATA_SIZE = 100
+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=DATA_SIZE)
+prompts = PromptLoader(seed=SEED).load_data(size=DATASET_SIZE)
 engine = PromptSearchEngine(prompts)
 
-# Initialize FastAPI
-app = FastAPI()
+# Initialize the API router
+router = APIRouter()
 
 
-# Request and Response Models
-class QueryRequest(BaseModel):
-    query: str
-    n: int = 5
+@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.
 
-class SimilarPrompt(BaseModel):
-    score: float
-    prompt: str
+    Returns:
+    QueryResponse: A response containing a list of similar prompts and their similarity scores.
 
-
-class QueryResponse(BaseModel):
-    similar_prompts: List[SimilarPrompt]
-
-
-# API endpoint
-@app.post("/most_similar", response_model=QueryResponse)
-async def get_most_similar(query_request: QueryRequest):
+    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
@@ -52,8 +48,14 @@ async def get_most_similar(query_request: QueryRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
 
-@app.get("/", response_class=HTMLResponse)
-async def home_page():
+@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(
         """
         <!DOCTYPE html>
@@ -77,7 +79,6 @@ async def home_page():
                 <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>

frontend/app_ui.py

@@ -1,26 +1,29 @@
-import argparse
+import os
 
 import requests
 import streamlit as st
 
+# Read API URL from environment variable
+API_URL = os.getenv("API_URL", "http://localhost:8000")
 
-def parse_arguments():
-    """Parse command-line arguments."""
-    parser = argparse.ArgumentParser(description="Prompt Similarity Finder")
-    parser.add_argument(
-        "--api_url",
-        type=str,
-        default="https://lazarr19-prompt-engine.hf.space",
-        help="The URL of the FastAPI service",
-    )
-    return parser.parse_args()
 
+def get_similar_prompts(query: str, n: int) -> dict:
+    """
+    Fetches similar prompts from the API based on the user query.
+
+    Args:
+    query (str): The user query for which similar prompts are to be retrieved.
+    n (int): The number of similar prompts to return.
+
+    Returns:
+    dict: A dictionary containing similar prompts, or None if there was an error.
 
-def get_similar_prompts(api_url, query, n):
-    """Fetch similar prompts from the FastAPI service."""
+    Raises:
+    requests.RequestException: If an HTTP error occurs during the request.
+    """
     try:
         response = requests.post(
-            f"{api_url}/most_similar", json={"query": query, "n": n}
+            f"{API_URL}/most_similar", json={"query": query, "n": n}
         )
         response.raise_for_status()  # Raise an exception for HTTP errors
         return response.json()
@@ -29,8 +32,16 @@ def get_similar_prompts(api_url, query, n):
         return None
 
 
-def get_color(score):
-    """Determine the color based on the score."""
+def get_color(score: float) -> str:
+    """
+    Determines the color based on the similarity score.
+
+    Args:
+    score (float): The similarity score of a prompt.
+
+    Returns:
+    str: The color representing the score, which could be "green", "orange", or "red".
+    """
     if score >= 0.8:
         return "green"
     elif score >= 0.5:
@@ -39,11 +50,13 @@ def get_color(score):
         return "red"
 
 
-def main(api_url):
-    """Main function to run the Streamlit app."""
+def main():
+    """
+    The main function for running the Streamlit app.
+    Sets up the UI for entering queries and retrieving similar prompts.
+    """
     st.title("Prompt Similarity Finder")
 
-    # User input for query
     query = st.text_input("Enter your query:", "")
     n = st.slider(
         "Number of similar prompts to retrieve:", min_value=1, max_value=40, value=5
@@ -52,7 +65,7 @@ def main(api_url):
     if st.button("Find Similar Prompts"):
         if query:
             with st.spinner("Fetching similar prompts..."):
-                result = get_similar_prompts(api_url, query, n)
+                result = get_similar_prompts(query, n)
                 if result:
                     similar_prompts = result.get("similar_prompts", [])
                     if similar_prompts:
@@ -60,7 +73,6 @@ def main(api_url):
                         for item in similar_prompts:
                             score = item["score"]
                             color = get_color(score)
-                            # Apply color only to the score part
                             st.markdown(
                                 f"<p><strong>Score:</strong> <span style='color:{color};'>{score:.2f}</span> <br> <strong>Prompt:</strong> {item['prompt']}</p>",
                                 unsafe_allow_html=True,
@@ -73,5 +85,4 @@ def main(api_url):
 
 
 if __name__ == "__main__":
-    args = parse_arguments()
-    main(args.api_url)
+    main()

run.py

@@ -3,7 +3,7 @@ import uvicorn
 
 def run_fastapi_app():
     uvicorn.run(
-        "main:app",  # Module name and app instance
+        "backend.main:app",  # Module name and app instance
         host="0.0.0.0",
         port=8000,
         reload=True,  # Enable auto-reload for development

src/prompt_loader.py

@@ -5,16 +5,46 @@ from datasets import load_dataset
 
 
 class PromptLoader:
+    """
+    A class for loading and sampling prompts from a dataset.
+    """
+
     def __init__(self, seed: int = 42) -> None:
+        """
+        Initializes the PromptLoader with a specified seed for random sampling.
+
+        Args:
+        seed (int): The seed value for the random number generator. Default is 42.
+        """
         self.randomizer = random.Random(seed)
         self.data: Optional[List[str]] = None
 
     def _load_data(self) -> None:
+        """
+        Loads the dataset of prompts and stores them in the `data` attribute.
+
+        This method uses the `datasets` library to load the dataset and extract prompts from the "train" split.
+        """
         self.data = load_dataset("daspartho/stable-diffusion-prompts")["train"][
             "prompt"
         ]
 
     def load_data(self, size: Optional[int] = None) -> List[str]:
+        """
+        Loads and samples prompts from the dataset.
+
+        If the dataset is not already loaded, it calls `_load_data()` to load it.
+
+        Args:
+        size (Optional[int]): The number of prompts to sample. If not specified, all loaded prompts are returned.
+
+        Returns:
+        List[str]: A list of sampled prompts. If `size` is specified, returns a random sample of the specified size.
+            If `size` is not specified, returns all loaded prompts.
+
+        Raises:
+        ValueError: If `size` is specified and is greater than the number of available prompts.
+        """
         if not self.data:
             self._load_data()