mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-06-20 04:08:06 +02:00
594deba73e
changes. Workspace support: - Support for separate workspaces - Addition of workspace CLI support for test purposes - Massive test update - Remove many 'user' references in services - workspace now provides the same separation - Update API
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
from . request_response_spec import RequestResponse, RequestResponseSpec
|
|
from .. schema import RowEmbeddingsRequest, RowEmbeddingsResponse
|
|
|
|
class RowEmbeddingsQueryClient(RequestResponse):
|
|
async def row_embeddings_query(
|
|
self, vector, schema_name, collection="default",
|
|
index_name=None, limit=10, timeout=600
|
|
):
|
|
request = RowEmbeddingsRequest(
|
|
vector=vector,
|
|
schema_name=schema_name,
|
|
collection=collection,
|
|
limit=limit
|
|
)
|
|
if index_name:
|
|
request.index_name = index_name
|
|
|
|
resp = await self.request(request, timeout=timeout)
|
|
|
|
if resp.error:
|
|
raise RuntimeError(resp.error.message)
|
|
|
|
# Return matches as list of dicts
|
|
return [
|
|
{
|
|
"index_name": match.index_name,
|
|
"index_value": match.index_value,
|
|
"text": match.text,
|
|
"score": match.score
|
|
}
|
|
for match in (resp.matches or [])
|
|
]
|
|
|
|
class RowEmbeddingsQueryClientSpec(RequestResponseSpec):
|
|
def __init__(
|
|
self, request_name, response_name,
|
|
):
|
|
super(RowEmbeddingsQueryClientSpec, self).__init__(
|
|
request_name = request_name,
|
|
request_schema = RowEmbeddingsRequest,
|
|
response_name = response_name,
|
|
response_schema = RowEmbeddingsResponse,
|
|
impl = RowEmbeddingsQueryClient,
|
|
)
|