fix: Reuse PG vector store to avoid recreating sqlalchemy engine (#95)

.changeset/slow-rivers-breathe.md

+---
+"create-llama": patch
+---
+
+Fix postgres connection leaking issue

templates/components/vectordbs/python/none/index.py

-import logging
 import os
+import logging
+from datetime import timedelta
 
+from cachetools import cached, TTLCache
 from llama_index.core.storage import StorageContext
 from llama_index.core.indices import load_index_from_storage
 
 logger = logging.getLogger("uvicorn")
 
 
+@cached(
+    TTLCache(maxsize=10, ttl=timedelta(minutes=5).total_seconds()),
+    key=lambda *args, **kwargs: "global_storage_context",
+)
+def get_storage_context(persist_dir: str) -> StorageContext:
+    return StorageContext.from_defaults(persist_dir=persist_dir)
+
+
 def get_index():
     storage_dir = os.getenv("STORAGE_DIR", "storage")
     # check if storage already exists
@@ -14,7 +24,7 @@ def get_index():
         return None
     # load the existing index
     logger.info(f"Loading index from {storage_dir}...")
-    storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
+    storage_context = get_storage_context(storage_dir)
     index = load_index_from_storage(storage_context)
     logger.info(f"Finished loading index from {storage_dir}")
     return index

templates/components/vectordbs/python/pg/vectordb.py

@@ -5,26 +5,33 @@ from urllib.parse import urlparse
 PGVECTOR_SCHEMA = "public"
 PGVECTOR_TABLE = "llamaindex_embedding"
 
+vector_store: PGVectorStore = None
+
 
 def get_vector_store():
-    original_conn_string = os.environ.get("PG_CONNECTION_STRING")
-    if original_conn_string is None or original_conn_string == "":
-        raise ValueError("PG_CONNECTION_STRING environment variable is not set.")
-
-    # The PGVectorStore requires both two connection strings, one for psycopg2 and one for asyncpg
-    # Update the configured scheme with the psycopg2 and asyncpg schemes
-    original_scheme = urlparse(original_conn_string).scheme + "://"
-    conn_string = original_conn_string.replace(
-        original_scheme, "postgresql+psycopg2://"
-    )
-    async_conn_string = original_conn_string.replace(
-        original_scheme, "postgresql+asyncpg://"
-    )
-
-    return PGVectorStore(
-        connection_string=conn_string,
-        async_connection_string=async_conn_string,
-        schema_name=PGVECTOR_SCHEMA,
-        table_name=PGVECTOR_TABLE,
-        embed_dim=int(os.environ.get("EMBEDDING_DIM", 1024)),
-    )
+    global vector_store
+
+    if vector_store is None:
+        original_conn_string = os.environ.get("PG_CONNECTION_STRING")
+        if original_conn_string is None or original_conn_string == "":
+            raise ValueError("PG_CONNECTION_STRING environment variable is not set.")
+
+        # The PGVectorStore requires both two connection strings, one for psycopg2 and one for asyncpg
+        # Update the configured scheme with the psycopg2 and asyncpg schemes
+        original_scheme = urlparse(original_conn_string).scheme + "://"
+        conn_string = original_conn_string.replace(
+            original_scheme, "postgresql+psycopg2://"
+        )
+        async_conn_string = original_conn_string.replace(
+            original_scheme, "postgresql+asyncpg://"
+        )
+
+        vector_store = PGVectorStore(
+            connection_string=conn_string,
+            async_connection_string=async_conn_string,
+            schema_name=PGVECTOR_SCHEMA,
+            table_name=PGVECTOR_TABLE,
+            embed_dim=int(os.environ.get("EMBEDDING_DIM", 1024)),
+        )
+
+    return vector_store

templates/types/streaming/fastapi/pyproject.toml

@@ -16,6 +16,7 @@ python-dotenv = "^1.0.0"
 aiostream = "^0.5.2"
 llama-index = "0.10.28"
 llama-index-core = "0.10.28"
+cachetools = "^5.3.3"
 
 [build-system]
 requires = ["poetry-core"]