Structured query support (#492)

* Tweak the structured query schema

* Structure query service

* Gateway support for nlp-query and structured-query

* API support

* Added CLI

* Update tests

* More tests
This commit is contained in:
cybermaggedon 2025-09-04 16:06:18 +01:00 committed by GitHub
parent 8d4aa0069c
commit a6d9f5e849
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 2813 additions and 31 deletions

View file

@ -426,3 +426,62 @@ class FlowInstance:
return result
def nlp_query(self, question, max_results=100):
"""
Convert a natural language question to a GraphQL query.
Args:
question: Natural language question
max_results: Maximum number of results to return (default: 100)
Returns:
dict with graphql_query, variables, detected_schemas, confidence
"""
input = {
"question": question,
"max_results": max_results
}
response = self.request(
"service/nlp-query",
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
def structured_query(self, question):
"""
Execute a natural language question against structured data.
Combines NLP query conversion and GraphQL execution.
Args:
question: Natural language question
Returns:
dict with data and optional errors
"""
input = {
"question": question
}
response = self.request(
"service/structured-query",
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

View file

@ -22,6 +22,8 @@ from .translators.embeddings_query import (
GraphEmbeddingsRequestTranslator, GraphEmbeddingsResponseTranslator
)
from .translators.objects_query import ObjectsQueryRequestTranslator, ObjectsQueryResponseTranslator
from .translators.nlp_query import QuestionToStructuredQueryRequestTranslator, QuestionToStructuredQueryResponseTranslator
from .translators.structured_query import StructuredQueryRequestTranslator, StructuredQueryResponseTranslator
# Register all service translators
TranslatorRegistry.register_service(
@ -114,6 +116,18 @@ TranslatorRegistry.register_service(
ObjectsQueryResponseTranslator()
)
TranslatorRegistry.register_service(
"nlp-query",
QuestionToStructuredQueryRequestTranslator(),
QuestionToStructuredQueryResponseTranslator()
)
TranslatorRegistry.register_service(
"structured-query",
StructuredQueryRequestTranslator(),
StructuredQueryResponseTranslator()
)
# Register single-direction translators for document loading
TranslatorRegistry.register_request("document", DocumentTranslator())
TranslatorRegistry.register_request("text-document", TextDocumentTranslator())

View file

@ -0,0 +1,47 @@
from typing import Dict, Any, Tuple
from ...schema import QuestionToStructuredQueryRequest, QuestionToStructuredQueryResponse
from .base import MessageTranslator
class QuestionToStructuredQueryRequestTranslator(MessageTranslator):
"""Translator for QuestionToStructuredQueryRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> QuestionToStructuredQueryRequest:
return QuestionToStructuredQueryRequest(
question=data.get("question", ""),
max_results=data.get("max_results", 100)
)
def from_pulsar(self, obj: QuestionToStructuredQueryRequest) -> Dict[str, Any]:
return {
"question": obj.question,
"max_results": obj.max_results
}
class QuestionToStructuredQueryResponseTranslator(MessageTranslator):
"""Translator for QuestionToStructuredQueryResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> QuestionToStructuredQueryResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: QuestionToStructuredQueryResponse) -> Dict[str, Any]:
result = {
"graphql_query": obj.graphql_query,
"variables": dict(obj.variables) if obj.variables else {},
"detected_schemas": list(obj.detected_schemas) if obj.detected_schemas else [],
"confidence": obj.confidence
}
# Handle system-level error
if obj.error:
result["error"] = {
"type": obj.error.type,
"message": obj.error.message
}
return result
def from_response_with_completion(self, obj: QuestionToStructuredQueryResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -0,0 +1,56 @@
from typing import Dict, Any, Tuple
from ...schema import StructuredQueryRequest, StructuredQueryResponse
from .base import MessageTranslator
import json
class StructuredQueryRequestTranslator(MessageTranslator):
"""Translator for StructuredQueryRequest schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> StructuredQueryRequest:
return StructuredQueryRequest(
question=data.get("question", "")
)
def from_pulsar(self, obj: StructuredQueryRequest) -> Dict[str, Any]:
return {
"question": obj.question
}
class StructuredQueryResponseTranslator(MessageTranslator):
"""Translator for StructuredQueryResponse schema objects"""
def to_pulsar(self, data: Dict[str, Any]) -> StructuredQueryResponse:
raise NotImplementedError("Response translation to Pulsar not typically needed")
def from_pulsar(self, obj: StructuredQueryResponse) -> Dict[str, Any]:
result = {}
# Handle structured query response data
if obj.data:
try:
result["data"] = json.loads(obj.data)
except json.JSONDecodeError:
result["data"] = obj.data
else:
result["data"] = None
# Handle errors (array of strings)
if obj.errors:
result["errors"] = list(obj.errors)
else:
result["errors"] = []
# Handle system-level error
if obj.error:
result["error"] = {
"type": obj.error.type,
"message": obj.error.message
}
return result
def from_response_with_completion(self, obj: StructuredQueryResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)"""
return self.from_pulsar(obj), True

View file

@ -8,13 +8,11 @@ from ..core.topic import topic
# Structured Query Service - executes GraphQL queries
class StructuredQueryRequest(Record):
query = String() # GraphQL query
variables = Map(String()) # GraphQL variables
operation_name = String() # Optional operation name for multi-operation documents
question = String()
class StructuredQueryResponse(Record):
error = Error()
data = String() # JSON-encoded GraphQL response data
errors = Array(String()) # GraphQL errors if any
############################################################################
############################################################################