Skip to content

Commit c33b445

Browse files
ShenKSPZShenKSPZ
committed
Fix: fix image generation MCP
1 parent de0c700 commit c33b445

File tree

7 files changed

+384
-77
lines changed

7 files changed

+384
-77
lines changed

‎src/components/chat/MarkdownContent.tsx

Lines changed: 108 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import type { HTMLProps } from 'react';
1111
import { MessageContent, MessageContentType } from '../../types/chat';
1212
import { MessageHelper } from '../../services/message-helper';
1313
import FileAttachmentDisplay from './FileAttachmentDisplay';
14-
import { Loader2 } from 'lucide-react';
14+
import { Loader2, ServerCog, AlertCircle } from 'lucide-react';
1515
import { useTranslation } from '../../hooks/useTranslation';
1616

1717
interface MarkdownContentProps {
@@ -33,6 +33,7 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
3333
const [imageContents, setImageContents] = useState<MessageContent[]>([]);
3434
const [isProcessingImage, setIsProcessingImage] = useState(false);
3535
const [imageGenerationError, setImageGenerationError] = useState<string | null>(null);
36+
const [customToolCalls, setCustomToolCalls] = useState<Array<{name: string, status: string}>>([]);
3637

3738
// Process content and check for thinking blocks, files, and image generation
3839
useEffect(() => {
@@ -67,10 +68,48 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
6768

6869
let processed = MessageHelper.MessageContentToText(textContents);
6970

71+
// Track custom tool calls
72+
const toolCalls: Array<{name: string, status: string}> = [];
73+
74+
// Detect tool calls in progress
75+
const customToolCallRegex = /\bExecuting tool call:\s+([a-zA-Z0-9_-]+)/g;
76+
let match: RegExpExecArray | null;
77+
while ((match = customToolCallRegex.exec(processed)) !== null) {
78+
const toolName = match[1];
79+
if (toolName !== 'generate_image') {
80+
toolCalls.push({ name: toolName, status: 'in_progress' });
81+
}
82+
}
83+
84+
// Detect tool results
85+
const toolResultRegex = /\bTool result:\s+({.+})/g;
86+
while ((match = toolResultRegex.exec(processed)) !== null) {
87+
// Mark the last tool as completed if it exists
88+
if (toolCalls.length > 0) {
89+
const lastTool = toolCalls[toolCalls.length - 1];
90+
lastTool.status = 'completed';
91+
}
92+
}
93+
94+
// Detect tool errors
95+
const toolErrorRegex = /\bError in tool call\s+([a-zA-Z0-9_-]+):\s+(.+)/g;
96+
while ((match = toolErrorRegex.exec(processed)) !== null) {
97+
const toolName = match[1];
98+
// Check if we already have this tool
99+
const existingTool = toolCalls.find(tool => tool.name === toolName);
100+
if (existingTool) {
101+
existingTool.status = 'error';
102+
} else {
103+
toolCalls.push({ name: toolName, status: 'error' });
104+
}
105+
}
106+
107+
setCustomToolCalls(toolCalls);
108+
70109
// Detect image generation in progress
71110
const imageGenInProgressMatch = processed.match(
72111
/(?:generating|creating|processing)\s+(?:an\s+)?image(?:s)?\s+(?:with|using|for|from)?(?:\s+prompt)?(?::|;)?\s*["']?([^"']+)["']?/i
73-
) || processed.match(/\bimage\s+generation\s+in\s+progress\b/i);
112+
) || processed.match(/\bimage\s+generation\s+in\s+progress\b/i) || processed.match(/\bGenerating image\b/i);
74113

75114
if ((imageGenInProgressMatch && images.length === 0) ||
76115
(processed.includes('generate_image') && processed.includes('tool call') && images.length === 0)) {
@@ -83,7 +122,7 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
83122
// Detect image generation errors
84123
const imageGenErrorMatch = processed.match(
85124
/(?:error|failed|couldn't|unable)\s+(?:in\s+)?(?:generating|creating|processing)\s+(?:an\s+)?image(?:s)?(?::|;)?\s*["']?([^"']+)["']?/i
86-
) || processed.match(/\bimage\s+generation\s+(?:error|failed)\b:?\s*["']?([^"']+)["']?/i);
125+
) || processed.match(/\bimage\s+generation\s+(?:error|failed)\b:?\s*["']?([^"']+)["']?/i) || processed.match(/\bError generating image\b:?\s*([^"\n]+)/i);
87126

88127
if (imageGenErrorMatch || (processed.includes('error') && processed.includes('generate_image'))) {
89128
const errorMessage = imageGenErrorMatch ? (imageGenErrorMatch[1] || "Unknown error occurred") : "Failed to generate image";
@@ -93,6 +132,33 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
93132
setImageGenerationError(null);
94133
}
95134

135+
// Process the content - remove the tool call info to make the response cleaner
136+
if (processed.includes('Executing tool call:') || processed.includes('Tool result:') ||
137+
processed.includes('Generating image') || processed.includes('Error generating image') ||
138+
processed.includes('Error in tool call')) {
139+
140+
// Split by the first occurrence of a tool-related message
141+
const toolMarkers = [
142+
'Executing tool call:',
143+
'Tool result:',
144+
'Generating image',
145+
'Error generating image',
146+
'Error in tool call'
147+
];
148+
149+
let firstToolIndex = processed.length;
150+
for (const marker of toolMarkers) {
151+
const index = processed.indexOf(marker);
152+
if (index > -1 && index < firstToolIndex) {
153+
firstToolIndex = index;
154+
}
155+
}
156+
157+
if (firstToolIndex < processed.length) {
158+
processed = processed.substring(0, firstToolIndex);
159+
}
160+
}
161+
96162
// Check if content contains thinking block
97163
const thinkMatch = processed.match(/<think>([\s\S]*?)<\/think>([\s\S]*)/);
98164

@@ -172,9 +238,47 @@ export const MarkdownContent: React.FC<MarkdownContentProps> = ({ content, isUse
172238
</div>
173239
)}
174240

241+
{/* Custom Tool Calls */}
242+
{customToolCalls.length > 0 && (
243+
<div className="mb-3 space-y-2">
244+
{customToolCalls.map((toolCall, index) => (
245+
<div
246+
key={index}
247+
className={`p-3 border rounded-md flex items-center gap-2 ${
248+
toolCall.status === 'error'
249+
? 'border-red-200 bg-red-50'
250+
: toolCall.status === 'completed'
251+
? 'border-green-200 bg-green-50'
252+
: 'border-blue-200 bg-blue-50'
253+
}`}
254+
>
255+
{toolCall.status === 'in_progress' && (
256+
<Loader2 size={18} className="text-blue-500 animate-spin" />
257+
)}
258+
{toolCall.status === 'completed' && (
259+
<ServerCog size={18} className="text-green-600" />
260+
)}
261+
{toolCall.status === 'error' && (
262+
<AlertCircle size={18} className="text-red-600" />
263+
)}
264+
<div>
265+
<div className="font-medium">
266+
{toolCall.status === 'in_progress' && t('tools.executing')}
267+
{toolCall.status === 'completed' && t('tools.executedSuccessfully')}
268+
{toolCall.status === 'error' && t('tools.executionFailed')}
269+
</div>
270+
<div className="text-sm">
271+
{t('tools.toolName')}: {toolCall.name}
272+
</div>
273+
</div>
274+
</div>
275+
))}
276+
</div>
277+
)}
278+
175279
{/* Image Generation In Progress */}
176280
{isProcessingImage && (
177-
<div className="p-4 mb-3 border border-gray-200 rounded-md">
281+
<div className="p-4 mb-3 border border-blue-200 rounded-md bg-blue-50">
178282
<div className="flex items-center gap-2 mb-2">
179283
<Loader2 size={20} className="text-blue-500 animate-spin" />
180284
<span className="font-medium">{t('imageGeneration.generating')}</span>

‎src/locales/en/translation.json

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,14 @@
3939
"selectImageProvider": "Select Image Provider",
4040
"noImageProvidersAvailable": "No image providers available"
4141
},
42+
"tools": {
43+
"executing": "Executing tool...",
44+
"executedSuccessfully": "Tool execution completed",
45+
"executionFailed": "Tool execution failed",
46+
"toolName": "Tool name",
47+
"result": "Result",
48+
"error": "Error"
49+
},
4250
"translation": {
4351
"title": "Translation",
4452
"sourceLanguageWithAutoSelected": "Source Language will be auto detected",

‎src/services/ai-service.ts

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { StreamControlHandler } from './streaming-control';
55
import { SETTINGS_CHANGE_EVENT, SettingsService } from './settings-service';
66
import { MCPToolAdapter } from './mcp-tool-adapter';
77
import { MCPService } from './mcp-service';
8+
import { AIServiceCapability } from '../types/capabilities';
89

910
export interface ModelOption {
1011
id: string;
@@ -199,18 +200,34 @@ export class AIService {
199200
throw new Error(`Provider ${providerName} not available`);
200201
}
201202

202-
// If MCP tools are specified, get them and add to options
203-
if (mcpTools && mcpTools.length > 0) {
203+
// Check if the selected model supports tool calls
204+
const modelCapabilities = provider.getModelCapabilities(modelName);
205+
const supportsTools = modelCapabilities.includes(AIServiceCapability.ToolUsage) ||
206+
modelCapabilities.includes(AIServiceCapability.FunctionCalling) ||
207+
modelCapabilities.includes(AIServiceCapability.MCPServer);
208+
209+
// If MCP tools are specified and model supports tools, get them and add to options
210+
if (mcpTools && mcpTools.length > 0 && supportsTools) {
204211
const mcpToolAdapter = MCPToolAdapter.getInstance();
205-
const allTools: Record<string, any> = {};
212+
const allTools: Record<string, unknown> = {};
206213

207214
for (const mcpToolId of mcpTools) {
208-
const tools = await mcpToolAdapter.getToolsForServer(mcpToolId, streamController);
209-
Object.assign(allTools, tools);
215+
try {
216+
const tools = await mcpToolAdapter.getToolsForServer(mcpToolId, streamController);
217+
Object.assign(allTools, tools);
218+
} catch (error) {
219+
console.error(`Error loading tools for MCP server ${mcpToolId}:`, error);
220+
// Continue with other tools if one fails
221+
}
210222
}
211223

212-
// Add tools to options
213-
options.tools = allTools;
224+
// Add tools to options if any were loaded
225+
if (Object.keys(allTools).length > 0) {
226+
console.log('Adding tools to request:', Object.keys(allTools));
227+
options.tools = allTools;
228+
}
229+
} else if (mcpTools && mcpTools.length > 0 && !supportsTools) {
230+
console.warn(`Model ${modelName} does not support tools, but tools were requested`);
214231
}
215232

216233
const result = await provider.getChatCompletion(
@@ -426,7 +443,7 @@ export class AIService {
426443
/**
427444
* Get all available MCP servers
428445
*/
429-
public getMCPServers(): Record<string, any> {
446+
public getMCPServers(): Record<string, unknown> {
430447
return MCPService.getInstance().getMCPServers();
431448
}
432449
}

‎src/services/chat-service.ts

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -900,7 +900,7 @@ export class ChatService {
900900
/**
901901
* Get available MCP servers
902902
*/
903-
public getAvailableMCPServers(): Record<string, any> {
903+
public getAvailableMCPServers(): Record<string, unknown> {
904904
return MCPService.getInstance().getMCPServers();
905905
}
906906

@@ -944,6 +944,28 @@ export class ChatService {
944944
const provider = settingsService.getSelectedProvider();
945945
const model = settingsService.getSelectedModel();
946946

947+
// Check if any of the selected MCP servers is for image generation
948+
const mcpService = MCPService.getInstance();
949+
const selectedMcpServers = mcpTools.map(id => mcpService.getMCPServer(id)).filter(Boolean);
950+
const hasImageGenerationTool = selectedMcpServers.some(server => server?.isImageGeneration);
951+
952+
// Check if the message content starts with image generation command
953+
const isImageGenerationRequest = content.toLowerCase().startsWith('/image')
954+
|| content.toLowerCase().includes('generate an image')
955+
|| content.toLowerCase().includes('create an image');
956+
957+
// If user is requesting image generation but hasn't selected the image tool, add it automatically
958+
if (isImageGenerationRequest && !hasImageGenerationTool) {
959+
// Find the default image generation server
960+
const imageServer = Object.values(mcpService.getMCPServers())
961+
.find(server => server.isImageGeneration);
962+
963+
if (imageServer) {
964+
console.log('Adding image generation tool automatically');
965+
mcpTools.push(imageServer.id);
966+
}
967+
}
968+
947969
//#region Save user message to database and update title
948970
// eslint-disable-next-line prefer-const
949971
let {conversation: updatedConversation, message: userMessage} = await MessageHelper.addUserMessageToConversation(content, currentConversation);

0 commit comments

Comments
 (0)