Skip to content

Commit 8533a54

Browse files
committed
feature: 支持用户传入http client 到 OpenAIModel 中来控制连接的生命周期
- 修改trace的调用方式避免上传trace不全
1 parent 32d7af9 commit 8533a54

6 files changed

Lines changed: 343 additions & 21 deletions

File tree

Lines changed: 326 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,326 @@
1+
# Tencent is pleased to support the open source community by making tRPC-Agent-Python available.
2+
#
3+
# Copyright (C) 2026 Tencent. All rights reserved.
4+
#
5+
# tRPC-Agent-Python is licensed under Apache-2.0.
6+
"""Langfuse reporting regression tests tied to the active pytest interpreter.
7+
8+
Expected behaviour (same test file, different venv):
9+
- ``./venv/bin/pytest tests/langfuse/tracing/test_langfuse_reporting_fixtures.py`` → PASS
10+
(dev env: no broken site-packages copy, or fixed source).
11+
- ``./examples/quickstart/venv/bin/pytest ...`` → FAIL
12+
(quickstart venv ships an older ``trpc_agent_sdk`` in site-packages with detached spans).
13+
14+
The probe reads span-creation code from **site-packages when present**, matching
15+
``run_agent.py`` under ``examples/quickstart/``. Repo-root ``sys.path`` is ignored
16+
for that detection so pytest at repo root still exercises the installed wheel.
17+
"""
18+
19+
from __future__ import annotations
20+
21+
import json
22+
import sys
23+
from pathlib import Path
24+
from typing import Callable
25+
from unittest.mock import MagicMock
26+
27+
import pytest
28+
from opentelemetry import trace
29+
from opentelemetry.sdk.trace import TracerProvider
30+
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
31+
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
32+
33+
import trpc_agent_sdk.server.langfuse.tracing.opentelemetry as otel_module
34+
from trpc_agent_sdk.events import Event
35+
from trpc_agent_sdk.server.langfuse.tracing.opentelemetry import LangfuseConfig, _LangfuseMixin
36+
from trpc_agent_sdk.telemetry._trace import trace_agent
37+
from trpc_agent_sdk.telemetry._trace import trace_runner
38+
from trpc_agent_sdk.telemetry._trace import tracer
39+
from trpc_agent_sdk.types import Content
40+
from trpc_agent_sdk.types import Part
41+
42+
SPAN_PREFIX = "trpc.python.agent"
43+
44+
SYSTEM_INSTRUCTION = (
45+
"You are an agent who's name is [assistant].\n\n"
46+
"You are a helpful assistant for query weather."
47+
)
48+
TOOLS = [
49+
{
50+
"function_declarations": [
51+
{
52+
"description": "get weather information for the specified city",
53+
"name": "get_weather_report",
54+
"parameters": {
55+
"properties": {"city": {"type": "STRING"}},
56+
"type": "OBJECT",
57+
},
58+
}
59+
]
60+
}
61+
]
62+
LLM_REQUEST = {
63+
"model": "glm-5.0-w4afp8",
64+
"config": {
65+
"system_instruction": SYSTEM_INSTRUCTION,
66+
"tools": TOOLS,
67+
},
68+
"contents": [
69+
{
70+
"parts": [{"text": "What's the weather like today?"}],
71+
"role": "user",
72+
}
73+
],
74+
}
75+
LLM_RESPONSE = {
76+
"content": {
77+
"parts": [
78+
{"text": "assistant reply", "thought": False},
79+
],
80+
"role": "model",
81+
},
82+
"partial": False,
83+
"usage_metadata": {
84+
"candidates_token_count": 107,
85+
"prompt_token_count": 185,
86+
"total_token_count": 292,
87+
},
88+
}
89+
90+
91+
def _iter_site_packages() -> list[Path]:
92+
paths: list[Path] = []
93+
prefix = Path(sys.prefix)
94+
for lib_dir in ("lib64", "lib"):
95+
base = prefix / lib_dir
96+
if not base.is_dir():
97+
continue
98+
for child in sorted(base.glob("python*/site-packages")):
99+
if child.is_dir():
100+
paths.append(child)
101+
return paths
102+
103+
104+
def _resolve_sdk_file(relative: str) -> Path:
105+
"""Resolve a SDK file, preferring site-packages over repo-source import."""
106+
for site in _iter_site_packages():
107+
candidate = site / "trpc_agent_sdk" / relative
108+
if candidate.is_file():
109+
return candidate
110+
module_path = "trpc_agent_sdk." + relative.replace("/", ".").removesuffix(".py")
111+
module = __import__(module_path, fromlist=["_"])
112+
return Path(module.__file__)
113+
114+
115+
def _span_pattern_from_source(
116+
source: str,
117+
*,
118+
detached_needle: str,
119+
current_needle: str,
120+
) -> str:
121+
"""Return ``current`` or ``detached`` based on how a span is opened in source text."""
122+
has_current = current_needle in source
123+
has_detached = detached_needle in source
124+
if has_current and not has_detached:
125+
return "current"
126+
if has_detached and not has_current:
127+
return "detached"
128+
if has_current:
129+
return "current"
130+
if has_detached:
131+
return "detached"
132+
raise AssertionError(
133+
f"cannot detect span pattern (needles detached={detached_needle!r} "
134+
f"current={current_needle!r})"
135+
)
136+
137+
138+
def _agent_run_span_pattern() -> str:
139+
source = _resolve_sdk_file("agents/_base_agent.py").read_text(encoding="utf-8")
140+
return _span_pattern_from_source(
141+
source,
142+
detached_needle='span = tracer.start_span(f"agent_run',
143+
current_needle='with tracer.start_as_current_span(f"agent_run',
144+
)
145+
146+
147+
def _invocation_span_pattern() -> str:
148+
source = _resolve_sdk_file("runners.py").read_text(encoding="utf-8")
149+
return _span_pattern_from_source(
150+
source,
151+
detached_needle='span = tracer.start_span("invocation")',
152+
current_needle='with tracer.start_as_current_span(f"invocation")',
153+
)
154+
155+
156+
_EXPORTER: InMemorySpanExporter | None = None
157+
158+
159+
@pytest.fixture(scope="session", autouse=True)
160+
def _init_otel_tracer_once():
161+
"""OpenTelemetry allows only one TracerProvider; share it across this module."""
162+
global _EXPORTER # noqa: PLW0603
163+
_EXPORTER = InMemorySpanExporter()
164+
provider = TracerProvider()
165+
provider.add_span_processor(SimpleSpanProcessor(_EXPORTER))
166+
trace.set_tracer_provider(provider)
167+
yield
168+
_EXPORTER = None
169+
170+
171+
def _clear_finished_spans() -> None:
172+
assert _EXPORTER is not None
173+
_EXPORTER.clear()
174+
175+
176+
def _run_with_span_pattern(pattern: str, span_name: str, callback: Callable[[], None]) -> None:
177+
if pattern == "current":
178+
with tracer.start_as_current_span(span_name):
179+
callback()
180+
return
181+
if pattern == "detached":
182+
span = tracer.start_span(span_name)
183+
try:
184+
callback()
185+
finally:
186+
span.end()
187+
return
188+
raise ValueError(f"unknown span pattern: {pattern}")
189+
190+
191+
def _finished_span_attributes(name_substring: str) -> dict:
192+
assert _EXPORTER is not None
193+
spans = _EXPORTER.get_finished_spans()
194+
matched = [span for span in spans if name_substring in span.name]
195+
assert matched, (
196+
f"no finished span containing {name_substring!r}; "
197+
f"got span names: {[span.name for span in spans]}"
198+
)
199+
return dict(matched[0].attributes or {})
200+
201+
202+
def _map_to_langfuse(raw_attributes: dict) -> dict:
203+
mixin = _LangfuseMixin()
204+
otel_module._langfuse_config = LangfuseConfig()
205+
return mixin._map_attributes_to_langfuse(raw_attributes)
206+
207+
208+
def _make_invocation_context() -> MagicMock:
209+
ctx = MagicMock()
210+
ctx.agent.name = "assistant"
211+
ctx.user_content = Content(role="user", parts=[Part(text="What's the weather like today?")])
212+
ctx.override_messages = None
213+
ctx.session.id = "a252d252-4b55-4713-80e4-90abb177c433"
214+
ctx.session.user_id = "demo_user"
215+
ctx.invocation_id = "e-d5a9872c-80e3-43ea-b2a8-0091257a1616"
216+
return ctx
217+
218+
219+
def probe_agent_run_langfuse_mapping() -> dict:
220+
_clear_finished_spans()
221+
ctx = _make_invocation_context()
222+
pattern = _agent_run_span_pattern()
223+
224+
def _callback() -> None:
225+
trace_agent(
226+
invocation_context=ctx,
227+
agent_action="Could you please tell me the city you're interested in?",
228+
state_begin={"user_name": "demo_user"},
229+
state_end={"user_name": "demo_user"},
230+
)
231+
232+
_run_with_span_pattern(pattern, "agent_run [assistant]", _callback)
233+
return _map_to_langfuse(_finished_span_attributes("agent_run"))
234+
235+
236+
def probe_invocation_langfuse_mapping() -> dict:
237+
_clear_finished_spans()
238+
ctx = _make_invocation_context()
239+
pattern = _invocation_span_pattern()
240+
user_message = Content(role="user", parts=[Part(text="What's the weather like today?")])
241+
last_event = Event(
242+
content=Content(role="model", parts=[Part(text="Could you please tell me the city you're interested in?")]),
243+
)
244+
245+
def _callback() -> None:
246+
trace_runner(
247+
app_name="weather_agent_demo",
248+
user_id="demo_user",
249+
session_id="a252d252-4b55-4713-80e4-90abb177c433",
250+
invocation_context=ctx,
251+
new_message=user_message,
252+
last_event=last_event,
253+
state_begin={"user_name": "demo_user"},
254+
state_end={"user_name": "demo_user"},
255+
)
256+
257+
_run_with_span_pattern(pattern, "invocation", _callback)
258+
return _map_to_langfuse(_finished_span_attributes("invocation"))
259+
260+
261+
def assert_valid_call_llm_langfuse_mapping(result: dict) -> None:
262+
assert result["langfuse.observation.type"] == "generation", result
263+
llm_input = json.loads(result["langfuse.observation.input"])
264+
config = llm_input.get("config", {})
265+
assert config.get("system_instruction"), result
266+
assert config.get("tools"), result
267+
model_params = json.loads(result["langfuse.observation.model.parameters"])
268+
assert model_params.get("system_instruction"), result
269+
assert model_params.get("tools"), result
270+
assert result["langfuse.observation.output"] != "unknown", result
271+
272+
273+
def assert_valid_run_agent_langfuse_mapping(result: dict) -> None:
274+
assert result.get("langfuse.observation.type") == "span", result
275+
assert result.get("langfuse.observation.input") == "What's the weather like today?", result
276+
assert result.get("langfuse.observation.output"), result
277+
278+
279+
def assert_valid_run_runner_langfuse_mapping(result: dict) -> None:
280+
assert result.get("langfuse.trace.name") == "[trpc-agent]: weather_agent_demo/assistant", result
281+
assert result.get("langfuse.user.id") == "demo_user", result
282+
assert result.get("langfuse.session.id") == "a252d252-4b55-4713-80e4-90abb177c433", result
283+
assert result.get("langfuse.observation.input") == "What's the weather like today?", result
284+
assert result.get("langfuse.observation.output"), result
285+
assert "langfuse.trace.metadata" in result, result
286+
287+
288+
@pytest.fixture(autouse=True)
289+
def _langfuse_config():
290+
original = otel_module._langfuse_config
291+
otel_module._langfuse_config = LangfuseConfig()
292+
yield
293+
otel_module._langfuse_config = original
294+
295+
296+
@pytest.fixture
297+
def mixin():
298+
return _LangfuseMixin()
299+
300+
301+
class TestLangfuseReportingSpanContext:
302+
"""End-to-end: telemetry must land on the span Langfuse exports (ok.txt vs error.txt)."""
303+
304+
def test_trace_agent_reaches_agent_run_span(self):
305+
result = probe_agent_run_langfuse_mapping()
306+
assert_valid_run_agent_langfuse_mapping(result)
307+
308+
def test_trace_runner_reaches_invocation_span(self):
309+
result = probe_invocation_langfuse_mapping()
310+
assert_valid_run_runner_langfuse_mapping(result)
311+
312+
313+
class TestLangfuseReportingCallLlmMapping:
314+
"""call_llm mapping must always include system prompt and tools (ok.txt generation)."""
315+
316+
def test_call_llm_mapping_includes_system_instruction_and_tools(self, mixin):
317+
attrs = {
318+
"gen_ai.operation.name": "call_llm",
319+
f"{SPAN_PREFIX}.llm_request": json.dumps(LLM_REQUEST, ensure_ascii=False),
320+
f"{SPAN_PREFIX}.llm_response": json.dumps(LLM_RESPONSE, ensure_ascii=False),
321+
"gen_ai.usage.input_tokens": 185,
322+
"gen_ai.usage.output_tokens": 107,
323+
"gen_ai.request.model": "glm-5.0-w4afp8",
324+
}
325+
result = mixin._map_attributes_to_langfuse(attrs)
326+
assert_valid_call_llm_langfuse_mapping(result)

