agentmon/internal/sdk/events.go

package sdk

import (
	"time"
)

// NewSessionStart creates a session.start event.
func NewSessionStart(sessionID string, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "session.start",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewSessionEnd creates a session.end event.
func NewSessionEnd(sessionID string, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "session.end",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewRunStart creates a run.start event.
func NewRunStart(sessionID, runID string, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "run.start",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewRunEnd creates a run.end event.
func NewRunEnd(sessionID, runID string, status string, durationMs int64, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "run.end",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
		},
		"payload": map[string]any{
			"status":      status,
			"duration_ms": durationMs,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewSpanStart creates a span.start event.
func NewSpanStart(sessionID, runID, traceID, spanID string, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "span.start",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
			"trace_id":   traceID,
			"span_id":    spanID,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewSpanEnd creates a span.end event.
func NewSpanEnd(sessionID, runID, traceID, spanID string, status string, durationMs int64, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "span.end",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
			"trace_id":   traceID,
			"span_id":    spanID,
		},
		"payload": map[string]any{
			"status":      status,
			"duration_ms": durationMs,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewError creates an error event.
func NewError(sessionID, runID, traceID, spanID string, errType, message string, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "error",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
			"trace_id":   traceID,
			"span_id":    spanID,
		},
		"payload": map[string]any{
			"error": map[string]any{
				"type":    errType,
				"message": message,
			},
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// NewMetricSnapshot creates a metric.snapshot event.
func NewMetricSnapshot(sessionID, runID string, metrics map[string]any, opts ...EventOption) Event {
	now := time.Now()
	event := map[string]any{
		"schema": map[string]any{
			"name":    schemaName,
			"version": schemaVersion,
		},
		"event": map[string]any{
			"id":   generateID(),
			"type": "metric.snapshot",
			"ts":   now.UTC().Format(time.RFC3339Nano),
		},
		"correlation": map[string]any{
			"session_id": sessionID,
			"run_id":     runID,
		},
		"payload": map[string]any{
			"metrics": metrics,
		},
	}

	for _, opt := range opts {
		opt(event)
	}

	return event
}

// EventOption is a function that modifies an event.
type EventOption func(Event)

// WithSource sets the source information on an event.
func WithSource(emitter *Emitter) EventOption {
	return func(e Event) {
		if event, ok := e["event"].(map[string]any); ok {
			event["source"] = map[string]any{
				"framework": emitter.config.Framework,
				"client_id": emitter.config.ClientID,
				"host":      emitter.config.Host,
			}
		}
	}
}

// WithAttributes adds attributes to an event.
func WithAttributes(attrs map[string]any) EventOption {
	return func(e Event) {
		if _, ok := e["attributes"]; !ok {
			e["attributes"] = make(map[string]any)
		}
		if attrsMap, ok := e["attributes"].(map[string]any); ok {
			for k, v := range attrs {
				attrsMap[k] = v
			}
		}
	}
}

// WithSpanKind sets the span_kind attribute.
func WithSpanKind(kind string) EventOption {
	return WithAttributes(map[string]any{"span_kind": kind})
}

// WithName sets the name attribute.
func WithName(name string) EventOption {
	return WithAttributes(map[string]any{"name": name})
}

// WithParentSpanID sets the parent_span_id in correlation.
func WithParentSpanID(parentID string) EventOption {
	return func(e Event) {
		if corr, ok := e["correlation"].(map[string]any); ok {
			corr["parent_span_id"] = parentID
		}
	}
}

// WithPayload sets the payload on an event.
func WithPayload(payload map[string]any) EventOption {
	return func(e Event) {
		e["payload"] = payload
	}
}

// WithSeq sets the sequence number on an event.
func WithSeq(seq int) EventOption {
	return func(e Event) {
		if event, ok := e["event"].(map[string]any); ok {
			event["seq"] = seq
		}
	}
}

// WithLLMUsage adds LLM usage information to a span.end or run.end payload.
func WithLLMUsage(model string, inputTokens, outputTokens int, costUSD float64) EventOption {
	return func(e Event) {
		if payload, ok := e["payload"].(map[string]any); ok {
			if _, ok := payload["llm"]; !ok {
				payload["llm"] = make(map[string]any)
			}
			if llm, ok := payload["llm"].(map[string]any); ok {
				llm["model"] = model
				llm["usage"] = map[string]any{
					"input_tokens":  inputTokens,
					"output_tokens": outputTokens,
				}
				llm["cost"] = map[string]any{
					"total_usd": costUSD,
				}
			}
		}
	}
}

// WithErrorDetails adds error details to a payload.
func WithErrorDetails(code string, retryable bool) EventOption {
	return func(e Event) {
		if payload, ok := e["payload"].(map[string]any); ok {
			if err, ok := payload["error"].(map[string]any); ok {
				err["code"] = code
				err["retryable"] = retryable
			}
		}
	}
}