feat: add SetOriginalModel method to ImageGenerationMetrics interface

- Add SetOriginalModel method to ImageGenerationMetrics interface to track original model from request body
- Implement SetOriginalModel in imageGeneration struct to delegate to baseMetrics
- Update tests to use separate SetOriginalModel, SetRequestModel, SetResponseModel calls instead of SetModel

Signed-off-by: Hrushikesh Patil <[email protected]>

Author: nutanix-Hrushikesh

internal/metrics/image_generation_metrics.go

@@ -25,6 +25,9 @@ type imageGeneration struct {
 type ImageGenerationMetrics interface {
 	// StartRequest initializes timing for a new request.
 	StartRequest(headers map[string]string)
+	// SetOriginalModel sets the original model from the incoming request body before any virtualization applies.
+	// This is usually called after parsing the request body. Example: dall-e-3
+	SetOriginalModel(originalModel internalapi.OriginalModel)
 	// SetRequestModel sets the request model name.
 	SetRequestModel(requestModel internalapi.RequestModel)
 	// SetResponseModel sets the response model name.
@@ -53,6 +56,11 @@ func (i *imageGeneration) StartRequest(headers map[string]string) {
 	i.baseMetrics.StartRequest(headers)
 }
 
+// SetOriginalModel sets the original model from the incoming request body before any virtualization applies.
+func (i *imageGeneration) SetOriginalModel(originalModel internalapi.OriginalModel) {
+	i.baseMetrics.SetOriginalModel(originalModel)
+}
+
 // SetRequestModel sets the request model for the request.
 func (i *imageGeneration) SetRequestModel(requestModel internalapi.RequestModel) {
 	i.baseMetrics.SetRequestModel(requestModel)

internal/metrics/image_generation_metrics_test.go

@@ -28,6 +28,7 @@ func TestImageGeneration_RecordTokenUsage(t *testing.T) {
 		attrsBase = []attribute.KeyValue{
 			attribute.Key(genaiAttributeOperationName).String(genaiOperationImageGeneration),
 			attribute.Key(genaiAttributeProviderName).String(genaiProviderOpenAI),
+			attribute.Key(genaiAttributeOriginalModel).String("test-model"),
 			attribute.Key(genaiAttributeRequestModel).String("test-model"),
 			attribute.Key(genaiAttributeResponseModel).String("test-model"),
 		}
@@ -36,7 +37,9 @@ func TestImageGeneration_RecordTokenUsage(t *testing.T) {
 	)
 
 	// Set labels and record usage.
-	im.SetModel("test-model", "test-model")
+	im.SetOriginalModel("test-model")
+	im.SetRequestModel("test-model")
+	im.SetResponseModel("test-model")
 	im.SetBackend(&filterapi.Backend{Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}})
 	im.RecordTokenUsage(t.Context(), 3, 7, nil)
 
@@ -60,6 +63,7 @@ func TestImageGeneration_RecordImageGeneration(t *testing.T) {
 		attrs := attribute.NewSet(
 			attribute.Key(genaiAttributeOperationName).String(genaiOperationImageGeneration),
 			attribute.Key(genaiAttributeProviderName).String(genaiProviderOpenAI),
+			attribute.Key(genaiAttributeOriginalModel).String("img-model"),
 			attribute.Key(genaiAttributeRequestModel).String("img-model"),
 			attribute.Key(genaiAttributeResponseModel).String("img-model"),
 			attribute.Key("gen_ai.image.count").Int(2),
@@ -68,7 +72,9 @@ func TestImageGeneration_RecordImageGeneration(t *testing.T) {
 		)
 
 		im.StartRequest(nil)
-		im.SetModel("img-model", "img-model")
+		im.SetOriginalModel("img-model")
+		im.SetRequestModel("img-model")
+		im.SetResponseModel("img-model")
 		im.SetBackend(&filterapi.Backend{Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}})
 
 		time.Sleep(10 * time.Millisecond)
@@ -94,13 +100,16 @@ func TestImageGeneration_HeaderLabelMapping(t *testing.T) {
 		"x-org-id":  "org456",
 	}
 
-	im.SetModel("test-model", "test-model")
+	im.SetOriginalModel("test-model")
+	im.SetRequestModel("test-model")
+	im.SetResponseModel("test-model")
 	im.SetBackend(&filterapi.Backend{Schema: filterapi.VersionedAPISchema{Name: filterapi.APISchemaOpenAI}})
 	im.RecordTokenUsage(t.Context(), 5, 0, requestHeaders)
 
 	attrs := attribute.NewSet(
 		attribute.Key(genaiAttributeOperationName).String(genaiOperationImageGeneration),
 		attribute.Key(genaiAttributeProviderName).String(genaiProviderOpenAI),
+		attribute.Key(genaiAttributeOriginalModel).String("test-model"),
 		attribute.Key(genaiAttributeRequestModel).String("test-model"),
 		attribute.Key(genaiAttributeResponseModel).String("test-model"),
 		attribute.Key(genaiAttributeTokenType).String(genaiTokenTypeInput),