From ccd41220cd8ac470f0b0500b394157e20b691d69 Mon Sep 17 00:00:00 2001
From: "dristy.cd" <dristy@clouddefense.io>
Date: Wed, 10 Jul 2024 10:37:38 +0530
Subject: [PATCH] Updating variable content_checksum to pb_checksum

---
 .../langchain_community/chains/pebblo_retrieval/models.py | 2 +-
 .../langchain_community/document_loaders/pebblo.py        | 8 +++-----
 .../tests/unit_tests/document_loaders/test_pebblo.py      | 6 ++++--
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/libs/community/langchain_community/chains/pebblo_retrieval/models.py b/libs/community/langchain_community/chains/pebblo_retrieval/models.py
index 9283a3a12e3201..13a54537b91699 100644
--- 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
index 41ff9b3aceff46..2e31b370cbd231 100644
--- 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
index 922acbc8818e67..d0a71faae7a8e8 100644
--- 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",
         ),