From 28ca00d0d94cfe51a872ccac2de9f4dac311f4f4 Mon Sep 17 00:00:00 2001
From: Ibrahim Haroon <99413953+Ibrahim-Haroon@users.noreply.github.com>
Date: Sat, 7 Jun 2025 16:31:30 -0400
Subject: [PATCH] fix(pgvector): handle case where distance is 0 by setting
 score to infinity (#2416)

# What does this PR do?
Fixes provider pgvector `query_vector` function for when the distance
between the query embedding and an embedding within the vector db is 0
(identical vectors). Catches `ZeroDivisionError` and then sets `score`
to infinity, which represent maximum similarity.

<!-- If resolving an issue, uncomment and update the line below -->
Closes [#2381]

## Test Plan
Checkout this PR

Execute this code and there will no longer be a `ZeroDivisionError`
exception
```
from llama_stack_client import LlamaStackClient

base_url = "http://localhost:8321"
client = LlamaStackClient(base_url=base_url)

models = client.models.list()
embedding_model = (
    em := next(m for m in models if m.model_type == "embedding")
).identifier
embedding_dimension = 384

_ = client.vector_dbs.register(
    vector_db_id="foo_db",
    embedding_model=embedding_model,
    embedding_dimension=embedding_dimension,
    provider_id="pgvector",
)

chunk = {
    "content": "foo",
    "mime_type": "text/plain",
    "metadata": {
        "document_id": "foo-id"
    }
}

client.vector_io.insert(vector_db_id="foo_db", chunks=[chunk])
client.vector_io.query(vector_db_id="foo_db", query="foo")
```
---
 llama_stack/providers/remote/vector_io/pgvector/pgvector.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py
index ea918c552..7d58a49f3 100644
--- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py
+++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py
@@ -116,7 +116,7 @@ class PGVectorIndex(EmbeddingIndex):
             scores = []
             for doc, dist in results:
                 chunks.append(Chunk(**doc))
-                scores.append(1.0 / float(dist))
+                scores.append(1.0 / float(dist) if dist != 0 else float("inf"))
 
             return QueryChunksResponse(chunks=chunks, scores=scores)