mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-26 00:46:22 +02:00
f22bf13aa6
* Collection+user fields in structured query * User/collection in structured query & agent
35 lines
No EOL
1.2 KiB
Python
35 lines
No EOL
1.2 KiB
Python
from . request_response_spec import RequestResponse, RequestResponseSpec
|
|
from .. schema import StructuredQueryRequest, StructuredQueryResponse
|
|
|
|
class StructuredQueryClient(RequestResponse):
|
|
async def structured_query(self, question, user="trustgraph", collection="default", timeout=600):
|
|
resp = await self.request(
|
|
StructuredQueryRequest(
|
|
question = question,
|
|
user = user,
|
|
collection = collection
|
|
),
|
|
timeout=timeout
|
|
)
|
|
|
|
if resp.error:
|
|
raise RuntimeError(resp.error.message)
|
|
|
|
# Return the full response structure for the tool to handle
|
|
return {
|
|
"data": resp.data,
|
|
"errors": resp.errors if resp.errors else [],
|
|
"error": resp.error
|
|
}
|
|
|
|
class StructuredQueryClientSpec(RequestResponseSpec):
|
|
def __init__(
|
|
self, request_name, response_name,
|
|
):
|
|
super(StructuredQueryClientSpec, self).__init__(
|
|
request_name = request_name,
|
|
request_schema = StructuredQueryRequest,
|
|
response_name = response_name,
|
|
response_schema = StructuredQueryResponse,
|
|
impl = StructuredQueryClient,
|
|
) |