Update backend settings

Author: merfanian

backend/core/embedders.py

@@ -69,7 +69,7 @@ def weight(self, value):
 class EmbedderManager:
     def __init__(self):
         self._device = torch.device(
-            "cuda" if torch.cuda.is_available() and settings.app.use_cuda else "cpu")
+            "cuda" if torch.cuda.is_available() and settings.service.use_cuda else "cpu")
         self._image_embedders = {}
         for embedder_config in settings.image_embedders:
             self._image_embedders[embedder_config.name] = ImageEmbedder(name=embedder_config.name,

backend/main.py

@@ -5,9 +5,9 @@
 import fastapi
 from fastapi import FastAPI, HTTPException, Body, Request
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-from fastapi.responses import FileResponse
 from pymilvus import Collection
 
 from core import embedder_manager, image_generator, query_manager
@@ -111,14 +111,13 @@ async def create_query(q: str = Body(..., embed=True)):
 
 
 @app.get("/search/{qid}")
-async def search(qid: int, n: int = 20, k: int = 1, image_size: int = 512,
+async def search(qid: int, n: int = settings.query.default_num_images_to_retrieve,
+                 k: int = settings.query.default_num_images_to_generate,
+                 image_size: int = settings.query.default_generated_image_size,
+                 include_base_images_in_preview: bool = settings.query.include_base_images_in_preview,
                  generator_engines: List[str] = fastapi.Query(None),
-                 include_base_images: bool = False,
                  request: Request = None
                  ):
-    if generator_engines is None:
-        generator_engines = [settings.generators.default_engine]
-
     query_object = query_manager.get_query(qid)
     query = query_object.query
 
@@ -192,16 +191,13 @@ async def search(qid: int, n: int = 20, k: int = 1, image_size: int = 512,
         "preview_url": str(request.url_for("gallery", qid=qid))
     }
 
-    if include_base_images:
+    if include_base_images_in_preview:
         res["base_images"] = [pil_image_to_base64(image) for image in generated_images]
 
     return res
 
 
-
-
-
-@app.get("/file", name="get_file")
+@app.get("/file")
 async def get_file(file_path: str):
     if not os.path.isfile(file_path):
         raise HTTPException(status_code=404, detail="File not found")

backend/monitoring/directory_watcher.py

@@ -82,7 +82,7 @@ def __init__(self):
         self.observer = Observer()
         self.handlers = {}
         self.watches = {}
-        self.executor = ThreadPoolExecutor(max_workers=settings.app.num_watcher_workers)
+        self.executor = ThreadPoolExecutor(max_workers=settings.directory.num_watcher_workers)
 
     def add_directory(self, path):
         session = SessionLocal()
@@ -104,7 +104,7 @@ def add_directory(self, path):
 
         handler = ImageChangeHandler(directory_id)
         self.handlers[renamed_path] = handler
-        watch = self.observer.schedule(handler, path, recursive=settings.app.recursive_indexing)
+        watch = self.observer.schedule(handler, path, recursive=settings.directory.recursive_indexing)
         self.watches[renamed_path] = watch
 
         # Start the background indexing
@@ -130,7 +130,7 @@ def _process_directory_in_background(self, path, directory_id):
         )
 
         to_add = []
-        if settings.app.recursive_indexing:  # If recursive search is required
+        if settings.directory.recursive_indexing:  # If recursive search is required
             for root, _, files in os.walk(path):
                 for file in files:
                     if file.lower().endswith(('.png', '.jpg', '.jpeg')):
@@ -151,7 +151,7 @@ def _process_directory_in_background(self, path, directory_id):
         # 2. Retrieve all not indexed images
         embedders = embedder_manager.get_image_embedders()
 
-        batch_size = settings.app.batch_size
+        batch_size = settings.directory.batch_size
         while True:
             not_indexed_images = session.query(Image).filter(Image.directory_id == directory_id, Image.is_indexed == False).all()
             if not not_indexed_images:
@@ -164,7 +164,7 @@ def _process_directory_in_background(self, path, directory_id):
             batch = not_indexed_images[:batch_size]
 
             # 3. Embed images in parallel for this batch
-            with ThreadPoolExecutor(max_workers=settings.app.num_embedding_workers) as pool:
+            with ThreadPoolExecutor(max_workers=settings.directory.num_embedding_workers) as pool:
                 futures = {pool.submit(self._embed_image, img_record.path, embedders): img_record.path for img_record in batch}
 
                 embeddings_batch = []

backend/settings/readonly.py

@@ -23,7 +23,9 @@ def _log_settings(cls):
         logger.info("Settings Loaded:")
         logger.info(f"Postgres Settings: {cls._instance._settings.postgres}")
         logger.info(f"Milvus Settings: {cls._instance._settings.milvus}")
-        logger.info(f"App Settings: {cls._instance._settings.app}")
+        logger.info(f"Service Settings: {cls._instance._settings.service}")
+        logger.info(f"Directory Settings: {cls._instance._settings.directory}")
+        logger.info(f"Query Settings: {cls._instance._settings.query}")
         logger.info(f"Generators Settings: {cls._instance._settings.generators}")
         if cls._instance._settings.json_config:
             logger.info(f"JSON Config: {cls._instance._settings.json_config.dict()}")
@@ -37,8 +39,16 @@ def milvus(self):
         return self._settings.milvus
 
     @property
-    def app(self):
-        return self._settings.app
+    def service(self):
+        return self._settings.service
+
+    @property
+    def query(self):
+        return self._settings.query
+
+    @property
+    def directory(self):
+        return self._settings.directory
 
     @property
     def generators(self):
@@ -55,9 +65,8 @@ def image_embedders(self):
         return self.json_config.image_embedders
 
     @property
-    def weights_path(self) :
-        return os.path.join(self._settings.app.embedders_config_dir_path, "weights.json")
-
+    def weights_path(self):
+        return os.path.join(self._settings.service.embedders_config_dir_path, "weights.json")
 
     def get_image_embedder_details(self, name: str):
         embedder = next(

backend/settings/settings_model.py

@@ -36,24 +36,28 @@ def uri(self) -> str:
         return f"{self.host}:{self.port}"
 
 
-class AppSettings(BaseModel):
-    embedders_config_dir_path: str = Field("./configs/balanced/")
-    use_cuda: bool = Field(False)
-    recursive_indexing: bool = Field(False)
-    batch_size: int = Field(100)
-    num_embedding_workers: int = Field(4)
+class QuerySettings(BaseModel):
+    default_num_images_to_retrieve: int = Field(20)
+    default_num_images_to_generate: int = Field(4)
+    default_generated_image_size: int = Field(512)
+    include_base_images_in_preview: bool = Field(False)
+
+
+class DirectorySettings(BaseModel):
     num_watcher_workers: int = Field(4)
+    num_embedding_workers: int = Field(4)
+    batch_size: int = Field(50)
+    recursive_indexing: bool = Field(False)
 
 
-class IndexSettings(BaseModel):
-    index_type: str = Field("HNSW")
-    metric_type: str = Field("COSINE")
+class ServiceSettings(BaseModel):
+    embedders_config_dir_path: str = Field("./configs/balanced/")
+    use_cuda: bool = Field(False)
 
 
-class GeneratorsSettings(BaseModel):
+class ImageGeneratorSettings(BaseModel):
     host: str = Field("0.0.0.0")
     port: int = Field(8001)
-    default_engine: str = Field("sdxl_lightning")
 
     @property
     def url(self) -> str:
@@ -64,8 +68,10 @@ class Settings(BaseSettings):
     # Environment-based settings
     postgres: PostgresSettings = PostgresSettings()
     milvus: MilvusSettings = MilvusSettings()
-    app: AppSettings = AppSettings()
-    generators: GeneratorsSettings = GeneratorsSettings()
+    service: ServiceSettings = ServiceSettings()
+    generators: ImageGeneratorSettings = ImageGeneratorSettings()
+    directory: DirectorySettings = DirectorySettings()
+    query: QuerySettings = QuerySettings()
 
     # JSON config
     json_config: Optional[JSONConfig] = None
@@ -78,7 +84,7 @@ def load_json_config(self):
         """
         Load and parse the JSON configuration file specified in app.config_path.
         """
-        config_path = Path(self.app.embedders_config_dir_path, "config.json")
+        config_path = Path(self.service.embedders_config_dir_path, "config.json")
         if config_path.exists():
             with open(config_path, "r") as file:
                 json_data = json.load(file)

needlectl/cli/directory.py

@@ -23,16 +23,8 @@ def add_directory(
     output_format = ctx.obj["output"]
     abs_path = os.path.abspath(path)
 
-    # manager = DockerComposeManager()
-    # manager.add_volume(service_name="backend", volume_path=abs_path)
-    #
-    # typer.echo("Restarting backend container...")
-    # manager.start_containers()
-
     client = BackendClient(api_url)
 
-    typer.echo("Waiting for API to become available...")
-
     try:
         client.wait_for_api()
         resp = client.add_directory(abs_path)