Merge pull request #309 from TogetherCrew/feat/302-telegram-raw-vecto…

…rize fix: Qdrant get_latest_point, added timestamp conversion support!
TogetherCrew · Oct 22, 2024 · 4e8fea8 · 4e8fea8
2 parents 6ab0726 + 439ca37
commit 4e8fea8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/dags/hivemind_etl_helpers/ingestion_pipeline.py b/dags/hivemind_etl_helpers/ingestion_pipeline.py
@@ -165,7 +165,15 @@ def get_latest_document_date(
                     logging.info("No documents found in the collection.")
                     latest_date = None
                 else:
-                    latest_date = parse(latest_document[0][0].payload[field_name])
+                    date_field = latest_document[0][0].payload[field_name]
+
+                    # if it was float timestamp
+                    if field_schema == models.PayloadSchemaType.FLOAT:
+                        latest_date = datetime.fromtimestamp(date_field)
+
+                    # it should be datetime in any other case
+                    else:
+                        latest_date = parse(date_field)
 
             else:
                 raise ValueError(

diff --git a/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py b/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py
@@ -106,11 +106,12 @@ def test_extract_chats_multiple_chats(self):
             )
 
         chat_ids = self.tc_chats.extract_chats()
-        self.assertEqual(
-            chat_ids,
-            [
-                (100001, "test chat"),
-                (100002, "test chat 2"),
-                (100003, "test chat 3"),
-            ],
-        )
+
+        self.assertEqual(len(chat_ids), 3)
+        expected_chats = [
+            (100001, "test chat"),
+            (100002, "test chat 2"),
+            (100003, "test chat 3"),
+        ]
+        for id in chat_ids:
+            self.assertIn(id, expected_chats)