Batch Ingestion

import json
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import islice
from pathlib import Path
from threading import Lock
from typing import Iterator

from dakera import DakeraClient

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(message)s')
logger = logging.getLogger(__name__)

client = DakeraClient(base_url="http://localhost:3000", api_key="dk-...")

# Configuration
AGENT_ID = "knowledge-base"
MAX_WORKERS = 20
RATE_LIMIT_RPS = 20        # requests per second, shared by all workers
CHUNK_SIZE = 800           # max characters per chunk (a longer sentence stays whole)
CHUNK_OVERLAP = 100        # overlap between consecutive chunks
CHECKPOINT_INTERVAL = 500  # source records between checkpoints
MAX_ATTEMPTS = 3           # attempts per chunk before a permanent failure
BATCH_ID = datetime.now(timezone.utc).isoformat()
BATCH_TAG = f"batch:{BATCH_ID}"  # one tag per run, for rollback with batch_forget

@dataclass
class Chunk:
    content: str
    doc_id: str
    source: str
    chunk_index: int
    total_chunks: int
    category: str
    importance: float = 0.80

def chunk_document(text: str, doc_id: str, source: str, category: str) -> list[Chunk]:
    """Split a document into sentence-aware overlapping chunks."""
    if len(text) <= CHUNK_SIZE:
        return [Chunk(text, doc_id, source, 0, 1, category)]

    chunks = []
    sentences = text.replace('. ', '.|').replace('! ', '!|').replace('? ', '?|').split('|')
    current_chunk = ""

    for sentence in sentences:
        if len(current_chunk) + len(sentence) + 1 > CHUNK_SIZE and current_chunk:
            chunks.append(current_chunk.strip())
            # Overlap: carry the last 100 characters into the next chunk
            current_chunk = current_chunk[-CHUNK_OVERLAP:] + " " + sentence
        else:
            current_chunk += (" " if current_chunk else "") + sentence

    if current_chunk.strip():
        chunks.append(current_chunk.strip())

    total = len(chunks)
    return [
        Chunk(content=c, doc_id=doc_id, source=source, chunk_index=i,
              total_chunks=total, category=category)
        for i, c in enumerate(chunks)
    ]

def load_checkpoint(checkpoint_file: str) -> int:
    """Number of source records a previous run finished (0 if none)."""
    p = Path(checkpoint_file)
    if p.exists():
        data = json.loads(p.read_text())
        logger.info(f"Resuming from checkpoint: record {data['last_index']}")
        return data["last_index"]
    return 0

def save_checkpoint(checkpoint_file: str, index: int, ingested: int, failed: int) -> None:
    """Record that every source record before `index` is done."""
    Path(checkpoint_file).write_text(json.dumps({
        "last_index": index,
        "ingested": ingested,
        "failed": failed,
        "timestamp": datetime.now(timezone.utc).isoformat()
    }, indent=2))

class RateLimiter:
    """Spaces requests so all workers together send at most `rps` per second."""
    def __init__(self, rps: float):
        self.min_interval = 1.0 / rps
        self.next_slot = time.monotonic()
        self.lock = Lock()

    def wait(self) -> None:
        with self.lock:
            now = time.monotonic()
            slot = max(now, self.next_slot)
            self.next_slot = slot + self.min_interval
        time.sleep(max(0.0, slot - now))

rate_limiter = RateLimiter(RATE_LIMIT_RPS)

def ingest_chunk(chunk: Chunk) -> bool:
    """Store one chunk, retrying with exponential backoff. Returns True on success."""
    for attempt in range(MAX_ATTEMPTS):
        try:
            rate_limiter.wait()
            client.store_memory(
                agent_id=AGENT_ID,
                content=chunk.content,
                importance=chunk.importance,
                memory_type="semantic",
                tags=[f"source:{chunk.source}", f"category:{chunk.category}", BATCH_TAG],
                metadata={
                    "source": chunk.source,
                    "doc_id": chunk.doc_id,
                    "chunk_index": chunk.chunk_index,
                    "total_chunks": chunk.total_chunks,
                    "category": chunk.category,
                    "ingestion_batch": BATCH_ID
                }
            )
            return True
        except Exception as e:
            if attempt < MAX_ATTEMPTS - 1:
                backoff = 2 ** attempt  # 1 s, then 2 s
                logger.warning(f"Retry {attempt + 1}/{MAX_ATTEMPTS - 1} for {chunk.doc_id}[{chunk.chunk_index}]: {e}. Waiting {backoff}s")
                time.sleep(backoff)
            else:
                logger.error(f"PERMANENT FAILURE: {chunk.doc_id}[{chunk.chunk_index}]: {e}")
    return False

