Skip to content

Commit 327eee2

Browse files
whyiugMittaPei
whyiug
and
MittaPei
authored
[HER Hack-Astron #2] Langfuse integration + end-to-end agent traces (#1607)
* feat(observability): add Langfuse tracing Signed-off-by: whyiug <whyiug@hotmail.com> * fix(agent): capture streaming token usage Signed-off-by: whyiug <whyiug@hotmail.com> * docs(observability): add judge evidence Signed-off-by: whyiug <whyiug@hotmail.com> * fix(observability): harden Langfuse production paths Signed-off-by: whyiug <whyiug@hotmail.com> * fix(observability): strip tracestate from Langfuse export Signed-off-by: MittaPei <315415437+MittaPei@users.noreply.github.com> * fix(observability): close Langfuse compatibility gaps Signed-off-by: whyiug <whyiug@hotmail.com> * fix(observability): harden trusted trace handoff Signed-off-by: whyiug <whyiug@hotmail.com> * test(observability): cover production trace handoff Signed-off-by: whyiug <whyiug@hotmail.com> --------- Signed-off-by: whyiug <whyiug@hotmail.com> Signed-off-by: MittaPei <315415437+MittaPei@users.noreply.github.com> Co-authored-by: MittaPei <315415437+MittaPei@users.noreply.github.com>
1 parent e03019c commit 327eee2

53 files changed

Lines changed: 6798 additions & 527 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎core/agent/api/v1/base_api.py‎

Lines changed: 123 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,21 @@
1+
import inspect
12
import json
23
import os
34
import time
45
import traceback
56
from abc import ABC, abstractmethod
7+
from contextlib import aclosing
68
from dataclasses import dataclass
79
from typing import Any, AsyncGenerator, List
810

911
# Use unified common package import module
1012
from common.exceptions.base import BaseExc
1113
from common.otlp.log_trace.node_trace_log import NodeTraceLog, Status
1214
from common.otlp.metrics.meter import Meter
15+
from common.otlp.trace.langfuse import langfuse_enabled
1316
from common.otlp.trace.span import Span
17+
from opentelemetry.trace import Status as OtelStatus
18+
from opentelemetry.trace import StatusCode
1419
from pydantic import BaseModel, ConfigDict
1520

1621
from agent.api.schemas.base_inputs import BaseInputs
@@ -102,104 +107,133 @@ async def _process_chunk(
102107
chunk_logs.append(chunk.model_dump_json())
103108
yield await self.create_chunk(chunk)
104109

110+
@staticmethod
111+
def _attach_total_usage(
112+
stop_chunk: ReasonChatCompletionChunk, node_trace_log: NodeTraceLog
113+
) -> None:
114+
"""Attach aggregated model usage to the terminal SSE chunk."""
115+
if not node_trace_log.trace:
116+
return
117+
118+
from openai.types.completion_usage import CompletionUsage
119+
120+
total_usage = {
121+
"completion_tokens": 0,
122+
"prompt_tokens": 0,
123+
"total_tokens": 0,
124+
}
125+
for node in node_trace_log.trace:
126+
if hasattr(node, "data") and hasattr(node.data, "usage"):
127+
total_usage["completion_tokens"] += node.data.usage.completion_tokens
128+
total_usage["prompt_tokens"] += node.data.usage.prompt_tokens
129+
total_usage["total_tokens"] += node.data.usage.total_tokens
130+
131+
if total_usage["total_tokens"] > 0:
132+
stop_chunk.usage = CompletionUsage(
133+
completion_tokens=total_usage["completion_tokens"],
134+
prompt_tokens=total_usage["prompt_tokens"],
135+
total_tokens=total_usage["total_tokens"],
136+
)
137+
138+
async def _terminal_chunks(self, context: RunContext) -> tuple[str, str]:
139+
"""Build normal terminal frames before control enters final cleanup."""
140+
if context.error.c != 0:
141+
context.error.m += f",{context.span.sid}"
142+
context.span.add_error_events({"traceback": context.error_log})
143+
144+
stop_chunk = await self.create_stop(context.span, context.error)
145+
self._attach_total_usage(stop_chunk, context.node_trace_log)
146+
context.chunk_logs.append(stop_chunk.model_dump_json())
147+
for chunk_log in context.chunk_logs:
148+
context.span.add_info_events({"response-chunk": chunk_log})
149+
return await self.create_chunk(stop_chunk), await self.create_done()
150+
151+
def _finalize_run(self, context: RunContext) -> None:
152+
"""Perform output-free cleanup, including on cancellation/aclose."""
153+
if os.getenv("UPLOAD_METRICS"):
154+
context.meter.in_error_count(context.error.c)
155+
attributes: dict[str, Any] = {"code": context.error.c}
156+
if langfuse_enabled():
157+
attributes["astron.agent.error_code"] = context.error.c
158+
if context.error.c != 0 and langfuse_enabled():
159+
# The application protocol reports errors as terminal SSE frames,
160+
# so no exception escapes this span context automatically. Mark
161+
# the swallowed failure explicitly without exporting its possibly
162+
# sensitive message to Langfuse.
163+
context.span.set_status(OtelStatus(StatusCode.ERROR))
164+
attributes.update(
165+
{
166+
"langfuse.observation.level": "ERROR",
167+
"langfuse.observation.status_message": "Agent execution failed",
168+
}
169+
)
170+
context.span.set_attributes(attributes=attributes)
171+
context.span.add_info_events({"message": context.error.m})
172+
context.node_trace_log.record_end()
173+
if os.getenv("UPLOAD_NODE_TRACE"):
174+
node_trace_log = context.node_trace_log.upload(
175+
status=Status(code=context.error.c, message=context.error.m),
176+
log_caller=self.log_caller,
177+
span=context.span,
178+
)
179+
context.span.add_info_events(
180+
{
181+
"node-trace": json.dumps(
182+
node_trace_log,
183+
ensure_ascii=False,
184+
default=json_serializer,
185+
)
186+
}
187+
)
188+
105189
async def run_runner(
106190
self, node_trace_log: NodeTraceLog, meter: Meter, span: Span
107191
) -> AsyncGenerator[str, None]:
108192

109193
with span.start("RunRunner") as sp:
110-
error: BaseExc = AgentNormalExc()
111-
error_log: str = ""
112194
chunk_logs: List[str] = []
195+
context = RunContext(
196+
error=AgentNormalExc(),
197+
error_log="",
198+
chunk_logs=chunk_logs,
199+
span=sp,
200+
node_trace_log=node_trace_log,
201+
meter=meter,
202+
)
113203

114204
try:
115-
runner = await self.build_runner(sp)
116-
if runner is None:
117-
raise AgentInternalExc("Failed to build runner")
118-
119-
async for chunk in runner.run(span=sp, node_trace_log=node_trace_log):
120-
chunk.id = span.sid
121-
async for processed_chunk in self._process_chunk(chunk, chunk_logs):
122-
yield processed_chunk
123-
124-
except BaseExc as e:
125-
error = e
126-
error_log = traceback.format_exc()
127-
except Exception as e: # pylint: disable=broad-exception-caught
128-
error = AgentInternalExc(str(e))
129-
error_log = traceback.format_exc()
130-
205+
try:
206+
runner = await self.build_runner(sp)
207+
if runner is None:
208+
raise AgentInternalExc("Failed to build runner")
209+
210+
runner_stream = runner.run(span=sp, node_trace_log=node_trace_log)
211+
if inspect.isawaitable(runner_stream):
212+
runner_stream = await runner_stream
213+
async with aclosing(runner_stream):
214+
async for chunk in runner_stream:
215+
chunk.id = span.sid
216+
processed_stream = self._process_chunk(chunk, chunk_logs)
217+
async with aclosing(processed_stream):
218+
async for processed_chunk in processed_stream:
219+
yield processed_chunk
220+
221+
except BaseExc as exc:
222+
context.error = exc
223+
context.error_log = traceback.format_exc()
224+
except Exception as exc: # pylint: disable=broad-exception-caught
225+
context.error = AgentInternalExc(str(exc))
226+
context.error_log = traceback.format_exc()
227+
228+
stop_frame, done_frame = await self._terminal_chunks(context)
229+
yield stop_frame
230+
yield done_frame
131231
finally:
132-
context = RunContext(
133-
error=error,
134-
error_log=error_log,
135-
chunk_logs=chunk_logs,
136-
span=sp,
137-
node_trace_log=node_trace_log,
138-
meter=meter,
139-
)
140-
"""Cleanup work after completing the run"""
141-
if context.error.c != 0:
142-
context.error.m += f",{context.span.sid}"
143-
context.span.add_error_events({"traceback": context.error_log})
144-
145-
stop_chunk = await self.create_stop(context.span, context.error)
146-
# Attach usage from node_trace if available
147-
if context.node_trace_log.trace:
148-
from openai.types.completion_usage import CompletionUsage
149-
150-
total_usage = {
151-
"completion_tokens": 0,
152-
"prompt_tokens": 0,
153-
"total_tokens": 0,
154-
}
155-
for node in context.node_trace_log.trace:
156-
if hasattr(node, "data") and hasattr(node.data, "usage"):
157-
total_usage[
158-
"completion_tokens"
159-
] += node.data.usage.completion_tokens
160-
total_usage[
161-
"prompt_tokens"
162-
] += node.data.usage.prompt_tokens
163-
total_usage["total_tokens"] += node.data.usage.total_tokens
164-
165-
if total_usage["total_tokens"] > 0:
166-
stop_chunk.usage = CompletionUsage(
167-
completion_tokens=total_usage["completion_tokens"],
168-
prompt_tokens=total_usage["prompt_tokens"],
169-
total_tokens=total_usage["total_tokens"],
170-
)
171-
172-
context.chunk_logs.append(stop_chunk.model_dump_json())
173-
174-
for chunk_log in context.chunk_logs:
175-
context.span.add_info_events({"response-chunk": chunk_log})
176-
177-
yield await self.create_chunk(stop_chunk)
178-
yield await self.create_done()
179-
180-
if os.getenv("UPLOAD_METRICS"):
181-
context.meter.in_error_count(context.error.c)
182-
# context.meter.in_error_count(
183-
# context.error.c, lables={"msg": context.error.m}
184-
# )
185-
context.span.set_attributes(attributes={"code": context.error.c})
186-
context.span.add_info_events({"message": context.error.m})
187-
context.node_trace_log.record_end()
188-
if os.getenv("UPLOAD_NODE_TRACE"):
189-
node_trace_log = context.node_trace_log.upload(
190-
status=Status(code=context.error.c, message=context.error.m),
191-
log_caller=self.log_caller,
192-
span=context.span,
193-
)
194-
context.span.add_info_events(
195-
{
196-
"node-trace": json.dumps(
197-
node_trace_log,
198-
ensure_ascii=False,
199-
default=json_serializer,
200-
)
201-
}
202-
)
232+
# Never yield from a generator finalizer. Starlette closes the
233+
# response iterator when a client disconnects; yielding here
234+
# raises "async generator ignored GeneratorExit" and prevents
235+
# the active trace spans from ending.
236+
self._finalize_run(context)
203237

204238
@staticmethod
205239
async def create_chunk(chunk: Any) -> str:

0 commit comments

Comments
 (0)