Retrieval-augmented generation (RAG) in Go, with the standard library

$ go run ./rag

index: 140 chunks, $0.00028 to embed

approach  correct  calls  tokens in      out       cost
rag         16/16     32       8714     3685    $0.0064
rag-tool    16/16     72      25834     6103    $0.0128
files       16/16     54      43201     6098    $0.0151

$ go run ./rag -ask rag "How do I make redirects permanent?"

// Chunk is one piece of one document: the unit that is embedded,
// searched, and handed to the model.
type Chunk struct {
	Doc     string `json:"doc"`     // the file it came from
	Heading string `json:"heading"` // the nearest heading above it
	Text    string `json:"text"`
}

// chunkDocs splits every Markdown file under docs into chunks, one per
// section: the text under each heading, with long sections cut at
// paragraph breaks so no chunk passes maxChunk bytes.
func chunkDocs(docs fs.FS, maxChunk int) ([]Chunk, error) {
	var out []Chunk
	err := fs.WalkDir(docs, ".", func(p string, d fs.DirEntry,
		err error) error {
		if err != nil || d.IsDir() || !strings.HasSuffix(p, ".md") {
			return err
		}
		data, err := fs.ReadFile(docs, p)
		if err != nil {
			return fmt.Errorf("read %s: %w", p, err)
		}
		heading, body := "", ""
		flush := func() {
			for _, part := range splitParas(body, maxChunk) {
				out = append(out, Chunk{Doc: p, Heading: heading, Text: part})
			}
			body = ""
		}
		for line := range strings.SplitSeq(string(data), "\n") {
			if h, ok := strings.CutPrefix(line, "#"); ok {
				flush()
				heading = strings.TrimSpace(strings.TrimLeft(h, "#"))
				continue
			}
			body += line + "\n"
		}
		flush()
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("chunk docs: %w", err)
	}
	return out, nil
}

[docs/deploy.md: Behind a reverse proxy]
When shortlink runs behind a reverse proxy such as nginx, every request
appears to come from the proxy's address, so every user shares one rate
limit. ...

// Embed turns texts into vectors with an embeddings model, through
// the same API and key as the chat calls.
func (c *Client) Embed(ctx context.Context, model string,
	texts []string) ([][]float64, error) {
	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
	defer cancel()
	body, err := json.Marshal(map[string]any{
		"model": model, "input": texts, "dimensions": embedDims,
	})
	if err != nil {
		return nil, fmt.Errorf("encode embeddings request: %w", err)
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		c.BaseURL+"/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("build embeddings request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.httpClient().Do(req)
	if err != nil {
		return nil, fmt.Errorf("send embeddings request: %w", err)
	}
	defer resp.Body.Close()
	return c.readEmbeddings(resp, len(texts))
}

// Index holds every chunk with its vector, and finds the chunks
// closest to a question by cosine similarity: an in-memory vector
// search, which for a few thousand chunks needs no database.
type Index struct {
	Chunks  []Chunk
	Vectors [][]float64
}

// Hit is one search result.
type Hit struct {
	Chunk
	Score float64 // cosine similarity, from -1 to 1
}

// Search returns the k chunks most similar to the question's vector,
// best first.
func (ix *Index) Search(q []float64, k int) []Hit {
	hits := make([]Hit, len(ix.Chunks))
	for i, v := range ix.Vectors {
		hits[i] = Hit{ix.Chunks[i], cosine(q, v)}
	}
	slices.SortFunc(hits, func(a, b Hit) int {
		return cmp.Compare(b.Score, a.Score)
	})
	return hits[:min(k, len(hits))]
}

func cosine(a, b []float64) float64 {
	var dot, na, nb float64
	for i := range min(len(a), len(b)) {
		dot += a[i] * b[i]
		na += a[i] * a[i]
		nb += b[i] * b[i]
	}
	if na == 0 || nb == 0 {
		return 0
	}
	return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

const ragPrompt = `Answer the question using only the passages from
shortlink's documentation that come with it. Name the document each
fact comes from. If the passages don't answer the question, say that
the documentation doesn't say.`

// askRAG is retrieval-augmented generation in one call: embed the
// question, find the closest chunks, and send them with the question.
func askRAG(ctx context.Context, c *Client, ix *Index, embedModel,
	question string) (string, error) {
	qv, err := c.Embed(ctx, embedModel, []string{question})
	if err != nil {
		return "", err
	}
	var passages []string
	for _, h := range ix.Search(qv[0], topK) {
		passages = append(passages, h.passage())
	}
	msgs := []Message{
		{Role: "system", Content: ragPrompt},
		{Role: "user", Content: "Passages:\n\n" +
			strings.Join(passages, "\n\n") + "\n\nQuestion: " + question},
	}
	reply, err := c.Chat(ctx, msgs, nil)
	if err != nil {
		return "", err
	}
	if reply.FinishReason == "length" {
		return "", errCutOff
	}
	return reply.Message.Content, nil
}

The documentation doesn't say. None of the provided passages mention
QR codes ...

const toolPrompt = `You answer questions about shortlink from its
documentation. Use search_docs to find passages; if the first results
don't answer the question, search again with other words. Answer only
from what you found, name the documents you used, and if the
documentation doesn't answer the question, say so.`

// searchTool is retrieval as a tool: the model decides when to search
// and what to search for, and can search again.
func searchTool(c *Client, ix *Index, embedModel string) Tool {
	return Tool{
		Name: "search_docs",
		Description: "Search shortlink's documentation. Returns the " +
			"passages closest in meaning to the query, best first.",
		Parameters: []byte(`{"type": "object", "properties": {"query":
			{"type": "string"}}, "required": ["query"]}`),
		Run: func(ctx context.Context, args string) (string, error) {
			var a struct {
				Query string `json:"query"`
			}
			if err := json.Unmarshal([]byte(args), &a); err != nil {
				return "", fmt.Errorf("arguments: %w", err)
			}
			qv, err := c.Embed(ctx, embedModel, []string{a.Query})
			if err != nil {
				return "", err
			}
			var out []string
			for _, h := range ix.Search(qv[0], topK) {
				out = append(out, h.passage())
			}
			return strings.Join(out, "\n\n"), nil
		},
	}
}