using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Net.WebSockets;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using FreneticUtilities.FreneticExtensions;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Text2Image;
using SwarmUI.Utils;
using SwarmUI.WebAPI;
namespace Mrleo1nid.SwarmAssistent;
/// Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches + Generate/Civitai.
public class SwarmAssistentExtension : Extension
{
public static PermInfo PermUse = Permissions.Register(new(
"swarm_assistent_use",
"[Swarm Assistent] Use",
"Allows using the Swarm Assistent chat (Ollama proxy).",
PermissionDefault.USER,
Permissions.GroupUser));
public static HttpClient HttpClient;
public AssistentConfig Config;
public AssistentMemory Memory;
const int MaxCivitaiHopsFallback = 2;
const int MaxLorasInInventoryFallback = 150;
const int MaxWildcardsInInventoryFallback = 80;
const int MaxCheckpointsInInventoryFallback = 60;
const int InventoryBlurbMaxFallback = 140;
const int DefaultNumCtxFallback = 16384;
static readonly Regex JsonFenceRe = new(@"```(?:json)?\s*([\s\S]*?)```", RegexOptions.IgnoreCase | RegexOptions.Compiled);
public override void OnPreInit()
{
ScriptFiles.Add("Assets/assistent.js");
StyleSheetFiles.Add("Assets/assistent.css");
ExtensionAuthor = "mrleo1nid";
Description = "Collaborative Krea 2 assistant: Ollama chat, persona presets, vector memory, model cards, Generate loop.";
License = "MIT";
Version = "0.7.3";
Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint", "memory"];
}
public override void OnInit()
{
HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
Config = new AssistentConfig(FilePath, DataRoot());
Memory = new AssistentMemory(DataRoot(), HttpClient, Config.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString() ?? "nomic-embed-text");
API.RegisterAPICall(AssistentListModels, false, PermUse);
API.RegisterAPICall(AssistentGetPacks, false, PermUse);
API.RegisterAPICall(AssistentListPersonas, false, PermUse);
API.RegisterAPICall(AssistentGetConfig, false, PermUse);
API.RegisterAPICall(AssistentGetSettings, false, PermUse);
API.RegisterAPICall(AssistentSaveSettings, true, PermUse);
API.RegisterAPICall(AssistentListInventory, false, PermUse);
API.RegisterAPICall(AssistentGetCard, false, PermUse);
API.RegisterAPICall(AssistentSaveCard, true, PermUse);
API.RegisterAPICall(AssistentEnqueueWanted, true, PermUse);
API.RegisterAPICall(AssistentGetCardMeta, false, PermUse);
API.RegisterAPICall(AssistentSearchCivitai, false, PermUse);
API.RegisterAPICall(AssistentGetTaste, false, PermUse);
API.RegisterAPICall(AssistentSaveTaste, true, PermUse);
API.RegisterAPICall(AssistentChat, true, PermUse);
API.RegisterAPICall(AssistentChatWS, true, PermUse);
Logs.Init("Swarm Assistent extension loaded (Config presets + vector memory)");
}
int CfgInt(string key, int fallback)
{
try
{
return Config?.LoadAssistant(Config.DefaultPersonaId())[key]?.Value() ?? fallback;
}
catch
{
return fallback;
}
}
static string Clip(string text, int max)
{
if (string.IsNullOrEmpty(text) || text.Length <= max)
{
return text ?? "";
}
return text[..max] + "…";
}
public static string NormalizeBaseUrl(string raw)
{
string url = (raw ?? "").Trim();
if (string.IsNullOrWhiteSpace(url))
{
url = "http://127.0.0.1:11434";
}
return url.TrimEnd('/');
}
public string ReadPackFile(string name)
{
return Config?.LoadPackPrompt(Config.DefaultPersonaId(), name);
}
public async Task AssistentListModels(Session session, string baseUrl)
{
string root = NormalizeBaseUrl(baseUrl);
try
{
using HttpResponseMessage resp = await HttpClient.GetAsync($"{root}/api/tags");
string body = await resp.Content.ReadAsStringAsync();
if (!resp.IsSuccessStatusCode)
{
return new JObject { ["error"] = $"Ollama /api/tags HTTP {(int)resp.StatusCode}: {Clip(body, 400)}" };
}
JObject parsed = JObject.Parse(body);
JArray all = [];
foreach (JToken m in parsed["models"] as JArray ?? [])
{
string name = m["name"]?.ToString() ?? m["model"]?.ToString() ?? "";
if (!string.IsNullOrWhiteSpace(name))
{
all.Add(name);
}
}
JObject roles = Config?.LoadOllamaRoles() ?? new JObject();
HashSet chatSet = new(
(roles["chat"] as JArray)?.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)) ?? [],
StringComparer.OrdinalIgnoreCase);
HashSet memSet = new(
(roles["memory"] as JArray)?.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)) ?? [],
StringComparer.OrdinalIgnoreCase);
// Heuristic fallbacks when sidecar missing
if (chatSet.Count == 0 && memSet.Count == 0)
{
foreach (JToken t in all)
{
string n = t.ToString();
if (LooksLikeEmbedModel(n))
{
memSet.Add(n);
}
else
{
chatSet.Add(n);
}
}
}
else
{
// Keep only tags that exist; anything unlabeled goes to chat if not memory
foreach (JToken t in all)
{
string n = t.ToString();
if (memSet.Contains(n) || LooksLikeEmbedModel(n))
{
memSet.Add(n);
chatSet.Remove(n);
}
else if (chatSet.Count == 0 || chatSet.Contains(n))
{
chatSet.Add(n);
}
else if (!memSet.Contains(n))
{
chatSet.Add(n);
}
}
}
JArray models = new(all.Select(t => t.ToString()).Where(n => chatSet.Contains(n) && !memSet.Contains(n) && !LooksLikeEmbedModel(n)));
JArray memoryModels = new(all.Select(t => t.ToString()).Where(n => memSet.Contains(n) || LooksLikeEmbedModel(n)).Distinct(StringComparer.OrdinalIgnoreCase).ToList());
if (memoryModels.Count == 0)
{
string fallback = Config?.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString() ?? "nomic-embed-text";
if (all.Any(t => string.Equals(t.ToString(), fallback, StringComparison.OrdinalIgnoreCase)
|| t.ToString().StartsWith(fallback.Split(':')[0], StringComparison.OrdinalIgnoreCase)))
{
memoryModels.Add(all.Select(t => t.ToString()).First(n =>
string.Equals(n, fallback, StringComparison.OrdinalIgnoreCase)
|| n.StartsWith(fallback.Split(':')[0], StringComparison.OrdinalIgnoreCase)));
}
}
return new JObject
{
["success"] = true,
["base_url"] = root,
["models"] = models,
["memory_models"] = memoryModels,
};
}
catch (Exception ex)
{
return new JObject { ["error"] = $"Ollama unreachable at {root}: {ex.Message}" };
}
}
static bool LooksLikeEmbedModel(string name)
{
string n = (name ?? "").ToLowerInvariant();
return n.Contains("embed") || n.Contains("nomic") || n.Contains("bge-") || n.Contains("minilm") || n.Contains("e5-");
}
public async Task AssistentGetPacks(Session session, string persona = null)
{
await Task.CompletedTask;
string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
JObject packs = new();
JArray order = [];
foreach (var p in Config.ListPacks(pid))
{
string text = Config.LoadPackPrompt(pid, p.id);
if (text is not null)
{
packs[p.id] = text;
}
order.Add(p.id);
}
return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = order, ["persona"] = pid };
}
public async Task AssistentGetConfig(Session session, string persona = null)
{
await Task.CompletedTask;
string pid = AssistentConfig.SafeId(persona) ?? Config.DefaultPersonaId();
return Config.BuildMergedConfigPayload(pid);
}
public async Task AssistentGetSettings(Session session)
{
await Task.CompletedTask;
return new JObject { ["success"] = true, ["settings"] = Config.LoadSettings() };
}
public async Task AssistentSaveSettings(Session session, JObject settings)
{
await Task.CompletedTask;
if (settings is null)
{
return new JObject { ["error"] = "settings required" };
}
string prevEmbed = Config.LoadSettings()["embed_model"]?.ToString();
Config.SaveSettings(settings);
string nextEmbed = settings["embed_model"]?.ToString();
if (!string.IsNullOrWhiteSpace(nextEmbed) && !string.Equals(prevEmbed, nextEmbed, StringComparison.OrdinalIgnoreCase))
{
try
{
await Memory.ReembedAllAsync(NormalizeBaseUrl(settings["base_url"]?.ToString()), nextEmbed);
}
catch (Exception ex)
{
Logs.Debug($"AssistentSaveSettings reembed: {ex.Message}");
}
}
return new JObject { ["success"] = true, ["path"] = Path.Combine(Config.OverlayRoot, "settings.json") };
}
static string DataRoot()
{
if (Directory.Exists("/mnt/swarm_data"))
{
return "/mnt/swarm_data";
}
try
{
string models = Path.GetFullPath(Path.Combine(Environment.CurrentDirectory, "Models"));
if (Directory.Exists(models))
{
return Path.GetDirectoryName(models) ?? Environment.CurrentDirectory;
}
}
catch
{
// ignore
}
return Environment.CurrentDirectory;
}
string PersonasOverlayJsonPath() => Path.Combine(DataRoot(), "Assistent", "personas.json");
string WantedModelsPath() => Path.Combine(DataRoot(), ".gpu-rent-wanted-models.yaml");
string WantedCardsDir() => Path.Combine(DataRoot(), ".gpu-rent-wanted-cards");
public async Task AssistentListPersonas(Session session)
{
await Task.CompletedTask;
var catalog = Config.ListPersonaCatalog();
JArray list = [];
foreach (var p in catalog)
{
list.Add(new JObject
{
["id"] = p.id,
["title"] = p.title,
["accent"] = p.accent,
["prompt"] = Config.RenderIdentityBlock(p.id),
["source"] = p.source,
});
}
return new JObject
{
["success"] = true,
["default"] = Config.DefaultPersonaId(),
["personas"] = list,
};
}
static string ModelWeightPath(string setName, string modelName)
{
if (string.IsNullOrWhiteSpace(modelName) || !Program.T2IModelSets.TryGetValue(setName, out T2IModelHandler handler))
{
return null;
}
if (!handler.Models.TryGetValue(modelName, out T2IModel model) && !handler.Models.TryGetValue(modelName.Replace('\\', '/'), out model))
{
// Try suffix match
model = handler.Models.Values.FirstOrDefault(m =>
string.Equals(m.Name, modelName, StringComparison.OrdinalIgnoreCase)
|| m.Name.EndsWith("/" + modelName, StringComparison.OrdinalIgnoreCase)
|| Path.GetFileNameWithoutExtension(m.Name) == Path.GetFileNameWithoutExtension(modelName));
}
if (model is null)
{
return null;
}
try
{
// SwarmUI T2IModel exposes RawFilePath in recent builds.
return model.RawFilePath;
}
catch
{
return null;
}
}
static string CardPathForWeight(string weightPath)
{
if (string.IsNullOrWhiteSpace(weightPath))
{
return null;
}
string dir = Path.GetDirectoryName(weightPath);
string stem = Path.GetFileNameWithoutExtension(weightPath);
if (string.IsNullOrWhiteSpace(dir) || string.IsNullOrWhiteSpace(stem))
{
return null;
}
return Path.Combine(dir, $"{stem}.assistent.json");
}
static string SetNameForKind(string kind)
{
return (kind ?? "").Trim().ToLowerInvariant() switch
{
"lora" => "LoRA",
"checkpoint" or "ckpt" or "stable-diffusion" => "Stable-Diffusion",
_ => null,
};
}
JObject ReadCardObject(string kind, string name)
{
string set = SetNameForKind(kind);
string weight = ModelWeightPath(set, name);
string card = CardPathForWeight(weight);
if (card is null || !File.Exists(card))
{
return null;
}
try
{
return JObject.Parse(File.ReadAllText(card, Encoding.UTF8));
}
catch
{
return null;
}
}
public async Task AssistentGetCard(Session session, string kind, string name)
{
await Task.CompletedTask;
if (string.IsNullOrWhiteSpace(kind) || string.IsNullOrWhiteSpace(name))
{
return new JObject { ["error"] = "kind and name required" };
}
JObject card = ReadCardObject(kind, name);
string set = SetNameForKind(kind);
string weight = ModelWeightPath(set, name);
return new JObject
{
["success"] = true,
["kind"] = kind,
["name"] = name,
["has_card"] = card is not null,
["weight_path"] = weight,
["card"] = card,
};
}
public async Task AssistentSaveCard(Session session, string kind, string name, JObject card, bool enqueue_wanted = false)
{
await Task.CompletedTask;
if (card is null)
{
return new JObject { ["error"] = "card required" };
}
kind = (kind ?? card["kind"]?.ToString() ?? "").Trim();
name = (name ?? card["name"]?.ToString() ?? "").Trim();
if (string.IsNullOrWhiteSpace(kind) || string.IsNullOrWhiteSpace(name))
{
return new JObject { ["error"] = "kind and name required" };
}
card["kind"] = kind;
card["name"] = name;
string set = SetNameForKind(kind);
string weight = ModelWeightPath(set, name);
if (!string.IsNullOrWhiteSpace(weight) && File.Exists(weight))
{
string path = CardPathForWeight(weight);
File.WriteAllText(path, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
_ = IngestCardToMemory(card, name);
return new JObject { ["success"] = true, ["path"] = path, ["installed"] = true };
}
// Not installed — draft into wanted-cards + optionally enqueue download for next up.
Directory.CreateDirectory(WantedCardsDir());
string rawVid = card["version_id"]?.ToString() ?? "draft";
string vid = Regex.IsMatch(rawVid, @"^\d+$") ? rawVid : "draft";
string draft = Path.Combine(WantedCardsDir(), $"{vid}.assistent.json");
File.WriteAllText(draft, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
if (enqueue_wanted || !string.IsNullOrWhiteSpace(card["civitai_url"]?.ToString()))
{
await AssistentEnqueueWanted(session, kind, card["civitai_url"]?.ToString(), card["version_id"]?.Value() ?? 0, card["title"]?.ToString() ?? name, card);
}
_ = IngestCardToMemory(card, name);
return new JObject { ["success"] = true, ["path"] = draft, ["installed"] = false, ["wanted"] = true };
}
async Task IngestCardToMemory(JObject card, string name)
{
if (Memory is null || card is null)
{
return;
}
try
{
string kind = (card["kind"]?.ToString() ?? "lora").Trim().ToLowerInvariant();
string key = (card["name"]?.ToString() ?? name ?? "").Trim();
List bits = [];
foreach (string field in new[] { "when", "avoid", "prompt_hint", "notes" })
{
string v = card[field]?.ToString();
if (!string.IsNullOrWhiteSpace(v))
{
bits.Add($"{field}: {v.Trim()}");
}
}
if (card["triggers"] is JArray tr)
{
string joined = string.Join(", ", tr.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)));
if (!string.IsNullOrWhiteSpace(joined))
{
bits.Add("triggers: " + joined);
}
}
if (bits.Count == 0 || string.IsNullOrWhiteSpace(key))
{
return;
}
string text = $"{kind} {key}. " + string.Join(" ", bits);
string baseUrl = NormalizeBaseUrl(Config.LoadSettings()["base_url"]?.ToString());
string embedModel = Config.LoadSettings()["embed_model"]?.ToString()
?? Config.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString();
await Memory.UpsertTextAsync(baseUrl, "card", key, text, "user", card, embedModel);
}
catch (Exception ex)
{
Logs.Debug($"IngestCardToMemory: {ex.Message}");
}
}
public async Task AssistentEnqueueWanted(Session session, string kind, string url, int version_id = 0, string title = null, JObject card = null)
{
await Task.CompletedTask;
kind = (kind ?? "lora").Trim().ToLowerInvariant();
if (kind is not ("lora" or "checkpoint" or "vae" or "embedding" or "controlnet" or "upscaler" or "clip"))
{
kind = "lora";
}
url = (url ?? "").Trim();
if (string.IsNullOrWhiteSpace(url) && version_id > 0)
{
url = $"https://civitai.red/models/0?modelVersionId={version_id}";
}
if (string.IsNullOrWhiteSpace(url))
{
return new JObject { ["error"] = "url or version_id required" };
}
if (version_id <= 0)
{
Match m = Regex.Match(url, @"modelVersionId=(\d+)", RegexOptions.IgnoreCase);
if (m.Success)
{
version_id = int.Parse(m.Groups[1].Value);
}
}
string path = WantedModelsPath();
Directory.CreateDirectory(Path.GetDirectoryName(path) ?? DataRoot());
Dictionary> sections = LoadWantedYaml(File.Exists(path) ? File.ReadAllText(path, Encoding.UTF8) : "");
if (version_id > 0)
{
foreach (List list in sections.Values)
{
if (list.Any(e => e.VersionId == version_id))
{
return new JObject { ["success"] = true, ["already"] = true, ["path"] = path, ["version_id"] = version_id };
}
}
}
else
{
foreach (List list in sections.Values)
{
if (list.Any(e => string.Equals(e.Url, url, StringComparison.OrdinalIgnoreCase)))
{
return new JObject { ["success"] = true, ["already"] = true, ["path"] = path };
}
}
}
if (!sections.TryGetValue(kind, out List bucket))
{
bucket = [];
sections[kind] = bucket;
}
bucket.Add(new WantedEntry { Url = url, Title = title, VersionId = version_id });
File.WriteAllText(path, WriteWantedYaml(sections), Encoding.UTF8);
if (card is not null && version_id > 0)
{
Directory.CreateDirectory(WantedCardsDir());
string draft = Path.Combine(WantedCardsDir(), $"{version_id}.assistent.json");
File.WriteAllText(draft, card.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
}
return new JObject { ["success"] = true, ["path"] = path, ["version_id"] = version_id };
}
sealed class WantedEntry
{
public string Url;
public string Title;
public int VersionId;
}
static Dictionary> LoadWantedYaml(string raw)
{
Dictionary> sections = new(StringComparer.OrdinalIgnoreCase);
string currentKind = null;
WantedEntry cur = null;
void Flush()
{
if (cur is null || string.IsNullOrWhiteSpace(cur.Url) || string.IsNullOrWhiteSpace(currentKind))
{
cur = null;
return;
}
if (!sections.TryGetValue(currentKind, out List list))
{
list = [];
sections[currentKind] = list;
}
list.Add(cur);
cur = null;
}
foreach (string line in (raw ?? "").Split('\n'))
{
string t = line.TrimEnd();
if (string.IsNullOrWhiteSpace(t) || t.TrimStart().StartsWith('#'))
{
continue;
}
Match kindLine = Regex.Match(t, @"^([A-Za-z0-9_-]+):\s*$");
if (kindLine.Success && !t.TrimStart().StartsWith('-'))
{
Flush();
currentKind = kindLine.Groups[1].Value.Trim().ToLowerInvariant();
continue;
}
Match urlLine = Regex.Match(t, @"^\s*-\s*url:\s*[""']?(.+?)[""']?\s*$");
if (urlLine.Success)
{
Flush();
cur = new WantedEntry { Url = urlLine.Groups[1].Value.Trim() };
continue;
}
if (cur is null)
{
continue;
}
Match titleLine = Regex.Match(t, @"^\s*title:\s*[""']?(.+?)[""']?\s*$");
if (titleLine.Success)
{
cur.Title = titleLine.Groups[1].Value.Trim();
continue;
}
Match vidLine = Regex.Match(t, @"^\s*version_id:\s*(\d+)\s*$");
if (vidLine.Success && int.TryParse(vidLine.Groups[1].Value, out int vid))
{
cur.VersionId = vid;
}
}
Flush();
return sections;
}
static string WriteWantedYaml(Dictionary> sections)
{
StringBuilder sb = new();
sb.AppendLine("# Assistent wanted queue — merged into local models.yaml on gpu-rent up/capture");
string[] order = ["checkpoint", "lora", "vae", "embedding", "controlnet", "upscaler", "clip"];
HashSet seen = new(StringComparer.OrdinalIgnoreCase);
foreach (string kind in order.Concat(sections.Keys.OrderBy(k => k, StringComparer.OrdinalIgnoreCase)))
{
if (!seen.Add(kind) || !sections.TryGetValue(kind, out List list) || list.Count == 0)
{
continue;
}
sb.AppendLine($"{kind}:");
foreach (WantedEntry e in list)
{
sb.AppendLine($" - url: \"{e.Url.Replace("\"", "%22")}\"");
if (!string.IsNullOrWhiteSpace(e.Title))
{
sb.AppendLine($" title: \"{e.Title.Replace("\"", "'")}\"");
}
if (e.VersionId > 0)
{
sb.AppendLine($" version_id: {e.VersionId}");
}
}
}
return sb.ToString();
}
string TasteJsonPath() => Path.Combine(DataRoot(), "Assistent", "taste.json");
public async Task AssistentGetTaste(Session session)
{
await Task.CompletedTask;
string path = TasteJsonPath();
if (!File.Exists(path))
{
return new JObject { ["success"] = true, ["taste"] = null };
}
try
{
JObject taste = JObject.Parse(File.ReadAllText(path, Encoding.UTF8));
return new JObject { ["success"] = true, ["taste"] = taste };
}
catch (Exception ex)
{
return new JObject { ["error"] = $"taste.json: {ex.Message}" };
}
}
public async Task AssistentSaveTaste(Session session, JObject taste)
{
await Task.CompletedTask;
if (taste is null)
{
return new JObject { ["error"] = "taste required" };
}
if (taste["updated"] == null)
{
taste["updated"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
}
string dir = Path.Combine(DataRoot(), "Assistent");
Directory.CreateDirectory(dir);
string path = TasteJsonPath();
File.WriteAllText(path, taste.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
return new JObject { ["success"] = true, ["path"] = path };
}
public async Task AssistentGetCardMeta(Session session, string kind, string name, int version_id = 0, bool fetch = false)
{
string set = SetNameForKind(kind);
string weight = ModelWeightPath(set, name);
JObject civitai = null;
JArray exampleUrls = [];
JArray previewUrls = [];
bool hasSidecar = false;
string fetchError = null;
bool fetched = false;
if (!string.IsNullOrWhiteSpace(weight))
{
string stem = Path.GetFileNameWithoutExtension(weight);
string dir = Path.GetDirectoryName(weight);
string side = Path.Combine(dir ?? "", $"{stem}.civitai.json");
if (File.Exists(side))
{
hasSidecar = true;
try
{
civitai = JObject.Parse(File.ReadAllText(side, Encoding.UTF8));
}
catch
{
// ignore
}
}
foreach (string suffix in new[] { ".preview.jpg", ".preview.png", ".preview.jpeg", ".jpg", ".png", ".webp" })
{
string prev = Path.Combine(dir ?? "", stem + suffix);
if (File.Exists(prev))
{
// Swarm View path — relative URL works in the same origin browser session.
previewUrls.Add($"View/Models/{(kind == "lora" ? "Lora" : "Stable-Diffusion")}/{Path.GetFileName(prev)}");
break;
}
}
}
if (civitai is not null)
{
if (version_id <= 0)
{
version_id = civitai["id"]?.Value() ?? 0;
}
CollectExampleUrls(civitai, exampleUrls);
}
string hash = null;
string trigger = null;
try
{
if (Program.T2IModelSets.TryGetValue(set, out T2IModelHandler h)
&& (h.Models.TryGetValue(name, out T2IModel m)
|| h.Models.TryGetValue(name.Replace('\\', '/'), out m)))
{
trigger = m.Metadata?.TriggerPhrase;
hash = m.Metadata?.Hash;
}
}
catch
{
// ignore
}
if (fetch && civitai is null)
{
string apiKey = session.User.GetGenericData("civitai_api", "key") ?? "";
if (string.IsNullOrWhiteSpace(apiKey))
{
fetchError = "Civitai: нет ключа в User Settings";
}
else
{
try
{
JObject remote = null;
if (version_id > 0)
{
remote = await FetchCivitaiModelVersion(apiKey, version_id);
}
if (remote is null && !string.IsNullOrWhiteSpace(hash))
{
string sha = hash.Trim().ToLowerInvariant();
if (sha.StartsWith("sha256:"))
{
sha = sha["sha256:".Length..];
}
if (sha.Length == 64)
{
remote = await FetchCivitaiByHash(apiKey, sha);
}
else
{
fetchError ??= "Civitai: хеш модели не SHA256";
}
}
if (remote is not null)
{
civitai = remote;
fetched = true;
version_id = remote["id"]?.Value() ?? version_id;
CollectExampleUrls(remote, exampleUrls);
if (!string.IsNullOrWhiteSpace(weight))
{
try
{
string stem = Path.GetFileNameWithoutExtension(weight);
string dir = Path.GetDirectoryName(weight);
string side = Path.Combine(dir ?? "", $"{stem}.civitai.json");
File.WriteAllText(side, remote.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
hasSidecar = true;
}
catch (Exception ex)
{
Logs.Debug($"AssistentGetCardMeta write sidecar: {ex.Message}");
}
}
}
else if (fetchError is null)
{
fetchError = string.IsNullOrWhiteSpace(hash)
? "Civitai: нет hash и version_id"
: "Хеш не найден на Civitai";
}
}
catch (Exception ex)
{
fetchError = $"Civitai: {ex.Message}";
}
}
}
JObject card = ReadCardObject(kind, name);
return new JObject
{
["success"] = true,
["kind"] = kind,
["name"] = name,
["version_id"] = version_id,
["trigger_phrase"] = trigger,
["has_card"] = card is not null,
["has_sidecar"] = hasSidecar,
["fetched"] = fetched,
["fetch_error"] = fetchError,
["card"] = card,
["civitai"] = civitai,
["example_urls"] = exampleUrls,
["preview_urls"] = previewUrls,
["weight_path"] = weight,
["hash"] = hash,
};
}
static void CollectExampleUrls(JObject civitai, JArray exampleUrls)
{
if (civitai?["images"] is not JArray imgs)
{
return;
}
foreach (JToken img in imgs.Take(6))
{
string u = img?["url"]?.ToString();
if (!string.IsNullOrWhiteSpace(u))
{
exampleUrls.Add(u);
}
}
}
async Task FetchCivitaiByHash(string apiKey, string sha)
{
string[] hosts = ["civitai.red", "civitai.com"];
Exception last = null;
foreach (string host in hosts)
{
try
{
string url = $"https://{host}/api/v1/model-versions/by-hash/{sha}";
using HttpRequestMessage req = new(HttpMethod.Get, url);
if (!string.IsNullOrWhiteSpace(apiKey))
{
req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim());
}
using HttpResponseMessage resp = await HttpClient.SendAsync(req);
string body = await resp.Content.ReadAsStringAsync();
if (resp.StatusCode == System.Net.HttpStatusCode.NotFound)
{
continue;
}
if (!resp.IsSuccessStatusCode)
{
last = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 160)}");
if ((int)resp.StatusCode is 401 or 403)
{
throw last;
}
continue;
}
return JObject.Parse(body);
}
catch (Exception ex) when (ex is not HttpRequestException && ex.Message.Contains("401"))
{
throw;
}
catch (Exception ex)
{
last = ex;
}
}
if (last is not null)
{
throw last;
}
return null;
}
async Task FetchCivitaiModelVersion(string apiKey, int versionId)
{
string[] hosts = ["civitai.red", "civitai.com"];
Exception last = null;
foreach (string host in hosts)
{
try
{
string url = $"https://{host}/api/v1/model-versions/{versionId}";
using HttpRequestMessage req = new(HttpMethod.Get, url);
if (!string.IsNullOrWhiteSpace(apiKey))
{
req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim());
}
using HttpResponseMessage resp = await HttpClient.SendAsync(req);
string body = await resp.Content.ReadAsStringAsync();
if (!resp.IsSuccessStatusCode)
{
last = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 160)}");
if ((int)resp.StatusCode is 401 or 403)
{
throw last;
}
continue;
}
return JObject.Parse(body);
}
catch (Exception ex)
{
last = ex;
if (ex.Message.Contains("401") || ex.Message.Contains("403"))
{
throw;
}
}
}
if (last is not null)
{
throw last;
}
return null;
}
/// Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).
/// Pass rescan=true after downloads so new files appear (calls Program.RefreshAllModelSets).
public async Task AssistentListInventory(Session session, bool rescan = false)
{
await Task.CompletedTask;
if (rescan)
{
try
{
Program.RefreshAllModelSets();
}
catch (Exception ex)
{
Logs.Debug($"AssistentListInventory rescan: {ex.Message}");
try
{
Program.ModelRefreshEvent?.Invoke();
}
catch (Exception ex2)
{
Logs.Debug($"AssistentListInventory ModelRefreshEvent: {ex2.Message}");
}
}
}
JArray loras = [];
JArray checkpoints = [];
JArray wildcards = [];
if (Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler loraHandler))
{
foreach (T2IModel model in loraHandler.Models.Values
.OrderByDescending(m => LooksLikeKreaArch(m))
.ThenBy(m => m.Name)
.Take(CfgInt("max_loras_inventory", MaxLorasInInventoryFallback)))
{
loras.Add(BuildInventoryModelEntry(model, "lora"));
}
}
if (Program.T2IModelSets.TryGetValue("Stable-Diffusion", out T2IModelHandler ckptHandler))
{
foreach (T2IModel model in ckptHandler.Models.Values
.OrderByDescending(m => LooksLikeKreaArch(m))
.ThenBy(m => m.Name)
.Take(CfgInt("max_checkpoints_inventory", MaxCheckpointsInInventoryFallback)))
{
checkpoints.Add(BuildInventoryModelEntry(model, "checkpoint"));
}
}
try
{
foreach (string name in WildcardsHelper.ListFiles.OrderBy(n => n).Take(CfgInt("max_wildcards_inventory", MaxWildcardsInInventoryFallback)))
{
wildcards.Add(new JObject { ["name"] = name });
}
}
catch (Exception ex)
{
Logs.Debug($"AssistentListInventory wildcards: {ex.Message}");
}
bool hasCivitaiKey = !string.IsNullOrWhiteSpace(session.User.GetGenericData("civitai_api", "key"));
return new JObject
{
["success"] = true,
["loras"] = loras,
["checkpoints"] = checkpoints,
["wildcards"] = wildcards,
["has_civitai_key"] = hasCivitaiKey,
["rescanned"] = rescan,
["inventory_at"] = DateTimeOffset.UtcNow.ToUnixTimeSeconds(),
};
}
static bool LooksLikeKreaArch(T2IModel model)
{
string arch = model?.ModelClass?.ID ?? "";
string compat = model?.ModelClass?.CompatClass?.ID ?? "";
string name = model?.Name ?? "";
string blob = $"{arch} {compat} {name}".ToLowerInvariant();
return blob.Contains("krea");
}
JObject BuildInventoryModelEntry(T2IModel model, string kind)
{
string weight = null;
try { weight = model.RawFilePath; } catch { /* ignore */ }
string cardPath = CardPathForWeight(weight);
bool hasCard = !string.IsNullOrWhiteSpace(cardPath) && File.Exists(cardPath);
string usage = model.Metadata?.UsageHint;
string desc = model.Metadata?.Description;
try
{
if (string.IsNullOrWhiteSpace(desc) && !string.IsNullOrWhiteSpace(model.Description))
{
desc = model.Description;
}
}
catch
{
// older Swarm builds
}
string blurb = null;
if (hasCard)
{
try
{
JObject card = JObject.Parse(File.ReadAllText(cardPath, Encoding.UTF8));
string fromCard = (card["notes"] ?? card["when"] ?? card["prompt_hint"])?.ToString();
if (!string.IsNullOrWhiteSpace(fromCard))
{
blurb = Clip(fromCard.Trim(), CfgInt("inventory_blurb_max", InventoryBlurbMaxFallback));
}
}
catch
{
// ignore bad card json
}
}
if (string.IsNullOrWhiteSpace(blurb))
{
string raw = !string.IsNullOrWhiteSpace(usage) ? usage : desc;
if (!string.IsNullOrWhiteSpace(raw))
{
blurb = Clip(CollapseWs(raw), CfgInt("inventory_blurb_max", InventoryBlurbMaxFallback));
}
}
JArray tags = null;
if (model.Metadata?.Tags is { Length: > 0 } tagArr)
{
tags = new JArray(tagArr.Where(t => !string.IsNullOrWhiteSpace(t)).Take(8));
}
string trigger = model.Metadata?.TriggerPhrase;
JArray triggers = null;
if (!string.IsNullOrWhiteSpace(trigger))
{
triggers = new JArray(trigger.Split(new[] { ',', ';' }, StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries).Take(12));
}
JObject entry = new()
{
["name"] = model.Name,
["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
["kind"] = kind,
["trigger_phrase"] = trigger,
["architecture"] = model.ModelClass?.ID,
["compat_class"] = model.ModelClass?.CompatClass?.ID,
["hash"] = model.Metadata?.Hash ?? "",
["has_card"] = hasCard,
["krea_likely"] = LooksLikeKreaArch(model),
};
if (!string.IsNullOrWhiteSpace(weight))
{
string stem = Path.GetFileNameWithoutExtension(weight);
string dir = Path.GetDirectoryName(weight);
string side = Path.Combine(dir ?? "", $"{stem}.civitai.json");
entry["has_sidecar"] = File.Exists(side);
foreach (string suffix in new[] { ".preview.jpg", ".preview.png", ".preview.jpeg", ".jpg", ".png", ".webp" })
{
string prev = Path.Combine(dir ?? "", stem + suffix);
if (File.Exists(prev))
{
string folder = kind == "lora" ? "Lora" : "Stable-Diffusion";
entry["preview_url"] = $"View/Models/{folder}/{Path.GetFileName(prev)}";
break;
}
}
}
else
{
entry["has_sidecar"] = false;
}
if (triggers is not null && triggers.Count > 0)
{
entry["triggers"] = triggers;
}
if (!string.IsNullOrWhiteSpace(blurb))
{
entry["blurb"] = blurb;
}
if (!string.IsNullOrWhiteSpace(usage))
{
entry["usage_hint"] = Clip(CollapseWs(usage), 120);
}
if (tags is not null && tags.Count > 0)
{
entry["tags"] = tags;
}
string defW = model.Metadata?.LoraDefaultWeight;
if (!string.IsNullOrWhiteSpace(defW) && kind == "lora")
{
entry["default_weight"] = defW;
}
return entry;
}
static string CollapseWs(string text)
{
if (string.IsNullOrWhiteSpace(text))
{
return "";
}
return Regex.Replace(text.Trim(), @"\s+", " ");
}
/// Search Civitai for LoRAs (prefers Krea 2 base). Uses Swarm-stored civitai_api key.
public async Task AssistentSearchCivitai(Session session, string query, int limit = 8)
{
string q = (query ?? "").Trim();
if (string.IsNullOrWhiteSpace(q))
{
return new JObject { ["error"] = "query is required" };
}
limit = Math.Clamp(limit, 1, 20);
string apiKey = session.User.GetGenericData("civitai_api", "key") ?? "";
HashSet installedNames = CollectInstalledLoraNames();
HashSet installedHashes = CollectInstalledLoraHashes();
string[] hosts = ["civitai.red", "civitai.com"];
Exception lastEx = null;
foreach (string host in hosts)
{
try
{
string url = $"https://{host}/api/v1/models?limit={limit}&types=LORA&query={Uri.EscapeDataString(q)}";
using HttpRequestMessage req = new(HttpMethod.Get, url);
if (!string.IsNullOrWhiteSpace(apiKey))
{
req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim());
}
using HttpResponseMessage resp = await HttpClient.SendAsync(req);
string body = await resp.Content.ReadAsStringAsync();
if (!resp.IsSuccessStatusCode)
{
lastEx = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 200)}");
continue;
}
JObject parsed = JObject.Parse(body);
JArray items = parsed["items"] as JArray ?? [];
JArray results = [];
foreach (JToken item in items)
{
if (item is not JObject mo)
{
continue;
}
JObject card = BuildCivitaiCard(mo, installedNames, installedHashes);
if (card is not null)
{
results.Add(card);
}
}
// Prefer Krea-compatible first
JArray sorted = new(results.OrderByDescending(t => LooksLikeKrea(t["base_model"]?.ToString())).ThenBy(t => t["name"]?.ToString()));
return new JObject
{
["success"] = true,
["query"] = q,
["host"] = host,
["results"] = sorted,
["has_civitai_key"] = !string.IsNullOrWhiteSpace(apiKey),
};
}
catch (Exception ex)
{
lastEx = ex;
}
}
return new JObject { ["error"] = $"Civitai search failed: {lastEx?.Message ?? "unknown"}" };
}
static bool LooksLikeKrea(string text) => !string.IsNullOrEmpty(text) && Regex.IsMatch(text, @"krea", RegexOptions.IgnoreCase);
static HashSet CollectInstalledLoraNames()
{
HashSet names = new(StringComparer.OrdinalIgnoreCase);
if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler))
{
return names;
}
foreach (T2IModel m in handler.Models.Values)
{
names.Add(m.Name);
string leaf = m.Name.Replace('\\', '/').AfterLast('/');
if (!string.IsNullOrEmpty(leaf))
{
names.Add(leaf);
names.Add(Path.GetFileNameWithoutExtension(leaf));
}
}
return names;
}
static HashSet CollectInstalledLoraHashes()
{
HashSet hashes = new(StringComparer.OrdinalIgnoreCase);
if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler))
{
return hashes;
}
foreach (T2IModel m in handler.Models.Values)
{
string h = m.Metadata?.Hash;
if (!string.IsNullOrWhiteSpace(h))
{
hashes.Add(h.Trim().ToLowerInvariant());
}
}
return hashes;
}
static JObject BuildCivitaiCard(JObject model, HashSet installedNames, HashSet installedHashes)
{
string name = model["name"]?.ToString() ?? "";
JArray versions = model["modelVersions"] as JArray;
JObject ver = versions?.FirstOrDefault() as JObject;
if (ver is null)
{
return null;
}
string baseModel = ver["baseModel"]?.ToString() ?? "";
JArray trained = ver["trainedWords"] as JArray ?? [];
List triggers = trained.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)).Take(8).ToList();
JObject file = null;
foreach (JToken f in ver["files"] as JArray ?? [])
{
if (f is JObject fo && (fo["primary"]?.Value() == true || (fo["name"]?.ToString() ?? "").EndsWith(".safetensors", StringComparison.OrdinalIgnoreCase)))
{
file = fo;
break;
}
}
file ??= (ver["files"] as JArray)?.FirstOrDefault() as JObject;
string downloadUrl = file?["downloadUrl"]?.ToString() ?? ver["downloadUrl"]?.ToString() ?? "";
string fileName = file?["name"]?.ToString() ?? "";
string sha = file?["hashes"]?["SHA256"]?.ToString() ?? file?["hashes"]?["AutoV2"]?.ToString() ?? "";
string saveName = string.IsNullOrWhiteSpace(fileName)
? Regex.Replace(name, @"[^\w\-.]+", "_").Trim('_')
: Path.GetFileNameWithoutExtension(fileName);
bool already = false;
if (!string.IsNullOrWhiteSpace(sha) && installedHashes.Contains(sha.Trim().ToLowerInvariant()))
{
already = true;
}
else if (installedNames.Contains(saveName) || installedNames.Contains(name) || installedNames.Contains(fileName))
{
already = true;
}
return new JObject
{
["id"] = model["id"],
["version_id"] = ver["id"],
["name"] = name,
["base_model"] = baseModel,
["krea_likely"] = LooksLikeKrea(baseModel),
["triggers"] = new JArray(triggers),
["download_url"] = downloadUrl,
["file_name"] = saveName,
["sha256"] = sha,
["already_installed"] = already,
["n_sfw"] = model["nsfw"]?.Value() ?? false,
};
}
List BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null, string personaId = null, IEnumerable skillIds = null)
{
List 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) && packName != "base_krea2" && packName != "core")
{
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;
}
string ResolvePersonaPrompt(string personaId) => Config.RenderIdentityBlock(personaId);
static JObject TryParsePatch(string reply)
{
if (string.IsNullOrWhiteSpace(reply))
{
return null;
}
foreach (Match match in JsonFenceRe.Matches(reply))
{
string raw = match.Groups[1].Value.Trim();
try
{
JObject obj = JObject.Parse(raw);
if (obj is not null && (obj["prompt"] != null || obj["loras"] != null || obj["width"] != null
|| obj["height"] != null || obj["steps"] != null || obj["cfg"] != null
|| obj["seed"] != null || obj["sigma_shift"] != null || obj["sampler"] != null
|| obj["actions"] != null || obj["search_query"] != null || obj["civitai_query"] != null
|| obj["use_init_image"] != null || obj["clear_init_image"] != null
|| obj["init_creativity"] != null || obj["denoise"] != null
|| obj["use_mask_image"] != null || obj["clear_mask_image"] != null
|| obj["mask_blur"] != null || obj["mask_grow"] != null
|| obj["look_at"] != null || obj["vision_from"] != null || obj["vision_slots"] != null
|| obj["slot_to_init"] != null || obj["slot_to_mask"] != null
|| obj["snapshot_generate"] != null || obj["select_slot"] != null
|| obj["aspect"] != null || obj["images"] != null || obj["batch"] != null
|| obj["vary"] != null || obj["lock_seed"] != null
|| obj["creativity"] != null || obj["intensity"] != null
|| obj["complexity"] != null || obj["movement"] != null
|| obj["clear_prompt_images"] != null || obj["slot_to_prompt_image"] != null
|| obj["pack"] != null || obj["memories"] != null || obj["memory"] != null))
{
return obj;
}
}
catch
{
// not json
}
}
return null;
}
static string ExtractSearchQuery(JObject patch)
{
if (patch is null)
{
return null;
}
string q = (patch["search_query"] ?? patch["civitai_query"])?.ToString()?.Trim();
if (!string.IsNullOrWhiteSpace(q))
{
return q;
}
if (patch["actions"] is JArray acts)
{
foreach (JToken a in acts)
{
if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
{
return q; // may still be null — caller checks
}
}
}
return null;
}
static bool WantsCivitaiSearch(JObject patch)
{
if (patch is null)
{
return false;
}
return !string.IsNullOrWhiteSpace(ExtractSearchQuery(patch));
}
static void ExtractChatPayload(JObject raw, ref string baseUrl, ref string model, ref string pack, ref bool includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills)
{
JObject whole = raw ?? [];
JObject nested = whole["raw"] as JObject;
if (string.IsNullOrWhiteSpace(baseUrl))
{
baseUrl = whole["base_url"]?.ToString()
?? whole["baseUrl"]?.ToString()
?? nested?["base_url"]?.ToString()
?? nested?["baseUrl"]?.ToString();
}
if (string.IsNullOrWhiteSpace(model))
{
model = whole["model"]?.ToString() ?? nested?["model"]?.ToString();
}
if (string.IsNullOrWhiteSpace(pack))
{
pack = whole["pack"]?.ToString() ?? nested?["pack"]?.ToString();
}
if (whole["includeBase"] is not null)
{
includeBase = whole.Value("includeBase") ?? includeBase;
}
userMessages = (whole["messages"] as JArray) ?? (nested?["messages"] as JArray);
contextJson = whole["context_json"]?.ToString() ?? nested?["context_json"]?.ToString();
persona = whole["persona"]?.ToString() ?? nested?["persona"]?.ToString() ?? "neutral";
skills = (whole["skills"] as JArray) ?? (nested?["skills"] as JArray);
}
async Task<(string reply, JObject raw, JArray civitaiResults)> RunChatWithHops(
Session session,
string root,
string modelName,
string packName,
bool includeBase,
string contextJson,
JArray userMessages,
Func onDelta = null,
Func onHopStart = null,
string personaId = null,
JArray skillIds = null,
string embedModel = null)
{
string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
List 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() ?? 10;
hits = await Memory.RetrieveAsync(root, retrieveQuery, topK, embed);
}
catch (Exception ex)
{
Logs.Debug($"Assistent memory retrieve: {ex.Message}");
}
string enrichedContext = InjectMemoryHits(contextJson, hits);
List 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);
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 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() == 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);
}
async Task ApplyMemoryActions(string root, JObject patch, string embedModel)
{
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() ?? "";
try
{
if (forget && string.IsNullOrWhiteSpace(text))
{
Memory.Forget(kind, key);
}
else if (upsert || !string.IsNullOrWhiteSpace(text))
{
await Memory.UpsertTextAsync(root, kind, key, text, "user", mo, embedModel);
}
}
catch (Exception ex)
{
Logs.Debug($"ApplyMemoryActions: {ex.Message}");
}
}
}
async Task<(string reply, JObject raw)> CallOllamaChat(
string root,
string modelName,
List ollamaMessages,
bool stream,
Func onDelta,
string personaId = null)
{
int numCtx = Config.LoadAssistant(AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId())["num_ctx"]?.Value()
?? DefaultNumCtxFallback;
JObject payload = new()
{
["model"] = modelName,
["stream"] = stream,
["messages"] = new JArray(ollamaMessages),
["options"] = new JObject
{
["num_ctx"] = numCtx,
},
["keep_alive"] = "15m",
};
using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
using HttpRequestMessage req = new(HttpMethod.Post, $"{root}/api/chat") { Content = content };
using HttpResponseMessage resp = await HttpClient.SendAsync(req, stream
? HttpCompletionOption.ResponseHeadersRead
: HttpCompletionOption.ResponseContentRead);
if (!resp.IsSuccessStatusCode)
{
string errBody = await resp.Content.ReadAsStringAsync();
throw new Exception($"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(errBody, 800)}");
}
if (!stream)
{
string body = await resp.Content.ReadAsStringAsync();
JObject parsed = JObject.Parse(body);
string reply = parsed["message"]?["content"]?.ToString() ?? parsed["response"]?.ToString() ?? "";
return (reply, parsed);
}
StringBuilder full = new();
await using Stream streamBody = await resp.Content.ReadAsStreamAsync();
using StreamReader reader = new(streamBody, Encoding.UTF8);
JObject last = null;
while (true)
{
string line = await reader.ReadLineAsync();
if (line is null)
{
break;
}
if (string.IsNullOrWhiteSpace(line))
{
continue;
}
JObject chunk = JObject.Parse(line);
last = chunk;
string delta = chunk["message"]?["content"]?.ToString() ?? "";
if (!string.IsNullOrEmpty(delta))
{
full.Append(delta);
if (onDelta is not null)
{
await onDelta(delta);
}
}
if (chunk["done"]?.Value() == true)
{
break;
}
}
return (full.ToString(), last ?? new JObject());
}
///
/// SwarmUI passes the whole request as the JObject param (not only a nested key).
/// Support both flat fields and legacy nested raw.
///
// ExtractChatPayload defined above
/// Proxy to Ollama /api/chat (non-stream), with optional Civitai search hop.
public async Task AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
{
ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills);
string root = NormalizeBaseUrl(baseUrl);
string modelName = (model ?? "").Trim();
if (string.IsNullOrWhiteSpace(modelName))
{
return new JObject { ["error"] = "model is required" };
}
if (userMessages is null || userMessages.Count == 0)
{
return new JObject { ["error"] = "messages required" };
}
string packName = (pack ?? "write_prompt").Trim();
string embedModel = raw?["embed_model"]?.ToString() ?? Config.LoadSettings()["embed_model"]?.ToString();
try
{
(string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
session, root, modelName, packName, includeBase, contextJson, userMessages, personaId: persona, skillIds: skills, embedModel: embedModel);
return new JObject
{
["success"] = true,
["reply"] = reply,
["model"] = modelName,
["pack"] = packName,
["persona"] = persona,
["raw"] = parsed,
["civitai_results"] = civitai,
};
}
catch (Exception ex)
{
return new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" };
}
}
/// WebSocket streaming chat (Ollama stream:true) + Civitai hops.
public async Task AssistentChatWS(Session session, WebSocket ws, string baseUrl, string model, string pack, bool includeBase, JObject raw)
{
ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills);
string root = NormalizeBaseUrl(baseUrl);
string modelName = (model ?? "").Trim();
if (string.IsNullOrWhiteSpace(modelName))
{
await ws.SendJson(new JObject { ["error"] = "model is required" }, API.WebsocketTimeout);
return null;
}
if (userMessages is null || userMessages.Count == 0)
{
await ws.SendJson(new JObject { ["error"] = "messages required" }, API.WebsocketTimeout);
return null;
}
string packName = (pack ?? "write_prompt").Trim();
string embedModel = raw?["embed_model"]?.ToString() ?? Config.LoadSettings()["embed_model"]?.ToString();
try
{
if (ws.State == WebSocketState.Open)
{
await ws.SendJson(new JObject
{
["phase"] = "waiting_ollama",
["notice"] = "Loading model into GPU…",
}, API.WebsocketTimeout);
}
async Task OnDelta(string delta)
{
if (ws.State == WebSocketState.Open)
{
await ws.SendJson(new JObject { ["delta"] = delta }, API.WebsocketTimeout);
}
}
async Task OnHopStart(int hop)
{
if (ws.State == WebSocketState.Open && hop > 0)
{
await ws.SendJson(new JObject
{
["clear_stream"] = true,
["hop"] = hop + 1,
["notice"] = "Civitai search done — refining…",
}, API.WebsocketTimeout);
}
}
(string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
session, root, modelName, packName, includeBase, contextJson, userMessages, OnDelta, OnHopStart, persona, skills, embedModel);
await ws.SendJson(new JObject
{
["success"] = true,
["done"] = true,
["reply"] = reply,
["model"] = modelName,
["pack"] = packName,
["persona"] = persona,
["raw"] = parsed,
["civitai_results"] = civitai,
}, API.WebsocketTimeout);
}
catch (Exception ex)
{
await ws.SendJson(new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" }, API.WebsocketTimeout);
}
return null;
}
}