mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 04:31:02 +02:00
Agent streaming
This commit is contained in:
Cyber MacGeddon
parent
5d95c90120
commit
83e71db3ab
4 changed files with 170 additions and 45 deletions
|
|
@ -12,16 +12,18 @@ class AgentRequestTranslator(MessageTranslator):
|
|||
state=data.get("state", None),
|
||||
group=data.get("group", None),
|
||||
history=data.get("history", []),
|
||||
user=data.get("user", "trustgraph")
|
||||
user=data.get("user", "trustgraph"),
|
||||
streaming=data.get("streaming", False)
|
||||
)
|
||||
|
||||
|
||||
def from_pulsar(self, obj: AgentRequest) -> Dict[str, Any]:
|
||||
return {
|
||||
"question": obj.question,
|
||||
"state": obj.state,
|
||||
"group": obj.group,
|
||||
"history": obj.history,
|
||||
"user": obj.user
|
||||
"user": obj.user,
|
||||
"streaming": getattr(obj, "streaming", False)
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -33,14 +35,36 @@ class AgentResponseTranslator(MessageTranslator):
|
|||
|
||||
def from_pulsar(self, obj: AgentResponse) -> Dict[str, Any]:
|
||||
result = {}
|
||||
if obj.answer:
|
||||
result["answer"] = obj.answer
|
||||
if obj.thought:
|
||||
result["thought"] = obj.thought
|
||||
if obj.observation:
|
||||
result["observation"] = obj.observation
|
||||
|
||||
# Check if this is a streaming response (has chunk_type)
|
||||
if hasattr(obj, 'chunk_type') and obj.chunk_type:
|
||||
result["chunk_type"] = obj.chunk_type
|
||||
if obj.content:
|
||||
result["content"] = obj.content
|
||||
result["end_of_message"] = getattr(obj, "end_of_message", False)
|
||||
result["end_of_dialog"] = getattr(obj, "end_of_dialog", False)
|
||||
else:
|
||||
# Legacy format
|
||||
if obj.answer:
|
||||
result["answer"] = obj.answer
|
||||
if obj.thought:
|
||||
result["thought"] = obj.thought
|
||||
if obj.observation:
|
||||
result["observation"] = obj.observation
|
||||
|
||||
# Always include error if present
|
||||
if hasattr(obj, 'error') and obj.error and obj.error.message:
|
||||
result["error"] = {"message": obj.error.message, "code": obj.error.code}
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def from_response_with_completion(self, obj: AgentResponse) -> Tuple[Dict[str, Any], bool]:
|
||||
"""Returns (response_dict, is_final)"""
|
||||
return self.from_pulsar(obj), (obj.answer is not None)
|
||||
# For streaming responses, check end_of_dialog
|
||||
if hasattr(obj, 'chunk_type') and obj.chunk_type:
|
||||
is_final = getattr(obj, 'end_of_dialog', False)
|
||||
else:
|
||||
# For legacy responses, check if answer is present
|
||||
is_final = (obj.answer is not None)
|
||||
|
||||
return self.from_pulsar(obj), is_final
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
|
||||
from pulsar.schema import Record, String, Array, Map
|
||||
from pulsar.schema import Record, String, Array, Map, Boolean
|
||||
|
||||
from ..core.topic import topic
|
||||
from ..core.primitives import Error
|
||||
|
|
@ -21,8 +21,16 @@ class AgentRequest(Record):
|
|||
group = Array(String())
|
||||
history = Array(AgentStep())
|
||||
user = String() # User context for multi-tenancy
|
||||
streaming = Boolean() # NEW: Enable streaming response delivery (default false)
|
||||
|
||||
class AgentResponse(Record):
|
||||
# Streaming-first design
|
||||
chunk_type = String() # "thought", "action", "observation", "answer", "error"
|
||||
content = String() # The actual content (interpretation depends on chunk_type)
|
||||
end_of_message = Boolean() # Current chunk type (thought/action/etc.) is complete
|
||||
end_of_dialog = Boolean() # Entire agent dialog is complete
|
||||
|
||||
# Legacy fields (deprecated but kept for backward compatibility)
|
||||
answer = String()
|
||||
error = Error()
|
||||
thought = String()
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ def output(text, prefix="> ", width=78):
|
|||
|
||||
async def question(
|
||||
url, question, flow_id, user, collection,
|
||||
plan=None, state=None, group=None, verbose=False
|
||||
plan=None, state=None, group=None, verbose=False, streaming=True
|
||||
):
|
||||
|
||||
if not url.endswith("/"):
|
||||
|
|
@ -62,16 +62,17 @@ async def question(
|
|||
"request": {
|
||||
"question": question,
|
||||
"user": user,
|
||||
"history": []
|
||||
"history": [],
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# Only add optional fields if they have values
|
||||
if state is not None:
|
||||
req["request"]["state"] = state
|
||||
if group is not None:
|
||||
req["request"]["group"] = group
|
||||
|
||||
|
||||
req = json.dumps(req)
|
||||
|
||||
await ws.send(req)
|
||||
|
|
@ -89,14 +90,34 @@ async def question(
|
|||
print("Ignore message")
|
||||
continue
|
||||
|
||||
if "thought" in obj["response"]:
|
||||
think(obj["response"]["thought"])
|
||||
response = obj["response"]
|
||||
|
||||
if "observation" in obj["response"]:
|
||||
observe(obj["response"]["observation"])
|
||||
# Handle streaming format (new format with chunk_type)
|
||||
if "chunk_type" in response:
|
||||
chunk_type = response["chunk_type"]
|
||||
content = response.get("content", "")
|
||||
|
||||
if "answer" in obj["response"]:
|
||||
print(obj["response"]["answer"])
|
||||
if chunk_type == "thought":
|
||||
think(content)
|
||||
elif chunk_type == "observation":
|
||||
observe(content)
|
||||
elif chunk_type == "answer":
|
||||
print(content)
|
||||
elif chunk_type == "error":
|
||||
raise RuntimeError(content)
|
||||
else:
|
||||
# Handle legacy format (backward compatibility)
|
||||
if "thought" in response:
|
||||
think(response["thought"])
|
||||
|
||||
if "observation" in response:
|
||||
observe(response["observation"])
|
||||
|
||||
if "answer" in response:
|
||||
print(response["answer"])
|
||||
|
||||
if "error" in response:
|
||||
raise RuntimeError(response["error"])
|
||||
|
||||
if obj["complete"]: break
|
||||
|
||||
|
|
@ -161,6 +182,12 @@ def main():
|
|||
help=f'Output thinking/observations'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--no-streaming',
|
||||
action="store_true",
|
||||
help=f'Disable streaming (use legacy mode)'
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
|
|
@ -176,6 +203,7 @@ def main():
|
|||
state = args.state,
|
||||
group = args.group,
|
||||
verbose = args.verbose,
|
||||
streaming = not args.no_streaming,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -191,6 +191,9 @@ class Processor(AgentService):
|
|||
|
||||
try:
|
||||
|
||||
# Check if streaming is enabled
|
||||
streaming = getattr(request, 'streaming', False)
|
||||
|
||||
if request.history:
|
||||
history = [
|
||||
Action(
|
||||
|
|
@ -215,12 +218,27 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug(f"Think: {x}")
|
||||
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="thought",
|
||||
content=x,
|
||||
end_of_message=True,
|
||||
end_of_dialog=False,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=x,
|
||||
observation=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -228,12 +246,27 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug(f"Observe: {x}")
|
||||
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="observation",
|
||||
content=x,
|
||||
end_of_message=True,
|
||||
end_of_dialog=False,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
answer=None,
|
||||
error=None,
|
||||
thought=None,
|
||||
observation=x,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -287,11 +320,25 @@ class Processor(AgentService):
|
|||
else:
|
||||
f = json.dumps(act.final)
|
||||
|
||||
r = AgentResponse(
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
if streaming:
|
||||
# Streaming format - mark as final dialog
|
||||
r = AgentResponse(
|
||||
chunk_type="answer",
|
||||
content=f,
|
||||
end_of_message=True,
|
||||
end_of_dialog=True,
|
||||
# Legacy fields for backward compatibility
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
answer=act.final,
|
||||
error=None,
|
||||
thought=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
|
|
@ -336,14 +383,32 @@ class Processor(AgentService):
|
|||
|
||||
logger.debug("Send error response...")
|
||||
|
||||
r = AgentResponse(
|
||||
error=Error(
|
||||
type = "agent-error",
|
||||
message = str(e),
|
||||
),
|
||||
response=None,
|
||||
error_obj = Error(
|
||||
type = "agent-error",
|
||||
message = str(e),
|
||||
)
|
||||
|
||||
# Check if streaming was enabled (may not be set if error occurred early)
|
||||
streaming = getattr(request, 'streaming', False) if 'request' in locals() else False
|
||||
|
||||
if streaming:
|
||||
# Streaming format
|
||||
r = AgentResponse(
|
||||
chunk_type="error",
|
||||
content=str(e),
|
||||
end_of_message=True,
|
||||
end_of_dialog=True,
|
||||
# Legacy fields for backward compatibility
|
||||
error=error_obj,
|
||||
response=None,
|
||||
)
|
||||
else:
|
||||
# Legacy format
|
||||
r = AgentResponse(
|
||||
error=error_obj,
|
||||
response=None,
|
||||
)
|
||||
|
||||
await respond(r)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue