llmgw/internal/control/adminapi/kbdataset.go

package adminapi

import (
	"encoding/json"
	"fmt"
	"net/http"
	"regexp"
	"strings"
)

// Training data from a collection (BACKLOG F4): questions a person might
// ask of each chunk and the answers it holds, written by a model of the
// cluster, for people to read through before an adapter is trained on them.
// What an adapter learns well is how the organisation's answers sound; what
// they say is still the knowledge base's to give.

// KBDatasetRequest asks for training data from a collection.
type KBDatasetRequest struct {
	// Model writes the questions and answers: a chat model of the catalog.
	Model string `json:"model"`
	// PerChunk is how many pairs each chunk gives; 8 when 0.
	PerChunk int `json:"per_chunk,omitempty"`
	// MaxChunks bounds how many chunks are read; 500 when 0.
	MaxChunks int `json:"max_chunks,omitempty"`
}

// KBDatasetLine is one line of the answer: a pair, or why a chunk gave none.
type KBDatasetLine struct {
	Q   string `json:"q,omitempty"`
	A   string `json:"a,omitempty"`
	Doc string `json:"doc"`
	Ord int    `json:"ord"`
	// Collection is where the pair came from: an adapter trained on it
	// belongs to a part of the tree that may read it.
	Collection string `json:"collection,omitempty"`
	Error      string `json:"error,omitempty"`
}

const kbDatasetPrompt = `From the passage of an organisation's document below, write %d different questions an employee might ask, each with a short, correct answer.
Rules: use only what the passage says. Phrase the questions differently from each other (formal, everyday, short, indirect). Answers are one or two sentences and keep the passage's numbers exactly. Leave out e-mail addresses and people's names. Write in the passage's own language.
Output: one JSON object per line, {"q": "...", "a": "..."}, and nothing else.

DOCUMENT %s:
%s`

var thinkBlock = regexp.MustCompile(`(?s)<think>.*?</think>`)

func (a *API) kbDataset(w http.ResponseWriter, r *http.Request) {
	c, ok := a.kbCollection(w, r)
	if !ok {
		return
	}
	if a.replay == nil {
		writeError(w, http.StatusNotImplemented, "no router here to ask a model")
		return
	}
	var in KBDatasetRequest
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	m, found := a.ctl.Catalog().Resolve(in.Model)
	if !found {
		writeError(w, http.StatusBadRequest, "no model "+in.Model+" in the catalog")
		return
	}
	if m.Capabilities.Embeddings {
		writeError(w, http.StatusBadRequest, m.Name+" embeds; questions are written by a chat model")
		return
	}
	per := in.PerChunk
	if per <= 0 || per > 30 {
		per = 8
	}
	limit := in.MaxChunks
	if limit <= 0 || limit > 5000 {
		limit = 500
	}
	chunks, err := a.kb.Store.CollectionChunks(r.Context(), c, limit)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	w.Header().Set("Content-Type", "application/x-ndjson")
	w.WriteHeader(http.StatusOK)
	enc := json.NewEncoder(w)
	flush := func() {
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
	}
	for _, ch := range chunks {
		if r.Context().Err() != nil {
			return
		}
		body, _ := json.Marshal(map[string]any{"model": m.Name, "temperature": 0.7, "max_tokens": 200 * per,
			"messages": []map[string]string{{"role": "user", "content": fmt.Sprintf(kbDatasetPrompt, per, ch.Document, ch.Body)}}})
		status, out, _, err := a.replay(r.Context(), "/v1/chat/completions", m.Name, body)
		line := KBDatasetLine{Doc: ch.Document, Ord: ch.Ord, Collection: c.ID}
		if err == nil && status != http.StatusOK {
			err = fmt.Errorf("%d: %s", status, strings.TrimSpace(string(out)))
		}
		if err != nil {
			line.Error = err.Error()
			_ = enc.Encode(line)
			flush()
			continue
		}
		pairs := datasetPairs(out)
		if len(pairs) == 0 {
			line.Error = "the model gave no pairs it could read"
			_ = enc.Encode(line)
		}
		for _, p := range pairs {
			p.Doc, p.Ord, p.Collection = ch.Document, ch.Ord, c.ID
			_ = enc.Encode(p)
		}
		flush()
	}
}

// datasetPairs reads the {"q","a"} lines out of a chat completion.
func datasetPairs(completion []byte) []KBDatasetLine {
	var res struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if json.Unmarshal(completion, &res) != nil || len(res.Choices) == 0 {
		return nil
	}
	text := thinkBlock.ReplaceAllString(res.Choices[0].Message.Content, "")
	var out []KBDatasetLine
	for _, line := range strings.Split(text, "\n") {
		line = strings.TrimSuffix(strings.TrimSpace(line), ",")
		if !strings.HasPrefix(line, "{") || !strings.HasSuffix(line, "}") {
			continue
		}
		var p struct{ Q, A string }
		if json.Unmarshal([]byte(line), &p) == nil && strings.TrimSpace(p.Q) != "" && strings.TrimSpace(p.A) != "" {
			out = append(out, KBDatasetLine{Q: strings.TrimSpace(p.Q), A: strings.TrimSpace(p.A)})
		}
	}
	return out
}