firebase-tools
Version:
Command-Line Interface for Firebase
193 lines (192 loc) • 7.15 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.PROMPT_GENERATE_SEED_DATA = exports.PROMPT_GENERATE_CONNECTOR = void 0;
exports.generateSchema = generateSchema;
exports.generateOperation = generateOperation;
exports.extractCodeBlock = extractCodeBlock;
const apiv2_1 = require("../apiv2");
const api_1 = require("../api");
const error_1 = require("../error");
const logger_1 = require("../logger");
const apiClient = new apiv2_1.Client({ urlPrefix: (0, api_1.dataconnectOrigin)(), auth: true });
exports.PROMPT_GENERATE_CONNECTOR = "Create 4 operations for an app using the instance schema with proper authentication.";
exports.PROMPT_GENERATE_SEED_DATA = "Create a mutation to populate the database with some seed data.";
function logCurl(method, path, body) {
const url = `${(0, api_1.dataconnectOrigin)()}${path}`;
const headers = [
'-H "Content-Type: application/json"',
'-H "Authorization: Bearer $(gcloud auth print-access-token)"',
].join(" ");
const curl = `curl -X ${method} "${url}" ${headers} -d '${JSON.stringify(body)}'`;
logger_1.logger.debug(`[Agent Service] Reusable cURL command:\\n${curl}`);
}
async function generateSchema(prompt, project, location, onStatus) {
const path = `/v1/projects/${project}/locations/${location}/services/-:generateSchema`;
const body = {
name: `projects/${project}/locations/${location}/services/-`,
prompt,
};
logCurl("POST", path, body);
const res = await apiClient.request({
method: "POST",
path,
body,
responseType: "stream",
resolveOnHTTPError: true,
});
if (res.status >= 400) {
const errorText = await readStream(res.body);
throw new error_1.FirebaseError(`Failed to generate schema. Status: ${res.status}, Message: ${errorText}`);
}
return consumeStream(res.body, onStatus);
}
async function generateOperation(prompt, service, project, schemas, onStatus) {
let location = "us-central1";
let serviceId = service;
if (service.startsWith("projects/")) {
const parts = service.split("/");
project = parts[1];
location = parts[3];
serviceId = parts[5];
}
if (schemas && schemas.length > 0) {
serviceId = "-";
}
const path = `/v1/projects/${project}/locations/${location}/services/${serviceId}:generateQuery`;
const body = {
name: `projects/${project}/locations/${location}/services/${serviceId}`,
prompt,
schemas,
};
logCurl("POST", path, body);
const res = await apiClient.request({
method: "POST",
path,
body,
responseType: "stream",
resolveOnHTTPError: true,
});
if (res.status >= 400) {
const errorText = await readStream(res.body);
throw new error_1.FirebaseError(`Failed to generate operation. Status: ${res.status}, Message: ${errorText}`);
}
return consumeStream(res.body, onStatus);
}
async function readStream(stream) {
return new Promise((resolve, reject) => {
let data = "";
stream.on("data", (chunk) => {
data += chunk.toString();
});
stream.on("end", () => {
resolve(data);
});
stream.on("error", (err) => {
reject(err);
});
});
}
async function consumeStream(stream, onStatus) {
return new Promise((resolve, reject) => {
let buffer = "";
let fullText = "";
stream.on("data", (chunk) => {
const text = chunk.toString();
fullText += text;
buffer += text;
let newlineIndex;
while ((newlineIndex = buffer.indexOf("\n")) !== -1) {
const line = buffer.substring(0, newlineIndex).trim();
buffer = buffer.substring(newlineIndex + 1);
if (line) {
try {
const obj = JSON.parse(line);
if (obj.status && onStatus) {
onStatus(obj.status);
}
}
catch (err) {
}
}
}
});
stream.on("end", () => {
try {
const response = JSON.parse(fullText);
if (Array.isArray(response)) {
let code = "";
for (const item of response) {
if (item.status && onStatus) {
onStatus(item.status);
}
if (item.part?.textChunk?.text) {
code += item.part.textChunk.text;
}
if (item.part?.codeChunk?.code) {
code += item.part.codeChunk.code;
}
}
if (code) {
resolve(extractCodeBlock(code));
}
else {
resolve(fullText);
}
}
else {
const resObj = response;
if (resObj.part?.codeChunk?.code) {
resolve(extractCodeBlock(resObj.part.codeChunk.code));
}
else if (resObj.part?.textChunk?.text) {
resolve(extractCodeBlock(resObj.part.textChunk.text));
}
else {
resolve(fullText);
}
}
}
catch (e) {
const lines = fullText.trim().split("\n");
let code = "";
for (const line of lines) {
try {
const obj = JSON.parse(line);
if (obj.part?.codeChunk?.code) {
code += obj.part.codeChunk.code;
}
else if (obj.part?.textChunk?.text) {
code += obj.part.textChunk.text;
}
if (obj.status && onStatus) {
onStatus(obj.status);
}
}
catch (err) {
logger_1.logger.error("Failed to parse FSQL Generate response: ", err);
}
}
if (code) {
resolve(extractCodeBlock(code));
}
else {
resolve(fullText);
}
}
});
stream.on("error", (err) => {
reject(err);
});
});
}
function extractCodeBlock(text) {
const regex = /```(?:[a-z]+\n)?([\s\S]*?)```/m;
const match = regex.exec(text);
if (match && match[1]) {
return match[1].trim();
}
if (!text.includes("{")) {
logger_1.logger.warn("[Agent Service] Response seems to be plain text, no GraphQL code block found.");
}
return text.trim();
}