88 lines
3.0 KiB
JavaScript
88 lines
3.0 KiB
JavaScript
const axios = require('axios');
|
|
const { GoogleGenerativeAI } = require('@google/generative-ai');
|
|
const ai = new GoogleGenerativeAI({ apiKey: process.env.GEMINI_API_KEY });
|
|
|
|
const fetchWrenData = async (prompt, tenantId) => {
|
|
try {
|
|
const url = process.env.WREN_AI_URL || 'http://localhost:5555/api/v1/text-to-sql';
|
|
|
|
const response = await axios.post(url,
|
|
{ prompt: prompt },
|
|
{
|
|
headers: {
|
|
|
|
'X-Wren-Session-Properties': `@user_org_id=${tenantId}`,
|
|
'Content-Type': 'application/json'
|
|
}
|
|
}
|
|
);
|
|
return response.data;
|
|
} catch (error) {
|
|
console.error('Wren AI Integration Error:', error.message);
|
|
throw new Error('Failed to fetch data payload from Wren AI endpoint');
|
|
}
|
|
};
|
|
|
|
const generateVegaSchema = async (question, dataArray) => {
|
|
try {
|
|
const model = ai.getGenerativeModel({
|
|
model: 'gemini-2.5-flash',
|
|
generationConfig: { responseMimeType: 'application/json' }
|
|
});
|
|
|
|
const systemPrompt = `You are a data visualization expert. I will provide a user's question and a JSON array of data.
|
|
Your task is to generate a strictly valid Vega-Lite JSON specification to visualize this data.
|
|
The data array will be provided to the Vega spec internally. Map the JSON keys to the correct x, y, and color axes.
|
|
Choose the best chart type (bar, line, arc) based on the question.`;
|
|
|
|
const userPrompt = `User Question: "${question}"\nData JSON: ${JSON.stringify(dataArray)}`;
|
|
|
|
const result = await model.generateContent({
|
|
contents: [{ role: 'user', parts: [{ text: `${systemPrompt}\n\n${userPrompt}` }] }]
|
|
});
|
|
|
|
return JSON.parse(result.response.text());
|
|
} catch (error) {
|
|
console.error('Gemini Engine Error:', error.message);
|
|
throw new Error('Failed to transform data architecture into valid Vega-Lite spec');
|
|
}
|
|
};
|
|
|
|
|
|
const handleOrchestration = async (req, res) => {
|
|
try {
|
|
const { prompt } = req.body;
|
|
|
|
|
|
const tenant_id = req.user ? req.user.client_id : null;
|
|
return res.status(200).json("ok");
|
|
|
|
if (!prompt) {
|
|
return res.status(400).json({ error: 'Prompt field is required in request body' });
|
|
}
|
|
if (!tenant_id) {
|
|
return res.status(400).json({ error: 'Tenant context (client_id) missing from auth token' });
|
|
}
|
|
|
|
|
|
const wrenData = await fetchWrenData(prompt, tenant_id);
|
|
|
|
|
|
const vegaSchema = await generateVegaSchema(prompt, wrenData);
|
|
|
|
|
|
if (!vegaSchema || !vegaSchema.$schema) {
|
|
return res.status(500).json({ error: 'Egress Pipeline Validation Failure: Output missing standard Vega $schema identifier' });
|
|
}
|
|
|
|
|
|
return res.status(200).json(vegaSchema);
|
|
|
|
} catch (error) {
|
|
console.error('API Gateway Orchestrator Crash:', error.message);
|
|
return res.status(500).json({ error: error.message });
|
|
}
|
|
};
|
|
|
|
module.exports = { handleOrchestration };
|