Skip to content

Commit a6e4630

Browse files
umaannamalaiumaannamalai
authored
Merge branch 'main' into llm-response-token-counts
2 parents ce76038 + 0aeec05 commit a6e4630

13 files changed

Lines changed: 581 additions & 194 deletions

File tree

‎.github/workflows/deploy.yml‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,14 @@ jobs:
6363
- wheel: cp314t-musllinux
6464
os: ubuntu-24.04
6565
# Windows
66-
# Windows wheels won't but published until the full release announcement.
67-
# - wheel: cp313-win
68-
# os: windows-2025
69-
# - wheel: cp314-win
70-
# os: windows-2025
66+
- wheel: cp313-win
67+
os: windows-2025
68+
- wheel: cp313t-win
69+
os: windows-2025
70+
- wheel: cp314-win
71+
os: windows-2025
72+
- wheel: cp314t-win
73+
os: windows-2025
7174

7275
name: Build wheels for ${{ matrix.wheel }}
7376
runs-on: ${{ matrix.os }}
@@ -94,7 +97,7 @@ jobs:
9497
CIBW_ARCHS_WINDOWS: AMD64 ARM64
9598
CIBW_ENVIRONMENT_LINUX: "LD_LIBRARY_PATH=/opt/rh/devtoolset-8/root/usr/lib64:/opt/rh/devtoolset-8/root/usr/lib:/opt/rh/devtoolset-8/root/usr/lib64/dyninst:/opt/rh/devtoolset-8/root/usr/lib/dyninst:/usr/local/lib64:/usr/local/lib"
9699
CIBW_ENABLE: cpython-freethreading
97-
CIBW_TEST_REQUIRES: pytest
100+
CIBW_TEST_REQUIRES: pytest certifi
98101
CIBW_TEST_COMMAND_LINUX: "export PYTHONPATH={project}/tests; pytest {project}/tests/agent_unittests -vx"
99102
CIBW_TEST_COMMAND_MACOS: "export PYTHONPATH={project}/tests; pytest {project}/tests/agent_unittests -vx"
100103
CIBW_TEST_COMMAND_WINDOWS: "set PYTHONPATH={project}/tests; pytest {project}/tests/agent_unittests -vx"

‎newrelic/common/llm_utils.py‎

Lines changed: 74 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -38,71 +38,125 @@ def _get_llm_metadata(transaction):
3838
return llm_metadata_dict
3939

4040

41-
class GeneratorProxy(ObjectProxy):
41+
class LLMStreamProxy(ObjectProxy):
4242
def __init__(self, wrapped, on_stop_iteration, on_error):
4343
super().__init__(wrapped)
4444
self._nr_on_stop_iteration = on_stop_iteration
4545
self._nr_on_error = on_error
46+
# Track if we've sent the LLM events yet to avoid sending them multiple times
47+
self._nr_closed = False
4648

4749
def __iter__(self):
4850
self._nr_wrapped_iter = self.__wrapped__.__iter__()
4951
return self
5052

5153
def __next__(self):
52-
transaction = current_transaction()
53-
if not transaction:
54-
return self._nr_wrapped_iter.__next__()
55-
56-
return_val = None
5754
try:
5855
return_val = self._nr_wrapped_iter.__next__()
5956
except StopIteration:
60-
self._nr_on_stop_iteration(self, transaction)
57+
transaction = current_transaction()
58+
if transaction:
59+
self._nr_closed = True
60+
self._nr_on_stop_iteration(self, transaction)
6161
raise
6262
except Exception:
63-
self._nr_on_error(self, transaction)
63+
transaction = current_transaction()
64+
if transaction:
65+
self._nr_closed = True
66+
self._nr_on_error(self, transaction)
6467
raise
65-
return return_val
68+
else:
69+
return return_val
6670

6771
def close(self):
72+
if self._nr_closed:
73+
# If we already sent the related events, we can just call close as there's nothing left to do.
74+
return self.__wrapped__.close()
75+
76+
transaction = current_transaction()
77+
if transaction:
78+
# Send the events as if we were hitting StopIteration.
79+
self._nr_closed = True
80+
self._nr_on_stop_iteration(self, transaction)
81+
6882
return self.__wrapped__.close()
6983

84+
def throw(self, *args):
85+
if self._nr_closed:
86+
# If we already sent the related events, we can just call throw as there's nothing left to do.
87+
return self.__wrapped__.throw(*args)
88+
89+
transaction = current_transaction()
90+
if transaction:
91+
# Send the events as if we were hitting an exception.
92+
self._nr_closed = True
93+
self._nr_on_error(self, transaction)
94+
95+
return self.__wrapped__.throw(*args)
96+
7097
def __copy__(self):
7198
# Required to properly interface with itertool.tee, which can be called by LangChain on generators
7299
self.__wrapped__, copy = itertools.tee(self.__wrapped__, 2)
73-
return GeneratorProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)
100+
return LLMStreamProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)
74101

