Long-Term Memory: pgvector, Embeddings, Chunking, the memory_search Tool

const EMBEDDING_URL = 'https://openrouter.ai/api/v1/embeddings'
const EMBEDDING_MODEL = 'openai/text-embedding-3-small' // fixed at 1536 dimensions
 
// This is not a model call: no temperature, no streaming, no improvisation.
// The same text in always gives the same 1536 floats out.
export async function embed(texts) {
  const res = await fetch(EMBEDDING_URL, {
    method: 'POST',
    headers: {
      Authorization: `Bearer ${process.env.OPENROUTER_API_KEY}`,
      'Content-Type': 'application/json',
    },
    // Send a batch: one request computes N, saving round-trip latency rather than money -
    // billing is per token and independent of batching
    body: JSON.stringify({ model: EMBEDDING_MODEL, input: texts }),
  })
  if (!res.ok) throw new Error(`embedding call failed: ${res.status}`)
  const json = await res.json()
  return json.data.map((item) => item.embedding)
}
 
// 0.02 dollars per million tokens; characters estimated 1:1 as tokens (D6's conservative convention)
export const embeddingCostUsd = (texts) =>
  (texts.reduce((sum, t) => sum + t.length, 0) / 1_000_000) * 0.02

import os
 
import httpx
 
EMBEDDING_URL = "https://openrouter.ai/api/v1/embeddings"
EMBEDDING_MODEL = "openai/text-embedding-3-small"  # fixed at 1536 dimensions
 
 
async def embed(texts: list[str]) -> list[list[float]]:
    """Not a model call: no temperature, no streaming, and the same text always
    yields the same vector."""
    async with httpx.AsyncClient(timeout=30) as client:
        res = await client.post(
            EMBEDDING_URL,
            headers={"Authorization": f"Bearer {os.environ['OPENROUTER_API_KEY']}"},
            # Send a batch: this saves round-trip latency, not money - billing is per token
            json={"model": EMBEDDING_MODEL, "input": texts},
        )
    res.raise_for_status()
    return [item["embedding"] for item in res.json()["data"]]
 
 
def embedding_cost_usd(texts: list[str]) -> float:
    # 0.02 dollars per million tokens; characters estimated 1:1 as tokens (D6's convention)
    return sum(len(t) for t in texts) / 1_000_000 * 0.02

// Dependencies: java.net.http (JDK 11+) plus Jackson
static final String EMBEDDING_URL = "https://openrouter.ai/api/v1/embeddings";
static final String EMBEDDING_MODEL = "openai/text-embedding-3-small"; // 1536 dimensions
static final ObjectMapper MAPPER = new ObjectMapper();
 
// Not a model call: no temperature, no streaming, and the same text always yields the same
// vector. Store vectors as float[]: each pgvector dimension is a 4-byte float, so double[]
// doubles memory for nothing
static List<float[]> embed(List<String> texts) throws Exception {
    var payload = MAPPER.createObjectNode().put("model", EMBEDDING_MODEL);
    payload.set("input", MAPPER.valueToTree(texts));
    var request = HttpRequest.newBuilder(URI.create(EMBEDDING_URL))
        .header("Authorization", "Bearer " + System.getenv("OPENROUTER_API_KEY"))
        .header("Content-Type", "application/json")
        .POST(HttpRequest.BodyPublishers.ofString(payload.toString()))
        .build();
    var res = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
    if (res.statusCode() != 200) throw new IOException("embedding call failed: " + res.statusCode());
 
    var vectors = new ArrayList<float[]>();
    for (JsonNode item : MAPPER.readTree(res.body()).get("data")) {
        var node = item.get("embedding");
        var vector = new float[node.size()];
        for (int i = 0; i < node.size(); i++) vector[i] = (float) node.get(i).asDouble();
        vectors.add(vector);
    }
    return vectors;
}
 
// 0.02 dollars per million tokens; characters estimated 1:1 as tokens (D6's convention)
static double embeddingCostUsd(List<String> texts) {
    return texts.stream().mapToInt(String::length).sum() / 1_000_000.0 * 0.02;
}

// Declaring both request and response as Codable types makes a misspelled field a
// compile error
struct EmbeddingRequest: Encodable {
    let model = "openai/text-embedding-3-small" // fixed at 1536 dimensions
    let input: [String]
}
 
struct EmbeddingResponse: Decodable {
    struct Item: Decodable { let embedding: [Float] }
    let data: [Item]
}
 
