mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-09 05:12:12 +02:00
Embeddings API scores (#671)
- Put scores in all responses - Remove unused 'middle' vector layer. Vector of texts -> vector of (vector embedding)
This commit is contained in:
cybermaggedon
GitHub
parent
4fa7cc7d7c
commit
f2ae0e8623
65 changed files with 1339 additions and 1292 deletions
|
|
@ -7,7 +7,7 @@ of chunk_ids
|
|||
import logging
|
||||
|
||||
from .... direct.milvus_doc_embeddings import DocVectors
|
||||
from .... schema import DocumentEmbeddingsResponse
|
||||
from .... schema import DocumentEmbeddingsResponse, ChunkMatch
|
||||
from .... schema import Error
|
||||
from .... base import DocumentEmbeddingsQueryService
|
||||
|
||||
|
|
@ -35,26 +35,33 @@ class Processor(DocumentEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Handle zero limit case
|
||||
if msg.limit <= 0:
|
||||
return []
|
||||
|
||||
chunk_ids = []
|
||||
resp = self.vecstore.search(
|
||||
vec,
|
||||
msg.user,
|
||||
msg.collection,
|
||||
limit=msg.limit
|
||||
)
|
||||
|
||||
for vec in msg.vectors:
|
||||
chunks = []
|
||||
for r in resp:
|
||||
chunk_id = r["entity"]["chunk_id"]
|
||||
# Milvus returns distance, convert to similarity score
|
||||
distance = r.get("distance", 0.0)
|
||||
score = 1.0 - distance if distance else 0.0
|
||||
chunks.append(ChunkMatch(
|
||||
chunk_id=chunk_id,
|
||||
score=score,
|
||||
))
|
||||
|
||||
resp = self.vecstore.search(
|
||||
vec,
|
||||
msg.user,
|
||||
msg.collection,
|
||||
limit=msg.limit
|
||||
)
|
||||
|
||||
for r in resp:
|
||||
chunk_id = r["entity"]["chunk_id"]
|
||||
chunk_ids.append(chunk_id)
|
||||
|
||||
return chunk_ids
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import os
|
|||
from pinecone import Pinecone, ServerlessSpec
|
||||
from pinecone.grpc import PineconeGRPC, GRPCClientConfig
|
||||
|
||||
from .... schema import ChunkMatch
|
||||
from .... base import DocumentEmbeddingsQueryService
|
||||
|
||||
# Module logger
|
||||
|
|
@ -51,38 +52,43 @@ class Processor(DocumentEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Handle zero limit case
|
||||
if msg.limit <= 0:
|
||||
return []
|
||||
|
||||
chunk_ids = []
|
||||
dim = len(vec)
|
||||
|
||||
for vec in msg.vectors:
|
||||
# Use dimension suffix in index name
|
||||
index_name = f"d-{msg.user}-{msg.collection}-{dim}"
|
||||
|
||||
dim = len(vec)
|
||||
# Check if index exists - return empty if not
|
||||
if not self.pinecone.has_index(index_name):
|
||||
logger.info(f"Index {index_name} does not exist")
|
||||
return []
|
||||
|
||||
# Use dimension suffix in index name
|
||||
index_name = f"d-{msg.user}-{msg.collection}-{dim}"
|
||||
index = self.pinecone.Index(index_name)
|
||||
|
||||
# Check if index exists - skip if not
|
||||
if not self.pinecone.has_index(index_name):
|
||||
logger.info(f"Index {index_name} does not exist, skipping this vector")
|
||||
continue
|
||||
results = index.query(
|
||||
vector=vec,
|
||||
top_k=msg.limit,
|
||||
include_values=False,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
index = self.pinecone.Index(index_name)
|
||||
chunks = []
|
||||
for r in results.matches:
|
||||
chunk_id = r.metadata["chunk_id"]
|
||||
score = r.score if hasattr(r, 'score') else 0.0
|
||||
chunks.append(ChunkMatch(
|
||||
chunk_id=chunk_id,
|
||||
score=score,
|
||||
))
|
||||
|
||||
results = index.query(
|
||||
vector=vec,
|
||||
top_k=msg.limit,
|
||||
include_values=False,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
for r in results.matches:
|
||||
chunk_id = r.metadata["chunk_id"]
|
||||
chunk_ids.append(chunk_id)
|
||||
|
||||
return chunk_ids
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from qdrant_client import QdrantClient
|
|||
from qdrant_client.models import PointStruct
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
|
||||
from .... schema import DocumentEmbeddingsResponse
|
||||
from .... schema import DocumentEmbeddingsResponse, ChunkMatch
|
||||
from .... schema import Error
|
||||
from .... base import DocumentEmbeddingsQueryService
|
||||
|
||||
|
|
@ -69,31 +69,36 @@ class Processor(DocumentEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
chunk_ids = []
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
for vec in msg.vectors:
|
||||
# Use dimension suffix in collection name
|
||||
dim = len(vec)
|
||||
collection = f"d_{msg.user}_{msg.collection}_{dim}"
|
||||
|
||||
# Use dimension suffix in collection name
|
||||
dim = len(vec)
|
||||
collection = f"d_{msg.user}_{msg.collection}_{dim}"
|
||||
# Check if collection exists - return empty if not
|
||||
if not self.collection_exists(collection):
|
||||
logger.info(f"Collection {collection} does not exist, returning empty results")
|
||||
return []
|
||||
|
||||
# Check if collection exists - return empty if not
|
||||
if not self.collection_exists(collection):
|
||||
logger.info(f"Collection {collection} does not exist, returning empty results")
|
||||
continue
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vec,
|
||||
limit=msg.limit,
|
||||
with_payload=True,
|
||||
).points
|
||||
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vec,
|
||||
limit=msg.limit,
|
||||
with_payload=True,
|
||||
).points
|
||||
chunks = []
|
||||
for r in search_result:
|
||||
chunk_id = r.payload["chunk_id"]
|
||||
score = r.score if hasattr(r, 'score') else 0.0
|
||||
chunks.append(ChunkMatch(
|
||||
chunk_id=chunk_id,
|
||||
score=score,
|
||||
))
|
||||
|
||||
for r in search_result:
|
||||
chunk_id = r.payload["chunk_id"]
|
||||
chunk_ids.append(chunk_id)
|
||||
|
||||
return chunk_ids
|
||||
return chunks
|
||||
|
||||
except Exception as e:
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ entities
|
|||
import logging
|
||||
|
||||
from .... direct.milvus_graph_embeddings import EntityVectors
|
||||
from .... schema import GraphEmbeddingsResponse
|
||||
from .... schema import GraphEmbeddingsResponse, EntityMatch
|
||||
from .... schema import Error, Term, IRI, LITERAL
|
||||
from .... base import GraphEmbeddingsQueryService
|
||||
|
||||
|
|
@ -41,42 +41,41 @@ class Processor(GraphEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
entity_set = set()
|
||||
entities = []
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Handle zero limit case
|
||||
if msg.limit <= 0:
|
||||
return []
|
||||
|
||||
for vec in msg.vectors:
|
||||
resp = self.vecstore.search(
|
||||
vec,
|
||||
msg.user,
|
||||
msg.collection,
|
||||
limit=msg.limit * 2
|
||||
)
|
||||
|
||||
resp = self.vecstore.search(
|
||||
vec,
|
||||
msg.user,
|
||||
msg.collection,
|
||||
limit=msg.limit * 2
|
||||
)
|
||||
entity_set = set()
|
||||
entities = []
|
||||
|
||||
for r in resp:
|
||||
ent = r["entity"]["entity"]
|
||||
|
||||
# De-dupe entities
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(ent)
|
||||
for r in resp:
|
||||
ent = r["entity"]["entity"]
|
||||
# Milvus returns distance, convert to similarity score
|
||||
distance = r.get("distance", 0.0)
|
||||
score = 1.0 - distance if distance else 0.0
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
# De-dupe entities, keep highest score
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(EntityMatch(
|
||||
entity=self.create_value(ent),
|
||||
score=score,
|
||||
))
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
|
||||
ents2 = []
|
||||
|
||||
for ent in entities:
|
||||
ents2.append(self.create_value(ent))
|
||||
|
||||
entities = ents2
|
||||
if len(entities) >= msg.limit:
|
||||
break
|
||||
|
||||
logger.debug("Send response...")
|
||||
return entities
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import os
|
|||
from pinecone import Pinecone, ServerlessSpec
|
||||
from pinecone.grpc import PineconeGRPC, GRPCClientConfig
|
||||
|
||||
from .... schema import GraphEmbeddingsResponse
|
||||
from .... schema import GraphEmbeddingsResponse, EntityMatch
|
||||
from .... schema import Error, Term, IRI, LITERAL
|
||||
from .... base import GraphEmbeddingsQueryService
|
||||
|
||||
|
|
@ -59,57 +59,53 @@ class Processor(GraphEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Handle zero limit case
|
||||
if msg.limit <= 0:
|
||||
return []
|
||||
|
||||
dim = len(vec)
|
||||
|
||||
# Use dimension suffix in index name
|
||||
index_name = f"t-{msg.user}-{msg.collection}-{dim}"
|
||||
|
||||
# Check if index exists - return empty if not
|
||||
if not self.pinecone.has_index(index_name):
|
||||
logger.info(f"Index {index_name} does not exist")
|
||||
return []
|
||||
|
||||
index = self.pinecone.Index(index_name)
|
||||
|
||||
# Heuristic hack, get (2*limit), so that we have more chance
|
||||
# of getting (limit) unique entities
|
||||
results = index.query(
|
||||
vector=vec,
|
||||
top_k=msg.limit * 2,
|
||||
include_values=False,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
entity_set = set()
|
||||
entities = []
|
||||
|
||||
for vec in msg.vectors:
|
||||
for r in results.matches:
|
||||
ent = r.metadata["entity"]
|
||||
score = r.score if hasattr(r, 'score') else 0.0
|
||||
|
||||
dim = len(vec)
|
||||
|
||||
# Use dimension suffix in index name
|
||||
index_name = f"t-{msg.user}-{msg.collection}-{dim}"
|
||||
|
||||
# Check if index exists - skip if not
|
||||
if not self.pinecone.has_index(index_name):
|
||||
logger.info(f"Index {index_name} does not exist, skipping this vector")
|
||||
continue
|
||||
|
||||
index = self.pinecone.Index(index_name)
|
||||
|
||||
# Heuristic hack, get (2*limit), so that we have more chance
|
||||
# of getting (limit) entities
|
||||
results = index.query(
|
||||
vector=vec,
|
||||
top_k=msg.limit * 2,
|
||||
include_values=False,
|
||||
include_metadata=True
|
||||
)
|
||||
|
||||
for r in results.matches:
|
||||
|
||||
ent = r.metadata["entity"]
|
||||
|
||||
# De-dupe entities
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(ent)
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
# De-dupe entities, keep highest score
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(EntityMatch(
|
||||
entity=self.create_value(ent),
|
||||
score=score,
|
||||
))
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
|
||||
ents2 = []
|
||||
|
||||
for ent in entities:
|
||||
ents2.append(self.create_value(ent))
|
||||
|
||||
entities = ents2
|
||||
if len(entities) >= msg.limit:
|
||||
break
|
||||
|
||||
return entities
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from qdrant_client import QdrantClient
|
|||
from qdrant_client.models import PointStruct
|
||||
from qdrant_client.models import Distance, VectorParams
|
||||
|
||||
from .... schema import GraphEmbeddingsResponse
|
||||
from .... schema import GraphEmbeddingsResponse, EntityMatch
|
||||
from .... schema import Error, Term, IRI, LITERAL
|
||||
from .... base import GraphEmbeddingsQueryService
|
||||
|
||||
|
|
@ -75,49 +75,46 @@ class Processor(GraphEmbeddingsQueryService):
|
|||
|
||||
try:
|
||||
|
||||
vec = msg.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Use dimension suffix in collection name
|
||||
dim = len(vec)
|
||||
collection = f"t_{msg.user}_{msg.collection}_{dim}"
|
||||
|
||||
# Check if collection exists - return empty if not
|
||||
if not self.collection_exists(collection):
|
||||
logger.info(f"Collection {collection} does not exist")
|
||||
return []
|
||||
|
||||
# Heuristic hack, get (2*limit), so that we have more chance
|
||||
# of getting (limit) unique entities
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vec,
|
||||
limit=msg.limit * 2,
|
||||
with_payload=True,
|
||||
).points
|
||||
|
||||
entity_set = set()
|
||||
entities = []
|
||||
|
||||
for vec in msg.vectors:
|
||||
for r in search_result:
|
||||
ent = r.payload["entity"]
|
||||
score = r.score if hasattr(r, 'score') else 0.0
|
||||
|
||||
# Use dimension suffix in collection name
|
||||
dim = len(vec)
|
||||
collection = f"t_{msg.user}_{msg.collection}_{dim}"
|
||||
|
||||
# Check if collection exists - return empty if not
|
||||
if not self.collection_exists(collection):
|
||||
logger.info(f"Collection {collection} does not exist, skipping this vector")
|
||||
continue
|
||||
|
||||
# Heuristic hack, get (2*limit), so that we have more chance
|
||||
# of getting (limit) entities
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=collection,
|
||||
query=vec,
|
||||
limit=msg.limit * 2,
|
||||
with_payload=True,
|
||||
).points
|
||||
|
||||
for r in search_result:
|
||||
ent = r.payload["entity"]
|
||||
|
||||
# De-dupe entities
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(ent)
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
# De-dupe entities, keep highest score
|
||||
if ent not in entity_set:
|
||||
entity_set.add(ent)
|
||||
entities.append(EntityMatch(
|
||||
entity=self.create_value(ent),
|
||||
score=score,
|
||||
))
|
||||
|
||||
# Keep adding entities until limit
|
||||
if len(entity_set) >= msg.limit: break
|
||||
|
||||
ents2 = []
|
||||
|
||||
for ent in entities:
|
||||
ents2.append(self.create_value(ent))
|
||||
|
||||
entities = ents2
|
||||
if len(entities) >= msg.limit:
|
||||
break
|
||||
|
||||
logger.debug("Send response...")
|
||||
return entities
|
||||
|
|
|
|||
|
|
@ -93,7 +93,9 @@ class Processor(FlowProcessor):
|
|||
async def query_row_embeddings(self, request: RowEmbeddingsRequest):
|
||||
"""Execute row embeddings query"""
|
||||
|
||||
matches = []
|
||||
vec = request.vector
|
||||
if not vec:
|
||||
return []
|
||||
|
||||
# Find the collection for this user/collection/schema
|
||||
qdrant_collection = self.find_collection(
|
||||
|
|
@ -105,47 +107,47 @@ class Processor(FlowProcessor):
|
|||
f"No Qdrant collection found for "
|
||||
f"{request.user}/{request.collection}/{request.schema_name}"
|
||||
)
|
||||
return []
|
||||
|
||||
try:
|
||||
# Build optional filter for index_name
|
||||
query_filter = None
|
||||
if request.index_name:
|
||||
query_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="index_name",
|
||||
match=MatchValue(value=request.index_name)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Query Qdrant
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=qdrant_collection,
|
||||
query=vec,
|
||||
limit=request.limit,
|
||||
with_payload=True,
|
||||
query_filter=query_filter,
|
||||
).points
|
||||
|
||||
# Convert to RowIndexMatch objects
|
||||
matches = []
|
||||
for point in search_result:
|
||||
payload = point.payload or {}
|
||||
match = RowIndexMatch(
|
||||
index_name=payload.get("index_name", ""),
|
||||
index_value=payload.get("index_value", []),
|
||||
text=payload.get("text", ""),
|
||||
score=point.score if hasattr(point, 'score') else 0.0
|
||||
)
|
||||
matches.append(match)
|
||||
|
||||
return matches
|
||||
|
||||
for vec in request.vectors:
|
||||
try:
|
||||
# Build optional filter for index_name
|
||||
query_filter = None
|
||||
if request.index_name:
|
||||
query_filter = Filter(
|
||||
must=[
|
||||
FieldCondition(
|
||||
key="index_name",
|
||||
match=MatchValue(value=request.index_name)
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Query Qdrant
|
||||
search_result = self.qdrant.query_points(
|
||||
collection_name=qdrant_collection,
|
||||
query=vec,
|
||||
limit=request.limit,
|
||||
with_payload=True,
|
||||
query_filter=query_filter,
|
||||
).points
|
||||
|
||||
# Convert to RowIndexMatch objects
|
||||
for point in search_result:
|
||||
payload = point.payload or {}
|
||||
match = RowIndexMatch(
|
||||
index_name=payload.get("index_name", ""),
|
||||
index_value=payload.get("index_value", []),
|
||||
text=payload.get("text", ""),
|
||||
score=point.score if hasattr(point, 'score') else 0.0
|
||||
)
|
||||
matches.append(match)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to query Qdrant: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
return matches
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to query Qdrant: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
async def on_message(self, msg, consumer, flow):
|
||||
"""Handle incoming query request"""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue