Build a RAG Server

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/pgvector/pgvector-go"
	"google.golang.org/genai"
)

const (
	embedModel = "gemini-embedding-001"
	chatModel  = "gemini-2.5-flash"
	embedDims  = 768
	topK       = 4
)

const schema = `
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS chunks (
    id        BIGSERIAL PRIMARY KEY,
    source    TEXT        NOT NULL,
    body      TEXT        NOT NULL,
    embedding vector(768) NOT NULL
);
CREATE INDEX IF NOT EXISTS chunks_embedding_idx
    ON chunks USING hnsw (embedding vector_cosine_ops);
`

func Chunk(text string, maxRunes int) []string {
	var chunks []string
	var cur strings.Builder
	for _, p := range strings.Split(text, "\n\n") {
		p = strings.TrimSpace(p)
		if p == "" {
			continue
		}
		if cur.Len() > 0 && utf8.RuneCountInString(cur.String())+utf8.RuneCountInString(p) > maxRunes {
			chunks = append(chunks, cur.String())
			cur.Reset()
		}
		if cur.Len() > 0 {
			cur.WriteString("\n\n")
		}
		cur.WriteString(p)
	}
	if cur.Len() > 0 {
		chunks = append(chunks, cur.String())
	}
	return chunks
}

type Server struct {
	ai *genai.Client
	db *pgxpool.Pool
}

func (s *Server) embed(ctx context.Context, texts []string, task string) ([][]float32, error) {
	contents := make([]*genai.Content, len(texts))
	for i, t := range texts {
		contents[i] = genai.NewContentFromText(t, genai.RoleUser)
	}
	dims := int32(embedDims)
	resp, err := s.ai.Models.EmbedContent(ctx, embedModel, contents, &genai.EmbedContentConfig{
		TaskType:             task,
		OutputDimensionality: &dims,
	})
	if err != nil {
		return nil, err
	}
	out := make([][]float32, len(resp.Embeddings))
	for i, e := range resp.Embeddings {
		out[i] = e.Values
	}
	return out, nil
}

func (s *Server) handleIngest(w http.ResponseWriter, r *http.Request) {
	var in struct {
		Source string `json:"source"`
		Text   string `json:"text"`
	}
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	chunks := Chunk(in.Text, 900)
	if len(chunks) == 0 {
		http.Error(w, "no text", http.StatusBadRequest)
		return
	}
	vecs, err := s.embed(r.Context(), chunks, "RETRIEVAL_DOCUMENT")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	batch := make([][]any, len(chunks))
	for i := range chunks {
		batch[i] = []any{in.Source, chunks[i], pgvector.NewVector(vecs[i])}
	}
	_, err = s.db.CopyFrom(r.Context(),
		pgx.Identifier{"chunks"},
		[]string{"source", "body", "embedding"},
		pgx.CopyFromRows(batch),
	)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	writeJSON(w, map[string]int{"chunks": len(chunks)})
}

type hit struct {
	Source string  `json:"source"`
	Body   string  `json:"-"`
	Score  float64 `json:"score"`
}

func (s *Server) handleAsk(w http.ResponseWriter, r *http.Request) {
	var in struct {
		Question string `json:"question"`
	}
	if err := json.NewDecoder(r.Body).Decode(&in); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	qv, err := s.embed(r.Context(), []string{in.Question}, "RETRIEVAL_QUERY")
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	rows, err := s.db.Query(r.Context(), `
        SELECT source, body, 1 - (embedding <=> $1) AS score
        FROM chunks
        ORDER BY embedding <=> $1
        LIMIT $2`, pgvector.NewVector(qv[0]), topK)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var hits []hit
	var ctxBuf strings.Builder
	for rows.Next() {
		var h hit
		if err := rows.Scan(&h.Source, &h.Body, &h.Score); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		fmt.Fprintf(&ctxBuf, "[%d] (%s) %s\n\n", len(hits)+1, h.Source, h.Body)
		hits = append(hits, h)
	}
	if err := rows.Err(); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	prompt := fmt.Sprintf("Context:\n%s\nQuestion: %s", ctxBuf.String(), in.Question)
	resp, err := s.ai.Models.GenerateContent(r.Context(), chatModel,
		genai.Text(prompt),
		&genai.GenerateContentConfig{
			SystemInstruction: genai.NewContentFromText(systemPrompt, genai.RoleUser),
			Temperature:       genai.Ptr[float32](0),
		})
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	writeJSON(w, map[string]any{"answer": resp.Text(), "sources": hits})
}

const systemPrompt = `Answer using only the numbered context provided.
If the context does not contain the answer, say you do not know.
Cite the numbers you used.`

func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(v)
}

func main() {
	ctx := context.Background()

	ai, err := genai.NewClient(ctx, &genai.ClientConfig{
		APIKey:  os.Getenv("GEMINI_API_KEY"),
		Backend: genai.BackendGeminiAPI,
	})
	if err != nil {
		log.Fatal(err)
	}

	db, err := pgxpool.New(ctx, os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	if _, err := db.Exec(ctx, schema); err != nil {
		log.Fatal(err)
	}

	srv := &Server{ai: ai, db: db}
	mux := http.NewServeMux()
	mux.HandleFunc("POST /ingest", srv.handleIngest)
	mux.HandleFunc("POST /ask", srv.handleAsk)

	s := &http.Server{Addr: ":8080", Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	log.Fatal(s.ListenAndServe())
}