letta-nightly 0.8.15.dev20250720104313__py3-none-any.whl → 0.8.16.dev20250721104533__py3-none-any.whl

letta/agents/letta_agent.py

@@ -1,4 +1,3 @@
-import asyncio
 import json
 import uuid
 from collections.abc import AsyncGenerator
@@ -44,6 +43,7 @@ from letta.schemas.llm_config import LLMConfig
 from letta.schemas.message import Message, MessageCreate
 from letta.schemas.openai.chat_completion_response import ToolCall, UsageStatistics
 from letta.schemas.provider_trace import ProviderTraceCreate
+from letta.schemas.step import StepProgression
 from letta.schemas.tool_execution_result import ToolExecutionResult
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User
@@ -239,100 +239,164 @@ class LettaAgent(BaseAgent):
             agent_step_span = tracer.start_span("agent_step", start_time=step_start)
             agent_step_span.set_attributes({"step_id": step_id})
 
-            request_data, response_data, current_in_context_messages, new_in_context_messages, valid_tool_names = (
-                await self._build_and_request_from_llm(
-                    current_in_context_messages,
-                    new_in_context_messages,
-                    agent_state,
-                    llm_client,
-                    tool_rules_solver,
-                    agent_step_span,
+            step_progression = StepProgression.START
+            should_continue = False
+            try:
+                request_data, response_data, current_in_context_messages, new_in_context_messages, valid_tool_names = (
+                    await self._build_and_request_from_llm(
+                        current_in_context_messages,
+                        new_in_context_messages,
+                        agent_state,
+                        llm_client,
+                        tool_rules_solver,
+                        agent_step_span,
+                    )
                 )
-            )
-            in_context_messages = current_in_context_messages + new_in_context_messages
-
-            log_event("agent.stream_no_tokens.llm_response.received")  # [3^]
-
-            response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
-
-            # update usage
-            usage.step_count += 1
-            usage.completion_tokens += response.usage.completion_tokens
-            usage.prompt_tokens += response.usage.prompt_tokens
-            usage.total_tokens += response.usage.total_tokens
-            MetricRegistry().message_output_tokens.record(
-                response.usage.completion_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
-            )
+                in_context_messages = current_in_context_messages + new_in_context_messages
 
-            if not response.choices[0].message.tool_calls:
-                # TODO: make into a real error
-                raise ValueError("No tool calls found in response, model must make a tool call")
-            tool_call = response.choices[0].message.tool_calls[0]
-            if response.choices[0].message.reasoning_content:
-                reasoning = [
-                    ReasoningContent(
-                        reasoning=response.choices[0].message.reasoning_content,
-                        is_native=True,
-                        signature=response.choices[0].message.reasoning_content_signature,
-                    )
-                ]
-            elif response.choices[0].message.omitted_reasoning_content:
-                reasoning = [OmittedReasoningContent()]
-            elif response.choices[0].message.content:
-                reasoning = [TextContent(text=response.choices[0].message.content)]  # reasoning placed into content for legacy reasons
-            else:
-                self.logger.info("No reasoning content found.")
-                reasoning = None
-
-            persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
-                tool_call,
-                valid_tool_names,
-                agent_state,
-                tool_rules_solver,
-                response.usage,
-                reasoning_content=reasoning,
-                step_id=step_id,
-                initial_messages=initial_messages,
-                agent_step_span=agent_step_span,
-                is_final_step=(i == max_steps - 1),
-            )
+                step_progression = StepProgression.RESPONSE_RECEIVED
+                log_event("agent.stream_no_tokens.llm_response.received")  # [3^]
 
-            # TODO (cliandy): handle message contexts with larger refactor and dedupe logic
-            new_message_idx = len(initial_messages) if initial_messages else 0
-            self.response_messages.extend(persisted_messages[new_message_idx:])
-            new_in_context_messages.extend(persisted_messages[new_message_idx:])
-            initial_messages = None
-            log_event("agent.stream_no_tokens.llm_response.processed")  # [4^]
+                response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
 
-            # log step time
-            now = get_utc_timestamp_ns()
-            step_ns = now - step_start
-            agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
-            agent_step_span.end()
+                # update usage
+                usage.step_count += 1
+                usage.completion_tokens += response.usage.completion_tokens
+                usage.prompt_tokens += response.usage.prompt_tokens
+                usage.total_tokens += response.usage.total_tokens
+                MetricRegistry().message_output_tokens.record(
+                    response.usage.completion_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
+                )
 
