feat: unify settings for vector and nodestore connections to PostgreSQL (#1730)

* Unify pgvector and postgres connection settings

* Remove local changes

* Update file pgvector->postgres
This commit is contained in:
Brett England 2024-03-15 04:55:17 -04:00 committed by GitHub
parent 68b3a34b03
commit 63de7e4930
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 39 additions and 45 deletions

View file

@ -38,7 +38,7 @@ class VectorStoreComponent:
def __init__(self, settings: Settings) -> None:
self.settings = settings
match settings.vectorstore.database:
case "pgvector":
case "postgres":
try:
from llama_index.vector_stores.postgres import ( # type: ignore
PGVectorStore,
@ -48,15 +48,17 @@ class VectorStoreComponent:
"Postgres dependencies not found, install with `poetry install --extras vector-stores-postgres`"
) from e
if settings.pgvector is None:
if settings.postgres is None:
raise ValueError(
"PGVectorStore settings not found. Please provide settings."
"Postgres settings not found. Please provide settings."
)
self.vector_store = typing.cast(
VectorStore,
PGVectorStore.from_params(
**settings.pgvector.model_dump(exclude_none=True)
**settings.postgres.model_dump(exclude_none=True),
table_name="embeddings",
embed_dim=settings.embedding.embed_dim,
),
)

View file

@ -105,7 +105,7 @@ class LLMSettings(BaseModel):
class VectorstoreSettings(BaseModel):
database: Literal["chroma", "qdrant", "pgvector"]
database: Literal["chroma", "qdrant", "postgres"]
class NodeStoreSettings(BaseModel):
@ -177,6 +177,10 @@ class EmbeddingSettings(BaseModel):
"Do not set it higher than your number of threads of your CPU."
),
)
embed_dim: int = Field(
384,
description="The dimension of the embeddings stored in the Postgres database",
)
class SagemakerSettings(BaseModel):
@ -280,17 +284,6 @@ class PostgresSettings(BaseModel):
)
class PGVectorSettings(PostgresSettings):
embed_dim: int = Field(
384,
description="The dimension of the embeddings stored in the Postgres database",
)
table_name: str = Field(
"embeddings",
description="The name of the table in the Postgres database where the embeddings are stored",
)
class QdrantSettings(BaseModel):
location: str | None = Field(
None,
@ -360,7 +353,6 @@ class Settings(BaseModel):
nodestore: NodeStoreSettings
qdrant: QdrantSettings | None = None
postgres: PostgresSettings | None = None
pgvector: PGVectorSettings | None = None
"""