Feat/vector db pgvector (#3879)

2024-05-10 17:20:30 +08:00 · 2024-05-10 17:20:30 +08:00 · 875249eb00
commit 875249eb00
parent 4d5a4e4cef
13 changed files with 316 additions and 9 deletions
--- a/api/commands.py
+++ b/api/commands.py
@ -305,6 +305,14 @@ def migrate_knowledge_vector_database():
                        "vector_store": {"class_prefix": collection_name}
                    }
                    dataset.index_struct = json.dumps(index_struct_dict)
+                elif vector_type == "pgvector":
+                    dataset_id = dataset.id
+                    collection_name = Dataset.gen_collection_name_by_id(dataset_id)
+                    index_struct_dict = {
+                        "type": 'pgvector',
+                        "vector_store": {"class_prefix": collection_name}
+                    }
+                    dataset.index_struct = json.dumps(index_struct_dict)
                else:
                    raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.")