-            # Log LLM Trace
-            await self.telemetry_manager.create_provider_trace_async(
-                actor=self.actor,
-                provider_trace_create=ProviderTraceCreate(
-                    request_json=request_data,
-                    response_json=response_data,
+                if not response.choices[0].message.tool_calls:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.no_tool_call.value)
+                    raise ValueError("No tool calls found in response, model must make a tool call")
+                tool_call = response.choices[0].message.tool_calls[0]
+                if response.choices[0].message.reasoning_content:
+                    reasoning = [
+                        ReasoningContent(
+                            reasoning=response.choices[0].message.reasoning_content,
+                            is_native=True,
+                            signature=response.choices[0].message.reasoning_content_signature,
+                        )
+                    ]
+                elif response.choices[0].message.omitted_reasoning_content:
+                    reasoning = [OmittedReasoningContent()]
+                elif response.choices[0].message.content:
+                    reasoning = [TextContent(text=response.choices[0].message.content)]  # reasoning placed into content for legacy reasons
+                else:
+                    self.logger.info("No reasoning content found.")
+                    reasoning = None
+
+                persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
+                    tool_call,
+                    valid_tool_names,
+                    agent_state,
+                    tool_rules_solver,
+                    response.usage,
+                    reasoning_content=reasoning,
                     step_id=step_id,
-                    organization_id=self.actor.organization_id,
-                ),
-            )
+                    initial_messages=initial_messages,
+                    agent_step_span=agent_step_span,
+                    is_final_step=(i == max_steps - 1),
+                )
+                step_progression = StepProgression.STEP_LOGGED
+
+                # TODO (cliandy): handle message contexts with larger refactor and dedupe logic
+                new_message_idx = len(initial_messages) if initial_messages else 0
+                self.response_messages.extend(persisted_messages[new_message_idx:])
+                new_in_context_messages.extend(persisted_messages[new_message_idx:])
+                initial_messages = None
+                log_event("agent.stream_no_tokens.llm_response.processed")  # [4^]
+
+                # log step time
+                now = get_utc_timestamp_ns()
+                step_ns = now - step_start
+                agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
+                agent_step_span.end()
+
+                # Log LLM Trace
+                await self.telemetry_manager.create_provider_trace_async(
+                    actor=self.actor,
+                    provider_trace_create=ProviderTraceCreate(
+                        request_json=request_data,
+                        response_json=response_data,
+                        step_id=step_id,
+                        organization_id=self.actor.organization_id,
+                    ),
+                )
+                step_progression = StepProgression.LOGGED_TRACE
 
-            # stream step
-            # TODO: improve TTFT
-            filter_user_messages = [m for m in persisted_messages if m.role != "user"]
-            letta_messages = Message.to_letta_messages_from_list(
-                filter_user_messages, use_assistant_message=use_assistant_message, reverse=False
-            )
+                # stream step
+                # TODO: improve TTFT
+                filter_user_messages = [m for m in persisted_messages if m.role != "user"]
+                letta_messages = Message.to_letta_messages_from_list(
+                    filter_user_messages, use_assistant_message=use_assistant_message, reverse=False
+                )
 
-            for message in letta_messages:
-                if include_return_message_types is None or message.message_type in include_return_message_types:
-                    yield f"data: {message.model_dump_json()}\n\n"
+                for message in letta_messages:
+                    if include_return_message_types is None or message.message_type in include_return_message_types:
+                        yield f"data: {message.model_dump_json()}\n\n"
 
-            MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                step_progression = StepProgression.FINISHED
+            except Exception as e:
+                # Handle any unexpected errors during step processing
+                self.logger.error(f"Error during step processing: {e}")
+
+                # This indicates we failed after we decided to stop stepping, which indicates a bug with our flow.
+                if not stop_reason:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                elif stop_reason.stop_reason in (StopReasonType.end_turn, StopReasonType.max_steps, StopReasonType.tool_rule):
+                    self.logger.error("Error occurred during step processing, with valid stop reason: %s", stop_reason.stop_reason)
+                elif stop_reason.stop_reason not in (StopReasonType.no_tool_call, StopReasonType.invalid_tool_call):
+                    raise ValueError(f"Invalid Stop Reason: {stop_reason}")
+
+                # Send error stop reason to client and re-raise
+                yield f"data: {stop_reason.model_dump_json()}\n\n", 500
+                raise
+
+            # Update step if it needs to be updated
+            finally:
+                if settings.track_stop_reason:
+                    self.logger.info("Running final update. Step Progression: %s", step_progression)
+                    try:
+                        if step_progression < StepProgression.STEP_LOGGED:
+                            await self.step_manager.log_step_async(
+                                actor=self.actor,
+                                agent_id=agent_state.id,
+                                provider_name=agent_state.llm_config.model_endpoint_type,
+                                provider_category=agent_state.llm_config.provider_category or "base",
+                                model=agent_state.llm_config.model,
+                                model_endpoint=agent_state.llm_config.model_endpoint,
+                                context_window_limit=agent_state.llm_config.context_window,
+                                usage=UsageStatistics(completion_tokens=0, prompt_tokens=0, total_tokens=0),
+                                provider_id=None,
+                                job_id=self.current_run_id if self.current_run_id else None,
+                                step_id=step_id,
+                                project_id=agent_state.project_id,
+                                stop_reason=stop_reason,
+                            )
+                        if step_progression <= StepProgression.RESPONSE_RECEIVED:
+                            # TODO (cliandy): persist response if we get it back
+                            if settings.track_errored_messages:
+                                for message in initial_messages:
+                                    message.is_err = True
+                                    message.step_id = step_id
+                                await self.message_manager.create_many_messages_async(initial_messages, actor=self.actor)
+                        elif step_progression <= StepProgression.LOGGED_TRACE:
+                            if stop_reason is None:
+                                self.logger.error("Error in step after logging step")
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        elif step_progression == StepProgression.FINISHED and not should_continue:
+                            if stop_reason is None:
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.end_turn.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        else:
+                            self.logger.error("Invalid StepProgression value")
+                    except Exception as e:
+                        self.logger.error("Failed to update step: %s", e)
 
             if not should_continue:
                 break
