Preparing Data for the RAG Server

# httpx for the Ollama API, and psycopg 3 with the pre-built binary
# package that bundles libpq.
httpx>=0.28,<1
psycopg[binary]>=3.2,<4

{
  "name": "prepare-data",
  "private": true,
  "type": "module",
  "scripts": {
    "build": "tsc",
    "start": "node dist/main.js"
  },
  "dependencies": {
    "pg": "^8.23.1"
  },
  "devDependencies": {
    "@types/node": "^22.0.0",
    "@types/pg": "^8.23.1",
    "typescript": "^5.9.0"
  }
}

<dependencies>
  <!-- Jackson, for reading and writing JSON; the HTTP client is part of
       the JDK. -->
  <dependency>
    <groupId>tools.jackson.core</groupId>
    <artifactId>jackson-databind</artifactId>
    <version>3.2.4</version>
  </dependency>
  <!-- The PostgreSQL JDBC driver (pgJDBC). -->
  <dependency>
    <groupId>org.postgresql</groupId>
    <artifactId>postgresql</artifactId>
    <version>42.7.14</version>
  </dependency>
</dependencies>

<ItemGroup>
  <!-- Npgsql, the ADO.NET data provider for PostgreSQL. HttpClient and
       System.Text.Json are part of .NET. -->
  <PackageReference Include="Npgsql" Version="10.0.3" />
</ItemGroup>

// pgx is a pure Go driver and toolkit for PostgreSQL; the Ollama API
// needs only the standard library.
require github.com/jackc/pgx/v5 v5.12.0

require (
    github.com/jackc/pgpassfile v1.0.0 // indirect
    github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
    golang.org/x/text v0.41.0 // indirect
)

[dependencies]
# The blocking client, with TLS from rustls and request and response bodies
# as JSON; the default features add HTTP/2 and other things this example
# does not use.
reqwest = { version = "0.13", default-features = false, features = ["blocking", "json", "rustls"] }
serde_json = "1"
# The synchronous PostgreSQL client; tokio-postgres is the asynchronous
# equivalent.
postgres = "0.19"

source "https://rubygems.org"

# Net::HTTP and json are part of the standard library. The pg gem wraps
# libpq, so libpq and its headers must be installed.
gem "pg", "~> 1.6"

import os
import sys

import httpx
import psycopg

import pg from "pg";

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpTimeoutException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.StringJoiner;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import tools.jackson.core.JacksonException;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.ObjectMapper;
import tools.jackson.databind.json.JsonMapper;
import tools.jackson.databind.node.ObjectNode;

using System.Globalization;
using System.Net.Http.Json;
using System.Text.Json;
using Npgsql;

import (
    "bytes"
    "context"
    "encoding/json"
    "fmt"
    "io"
    "net"
    "net/http"
    "net/url"
    "os"
    "strconv"
    "strings"
    "time"

    "github.com/jackc/pgx/v5"
)

use std::env;
use std::error::Error;
use std::process::ExitCode;
use std::time::Duration;

use postgres::{Client, Config, NoTls, Transaction};
use serde_json::{json, Value};

require "json"
require "net/http"
require "pg"
require "uri"

# The or operator treats an empty variable as unset, as libpq does.
OLLAMA_URL = (os.environ.get("OLLAMA_URL") or "http://localhost:11434").rstrip("/")
EMBED_MODEL = os.environ.get("EMBED_MODEL") or "nomic-embed-text"

# This recipe uses its own table, so that it never disturbs the table that
# a running pipeline reads.
TABLE = "cookbook_prepare_data"

# The documents to store; in a real application they come from your own
# data. The text is fictional.
DOCUMENTS = [
    (1, "Charging the Acme Gadget",
     "The Acme Gadget charges through its USB-C port. A full charge takes "
     "two hours, and a fully charged Acme Gadget runs for ten hours."),
    (2, "Resetting the Acme Gadget",
     "To reset the Acme Gadget to its factory settings, hold down the "
     "power button for eight seconds until the status light flashes red."),
    (3, "Warranty for the Acme Gadget",
     "The Acme Gadget comes with a two-year warranty that covers "
     "manufacturing defects but not accidental damage."),
]


def connection_settings() -> dict:
    """Reads the connection settings from the standard libpq environment
    variables, with defaults suitable for a local development server."""
    return {
        "host": os.environ.get("PGHOST") or "localhost",
        "port": os.environ.get("PGPORT") or "5432",
        "dbname": os.environ.get("PGDATABASE") or "postgres",
        "user": os.environ.get("PGUSER") or "postgres",
        "password": os.environ.get("PGPASSWORD") or None,
    }

// The || operator, unlike ??, also treats an empty variable as unset, as
// libpq does.
const ollamaUrl = (process.env.OLLAMA_URL || "http://localhost:11434").replace(/\/+$/, "");
const embedModel = process.env.EMBED_MODEL || "nomic-embed-text";

// This recipe uses its own table, so that it never disturbs the table that
// a running pipeline reads.
const table = "cookbook_prepare_data";

// The documents to store; in a real application they come from your own
// data. The text is fictional.
const documents = [
  {
    id: 1,
    title: "Charging the Acme Gadget",
    content:
      "The Acme Gadget charges through its USB-C port. A full charge takes " +
      "two hours, and a fully charged Acme Gadget runs for ten hours.",
  },
  {
    id: 2,
    title: "Resetting the Acme Gadget",
    content:
      "To reset the Acme Gadget to its factory settings, hold down the " +
      "power button for eight seconds until the status light flashes red.",
  },
  {
    id: 3,
    title: "Warranty for the Acme Gadget",
    content:
      "The Acme Gadget comes with a two-year warranty that covers " +
      "manufacturing defects but not accidental damage.",
  },
];

const dbConfig: pg.ClientConfig = {
  host: process.env.PGHOST || "localhost",
  port: Number(process.env.PGPORT || "5432"),
  database: process.env.PGDATABASE || "postgres",
  user: process.env.PGUSER || "postgres",
  password: process.env.PGPASSWORD || undefined,
};

/** Returns the value of an environment variable, or a fallback. */
static String envOr(String name, String fallback) {
    String value = System.getenv(name);
    return (value != null && !value.isEmpty()) ? value : fallback;
}

static final String OLLAMA_URL =
        envOr("OLLAMA_URL", "http://localhost:11434").replaceAll("/+$", "");
static final String EMBED_MODEL = envOr("EMBED_MODEL", "nomic-embed-text");

// This recipe uses its own table, so that it never disturbs the table
// that a running pipeline reads.
static final String TABLE = "cookbook_prepare_data";

/** A row to store. */
record Document(int id, String title, String content) {}

// The documents to store; in a real application they come from your
// own data. The text is fictional.
static final List<Document> DOCUMENTS = List.of(
        new Document(1, "Charging the Acme Gadget",
                "The Acme Gadget charges through its USB-C port. A full charge takes "
                + "two hours, and a fully charged Acme Gadget runs for ten hours."),
        new Document(2, "Resetting the Acme Gadget",
                "To reset the Acme Gadget to its factory settings, hold down the "
                + "power button for eight seconds until the status light flashes red."),
        new Document(3, "Warranty for the Acme Gadget",
                "The Acme Gadget comes with a two-year warranty that covers "
                + "manufacturing defects but not accidental damage."));

/**
 * Returns the JDBC URL that identifies the server and database, built
 * from the standard libpq environment variables.
 */
static String connectionUrl() {
    return "jdbc:postgresql://"
            + envOr("PGHOST", "localhost") + ":"
            + envOr("PGPORT", "5432") + "/"
            + envOr("PGDATABASE", "postgres");
}

/**
 * Returns the user name and password as connection properties; the
 * password is omitted if PGPASSWORD is unset or empty.
 */
static Properties connectionProperties() {
    Properties props = new Properties();
    props.setProperty("user", envOr("PGUSER", "postgres"));
    String password = envOr("PGPASSWORD", null);
    if (password != null) {
        props.setProperty("password", password);
    }
    return props;
}

// Returns the value of an environment variable, or null if the variable
// is unset or empty; libpq also treats an empty variable as unset.
static string? EnvValue(string name) =>
    Environment.GetEnvironmentVariable(name) is { Length: > 0 } value
        ? value
        : null;

// Returns the value of an environment variable, or a fallback if unset.
static string EnvOr(string name, string fallback) =>
    EnvValue(name) ?? fallback;

static string OllamaUrl() =>
    EnvOr("OLLAMA_URL", "http://localhost:11434").TrimEnd('/');
static string EmbedModel() => EnvOr("EMBED_MODEL", "nomic-embed-text");

// This recipe uses its own table, so that it never disturbs the table that
// a running pipeline reads.
const string Table = "cookbook_prepare_data";

// The documents to store; in a real application they come from your own
// data. The text is fictional.
Document[] documents =
[
    new(1, "Charging the Acme Gadget",
        "The Acme Gadget charges through its USB-C port. A full charge takes " +
        "two hours, and a fully charged Acme Gadget runs for ten hours."),
    new(2, "Resetting the Acme Gadget",
        "To reset the Acme Gadget to its factory settings, hold down the " +
        "power button for eight seconds until the status light flashes red."),
    new(3, "Warranty for the Acme Gadget",
        "The Acme Gadget comes with a two-year warranty that covers " +
        "manufacturing defects but not accidental damage."),
];

// Builds a connection string from the standard libpq environment
// variables, with defaults suitable for a local development server. The
// builder takes care of quoting any special characters.
static string ConnectionString() =>
    new NpgsqlConnectionStringBuilder
    {
        Host = EnvOr("PGHOST", "localhost"),
        Port = int.Parse(EnvOr("PGPORT", "5432")),
        Database = EnvOr("PGDATABASE", "postgres"),
        Username = EnvOr("PGUSER", "postgres"),
        Password = EnvValue("PGPASSWORD"),
    }.ConnectionString;

// envOr returns the value of an environment variable, or a fallback if
// the variable is unset or empty.
func envOr(name, fallback string) string {
    if value := os.Getenv(name); value != "" {
        return value
    }
    return fallback
}

var (
    ollamaURL  = strings.TrimRight(envOr("OLLAMA_URL", "http://localhost:11434"), "/")
    embedModel = envOr("EMBED_MODEL", "nomic-embed-text")
)

// This recipe uses its own table, so that it never disturbs the table that
// a running pipeline reads.
const table = "cookbook_prepare_data"

// document is a row to store. The text is fictional; in a real
// application the documents come from your own data.
type document struct {
    ID      int
    Title   string
    Content string
}

var documents = []document{
    {1, "Charging the Acme Gadget",
        "The Acme Gadget charges through its USB-C port. A full charge takes " +
            "two hours, and a fully charged Acme Gadget runs for ten hours."},
    {2, "Resetting the Acme Gadget",
        "To reset the Acme Gadget to its factory settings, hold down the " +
            "power button for eight seconds until the status light flashes red."},
    {3, "Warranty for the Acme Gadget",
        "The Acme Gadget comes with a two-year warranty that covers " +
            "manufacturing defects but not accidental damage."},
}

// connString builds a postgres:// URL from the standard libpq environment
// variables. net/url takes care of escaping any special characters in
// the user name, password, or database name.
func connString() string {
    u := url.URL{
        Scheme: "postgres",
        Host:   net.JoinHostPort(envOr("PGHOST", "localhost"), envOr("PGPORT", "5432")),
        Path:   "/" + envOr("PGDATABASE", "postgres"),
    }
    if password := os.Getenv("PGPASSWORD"); password != "" {
        u.User = url.UserPassword(envOr("PGUSER", "postgres"), password)
    } else {
        u.User = url.User(envOr("PGUSER", "postgres"))
    }
    return u.String()
}

/// Returns the value of an environment variable, if it is set and not
/// empty; libpq also treats an empty variable as unset.
fn env_value(name: &str) -> Option<String> {
    env::var(name).ok().filter(|value| !value.is_empty())
}

/// Returns the value of an environment variable, or a fallback if the
/// variable is unset or empty.
fn env_or(name: &str, fallback: &str) -> String {
    env_value(name).unwrap_or_else(|| fallback.to_string())
}

fn ollama_url() -> String {
    env_or("OLLAMA_URL", "http://localhost:11434")
        .trim_end_matches('/')
        .to_string()
}

fn embed_model() -> String {
    env_or("EMBED_MODEL", "nomic-embed-text")
}

/// This recipe uses its own table, so that it never disturbs the table
/// that a running pipeline reads.
const TABLE: &str = "cookbook_prepare_data";

/// The documents to store, as (id, title, content); in a real application
/// they come from your own data. The text is fictional.
const DOCUMENTS: [(i32, &str, &str); 3] = [
    (
        1,
        "Charging the Acme Gadget",
        "The Acme Gadget charges through its USB-C port. A full charge takes \
         two hours, and a fully charged Acme Gadget runs for ten hours.",
    ),
    (
        2,
        "Resetting the Acme Gadget",
        "To reset the Acme Gadget to its factory settings, hold down the \
         power button for eight seconds until the status light flashes red.",
    ),
    (
        3,
        "Warranty for the Acme Gadget",
        "The Acme Gadget comes with a two-year warranty that covers \
         manufacturing defects but not accidental damage.",
    ),
];

/// The errors that the functions below return; a string converts into one
/// with the ? operator or into().
type BoxError = Box<dyn Error>;

/// Builds the connection configuration from the standard libpq
/// environment variables, with defaults suitable for a local server.
fn config() -> Result<Config, BoxError> {
    let mut config = Config::new();
    config
        .host(&env_or("PGHOST", "localhost"))
        .port(env_or("PGPORT", "5432").parse()?)
        .dbname(&env_or("PGDATABASE", "postgres"))
        .user(&env_or("PGUSER", "postgres"));
    if let Some(password) = env_value("PGPASSWORD") {
        config.password(&password);
    }
    Ok(config)
}

# Returns the value of an environment variable, or a fallback if the
# variable is unset or empty; libpq also treats an empty variable as unset.
def env_or(name, fallback)
  value = ENV[name]
  value.nil? || value.empty? ? fallback : value
end

OLLAMA_URL = env_or("OLLAMA_URL", "http://localhost:11434").sub(%r{/+\z}, "")
EMBED_MODEL = env_or("EMBED_MODEL", "nomic-embed-text")

# This recipe uses its own table, so that it never disturbs the table that
# a running pipeline reads.
TABLE = "cookbook_prepare_data"

# The documents to store, as [id, title, content]; in a real application
# they come from your own data. The text is fictional.
DOCUMENTS = [
  [1, "Charging the Acme Gadget",
   "The Acme Gadget charges through its USB-C port. A full charge takes " \
   "two hours, and a fully charged Acme Gadget runs for ten hours."],
  [2, "Resetting the Acme Gadget",
   "To reset the Acme Gadget to its factory settings, hold down the " \
   "power button for eight seconds until the status light flashes red."],
  [3, "Warranty for the Acme Gadget",
   "The Acme Gadget comes with a two-year warranty that covers " \
   "manufacturing defects but not accidental damage."]
].freeze

# Returns the connection settings from the standard libpq environment
# variables, with defaults suitable for a local development server.
# compact removes the password if PGPASSWORD is not set.
def connection_params
  {
    host: env_or("PGHOST", "localhost"),
    port: env_or("PGPORT", "5432"),
    dbname: env_or("PGDATABASE", "postgres"),
    user: env_or("PGUSER", "postgres"),
    password: env_or("PGPASSWORD", nil)
  }.compact
end

def embed(texts: list[str]) -> list[list[float]]:
    """Asks Ollama for the embedding of each text, in the same order."""
    # The embed endpoint accepts a list of texts, so one request embeds
    # every document.
    response = httpx.post(
        f"{OLLAMA_URL}/api/embed",
        json={"model": EMBED_MODEL, "input": texts},
        timeout=300.0,
    )
    response.raise_for_status()

    embeddings = response.json()["embeddings"]
    if len(embeddings) != len(texts):
        raise ValueError(
            f"Expected {len(texts)} embeddings, got {len(embeddings)}")
    return embeddings

// Asks Ollama for the embedding of each text, in the same order.
async function embed(texts: string[]): Promise<number[][]> {
  // The embed endpoint accepts a list of texts, so one request embeds
  // every document.
  const response = await fetch(`${ollamaUrl}/api/embed`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ model: embedModel, input: texts }),
    signal: AbortSignal.timeout(300_000),
  });
  // fetch only rejects when it cannot reach the server, so check the
  // status explicitly.
  if (!response.ok) {
    throw new Error(`Ollama returned HTTP ${response.status}: ${(await response.text()).slice(0, 200)}`);
  }

  const { embeddings } = (await response.json()) as { embeddings?: number[][] };
  if (!Array.isArray(embeddings) || embeddings.length !== texts.length) {
    throw new Error(`Expected ${texts.length} embeddings from Ollama`);
  }
  return embeddings;
}

// A mapper is safe to share, so the program creates only one.
static final ObjectMapper MAPPER = JsonMapper.builder().build();

