llmgw/internal/router/rag.go

package router

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	storepkg "gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Answering from knowledge bases (BACKLOG F3): a chat request that names
// collections in X-Llmgw-Knowledge has the passages its last question needs
// found in them and put in front of the model, with their sources, before it
// is judged and sent. Which collections it may name is decided by the key's
// place in the tree, exactly as llmgw-kb's own requests are; one it may not
// read is answered as one that does not exist.

const (
	// knowledgeHeader names the collections, by name or id, comma separated.
	knowledgeHeader = "X-Llmgw-Knowledge"
	// knowledgeKHeader is how many passages; kbDefaultK when absent.
	knowledgeKHeader = "X-Llmgw-Knowledge-K"
	// sourcesHeader answers which passages were given: document#chunk, each
	// path-escaped, in the order they were numbered.
	sourcesHeader = "x-llmgw-sources"

	kbDefaultK = 5
	kbMaxK     = 20
)

// kbContextIntro is what the model is told about the passages. In English,
// which every model reads; it is asked to answer in the question's language.
const kbContextIntro = "Below are passages from the organisation's own documents, found for the user's last " +
	"question. Answer from them, and cite the ones you use as [1], [2] and so on. If they do not hold the " +
	"answer, say so rather than guessing. Answer in the language of the question."

const kbNothingFound = "The user asked a question meant to be answered from the organisation's documents, " +
	"and nothing in them matched it. Say so, rather than answering from general knowledge as if it were " +
	"the organisation's."

// kbEffective is a caller's place in the tree for reading collections: a
// personal key has one only in a project it names and belongs to.
func (r *Router) kbEffective(ctx context.Context, p auth.Principal, project string) (storepkg.Effective, error) {
	e := storepkg.Effective{KeyID: p.KeyID, UserID: p.UserID, ProjectID: p.ProjectID, TeamID: p.TeamID, OrgID: p.OrgID}
	project = strings.TrimSpace(project)
	if project == "" || p.ProjectID != "" || p.UserID == "" {
		return e, nil
	}
	may, err := r.store.CanUseProject(ctx, p.UserID, project)
	if err != nil {
		return e, err
	}
	if !may {
		return e, &kbError{status: http.StatusForbidden, kind: "permission_error", msg: "you are not a member of project " + project}
	}
	eff, err := r.store.ResolveFor(ctx, p.KeyID, p.UserID, project)
	if err != nil {
		return e, err
	}
	e.ProjectID, e.TeamID, e.OrgID = eff.ProjectID, eff.TeamID, eff.OrgID
	return e, nil
}

// kbError is a refusal with the status and kind a client is answered with.
type kbError struct {
	status     int
	kind, msg  string
	retryAfter string
}

func (e *kbError) Error() string { return e.msg }

