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

Code before:


def check_extension_version(cur):
    cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'vector'")
    result = cur.fetchone()
    return result[0] if result else None


def upsert_models(conn, keys_models: list[tuple[str, BaseModel]]):
    with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
        query = sql.SQL(
            """
            INSERT INTO metadata_store (key, data)
            VALUES %s
            ON CONFLICT (key) DO UPDATE
            SET data = EXCLUDED.data
        """
        )

        values = [(key, Json(model.model_dump())) for key, model in keys_models]
        execute_values(cur, query, values, template="(%s, %s)")


def remove_vector_store_metadata(conn: psycopg2.extensions.connection, vector_store_id: str) -> None:
    """
    Performs removal of vector store metadata from PGVector metadata_store table when unregister vector store

    Args:
        conn: active PostgreSQL connection
        vector_store_id: identifier of VectorStore resource
    """
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM metadata_store WHERE key = %s", (vector_store_id,))
            if cur.rowcount > 0:
                log.info(f"Removed metadata for vector store '{vector_store_id}' from PGVector metadata_store table.")
            else:
                log.debug(f"No metadata found for vector store '{vector_store_id}' in PGVector metadata_store table.")

    except Exception as e:
        raise RuntimeError(
            f"Error removing metadata from PGVector metadata_store for vector_store: {vector_store_id}"
        ) from e


def load_models(cur, cls):
    cur.execute("SELECT key, data FROM metadata_store")
    rows = cur.fetchall()
    return [TypeAdapter(cls).validate_python(row["data"]) for row in rows]



Review comment (suggestion):
nit for clarity:
```suggestion
    Performs removal of vector store metadata from PGVector metadata_store table when vector store is unregistered
```

Code after:


def check_extension_version(cur):
    cur.execute("SELECT extversion FROM pg_extension WHERE extname = 'vector'")
    result = cur.fetchone()
    return result[0] if result else None


def upsert_models(conn, keys_models: list[tuple[str, BaseModel]]):
    with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
        query = sql.SQL(
            """
            INSERT INTO metadata_store (key, data)
            VALUES %s
            ON CONFLICT (key) DO UPDATE
            SET data = EXCLUDED.data
        """
        )

        values = [(key, Json(model.model_dump())) for key, model in keys_models]
        execute_values(cur, query, values, template="(%s, %s)")


def remove_vector_store_metadata(conn: psycopg2.extensions.connection, vector_store_id: str) -> None:
    """
    Performs removal of vector store metadata from PGVector metadata_store table when vector store is unregistered

    Args:
        conn: active PostgreSQL connection
        vector_store_id: identifier of VectorStore resource
    """
    try:
        with conn.cursor() as cur:
            cur.execute("DELETE FROM metadata_store WHERE key = %s", (vector_store_id,))
            if cur.rowcount > 0:
                log.info(f"Removed metadata for vector store '{vector_store_id}' from PGVector metadata_store table.")

    except Exception as e:
        raise RuntimeError(
            f"Error removing metadata from PGVector metadata_store for vector_store: {vector_store_id}"
        ) from e


def load_models(cur, cls):
    cur.execute("SELECT key, data FROM metadata_store")
    rows = cur.fetchall()
    return [TypeAdapter(cls).validate_python(row["data"]) for row in rows]


class PGVectorIndex(EmbeddingIndex):
    # reference: https://github.com/pgvector/pgvector?tab=readme-ov-file#querying