/// Not a model call: no temperature, no streaming, and the same text always yields
/// the same vector
func embed(_ texts: [String]) async throws -> [[Float]] {
    let key = ProcessInfo.processInfo.environment["OPENROUTER_API_KEY"] ?? ""
    var request = URLRequest(url: URL(string: "https://openrouter.ai/api/v1/embeddings")!)
    request.httpMethod = "POST"
    request.setValue("Bearer \(key)", forHTTPHeaderField: "Authorization")
    request.setValue("application/json", forHTTPHeaderField: "Content-Type")
    // Send a batch: this saves round-trip latency, not money - billing is per token
    request.httpBody = try JSONEncoder().encode(EmbeddingRequest(input: texts))
 
    let (data, response) = try await URLSession.shared.data(for: request)
    guard (response as? HTTPURLResponse)?.statusCode == 200 else {
        throw URLError(.badServerResponse)
    }
    return try JSONDecoder().decode(EmbeddingResponse.self, from: data).data.map(\.embedding)
}
 
// 0.02 dollars per million tokens; characters estimated 1:1 as tokens (D6's convention)
func embeddingCostUsd(_ texts: [String]) -> Double {
    Double(texts.reduce(0) { $0 + $1.count }) / 1_000_000 * 0.02
}

const CHUNK_SIZE = 400 // target 400 characters: one card holds one thing
const OVERLAP = 80 // 80 characters of overlap: a split sentence stays intact in the other chunk
 
// Find the nearest natural boundary in the back half to finish on; 400 is a target,
// not a hard ceiling. The full stop below is the CJK one because the lab's fixtures are
// Chinese; for English text add '.' and '?' to the same search
export function chunkText(text) {
  const chunks = []
  let start = 0
  while (start < text.length) {
    let end = Math.min(start + CHUNK_SIZE, text.length)
    if (end < text.length) {
      const half = start + Math.floor(CHUNK_SIZE / 2)
      const window = text.slice(half, end)
      const hit = Math.max(window.lastIndexOf('。'), window.lastIndexOf('\n'))
      if (hit >= 0) end = half + hit + 1
    }
    const piece = text.slice(start, end).trim()
    if (piece) chunks.push(piece)
    if (end >= text.length) break
    start = Math.max(end - OVERLAP, start + 1) // step back to create overlap, always advancing
  }
  return chunks
}

CHUNK_SIZE = 400  # target 400 characters: one card holds one thing
OVERLAP = 80  # 80 characters of overlap: a split sentence stays intact in the other chunk
 
 
def chunk_text(text: str) -> list[str]:
    """Finish at the nearest natural boundary in the back half; 400 is a target,
    not a hard ceiling."""
    chunks: list[str] = []
    start = 0
    while start < len(text):
        end = min(start + CHUNK_SIZE, len(text))
        if end < len(text):
            half = start + CHUNK_SIZE // 2
            window = text[half:end]
            hit = max(window.rfind("。"), window.rfind("\n"))
            if hit >= 0:
                end = half + hit + 1
        piece = text[start:end].strip()
        if piece:
            chunks.append(piece)
        if end >= len(text):
            break
        start = max(end - OVERLAP, start + 1)  # step back to overlap, always advancing
    return chunks

// Dependencies: java.util. String.lastIndexOf is enough; no regex needed for boundaries
static final int CHUNK_SIZE = 400; // target 400 characters: one card holds one thing
static final int OVERLAP = 80;     // 80 characters of overlap keeps a split sentence intact
 
static List<String> chunkText(String text) {
    var chunks = new ArrayList<String>();
    int start = 0;
    while (start < text.length()) {
        int end = Math.min(start + CHUNK_SIZE, text.length());
        if (end < text.length()) {
            int half = start + CHUNK_SIZE / 2;
            String window = text.substring(half, end);
            int hit = Math.max(window.lastIndexOf('。'), window.lastIndexOf('\n'));
            if (hit >= 0) end = half + hit + 1;
        }
        String piece = text.substring(start, end).strip();
        if (!piece.isEmpty()) chunks.add(piece);
        if (end >= text.length()) break;
        start = Math.max(end - OVERLAP, start + 1); // step back to overlap, always advancing
    }
    return List.copyOf(chunks);
}

let chunkSize = 400 // target 400 characters: one card holds one thing
let overlap = 80 // 80 characters of overlap: a split sentence stays intact in the other chunk
 
// Swift's String indices are not integers, so flatten to [Character] and slice by index:
// this dodges index-offset traps and keeps "one character is one character" consistent
// with the other three versions
func chunkText(_ text: String) -> [String] {
    let chars = Array(text)
    var chunks: [String] = []
    var start = 0
    while start < chars.count {
        var end = min(start + chunkSize, chars.count)
        if end < chars.count {
            let half = start + chunkSize / 2
            if let hit = chars[half..<end].lastIndex(where: { $0 == "。" || $0 == "\n" }) {
                end = hit + 1
            }
        }
        let piece = String(chars[start..<end]).trimmingCharacters(in: .whitespacesAndNewlines)
        if !piece.isEmpty { chunks.append(piece) }
        if end >= chars.count { break }
        start = max(end - overlap, start + 1) // step back to overlap, always advancing
    }
    return chunks
}