// retrieve finds and adds the passages a request asked for. It answers false
// when it has answered the client itself.
func (r *Router) retrieve(ctx context.Context, w http.ResponseWriter, req *parsedRequest, start time.Time) bool {
	names := splitList(req.Header.Get(knowledgeHeader))
	if len(names) == 1 && strings.EqualFold(names[0], "none") {
		// Asked for no passages, whatever the default.
		return true
	}
	// Without the header, the default of the caller's part of the tree, if
	// it has one. A default is the administrator's choice, not the
	// caller's, so it never refuses a request: a collection the caller
	// cannot read is left out, and a search that fails is answered without
	// passages.
	explicit := len(names) > 0
	if !explicit {
		if req.Path != "/v1/chat/completions" || req.internal || r.store == nil || r.store.KnowledgeAvailable() != nil {
			return true
		}
		p, ok := auth.FromContext(ctx)
		if !ok {
			return true
		}
		if names = r.kbDefaults(ctx, p, req.Header.Get(projectHeader)); len(names) == 0 {
			return true
		}
	}
	fail := func(err error) bool {
		if !explicit {
			r.log.Warn("a knowledge-base default could not be searched; answering without it", "model", req.Model, "err", err)
			r.m.kbRetrievals.WithLabelValues("skipped").Inc()
			return true
		}
		var ke *kbError
		if !errors.As(err, &ke) {
			ke = &kbError{status: http.StatusInternalServerError, kind: "server_error", msg: err.Error()}
		}
		if ke.retryAfter != "" {
			w.Header().Set("Retry-After", ke.retryAfter)
		}
		writeError(w, ke.status, ke.kind, ke.msg)
		r.m.kbRetrievals.WithLabelValues("refused").Inc()
		r.m.observeRequest(r.metricModel(req.Model), "", ke.status, time.Since(start), nil, 0)
		return false
	}
	if req.Path != "/v1/chat/completions" {
		return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
			msg: knowledgeHeader + " is for chat completions"})
	}
	p, ok := auth.FromContext(ctx)
	if !ok || r.store == nil {
		return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error",
			msg: "knowledge bases need authentication and the store, which this gateway runs without"})
	}
	if err := r.store.KnowledgeAvailable(); err != nil {
		return fail(&kbError{status: http.StatusNotImplemented, kind: "invalid_request_error", msg: err.Error()})
	}
	k := kbDefaultK
	if v := req.Header.Get(knowledgeKHeader); v != "" {
		n, err := strconv.Atoi(v)
		if err != nil || n < 1 || n > kbMaxK {
			return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
				msg: fmt.Sprintf("%s is 1 to %d", knowledgeKHeader, kbMaxK)})
		}
		k = n
	}
	began := time.Now()
	e, err := r.kbEffective(ctx, p, req.Header.Get(projectHeader))
	if err != nil {
		return fail(err)
	}
	allowed, err := r.store.CollectionsFor(ctx, e, false)
	if err != nil {
		return fail(err)
	}
	var colls []storepkg.KBCollection
	for _, n := range names {
		c, ok := pickCollection(allowed, n)
		if !ok && !explicit {
			continue
		}
		if !ok {
			msg := "no collection " + n
			if e.OrgID == "" {
				msg += " (a personal key names its project in " + projectHeader + ")"
			}
			return fail(&kbError{status: http.StatusNotFound, kind: "not_found_error", msg: msg})
		}
		colls = append(colls, c)
	}
	if len(colls) == 0 {
		return true
	}

	question := lastUserText(req.UpstreamBody)
	if strings.TrimSpace(question) == "" && !explicit {
		return true
	}
	if strings.TrimSpace(question) == "" {
		return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error",
			msg: knowledgeHeader + " needs a last user message with text to search for"})
	}
	embeddings := map[string][]float32{}
	for _, c := range colls {
		if _, done := embeddings[c.EmbedModel]; done {
			continue
		}
		v, err := r.embedFor(ctx, p, c.EmbedModel, question)
		if err != nil {
			return fail(err)
		}
		embeddings[c.EmbedModel] = v
	}
	hits, err := r.store.SearchChunks(ctx, storepkg.KBSearch{OrgID: e.OrgID, Collections: colls,
		Embeddings: embeddings, Text: question, K: k})
	if err != nil {
		return fail(err)
	}

	// What fits: the context the model has, less what the request already
	// holds and the answer it asked room for — the last passages go first.
	room := -1
	if _, entry := r.resolve(req.Model); entry.ContextLength > 0 {
		room = (entry.ContextLength-req.MaxTokens)*4 - len(req.UpstreamBody) - len(kbContextIntro) - 64
	}
	var b strings.Builder
	var sources []string
	for i, h := range hits {
		block := fmt.Sprintf("\n\n[%d] %s\n%s", i+1, h.Document, strings.TrimSpace(h.Body))
		if room >= 0 && b.Len()+len(block) > room {
			break
		}
		b.WriteString(block)
		sources = append(sources, url.PathEscape(h.Document)+"#"+strconv.Itoa(h.Ord))
	}
	text := kbNothingFound
	if len(sources) > 0 {
		text = kbContextIntro + b.String()
	}
	body, err := withSystemText(req.UpstreamBody, text)
	if err != nil {
		return fail(&kbError{status: http.StatusBadRequest, kind: "invalid_request_error", msg: err.Error()})
	}
	req.UpstreamBody = body
	req.PromptTokens = len(body) / 4
	req.knowledge = true
	// The collections searched, which is who else could have had these
	// passages added: the prefix cache is shared among exactly them.
	ids := make([]string, len(colls))
	for i, c := range colls {
		ids[i] = c.ID
	}
	sort.Strings(ids)
	req.kbSalt = strings.Join(ids, ",")
	// Set now: the proxy writes the headers when the answer starts.
	w.Header().Set(sourcesHeader, strings.Join(sources, ","))
	result := "found"
	if len(sources) == 0 {
		result = "nothing"
	}
	r.m.kbRetrievals.WithLabelValues(result).Inc()
	r.m.kbRetrieval.Observe(time.Since(began).Seconds())
	return true
}

