descrybe/apps/api/cmd/mock-llm/main.go

// Command mock-llm serves a tiny OpenAI-compatible Chat Completions API for
// local/CI processing proofs. It reuses processing.HeuristicCompleter so
// enhance JSON shapes match the offline fallback, while exercising the real
// OpenAIClient HTTP path (no production API keys).
//
// Usage:
//
//	go run ./cmd/mock-llm -addr 127.0.0.1:18767
//	go run ./cmd/mock-llm -addr 127.0.0.1:18767 -key local-test -model mock-llm
//
// Then point platform env (or /integrations/ai) at:
//
//	OPENAI_API_KEY=local-test
//	OPENAI_BASE_URL=http://127.0.0.1:18767/v1
//	OPENAI_MODEL=mock-llm
package main

import (
	"context"
	"encoding/json"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync/atomic"
	"time"

	"github.com/descrybe/descrybe-v2/apps/api/internal/processing"
)

type server struct {
	apiKey string
	model  string
	logReq atomic.Int64
}

type chatRequest struct {
	Model       string `json:"model"`
	Messages    []struct {
		Role    string `json:"role"`
		Content string `json:"content"`
	} `json:"messages"`
	Temperature float64 `json:"temperature"`
	MaxTokens   int     `json:"max_tokens"`
}

func main() {
	addr := flag.String("addr", "127.0.0.1:18767", "listen address")
	key := flag.String("key", envOr("MOCK_LLM_API_KEY", "local-test"), "Bearer API key (non-empty placeholder)")
	model := flag.String("model", envOr("MOCK_LLM_MODEL", "mock-llm"), "model id returned by /v1/models")
	flag.Parse()

	s := &server{
		apiKey: strings.TrimSpace(*key),
		model:  strings.TrimSpace(*model),
	}
	if s.apiKey == "" {
		log.Fatal("mock-llm: API key must be non-empty (Descrybe Completer.Enabled requires it)")
	}
	if s.model == "" {
		s.model = "mock-llm"
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", s.handleHealth)
	mux.HandleFunc("/v1/models", s.handleModels)
	mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions)
	mux.HandleFunc("/v1/embeddings", s.handleEmbeddings)

	log.Printf("mock-llm listening on http://%s", *addr)
	log.Printf("OpenAI base: http://%s/v1  model=%s  key=<redacted>", *addr, s.model)
	log.Printf("Wire: OPENAI_BASE_URL=http://%s/v1 OPENAI_API_KEY=<redacted> OPENAI_MODEL=%s", *addr, s.model)
	if err := http.ListenAndServe(*addr, mux); err != nil {
		log.Fatal(err)
	}
}

func envOr(k, def string) string {
	if v := strings.TrimSpace(os.Getenv(k)); v != "" {
		return v
	}
	return def
}

func (s *server) handleHealth(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{
		"status":  "ok",
		"service": "mock-llm",
		"model":   s.model,
	})
}

func (s *server) authOK(r *http.Request) bool {
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, "Bearer ") {
		return false
	}
	token := strings.TrimSpace(strings.TrimPrefix(h, "Bearer "))
	return token == s.apiKey
}

func (s *server) handleModels(w http.ResponseWriter, r *http.Request) {
	s.logReq.Add(1)
	if r.Method != http.MethodGet {
		http.Error(w, `{"error":{"message":"method not allowed"}}`, http.StatusMethodNotAllowed)
		return
	}
	if !s.authOK(r) {
		http.Error(w, `{"error":{"message":"unauthorized"}}`, http.StatusUnauthorized)
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{
		"object": "list",
		"data": []map[string]any{
			{"id": s.model, "object": "model", "owned_by": "descrybe-mock"},
		},
	})
}

func (s *server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	s.logReq.Add(1)
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":{"message":"method not allowed"}}`, http.StatusMethodNotAllowed)
		return
	}
	if !s.authOK(r) {
		http.Error(w, `{"error":{"message":"unauthorized"}}`, http.StatusUnauthorized)
		return
	}
	body, err := io.ReadAll(io.LimitReader(r.Body, 2<<20))
	if err != nil {
		http.Error(w, `{"error":{"message":"read body"}}`, http.StatusBadRequest)
		return
	}
	var req chatRequest
	if err := json.Unmarshal(body, &req); err != nil {
		http.Error(w, `{"error":{"message":"invalid json"}}`, http.StatusBadRequest)
		return
	}
	system, user := splitMessages(req.Messages)
	comp, err := processing.HeuristicCompleter{}.Complete(context.Background(), system, user)
	if err != nil {
		http.Error(w, `{"error":{"message":"completer failed"}}`, http.StatusInternalServerError)
		return
	}
	model := strings.TrimSpace(req.Model)
	if model == "" {
		model = s.model
	}
	promptTokens := estimateTokens(system) + estimateTokens(user)
	completionTokens := estimateTokens(comp.Text)
	writeJSON(w, http.StatusOK, map[string]any{
		"id":      "chatcmpl-mock",
		"object":  "chat.completion",
		"created": time.Now().Unix(),
		"model":   model,
		"choices": []map[string]any{
			{
				"index": 0,
				"message": map[string]any{
					"role":    "assistant",
					"content": comp.Text,
				},
				"finish_reason": "stop",
			},
		},
		"usage": map[string]any{
			"prompt_tokens":     promptTokens,
			"completion_tokens": completionTokens,
			"total_tokens":      promptTokens + completionTokens,
		},
	})
}

func (s *server) handleEmbeddings(w http.ResponseWriter, r *http.Request) {
	s.logReq.Add(1)
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":{"message":"method not allowed"}}`, http.StatusMethodNotAllowed)
		return
	}
	if !s.authOK(r) {
		http.Error(w, `{"error":{"message":"unauthorized"}}`, http.StatusUnauthorized)
		return
	}
	// Tiny fixed vector — enough for platform role probe / CI smoke.
	writeJSON(w, http.StatusOK, map[string]any{
		"object": "list",
		"model":  s.model + "-embed",
		"data": []map[string]any{
			{"object": "embedding", "index": 0, "embedding": []float32{0.01, 0.02, 0.03, 0.04}},
		},
		"usage": map[string]any{"prompt_tokens": 1, "total_tokens": 1},
	})
}

func splitMessages(msgs []struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}) (system, user string) {
	var users []string
	for _, m := range msgs {
		switch strings.ToLower(strings.TrimSpace(m.Role)) {
		case "system":
			if system == "" {
				system = m.Content
			} else {
				system += "\n" + m.Content
			}
		case "user":
			users = append(users, m.Content)
		case "assistant":
			// ignore prior assistant turns in this stub
		}
	}
	return system, strings.Join(users, "\n")
}

func estimateTokens(s string) int {
	n := len(strings.Fields(s))
	if n < 1 && strings.TrimSpace(s) != "" {
		return 1
	}
	return n
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		log.Printf("encode: %v", err)
	}
}