query-agent/index.js

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

import { Server } from "socket.io";
import { runAIControlledWorkflow, agentLog } from "./work.js";

function returnHtml(stringMessage) {
  return `<div style="background-color: transparent; color: white;">${stringMessage}</div>`;
}

// Main init function that adds Query Agent to existing Express server
function init(serverOrIo, options = {}) {
  if (!serverOrIo) {
    throw new Error("HTTP server or Socket.IO instance is required");
  }

  if (
    !options?.executeSQLQuery &&
    typeof options?.executeSQLQuery !== "function"
  ) {
    throw new Error("executeSQLQuery is required and must be a function");
  }

  const socketPath = "/query-agent";
  const routePrefix = socketPath;
  const { corsOptions = { origin: "*" } } = options;

  // Initialize Socket.IO with custom path and CORS
  let io;
  agentLog("serverOrIo", serverOrIo && serverOrIo.on ? "true" : "false");
  if (serverOrIo instanceof Server) {
    // Existing Socket.IO instance passed
    io = serverOrIo;
  } else {
    agentLog("corsOptions", corsOptions);
    io = new Server(serverOrIo, {
      cors: corsOptions,
    });
  }

  // Socket.IO namespace for /query-agent
  const queryAgentNamespace = io.of(socketPath);
  //   let queryAgentNamespace = io;

  queryAgentNamespace.use((socket, next) => {
    const token =
      socket.handshake.auth.token || socket.handshake.headers["authorization"];

    if (!token) {
      return next(new Error("Unauthorized: No token provided"));
    }

    if (token !== process.env.AgentToken) {
      return next(new Error("Unauthorized: Invalid token"));
    }

    try {
      next();
    } catch (err) {
      return next(new Error("Unauthorized: Invalid token"));
    }
  });

  queryAgentNamespace.on("connection", (socket) => {
    agentLog(`🔌 Client connected to /query-agent namespace: ${socket?.id}`);

    // Handle query events
    socket.on("query", async (data) => {
      try {
        agentLog(`📥 Received query from ${socket.id}:`, data);
        agentLog("query data:", data);
        const {
          userQuery,
          dbType = "MySQL",
          otherDetails = "",
          model,
          apiKey,
        } = data;

        if (!userQuery) {
          queryAgentNamespace.to(socket.id).emit("response", {
            status: "failed",
            message: returnHtml("User query is required"),
          });
          return;
        }

        if (!model) {
          queryAgentNamespace.to(socket.id).emit("response", {
            status: "failed",
            message: returnHtml("Model is required"),
          });
          return;
        }

        if (!apiKey) {
          queryAgentNamespace.to(socket.id).emit("response", {
            status: "failed",
            message: returnHtml("Auth key is required"),
          });
          return;
        }

        // Emit status updates during processing
        queryAgentNamespace.to(socket.id).emit("PROCESSING", {
          status: "processing",
          message: returnHtml("Starting AI workflow..."),
        });

        await runAIControlledWorkflow({
          userQuery,
          dbType,
          otherDetails,
          queryAgentNamespace,
          socketId: socket.id,
          executeSQLQuery: options?.executeSQLQuery,
          model: model,
          authKey: apiKey,
        });
      } catch (error) {
        agentLog("Socket error:", error);
        queryAgentNamespace.to(socket.id).emit("response", {
          status: "failed",
          message: returnHtml(
            error.message || "Socket error occurred, please try again"
          ),
        });
      }
    });

    socket.on("connect", () => {
      agentLog(`🔌 Client connected to /query-agent namespace: ${socket?.id}`);
    });

    socket.on("disconnect", () => {
      agentLog(
        `🔌 Client disconnected from /query-agent namespace: ${socket?.id}`
      );
    });
  });

  queryAgentNamespace.on("error", (error) => {
    agentLog("Socket error:", error);
  });

  queryAgentNamespace.on("disconnect", () => {
    agentLog("Client disconnected from /query-agent namespace");
  });

  console.log(`✅ Query Agent initialized:`);
  console.log(`📡 Socket.IO namespace: /query-agent`);

  return { io, queryAgentNamespace };
}

export default init;