Skip to content

Commit dd29433

Browse files
umaannamalaiumaannamalai
committed
Add openai sync instrumentation.
1 parent 2834663 commit dd29433

File tree

4 files changed

+318
-4
lines changed

4 files changed

+318
-4
lines changed

‎newrelic/config.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2037,6 +2037,16 @@ def _process_trace_cache_import_hooks():
20372037

20382038

20392039
def _process_module_builtin_defaults():
2040+
_process_module_definition(
2041+
"openai.api_resources.chat_completion",
2042+
"newrelic.hooks.mlmodel_openai",
2043+
"instrument_openai_api_resources_chat_completion",
2044+
)
2045+
_process_module_definition(
2046+
"openai.util",
2047+
"newrelic.hooks.mlmodel_openai",
2048+
"instrument_openai_util",
2049+
)
20402050
_process_module_definition(
20412051
"asyncio.base_events",
20422052
"newrelic.hooks.coroutines_asyncio",

‎newrelic/hooks/mlmodel_openai.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
2+
# Copyright 2010 New Relic, Inc.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
import openai
16+
import uuid
17+
from newrelic.api.function_trace import FunctionTrace
18+
from newrelic.common.object_wrapper import wrap_function_wrapper
19+
from newrelic.api.transaction import current_transaction
20+
from newrelic.api.time_trace import get_trace_linking_metadata
21+
from newrelic.core.config import global_settings
22+
from newrelic.common.object_names import callable_name
23+
from newrelic.core.attribute import MAX_LOG_MESSAGE_LENGTH
24+
25+
26+
def wrap_chat_completion_create(wrapped, instance, args, kwargs):
27+
transaction = current_transaction()
28+
29+
if not transaction:
30+
return
31+
32+
ft_name = callable_name(wrapped)
33+
with FunctionTrace(ft_name) as ft:
34+
response = wrapped(*args, **kwargs)
35+
36+
if not response:
37+
return
38+
39+
custom_attrs_dict = transaction._custom_params
40+
conversation_id = custom_attrs_dict["conversation_id"] if "conversation_id" in custom_attrs_dict.keys() else str(uuid.uuid4())
41+
42+
chat_completion_id = str(uuid.uuid4())
43+
available_metadata = get_trace_linking_metadata()
44+
span_id = available_metadata.get("span.id", "")
45+
trace_id = available_metadata.get("trace.id", "")
46+
47+
response_headers = getattr(response, "_nr_response_headers")
48+
response_model = response.model
49+
settings = transaction.settings if transaction.settings is not None else global_settings()
50+
51+
chat_completion_summary_dict = {
52+
"id": chat_completion_id,
53+
"appName": settings.app_name,
54+
"conversation_id": conversation_id,
55+
"span_id": span_id,
56+
"trace_id": trace_id,
57+
"transaction_id": transaction._transaction_id,
58+
"api_key_last_four_digits": f"sk-{response.api_key[-4:]}",
59+
"response_time": ft.duration,
60+
"request.model": kwargs.get("model") or kwargs.get("engine"),
61+
"response.model": response_model,
62+
"response.organization": response.organization,
63+
"response.usage.completion_tokens": response.usage.completion_tokens,
64+
"response.usage.total_tokens": response.usage.total_tokens,
65+
"response.usage.prompt_tokens": response.usage.prompt_tokens,
66+
"request.temperature": kwargs.get("temperature"),
67+
"request.max_tokens": kwargs.get("max_tokens"),
68+
"response.choices.finish_reason": response.choices[0].finish_reason,
69+
"response.api_type": response.api_type,
70+
"response.headers.llmVersion": response_headers.get("openai-version"),
71+
"response.headers.ratelimitLimitRequests": check_rate_limit_header(response_headers, "x-ratelimit-limit-requests", True),
72+
"response.headers.ratelimitLimitTokens": check_rate_limit_header(response_headers, "x-ratelimit-limit-tokens", True),
73+
"response.headers.ratelimitResetTokens": check_rate_limit_header(response_headers, "x-ratelimit-reset-tokens", False),
74+
"response.headers.ratelimitResetRequests": check_rate_limit_header(response_headers, "x-ratelimit-reset-requests", False),
75+
"response.headers.ratelimitRemainingTokens": check_rate_limit_header(response_headers, "x-ratelimit-remaining-tokens", True),
76+
"response.headers.ratelimitRemainingRequests": check_rate_limit_header(response_headers, "x-ratelimit-remaining-requests", True),
77+
"vendor": "openAI",
78+
"ingest_source": "Python",
79+
"number_of_messages": len(kwargs.get("messages", [])) + len(response.choices),
80+
"api_version": response_headers.get("openai-version")
81+
}
82+
83+
transaction.record_ml_event("LlmChatCompletionSummary", chat_completion_summary_dict)
84+
message_list = list(kwargs.get("messages", [])) + [response.choices[0].message]
85+
86+
create_chat_completion_message_event(transaction, message_list, chat_completion_id, span_id, trace_id, response_model)
87+
88+
return response
89+
90+
91+
def check_rate_limit_header(response_headers, header_name, is_int):
92+
if header_name in response_headers:
93+
header_value = response_headers.get(header_name)
94+
if is_int:
95+
header_value = int(header_value)
96+
return header_value
97+
else:
98+
return None
99+
100+
101+
def create_chat_completion_message_event(transaction, message_list, chat_completion_id, span_id, trace_id, response_model):
102+
if not transaction:
103+
return
104+
105+
for index, message in enumerate(message_list):
106+
chat_completion_message_dict = {
107+
"id": str(uuid.uuid4()),
108+
"span_id": span_id,
109+
"trace_id": trace_id,
110+
"transaction_id": transaction._transaction_id,
111+
"content": message.get("content", "")[:MAX_LOG_MESSAGE_LENGTH],
112+
"role": message.get("role"),
113+
"completion_id": chat_completion_id,
114+
"sequence": index,
115+
"model": response_model,
116+
"vendor": "openAI",
117+
"ingest_source": "Python",
118+
}
119+
transaction.record_ml_event("LlmChatCompletionMessage", chat_completion_message_dict)
120+
121+
122+
def wrap_convert_to_openai_object(wrapped, instance, args, kwargs):
123+
resp = args[0]
124+
returned_response = wrapped(*args, **kwargs)
125+
126+
if isinstance(resp, openai.openai_response.OpenAIResponse):
127+
setattr(returned_response, "_nr_response_headers", getattr(resp, "_headers", {}))
128+
129+
return returned_response
130+
131+
132+
def instrument_openai_api_resources_chat_completion(module):
133+
if hasattr(module.ChatCompletion, "create"):
134+
wrap_function_wrapper(module, "ChatCompletion.create", wrap_chat_completion_create)
135+
136+
137+
def instrument_openai_util(module):
138+
wrap_function_wrapper(module, "convert_to_openai_object", wrap_convert_to_openai_object)

‎tests/mlmodel_openai/conftest.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,11 @@
3636
"transaction_tracer.stack_trace_threshold": 0.0,
3737
"debug.log_data_collector_payloads": True,
3838
"debug.record_transaction_failure": True,
39-
"ml_insights_event.enabled": True,
39+
"machine_learning.enabled": True,
40+
#"machine_learning.inference_events_value.enabled": True,
41+
"ml_insights_events.enabled": True
4042
}
43+
4144
collector_agent_registration = collector_agent_registration_fixture(
4245
app_name="Python Agent Test (mlmodel_openai)",
4346
default_settings=_default_settings,

‎tests/mlmodel_openai/test_chat_completion.py‎

Lines changed: 166 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,182 @@
1313
# limitations under the License.
1414

1515
import openai
16+
import pytest
17+
from testing_support.fixtures import ( # function_not_called,; override_application_settings,
18+
function_not_called,
19+
override_application_settings,
20+
reset_core_stats_engine,
21+
)
22+
23+
from newrelic.api.time_trace import current_trace
24+
from newrelic.api.transaction import current_transaction, add_custom_attribute
25+
from testing_support.validators.validate_ml_event_count import validate_ml_event_count
26+
from testing_support.validators.validate_ml_event_payload import (
27+
validate_ml_event_payload,
28+
)
29+
from testing_support.validators.validate_ml_events import validate_ml_events
30+
from testing_support.validators.validate_ml_events_outside_transaction import (
31+
validate_ml_events_outside_transaction,
32+
)
33+
34+
import newrelic.core.otlp_utils
35+
from newrelic.api.application import application_instance as application
36+
from newrelic.api.background_task import background_task
37+
from newrelic.api.transaction import record_ml_event
38+
from newrelic.core.config import global_settings
39+
from newrelic.packages import six
1640

1741
_test_openai_chat_completion_sync_messages = (
1842
{"role": "system", "content": "You are a scientist."},
19-
{"role": "user", "content": "What is the boiling point of water?"},
20-
{"role": "assistant", "content": "The boiling point of water is 212 degrees Fahrenheit."},
43+
#{"role": "user", "content": "What is the boiling point of water?"},
44+
#{"role": "assistant", "content": "The boiling point of water is 212 degrees Fahrenheit."},
2145
{"role": "user", "content": "What is 212 degrees Fahrenheit converted to Celsius?"},
2246
)
2347

2448

25-
def test_openai_chat_completion_sync():
49+
def set_trace_info():
50+
txn = current_transaction()
51+
if txn:
52+
txn._trace_id = "trace-id"
53+
trace = current_trace()
54+
if trace:
55+
trace.guid = "span-id"
56+
57+
58+
sync_chat_completion_recorded_events = [
59+
(
60+
{'type': 'LlmChatCompletionSummary'},
61+
{
62+
'id': None, # UUID that varies with each run
63+
'appName': 'Python Agent Test (mlmodel_openai)',
64+
'conversation_id': 'my-awesome-id',
65+
'transaction_id': None,
66+
'span_id': "span-id",
67+
'trace_id': "trace-id",
68+
'api_key_last_four_digits': 'sk-CRET',
69+
'response_time': None, # Response time varies each test run
70+
'request.model': 'gpt-3.5-turbo',
71+
'response.model': 'gpt-3.5-turbo-0613',
72+
'response.organization': 'new-relic-nkmd8b',
73+
'response.usage.completion_tokens': 11,
74+
'response.usage.total_tokens': 64,
75+
'response.usage.prompt_tokens': 53,
76+
'request.temperature': 0.7,
77+
'request.max_tokens': 100,
78+
'response.choices.finish_reason': 'stop',
79+
'response.api_type': 'None',
80+
'response.headers.llmVersion': '2020-10-01',
81+
'response.headers.ratelimitLimitRequests': 200,
82+
'response.headers.ratelimitLimitTokens': 40000,
83+
'response.headers.ratelimitResetTokens': "90ms",
84+
'response.headers.ratelimitResetRequests': "7m12s",
85+
'response.headers.ratelimitRemainingTokens': 39940,
86+
'response.headers.ratelimitRemainingRequests': 199,
87+
'vendor': 'openAI',
88+
'ingest_source': 'Python',
89+
'number_of_messages': 3,
90+
'api_version': '2020-10-01'
91+
},
92+
),
93+
(
94+
{'type': 'LlmChatCompletionMessage'},
95+
{
96+
'id': None,
97+
'span_id': "span-id",
98+
'trace_id': "trace-id",
99+
'transaction_id': None,
100+
'content': 'You are a scientist.',
101+
'role': 'system',
102+
'completion_id': None,
103+
'sequence': 0,
104+
'model': 'gpt-3.5-turbo-0613',
105+
'vendor': 'openAI',
106+
'ingest_source': 'Python'
107+
},
108+
),
109+
(
110+
{'type': 'LlmChatCompletionMessage'},
111+
{
112+
'id': None,
113+
'span_id': "span-id",
114+
'trace_id': "trace-id",
115+
'transaction_id': None,
116+
'content': 'What is 212 degrees Fahrenheit converted to Celsius?',
117+
'role': 'user',
118+
'completion_id': None,
119+
'sequence': 1,
120+
'model': 'gpt-3.5-turbo-0613',
121+
'vendor': 'openAI',
122+
'ingest_source': 'Python'
123+
},
124+
),
125+
(
126+
{'type': 'LlmChatCompletionMessage'},
127+
{
128+
'id': None,
129+
'span_id': "span-id",
130+
'trace_id': "trace-id",
131+
'transaction_id': None,
132+
'content': '212 degrees Fahrenheit is equal to 100 degrees Celsius.',
133+
'role': 'assistant',
134+
'completion_id': None,
135+
'sequence': 2,
136+
'model': 'gpt-3.5-turbo-0613',
137+
'vendor': 'openAI',
138+
'ingest_source': 'Python'
139+
}
140+
),
141+
]
142+
143+
144+
@reset_core_stats_engine()
145+
@validate_ml_events(sync_chat_completion_recorded_events)
146+
# One summary event, one system message, one user message, and one response message from the assistant
147+
@validate_ml_event_count(count=4)
148+
@background_task()
149+
def test_openai_chat_completion_sync_in_txn():
150+
set_trace_info()
151+
add_custom_attribute("conversation_id", "my-awesome-id")
152+
openai.ChatCompletion.create(
153+
model="gpt-3.5-turbo",
154+
messages=_test_openai_chat_completion_sync_messages,
155+
temperature=0.7,
156+
max_tokens=100
157+
)
158+
159+
160+
@reset_core_stats_engine()
161+
# One summary event, one system message, one user message, and one response message from the assistant
162+
@validate_ml_event_count(count=0)
163+
def test_openai_chat_completion_sync_outside_txn():
164+
set_trace_info()
165+
add_custom_attribute("conversation_id", "my-awesome-id")
166+
openai.ChatCompletion.create(
167+
model="gpt-3.5-turbo",
168+
messages=_test_openai_chat_completion_sync_messages,
169+
temperature=0.7,
170+
max_tokens=100
171+
)
172+
173+
174+
disabled_ml_settings = {
175+
"machine_learning.enabled": False,
176+
"machine_learning.inference_events_value.enabled": False,
177+
"ml_insights_events.enabled": False
178+
}
179+
180+
181+
@override_application_settings(disabled_ml_settings)
182+
@reset_core_stats_engine()
183+
# One summary event, one system message, one user message, and one response message from the assistant
184+
@validate_ml_event_count(count=0)
185+
def test_openai_chat_completion_sync_disabled_settings():
186+
set_trace_info()
26187
openai.ChatCompletion.create(
27188
model="gpt-3.5-turbo",
28189
messages=_test_openai_chat_completion_sync_messages,
190+
temperature=0.7,
191+
max_tokens=100
29192
)
30193

31194

0 commit comments

Comments
 (0)