Merge pull request #16 from epinzur/async-ingest

added async support for ingest
epinzur · Jun 14, 2024 · bc107f4 · bc107f4
2 parents c7ee688 + e5500db
commit bc107f4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/ragulate/pipelines/ingest_pipeline.py b/ragulate/pipelines/ingest_pipeline.py
@@ -1,3 +1,4 @@
+import asyncio
 from typing import Any, Dict
 
 from tqdm import tqdm
@@ -24,4 +25,7 @@ def ingest(self):
             source_files.extend(dataset.get_source_file_paths())
 
         for source_file in tqdm(source_files):
-            ingest_method(file_path=source_file, **params)
+            if asyncio.iscoroutinefunction(ingest_method):
+                asyncio.run(ingest_method(file_path=source_file, **params))
+            else:
+                ingest_method(file_path=source_file, **params)
diff --git a/ragulate/pipelines/query_pipeline.py b/ragulate/pipelines/query_pipeline.py
@@ -149,7 +149,9 @@ def query(self):
                         pipeline.invoke(query)
                 except Exception as e:
                     # TODO: figure out why the logger isn't working after tru-lens starts. For now use print()
-                    print(f"ERROR: Query: '{query}' caused exception, skipping. Exception {e}")
+                    print(
+                        f"ERROR: Query: '{query}' caused exception, skipping. Exception {e}"
+                    )
                     logger.error(f"Query: '{query}' caused exception, skipping.")
                 finally:
                     self.update_progress(query_change=1)