Skip to content

Commit a8674d2

Browse files
lalalunelalalune
committed
fix responses with localai (and prolly elsewhere in gui)
1 parent 777ee94 commit a8674d2

File tree

3 files changed

+272
-21
lines changed

3 files changed

+272
-21
lines changed

‎packages/cli/src/server/api/agent.ts

+18-11
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,24 @@
1-
import fs from "node:fs";
2-
import path from "node:path";
3-
import { Readable } from "node:stream";
41
import type {
52
Agent,
63
Character,
74
Content,
85
IAgentRuntime,
9-
Media,
106
Memory,
11-
UUID,
7+
UUID
128
} from "@elizaos/core";
139
import {
1410
ChannelType,
1511
ModelType,
1612
composePrompt,
13+
composePromptFromState,
1714
createUniqueUuid,
15+
logger,
1816
messageHandlerTemplate,
19-
parseJSONObjectFromText,
2017
validateUuid
2118
} from "@elizaos/core";
22-
import { logger} from "@elizaos/core";
2319
import express from "express";
20+
import fs from "node:fs";
21+
import { Readable } from "node:stream";
2422
import type { AgentServer } from "..";
2523
import { upload } from "../loader";
2624

@@ -1473,18 +1471,27 @@ export function agentRouter(
14731471
createdAt: Date.now(),
14741472
};
14751473

1476-
let state = await runtime.composeState(userMessage);
1474+
// save message
1475+
await runtime.createMemory(memory, "messages");
1476+
1477+
console.log("*** memory", memory);
14771478