/** Thrown when Ollama rejects a request or returns an unexpected reply. */
static class EmbedException extends Exception {
    EmbedException(String message) {
        super(message);
    }
}

/**
 * Sends a request and returns the response with its whole body. The
 * request timeout stops applying once the headers arrive, so a server
 * that stalls part of the way through the body would hang send(); this
 * applies the same timeout to the whole exchange instead.
 */
static HttpResponse<String> send(HttpClient client, HttpRequest request)
        throws IOException, InterruptedException {
    Duration timeout = request.timeout().orElseThrow();
    CompletableFuture<HttpResponse<String>> pending =
            client.sendAsync(request, HttpResponse.BodyHandlers.ofString());
    try {
        return pending.get(timeout.toMillis(), TimeUnit.MILLISECONDS);
    } catch (TimeoutException e) {
        pending.cancel(true);
        throw new HttpTimeoutException("request timed out");
    } catch (ExecutionException e) {
        if (e.getCause() instanceof IOException cause) {
            throw cause;
        }
        throw new IOException(e.getCause());
    }
}

/**
 * Asks Ollama for the embedding of each text, in the same order.
 * pgvector stores single-precision values, so a float loses nothing
 * that the table would keep.
 */
static List<float[]> embed(List<String> texts)
        throws IOException, InterruptedException, EmbedException {
    // The embed endpoint accepts a list of texts, so one request embeds
    // every document.
    ObjectNode payload = MAPPER.createObjectNode().put("model", EMBED_MODEL);
    texts.forEach(payload.putArray("input")::add);

    HttpRequest request = HttpRequest.newBuilder(URI.create(OLLAMA_URL + "/api/embed"))
            .timeout(Duration.ofSeconds(300))
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(MAPPER.writeValueAsString(payload)))
            .build();
    HttpResponse<String> response = send(HttpClient.newHttpClient(), request);

    // The client treats an error status as a normal response, so check
    // it here.
    if (response.statusCode() != 200) {
        String body = response.body().strip();
        throw new EmbedException("Ollama returned HTTP " + response.statusCode()
                + ": " + body.substring(0, Math.min(body.length(), 200)));
    }

    // path() returns a missing node, rather than null, for a field that
    // is not there, and a missing node holds no elements.
    List<float[]> embeddings = new ArrayList<>();
    for (JsonNode node : MAPPER.readTree(response.body()).path("embeddings")) {
        float[] embedding = new float[node.size()];
        for (int i = 0; i < embedding.length; i++) {
            embedding[i] = node.get(i).floatValue();
        }
        embeddings.add(embedding);
    }
    if (embeddings.size() != texts.size()) {
        throw new EmbedException("Expected " + texts.size()
                + " embeddings, got " + embeddings.size());
    }
    return embeddings;
}

// Asks Ollama for the embedding of each text, in the same order. pgvector
// stores single-precision values, so a float loses nothing that the table
// would keep.
static async Task<float[][]> EmbedAsync(string[] texts)
{
    using var client = new HttpClient { Timeout = TimeSpan.FromSeconds(300) };
    // The embed endpoint accepts a list of texts, so one request embeds
    // every document.
    using var response = await client.PostAsJsonAsync(
        $"{OllamaUrl()}/api/embed", new { model = EmbedModel(), input = texts });

    // HttpClient treats an error status as a normal response, so check it
    // here.
    if (!response.IsSuccessStatusCode)
    {
        var body = (await response.Content.ReadAsStringAsync()).Trim();
        throw new EmbedException(
            $"Ollama returned HTTP {(int)response.StatusCode}: " +
            body[..Math.Min(body.Length, 200)]);
    }

    // ReadFromJsonAsync matches the embeddings field to the Embeddings
    // property without regard to case, and leaves it null if it is missing.
    var reply = await response.Content.ReadFromJsonAsync<EmbedReply>();
    if (reply?.Embeddings is not { } embeddings || embeddings.Length != texts.Length)
    {
        throw new EmbedException(
            $"Expected {texts.Length} embeddings, got {reply?.Embeddings?.Length ?? 0}");
    }
    return embeddings;
}

// embed asks Ollama for the embedding of each text, in the same order.
func embed(texts []string) ([][]float32, error) {
    // The embed endpoint accepts a list of texts, so one request embeds
    // every document.
    payload, err := json.Marshal(map[string]any{"model": embedModel, "input": texts})
    if err != nil {
        return nil, err
    }

    client := &http.Client{Timeout: 300 * time.Second}
    response, err := client.Post(ollamaURL+"/api/embed", "application/json", bytes.NewReader(payload))
    if err != nil {
        return nil, err
    }
    defer response.Body.Close()

    body, err := io.ReadAll(response.Body)
    if err != nil {
        return nil, err
    }
    // Go treats an error status as a normal response, so check it here.
    if response.StatusCode != http.StatusOK {
        return nil, fmt.Errorf("ollama returned HTTP %d: %.200s", response.StatusCode, body)
    }

    var result struct {
        Embeddings [][]float32 `json:"embeddings"`
    }
    if err := json.Unmarshal(body, &result); err != nil {
        return nil, fmt.Errorf("unexpected response from ollama: %w", err)
    }
    if len(result.Embeddings) != len(texts) {
        return nil, fmt.Errorf("expected %d embeddings, got %d", len(texts), len(result.Embeddings))
    }
    return result.Embeddings, nil
}

