server.js

/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import express from 'express';
import path from 'path';
import { fileURLToPath } from 'url';
import { GoogleGenAI, Type } from "@google/genai";
import 'dotenv/config';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const app = express();
app.use(express.json());
const port = process.env.PORT || 8080;

const responseSchema = {
    type: Type.OBJECT,
    properties: {
        nodes: {
            type: Type.ARRAY,
            items: {
                type: Type.OBJECT,
                properties: {
                    id: { type: Type.STRING, description: "Unique identifier for the node, derived from the agent name." },
                    type: { type: Type.STRING, enum: ['agentNode'] },
                    data: {
                        type: Type.OBJECT,
                        properties: {
                            label: { type: Type.STRING, description: "The agent's name from .name()" },
                            agentType: { type: Type.STRING, enum: ['LLM', 'SEQUENTIAL', 'PARALLEL', 'LOOP'], description: "The type of the agent." },
                            model: { type: Type.STRING, description: "The model used by the agent, from .model(). Can be null." },
                            instruction: { type: Type.STRING, description: "The instruction prompt for the agent, from .instruction(). Can be null." },
                            outputKey: { type: Type.STRING, description: "The key for the agent's output, from .outputKey(). Can be null." },
                            tools: { type: Type.ARRAY, items: { type: Type.STRING }, description: "List of tools used. Can be null or empty." },
                            maxIterations: { type: Type.NUMBER, description: "The maximum number of iterations for a LoopAgent, from .maxIterations(). Can be null." },
                            callbacks: { type: Type.ARRAY, items: { type: Type.STRING }, description: "List of callbacks defined on the agent (e.g., 'afterAgentCallback'). Can be null or empty." }
                        },
                        required: ['label', 'agentType']
                    },
                    position: {
                        type: Type.OBJECT,
                        properties: {
                            x: { type: Type.NUMBER },
                            y: { type: Type.NUMBER }
                        },
                        required: ['x', 'y']
                    }
                },
                required: ['id', 'type', 'data', 'position']
            }
        },
        edges: {
            type: Type.ARRAY,
            items: {
                type: Type.OBJECT,
                properties: {
                    id: { type: Type.STRING, description: "Unique identifier for the edge, e.g., 'source->target'." },
                    source: { type: Type.STRING, description: "The id of the source node." },
                    target: { type: Type.STRING, description: "The id of the target node." },
                    animated: { type: Type.BOOLEAN, description: "Whether the edge should be animated." },
                    label: { type: Type.STRING, description: "Label for the edge, used for outputKey connections. Can be null." },
                    data: {
                        type: Type.OBJECT,
                        properties: {
                           edgeType: { type: Type.STRING, enum: ['structural', 'dataFlow', 'toolUsage'], description: "The type of the connection." }
                        },
                        required: ['edgeType']
                    }
                },
                required: ['id', 'source', 'target', 'data']
            }
        }
    },
    required: ['nodes', 'edges']
};

app.post('/api/generate', async (req, res) => {
    const { code } = req.body;
    const apiKey = process.env.GOOGLE_API_KEY;

    if (!apiKey) {
        return res.status(500).json({ error: 'API key not configured' });
    }
    if (!code) {
        return res.status(400).json({ error: 'Code is required' });
    }

    try {
        const ai = new GoogleGenAI({ apiKey });
        const prompt = `
You are an expert code parser for an AI Agent Development Kit. Your task is to analyze the provided Java agent definition code and transform it into a JSON structure for React Flow.

Extract details for each agent:
- \`name\`: From \`.name("..."})\`. Use this for the node \`id\` and \`data.label\`.
- \`type\`: \`LlmAgent\` -> "LLM", \`SequentialAgent\` -> "SEQUENTIAL", \`ParallelAgent\` -> "PARALLEL", \`LoopAgent\` -> "LOOP". Use this for \`data.agentType\`.
- \`model\`: From \`.model("..."})\`.
- \`instruction\`: From \`.instruction("""...""")\`.
- \`outputKey\`: From \`.outputKey("..."})\`.
- \`tools\`: From any \`.tools(...)\` calls. Extract the name of the tool. If an agent is used as a tool (e.g., \`AgentTool.from(someAgentVariable)\` or \`AgentTool.create(someAgentVariable)\`), you MUST find the definition of \`someAgentVariable\`, get its name from its \`.name("..."})\` call, and use that agent name as the tool name in the \`tools\` array. For example, if a variable \`searchAgent\` points to an agent named "google-search-agent", the tools list should contain "google-search-agent", not "searchAgent".
- \`maxIterations\`: For \`LoopAgent\` only, from \`.maxIterations(...)\`. This should be a number.
- \`callbacks\`: From any \`.*Callback(...)\` calls (e.g., \`.beforeToolCallback(...)\`, \`.afterAgentCallback(...)\`). Extract the name of the callback method itself (e.g., "beforeToolCallback") and add it to the \`callbacks\` array.

Create three types of edges:
1.  **Structural Edges**: When an agent's builder calls \`.subAgents(...)\`, it is a parent. Create an edge from the parent agent to each agent variable passed into \`subAgents\`. Set \`data: { "edgeType": "structural" }\` and \`animated: true\`.
2.  **Data Flow Edges**: An agent's \`instruction\` might use a placeholder like \`{some_key}\`. If \`some_key\` matches the \`outputKey\` of another agent, create an edge from the agent with that \`outputKey\` to the agent using the placeholder. The \`label\` for this edge MUST be the key itself (e.g., "some_key"). Set \`data: { "edgeType": "dataFlow" }\` and \`animated: true\`.
3.  **Tool Usage Edges**: If an agent's builder calls \`.tools(AgentTool.from(someAgentVariable))\`, create an edge from the agent using the tool to the agent represented by \`someAgentVariable\`. Set \`data: { "edgeType": "toolUsage" }\` and \`animated: false\`.

Analyze the code below and generate a JSON object matching the provided schema. Ensure all agents mentioned are included as nodes.

Code to analyze:
---
${code}
---

Your output must be a valid JSON object.
- For nodes, populate \`id\`, \`type: 'agentNode'\`, \`data\` (with all extracted fields), and a placeholder \`position: { x: 0, y: 0 }\`.
- For edges, populate \`id\`, \`source\`, \`target\`, \`data\`, and the \`animated\` property as instructed. For data flow edges, also set the \`label\`.
`;
        const response = await ai.models.generateContent({
            model: "gemini-2.5-flash",
            contents: prompt,
            config: {
                responseMimeType: "application/json",
                responseSchema: responseSchema,
            },
        });

        const jsonText = response.text.trim();
        const parsedData = JSON.parse(jsonText);
        res.json(parsedData);
    } catch (error) {
        console.error('Error generating graph:', error);
        res.status(500).json({ error: 'Failed to generate graph' });
    }
});

// Serve static files from the React app
app.use(express.static(path.join(__dirname, 'dist')));

// The "catchall" handler: for any request that doesn't
// match one above, send back React's index.html file.
app.get('*', (req, res) => {
  res.sendFile(path.join(__dirname, 'dist', 'index.html'));
});

app.listen(port, () => {
  console.log(`Server listening on port ${port}`);
});