@@ -397,17 +461,6 @@ class LettaAgent(BaseAgent):
         stop_reason = None
         usage = LettaUsageStatistics()
         for i in range(max_steps):
-            # Check for job cancellation at the start of each step
-            if await self._check_run_cancellation():
-                stop_reason = LettaStopReason(stop_reason=StopReasonType.cancelled.value)
-                logger.info(f"Agent execution cancelled for run {self.current_run_id}")
-                break
-
-            step_id = generate_step_id()
-            step_start = get_utc_timestamp_ns()
-            agent_step_span = tracer.start_span("agent_step", start_time=step_start)
-            agent_step_span.set_attributes({"step_id": step_id})
-
             # If dry run, build request data and return it without making LLM call
             if dry_run:
                 request_data, valid_tool_names = await self._create_llm_request_data_async(
@@ -418,84 +471,159 @@ class LettaAgent(BaseAgent):
                 )
                 return request_data
 
-            request_data, response_data, current_in_context_messages, new_in_context_messages, valid_tool_names = (
-                await self._build_and_request_from_llm(
-                    current_in_context_messages, new_in_context_messages, agent_state, llm_client, tool_rules_solver, agent_step_span
-                )
-            )
-            in_context_messages = current_in_context_messages + new_in_context_messages
-
-            log_event("agent.step.llm_response.received")  # [3^]
+            # Check for job cancellation at the start of each step
+            if await self._check_run_cancellation():
+                stop_reason = LettaStopReason(stop_reason=StopReasonType.cancelled.value)
+                logger.info(f"Agent execution cancelled for run {self.current_run_id}")
+                break
 
-            response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
+            step_id = generate_step_id()
+            step_start = get_utc_timestamp_ns()
+            agent_step_span = tracer.start_span("agent_step", start_time=step_start)
+            agent_step_span.set_attributes({"step_id": step_id})
 
-            usage.step_count += 1
-            usage.completion_tokens += response.usage.completion_tokens
-            usage.prompt_tokens += response.usage.prompt_tokens
-            usage.total_tokens += response.usage.total_tokens
-            usage.run_ids = [run_id] if run_id else None
-            MetricRegistry().message_output_tokens.record(
-                response.usage.completion_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
-            )
+            step_progression = StepProgression.START
+            should_continue = False
 
-            if not response.choices[0].message.tool_calls:
-                # TODO: make into a real error
-                raise ValueError("No tool calls found in response, model must make a tool call")
-            tool_call = response.choices[0].message.tool_calls[0]
-            if response.choices[0].message.reasoning_content:
-                reasoning = [
-                    ReasoningContent(
-                        reasoning=response.choices[0].message.reasoning_content,
-                        is_native=True,
-                        signature=response.choices[0].message.reasoning_content_signature,
+            try:
+                request_data, response_data, current_in_context_messages, new_in_context_messages, valid_tool_names = (
+                    await self._build_and_request_from_llm(
+                        current_in_context_messages, new_in_context_messages, agent_state, llm_client, tool_rules_solver, agent_step_span
                     )
-                ]
-            elif response.choices[0].message.content:
-                reasoning = [TextContent(text=response.choices[0].message.content)]  # reasoning placed into content for legacy reasons
-            elif response.choices[0].message.omitted_reasoning_content:
-                reasoning = [OmittedReasoningContent()]
-            else:
-                self.logger.info("No reasoning content found.")
-                reasoning = None
-
-            persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
-                tool_call,
-                valid_tool_names,
-                agent_state,
-                tool_rules_solver,
-                response.usage,
-                reasoning_content=reasoning,
-                step_id=step_id,
-                initial_messages=initial_messages,
-                agent_step_span=agent_step_span,
-                is_final_step=(i == max_steps - 1),
-                run_id=run_id,
-            )
-            new_message_idx = len(initial_messages) if initial_messages else 0
-            self.response_messages.extend(persisted_messages[new_message_idx:])
-            new_in_context_messages.extend(persisted_messages[new_message_idx:])
+                )
+                in_context_messages = current_in_context_messages + new_in_context_messages
 
