Agent streaming

This commit is contained in:
Cyber MacGeddon 2025-11-25 20:16:13 +00:00
parent 5d95c90120
commit 83e71db3ab
4 changed files with 170 additions and 45 deletions

View file

@ -12,7 +12,8 @@ class AgentRequestTranslator(MessageTranslator):
state=data.get("state", None), state=data.get("state", None),
group=data.get("group", None), group=data.get("group", None),
history=data.get("history", []), 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]: def from_pulsar(self, obj: AgentRequest) -> Dict[str, Any]:
@ -21,7 +22,8 @@ class AgentRequestTranslator(MessageTranslator):
"state": obj.state, "state": obj.state,
"group": obj.group, "group": obj.group,
"history": obj.history, "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]: def from_pulsar(self, obj: AgentResponse) -> Dict[str, Any]:
result = {} result = {}
# 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: if obj.answer:
result["answer"] = obj.answer result["answer"] = obj.answer
if obj.thought: if obj.thought:
result["thought"] = obj.thought result["thought"] = obj.thought
if obj.observation: if obj.observation:
result["observation"] = 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 return result
def from_response_with_completion(self, obj: AgentResponse) -> Tuple[Dict[str, Any], bool]: def from_response_with_completion(self, obj: AgentResponse) -> Tuple[Dict[str, Any], bool]:
"""Returns (response_dict, is_final)""" """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

View file

@ -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.topic import topic
from ..core.primitives import Error from ..core.primitives import Error
@ -21,8 +21,16 @@ class AgentRequest(Record):
group = Array(String()) group = Array(String())
history = Array(AgentStep()) history = Array(AgentStep())
user = String() # User context for multi-tenancy user = String() # User context for multi-tenancy
streaming = Boolean() # NEW: Enable streaming response delivery (default false)
class AgentResponse(Record): 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() answer = String()
error = Error() error = Error()
thought = String() thought = String()

View file

@ -29,7 +29,7 @@ def output(text, prefix="> ", width=78):
async def question( async def question(
url, question, flow_id, user, collection, 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("/"): if not url.endswith("/"):
@ -62,7 +62,8 @@ async def question(
"request": { "request": {
"question": question, "question": question,
"user": user, "user": user,
"history": [] "history": [],
"streaming": streaming
} }
} }
@ -89,14 +90,34 @@ async def question(
print("Ignore message") print("Ignore message")
continue continue
if "thought" in obj["response"]: response = obj["response"]
think(obj["response"]["thought"])
if "observation" in obj["response"]: # Handle streaming format (new format with chunk_type)
observe(obj["response"]["observation"]) if "chunk_type" in response:
chunk_type = response["chunk_type"]
content = response.get("content", "")
if "answer" in obj["response"]: if chunk_type == "thought":
print(obj["response"]["answer"]) 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 if obj["complete"]: break
@ -161,6 +182,12 @@ def main():
help=f'Output thinking/observations' help=f'Output thinking/observations'
) )
parser.add_argument(
'--no-streaming',
action="store_true",
help=f'Disable streaming (use legacy mode)'
)
args = parser.parse_args() args = parser.parse_args()
try: try:
@ -176,6 +203,7 @@ def main():
state = args.state, state = args.state,
group = args.group, group = args.group,
verbose = args.verbose, verbose = args.verbose,
streaming = not args.no_streaming,
) )
) )

View file

@ -191,6 +191,9 @@ class Processor(AgentService):
try: try:
# Check if streaming is enabled
streaming = getattr(request, 'streaming', False)
if request.history: if request.history:
history = [ history = [
Action( Action(
@ -215,6 +218,21 @@ class Processor(AgentService):
logger.debug(f"Think: {x}") logger.debug(f"Think: {x}")
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( r = AgentResponse(
answer=None, answer=None,
error=None, error=None,
@ -228,6 +246,21 @@ class Processor(AgentService):
logger.debug(f"Observe: {x}") logger.debug(f"Observe: {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( r = AgentResponse(
answer=None, answer=None,
error=None, error=None,
@ -287,6 +320,20 @@ class Processor(AgentService):
else: else:
f = json.dumps(act.final) f = json.dumps(act.final)
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( r = AgentResponse(
answer=act.final, answer=act.final,
error=None, error=None,
@ -336,11 +383,29 @@ class Processor(AgentService):
logger.debug("Send error response...") logger.debug("Send error response...")
r = AgentResponse( error_obj = Error(
error=Error(
type = "agent-error", type = "agent-error",
message = str(e), 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, response=None,
) )