Skip to content

Commit

Permalink
bug fixes and unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@tippmar-nr
tippmar-nr committed Feb 10, 2025
1 parent a4ac218 commit b80e686
Show file tree
Hide file tree
Showing 4 changed files with 175 additions and 43 deletions.
2 changes: 1 addition & 1 deletion ...lic/Agent/Extensions/NewRelic.Agent.Extensions/Llm/OpenAi/OpenAiHeaderDictionaryHelper.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ public static IDictionary<string, string> GetOpenAiHeaders(this IDictionary<stri
case "openai-version":
llmHeaders.Add("llmVersion", header.Value);
break;
case "x-ratelimit-limit-requests:":
case "x-ratelimit-limit-requests":
llmHeaders.Add("ratelimitLimitRequests", header.Value);
break;
case "x-ratelimit-limit-tokens":
Expand Down
2 changes: 1 addition & 1 deletion src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/OpenAI/OpenAiHeaderDictionaryHelper.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ public static IDictionary<string, string> GetOpenAiHeaders(this IDictionary<stri
case "openai-version":
llmHeaders.Add("llmVersion", header.Value);
break;
case "x-ratelimit-limit-requests:":
case "x-ratelimit-limit-requests":
llmHeaders.Add("ratelimitLimitRequests", header.Value);
break;
case "x-ratelimit-limit-tokens":
Expand Down
140 changes: 99 additions & 41 deletions tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Llm/EventHelperTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,12 @@ public void Setup()
}

[Test]
public void CreateChatCompletionEvent_ShouldRecordLlmChatCompletionSummaryEvent()
[TestCase(null, null, null)]
[TestCase(100, 0.5f, "organization")]
public void CreateChatCompletionEvent_ShouldRecordLlmChatCompletionSummaryEvent(int? maxTokens, float? temperature, string organization)
{
// Arrange
var requestId = "123";
var temperature = 98.6f;
var maxTokens = 100;
var requestModel = "model1";
var responseModel = "model2";
var numMessages = 5;
Expand Down Expand Up @@ -68,46 +68,78 @@ public void CreateChatCompletionEvent_ShouldRecordLlmChatCompletionSummaryEvent(
});

// Act
var completionId = EventHelper.CreateChatCompletionEvent(_agent, _segment, requestId, temperature, maxTokens, requestModel, responseModel, numMessages, finishReason, vendor, false, headers, null);
var completionId = EventHelper.CreateChatCompletionEvent(_agent, _segment, requestId, temperature, maxTokens, requestModel, responseModel, numMessages, finishReason, vendor, false, headers, null, organization);

// Assert
Assert.That(completionId, Is.Not.Null);

// assert that _agent.RecordLlmEvent was called one time
Mock.Assert(() => _agent.RecordLlmEvent("LlmChatCompletionSummary", Arg.IsAny<Dictionary<string, object>>()), Occurs.Once());

var expectedAttributeCount = 21;

Assert.Multiple(() =>
{
// assert that the attributes passed to _agent.RecordLlmEvent are correct
Assert.That(llmAttributes, Is.Not.Null);
Assert.That(llmAttributes.Count, Is.EqualTo(23));
Assert.That(llmAttributes["id"], Is.EqualTo(completionId));
Assert.That(llmAttributes["request_id"], Is.EqualTo(requestId));
Assert.That(llmAttributes["span_id"], Is.EqualTo(_segment.SpanId));
Assert.That(llmAttributes["trace_id"], Is.EqualTo(_agent.GetLinkingMetadata()["trace.id"]));
Assert.That(llmAttributes["request.temperature"], Is.EqualTo(temperature));
Assert.That(llmAttributes["request.max_tokens"], Is.EqualTo(maxTokens));
Assert.That(llmAttributes["request.model"], Is.EqualTo(requestModel));
Assert.That(llmAttributes["response.model"], Is.EqualTo(responseModel));
Assert.That(llmAttributes["response.number_of_messages"], Is.EqualTo(numMessages));
Assert.That(llmAttributes["response.choices.finish_reason"], Is.EqualTo(finishReason));
Assert.That(llmAttributes["vendor"], Is.EqualTo(vendor));
Assert.That(llmAttributes["ingest_source"], Is.EqualTo("DotNet"));
Assert.That(llmAttributes["duration"], Is.EqualTo((float)_segment.DurationOrZero.TotalMilliseconds));
Assert.That(llmAttributes["llmVersion"], Is.EqualTo("1.0"));
Assert.That(llmAttributes["ratelimitLimitRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitLimitTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitResetTokens"], Is.EqualTo("ratelimitResetTokens"));
Assert.That(llmAttributes["ratelimitResetRequests"], Is.EqualTo("ratelimitResetRequests"));
Assert.That(llmAttributes["ratelimitRemainingTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitRemainingRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitLimitTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitResetTokensUsageBased"], Is.EqualTo("ratelimitResetTokensUsageBased"));
Assert.That(llmAttributes["ratelimitRemainingTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.llmVersion"], Is.EqualTo("1.0"));
Assert.That(llmAttributes["response.headers.ratelimitLimitRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitLimitTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitResetTokens"], Is.EqualTo("ratelimitResetTokens"));
Assert.That(llmAttributes["response.headers.ratelimitResetRequests"], Is.EqualTo("ratelimitResetRequests"));
Assert.That(llmAttributes["response.headers.ratelimitRemainingTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitRemainingRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitLimitTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitResetTokensUsageBased"], Is.EqualTo("ratelimitResetTokensUsageBased"));
Assert.That(llmAttributes["response.headers.ratelimitRemainingTokensUsageBased"], Is.EqualTo(99));

Assert.That(llmAttributes, Does.Not.ContainKey("error"));
Assert.That(llmAttributes, Does.Not.ContainKey("header1"));
Assert.That(llmAttributes, Does.Not.ContainKey("header2"));

if (maxTokens.HasValue)
{
Assert.That(llmAttributes["request.max_tokens"], Is.EqualTo(maxTokens));
++expectedAttributeCount;
}
else
{
Assert.That(llmAttributes, Does.Not.ContainKey("request.max_tokens"));
}

if (temperature.HasValue)
{
Assert.That(llmAttributes["request.temperature"], Is.EqualTo(temperature));
++expectedAttributeCount;
}
else
{
Assert.That(llmAttributes, Does.Not.ContainKey("request.temperature"));
}

if (!string.IsNullOrEmpty(organization))
{
Assert.That(llmAttributes["response.organization"], Is.EqualTo(organization));
++expectedAttributeCount;
}
else
{
Assert.That(llmAttributes, Does.Not.ContainKey("response.organization"));
}

Assert.That(llmAttributes.Count, Is.EqualTo(expectedAttributeCount));

});
}