1478-
const prompt = composePrompt({
1479+
let state = await runtime.composeState(memory);
1480+
1481+
console.log("*** state", state);
1482+
1483+
const prompt = composePromptFromState({
14791484
state,
14801485
template: messageHandlerTemplate,
14811486
});
14821487

1483-
const responseText = await runtime.useModel(ModelType.TEXT_LARGE, {
1488+
console.log("*** prompt", prompt);
1489+
1490+
const response = await runtime.useModel(ModelType.OBJECT_LARGE, {
14841491
prompt,
14851492
});
14861493

1487-
const response = parseJSONObjectFromText(responseText) as Content;
1494+
console.log("*** response", response);
14881495

14891496
if (!response) {
14901497
res.status(500).json({

‎packages/core/src/actions/reply.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ Your response should include the valid JSON block and nothing else.`;
4949
*/
5050
export const replyAction = {
5151
name: "REPLY",
52-
similes: ["GREET", "REPLY_TO_MESSAGE", "SEND_REPLY", "RESPOND"],
52+
similes: ["GREET", "REPLY_TO_MESSAGE", "SEND_REPLY", "RESPOND", "RESPONSE"],
5353
description:
5454
"Replies to the current conversation with the text from the generated message. Default if the agent is responding with a message and no other action. Use REPLY at the beginning of a chain of actions as an acknowledgement, and at the end of a chain of actions as a final response.",
5555
validate: async (_runtime: IAgentRuntime) => {

‎packages/plugin-local-ai/src/index.ts

+253-9
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import fs from "node:fs";
22
import path from "node:path";
33
import { Readable } from "node:stream";
44
import { fileURLToPath } from "node:url";
5-
import type { GenerateTextParams, ModelTypeName, TextEmbeddingParams } from "@elizaos/core";
5+
import type { GenerateTextParams, ModelTypeName, TextEmbeddingParams, ObjectGenerationParams } from "@elizaos/core";
66
import {
77
type IAgentRuntime,
88
ModelType,
@@ -370,19 +370,18 @@ class LocalAIManager {
370370
*
371371
* @returns A Promise that resolves to a boolean indicating whether the model download was successful.
372372
*/
373-
private async downloadModel(): Promise<boolean> {
373+
private async downloadModel(modelType: ModelTypeName): Promise<boolean> {
374+
const modelSpec = modelType === ModelType.TEXT_LARGE ? MODEL_SPECS.medium : MODEL_SPECS.small;
375+
const modelPath = modelType === ModelType.TEXT_LARGE ? this.mediumModelPath : this.modelPath;
374376
try {
375-
// Determine which model to download based on current modelPath
376-
const isLargeModel = this.modelPath === this.mediumModelPath;
377-
const modelSpec = isLargeModel ? MODEL_SPECS.medium : MODEL_SPECS.small;
378377
return await this.downloadManager.downloadModel(
379378
modelSpec,
380-
this.modelPath,
379+
modelPath,
381380
);
382381
} catch (error) {
383382
logger.error("Model download failed:", {
384383
error: error instanceof Error ? error.message : String(error),
385-
modelPath: this.modelPath,
384+
modelPath,
386385
});
387386
throw error;
388387
}
@@ -865,7 +864,7 @@ class LocalAIManager {
865864
await this.checkPlatformCapabilities();
866865

867866
// Download model if needed
868-
await this.downloadModel();
867+
await this.downloadModel(ModelType.TEXT_SMALL);
869868

870869
// Initialize Llama and small model
871870
try {
@@ -912,6 +911,8 @@ class LocalAIManager {
912911
await this.lazyInitSmallModel();
913912
}
914913

914+
await this.downloadModel(ModelType.TEXT_LARGE);
915+
915916
// Initialize medium model
916917
try {
917918
const mediumModel = await this.llama!.loadModel(
@@ -1166,7 +1167,7 @@ export const localAIPlugin: Plugin = {
11661167
_runtime: IAgentRuntime,
11671168
params: TextEmbeddingParams
11681169
) => {
1169-
const { text } = params;
1170+
const text = params?.text;
11701171
try {
11711172
// Add detailed logging of the input text and its structure
11721173
logger.info("TEXT_EMBEDDING handler - Initial input:", {
@@ -1206,6 +1207,249 @@ export const localAIPlugin: Plugin = {
12061207
}
12071208
},
12081209

1210+
[ModelType.OBJECT_SMALL]: async (
1211+
runtime: IAgentRuntime,
1212+
params: ObjectGenerationParams
1213+
) => {
1214+
try {
1215+
logger.info("OBJECT_SMALL handler - Processing request:", {
1216+
prompt: params.prompt,
1217+
hasSchema: !!params.schema,
1218+
temperature: params.temperature,
1219+
});
1220+
1221+
// Enhance the prompt to request JSON output
1222+
let jsonPrompt = params.prompt;
1223+
if (!jsonPrompt.includes("```json") && !jsonPrompt.includes("respond with valid JSON")) {
1224+
jsonPrompt += "\nPlease respond with valid JSON only, without any explanations, markdown formatting, or additional text.";
1225+
}
1226+
1227+
const modelConfig = localAIManager.getTextModelSource();
1228+
1229+
// Generate text based on the configured model source
1230+
let textResponse: string;
1231+
if (modelConfig.source !== "local") {
1232+
textResponse = await localAIManager.generateTextOllamaStudio({
1233+
prompt: jsonPrompt,
1234+
stopSequences: params.stopSequences,
1235+
runtime,
1236+
modelType: ModelType.TEXT_SMALL,
1237+
});
1238+
} else {
1239+
textResponse = await localAIManager.generateText({
1240+
prompt: jsonPrompt,
1241+
stopSequences: params.stopSequences,
1242+
runtime,
1243+
modelType: ModelType.TEXT_SMALL,
1244+
});
1245+
}
1246+
1247+
// Extract and parse JSON from the text response
1248+
try {
1249+
// Function to extract JSON content from text
1250+
const extractJSON = (text: string): string => {
1251+
// Try to find content between JSON codeblocks or markdown blocks
1252+
const jsonBlockRegex = /```(?:json)?\s*([\s\S]*?)\s*```/;
1253+
const match = text.match(jsonBlockRegex);
1254+
1255+
if (match && match[1]) {
1256+
return match[1].trim();
1257+
}
1258+
1259+
// If no code blocks, try to find JSON-like content
1260+
// This regex looks for content that starts with { and ends with }
1261+
const jsonContentRegex = /\s*(\{[\s\S]*\})\s*$/;
1262+
const contentMatch = text.match(jsonContentRegex);
1263+
1264+
if (contentMatch && contentMatch[1]) {
1265+
return contentMatch[1].trim();
1266+
}
1267+
1268+
// If no JSON-like content found, return the original text
1269+
return text.trim();
1270+
};
1271+
1272+
const extractedJsonText = extractJSON(textResponse);
1273+
logger.debug("Extracted JSON text:", extractedJsonText);
1274+
1275+
let jsonObject;
1276+
try {
1277+
jsonObject = JSON.parse(extractedJsonText);
1278+
} catch (parseError) {
1279+
// Try fixing common JSON issues
1280+
logger.debug("Initial JSON parse failed, attempting to fix common issues");
1281+
1282+
// Replace any unescaped newlines in string values
1283+
const fixedJson = extractedJsonText
1284+
.replace(/:\s*"([^"]*)(?:\n)([^"]*)"/g, ': "$1\\n$2"')
1285+
// Remove any non-JSON text that might have gotten mixed into string values
1286+
.replace(/"([^"]*?)[^a-zA-Z0-9\s\.,;:\-_\(\)"'\[\]{}]([^"]*?)"/g, '"$1$2"')
1287+
// Fix missing quotes around property names
1288+
.replace(/(\s*)(\w+)(\s*):/g, '$1"$2"$3:')
1289+
// Fix trailing commas in arrays and objects
1290+
.replace(/,(\s*[\]}])/g, '$1');
1291+
1292+
try {
1293+
jsonObject = JSON.parse(fixedJson);
1294+
} catch (finalError) {
1295+
logger.error("Failed to parse JSON after fixing:", finalError);
1296+
throw new Error("Invalid JSON returned from model");
1297+
}
1298+
}
1299+
1300+
// Validate against schema if provided
1301+
if (params.schema) {
1302+
try {
1303+
// Simplistic schema validation - check if all required properties exist
1304+
for (const key of Object.keys(params.schema)) {
1305+
if (!(key in jsonObject)) {
1306+
jsonObject[key] = null; // Add missing properties with null value
1307+
}
1308+
}
1309+
} catch (schemaError) {
1310+
logger.error("Schema validation failed:", schemaError);
1311+
}
1312+
}
1313+
1314+
return jsonObject;
1315+
} catch (parseError) {
1316+
logger.error("Failed to parse JSON:", parseError);
1317+
logger.error("Raw response:", textResponse);
1318+
throw new Error("Invalid JSON returned from model");
1319+
}
1320+
} catch (error) {
1321+
logger.error("Error in OBJECT_SMALL handler:", error);
1322+
throw error;
1323+
}
1324+
},
1325+
1326+
[ModelType.OBJECT_LARGE]: async (
1327+
runtime: IAgentRuntime,
1328+
params: ObjectGenerationParams
1329+
) => {
1330+
try {
1331+
logger.info("OBJECT_LARGE handler - Processing request:", {
1332+
prompt: params.prompt,
1333+
hasSchema: !!params.schema,
1334+
temperature: params.temperature,
1335+
});
1336+
1337+
// Enhance the prompt to request JSON output
1338+
let jsonPrompt = params.prompt;
1339+
if (!jsonPrompt.includes("```json") && !jsonPrompt.includes("respond with valid JSON")) {
1340+
jsonPrompt += "\nPlease respond with valid JSON only, without any explanations, markdown formatting, or additional text.";
1341+
}
1342+
1343+
const modelConfig = localAIManager.getTextModelSource();
1344+
1345+
// Generate text based on the configured model source
1346+
let textResponse: string;
1347+
if (modelConfig.source !== "local") {
1348+
textResponse = await localAIManager.generateTextOllamaStudio({
1349+
prompt: jsonPrompt,
1350+
stopSequences: params.stopSequences,
1351+
runtime,
1352+
modelType: ModelType.TEXT_LARGE,
1353+
});
1354+
} else {
1355+
textResponse = await localAIManager.generateText({
1356+
prompt: jsonPrompt,
1357+
stopSequences: params.stopSequences,
1358+
runtime,
1359+
modelType: ModelType.TEXT_LARGE,
1360+
});
1361+
}
1362+
1363+
// Extract and parse JSON from the text response
1364+
try {
1365+
// Function to extract JSON content from text
1366+
const extractJSON = (text: string): string => {
1367+
// Try to find content between JSON codeblocks or markdown blocks
1368+
const jsonBlockRegex = /```(?:json)?\s*([\s\S]*?)\s*```/;
1369+
const match = text.match(jsonBlockRegex);
1370+
1371+
if (match && match[1]) {
1372+
return match[1].trim();
1373+
}
1374+
1375+
// If no code blocks, try to find JSON-like content
1376+
// This regex looks for content that starts with { and ends with }
1377+
const jsonContentRegex = /\s*(\{[\s\S]*\})\s*$/;
1378+
const contentMatch = text.match(jsonContentRegex);
1379+
1380+
if (contentMatch && contentMatch[1]) {
1381+
return contentMatch[1].trim();
1382+
}
1383+
1384+
// If no JSON-like content found, return the original text
1385+
return text.trim();
1386+
};
1387+
1388+
// Clean up the extracted JSON to handle common formatting issues
1389+
const cleanupJSON = (jsonText: string): string => {
1390+
// Remove common logging/debugging patterns that might get mixed into the JSON
1391+
return jsonText
1392+
// Remove any lines that look like log statements
1393+
.replace(/\[DEBUG\].*?(\n|$)/g, '\n')
1394+
.replace(/\[LOG\].*?(\n|$)/g, '\n')
1395+
.replace(/console\.log.*?(\n|$)/g, '\n');
1396+
};
1397+
1398+
const extractedJsonText = extractJSON(textResponse);
1399+
const cleanedJsonText = cleanupJSON(extractedJsonText);
1400+
logger.debug("Extracted JSON text:", cleanedJsonText);
1401+
1402+
let jsonObject;
1403+
try {
1404+
jsonObject = JSON.parse(cleanedJsonText);
1405+
} catch (parseError) {
1406+
// Try fixing common JSON issues
1407+
logger.debug("Initial JSON parse failed, attempting to fix common issues");
1408+
1409+
// Replace any unescaped newlines in string values
1410+
const fixedJson = cleanedJsonText
1411+
.replace(/:\s*"([^"]*)(?:\n)([^"]*)"/g, ': "$1\\n$2"')
1412+
// Remove any non-JSON text that might have gotten mixed into string values
1413+
.replace(/"([^"]*?)[^a-zA-Z0-9\s\.,;:\-_\(\)"'\[\]{}]([^"]*?)"/g, '"$1$2"')
1414+
// Fix missing quotes around property names
1415+
.replace(/(\s*)(\w+)(\s*):/g, '$1"$2"$3:')
1416+
// Fix trailing commas in arrays and objects
1417+
.replace(/,(\s*[\]}])/g, '$1');
1418+
1419+
try {
1420+
jsonObject = JSON.parse(fixedJson);
1421+
} catch (finalError) {
1422+
logger.error("Failed to parse JSON after fixing:", finalError);
1423+
throw new Error("Invalid JSON returned from model");
1424+
}
1425+
}
1426+
1427+
// Validate against schema if provided
1428+
if (params.schema) {
1429+
try {
1430+
// Simplistic schema validation - check if all required properties exist
1431+
for (const key of Object.keys(params.schema)) {
1432+
if (!(key in jsonObject)) {
1433+
jsonObject[key] = null; // Add missing properties with null value
1434+
}
1435+
}
1436+
} catch (schemaError) {
1437+
logger.error("Schema validation failed:", schemaError);
1438+
}
1439+
}
1440+
1441+
return jsonObject;
1442+
} catch (parseError) {
1443+
logger.error("Failed to parse JSON:", parseError);
1444+
logger.error("Raw response:", textResponse);
1445+
throw new Error("Invalid JSON returned from model");
1446+
}
1447+
} catch (error) {
1448+
logger.error("Error in OBJECT_LARGE handler:", error);
1449+
throw error;
1450+
}
1451+
},
1452+
12091453
[ModelType.TEXT_TOKENIZER_ENCODE]: async (
12101454
_runtime: IAgentRuntime,
12111455
{ text }: { text: string },

0 commit comments

Comments
 (0)