-            initial_messages = None
-            log_event("agent.step.llm_response.processed")  # [4^]
+                step_progression = StepProgression.RESPONSE_RECEIVED
+                log_event("agent.step.llm_response.received")  # [3^]
 
-            # log step time
-            now = get_utc_timestamp_ns()
-            step_ns = now - step_start
-            agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
-            agent_step_span.end()
+                response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
 
-            # Log LLM Trace
-            await self.telemetry_manager.create_provider_trace_async(
-                actor=self.actor,
-                provider_trace_create=ProviderTraceCreate(
-                    request_json=request_data,
-                    response_json=response_data,
+                usage.step_count += 1
+                usage.completion_tokens += response.usage.completion_tokens
+                usage.prompt_tokens += response.usage.prompt_tokens
+                usage.total_tokens += response.usage.total_tokens
+                usage.run_ids = [run_id] if run_id else None
+                MetricRegistry().message_output_tokens.record(
+                    response.usage.completion_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
+                )
+
+                if not response.choices[0].message.tool_calls:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.no_tool_call.value)
+                    raise ValueError("No tool calls found in response, model must make a tool call")
+                tool_call = response.choices[0].message.tool_calls[0]
+                if response.choices[0].message.reasoning_content:
+                    reasoning = [
+                        ReasoningContent(
+                            reasoning=response.choices[0].message.reasoning_content,
+                            is_native=True,
+                            signature=response.choices[0].message.reasoning_content_signature,
+                        )
+                    ]
+                elif response.choices[0].message.content:
+                    reasoning = [TextContent(text=response.choices[0].message.content)]  # reasoning placed into content for legacy reasons
+                elif response.choices[0].message.omitted_reasoning_content:
+                    reasoning = [OmittedReasoningContent()]
+                else:
+                    self.logger.info("No reasoning content found.")
+                    reasoning = None
+
+                persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
+                    tool_call,
+                    valid_tool_names,
+                    agent_state,
+                    tool_rules_solver,
+                    response.usage,
+                    reasoning_content=reasoning,
                     step_id=step_id,
-                    organization_id=self.actor.organization_id,
-                ),
-            )
+                    initial_messages=initial_messages,
+                    agent_step_span=agent_step_span,
+                    is_final_step=(i == max_steps - 1),
+                    run_id=run_id,
+                )
+                step_progression = StepProgression.STEP_LOGGED
+
+                new_message_idx = len(initial_messages) if initial_messages else 0
+                self.response_messages.extend(persisted_messages[new_message_idx:])
+                new_in_context_messages.extend(persisted_messages[new_message_idx:])
+
+                initial_messages = None
+                log_event("agent.step.llm_response.processed")  # [4^]
+
+                # log step time
+                now = get_utc_timestamp_ns()
+                step_ns = now - step_start
+                agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
+                agent_step_span.end()
+
+                # Log LLM Trace
+                await self.telemetry_manager.create_provider_trace_async(
+                    actor=self.actor,
+                    provider_trace_create=ProviderTraceCreate(
+                        request_json=request_data,
+                        response_json=response_data,
+                        step_id=step_id,
+                        organization_id=self.actor.organization_id,
+                    ),
+                )
 
-            MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                step_progression = StepProgression.LOGGED_TRACE
+                MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                step_progression = StepProgression.FINISHED
+
+            except Exception as e:
+                # Handle any unexpected errors during step processing
+                self.logger.error(f"Error during step processing: {e}")
+
+                # This indicates we failed after we decided to stop stepping, which indicates a bug with our flow.
+                if not stop_reason:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                elif stop_reason.stop_reason in (StopReasonType.end_turn, StopReasonType.max_steps, StopReasonType.tool_rule):
+                    self.logger.error("Error occurred during step processing, with valid stop reason: %s", stop_reason.stop_reason)
+                elif stop_reason.stop_reason not in (StopReasonType.no_tool_call, StopReasonType.invalid_tool_call):
+                    raise ValueError(f"Invalid Stop Reason: {stop_reason}")
+                raise
+
+                # Update step if it needs to be updated
+            finally:
+                if settings.track_stop_reason:
+                    self.logger.info("Running final update. Step Progression: %s", step_progression)
+                    try:
+                        if step_progression < StepProgression.STEP_LOGGED:
+                            await self.step_manager.log_step_async(
+                                actor=self.actor,
+                                agent_id=agent_state.id,
+                                provider_name=agent_state.llm_config.model_endpoint_type,
+                                provider_category=agent_state.llm_config.provider_category or "base",
+                                model=agent_state.llm_config.model,
+                                model_endpoint=agent_state.llm_config.model_endpoint,
+                                context_window_limit=agent_state.llm_config.context_window,
+                                usage=UsageStatistics(completion_tokens=0, prompt_tokens=0, total_tokens=0),
+                                provider_id=None,
+                                job_id=self.current_run_id if self.current_run_id else None,
+                                step_id=step_id,
+                                project_id=agent_state.project_id,
+                                stop_reason=stop_reason,
+                            )
+                        if step_progression <= StepProgression.RESPONSE_RECEIVED:
+                            # TODO (cliandy): persist response if we get it back
+                            if settings.track_errored_messages:
+                                for message in initial_messages:
+                                    message.is_err = True
+                                    message.step_id = step_id
+                                await self.message_manager.create_many_messages_async(initial_messages, actor=self.actor)
+                        elif step_progression <= StepProgression.LOGGED_TRACE:
+                            if stop_reason is None:
+                                self.logger.error("Error in step after logging step")
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        elif step_progression == StepProgression.FINISHED and not should_continue:
+                            if stop_reason is None:
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.end_turn.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        else:
+                            self.logger.error("Invalid StepProgression value")
+                    except Exception as e:
+                        self.logger.error("Failed to update step: %s", e)
 
             if not should_continue:
                 break