def load_jsonl(file_path: str) -> Iterator[dict]:
    """Stream documents from a JSONL file."""
    with open(file_path) as f:
        for line in f:
            line = line.strip()
            if line:
                yield json.loads(line)

def batch_ingest(
    source_file: str,
    checkpoint_file: str = ".ingest_checkpoint.json",
    dry_run: bool = False
) -> dict:
    """Ingest a JSONL file in windows of CHECKPOINT_INTERVAL records, resumable after a crash."""
    start_index = load_checkpoint(checkpoint_file)
    records = enumerate(load_jsonl(source_file))
    for _ in islice(records, start_index):
        pass  # skip records a previous run finished

    ingested, failed, planned = 0, 0, 0
    failed_chunks: list[Chunk] = []
    start_time = time.monotonic()

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        while window := list(islice(records, CHECKPOINT_INTERVAL)):
            chunks = [
                chunk
                for i, doc in window
                for chunk in chunk_document(
                    text=doc["content"],
                    doc_id=doc.get("id", f"doc_{i:06d}"),
                    source=doc.get("source", source_file),
                    category=doc.get("category", "general"),
                )
            ]
            planned += len(chunks)
            if dry_run:
                continue

            # Every chunk of the window finishes (stored or failed) before the checkpoint
            for chunk, ok in zip(chunks, executor.map(ingest_chunk, chunks)):
                if ok:
                    ingested += 1
                else:
                    failed += 1
                    failed_chunks.append(chunk)

            done = window[-1][0] + 1
            save_checkpoint(checkpoint_file, done, ingested, failed)
            elapsed = time.monotonic() - start_time
            logger.info(f"Progress: {done} records | {ingested} chunks ok, {failed} failed | "
                        f"{ingested / max(elapsed, 1):.1f} chunks/s")

    if dry_run:
        logger.info(f"DRY RUN: would ingest {planned} chunks. Nothing was stored.")
        return {"dry_run": True, "chunks": planned}

    # Write failed chunks for manual review; never skip them silently
    if failed_chunks:
        failed_file = f"failed_chunks_{BATCH_ID[:10]}.jsonl"
        with open(failed_file, 'w') as f:
            for c in failed_chunks:
                f.write(json.dumps({"doc_id": c.doc_id, "chunk_index": c.chunk_index, "content": c.content[:200]}) + '\n')
        logger.warning(f"Wrote {len(failed_chunks)} failed chunks to {failed_file}")

    elapsed = time.monotonic() - start_time
    result = {
        "total_chunks": planned,
        "ingested": ingested,
        "failed": failed,
        "elapsed_seconds": round(elapsed, 1),
        "throughput_per_min": round(ingested / max(elapsed / 60, 1e-9)),
        "batch_tag": BATCH_TAG
    }
    logger.info(f"Batch complete: {result}")
    return result

def verify(sample_queries: list[str]) -> None:
    """Stage 5: count the agent's memories and spot-check recall."""
    stats = client.agent_stats(AGENT_ID)
    logger.info(f"Memories stored for {AGENT_ID}: {stats['total_memories']}")
    for q in sample_queries:
        hits = client.recall(agent_id=AGENT_ID, query=q, top_k=5)
        logger.info(f"{q!r}: {[m.content[:60] for m in hits.memories]}")

# Run the ingestion, then verify
result = batch_ingest("knowledge-base.jsonl", dry_run=False)
print(f"Ingested {result['ingested']:,} chunks in {result['elapsed_seconds']}s ({result['throughput_per_min']}/min)")
verify(["How do I configure SSO?", "What does the SAML setup need?"])

