api_ingestion.go

package langfuse

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

const (
	ingestionPath  = "/api/public/ingestion"
	maxRetries     = 3
	retryBaseDelay = 200 * time.Millisecond
	flushTimeout   = 15 * time.Second
)

// --- ingestion event types ---

type ingestionEvent struct {
	ID        string `json:"id"`
	Type      string `json:"type"`
	Timestamp string `json:"timestamp"`
	Body      any    `json:"body"`
}

type ingestionRequest struct {
	Batch []ingestionEvent `json:"batch"`
}

type traceBody struct {
	ID        string         `json:"id"`
	Timestamp string         `json:"timestamp,omitempty"`
	Name      string         `json:"name,omitempty"`
	UserID    string         `json:"userId,omitempty"`
	SessionID string         `json:"sessionId,omitempty"`
	Input     any            `json:"input,omitempty"`
	Output    any            `json:"output,omitempty"`
	Metadata  map[string]any `json:"metadata,omitempty"`
	Tags      []string       `json:"tags,omitempty"`
	Public    bool           `json:"public,omitempty"`
	Release   string         `json:"release,omitempty"`
	Version   string         `json:"version,omitempty"`
}

type spanBody struct {
	ID                  string           `json:"id"`
	TraceID             string           `json:"traceId"`
	ParentObservationID string           `json:"parentObservationId,omitempty"`
	Name                string           `json:"name,omitempty"`
	StartTime           string           `json:"startTime,omitempty"`
	EndTime             string           `json:"endTime,omitempty"`
	Input               any              `json:"input,omitempty"`
	Output              any              `json:"output,omitempty"`
	Metadata            map[string]any   `json:"metadata,omitempty"`
	Level               ObservationLevel `json:"level,omitempty"`
	StatusMessage       string           `json:"statusMessage,omitempty"`
	Version             string           `json:"version,omitempty"`
}

type generationBody struct {
	ID                  string            `json:"id"`
	TraceID             string            `json:"traceId"`
	ParentObservationID string            `json:"parentObservationId,omitempty"`
	Name                string            `json:"name,omitempty"`
	StartTime           string            `json:"startTime,omitempty"`
	EndTime             string            `json:"endTime,omitempty"`
	CompletionStartTime string            `json:"completionStartTime,omitempty"`
	Model               string            `json:"model,omitempty"`
	ModelParameters     map[string]string `json:"modelParameters,omitempty"`
	Input               any               `json:"input,omitempty"`
	Output              any               `json:"output,omitempty"`
	Usage               *Usage            `json:"usage,omitempty"`
	Metadata            map[string]any    `json:"metadata,omitempty"`
	Level               ObservationLevel  `json:"level,omitempty"`
	StatusMessage       string            `json:"statusMessage,omitempty"`
	Version             string            `json:"version,omitempty"`
	PromptName          string            `json:"promptName,omitempty"`
	PromptVersion       int               `json:"promptVersion,omitempty"`
}

type eventBody struct {
	ID                  string           `json:"id"`
	TraceID             string           `json:"traceId"`
	ParentObservationID string           `json:"parentObservationId,omitempty"`
	Name                string           `json:"name,omitempty"`
	StartTime           string           `json:"startTime,omitempty"`
	Input               any              `json:"input,omitempty"`
	Output              any              `json:"output,omitempty"`
	Metadata            map[string]any   `json:"metadata,omitempty"`
	Level               ObservationLevel `json:"level,omitempty"`
	StatusMessage       string           `json:"statusMessage,omitempty"`
	Version             string           `json:"version,omitempty"`
}

type scoreBody struct {
	ID            string  `json:"id"`
	TraceID       string  `json:"traceId"`
	Name          string  `json:"name"`
	Value         float64 `json:"value"`
	ObservationID string  `json:"observationId,omitempty"`
	Comment       string  `json:"comment,omitempty"`
	DataType      string  `json:"dataType,omitempty"`
}

// --- buffering & flush ---

// Flush sends all buffered ingestion events immediately. It blocks until the send completes.
func (c *Client) Flush(ctx context.Context) error {
	c.mu.Lock()
	batch := c.takeBufLocked()
	c.mu.Unlock()

	if len(batch) == 0 {
		return nil
	}
	return c.sendBatch(ctx, batch)
}

// Shutdown flushes all pending ingestion events and releases resources.
func (c *Client) Shutdown(ctx context.Context) error {
	c.mu.Lock()
	c.closed = true
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	batch := c.takeBufLocked()
	c.mu.Unlock()

	// Wait for any in-flight sends.
	c.wg.Wait()

	if len(batch) == 0 {
		return nil
	}
	return c.sendBatch(ctx, batch)
}