/// Asks Ollama for the embedding of each text, in the same order.
/// pgvector stores single-precision values, so f32 loses nothing that the
/// table would keep.
fn embed(texts: &[&str]) -> Result<Vec<Vec<f32>>, BoxError> {
    let client = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(300))
        .build()?;
    // The embed endpoint accepts a list of texts, so one request embeds
    // every document.
    let response = client
        .post(format!("{}/api/embed", ollama_url()))
        .json(&json!({ "model": embed_model(), "input": texts }))
        .send()?;

    // reqwest treats an error status as a normal response, so check it
    // here.
    let status = response.status();
    if !status.is_success() {
        let body = response.text()?;
        let body: String = body.trim().chars().take(200).collect();
        return Err(format!("Ollama returned HTTP {}: {body}", status.as_u16()).into());
    }

    let mut reply: Value = response.json()?;
    let embeddings: Vec<Vec<f32>> = serde_json::from_value(reply["embeddings"].take())?;
    if embeddings.len() != texts.len() {
        return Err(format!("Expected {} embeddings, got {}", texts.len(), embeddings.len()).into());
    }
    Ok(embeddings)
}

# Raised when Ollama rejects a request or returns an unexpected reply.
class EmbedError < StandardError; end

# Asks Ollama for the embedding of each text, in the same order.
def embed(texts)
  uri = URI("#{OLLAMA_URL}/api/embed")
  request = Net::HTTP::Post.new(uri, "Content-Type" => "application/json")
  # The embed endpoint accepts a list of texts, so one request embeds
  # every document.
  request.body = JSON.generate({ model: EMBED_MODEL, input: texts })

  # start opens a connection, and closes it when the block exits, whether
  # or not an error occurred.
  response = Net::HTTP.start(
    uri.host, uri.port,
    use_ssl: uri.scheme == "https", open_timeout: 300, read_timeout: 300
  ) { |http| http.request(request) }
  # Net::HTTP treats an error status as a normal response, so check it
  # here.
  unless response.is_a?(Net::HTTPSuccess)
    raise EmbedError, "Ollama returned HTTP #{response.code}: #{response.body.to_s.strip[0, 200]}"
  end

  # fetch raises KeyError if the reply has no embeddings.
  embeddings = JSON.parse(response.body).fetch("embeddings")
  unless embeddings.is_a?(Array) && embeddings.length == texts.length
    raise EmbedError, "Expected #{texts.length} embeddings, got #{embeddings.inspect[0, 200]}"
  end

  embeddings
end

def create_table(conn: psycopg.Connection, dimensions: int) -> None:
    """Creates the table, replacing any earlier copy. The vector column
    must have the same number of dimensions as the embedding model."""
    conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
    conn.execute(f"DROP TABLE IF EXISTS {TABLE}")
    # The table name is a constant and the dimensions are an integer, so
    # building them into the statement is safe; a statement cannot take
    # either as a parameter.
    conn.execute(
        f"""
        CREATE TABLE {TABLE} (
            id        integer PRIMARY KEY,
            title     text NOT NULL,
            content   text NOT NULL,
            embedding vector({int(dimensions)}) NOT NULL
        )
        """
    )

// Creates the table, replacing any earlier copy. The vector column must
// have the same number of dimensions as the embedding model.
async function createTable(client: pg.Client, dimensions: number): Promise<void> {
  await client.query("CREATE EXTENSION IF NOT EXISTS vector");
  await client.query(`DROP TABLE IF EXISTS ${table}`);
  // The table name is a constant and the dimensions are an integer, so
  // building them into the statement is safe; a statement cannot take
  // either as a parameter.
  await client.query(`
    CREATE TABLE ${table} (
      id        integer PRIMARY KEY,
      title     text NOT NULL,
      content   text NOT NULL,
      embedding vector(${Math.trunc(dimensions)}) NOT NULL
    )`);
}

/**
 * Creates the table, replacing any earlier copy. The vector column
 * must have the same number of dimensions as the embedding model.
 */
static void createTable(Connection conn, int dimensions) throws SQLException {
    try (Statement stmt = conn.createStatement()) {
        stmt.execute("CREATE EXTENSION IF NOT EXISTS vector");
        stmt.execute("DROP TABLE IF EXISTS " + TABLE);
        // The table name is a constant and the dimensions are an int,
        // so building them into the statement is safe; a statement
        // cannot take either as a parameter.
        stmt.execute("""
                CREATE TABLE %s (
                    id        integer PRIMARY KEY,
                    title     text NOT NULL,
                    content   text NOT NULL,
                    embedding vector(%d) NOT NULL
                )""".formatted(TABLE, dimensions));
    }
}

// Creates the table, replacing any earlier copy. The vector column must
// have the same number of dimensions as the embedding model.
static async Task CreateTableAsync(NpgsqlConnection conn, int dimensions)
{
    await using var cmd = conn.CreateCommand();
    cmd.CommandText = "CREATE EXTENSION IF NOT EXISTS vector";
    await cmd.ExecuteNonQueryAsync();
    cmd.CommandText = $"DROP TABLE IF EXISTS {Table}";
    await cmd.ExecuteNonQueryAsync();
    // The table name is a constant and the dimensions are an int, so
    // building them into the statement is safe; a statement cannot take
    // either as a parameter.
    cmd.CommandText = $"""
        CREATE TABLE {Table} (
            id        integer PRIMARY KEY,
            title     text NOT NULL,
            content   text NOT NULL,
            embedding vector({dimensions}) NOT NULL
        )
        """;
    await cmd.ExecuteNonQueryAsync();
}

// createTable creates the table, replacing any earlier copy. The vector
// column must have the same number of dimensions as the embedding model.
func createTable(ctx context.Context, tx pgx.Tx, dimensions int) error {
    if _, err := tx.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
        return err
    }
    if _, err := tx.Exec(ctx, "DROP TABLE IF EXISTS "+table); err != nil {
        return err
    }
    // The table name is a constant and the dimensions are an integer, so
    // building them into the statement is safe; a statement cannot take
    // either as a parameter.
    _, err := tx.Exec(ctx, `
        CREATE TABLE `+table+` (
            id        integer PRIMARY KEY,
            title     text NOT NULL,
            content   text NOT NULL,
            embedding vector(`+strconv.Itoa(dimensions)+`) NOT NULL
        )`)
    return err
}

