master-ai/vibn-frontend/lib/ai/gemini-chat.ts

import { GoogleGenAI } from '@google/genai';

const GEMINI_API_KEY = process.env.GOOGLE_API_KEY || "";
const GEMINI_MODEL = process.env.VIBN_CHAT_MODEL || "gemini-3.1-pro-preview";

if (!GEMINI_API_KEY) {
  console.warn(`[GeminiChat] WARNING: GOOGLE_API_KEY is not set. Chat stream will fail with 403 Forbidden.`);
}

const ai = new GoogleGenAI({ apiKey: GEMINI_API_KEY });

export interface ChatMessage {
  role: "user" | "assistant" | "tool";
  content: string;
  toolCalls?: ToolCall[];
  toolCallId?: string;
  toolName?: string;
  thoughtSignature?: string;
}

export interface ToolCall {
  id: string;
  name: string;
  args: Record<string, unknown>;
  thoughtSignature?: string;
}

export interface ToolDefinition {
  name: string;
  description: string;
  parameters: Record<string, unknown>;
}

export interface ChatChunk {
  type: "text" | "thinking" | "tool_call" | "done" | "error";
  text?: string;
  toolCall?: ToolCall;
  error?: string;
}

function toGeminiContents(messages: ChatMessage[]) {
  const contents: any[] = [];

  for (const msg of messages) {
    if (msg.role === "user") {
      contents.push({ role: "user", parts: [{ text: msg.content }] });
    } else if (msg.role === "assistant") {
      const parts: any[] = [];
      if (msg.content) parts.push({ text: msg.content });
      if (msg.toolCalls?.length) {
        for (const tc of msg.toolCalls) {
          const part: any = {
            functionCall: { name: tc.name, args: tc.args },
          };
          if (tc.thoughtSignature) {
            part.thoughtSignature = tc.thoughtSignature;
          }
          parts.push(part);
        }
      }
      if (parts.length) contents.push({ role: "model", parts });
    } else if (msg.role === "tool") {
      const part = {
        functionResponse: {
          name: msg.toolName || "unknown",
          response: { name: msg.toolName || "unknown", content: msg.content },
        },
      };
      const last = contents[contents.length - 1];
      if (last?.role === "user") {
        last.parts.push(part);
      } else {
        contents.push({ role: "user", parts: [part] });
      }
    }
  }
  return contents;
}

function toGeminiFunctions(tools: ToolDefinition[]) {
  if (!tools.length) return undefined;
  return [
    {
      functionDeclarations: tools.map((t) => ({
        name: t.name,
        description: t.description,
        parameters: t.parameters as any,
      })),
    },
  ];
}

export async function callGeminiChat(opts: {
  systemPrompt: string;
  messages: ChatMessage[];
  tools?: ToolDefinition[];
  temperature?: number;
  includeThoughts?: boolean;
}): Promise<{
  text: string;
  thoughts: string;
  toolCalls: ToolCall[];
  finishReason?: string;
  error?: string;
}> {
  try {
    const config: any = {
      temperature: opts.temperature ?? 0.7,
      maxOutputTokens: 8192,
    };
    
    if (opts.systemPrompt) {
        config.systemInstruction = opts.systemPrompt;
    }

    if (opts.includeThoughts) {
        config.thinkingConfig = { thinkingBudget: 1024 };
    }

    const fns = toGeminiFunctions(opts.tools ?? []);
    if (fns) config.tools = fns;

    const response = await ai.models.generateContent({
      model: GEMINI_MODEL,
      contents: toGeminiContents(opts.messages),
      config
    });

    let text = "";
    let thoughts = "";
    const toolCalls: ToolCall[] = [];
    
    const parts = response.candidates?.[0]?.content?.parts ?? [];

    for (const part of parts) {
      if (part.text) {
        if ((part as any).thought) thoughts += part.text;
        else text += part.text;
      }
      if (part.functionCall) {
        toolCalls.push({
          id: `tc-${Date.now()}-${Math.random().toString(36).slice(2)}`,
          name: part.functionCall.name,
          args: part.functionCall.args as Record<string, unknown> ?? {},
          thoughtSignature: (part as any).thoughtSignature,
        });
      }
    }

    return { 
      text, 
      thoughts, 
      toolCalls, 
      finishReason: response.candidates?.[0]?.finishReason 
    };

  } catch (error) {
    return {
      text: "",
      thoughts: "",
      toolCalls: [],
      error: `GoogleGenAI error: ${error instanceof Error ? error.message : String(error)}`,
    };
  }
}

export async function* streamGeminiChat(opts: {
  systemPrompt: string;
  messages: ChatMessage[];
  tools?: ToolDefinition[];
  temperature?: number;
}): AsyncGenerator<ChatChunk> {
  try {
    const config: any = {
      temperature: opts.temperature ?? 0.7,
      maxOutputTokens: 8192,
      thinkingConfig: { thinkingBudget: 1024 },
    };

    if (opts.systemPrompt) {
        config.systemInstruction = opts.systemPrompt;
    }

    const fns = toGeminiFunctions(opts.tools ?? []);
    if (fns) config.tools = fns;

    const streamResult = await ai.models.generateContentStream({
      model: GEMINI_MODEL,
      contents: toGeminiContents(opts.messages),
      config
    });

    for await (const chunk of streamResult) {
      const parts = chunk.candidates?.[0]?.content?.parts ?? [];
      for (const part of parts) {
        if (part.text) {
          yield (part as any).thought
            ? { type: "thinking", text: part.text }
            : { type: "text", text: part.text };
        }
      }
    }

    yield { type: "done" };

  } catch (error) {
    yield {
      type: "error",
      error: `GoogleGenAI error: ${error instanceof Error ? error.message : String(error)}`,
    };
  }
}