query-agent/work.js

An AI-powered database query agent that integrates with existing Express apps using Socket.IO and HTTP routes

import { GoogleGenerativeAI } from "@google/generative-ai";

// Logging function that only prints when agentLog=1
export const agentLog = (message, ...args) => {
  if (process.env.agentLog === "1") {
    console.log(message, ...args);
  }
};

// Function declaration for the Gemini API
const sqlQueryFunctionDeclaration = {
  name: "execute_sql_query",
  description:
    "Execute a SELECT SQL query to retrieve data from the database. Only SELECT queries are allowed for security reasons.",
  parameters: {
    type: "object",
    properties: {
      query: {
        type: "string",
        description:
          "The SQL SELECT query to execute. Must be properly formatted for the database type being used.",
      },
      reason: {
        type: "string",
        description:
          "Brief explanation of why this query is needed to answer the user's question.",
      },
    },
    required: ["query", "reason"],
  },
};

// Generate system instruction based on database type
const generateSystemInstruction = (dbType, otherDetails) => {
  return `You are a helpful database assistant that helps query the ${dbType} database.

IMPORTANT GUIDELINES:
1. You can only execute SELECT queries - no INSERT, UPDATE, DELETE, or DDL operations
2. Before generating any final query, explore the database schema step by step:
   - First, understand whole databases/schemas are available
   - Then explore relevant tables and their structure
   - Check column names, data types, and relationships
   - Only then generate the final query to answer the user's question

3. Always explain your reasoning and break down complex queries
4. If you encounter errors, analyze them and correct your approach
5. Present final results in a clear, formatted HTML response
6. If the user asks about tables/columns that don't exist, show available options
7. Table name and column names etc, are case-sensitive, and there's a high chance they may not exactly match the user's input. You should cross-check the details to create query and call the execute_sql_query function call and extract the actual table name and column names etc. from the model's answer.
8. join multiple tables, if needed, to answer the user's question.
9. for Html response prefer dart mode theme, background must be transparent, and font color must be white.
10. Don't Add any comment in HTMl response Code.
11. Don't use "vh" and "vw" to give height and width in the response html code use % instead.
12. If Response is Too Long then don't send partial html output that cut from in the middle. Say "Response is Too Long"
13. Don't send initial message of planning start direct start executing the query.

PROCESS :
- Use the execute_sql_query function to explore the database incrementally
- Start with schema exploration queries if needed
- Build up to the final query that answers the user's question
- Format the final response as HTML for better readability

IMPORTANT GUIDELINES AND PROCESS FOR CHART (if user ask about any chart or graph):
1. you should use the execute_sql_query function and follow above PROCESS section to get the final data. if final data is empty then prepare response like "Not enough data to render chart" else follow below process.
2. Collect Data and formate data structure like below:
  type: {type of Chart}
  labels: string[];
  datasets: {
    label: string;
    data: number[];
    backgroundColor?: string;
    borderColor?: string;
    borderWidth?: number;
  }[];
3. must provide pure json output format.
4. Don't add any comment in the json output.
5. No extra space or any other text in the json and outside the json output.



Database Type: ${dbType}
${otherDetails ? `Schema/Additional Context: ${otherDetails}` : ""}

IMPORTANT FOR POSTGRESQL:
- POSTGRESQL uses SCHEMAS instead of SCHEMATA. 
- Use proper schema prefixes: ${
    otherDetails
      ? `${otherDetails.replace("schema is ", "")}.table_name`
      : "schema.table_name"
  }
- Remember PostgreSQL is case-sensitive for quoted identifiers
- Use LIMIT instead of TOP for row limiting

IMPORTANT FOR MySQL:
- MySQL uses SCHEMATA instead of SCHEMAS.

Remember: Always put your main query/response at the end for optimal performance.`;
};

// Get model and system instruction (simplified approach)
const getModelAndSystemInstruction = ({
  dbType,
  otherDetails,
  modelName = "gemini-2.5-flash",
  authKey,
}) => {
  agentLog("modelName", modelName);
  const genAI = new GoogleGenerativeAI(authKey);
  const model = genAI.getGenerativeModel({
    model: modelName, // Using the updated model
    tools: [
      {
        functionDeclarations: [sqlQueryFunctionDeclaration],
      },
    ],
    generationConfig: {
      temperature: 0.1,
    },
  });

  const systemInstruction = generateSystemInstruction(dbType, otherDetails);

  return {
    model: model,
    instruction: systemInstruction,
  };
};

// Process function calls and execute SQL queries
const processFunctionCall = async (
  functionCall,
  queryAgentNamespace,
  socketId,
  executeSQLQuery
) => {
  const { name, args } = functionCall;
  queryAgentNamespace.to(socketId).emit("PROCESSING", functionCall);
  if (name === "execute_sql_query") {
    agentLog(`🔍 Executing SQL query: ${args.query}`);
    agentLog(`💭 Reason: ${args.reason}`);

    const result = await executeSQLQuery(args.query);
    agentLog(`📊 Query result:`, {
      success: result.success,
      rowCount: result.rowCount || 0,
      hasError: !!result.error,
      dataPreview:
        result.success && result.data
          ? `${result.data.length} rows`
          : result.error || "No data",
    });

    queryAgentNamespace.to(socketId).emit("PROCESSING", {
      name: "execute_sql_query_result",
      args: {
        query: args.query,
        reason: args.reason,
      },
      result: result,
    });

    return {
      name: name,
      response: {
        result: result,
      },
    };
  }

  throw new Error(`Unknown function: ${name}`);
};

// Main workflow function with context caching optimization
export async function runAIControlledWorkflow({
  userQuery,
  dbType = "MySQL",
  otherDetails = "",
  queryAgentNamespace,
  socketId,
  executeSQLQuery,
  model: modelName,
  authKey,
}) {
  if (!userQuery) {
    queryAgentNamespace.to(socketId).emit("response", {
      status: "failed",
      message: "User query is required",
    });
    return;
  }

  try {
    // Get model and system instruction
    const { model, instruction } = getModelAndSystemInstruction({
      dbType,
      otherDetails,
      modelName,
      authKey,
    });

    // Start conversation with system instruction
    let contents = [
      {
        role: "user",
        parts: [{ text: instruction }],
      },
    ];

    agentLog("Step 1: Sending system instruction...");
    const systemResponse = await model.generateContent({
      contents: contents,
    });

    agentLog(
      "✅ System instruction acknowledged:",
      systemResponse.response.text().substring(0, 100) + "..."
    );

    // Add system response to conversation
    contents.push({
      role: "model",
      parts: [{ text: systemResponse.response.text() }],
    });

    // Add user query
    contents.push({
      role: "user",
      parts: [{ text: `User Question: ${userQuery}` }],
    });

    agentLog("Step 2: Sending user query:", userQuery);

    // Function calling loop
    let maxIterations = 10; // Prevent infinite loops
    let iteration = 0;

    while (iteration < maxIterations) {
      iteration++;
      agentLog(`Iteration ${iteration}`);
      agentLog(contents);
      // Generate response
      const result = await model.generateContent({
        contents: contents,
      });

      const response = result.response;

      const functionCalls = response.functionCalls?.();

      // Log function calls for debugging (optional)
      if (functionCalls && functionCalls.length > 0) {
        agentLog(`📞 Function calls: ${functionCalls.length}`);
      }

      // Check if there are function calls
      if (functionCalls && functionCalls.length > 0) {
        // Process each function call
        for (const functionCall of functionCalls) {
          const functionResponse = await processFunctionCall(
            functionCall,
            queryAgentNamespace,
            socketId,
            executeSQLQuery
          );

          // Add the model's response with function call to conversation
          contents.push({
            role: "model",
            parts: [{ functionCall: functionCall }],
          });

          // Add the function response to conversation
          contents.push({
            role: "user",
            parts: [{ functionResponse: functionResponse }],
          });
        }
      } else {
        // No more function calls, we have the final response
        const finalText = response.text();
        agentLog(
          "🏁 Final response generated:",
          finalText ? `${finalText.length} characters` : "EMPTY!"
        );
        agentLog("📄 Final text preview:", finalText);

        if (!finalText || finalText.trim().length === 0) {
          agentLog("❌ Empty final response - something went wrong!");
          const debugResponse = `
                    <div style="font-family: Arial, sans-serif; padding: 20px;">
                        <h3 style="color: red;">Debug: Empty Response Detected</h3>
                        <div style="background: #3f3f46; padding: 15px; border-radius: 5px; border: 1px solid #555; color: #fff;">
                            <p><strong>Issue:</strong> The AI completed processing but returned an empty response.</p>
                            <p><strong>Iterations completed:</strong> ${iteration}</p>
                            <p><strong>Check server logs for function call details.</strong></p>
                            <p><strong>Your query:</strong> ${userQuery}</p>
                            <p><strong>Database:</strong> ${dbType} (${otherDetails})</p>
                        </div>
                    </div>`;
          queryAgentNamespace.to(socketId).emit("response", {
            status: "failed",
            message: debugResponse,
          });
          return;
        }

        // Check if response looks like HTML
        if (finalText.includes("```json")) {
          queryAgentNamespace.to(socketId).emit("response", {
            status: "success",
            message: finalText,
          });
          return;
        } else if (finalText.includes("<") && finalText.includes(">")) {
          queryAgentNamespace.to(socketId).emit("response", {
            status: "success",
            message: finalText,
          });
          return;
        } else {
          // Wrap in basic HTML structure
          const htmlResponse = `
                    <div style="font-family: Arial, sans-serif; padding: 20px;">
                        <h3>Query Result</h3>
                        <div style="background: #3f3f46; padding: 15px; border-radius: 5px; color: #fff;">
                            ${finalText.replace(/\n/g, "<br>")}
                        </div>
                    </div>`;
          queryAgentNamespace.to(socketId).emit("response", {
            status: "success",
            message: htmlResponse,
          });
          return;
        }
      }
    }

    // If we reach max iterations, return what we have
    queryAgentNamespace.to(socketId).emit("response", {
      status: "partial_success",
      message: "Reached maximum iterations. Here's what I found:",
    });
    return;
  } catch (error) {
    agentLog("Error in AI workflow:", error);
    queryAgentNamespace.to(socketId).emit("response", {
      status: "error",
      message:
        error?.message || "An error occurred while processing your request",
    });
    return;
  }
}