Maven/internal/llm/client.go

// Package llm is the shared llama-server completion client — one seam both the
// phraser (talking back) and the router (routing) call. It does NOT spawn the
// server; the daemon owns one llama-server (spawned by the phraser) and hands
// its base URL here, so a single resident model serves both callers.
package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"
)

// Gate — admission control for a completion. Enter blocks or refuses while the
// resident model is being swapped, and the returned release says the request is
// done. The phraser implements it: a swap kills the running llama-server, so
// every holder of a base URL has to be counted before the kill, not just the
// phrasing paths.
//
// Without this the drain saw only the phraser's own calls. The LLM router, the
// replier, the mail extractor and the memory evaluator all reach llama-server
// through this client, so a swap could report zero requests in flight and kill
// the server out from under a routing decision. The turn then finished on the
// new model, which is the "half of one model and half of another" the swap is
// supposed to make impossible.
type Gate interface {
	Enter() (release func(), err error)
}

type Client struct {
	// mu guards base and gate. The base URL can change when the daemon swaps the
	// resident model (Vikunja #250) and every holder of this client — the LLM
	// router, the replier, the mail extractor — must follow without being
	// rebuilt. A swap re-points the client, it does not replace it.
	//
	// On the deploy shape the new server binds the same fixed port the killed
	// one released (startLlamaProc passes the port out of phraser.listen), so
	// SetBaseURL is normally a no-op and the gate is the part doing the work.
	// The re-pointing stays because nothing guarantees the port: a phraser
	// listening on :0, or a future swap that moves the server, changes the base.
	mu   sync.RWMutex
	base string
	gate Gate
	http *http.Client
}

// SetGate installs the admission gate. Nil (the default, and what the eval
// harness and the tests use) means no gating.
func (c *Client) SetGate(g Gate) {
	c.mu.Lock()
	c.gate = g
	c.mu.Unlock()
}

func (c *Client) enter() (func(), error) {
	c.mu.RLock()
	g := c.gate
	c.mu.RUnlock()
	if g == nil {
		return func() {}, nil
	}
	return g.Enter()
}

func New(baseURL string, timeout time.Duration) *Client {
	return &Client{base: baseURL, http: &http.Client{Timeout: timeout}}
}

// SetBaseURL re-points the client at another llama-server. Safe to call while
// requests are in flight: a request that already read the old base finishes
// against the old base (or fails, and every caller of Complete has a fallback),
// and the next one uses the new base. It is deliberately NOT a queue-and-retry —
// the phraser quiesces around a swap, so the window is small and a lost turn
// degrades to the classifier rather than hanging.
func (c *Client) SetBaseURL(base string) {
	c.mu.Lock()
	c.base = base
	c.mu.Unlock()
}

// BaseURL is the server this client currently talks to.
func (c *Client) BaseURL() string {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.base
}

type Req struct {
	System    string
	User      string
	Grammar   string // GBNF; empty ⇒ unconstrained
	MaxTokens int
	// RepeatPenalty > 0 ⇒ penalize token repetition (curbs the sub-1B "тоже
	// тоже тоже" loop). 0 ⇒ server default (no extra penalty).
	RepeatPenalty float64
	// Stop — sequences that end generation early (e.g. newline for a one-liner).
	Stop []string
}

type msg struct {
	Role             string `json:"role"`
	Content          string `json:"content"`
	Reasoning        string `json:"reasoning,omitempty"`
	ReasoningContent string `json:"reasoning_content,omitempty"`
}
type body struct {
	Messages      []msg    `json:"messages"`
	MaxTokens     int      `json:"max_tokens,omitempty"`
	Grammar       string   `json:"grammar,omitempty"`
	Temp          float64  `json:"temperature"`
	RepeatPenalty float64  `json:"repeat_penalty,omitempty"`
	Stop          []string `json:"stop,omitempty"`
}
type resp struct {
	Choices []struct {
		Message msg `json:"message"`
	} `json:"choices"`
}

func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
	release, err := c.enter()
	if err != nil {
		return "", err
	}
	defer release()
	b, _ := json.Marshal(body{
		Messages:      []msg{{Role: "system", Content: r.System}, {Role: "user", Content: r.User}},
		MaxTokens:     r.MaxTokens,
		Grammar:       r.Grammar,
		Temp:          0,
		RepeatPenalty: r.RepeatPenalty,
		Stop:          r.Stop,
	})
	req, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL()+"/v1/chat/completions", bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	httpResp, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer httpResp.Body.Close()
	if httpResp.StatusCode != 200 {
		return "", fmt.Errorf("llm: status %d", httpResp.StatusCode)
	}
	var out resp
	if err := json.NewDecoder(httpResp.Body).Decode(&out); err != nil {
		return "", err
	}
	if len(out.Choices) == 0 {
		return "", fmt.Errorf("llm: no choices")
	}
	content := out.Choices[0].Message.Content
	if content == "" {
		content = out.Choices[0].Message.ReasoningContent
	}
	return content, nil
}