// The tool definition follows D5's rules: format notes, a valid example, and both bounds
// in the schema
export const memorySearch = {
  name: 'memory_search',
  description:
    "Search this user's long-term memory semantically for settled facts and preferences. " +
    'Searches cross-session long-term memory only; for anything said in this session, read ' +
    'the context above rather than using this tool.',
  parameters: {
    type: 'object',
    properties: {
      query: {
        type: 'string',
        description: 'One natural-language sentence describing what you want to recall, for example the user food restrictions',
      },
      limit: {
        type: 'integer',
        description: 'How many rows to return, an integer between 1 and 10, default 5',
        minimum: 1,
        maximum: 10,
      },
    },
    required: ['query'],
    // No user_id here: identity comes only from the session, never from the model
  },
}
 
const MIN_SCORE = 0.3 // below this, return nothing; recalibrate when changing embedding models
 
export async function runMemorySearch(args, ctx) {
  const errors = validate(memorySearch, args) // D5's validator, reporting all at once
  if (errors.length > 0) {
    return `The call failed. ${errors.join('; ')}. Fix it and call memory_search again.`
  }
  const vector = await embedOne(args.query)
  // userId comes from the session context, not from an argument
  const hits = (await ctx.store.search(ctx.userId, vector, args.limit ?? 5)).filter(
    (hit) => hit.score >= MIN_SCORE
  )
  if (hits.length === 0) {
    // An empty result has to be said out loud: given an empty string the model treats it
    // as "no constraints" and invents one
    return 'No relevant long-term memory found. Do not speculate about the user preferences.'
  }
  return hits
    .map((hit, i) => `${i + 1}. [score ${hit.score.toFixed(2)} | ${hit.createdAt}] ${hit.content}`)
    .join('\n')
}

MEMORY_SEARCH = {
    "name": "memory_search",
    "description": (
        "Search this user's long-term memory semantically for settled facts and preferences. "
        "Searches cross-session long-term memory only; for anything said in this session, "
        "read the context above rather than using this tool."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "One natural-language sentence describing what you want to recall, for example the user food restrictions",
            },
            "limit": {
                "type": "integer",
                "description": "How many rows to return, an integer between 1 and 10, default 5",
                "minimum": 1,
                "maximum": 10,
            },
        },
        "required": ["query"],
        # No user_id here: identity comes only from the session, never from the model
    },
}
 
MIN_SCORE = 0.3  # below this, return nothing; recalibrate when changing embedding models
 
 
async def run_memory_search(args: dict, ctx: Context) -> str:
    errors = validate(MEMORY_SEARCH, args)  # D5's validator, reporting all at once
    if errors:
        return "The call failed. " + "; ".join(errors) + ". Fix it and call memory_search again."
 
    vector = await embed_one(args["query"])
    # user_id comes from the session context, not from an argument
    hits = [
        hit
        for hit in await ctx.store.search(ctx.user_id, vector, args.get("limit", 5))
        if hit.score >= MIN_SCORE
    ]
    if not hits:
        # An empty result has to be said out loud, or the model invents a preference
        return "No relevant long-term memory found. Do not speculate about the user preferences."
    return "\n".join(
        f"{i}. [score {hit.score:.2f} | {hit.created_at}] {hit.content}"
        for i, hit in enumerate(hits, start=1)
    )

// Dependencies: Jackson. The schema is a text block, sparing hand-escaped JSON
static final double MIN_SCORE = 0.3; // below this, return nothing; recalibrate per model
 
// readTree throws a checked exception and a static initializer has nowhere to declare
// throws. The schema is a hard-coded literal, so a parse failure can only mean you
// mistyped the JSON - wrap it and fail fast
private static JsonNode schema(String json) {
    try {
        return MAPPER.readTree(json);
    } catch (JsonProcessingException e) {
        throw new IllegalStateException("the built-in schema is not valid JSON", e);
    }
}
 