@@ -577,6 +705,7 @@ class LettaAgent(BaseAgent):
             request_span.set_attributes({f"llm_config.{k}": v for k, v in agent_state.llm_config.model_dump().items() if v is not None})
 
         for i in range(max_steps):
+            step_id = generate_step_id()
             # Check for job cancellation at the start of each step
             if await self._check_run_cancellation():
                 stop_reason = LettaStopReason(stop_reason=StopReasonType.cancelled.value)
@@ -584,163 +713,230 @@ class LettaAgent(BaseAgent):
                 yield f"data: {stop_reason.model_dump_json()}\n\n"
                 break
 
-            step_id = generate_step_id()
             step_start = get_utc_timestamp_ns()
             agent_step_span = tracer.start_span("agent_step", start_time=step_start)
             agent_step_span.set_attributes({"step_id": step_id})
 
-            (
-                request_data,
-                stream,
-                current_in_context_messages,
-                new_in_context_messages,
-                valid_tool_names,
-                provider_request_start_timestamp_ns,
-            ) = await self._build_and_request_from_llm_streaming(
-                first_chunk,
-                agent_step_span,
-                request_start_timestamp_ns,
-                current_in_context_messages,
-                new_in_context_messages,
-                agent_state,
-                llm_client,
-                tool_rules_solver,
-            )
-            log_event("agent.stream.llm_response.received")  # [3^]
-
-            # TODO: THIS IS INCREDIBLY UGLY
-            # TODO: THERE ARE MULTIPLE COPIES OF THE LLM_CONFIG EVERYWHERE THAT ARE GETTING MANIPULATED
-            if agent_state.llm_config.model_endpoint_type in [ProviderType.anthropic, ProviderType.bedrock]:
-                interface = AnthropicStreamingInterface(
-                    use_assistant_message=use_assistant_message,
-                    put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
-                )
-            elif agent_state.llm_config.model_endpoint_type == ProviderType.openai:
-                interface = OpenAIStreamingInterface(
-                    use_assistant_message=use_assistant_message,
-                    put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
-                )
-            else:
-                raise ValueError(f"Streaming not supported for {agent_state.llm_config}")
-
-            async for chunk in interface.process(
-                stream,
-                ttft_span=request_span,
-                provider_request_start_timestamp_ns=provider_request_start_timestamp_ns,
-            ):
-                # Measure time to first token
-                if first_chunk and request_span is not None:
-                    now = get_utc_timestamp_ns()
-                    ttft_ns = now - request_start_timestamp_ns
-                    request_span.add_event(name="time_to_first_token_ms", attributes={"ttft_ms": ns_to_ms(ttft_ns)})
-                    metric_attributes = get_ctx_attributes()
-                    metric_attributes["model.name"] = agent_state.llm_config.model
-                    MetricRegistry().ttft_ms_histogram.record(ns_to_ms(ttft_ns), metric_attributes)
-                    first_chunk = False
-
-                if include_return_message_types is None or chunk.message_type in include_return_message_types:
-                    # filter down returned data
-                    yield f"data: {chunk.model_dump_json()}\n\n"
-
-            stream_end_time_ns = get_utc_timestamp_ns()
-
-            # update usage
-            usage.step_count += 1
-            usage.completion_tokens += interface.output_tokens
-            usage.prompt_tokens += interface.input_tokens
-            usage.total_tokens += interface.input_tokens + interface.output_tokens
-            MetricRegistry().message_output_tokens.record(
-                interface.output_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
-            )
-
-            # log LLM request time
-            llm_request_ms = ns_to_ms(stream_end_time_ns - provider_request_start_timestamp_ns)
-            agent_step_span.add_event(name="llm_request_ms", attributes={"duration_ms": llm_request_ms})
-            MetricRegistry().llm_execution_time_ms_histogram.record(
-                llm_request_ms,
-                dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model}),
-            )
-
-            # Process resulting stream content
+            step_progression = StepProgression.START
+            should_continue = False
             try:
