From 19a38d9310dd884a45741e1619402811fe4138fb Mon Sep 17 00:00:00 2001 From: yangdx Date: Mon, 21 Jul 2025 01:46:41 +0800 Subject: [PATCH] Feat: add PostgreSQL extensions for vector and AGE - Ensure VECTOR extension is available when PostgreSQL init - Ensure AGE extension is available when PGGraphStorage init --- lightrag/kg/postgres_impl.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index 43825128..842e1e54 100644 --- a/lightrag/kg/postgres_impl.py +++ b/lightrag/kg/postgres_impl.py @@ -73,6 +73,10 @@ class PostgreSQLDB: max_size=self.max, ) + # Ensure VECTOR extension is available + async with self.pool.acquire() as connection: + await self.configure_vector_extension(connection) + logger.info( f"PostgreSQL, Connected to database at {self.host}:{self.port}/{self.database}" ) @@ -82,6 +86,26 @@ class PostgreSQLDB: ) raise + @staticmethod + async def configure_vector_extension(connection: asyncpg.Connection) -> None: + """Create VECTOR extension if it doesn't exist for vector similarity operations.""" + try: + await connection.execute("CREATE EXTENSION IF NOT EXISTS vector") # type: ignore + logger.info("VECTOR extension ensured for PostgreSQL") + except Exception as e: + logger.warning(f"Could not create VECTOR extension: {e}") + # Don't raise - let the system continue without vector extension + + @staticmethod + async def configure_age_extension(connection: asyncpg.Connection) -> None: + """Create AGE extension if it doesn't exist for graph operations.""" + try: + await connection.execute("CREATE EXTENSION IF NOT EXISTS age") # type: ignore + logger.info("AGE extension ensured for PostgreSQL") + except Exception as e: + logger.warning(f"Could not create AGE extension: {e}") + # Don't raise - let the system continue without AGE extension + @staticmethod async def configure_age(connection: asyncpg.Connection, graph_name: str) -> None: """Set the Apache AGE environment and creates a graph if it does not exist. @@ -1851,6 +1875,11 @@ class PGGraphStorage(BaseGraphStorage): f"PostgreSQL Graph initialized: workspace='{self.workspace}', graph_name='{self.graph_name}'" ) + # Create AGE extension and configure graph environment once at initialization + async with self.db.pool.acquire() as connection: + # First ensure AGE extension is created + await PostgreSQLDB.configure_age_extension(connection) + # Execute each statement separately and ignore errors queries = [ f"SELECT create_graph('{self.graph_name}')",