Context Window Management

from dakera import DakeraClient
from datetime import datetime, timezone
from typing import NamedTuple
import math

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

# Token budget configuration for 128k context model
TOKEN_BUDGET = {
    "system_prompt": 8_000,
    "pinned_memories": 12_000,
    "scored_memories": 20_000,  # ← Our managed budget
    "conversation": 18_000,
    "user_message": 3_000,
    "response_reserved": 67_000,
}

# Scoring weights — tune for your use case
SCORING_WEIGHTS = {
    "relevance": 0.50,   # Dakera semantic similarity score
    "importance": 0.30,  # User-defined importance 0.0-1.0
    "recency": 0.20,     # Decay factor from created_at
}
RECENCY_DECAY_HOURS = 168  # 7 days: memories older than this get near-zero recency

class ScoredMemory(NamedTuple):
    content: str
    composite_score: float
    token_estimate: int
    memory_id: str

def estimate_tokens(text: str) -> int:
    """Conservative character-based token estimate. Use tiktoken for precision."""
    return max(1, len(text) // 4)

def compute_recency_factor(created_at_iso: str | None) -> float:
    """Returns 1.0 for fresh memories, decaying toward 0.0 over RECENCY_DECAY_HOURS."""
    if not created_at_iso:
        return 0.5  # Unknown age — neutral recency
    try:
        created = datetime.fromisoformat(created_at_iso.replace('Z', '+00:00'))
        now = datetime.now(timezone.utc)
        hours_old = (now - created).total_seconds() / 3600
        return max(0.0, 1.0 - hours_old / RECENCY_DECAY_HOURS)
    except (ValueError, TypeError):
        return 0.5

def score_memory(mem: dict, relevance_score: float) -> ScoredMemory:
    """Compute composite score for a single recalled memory."""
    importance = mem.get("importance", 0.5)
    recency = compute_recency_factor(mem.get("created_at"))

    composite = (
        relevance_score * SCORING_WEIGHTS["relevance"]
        + importance * SCORING_WEIGHTS["importance"]
        + recency * SCORING_WEIGHTS["recency"]
    )

    return ScoredMemory(
        content=mem["content"],
        composite_score=round(composite, 4),
        token_estimate=estimate_tokens(mem["content"]),
        memory_id=mem.get("id", "")
    )

def build_memory_context(
    agent_id: str,
    query: str,
    memory_budget_tokens: int = TOKEN_BUDGET["scored_memories"]
) -> dict:
    """
    Recall, score, and trim memories to fit within the token budget.
    Returns context string + usage statistics.
    """
    # Step 1: Pin critical memories (goals, lessons, identity)
    pinned_result = client.recall(
        agent_id=agent_id,
        query="active goals lessons critical identity constraints",
        top_k=10,
        min_importance=0.90
    )
    pinned_memories = pinned_result.get("memories", [])
    pinned_tokens = sum(estimate_tokens(m["content"]) for m in pinned_memories)

    # Step 2: Broad recall for scored candidates
    candidate_result = client.recall(
        agent_id=agent_id,
        query=query,
        top_k=50,
        min_importance=0.30
    )

    # Step 3: Score all candidates
    # Dakera returns memories in relevance order -- index 0 is most relevant
    # Approximate relevance scores: 1.0 for rank 1, decaying by rank
    scored: list[ScoredMemory] = []
    total_memories = len(candidate_result.get("memories", []))
    for i, mem in enumerate(candidate_result.get("memories", [])):
        approx_relevance = 1.0 - (i / max(total_memories, 1)) * 0.6
        scored.append(score_memory(mem, approx_relevance))

    # Step 4: Sort by composite score
    scored.sort(key=lambda m: m.composite_score, reverse=True)

    # Step 5: Fill budget greedily
    selected: list[ScoredMemory] = []
    tokens_used = 0
    trimmed_count = 0

    for mem in scored:
        if tokens_used + mem.token_estimate > memory_budget_tokens:
            trimmed_count += 1
            continue
        selected.append(mem)
        tokens_used += mem.token_estimate

    # Step 6: Build context strings
    pinned_context = "
".join(m["content"] for m in pinned_memories)
    scored_context = "
".join(m.content for m in selected)

    full_context = ""
    if pinned_context:
        full_context += f"[Critical context]
{pinned_context}

"
    if scored_context:
        full_context += f"[Relevant context]
{scored_context}"

    return {
        "context": full_context,
        "stats": {
            "pinned_memories": len(pinned_memories),
            "pinned_tokens": pinned_tokens,
            "scored_selected": len(selected),
            "scored_trimmed": trimmed_count,
            "scored_tokens": tokens_used,
            "total_tokens": pinned_tokens + tokens_used,
            "budget_utilization": round(tokens_used / memory_budget_tokens, 3),
        }
    }

# --- Usage in a long document analysis agent ---
def analyze_document(agent_id: str, document_text: str, question: str) -> str:
    # Build memory context scoped to the question
    context = build_memory_context(agent_id, query=question, memory_budget_tokens=18_000)

    stats = context["stats"]
    print(f"Memory context: {stats['scored_selected']} memories, "
          f"{stats['total_tokens']:,} tokens, "
          f"{stats['budget_utilization']:.0%} budget used, "
          f"{stats['scored_trimmed']} memories trimmed")

    system_prompt = f"""You are a document analysis assistant.

{context['context']}

Document excerpt:
{document_text[:10000]}"""

    # return llm.complete(system=system_prompt, user=question)
    return f"[Context loaded: {stats['total_tokens']} tokens] Ready to answer: {question}"

result = analyze_document("doc-analysis-agent", "...", "What are the key risk factors?")
print(result)

import { DakeraClient } from '@dakera-ai/dakera';

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

const TOKEN_BUDGET = {
  systemPrompt: 8_000,
  pinnedMemories: 12_000,
  scoredMemories: 20_000,
  conversation: 18_000,
  userMessage: 3_000,
  responseReserved: 67_000,
};

const SCORING_WEIGHTS = { relevance: 0.5, importance: 0.3, recency: 0.2 };
const RECENCY_DECAY_HOURS = 168; // 7 days

function estimateTokens(text: string): number {
  return Math.max(1, Math.floor(text.length / 4));
}

function computeRecencyFactor(createdAt?: string): number {
  if (!createdAt) return 0.5;
  const hoursOld = (Date.now() - new Date(createdAt).getTime()) / 3_600_000;
  return Math.max(0, 1 - hoursOld / RECENCY_DECAY_HOURS);
}

interface ScoredMemory {
  content: string;
  compositeScore: number;
  tokenEstimate: number;
  memoryId: string;
}

function scoreMemory(mem: { content: string; importance?: number; created_at?: string; id?: string }, rankPosition: number, totalMemories: number): ScoredMemory {
  const approxRelevance = 1.0 - (rankPosition / Math.max(totalMemories, 1)) * 0.6;
  const importance = mem.importance ?? 0.5;
  const recency = computeRecencyFactor(mem.created_at);

  const composite =
    approxRelevance * SCORING_WEIGHTS.relevance +
    importance * SCORING_WEIGHTS.importance +
    recency * SCORING_WEIGHTS.recency;

  return {
    content: mem.content,
    compositeScore: Math.round(composite * 10000) / 10000,
    tokenEstimate: estimateTokens(mem.content),
    memoryId: mem.id ?? '',
  };
}

async function buildMemoryContext(
  agentId: string,
  query: string,
  memoryBudget = TOKEN_BUDGET.scoredMemories
): Promise<{ context: string; stats: Record<string, number> }> {
  // Pinned memories (always included)
  const pinnedResult = await client.recall(agentId, 'active goals lessons critical identity', {
    top_k: 10,
    min_importance: 0.9,
  });
  const pinnedTokens = pinnedResult.memories.reduce((sum, m) => sum + estimateTokens(m.content), 0);

  // Broad candidate recall
  const candidates = await client.recall(agentId, query, { top_k: 50, min_importance: 0.3 });

  // Score and sort
  const scored: ScoredMemory[] = candidates.memories
    .map((m, i) => scoreMemory(m, i, candidates.memories.length))
    .sort((a, b) => b.compositeScore - a.compositeScore);

  // Budget fill
  const selected: ScoredMemory[] = [];
  let tokensUsed = 0;
  let trimmedCount = 0;

  for (const mem of scored) {
    if (tokensUsed + mem.tokenEstimate > memoryBudget) { trimmedCount++; continue; }
    selected.push(mem);
    tokensUsed += mem.tokenEstimate;
  }

  const pinnedBlock = pinnedResult.memories.map(m => m.content).join('
');
  const scoredBlock = selected.map(m => m.content).join('
');
  const context = [
    pinnedBlock ? `[Critical context]
${pinnedBlock}` : '',
    scoredBlock ? `[Relevant context]
${scoredBlock}` : '',
  ].filter(Boolean).join('

');

  return {
    context,
    stats: {
      pinnedCount: pinnedResult.memories.length,
      pinnedTokens,
      selectedCount: selected.length,
      trimmedCount,
      scoredTokens: tokensUsed,
      totalTokens: pinnedTokens + tokensUsed,
      budgetUtilization: Math.round((tokensUsed / memoryBudget) * 1000) / 1000,
    },
  };
}

// Usage
const { context, stats } = await buildMemoryContext('doc-analysis-agent', 'key risk factors in contract', 18_000);
console.log(`Loaded ${stats.selectedCount} memories, ${stats.totalTokens} tokens (${(stats.budgetUtilization * 100).toFixed(0)}% of budget)`);

use dakera_rs::{Client, RecallRequest};
use std::time::{SystemTime, UNIX_EPOCH};

let client = Client::new("http://localhost:3300", "dk-...");

const MEMORY_BUDGET_TOKENS: usize = 20_000;
const RECENCY_DECAY_HOURS: f64 = 168.0;

fn estimate_tokens(text: &str) -> usize {
    (text.len() / 4).max(1)
}

fn compute_recency(hours_old: f64) -> f64 {
    (1.0 - hours_old / RECENCY_DECAY_HOURS).max(0.0)
}

// Broad candidate recall
let candidates = client.recall("doc-analysis-agent", RecallRequest {
    query: "document analysis context user preferences".into(),
    top_k: Some(50),
    min_importance: Some(0.3),
    ..Default::default()
}).await?;

// Score and rank memories
let mut scored: Vec<(f64, &str, usize)> = candidates.memories.iter()
    .enumerate()
    .map(|(i, mem)| {
        let total = candidates.memories.len() as f64;
        let relevance = 1.0 - (i as f64 / total.max(1.0)) * 0.6;
        let importance = mem.importance.unwrap_or(0.5) as f64;
        let recency = compute_recency(24.0); // Simplified -- use actual created_at in production
        let score = relevance * 0.5 + importance * 0.3 + recency * 0.2;
        let tokens = estimate_tokens(&mem.content);
        (score, mem.content.as_str(), tokens)
    })
    .collect();

scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());

// Budget fill
let mut context_parts: Vec<&str> = Vec::new();
let mut tokens_used = 0usize;
let mut trimmed = 0usize;

for (_, content, tokens) in &scored {
    if tokens_used + tokens > MEMORY_BUDGET_TOKENS {
        trimmed += 1;
        continue;
    }
    context_parts.push(content);
    tokens_used += tokens;
}

let memory_context = context_parts.join("
");
println!("Selected {} memories, {} tokens, {} trimmed", context_parts.len(), tokens_used, trimmed);

client := dakera.NewClient("http://localhost:3300", "dk-...")
ctx := context.Background()

const memoryBudgetTokens = 20_000

func estimateTokens(text string) int {
    t := len(text) / 4
    if t < 1 { return 1 }
    return t
}

// Broad recall
candidates, _ := client.Recall(ctx, "doc-analysis-agent", dakera.RecallRequest{
    Query:         "document analysis context user preferences",
    TopK:          50,
    MinImportance: 0.3,
})

type scoredMem struct {
    content string
    score   float64
    tokens  int
}

total := float64(len(candidates.Memories))
scored := make([]scoredMem, len(candidates.Memories))
for i, mem := range candidates.Memories {
    relevance := 1.0 - float64(i)/math.Max(total, 1)*0.6
    importance := mem.Importance
    recency := 0.7 // Simplified -- compute from created_at in production
    compositeScore := relevance*0.5 + importance*0.3 + recency*0.2
    scored[i] = scoredMem{mem.Content, compositeScore, estimateTokens(mem.Content)}
}

sort.Slice(scored, func(i, j int) bool { return scored[i].score > scored[j].score })

// Budget fill
var selected []string
tokensUsed, trimmed := 0, 0
for _, m := range scored {
    if tokensUsed+m.tokens > memoryBudgetTokens {
        trimmed++
        continue
    }
    selected = append(selected, m.content)
    tokensUsed += m.tokens
}

memoryContext := strings.Join(selected, "
")
fmt.Printf("Selected %d memories, %d tokens, %d trimmed
", len(selected), tokensUsed, trimmed)