# Roll the whole run back if it went wrong:
# from dakera import BatchForgetRequest, BatchMemoryFilter
# client.batch_forget(BatchForgetRequest(agent_id=AGENT_ID, filter=BatchMemoryFilter(tags=[result["batch_tag"]])))

import { DakeraClient } from '@dakera-ai/dakera';
import { createReadStream, existsSync, readFileSync, writeFileSync } from 'fs';
import { createInterface } from 'readline';

const client = new DakeraClient({ baseUrl: 'http://localhost:3000', apiKey: 'dk-...' });

const CONFIG = {
  agentId: 'knowledge-base',
  maxWorkers: 20,
  rateLimitRps: 20,        // shared by all workers
  chunkSize: 800,
  chunkOverlap: 100,
  checkpointInterval: 500, // source records between checkpoints
  maxAttempts: 3,
  batchId: new Date().toISOString(),
};
const BATCH_TAG = `batch:${CONFIG.batchId}`;

interface Chunk {
  content: string;
  docId: string;
  source: string;
  chunkIndex: number;
  totalChunks: number;
  category: string;
}

function chunkDocument(text: string, docId: string, source: string, category: string): Chunk[] {
  if (text.length <= CONFIG.chunkSize) {
    return [{ content: text, docId, source, chunkIndex: 0, totalChunks: 1, category }];
  }
  const sentences = text.replace(/\. /g, '.|').replace(/! /g, '!|').replace(/\? /g, '?|').split('|');
  const rawChunks: string[] = [];
  let current = '';
  for (const sentence of sentences) {
    if (current.length + sentence.length + 1 > CONFIG.chunkSize && current) {
      rawChunks.push(current.trim());
      current = current.slice(-CONFIG.chunkOverlap) + ' ' + sentence;
    } else {
      current += (current ? ' ' : '') + sentence;
    }
  }
  if (current.trim()) rawChunks.push(current.trim());
  return rawChunks.map((content, i) => ({
    content, docId, source, chunkIndex: i, totalChunks: rawChunks.length, category,
  }));
}

// Spaces requests so all workers together send at most rateLimitRps per second
let nextSlot = Date.now();
async function rateLimit(): Promise<void> {
  const slot = Math.max(Date.now(), nextSlot);
  nextSlot = slot + 1000 / CONFIG.rateLimitRps;
  await new Promise(r => setTimeout(r, slot - Date.now()));
}

async function ingestChunk(chunk: Chunk): Promise<boolean> {
  for (let attempt = 0; attempt < CONFIG.maxAttempts; attempt++) {
    try {
      await rateLimit();
      await client.storeMemory(CONFIG.agentId, {
        content: chunk.content,
        importance: 0.80,
        memory_type: 'semantic',
        tags: [`source:${chunk.source}`, `category:${chunk.category}`, BATCH_TAG],
        metadata: {
          source: chunk.source,
          doc_id: chunk.docId,
          chunk_index: chunk.chunkIndex,
          total_chunks: chunk.totalChunks,
          category: chunk.category,
          ingestion_batch: CONFIG.batchId,
        },
      });
      return true;
    } catch (err) {
      if (attempt < CONFIG.maxAttempts - 1) {
        const backoff = 2 ** attempt * 1000; // 1 s, then 2 s
        console.warn(`Retry ${attempt + 1} for ${chunk.docId}[${chunk.chunkIndex}]: waiting ${backoff}ms`);
        await new Promise(r => setTimeout(r, backoff));
      } else {
        console.error(`PERMANENT FAILURE: ${chunk.docId}[${chunk.chunkIndex}]: ${err}`);
      }
    }
  }
  return false;
}

// Run tasks with at most `limit` in flight; results keep their order
async function pool<T, R>(items: T[], limit: number, fn: (item: T) => Promise<R>): Promise<R[]> {
  const results: R[] = new Array(items.length);
  let next = 0;
  const worker = async () => {
    while (next < items.length) {
      const i = next++;
      results[i] = await fn(items[i]);
    }
  };
  await Promise.all(Array.from({ length: limit }, worker));
  return results;
}

