Leonid PershinandCursor
880e2dbea2
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>
386 lines
14 KiB
C#
386 lines
14 KiB
C#
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}");
|
|
}
|
|
}
|
|
}
|
|
}
|