@mixio-pro/kalaasetu-mcp 1.2.2 → 2.0.2-beta

package/src/tools/fal/models.ts

@@ -1,54 +1,159 @@
+import { z } from "zod";
+import {
+  loadFalConfig,
+  saveFalConfig,
+  FAL_BASE_URL,
+  DEFAULT_TIMEOUT,
+} from "./config";
+import { safeToolExecute } from "../../utils/tool-wrapper";
+
 /**
- * Models module for fal.ai MCP server.
- * Provides tools for retrieving information about configured fal.ai models.
+ * Extract simplified input schema from FAL OpenAPI response.
+ * Returns only the properties object from the input schema.
  */
+function extractInputSchema(openApiSchema: any): Record<string, any> | null {
+  try {
+    const schemas = openApiSchema?.components?.schemas;
+    if (!schemas) return null;
 
-import { z } from "zod";
-import { loadFalConfig, FAL_BASE_URL, DEFAULT_TIMEOUT } from "./config";
-import { safeToolExecute } from "../../utils/tool-wrapper";
+    // Find the input schema - usually named like "Ltx2ImageToVideoInput" or similar
+    const inputSchemaKey = Object.keys(schemas).find(
+      (key) =>
+        key.toLowerCase().includes("input") &&
+        !key.toLowerCase().includes("output")
+    );
+
+    if (!inputSchemaKey) return null;
+
+    const inputSchema = schemas[inputSchemaKey];
+    if (!inputSchema?.properties) return null;
+
+    // Extract simplified properties
+    const simplified: Record<string, any> = {};
+    for (const [propName, propDef] of Object.entries(
+      inputSchema.properties as Record<string, any>
+    )) {
+      simplified[propName] = {
+        type: propDef.type,
+        description: propDef.description,
+        ...(propDef.enum && { enum: propDef.enum }),
+        ...(propDef.default !== undefined && { default: propDef.default }),
+        ...(propDef.examples && { example: propDef.examples[0] }),
+      };
+    }
+
+    return simplified;
+  } catch (e) {
+    return null;
+  }
+}
 
 /**
  * Tool to list available generation presets and their intents.
+ * Dynamically fetches and caches input schemas from FAL API.
  */
 export const falListPresets = {
   name: "fal_list_presets",
   description:
     "The entry point for discovering fal.ai capabilities on this server. " +
     "Lists all available generation presets, including their high-level 'intent' (e.g., 'Generate cinematic video'), " +
-    "and the types of input/output they support. Call this first when you need to perform an AI generation task.",
-  parameters: z.object({}),
-  timeoutMs: 30000,
-  execute: async () => {
+    "input/output types, and INPUT SCHEMA with parameter details. Call this first when you need to perform an AI generation task.",
+  parameters: z.object({
+    refresh_schemas: z
+      .boolean()
+      .optional()
+      .describe("If true, re-fetch schemas from FAL API even if cached."),
+  }),
+  timeoutMs: 60000, // Allow more time for schema fetching
+  execute: async (args: { refresh_schemas?: boolean }) => {
     return safeToolExecute(async () => {
       const config = loadFalConfig();
+      let updated = false;
+
+      // Fetch schemas for presets that don't have them (or if refresh requested)
+      for (const preset of config.presets) {
+        const shouldFetch = !preset.input_schema || args.refresh_schemas;
+        console.log(
+          `[fal_list_presets] ${
+            preset.presetName
+          }: shouldFetch=${shouldFetch}, hasSchema=${!!preset.input_schema}, refresh=${
+            args.refresh_schemas
+          }`
+        );
+
+        if (shouldFetch) {
+          try {
+            const url = `https://fal.ai/api/openapi/queue/openapi.json?endpoint_id=${preset.modelId}`;
+            console.log(`[fal_list_presets] Fetching schema from: ${url}`);
+            const response = await fetch(url, {
+              method: "GET",
+              signal: AbortSignal.timeout(10000),
+            });
+
+            if (response.ok) {
+              const openApiSchema = await response.json();
+              const simplified = extractInputSchema(openApiSchema);
+              console.log(
+                `[fal_list_presets] Extracted schema for ${preset.presetName}:`,
+                simplified ? Object.keys(simplified) : null
+              );
+
+              if (simplified) {
+                preset.input_schema = simplified;
+                updated = true;
+              }
+            } else {
+              console.log(
+                `[fal_list_presets] Fetch failed: ${response.status}`
+              );
+            }
+          } catch (e: any) {
+            console.log(
+              `[fal_list_presets] Error fetching schema for ${preset.presetName}:`,
+              e.message
+            );
+          }
+        }
+      }
+
+      // Save updated config if schemas were fetched
+      if (updated) {
+        console.log(`[fal_list_presets] Saving updated config...`);
+        const saved = saveFalConfig(config);
+        console.log(`[fal_list_presets] Config saved: ${saved}`);
+      }
+
+      // Return enriched preset list
       const summary = config.presets.map((p) => ({
         presetName: p.presetName,
         intent: p.intent,
         inputType: p.inputType,
         outputType: p.outputType,
         description: p.description,
+        inputSchema: p.input_schema,
       }));
+
       return JSON.stringify(summary, null, 2);
     }, "fal_list_presets");
   },
 };
 
 /**
- * Tool to get full details for a specific preset, including default parameters.
+ * Tool to get full details for a specific preset, including default parameters
+ * and real-time model metadata from fal.ai.
  */
 export const falGetPresetDetails = {
   name: "fal_get_preset_details",
   description:
     "Retrieve full details for a specific generation preset. " +
-    "Use this to see the 'modelId' being used and, most importantly, the 'defaultParams'. " +
-    "The default parameters shown here can be overridden in the 'parameters' argument of 'fal_generate'. " +
+    "This tool fetches both the local preset configuration (like 'defaultParams') " +
+    "and the live model metadata (schema, benchmarks, etc.) from fal.ai. " +
+    "Use this to understand the full capabilities and constraints of a model. " +
     "ONLY USE WHEN WORKING WITH FAL MODELS/PRESETS.",
   parameters: z.object({
     preset_name: z
       .string()
-      .describe(
-        "The name of the preset to inspect (e.g., 'ltx_image_to_video')."
-      ),
+      .describe("The name of the preset to inspect (e.g., 'cinematic_image')."),
   }),
   timeoutMs: 30000,
   execute: async (args: { preset_name: string }) => {
@@ -60,7 +165,43 @@ export const falGetPresetDetails = {
       if (!preset) {
         throw new Error(`Preset '${args.preset_name}' not found.`);
       }
-      return JSON.stringify(preset, null, 2);
+
+      // Fetch live model metadata/schema from fal.ai API
+      // Based on: https://fal.ai/api/openapi/queue/openapi.json?endpoint_id={model_id}
+      let modelMetadata: any = null;
+      let schemaSource = "none";
+
+      try {
+        const url = `https://fal.ai/api/openapi/queue/openapi.json?endpoint_id=${preset.modelId}`;
+        const schema = await publicRequest(url);
+
+        // Extract relevant schema parts if possible, or return full schema
+        if (schema) {
+          modelMetadata = schema;
+          schemaSource = "api";
+        }
+      } catch (e) {
+        // console.error(`Failed to fetch schema for ${preset.modelId}:`, e);
+        // Fallback to locally defined schema if available
+      }
+
+      // Fallback: If API failed or returned nothing, use manual schema from config
+      if (!modelMetadata && preset.input_schema) {
+        modelMetadata = preset.input_schema;
+        schemaSource = "local_config";
+      }
+
+      return JSON.stringify(
+        {
+          preset,
+          modelMetadata,
+          _meta: {
+            schemaSource,
+          },
+        },
+        null,
+        2
+      );
     }, "fal_get_preset_details");
   },
 };
