Skip to content

Commit

Permalink
Add request/ response IDs.
Browse files Browse the repository at this point in the history
  • Loading branch information
@umaannamalai
umaannamalai committed Oct 13, 2023
1 parent f6d13f8 commit 182c7a8
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 17 deletions.
23 changes: 14 additions & 9 deletions newrelic/hooks/mlmodel_openai.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):
return

custom_attrs_dict = transaction._custom_params
conversation_id = custom_attrs_dict["conversation_id"] if "conversation_id" in custom_attrs_dict.keys() else str(uuid.uuid4())
conversation_id = custom_attrs_dict.get("conversation_id", str(uuid.uuid4()))

chat_completion_id = str(uuid.uuid4())
available_metadata = get_trace_linking_metadata()
Expand All @@ -46,6 +46,8 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):

response_headers = getattr(response, "_nr_response_headers", None)
response_model = response.model
response_id = response.get("id", "")
request_id = response_headers.get("x-request-id", "")
settings = transaction.settings if transaction.settings is not None else global_settings()

chat_completion_summary_dict = {
Expand All @@ -55,8 +57,9 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):
"span_id": span_id,
"trace_id": trace_id,
"transaction_id": transaction._transaction_id,
"request_id": request_id,
"api_key_last_four_digits": f"sk-{response.api_key[-4:]}",
"response_time": ft.duration,
"duration": ft.duration,
"request.model": kwargs.get("model") or kwargs.get("engine"),
"response.model": response_model,
"response.organization": response.organization,
Expand All @@ -77,44 +80,46 @@ def wrap_chat_completion_create(wrapped, instance, args, kwargs):
"vendor": "openAI",
"ingest_source": "Python",
"number_of_messages": len(kwargs.get("messages", [])) + len(response.choices),
"api_version": response_headers.get("openai-version", "")
}

transaction.record_ml_event("LlmChatCompletionSummary", chat_completion_summary_dict)
message_list = list(kwargs.get("messages", [])) + [response.choices[0].message]

create_chat_completion_message_event(transaction, message_list, chat_completion_id, span_id, trace_id, response_model)
create_chat_completion_message_event(transaction, settings.app_name, message_list, chat_completion_id, span_id, trace_id, response_model, response_id, request_id)

return response


def check_rate_limit_header(response_headers, header_name, is_int):
if not response_headers:
return None
return ""

if header_name in response_headers:
header_value = response_headers.get(header_name)
if is_int:
header_value = int(header_value)
return header_value
else:
return None
return ""


def create_chat_completion_message_event(transaction, message_list, chat_completion_id, span_id, trace_id, response_model):
def create_chat_completion_message_event(transaction, app_name, message_list, chat_completion_id, span_id, trace_id, response_model, response_id, request_id):
if not transaction:
return

for index, message in enumerate(message_list):
chat_completion_message_dict = {
"id": str(uuid.uuid4()),
"id": "%s-%s" % (response_id, index),
"appName": app_name,
"request_id": request_id,
"span_id": span_id,
"trace_id": trace_id,
"transaction_id": transaction._transaction_id,
"content": message.get("content", "")[:MAX_LOG_MESSAGE_LENGTH],
"role": message.get("role"),
"completion_id": chat_completion_id,
"sequence": index,
"model": response_model,
"response.model": response_model,
"vendor": "openAI",
"ingest_source": "Python",
}
Expand Down
22 changes: 14 additions & 8 deletions tests/mlmodel_openai/test_chat_completion.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,9 @@ def set_trace_info():
'transaction_id': None,
'span_id': "span-id",
'trace_id': "trace-id",
'request_id': "49dbbffbd3c3f4612aa48def69059ccd",
'api_key_last_four_digits': 'sk-CRET',
'response_time': None, # Response time varies each test run
'duration': None, # Response time varies each test run
'request.model': 'gpt-3.5-turbo',
'response.model': 'gpt-3.5-turbo-0613',
'response.organization': 'new-relic-nkmd8b',
Expand All @@ -86,53 +87,58 @@ def set_trace_info():
'vendor': 'openAI',
'ingest_source': 'Python',
'number_of_messages': 3,
'api_version': '2020-10-01'
},
),
(
{'type': 'LlmChatCompletionMessage'},
{
'id': None,
'id': "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-0",
'appName': 'Python Agent Test (mlmodel_openai)',
'request_id': "49dbbffbd3c3f4612aa48def69059ccd",
'span_id': "span-id",
'trace_id': "trace-id",
'transaction_id': None,
'content': 'You are a scientist.',
'role': 'system',
'completion_id': None,
'sequence': 0,
'model': 'gpt-3.5-turbo-0613',
'response.model': 'gpt-3.5-turbo-0613',
'vendor': 'openAI',
'ingest_source': 'Python'
},
),
(
{'type': 'LlmChatCompletionMessage'},
{
'id': None,
'id': "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-1",
'appName': 'Python Agent Test (mlmodel_openai)',
'request_id': "49dbbffbd3c3f4612aa48def69059ccd",
'span_id': "span-id",
'trace_id': "trace-id",
'transaction_id': None,
'content': 'What is 212 degrees Fahrenheit converted to Celsius?',
'role': 'user',
'completion_id': None,
'sequence': 1,
'model': 'gpt-3.5-turbo-0613',
'response.model': 'gpt-3.5-turbo-0613',
'vendor': 'openAI',
'ingest_source': 'Python'
},
),
(
{'type': 'LlmChatCompletionMessage'},
{
'id': None,
'id': "chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTemv-2",
'appName': 'Python Agent Test (mlmodel_openai)',
'request_id': "49dbbffbd3c3f4612aa48def69059ccd",
'span_id': "span-id",
'trace_id': "trace-id",
'transaction_id': None,
'content': '212 degrees Fahrenheit is equal to 100 degrees Celsius.',
'role': 'assistant',
'completion_id': None,
'sequence': 2,
'model': 'gpt-3.5-turbo-0613',
'response.model': 'gpt-3.5-turbo-0613',
'vendor': 'openAI',
'ingest_source': 'Python'
}
Expand Down

0 comments on commit 182c7a8

Please sign in to comment.