1467 lines
54 KiB
C#
1467 lines
54 KiB
C#
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;
|
|
|
|
/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches + Generate/Civitai.</summary>
|
|
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 static readonly string[] PackNames =
|
|
[
|
|
"base_krea2",
|
|
"write_prompt",
|
|
"critique_image",
|
|
"compose_scene",
|
|
"fix_params",
|
|
"inpaint_edit",
|
|
"describe_ref",
|
|
"catalog_card",
|
|
];
|
|
|
|
public static readonly string[] DefaultPersonaIds = ["neutral", "lewd", "aggressive"];
|
|
|
|
const int MaxCivitaiHops = 2;
|
|
const int MaxLorasInInventory = 150;
|
|
const int MaxWildcardsInInventory = 80;
|
|
const int MaxCheckpointsInInventory = 60;
|
|
const int InventoryBlurbMax = 140;
|
|
/// <summary>Ollama default num_ctx is 4096; Assistent system+inventory+vision exceeds that.</summary>
|
|
const int DefaultNumCtx = 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, multi-window board, personas, model cards, Generate loop, Civitai Confirm.";
|
|
License = "MIT";
|
|
Version = "0.5.3";
|
|
Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint"];
|
|
}
|
|
|
|
public override void OnInit()
|
|
{
|
|
HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
|
|
API.RegisterAPICall(AssistentListModels, false, PermUse);
|
|
API.RegisterAPICall(AssistentGetPacks, false, PermUse);
|
|
API.RegisterAPICall(AssistentListPersonas, false, 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(AssistentChat, true, PermUse);
|
|
API.RegisterAPICall(AssistentChatWS, true, PermUse);
|
|
Logs.Init("Swarm Assistent extension loaded (Ollama proxy + personas + model cards)");
|
|
}
|
|
|
|
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)
|
|
{
|
|
string safe = name.Replace('\\', '/').AfterLast('/').Replace("..", "");
|
|
if (!PackNames.Contains(safe))
|
|
{
|
|
return null;
|
|
}
|
|
string path = Path.Combine(FilePath, "Prompts", $"{safe}.md");
|
|
if (!File.Exists(path))
|
|
{
|
|
return null;
|
|
}
|
|
return File.ReadAllText(path, Encoding.UTF8);
|
|
}
|
|
|
|
public async Task<JObject> 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 models = [];
|
|
foreach (JToken m in parsed["models"] as JArray ?? [])
|
|
{
|
|
models.Add(m["name"]?.ToString() ?? "");
|
|
}
|
|
return new JObject { ["success"] = true, ["base_url"] = root, ["models"] = models };
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
return new JObject { ["error"] = $"Ollama unreachable at {root}: {ex.Message}" };
|
|
}
|
|
}
|
|
|
|
public async Task<JObject> AssistentGetPacks(Session session)
|
|
{
|
|
JObject packs = new();
|
|
foreach (string name in PackNames)
|
|
{
|
|
string text = ReadPackFile(name);
|
|
if (text is not null)
|
|
{
|
|
packs[name] = text;
|
|
}
|
|
}
|
|
return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = new JArray(PackNames) };
|
|
}
|
|
|
|
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 string ReadPersonaFile(string id)
|
|
{
|
|
string safe = (id ?? "").Replace('\\', '/').AfterLast('/').Replace("..", "");
|
|
if (string.IsNullOrWhiteSpace(safe))
|
|
{
|
|
return null;
|
|
}
|
|
string path = Path.Combine(FilePath, "Personas", $"{safe}.md");
|
|
if (!File.Exists(path))
|
|
{
|
|
return null;
|
|
}
|
|
return File.ReadAllText(path, Encoding.UTF8);
|
|
}
|
|
|
|
public async Task<JObject> AssistentListPersonas(Session session)
|
|
{
|
|
await Task.CompletedTask;
|
|
Dictionary<string, JObject> byId = new(StringComparer.OrdinalIgnoreCase);
|
|
string def = "neutral";
|
|
|
|
foreach (string id in DefaultPersonaIds)
|
|
{
|
|
string text = ReadPersonaFile(id);
|
|
if (string.IsNullOrWhiteSpace(text))
|
|
{
|
|
continue;
|
|
}
|
|
byId[id] = new JObject
|
|
{
|
|
["id"] = id,
|
|
["title"] = id switch
|
|
{
|
|
"lewd" => "Пошляк",
|
|
"aggressive" => "Агрессивный",
|
|
_ => "Нейтральный",
|
|
},
|
|
["prompt"] = text,
|
|
["source"] = "bundled",
|
|
};
|
|
}
|
|
|
|
string overlay = PersonasOverlayJsonPath();
|
|
if (File.Exists(overlay))
|
|
{
|
|
try
|
|
{
|
|
JObject parsed = JObject.Parse(File.ReadAllText(overlay, Encoding.UTF8));
|
|
if (parsed["default"] != null)
|
|
{
|
|
def = parsed["default"]?.ToString() ?? def;
|
|
}
|
|
if (parsed["personas"] is JArray arr)
|
|
{
|
|
foreach (JToken t in arr)
|
|
{
|
|
if (t is not JObject po)
|
|
{
|
|
continue;
|
|
}
|
|
string id = (po["id"]?.ToString() ?? "").Trim();
|
|
if (string.IsNullOrWhiteSpace(id))
|
|
{
|
|
continue;
|
|
}
|
|
byId[id] = new JObject
|
|
{
|
|
["id"] = id,
|
|
["title"] = po["title"]?.ToString() ?? id,
|
|
["prompt"] = po["prompt"]?.ToString() ?? "",
|
|
["source"] = "overlay",
|
|
};
|
|
}
|
|
}
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
Logs.Debug($"AssistentListPersonas overlay: {ex.Message}");
|
|
}
|
|
}
|
|
|
|
JArray list = [];
|
|
foreach (JObject p in byId.Values.OrderBy(p => p["id"]?.ToString()))
|
|
{
|
|
list.Add(p);
|
|
}
|
|
if (!byId.ContainsKey(def) && list.Count > 0)
|
|
{
|
|
def = list[0]?["id"]?.ToString() ?? "neutral";
|
|
}
|
|
return new JObject
|
|
{
|
|
["success"] = true,
|
|
["default"] = def,
|
|
["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<JObject> 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<JObject> 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);
|
|
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 vid = card["version_id"]?.ToString() ?? "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<int?>() ?? 0, card["title"]?.ToString() ?? name, card);
|
|
}
|
|
return new JObject { ["success"] = true, ["path"] = draft, ["installed"] = false, ["wanted"] = true };
|
|
}
|
|
|
|
public async Task<JObject> 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<string, List<WantedEntry>> sections = LoadWantedYaml(File.Exists(path) ? File.ReadAllText(path, Encoding.UTF8) : "");
|
|
|
|
if (version_id > 0)
|
|
{
|
|
foreach (List<WantedEntry> 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<WantedEntry> 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<WantedEntry> 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<string, List<WantedEntry>> LoadWantedYaml(string raw)
|
|
{
|
|
Dictionary<string, List<WantedEntry>> 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<WantedEntry> 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<string, List<WantedEntry>> 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<string> 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<WantedEntry> 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();
|
|
}
|
|
|
|
public async Task<JObject> AssistentGetCardMeta(Session session, string kind, string name, int version_id = 0)
|
|
{
|
|
// Pull Civitai sidecar next to weight + optional API version for examples.
|
|
await Task.CompletedTask;
|
|
string set = SetNameForKind(kind);
|
|
string weight = ModelWeightPath(set, name);
|
|
JObject civitai = null;
|
|
JArray exampleUrls = [];
|
|
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))
|
|
{
|
|
try
|
|
{
|
|
civitai = JObject.Parse(File.ReadAllText(side, Encoding.UTF8));
|
|
}
|
|
catch
|
|
{
|
|
// ignore
|
|
}
|
|
}
|
|
}
|
|
if (civitai is not null)
|
|
{
|
|
if (version_id <= 0)
|
|
{
|
|
version_id = civitai["id"]?.Value<int?>() ?? 0;
|
|
}
|
|
if (civitai["images"] is JArray imgs)
|
|
{
|
|
foreach (JToken img in imgs.Take(3))
|
|
{
|
|
string u = img?["url"]?.ToString();
|
|
if (!string.IsNullOrWhiteSpace(u))
|
|
{
|
|
exampleUrls.Add(u);
|
|
}
|
|
}
|
|
}
|
|
if (civitai["trainedWords"] is null && civitai["model"] is JObject)
|
|
{
|
|
// keep as-is
|
|
}
|
|
}
|
|
JObject card = ReadCardObject(kind, name);
|
|
string trigger = null;
|
|
try
|
|
{
|
|
if (Program.T2IModelSets.TryGetValue(set, out T2IModelHandler h)
|
|
&& h.Models.TryGetValue(name, out T2IModel m))
|
|
{
|
|
trigger = m.Metadata?.TriggerPhrase;
|
|
}
|
|
}
|
|
catch
|
|
{
|
|
// ignore
|
|
}
|
|
return new JObject
|
|
{
|
|
["success"] = true,
|
|
["kind"] = kind,
|
|
["name"] = name,
|
|
["version_id"] = version_id,
|
|
["trigger_phrase"] = trigger,
|
|
["has_card"] = card is not null,
|
|
["card"] = card,
|
|
["civitai"] = civitai,
|
|
["example_urls"] = exampleUrls,
|
|
["weight_path"] = weight,
|
|
};
|
|
}
|
|
|
|
/// <summary>Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).
|
|
/// Pass rescan=true after downloads so new files appear (calls Program.RefreshAllModelSets).</summary>
|
|
public async Task<JObject> 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(MaxLorasInInventory))
|
|
{
|
|
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(MaxCheckpointsInInventory))
|
|
{
|
|
checkpoints.Add(BuildInventoryModelEntry(model, "checkpoint"));
|
|
}
|
|
}
|
|
|
|
try
|
|
{
|
|
foreach (string name in WildcardsHelper.ListFiles.OrderBy(n => n).Take(MaxWildcardsInInventory))
|
|
{
|
|
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(), InventoryBlurbMax);
|
|
}
|
|
}
|
|
catch
|
|
{
|
|
// ignore bad card json
|
|
}
|
|
}
|
|
if (string.IsNullOrWhiteSpace(blurb))
|
|
{
|
|
string raw = !string.IsNullOrWhiteSpace(usage) ? usage : desc;
|
|
if (!string.IsNullOrWhiteSpace(raw))
|
|
{
|
|
blurb = Clip(CollapseWs(raw), InventoryBlurbMax);
|
|
}
|
|
}
|
|
|
|
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 (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+", " ");
|
|
}
|
|
|
|
/// <summary>Search Civitai for LoRAs (prefers Krea 2 base). Uses Swarm-stored civitai_api key.</summary>
|
|
public async Task<JObject> 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<string> installedNames = CollectInstalledLoraNames();
|
|
HashSet<string> 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<string> CollectInstalledLoraNames()
|
|
{
|
|
HashSet<string> 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<string> CollectInstalledLoraHashes()
|
|
{
|
|
HashSet<string> 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<string> installedNames, HashSet<string> 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<string> 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<bool>() == 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<bool>() ?? false,
|
|
};
|
|
}
|
|
|
|
List<JObject> BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null, string personaId = null)
|
|
{
|
|
List<JObject> ollamaMessages = [];
|
|
StringBuilder system = new();
|
|
if (includeBase)
|
|
{
|
|
string basePack = ReadPackFile("base_krea2");
|
|
if (!string.IsNullOrWhiteSpace(basePack))
|
|
{
|
|
system.AppendLine(basePack);
|
|
}
|
|
}
|
|
string personaPrompt = ResolvePersonaPrompt(personaId);
|
|
if (!string.IsNullOrWhiteSpace(personaPrompt))
|
|
{
|
|
system.AppendLine();
|
|
system.AppendLine($"## Persona: {personaId ?? "neutral"}");
|
|
system.AppendLine(personaPrompt);
|
|
}
|
|
if (!string.IsNullOrWhiteSpace(packName) && packName != "base_krea2")
|
|
{
|
|
string situational = ReadPackFile(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)
|
|
{
|
|
string id = (personaId ?? "neutral").Trim();
|
|
if (string.IsNullOrWhiteSpace(id))
|
|
{
|
|
id = "neutral";
|
|
}
|
|
string overlay = PersonasOverlayJsonPath();
|
|
if (File.Exists(overlay))
|
|
{
|
|
try
|
|
{
|
|
JObject parsed = JObject.Parse(File.ReadAllText(overlay, Encoding.UTF8));
|
|
if (parsed["personas"] is JArray arr)
|
|
{
|
|
foreach (JToken t in arr)
|
|
{
|
|
if (t is JObject po && string.Equals(po["id"]?.ToString(), id, StringComparison.OrdinalIgnoreCase))
|
|
{
|
|
string p = po["prompt"]?.ToString();
|
|
if (!string.IsNullOrWhiteSpace(p))
|
|
{
|
|
return p;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
catch
|
|
{
|
|
// fall through to bundled
|
|
}
|
|
}
|
|
return ReadPersonaFile(id);
|
|
}
|
|
|
|
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))
|
|
{
|
|
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;
|
|
}
|
|
if (!string.IsNullOrWhiteSpace(ExtractSearchQuery(patch)))
|
|
{
|
|
return true;
|
|
}
|
|
if (patch["actions"] is JArray acts)
|
|
{
|
|
foreach (JToken a in acts)
|
|
{
|
|
if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
|
|
{
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
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)
|
|
{
|
|
List<JObject> messages = BuildOllamaMessages(packName, includeBase, contextJson, userMessages, personaId: personaId);
|
|
JArray civitaiResults = [];
|
|
string reply = "";
|
|
JObject lastRaw = null;
|
|
for (int hop = 0; hop < MaxCivitaiHops; hop++)
|
|
{
|
|
if (onHopStart is not null)
|
|
{
|
|
await onHopStart(hop);
|
|
}
|
|
(reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta);
|
|
JObject patch = TryParsePatch(reply);
|
|
if (hop + 1 >= MaxCivitaiHops || !WantsCivitaiSearch(patch))
|
|
{
|
|
break;
|
|
}
|
|
string query = ExtractSearchQuery(patch);
|
|
if (string.IsNullOrWhiteSpace(query))
|
|
{
|
|
query = userMessages.LastOrDefault(m => m["role"]?.ToString() == "user")?["content"]?.ToString() ?? "";
|
|
}
|
|
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);
|
|
}
|
|
|
|
async Task<(string reply, JObject raw)> CallOllamaChat(
|
|
string root,
|
|
string modelName,
|
|
List<JObject> ollamaMessages,
|
|
bool stream,
|
|
Func<string, Task> onDelta)
|
|
{
|
|
JObject payload = new()
|
|
{
|
|
["model"] = modelName,
|
|
["stream"] = stream,
|
|
["messages"] = new JArray(ollamaMessages),
|
|
["options"] = new JObject
|
|
{
|
|
["num_ctx"] = DefaultNumCtx,
|
|
},
|
|
};
|
|
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<bool>() == true)
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
return (full.ToString(), last ?? new JObject());
|
|
}
|
|
|
|
/// <summary>
|
|
/// SwarmUI passes the whole request as the JObject param (not only a nested key).
|
|
/// Support both flat fields and legacy nested <c>raw</c>.
|
|
/// </summary>
|
|
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)
|
|
{
|
|
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<bool?>("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";
|
|
}
|
|
|
|
/// <summary>Proxy to Ollama /api/chat (non-stream), with optional Civitai search hop.</summary>
|
|
public async Task<JObject> 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);
|
|
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();
|
|
try
|
|
{
|
|
(string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
|
|
session, root, modelName, packName, includeBase, contextJson, userMessages, personaId: persona);
|
|
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}" };
|
|
}
|
|
}
|
|
|
|
/// <summary>WebSocket streaming chat (Ollama stream:true) + Civitai hops.</summary>
|
|
public async Task<JObject> 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);
|
|
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();
|
|
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);
|
|
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;
|
|
}
|
|
}
|