Background Index Attempt Creation (#1010)

This commit is contained in:
Yuhong Sun
2024-01-28 23:14:20 -08:00
committed by GitHub
parent c0c9c67534
commit 4b45164496
35 changed files with 1022 additions and 370 deletions

View File

@@ -7,7 +7,13 @@ Create Date: 2024-01-25 17:12:31.813160
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy import table, column, String, Integer, Boolean
from danswer.configs.model_configs import DOCUMENT_ENCODER_MODEL
from danswer.configs.model_configs import DOC_EMBEDDING_DIM
from danswer.configs.model_configs import NORMALIZE_EMBEDDINGS
from danswer.configs.model_configs import ASYM_QUERY_PREFIX
from danswer.configs.model_configs import ASYM_PASSAGE_PREFIX
from danswer.db.models import IndexModelStatus
# revision identifiers, used by Alembic.
@@ -26,6 +32,7 @@ def upgrade() -> None:
sa.Column("normalize", sa.Boolean(), nullable=False),
sa.Column("query_prefix", sa.String(), nullable=False),
sa.Column("passage_prefix", sa.String(), nullable=False),
sa.Column("index_name", sa.String(), nullable=False),
sa.Column(
"status",
sa.Enum(IndexModelStatus, native=False),
@@ -33,10 +40,46 @@ def upgrade() -> None:
),
sa.PrimaryKeyConstraint("id"),
)
EmbeddingModel = table(
"embedding_model",
column("id", Integer),
column("model_name", String),
column("model_dim", Integer),
column("normalize", Boolean),
column("query_prefix", String),
column("passage_prefix", String),
column("index_name", String),
column(
"status", sa.Enum(IndexModelStatus, name="indexmodelstatus", native=False)
),
)
op.bulk_insert(
EmbeddingModel,
[
{
"model_name": DOCUMENT_ENCODER_MODEL,
"model_dim": DOC_EMBEDDING_DIM,
"normalize": NORMALIZE_EMBEDDINGS,
"query_prefix": ASYM_QUERY_PREFIX,
"passage_prefix": ASYM_PASSAGE_PREFIX,
"index_name": "danswer_chunk",
"status": IndexModelStatus.PRESENT,
}
],
)
op.add_column(
"index_attempt",
sa.Column("embedding_model_id", sa.Integer(), nullable=True),
)
op.execute(
"UPDATE index_attempt SET embedding_model_id=1 WHERE embedding_model_id IS NULL"
)
op.alter_column(
"index_attempt",
"embedding_model_id",
existing_type=sa.Integer(),
nullable=False,
)
op.create_foreign_key(
"index_attempt__embedding_model_fk",
"index_attempt",