src/dynamicRetrieval.ts

import { ScoredMemory, MemoryConfig, DEFAULT_MEMORY_CONFIG, MemoryType } from './types';
import { applyDecayWeighting } from './memoryDecay';
import { formatMemoriesByType } from './typedMemory';

type QueryComplexity = 'simple' | 'moderate' | 'complex';

export function estimateQueryComplexity(query: string): QueryComplexity {
  const wc = query.split(/\s+/).length;
  if (wc <= 5) return 'simple';
  if (/\b(what('s|\s+is)\s+my|am\s+i|do\s+i|where\s+do\s+i)\b/i.test(query)) return 'simple';
  if (wc >= 15) return 'complex';
  if (/\b(plan|schedule|organize|considering|taking\s+into\s+account|based\s+on)\b/i.test(query)) return 'complex';
  if (/\band\b.*\band\b/i.test(query)) return 'complex';
  return 'moderate';
}

export function inferRelevantMemoryTypes(query: string): MemoryType[] {
  const types: MemoryType[] = []; const q = query.toLowerCase();
  if (/\b(yesterday|last|when|recently|remember\s+when|did\s+we|did\s+i|what\s+happened)\b/.test(q)) types.push('episodic');
  if (/\b(how\s+(should|do)|what's\s+the\s+rule|format|style|remind|procedure)\b/.test(q)) types.push('procedural');
  if (/\b(what('s|\s+is)\s+my|suggest|recommend|prefer|like|allergic|favorite|do\s+i\s+(like|have|live))\b/.test(q)) types.push('semantic');
  return types.length ? types : ['semantic', 'episodic', 'procedural'];
}

export interface RetrievalParams { topK: number; similarityThreshold: number; maxTokenBudget: number; priorityTypes: MemoryType[]; boostRecent: boolean; }

export function computeRetrievalParams(query: string, config: MemoryConfig = DEFAULT_MEMORY_CONFIG): RetrievalParams {
  const complexity = estimateQueryComplexity(query);
  const priorityTypes = inferRelevantMemoryTypes(query);
  const boostRecent = /\b(today|just|now|this\s+(morning|afternoon|evening)|recently|latest)\b/i.test(query);
  switch (complexity) {
    case 'simple': return { topK: 3, similarityThreshold: 0.50, maxTokenBudget: 100, priorityTypes, boostRecent };
    case 'moderate': return { topK: config.defaultTopK, similarityThreshold: config.minSimilarityThreshold, maxTokenBudget: config.maxContextTokens, priorityTypes, boostRecent };
    case 'complex': return { topK: 12, similarityThreshold: 0.35, maxTokenBudget: 500, priorityTypes, boostRecent };
  }
}

export function buildDynamicRetrievalQuery(embedding: number[], params: RetrievalParams) {
  const typeFilter = params.priorityTypes.length < 3 ? `AND type IN (${params.priorityTypes.map(() => '?').join(',')})` : '';
  const orderClause = params.boostRecent
    ? `ORDER BY (1 - vec_distance_cosine(embedding, vec_f32(?))) * (1 + 0.2 * CASE WHEN (julianday('now') - julianday(last_accessed_at / 1000, 'unixepoch')) < 1 THEN 1 ELSE 0 END) DESC`
    : `ORDER BY vec_distance_cosine(embedding, vec_f32(?)) ASC`;
  const sql = `SELECT id, text, embedding, type, source, created_at, last_accessed_at, access_count, importance, (1 - vec_distance_cosine(embedding, vec_f32(?))) AS similarity FROM memories WHERE (1 - vec_distance_cosine(embedding, vec_f32(?))) >= ? ${typeFilter} ${orderClause} LIMIT ?`;
  const queryParams: unknown[] = [JSON.stringify(embedding), JSON.stringify(embedding), params.similarityThreshold, ...params.priorityTypes.length < 3 ? params.priorityTypes : [], JSON.stringify(embedding), params.topK];
  return { sql, queryParams };
}

export function processRetrievedMemories(rawResults: Array<{ record: any; cosineSimilarity: number }>, params: RetrievalParams, config: MemoryConfig = DEFAULT_MEMORY_CONFIG) {
  const scored = applyDecayWeighting(rawResults, config);
  let tokenCount = 0; const selected: ScoredMemory[] = [];
  for (const m of scored) { const t = Math.ceil(m.text.length / 4); if (tokenCount + t > params.maxTokenBudget) break; selected.push(m); tokenCount += t; }
  return { formattedContext: formatMemoriesByType(selected.map(m => ({ text: m.text, type: m.type, source: m.source }))), memoriesUsed: selected.map(m => m.id) };
}