Batch Ingestion

import asyncio
import json
import logging
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterator
import concurrent.futures

from dakera import DakeraClient

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

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

# Configuration
AGENT_ID = "knowledge-base"
MAX_WORKERS = 20
RATE_LIMIT_RPS = 20        # requests per second
CHUNK_SIZE = 800            # max characters per chunk
CHUNK_OVERLAP = 100         # overlap between consecutive chunks
CHECKPOINT_INTERVAL = 500  # save progress every N records
MAX_RETRIES = 3
BATCH_ID = datetime.now(timezone.utc).isoformat()

@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 = ""
    chunk_idx = 0

    for sentence in sentences:
        if len(current_chunk) + len(sentence) + 1 > CHUNK_SIZE and current_chunk:
            chunks.append(current_chunk.strip())
            # Overlap: keep last 100 chars of current chunk
            overlap_start = max(0, len(current_chunk) - CHUNK_OVERLAP)
            current_chunk = current_chunk[overlap_start:] + " " + sentence
            chunk_idx += 1
        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:
    """Load last successfully ingested index from checkpoint file."""
    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:
    """Write checkpoint for resume capability."""
    Path(checkpoint_file).write_text(json.dumps({
        "last_index": index,
        "ingested": ingested,
        "failed": failed,
        "timestamp": datetime.now(timezone.utc).isoformat()
    }, indent=2))

class RateLimiter:
    """Token bucket rate limiter."""
    def __init__(self, rps: float):
        self.min_interval = 1.0 / rps
        self.last_call = 0.0

    def wait(self) -> None:
        now = time.monotonic()
        elapsed = now - self.last_call
        wait_time = self.min_interval - elapsed
        if wait_time > 0:
            time.sleep(wait_time)
        self.last_call = time.monotonic()

rate_limiter = RateLimiter(RATE_LIMIT_RPS)

def ingest_chunk(chunk: Chunk) -> bool:
    """Ingest a single chunk with retry logic. Returns True on success."""
    for attempt in range(MAX_RETRIES):
        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}",
                    f"batch:{BATCH_ID[:10]}"
                ],
                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_RETRIES - 1:
                backoff = 2 ** attempt
                logger.warning(f"Retry {attempt + 1}/{MAX_RETRIES} 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
    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:
    """Main batch ingestion function with progress tracking and resume support."""

    # Resume from checkpoint if available
    start_index = load_checkpoint(checkpoint_file)
    ingested, failed = 0, 0
    start_time = time.monotonic()

    # Collect all chunks to ingest (skipping already-processed records)
    all_chunks: list[Chunk] = []
    for i, doc in enumerate(load_jsonl(source_file)):
        if i < start_index:
            continue  # Skip already processed
        doc_chunks = 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")
        )
        all_chunks.extend(doc_chunks)

    total = len(all_chunks)
    logger.info(f"Ingesting {total} chunks from {source_file} (starting at record {start_index})")

    if dry_run:
        logger.info(f"DRY RUN: would ingest {total} chunks. Exiting.")
        return {"dry_run": True, "chunks": total}

    failed_chunks: list[Chunk] = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(ingest_chunk, chunk): chunk for chunk in all_chunks}

        for i, future in enumerate(concurrent.futures.as_completed(futures)):
            chunk = futures[future]
            if future.result():
                ingested += 1
            else:
                failed += 1
                failed_chunks.append(chunk)

            # Progress logging and checkpoint
            if (i + 1) % CHECKPOINT_INTERVAL == 0:
                elapsed = time.monotonic() - start_time
                rate = ingested / elapsed if elapsed > 0 else 0
                eta_min = (total - i - 1) / max(rate * 60, 1)
                logger.info(
                    f"Progress: {i+1}/{total} chunks | "
                    f"{ingested} ok, {failed} failed | "
                    f"{rate:.1f} chunks/s | ETA: {eta_min:.1f} min"
                )
                save_checkpoint(checkpoint_file, start_index + i + 1, ingested, failed)

    # Final checkpoint
    save_checkpoint(checkpoint_file, start_index + total, ingested, failed)

    # Write failed records for manual review
    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]}) + '
