gemini-compat-endpoint

// Tool-related helpers extracted from generator.ts to keep behavior unchanged.
// Includes schema conversion, tool gating, name sanitization, and response sanitization.

import { type GeneratorController } from "@lmstudio/sdk";
import { z } from "zod";

// Local types (scoped to tool schema transformation)
export type JsonSchema = any;
export type GeminiDataSchema = {
  type: string;
  properties?: Record<string, GeminiDataSchema>;
  required?: string[];
  items?: GeminiDataSchema;
  description?: string;
  enum?: string[] | number[] | boolean[];
};

export function convertJsonSchemaToGemini(schema: JsonSchema): GeminiDataSchema {
  if (!schema || typeof schema !== "object") return { type: "OBJECT" };
  const t = (schema.type || "object").toString().toLowerCase();
  switch (t) {
    case "string": {
      const ds: GeminiDataSchema = { type: "STRING" };
      if (schema.description) ds.description = schema.description;
      if (Array.isArray(schema.enum)) ds.enum = schema.enum;
      return ds;
    }
    case "number": {
      const ds: GeminiDataSchema = { type: "NUMBER" };
      if (schema.description) ds.description = schema.description;
      return ds;
    }
    case "integer": {
      const ds: GeminiDataSchema = { type: "INTEGER" };
      if (schema.description) ds.description = schema.description;
      return ds;
    }
    case "boolean": {
      const ds: GeminiDataSchema = { type: "BOOLEAN" };
      if (schema.description) ds.description = schema.description;
      return ds;
    }
    case "array": {
      const ds: GeminiDataSchema = { type: "ARRAY" };
      if (schema.description) ds.description = schema.description;
      if (schema.items) ds.items = convertJsonSchemaToGemini(schema.items);
      return ds;
    }
    case "object":
    default: {
      const ds: GeminiDataSchema = { type: "OBJECT" };
      if (schema.description) ds.description = schema.description;
      const props: Record<string, GeminiDataSchema> = {};
      const required: string[] = Array.isArray(schema.required) ? [...schema.required] : [];
      if (schema.properties && typeof schema.properties === "object") {
        for (const [k, v] of Object.entries(schema.properties)) {
          props[k] = convertJsonSchemaToGemini(v as JsonSchema);
        }
      }
      if (Object.keys(props).length) ds.properties = props;
      if (required.length) ds.required = required;
      return ds;
    }
  }
}

export function sanitizeToolName(name: string): string {
  // Gemini: must start with a letter, then letters/digits/_; max length 64
  let n = name || "tool";
  n = n.replace(/[^A-Za-z0-9_]/g, "_");
  if (!/^[A-Za-z]/.test(n)) n = "t_" + n;
  if (n.length > 64) n = n.slice(0, 64);
  return n;
}

export function pickRelevantToolDefs(defs: any[] | undefined, userText: string): any[] {
  if (!Array.isArray(defs) || !userText) return [];
  const t = userText.toLowerCase();
  const picked: any[] = [];
  for (const d of defs) {
    const name = String(d?.name ?? d?.toolName ?? "").toLowerCase();
    const desc = String(d?.description ?? "").toLowerCase();
    const hay = `${name} ${desc}`;
    // Simple conservative gating heuristics
    const img = /(image|bild|picture|photo|render|zeichnen|draw|paint|image2image)/i;
    const search = /(search|find|lookup|look\s*up|suche)/i;
    const http = /(fetch|download|http|https|request|url)/i;
    const translate = /(translate|übersetz|uebersetz)/i;
    const file = /(file|read|write|save|fs|filesystem)/i;

    let match = false;
    if (img.test(t) && /image|img|draw|render/.test(hay)) match = true;
    else if (search.test(t) && /search|find/.test(hay)) match = true;
    else if (http.test(t) && /fetch|http|request|download|url/.test(hay)) match = true;
    else if (translate.test(t) && /translat|uebersetz|übersetz/.test(hay)) match = true;
    else if (file.test(t) && /file|fs|filesystem|read|write/.test(hay)) match = true;

    if (match) picked.push(d);
  }
  return picked;
}

