From 1f818b7a4d2de8601be165cd7615e254112ff868 Mon Sep 17 00:00:00 2001 From: Qunfei Wu Date: Fri, 22 Nov 2024 14:06:51 +0100 Subject: [PATCH] Adding a feature called create_extension in init method. --- .../document_stores/pgvector/document_store.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py index feb99b1f4..6301ced27 100644 --- a/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py +++ b/integrations/pgvector/src/haystack_integrations/document_stores/pgvector/document_store.py @@ -90,6 +90,7 @@ def __init__( hnsw_index_name: str = "haystack_hnsw_index", hnsw_ef_search: Optional[int] = None, keyword_index_name: str = "haystack_keyword_index", + create_extension: bool = True, ): """ Creates a new PgvectorDocumentStore instance. @@ -153,6 +154,7 @@ def __init__( self.hnsw_ef_search = hnsw_ef_search self.keyword_index_name = keyword_index_name self.language = language + self.create_extension = create_extension self._connection = None self._cursor = None self._dict_cursor = None @@ -182,6 +184,8 @@ def _create_connection(self): conn_str = self.connection_string.resolve_value() or "" connection = connect(conn_str) connection.autocommit = True + if self.create_extension: + connection.execute("CREATE EXTENSION IF NOT EXISTS vector") register_vector(connection) # Note: this must be called before creating the cursors. self._connection = connection