feat: [Sidecar] Add ability to remove one or all references (#193)

gjreda · web-flow · commit 7f9cc08e06e5 · 2023-06-26T17:32:10.000+01:00
* Add ability to delete references

* types

* updates

* fix comment

* fix type

* Address PR comments

* types

* readme
diff --git a/python/README.md b/python/README.md
@@ -39,6 +39,7 @@ The application current has three main functions:
 2. Ingest Status
 3. Rewrite
 4. Chat
+5. Delete (References)
 
 ### Ingest
 `Ingest` is called when a Refstudio user uploads Reference documents.
@@ -182,4 +183,34 @@ $ poetry run python main.py --text "What can you tell me about hidden feedback l
     "text": "Hidden feedback loops in machine learning refer to situations where two systems indirectly influence each other through the world, leading to changes in behavior that may not be immediately visible. These loops may exist between completely disjoint systems and can make analyzing the effect of proposed changes extremely difficult, adding cost to even simple improvements. It is recommended to look carefully for hidden feedback loops and remove them whenever feasible."
   }
 ]
+```
+
+### Delete (References)
+`delete` removes a Reference from project storage. It is called when a Reference is deleted from the UI.
+
+It takes a list of `source_filenames` as input and returns a response status. Alternatively, you can use the `--all` parameter to delete all References in storage.
+
+To run `delete`:
+
+```bash
+$ poetry run python main.py delete --source_filenames grobid-fails.pdf "Machine Learning at Scale.pdf"
+
+
+# Example:
+$ poetry run python main.py delete --source_filenames grobid-fails.pdf "Machine Learning at Scale.pdf" | jq
+
+# Response:
+{
+  "status": "ok",
+  "message": ""
+}
+
+# Another Example (error):
+$ poetry run python main.py delete --source_filenames file-does-not-exist.pdf | jq
+
+# Error Response:
+{
+  "status": "error",
+  "message": "Unable to delete file-does-not-exist.pdf: not found in storage"
+}
 ```
diff --git a/python/cli.schema.json b/python/cli.schema.json
@@ -19,13 +19,17 @@
       "items": {
         "$ref": "#/definitions/ChatResponseChoice"
       }
+    },
+    "delete": {
+      "$ref": "#/definitions/DeleteStatusResponse"
     }
   },
   "required": [
     "ingest",
     "ingest_status",
     "rewrite",
-    "chat"
+    "chat",
+    "delete"
   ],
   "definitions": {
     "IngestStatus": {
@@ -231,6 +235,22 @@
         "index",
         "text"
       ]
+    },
+    "DeleteStatusResponse": {
+      "title": "DeleteStatusResponse",
+      "type": "object",
+      "properties": {
+        "status": {
+          "$ref": "#/definitions/ResponseStatus"
+        },
+        "message": {
+          "type": "string"
+        }
+      },
+      "required": [
+        "status",
+        "message"
+      ]
     }
   }
 }
diff --git a/python/main.py b/python/main.py
@@ -1,4 +1,4 @@
-from sidecar import chat, cli, ingest, rewrite
+from sidecar import chat, cli, ingest, rewrite, storage
 
 if __name__ == '__main__':
     parser = cli.get_arg_parser()
@@ -19,3 +19,9 @@
     if args.command == "chat":
         chat.ask_question(args.text)
 
+    if args.command == "delete":
+        if args.all:
+            storage.delete_references(all_=args.all)
+        else:
+            storage.delete_references(source_filenames=args.source_filenames)
+        
diff --git a/python/sidecar/cli.py b/python/sidecar/cli.py
@@ -43,4 +43,18 @@ def get_arg_parser():
         "--text",
         type=str,
     )
+
+    delete_parser = subparsers.add_parser(
+        "delete",
+        description="Deletes a Reference"
+    )
+    delete_parser.add_argument(
+        "--source_filenames",
+        nargs="*",
+        type=str
+    )
+    delete_parser.add_argument(
+        "--all",
+        action="store_true"
+    )
     return parser
diff --git a/python/sidecar/ingest.py b/python/sidecar/ingest.py
@@ -10,6 +10,7 @@
 from grobid_client.grobid_client import GrobidClient
 from sidecar import shared, typing
 
+from .settings import REFERENCES_JSON_PATH, UPLOADS_DIR
 from .shared import HiddenPrints, chunk_text, get_filename_md5
 from .storage import JsonStorage
 from .typing import Author, IngestResponse, Reference
@@ -35,21 +36,6 @@
 GROBID_SERVER_URL = "https://kermitt2-grobid.hf.space"
 GROBID_TIMEOUT = 60 * 5
 
-APPDATA_DIR = Path(
-    os.environ.get('APP_DATA_DIR', '/tmp')
-)
-PROJECT_NAME = Path(
-    os.environ.get('PROJECT_NAME', 'project-x')
-)
-UPLOADS_DIR = Path(
-    os.path.join(APPDATA_DIR, PROJECT_NAME, 'uploads')
-)
-REFERENCES_JSON_PATH = Path(
-    os.path.join(
-        APPDATA_DIR, PROJECT_NAME, '.storage', 'references.json'
-    )
-)
-
 
 class PDFIngestion:
     def __init__(self, input_dir: Path):
diff --git a/python/sidecar/settings.py b/python/sidecar/settings.py
@@ -0,0 +1,22 @@
+import os
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+APPDATA_DIR = Path(
+    os.environ.get('APP_DATA_DIR', '/tmp')
+)
+PROJECT_NAME = Path(
+    os.environ.get('PROJECT_NAME', 'project-x')
+)
+UPLOADS_DIR = Path(
+    os.path.join(APPDATA_DIR, PROJECT_NAME, 'uploads')
+)
+REFERENCES_JSON_PATH = Path(
+    os.path.join(
+        APPDATA_DIR, PROJECT_NAME, '.storage', 'references.json'
+    )
+)
diff --git a/python/sidecar/storage.py b/python/sidecar/storage.py
@@ -1,8 +1,9 @@
 import json
 import logging
 import os
+import sys
 
-from sidecar import typing
+from sidecar import settings, typing
 
 logging.root.setLevel(logging.NOTSET)
 
@@ -21,6 +22,12 @@
 logger.disabled = os.environ.get("SIDECAR_ENABLE_LOGGING", "false").lower() == "true"
 
 
+def delete_references(source_filenames: list[str] = [], all_: bool = False):
+    storage = JsonStorage(settings.REFERENCES_JSON_PATH)
+    storage.load()
+    storage.delete(source_filenames=source_filenames, all_=all_)
+
+
 class JsonStorage:
     def __init__(self, filepath: str):
         self.filepath = filepath
@@ -52,6 +59,53 @@ def save(self):
         with open(self.filepath, 'w') as f:
             json.dump(contents, f, indent=2, default=str)
     
+    def delete(self, source_filenames: list[str] = [], all_: bool = False):
+        """
+        Delete one or more References from storage.
+        
+        Parameters
+        ----------
+        source_filenames : list[str]
+            List of source filenames to be deleted
+        all_ : bool, default False
+            Delete all References from storage
+        """
+        if not source_filenames and not all_:
+            msg = ("`delete` operation requires one of "
+                   "`source_filenames` or `all_` input parameters")
+            raise ValueError(msg)
+    
+        # preprocess references into a dict of source_filename: Reference
+        # so that we can simply do `del refs[filename]`
+        refs = {
+            ref.source_filename: ref for ref in self.references
+        }
+
+        if all_:
+            source_filenames = list(refs.keys())
+
+        for filename in source_filenames:
+            try:
+                del refs[filename]
+            except KeyError:
+                msg = f"Unable to delete {filename}: not found in storage"
+                logger.warning(msg)
+                response = typing.DeleteStatusResponse(
+                    status=typing.ResponseStatus.ERROR,
+                    message=msg
+                )
+                sys.stdout.write(response.json())
+                return
+        
+        self.references = list(refs.values())
+        self.save()
+        
+        response = typing.DeleteStatusResponse(
+            status=typing.ResponseStatus.OK,
+            message=""
+        )
+        sys.stdout.write(response.json())
+    
     def update(self, target: typing.Reference):
         """
         Update a Reference in storage with the target reference.
diff --git a/python/sidecar/typing.py b/python/sidecar/typing.py
@@ -43,6 +43,10 @@ class Reference(RefStudioModel):
     metadata: dict[str, Any] = {}
 
 
+class ReferenceDelete(RefStudioModel):
+    source_filenames: list[str]
+
+
 class Author(RefStudioModel):
     full_name: str
     given_name: str | None = None
@@ -71,6 +75,11 @@ class IngestStatusResponse(RefStudioModel):
     reference_statuses: list[ReferenceStatus]
 
 
+class DeleteStatusResponse(RefStudioModel):
+    status: ResponseStatus
+    message: str
+
+
 class RewriteChoice(RefStudioModel):
     index: int
     text: str
@@ -86,6 +95,7 @@ class CliCommands(RefStudioModel):
     ingest_status: IngestStatusResponse
     rewrite: list[RewriteChoice]
     chat: list[ChatResponseChoice]
+    delete: DeleteStatusResponse
 
 
 Reference.update_forward_refs()
diff --git a/python/tests/test_storage.py b/python/tests/test_storage.py
@@ -1,3 +1,4 @@
+import json
 from pathlib import Path
 
 from sidecar import storage
@@ -69,4 +70,74 @@ def test_json_storage_update(monkeypatch, tmp_path):
     assert jstore.references[1].filename_md5 == "abcdefg123"
     assert jstore.references[1].citation_key is None
     assert len(jstore.references[1].authors) == 2
-    assert len(jstore.references[1].chunks) == 6
+    assert len(jstore.references[1].chunks) == 6
+
+
+def test_storage_delete_references(monkeypatch, tmp_path, capsys):
+    fp = Path(__file__).parent.joinpath("fixtures/data/references.json")
+    jstore = storage.JsonStorage(filepath=fp)
+    jstore.load()
+
+    # mock the filepath before save
+    # we do not want modify our test reference data for other tests
+    savepath = tmp_path.joinpath("references.json")
+    monkeypatch.setattr(jstore, "filepath", savepath)
+
+    # -------------
+
+    # test: delete with `all_=True`
+    # expect: all References should be deleted, json response in stdout = OK
+    jstore.delete(all_=True)
+
+    captured = capsys.readouterr()
+    output = json.loads(captured.out)
+
+    assert output['status'] == 'ok'
+    assert output['message'] == ""
+
+    # reload from `savepath` to check that the delete was successful
+    jstore = storage.JsonStorage(filepath=savepath)
+    jstore.load()
+
+    source_filenames_from_storage = [ref.source_filename for ref in jstore.references]
+    assert len(source_filenames_from_storage) == 0
+
+    # -------------
+
+    jstore = storage.JsonStorage(filepath=fp)
+    jstore.load()
+    monkeypatch.setattr(jstore, "filepath", savepath)
+
+    # test: delete source_filename that is present in references.json
+    # expect: corresponding Reference is deleted
+    #   json response in stdout should be of status: ok
+    to_be_deleted = ["some_file.pdf"]
+
+    jstore.delete(source_filenames=to_be_deleted)
+    
+    captured = capsys.readouterr()
+    output = json.loads(captured.out)
+
+    assert output['status'] == 'ok'
+    assert output['message'] == ""
+
+    # reload from `savepath` to check that the delete was successful
+    jstore = storage.JsonStorage(filepath=savepath)
+    jstore.load()
+
+    source_filenames_from_storage = [ref.source_filename for ref in jstore.references]
+    assert "some_file.pdf" not in source_filenames_from_storage
+
+    # -------------
+
+    # test: delete source_filename that is NOT present in references.json
+    # expect: json response in stdout should be of status: error
+    to_be_deleted = ["not_in_references_storage.pdf"]
+
+    jstore.delete(to_be_deleted)
+    
+    captured = capsys.readouterr()
+    output = json.loads(captured.out)
+
+    assert output['status'] == 'error'
+    assert output['message'] != ""
diff --git a/src/api/types.ts b/src/api/types.ts
@@ -19,6 +19,7 @@ export interface CliCommands {
   ingest_status: IngestStatusResponse;
   rewrite: RewriteChoice[];
   chat: ChatResponseChoice[];
+  delete: DeleteStatusResponse;
 }
 export interface IngestResponse {
   project_name: string;
@@ -67,3 +68,7 @@ export interface ChatResponseChoice {
   index: number;
   text: string;
 }
+export interface DeleteStatusResponse {
+  status: ResponseStatus;
+  message: string;
+}