// enqueue adds an event to the buffer and triggers a flush if needed.
func (c *Client) enqueue(e ingestionEvent) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		log.Printf("langfuse: warning: event enqueued after Shutdown, dropping %s", e.Type)
		return
	}

	c.buf = append(c.buf, e)

	// Drop oldest events if buffer exceeds max size.
	if len(c.buf) > c.cfg.MaxBufferSize {
		c.buf = c.buf[len(c.buf)-c.cfg.MaxBufferSize:]
		log.Printf("langfuse: buffer full (%d), dropped oldest event", c.cfg.MaxBufferSize)
	}

	if len(c.buf) >= c.cfg.FlushBatch {
		batch := c.takeBufLocked()
		c.flushAsync(batch)
		return
	}

	// Start a flush timer if not already running.
	if c.timer == nil {
		c.timer = time.AfterFunc(c.cfg.FlushInterval, func() {
			c.mu.Lock()
			batch := c.takeBufLocked()
			if len(batch) > 0 {
				// Increment wg under the lock so Shutdown's wg.Wait()
				// is guaranteed to see it before returning.
				c.wg.Add(1)
			}
			c.mu.Unlock()
			if len(batch) > 0 {
				go func() {
					defer c.wg.Done()
					c.flushSem <- struct{}{}
					defer func() { <-c.flushSem }()
					ctx, cancel := context.WithTimeout(context.Background(), flushTimeout)
					defer cancel()
					if err := c.sendBatch(ctx, batch); err != nil {
						log.Printf("langfuse: flush error: %v", err)
					}
				}()
			}
		})
	}
}

// takeBufLocked swaps out the buffer and stops the flush timer.
// Must be called with c.mu held.
func (c *Client) takeBufLocked() []ingestionEvent {
	if len(c.buf) == 0 {
		return nil
	}
	batch := c.buf
	c.buf = nil
	if c.timer != nil {
		c.timer.Stop()
		c.timer = nil
	}
	return batch
}

// flushAsync sends a batch in a background goroutine. Must be called with c.mu held
// so that wg.Add is visible to Shutdown before the lock is released.
func (c *Client) flushAsync(batch []ingestionEvent) {
	c.wg.Add(1)
	go func() {
		defer c.wg.Done()
		c.flushSem <- struct{}{}
		defer func() { <-c.flushSem }()
		ctx, cancel := context.WithTimeout(context.Background(), flushTimeout)
		defer cancel()
		if err := c.sendBatch(ctx, batch); err != nil {
			log.Printf("langfuse: flush error: %v", err)
		}
	}()
}

// sendBatch sends a batch of ingestion events with retries for transient errors.
func (c *Client) sendBatch(ctx context.Context, batch []ingestionEvent) error {
	body, err := json.Marshal(ingestionRequest{Batch: batch})
	if err != nil {
		return fmt.Errorf("marshal: %w", err)
	}

	var lastErr error
	for attempt := 0; attempt < maxRetries; attempt++ {
		if attempt > 0 {
			delay := retryBaseDelay * time.Duration(1<<(attempt-1))
			select {
			case <-ctx.Done():
				return fmt.Errorf("retry cancelled: %w (last error: %v)", ctx.Err(), lastErr)
			case <-time.After(delay):
			}
			log.Printf("langfuse: retrying send (attempt %d/%d)", attempt+1, maxRetries)
		}

		lastErr = c.doSend(ctx, body, len(batch))
		if lastErr == nil {
			return nil
		}
		if !isRetriable(lastErr) {
			return lastErr
		}
	}
	return fmt.Errorf("failed after %d attempts: %w", maxRetries, lastErr)
}

// doSend performs a single HTTP POST to the ingestion endpoint.
func (c *Client) doSend(ctx context.Context, body []byte, batchLen int) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.Host+ingestionPath, bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("new request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	req.SetBasicAuth(c.cfg.PublicKey, c.cfg.SecretKey)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return &retriableError{err: fmt.Errorf("http: %w", err)}
	}
	defer resp.Body.Close()

	respBody, _ := io.ReadAll(resp.Body)

	if resp.StatusCode >= 500 {
		return &retriableError{err: fmt.Errorf("langfuse API returned %d: %s", resp.StatusCode, truncate(respBody, 512))}
	}
	if resp.StatusCode == http.StatusTooManyRequests {
		return &retriableError{err: fmt.Errorf("langfuse API returned 429: %s", truncate(respBody, 512))}
	}
	if resp.StatusCode >= 400 {
		return fmt.Errorf("langfuse API returned %d: %s", resp.StatusCode, truncate(respBody, 512))
	}

	// Log partial failures from 207 Multi-Status responses.
	if resp.StatusCode == http.StatusMultiStatus && len(respBody) > 0 {
		var result struct {
			Errors []any `json:"errors"`
		}
		if json.Unmarshal(respBody, &result) == nil && len(result.Errors) > 0 {
			log.Printf("langfuse: ingestion partial failure (%d events, %d errors): %s",
				batchLen, len(result.Errors), truncate(respBody, 512))
		}
	}

	return nil
}

// --- retriable error ---

// retriableError wraps transient errors (5xx, 429, network) to signal that
// sendBatch should retry the request.
type retriableError struct {
	err error
}

func (e *retriableError) Error() string { return e.err.Error() }
func (e *retriableError) Unwrap() error { return e.err }

// isRetriable reports whether err is a transient failure worth retrying.
func isRetriable(err error) bool {
	var re *retriableError
	return errors.As(err, &re)
}

// event wraps a body into an ingestion event and enqueues it for sending.
func (c *Client) event(typ string, body any) {
	c.enqueue(ingestionEvent{
		ID:        newID(),
		Type:      typ,
		Timestamp: now(),
		Body:      body,
	})
}