async function batchIngest(sourceFile: string, checkpointFile = '.ingest_checkpoint.json'): Promise<void> {
  const startIndex = existsSync(checkpointFile)
    ? JSON.parse(readFileSync(checkpointFile, 'utf8')).last_index as number
    : 0;
  const start = Date.now();
  let ingested = 0, failed = 0, index = 0;
  let window: { index: number; doc: any }[] = [];
  const failedChunks: Chunk[] = [];

  const flush = async () => {
    const chunks = window.flatMap(({ index: i, doc }) =>
      chunkDocument(doc.content, doc.id ?? `doc_${i}`, doc.source ?? sourceFile, doc.category ?? 'general'));
    const results = await pool(chunks, CONFIG.maxWorkers, ingestChunk);
    results.forEach((ok, i) => ok ? ingested++ : (failed++, failedChunks.push(chunks[i])));
    const done = window[window.length - 1].index + 1;
    writeFileSync(checkpointFile, JSON.stringify({ last_index: done, ingested, failed }));
    const rate = ingested / Math.max((Date.now() - start) / 1000, 1);
    console.log(`Progress: ${done} records | ${ingested} chunks ok, ${failed} failed | ${rate.toFixed(1)} chunks/s`);
    window = [];
  };

  const rl = createInterface({ input: createReadStream(sourceFile) });
  for await (const line of rl) {
    if (!line.trim()) continue;
    const i = index++;
    if (i < startIndex) continue; // finished by a previous run
    window.push({ index: i, doc: JSON.parse(line) });
    if (window.length === CONFIG.checkpointInterval) await flush();
  }
  if (window.length) await flush();

  if (failedChunks.length) {
    writeFileSync(`failed_chunks_${CONFIG.batchId.slice(0, 10)}.jsonl`,
      failedChunks.map(c => JSON.stringify({ doc_id: c.docId, chunk_index: c.chunkIndex })).join('\n'));
  }
  const elapsed = (Date.now() - start) / 1000;
  console.log(`Complete: ${ingested} ingested, ${failed} failed in ${elapsed.toFixed(1)}s (tag ${BATCH_TAG})`);
}

await batchIngest('knowledge-base.jsonl');

// Verify: count, then spot-check recall
const stats = await client.agentStats(CONFIG.agentId);
console.log(`Memories stored: ${stats.total_memories}`);
const { memories } = await client.recall(CONFIG.agentId, 'How do I configure SSO?', { top_k: 5 });
memories.forEach(m => console.log(m.score.toFixed(2), m.content.slice(0, 60)));

use std::sync::Arc;
use dakera_client::{DakeraClient, MemoryType, StoreMemoryRequest};
use tokio::{sync::Semaphore, task::JoinSet, time::{interval, sleep, Duration}};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let client = Arc::new(
        DakeraClient::builder("http://localhost:3000").api_key("dk-...").build()?,
    );
    let semaphore = Arc::new(Semaphore::new(20));          // 20 concurrent workers
    let mut ticker = interval(Duration::from_millis(50));  // 20 requests per second
    let batch_tag = "batch:2024-08-01T10:00:00Z".to_string(); // one tag per run

    // Output of the load, chunk and tag stages (see the Python tab)
    let chunks: Vec<String> = load_chunks("knowledge-base.jsonl");
    println!("Ingesting {} chunks", chunks.len());

    let mut set = JoinSet::new();
    for (i, content) in chunks.into_iter().enumerate() {
        ticker.tick().await;
        let permit = semaphore.clone().acquire_owned().await?;
        let client = client.clone();
        let tags = vec!["source:docs".to_string(), "category:general".to_string(), batch_tag.clone()];

        set.spawn(async move {
            let _permit = permit; // released when the task ends
            for attempt in 0..3u32 {
                let request = StoreMemoryRequest::new("knowledge-base", content.clone())
                    .with_type(MemoryType::Semantic)
                    .with_importance(0.80)
                    .with_tags(tags.clone());
                match client.store_memory(request).await {
                    Ok(_) => return (i, true),
                    Err(e) if attempt < 2 => {
                        eprintln!("Retry {} for chunk {}: {:?}", attempt + 1, i, e);
                        sleep(Duration::from_secs(1 << attempt)).await; // 1 s, then 2 s
                    }
                    Err(e) => eprintln!("Permanent failure at chunk {}: {:?}", i, e),
                }
            }
            (i, false)
        });
    }

    let (mut ingested, mut failed) = (0u64, 0u64);
    while let Some(result) = set.join_next().await {
        match result {
            Ok((_, true)) => ingested += 1,
            _ => failed += 1,
        }
    }
    println!("Batch complete: {} ingested, {} failed", ingested, failed);
    Ok(())
}

