diff --git a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py index 3aba946c88..4f0f67096e 100644 --- a/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py +++ b/src/praisonai-agents/praisonaiagents/agent/chat_mixin.py @@ -29,6 +29,9 @@ # into tool worker threads with the context, so each turn has its own counters and # overlapping turns on a shared agent never exhaust each other's limits. _turn_loop_guard = contextvars.ContextVar("praisonai_turn_loop_guard", default=None) +# Agent whose streamed turn already records agent_start/agent_end, set while its +# non-streaming chat() fallback runs so that turn is not recorded twice. +_stream_traced_agent = contextvars.ContextVar("praisonai_stream_traced_agent", default=None) def _agent_end_status(db_status: str) -> str: @@ -3554,7 +3557,9 @@ def chat(self, prompt: str, temperature: Optional[float] = None, tools: Optional # Emit context trace event (zero overhead when not set) from ..trace.context_events import get_context_emitter _trace_emitter = get_context_emitter() - _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal}) + _trace_turn = _stream_traced_agent.get() is not self + if _trace_turn: + _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal}) durable_context = None durable_token = None previous_run_state = self._get_db_run_state() @@ -3623,7 +3628,8 @@ def chat(self, prompt: str, temperature: Optional[float] = None, tools: Optional end_durable_run(durable_token) durable_context.close() try: - _trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)}) + if _trace_turn: + _trace_emitter.agent_end(self.name, {"status": _agent_end_status(db_status)}) finally: if getattr(self, "_active_turn_token", None) is _cancel: self._active_turn_token = None @@ -5431,9 +5437,26 @@ def _start_stream(self, prompt: str, **kwargs) -> Generator[str, None, None]: # via _enter_turn) so overlapping streams on a shared agent neither # exhaust nor inherit each other's per-turn limits. guard_token = _turn_loop_guard.set([self, None]) + # Record the streamed turn like chat()/achat() do, so a live replay can + # show the agent running while its words stream (no-op when tracing is off). + from ..trace.context_events import get_context_emitter + _trace_emitter = get_context_emitter() + status = "completed" + if _trace_emitter._enabled: + _trace_emitter.agent_start(self.name, {"role": self.role, "goal": self.goal}) try: - yield from self._start_stream_turn(prompt, **kwargs) + # _start_stream_turn returns "failed" when a guardrail or hook blocks + # the turn, matching chat(), which records a blocked turn as failed. + status = (yield from self._start_stream_turn(prompt, **kwargs)) or status + except (GeneratorExit, InterruptedError): + status = "cancelled" + raise + except BaseException: + status = "failed" + raise finally: + if _trace_emitter._enabled: + _trace_emitter.agent_end(self.name, {"status": status}) try: _turn_loop_guard.reset(guard_token) except (ValueError, LookupError): @@ -5458,7 +5481,7 @@ def _start_stream_turn(self, prompt: str, **kwargs) -> Generator[str, None, None if not _in_ok: logging.warning(f"Agent {getattr(self, 'name', '')}: input blocked by guardrail: {_in_err}") yield f"[Input blocked by guardrail: {_in_err}]" - return + return "failed" prompt = _in_prompt # Same pre-turn checks as chat()/achat(): a streamed turn must not @@ -5479,7 +5502,7 @@ def _start_stream_turn(self, prompt: str, **kwargs) -> Generator[str, None, None )) if self._hook_runner.is_blocked(hook_results): logging.warning(f"Agent {self.name} execution blocked by BEFORE_AGENT hook") - return + return "failed" for res in hook_results: if res.output and res.output.modified_input and "prompt" in res.output.modified_input: prompt = res.output.modified_input["prompt"] @@ -6299,6 +6322,9 @@ def _stream_fallback_chat(self, prompt: str, kwargs: Dict[str, Any], streaming_e else: fallback_kwargs.pop('stream', None) + # The streamed turn already records agent_start/agent_end; keep chat() + # from recording the same turn a second time. + _traced = _stream_traced_agent.set(self) try: return self.chat(prompt, **fallback_kwargs) except Exception as fallback_error: @@ -6311,6 +6337,8 @@ def _stream_fallback_chat(self, prompt: str, kwargs: Dict[str, Any], streaming_e except Exception: pass raise fallback_error from streaming_error + finally: + _stream_traced_agent.reset(_traced) def _create_llm_summarize_function(self): """ diff --git a/src/praisonai-agents/tests/unit/trace/test_unnamed_agent_trace.py b/src/praisonai-agents/tests/unit/trace/test_unnamed_agent_trace.py index f701299270..c3c486bab3 100644 --- a/src/praisonai-agents/tests/unit/trace/test_unnamed_agent_trace.py +++ b/src/praisonai-agents/tests/unit/trace/test_unnamed_agent_trace.py @@ -172,3 +172,89 @@ async def fake_dispatch(*args, **kwargs): assert d.called and s.get_events() == [] finally: reset_context_emitter(token) + + +# --- streamed turns (start(stream=True)) --------------------------------------- + +def _fake_stream(self, prompt, **kwargs): + yield "Solar " + yield "sails." + + +def test_streamed_turn_records_agent_start_and_end(sink): + from praisonaiagents import Agent + + agent = Agent(instructions="x") + with patch.object(type(agent), "_start_stream_turn", _fake_stream): + text = "".join(agent._start_stream("hi")) + assert text == "Solar sails." + kinds = [(e.event_type.value, e.node_id, e.data.get("status")) for e in sink.get_events()] + assert kinds == [("agent_start", "agent:Agent", None), ("agent_end", "agent:Agent", "completed")] + + +def test_streamed_turn_stopped_early_is_cancelled(sink): + from praisonaiagents import Agent + + agent = Agent(name="Writer", instructions="x") + with patch.object(type(agent), "_start_stream_turn", _fake_stream): + gen = agent._start_stream("hi") + next(gen) + gen.close() # the page stopped reading + assert _of(sink, "agent_end")[0].data == {"status": "cancelled"} + + +def test_streamed_turn_that_raises_is_failed(sink): + from praisonaiagents import Agent + + def boom(self, prompt, **kwargs): + yield "x" + raise RuntimeError("provider down") + + agent = Agent(name="Writer", instructions="x") + with patch.object(type(agent), "_start_stream_turn", boom): + with pytest.raises(RuntimeError): + list(agent._start_stream("hi")) + assert _of(sink, "agent_end")[0].data == {"status": "failed"} + + +def test_streamed_turn_records_nothing_when_tracing_is_off(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-not-used") + from praisonaiagents import Agent + + s = ContextListSink() + token = set_context_emitter(ContextTraceEmitter(sink=s, session_id="t", enabled=False)) + try: + agent = Agent(name="Writer", instructions="x") + with patch.object(type(agent), "_start_stream_turn", _fake_stream): + assert "".join(agent._start_stream("hi")) == "Solar sails." + assert s.get_events() == [] + finally: + reset_context_emitter(token) + + +def test_streamed_turn_blocked_by_guardrail_is_failed(sink): + from praisonaiagents import Agent + + agent = Agent(name="Writer", instructions="x") + with patch.object(agent, "_validate_input_with_guardrail", return_value=(False, None, "no")): + assert "blocked" in "".join(agent._start_stream("hi")) + assert _of(sink, "agent_end")[0].data == {"status": "failed"} + + +def test_streamed_turn_fallback_records_the_turn_once(sink): + from praisonaiagents import Agent + + def falls_back(self, prompt, **kwargs): + yield self._stream_fallback_chat(prompt, kwargs, RuntimeError("stream broke")) + + agent = Agent(name="Writer", instructions="x") + with patch.object(type(agent), "_start_stream_impl", falls_back), \ + patch.object(agent, "_chat_completion", return_value=_fake_completion()): + assert "".join(agent._start_stream("hi")) == "Solar sails ride sunlight." + kinds = [e.event_type.value for e in sink.get_events() if e.event_type.value.startswith("agent_")] + assert kinds == ["agent_start", "agent_end"] + assert _of(sink, "agent_end")[0].data == {"status": "completed"} + # chat() outside a stream still records its own turn. + with patch.object(agent, "_chat_completion", return_value=_fake_completion()): + agent.chat("again") + assert len(_of(sink, "agent_start")) == 2