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

Code before:
                return result[0]

            log.info(f"vector_store: {self.vector_store.identifier} currently doesn't have any records.")
            return 0

        except psycopg2.Error as e:
            raise RuntimeError(f"Failed to check if vector store has records in PGVector: {e}") from e

    async def create_gin_index(self, cur: cursor) -> None:
        """Create GIN index for full-text search performance

        Args:
            cur: PostgreSQL cursor

        Raises:
            RuntimeError: If the error occurred when creating GIN index
        """

        try:
            cur.execute(
                f"""
                CREATE INDEX IF NOT EXISTS {self.table_name}_content_gin_idx
                ON {self.table_name} USING GIN(tokenized_content)
            """
            )
            log.info(f"GIN index was successfully created for vector_store: {self.vector_store.identifier}.")

        except psycopg2.Error as e:
            raise RuntimeError(
                f"Failed to create GIN index for vector_store: {self.vector_store.identifier}: {e}"
            ) from e


class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtocolPrivate):
    def __init__(
        self, config: PGVectorVectorIOConfig, inference_api: Inference, files_api: Files | None = None
    ) -> None:
        super().__init__(inference_api=inference_api, files_api=files_api, kvstore=None)
        self.config = config
        self.conn = None
        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()


Review comment (suggestion):
I'm not too sure about this message. If the GIN index already exists, it will still log it was successfully created. It's a bit misleading. We have two options:
- We could make it generic like "GIN index verified for vector_store".
- OR check existence and assign to boolean, then log/execute based on that.

To avoid modifying production code, I think option 1 is fine. But, I'll leave it up to you :) 

Code after:
                return result[0]

            log.info(f"vector_store: {self.vector_store.identifier} currently doesn't have any records.")
            return 0

        except psycopg2.Error as e:
            raise RuntimeError(f"Failed to check if vector store has records in PGVector: {e}") from e

    async def create_gin_index(self, cur: cursor) -> None:
        """Create GIN index for full-text search performance

        Args:
            cur: PostgreSQL cursor

        Raises:
            RuntimeError: If the error occurred when creating GIN index
        """

        try:
            cur.execute(
                f"""
                CREATE INDEX IF NOT EXISTS {self.table_name}_content_gin_idx
                ON {self.table_name} USING GIN(tokenized_content)
            """
            )
            log.info(f"GIN index verified for vector_store: {self.vector_store.identifier}.")

        except psycopg2.Error as e:
            raise RuntimeError(
                f"Failed to create GIN index for vector_store: {self.vector_store.identifier}: {e}"
            ) from e


class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtocolPrivate):
    def __init__(
        self, config: PGVectorVectorIOConfig, inference_api: Inference, files_api: Files | None = None
    ) -> None:
        super().__init__(inference_api=inference_api, files_api=files_api, kvstore=None)
        self.config = config
        self.conn = None
        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()