poimen-workflows/internal/batching/llm_batch.go

package batching

import (
	"fmt"
	"sync"
	"time"
)

// LLMRequest represents a single LLM request to be batched
type LLMRequest struct {
	ID        string                 `json:"id"`
	Type      string                 `json:"type"` // "implementer", "judge", "planner"
	Model     string                 `json:"model"`
	Prompt    string                 `json:"prompt"`
	Metadata  map[string]interface{} `json:"metadata,omitempty"`
	Timestamp time.Time              `json:"timestamp"`
	ResultCh  chan *LLMResult        `json:"-"`
}

// LLMResult represents the result of a single LLM request
type LLMResult struct {
	RequestID  string                 `json:"request_id"`
	Response   string                 `json:"response"`
	Error      error                  `json:"error,omitempty"`
	Duration   time.Duration          `json:"duration"`
	TokenCount int                    `json:"token_count"`
	Metadata   map[string]interface{} `json:"metadata,omitempty"`
}

// LLMBatch represents a batch of LLM requests
type LLMBatch struct {
	ID           string
	Requests     []*LLMRequest
	Model        string
	Type         string
	CreatedAt    time.Time
	ExecutedAt   time.Time
	Status       string // "pending", "executing", "completed", "failed"
	Error        error
	Results      map[string]*LLMResult
	ExecutionTime time.Duration
}

// LLMBatcher batches LLM requests for efficient API usage
type LLMBatcher struct {
	mu              sync.RWMutex
	queue           []*LLMRequest
	maxBatchSize    int
	maxBatchAge     time.Duration
	lastFlushTime   time.Time
	executedBatches []*LLMBatch
	pendingBatches  []*LLMBatch
	stats           *LLMBatchStats
}

// LLMBatchStats tracks LLM batching statistics
type LLMBatchStats struct {
	TotalRequests      int
	TotalBatches       int
	AvgRequestsPerBatch float64
	APICallsSaved      int // Total API calls saved (individual requests - batches)
	TotalTokens        int
	TotalExecutionTime time.Duration
}

// NewLLMBatcher creates a new LLM batcher
func NewLLMBatcher(maxBatchSize int, maxBatchAge time.Duration) *LLMBatcher {
	if maxBatchSize <= 0 {
		maxBatchSize = 10
	}
	if maxBatchAge <= 0 {
		maxBatchAge = 2 * time.Second
	}

	return &LLMBatcher{
		queue:           make([]*LLMRequest, 0),
		maxBatchSize:    maxBatchSize,
		maxBatchAge:     maxBatchAge,
		lastFlushTime:   time.Now(),
		executedBatches: make([]*LLMBatch, 0),
		pendingBatches:  make([]*LLMBatch, 0),
		stats: &LLMBatchStats{
			TotalRequests: 0,
			TotalBatches:  0,
		},
	}
}

// Enqueue adds an LLM request to the queue
func (lb *LLMBatcher) Enqueue(req *LLMRequest) {
	if req == nil {
		return
	}

	if req.ID == "" {
		req.ID = fmt.Sprintf("req-%d", time.Now().UnixNano())
	}

	req.Timestamp = time.Now()
	if req.ResultCh == nil {
		req.ResultCh = make(chan *LLMResult, 1)
	}

	lb.mu.Lock()
	defer lb.mu.Unlock()

	lb.queue = append(lb.queue, req)
	lb.stats.TotalRequests++

	// Auto-flush if batch is full
	if len(lb.queue) >= lb.maxBatchSize {
		lb.flushLocked()
	}
}

// flushLocked creates a batch from queued requests (must be called with lock held)
func (lb *LLMBatcher) flushLocked() {
	if len(lb.queue) == 0 {
		return
	}

	// Group by type and model
	groups := make(map[string][]*LLMRequest)
	for _, req := range lb.queue {
		key := fmt.Sprintf("%s:%s", req.Type, req.Model)
		groups[key] = append(groups[key], req)
	}

	// Create batch for each group
	for key, reqs := range groups {
		batch := &LLMBatch{
			ID:        fmt.Sprintf("batch-%d", lb.stats.TotalBatches),
			Requests:  reqs,
			Model:     reqs[0].Model,
			Type:      reqs[0].Type,
			CreatedAt: time.Now(),
			Status:    "pending",
			Results:   make(map[string]*LLMResult),
		}

		lb.pendingBatches = append(lb.pendingBatches, batch)
		lb.stats.TotalBatches++
		_ = key // Silence unused variable warning
	}

	lb.queue = make([]*LLMRequest, 0)
	lb.lastFlushTime = time.Now()
}

// Flush manually flushes the current queue
func (lb *LLMBatcher) Flush() {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	lb.flushLocked()
}

// GetPendingBatch returns the next pending batch without removing it
func (lb *LLMBatcher) GetPendingBatch() *LLMBatch {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	if len(lb.pendingBatches) == 0 {
		return nil
	}

	return lb.pendingBatches[0]
}