/// Creates the table, replacing any earlier copy. The vector column must
/// have the same number of dimensions as the embedding model.
fn create_table(tx: &mut Transaction, dimensions: usize) -> Result<(), postgres::Error> {
    tx.batch_execute("CREATE EXTENSION IF NOT EXISTS vector")?;
    tx.batch_execute(&format!("DROP TABLE IF EXISTS {TABLE}"))?;
    // The table name is a constant and the dimensions are an integer, so
    // building them into the statement is safe; a statement cannot take
    // either as a parameter.
    tx.batch_execute(&format!(
        "CREATE TABLE {TABLE} (
            id        integer PRIMARY KEY,
            title     text NOT NULL,
            content   text NOT NULL,
            embedding vector({dimensions}) NOT NULL
        )"
    ))
}

# Creates the table, replacing any earlier copy. The vector column must
# have the same number of dimensions as the embedding model.
def create_table(conn, dimensions)
  conn.exec("CREATE EXTENSION IF NOT EXISTS vector")
  conn.exec("DROP TABLE IF EXISTS #{TABLE}")
  # The table name is a constant and Integer() accepts only an integer, so
  # building them into the statement is safe; a statement cannot take
  # either as a parameter.
  conn.exec(<<~SQL)
    CREATE TABLE #{TABLE} (
        id        integer PRIMARY KEY,
        title     text NOT NULL,
        content   text NOT NULL,
        embedding vector(#{Integer(dimensions)}) NOT NULL
    )
  SQL
end

def insert_documents(conn: psycopg.Connection,
                     embeddings: list[list[float]]) -> None:
    """Stores each document with its embedding."""
    rows = []
    for (doc_id, title, content), embedding in zip(DOCUMENTS, embeddings):
        # pgvector reads a vector from text such as '[0.1,0.2,0.3]', and
        # the cast tells the server how to interpret the parameter.
        vector = "[" + ",".join(str(value) for value in embedding) + "]"
        rows.append((doc_id, title, content, vector))

    with conn.cursor() as cur:
        cur.executemany(
            f"INSERT INTO {TABLE} (id, title, content, embedding) "
            "VALUES (%s, %s, %s, %s::vector)",
            rows,
        )

// Stores each document with its embedding.
async function insertDocuments(client: pg.Client, embeddings: number[][]): Promise<void> {
  for (const [i, doc] of documents.entries()) {
    // pgvector reads a vector from text such as '[0.1,0.2,0.3]', and the
    // cast tells the server how to interpret the parameter.
    const vector = `[${embeddings[i].join(",")}]`;
    await client.query(
      `INSERT INTO ${table} (id, title, content, embedding)
       VALUES ($1, $2, $3, $4::vector)`,
      [doc.id, doc.title, doc.content, vector],
    );
  }
}

/** Stores each document with its embedding, in a single batch. */
static void insertDocuments(Connection conn, List<float[]> embeddings)
        throws SQLException {
    String sql = "INSERT INTO " + TABLE + " (id, title, content, embedding) "
            + "VALUES (?, ?, ?, ?::vector)";
    try (PreparedStatement stmt = conn.prepareStatement(sql)) {
        for (int i = 0; i < DOCUMENTS.size(); i++) {
            Document doc = DOCUMENTS.get(i);
            // pgvector reads a vector from text such as
            // '[0.1,0.2,0.3]', and the cast tells the server how to
            // interpret the parameter.
            StringJoiner vector = new StringJoiner(",", "[", "]");
            for (float value : embeddings.get(i)) {
                vector.add(Float.toString(value));
            }
            stmt.setInt(1, doc.id());
            stmt.setString(2, doc.title());
            stmt.setString(3, doc.content());
            stmt.setString(4, vector.toString());
            stmt.addBatch();
        }
        stmt.executeBatch();
    }
}

// Stores each document with its embedding.
static async Task InsertDocumentsAsync(
    NpgsqlConnection conn, Document[] documents, float[][] embeddings)
{
    for (var i = 0; i < documents.Length; i++)
    {
        // pgvector reads a vector from text such as '[0.1,0.2,0.3]'. The
        // invariant culture always writes a full stop as the decimal
        // separator, and the cast tells the server how to interpret the
        // parameter.
        var vector = "[" + string.Join(",", embeddings[i].Select(
            value => value.ToString(CultureInfo.InvariantCulture))) + "]";
        await using var cmd = new NpgsqlCommand(
            $"INSERT INTO {Table} (id, title, content, embedding) " +
            "VALUES ($1, $2, $3, $4::vector)", conn)
        {
            Parameters =
            {
                new() { Value = documents[i].Id },
                new() { Value = documents[i].Title },
                new() { Value = documents[i].Content },
                new() { Value = vector },
            },
        };
        await cmd.ExecuteNonQueryAsync();
    }
}

// vectorText formats an embedding as the text that pgvector reads, such as
// '[0.1,0.2,0.3]'.
func vectorText(embedding []float32) string {
    parts := make([]string, len(embedding))
    for i, value := range embedding {
        parts[i] = strconv.FormatFloat(float64(value), 'g', -1, 32)
    }
    return "[" + strings.Join(parts, ",") + "]"
}

// insertDocuments stores each document with its embedding.
func insertDocuments(ctx context.Context, tx pgx.Tx, embeddings [][]float32) error {
    for i, doc := range documents {
        // The cast tells the server how to interpret the text parameter.
        _, err := tx.Exec(ctx,
            "INSERT INTO "+table+" (id, title, content, embedding) VALUES ($1, $2, $3, $4::vector)",
            doc.ID, doc.Title, doc.Content, vectorText(embeddings[i]))
        if err != nil {
            return err
        }
    }
    return nil
}

