Repository: meta-llama/llama-stack
File: src/llama_stack/providers/remote/vector_io/pgvector/pgvector.py (Python)

Code before:
            index = VectorStoreWithIndex(vector_store, index=pgvector_index, inference_api=self.inference_api)
            self.cache[vector_store.identifier] = index

    async def shutdown(self) -> None:
        if self.conn is not None:
            self.conn.close()
            log.info("Connection to PGVector database server closed")
        # Clean up mixin resources (file batch tasks)
        await super().shutdown()

    async def register_vector_store(self, vector_store: VectorStore) -> None:
        # Persist vector DB metadata in the KV store
        if self.kvstore is None:
            raise RuntimeError("KVStore not initialized. Call initialize() before registering vector stores.")

        # Save to kvstore for persistence
        key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
        await self.kvstore.set(key=key, value=vector_store.model_dump_json())

        # Upsert model metadata in Postgres
        upsert_models(self.conn, [(vector_store.identifier, vector_store)])

        # Create and cache the PGVector index table for the vector DB
        pgvector_index = PGVectorIndex(
            vector_store=vector_store, dimension=vector_store.embedding_dimension, conn=self.conn, kvstore=self.kvstore
        )
        await pgvector_index.initialize()
        index = VectorStoreWithIndex(vector_store, index=pgvector_index, inference_api=self.inference_api)
        self.cache[vector_store.identifier] = index

    async def unregister_vector_store(self, vector_store_id: str) -> None:
        # Remove provider index and cache
        if vector_store_id in self.cache:
            await self.cache[vector_store_id].index.delete()
            del self.cache[vector_store_id]

        # Delete vector DB metadata from KV store
        if self.kvstore is None:
            raise RuntimeError("KVStore not initialized. Call initialize() before unregistering vector stores.")
        await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")

    async def insert_chunks(
        self, vector_store_id: str, chunks: list[EmbeddedChunk], ttl_seconds: int | None = None
    ) -> None:
        index = await self._get_and_cache_vector_store_index(vector_store_id)
        await index.insert_chunks(chunks)

    async def query_chunks(
        self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None
    ) -> QueryChunksResponse:
        index = await self._get_and_cache_vector_store_index(vector_store_id)

Review comment (suggestion):
how about -

```
   execute("CREATE EXTENSION IF NOT EXISTS vector;")
   log.info(f"...{check_extension_version}")
```


side note - @ashwinb and i had a discussion about resource ownership wrt the s3 files provider. we landed on stack should not own the external resource, it should use the external resource, but there's a config option to allow ownership. in this case stack is already owning (it creates the tables) so this is fine.

Code after:
        self.cache = {}
        self.vector_store_table = None
        self.metadata_collection_name = "openai_vector_stores_metadata"

    async def initialize(self) -> None:
        # Create a safe config representation with masked password for logging
        safe_config = {**self.config.model_dump(exclude={"password"}), "password": "******"}
        log.info(f"Initializing PGVector memory adapter with config: {safe_config}")
        self.kvstore = await kvstore_impl(self.config.persistence)
        await self.initialize_openai_vector_stores()

        try:
            self.conn = psycopg2.connect(
                host=self.config.host,
                port=self.config.port,
                database=self.config.db,
                user=self.config.user,
                password=self.config.password,
            )
            self.conn.autocommit = True
            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
                version = check_extension_version(cur)
                if version:
                    log.info(f"Vector extension version: {version}")
                else:
                    create_vector_extension(cur)

                cur.execute(
                    """
                    CREATE TABLE IF NOT EXISTS metadata_store (
                        key TEXT PRIMARY KEY,
                        data JSONB
                    )
                """
                )
        except Exception as e:
            log.exception("Could not connect to PGVector database server")
            raise RuntimeError("Could not connect to PGVector database server") from e

        # Load existing vector stores from KV store into cache
        start_key = VECTOR_DBS_PREFIX
        end_key = f"{VECTOR_DBS_PREFIX}\xff"
        stored_vector_stores = await self.kvstore.values_in_range(start_key, end_key)
        for vector_store_data in stored_vector_stores:
            vector_store = VectorStore.model_validate_json(vector_store_data)
            pgvector_index = PGVectorIndex(
                vector_store=vector_store,
                dimension=vector_store.embedding_dimension,
                conn=self.conn,
                kvstore=self.kvstore,
            )