@@ -82,22 +223,6 @@ async function publicRequest(url: string): Promise<any> {
   return response.json();
 }
 
-/**
- * Tool to retrieve the current full FAL configuration.
- */
-export const falGetConfig = {
-  name: "fal_get_config",
-  description: "Retrieve the full FAL configuration JSON file content.",
-  parameters: z.object({}),
-  timeoutMs: 30000,
-  execute: async () => {
-    return safeToolExecute(async () => {
-      const config = loadFalConfig();
-      return JSON.stringify(config, null, 2);
-    }, "fal_get_config");
-  },
-};
-
 /* 
 // ORIGINAL UNRESTRICTED TOOLS - Commented out for reference
 

package/src/tools/gemini.ts

@@ -12,6 +12,10 @@ import { PassThrough } from "stream";
 import { getStorage } from "../storage";
 import { generateTimestampedFilename } from "../utils/filename";
 import { safeToolExecute } from "../utils/tool-wrapper";
+import {
+  resolveEnhancer,
+  listImageEnhancerPresets,
+} from "../utils/prompt-enhancer-presets";
 
 const ai = new GoogleGenAI({
   apiKey: process.env.GEMINI_API_KEY || "",
@@ -232,6 +236,13 @@ export const geminiTextToImage = {
       .describe(
         "Optional: local paths or URLs of images to use as visual references for style or composition."
       ),
+    enhancer_preset: z
+      .string()
+      .optional()
+      .describe(
+        "Optional: Name of a prompt enhancer preset to apply (e.g., 'cinematic', 'photorealistic', 'anime'). " +
+          "Automatically enhances the prompt with professional style modifiers."
+      ),
   }),
   timeoutMs: 300000,
   execute: async (args: {
@@ -239,10 +250,20 @@ export const geminiTextToImage = {
     aspect_ratio?: string;
     output_path?: string;
     reference_images?: string[];
+    enhancer_preset?: string;
   }) => {
     return safeToolExecute(async () => {
       try {
-        const contents: any[] = [args.prompt];
+        // Apply prompt enhancement if preset specified
+        let enhancedPrompt = args.prompt;
+        if (args.enhancer_preset) {
+          const enhancer = resolveEnhancer(args.enhancer_preset);
+          if (enhancer.hasTransformations()) {
+            enhancedPrompt = enhancer.enhance(args.prompt);
+          }
+        }
+
+        const contents: any[] = [enhancedPrompt];
 
         if (args.reference_images && Array.isArray(args.reference_images)) {
           for (const refPath of args.reference_images) {
@@ -337,6 +358,13 @@ export const geminiEditImage = {
       .describe(
         "Optional: additional images to guide the edit (e.g., to reference a specific character or object style)."
       ),
+    enhancer_preset: z
+      .string()
+      .optional()
+      .describe(
+        "Optional: Name of a prompt enhancer preset to apply (e.g., 'cinematic', 'photorealistic'). " +
+          "Enhances the edit instructions with professional style modifiers."
+      ),
   }),
   timeoutMs: 300000,
   execute: async (args: {
@@ -344,11 +372,21 @@ export const geminiEditImage = {
     prompt: string;
     output_path?: string;
     reference_images?: string[];
+    enhancer_preset?: string;
   }) => {
     return safeToolExecute(async () => {
       try {
+        // Apply prompt enhancement if preset specified
+        let enhancedPrompt = args.prompt;
+        if (args.enhancer_preset) {
+          const enhancer = resolveEnhancer(args.enhancer_preset);
+          if (enhancer.hasTransformations()) {
+            enhancedPrompt = enhancer.enhance(args.prompt);
+          }
+        }
+
         const imagePart = await fileToGenerativePart(args.image_path);
-        const contents: any[] = [args.prompt, imagePart];
+        const contents: any[] = [enhancedPrompt, imagePart];
 
         if (args.reference_images) {
           for (const refPath of args.reference_images) {

package/src/tools/get-status.ts

@@ -0,0 +1,174 @@
+/**
+ * Unified status checking tool for FAL and Vertex AI generation operations.
+ * This tool allows resuming/checking status of operations that timed out.
+ */
+
+import { z } from "zod";
+import { safeToolExecute } from "../utils/tool-wrapper";
+import { getGoogleAccessToken } from "../utils/google-auth";
+
+const FAL_KEY = process.env.FAL_KEY;
+
+/**
+ * Check FAL generation status using the status URL
+ */
+async function checkFalStatus(statusUrl: string): Promise<any> {
+  if (!FAL_KEY) {
+    throw new Error("FAL_KEY environment variable not set");
+  }
+
+  const response = await fetch(statusUrl, {
+    method: "GET",
+    headers: {
+      Authorization: `Key ${FAL_KEY}`,
+      "Content-Type": "application/json",
+    },
+  });
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    throw new Error(`FAL API error [${response.status}]: ${errorText}`);
+  }
+
+  const statusResult = (await response.json()) as { status?: string };
+
+  if (statusResult.status === "COMPLETED") {
+    // Fetch the actual result
+    const responseUrl = statusUrl.replace(/\/status$/, "");
+    const resultResponse = await fetch(responseUrl, {
+      method: "GET",
+      headers: {
+        Authorization: `Key ${FAL_KEY}`,
+        "Content-Type": "application/json",
+      },
+    });
+
+    if (resultResponse.ok) {
+      return await resultResponse.json();
+    }
+  }
+
+  return statusResult;
+}
+
+/**
+ * Check Vertex AI operation status
+ */
+async function checkVertexStatus(
+  operationName: string,
+  projectId: string,
+  locationId: string
+): Promise<any> {
+  const accessToken = await getGoogleAccessToken();
+
+  const operationsUrl = `https://${locationId}-aiplatform.googleapis.com/v1/projects/${projectId}/locations/${locationId}/publishers/google/models/veo-3.1-generate-preview/operations/${operationName}`;
+
+  const response = await fetch(operationsUrl, {
+    method: "GET",
+    headers: {
+      Authorization: `Bearer ${accessToken}`,
+      "Content-Type": "application/json",
+    },
+  });
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    throw new Error(`Vertex AI API error [${response.status}]: ${errorText}`);
+  }
+
+  return await response.json();
+}
+
+export const getGenerationStatus = {
+  name: "get_generation_status",
+  description:
+    "Check the status or retrieve the result of a generation operation that was started by 'fal_generate' or 'generateVideoi2v'. " +
+    "Use this when the original generation tool returned an 'IN_PROGRESS' status with a 'resume_id'. " +
+    "Pass the resume_id exactly as it was returned. " +
+    "For FAL operations, the resume_id is a full URL. " +
+    "For Vertex AI operations, the resume_id is an operation name.",
+  parameters: z.object({
+    resume_id: z
+      .string()
+      .describe(
+        "The resume_id returned by the original generation tool. " +
+          "For FAL: This is a full URL (starts with 'https://queue.fal.run/...'). " +
+          "For Vertex AI: This is an operation name."
+      ),
+    source: z
+      .enum(["fal", "vertex", "auto"])
+      .optional()
+      .default("auto")
+      .describe(
+        "Source of the operation: 'fal' for FAL AI, 'vertex' for Google Vertex AI, or 'auto' to auto-detect based on resume_id format."
+      ),
+    project_id: z
+      .string()
+      .optional()
+      .default("mixio-pro")
+      .describe("GCP Project ID (only needed for Vertex AI operations)."),
+    location_id: z
+      .string()
+      .optional()
+      .default("us-central1")
+      .describe("GCP region (only needed for Vertex AI operations)."),
+  }),
+  timeoutMs: 30000, // 30 seconds for status check
+  execute: async (args: {
+    resume_id: string;
+    source?: "fal" | "vertex" | "auto";
+    project_id?: string;
+    location_id?: string;
+  }) => {
+    return safeToolExecute(async () => {
+      const {
+        resume_id,
+        source = "auto",
+        project_id = "mixio-pro",
+        location_id = "us-central1",
+      } = args;
+
+      // Auto-detect source based on resume_id format
+      let detectedSource = source;
+      if (source === "auto") {
+        if (
+          resume_id.startsWith("https://queue.fal.run") ||
+          resume_id.startsWith("https://fal.run")
+        ) {
+          detectedSource = "fal";
+        } else {
+          detectedSource = "vertex";
+        }
+      }
+
+      let result: any;
+
+      if (detectedSource === "fal") {
+        result = await checkFalStatus(resume_id);
+      } else {
+        result = await checkVertexStatus(resume_id, project_id, location_id);
+      }
+
+      // Normalize the response
+      const status =
+        result.status || (result.done ? "COMPLETED" : "IN_PROGRESS");
+
+      return JSON.stringify(
+        {
+          source: detectedSource,
+          status,
+          resume_id,
+          result,
+          message:
+            status === "COMPLETED"
+              ? "Generation completed! The result is included in the 'result' field."
+              : status === "FAILED"
+              ? "Generation failed. Check the 'result' field for error details."
+              : "Generation is still in progress. Call this tool again with the same resume_id to check later.",
+        },
+        null,
+        2
+      );
+    }, "get_generation_status");
+  },
+};