// MarkBatchExecuting marks a batch as executing
func (lb *LLMBatcher) MarkBatchExecuting(batchID string) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	for _, batch := range lb.pendingBatches {
		if batch.ID == batchID {
			batch.Status = "executing"
			break
		}
	}
}

// MarkBatchCompleted marks a batch as completed and delivers results
func (lb *LLMBatcher) MarkBatchCompleted(batchID string, results map[string]*LLMResult) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	var idx int
	var found *LLMBatch
	for i, batch := range lb.pendingBatches {
		if batch.ID == batchID {
			idx = i
			found = batch
			break
		}
	}

	if found != nil {
		found.Status = "completed"
		found.ExecutedAt = time.Now()
		found.ExecutionTime = found.ExecutedAt.Sub(found.CreatedAt)
		found.Results = results

		// Deliver results to request channels
		for _, req := range found.Requests {
			if result, exists := results[req.ID]; exists {
				select {
				case req.ResultCh <- result:
				default:
					// Channel not ready or closed
				}
			}
		}

		// Update stats
		lb.stats.TotalTokens += countTokensInBatch(found)
		lb.stats.TotalExecutionTime += found.ExecutionTime

		// Move to executed batches
		lb.executedBatches = append(lb.executedBatches, found)
		lb.pendingBatches = append(lb.pendingBatches[:idx], lb.pendingBatches[idx+1:]...)
	}
}

// MarkBatchFailed marks a batch as failed
func (lb *LLMBatcher) MarkBatchFailed(batchID string, err error) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	var found *LLMBatch
	for _, batch := range lb.pendingBatches {
		if batch.ID == batchID {
			found = batch
			break
		}
	}

	if found != nil {
		found.Status = "failed"
		found.Error = err
		found.ExecutedAt = time.Now()

		// Deliver errors to request channels
		for _, req := range found.Requests {
			result := &LLMResult{
				RequestID: req.ID,
				Error:     err,
			}

			select {
			case req.ResultCh <- result:
			default:
				// Channel not ready or closed
			}
		}
	}
}

// GetStats returns batching statistics
func (lb *LLMBatcher) GetStats() *LLMBatchStats {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	stats := *lb.stats
	if stats.TotalBatches > 0 {
		stats.AvgRequestsPerBatch = float64(stats.TotalRequests) / float64(stats.TotalBatches)
		// API calls saved: total requests - total batches
		stats.APICallsSaved = stats.TotalRequests - stats.TotalBatches
	}

	return &stats
}

// QueueSize returns current queue size
func (lb *LLMBatcher) QueueSize() int {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	return len(lb.queue)
}

// PendingBatchCount returns number of pending batches
func (lb *LLMBatcher) PendingBatchCount() int {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	return len(lb.pendingBatches)
}

// GetBatchByID returns a batch by ID
func (lb *LLMBatcher) GetBatchByID(batchID string) *LLMBatch {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	for _, batch := range lb.pendingBatches {
		if batch.ID == batchID {
			return batch
		}
	}

	for _, batch := range lb.executedBatches {
		if batch.ID == batchID {
			return batch
		}
	}

	return nil
}

// TimeSinceLastFlush returns time since last flush
func (lb *LLMBatcher) TimeSinceLastFlush() time.Duration {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	return time.Since(lb.lastFlushTime)
}

// ShouldFlush checks if queue should be flushed based on age
func (lb *LLMBatcher) ShouldFlush() bool {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	if len(lb.queue) == 0 {
		return false
	}

	return time.Since(lb.lastFlushTime) >= lb.maxBatchAge
}

// GetExecutedBatches returns all executed batches
func (lb *LLMBatcher) GetExecutedBatches() []*LLMBatch {
	lb.mu.RLock()
	defer lb.mu.RUnlock()

	result := make([]*LLMBatch, len(lb.executedBatches))
	copy(result, lb.executedBatches)

	return result
}

// Clear clears all pending operations
func (lb *LLMBatcher) Clear() {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	lb.queue = make([]*LLMRequest, 0)
	lb.pendingBatches = make([]*LLMBatch, 0)
	lb.executedBatches = make([]*LLMBatch, 0)
}

// countTokensInBatch counts total tokens in a batch
func countTokensInBatch(batch *LLMBatch) int {
	total := 0
	for _, result := range batch.Results {
		total += result.TokenCount
	}
	return total
}

// GetBatchInfo returns human-readable batch information
func (batch *LLMBatch) GetInfo() map[string]interface{} {
	return map[string]interface{}{
		"id":            batch.ID,
		"type":          batch.Type,
		"model":         batch.Model,
		"status":        batch.Status,
		"request_count": len(batch.Requests),
		"created_at":    batch.CreatedAt,
		"executed_at":   batch.ExecutedAt,
		"duration":      batch.ExecutionTime,
		"error":         batch.Error,
	}
}