// A factory method as in Swift: D5's record ToolSpec requires a fourth argument, execute,
// and execute needs the session's ctx, so it cannot be a static final constant
static ToolSpec memorySearch(AgentContext ctx) {
    return new ToolSpec(
        "memory_search",
        "Search this user's long-term memory semantically for settled facts and preferences. "
            + "Searches cross-session long-term memory only; for anything said in this session, "
            + "read the context above rather than using this tool.",
        schema("""
        {
          "type": "object",
          "properties": {
            "query": {
              "type": "string",
              "description": "One natural-language sentence describing what you want to recall, for example the user food restrictions"
            },
            "limit": {
              "type": "integer",
              "description": "How many rows to return, an integer between 1 and 10, default 5",
              "minimum": 1,
              "maximum": 10
            }
          },
          "required": ["query"]
        }
            """), // No user_id here: identity comes only from the session
        // ToolImpl.apply throws no checked exception, so a retrieval failure is folded
        // here into one sentence for the model
        args -> {
            try {
                return runMemorySearch(args, ctx);
            } catch (Exception e) {
                return "Long-term memory is temporarily unavailable. Do not speculate about the user preferences.";
            }
        });
}
 
// ctx carries the session's userId and store, neither of which the model can touch
static String runMemorySearch(JsonNode args, AgentContext ctx) throws Exception {
    // Validation uses the same schema: rebuilding a ToolSpec only reads its parameters
    // and never runs execute
    var errors = validate(memorySearch(ctx), args); // D5's validator, all at once
    if (!errors.isEmpty()) {
        return "The call failed. " + String.join("; ", errors) + ". Fix it and call memory_search again.";
    }
    var vector = embedOne(args.get("query").asText());
    var hits = ctx.store().search(ctx.userId(), vector, args.path("limit").asInt(5)).stream()
        .filter(hit -> hit.score() >= MIN_SCORE)
        .toList();
    if (hits.isEmpty()) {
        // An empty result has to be said out loud, or the model invents a preference
        return "No relevant long-term memory found. Do not speculate about the user preferences.";
    }
    return IntStream.range(0, hits.size())
        .mapToObj(i -> "%d. [score %.2f | %s] %s".formatted(
            i + 1, hits.get(i).score(), hits.get(i).createdAt(), hits.get(i).content()))
        .collect(Collectors.joining("\n"));
}

// D5's ToolSpec binds the manual and the local implementation to one type, and execute is
// required - so Swift and Java must hand over the implementation too, unlike the JS version
// which can pass a schema object alone. That is a real difference between the four
// languages, not one being more verbose
func memorySearchTool(ctx: AgentContext) -> ToolSpec {
    ToolSpec(
        name: "memory_search",
        description: """
            Search this user's long-term memory semantically for settled facts and preferences. \
            Searches cross-session long-term memory only; for anything said in this session, \
            read the context above rather than using this tool.
            """,
        parameters: .init(
            properties: [
                "query": ParameterSpec(
                    type: "string",
                    description: "One natural-language sentence describing what you want to recall, for example the user food restrictions"),
                "limit": ParameterSpec(
                    type: "integer",
                    description: "How many rows to return, an integer between 1 and 10, default 5",
                    minimum: 1, maximum: 10),
            ],
            required: ["query"]), // No user_id here: identity comes only from the session
        // execute's signature does not throw, so a retrieval failure is folded here into
        // one sentence for the model rather than escaping and breaking the whole round
        execute: { args in
            (try? await runMemorySearch(args, ctx: ctx))
                ?? "Long-term memory is temporarily unavailable. Do not speculate about the user preferences."
        })
}
 
let minScore: Float = 0.3 // below this, return nothing; recalibrate per embedding model
 
func runMemorySearch(_ args: [String: JSONValue], ctx: AgentContext) async throws -> String {
    // Validation uses the same schema: rebuilding a ToolSpec here only reads its parameters
    if case .failed(let errors) = validate(memorySearchTool(ctx: ctx), args: args) { // D5's validator
        return "The call failed. \(errors.joined(separator: "; ")). Fix it and call memory_search again."
    }
    guard case .string(let query)? = args["query"] else {
        return "The call failed. Required argument query is missing. Fix it and call memory_search again."
    }
    var limit = 5
    if case .number(let value)? = args["limit"] { limit = Int(value) }
 
    // userID comes from the session context, not from an argument
    let vector = try await embedOne(query)
    let hits = try await ctx.store.search(userID: ctx.userID, vector: vector, limit: limit)
        .filter { $0.score >= minScore }
    guard !hits.isEmpty else {
        // An empty result has to be said out loud, or the model invents a preference
        return "No relevant long-term memory found. Do not speculate about the user preferences."
    }
    return hits.enumerated()
        .map { index, hit in
            let score = String(format: "%.2f", hit.score)
            return "\(index + 1). [score \(score) | \(hit.createdAt)] \(hit.content)"
        }
        .joined(separator: "\n")
}