feat(go): update telemetry for compatibility with Firebase AIM

go/ai/action_test.go

@@ -140,7 +140,11 @@ func TestGenerateAction(t *testing.T) {
 					t.Errorf("chunks mismatch (-want +got):\n%s", diff)
 				}
 
-				if diff := cmp.Diff(tc.ExpectResponse, resp, cmp.Options{cmpopts.EquateEmpty()}); diff != "" {
+				if diff := cmp.Diff(tc.ExpectResponse, resp, cmp.Options{
+					cmpopts.EquateEmpty(),
+					cmpopts.IgnoreFields(ModelResponse{}, "LatencyMs"),
+					cmpopts.IgnoreFields(GenerationUsage{}, "InputCharacters", "OutputCharacters"),
+				}); diff != "" {
 					t.Errorf("response mismatch (-want +got):\n%s", diff)
 				}
 			} else {
@@ -149,7 +153,11 @@ func TestGenerateAction(t *testing.T) {
 					t.Fatalf("action failed: %v", err)
 				}
 
-				if diff := cmp.Diff(tc.ExpectResponse, resp, cmp.Options{cmpopts.EquateEmpty()}); diff != "" {
+				if diff := cmp.Diff(tc.ExpectResponse, resp, cmp.Options{
+					cmpopts.EquateEmpty(),
+					cmpopts.IgnoreFields(ModelResponse{}, "LatencyMs"),
+					cmpopts.IgnoreFields(GenerationUsage{}, "InputCharacters", "OutputCharacters"),
+				}); diff != "" {
 					t.Errorf("response mismatch (-want +got):\n%s", diff)
 				}
 			}

go/ai/document.go

@@ -19,6 +19,7 @@ package ai
 import (
 	"encoding/json"
 	"fmt"
+	"strings"
 )
 
 // A Document is a piece of data that can be embedded, indexed, or retrieved.
@@ -160,6 +161,33 @@ func (p *Part) IsReasoning() bool {
 	return p.Kind == PartReasoning
 }
 
+// IsImage reports whether the [Part] contains an image.
+func (p *Part) IsImage() bool {
+	if p == nil || !p.IsMedia() {
+		return false
+	}
+	return IsImageContentType(p.ContentType) ||
+		strings.HasPrefix(p.Text, "data:image/")
+}
+
+// IsVideo reports whether the [Part] contains a video.
+func (p *Part) IsVideo() bool {
+	if p == nil || !p.IsMedia() {
+		return false
+	}
+	return IsVideoContentType(p.ContentType) ||
+		strings.HasPrefix(p.Text, "data:video/")
+}
+
+// IsAudio reports whether the [Part] contains an audio file.
+func (p *Part) IsAudio() bool {
+	if p == nil || !p.IsMedia() {
+		return false
+	}
+	return IsAudioContentType(p.ContentType) ||
+		strings.HasPrefix(p.Text, "data:audio/")
+}
+
 // IsResource reports whether the [Part] contains a resource reference.
 func (p *Part) IsResource() bool {
 	return p.Kind == PartResource
@@ -313,3 +341,21 @@ func DocumentFromText(text string, metadata map[string]any) *Document {
 		Metadata: metadata,
 	}
 }
+
+// IsImageContentType checks if the content type represents an image.
+func IsImageContentType(contentType string) bool {
+	return strings.HasPrefix(contentType, "image/") ||
+		strings.HasPrefix(contentType, "data:image/")
+}
+
+// IsVideoContentType checks if the content type represents a video.
+func IsVideoContentType(contentType string) bool {
+	return strings.HasPrefix(contentType, "video/") ||
+		strings.HasPrefix(contentType, "data:video/")
+}
+
+// IsAudioContentType checks if the content type represents an audio file.
+func IsAudioContentType(contentType string) bool {
+	return strings.HasPrefix(contentType, "audio/") ||
+		strings.HasPrefix(contentType, "data:audio/")
+}

go/ai/evaluator.go

@@ -196,7 +196,12 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat
 				if datapoint.TestCaseId == "" {
 					datapoint.TestCaseId = uuid.New().String()
 				}
-				_, err := tracing.RunInNewSpan(ctx, fmt.Sprintf("TestCase %s", datapoint.TestCaseId), "evaluator", false, datapoint,
+				spanMetadata := &tracing.SpanMetadata{
+					Name:    fmt.Sprintf("TestCase %s", datapoint.TestCaseId),
+					Type:    "evaluator",
+					Subtype: "evaluator",
+				}
+				_, err := tracing.RunInNewSpan(ctx, spanMetadata, datapoint,
 					func(ctx context.Context, input *Example) (*EvaluatorCallbackResponse, error) {
 						traceId := trace.SpanContextFromContext(ctx).TraceID().String()
 						spanId := trace.SpanContextFromContext(ctx).SpanID().String()

go/ai/generate.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"slices"
 	"strings"
+	"time"
 
 	"github.com/firebase/genkit/go/core"
 	"github.com/firebase/genkit/go/core/api"
@@ -122,7 +123,12 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction {
 					"err", err)
 			}()
 
-			return tracing.RunInNewSpan(ctx, "generate", "util", false, actionOpts,
+			spanMetadata := &tracing.SpanMetadata{
+				Name:    "generate",
+				Type:    "util",
+				Subtype: "util",
+			}
+			return tracing.RunInNewSpan(ctx, spanMetadata, actionOpts,
 				func(ctx context.Context, actionOpts *GenerateActionOptions) (*ModelResponse, error) {
 					return GenerateWithRequest(ctx, r, actionOpts, nil, cb)
 				})
@@ -176,6 +182,7 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model {
 		simulateSystemPrompt(opts, nil),
 		augmentWithContext(opts, nil),
 		validateSupport(name, opts),
+		addAutomaticTelemetry(),
 	}
 	fn = core.ChainMiddleware(mws...)(fn)
 
@@ -1073,6 +1080,209 @@ func handleResumeOption(ctx context.Context, r api.Registry, genOpts *GenerateAc
 	}, nil
 }
 
+// addAutomaticTelemetry creates middleware that automatically measures latency and calculates character and media counts.
+func addAutomaticTelemetry() ModelMiddleware {
+	return func(fn ModelFunc) ModelFunc {
+		return func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
+			startTime := time.Now()
+
+			// Call the underlying model function
+			resp, err := fn(ctx, req, cb)
+			if err != nil {
+				return nil, err
+			}
+
+			// Calculate latency
+			latencyMs := float64(time.Since(startTime).Nanoseconds()) / 1e6
+			if resp.LatencyMs == 0 {
+				resp.LatencyMs = latencyMs
+			}
+
+			// Calculate character and media counts automatically if Usage is available
+			if resp.Usage != nil {
+				if resp.Usage.InputCharacters == 0 {
+					resp.Usage.InputCharacters = calculateInputCharacters(req)
+				}
+				if resp.Usage.OutputCharacters == 0 {
+					resp.Usage.OutputCharacters = calculateOutputCharacters(resp)
+				}
+				if resp.Usage.InputImages == 0 {
+					resp.Usage.InputImages = calculateInputImages(req)
+				}
+				if resp.Usage.OutputImages == 0 {
+					resp.Usage.OutputImages = calculateOutputImages(resp)
+				}
+				if resp.Usage.InputVideos == 0 {
+					resp.Usage.InputVideos = float64(calculateInputVideos(req))
+				}
+				if resp.Usage.OutputVideos == 0 {
+					resp.Usage.OutputVideos = float64(calculateOutputVideos(resp))
+				}
+				if resp.Usage.InputAudioFiles == 0 {
+					resp.Usage.InputAudioFiles = float64(calculateInputAudio(req))
+				}
+				if resp.Usage.OutputAudioFiles == 0 {
+					resp.Usage.OutputAudioFiles = float64(calculateOutputAudio(resp))
+				}
+			} else {
+				// Create GenerationUsage if it doesn't exist
+				resp.Usage = &GenerationUsage{
+					InputCharacters:  calculateInputCharacters(req),
+					OutputCharacters: calculateOutputCharacters(resp),
+					InputImages:      calculateInputImages(req),
+					OutputImages:     calculateOutputImages(resp),
+					InputVideos:      float64(calculateInputVideos(req)),
+					OutputVideos:     float64(calculateOutputVideos(resp)),
+					InputAudioFiles:  float64(calculateInputAudio(req)),
+					OutputAudioFiles: float64(calculateOutputAudio(resp)),
+				}
+			}
+
+			return resp, nil
+		}
+	}
+}
+
+// calculateInputCharacters counts the total characters in the input request.
+func calculateInputCharacters(req *ModelRequest) int {
+	if req == nil {
+		return 0
+	}
+
+	totalChars := 0
+	for _, msg := range req.Messages {
+		if msg == nil {
+			continue
+		}
+		for _, part := range msg.Content {
+			if part != nil && part.Text != "" {
+				totalChars += len(part.Text)
+			}
+		}
+	}
+	return totalChars
+}
+
+// calculateOutputCharacters counts the total characters in the output response.
+func calculateOutputCharacters(resp *ModelResponse) int {
+	if resp == nil || resp.Message == nil {
+		return 0
+	}
+
+	totalChars := 0
+	for _, part := range resp.Message.Content {
+		if part != nil && part.Text != "" {
+			totalChars += len(part.Text)
+		}
+	}
+	return totalChars
+}
+
+// calculateInputImages counts the total number of images in the input request.
+func calculateInputImages(req *ModelRequest) int {
+	if req == nil {
+		return 0
+	}
+
+	imageCount := 0
+	for _, msg := range req.Messages {
+		if msg == nil {
+			continue
+		}
+		for _, part := range msg.Content {
+			if part != nil && part.IsImage() {
+				imageCount++
+			}
+		}
+	}
+	return imageCount
+}
+
+// calculateOutputImages counts the total number of images in the output response.
+func calculateOutputImages(resp *ModelResponse) int {
+	if resp == nil || resp.Message == nil {
+		return 0
+	}
+
+	imageCount := 0
+	for _, part := range resp.Message.Content {
+		if part != nil && part.IsImage() {
+			imageCount++
+		}
+	}
+	return imageCount
+}
+
+// calculateInputVideos counts the total number of videos in the input request.
+func calculateInputVideos(req *ModelRequest) int {
+	if req == nil {
+		return 0
+	}
+
+	videoCount := 0
+	for _, msg := range req.Messages {
+		if msg == nil {
+			continue
+		}
+		for _, part := range msg.Content {
+			if part != nil && part.IsVideo() {
+				videoCount++
+			}
+		}
+	}
+	return videoCount
+}
+
+// calculateOutputVideos counts the total number of videos in the output response.
+func calculateOutputVideos(resp *ModelResponse) int {
+	if resp == nil || resp.Message == nil {
+		return 0
+	}
+
+	videoCount := 0
+	for _, part := range resp.Message.Content {
+		if part != nil && part.IsVideo() {
+			videoCount++
+		}
+	}
+	return videoCount
+}
+
+// calculateInputAudio counts the total number of audio files in the input request.
+func calculateInputAudio(req *ModelRequest) int {
+	if req == nil {
+		return 0
+	}
+
+	audioCount := 0
+	for _, msg := range req.Messages {
+		if msg == nil {
+			continue
+		}
+		for _, part := range msg.Content {
+			if part != nil && part.IsAudio() {
+				audioCount++
+			}
+		}
+	}
+	return audioCount
+}
+
+// calculateOutputAudio counts the total number of audio files in the output response.
+func calculateOutputAudio(resp *ModelResponse) int {
+	if resp == nil || resp.Message == nil {
+		return 0
+	}
+
+	audioCount := 0
+	for _, part := range resp.Message.Content {
+		if part != nil && part.IsAudio() {
+			audioCount++
+		}
+	}
+	return audioCount
+}
+
 // processResources processes messages to replace resource parts with actual content.
 func processResources(ctx context.Context, r api.Registry, messages []*Message) ([]*Message, error) {
 	processedMessages := make([]*Message, len(messages))

go/ai/generate_test.go

@@ -38,6 +38,8 @@ var r = registry.New()
 func init() {
 	// Set up default formats
 	ConfigureFormats(r)
+	// Register the generate action that Generate() function expects
+	DefineGenerateAction(context.Background(), r)
 }
 
 // echoModel attributes

go/core/action.go

@@ -155,7 +155,6 @@ func newAction[In, Out, Stream any](
 
 	return &ActionDef[In, Out, Stream]{
 		fn: func(ctx context.Context, input In, cb StreamCallback[Stream]) (Out, error) {
-			tracing.SetCustomMetadataAttr(ctx, "subtype", string(atype))
 			return fn(ctx, input, cb)
 		},
 		desc: &api.ActionDesc{
@@ -185,7 +184,23 @@ func (a *ActionDef[In, Out, Stream]) Run(ctx context.Context, input In, cb Strea
 			"err", err)
 	}()
 
-	return tracing.RunInNewSpan(ctx, a.desc.Name, "action", false, input,
+	// Create span metadata and inject flow name if we're in a flow context
+	spanMetadata := &tracing.SpanMetadata{
+		Name:    a.desc.Name,
+		Type:    "action",
+		Subtype: string(a.desc.Type), // The actual action type becomes the subtype
+		// IsRoot will be automatically determined in tracing.go based on parent span presence
+	}
+
+	// Auto-inject flow name if we're in a flow context
+	if flowName := FlowNameFromContext(ctx); flowName != "" {
+		if spanMetadata.Metadata == nil {
+			spanMetadata.Metadata = make(map[string]string)
+		}
+		spanMetadata.Metadata["flow:name"] = flowName
+	}
+
+	return tracing.RunInNewSpan(ctx, spanMetadata, input,
 		func(ctx context.Context, input In) (Out, error) {
 			start := time.Now()
 			var err error
@@ -220,9 +235,7 @@ func (a *ActionDef[In, Out, Stream]) RunJSON(ctx context.Context, input json.Raw
 	}
 	var in In
 	if input != nil {
-		if err := json.Unmarshal(input, &in); err != nil {
-			return nil, err
-		}
+		json.Unmarshal(input, &in)
 	}
 	var callback func(context.Context, Stream) error
 	if cb != nil {