trpc_agent_sdk/agents/_base_agent.py

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -267,12 +267,7 @@ async def run_async(
267267
# because __aexit__ of the context manager is not guaranteed to run when
268268
# an async generator is cancelled, but try/finally always executes
269269
# even under CancelledError (PEP 492).
270-
from opentelemetry import context as context_api
271-
from opentelemetry.trace import set_span_in_context
272-
273-
span = tracer.start_span(f"agent_run [{self.name}]")
274-
_ctx_token = context_api.attach(set_span_in_context(span, context_api.get_current()))
275-
try:
270+
with tracer.start_as_current_span(f"agent_run [{self.name}]"):
276271
ctx = self._create_invocation_context(parent_context)
277272
if ctx.agent_context is None:
278273
ctx.agent_context = create_agent_context()
@@ -333,9 +328,6 @@ async def run_async(
333328

334329
# avoid memory leak
335330
reset_invocation_ctx(token)
336-
finally:
337-
context_api.detach(_ctx_token)
338-
span.end()
339331

340332
@abstractmethod
341333
async def _run_async_impl(self, ctx: InvocationContext) -> AsyncGenerator[Event, None]:

trpc_agent_sdk/models/_openai_model.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,12 @@ def __init__(
171171
# Extract OpenAI-specific config
172172
self.organization: str = kwargs.get(const.ORGANIZATION, "")
173173
self.client_args = kwargs.get(const.CLIENT_ARGS, {})
174+
# Allow callers (e.g. trpc-claw) to inject a tuned httpx client so the
175+
# underlying openai.AsyncOpenAI honors connection-pool settings such as
176+
# keepalive_expiry / max_keepalive_connections (avoids reusing stale
177+
# keep-alive sockets that gateways close earlier than httpx). Accept it
178+
# either as a top-level ``http_client`` kwarg or inside ``client_args``.
179+
self._http_client = self.client_args.get("http_client", None)
174180

175181
# Tool prompt configuration
176182
self.add_tools_to_prompt = add_tools_to_prompt
@@ -1109,7 +1115,10 @@ async def _generate_single(self,
11091115
else:
11101116
return self._create_response_without_content(response_dict)
11111117
finally:
1112-
await client.close()
1118+
# An externally-injected http client is owned by the caller and
1119+
# reused across requests; only close clients we created ourselves.
1120+
if self._http_client is None:
1121+
await client.close()
11131122

11141123
def _convert_tools_to_openai_format(self, tools: List[Tool]) -> List[Dict[str, Any]]:
11151124
"""Convert Google GenAI tools format to OpenAI tools format.
@@ -1772,4 +1781,7 @@ async def _generate_stream(self,
17721781
custom_metadata={"error": str(ex)},
17731782
)
17741783
finally:
1775-
await client.close()
1784+
# An externally-injected http client is owned by the caller and
1785+
# reused across requests; only close clients we created ourselves.
1786+
if self._http_client is None:
1787+
await client.close()

0 commit comments

Comments
 (0)