Live retrieval API with index API and retriever API (#94)

* Added retrieving among a set of specified documents

* Adapted run_retriever

* simplified code in retriever

* simplified run_retriever.py

* misc style changes

* adapted retriever to rag api

* cherry picking from RAG interface is now possible by providing document_ids in the json file

* Preparation for the retriever api

* basic retriever api

* Updated retrieval api + documentation

Co-authored-by: MikaelKalajdzic <mikael@kalajdzic.ch>
Co-authored-by: fabnemEPFL <fabrice.nemo@epfl.ch>

* adapted retriever to effectively use the arguments minSimilarity and maxMatches

* execution state shut down at the end of processing

* import fixes on run_retriever.py

* removed useless enumerate

* retriever config and fixes with retriever_api

* fixed mmore retriever

* fixed default retriever_api config

* fix default DB uri

* fixed error 404 if deletion of file already absent

* new fixes of HTTP errors

* changed retriever endpoint to post

* making a single FastAPI combining index-api and retriever-api

* various fixes

* more fixes

* retrieve the original name of the file sent on the download endpoit

* made one shady line of the code clearer

* fixed the logic of retriever

* cosmetic changes

* moved utilities to utils.py

* added live-retrieval in the module

* cosmetic changes

* removed the retriever endpoint from run_index_api

* removed useless import

* initializing the indexer and processors on starttime of the live retrieval

* small fix

* added a tag to the retriever endpoint

* small changes

* fixes

* black and isort fixes

* logic fix

* cosmetic changes

* fixes

* pre-loading of processors

* chunking documents in index_api

* attempt of fix

* smol fix

* changed default LLM for RAG

* initialize PDF processor on startup

* split document id and chunk id

* fixed typo

* fix

* changed default values

* fix

* ruff formatting

* removed useless imports

* reformatted tests

* pyright compliance and reformatting

* ruff compliance

* ruff fixes

* small fixes

* renamed a function

* again

* manageable preloading of processor models

* changed default model for media processor

---------

Co-authored-by: leagrieder <lea@grieder.org>
Co-authored-by: MikaelKalajdzic <mikael@kalajdzic.ch>

Author: 3 people

docs/index_api.md

@@ -173,7 +173,7 @@ Returns the file with binary content.
 - File types supported:
 
     ```
-    .pdf, .docx, .pptx, .md, .txt, .xlsx, .xls, .csv, .mp4, .avi, .mov, .mkv, .mp3, .wav, .aac, .eml, .html
+    .pdf, .docx, .pptx, .md, .txt, .xlsx, .xls, .csv, .mp4, .avi, .mov, .mkv, .mp3, .wav, .aac, .eml, .html, .htm
     ```
     
 

examples/rag/config.yaml

@@ -1,7 +1,7 @@
 rag:
   llm: 
     llm_name: OpenMeditron/meditron3-8b
-    max_new_tokens: 250
+    max_new_tokens: 1200
   retriever:
     db:
       uri: ./proc_demo.db

examples/rag/config_api.yaml

@@ -2,8 +2,8 @@
 rag: 
   # LLM Config
   llm: 
-    llm_name: "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # "epfl-llm/meditron-70b" # "gpt-4o-mini" # Anything supported
-    max_new_tokens: 100
+    llm_name: Qwen/Qwen3-8B # "epfl-llm/meditron-70b" # "gpt-4o-mini" # Anything supported
+    max_new_tokens: 1200
     temperature: 0.8
   # Retriever Config
   retriever:

examples/retriever_api/config.yaml

@@ -0,0 +1,6 @@
+db:
+  uri: ./proc_demo.db
+  name: my_db
+hybrid_search_weight: 0.5
+k: 5
+collection_name: my_docs

pyproject.toml

@@ -164,7 +164,7 @@ exclude = ["src/mmore/run_retriever.py"]    # TODO: add back when GH CI bug is f
 
 [tool.ruff.lint]
 select = ["E", "F", "W", "I", "N"]
-ignore = ["E501"]   # Avoid enforcing line-length violations (`E501`)
+ignore = ["E501", "E402"]   # Avoid enforcing line-length violations (`E501`) and Module level import not at top of file (`E402`)
 
 [tool.ruff.lint.per-file-ignores]
 "__init__.py" = ["F401"]

src/mmore/cli.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import click
 
 
@@ -97,36 +99,82 @@ def index(config_file: str, documents_path: str, collection_name: str):
     "-c",
     type=str,
     required=True,
-    help="Dispatcher configuration file path.",
+    help="Retriever configuration file path.",
 )
 @click.option(
     "--input-file",
     "-f",
-    type=str,
-    required=True,
+    type=Optional[str],
+    required=False,
+    default=None,
     help="Path to the JSONL file of the input queries.",
 )
 @click.option(
     "--output-file",
     "-o",
-    type=str,
-    required=True,
+    type=Optional[str],
+    required=False,
+    default=None,
     help="Path to which save the results of the retriever as a JSON.",
 )
-def retrieve(config_file: str, input_file: str, output_file: str):
+@click.option(
+    "--host", type=str, default="0.0.0.0", help="Host on which the API should be run."
+)
+@click.option(
+    "--port", type=int, default=8001, help="Port on which the API should be run."
+)
+def retrieve(
+    config_file: str,
+    input_file: Optional[str],
+    output_file: Optional[str],
+    host: str,
+    port: int,
+):
     """Retrieve documents for specified queries.
 
     Args:
-      config_file: path to the config file for the retriver.
+      config_file: path to the config file for the retriever.
       input_file: path to the JSONL file of the input queries.
       output_file: path to which save the results of the retriever as a JSON.
 
     Returns:
 
     """
     from .run_retriever import retrieve as run_retrieve