')
        logger.warning(f"Wrote {len(failed_chunks)} failed chunks to {failed_file}")

    elapsed = time.monotonic() - start_time
    result = {
        "total_chunks": total,
        "ingested": ingested,
        "failed": failed,
        "elapsed_seconds": round(elapsed, 1),
        "throughput_per_min": round(ingested / (elapsed / 60), 0),
        "success_rate": round(ingested / max(total, 1), 4)
    }
    logger.info(f"Batch complete: {result}")
    return result

# Run the ingestion
result = batch_ingest("knowledge-base.jsonl", dry_run=False)
print(f"Ingested {result['ingested']:,} chunks in {result['elapsed_seconds']}s ({result['throughput_per_min']:.0f}/min)")

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

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

const CONFIG = {
  agentId: 'knowledge-base',
  maxWorkers: 20,
  rateLimitRps: 20,
  chunkSize: 800,
  chunkOverlap: 100,
  checkpointInterval: 500,
  maxRetries: 3,
  batchId: new Date().toISOString(),
};

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

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());
      const overlapStart = Math.max(0, current.length - CONFIG.chunkOverlap);
      current = current.slice(overlapStart) + ' ' + 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,
  }));
}

class RateLimiter {
  private lastCall = 0;
  private readonly minInterval: number;
  constructor(rps: number) { this.minInterval = 1000 / rps; }

  async wait(): Promise<void> {
    const now = Date.now();
    const elapsed = now - this.lastCall;
    if (elapsed < this.minInterval) {
      await new Promise(r => setTimeout(r, this.minInterval - elapsed));
    }
    this.lastCall = Date.now();
  }
}

const rateLimiter = new RateLimiter(CONFIG.rateLimitRps);

async function ingestChunk(chunk: Chunk): Promise<boolean> {
  for (let attempt = 0; attempt < CONFIG.maxRetries; attempt++) {
    try {
      await rateLimiter.wait();
      await client.storeMemory(CONFIG.agentId, {
        content: chunk.content,
        importance: chunk.importance ?? 0.80,
        memoryType: 'semantic',
        tags: [`source:${chunk.source}`, `category:${chunk.category}`, `batch:${CONFIG.batchId.slice(0, 10)}`],
      });
      return true;
    } catch (err) {
      if (attempt < CONFIG.maxRetries - 1) {
        const backoff = Math.pow(2, attempt) * 1000;
        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;
      }
    }
  }
  return false;
}

async function batchIngest(sourceFile: string, checkpointFile = '.ingest_checkpoint.json'): Promise<void> {
  const start = Date.now();
  let ingested = 0, failed = 0;

  // Read all documents and chunk them
  const allChunks: Chunk[] = [];
  const rl = createInterface({ input: createReadStream(sourceFile) });
  let lineIndex = 0;
  for await (const line of rl) {
    if (!line.trim()) continue;
    const doc = JSON.parse(line);
    const chunks = chunkDocument(doc.content, doc.id ?? `doc_${lineIndex}`, doc.source ?? sourceFile, doc.category ?? 'general');
    allChunks.push(...chunks);
    lineIndex++;
  }

  console.log(`Ingesting ${allChunks.length} chunks from ${lineIndex} documents`);

  // Process in batches of maxWorkers
  const BATCH = CONFIG.maxWorkers;
  for (let i = 0; i < allChunks.length; i += BATCH) {
    const batch = allChunks.slice(i, i + BATCH);
    const results = await Promise.all(batch.map(c => ingestChunk(c)));
    ingested += results.filter(Boolean).length;
    failed += results.filter(r => !r).length;

    if ((i + BATCH) % CONFIG.checkpointInterval < BATCH) {
      const elapsed = (Date.now() - start) / 1000;
      const rate = ingested / elapsed;
      const eta = (allChunks.length - i - BATCH) / Math.max(rate * 60, 1);
      console.log(`Progress: ${i + BATCH}/${allChunks.length} | ${ingested} ok | ${rate.toFixed(1)} chunks/s | ETA: ${eta.toFixed(1)} min`);
      writeFileSync(checkpointFile, JSON.stringify({ index: i + BATCH, ingested, failed }));
    }
  }

  const elapsed = (Date.now() - start) / 1000;
  console.log(`Complete: ${ingested} ingested, ${failed} failed in ${elapsed.toFixed(1)}s (${(ingested / (elapsed / 60)).toFixed(0)}/min)`);
}

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

