feat: add agents and log error properly (#8914)

* add agents and log error properly

* fix llm stream adapter
This commit is contained in:
Kian Jones
2026-01-18 18:50:47 -08:00
committed by Sarah Wooders
parent 6472834130
commit 81b5d71889
11 changed files with 52 additions and 7 deletions

View File

@@ -25,11 +25,13 @@ class LettaLLMAdapter(ABC):
llm_client: LLMClientBase,
llm_config: LLMConfig,
agent_id: str | None = None,
agent_tags: list[str] | None = None,
run_id: str | None = None,
) -> None:
self.llm_client: LLMClientBase = llm_client
self.llm_config: LLMConfig = llm_config
self.agent_id: str | None = agent_id
self.agent_tags: list[str] | None = agent_tags
self.run_id: str | None = run_id
self.message_id: str | None = None
self.request_data: dict | None = None

View File

@@ -125,6 +125,7 @@ class LettaLLMRequestAdapter(LettaLLMAdapter):
response_json=self.response_data,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=self.agent_tags,
run_id=self.run_id,
),
),

View File

@@ -26,8 +26,15 @@ class LettaLLMStreamAdapter(LettaLLMAdapter):
specific streaming formats.
"""
def __init__(self, llm_client: LLMClientBase, llm_config: LLMConfig, agent_id: str | None = None, run_id: str | None = None) -> None:
super().__init__(llm_client, llm_config, agent_id=agent_id, run_id=run_id)
def __init__(
self,
llm_client: LLMClientBase,
llm_config: LLMConfig,
agent_id: str | None = None,
agent_tags: list[str] | None = None,
run_id: str | None = None,
) -> None:
super().__init__(llm_client, llm_config, agent_id=agent_id, agent_tags=agent_tags, run_id=run_id)
self.interface: OpenAIStreamingInterface | AnthropicStreamingInterface | None = None
async def invoke_llm(
@@ -227,6 +234,7 @@ class LettaLLMStreamAdapter(LettaLLMAdapter):
response_json=response_json,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=self.agent_tags,
run_id=self.run_id,
),
),

View File

@@ -43,6 +43,7 @@ class SimpleLLMRequestAdapter(LettaLLMRequestAdapter):
telemetry_manager=self.telemetry_manager,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=self.agent_tags,
run_id=self.run_id,
call_type="agent_step",
)

View File

@@ -281,6 +281,7 @@ class SimpleLLMStreamAdapter(LettaLLMStreamAdapter):
response_json=response_json,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=self.agent_tags,
run_id=self.run_id,
),
),

View File

@@ -91,6 +91,7 @@ class EphemeralSummaryAgent(BaseAgent):
llm_client.set_telemetry_context(
telemetry_manager=TelemetryManager(),
agent_id=self.agent_id,
agent_tags=agent_state.tags,
call_type="summarization",
)
response_data = await llm_client.request_async_with_telemetry(request_data, agent_state.llm_config)

View File

@@ -416,6 +416,7 @@ class LettaAgent(BaseAgent):
response_json=response_data,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
run_id=self.current_run_id,
),
)
@@ -763,6 +764,7 @@ class LettaAgent(BaseAgent):
response_json=response_data,
step_id=step_id,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
run_id=self.current_run_id,
),
)
@@ -1230,6 +1232,7 @@ class LettaAgent(BaseAgent):
},
step_id=step_id,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
run_id=self.current_run_id,
),
)
@@ -1456,6 +1459,7 @@ class LettaAgent(BaseAgent):
llm_client.set_telemetry_context(
telemetry_manager=self.telemetry_manager,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
run_id=self.current_run_id,
call_type="agent_step",
)
@@ -1524,6 +1528,7 @@ class LettaAgent(BaseAgent):
llm_client.set_telemetry_context(
telemetry_manager=self.telemetry_manager,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
run_id=self.current_run_id,
call_type="agent_step",
)

View File

@@ -155,7 +155,9 @@ class LettaAgentV2(BaseAgentV2):
response = self._step(
run_id=None,
messages=in_context_messages + input_messages_to_persist,
llm_adapter=LettaLLMRequestAdapter(llm_client=self.llm_client, llm_config=self.agent_state.llm_config),
llm_adapter=LettaLLMRequestAdapter(
llm_client=self.llm_client, llm_config=self.agent_state.llm_config, agent_tags=self.agent_state.tags
),
dry_run=True,
enforce_run_id_set=False,
)
@@ -206,7 +208,11 @@ class LettaAgentV2(BaseAgentV2):
messages=in_context_messages + self.response_messages,
input_messages_to_persist=input_messages_to_persist,
llm_adapter=LettaLLMRequestAdapter(
llm_client=self.llm_client, llm_config=self.agent_state.llm_config, agent_id=self.agent_state.id, run_id=run_id
llm_client=self.llm_client,
llm_config=self.agent_state.llm_config,
agent_id=self.agent_state.id,
agent_tags=self.agent_state.tags,
run_id=run_id,
),
run_id=run_id,
use_assistant_message=use_assistant_message,
@@ -289,6 +295,7 @@ class LettaAgentV2(BaseAgentV2):
llm_client=self.llm_client,
llm_config=self.agent_state.llm_config,
agent_id=self.agent_state.id,
agent_tags=self.agent_state.tags,
run_id=run_id,
)
else:
@@ -296,6 +303,7 @@ class LettaAgentV2(BaseAgentV2):
llm_client=self.llm_client,
llm_config=self.agent_state.llm_config,
agent_id=self.agent_state.id,
agent_tags=self.agent_state.tags,
run_id=run_id,
)

View File

@@ -39,6 +39,7 @@ class LLMClientBase:
self.use_tool_naming = use_tool_naming
self._telemetry_manager: Optional["TelemetryManager"] = None
self._telemetry_agent_id: Optional[str] = None
self._telemetry_agent_tags: Optional[List[str]] = None
self._telemetry_run_id: Optional[str] = None
self._telemetry_step_id: Optional[str] = None
self._telemetry_call_type: Optional[str] = None
@@ -47,6 +48,7 @@ class LLMClientBase:
self,
telemetry_manager: Optional["TelemetryManager"] = None,
agent_id: Optional[str] = None,
agent_tags: Optional[List[str]] = None,
run_id: Optional[str] = None,
step_id: Optional[str] = None,
call_type: Optional[str] = None,
@@ -54,6 +56,7 @@ class LLMClientBase:
"""Set telemetry context for provider trace logging."""
self._telemetry_manager = telemetry_manager
self._telemetry_agent_id = agent_id
self._telemetry_agent_tags = agent_tags
self._telemetry_run_id = run_id
self._telemetry_step_id = step_id
self._telemetry_call_type = call_type
@@ -68,11 +71,13 @@ class LLMClientBase:
logger = get_logger(__name__)
response_data = None
error_msg = None
error_type = None
try:
response_data = await self.request_async(request_data, llm_config)
return response_data
except Exception as e:
error_msg = str(e)
error_type = type(e).__name__
raise
finally:
if self._telemetry_manager and settings.track_provider_trace:
@@ -85,9 +90,10 @@ class LLMClientBase:
actor=pydantic_actor,
provider_trace=ProviderTrace(
request_json=request_data,
response_json=response_data if response_data else {"error": error_msg},
response_json=response_data if response_data else {"error": error_msg, "error_type": error_type},
step_id=self._telemetry_step_id,
agent_id=self._telemetry_agent_id,
agent_tags=self._telemetry_agent_tags,
run_id=self._telemetry_run_id,
call_type=self._telemetry_call_type,
),
@@ -128,6 +134,7 @@ class LLMClientBase:
response_json=response_json,
step_id=self._telemetry_step_id,
agent_id=self._telemetry_agent_id,
agent_tags=self._telemetry_agent_tags,
run_id=self._telemetry_run_id,
call_type=self._telemetry_call_type,
),

View File

@@ -70,8 +70,15 @@ class SocketProviderTraceBackend(ProviderTraceBackendClient):
response = provider_trace.response_json or {}
request = provider_trace.request_json or {}
# Extract error if present
error = response.get("error", {}).get("message") if isinstance(response.get("error"), dict) else None
# Extract error if present - handles both {"error": "msg"} and {"error": {"message": "msg"}}
raw_error = response.get("error")
if isinstance(raw_error, dict):
error = raw_error.get("message")
elif isinstance(raw_error, str):
error = raw_error
else:
error = None
error_type = response.get("error_type")
record = {
"protocol_version": PROTOCOL_VERSION,
@@ -84,6 +91,7 @@ class SocketProviderTraceBackend(ProviderTraceBackendClient):
"request": request,
"response": response if not error else None,
"error": error,
"error_type": error_type,
"timestamp": datetime.now(timezone.utc).isoformat(),
}

View File

@@ -172,6 +172,7 @@ class Summarizer:
actor=self.actor,
include_ack=True,
agent_id=self.agent_id,
agent_tags=agent_state.tags,
)
# TODO add counts back
@@ -429,6 +430,7 @@ async def simple_summary(
prompt: str | None = None,
telemetry_manager: "TelemetryManager | None" = None,
agent_id: str | None = None,
agent_tags: List[str] | None = None,
run_id: str | None = None,
) -> str:
"""Generate a simple summary from a list of messages.
@@ -450,6 +452,7 @@ async def simple_summary(
llm_client.set_telemetry_context(
telemetry_manager=tm,
agent_id=agent_id,
agent_tags=agent_tags,
run_id=run_id,
call_type="summarization",
)