Feat: add PostgreSQL extensions for vector and AGE
Browse files- Ensure VECTOR extension is available when PostgreSQL init
- Ensure AGE extension is available when PGGraphStorage init
- lightrag/kg/postgres_impl.py +29 -0
lightrag/kg/postgres_impl.py
CHANGED
@@ -73,6 +73,10 @@ class PostgreSQLDB:
|
|
73 |
max_size=self.max,
|
74 |
)
|
75 |
|
|
|
|
|
|
|
|
|
76 |
logger.info(
|
77 |
f"PostgreSQL, Connected to database at {self.host}:{self.port}/{self.database}"
|
78 |
)
|
@@ -82,6 +86,26 @@ class PostgreSQLDB:
|
|
82 |
)
|
83 |
raise
|
84 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
85 |
@staticmethod
|
86 |
async def configure_age(connection: asyncpg.Connection, graph_name: str) -> None:
|
87 |
"""Set the Apache AGE environment and creates a graph if it does not exist.
|
@@ -1851,6 +1875,11 @@ class PGGraphStorage(BaseGraphStorage):
|
|
1851 |
f"PostgreSQL Graph initialized: workspace='{self.workspace}', graph_name='{self.graph_name}'"
|
1852 |
)
|
1853 |
|
|
|
|
|
|
|
|
|
|
|
1854 |
# Execute each statement separately and ignore errors
|
1855 |
queries = [
|
1856 |
f"SELECT create_graph('{self.graph_name}')",
|
|
|
73 |
max_size=self.max,
|
74 |
)
|
75 |
|
76 |
+
# Ensure VECTOR extension is available
|
77 |
+
async with self.pool.acquire() as connection:
|
78 |
+
await self.configure_vector_extension(connection)
|
79 |
+
|
80 |
logger.info(
|
81 |
f"PostgreSQL, Connected to database at {self.host}:{self.port}/{self.database}"
|
82 |
)
|
|
|
86 |
)
|
87 |
raise
|
88 |
|
89 |
+
@staticmethod
|
90 |
+
async def configure_vector_extension(connection: asyncpg.Connection) -> None:
|
91 |
+
"""Create VECTOR extension if it doesn't exist for vector similarity operations."""
|
92 |
+
try:
|
93 |
+
await connection.execute("CREATE EXTENSION IF NOT EXISTS vector") # type: ignore
|
94 |
+
logger.info("VECTOR extension ensured for PostgreSQL")
|
95 |
+
except Exception as e:
|
96 |
+
logger.warning(f"Could not create VECTOR extension: {e}")
|
97 |
+
# Don't raise - let the system continue without vector extension
|
98 |
+
|
99 |
+
@staticmethod
|
100 |
+
async def configure_age_extension(connection: asyncpg.Connection) -> None:
|
101 |
+
"""Create AGE extension if it doesn't exist for graph operations."""
|
102 |
+
try:
|
103 |
+
await connection.execute("CREATE EXTENSION IF NOT EXISTS age") # type: ignore
|
104 |
+
logger.info("AGE extension ensured for PostgreSQL")
|
105 |
+
except Exception as e:
|
106 |
+
logger.warning(f"Could not create AGE extension: {e}")
|
107 |
+
# Don't raise - let the system continue without AGE extension
|
108 |
+
|
109 |
@staticmethod
|
110 |
async def configure_age(connection: asyncpg.Connection, graph_name: str) -> None:
|
111 |
"""Set the Apache AGE environment and creates a graph if it does not exist.
|
|
|
1875 |
f"PostgreSQL Graph initialized: workspace='{self.workspace}', graph_name='{self.graph_name}'"
|
1876 |
)
|
1877 |
|
1878 |
+
# Create AGE extension and configure graph environment once at initialization
|
1879 |
+
async with self.db.pool.acquire() as connection:
|
1880 |
+
# First ensure AGE extension is created
|
1881 |
+
await PostgreSQLDB.configure_age_extension(connection)
|
1882 |
+
|
1883 |
# Execute each statement separately and ignore errors
|
1884 |
queries = [
|
1885 |
f"SELECT create_graph('{self.graph_name}')",
|