-                tool_call = interface.get_tool_call_object()
-            except ValueError as e:
-                stop_reason = LettaStopReason(stop_reason=StopReasonType.no_tool_call.value)
-                yield f"data: {stop_reason.model_dump_json()}\n\n"
-                raise e
-            except Exception as e:
-                stop_reason = LettaStopReason(stop_reason=StopReasonType.invalid_tool_call.value)
-                yield f"data: {stop_reason.model_dump_json()}\n\n"
-                raise e
-            reasoning_content = interface.get_reasoning_content()
-            persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
-                tool_call,
-                valid_tool_names,
-                agent_state,
-                tool_rules_solver,
-                UsageStatistics(
-                    completion_tokens=interface.output_tokens,
-                    prompt_tokens=interface.input_tokens,
-                    total_tokens=interface.input_tokens + interface.output_tokens,
-                ),
-                reasoning_content=reasoning_content,
-                pre_computed_assistant_message_id=interface.letta_message_id,
-                step_id=step_id,
-                initial_messages=initial_messages,
-                agent_step_span=agent_step_span,
-                is_final_step=(i == max_steps - 1),
-            )
-            new_message_idx = len(initial_messages) if initial_messages else 0
-            self.response_messages.extend(persisted_messages[new_message_idx:])
-            new_in_context_messages.extend(persisted_messages[new_message_idx:])
+                (
+                    request_data,
+                    stream,
+                    current_in_context_messages,
+                    new_in_context_messages,
+                    valid_tool_names,
+                    provider_request_start_timestamp_ns,
+                ) = await self._build_and_request_from_llm_streaming(
+                    first_chunk,
+                    agent_step_span,
+                    request_start_timestamp_ns,
+                    current_in_context_messages,
+                    new_in_context_messages,
+                    agent_state,
+                    llm_client,
+                    tool_rules_solver,
+                )
 
-            initial_messages = None
+                step_progression = StepProgression.STREAM_RECEIVED
+                log_event("agent.stream.llm_response.received")  # [3^]
 
-            # log total step time
-            now = get_utc_timestamp_ns()
-            step_ns = now - step_start
-            agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
-            agent_step_span.end()
+                # TODO: THIS IS INCREDIBLY UGLY
+                # TODO: THERE ARE MULTIPLE COPIES OF THE LLM_CONFIG EVERYWHERE THAT ARE GETTING MANIPULATED
+                if agent_state.llm_config.model_endpoint_type in [ProviderType.anthropic, ProviderType.bedrock]:
+                    interface = AnthropicStreamingInterface(
+                        use_assistant_message=use_assistant_message,
+                        put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
+                    )
+                elif agent_state.llm_config.model_endpoint_type == ProviderType.openai:
+                    interface = OpenAIStreamingInterface(
+                        use_assistant_message=use_assistant_message,
+                        put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
+                    )
+                else:
+                    raise ValueError(f"Streaming not supported for {agent_state.llm_config}")
+
+                async for chunk in interface.process(
+                    stream,
+                    ttft_span=request_span,
+                    provider_request_start_timestamp_ns=provider_request_start_timestamp_ns,
+                ):
+                    # Measure time to first token
+                    if first_chunk and request_span is not None:
+                        now = get_utc_timestamp_ns()
+                        ttft_ns = now - request_start_timestamp_ns
+                        request_span.add_event(name="time_to_first_token_ms", attributes={"ttft_ms": ns_to_ms(ttft_ns)})
+                        metric_attributes = get_ctx_attributes()
+                        metric_attributes["model.name"] = agent_state.llm_config.model
+                        MetricRegistry().ttft_ms_histogram.record(ns_to_ms(ttft_ns), metric_attributes)
+                        first_chunk = False
+
+                    if include_return_message_types is None or chunk.message_type in include_return_message_types:
+                        # filter down returned data
+                        yield f"data: {chunk.model_dump_json()}\n\n"
+
+                stream_end_time_ns = get_utc_timestamp_ns()
+
+                # update usage
+                usage.step_count += 1
+                usage.completion_tokens += interface.output_tokens
+                usage.prompt_tokens += interface.input_tokens
+                usage.total_tokens += interface.input_tokens + interface.output_tokens
+                MetricRegistry().message_output_tokens.record(
+                    interface.output_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
+                )
 
-            # TODO (cliandy): the stream POST request span has ended at this point, we should tie this to the stream
-            # log_event("agent.stream.llm_response.processed") # [4^]
+                # log LLM request time
+                llm_request_ms = ns_to_ms(stream_end_time_ns - provider_request_start_timestamp_ns)
+                agent_step_span.add_event(name="llm_request_ms", attributes={"duration_ms": llm_request_ms})
+                MetricRegistry().llm_execution_time_ms_histogram.record(
+                    llm_request_ms,
+                    dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model}),
+                )
 
