mirror of
https://github.com/vndangkhoa/Sys-Arc-Visl.git
synced 2026-04-05 01:17:57 +07:00
543 lines
19 KiB
TypeScript
543 lines
19 KiB
TypeScript
export interface NodeMetadata {
|
|
techStack: string[];
|
|
role: string;
|
|
description: string;
|
|
}
|
|
|
|
export interface AIResponse {
|
|
success: boolean;
|
|
mermaidCode?: string;
|
|
metadata?: Record<string, NodeMetadata>; // Keyed by node label
|
|
analysis?: any; // For visual analysis results
|
|
error?: string;
|
|
}
|
|
|
|
import { webLlmService } from './webLlmService';
|
|
import { visionService } from './visionService';
|
|
|
|
const SYSTEM_PROMPT = `You are an expert System Architect AI. Your task is to transform technical descriptions into precise Mermaid.js diagrams and high-fidelity metadata.
|
|
|
|
Return ONLY a strictly valid JSON object. No Markdown headers, no preamble.
|
|
|
|
EXPECTED JSON FORMAT:
|
|
{
|
|
"mermaidCode": "flowchart TD\\n A[Client] --> B[API Server]\\n ...",
|
|
"metadata": {
|
|
"NodeLabel": {
|
|
"techStack": ["Technologies used"],
|
|
"role": "Specific role (e.g., Load Balancer, Cache, Database)",
|
|
"description": "Concise technical responsibility"
|
|
}
|
|
}
|
|
}
|
|
|
|
ARCHITECTURAL RULES:
|
|
1. Use 'flowchart TD' or 'flowchart LR'.
|
|
2. Use descriptive but concise ID/Labels (e.g., 'API', 'DB_PROD').
|
|
3. Labels must match the keys in the 'metadata' object EXACTLY.
|
|
4. If input is already Mermaid code, wrap it in the JSON 'mermaidCode' field and infer metadata.
|
|
5. Identify semantic roles: use keywords like 'Client', 'Server', 'Worker', 'Database', 'Queue' in labels.
|
|
6. Escape double quotes inside the mermaid string correctly.
|
|
|
|
DIAGRAM QUALITY RULES:
|
|
7. NEVER use HTML tags (like <br/>, <b>, etc.) in node labels. Use short, clean text only.
|
|
8. Use DIAMOND shapes {Decision} for review, approval, or decision steps (e.g., A{Approve?}).
|
|
9. Use CYLINDER shapes [(Database)] for data stores.
|
|
10. Use ROUNDED shapes (Process) for human/manual tasks.
|
|
11. For complex workflows, add step numbers as edge labels: A -->|1| B -->|2| C.
|
|
12. Include feedback loops where logical (e.g., connect "Collect Feedback" back to analysis nodes).
|
|
13. Use subgraphs/swimlanes to group related components by team, role, or domain.
|
|
14. Ensure consistent node shapes within each swimlane/subgraph.`;
|
|
|
|
|
|
/**
|
|
* Exponential backoff retry wrapper for fetch.
|
|
*/
|
|
async function fetchWithRetry(
|
|
url: string,
|
|
options: RequestInit,
|
|
retries: number = 2,
|
|
backoff: number = 1000
|
|
): Promise<Response> {
|
|
try {
|
|
const response = await fetch(url, options);
|
|
// Retry on 429 (Too Many Requests) or 5xx (Server Errors)
|
|
if (!response.ok && (response.status === 429 || response.status >= 500)) {
|
|
throw new Error(`Retryable error: ${response.status} ${response.statusText}`);
|
|
}
|
|
return response;
|
|
} catch (error) {
|
|
if (retries > 0) {
|
|
console.warn(`Fetch failed, retrying in ${backoff}ms... (${retries} retries left)`);
|
|
await new Promise(resolve => setTimeout(resolve, backoff));
|
|
return fetchWithRetry(url, options, retries - 1, backoff * 2);
|
|
}
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
async function callOnlineAI(
|
|
provider: 'openai' | 'gemini' | 'ollama-cloud',
|
|
apiKey: string,
|
|
ollamaUrl: string,
|
|
messages: any[],
|
|
customSystemPrompt?: string
|
|
): Promise<AIResponse> {
|
|
const activePrompt = customSystemPrompt || SYSTEM_PROMPT;
|
|
try {
|
|
let url = '';
|
|
let headers: Record<string, string> = { 'Content-Type': 'application/json' };
|
|
let body: any = {};
|
|
|
|
if (provider === 'openai') {
|
|
url = 'https://api.openai.com/v1/chat/completions';
|
|
headers['Authorization'] = `Bearer ${apiKey}`;
|
|
body = {
|
|
model: 'gpt-4o',
|
|
messages: [{ role: 'system', content: activePrompt }, ...messages],
|
|
response_format: { type: 'json_object' }
|
|
};
|
|
} else if (provider === 'gemini') {
|
|
// Simple Gemini API call (v1beta)
|
|
url = `https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key=${apiKey}`;
|
|
body = {
|
|
contents: [{
|
|
parts: [{
|
|
text: `${activePrompt}\n\nTask: ${messages[messages.length - 1].content}`
|
|
}]
|
|
}],
|
|
generationConfig: { responseMimeType: 'application/json' }
|
|
};
|
|
} else if (provider === 'ollama-cloud') {
|
|
url = `${ollamaUrl}/api/chat`;
|
|
body = {
|
|
model: 'llava', // Default for vision
|
|
messages: [{ role: 'system', content: activePrompt }, ...messages],
|
|
format: 'json',
|
|
stream: false
|
|
};
|
|
}
|
|
|
|
const response = await fetchWithRetry(url, {
|
|
method: 'POST',
|
|
headers,
|
|
body: JSON.stringify(body)
|
|
});
|
|
|
|
if (!response.ok) {
|
|
let errorMsg = `${provider} error: ${response.status} ${response.statusText}`;
|
|
if (response.status === 401) errorMsg = 'Invalid API Key. Please check your settings.';
|
|
if (response.status === 429) errorMsg = 'Rate limit exceeded. Please try again later.';
|
|
throw new Error(errorMsg);
|
|
}
|
|
|
|
const data = await response.json();
|
|
let content = '{}';
|
|
|
|
if (provider === 'openai') {
|
|
content = data.choices[0].message.content;
|
|
} else if (provider === 'gemini') {
|
|
content = data.candidates[0].content.parts[0].text;
|
|
} else {
|
|
content = data.message?.content || '{}';
|
|
}
|
|
|
|
const parsed = JSON.parse(content);
|
|
return {
|
|
success: true,
|
|
mermaidCode: parsed.mermaidCode,
|
|
metadata: parsed.metadata
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : 'Online model failure'
|
|
};
|
|
}
|
|
}
|
|
|
|
async function callLocalAI(
|
|
ollamaUrl: string,
|
|
model: string,
|
|
messages: any[],
|
|
customSystemPrompt?: string
|
|
): Promise<AIResponse> {
|
|
const activePrompt = customSystemPrompt || SYSTEM_PROMPT;
|
|
try {
|
|
const response = await fetchWithRetry(`${ollamaUrl}/api/chat`, {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
model: model,
|
|
messages: [
|
|
{ role: 'system', content: activePrompt },
|
|
...messages
|
|
],
|
|
format: 'json',
|
|
stream: false,
|
|
options: {
|
|
temperature: 0.2,
|
|
num_predict: 2000,
|
|
}
|
|
}),
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const errorText = await response.text();
|
|
throw new Error(`Local model error (${response.status}): ${errorText}`);
|
|
}
|
|
|
|
const data = await response.json();
|
|
const content = data.message?.content || '{}';
|
|
|
|
try {
|
|
// Strip markdown code blocks if present (```json ... ```)
|
|
let cleanContent = content.trim();
|
|
if (cleanContent.startsWith('```')) {
|
|
cleanContent = cleanContent.replace(/^```(?:json)?\s*\n?/, '').replace(/\n?```\s*$/, '');
|
|
}
|
|
|
|
const parsed = JSON.parse(cleanContent);
|
|
return {
|
|
success: true,
|
|
mermaidCode: parsed.mermaidCode,
|
|
metadata: parsed.metadata,
|
|
};
|
|
} catch (e) {
|
|
console.error('Failed to parse model response:', content);
|
|
return {
|
|
success: false,
|
|
error: 'Model returned invalid JSON format',
|
|
};
|
|
}
|
|
} catch (error) {
|
|
let msg = error instanceof Error ? error.message : 'Unknown local AI error';
|
|
if (msg.includes('Failed to fetch') || msg.includes('ECONNREFUSED')) {
|
|
msg = 'Could not connect to Ollama. Is it running locally on the specified URL?';
|
|
}
|
|
return {
|
|
success: false,
|
|
error: msg,
|
|
};
|
|
}
|
|
}
|
|
|
|
|
|
async function callBrowserAI(
|
|
messages: any[],
|
|
customSystemPrompt?: string
|
|
): Promise<AIResponse> {
|
|
const activePrompt = customSystemPrompt || SYSTEM_PROMPT;
|
|
try {
|
|
if (!webLlmService.getStatus().isReady) {
|
|
throw new Error('Browser model is not loaded. Please initialize it in Settings.');
|
|
}
|
|
|
|
// --- Vision Processing Pipeline ---
|
|
let processedMessages = [];
|
|
|
|
for (const msg of messages) {
|
|
let content = msg.content;
|
|
|
|
if (msg.images && msg.images.length > 0) {
|
|
if (!visionService.getStatus().isReady) {
|
|
throw new Error('Vision model is not loaded. Please initialize it in Settings (Browser Tab).');
|
|
}
|
|
|
|
// Analyze the first image
|
|
// Assuming msg.images[0] is base64 string
|
|
const imageDescription = await visionService.analyzeImage(msg.images[0]);
|
|
|
|
// Augment the prompt with the description
|
|
content = `${content}\n\n[VISUAL CONTEXT FROM IMAGE]:\n${imageDescription}\n\n(Use this visual description to generate the Mermaid code.)`;
|
|
}
|
|
|
|
processedMessages.push({
|
|
role: (msg.role === 'user' || msg.role === 'assistant' || msg.role === 'system') ? msg.role : 'user',
|
|
content: content
|
|
});
|
|
}
|
|
|
|
const fullMessages = [
|
|
{ role: 'system' as const, content: activePrompt },
|
|
...processedMessages
|
|
];
|
|
|
|
const generator = await webLlmService.chat(fullMessages);
|
|
let fullContent = "";
|
|
|
|
for await (const chunk of generator) {
|
|
fullContent += chunk;
|
|
}
|
|
|
|
// Parse JSON
|
|
let cleanContent = fullContent.trim();
|
|
if (cleanContent.startsWith('```')) {
|
|
cleanContent = cleanContent.replace(/^```(?:json)?\s*\n?/, '').replace(/\n?```\s*$/, '');
|
|
}
|
|
|
|
const parsed = JSON.parse(cleanContent);
|
|
return {
|
|
success: true,
|
|
mermaidCode: parsed.mermaidCode,
|
|
metadata: parsed.metadata,
|
|
analysis: parsed.analysis // Forward analysis field if present
|
|
};
|
|
|
|
} catch (error) {
|
|
return {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : 'Browser model logic failed'
|
|
};
|
|
}
|
|
}
|
|
|
|
const SYSTEM_PROMPT_SIMPLE = `You are a System Architect. Create a SIMPLE, high-level Mermaid diagram.
|
|
- Focus on the main data flow (max 5-7 nodes).
|
|
- Use simple labels (e.g., "User", "API", "DB"). NEVER use HTML tags like <br/>.
|
|
- Use diamond {Decision?} for approval/review steps.
|
|
- Minimal metadata (role only).`;
|
|
|
|
const SYSTEM_PROMPT_COMPLEX = `You are an Expert Solution Architect. Create a DETAILED, comprehensive Mermaid diagram.
|
|
- Include all subsystems, queues, workers, and external services.
|
|
- Use swimlanes (subgraphs) to group components by role/domain.
|
|
- NEVER use HTML tags like <br/> in labels. Keep labels clean and concise.
|
|
- Use diamond {Decision?} for approval/review steps.
|
|
- Use cylinder [(DB)] for databases, rounded (Task) for human tasks.
|
|
- Add step numbers as edge labels for complex flows: A -->|1| B -->|2| C.
|
|
- Include feedback loops connecting outputs back to inputs where logical.
|
|
- Detailed metadata (techStack, role, description).`;
|
|
|
|
function getSystemPrompt(complexity: 'simple' | 'complex') {
|
|
return complexity === 'simple' ? SYSTEM_PROMPT_SIMPLE : SYSTEM_PROMPT_COMPLEX;
|
|
}
|
|
|
|
export async function analyzeImage(
|
|
imageBase64: string,
|
|
ollamaUrl: string,
|
|
model: string,
|
|
aiMode: 'online' | 'offline' | 'browser' = 'offline',
|
|
onlineProvider?: 'openai' | 'gemini' | 'ollama-cloud' | 'browser',
|
|
apiKey?: string,
|
|
complexity: 'simple' | 'complex' = 'simple'
|
|
): Promise<AIResponse> {
|
|
const rawBase64 = imageBase64.includes(',') ? imageBase64.split(',')[1] : imageBase64;
|
|
const prompt = getSystemPrompt(complexity) + "\n" + SYSTEM_PROMPT.split('\n').slice(16).join('\n'); // Append JSON rules
|
|
|
|
const messages = [{
|
|
role: 'user',
|
|
content: 'Analyze this system design diagram. Return MermaidJS and technical metadata.',
|
|
images: [rawBase64]
|
|
}];
|
|
|
|
if (aiMode === 'online' && onlineProvider) {
|
|
return callOnlineAI(onlineProvider as any, apiKey || '', ollamaUrl, messages, prompt);
|
|
}
|
|
if (aiMode === 'browser') {
|
|
return callBrowserAI(messages, prompt);
|
|
}
|
|
return callLocalAI(ollamaUrl, model, messages, prompt);
|
|
}
|
|
|
|
export async function interpretText(
|
|
text: string,
|
|
ollamaUrl: string,
|
|
model: string,
|
|
aiMode: 'online' | 'offline' | 'browser' = 'offline',
|
|
onlineProvider?: 'openai' | 'gemini' | 'ollama-cloud' | 'browser',
|
|
apiKey?: string,
|
|
complexity: 'simple' | 'complex' = 'simple'
|
|
): Promise<AIResponse> {
|
|
const prompt = getSystemPrompt(complexity) + "\n" + SYSTEM_PROMPT.split('\n').slice(16).join('\n'); // Append JSON rules
|
|
|
|
const messages = [{
|
|
role: 'user',
|
|
content: `Create a system design based on this description: ${text}`,
|
|
}];
|
|
|
|
if (aiMode === 'online' && onlineProvider) {
|
|
return callOnlineAI(onlineProvider as any, apiKey || '', ollamaUrl, messages, prompt);
|
|
}
|
|
if (aiMode === 'browser') {
|
|
return callBrowserAI(messages, prompt);
|
|
}
|
|
return callLocalAI(ollamaUrl, model, messages, prompt);
|
|
}
|
|
|
|
export async function analyzeSVG(
|
|
svgContent: string,
|
|
ollamaUrl: string,
|
|
model: string,
|
|
aiMode: 'online' | 'offline' | 'browser' = 'offline',
|
|
onlineProvider?: 'openai' | 'gemini' | 'ollama-cloud' | 'browser',
|
|
apiKey?: string
|
|
): Promise<AIResponse> {
|
|
const messages = [{
|
|
role: 'user',
|
|
content: `Analyze this SVG architecture: ${svgContent}`,
|
|
}];
|
|
|
|
if (aiMode === 'online' && onlineProvider) {
|
|
return callOnlineAI(onlineProvider as any, apiKey || '', ollamaUrl, messages);
|
|
}
|
|
if (aiMode === 'browser') {
|
|
return callBrowserAI(messages);
|
|
}
|
|
return callLocalAI(ollamaUrl, model, messages);
|
|
}
|
|
|
|
const SUGGEST_PROMPT = `You are a Mermaid.js syntax and logic expert.
|
|
Your task is to analyze the provided Mermaid flowchart code and either:
|
|
1. Fix any syntax errors that prevent it from rendering.
|
|
2. Improve the logical flow or visual clarity if the syntax is already correct.
|
|
|
|
Return ONLY a strictly valid JSON object. No Markdown headers, no preamble.
|
|
|
|
EXPECTED JSON FORMAT:
|
|
{
|
|
"mermaidCode": "flowchart TD\\n A[Fixed/Improved] --> B[Nodes]",
|
|
"explanation": "Briefly explain what was fixed or improved."
|
|
}
|
|
|
|
RULES:
|
|
- Maintain the original intent of the diagram.
|
|
- Use best practices for Mermaid layout and labeling.
|
|
- If the code is already perfect, return it as is but provide a positive explanation.`;
|
|
|
|
const VISUAL_ANALYSIS_PROMPT = `You are a Visualization and UX Expert specialized in node-graph diagrams.
|
|
Your task is to analyze the provided graph structure and metrics to suggest specific improvements for layout, grouping, and visual clarity.
|
|
|
|
Return ONLY a strictly valid JSON object. Do not include markdown formatting like \`\`\`json.
|
|
|
|
EXPECTED JSON FORMAT:
|
|
{
|
|
"analysis": {
|
|
"suggestions": [
|
|
{
|
|
"id": "unique-id",
|
|
"title": "Short title",
|
|
"description": "Detailed explanation",
|
|
"type": "spacing" | "grouping" | "routing" | "hierarchy" | "style",
|
|
"impact": "high" | "medium" | "low",
|
|
"fix_strategy": "algorithmic_spacing" | "algorithmic_routing" | "group_nodes" | "unknown"
|
|
}
|
|
],
|
|
"summary": {
|
|
"critique": "Overall analysis",
|
|
"score": 0-100
|
|
}
|
|
}
|
|
}
|
|
|
|
EXAMPLE RESPONSE:
|
|
{
|
|
"analysis": {
|
|
"suggestions": [
|
|
{
|
|
"id": "sug-1",
|
|
"title": "Group Database Nodes",
|
|
"description": "Several database nodes are scattered. Group them for better logical separation.",
|
|
"type": "grouping",
|
|
"impact": "high",
|
|
"fix_strategy": "group_nodes"
|
|
}
|
|
],
|
|
"summary": {
|
|
"critique": "The flow is generally good but lacks logical grouping for backend services.",
|
|
"score": 75
|
|
}
|
|
}
|
|
}
|
|
|
|
RULES:
|
|
- Focus on readability, flow, and logical grouping.
|
|
- Identify if nodes are too cluttered or if the flow is confusing.
|
|
- Suggest grouping for nodes that appear related based on their labels.`;
|
|
|
|
export async function suggestFix(
|
|
code: string,
|
|
ollamaUrl: string,
|
|
modelName: string,
|
|
aiMode: 'online' | 'offline' | 'browser',
|
|
onlineProvider: 'openai' | 'gemini' | 'ollama-cloud' | 'browser',
|
|
apiKey: string
|
|
): Promise<{ success: boolean; mermaidCode?: string; explanation?: string; error?: string }> {
|
|
const messages = [{
|
|
role: 'user',
|
|
content: `CURRENT MERMAID CODE:\n${code}\n\nPlease analyze and provide a fix or improvement.`,
|
|
}];
|
|
|
|
try {
|
|
let response: AIResponse;
|
|
if (aiMode === 'online' && onlineProvider) {
|
|
response = await callOnlineAI(onlineProvider as any, apiKey || '', ollamaUrl, messages, SUGGEST_PROMPT);
|
|
} else if (aiMode === 'browser') {
|
|
response = await callBrowserAI(messages, SUGGEST_PROMPT);
|
|
} else {
|
|
response = await callLocalAI(ollamaUrl, modelName, messages, SUGGEST_PROMPT);
|
|
}
|
|
|
|
if (!response.success) throw new Error(response.error);
|
|
|
|
return {
|
|
success: true,
|
|
mermaidCode: response.mermaidCode,
|
|
explanation: (response as any).explanation || 'Code improved.'
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : 'Fix suggestion failed'
|
|
};
|
|
}
|
|
}
|
|
|
|
export async function analyzeVisualLayout(
|
|
nodes: any[],
|
|
edges: any[],
|
|
metrics: any,
|
|
ollamaUrl: string,
|
|
modelName: string,
|
|
aiMode: 'online' | 'offline' | 'browser',
|
|
onlineProvider: 'openai' | 'gemini' | 'ollama-cloud' | 'browser',
|
|
apiKey: string
|
|
): Promise<{ success: boolean; analysis?: any; error?: string }> {
|
|
const context = {
|
|
nodeCount: nodes.length,
|
|
edgeCount: edges.length,
|
|
nodeLabels: nodes.map(n => n.data?.label || n.id),
|
|
metrics: metrics
|
|
};
|
|
|
|
const messages = [{
|
|
role: 'user',
|
|
content: `ANALYZE THIS DIAGRAM LAYOUT:\n${JSON.stringify(context, null, 2)}\n\nProvide visual improvement suggestions.`,
|
|
}];
|
|
|
|
try {
|
|
let response: AIResponse;
|
|
if (aiMode === 'online' && onlineProvider) {
|
|
response = await callOnlineAI(onlineProvider as any, apiKey || '', ollamaUrl, messages, VISUAL_ANALYSIS_PROMPT);
|
|
} else if (aiMode === 'browser') {
|
|
response = await callBrowserAI(messages, VISUAL_ANALYSIS_PROMPT);
|
|
} else {
|
|
response = await callLocalAI(ollamaUrl, modelName, messages, VISUAL_ANALYSIS_PROMPT);
|
|
}
|
|
|
|
if (!response.success) throw new Error(response.error);
|
|
|
|
// The AI response parsing logic in callLocalAI/callOnlineAI assigns unknown JSON fields to the object
|
|
// We expect 'analysis' field in the JSON
|
|
const analysis = (response as any).analysis || (response as any).visualAnalysis;
|
|
|
|
return {
|
|
success: true,
|
|
analysis: analysis
|
|
};
|
|
} catch (error) {
|
|
return {
|
|
success: false,
|
|
error: error instanceof Error ? error.message : 'Visual analysis failed'
|
|
};
|
|
}
|
|
}
|