Fixes rag & retriever API (#93)

* 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

* recovered the right version of retriever.py

* enhanced rag docs

* removed duplicates in run_retriever.py

* changed the example for a public version of meditron

* small fix in docs

* truncate the output of the model in the RAG pipeline

* upgraded the requirement of transformers

* fix

* reformatting using black

---------

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

Author: 3 people

docs/rag.md

@@ -12,21 +12,21 @@
 
 Here is a minimal example to create a RAG pipeline hosted through [LangServe](https://python.langchain.com/docs/langserve/) servers.
 
-1. Create your RAG Inference config file based on the [local example](/examples/rag/config.yaml) or the [API example](/examples/rag/config_api.yaml).
+1. Create your RAG Inference config file based on the [local example](/examples/rag/config.yaml) or the [API example](/examples/rag/config_api.yaml). You can check the structure of the configuration file with the dataclass [RAGConfig](/src/mmore/rag/pipeline.py).
 
 2. Start your RAG pipeline using the `run_rag.py` script and your config file
     ```bash
     python3 -m mmore rag --config_file /path/to/config.yaml
     ```
 
-3. Query the server like any other LangServe server
+3. In API mode, query the server like any other LangServe server:
     ```bash
     curl --location --request POST http://localhost:8000/rag/invoke \
     -H 'Content-Type: application/json' \
     -d '{
         "input": {
             "input": "What is Meditron?",
-            "collection_name": "med_docs"
+            "collection_name": "my_docs"
         }
     }'
     ```
@@ -36,6 +36,8 @@ Here is a minimal example to create a RAG pipeline hosted through [LangServe](ht
     -H 'Content-Type: application/json' 
     ```
 
+    In local mode, the pipeline is run directly with the input data specified in the configuration file and the result is saved at the specified path.
+
 See [`examples/rag`](/examples/rag/) for other use cases.
 
 ## :mag: Modules

docs/retriever_api_specs.yaml

@@ -0,0 +1,183 @@
+openapi: 3.1.1
+info:
+  title: mmore Retriever API
+  description: |
+    This API is based on the OpenAPI 3.1 specification. You can find out more about Swagger at [https://swagger.io](https://swagger.io).
+
+    ## Overview
+
+    This API defines the retriever API of mmore, handling:
+
+    1. **File Operations** - Direct file management within mmore.
+    2. **Context Retrieval** - Semantic search based on the subset of documents that the user wants.
+
+    ## API Versioning
+
+    All requests must be prefixed with `/v1` for this current version of the API.
+
+    ## Roadmap & Considerations
+    - Authorization layer for mmore operations
+    - Permission control on some operation such as file deletion
+    - Advanced query for retrieval (send more sophisticated object, for instance representing a whole conversation, rather than one string)
+
+  version: 1.0.0
+servers:
+  - url: /v1
+    description: API server
+
+paths:
+  /files:
+    post:
+      tags:
+        - File Operations
+      summary: Upload a file
+      description: |
+        Upload a new file
+
+         **Requirements**:
+        - Unique fileId
+      requestBody:
+        content:
+          multipart/form-data:
+            schema:
+              type: object
+              properties:
+                fileId:
+                  type: string
+                  description: Unique identifier for the file
+                file:
+                  type: string
+                  format: binary
+                  description: The file content
+              required:
+                - fileId
+                - file
+      responses:
+        '201':
+          description: File successfully uploaded
+
+  /files/{id}:
+    put:
+      tags:
+        - File Operations
+      summary: Update a file
+      description: Replace an existing file with a new version
+      parameters:
+        - name: id
+          in: path
+          required: true
+          schema:
+            type: string
+      requestBody:
+        content:
+          multipart/form-data:
+            schema:
+              type: object
+              properties:
+                file:
+                  type: string
+                  format: binary
+                  description: The new file content
+              required:
+                - file
+      responses:
+        '200':
+          description: File successfully updated
+
+    delete:
+      tags:
+        - File Operations
+      summary: Remove a file
+      description: |
+        Delete a file from the system.
+
+        **Warning**: Not limited operation.
+      parameters:
+        - name: id
+          in: path
+          required: true
+          schema:
+            type: string
+      responses:
+        '200':
+          description: File successfully deleted
+
+    get:
+      tags:
+        - File Operations
+      summary: Download a file
+      description: Download a file from the system
+      parameters:
+        - name: id
+          in: path
+          required: true
+          schema:
+            type: string
+      responses:
+        '200':
+          description: File content
+          content:
+            application/octet-stream:
+              schema:
+                type: string
+                format: binary
+
+  /retrieve:
+    post:
+      tags:
+        - Context Retrieval
+      summary: Retrieve on MMORE
+      description: |
+        Search for files based on content similarity.
+      
+        This request is meant to be called by the Gateway to MMORE service upon chat completions request from Moove.
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                fileIds:
+                  type: array
+                  items:
+                    type: string
+                  description: List of file IDs to search within
+                maxMatches:
+                  type: integer
+                  minimum: 1
+                  description: Maximum number of matches to return
+                minSimilarity:
+                  type: number
+                  format: float
+                  minimum: -1.0
+                  maximum: 1.0
+                  default: 0.0
+                  description: Minimum similarity score for results (-1.0 to 1.0)
+                query:
+                  type: string
+                  description: Search query
+              required:
+                - fileIds
+                - query
+                - maxMatches
+      responses:
+        '200':
+          description: List of matching files with related content, sorted by highest similarity first.
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  type: object
+                  properties:
+                    fileId:
+                      type: string
+                    content:
+                      type: string
+                    similarity:
+                      type: number
+                      format: float
+                      minimum: -1.0
+                      maximum: 1.0
+                description: Results sorted by similarity (highest first)

examples/rag/config_api.yaml

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

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "numpy==1.26.3",
     "pandas==2.2.3",
     "datasets==2.19.1",
-    "transformers==4.47.0",
+    "transformers==4.52",
     "fastapi[standard]",
     "fastapi==0.115.5",
     "fasteners==0.19",

src/mmore/cli.py

@@ -33,7 +33,10 @@ def process(config_file: str):
     help="Path to the config file for post-processing.",
 )
 @click.option(
-    "--input-data", type=str, required=True, help="Path to the input JSONL file of documents."
+    "--input-data",
+    type=str,
+    required=True,
+    help="Path to the input JSONL file of documents.",
 )
 def postprocess(config_file: str, input_data: str):
     """Run the post-processors pipeline.
@@ -59,7 +62,11 @@ def postprocess(config_file: str, input_data: str):
     help="Path to the config file for indexing.",
 )
 @click.option(
-    "--documents-path", "-f", type=str, required=False, help="Path to the JSONL file of the (post)processed documents."
+    "--documents-path",
+    "-f",
+    type=str,
+    required=False,
+    help="Path to the JSONL file of the (post)processed documents.",
 )
 @click.option(
     "--collection-name",
@@ -93,10 +100,18 @@ def index(config_file: str, documents_path: str, collection_name: str):
     help="Dispatcher configuration file path.",
 )
 @click.option(
-    "--input-file", "-f", type=str, required=True, help="Path to the JSONL file of the input queries."
+    "--input-file",
+    "-f",
+    type=str,
+    required=True,
+    help="Path to the JSONL file of the input queries.",
 )
 @click.option(
-    "--output-file", "-o", type=str, required=True, help="Path to which save the results of the retriever as a JSON."
+    "--output-file",
+    "-o",
+    type=str,
+    required=True,
+    help="Path to which save the results of the retriever as a JSON.",
 )
 def retrieve(config_file: str, input_file: str, output_file: str):
     """Retrieve documents for specified queries.
@@ -171,8 +186,8 @@ def dashboard_backend(host, port):
     """Run the dashboard backend.
 
     Args:
-      host: 
-      port: 
+      host:
+      port:
 
     Returns:
 

src/mmore/process/dispatcher.py

@@ -342,6 +342,8 @@ def batch_list(
         else:
             results = list(self._dispatch_local(task_lists))
 
+        ExecutionState.shutdown()
+
         return results
 
     def __call__(self) -> List[List[MultimodalSample]]:

src/mmore/process/execution_state.py

@@ -47,6 +47,12 @@ def initialize(distributed_mode=False, client=None):
             ExecutionState._local_state = False
             logger.info("Execution state initialized (local mode)")
 
+    @staticmethod
+    def shutdown():
+        ExecutionState._use_dask = None
+        ExecutionState._dask_var = None
+        ExecutionState._local_state = False
+
     @staticmethod
     def get_should_stop_execution() -> bool:
         """Returns the global execution state (True if it should stop)"""

src/mmore/rag/pipeline.py

@@ -87,9 +87,15 @@ def _build_chain(retriever, format_docs, prompt, llm) -> Runnable:
         validate_input = RunnableLambda(
             lambda x: MMOREInput.model_validate(x).model_dump()
         )
-        validate_output = RunnableLambda(
-            lambda x: MMOREOutput.model_validate(x).model_dump()
-        )
+
+        def make_output(x):
+            """Validate the output of the LLM and keep only the actual answer of the assistant"""
+            res_dict = MMOREOutput.model_validate(x).model_dump()
+            res_dict["answer"] = res_dict["answer"].split("<|assistant|>\n")[-1]
+
+            return res_dict
+
+        validate_output = RunnableLambda(make_output)
 
         rag_chain_from_docs = prompt | llm | StrOutputParser()
 

src/mmore/rag/types.py

@@ -17,6 +17,10 @@ class MMOREInput(BaseModel):
         ...,
         description="The collection",
     )
+    document_ids: List[str] = Field(
+        default_factory=list,  # Set default to an empty list if not provided
+        description="List of document IDs, defaults to an empty list if not provided.",
+    )
 
 
 # ------------------------------- Simple Output ------------------------------ #

src/mmore/run_index_api.py

@@ -131,7 +131,7 @@ async def upload_files(
         with tempfile.TemporaryDirectory() as temp_dir:
             logging.info(f"Starting to process {len(files)} files with custom IDs")
 
-            for i, (file, file_id) in enumerate(zip(files, listIds)):
+            for file, file_id in zip(files, listIds):
                 if file.filename is None:
                     raise HTTPException(
                         status_code=422,