feat: Add config support for pgvector

Added config support for using pgvector

Feature #COG-170
This commit is contained in:
Igor Ilic 2024-10-11 13:23:11 +02:00
parent 028f0b5100
commit d68a3be32e
7 changed files with 13 additions and 8 deletions

View file

@ -7,19 +7,19 @@ GRAPHISTRY_PASSWORD=
SENTRY_REPORTING_URL= SENTRY_REPORTING_URL=
GRAPH_DATABASE_PROVIDER="neo4j" # or "networkx" GRAPH_DATABASE_PROVIDER="neo4j" # "neo4j" or "networkx"
# Not needed if using networkx # Not needed if using networkx
GRAPH_DATABASE_URL= GRAPH_DATABASE_URL=
GRAPH_DATABASE_USERNAME= GRAPH_DATABASE_USERNAME=
GRAPH_DATABASE_PASSWORD= GRAPH_DATABASE_PASSWORD=
VECTOR_ENGINE_PROVIDER="qdrant" # or "weaviate" or "lancedb" VECTOR_ENGINE_PROVIDER="qdrant" # or "qdrant", "pgvector", "weaviate" or "lancedb"
# Not needed if using "lancedb" # Not needed if using "lancedb"
VECTOR_DB_URL= VECTOR_DB_URL=
VECTOR_DB_KEY= VECTOR_DB_KEY=
# Database provider # Database provider
DB_PROVIDER="sqlite" # or "postgres" DB_PROVIDER="sqlite" # "sqlite" or "postgres"
# Database name # Database name
DB_NAME=cognee_db DB_NAME=cognee_db

View file

@ -190,11 +190,11 @@ Cognee supports a variety of tools and services for different operations:
- **Local Setup**: By default, LanceDB runs locally with NetworkX and OpenAI. - **Local Setup**: By default, LanceDB runs locally with NetworkX and OpenAI.
- **Vector Stores**: Cognee supports Qdrant and Weaviate for vector storage. - **Vector Stores**: Cognee supports LanceDB, Qdrant, PGVector and Weaviate for vector storage.
- **Language Models (LLMs)**: You can use either Anyscale or Ollama as your LLM provider. - **Language Models (LLMs)**: You can use either Anyscale or Ollama as your LLM provider.
- **Graph Stores**: In addition to LanceDB, Neo4j is also supported for graph storage. - **Graph Stores**: In addition to NetworkX, Neo4j is also supported for graph storage.
- **User management**: Create individual user graphs and manage permissions - **User management**: Create individual user graphs and manage permissions

View file

@ -362,7 +362,7 @@ class LLMConfig(BaseModel):
apiKey: str apiKey: str
class VectorDBConfig(BaseModel): class VectorDBConfig(BaseModel):
provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"]] provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
url: str url: str
apiKey: str apiKey: str

View file

@ -26,6 +26,8 @@ def create_vector_engine(config: VectorConfig, embedding_engine):
api_key = config["vector_db_key"], api_key = config["vector_db_key"],
embedding_engine = embedding_engine embedding_engine = embedding_engine
) )
elif config["vector_db_provider"] == "pgvector":
pass
else: else:
from .lancedb.LanceDBAdapter import LanceDBAdapter from .lancedb.LanceDBAdapter import LanceDBAdapter

View file

@ -13,6 +13,9 @@ def get_settings():
}, { }, {
"value": "lancedb", "value": "lancedb",
"label": "LanceDB", "label": "LanceDB",
}, {
"value": "pgvector",
"label": "PGVector",
}] }]
vector_config = get_vectordb_config() vector_config = get_vectordb_config()

View file

@ -5,7 +5,7 @@ from cognee.infrastructure.databases.vector import get_vectordb_config
class VectorDBConfig(BaseModel): class VectorDBConfig(BaseModel):
url: str url: str
apiKey: str apiKey: str
provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"]] provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
async def save_vector_db_config(vector_db_config: VectorDBConfig): async def save_vector_db_config(vector_db_config: VectorDBConfig):
vector_config = get_vectordb_config() vector_config = get_vectordb_config()

View file

@ -62,7 +62,7 @@ services:
- cognee-network - cognee-network
postgres: postgres:
image: postgres:latest image: pgvector/pgvector:pg17
container_name: postgres container_name: postgres
environment: environment:
POSTGRES_USER: cognee POSTGRES_USER: cognee