Expand Down Expand Up @@ -202,10 +234,10 @@ public void CreateChatCompletionEvent_ShouldNoticeError_WhenIsErrorIsTrue()
});
}

[TestCase(false)]
[TestCase(true)]
[Test]
public void CreateChatMessageEvent_ShouldRecordLlmChatCompletionEvent(bool isResponse)
[TestCase(false, null)]
[TestCase(true, 123)]
public void CreateChatMessageEvent_ShouldRecordLlmChatCompletionEvent(bool isResponse, int? tokenCount)
{
// Arrange
var requestId = "123";
Expand All @@ -232,7 +264,10 @@ public void CreateChatMessageEvent_ShouldRecordLlmChatCompletionEvent(bool isRes

// Act
var responseId = Guid.NewGuid().ToString();
EventHelper.CreateChatMessageEvent(_agent, _segment, requestId, responseId, responseModel, content, role, sequence, completionId, isResponse, vendor);
EventHelper.CreateChatMessageEvent(_agent, _segment, requestId, responseId, responseModel, content, role, sequence, completionId, isResponse, vendor, tokenCount);


var expectedAttributeCount = 11;

// Assert
Mock.Assert(() => _agent.RecordLlmEvent("LlmChatCompletionMessage", Arg.IsAny<Dictionary<string, object>>()), Occurs.Once());
Expand All @@ -241,7 +276,6 @@ public void CreateChatMessageEvent_ShouldRecordLlmChatCompletionEvent(bool isRes
{
// assert that the attributes passed to _agent.RecordLlmEvent are correct
Assert.That(llmAttributes, Is.Not.Null);
Assert.That(llmAttributes.Count, Is.EqualTo(isResponse ? 12 : 11));
Assert.That(llmAttributes["id"], Is.EqualTo(responseId + "-" + sequence));
Assert.That(llmAttributes["request_id"], Is.EqualTo(requestId));
Assert.That(llmAttributes["span_id"], Is.EqualTo(_segment.SpanId));
Expand All @@ -253,17 +287,33 @@ public void CreateChatMessageEvent_ShouldRecordLlmChatCompletionEvent(bool isRes
Assert.That(llmAttributes["role"], Is.EqualTo(role));
Assert.That(llmAttributes["sequence"], Is.EqualTo(sequence));
Assert.That(llmAttributes["completion_id"], Is.EqualTo(completionId));
});

if (isResponse)
Assert.That(llmAttributes["is_response"], Is.True);
else
Assert.That(llmAttributes, Does.Not.ContainKey("is_response"));
if (isResponse)
{
Assert.That(llmAttributes["is_response"], Is.True);
++expectedAttributeCount;
}
else
Assert.That(llmAttributes, Does.Not.ContainKey("is_response"));

if (tokenCount.HasValue)
{
Assert.That(llmAttributes["token_count"], Is.EqualTo(tokenCount));
++expectedAttributeCount;
}
else
Assert.That(llmAttributes, Does.Not.ContainKey("token_count"));

Assert.That(llmAttributes.Count, Is.EqualTo(expectedAttributeCount));
});

}

