mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 04:31:02 +02:00
aecf00f040
Update RAG and Agent clients for streaming message handling GraphRAG now sends multiple message types in a stream: - 'explain' messages with explain_id and explain_graph for provenance - 'chunk' messages with response text fragments - end_of_session marker for stream completion Updated all clients to handle this properly: CLI clients (trustgraph-base/trustgraph/clients/): - graph_rag_client.py: Added chunk_callback and explain_callback - document_rag_client.py: Added chunk_callback and explain_callback - agent_client.py: Added think, observe, answer_callback, error_callback Internal clients (trustgraph-base/trustgraph/base/): - graph_rag_client.py: Async callbacks for streaming - agent_client.py: Async callbacks for streaming All clients now: - Route messages by chunk_type/message_type - Stream via optional callbacks for incremental delivery - Wait for proper completion signals (end_of_dialog/end_of_session/end_of_stream) - Accumulate and return complete response for callers not using callbacks Updated callers: - extract/kg/agent/extract.py: Uses new invoke(question=...) API - tests/integration/test_agent_kg_extraction_integration.py: Updated mocks This fixes the agent infinite loop issue where knowledge_query was returning the first 'explain' message (empty response) instead of waiting for the actual answer chunks. Concurrency in triples query
81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
|
|
from . request_response_spec import RequestResponse, RequestResponseSpec
|
|
from .. schema import AgentRequest, AgentResponse
|
|
from .. knowledge import Uri, Literal
|
|
|
|
class AgentClient(RequestResponse):
|
|
async def invoke(self, question, plan=None, state=None,
|
|
history=[], think=None, observe=None, answer_callback=None,
|
|
timeout=300):
|
|
"""
|
|
Invoke the agent with optional streaming callbacks.
|
|
|
|
Args:
|
|
question: The question to ask
|
|
plan: Optional plan context
|
|
state: Optional state context
|
|
history: Conversation history
|
|
think: Optional async callback(content, end_of_message) for thought chunks
|
|
observe: Optional async callback(content, end_of_message) for observation chunks
|
|
answer_callback: Optional async callback(content, end_of_message) for answer chunks
|
|
timeout: Request timeout in seconds
|
|
|
|
Returns:
|
|
Complete answer text (accumulated from all answer chunks)
|
|
"""
|
|
accumulated_answer = []
|
|
|
|
async def recipient(resp):
|
|
if resp.error:
|
|
raise RuntimeError(resp.error.message)
|
|
|
|
# Handle thought chunks
|
|
if resp.chunk_type == 'thought':
|
|
if think:
|
|
await think(resp.content, resp.end_of_message)
|
|
return False # Continue receiving
|
|
|
|
# Handle observation chunks
|
|
if resp.chunk_type == 'observation':
|
|
if observe:
|
|
await observe(resp.content, resp.end_of_message)
|
|
return False # Continue receiving
|
|
|
|
# Handle answer chunks
|
|
if resp.chunk_type == 'answer':
|
|
if resp.content:
|
|
accumulated_answer.append(resp.content)
|
|
if answer_callback:
|
|
await answer_callback(resp.content, resp.end_of_message)
|
|
|
|
# Complete when dialog ends
|
|
if resp.end_of_dialog:
|
|
return True
|
|
|
|
return False # Continue receiving
|
|
|
|
await self.request(
|
|
AgentRequest(
|
|
question = question,
|
|
plan = plan,
|
|
state = state,
|
|
history = history,
|
|
),
|
|
recipient=recipient,
|
|
timeout=timeout,
|
|
)
|
|
|
|
return "".join(accumulated_answer)
|
|
|
|
class AgentClientSpec(RequestResponseSpec):
|
|
def __init__(
|
|
self, request_name, response_name,
|
|
):
|
|
super(AgentClientSpec, self).__init__(
|
|
request_name = request_name,
|
|
request_schema = AgentRequest,
|
|
response_name = response_name,
|
|
response_schema = AgentResponse,
|
|
impl = AgentClient,
|
|
)
|
|
|