-            # Log LLM Trace
-            # TODO (cliandy): we are piecing together the streamed response here. Content here does not match the actual response schema.
-            await self.telemetry_manager.create_provider_trace_async(
-                actor=self.actor,
-                provider_trace_create=ProviderTraceCreate(
-                    request_json=request_data,
-                    response_json={
-                        "content": {
-                            "tool_call": tool_call.model_dump_json(),
-                            "reasoning": [content.model_dump_json() for content in reasoning_content],
-                        },
-                        "id": interface.message_id,
-                        "model": interface.model,
-                        "role": "assistant",
-                        # "stop_reason": "",
-                        # "stop_sequence": None,
-                        "type": "message",
-                        "usage": {"input_tokens": interface.input_tokens, "output_tokens": interface.output_tokens},
-                    },
+                # Process resulting stream content
+                try:
+                    tool_call = interface.get_tool_call_object()
+                except ValueError as e:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.no_tool_call.value)
+                    raise e
+                except Exception as e:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.invalid_tool_call.value)
+                    raise e
+                reasoning_content = interface.get_reasoning_content()
+                persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
+                    tool_call,
+                    valid_tool_names,
+                    agent_state,
+                    tool_rules_solver,
+                    UsageStatistics(
+                        completion_tokens=interface.output_tokens,
+                        prompt_tokens=interface.input_tokens,
+                        total_tokens=interface.input_tokens + interface.output_tokens,
+                    ),
+                    reasoning_content=reasoning_content,
+                    pre_computed_assistant_message_id=interface.letta_message_id,
                     step_id=step_id,
-                    organization_id=self.actor.organization_id,
-                ),
-            )
+                    initial_messages=initial_messages,
+                    agent_step_span=agent_step_span,
+                    is_final_step=(i == max_steps - 1),
+                )
+                step_progression = StepProgression.STEP_LOGGED
+
+                new_message_idx = len(initial_messages) if initial_messages else 0
+                self.response_messages.extend(persisted_messages[new_message_idx:])
+                new_in_context_messages.extend(persisted_messages[new_message_idx:])
+
+                initial_messages = None
+
+                # log total step time
+                now = get_utc_timestamp_ns()
+                step_ns = now - step_start
+                agent_step_span.add_event(name="step_ms", attributes={"duration_ms": ns_to_ms(step_ns)})
+                agent_step_span.end()
+
+                # TODO (cliandy): the stream POST request span has ended at this point, we should tie this to the stream
+                # log_event("agent.stream.llm_response.processed") # [4^]
+
+                # Log LLM Trace
+                # We are piecing together the streamed response here.
+                # Content here does not match the actual response schema as streams come in chunks.
+                await self.telemetry_manager.create_provider_trace_async(
+                    actor=self.actor,
+                    provider_trace_create=ProviderTraceCreate(
+                        request_json=request_data,
+                        response_json={
+                            "content": {
+                                "tool_call": tool_call.model_dump_json(),
+                                "reasoning": [content.model_dump_json() for content in reasoning_content],
+                            },
+                            "id": interface.message_id,
+                            "model": interface.model,
+                            "role": "assistant",
+                            # "stop_reason": "",
+                            # "stop_sequence": None,
+                            "type": "message",
+                            "usage": {
+                                "input_tokens": interface.input_tokens,
+                                "output_tokens": interface.output_tokens,
+                            },
+                        },
+                        step_id=step_id,
+                        organization_id=self.actor.organization_id,
+                    ),
+                )
+                step_progression = StepProgression.LOGGED_TRACE
 
-            tool_return = [msg for msg in persisted_messages if msg.role == "tool"][-1].to_letta_messages()[0]
-            if not (use_assistant_message and tool_return.name == "send_message"):
-                # Apply message type filtering if specified
-                if include_return_message_types is None or tool_return.message_type in include_return_message_types:
-                    yield f"data: {tool_return.model_dump_json()}\n\n"
+                # yields tool response as this is handled from Letta and not the response from the LLM provider
+                tool_return = [msg for msg in persisted_messages if msg.role == "tool"][-1].to_letta_messages()[0]
+                if not (use_assistant_message and tool_return.name == "send_message"):
+                    # Apply message type filtering if specified
+                    if include_return_message_types is None or tool_return.message_type in include_return_message_types:
+                        yield f"data: {tool_return.model_dump_json()}\n\n"
 
