From a400c80008c0492c5f979e8a583224b7b3eeba90 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 23 Feb 2026 16:36:25 +0000 Subject: [PATCH] Added row embeddings API and CLI support --- docs/tech-specs/structured-data-2.md | 146 ++++++++++++++++++ trustgraph-base/trustgraph/api/async_flow.py | 60 +++++++ .../trustgraph/api/async_socket_client.py | 23 +++ trustgraph-base/trustgraph/api/flow.py | 75 +++++++++ .../trustgraph/api/socket_client.py | 70 +++++++++ .../clients/row_embeddings_client.py | 60 +++++++ .../trustgraph/messaging/__init__.py | 13 +- .../messaging/translators/__init__.py | 3 +- .../messaging/translators/embeddings_query.py | 62 +++++++- trustgraph-cli/pyproject.toml | 1 + .../trustgraph/cli/invoke_row_embeddings.py | 126 +++++++++++++++ .../trustgraph/gateway/dispatch/manager.py | 2 + .../gateway/dispatch/row_embeddings_query.py | 31 ++++ 13 files changed, 667 insertions(+), 5 deletions(-) create mode 100644 trustgraph-base/trustgraph/clients/row_embeddings_client.py create mode 100644 trustgraph-cli/trustgraph/cli/invoke_row_embeddings.py create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/row_embeddings_query.py diff --git a/docs/tech-specs/structured-data-2.md b/docs/tech-specs/structured-data-2.md index 10bbdae9..1b70a6c3 100644 --- a/docs/tech-specs/structured-data-2.md +++ b/docs/tech-specs/structured-data-2.md @@ -368,8 +368,154 @@ This separation keeps concerns clean: - Embeddings API handles semantic similarity - User workflow: fuzzy search via embeddings to find candidates, then exact query to get full row data +#### Request/Response Schema + +```python +@dataclass +class RowEmbeddingsRequest: + vectors: list[list[float]] # Query vectors (pre-computed embeddings) + user: str = "" + collection: str = "" + schema_name: str = "" + index_name: str = "" # Optional: filter to specific index + limit: int = 10 # Max results per vector + +@dataclass +class RowIndexMatch: + index_name: str = "" # The matched index field(s) + index_value: list[str] = [] # The matched value(s) + text: str = "" # Original text that was embedded + score: float = 0.0 # Similarity score + +@dataclass +class RowEmbeddingsResponse: + error: Error | None = None + matches: list[RowIndexMatch] = [] +``` + +#### Query Processor + Module: `trustgraph-flow/trustgraph/query/row_embeddings/qdrant` +Entry point: `row-embeddings-query-qdrant` + +The processor: +1. Receives `RowEmbeddingsRequest` with query vectors +2. Finds the appropriate Qdrant collection by prefix matching +3. Searches for nearest vectors with optional `index_name` filter +4. Returns `RowEmbeddingsResponse` with matching index information + +#### API Gateway Integration + +The gateway exposes row embeddings queries via the standard request/response pattern: + +| Component | Location | +|-----------|----------| +| Dispatcher | `trustgraph-flow/trustgraph/gateway/dispatch/row_embeddings_query.py` | +| Registration | Add `"row-embeddings"` to `request_response_dispatchers` in `manager.py` | + +Flow interface name: `row-embeddings` + +Interface definition in flow blueprint: +```json +{ + "interfaces": { + "row-embeddings": { + "request": "non-persistent://tg/request/row-embeddings:{id}", + "response": "non-persistent://tg/response/row-embeddings:{id}" + } + } +} +``` + +#### Python SDK Support + +The SDK provides methods for row embeddings queries: + +```python +# Flow-scoped query (preferred) +api = Api(url) +flow = api.flow().id("default") + +# Query with text (SDK computes embeddings) +matches = flow.row_embeddings_query( + text="Chestnut Street", + collection="my_collection", + schema_name="addresses", + index_name="street_name", # Optional filter + limit=10 +) + +# Query with pre-computed vectors +matches = flow.row_embeddings_query( + vectors=[[0.1, 0.2, ...]], + collection="my_collection", + schema_name="addresses" +) + +# Each match contains: +for match in matches: + print(match.index_name) # e.g., "street_name" + print(match.index_value) # e.g., ["CHESTNUT ST"] + print(match.text) # e.g., "CHESTNUT ST" + print(match.score) # e.g., 0.95 +``` + +#### CLI Utility + +Command: `tg-invoke-row-embeddings` + +```bash +# Query by text (computes embedding automatically) +tg-invoke-row-embeddings \ + --text "Chestnut Street" \ + --collection my_collection \ + --schema addresses \ + --index street_name \ + --limit 10 + +# Query by vector file +tg-invoke-row-embeddings \ + --vectors vectors.json \ + --collection my_collection \ + --schema addresses + +# Output formats +tg-invoke-row-embeddings --text "..." --format json +tg-invoke-row-embeddings --text "..." --format table +``` + +#### Typical Usage Pattern + +The row embeddings query is typically used as part of a fuzzy-to-exact lookup flow: + +```python +# Step 1: Fuzzy search via embeddings +matches = flow.row_embeddings_query( + text="chestnut street", + collection="geo", + schema_name="streets" +) + +# Step 2: Exact lookup via GraphQL for full row data +for match in matches: + query = f''' + query {{ + streets(where: {{ {match.index_name}: {{ eq: "{match.index_value[0]}" }} }}) {{ + street_name + city + zip_code + }} + }} + ''' + rows = flow.rows_query(query, collection="geo") +``` + +This two-step pattern enables: +- Finding "CHESTNUT ST" when user searches for "Chestnut Street" +- Retrieving complete row data with all fields +- Combining semantic similarity with structured data access + ### Row Data Ingestion Deferred to a subsequent phase. Will be designed alongside other ingestion changes. diff --git a/trustgraph-base/trustgraph/api/async_flow.py b/trustgraph-base/trustgraph/api/async_flow.py index 2cf1bedf..440cebae 100644 --- a/trustgraph-base/trustgraph/api/async_flow.py +++ b/trustgraph-base/trustgraph/api/async_flow.py @@ -766,3 +766,63 @@ class AsyncFlowInstance: request_data.update(kwargs) return await self.request("rows", request_data) + + async def row_embeddings_query( + self, text: str, schema_name: str, user: str = "trustgraph", + collection: str = "default", index_name: Optional[str] = None, + limit: int = 10, **kwargs: Any + ): + """ + Query row embeddings for semantic search on structured data. + + Performs semantic search over row index embeddings to find rows whose + indexed field values are most similar to the input text. Enables + fuzzy/semantic matching on structured data. + + Args: + text: Query text for semantic search + schema_name: Schema name to search within + user: User identifier (default: "trustgraph") + collection: Collection identifier (default: "default") + index_name: Optional index name to filter search to specific index + limit: Maximum number of results to return (default: 10) + **kwargs: Additional service-specific parameters + + Returns: + dict: Response containing matches with index_name, index_value, + text, and score + + Example: + ```python + async_flow = await api.async_flow() + flow = async_flow.id("default") + + # Search for customers by name similarity + results = await flow.row_embeddings_query( + text="John Smith", + schema_name="customers", + user="trustgraph", + collection="sales", + limit=5 + ) + + for match in results.get("matches", []): + print(f"{match['index_name']}: {match['index_value']} (score: {match['score']})") + ``` + """ + # First convert text to embeddings vectors + emb_result = await self.embeddings(text=text) + vectors = emb_result.get("vectors", []) + + request_data = { + "vectors": vectors, + "schema_name": schema_name, + "user": user, + "collection": collection, + "limit": limit + } + if index_name: + request_data["index_name"] = index_name + request_data.update(kwargs) + + return await self.request("row-embeddings", request_data) diff --git a/trustgraph-base/trustgraph/api/async_socket_client.py b/trustgraph-base/trustgraph/api/async_socket_client.py index ac83876b..3241e0f7 100644 --- a/trustgraph-base/trustgraph/api/async_socket_client.py +++ b/trustgraph-base/trustgraph/api/async_socket_client.py @@ -345,3 +345,26 @@ class AsyncSocketFlowInstance: request.update(kwargs) return await self.client._send_request("mcp-tool", self.flow_id, request) + + async def row_embeddings_query( + self, text: str, schema_name: str, user: str = "trustgraph", + collection: str = "default", index_name: Optional[str] = None, + limit: int = 10, **kwargs + ): + """Query row embeddings for semantic search on structured data""" + # First convert text to embeddings vectors + emb_result = await self.embeddings(text=text) + vectors = emb_result.get("vectors", []) + + request = { + "vectors": vectors, + "schema_name": schema_name, + "user": user, + "collection": collection, + "limit": limit + } + if index_name: + request["index_name"] = index_name + request.update(kwargs) + + return await self.client._send_request("row-embeddings", self.flow_id, request) diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index e8da1522..cc07f794 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/trustgraph-base/trustgraph/api/flow.py @@ -1297,3 +1297,78 @@ class FlowInstance: return response["schema-matches"] + def row_embeddings_query( + self, text, schema_name, user="trustgraph", collection="default", + index_name=None, limit=10 + ): + """ + Query row data using semantic similarity on indexed fields. + + Finds rows whose indexed field values are semantically similar to the + input text, using vector embeddings. This enables fuzzy/semantic matching + on structured data. + + Args: + text: Query text for semantic search + schema_name: Schema name to search within + user: User/keyspace identifier (default: "trustgraph") + collection: Collection identifier (default: "default") + index_name: Optional index name to filter search to specific index + limit: Maximum number of results (default: 10) + + Returns: + dict: Query results with matches containing index_name, index_value, + text, and score + + Example: + ```python + flow = api.flow().id("default") + + # Search for customers by name similarity + results = flow.row_embeddings_query( + text="John Smith", + schema_name="customers", + user="trustgraph", + collection="sales", + limit=5 + ) + + # Filter to specific index + results = flow.row_embeddings_query( + text="machine learning engineer", + schema_name="employees", + index_name="job_title", + limit=10 + ) + ``` + """ + + # First convert text to embeddings vectors + emb_result = self.embeddings(text=text) + vectors = emb_result.get("vectors", []) + + # Query row embeddings for semantic search + input = { + "vectors": vectors, + "schema_name": schema_name, + "user": user, + "collection": collection, + "limit": limit + } + + if index_name: + input["index_name"] = index_name + + response = self.request( + "service/row-embeddings", + input + ) + + # Check for system-level error + if "error" in response and response["error"]: + error_type = response["error"].get("type", "unknown") + error_message = response["error"].get("message", "Unknown error") + raise ProtocolException(f"{error_type}: {error_message}") + + return response + diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index c0246612..e8de442a 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -881,3 +881,73 @@ class SocketFlowInstance: request.update(kwargs) return self.client._send_request_sync("mcp-tool", self.flow_id, request, False) + + def row_embeddings_query( + self, + text: str, + schema_name: str, + user: str = "trustgraph", + collection: str = "default", + index_name: Optional[str] = None, + limit: int = 10, + **kwargs: Any + ) -> Dict[str, Any]: + """ + Query row data using semantic similarity on indexed fields. + + Finds rows whose indexed field values are semantically similar to the + input text, using vector embeddings. This enables fuzzy/semantic matching + on structured data. + + Args: + text: Query text for semantic search + schema_name: Schema name to search within + user: User/keyspace identifier (default: "trustgraph") + collection: Collection identifier (default: "default") + index_name: Optional index name to filter search to specific index + limit: Maximum number of results (default: 10) + **kwargs: Additional parameters passed to the service + + Returns: + dict: Query results with matches containing index_name, index_value, + text, and score + + Example: + ```python + socket = api.socket() + flow = socket.flow("default") + + # Search for customers by name similarity + results = flow.row_embeddings_query( + text="John Smith", + schema_name="customers", + user="trustgraph", + collection="sales", + limit=5 + ) + + # Filter to specific index + results = flow.row_embeddings_query( + text="machine learning engineer", + schema_name="employees", + index_name="job_title", + limit=10 + ) + ``` + """ + # First convert text to embeddings vectors + emb_result = self.embeddings(text=text) + vectors = emb_result.get("vectors", []) + + request = { + "vectors": vectors, + "schema_name": schema_name, + "user": user, + "collection": collection, + "limit": limit + } + if index_name: + request["index_name"] = index_name + request.update(kwargs) + + return self.client._send_request_sync("row-embeddings", self.flow_id, request, False) diff --git a/trustgraph-base/trustgraph/clients/row_embeddings_client.py b/trustgraph-base/trustgraph/clients/row_embeddings_client.py new file mode 100644 index 00000000..4f911e3c --- /dev/null +++ b/trustgraph-base/trustgraph/clients/row_embeddings_client.py @@ -0,0 +1,60 @@ + +import _pulsar + +from .. schema import RowEmbeddingsRequest, RowEmbeddingsResponse +from .. schema import row_embeddings_request_queue +from .. schema import row_embeddings_response_queue +from . base import BaseClient + +# Ugly +ERROR=_pulsar.LoggerLevel.Error +WARN=_pulsar.LoggerLevel.Warn +INFO=_pulsar.LoggerLevel.Info +DEBUG=_pulsar.LoggerLevel.Debug + +class RowEmbeddingsClient(BaseClient): + + def __init__( + self, log_level=ERROR, + subscriber=None, + input_queue=None, + output_queue=None, + pulsar_host="pulsar://pulsar:6650", + pulsar_api_key=None, + ): + + if input_queue == None: + input_queue = row_embeddings_request_queue + + if output_queue == None: + output_queue = row_embeddings_response_queue + + super(RowEmbeddingsClient, self).__init__( + log_level=log_level, + subscriber=subscriber, + input_queue=input_queue, + output_queue=output_queue, + pulsar_host=pulsar_host, + pulsar_api_key=pulsar_api_key, + input_schema=RowEmbeddingsRequest, + output_schema=RowEmbeddingsResponse, + ) + + def request( + self, vectors, schema_name, user="trustgraph", collection="default", + index_name=None, limit=10, timeout=300 + ): + kwargs = dict( + user=user, collection=collection, + vectors=vectors, schema_name=schema_name, + limit=limit, timeout=timeout + ) + if index_name: + kwargs["index_name"] = index_name + + response = self.call(**kwargs) + + if response.error: + raise RuntimeError(f"{response.error.type}: {response.error.message}") + + return response.matches diff --git a/trustgraph-base/trustgraph/messaging/__init__.py b/trustgraph-base/trustgraph/messaging/__init__.py index 4d4e3c84..9fbcbf16 100644 --- a/trustgraph-base/trustgraph/messaging/__init__.py +++ b/trustgraph-base/trustgraph/messaging/__init__.py @@ -19,7 +19,8 @@ from .translators.prompt import PromptRequestTranslator, PromptResponseTranslato from .translators.tool import ToolRequestTranslator, ToolResponseTranslator from .translators.embeddings_query import ( DocumentEmbeddingsRequestTranslator, DocumentEmbeddingsResponseTranslator, - GraphEmbeddingsRequestTranslator, GraphEmbeddingsResponseTranslator + GraphEmbeddingsRequestTranslator, GraphEmbeddingsResponseTranslator, + RowEmbeddingsRequestTranslator, RowEmbeddingsResponseTranslator ) from .translators.rows_query import RowsQueryRequestTranslator, RowsQueryResponseTranslator from .translators.nlp_query import QuestionToStructuredQueryRequestTranslator, QuestionToStructuredQueryResponseTranslator @@ -107,11 +108,17 @@ TranslatorRegistry.register_service( ) TranslatorRegistry.register_service( - "graph-embeddings-query", - GraphEmbeddingsRequestTranslator(), + "graph-embeddings-query", + GraphEmbeddingsRequestTranslator(), GraphEmbeddingsResponseTranslator() ) +TranslatorRegistry.register_service( + "row-embeddings-query", + RowEmbeddingsRequestTranslator(), + RowEmbeddingsResponseTranslator() +) + TranslatorRegistry.register_service( "rows-query", RowsQueryRequestTranslator(), diff --git a/trustgraph-base/trustgraph/messaging/translators/__init__.py b/trustgraph-base/trustgraph/messaging/translators/__init__.py index 189265e1..5b5820fa 100644 --- a/trustgraph-base/trustgraph/messaging/translators/__init__.py +++ b/trustgraph-base/trustgraph/messaging/translators/__init__.py @@ -15,7 +15,8 @@ from .flow import FlowRequestTranslator, FlowResponseTranslator from .prompt import PromptRequestTranslator, PromptResponseTranslator from .embeddings_query import ( DocumentEmbeddingsRequestTranslator, DocumentEmbeddingsResponseTranslator, - GraphEmbeddingsRequestTranslator, GraphEmbeddingsResponseTranslator + GraphEmbeddingsRequestTranslator, GraphEmbeddingsResponseTranslator, + RowEmbeddingsRequestTranslator, RowEmbeddingsResponseTranslator ) from .rows_query import RowsQueryRequestTranslator, RowsQueryResponseTranslator from .diagnosis import StructuredDataDiagnosisRequestTranslator, StructuredDataDiagnosisResponseTranslator diff --git a/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py b/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py index a08f9b6c..141a7330 100644 --- a/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py +++ b/trustgraph-base/trustgraph/messaging/translators/embeddings_query.py @@ -1,7 +1,8 @@ from typing import Dict, Any, Tuple from ...schema import ( DocumentEmbeddingsRequest, DocumentEmbeddingsResponse, - GraphEmbeddingsRequest, GraphEmbeddingsResponse + GraphEmbeddingsRequest, GraphEmbeddingsResponse, + RowEmbeddingsRequest, RowEmbeddingsResponse, RowIndexMatch ) from .base import MessageTranslator from .primitives import ValueTranslator @@ -92,3 +93,62 @@ class GraphEmbeddingsResponseTranslator(MessageTranslator): def from_response_with_completion(self, obj: GraphEmbeddingsResponse) -> Tuple[Dict[str, Any], bool]: """Returns (response_dict, is_final)""" return self.from_pulsar(obj), True + + +class RowEmbeddingsRequestTranslator(MessageTranslator): + """Translator for RowEmbeddingsRequest schema objects""" + + def to_pulsar(self, data: Dict[str, Any]) -> RowEmbeddingsRequest: + return RowEmbeddingsRequest( + vectors=data["vectors"], + limit=int(data.get("limit", 10)), + user=data.get("user", "trustgraph"), + collection=data.get("collection", "default"), + schema_name=data.get("schema_name", ""), + index_name=data.get("index_name") + ) + + def from_pulsar(self, obj: RowEmbeddingsRequest) -> Dict[str, Any]: + result = { + "vectors": obj.vectors, + "limit": obj.limit, + "user": obj.user, + "collection": obj.collection, + "schema_name": obj.schema_name, + } + if obj.index_name: + result["index_name"] = obj.index_name + return result + + +class RowEmbeddingsResponseTranslator(MessageTranslator): + """Translator for RowEmbeddingsResponse schema objects""" + + def to_pulsar(self, data: Dict[str, Any]) -> RowEmbeddingsResponse: + raise NotImplementedError("Response translation to Pulsar not typically needed") + + def from_pulsar(self, obj: RowEmbeddingsResponse) -> Dict[str, Any]: + result = {} + + if obj.error is not None: + result["error"] = { + "type": obj.error.type, + "message": obj.error.message + } + + if obj.matches is not None: + result["matches"] = [ + { + "index_name": match.index_name, + "index_value": match.index_value, + "text": match.text, + "score": match.score + } + for match in obj.matches + ] + + return result + + def from_response_with_completion(self, obj: RowEmbeddingsResponse) -> Tuple[Dict[str, Any], bool]: + """Returns (response_dict, is_final)""" + return self.from_pulsar(obj), True diff --git a/trustgraph-cli/pyproject.toml b/trustgraph-cli/pyproject.toml index 49e24b8e..66df74f1 100644 --- a/trustgraph-cli/pyproject.toml +++ b/trustgraph-cli/pyproject.toml @@ -49,6 +49,7 @@ tg-invoke-document-embeddings = "trustgraph.cli.invoke_document_embeddings:main" tg-invoke-mcp-tool = "trustgraph.cli.invoke_mcp_tool:main" tg-invoke-nlp-query = "trustgraph.cli.invoke_nlp_query:main" tg-invoke-rows-query = "trustgraph.cli.invoke_rows_query:main" +tg-invoke-row-embeddings = "trustgraph.cli.invoke_row_embeddings:main" tg-invoke-prompt = "trustgraph.cli.invoke_prompt:main" tg-invoke-structured-query = "trustgraph.cli.invoke_structured_query:main" tg-load-doc-embeds = "trustgraph.cli.load_doc_embeds:main" diff --git a/trustgraph-cli/trustgraph/cli/invoke_row_embeddings.py b/trustgraph-cli/trustgraph/cli/invoke_row_embeddings.py new file mode 100644 index 00000000..7393b4c3 --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/invoke_row_embeddings.py @@ -0,0 +1,126 @@ +""" +Queries row data by text similarity using vector embeddings on indexed fields. +Returns matching rows with their index values and similarity scores. +""" + +import argparse +import os +from trustgraph.api import Api + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) + +def query(url, flow_id, query_text, schema_name, user, collection, index_name, limit, token=None): + + # Create API client + api = Api(url=url, token=token) + socket = api.socket() + flow = socket.flow(flow_id) + + try: + # Call row embeddings query service + result = flow.row_embeddings_query( + text=query_text, + schema_name=schema_name, + user=user, + collection=collection, + index_name=index_name, + limit=limit + ) + + matches = result.get("matches", []) + for match in matches: + print(f"Index: {match['index_name']}") + print(f" Values: {match['index_value']}") + print(f" Text: {match['text']}") + print(f" Score: {match['score']:.4f}") + print() + + finally: + # Clean up socket connection + socket.close() + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-invoke-row-embeddings', + description=__doc__, + ) + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + + parser.add_argument( + '-f', '--flow-id', + default="default", + help=f'Flow ID (default: default)' + ) + + parser.add_argument( + '-U', '--user', + default="trustgraph", + help='User/keyspace (default: trustgraph)', + ) + + parser.add_argument( + '-c', '--collection', + default="default", + help='Collection (default: default)', + ) + + parser.add_argument( + '-s', '--schema-name', + required=True, + help='Schema name to search within (required)', + ) + + parser.add_argument( + '-i', '--index-name', + default=None, + help='Index name to filter search (optional)', + ) + + parser.add_argument( + '-l', '--limit', + type=int, + default=10, + help='Maximum number of results (default: 10)', + ) + + parser.add_argument( + 'query', + nargs=1, + help='Query text to search for similar row index values', + ) + + args = parser.parse_args() + + try: + + query( + url=args.url, + flow_id=args.flow_id, + query_text=args.query[0], + schema_name=args.schema_name, + user=args.user, + collection=args.collection, + index_name=args.index_name, + limit=args.limit, + token=args.token, + ) + + except Exception as e: + + print("Exception:", e, flush=True) + +if __name__ == "__main__": + main() diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index d7d04f83..35edad76 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -27,6 +27,7 @@ from . structured_diag import StructuredDiagRequestor from . embeddings import EmbeddingsRequestor from . graph_embeddings_query import GraphEmbeddingsQueryRequestor from . document_embeddings_query import DocumentEmbeddingsQueryRequestor +from . row_embeddings_query import RowEmbeddingsQueryRequestor from . mcp_tool import McpToolRequestor from . text_load import TextLoad from . document_load import DocumentLoad @@ -62,6 +63,7 @@ request_response_dispatchers = { "nlp-query": NLPQueryRequestor, "structured-query": StructuredQueryRequestor, "structured-diag": StructuredDiagRequestor, + "row-embeddings": RowEmbeddingsQueryRequestor, } global_dispatchers = { diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/row_embeddings_query.py b/trustgraph-flow/trustgraph/gateway/dispatch/row_embeddings_query.py new file mode 100644 index 00000000..8b139fc2 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/row_embeddings_query.py @@ -0,0 +1,31 @@ + +from ... schema import RowEmbeddingsRequest, RowEmbeddingsResponse +from ... messaging import TranslatorRegistry + +from . requestor import ServiceRequestor + +class RowEmbeddingsQueryRequestor(ServiceRequestor): + def __init__( + self, backend, request_queue, response_queue, timeout, + consumer, subscriber, + ): + + super(RowEmbeddingsQueryRequestor, self).__init__( + backend=backend, + request_queue=request_queue, + response_queue=response_queue, + request_schema=RowEmbeddingsRequest, + response_schema=RowEmbeddingsResponse, + subscription = subscriber, + consumer_name = consumer, + timeout=timeout, + ) + + self.request_translator = TranslatorRegistry.get_request_translator("row-embeddings-query") + self.response_translator = TranslatorRegistry.get_response_translator("row-embeddings-query") + + def to_request(self, body): + return self.request_translator.to_pulsar(body) + + def from_response(self, message): + return self.response_translator.from_response_with_completion(message)