124 lines
4.3 KiB
TypeScript
124 lines
4.3 KiB
TypeScript
import fs from 'fs';
|
|
import dotenv from 'dotenv';
|
|
import { logger } from './logger';
|
|
|
|
dotenv.config();
|
|
|
|
const LMSTUDIO_BASE_URL = process.env.LMSTUDIO_BASE_URL;
|
|
const LMSTUDIO_API_KEY = process.env.LMSTUDIO_API_KEY;
|
|
const LMSTUDIO_MODEL = process.env.LMSTUDIO_MODEL;
|
|
|
|
async function callLmstudio(prompt: string): Promise<any> {
|
|
if (!LMSTUDIO_BASE_URL) {
|
|
throw new Error('LMSTUDIO_BASE_URL is not defined in the .env file');
|
|
}
|
|
|
|
for (let i = 0; i < 10; i++) {
|
|
let llmResponse = "";
|
|
|
|
try {
|
|
const response = await fetch(`${LMSTUDIO_BASE_URL}/chat/completions`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': `Bearer ${LMSTUDIO_API_KEY}`,
|
|
},
|
|
body: JSON.stringify({
|
|
model: LMSTUDIO_MODEL,
|
|
messages: [
|
|
{
|
|
role: 'user',
|
|
content: prompt,
|
|
},
|
|
],
|
|
temperature: 0.7,
|
|
}),
|
|
});
|
|
|
|
const data = await response.json();
|
|
if (data.choices && data.choices.length > 0) {
|
|
const content = data.choices[0].message.content;
|
|
llmResponse = content;
|
|
const jsonMatch = content.match(/\{[\s\S]*\}/);
|
|
if (jsonMatch) {
|
|
return JSON.parse(jsonMatch[0]);
|
|
} else {
|
|
const arrayMatch = content.match(/\[[\s\S]*\]/);
|
|
if (arrayMatch) {
|
|
return JSON.parse(arrayMatch[0]);
|
|
}
|
|
}
|
|
// If no JSON/array found, return the raw content
|
|
return content;
|
|
} else {
|
|
logger.error('Unexpected API response:', data);
|
|
}
|
|
} catch (error) {
|
|
logger.error(`Attempt ${i + 1} failed:`, error);
|
|
logger.debug(`LLM response: ${llmResponse}`)
|
|
}
|
|
}
|
|
|
|
throw new Error('Failed to get response from LLM after 10 attempts');
|
|
}
|
|
|
|
async function callLMStudioAPIWithFile(imagePath: string, prompt: string): Promise<any> {
|
|
if (!LMSTUDIO_BASE_URL) {
|
|
throw new Error('LMSTUDIO_BASE_URL is not defined in the .env file');
|
|
}
|
|
|
|
const imageBuffer = fs.readFileSync(imagePath);
|
|
const base64Image = imageBuffer.toString('base64');
|
|
|
|
for (let i = 0; i < 10; i++) {
|
|
let llmResponse = "";
|
|
|
|
try {
|
|
const response = await fetch(`${LMSTUDIO_BASE_URL}/chat/completions`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': `Bearer ${LMSTUDIO_API_KEY}`,
|
|
},
|
|
body: JSON.stringify({
|
|
model: LMSTUDIO_MODEL,
|
|
messages: [
|
|
{
|
|
role: 'user',
|
|
content: [
|
|
{ type: 'image_url', image_url: { url: `data:image/jpeg;base64,${base64Image}` } },
|
|
{ type: 'text', text: prompt },
|
|
],
|
|
},
|
|
],
|
|
temperature: 0.7,
|
|
}),
|
|
});
|
|
|
|
const data = await response.json();
|
|
if (data.choices && data.choices.length > 0) {
|
|
const content = data.choices[0].message.content;
|
|
llmResponse = content;
|
|
const jsonMatch = content.match(/\{[\s\S]*\}/);
|
|
if (jsonMatch) {
|
|
return JSON.parse(jsonMatch[0]);
|
|
} else {
|
|
const arrayMatch = content.match(/\[[\s\S]*\]/);
|
|
if (arrayMatch) {
|
|
return JSON.parse(arrayMatch[0]);
|
|
}
|
|
}
|
|
} else {
|
|
logger.error('Unexpected API response:', data);
|
|
}
|
|
} catch (error) {
|
|
logger.error(`Attempt ${i + 1} failed:`, error);
|
|
logger.debug(`LLM response: ${llmResponse}`)
|
|
}
|
|
}
|
|
|
|
throw new Error('Failed to describe image after 10 attempts');
|
|
}
|
|
|
|
export { callLmstudio, callLMStudioAPIWithFile };
|