using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Microsoft.Data.Sqlite;
using Newtonsoft.Json.Linq;
using SwarmUI.Utils;
namespace Mrleo1nid.SwarmAssistent;
/// Local SQLite vector memory with Ollama /api/embed.
/// Two layers: shared (persona='') is visible to every personality; personal (persona=id)
/// is not written back to shared. On retrieve, personal overwrites shared on the same kind+key.
public sealed partial class AssistentMemory : IDisposable
{
public const string SharedPersona = "";
public sealed class RetrieveOptions
{
public int TopK { get; set; } = 10;
public float MinScore { get; set; } = 0.32f;
public string KindFilter { get; set; }
public IReadOnlyDictionary Quotas { get; set; }
public bool ApplyQuotas { get; set; } = true;
}
static readonly Dictionary DefaultQuotas = new(StringComparer.OrdinalIgnoreCase)
{
["card"] = 3,
["lora"] = 3,
["pitfall"] = 3,
["path"] = 2,
["note"] = 4,
["model"] = 2,
["aspect"] = 1,
["heard"] = 3,
};
public static Dictionary CopyDefaultQuotas()
=> new(DefaultQuotas, StringComparer.OrdinalIgnoreCase);
readonly string _dataRoot;
readonly string _dbPath;
readonly HttpClient _http;
readonly object _lock = new();
SqliteConnection _conn;
string _embedModel;
int _dims;
int _seedVersion;
public AssistentMemory(string dataRoot, HttpClient http, string defaultEmbedModel = "nomic-embed-text")
{
_dataRoot = string.IsNullOrWhiteSpace(dataRoot) ? "." : dataRoot;
string dir = Path.Combine(_dataRoot, "Assistent", "memory");
Directory.CreateDirectory(dir);
_dbPath = Path.Combine(dir, "assistent.sqlite");
_http = http;
_embedModel = string.IsNullOrWhiteSpace(defaultEmbedModel) ? "nomic-embed-text" : defaultEmbedModel.Trim();
}
public string EmbedModel => _embedModel;
public int Dims => _dims;
public int SeedVersion => _seedVersion;
public static string NormalizePersona(string raw)
{
if (string.IsNullOrWhiteSpace(raw))
{
return SharedPersona;
}
string s = raw.Trim();
if (s is "_" or "*" or "shared" or "common" or "global" or "_shared")
{
return SharedPersona;
}
return AssistentConfig.SafeId(s) ?? SharedPersona;
}
public static bool IsShared(string persona) => string.IsNullOrEmpty(NormalizePersona(persona));
void EnsureOpen()
{
if (_conn is not null)
{
return;
}
AssistentSqliteBootstrap.EnsureInitialized();
_conn = new SqliteConnection($"Data Source={_dbPath}");
_conn.Open();
TryPragma("journal_mode=WAL");
TryPragma("busy_timeout=5000");
using (SqliteCommand cmd = _conn.CreateCommand())
{
cmd.CommandText =
"""
CREATE TABLE IF NOT EXISTS meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS memories (
id INTEGER PRIMARY KEY AUTOINCREMENT,
kind TEXT NOT NULL,
key TEXT NOT NULL,
persona TEXT NOT NULL DEFAULT '',
text TEXT NOT NULL,
source TEXT NOT NULL DEFAULT 'user',
meta_json TEXT,
embedding BLOB,
updated INTEGER NOT NULL,
UNIQUE(kind, key, source, persona)
);
""";
cmd.ExecuteNonQuery();
}
MigratePersonaColumn();
using (SqliteCommand idx = _conn.CreateCommand())
{
idx.CommandText =
"""
CREATE INDEX IF NOT EXISTS idx_memories_kind ON memories(kind);
CREATE INDEX IF NOT EXISTS idx_memories_persona ON memories(persona);
""";
idx.ExecuteNonQuery();
}
EnsureFts();
try
{
EnsureTagsSchema();
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory tags schema: {ex.Message}");
}
try
{
EnsureStoreSchema();
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory store schema: {ex.Message}");
}
try
{
EnsureTrainingSchema();
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory training schema: {ex.Message}");
}
try
{
EnsureUserPrefsSchema();
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory user_prefs schema: {ex.Message}");
}
_embedModel = GetMeta("embed_model") ?? _embedModel;
_ = int.TryParse(GetMeta("dims"), out _dims);
_ = int.TryParse(GetMeta("seed_version"), out _seedVersion);
}
bool HasColumn(string table, string column)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = $"PRAGMA table_info({table})";
using SqliteDataReader reader = cmd.ExecuteReader();
while (reader.Read())
{
if (string.Equals(reader.GetString(1), column, StringComparison.OrdinalIgnoreCase))
{
return true;
}
}
return false;
}
void MigratePersonaColumn()
{
if (HasColumn("memories", "persona"))
{
return;
}
using SqliteTransaction tx = _conn.BeginTransaction();
using (SqliteCommand cmd = _conn.CreateCommand())
{
cmd.Transaction = tx;
cmd.CommandText =
"""
CREATE TABLE memories_v2 (
id INTEGER PRIMARY KEY AUTOINCREMENT,
kind TEXT NOT NULL,
key TEXT NOT NULL,
persona TEXT NOT NULL DEFAULT '',
text TEXT NOT NULL,
source TEXT NOT NULL DEFAULT 'user',
meta_json TEXT,
embedding BLOB,
updated INTEGER NOT NULL,
UNIQUE(kind, key, source, persona)
);
INSERT INTO memories_v2 (kind, key, persona, text, source, meta_json, embedding, updated)
SELECT kind, key, '', text, source, meta_json, embedding, updated FROM memories;
DROP TABLE memories;
ALTER TABLE memories_v2 RENAME TO memories;
CREATE INDEX IF NOT EXISTS idx_memories_kind ON memories(kind);
CREATE INDEX IF NOT EXISTS idx_memories_persona ON memories(persona);
""";
cmd.ExecuteNonQuery();
}
tx.Commit();
SetMeta("schema_version", "2");
Logs.Debug("AssistentMemory: migrated sqlite to shared+personal persona column (existing rows → shared)");
}
void TryPragma(string pragma)
{
try
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "PRAGMA " + pragma;
_ = cmd.ExecuteScalar();
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory PRAGMA {pragma}: {ex.Message}");
}
}
bool TableExists(string name)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT 1 FROM sqlite_master WHERE type IN ('table','view') AND name = $n LIMIT 1";
cmd.Parameters.AddWithValue("$n", name);
return cmd.ExecuteScalar() is not null;
}
void Exec(string sql)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = sql;
cmd.ExecuteNonQuery();
}
void EnsureFts()
{
try
{
Exec(
"""
CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5(
key,
text,
tokenize = 'unicode61 remove_diacritics 2'
);
""");
Exec(
"""
CREATE TRIGGER IF NOT EXISTS memories_fts_ai AFTER INSERT ON memories BEGIN
INSERT INTO memories_fts(rowid, key, text) VALUES (new.id, new.key, new.text);
END;
""");
Exec(
"""
CREATE TRIGGER IF NOT EXISTS memories_fts_ad AFTER DELETE ON memories BEGIN
INSERT INTO memories_fts(memories_fts, rowid) VALUES('delete', old.id);
END;
""");
Exec(
"""
CREATE TRIGGER IF NOT EXISTS memories_fts_au AFTER UPDATE ON memories BEGIN
INSERT INTO memories_fts(memories_fts, rowid) VALUES('delete', old.id);
INSERT INTO memories_fts(rowid, key, text) VALUES (new.id, new.key, new.text);
END;
""");
int mem = 0, fts = 0;
using (SqliteCommand c = _conn.CreateCommand())
{
c.CommandText = "SELECT COUNT(*) FROM memories";
mem = Convert.ToInt32(c.ExecuteScalar());
}
using (SqliteCommand c = _conn.CreateCommand())
{
c.CommandText = "SELECT COUNT(*) FROM memories_fts";
fts = Convert.ToInt32(c.ExecuteScalar());
}
if (mem != fts)
{
Exec("DELETE FROM memories_fts");
Exec("INSERT INTO memories_fts(rowid, key, text) SELECT id, key, text FROM memories");
}
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory FTS5 unavailable, cosine-only: {ex.Message}");
}
}
string GetMeta(string key)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT value FROM meta WHERE key = $k";
cmd.Parameters.AddWithValue("$k", key);
return cmd.ExecuteScalar()?.ToString();
}
void SetMeta(string key, string value)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "INSERT INTO meta(key, value) VALUES($k, $v) ON CONFLICT(key) DO UPDATE SET value = excluded.value";
cmd.Parameters.AddWithValue("$k", key);
cmd.Parameters.AddWithValue("$v", value ?? "");
cmd.ExecuteNonQuery();
}
static byte[] FloatsToBytes(float[] v)
{
byte[] bytes = new byte[v.Length * 4];
Buffer.BlockCopy(v, 0, bytes, 0, bytes.Length);
return bytes;
}
static float[] BytesToFloats(byte[] bytes)
{
if (bytes is null || bytes.Length < 4 || bytes.Length % 4 != 0)
{
return Array.Empty();
}
float[] v = new float[bytes.Length / 4];
Buffer.BlockCopy(bytes, 0, v, 0, bytes.Length);
return v;
}
static float Cosine(float[] a, float[] b)
{
if (a.Length == 0 || a.Length != b.Length)
{
return float.NegativeInfinity;
}
double dot = 0, na = 0, nb = 0;
for (int i = 0; i < a.Length; i++)
{
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
if (na <= 0 || nb <= 0)
{
return float.NegativeInfinity;
}
return (float)(dot / (Math.Sqrt(na) * Math.Sqrt(nb)));
}
static int SourceRank(string source)
{
return string.Equals(source, "user", StringComparison.OrdinalIgnoreCase) ? 1 : 0;
}
public async Task EmbedAsync(string baseUrl, string model, string text, string keepAlive = "60m")
{
string root = (baseUrl ?? "http://127.0.0.1:11434").TrimEnd('/');
string m = string.IsNullOrWhiteSpace(model) ? _embedModel : model.Trim();
JObject payload = new()
{
["model"] = m,
["input"] = text ?? "",
["keep_alive"] = keepAlive,
};
using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
using HttpResponseMessage resp = await _http.PostAsync($"{root}/api/embed", content);
string body = await resp.Content.ReadAsStringAsync();
if (!resp.IsSuccessStatusCode)
{
throw new Exception($"Ollama /api/embed HTTP {(int)resp.StatusCode}: {body[..Math.Min(body.Length, 200)]}");
}
JObject parsed = JObject.Parse(body);
JArray embeddings = parsed["embeddings"] as JArray;
JToken first = embeddings?.FirstOrDefault() ?? parsed["embedding"];
if (first is not JArray vec)
{
throw new Exception("Ollama /api/embed: no embeddings in response");
}
float[] floats = vec.Select(t => t.Value()).ToArray();
return floats;
}
bool BundledExists(string kind, string key, string persona)
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT 1 FROM memories WHERE source = 'bundled' AND kind = $kind AND key = $key AND persona = $persona LIMIT 1";
cmd.Parameters.AddWithValue("$kind", kind);
cmd.Parameters.AddWithValue("$key", key);
cmd.Parameters.AddWithValue("$persona", persona ?? SharedPersona);
return cmd.ExecuteScalar() is not null;
}
public async Task EnsureSeedAsync(string baseUrl, AssistentConfig config, string modelOverride = null)
{
lock (_lock)
{
EnsureOpen();
}
JObject assistant = config.LoadAssistant(config.DefaultPersonaId());
int wantVersion = assistant["seed_version"]?.Value() ?? 1;
string wantModel = string.IsNullOrWhiteSpace(modelOverride)
? (assistant["embed_model"]?.ToString() ?? _embedModel)
: modelOverride.Trim();
bool needReseed = _seedVersion != wantVersion || !string.Equals(_embedModel, wantModel, StringComparison.OrdinalIgnoreCase);
List docs = config.LoadMemorySeedDocs();
if (docs.Count == 0)
{
TryIndexTags();
TryIndexExamples();
return;
}
if (!needReseed)
{
bool missing;
lock (_lock)
{
EnsureOpen();
missing = docs.Any(d =>
{
string kind = (d["kind"]?.ToString() ?? "note").Trim().ToLowerInvariant();
string key = (d["key"]?.ToString() ?? "").Trim();
string persona = NormalizePersona(d["persona"]?.ToString());
return !string.IsNullOrWhiteSpace(key) && !BundledExists(kind, key, persona);
});
}
if (!missing)
{
TryIndexTags();
TryIndexExamples();
return;
}
}
float[] probe;
try
{
probe = await EmbedAsync(baseUrl, wantModel, docs[0]["text"]?.ToString() ?? "seed");
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory seed defer (embed unavailable): {ex.Message}");
TryIndexTags();
TryIndexExamples();
return;
}
lock (_lock)
{
EnsureOpen();
if (!string.Equals(_embedModel, wantModel, StringComparison.OrdinalIgnoreCase) || (_dims > 0 && _dims != probe.Length))
{
using SqliteCommand clear = _conn.CreateCommand();
clear.CommandText = "DELETE FROM memories";
clear.ExecuteNonQuery();
}
else if (needReseed)
{
using SqliteCommand clearBundled = _conn.CreateCommand();
clearBundled.CommandText = "DELETE FROM memories WHERE source = 'bundled'";
clearBundled.ExecuteNonQuery();
}
_embedModel = wantModel;
_dims = probe.Length;
_seedVersion = wantVersion;
SetMeta("embed_model", _embedModel);
SetMeta("dims", _dims.ToString());
SetMeta("seed_version", _seedVersion.ToString());
}
foreach (JObject doc in docs)
{
string kind = (doc["kind"]?.ToString() ?? "note").Trim();
string key = (doc["key"]?.ToString() ?? "").Trim();
string text = (doc["text"]?.ToString() ?? "").Trim();
string persona = NormalizePersona(doc["persona"]?.ToString());
if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
{
continue;
}
if (!needReseed)
{
lock (_lock)
{
EnsureOpen();
if (BundledExists(kind.ToLowerInvariant(), key, persona))
{
continue;
}
}
}
try
{
float[] vec = await EmbedAsync(baseUrl, wantModel, text);
Upsert(kind, key, text, "bundled", doc["tags"], vec, persona);
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory seed item {kind}/{key}: {ex.Message}");
}
}
TryIndexTags();
TryIndexExamples();
}
public void Upsert(string kind, string key, string text, string source, JToken meta, float[] embedding, string persona = null)
{
lock (_lock)
{
EnsureOpen();
kind = (kind ?? "note").Trim().ToLowerInvariant();
key = (key ?? "").Trim();
text = (text ?? "").Trim();
source = string.IsNullOrWhiteSpace(source) ? "user" : source.Trim();
persona = NormalizePersona(persona);
if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(text))
{
return;
}
if (embedding is { Length: > 0 })
{
if (_dims <= 0)
{
_dims = embedding.Length;
SetMeta("dims", _dims.ToString());
}
}
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText =
"""
INSERT INTO memories(kind, key, persona, text, source, meta_json, embedding, updated)
VALUES($kind, $key, $persona, $text, $source, $meta, $emb, $upd)
ON CONFLICT(kind, key, source, persona) DO UPDATE SET
text = excluded.text,
meta_json = excluded.meta_json,
embedding = excluded.embedding,
updated = excluded.updated
""";
cmd.Parameters.AddWithValue("$kind", kind);
cmd.Parameters.AddWithValue("$key", key);
cmd.Parameters.AddWithValue("$persona", persona);
cmd.Parameters.AddWithValue("$text", text);
cmd.Parameters.AddWithValue("$source", source);
cmd.Parameters.AddWithValue("$meta", meta?.ToString(Newtonsoft.Json.Formatting.None) ?? "");
cmd.Parameters.AddWithValue("$emb", embedding is null ? (object)DBNull.Value : FloatsToBytes(embedding));
cmd.Parameters.AddWithValue("$upd", DateTimeOffset.UtcNow.ToUnixTimeSeconds());
cmd.ExecuteNonQuery();
}
}
/// Delete a row in one layer. Default (persona=current) never touches shared.
/// Pass SharedPersona to forget a shared fact. Bundled rows are kept unless source is set.
public void Forget(string kind, string key, string source = null, string persona = null)
{
lock (_lock)
{
EnsureOpen();
persona = NormalizePersona(persona);
using SqliteCommand cmd = _conn.CreateCommand();
if (string.IsNullOrWhiteSpace(source))
{
cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND persona = $persona AND source != 'bundled'";
}
else
{
cmd.CommandText = "DELETE FROM memories WHERE kind = $kind AND key = $key AND persona = $persona AND source = $source";
cmd.Parameters.AddWithValue("$source", source);
}
cmd.Parameters.AddWithValue("$kind", (kind ?? "").Trim().ToLowerInvariant());
cmd.Parameters.AddWithValue("$key", (key ?? "").Trim());
cmd.Parameters.AddWithValue("$persona", persona);
cmd.ExecuteNonQuery();
}
}
static Dictionary PersonaRankMap(IEnumerable personaChain)
{
Dictionary rank = new(StringComparer.OrdinalIgnoreCase)
{
[SharedPersona] = 0,
};
int i = 1;
foreach (string id in personaChain ?? [])
{
string p = NormalizePersona(id);
if (p == SharedPersona)
{
continue;
}
rank[p] = i++;
}
return rank;
}
static string BuildFtsMatch(string query)
{
if (string.IsNullOrWhiteSpace(query))
{
return null;
}
HashSet tokens = new(StringComparer.OrdinalIgnoreCase);
foreach (Match m in Regex.Matches(query, @"[\p{L}\p{N}_-]{2,}"))
{
string t = m.Value.Replace("\"", "").Trim('_', '-');
if (t.Length >= 2)
{
tokens.Add(t);
}
string spaced = t.Replace('_', ' ').Replace('-', ' ');
if (!string.Equals(spaced, t, StringComparison.OrdinalIgnoreCase) && spaced.Length >= 2)
{
foreach (string p in spaced.Split(' ', StringSplitOptions.RemoveEmptyEntries))
{
if (p.Length >= 2)
{
tokens.Add(p);
}
}
}
if (tokens.Count >= 12)
{
break;
}
}
if (tokens.Count == 0)
{
return null;
}
return string.Join(" OR ", tokens.Select(t => $"\"{t}\""));
}
Dictionary FtsRowRanks(string match, int limit)
{
Dictionary ranks = [];
if (string.IsNullOrWhiteSpace(match) || !TableExists("memories_fts"))
{
return ranks;
}
try
{
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT rowid FROM memories_fts WHERE memories_fts MATCH $q ORDER BY rank LIMIT $lim";
cmd.Parameters.AddWithValue("$q", match);
cmd.Parameters.AddWithValue("$lim", limit);
using SqliteDataReader reader = cmd.ExecuteReader();
int r = 0;
while (reader.Read())
{
ranks[reader.GetInt64(0)] = r++;
}
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory FTS match: {ex.Message}");
}
return ranks;
}
static float KeyBoost(string query, string key)
{
if (string.IsNullOrWhiteSpace(query))
{
return 0;
}
string q = query.Replace('_', ' ');
string k = (key ?? "").Replace('_', ' ');
if (k.Length >= 3 && q.Contains(k, StringComparison.OrdinalIgnoreCase))
{
return 0.22f;
}
if (k.Length >= 3 && k.Contains(q.Trim(), StringComparison.OrdinalIgnoreCase) && q.Trim().Length >= 4)
{
return 0.12f;
}
return 0;
}
static JArray ApplyQuotas(IEnumerable ordered, int topK, IReadOnlyDictionary quotas)
{
Dictionary used = new(StringComparer.OrdinalIgnoreCase);
List picked = [];
List overflow = [];
foreach (JObject row in ordered)
{
string kind = row["kind"]?.ToString() ?? "note";
int cap = 2;
if (quotas is not null && quotas.TryGetValue(kind, out int q))
{
cap = q;
}
else if (DefaultQuotas.TryGetValue(kind, out int d))
{
cap = d;
}
used.TryGetValue(kind, out int n);
if (n < cap)
{
picked.Add(row);
used[kind] = n + 1;
}
else
{
overflow.Add(row);
}
if (picked.Count >= topK)
{
return new JArray(picked);
}
}
foreach (JObject row in overflow)
{
if (picked.Count >= topK)
{
break;
}
picked.Add(row);
}
return new JArray(picked);
}
/// Retrieve shared + the given persona chain. Hybrid FTS + cosine; personal overwrites shared on kind+key.
public async Task RetrieveAsync(string baseUrl, string query, int topK = 10, string modelOverride = null, IEnumerable personaChain = null, RetrieveOptions options = null)
{
options ??= new RetrieveOptions { TopK = topK };
if (options.TopK <= 0)
{
options.TopK = topK;
}
if (string.IsNullOrWhiteSpace(query))
{
return [];
}
lock (_lock)
{
EnsureOpen();
}
string model = string.IsNullOrWhiteSpace(modelOverride) ? _embedModel : modelOverride;
float[] qvec = null;
try
{
qvec = await EmbedAsync(baseUrl, model, query);
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory retrieve embed: {ex.Message}");
}
Dictionary rank = PersonaRankMap(personaChain);
string kindFilter = string.IsNullOrWhiteSpace(options.KindFilter) ? null : options.KindFilter.Trim().ToLowerInvariant();
string ftsMatch = BuildFtsMatch(query);
List<(float hybrid, float cosine, int personaRank, int sourceRank, bool fts, JObject row)> scored = [];
lock (_lock)
{
EnsureOpen();
Dictionary ftsRanks = FtsRowRanks(ftsMatch, Math.Max(40, options.TopK * 4));
using SqliteCommand cmd = _conn.CreateCommand();
List personaKeys = rank.Keys.ToList();
StringBuilder sql = new("SELECT id, kind, key, text, source, embedding, persona FROM memories WHERE persona IN (");
for (int i = 0; i < personaKeys.Count; i++)
{
if (i > 0)
{
sql.Append(',');
}
string pName = "$p" + i;
sql.Append(pName);
cmd.Parameters.AddWithValue(pName, personaKeys[i]);
}
sql.Append(')');
if (kindFilter is not null)
{
sql.Append(" AND kind = $kind");
cmd.Parameters.AddWithValue("$kind", kindFilter);
}
cmd.CommandText = sql.ToString();
using SqliteDataReader reader = cmd.ExecuteReader();
while (reader.Read())
{
long id = reader.GetInt64(0);
string kind = reader.GetString(1);
if (kindFilter is not null && !string.Equals(kind, kindFilter, StringComparison.OrdinalIgnoreCase))
{
continue;
}
string persona = reader.IsDBNull(6) ? SharedPersona : reader.GetString(6) ?? SharedPersona;
if (!rank.TryGetValue(persona, out int personaRank))
{
continue;
}
string key = reader.GetString(2);
string text = reader.GetString(3);
string source = reader.GetString(4);
float cosine = 0;
if (qvec is not null && !reader.IsDBNull(5))
{
cosine = Cosine(qvec, BytesToFloats((byte[])reader.GetValue(5)));
if (float.IsNegativeInfinity(cosine))
{
cosine = 0;
}
}
bool ftsHit = ftsRanks.TryGetValue(id, out int ftsRank);
float ftsBoost = ftsHit ? 0.28f * (1f - (ftsRank / 50f)) : 0;
float hybrid = cosine + ftsBoost + KeyBoost(query, key);
bool shared = persona == SharedPersona;
scored.Add((hybrid, cosine, personaRank, SourceRank(source), ftsHit, new JObject
{
["kind"] = kind,
["key"] = key,
["text"] = text,
["source"] = source,
["scope"] = shared ? "shared" : "personal",
["persona"] = shared ? "shared" : persona,
["score"] = Math.Round(hybrid, 4),
["cosine"] = Math.Round(cosine, 4),
["fts"] = ftsHit,
}));
}
}
Dictionary best = new(StringComparer.OrdinalIgnoreCase);
foreach (var item in scored)
{
string id = $"{item.row["kind"]}\n{item.row["key"]}";
if (best.TryGetValue(id, out var cur))
{
if (item.personaRank < cur.personaRank
|| (item.personaRank == cur.personaRank && item.sourceRank < cur.sourceRank)
|| (item.personaRank == cur.personaRank && item.sourceRank == cur.sourceRank && item.hybrid <= cur.hybrid))
{
continue;
}
}
best[id] = item;
}
float min = options.MinScore;
IEnumerable ordered = best.Values
.Where(s => s.fts || s.hybrid >= min || s.cosine >= min)
.OrderByDescending(s => s.hybrid)
.Select(s => s.row);
int k = Math.Clamp(options.TopK, 1, 30);
if (options.ApplyQuotas)
{
return ApplyQuotas(ordered, k, options.Quotas ?? DefaultQuotas);
}
return new JArray(ordered.Take(k));
}
/// Exact kind+key read with the same personal-over-shared overlay as retrieve.
public JObject Get(string kind, string key, IEnumerable personaChain = null)
{
kind = (kind ?? "note").Trim().ToLowerInvariant();
key = (key ?? "").Trim();
if (string.IsNullOrWhiteSpace(key))
{
return null;
}
Dictionary rank = PersonaRankMap(personaChain);
JObject best = null;
int bestPersona = -1, bestSource = -1;
lock (_lock)
{
EnsureOpen();
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT kind, key, text, source, persona, updated, meta_json FROM memories WHERE kind = $kind AND key = $key";
cmd.Parameters.AddWithValue("$kind", kind);
cmd.Parameters.AddWithValue("$key", key);
using SqliteDataReader reader = cmd.ExecuteReader();
while (reader.Read())
{
string persona = reader.IsDBNull(4) ? SharedPersona : reader.GetString(4) ?? SharedPersona;
if (!rank.TryGetValue(persona, out int personaRank))
{
continue;
}
int src = SourceRank(reader.GetString(3));
if (best is not null && (personaRank < bestPersona || (personaRank == bestPersona && src <= bestSource)))
{
continue;
}
bestPersona = personaRank;
bestSource = src;
bool shared = persona == SharedPersona;
best = new JObject
{
["kind"] = reader.GetString(0),
["key"] = reader.GetString(1),
["text"] = reader.GetString(2),
["source"] = reader.GetString(3),
["scope"] = shared ? "shared" : "personal",
["persona"] = shared ? "shared" : persona,
["updated"] = reader.IsDBNull(5) ? 0 : reader.GetInt64(5),
["meta_json"] = reader.IsDBNull(6) ? null : reader.GetString(6),
};
}
}
return best;
}
public async Task SearchAsync(string baseUrl, string query, string kind, int topK, string modelOverride, IEnumerable personaChain)
{
return await RetrieveAsync(baseUrl, query, topK, modelOverride, personaChain, new RetrieveOptions
{
TopK = Math.Clamp(topK, 1, 30),
MinScore = 0.18f,
KindFilter = kind,
ApplyQuotas = false,
});
}
public JArray ListAll(int limit = 200)
{
lock (_lock)
{
EnsureOpen();
List rows = [];
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT kind, key, text, source, persona, updated FROM memories ORDER BY updated DESC LIMIT $lim";
cmd.Parameters.AddWithValue("$lim", Math.Clamp(limit, 1, 2000));
using SqliteDataReader reader = cmd.ExecuteReader();
while (reader.Read())
{
string persona = reader.IsDBNull(4) ? SharedPersona : reader.GetString(4) ?? SharedPersona;
bool shared = string.IsNullOrEmpty(persona);
rows.Add(new JObject
{
["kind"] = reader.GetString(0),
["key"] = reader.GetString(1),
["text"] = reader.GetString(2),
["source"] = reader.GetString(3),
["scope"] = shared ? "shared" : "personal",
["persona"] = shared ? "shared" : persona,
["updated"] = reader.IsDBNull(5) ? 0 : reader.GetInt64(5),
});
}
return new JArray(rows);
}
}
public int CountAll()
{
lock (_lock)
{
EnsureOpen();
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT COUNT(*) FROM memories";
return Convert.ToInt32(cmd.ExecuteScalar());
}
}
public async Task UpsertTextAsync(string baseUrl, string kind, string key, string text, string source = "user", JToken meta = null, string modelOverride = null, string persona = null)
{
string model = string.IsNullOrWhiteSpace(modelOverride) ? _embedModel : modelOverride;
float[] vec = await EmbedAsync(baseUrl, model, text);
Upsert(kind, key, text, source, meta, vec, persona);
}
public async Task ReembedAllAsync(string baseUrl, string newModel)
{
List<(string kind, string key, string text, string source, string meta, string persona)> rows = [];
lock (_lock)
{
EnsureOpen();
using SqliteCommand cmd = _conn.CreateCommand();
cmd.CommandText = "SELECT kind, key, text, source, meta_json, persona FROM memories";
using SqliteDataReader reader = cmd.ExecuteReader();
while (reader.Read())
{
rows.Add((
reader.GetString(0),
reader.GetString(1),
reader.GetString(2),
reader.GetString(3),
reader.IsDBNull(4) ? "" : reader.GetString(4),
reader.FieldCount > 5 && !reader.IsDBNull(5) ? reader.GetString(5) : SharedPersona));
}
}
if (rows.Count == 0)
{
_embedModel = newModel;
lock (_lock)
{
EnsureOpen();
SetMeta("embed_model", _embedModel);
}
return;
}
float[] first = await EmbedAsync(baseUrl, newModel, rows[0].text);
lock (_lock)
{
EnsureOpen();
_embedModel = newModel;
_dims = first.Length;
SetMeta("embed_model", _embedModel);
SetMeta("dims", _dims.ToString());
}
foreach (var row in rows)
{
try
{
float[] vec = await EmbedAsync(baseUrl, newModel, row.text);
JToken meta = null;
if (!string.IsNullOrWhiteSpace(row.meta))
{
try { meta = JToken.Parse(row.meta); } catch { /* ignore */ }
}
Upsert(row.kind, row.key, row.text, row.source, meta, vec, row.persona);
}
catch (Exception ex)
{
Logs.Debug($"AssistentMemory reembed {row.kind}/{row.key}: {ex.Message}");
}
}
}
public void Dispose()
{
lock (_lock)
{
_conn?.Dispose();
_conn = null;
}
}
}