fix kdbi

medirag/index/runner.py

@@ -23,6 +23,6 @@ document_processor = DailyMedDocumentProcessor(extracted_dir=extracted_dir)
 documents = document_processor.load_documents()
 
 # Index and query documents
-indexer = KDBAIDailyMedIndexer()
+indexer = KDBAIDailyMedIndexer(table_name="daily_med_v2")
 indexer.load_index(documents=documents)
 print("done")

misc/create_kdbai_table.py

@@ -7,25 +7,25 @@ load_dotenv()
 
 session = kdbai.Session(api_key=os.getenv("KDBAI_API_KEY"), endpoint=os.getenv("KDBAI_ENDPOINT"))
 
-schema = dict(
-    columns=[
-        dict(name="document_id", pytype="bytes"),
-        dict(name="text", pytype="bytes"),
-        dict(
-            name="embedding",
-            vectorIndex=dict(type="flat", metric="CS", dims=768),
-        ),
-    ]
-)
+schema = [
+    {"name": "document_id", "type": "bytes"},
+    {"name": "text", "type": "bytes"},
+    {"name": "embedding", "type": "float32s"},
+]
 
-KDBAI_TABLE_NAME = "daily_med_v1"
+indexFlat = {
+    "name": "flat_index",
+    "type": "flat",
+    "column": "embedding",
+    "params": {"dims": 768, "metric": "CS"},
+}
 
-table = session.database("default").table(KDBAI_TABLE_NAME)
-print()
-# # First ensure the table does not already exist
-# if KDBAI_TABLE_NAME in session.list():
-#     pass
-#     # session.table(KDBAI_TABLE_NAME).drop()
-# else:
-#     # Create the table
-#     table = session.database("default").create_table(KDBAI_TABLE_NAME, schema=schema)
+KDBAI_TABLE_NAME = "daily_med_v2"
+database = session.database("default")
+# First ensure the table does not already exist
+for table in database.tables:
+    if table.name == KDBAI_TABLE_NAME:
+        ## table.drop()
+        print("exist")
+        break
+table = database.create_table(KDBAI_TABLE_NAME, schema=schema, indexes=[indexFlat])