Files
swarm-assistent/AssistentChatPipeline.cs
T
Leonid PershinandCursor cf89348f85 Ship persona shelves, Exact controls, and overlay author pipeline.
Add Leonid as a shelf-based example with preference_bias slider; support /persona new clone-to-overlay and UI-only delete.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-22 01:48:02 +03:00

628 lines
23 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;
const int MaxToolHopsFallback = 4;
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, packName);
JArray hits = [];
try
{
AssistentMemory.RetrieveOptions opt = MemoryRetrieveOptions(pid);
hits = await Memory.RetrieveAsync(root, retrieveQuery, opt.TopK, embed, Config.PersonaExtendsChain(pid), opt);
}
catch (Exception ex)
{
Logs.Debug($"Assistent memory retrieve: {ex.Message}");
}
string enrichedContext = InjectMemoryHits(contextJson, hits);
enrichedContext = EnrichPersonaContext(enrichedContext, pid, packName);
List<JObject> messages = BuildOllamaMessages(packName, includeBase, enrichedContext, userMessages, personaId: pid, skillIds: skills);
JArray civitaiResults = [];
string reply = "";
JObject lastRaw = null;
int maxHops = Math.Max(CfgInt("max_civitai_hops", MaxCivitaiHopsFallback), CfgInt("max_tool_hops", MaxToolHopsFallback));
HashSet<string> hopDone = new(StringComparer.OrdinalIgnoreCase);
var chain = Config.PersonaExtendsChain(pid);
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);
ApplyPersonaActions(patch, ref pid);
if (hop + 1 >= maxHops)
{
break;
}
string tool = NextToolHop(patch);
if (string.IsNullOrWhiteSpace(tool))
{
break;
}
(string follow, JArray civitaiHop) = await RunToolHop(session, root, embed, pid, chain, patch, tool, hopDone);
if (follow is null)
{
break;
}
if (civitaiHop is { Count: > 0 })
{
civitaiResults = civitaiHop;
}
messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
messages.Add(new JObject { ["role"] = "user", ["content"] = follow });
}
return (reply, lastRaw, civitaiResults);
}
static string BuildRetrieveQuery(JArray userMessages, string contextJson, string packName = null)
{
StringBuilder sb = new();
if (!string.IsNullOrWhiteSpace(packName))
{
sb.Append(packName).Append(' ');
}
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(12))
{
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(' ');
}
string aspect = ctx["aspect"]?.ToString();
if (!string.IsNullOrWhiteSpace(aspect))
{
sb.Append(aspect).Append(' ');
}
string prompt = ctx["prompt"]?.ToString();
if (!string.IsNullOrWhiteSpace(prompt))
{
sb.Append(prompt.Length > 400 ? prompt[..400] : prompt).Append(' ');
}
}
catch
{
// ignore
}
}
foreach (JToken msg in (userMessages ?? []).Reverse().Take(3))
{
if (msg is JObject mo && string.Equals(mo["role"]?.ToString(), "user", StringComparison.OrdinalIgnoreCase))
{
string c = mo["content"]?.ToString() ?? "";
sb.Append(c.Length > 500 ? c[..500] : c).Append(' ');
}
}
string q = CollapseWs(sb.ToString());
return string.IsNullOrWhiteSpace(q) ? "krea2 prompting" : q;
}
AssistentMemory.RetrieveOptions MemoryRetrieveOptions(string pid)
{
JObject a = Config.LoadAssistant(pid) ?? new JObject();
AssistentMemory.RetrieveOptions opt = new()
{
TopK = a["memory_top_k"]?.Value<int?>() ?? 10,
MinScore = a["memory_min_score"]?.Value<float?>() ?? 0.32f,
ApplyQuotas = true,
};
if (a["memory_quotas"] is JObject quotas)
{
Dictionary<string, int> d = new(StringComparer.OrdinalIgnoreCase);
foreach (JProperty p in quotas.Properties())
{
d[p.Name] = p.Value?.Value<int?>() ?? 2;
}
opt.Quotas = d;
}
return opt;
}
async Task<(string follow, JArray civitai)> RunToolHop(
Session session,
string root,
string embed,
string pid,
IEnumerable<string> chain,
JObject patch,
string tool,
HashSet<string> hopDone)
{
if (tool == "memory_get")
{
JArray got = [];
foreach (JToken t in patch["memories"] as JArray ?? [])
{
if (t is not JObject mo)
{
continue;
}
string kind = mo["kind"]?.ToString() ?? "note";
string key = mo["key"]?.ToString() ?? "";
if (string.IsNullOrWhiteSpace(key))
{
continue;
}
string sig = $"get:{kind}:{key}";
if (!hopDone.Add(sig))
{
continue;
}
JObject row = Memory.Get(kind, key, chain);
got.Add(row ?? new JObject { ["kind"] = kind, ["key"] = key, ["missing"] = true });
}
if (got.Count == 0)
{
return (null, null);
}
return (
"memory_get results (JSON). Use these facts; omit memory_get unless you need a different key.\n```json\n"
+ got.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
null);
}
if (tool == "memory_search")
{
string q = ExtractMemoryQuery(patch);
if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("search:" + q))
{
return (null, null);
}
string kind = patch["memory_kind"]?.ToString();
int topK = Config.LoadAssistant(pid)["memory_top_k"]?.Value<int?>() ?? 10;
JArray rows = await Memory.SearchAsync(root, q, kind, topK, embed, chain);
return (
"memory_search results (JSON, hybrid FTS+vector). Omit memory_search unless you need a different query.\n```json\n"
+ rows.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
null);
}
if (tool == "lookup_tags")
{
string q = ExtractTagQuery(patch);
if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("tags:" + q))
{
return (null, null);
}
int lim = Config.LoadAssistant(pid)["tag_lookup_limit"]?.Value<int?>() ?? 20;
JArray tags = Memory.LookupTags(q, lim);
return (
"lookup_tags results from Danbooru csv (canonical name, aliases, post_count). Krea prompts stay natural prose — use this to check spelling/aliases, do not dump tag soup.\n```json\n"
+ tags.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
null);
}
if (tool == "civitai")
{
string query = ExtractSearchQuery(patch);
if (string.IsNullOrWhiteSpace(query) || !hopDone.Add("civitai:" + query))
{
return (null, null);
}
JObject search = await AssistentSearchCivitai(session, query, 8);
if (search["error"] is not null)
{
return ($"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.", null);
}
JArray civitaiResults = search["results"] as JArray ?? [];
return (
"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```",
civitaiResults);
}
return (null, null);
}
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;
}
string EnrichPersonaContext(string contextJson, string personaId, string packName)
{
JObject ctx;
try
{
ctx = string.IsNullOrWhiteSpace(contextJson) ? new JObject() : JObject.Parse(contextJson);
}
catch
{
ctx = new JObject { ["_raw_context"] = contextJson };
}
string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
ctx["persona_source"] = Config.PersonaSource(pid);
JObject schema = Config.LoadControlsSchema(pid);
JObject values = Config.LoadControlValues(pid);
if (schema.Properties().Any())
{
ctx["persona_controls"] = new JObject
{
["schema"] = schema,
["values"] = values,
};
}
JArray catalog = [];
foreach (var p in Config.ListPersonaCatalog())
{
catalog.Add(new JObject
{
["id"] = p.id,
["title"] = p.title,
["source"] = p.source,
});
}
ctx["personas"] = catalog;
if (string.Equals(packName, "author_persona", StringComparison.OrdinalIgnoreCase)
|| string.Equals(packName, "persona", StringComparison.OrdinalIgnoreCase))
{
JObject shelves = Config.LoadIdentityParts(pid);
shelves.Remove("extra");
ctx["persona_shelves"] = shelves;
ctx["persona_controls_schema"] = schema;
}
return ctx.ToString(Newtonsoft.Json.Formatting.None);
}
/// <summary>Apply overlay persona clone/write from patch. Ignores persona_delete. Updates pid ref after switch.</summary>
void ApplyPersonaActions(JObject patch, ref string personaId)
{
if (patch is null || Config is null)
{
return;
}
// Never honor delete from the model.
bool wantClone = false, wantWrite = false;
if (patch["actions"] is JArray acts)
{
foreach (JToken a in acts)
{
string s = a?.ToString() ?? "";
if (string.Equals(s, "persona_clone", StringComparison.OrdinalIgnoreCase))
{
wantClone = true;
}
if (string.Equals(s, "persona_write", StringComparison.OrdinalIgnoreCase))
{
wantWrite = true;
}
}
}
if (patch["persona_clone"] is JObject)
{
wantClone = true;
}
if (patch["persona_shelves"] is JObject)
{
wantWrite = true;
}
try
{
if (wantClone && patch["persona_clone"] is JObject clone)
{
string from = AssistentConfig.SafeId(clone["from"]?.ToString()) ?? personaId;
string to = AssistentConfig.SafeId(clone["to"]?.ToString());
string title = clone["title"]?.ToString();
bool overwrite = clone["overwrite"]?.Value<bool?>() == true;
if (to is not null)
{
Config.ClonePersonaToOverlay(from, to, title, overwrite);
personaId = to;
patch["_persona_cloned"] = to;
}
}
if (wantWrite && patch["persona_shelves"] is JObject shelves)
{
string target = AssistentConfig.SafeId(patch["persona"]?.ToString())
?? AssistentConfig.SafeId(patch["persona_clone"]?["to"]?.ToString())
?? personaId;
if (target is not null)
{
Config.SavePersonaShelves(target, shelves);
patch["_persona_written"] = target;
}
}
// Control values from model patch (Exact).
if (patch["controls"] is JObject ctrlVals)
{
string ctrlPid = AssistentConfig.SafeId(patch["persona"]?.ToString()) ?? personaId;
Config.SaveControlValues(ctrlPid, ctrlVals);
patch["_controls_saved"] = true;
}
}
catch (Exception ex)
{
Logs.Warning($"Assistent persona actions: {ex.Message}");
patch["_persona_error"] = ex.Message;
}
}
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}");
}
}
}
}