192 lines
6.5 KiB
Go
192 lines
6.5 KiB
Go
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type server struct {
|
|
model string
|
|
delay time.Duration
|
|
streamDelay time.Duration
|
|
responseBytes int
|
|
promptTokens int64
|
|
outputTokens int64
|
|
failEvery int64
|
|
requests atomic.Int64
|
|
}
|
|
|
|
func main() {
|
|
listen := flag.String("listen", "127.0.0.1:11435", "listen address")
|
|
model := flag.String("model", "qwen3:8b", "mock model name")
|
|
delay := flag.Duration("delay", 0, "delay before response headers")
|
|
streamDelay := flag.Duration("stream-delay", 0, "delay between streaming chunks")
|
|
responseBytes := flag.Int("response-bytes", 128, "approximate generated content bytes")
|
|
promptTokens := flag.Int64("prompt-tokens", 16, "reported prompt token count")
|
|
outputTokens := flag.Int64("output-tokens", 8, "reported completion token count")
|
|
failEvery := flag.Int64("fail-every", 0, "return HTTP 503 for every Nth inference request; 0 disables")
|
|
flag.Parse()
|
|
if *responseBytes < 0 || *promptTokens < 0 || *outputTokens < 0 || *failEvery < 0 {
|
|
log.Fatal("numeric flags must be non-negative")
|
|
}
|
|
s := &server{model: *model, delay: *delay, streamDelay: *streamDelay, responseBytes: *responseBytes, promptTokens: *promptTokens, outputTokens: *outputTokens, failEvery: *failEvery}
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/ps", s.ps)
|
|
mux.HandleFunc("/api/tags", s.tags)
|
|
mux.HandleFunc("/api/show", s.show)
|
|
mux.HandleFunc("/api/chat", s.nativeChat)
|
|
mux.HandleFunc("/api/generate", s.nativeGenerate)
|
|
mux.HandleFunc("/v1/chat/completions", s.openAIChat)
|
|
mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNoContent) })
|
|
h := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 2 * time.Minute}
|
|
log.Printf("mock Ollama listening on http://%s model=%s delay=%s response_bytes=%d", *listen, *model, *delay, *responseBytes)
|
|
log.Fatal(h.ListenAndServe())
|
|
}
|
|
|
|
func (s *server) ps(w http.ResponseWriter, r *http.Request) {
|
|
writeJSON(w, map[string]any{"models": []map[string]any{{"name": s.model, "model": s.model, "size": 1 << 30, "size_vram": 1 << 30, "context_length": 32768}}})
|
|
}
|
|
|
|
func (s *server) tags(w http.ResponseWriter, r *http.Request) {
|
|
writeJSON(w, map[string]any{"models": []map[string]any{{"name": s.model, "model": s.model, "size": 1 << 30, "details": map[string]any{"parameter_size": "8B", "quantization_level": "Q4_K_M"}}}})
|
|
}
|
|
|
|
func (s *server) show(w http.ResponseWriter, r *http.Request) {
|
|
writeJSON(w, map[string]any{"capabilities": []string{"completion", "tools", "thinking", "vision"}, "model_info": map[string]any{"mock.context_length": 32768}})
|
|
}
|
|
|
|
func (s *server) shouldFail() bool {
|
|
n := s.requests.Add(1)
|
|
return s.failEvery > 0 && n%s.failEvery == 0
|
|
}
|
|
|
|
func (s *server) openAIChat(w http.ResponseWriter, r *http.Request) {
|
|
if s.delay > 0 {
|
|
time.Sleep(s.delay)
|
|
}
|
|
if s.shouldFail() {
|
|
writeJSONStatus(w, http.StatusServiceUnavailable, map[string]any{"error": map[string]any{"message": "mock failure", "type": "server_error"}})
|
|
return
|
|
}
|
|
var req struct {
|
|
Stream bool `json:"stream"`
|
|
}
|
|
_ = json.NewDecoder(io.LimitReader(r.Body, 8<<20)).Decode(&req)
|
|
content := strings.Repeat("x", s.responseBytes)
|
|
if req.Stream {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(http.StatusOK)
|
|
f, _ := w.(http.Flusher)
|
|
parts := split(content, 4)
|
|
for i, p := range parts {
|
|
chunk := map[string]any{"id": "chatcmpl-mock", "object": "chat.completion.chunk", "choices": []map[string]any{{"index": 0, "delta": map[string]any{"content": p}, "finish_reason": nil}}}
|
|
if i == len(parts)-1 {
|
|
chunk["choices"] = []map[string]any{{"index": 0, "delta": map[string]any{}, "finish_reason": "stop"}}
|
|
}
|
|
b, _ := json.Marshal(chunk)
|
|
fmt.Fprintf(w, "data: %s\n\n", b)
|
|
if f != nil {
|
|
f.Flush()
|
|
}
|
|
if s.streamDelay > 0 {
|
|
time.Sleep(s.streamDelay)
|
|
}
|
|
}
|
|
io.WriteString(w, "data: [DONE]\n\n")
|
|
if f != nil {
|
|
f.Flush()
|
|
}
|
|
return
|
|
}
|
|
writeJSON(w, map[string]any{
|
|
"id": "chatcmpl-mock", "object": "chat.completion", "created": time.Now().Unix(), "model": s.model,
|
|
"choices": []map[string]any{{"index": 0, "message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop"}},
|
|
"usage": map[string]any{"prompt_tokens": s.promptTokens, "completion_tokens": s.outputTokens, "total_tokens": s.promptTokens + s.outputTokens},
|
|
})
|
|
}
|
|
|
|
func (s *server) nativeChat(w http.ResponseWriter, r *http.Request) {
|
|
s.native(w, r, "message")
|
|
}
|
|
|
|
func (s *server) nativeGenerate(w http.ResponseWriter, r *http.Request) {
|
|
s.native(w, r, "response")
|
|
}
|
|
|
|
func (s *server) native(w http.ResponseWriter, r *http.Request, field string) {
|
|
if s.delay > 0 {
|
|
time.Sleep(s.delay)
|
|
}
|
|
if s.shouldFail() {
|
|
writeJSONStatus(w, http.StatusServiceUnavailable, map[string]any{"error": "mock failure"})
|
|
return
|
|
}
|
|
var req map[string]any
|
|
_ = json.NewDecoder(io.LimitReader(r.Body, 8<<20)).Decode(&req)
|
|
stream, _ := req["stream"].(bool)
|
|
content := strings.Repeat("x", s.responseBytes)
|
|
w.Header().Set("Content-Type", "application/x-ndjson")
|
|
bw := bufio.NewWriter(w)
|
|
if !stream {
|
|
doc := map[string]any{"model": s.model, "done": true, "prompt_eval_count": s.promptTokens, "eval_count": s.outputTokens}
|
|
if field == "message" {
|
|
doc[field] = map[string]any{"role": "assistant", "content": content}
|
|
} else {
|
|
doc[field] = content
|
|
}
|
|
_ = json.NewEncoder(bw).Encode(doc)
|
|
_ = bw.Flush()
|
|
return
|
|
}
|
|
f, _ := w.(http.Flusher)
|
|
for _, p := range split(content, 4) {
|
|
doc := map[string]any{"model": s.model, "done": false}
|
|
if field == "message" {
|
|
doc[field] = map[string]any{"role": "assistant", "content": p}
|
|
} else {
|
|
doc[field] = p
|
|
}
|
|
_ = json.NewEncoder(bw).Encode(doc)
|
|
_ = bw.Flush()
|
|
if f != nil {
|
|
f.Flush()
|
|
}
|
|
if s.streamDelay > 0 {
|
|
time.Sleep(s.streamDelay)
|
|
}
|
|
}
|
|
_ = json.NewEncoder(bw).Encode(map[string]any{"model": s.model, "done": true, "prompt_eval_count": s.promptTokens, "eval_count": s.outputTokens})
|
|
_ = bw.Flush()
|
|
if f != nil {
|
|
f.Flush()
|
|
}
|
|
}
|
|
|
|
func split(s string, n int) []string {
|
|
if n <= 1 || len(s) == 0 {
|
|
return []string{s}
|
|
}
|
|
out := make([]string, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
start := len(s) * i / n
|
|
end := len(s) * (i + 1) / n
|
|
out = append(out, s[start:end])
|
|
}
|
|
return out
|
|
}
|
|
|
|
func writeJSON(w http.ResponseWriter, v any) { writeJSONStatus(w, http.StatusOK, v) }
|
|
func writeJSONStatus(w http.ResponseWriter, status int, v any) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(status)
|
|
_ = json.NewEncoder(w).Encode(v)
|
|
}
|