Ship Assistent 0.8.1: split modules and shared+personal vector memory.
Personal RAG never leaks into the shared store; retrieve merges shared plus the persona chain, with personal overwrite on kind+key. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Leonid Pershin
Cursor
@@ -0,0 +1,328 @@
|
||||
using System;
|
||||
using System.Collections.Generic;
|
||||
using System.IO;
|
||||
using System.Linq;
|
||||
using System.Net.Http;
|
||||
using System.Net.WebSockets;
|
||||
using System.Text;
|
||||
using System.Threading.Tasks;
|
||||
using Newtonsoft.Json.Linq;
|
||||
using SwarmUI.Accounts;
|
||||
using SwarmUI.Utils;
|
||||
using SwarmUI.WebAPI;
|
||||
|
||||
namespace Mrleo1nid.SwarmAssistent;
|
||||
|
||||
/// <summary>Ollama transport: model listing, /api/chat calls and the two chat API endpoints.</summary>
|
||||
public partial class SwarmAssistentExtension
|
||||
{
|
||||
const int DefaultNumCtxFallback = 16384;
|
||||
|
||||
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 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<string> chatSet = new(
|
||||
(roles["chat"] as JArray)?.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)) ?? [],
|
||||
StringComparer.OrdinalIgnoreCase);
|
||||
HashSet<string> 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-");
|
||||
}
|
||||
|
||||
async Task<(string reply, JObject raw)> CallOllamaChat(
|
||||
string root,
|
||||
string modelName,
|
||||
List<JObject> ollamaMessages,
|
||||
bool stream,
|
||||
Func<string, Task> onDelta,
|
||||
string personaId = null)
|
||||
{
|
||||
int numCtx = Config.LoadAssistant(AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId())["num_ctx"]?.Value<int?>()
|
||||
?? 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<bool>() == true)
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
return (full.ToString(), last ?? new JObject());
|
||||
}
|
||||
|
||||
/// <summary>SwarmUI hands the whole request body over as the JObject param, so every field is read flat off it.</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, out JArray skills)
|
||||
{
|
||||
JObject whole = raw ?? [];
|
||||
if (string.IsNullOrWhiteSpace(baseUrl))
|
||||
{
|
||||
baseUrl = whole["base_url"]?.ToString() ?? whole["baseUrl"]?.ToString();
|
||||
}
|
||||
if (string.IsNullOrWhiteSpace(model))
|
||||
{
|
||||
model = whole["model"]?.ToString();
|
||||
}
|
||||
if (string.IsNullOrWhiteSpace(pack))
|
||||
{
|
||||
pack = whole["pack"]?.ToString();
|
||||
}
|
||||
if (whole["includeBase"] is not null)
|
||||
{
|
||||
includeBase = whole.Value<bool?>("includeBase") ?? includeBase;
|
||||
}
|
||||
userMessages = whole["messages"] as JArray;
|
||||
contextJson = whole["context_json"]?.ToString();
|
||||
persona = whole["persona"]?.ToString() ?? "neutral";
|
||||
skills = whole["skills"] as JArray;
|
||||
}
|
||||
|
||||
/// <summary>Shared validation for both chat endpoints. Returns an error message, or null when the request is usable.</summary>
|
||||
static string ValidateChatRequest(string modelName, JArray userMessages)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(modelName))
|
||||
{
|
||||
return "model is required";
|
||||
}
|
||||
if (userMessages is null || userMessages.Count == 0)
|
||||
{
|
||||
return "messages required";
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/// <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, out JArray skills);
|
||||
string root = NormalizeBaseUrl(baseUrl);
|
||||
string modelName = (model ?? "").Trim();
|
||||
string invalid = ValidateChatRequest(modelName, userMessages);
|
||||
if (invalid is not null)
|
||||
{
|
||||
return new JObject { ["error"] = invalid };
|
||||
}
|
||||
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}" };
|
||||
}
|
||||
}
|
||||
|
||||
/// <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, out JArray skills);
|
||||
string root = NormalizeBaseUrl(baseUrl);
|
||||
string modelName = (model ?? "").Trim();
|
||||
string invalid = ValidateChatRequest(modelName, userMessages);
|
||||
if (invalid is not null)
|
||||
{
|
||||
await ws.SendJson(new JObject { ["error"] = invalid }, 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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user