Ship Assistent 0.8.1: split modules and shared+personal vector memory.
Personal RAG never leaks into the shared store; retrieve merges shared plus the persona chain, with personal overwrite on kind+key. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Leonid Pershin
Cursor
@@ -0,0 +1,385 @@
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.Linq;
|
||||
using System.Text;
|
||||
using System.Threading.Tasks;
|
||||
using Newtonsoft.Json.Linq;
|
||||
using SwarmUI.Accounts;
|
||||
using SwarmUI.Utils;
|
||||
|
||||
namespace Mrleo1nid.SwarmAssistent;
|
||||
|
||||
/// <summary>Prompt assembly, memory retrieval/writeback and the Civitai search hop loop.</summary>
|
||||
public partial class SwarmAssistentExtension
|
||||
{
|
||||
const int MaxCivitaiHopsFallback = 2;
|
||||
|
||||
List<JObject> BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null, string personaId = null, IEnumerable<string> skillIds = null)
|
||||
{
|
||||
List<JObject> ollamaMessages = [];
|
||||
StringBuilder system = new();
|
||||
string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
|
||||
|
||||
if (includeBase)
|
||||
{
|
||||
string core = Config.LoadCorePrompt(pid);
|
||||
if (!string.IsNullOrWhiteSpace(core))
|
||||
{
|
||||
system.AppendLine(core);
|
||||
}
|
||||
}
|
||||
|
||||
JObject exact = Config.LoadExactForPrompt(pid);
|
||||
if (exact is not null && exact.Count > 0)
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine("## Exact memory (canonical KV defaults — prefer over RAG for numbers)");
|
||||
system.AppendLine("```json");
|
||||
system.AppendLine(exact.ToString(Newtonsoft.Json.Formatting.None));
|
||||
system.AppendLine("```");
|
||||
}
|
||||
|
||||
foreach (string skillId in skillIds ?? Config.ResolveEnabledSkills(pid, null))
|
||||
{
|
||||
string skillText = Config.LoadSkillPrompt(pid, skillId);
|
||||
if (!string.IsNullOrWhiteSpace(skillText))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine($"## Skill: {skillId}");
|
||||
system.AppendLine(skillText);
|
||||
}
|
||||
}
|
||||
|
||||
string identity = Config.RenderIdentityBlock(pid);
|
||||
if (!string.IsNullOrWhiteSpace(identity))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine(identity);
|
||||
}
|
||||
|
||||
if (!string.IsNullOrWhiteSpace(packName))
|
||||
{
|
||||
string situational = Config.LoadPackPrompt(pid, packName);
|
||||
if (!string.IsNullOrWhiteSpace(situational))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine($"## Active mode: {packName}");
|
||||
system.AppendLine(situational);
|
||||
}
|
||||
}
|
||||
if (!string.IsNullOrWhiteSpace(contextJson))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine("## Live SwarmUI context (JSON — trust this over guesses)");
|
||||
system.AppendLine("```json");
|
||||
system.AppendLine(contextJson);
|
||||
system.AppendLine("```");
|
||||
}
|
||||
if (!string.IsNullOrWhiteSpace(extraSystem))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine(extraSystem);
|
||||
}
|
||||
if (system.Length > 0)
|
||||
{
|
||||
ollamaMessages.Add(new JObject
|
||||
{
|
||||
["role"] = "system",
|
||||
["content"] = system.ToString(),
|
||||
});
|
||||
}
|
||||
foreach (JToken msg in userMessages ?? [])
|
||||
{
|
||||
if (msg is not JObject mo)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
JObject copy = new()
|
||||
{
|
||||
["role"] = mo["role"]?.ToString() ?? "user",
|
||||
["content"] = mo["content"]?.ToString() ?? "",
|
||||
};
|
||||
if (mo["images"] is JArray images && images.Count > 0)
|
||||
{
|
||||
copy["images"] = images;
|
||||
}
|
||||
ollamaMessages.Add(copy);
|
||||
}
|
||||
return ollamaMessages;
|
||||
}
|
||||
|
||||
async Task<(string reply, JObject raw, JArray civitaiResults)> RunChatWithHops(
|
||||
Session session,
|
||||
string root,
|
||||
string modelName,
|
||||
string packName,
|
||||
bool includeBase,
|
||||
string contextJson,
|
||||
JArray userMessages,
|
||||
Func<string, Task> onDelta = null,
|
||||
Func<int, Task> onHopStart = null,
|
||||
string personaId = null,
|
||||
JArray skillIds = null,
|
||||
string embedModel = null)
|
||||
{
|
||||
string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
|
||||
List<string> skills = Config.ResolveEnabledSkills(pid, skillIds);
|
||||
string embed = string.IsNullOrWhiteSpace(embedModel)
|
||||
? (Config.LoadSettings()["embed_model"]?.ToString()
|
||||
?? Config.LoadAssistant(pid)["embed_model"]?.ToString()
|
||||
?? "nomic-embed-text")
|
||||
: embedModel;
|
||||
|
||||
try
|
||||
{
|
||||
await Memory.EnsureSeedAsync(root, Config, embed);
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
Logs.Debug($"Assistent memory seed: {ex.Message}");
|
||||
}
|
||||
|
||||
string retrieveQuery = BuildRetrieveQuery(userMessages, contextJson);
|
||||
JArray hits = [];
|
||||
try
|
||||
{
|
||||
int topK = Config.LoadAssistant(pid)["memory_top_k"]?.Value<int?>() ?? 10;
|
||||
hits = await Memory.RetrieveAsync(root, retrieveQuery, topK, embed, Config.PersonaExtendsChain(pid));
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
Logs.Debug($"Assistent memory retrieve: {ex.Message}");
|
||||
}
|
||||
|
||||
string enrichedContext = InjectMemoryHits(contextJson, hits);
|
||||
List<JObject> messages = BuildOllamaMessages(packName, includeBase, enrichedContext, userMessages, personaId: pid, skillIds: skills);
|
||||
JArray civitaiResults = [];
|
||||
string reply = "";
|
||||
JObject lastRaw = null;
|
||||
int maxHops = CfgInt("max_civitai_hops", MaxCivitaiHopsFallback);
|
||||
for (int hop = 0; hop < maxHops; hop++)
|
||||
{
|
||||
if (onHopStart is not null)
|
||||
{
|
||||
await onHopStart(hop);
|
||||
}
|
||||
(reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta, pid);
|
||||
JObject patch = TryParsePatch(reply);
|
||||
await ApplyMemoryActions(root, patch, embed, pid);
|
||||
if (hop + 1 >= maxHops || !WantsCivitaiSearch(patch))
|
||||
{
|
||||
break;
|
||||
}
|
||||
string query = ExtractSearchQuery(patch);
|
||||
if (string.IsNullOrWhiteSpace(query))
|
||||
{
|
||||
break;
|
||||
}
|
||||
JObject search = await AssistentSearchCivitai(session, query, 8);
|
||||
if (search["error"] is not null)
|
||||
{
|
||||
messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
|
||||
messages.Add(new JObject
|
||||
{
|
||||
["role"] = "user",
|
||||
["content"] = $"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.",
|
||||
});
|
||||
continue;
|
||||
}
|
||||
civitaiResults = search["results"] as JArray ?? [];
|
||||
messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
|
||||
messages.Add(new JObject
|
||||
{
|
||||
["role"] = "user",
|
||||
["content"] =
|
||||
"Civitai search results (JSON). Prefer `krea_likely: true`. Do NOT download yourself — the UI shows Confirm cards. " +
|
||||
"Pick useful LoRAs from results or available_loras, emit a normal patch (prompt/loras). " +
|
||||
"Omit search_civitai from actions unless you need a different query.\n```json\n" +
|
||||
civitaiResults.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
|
||||
});
|
||||
}
|
||||
return (reply, lastRaw, civitaiResults);
|
||||
}
|
||||
|
||||
static string BuildRetrieveQuery(JArray userMessages, string contextJson)
|
||||
{
|
||||
StringBuilder sb = new();
|
||||
if (!string.IsNullOrWhiteSpace(contextJson))
|
||||
{
|
||||
try
|
||||
{
|
||||
JObject ctx = JObject.Parse(contextJson);
|
||||
string ckpt = ctx["checkpoint"]?.ToString() ?? ctx["current_model"]?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(ckpt))
|
||||
{
|
||||
sb.Append(ckpt).Append(' ');
|
||||
}
|
||||
if (ctx["enabled_loras"] is JArray en)
|
||||
{
|
||||
foreach (JToken t in en.Take(8))
|
||||
{
|
||||
string n = t?["name"]?.ToString() ?? t?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(n))
|
||||
{
|
||||
sb.Append(n).Append(' ');
|
||||
}
|
||||
}
|
||||
}
|
||||
if (ctx["krea_profile"] != null)
|
||||
{
|
||||
sb.Append("krea ").Append(ctx["krea_profile"]).Append(' ');
|
||||
}
|
||||
}
|
||||
catch
|
||||
{
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
foreach (JToken msg in (userMessages ?? []).Reverse().Take(2))
|
||||
{
|
||||
if (msg is JObject mo && string.Equals(mo["role"]?.ToString(), "user", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
sb.Append(mo["content"]?.ToString()).Append(' ');
|
||||
}
|
||||
}
|
||||
string q = CollapseWs(sb.ToString());
|
||||
return string.IsNullOrWhiteSpace(q) ? "krea2 prompting" : q;
|
||||
}
|
||||
|
||||
static string InjectMemoryHits(string contextJson, JArray hits, JObject exact = null)
|
||||
{
|
||||
JObject ctx;
|
||||
try
|
||||
{
|
||||
ctx = string.IsNullOrWhiteSpace(contextJson) ? new JObject() : JObject.Parse(contextJson);
|
||||
}
|
||||
catch
|
||||
{
|
||||
ctx = new JObject { ["_raw_context"] = contextJson };
|
||||
}
|
||||
ctx["memory_hits"] = hits ?? new JArray();
|
||||
// Never re-inject full Exact into live context (already in system prompt).
|
||||
ctx.Remove("exact");
|
||||
if (ctx["session_exact"] is null)
|
||||
{
|
||||
ctx["session_exact"] = new JObject();
|
||||
}
|
||||
// Slim inventory for LLM: keep enabled + current, drop full dump if present
|
||||
if (ctx["available_loras"] is JArray allLoras && allLoras.Count > 24)
|
||||
{
|
||||
HashSet<string> keep = new(StringComparer.OrdinalIgnoreCase);
|
||||
if (ctx["enabled_loras"] is JArray en)
|
||||
{
|
||||
foreach (JToken t in en)
|
||||
{
|
||||
string n = t?["name"]?.ToString() ?? t?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(n))
|
||||
{
|
||||
keep.Add(n);
|
||||
}
|
||||
}
|
||||
}
|
||||
foreach (JToken hit in hits ?? [])
|
||||
{
|
||||
if (string.Equals(hit?["kind"]?.ToString(), "lora", StringComparison.OrdinalIgnoreCase)
|
||||
|| string.Equals(hit?["kind"]?.ToString(), "card", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
string k = hit?["key"]?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(k))
|
||||
{
|
||||
keep.Add(k);
|
||||
}
|
||||
}
|
||||
}
|
||||
JArray slim = [];
|
||||
foreach (JToken t in allLoras)
|
||||
{
|
||||
string n = t?["name"]?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(n) && (keep.Contains(n) || slim.Count < 12))
|
||||
{
|
||||
if (keep.Contains(n) || t?["krea_likely"]?.Value<bool>() == true)
|
||||
{
|
||||
slim.Add(t);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (slim.Count == 0)
|
||||
{
|
||||
foreach (JToken t in allLoras.Take(12))
|
||||
{
|
||||
slim.Add(t);
|
||||
}
|
||||
}
|
||||
ctx["available_loras"] = slim;
|
||||
ctx["available_loras_truncated"] = true;
|
||||
ctx["available_loras_total"] = allLoras.Count;
|
||||
}
|
||||
return ctx.ToString(Newtonsoft.Json.Formatting.None);
|
||||
}
|
||||
|
||||
static string MemoryWritePersona(JObject mo, string currentPersonaId)
|
||||
{
|
||||
string scope = (mo?["scope"]?.ToString() ?? "").Trim().ToLowerInvariant();
|
||||
if (scope is "shared" or "common" or "global")
|
||||
{
|
||||
return AssistentMemory.SharedPersona;
|
||||
}
|
||||
// Personal only — never let the model write into another personality's store.
|
||||
return AssistentConfig.SafeId(currentPersonaId) ?? AssistentMemory.SharedPersona;
|
||||
}
|
||||
|
||||
async Task ApplyMemoryActions(string root, JObject patch, string embedModel, string personaId)
|
||||
{
|
||||
if (patch is null || Memory is null)
|
||||
{
|
||||
return;
|
||||
}
|
||||
bool upsert = false, forget = false;
|
||||
if (patch["actions"] is JArray acts)
|
||||
{
|
||||
foreach (JToken a in acts)
|
||||
{
|
||||
string s = a?.ToString() ?? "";
|
||||
if (string.Equals(s, "memory_upsert", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
upsert = true;
|
||||
}
|
||||
if (string.Equals(s, "memory_forget", StringComparison.OrdinalIgnoreCase))
|
||||
{
|
||||
forget = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
JArray memories = patch["memories"] as JArray;
|
||||
if (memories is null || memories.Count == 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
foreach (JToken t in memories)
|
||||
{
|
||||
if (t is not JObject mo)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
string kind = mo["kind"]?.ToString() ?? "note";
|
||||
string key = mo["key"]?.ToString() ?? "";
|
||||
string text = mo["text"]?.ToString() ?? "";
|
||||
string target = MemoryWritePersona(mo, personaId);
|
||||
try
|
||||
{
|
||||
if (forget && string.IsNullOrWhiteSpace(text))
|
||||
{
|
||||
Memory.Forget(kind, key, persona: target);
|
||||
}
|
||||
else if (upsert || !string.IsNullOrWhiteSpace(text))
|
||||
{
|
||||
await Memory.UpsertTextAsync(root, kind, key, text, "user", mo, embedModel, target);
|
||||
}
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
Logs.Debug($"ApplyMemoryActions: {ex.Message}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user