Updating variable content_checksum to pb_checksum

daxa-ai · Jul 10, 2024 · ccd4122 · ccd4122
1 parent ca2296b
commit ccd4122
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/libs/community/langchain_community/chains/pebblo_retrieval/models.py b/libs/community/langchain_community/chains/pebblo_retrieval/models.py
@@ -129,7 +129,7 @@ class Context(BaseModel):
     retrieved_from: Optional[str]
     doc: Optional[str]
     vector_db: str
-    content_checksum: Optional[str]
+    pb_checksum: Optional[str]
 
 
 class Prompt(BaseModel):

diff --git a/libs/community/langchain_community/document_loaders/pebblo.py b/libs/community/langchain_community/document_loaders/pebblo.py
@@ -255,9 +255,7 @@ def _classify_doc(self, loading_end: bool = False) -> dict:
                     classified_data = classified_docs.get(doc_data["pb_id"], {})
                     doc_data.update(
                         {
-                            "content_checksum": classified_data.get(
-                                "content_checksum", None
-                            ),
+                            "pb_checksum": classified_data.get("pb_checksum", None),
                             "loader_source_path": classified_data.get(
                                 "loader_source_path", None
                             ),
@@ -534,6 +532,6 @@ def _add_pebblo_specific_metadata(self, classified_docs: dict) -> None:
                 )
             )
             doc_metadata["pb_id"] = doc.pb_id
-            doc_metadata["content_checksum"] = classified_docs.get(doc.pb_id, {}).get(
-                "content_checksum", None
+            doc_metadata["pb_checksum"] = classified_docs.get(doc.pb_id, {}).get(
+                "pb_checksum", None
             )
diff --git a/libs/community/tests/unit_tests/document_loaders/test_pebblo.py b/libs/community/tests/unit_tests/document_loaders/test_pebblo.py
@@ -70,7 +70,8 @@ def test_csv_loader_load_valid_data(mocker: MockerFixture) -> None:
                 "row": 0,
                 "full_path": full_file_path,
                 "pb_id": "0",
-                "content_checksum": None,  # For UT as here we are not calculating checksum
+                # For UT as here we are not calculating checksum
+                "pb_checksum": None,
             },
             page_content="column1: value1\ncolumn2: value2\ncolumn3: value3",
         ),
@@ -80,7 +81,8 @@ def test_csv_loader_load_valid_data(mocker: MockerFixture) -> None:
                 "row": 1,
                 "full_path": full_file_path,
                 "pb_id": "1",
-                "content_checksum": None,  # For UT as here we are not calculating checksum
+                # For UT as here we are not calculating checksum
+                "pb_checksum": None,
             },
             page_content="column1: value4\ncolumn2: value5\ncolumn3: value6",
         ),