// pickCollection finds a collection by id, then by name.
func pickCollection(cs []storepkg.KBCollection, want string) (storepkg.KBCollection, bool) {
	for _, c := range cs {
		if c.ID == want {
			return c, true
		}
	}
	for _, c := range cs {
		if strings.EqualFold(c.Name, want) {
			return c, true
		}
	}
	return storepkg.KBCollection{}, false
}

func splitList(v string) []string {
	var out []string
	seen := map[string]bool{}
	for _, s := range strings.Split(v, ",") {
		if s = strings.TrimSpace(s); s != "" && !seen[s] {
			seen[s] = true
			out = append(out, s)
		}
	}
	return out
}

// embedFor has a model embed the question as the caller: through the
// ordinary path, so the caller's limits, budget and guardrails apply to it
// and it is counted to them.
func (r *Router) embedFor(ctx context.Context, p auth.Principal, model, text string) ([]float32, error) {
	// A question, not a document: what the model wants before a query.
	if _, m := r.resolve(model); m.QueryPrefix != "" {
		text = m.QueryPrefix + text
	}
	body, _ := json.Marshal(map[string]any{"model": model, "input": []string{text}})
	hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, "/v1/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	hreq.Header.Set("Content-Type", "application/json")
	w := &bufferWriter{header: http.Header{}}
	r.handleInference("/v1/embeddings")(w, hreq)
	status := w.status
	if status == 0 {
		status = http.StatusOK
	}
	if status != http.StatusOK {
		msg := errorMessageOf(w.buf.Bytes())
		ke := &kbError{status: status, kind: "server_error",
			msg:        fmt.Sprintf("the knowledge base's model %s could not embed the question: %s", model, msg),
			retryAfter: w.header.Get("Retry-After")}
		if status < 500 {
			ke.kind = "invalid_request_error"
		}
		return nil, ke
	}
	var res struct {
		Data []struct {
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.Unmarshal(w.buf.Bytes(), &res); err != nil || len(res.Data) != 1 {
		return nil, fmt.Errorf("%s answered no embedding", model)
	}
	return res.Data[0].Embedding, nil
}

func errorMessageOf(b []byte) string {
	var e struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if json.Unmarshal(b, &e) == nil && e.Error.Message != "" {
		return e.Error.Message
	}
	return strings.TrimSpace(string(b))
}

// lastUserText is the text of a chat request's last user message: its
// content, or the text parts of it.
func lastUserText(body []byte) string {
	var in struct {
		Messages []struct {
			Role    string          `json:"role"`
			Content json.RawMessage `json:"content"`
		} `json:"messages"`
	}
	if json.Unmarshal(body, &in) != nil {
		return ""
	}
	for i := len(in.Messages) - 1; i >= 0; i-- {
		m := in.Messages[i]
		if m.Role != "user" {
			continue
		}
		var s string
		if json.Unmarshal(m.Content, &s) == nil {
			return s
		}
		var parts []struct {
			Type string `json:"type"`
			Text string `json:"text"`
		}
		if json.Unmarshal(m.Content, &parts) == nil {
			var b strings.Builder
			for _, p := range parts {
				if p.Type == "text" {
					b.WriteString(p.Text + "\n")
				}
			}
			return b.String()
		}
		return ""
	}
	return ""
}

// withSystemText adds text to a chat request's system message: appended to
// the first one when it leads, or as a new one at the start. One system
// message, first, because some models' templates refuse anything else.
func withSystemText(body []byte, text string) ([]byte, error) {
	var top map[string]json.RawMessage
	if err := json.Unmarshal(body, &top); err != nil {
		return nil, fmt.Errorf("the body is not a JSON object: %w", err)
	}
	var msgs []map[string]json.RawMessage
	if err := json.Unmarshal(top["messages"], &msgs); err != nil {
		return nil, errors.New("the request has no messages")
	}
	added := false
	if len(msgs) > 0 {
		var role string
		_ = json.Unmarshal(msgs[0]["role"], &role)
		if role == "system" || role == "developer" {
			var s string
			if json.Unmarshal(msgs[0]["content"], &s) == nil {
				msgs[0]["content"], _ = json.Marshal(strings.TrimRight(s, "\n") + "\n\n" + text)
				added = true
			} else {
				var parts []json.RawMessage
				if json.Unmarshal(msgs[0]["content"], &parts) == nil {
					part, _ := json.Marshal(map[string]string{"type": "text", "text": "\n\n" + text})
					msgs[0]["content"], _ = json.Marshal(append(parts, part))
					added = true
				}
			}
		}
	}
	if !added {
		content, _ := json.Marshal(text)
		role, _ := json.Marshal("system")
		msgs = append([]map[string]json.RawMessage{{"role": role, "content": content}}, msgs...)
	}
	var err error
	if top["messages"], err = json.Marshal(msgs); err != nil {
		return nil, err
	}
	return json.Marshal(top)
}

// kbDefaultTTL is how long a caller's default is remembered: a change to one
// is in force within it, and a chat request does not cost a query.
const kbDefaultTTL = 30 * time.Second

type kbDefaultEntry struct {
	at    time.Time
	names []string
}

// kbDefaults is the default collections in force for a caller.
func (r *Router) kbDefaults(ctx context.Context, p auth.Principal, project string) []string {
	key := p.KeyID + "|" + p.ProjectID + "|" + project
	r.kbDefMu.Lock()
	e, ok := r.kbDefCache[key]
	r.kbDefMu.Unlock()
	if ok && time.Since(e.at) < kbDefaultTTL {
		return e.names
	}
	var names []string
	if eff, err := r.kbEffective(ctx, p, project); err == nil {
		names, err = r.store.KBDefaultFor(ctx, eff)
		if err != nil {
			r.log.Warn("could not read the knowledge-base defaults", "err", err)
		}
	}
	r.kbDefMu.Lock()
	if r.kbDefCache == nil || len(r.kbDefCache) > 10000 {
		r.kbDefCache = map[string]kbDefaultEntry{}
	}
	r.kbDefCache[key] = kbDefaultEntry{at: time.Now(), names: names}
	r.kbDefMu.Unlock()
	return names
}

// withCacheSalt sets vLLM's cache_salt to one derived from key — the
// collections a request's passages came from — so requests that can see the
// same passages still share a cache, and nobody else's do: not another
// organization's, nor another team's of the same one.
func withCacheSalt(body []byte, key string) []byte {
	var top map[string]json.RawMessage
	if json.Unmarshal(body, &top) != nil {
		return body
	}
	sum := sha256.Sum256([]byte("llmgw-kb\x00" + key))
	top["cache_salt"], _ = json.Marshal(hex.EncodeToString(sum[:16]))
	out, err := json.Marshal(top)
	if err != nil {
		return body
	}
	return out
}