/// Stores each document with its embedding.
fn insert_documents(tx: &mut Transaction, embeddings: &[Vec<f32>]) -> Result<(), postgres::Error> {
    // The driver checks each value against the type the server expects,
    // and a Rust string cannot be sent as a vector, so the parameter is
    // declared as text and then cast to a vector.
    let statement = tx.prepare(&format!(
        "INSERT INTO {TABLE} (id, title, content, embedding) \
         VALUES ($1, $2, $3, $4::text::vector)"
    ))?;
    for ((id, title, content), embedding) in DOCUMENTS.iter().zip(embeddings) {
        // pgvector reads a vector from text such as '[0.1,0.2,0.3]'.
        let values: Vec<String> = embedding.iter().map(|value| value.to_string()).collect();
        let vector = format!("[{}]", values.join(","));
        tx.execute(&statement, &[id, title, content, &vector])?;
    }
    Ok(())
}

# Stores each document with its embedding.
def insert_documents(conn, embeddings)
  DOCUMENTS.zip(embeddings).each do |(id, title, content), embedding|
    # pgvector reads a vector from text such as '[0.1,0.2,0.3]', and the
    # cast tells the server how to interpret the parameter.
    vector = "[#{embedding.join(",")}]"
    conn.exec_params(
      "INSERT INTO #{TABLE} (id, title, content, embedding) " \
      "VALUES ($1, $2, $3, $4::vector)",
      [id, title, content, vector]
    )
  end
end

def describe_table(conn: psycopg.Connection) -> tuple[int, int]:
    """Returns the number of rows, and the number of dimensions of their
    embeddings."""
    row = conn.execute(
        f"SELECT count(*), max(vector_dims(embedding)) FROM {TABLE}"
    ).fetchone()
    return row[0], row[1]

// Returns the number of rows, and the number of dimensions of their
// embeddings.
async function describeTable(client: pg.Client): Promise<{ rows: number; dimensions: number }> {
  const result = await client.query<{ rows: string; dimensions: number }>(
    `SELECT count(*) AS rows, max(vector_dims(embedding)) AS dimensions FROM ${table}`,
  );
  // count(*) is a bigint, which node-postgres returns as a string.
  return { rows: Number(result.rows[0].rows), dimensions: result.rows[0].dimensions };
}

/**
 * Returns the number of rows, and the number of dimensions of their
 * embeddings.
 */
static long[] describeTable(Connection conn) throws SQLException {
    try (Statement stmt = conn.createStatement();
         ResultSet rs = stmt.executeQuery(
                 "SELECT count(*), max(vector_dims(embedding)) FROM " + TABLE)) {
        rs.next();
        return new long[] {rs.getLong(1), rs.getLong(2)};
    }
}

// Returns the number of rows, and the number of dimensions of their
// embeddings. count(*) returns a bigint, which maps to long, and
// vector_dims() returns an integer, which maps to int.
static async Task<(long Rows, int Dimensions)> DescribeTableAsync(NpgsqlConnection conn)
{
    await using var cmd = new NpgsqlCommand(
        $"SELECT count(*), max(vector_dims(embedding)) FROM {Table}", conn);
    await using var reader = await cmd.ExecuteReaderAsync();
    await reader.ReadAsync();
    return (reader.GetInt64(0), reader.GetInt32(1));
}

// describeTable returns the number of rows, and the number of dimensions
// of their embeddings.
func describeTable(ctx context.Context, conn *pgx.Conn) (rows, dimensions int, err error) {
    err = conn.QueryRow(ctx,
        "SELECT count(*), max(vector_dims(embedding)) FROM "+table,
    ).Scan(&rows, &dimensions)
    return rows, dimensions, err
}

/// Returns the number of rows, and the number of dimensions of their
/// embeddings. count(*) returns a bigint, which maps to i64, and
/// vector_dims() returns an integer, which maps to i32.
fn describe_table(client: &mut Client) -> Result<(i64, i32), postgres::Error> {
    let row = client.query_one(
        &format!("SELECT count(*), max(vector_dims(embedding)) FROM {TABLE}"),
        &[],
    )?;
    Ok((row.get(0), row.get(1)))
}

# Returns the number of rows, and the number of dimensions of their
# embeddings. Results arrive as strings, so convert them to Integers.
def describe_table(conn)
  row = conn.exec("SELECT count(*), max(vector_dims(embedding)) FROM #{TABLE}")[0]
  [row["count"].to_i, row["max"].to_i]
end

def drop_table(conn: psycopg.Connection) -> None:
    """Removes the table. A real application keeps it, so that the RAG
    Server can read it."""
    conn.execute(f"DROP TABLE IF EXISTS {TABLE}")

// Removes the table. A real application keeps it, so that the RAG Server
// can read it.
async function dropTable(client: pg.Client): Promise<void> {
  await client.query(`DROP TABLE IF EXISTS ${table}`);
}

/**
 * Removes the table. A real application keeps it, so that the RAG
 * Server can read it.
 */
static void dropTable(Connection conn) throws SQLException {
    try (Statement stmt = conn.createStatement()) {
        stmt.execute("DROP TABLE IF EXISTS " + TABLE);
    }
}

// Removes the table. A real application keeps it, so that the RAG Server
// can read it.
static async Task DropTableAsync(NpgsqlConnection conn)
{
    await using var cmd = new NpgsqlCommand($"DROP TABLE IF EXISTS {Table}", conn);
    await cmd.ExecuteNonQueryAsync();
}

// dropTable removes the table. A real application keeps it, so that the
// RAG Server can read it.
func dropTable(ctx context.Context, conn *pgx.Conn) error {
    _, err := conn.Exec(ctx, "DROP TABLE IF EXISTS "+table)
    return err
}

