Merge pull request #72 from TamiTakamiya/TamiTakamiya/llamastack-pgvector

Llama Stack PGVector Support

Author: tisnik

README.md

@@ -74,6 +74,7 @@ You can generate the vector database either using
 2. [Llama-Index Postgres (PGVector) Vector Store](#postgres-pgvector-vector-store)
 3. [Llama-Stack Faiss Vector-IO](#llama-stack-faiss)
 4. [Llama-Stack SQLite-vec Vector-IO](#llama-stack-sqlite-vec)
+5. [Llama-Stack Postgres (PGVector) Vector Store](#llama-stack-postgres-pgvector-vector-store)
 
 Llama-Index approaches require you to download the embedding model, and we also
 recommend it for Llama-Stack targets even though it should work even without
@@ -320,7 +321,54 @@ python scripts/query_rag.py \
   -k 5 \
   -q "how can I configure a cinder backend"
 ```
+### Llama-Stack Postgres (PGVector) Vector Store
 
+To generate a vector database stored in Postgres (PGVector) for Llama-Stack, run the following
+commands:
+
+1. Start Postgres with the pgvector extension by running:
+
+    ```bash
+    make start-postgres-debug
+    ```
+
+    The `data` folder of Postgres is created at `./postgresql/data`. Note that this command
+    also creates `./output`, which is not used for the Llama-Stack version while it is used for Llama-Index version.
+
+2. Run:
+
+    ```bash
+    POSTGRES_USER=postgres \
+    POSTGRES_PASSWORD=somesecret \
+    POSTGRES_HOST=localhost \
+    POSTGRES_PORT=15432 \
+    POSTGRES_DATABASE=postgres \
+    uv run python ./custom_processor.py \
+     -o ./output \
+     -f custom_docs/0.1/ \
+     -md embeddings_model/ \
+     -mn sentence-transformers/all-mpnet-base-v2 \
+     -i custom_docs-0_1 \
+     --vector-store-type llamastack-pgvector
+    ```
+
+    Which generates embeddings on PostgreSQL, which can be used for RAG.
+
+3. When you run `query_rag.py` to check some results, specify these environment variables for database access:
+
+   ```bash
+    POSTGRES_USER=postgres \
+    POSTGRES_PASSWORD=somesecret \
+    POSTGRES_HOST=localhost \
+    POSTGRES_PORT=15432 \
+    POSTGRES_DATABASE=postgres \
+    uv run python scripts/query_rag.py \
+    -p vector_db/custom_docs/0.1 \
+    -x custom-docs-0_1 \
+    -m embeddings_model \
+    -k 5 \
+    -q "how can I configure a cinder backend"
+   ```
 ## Update lockfiles
 
 The lock file is used in this repository:

scripts/query_rag.py

@@ -194,8 +194,15 @@ def _llama_stack_query(args: argparse.Namespace) -> None:  # noqa: C901
             "mode": "vector",  # "vector", "keyword", or "hybrid". Default "vector"
             "score_threshold": 0,
         }
+        vector_stores = cfg.get("registered_resources", {}).get("vector_stores", [])
+        if not vector_stores:
+            logging.error("No vector stores found in configuration")
+            exit(1)
+        vector_store_id = vector_stores[0]["vector_store_id"]
         res = client.vector_io.query(
-            vector_store_id=args.product_index, query=args.query, params=query_cfg
+            vector_store_id=vector_store_id,
+            query=args.query,
+            params=query_cfg,
         )
 
         if len(res.chunks) == 0:
@@ -308,7 +315,9 @@ def _llama_stack_query(args: argparse.Namespace) -> None:  # noqa: C901
 
     vector_store_type = args.vector_store_type
     if args.vector_store_type == "auto":
-        if os.path.exists(os.path.join(args.db_path, "metadata.json")):
+        if os.environ.get("POSTGRES_DATABASE"):
+            args.vector_store_type = "llamastack-pgvector"
+        elif os.path.exists(os.path.join(args.db_path, "metadata.json")):
             args.vector_store_type = "faiss"
         elif os.path.exists(os.path.join(args.db_path, "sqlite-vec_store.db")):
             args.vector_store_type = "llamastack-sqlite-vec"

src/lightspeed_rag_content/document_processor.py

@@ -234,11 +234,9 @@ class _LlamaStackDB(_BaseDB):
     provider_type: inline::rag-runtime
   vector_io:
   - config:
-      persistence:
-        namespace: vector_io::{provider_type}
-        backend: kv_rag
+      {vector_io_cfg}
     provider_id: {index_id}
-    provider_type: inline::{provider_type}
+    provider_type: {provider_type_prefix}::{provider_type}
 storage:
   backends:
     kv_rag:
@@ -284,6 +282,19 @@ class _LlamaStackDB(_BaseDB):
     provider_id: {vector_io_provider_id}
     vector_store_id: {vector_store_id}"""
 
+    # Template for vector_io/config section
+    VECTOR_IO_CONFIG_TEMPLATE_FOR_SQLITE = """persistence:
+        namespace: vector_io::{provider_type}
+        backend: kv_rag"""
+    VECTOR_IO_CONFIG_TEMPLATE_FOR_PGVECTOR = """persistence:
+        namespace: vector_io::{provider_type}
+        backend: kv_default
+      host: ${{env.POSTGRES_HOST}}
+      port: ${{env.POSTGRES_PORT}}
+      db: ${{env.POSTGRES_DATABASE}}
+      user: ${{env.POSTGRES_USER}}
+      password: ${{env.POSTGRES_PASSWORD}}"""
+
     CFG_FILENAME = "llama-stack.yaml"
 
     def __init__(self, config: _Config):
@@ -317,6 +328,7 @@ def __init__(self, config: _Config):
         assert config.vector_store_type in (  # noqa: S101
             "llamastack-faiss",
             "llamastack-sqlite-vec",
+            "llamastack-pgvector",
         )
 
         super().__init__(config)
@@ -363,14 +375,33 @@ def write_yaml_config(
         self, index_id: str, filename: str, db_file: str, files_metadata_db_file: str
     ) -> None:
         """Write a llama-stack configuration file using class templates."""
-        if self.config.vector_store_type == "llamastack-faiss":
-            vector_io_cfg = ""
+        if self.config.vector_store_type == "llamastack-pgvector":
+            provider_type_prefix = "remote"
+            required_vars = [
+                "POSTGRES_USER",
+                "POSTGRES_PASSWORD",
+                "POSTGRES_HOST",
+                "POSTGRES_PORT",
+                "POSTGRES_DATABASE",
+            ]
+            missing = [v for v in required_vars if not os.getenv(v)]
+            if missing:
+                raise ValueError(
+                    f"Missing required environment variables: {', '.join(missing)}"
+                )
+            vector_io_cfg = self.VECTOR_IO_CONFIG_TEMPLATE_FOR_PGVECTOR.format(
+                provider_type=self.provider_type,
+            )
         else:
-            vector_io_cfg = "db_path: " + db_file
+            provider_type_prefix = "inline"
+            vector_io_cfg = self.VECTOR_IO_CONFIG_TEMPLATE_FOR_SQLITE.format(
+                provider_type=self.provider_type,
+            )
 
         with open(filename, "w", encoding="utf-8") as fd:
             data = self.TEMPLATE.format(
                 index_id=index_id,
+                provider_type_prefix=provider_type_prefix,
                 provider_type=self.provider_type,
                 vector_io_cfg=vector_io_cfg,
                 kv_db_path=db_file,

src/lightspeed_rag_content/utils.py

@@ -121,7 +121,13 @@ def get_common_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--vector-store-type",
         default="faiss",
-        choices=["faiss", "postgres", "llamastack-faiss", "llamastack-sqlite-vec"],
+        choices=[
+            "faiss",
+            "postgres",
+            "llamastack-faiss",
+            "llamastack-sqlite-vec",
+            "llamastack-pgvector",
+        ],
         help="vector store type to be used.",
     )
     parser.add_argument(