Use ragstack-colbert and ragstack-langchain types in ragulate (#632)

datastax · Jul 30, 2024 · 3d80907 · 3d80907
1 parent b4430e9
commit 3d80907
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/libs/ragulate/colbert_chunk_size_and_k.py b/libs/ragulate/colbert_chunk_size_and_k.py
@@ -25,8 +25,8 @@
 
 batch_size = 640
 
-astra_token = os.getenv("ASTRA_DB_TOKEN")
-database_id = os.getenv("ASTRA_DB_ID")
+astra_token = os.environ["ASTRA_DB_TOKEN"]
+database_id = os.environ["ASTRA_DB_ID"]
 keyspace = "colbert"
 
 
@@ -38,7 +38,7 @@
 
 
 def get_embedding_model(chunk_size: int) -> ColbertEmbeddingModel:
-    return ColbertEmbeddingModel(doc_maxlen=chunk_size, batch_size=batch_size)
+    return ColbertEmbeddingModel(doc_maxlen=chunk_size, chunk_batch_size=batch_size)
 
 
 def get_database(chunk_size: int) -> CassandraDatabase:

diff --git a/libs/ragulate/poetry.lock b/libs/ragulate/poetry.lock
diff --git a/libs/ragulate/pyproject.toml b/libs/ragulate/pyproject.toml
@@ -42,7 +42,7 @@ types-seaborn = "^0.13.2"
 types-setuptools = "^70.0.0"
 types-tqdm = ">=4.66.0"
 pandas-stubs = ">=2.2.2"
-ragstack-ai-colbert = "^1.0.5"
+ragstack-ai-colbert = "^1.0.6"
 ragstack-ai-langchain = "^1.4.0"
 
 [build-system]
@@ -62,6 +62,5 @@ show_error_codes = true
 show_error_context = true
 
 [[tool.mypy.overrides]]
-# TODO: remove ragstack exclusions once packages containing py.typed are published
-module = "trulens_eval.*,plotly.*,cerberus.*,transformers.*,ragstack_colbert.*,ragstack_langchain.*"
+module = "trulens_eval.*,plotly.*,cerberus.*,transformers.*"
 ignore_missing_imports = true