+    from .run_retriever import run_api
+
+    if input_file:
+        assert isinstance(output_file, str)
+        run_retrieve(config_file, input_file, output_file)
+    else:
+        run_api(config_file, host, port)
+
+
+@main.command()
+@click.option(
+    "--config-file",
+    "-c",
+    type=str,
+    required=True,
+    help="Retriever configuration file path.",
+)
+@click.option(
+    "--host", type=str, default="0.0.0.0", help="Host on which the API should be run."
+)
+@click.option(
+    "--port", type=int, default=8000, help="Port on which the API should be run."
+)
+def live_retrieval(config_file: str, host: str, port: int):
+    """API for live indexing and retrieval of documents.
+
+    Args:
+      config_file: Path to the retriever configuration file.
+      host: Host on which the API should be run.
+      port: Port on which the API should be run.
+    """
+    from .run_live_retrieval import run
 
-    run_retrieve(config_file, input_file, output_file)
+    run(config_file, host, port)
 
 
 @main.command()
@@ -148,16 +196,24 @@ def rag(config_file: str):
 
 
 @main.command()
+@click.option(
+    "--config-file",
+    "-c",
+    type=str,
+    required=True,
+    help="Retriever configuration file path.",
+)
 @click.option(
     "--host", type=str, default="0.0.0.0", help="Host on which the API should be run."
 )
 @click.option(
     "--port", type=int, default=8000, help="Port on which the API should be run."
 )
-def index_api(host, port):
+def index_api(config_file, host, port):
     """Run the Index API.
 
     Args:
+      config_file: Path to the retriever configuration file.
       host: Host on which the API should be run.
       port: Port on which the API should be run.
 
@@ -166,7 +222,7 @@ def index_api(host, port):
     """
     from .run_index_api import run_api
 
-    run_api(host, port)
+    run_api(config_file, host, port)
 
 
 @main.command()

src/mmore/index/indexer.py

@@ -23,7 +23,7 @@
 
 @dataclass
 class DBConfig:
-    uri: str = "demo.db"
+    uri: str = "./proc_demo.db"
     name: str = "my_db"
 
 
@@ -114,7 +114,8 @@ def _create_collection_with_schema(self, collection_name: str):
         fields = [
             FieldSchema(
                 name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=128
-            ),  # Add doc_id field
+            ),
+            FieldSchema(name="document_id", dtype=DataType.VARCHAR, max_length=128),
             FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535),
             FieldSchema(
                 name="dense_embedding",
@@ -192,6 +193,7 @@ def _index_documents(
             data = [
                 {
                     "id": sample.id,
+                    "document_id": sample.document_id,
                     "text": sample.text,
                     "dense_embedding": d,
                     "sparse_embedding": s.reshape(1, -1),
@@ -206,7 +208,7 @@ def _index_documents(
                 partition_name=partition_name,
             )
 
-            inserted += list(batch_inserted.values())[0]
+            inserted += batch_inserted["insert_count"]
 
         return inserted
 

src/mmore/process/config.py

@@ -50,7 +50,7 @@ def get_config_path():
         app_dir.mkdir(parents=True, exist_ok=True)
     except PermissionError as e:
         click.echo(f"Error creating config directory: {e}", err=True)
-        raise
+        raise e
     return app_dir / "config.yaml"
 
 

src/mmore/process/crawler.py

@@ -144,7 +144,7 @@ def _get_metadata_jsonl_path(self, results_jsonl_path):
                 if "metadata" in data and "file_path" in data["metadata"]:
                     file_paths.append(data["metadata"]["file_path"])
                 else:
-                    print(
+                    logger.error(
                         f"Warning file_path not found in metadate (line{i} of {results_jsonl_path})"
                     )
         return file_paths
@@ -224,9 +224,9 @@ def from_yaml(yaml_path: str):
             with open(yaml_path, "r") as file:
                 config = yaml.safe_load(file)
             return CrawlerConfig.from_dict(config)
-        except (FileNotFoundError, yaml.YAMLError):
+        except (FileNotFoundError, yaml.YAMLError) as e:
             logger.error(f"[Crawler] Error processing {yaml_path}.")
-            raise
+            raise e
 
     def to_dict(self):
         """

src/mmore/process/dispatcher.py

@@ -85,7 +85,7 @@ def from_dict(config: Dict) -> "DispatcherConfig":
             use_fast_processors=config.get("use_fast_processors", True),
             distributed=config.get("distributed", False),
             scheduler_file=config.get("scheduler_file"),
-            processor_config=config.get("processor"),
+            processor_config=config.get("processor_config"),
             process_batch_sizes=config.get("process_batch_sizes"),
             batch_multiplier=config.get("batch_multiplier", 1),
             extract_images=config.get("extract_images", False),
@@ -100,9 +100,9 @@ def from_yaml(yaml_path: str):
             with open(yaml_path, "r") as file:
                 config = yaml.safe_load(file)
             return DispatcherConfig.from_dict(config)
-        except (FileNotFoundError, yaml.YAMLError):
+        except (FileNotFoundError, yaml.YAMLError) as e:
             logger.error(f"[Dispatcher] Error processing file {yaml_path}")
-            raise
+            raise e
 
     def to_dict(self) -> Dict:
         """Convert the DispatcherConfig object to a dictionary."""
@@ -111,7 +111,7 @@ def to_dict(self) -> Dict:
             "distributed": self.distributed,
             "scheduler_file": self.scheduler_file,
             "output_path": self.output_path,
-            "processor": self.processor_config,
+            "processor_config": self.processor_config,
             "process_batch_sizes": self.process_batch_sizes,
             "batch_multiplier": self.batch_multiplier,
             "extract_images": self.extract_images,