diff --git a/chat2db/database/postgres.py b/chat2db/database/postgres.py index c1c2a163f65d82b751435d2b1d9d099dec06a503..ea4470d49368fe289d8fbc1e07498191aa8ec6a2 100644 --- a/chat2db/database/postgres.py +++ b/chat2db/database/postgres.py @@ -34,26 +34,15 @@ class TableInfo(Base): TIMESTAMP(timezone=True), server_default=func.current_timestamp(), onupdate=func.current_timestamp()) - if config['DATABASE_TYPE'].lower() == 'opengauss': - __table_args__ = ( - Index( - 'table_note_vector_index', - table_note_vector, - opengauss_using='hnsw', - opengauss_with={'m': 16, 'ef_construction': 200}, - opengauss_ops={'vector': 'vector_cosine_ops'} - ), - ) - else: - __table_args__ = ( - Index( - 'table_note_vector_index', - table_note_vector, - postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 200}, - postgresql_ops={'table_note_vector': 'vector_cosine_ops'} - ), - ) + __table_args__ = ( + Index( + 'table_note_vector_index', + table_note_vector, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 200}, + postgresql_ops={'table_note_vector': 'vector_cosine_ops'} + ), + ) class ColumnInfo(Base): @@ -78,26 +67,15 @@ class SqlExample(Base): TIMESTAMP(timezone=True), server_default=func.current_timestamp(), onupdate=func.current_timestamp()) - if config['DATABASE_TYPE'].lower() == 'opengauss': - __table_args__ = ( - Index( - 'question_vector_index', - question_vector, - opengauss_using='hnsw', - opengauss_with={'m': 16, 'ef_construction': 200}, - opengauss_ops={'vector': 'vector_cosine_ops'} - ), - ) - else: - __table_args__ = ( - Index( - 'question_vector_index', - question_vector, - postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 200}, - postgresql_ops={'question_vector': 'vector_cosine_ops'} - ), - ) + __table_args__ = ( + Index( + 'question_vector_index', + question_vector, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 200}, + postgresql_ops={'question_vector': 'vector_cosine_ops'} + ), + ) class PostgresDB: diff --git a/data_chain/stores/database/database.py b/data_chain/stores/database/database.py index abe2c81fcd994329de3223582f890f58f1bf5449..4e8ae10d581bdeb25159026b668404bb1f3f08db 100644 --- a/data_chain/stores/database/database.py +++ b/data_chain/stores/database/database.py @@ -282,26 +282,15 @@ class DocumentEntity(Base): server_default=func.current_timestamp(), onupdate=func.current_timestamp() ) - if config["DATABASE_TYPE"].lower() == 'opengauss': - __table_args__ = ( - Index( - 'abstract_vector_index', - abstract_vector, - opengauss_using='hnsw', - opengauss_with={'m': 16, 'ef_construction': 200}, - opengauss_ops={'abstract_vector': 'vector_cosine_ops'} - ), - ) - else: - __table_args__ = ( - Index( - 'abstract_vector_index', - abstract_vector, - postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 200}, - postgresql_ops={'abstract_vector': 'vector_cosine_ops'} - ), - ) + __table_args__ = ( + Index( + 'abstract_vector_index', + abstract_vector, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 200}, + postgresql_ops={'abstract_vector': 'vector_cosine_ops'} + ), + ) class ChunkEntity(Base): @@ -333,26 +322,15 @@ class ChunkEntity(Base): TIMESTAMP(timezone=True), server_default=func.current_timestamp(), onupdate=func.current_timestamp()) - if config["DATABASE_TYPE"].lower() == 'opengauss': - __table_args__ = ( - Index( - 'text_vector_index', - text_vector, - opengauss_using='hnsw', - opengauss_with={'m': 16, 'ef_construction': 200}, - opengauss_ops={'text_vector': 'vector_cosine_ops'} - ), - ) - else: - __table_args__ = ( - Index( - 'text_vector_index', - text_vector, - postgresql_using='hnsw', - postgresql_with={'m': 16, 'ef_construction': 200}, - postgresql_ops={'text_vector': 'vector_cosine_ops'} - ), - ) + __table_args__ = ( + Index( + 'text_vector_index', + text_vector, + postgresql_using='hnsw', + postgresql_with={'m': 16, 'ef_construction': 200}, + postgresql_ops={'text_vector': 'vector_cosine_ops'} + ), + ) class ImageEntity(Base): @@ -568,7 +546,7 @@ class DataBase: database_url = f"postgresql+asyncpg://{config['DATABASE_USER']}:{encoded_password}@{config['DATABASE_HOST']}:{config['DATABASE_PORT']}/{config['DATABASE_DB']}" engine = create_async_engine( database_url, - echo=True, + echo=False, pool_recycle=300, pool_pre_ping=True )