/// Removes the table. A real application keeps it, so that the RAG Server
/// can read it.
fn drop_table(client: &mut Client) -> Result<(), postgres::Error> {
    client.batch_execute(&format!("DROP TABLE IF EXISTS {TABLE}"))
}

# Removes the table. A real application keeps it, so that the RAG Server
# can read it.
def drop_table(conn)
  conn.exec("DROP TABLE IF EXISTS #{TABLE}")
end

def main() -> int:
    cleanup_failed = False
    try:
        embeddings = embed([content for _, _, content in DOCUMENTS])
        print(f"Embedded {len(embeddings)} documents with {EMBED_MODEL}")

        # Autocommit makes each statement take effect at once, so that the
        # transaction block below is the only place that groups them.
        with psycopg.connect(**connection_settings(), autocommit=True) as conn:
            try:
                with conn.transaction():
                    create_table(conn, len(embeddings[0]))
                    insert_documents(conn, embeddings)

                rows, dimensions = describe_table(conn)
                print(f"Stored {rows} rows with {dimensions}-dimension "
                      f"embeddings in {TABLE}")
            finally:
                # Remove the table whether or not a step failed. A failure
                # here is reported on its own, so that it cannot hide the
                # error that came first.
                try:
                    drop_table(conn)
                except psycopg.Error as e:
                    print(f"Cleanup error: {e}", file=sys.stderr)
                    cleanup_failed = True
    # httpx.HTTPError covers a failure to reach Ollama or an error status,
    # and psycopg.Error covers every error from the database driver.
    except (httpx.HTTPError, psycopg.Error, KeyError, ValueError) as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1

    return 1 if cleanup_failed else 0


if __name__ == "__main__":
    sys.exit(main())

// Returns a readable description of an error. When a host name resolves
// to several addresses, Node reports a failure to reach all of them as an
// AggregateError, whose own message is empty; fetch attaches the cause of
// a connection failure to a TypeError.
function describe(err: unknown): string {
  if (err instanceof AggregateError) {
    return err.errors.map(describe).join("; ");
  }
  if (err instanceof Error) {
    return err.cause ? `${err.message}: ${describe(err.cause)}` : err.message;
  }
  return String(err);
}

async function main(): Promise<number> {
  let client: pg.Client | undefined;
  let cleanupFailed = false;
  try {
    const embeddings = await embed(documents.map((doc) => doc.content));
    console.log(`Embedded ${embeddings.length} documents with ${embedModel}`);

    client = new pg.Client(dbConfig);
    await client.connect();
    try {
      await client.query("BEGIN");
      await createTable(client, embeddings[0].length);
      await insertDocuments(client, embeddings);
      await client.query("COMMIT");

      const { rows, dimensions } = await describeTable(client);
      console.log(`Stored ${rows} rows with ${dimensions}-dimension embeddings in ${table}`);
    } finally {
      // Roll back anything unfinished (a no-op after a commit), and remove
      // the table whether or not a step failed. A failure here is reported
      // on its own, so that it cannot hide the error that came first.
      try {
        await client.query("ROLLBACK");
        await dropTable(client);
      } catch (cleanupErr) {
        console.error(`Cleanup error: ${describe(cleanupErr)}`);
        cleanupFailed = true;
      }
    }
    return cleanupFailed ? 1 : 0;
  } catch (err) {
    console.error(`Error: ${describe(err)}`);
    return 1;
  } finally {
    // client is undefined if the program failed before it connected.
    await client?.end();
  }
}

process.exitCode = await main();

/**
 * Creates the table and inserts the documents in one transaction,
 * rolling it back if a step fails.
 */
static void store(Connection conn, List<float[]> embeddings) throws SQLException {
    conn.setAutoCommit(false);
    try {
        createTable(conn, embeddings.get(0).length);
        insertDocuments(conn, embeddings);
        conn.commit();
    } catch (SQLException e) {
        // Attach any failure to roll back to the first error, rather
        // than letting it replace that error.
        try {
            conn.rollback();
        } catch (SQLException rollbackError) {
            e.addSuppressed(rollbackError);
        }
        throw e;
    } finally {
        conn.setAutoCommit(true);
    }
}

public static void main(String[] args) {
    boolean cleanupFailed = false;
    try {
        List<float[]> embeddings =
                embed(DOCUMENTS.stream().map(Document::content).toList());
        System.out.println("Embedded " + embeddings.size()
                + " documents with " + EMBED_MODEL);

        // try-with-resources closes the connection when the block
        // exits, whether or not an error occurred.
        try (Connection conn = DriverManager.getConnection(
                connectionUrl(), connectionProperties())) {
            try {
                store(conn, embeddings);
                long[] summary = describeTable(conn);
                System.out.println("Stored " + summary[0] + " rows with "
                        + summary[1] + "-dimension embeddings in " + TABLE);
            } finally {
                // Remove the table whether or not a step failed. A
                // failure here is reported on its own, so that it
                // cannot hide the error that came first.
                try {
                    dropTable(conn);
                } catch (SQLException e) {
                    System.err.println("Cleanup error: " + e.getMessage());
                    cleanupFailed = true;
                }
            }
        }
    // IOException covers failures to reach Ollama, such as a refused
    // connection or a timeout. The message of a refused connection can
    // be empty, so print the exception itself, which includes its type.
    } catch (IOException | InterruptedException e) {
        System.err.println("Error: " + e);
        System.exit(1);
    // SQLException covers connection failures and errors reported by
    // the server, and JacksonException covers a reply that is not JSON.
    } catch (SQLException | EmbedException | JacksonException e) {
        System.err.println("Error: " + e.getMessage());
        System.exit(1);
    }
    if (cleanupFailed) {
        System.exit(1);
    }
}