mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-06-25 06:38:06 +02:00
Minor agent tweaks (#692)
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
This commit is contained in:
parent
45e6ad4abc
commit
aecf00f040
8 changed files with 246 additions and 58 deletions
|
|
@ -42,10 +42,50 @@ class GraphRagClient(BaseClient):
|
|||
|
||||
def request(
|
||||
self, query, user="trustgraph", collection="default",
|
||||
chunk_callback=None,
|
||||
explain_callback=None,
|
||||
timeout=500
|
||||
):
|
||||
"""
|
||||
Request a graph RAG query with optional streaming callbacks.
|
||||
|
||||
return self.call(
|
||||
user=user, collection=collection, query=query, timeout=timeout
|
||||
).response
|
||||
Args:
|
||||
query: The question to ask
|
||||
user: User identifier
|
||||
collection: Collection identifier
|
||||
chunk_callback: Optional callback(text, end_of_stream) for text chunks
|
||||
explain_callback: Optional callback(explain_id, explain_graph) for explain notifications
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
Complete response text (accumulated from all chunks)
|
||||
"""
|
||||
accumulated_response = []
|
||||
|
||||
def inspect(x):
|
||||
# Handle explain notifications
|
||||
if x.message_type == 'explain':
|
||||
if explain_callback and x.explain_id:
|
||||
explain_callback(x.explain_id, x.explain_graph)
|
||||
return False # Continue receiving
|
||||
|
||||
# Handle text chunks
|
||||
if x.message_type == 'chunk':
|
||||
if x.response:
|
||||
accumulated_response.append(x.response)
|
||||
if chunk_callback:
|
||||
chunk_callback(x.response, x.end_of_stream)
|
||||
|
||||
# Complete when session ends
|
||||
if x.end_of_session:
|
||||
return True
|
||||
|
||||
return False # Continue receiving
|
||||
|
||||
self.call(
|
||||
user=user, collection=collection, query=query,
|
||||
inspect=inspect, timeout=timeout
|
||||
)
|
||||
|
||||
return "".join(accumulated_response)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue