final version

.gitignore

@@ -0,0 +1 @@
+__pycache__

backend/routes.py

@@ -7,6 +7,11 @@ 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))

frontend/app_ui.py

@@ -3,6 +3,11 @@ import os
 import requests
 import streamlit as st
 
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
 # Read API URL from environment variable
 API_URL = os.getenv("API_URL", "http://localhost:8000")
 

poe/common-tasks.toml

@@ -34,7 +34,7 @@ cmd  = "ruff check ."
 
 [tool.poe.tasks.test]
 help = "Run unit tests"
-cmd  = "pytest -p no:cacheprovider"
+cmd  = "python -m unittest discover -s tests"
 
 [tool.poe.tasks.clean]
 help = "Remove automatically generated files"

poetry.lock

@@ -3718,4 +3718,4 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "38832b2f1f7e879f5efe88601e5ba8d8971bbbe8b4326625936762f860a7c128"
+content-hash = "14c56c888e2fbf236863e1a06b7a2a42c79377dea1917f6d7387ed106713abfd"

pyproject.toml

@@ -15,6 +15,7 @@ numpy = "1.26.4"
 fastapi = "^0.111.1"
 uvicorn = "^0.30.3"
 streamlit = "^1.37.0"
+python-dotenv = "^1.0.1"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.1.1"

src/prompt_loader.py

@@ -19,7 +19,7 @@ class PromptLoader:
         self.randomizer = random.Random(seed)
         self.data: Optional[List[str]] = None
 
-    def _load_data(self) -> None:
+    def _get_data(self) -> None:
         """
         Loads the dataset of prompts and stores them in the `data` attribute.
 
@@ -33,7 +33,7 @@ class PromptLoader:
         """
         Loads and samples prompts from the dataset.
 
-        If the dataset is not already loaded, it calls `_load_data()` to load it.
+        If the dataset is not already loaded, it calls `_get_data()` to load it.
 
         Args:
         size (Optional[int]): The number of prompts to sample. If not specified, all loaded prompts are returned.
@@ -46,7 +46,7 @@ class PromptLoader:
         ValueError: If `size` is specified and is greater than the number of available prompts.
         """
         if not self.data:
-            self._load_data()
+            self._get_data()
 
         if size:
             if size > len(self.data):

tests/test_load_data.py

@@ -0,0 +1,60 @@
+import unittest
+from unittest.mock import patch, MagicMock
+from src.prompt_loader import (
+    PromptLoader,
+)
+
+
+class TestPromptLoader(unittest.TestCase):
+
+    def setUp(self) -> None:
+        # Set up a mock dataset for testing
+        self.mock_data = {"train": {"prompt": ["prompt1", "prompt2", "prompt3"]}}
+        self.loader = PromptLoader(seed=42)
+
+    @patch("src.prompt_loader.load_dataset")
+    def test_load_data_without_size(self, mock_load_dataset: MagicMock) -> None:
+        mock_load_dataset.return_value = self.mock_data
+
+        self.loader.load_data()
+        self.assertEqual(self.loader.data, ["prompt1", "prompt2", "prompt3"])
+
+    @patch("src.prompt_loader.load_dataset")
+    def test_load_data_with_size(self, mock_load_dataset: MagicMock) -> None:
+        mock_load_dataset.return_value = self.mock_data
+        self.loader.load_data()
+        sampled_data = self.loader.load_data(size=2)
+
+        self.assertEqual(len(sampled_data), 2)
+        self.assertTrue(set(sampled_data).issubset({"prompt1", "prompt2", "prompt3"}))
+
+    @patch("src.prompt_loader.load_dataset")
+    def test_load_data_size_exceeds(self, mock_load_dataset: MagicMock) -> None:
+        mock_load_dataset.return_value = self.mock_data
+        self.loader.load_data()
+
+        with self.assertRaises(ValueError):
+            self.loader.load_data(size=10)
+
+    @patch("src.prompt_loader.load_dataset")
+    def test_data_loading_on_demand(self, mock_load_dataset: MagicMock) -> None:
+        mock_load_dataset.return_value = self.mock_data
+        mock_load_dataset.assert_not_called()
+
+        self.loader.load_data()
+        mock_load_dataset.assert_called_once()
+
+    @patch("src.prompt_loader.load_dataset")
+    def test_random_sampling(self, mock_load_dataset: MagicMock) -> None:
+        mock_load_dataset.return_value = self.mock_data
+
+        self.loader.load_data()
+        sample = self.loader.load_data(size=2)
+
+        self.assertEqual(len(sample), 2)
+        self.assertTrue(set(sample).issubset({"prompt1", "prompt2", "prompt3"}))
+        self.assertNotEqual(sample, ["prompt1", "prompt2"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_similar_prompts.py

@@ -0,0 +1,39 @@
+import unittest
+from unittest.mock import patch, Mock
+import requests
+
+# Assuming the function to be tested is in a module named `frontend.app_ui`
+from frontend.app_ui import get_similar_prompts
+
+
+class TestGetSimilarPrompts(unittest.TestCase):
+
+    @patch("frontend.app_ui.requests.post")
+    def test_get_similar_prompts_success(self, mock_post):
+        # Mock the response object to simulate a successful API call
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"prompts": ["prompt1", "prompt2", "prompt3"]}
+        mock_post.return_value = mock_response
+
+        # Call the function with a sample query and number
+        result = get_similar_prompts("test query", 3)
+
+        # Assertions
+        self.assertIsInstance(result, dict)
+        self.assertEqual(result, {"prompts": ["prompt1", "prompt2", "prompt3"]})
+
+    @patch("frontend.app_ui.requests.post")
+    def test_get_similar_prompts_failure(self, mock_post):
+        # Mock the response object to simulate a failed API call
+        mock_post.side_effect = requests.RequestException("Mock request exception")
+
+        # Call the function with a sample query and number
+        result = get_similar_prompts("test query", 3)
+
+        # Assertions
+        self.assertIsNone(result)
+
+
+if __name__ == "__main__":
+    unittest.main()