Expand Assistent to v0.3: Generate loop, Civitai Confirm, img2img/inpaint.
Server inventory and streaming chat, auto-apply/generate with Interrupt, Civitai search cards (Confirm-only download), plus Init/Mask wiring and inpaint_edit pack. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Leonid Pershin
Cursor
+553
-30
@@ -3,18 +3,22 @@ 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.</summary>
|
||||
/// <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(
|
||||
@@ -33,17 +37,24 @@ public class SwarmAssistentExtension : Extension
|
||||
"critique_image",
|
||||
"compose_scene",
|
||||
"fix_params",
|
||||
"inpaint_edit",
|
||||
];
|
||||
|
||||
const int MaxCivitaiHops = 2;
|
||||
const int MaxLorasInInventory = 120;
|
||||
const int MaxWildcardsInInventory = 80;
|
||||
|
||||
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 via Ollama: chat, vision, prompts, LoRA triggers, size patches.";
|
||||
Description = "Collaborative Krea 2 assistant via Ollama: chat, vision, img2img/inpaint, Generate loop, Civitai Confirm.";
|
||||
License = "MIT";
|
||||
Version = "0.2.0";
|
||||
Tags = ["tabs", "ui", "llm", "ollama", "krea"];
|
||||
Version = "0.3.1";
|
||||
Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint"];
|
||||
}
|
||||
|
||||
public override void OnInit()
|
||||
@@ -51,8 +62,11 @@ public class SwarmAssistentExtension : Extension
|
||||
HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
|
||||
API.RegisterAPICall(AssistentListModels, false, PermUse);
|
||||
API.RegisterAPICall(AssistentGetPacks, false, PermUse);
|
||||
API.RegisterAPICall(AssistentListInventory, false, PermUse);
|
||||
API.RegisterAPICall(AssistentSearchCivitai, false, PermUse);
|
||||
API.RegisterAPICall(AssistentChat, true, PermUse);
|
||||
Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs)");
|
||||
API.RegisterAPICall(AssistentChatWS, true, PermUse);
|
||||
Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs + inventory/Civitai)");
|
||||
}
|
||||
|
||||
static string Clip(string text, int max)
|
||||
@@ -128,24 +142,231 @@ public class SwarmAssistentExtension : Extension
|
||||
return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = new JArray(PackNames) };
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Proxy to Ollama /api/chat (non-stream).
|
||||
/// <paramref name="raw"/> must include messages (JArray) and optional context_json.
|
||||
/// </summary>
|
||||
public async Task<JObject> AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
|
||||
/// <summary>Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).</summary>
|
||||
public async Task<JObject> AssistentListInventory(Session session)
|
||||
{
|
||||
string root = NormalizeBaseUrl(baseUrl ?? raw?["base_url"]?.ToString());
|
||||
string modelName = (model ?? raw?["model"]?.ToString() ?? "").Trim();
|
||||
if (string.IsNullOrWhiteSpace(modelName))
|
||||
await Task.CompletedTask;
|
||||
JArray loras = [];
|
||||
JArray checkpoints = [];
|
||||
JArray wildcards = [];
|
||||
|
||||
if (Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler loraHandler))
|
||||
{
|
||||
return new JObject { ["error"] = "model is required" };
|
||||
}
|
||||
JArray userMessages = raw?["messages"] as JArray;
|
||||
if (userMessages is null || userMessages.Count == 0)
|
||||
{
|
||||
return new JObject { ["error"] = "messages required" };
|
||||
foreach (T2IModel model in loraHandler.Models.Values.OrderBy(m => m.Name).Take(MaxLorasInInventory))
|
||||
{
|
||||
loras.Add(new JObject
|
||||
{
|
||||
["name"] = model.Name,
|
||||
["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
|
||||
["trigger_phrase"] = model.Metadata?.TriggerPhrase,
|
||||
["architecture"] = model.ModelClass?.ID,
|
||||
["compat_class"] = model.ModelClass?.CompatClass?.ID,
|
||||
["hash"] = model.Metadata?.Hash ?? "",
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (Program.T2IModelSets.TryGetValue("Stable-Diffusion", out T2IModelHandler ckptHandler))
|
||||
{
|
||||
foreach (T2IModel model in ckptHandler.Models.Values.OrderBy(m => m.Name).Take(60))
|
||||
{
|
||||
checkpoints.Add(new JObject
|
||||
{
|
||||
["name"] = model.Name,
|
||||
["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
|
||||
["architecture"] = model.ModelClass?.ID,
|
||||
["compat_class"] = model.ModelClass?.CompatClass?.ID,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
};
|
||||
}
|
||||
|
||||
/// <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)
|
||||
{
|
||||
List<JObject> ollamaMessages = [];
|
||||
StringBuilder system = new();
|
||||
if (includeBase)
|
||||
@@ -156,7 +377,6 @@ public class SwarmAssistentExtension : Extension
|
||||
system.AppendLine(basePack);
|
||||
}
|
||||
}
|
||||
string packName = (pack ?? raw?["pack"]?.ToString() ?? "write_prompt").Trim();
|
||||
if (!string.IsNullOrWhiteSpace(packName) && packName != "base_krea2")
|
||||
{
|
||||
string situational = ReadPackFile(packName);
|
||||
@@ -167,7 +387,6 @@ public class SwarmAssistentExtension : Extension
|
||||
system.AppendLine(situational);
|
||||
}
|
||||
}
|
||||
string contextJson = raw?["context_json"]?.ToString();
|
||||
if (!string.IsNullOrWhiteSpace(contextJson))
|
||||
{
|
||||
system.AppendLine();
|
||||
@@ -176,6 +395,11 @@ public class SwarmAssistentExtension : Extension
|
||||
system.AppendLine(contextJson);
|
||||
system.AppendLine("```");
|
||||
}
|
||||
if (!string.IsNullOrWhiteSpace(extraSystem))
|
||||
{
|
||||
system.AppendLine();
|
||||
system.AppendLine(extraSystem);
|
||||
}
|
||||
if (system.Length > 0)
|
||||
{
|
||||
ollamaMessages.Add(new JObject
|
||||
@@ -184,7 +408,7 @@ public class SwarmAssistentExtension : Extension
|
||||
["content"] = system.ToString(),
|
||||
});
|
||||
}
|
||||
foreach (JToken msg in userMessages)
|
||||
foreach (JToken msg in userMessages ?? [])
|
||||
{
|
||||
if (msg is not JObject mo)
|
||||
{
|
||||
@@ -201,24 +425,264 @@ public class SwarmAssistentExtension : Extension
|
||||
}
|
||||
ollamaMessages.Add(copy);
|
||||
}
|
||||
return ollamaMessages;
|
||||
}
|
||||
|
||||
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))
|
||||
{
|
||||
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)
|
||||
{
|
||||
List<JObject> messages = BuildOllamaMessages(packName, includeBase, contextJson, userMessages);
|
||||
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"] = false,
|
||||
["stream"] = stream,
|
||||
["messages"] = new JArray(ollamaMessages),
|
||||
};
|
||||
try
|
||||
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)
|
||||
{
|
||||
using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
|
||||
using HttpResponseMessage resp = await HttpClient.PostAsync($"{root}/api/chat", content);
|
||||
string body = await resp.Content.ReadAsStringAsync();
|
||||
if (!resp.IsSuccessStatusCode)
|
||||
{
|
||||
return new JObject { ["error"] = $"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(body, 800)}" };
|
||||
}
|
||||
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)
|
||||
{
|
||||
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();
|
||||
}
|
||||
|
||||
/// <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);
|
||||
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);
|
||||
return new JObject
|
||||
{
|
||||
["success"] = true,
|
||||
@@ -226,6 +690,7 @@ public class SwarmAssistentExtension : Extension
|
||||
["model"] = modelName,
|
||||
["pack"] = packName,
|
||||
["raw"] = parsed,
|
||||
["civitai_results"] = civitai,
|
||||
};
|
||||
}
|
||||
catch (Exception ex)
|
||||
@@ -233,4 +698,62 @@ public class SwarmAssistentExtension : Extension
|
||||
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);
|
||||
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
|
||||
{
|
||||
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);
|
||||
await ws.SendJson(new JObject
|
||||
{
|
||||
["success"] = true,
|
||||
["done"] = true,
|
||||
["reply"] = reply,
|
||||
["model"] = modelName,
|
||||
["pack"] = packName,
|
||||
["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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user