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
	categorizeConfidence float64
	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")
	// Lets a local run reproduce a low-confidence taxonomy pick (see
	// processing.DefaultMinCategorizeConfidence).
	categorizeConfidence := flag.Float64("categorize-confidence", 0,
		"when > 0, answer categorize calls with this confidence score")
	flag.Parse()

	s := &server{
		apiKey:               strings.TrimSpace(*key),
		model:                strings.TrimSpace(*model),
		categorizeConfidence: *categorizeConfidence,
	}
	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 truncTrace(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n]
}

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)
	if os.Getenv("MOCK_LLM_TRACE") != "" {
		log.Printf("mock-llm: roles=%d system_len=%d user_len=%d system_head=%q",
			len(req.Messages), len(system), len(user), truncTrace(system, 60))
	}
	var comp processing.Completion
	// A category formula in the prompt is answered in that formula's shape, so a
	// local run proves the formula reached the model and the reply passes the
	// pipeline's formula gate. Everything else keeps the heuristic fallback.
	if text, ok := s.categorizeReply(system, user); ok {
		comp = processing.Completion{Text: text}
	} else if text, ok := buildFormulaReply(system, user); ok {
		comp = processing.Completion{Text: text}
	} else {
		var err error
		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},
	})
}

// categorizeReply overrides the heuristic taxonomy pick so a local run can
// reproduce a low-confidence answer.
func (s *server) categorizeReply(system, user string) (string, bool) {
	if s.categorizeConfidence <= 0 || !strings.Contains(strings.ToLower(system), "categoryid") {
		return "", false
	}
	id := firstCategoryIDFromPrompt(user)
	if id == "" {
		return "", false
	}
	b, err := json.Marshal(map[string]any{"categoryId": id, "confidence": s.categorizeConfidence})
	if err != nil {
		return "", false
	}
	return string(b), true
}

// firstCategoryIDFromPrompt mirrors the "(ID: …)" format of the categorize prompt.
func firstCategoryIDFromPrompt(user string) string {
	const marker = "(id:"
	lower := strings.ToLower(user)
	at := strings.Index(lower, marker)
	if at < 0 {
		return ""
	}
	rest := strings.TrimSpace(user[at+len(marker):])
	end := strings.Index(rest, ")")
	if end <= 0 {
		return ""
	}
	return strings.TrimSpace(rest[:end])
}

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)) {
		// Reasoning-model requests send the system message as "developer"
		// (see processing.OpenAIClient sysRole). Treating it as a user turn made the
		// mock mis-route every call whose branch depends on the system prompt.
		case "system", "developer":
			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)
	}
}