75102

76-
class AsyncGeneratorProxy(ObjectProxy):
103+
class AsyncLLMStreamProxy(ObjectProxy):
77104
def __init__(self, wrapped, on_stop_iteration, on_error):
78105
super().__init__(wrapped)
79106
self._nr_on_stop_iteration = on_stop_iteration
80107
self._nr_on_error = on_error
108+
# Track if we've sent the LLM events yet to avoid sending them multiple times
109+
self._nr_closed = False
81110

82111
def __aiter__(self):
83112
self._nr_wrapped_iter = self.__wrapped__.__aiter__()
84113
return self
85114

86115
async def __anext__(self):
87-
transaction = current_transaction()
88-
if not transaction:
89-
return await self._nr_wrapped_iter.__anext__()
90-
91-
return_val = None
92116
try:
93117
return_val = await self._nr_wrapped_iter.__anext__()
94118
except StopAsyncIteration:
95-
self._nr_on_stop_iteration(self, transaction)
119+
transaction = current_transaction()
120+
if transaction:
121+
self._nr_closed = True
122+
self._nr_on_stop_iteration(self, transaction)
96123
raise
97124
except Exception:
98-
self._nr_on_error(self, transaction)
125+
transaction = current_transaction()
126+
if transaction:
127+
self._nr_closed = True
128+
self._nr_on_error(self, transaction)
99129
raise
100-
return return_val
130+
else:
131+
return return_val
101132

102133
async def aclose(self):
134+
if self._nr_closed:
135+
# If we already sent the related events, we can just call aclose as there's nothing left to do.
136+
return await self.__wrapped__.aclose()
137+
138+
transaction = current_transaction()
139+
if transaction:
140+
# Send the events as if we were hitting StopAsyncIteration.
141+
self._nr_closed = True
142+
self._nr_on_stop_iteration(self, transaction)
143+
103144
return await self.__wrapped__.aclose()
104145

146+
async def athrow(self, *args):
147+
if self._nr_closed:
148+
# If we already sent the related events, we can just call athrow as there's nothing left to do.
149+
return await self.__wrapped__.athrow(*args)
150+
151+
transaction = current_transaction()
152+
if transaction:
153+
# Send the events as if we were hitting an exception.
154+
self._nr_closed = True
155+
self._nr_on_error(self, transaction)
156+
157+
return await self.__wrapped__.athrow(*args)
158+
105159
def __copy__(self):
106160
# Required to properly interface with itertool.tee, which can be called by LangChain on generators
107161
self.__wrapped__, copy = itertools.tee(self.__wrapped__, n=2)
108-
return AsyncGeneratorProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)
162+
return AsyncLLMStreamProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)

‎newrelic/hooks/external_botocore.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1158,15 +1158,15 @@ def converse_record_stream_chunk(self, event, transaction):
11581158

11591159
class EventStreamWrapper(ObjectProxy):
11601160
def __iter__(self):
1161-
g = GeneratorProxy(self.__wrapped__.__iter__())
1161+
g = LLMStreamProxy(self.__wrapped__.__iter__())
11621162
g._nr_ft = getattr(self, "_nr_ft", None)
11631163
g._nr_bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
11641164
g._nr_model_extractor = getattr(self, "_nr_model_extractor", NULL_EXTRACTOR)
11651165
g._nr_is_converse = getattr(self, "_nr_is_converse", False)
11661166
return g
11671167

11681168

1169-
class GeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
1169+
class LLMStreamProxy(BedrockRecordEventMixin, ObjectProxy):
11701170
def __init__(self, wrapped):
11711171
super().__init__(wrapped)
11721172
self._nr_request_timestamp = int(1000.0 * time.time())
@@ -1197,15 +1197,15 @@ def close(self):
11971197

11981198
class AsyncEventStreamWrapper(ObjectProxy):
11991199
def __aiter__(self):
1200-
g = AsyncGeneratorProxy(self.__wrapped__.__aiter__())
1200+
g = AsyncLLMStreamProxy(self.__wrapped__.__aiter__())
12011201
g._nr_ft = getattr(self, "_nr_ft", None)
12021202
g._nr_bedrock_attrs = getattr(self, "_nr_bedrock_attrs", {})
12031203
g._nr_model_extractor = getattr(self, "_nr_model_extractor", NULL_EXTRACTOR)
12041204
g._nr_is_converse = getattr(self, "_nr_is_converse", False)
12051205
return g
12061206

12071207

1208-
class AsyncGeneratorProxy(BedrockRecordEventMixin, ObjectProxy):
1208+
class AsyncLLMStreamProxy(BedrockRecordEventMixin, ObjectProxy):
12091209
def __init__(self, wrapped):
12101210
super().__init__(wrapped)
12111211
self._nr_request_timestamp = int(1000.0 * time.time())

0 commit comments

Comments
 (0)