[Test]
public void CreateEmbeddingEvent_ShouldRecordLlmEmbeddingEvent()
[TestCase("organization")]
[TestCase(null)]

public void CreateEmbeddingEvent_ShouldRecordLlmEmbeddingEvent(string organization)
{
// Arrange
var requestId = "123";
Expand Down Expand Up @@ -303,15 +353,15 @@ public void CreateEmbeddingEvent_ShouldRecordLlmEmbeddingEvent()
});

// Act
EventHelper.CreateEmbeddingEvent(_agent, _segment, requestId, input, requestModel, responseModel, vendor, false, headers, null);
EventHelper.CreateEmbeddingEvent(_agent, _segment, requestId, input, requestModel, responseModel, vendor, false, headers, null, organization);

// Assert
Mock.Assert(() => _agent.RecordLlmEvent("LlmEmbedding", Arg.IsAny<Dictionary<string, object>>()), Occurs.Once());

var expectedAttributeCount = 20;
Assert.Multiple(() =>
{
Assert.That(llmAttributes, Is.Not.Null);
Assert.That(llmAttributes.Count, Is.EqualTo(20));
Assert.That(llmAttributes.ContainsKey("id"));
Assert.That(llmAttributes["request_id"], Is.EqualTo(requestId));
Assert.That(llmAttributes["span_id"], Is.EqualTo(_segment.SpanId));
Expand All @@ -322,22 +372,30 @@ public void CreateEmbeddingEvent_ShouldRecordLlmEmbeddingEvent()
Assert.That(llmAttributes["vendor"], Is.EqualTo(vendor));
Assert.That(llmAttributes["ingest_source"], Is.EqualTo("DotNet"));
Assert.That(llmAttributes["duration"], Is.EqualTo((float)_segment.DurationOrZero.TotalMilliseconds));
Assert.That(llmAttributes["llmVersion"], Is.EqualTo("1.0"));

Assert.That(llmAttributes["ratelimitLimitRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitLimitTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitResetTokens"], Is.EqualTo("ratelimitResetTokens"));
Assert.That(llmAttributes["ratelimitResetRequests"], Is.EqualTo("ratelimitResetRequests"));
Assert.That(llmAttributes["ratelimitRemainingTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitRemainingRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitLimitTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["ratelimitResetTokensUsageBased"], Is.EqualTo("ratelimitResetTokensUsageBased"));
Assert.That(llmAttributes["ratelimitRemainingTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.llmVersion"], Is.EqualTo("1.0"));
Assert.That(llmAttributes["response.headers.ratelimitLimitRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitLimitTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitResetTokens"], Is.EqualTo("ratelimitResetTokens"));
Assert.That(llmAttributes["response.headers.ratelimitResetRequests"], Is.EqualTo("ratelimitResetRequests"));
Assert.That(llmAttributes["response.headers.ratelimitRemainingTokens"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitRemainingRequests"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitLimitTokensUsageBased"], Is.EqualTo(99));
Assert.That(llmAttributes["response.headers.ratelimitResetTokensUsageBased"], Is.EqualTo("ratelimitResetTokensUsageBased"));
Assert.That(llmAttributes["response.headers.ratelimitRemainingTokensUsageBased"], Is.EqualTo(99));

if (!string.IsNullOrEmpty(organization))
{
Assert.That(llmAttributes["response.organization"], Is.EqualTo(organization));
++expectedAttributeCount;
}
else
Assert.That(llmAttributes, Does.Not.ContainKey("response.organization"));

Assert.That(llmAttributes, Does.Not.ContainKey("error"));
Assert.That(llmAttributes, Does.Not.ContainKey("header1"));
Assert.That(llmAttributes, Does.Not.ContainKey("header2"));

Assert.That(llmAttributes.Count, Is.EqualTo(expectedAttributeCount));
});
}

Expand Down
74 changes: 74 additions & 0 deletions .../Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Llm/OpenAiHeaderDictionaryHelperTests.cs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
// Copyright 2020 New Relic, Inc. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

using NUnit.Framework;
using System.Collections.Generic;
using NewRelic.Agent.Extensions.Llm.OpenAi;

namespace NewRelic.Agent.Extensions.Tests.Llm.OpenAi
{
[TestFixture]
public class OpenAiHeaderDictionaryHelperTests
{
[Test]
public void GetOpenAiHeaders_ShouldReturnCorrectHeaders()
{
// Arrange
var headers = new Dictionary<string, string>
{
{ "openai-version", "1.0" },
{ "x-ratelimit-limit-requests", "1000" },
{ "x-ratelimit-limit-tokens", "5000" },
{ "x-ratelimit-remaining-requests", "900" },
{ "x-ratelimit-remaining-tokens", "4500" },
{ "x-ratelimit-reset-requests", "60" },
{ "x-ratelimit-reset-tokens", "30" }
};

// Act
var result = headers.GetOpenAiHeaders();

// Assert
Assert.Multiple(() =>
{
Assert.That(result, Has.Count.EqualTo(7));
Assert.That(result["llmVersion"], Is.EqualTo("1.0"));
Assert.That(result["ratelimitLimitRequests"], Is.EqualTo("1000"));
Assert.That(result["ratelimitLimitTokens"], Is.EqualTo("5000"));
Assert.That(result["ratelimitRemainingRequests"], Is.EqualTo("900"));
Assert.That(result["ratelimitRemainingTokens"], Is.EqualTo("4500"));
Assert.That(result["ratelimitResetRequests"], Is.EqualTo("60"));
Assert.That(result["ratelimitResetTokens"], Is.EqualTo("30"));
});
}

[Test]
public void TryGetOpenAiOrganization_ShouldReturnOrganization_WhenHeaderExists()
{
// Arrange
var headers = new Dictionary<string, string>
{
{ "openai-organization", "org-12345" }
};

// Act
var result = headers.TryGetOpenAiOrganization();

// Assert
Assert.That(result, Is.EqualTo("org-12345"));
}

[Test]
public void TryGetOpenAiOrganization_ShouldReturnNull_WhenHeaderDoesNotExist()
{
// Arrange
var headers = new Dictionary<string, string>();

// Act
var result = headers.TryGetOpenAiOrganization();

// Assert
Assert.That(result, Is.Null);
}
}
}

0 comments on commit b80e686

Please sign in to comment.