gzdaniel commited on
Commit
b09e70e
·
1 Parent(s): f91fc9e

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

Files changed (1) hide show
  1. 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}')",