export function buildGeminiTools(ctl: GeneratorController, userText: string) {
  const defs = (ctl as any).getToolDefinitions?.() as any[] | undefined;
  let gated = pickRelevantToolDefs(defs, userText);
  // Fallback: if gating found nothing, include all tool definitions
  if ((!gated || gated.length === 0) && Array.isArray(defs) && defs.length) {
    gated = defs;
  }

  const functionDecls: any[] = [];
  const originalToSafe = new Map<string, string>();
  const safeToOriginal = new Map<string, string>();
  const safeNames: string[] = [];
  if (Array.isArray(gated) && gated.length) {
    for (const d of gated) {
      const origName = String(d?.function?.name ?? d?.name ?? d?.toolName ?? "tool");
      let safeName = sanitizeToolName(origName);
      // ensure uniqueness
      let i = 1;
      while (safeToOriginal.has(safeName) && safeToOriginal.get(safeName) !== origName) {
        const base = safeName.slice(0, Math.max(0, 60));
        safeName = `${base}_${++i}`;
      }
      originalToSafe.set(origName, safeName);
      safeToOriginal.set(safeName, origName);
      const description = d?.function?.description ?? d?.description ?? "";
      const parametersSchema = convertJsonSchemaToGemini(
        d?.function?.parameters ?? d?.parameters ?? d?.schema ?? { type: "object", properties: {} }
      );
      functionDecls.push({
        name: safeName,
        description,
        parameters: parametersSchema,
      });
      safeNames.push(safeName);
    }
  }
  const tools = functionDecls.length ? [{ functionDeclarations: functionDecls }] : undefined;
  return { tools, originalToSafe, safeToOriginal, safeNames } as const;
}

/**
 * Pass-through helper for tool results.
 * IMPORTANT:
 * - If we inject image markdown into the chat ourselves, we must prevent double-rendering.
 *   In that case we remove tool-returned image objects (incl. $hint/markdown) from the
 *   payload that is fed back into the model context.
 * - We also strip purely logistical Preview/Original path lines.
 */
export function sanitizeToolResponseForModel(payload: any, mode: 'A' | 'B' | 'C', chatWd?: string): any {
  if (payload === null || payload === undefined) return {};

  // Robust sanitizer: tool payloads come in many shapes (array, object with content/result/output arrays, nested).
  // We remove explicit tool image objects (incl. $hint/markdown) and strip purely logistical Preview/Original lines.

  const stripLogisticalText = (item: any): boolean => {
    if (!item || typeof item !== 'object') return false;
    if (item.type !== 'text') return false;
    if (typeof item.text !== 'string') return false;
    const t = item.text.trim();
    // Strip "Original vN:" and "Preview vN:" lines (contain file:// paths)
    if (/^Original v\d+:/i.test(t) || /^Preview v\d+:/i.test(t)) return true;
    // Strip JSON metadata objects containing local paths or sensitive fields
    if (t.startsWith('{')) {
      try {
        const parsed = JSON.parse(t);
        const stringified = JSON.stringify(parsed);
        // Check for path patterns that leak local filesystem info
        if (/lmstudio_attachment:|file:\/\/\/|\/Users\/|\/home\/|C:\\Users\\|\.lmstudio\//i.test(stringified)) {
          return true;
        }
        // Check for known metadata fields that contain local paths
        if (parsed.original_png_url || parsed.analysis_preview_url || parsed.source?.includes?.('/Users/')) {
          return true;
        }
      } catch {
        // Not valid JSON - check raw string for path patterns
        if (/file:\/\/\/|\/Users\/|\/home\/|C:\\Users\\/i.test(t)) {
          return true;
        }
      }
    }
    return false;
  };

  const sanitizeArray = (arr: any[]): any[] => {
    const hasToolImageObjects = arr.some((entry: any) => entry && typeof entry === 'object' && entry.type === 'image');
    return arr
      .filter((entry: any) => {
        if (!entry || typeof entry !== 'object') return true;
        if (hasToolImageObjects && entry.type === 'image') return false;
        if (stripLogisticalText(entry)) return false;
        return true;
      })
      .map(sanitizeAny);
  };

  const sanitizeObject = (obj: Record<string, any>): Record<string, any> => {
    const out: Record<string, any> = Array.isArray(obj) ? [] as any : {};
    for (const [k, v] of Object.entries(obj)) {
      out[k] = sanitizeAny(v);
    }
    return out;
  };

  const sanitizeAny = (v: any): any => {
    if (v === null || v === undefined) return v;
    if (Array.isArray(v)) return sanitizeArray(v);
    if (typeof v === 'object') return sanitizeObject(v as any);
    // If it's a JSON-encoded string containing an array/object, try to sanitize it too.
    if (typeof v === 'string') {
      const s = v.trim();
      if ((s.startsWith('[') && s.endsWith(']')) || (s.startsWith('{') && s.endsWith('}'))) {
        try {
          const parsed = JSON.parse(s);
          const sanitized = sanitizeAny(parsed);
          return sanitized;
        } catch {
          return v;
        }
      }
    }
    return v;
  };

  try {
    return sanitizeAny(payload);
  } catch {
    return payload;
  }
}