mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 00:16:23 +02:00
903 lines
29 KiB
Python
903 lines
29 KiB
Python
"""
|
|
TrustGraph Synchronous WebSocket Client
|
|
|
|
This module provides synchronous WebSocket-based access to TrustGraph services with
|
|
streaming support for real-time responses from agents, RAG queries, and text completions.
|
|
|
|
Uses a persistent WebSocket connection with a background reader task that
|
|
multiplexes requests by ID.
|
|
"""
|
|
|
|
import json
|
|
import asyncio
|
|
import websockets
|
|
from typing import Optional, Dict, Any, Iterator, Union, List
|
|
from threading import Lock
|
|
|
|
from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk, ProvenanceEvent
|
|
from . exceptions import ProtocolException, raise_from_error_dict
|
|
|
|
|
|
def build_term(value: Any, term_type: Optional[str] = None,
|
|
datatype: Optional[str] = None, language: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Build wire-format Term dict from a value.
|
|
|
|
Auto-detection rules (when term_type is None):
|
|
- Already a dict with 't' key -> return as-is (already a Term)
|
|
- Starts with http://, https://, urn: -> IRI
|
|
- Wrapped in <> (e.g., <http://...>) -> IRI (angle brackets stripped)
|
|
- Anything else -> literal
|
|
|
|
Args:
|
|
value: The term value (string, dict, or None)
|
|
term_type: One of 'iri', 'literal', or None for auto-detect
|
|
datatype: Datatype for literal objects (e.g., xsd:integer)
|
|
language: Language tag for literal objects (e.g., en)
|
|
|
|
Returns:
|
|
dict: Wire-format Term dict, or None if value is None
|
|
"""
|
|
if value is None:
|
|
return None
|
|
|
|
# If already a Term dict, return as-is
|
|
if isinstance(value, dict) and "t" in value:
|
|
return value
|
|
|
|
# Convert to string for processing
|
|
value = str(value)
|
|
|
|
# Auto-detect type if not specified
|
|
if term_type is None:
|
|
if value.startswith("<") and value.endswith(">") and not value.startswith("<<"):
|
|
# Angle-bracket wrapped IRI: <http://...>
|
|
value = value[1:-1] # Strip < and >
|
|
term_type = "iri"
|
|
elif value.startswith(("http://", "https://", "urn:")):
|
|
term_type = "iri"
|
|
else:
|
|
term_type = "literal"
|
|
|
|
if term_type == "iri":
|
|
# Strip angle brackets if present
|
|
if value.startswith("<") and value.endswith(">"):
|
|
value = value[1:-1]
|
|
return {"t": "i", "i": value}
|
|
elif term_type == "literal":
|
|
result = {"t": "l", "v": value}
|
|
if datatype:
|
|
result["dt"] = datatype
|
|
if language:
|
|
result["ln"] = language
|
|
return result
|
|
else:
|
|
raise ValueError(f"Unknown term type: {term_type}")
|
|
|
|
|
|
class SocketClient:
|
|
"""
|
|
Synchronous WebSocket client with persistent connection.
|
|
|
|
Maintains a single websocket connection and multiplexes requests
|
|
by ID via a background reader task. Provides synchronous generators
|
|
for streaming responses.
|
|
"""
|
|
|
|
def __init__(self, url: str, timeout: int, token: Optional[str]) -> None:
|
|
self.url: str = self._convert_to_ws_url(url)
|
|
self.timeout: int = timeout
|
|
self.token: Optional[str] = token
|
|
self._request_counter: int = 0
|
|
self._lock: Lock = Lock()
|
|
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
self._socket = None
|
|
self._connect_cm = None
|
|
self._reader_task = None
|
|
self._pending: Dict[str, asyncio.Queue] = {}
|
|
self._connected: bool = False
|
|
|
|
def _convert_to_ws_url(self, url: str) -> str:
|
|
if url.startswith("http://"):
|
|
return url.replace("http://", "ws://", 1)
|
|
elif url.startswith("https://"):
|
|
return url.replace("https://", "wss://", 1)
|
|
elif url.startswith("ws://") or url.startswith("wss://"):
|
|
return url
|
|
else:
|
|
return f"ws://{url}"
|
|
|
|
def _build_ws_url(self):
|
|
ws_url = f"{self.url.rstrip('/')}/api/v1/socket"
|
|
if self.token:
|
|
ws_url = f"{ws_url}?token={self.token}"
|
|
return ws_url
|
|
|
|
def _get_loop(self):
|
|
"""Get or create the event loop, reusing across calls."""
|
|
if self._loop is None or self._loop.is_closed():
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
if loop.is_running():
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
self._loop = loop
|
|
return self._loop
|
|
|
|
async def _ensure_connected(self):
|
|
"""Lazily establish the persistent websocket connection."""
|
|
if self._connected:
|
|
return
|
|
|
|
ws_url = self._build_ws_url()
|
|
self._connect_cm = websockets.connect(
|
|
ws_url, ping_interval=20, ping_timeout=self.timeout
|
|
)
|
|
self._socket = await self._connect_cm.__aenter__()
|
|
self._connected = True
|
|
self._reader_task = asyncio.create_task(self._reader())
|
|
|
|
async def _reader(self):
|
|
"""Background task to read responses and route by request ID."""
|
|
try:
|
|
async for raw_message in self._socket:
|
|
response = json.loads(raw_message)
|
|
request_id = response.get("id")
|
|
|
|
if request_id and request_id in self._pending:
|
|
await self._pending[request_id].put(response)
|
|
|
|
except websockets.exceptions.ConnectionClosed:
|
|
pass
|
|
except Exception as e:
|
|
for queue in self._pending.values():
|
|
try:
|
|
await queue.put({"error": str(e)})
|
|
except:
|
|
pass
|
|
finally:
|
|
self._connected = False
|
|
|
|
def _next_request_id(self):
|
|
with self._lock:
|
|
self._request_counter += 1
|
|
return f"req-{self._request_counter}"
|
|
|
|
def flow(self, flow_id: str) -> "SocketFlowInstance":
|
|
return SocketFlowInstance(self, flow_id)
|
|
|
|
def _send_request_sync(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any],
|
|
streaming: bool = False,
|
|
streaming_raw: bool = False,
|
|
include_provenance: bool = False
|
|
) -> Union[Dict[str, Any], Iterator[StreamingChunk], Iterator[Dict[str, Any]]]:
|
|
"""Synchronous wrapper around async WebSocket communication."""
|
|
loop = self._get_loop()
|
|
|
|
if streaming_raw:
|
|
return self._streaming_generator_raw(service, flow, request, loop)
|
|
elif streaming:
|
|
return self._streaming_generator(service, flow, request, loop, include_provenance)
|
|
else:
|
|
return loop.run_until_complete(self._send_request_async(service, flow, request))
|
|
|
|
def _streaming_generator(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any],
|
|
loop: asyncio.AbstractEventLoop,
|
|
include_provenance: bool = False
|
|
) -> Iterator[StreamingChunk]:
|
|
"""Generator that yields streaming chunks."""
|
|
async_gen = self._send_request_async_streaming(service, flow, request, include_provenance)
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
chunk = loop.run_until_complete(async_gen.__anext__())
|
|
yield chunk
|
|
except StopAsyncIteration:
|
|
break
|
|
finally:
|
|
try:
|
|
loop.run_until_complete(async_gen.aclose())
|
|
except:
|
|
pass
|
|
|
|
def _streaming_generator_raw(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any],
|
|
loop: asyncio.AbstractEventLoop
|
|
) -> Iterator[Dict[str, Any]]:
|
|
"""Generator that yields raw response dicts."""
|
|
async_gen = self._send_request_async_streaming_raw(service, flow, request)
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
data = loop.run_until_complete(async_gen.__anext__())
|
|
yield data
|
|
except StopAsyncIteration:
|
|
break
|
|
finally:
|
|
try:
|
|
loop.run_until_complete(async_gen.aclose())
|
|
except:
|
|
pass
|
|
|
|
async def _send_request_async_streaming_raw(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any]
|
|
) -> Iterator[Dict[str, Any]]:
|
|
"""Async streaming that yields raw response dicts."""
|
|
await self._ensure_connected()
|
|
|
|
request_id = self._next_request_id()
|
|
queue = asyncio.Queue()
|
|
self._pending[request_id] = queue
|
|
|
|
try:
|
|
message = {
|
|
"id": request_id,
|
|
"service": service,
|
|
"request": request
|
|
}
|
|
if flow:
|
|
message["flow"] = flow
|
|
|
|
await self._socket.send(json.dumps(message))
|
|
|
|
while True:
|
|
response = await queue.get()
|
|
|
|
if "error" in response:
|
|
raise_from_error_dict(response["error"])
|
|
|
|
if "response" in response:
|
|
yield response["response"]
|
|
|
|
if response.get("complete"):
|
|
break
|
|
|
|
finally:
|
|
self._pending.pop(request_id, None)
|
|
|
|
async def _send_request_async(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
"""Async non-streaming request over persistent connection."""
|
|
await self._ensure_connected()
|
|
|
|
request_id = self._next_request_id()
|
|
queue = asyncio.Queue()
|
|
self._pending[request_id] = queue
|
|
|
|
try:
|
|
message = {
|
|
"id": request_id,
|
|
"service": service,
|
|
"request": request
|
|
}
|
|
if flow:
|
|
message["flow"] = flow
|
|
|
|
await self._socket.send(json.dumps(message))
|
|
|
|
response = await queue.get()
|
|
|
|
if "error" in response:
|
|
raise_from_error_dict(response["error"])
|
|
|
|
if "response" not in response:
|
|
raise ProtocolException("Missing response in message")
|
|
|
|
return response["response"]
|
|
|
|
finally:
|
|
self._pending.pop(request_id, None)
|
|
|
|
async def _send_request_async_streaming(
|
|
self,
|
|
service: str,
|
|
flow: Optional[str],
|
|
request: Dict[str, Any],
|
|
include_provenance: bool = False
|
|
) -> Iterator[StreamingChunk]:
|
|
"""Async streaming request over persistent connection."""
|
|
await self._ensure_connected()
|
|
|
|
request_id = self._next_request_id()
|
|
queue = asyncio.Queue()
|
|
self._pending[request_id] = queue
|
|
|
|
try:
|
|
message = {
|
|
"id": request_id,
|
|
"service": service,
|
|
"request": request
|
|
}
|
|
if flow:
|
|
message["flow"] = flow
|
|
|
|
await self._socket.send(json.dumps(message))
|
|
|
|
while True:
|
|
response = await queue.get()
|
|
|
|
if "error" in response:
|
|
raise_from_error_dict(response["error"])
|
|
|
|
if "response" in response:
|
|
resp = response["response"]
|
|
|
|
if "error" in resp:
|
|
raise_from_error_dict(resp["error"])
|
|
|
|
chunk = self._parse_chunk(resp, include_provenance=include_provenance)
|
|
if chunk is not None:
|
|
yield chunk
|
|
|
|
if resp.get("end_of_session") or resp.get("end_of_dialog") or response.get("complete"):
|
|
break
|
|
|
|
finally:
|
|
self._pending.pop(request_id, None)
|
|
|
|
def _parse_chunk(self, resp: Dict[str, Any], include_provenance: bool = False) -> Optional[StreamingChunk]:
|
|
"""Parse response chunk into appropriate type. Returns None for non-content messages."""
|
|
chunk_type = resp.get("chunk_type")
|
|
message_type = resp.get("message_type")
|
|
|
|
# Handle GraphRAG/DocRAG message format with message_type
|
|
if message_type == "explain":
|
|
if include_provenance:
|
|
return ProvenanceEvent(
|
|
explain_id=resp.get("explain_id", ""),
|
|
explain_graph=resp.get("explain_graph", "")
|
|
)
|
|
return None
|
|
|
|
# Handle Agent message format with chunk_type="explain"
|
|
if chunk_type == "explain":
|
|
if include_provenance:
|
|
return ProvenanceEvent(
|
|
explain_id=resp.get("explain_id", ""),
|
|
explain_graph=resp.get("explain_graph", "")
|
|
)
|
|
return None
|
|
|
|
if chunk_type == "thought":
|
|
return AgentThought(
|
|
content=resp.get("content", ""),
|
|
end_of_message=resp.get("end_of_message", False),
|
|
message_id=resp.get("message_id", ""),
|
|
)
|
|
elif chunk_type == "observation":
|
|
return AgentObservation(
|
|
content=resp.get("content", ""),
|
|
end_of_message=resp.get("end_of_message", False),
|
|
message_id=resp.get("message_id", ""),
|
|
)
|
|
elif chunk_type == "answer" or chunk_type == "final-answer":
|
|
return AgentAnswer(
|
|
content=resp.get("content", ""),
|
|
end_of_message=resp.get("end_of_message", False),
|
|
end_of_dialog=resp.get("end_of_dialog", False),
|
|
message_id=resp.get("message_id", ""),
|
|
)
|
|
elif chunk_type == "action":
|
|
return AgentThought(
|
|
content=resp.get("content", ""),
|
|
end_of_message=resp.get("end_of_message", False)
|
|
)
|
|
else:
|
|
content = resp.get("response", resp.get("chunk", resp.get("text", "")))
|
|
return RAGChunk(
|
|
content=content,
|
|
end_of_stream=resp.get("end_of_stream", False),
|
|
error=None
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""Close the persistent WebSocket connection."""
|
|
if self._loop and not self._loop.is_closed():
|
|
try:
|
|
self._loop.run_until_complete(self._close_async())
|
|
except:
|
|
pass
|
|
|
|
async def _close_async(self):
|
|
# Cancel reader task
|
|
if self._reader_task:
|
|
self._reader_task.cancel()
|
|
try:
|
|
await self._reader_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._reader_task = None
|
|
|
|
# Exit the websockets context manager — this cleanly shuts down
|
|
# the connection and its keepalive task
|
|
if self._connect_cm:
|
|
try:
|
|
await self._connect_cm.__aexit__(None, None, None)
|
|
except Exception:
|
|
pass
|
|
self._connect_cm = None
|
|
|
|
self._socket = None
|
|
self._connected = False
|
|
self._pending.clear()
|
|
|
|
|
|
class SocketFlowInstance:
|
|
"""
|
|
Synchronous WebSocket flow instance for streaming operations.
|
|
|
|
Provides the same interface as REST FlowInstance but with WebSocket-based
|
|
streaming support for real-time responses.
|
|
"""
|
|
|
|
def __init__(self, client: SocketClient, flow_id: str) -> None:
|
|
self.client: SocketClient = client
|
|
self.flow_id: str = flow_id
|
|
|
|
def agent(
|
|
self,
|
|
question: str,
|
|
user: str,
|
|
state: Optional[Dict[str, Any]] = None,
|
|
group: Optional[str] = None,
|
|
history: Optional[List[Dict[str, Any]]] = None,
|
|
streaming: bool = False,
|
|
**kwargs: Any
|
|
) -> Union[Dict[str, Any], Iterator[StreamingChunk]]:
|
|
"""Execute an agent operation with streaming support."""
|
|
request = {
|
|
"question": question,
|
|
"user": user,
|
|
"streaming": streaming
|
|
}
|
|
if state is not None:
|
|
request["state"] = state
|
|
if group is not None:
|
|
request["group"] = group
|
|
if history is not None:
|
|
request["history"] = history
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync("agent", self.flow_id, request, streaming=True)
|
|
|
|
def agent_explain(
|
|
self,
|
|
question: str,
|
|
user: str,
|
|
collection: str,
|
|
state: Optional[Dict[str, Any]] = None,
|
|
group: Optional[str] = None,
|
|
history: Optional[List[Dict[str, Any]]] = None,
|
|
**kwargs: Any
|
|
) -> Iterator[Union[StreamingChunk, ProvenanceEvent]]:
|
|
"""Execute an agent operation with explainability support."""
|
|
request = {
|
|
"question": question,
|
|
"user": user,
|
|
"collection": collection,
|
|
"streaming": True
|
|
}
|
|
if state is not None:
|
|
request["state"] = state
|
|
if group is not None:
|
|
request["group"] = group
|
|
if history is not None:
|
|
request["history"] = history
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync(
|
|
"agent", self.flow_id, request,
|
|
streaming=True, include_provenance=True
|
|
)
|
|
|
|
def text_completion(self, system: str, prompt: str, streaming: bool = False, **kwargs) -> Union[str, Iterator[str]]:
|
|
"""Execute text completion with optional streaming."""
|
|
request = {
|
|
"system": system,
|
|
"prompt": prompt,
|
|
"streaming": streaming
|
|
}
|
|
request.update(kwargs)
|
|
|
|
result = self.client._send_request_sync("text-completion", self.flow_id, request, streaming)
|
|
|
|
if streaming:
|
|
return self._text_completion_generator(result)
|
|
else:
|
|
return result.get("response", "")
|
|
|
|
def _text_completion_generator(self, result: Iterator[StreamingChunk]) -> Iterator[str]:
|
|
for chunk in result:
|
|
if hasattr(chunk, 'content'):
|
|
yield chunk.content
|
|
|
|
def graph_rag(
|
|
self,
|
|
query: str,
|
|
user: str,
|
|
collection: str,
|
|
entity_limit: int = 50,
|
|
triple_limit: int = 30,
|
|
max_subgraph_size: int = 1000,
|
|
max_path_length: int = 2,
|
|
edge_score_limit: int = 30,
|
|
edge_limit: int = 25,
|
|
streaming: bool = False,
|
|
**kwargs: Any
|
|
) -> Union[str, Iterator[str]]:
|
|
"""Execute graph-based RAG query with optional streaming."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection,
|
|
"entity-limit": entity_limit,
|
|
"triple-limit": triple_limit,
|
|
"max-subgraph-size": max_subgraph_size,
|
|
"max-path-length": max_path_length,
|
|
"edge-score-limit": edge_score_limit,
|
|
"edge-limit": edge_limit,
|
|
"streaming": streaming
|
|
}
|
|
request.update(kwargs)
|
|
|
|
result = self.client._send_request_sync("graph-rag", self.flow_id, request, streaming)
|
|
|
|
if streaming:
|
|
return self._rag_generator(result)
|
|
else:
|
|
return result.get("response", "")
|
|
|
|
def graph_rag_explain(
|
|
self,
|
|
query: str,
|
|
user: str,
|
|
collection: str,
|
|
entity_limit: int = 50,
|
|
triple_limit: int = 30,
|
|
max_subgraph_size: int = 1000,
|
|
max_path_length: int = 2,
|
|
edge_score_limit: int = 30,
|
|
edge_limit: int = 25,
|
|
**kwargs: Any
|
|
) -> Iterator[Union[RAGChunk, ProvenanceEvent]]:
|
|
"""Execute graph-based RAG query with explainability support."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection,
|
|
"entity-limit": entity_limit,
|
|
"triple-limit": triple_limit,
|
|
"max-subgraph-size": max_subgraph_size,
|
|
"max-path-length": max_path_length,
|
|
"edge-score-limit": edge_score_limit,
|
|
"edge-limit": edge_limit,
|
|
"streaming": True,
|
|
"explainable": True,
|
|
}
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync(
|
|
"graph-rag", self.flow_id, request,
|
|
streaming=True, include_provenance=True
|
|
)
|
|
|
|
def document_rag(
|
|
self,
|
|
query: str,
|
|
user: str,
|
|
collection: str,
|
|
doc_limit: int = 10,
|
|
streaming: bool = False,
|
|
**kwargs: Any
|
|
) -> Union[str, Iterator[str]]:
|
|
"""Execute document-based RAG query with optional streaming."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection,
|
|
"doc-limit": doc_limit,
|
|
"streaming": streaming
|
|
}
|
|
request.update(kwargs)
|
|
|
|
result = self.client._send_request_sync("document-rag", self.flow_id, request, streaming)
|
|
|
|
if streaming:
|
|
return self._rag_generator(result)
|
|
else:
|
|
return result.get("response", "")
|
|
|
|
def document_rag_explain(
|
|
self,
|
|
query: str,
|
|
user: str,
|
|
collection: str,
|
|
doc_limit: int = 10,
|
|
**kwargs: Any
|
|
) -> Iterator[Union[RAGChunk, ProvenanceEvent]]:
|
|
"""Execute document-based RAG query with explainability support."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection,
|
|
"doc-limit": doc_limit,
|
|
"streaming": True,
|
|
"explainable": True,
|
|
}
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync(
|
|
"document-rag", self.flow_id, request,
|
|
streaming=True, include_provenance=True
|
|
)
|
|
|
|
def _rag_generator(self, result: Iterator[StreamingChunk]) -> Iterator[str]:
|
|
for chunk in result:
|
|
if hasattr(chunk, 'content'):
|
|
yield chunk.content
|
|
|
|
def prompt(
|
|
self,
|
|
id: str,
|
|
variables: Dict[str, str],
|
|
streaming: bool = False,
|
|
**kwargs: Any
|
|
) -> Union[str, Iterator[str]]:
|
|
"""Execute a prompt template with optional streaming."""
|
|
request = {
|
|
"id": id,
|
|
"variables": variables,
|
|
"streaming": streaming
|
|
}
|
|
request.update(kwargs)
|
|
|
|
result = self.client._send_request_sync("prompt", self.flow_id, request, streaming)
|
|
|
|
if streaming:
|
|
return self._rag_generator(result)
|
|
else:
|
|
return result.get("response", "")
|
|
|
|
def graph_embeddings_query(
|
|
self,
|
|
text: str,
|
|
user: str,
|
|
collection: str,
|
|
limit: int = 10,
|
|
**kwargs: Any
|
|
) -> Dict[str, Any]:
|
|
"""Query knowledge graph entities using semantic similarity."""
|
|
emb_result = self.embeddings(texts=[text])
|
|
vector = emb_result.get("vectors", [[]])[0]
|
|
|
|
request = {
|
|
"vector": vector,
|
|
"user": user,
|
|
"collection": collection,
|
|
"limit": limit
|
|
}
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync("graph-embeddings", self.flow_id, request, False)
|
|
|
|
def document_embeddings_query(
|
|
self,
|
|
text: str,
|
|
user: str,
|
|
collection: str,
|
|
limit: int = 10,
|
|
**kwargs: Any
|
|
) -> Dict[str, Any]:
|
|
"""Query document chunks using semantic similarity."""
|
|
emb_result = self.embeddings(texts=[text])
|
|
vector = emb_result.get("vectors", [[]])[0]
|
|
|
|
request = {
|
|
"vector": vector,
|
|
"user": user,
|
|
"collection": collection,
|
|
"limit": limit
|
|
}
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync("document-embeddings", self.flow_id, request, False)
|
|
|
|
def embeddings(self, texts: list, **kwargs: Any) -> Dict[str, Any]:
|
|
"""Generate vector embeddings for one or more texts."""
|
|
request = {"texts": texts}
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync("embeddings", self.flow_id, request, False)
|
|
|
|
def triples_query(
|
|
self,
|
|
s: Optional[Union[str, Dict[str, Any]]] = None,
|
|
p: Optional[Union[str, Dict[str, Any]]] = None,
|
|
o: Optional[Union[str, Dict[str, Any]]] = None,
|
|
g: Optional[str] = None,
|
|
user: Optional[str] = None,
|
|
collection: Optional[str] = None,
|
|
limit: int = 100,
|
|
**kwargs: Any
|
|
) -> List[Dict[str, Any]]:
|
|
"""Query knowledge graph triples using pattern matching."""
|
|
request = {"limit": limit}
|
|
|
|
s_term = build_term(s)
|
|
p_term = build_term(p)
|
|
o_term = build_term(o)
|
|
|
|
if s_term is not None:
|
|
request["s"] = s_term
|
|
if p_term is not None:
|
|
request["p"] = p_term
|
|
if o_term is not None:
|
|
request["o"] = o_term
|
|
if g is not None:
|
|
request["g"] = g
|
|
if user is not None:
|
|
request["user"] = user
|
|
if collection is not None:
|
|
request["collection"] = collection
|
|
request.update(kwargs)
|
|
|
|
result = self.client._send_request_sync("triples", self.flow_id, request, False)
|
|
if isinstance(result, dict) and "response" in result:
|
|
return result["response"]
|
|
return result
|
|
|
|
def triples_query_stream(
|
|
self,
|
|
s: Optional[Union[str, Dict[str, Any]]] = None,
|
|
p: Optional[Union[str, Dict[str, Any]]] = None,
|
|
o: Optional[Union[str, Dict[str, Any]]] = None,
|
|
g: Optional[str] = None,
|
|
user: Optional[str] = None,
|
|
collection: Optional[str] = None,
|
|
limit: int = 100,
|
|
batch_size: int = 20,
|
|
**kwargs: Any
|
|
) -> Iterator[List[Dict[str, Any]]]:
|
|
"""Query knowledge graph triples with streaming batches."""
|
|
request = {
|
|
"limit": limit,
|
|
"streaming": True,
|
|
"batch-size": batch_size,
|
|
}
|
|
|
|
s_term = build_term(s)
|
|
p_term = build_term(p)
|
|
o_term = build_term(o)
|
|
|
|
if s_term is not None:
|
|
request["s"] = s_term
|
|
if p_term is not None:
|
|
request["p"] = p_term
|
|
if o_term is not None:
|
|
request["o"] = o_term
|
|
if g is not None:
|
|
request["g"] = g
|
|
if user is not None:
|
|
request["user"] = user
|
|
if collection is not None:
|
|
request["collection"] = collection
|
|
request.update(kwargs)
|
|
|
|
for response in self.client._send_request_sync("triples", self.flow_id, request, streaming_raw=True):
|
|
if isinstance(response, dict) and "response" in response:
|
|
yield response["response"]
|
|
else:
|
|
yield response
|
|
|
|
def sparql_query_stream(
|
|
self,
|
|
query: str,
|
|
user: str = "trustgraph",
|
|
collection: str = "default",
|
|
limit: int = 10000,
|
|
batch_size: int = 20,
|
|
**kwargs: Any
|
|
) -> Iterator[Dict[str, Any]]:
|
|
"""Execute a SPARQL query with streaming batches."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection,
|
|
"limit": limit,
|
|
"streaming": True,
|
|
"batch-size": batch_size,
|
|
}
|
|
request.update(kwargs)
|
|
|
|
for response in self.client._send_request_sync(
|
|
"sparql", self.flow_id, request, streaming_raw=True
|
|
):
|
|
yield response
|
|
|
|
def rows_query(
|
|
self,
|
|
query: str,
|
|
user: str,
|
|
collection: str,
|
|
variables: Optional[Dict[str, Any]] = None,
|
|
operation_name: Optional[str] = None,
|
|
**kwargs: Any
|
|
) -> Dict[str, Any]:
|
|
"""Execute a GraphQL query against structured rows."""
|
|
request = {
|
|
"query": query,
|
|
"user": user,
|
|
"collection": collection
|
|
}
|
|
if variables:
|
|
request["variables"] = variables
|
|
if operation_name:
|
|
request["operationName"] = operation_name
|
|
request.update(kwargs)
|
|
|
|
return self.client._send_request_sync("rows", self.flow_id, request, False)
|
|
|
|
def mcp_tool(
|
|
self,
|
|
name: str,
|
|
parameters: Dict[str, Any],
|
|
**kwargs: Any
|
|
) -> Dict[str, Any]:
|
|
"""Execute a Model Context Protocol (MCP) tool."""
|
|
request = {
|
|
"name": name,
|
|
"parameters": parameters
|
|
}
|
|
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."""
|
|
emb_result = self.embeddings(texts=[text])
|
|
vector = emb_result.get("vectors", [[]])[0]
|
|
|
|
request = {
|
|
"vector": vector,
|
|
"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)
|