use dakera_rs::{Client, StoreMemoryRequest};
use std::sync::Arc;
use tokio::{sync::Semaphore, task::JoinSet, time::{sleep, Duration}};

let client = Arc::new(Client::new("http://localhost:3300", "dk-..."));
let semaphore = Arc::new(Semaphore::new(20)); // 20 concurrent workers

let documents: Vec<String> = load_documents("knowledge-base.jsonl");
let total = documents.len();
println!("Ingesting {} documents", total);

let mut set = JoinSet::new();
let mut ingested = 0u64;
let mut failed = 0u64;

for (i, content) in documents.into_iter().enumerate() {
    let client = client.clone();
    let permit = semaphore.clone().acquire_owned().await.unwrap();

    set.spawn(async move {
        let _permit = permit; // Released when dropped
        let mut last_err = None;

        // Exponential backoff retry
        for attempt in 0..3u32 {
            match client.store_memory("knowledge-base", StoreMemoryRequest {
                content: content.clone(),
                importance: Some(0.80),
                memory_type: "semantic".into(),
                tags: vec!["source:docs".into(), "category:general".into()],
                ..Default::default()
            }).await {
                Ok(_) => return (i, true),
                Err(e) => {
                    last_err = Some(e);
                    if attempt < 2 {
                        sleep(Duration::from_secs(2u64.pow(attempt))).await;
                    }
                }
            }
        }
        eprintln!("Permanent failure at doc {}: {:?}", i, last_err);
        (i, false)
    });
}

while let Some(result) = set.join_next().await {
    match result {
        Ok((_, true)) => ingested += 1,
        Ok((_, false)) => failed += 1,
        Err(e) => { failed += 1; eprintln!("Join error: {:?}", e); }
    }
    if (ingested + failed) % 500 == 0 {
        println!("Progress: {} ok, {} failed / {} total", ingested, failed, total);
    }
}

println!("Batch complete: {} ingested, {} failed", ingested, failed);

package main

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

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

const (
    maxWorkers = 20
    maxRetries = 3
    rateLimit  = 20 // req/s
)

func ingestChunkWithRetry(ctx context.Context, client *dakera.Client, content, source, docID string) bool {
    for attempt := 0; attempt < maxRetries; attempt++ {
        _, err := client.StoreMemory(ctx, "knowledge-base", dakera.StoreMemoryRequest{
            Content:    content,
            Importance: 0.80,
            MemoryType: "semantic",
            Tags:       []string{"source:" + source, "category:general"},
            Metadata: map[string]interface{}{
                "source": source,
                "doc_id": docID,
            },
        })
        if err == nil {
            return true
        }
        if attempt < maxRetries-1 {
            backoff := time.Duration(1<<uint(attempt)) * time.Second
            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
}

func batchIngest(ctx context.Context, client *dakera.Client, sourceFile string) {
    f, _ := os.Open(sourceFile)
    defer f.Close()

    sem := make(chan struct{}, maxWorkers)
    rateTicker := time.NewTicker(time.Second / rateLimit)
    defer rateTicker.Stop()

    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)

        <-rateTicker.C // Rate limit
        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)
            }

            if (ingested.Load()+failed.Load())%500 == 0 {
                fmt.Printf("Progress: %d ok, %d failed
", ingested.Load(), failed.Load())
            }
        }(content, docID, source)
    }

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

func main() {
    client := dakera.NewClient("http://localhost:3300", "dk-...")
    batchIngest(context.Background(), client, "knowledge-base.jsonl")
}