package main

import (
    "bufio"
    "context"
    "encoding/json"
    "fmt"
    "log"
    "os"
    "sync"
    "sync/atomic"
    "time"

    dakera "github.com/dakera-ai/dakera-go"
)

const (
    agentID     = "knowledge-base"
    maxWorkers  = 20
    maxAttempts = 3
    rateLimit   = 20 // requests per second, shared by all workers
    batchTag    = "batch:2024-08-01T10:00:00Z"
)

// One tick per allowed request, shared by every worker
var limiter = time.Tick(time.Second / rateLimit)

func ingestChunkWithRetry(ctx context.Context, client *dakera.Client, content, source, docID string) bool {
    importance := float32(0.80)
    for attempt := 0; attempt < maxAttempts; attempt++ {
        <-limiter
        _, err := client.StoreMemory(ctx, agentID, dakera.StoreMemoryRequest{
            Content:    content,
            MemoryType: "semantic",
            Importance: &importance,
            Tags:       []string{"source:" + source, "category:general", batchTag},
            Metadata: map[string]interface{}{
                "source": source,
                "doc_id": docID,
            },
        })
        if err == nil {
            return true
        }
        if attempt < maxAttempts-1 {
            backoff := time.Duration(1<<uint(attempt)) * time.Second // 1 s, then 2 s
            log.Printf("Retry %d for %s: %v. Waiting %v", attempt+1, docID, err, backoff)
            time.Sleep(backoff)
        } else {
            log.Printf("PERMANENT FAILURE: %s: %v", docID, err)
        }
    }
    return false
}

// Reads one chunk per JSONL line (chunking happens upstream; see the Python tab)
func batchIngest(ctx context.Context, client *dakera.Client, sourceFile string) error {
    f, err := os.Open(sourceFile)
    if err != nil {
        return err
    }
    defer f.Close()

    sem := make(chan struct{}, maxWorkers)
    var wg sync.WaitGroup
    var ingested, failed atomic.Int64
    total := 0

    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        line := scanner.Text()
        if line == "" {
            continue
        }
        var doc map[string]interface{}
        if err := json.Unmarshal([]byte(line), &doc); err != nil {
            continue
        }

        total++
        content, _ := doc["content"].(string)
        docID, _ := doc["id"].(string)
        if docID == "" {
            docID = fmt.Sprintf("doc_%06d", total)
        }
        source, _ := doc["source"].(string)

        sem <- struct{}{}
        wg.Add(1)
        go func(c, d, s string) {
            defer wg.Done()
            defer func() { <-sem }()
            if ingestChunkWithRetry(ctx, client, c, s, d) {
                ingested.Add(1)
            } else {
                failed.Add(1)
            }
        }(content, docID, source)
    }

    wg.Wait()
    fmt.Printf("Complete: %d ingested, %d failed out of %d total\n", ingested.Load(), failed.Load(), total)
    return scanner.Err()
}

func main() {
    client := dakera.NewClientWithOptions(dakera.ClientOptions{
        BaseURL: "http://localhost:3000",
        APIKey:  "dk-...",
    })
    if err := batchIngest(context.Background(), client, "knowledge-base.jsonl"); err != nil {
        log.Fatal(err)
    }
}