-            # TODO (cliandy): consolidate and expand with trace
-            MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                # TODO (cliandy): consolidate and expand with trace
+                MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
+                step_progression = StepProgression.FINISHED
+
+            except Exception as e:
+                # Handle any unexpected errors during step processing
+                self.logger.error(f"Error during step processing: {e}")
+
+                # This indicates we failed after we decided to stop stepping, which indicates a bug with our flow.
+                if not stop_reason:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                elif stop_reason.stop_reason in (StopReasonType.end_turn, StopReasonType.max_steps, StopReasonType.tool_rule):
+                    self.logger.error("Error occurred during step processing, with valid stop reason: %s", stop_reason.stop_reason)
+                elif stop_reason.stop_reason not in (StopReasonType.no_tool_call, StopReasonType.invalid_tool_call):
+                    raise ValueError(f"Invalid Stop Reason: {stop_reason}")
+
+                # Send error stop reason to client and re-raise with expected response code
+                yield f"data: {stop_reason.model_dump_json()}\n\n", 500
+                raise
+
+            # Update step if it needs to be updated
+            finally:
+                if settings.track_stop_reason:
+                    self.logger.info("Running final update. Step Progression: %s", step_progression)
+                    try:
+                        if step_progression < StepProgression.STEP_LOGGED:
+                            await self.step_manager.log_step_async(
+                                actor=self.actor,
+                                agent_id=agent_state.id,
+                                provider_name=agent_state.llm_config.model_endpoint_type,
+                                provider_category=agent_state.llm_config.provider_category or "base",
+                                model=agent_state.llm_config.model,
+                                model_endpoint=agent_state.llm_config.model_endpoint,
+                                context_window_limit=agent_state.llm_config.context_window,
+                                usage=UsageStatistics(completion_tokens=0, prompt_tokens=0, total_tokens=0),
+                                provider_id=None,
+                                job_id=self.current_run_id if self.current_run_id else None,
+                                step_id=step_id,
+                                project_id=agent_state.project_id,
+                                stop_reason=stop_reason,
+                            )
+                        if step_progression <= StepProgression.STREAM_RECEIVED:
+                            if first_chunk and settings.track_errored_messages:
+                                for message in initial_messages:
+                                    message.is_err = True
+                                    message.step_id = step_id
+                                await self.message_manager.create_many_messages_async(initial_messages, actor=self.actor)
+                        elif step_progression <= StepProgression.LOGGED_TRACE:
+                            if stop_reason is None:
+                                self.logger.error("Error in step after logging step")
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.error.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        elif step_progression == StepProgression.FINISHED and not should_continue:
+                            if stop_reason is None:
+                                stop_reason = LettaStopReason(stop_reason=StopReasonType.end_turn.value)
+                            await self.step_manager.update_step_stop_reason(self.actor, step_id, stop_reason.stop_reason)
+                        else:
+                            self.logger.error("Invalid StepProgression value")
+                    except Exception as e:
+                        self.logger.error("Failed to update step: %s", e)
 
             if not should_continue:
                 break
-
         # Extend the in context message ids
         if not agent_state.message_buffer_autoclear:
             await self._rebuild_context_window(
@@ -919,7 +1115,7 @@ class LettaAgent(BaseAgent):
                 in_context_messages=in_context_messages,
                 new_letta_messages=new_letta_messages,
             )
-        await self.agent_manager.set_in_context_messages_async(
+        await self.agent_manager.update_message_ids_async(
             agent_id=self.agent_id,
             message_ids=[m.id for m in new_in_context_messages],
             actor=self.actor,
@@ -936,7 +1132,7 @@ class LettaAgent(BaseAgent):
         new_in_context_messages, updated = await self.summarizer.summarize(
             in_context_messages=in_context_messages, new_letta_messages=[], force=True
         )
-        return await self.agent_manager.set_in_context_messages_async(
+        return await self.agent_manager.update_message_ids_async(
             agent_id=self.agent_id, message_ids=[m.id for m in new_in_context_messages], actor=self.actor
         )
 
@@ -948,18 +1144,17 @@ class LettaAgent(BaseAgent):
         agent_state: AgentState,
         tool_rules_solver: ToolRulesSolver,
     ) -> tuple[dict, list[str]]:
-        self.num_messages, self.num_archival_memories = await asyncio.gather(
-            (
-                self.message_manager.size_async(actor=self.actor, agent_id=agent_state.id)
-                if self.num_messages is None
-                else asyncio.sleep(0, result=self.num_messages)
-            ),
-            (
-                self.passage_manager.agent_passage_size_async(actor=self.actor, agent_id=agent_state.id)
-                if self.num_archival_memories is None
-                else asyncio.sleep(0, result=self.num_archival_memories)
-            ),
-        )
+        if not self.num_messages:
+            self.num_messages = await self.message_manager.size_async(
+                agent_id=agent_state.id,
+                actor=self.actor,
+            )
+        if not self.num_archival_memories:
+            self.num_archival_memories = await self.passage_manager.agent_passage_size_async(
+                agent_id=agent_state.id,
+                actor=self.actor,
+            )
+
         in_context_messages = await self._rebuild_memory_async(
             in_context_messages,
             agent_state,
@@ -1108,6 +1303,7 @@ class LettaAgent(BaseAgent):
             job_id=run_id if run_id else self.current_run_id,
             step_id=step_id,
             project_id=agent_state.project_id,
+            stop_reason=stop_reason,
         )
 
         tool_call_messages = create_letta_messages_from_llm_response(