@ssweens/pi-vertex 1.0.1 → 1.1.1

package/streaming/gemini.ts

@@ -1,20 +1,49 @@
 /**
  * Gemini streaming handler using @google/genai SDK
+ *
+ * Aligned with pi-mono's google-vertex.ts for consistent handling of:
+ * - Thinking content (thought blocks with signatures)
+ * - Tool calls with unique IDs and deduplication
+ * - Thinking configuration (levels for Gemini 3, budgets for Gemini 2.5)
+ * - Usage tracking including thinking tokens
  */
 
-import { GoogleGenAI } from "@google/genai";
-import type { VertexModelConfig, Context, StreamOptions } from "../types.js";
+import { GoogleGenAI, FinishReason, ThinkingLevel } from "@google/genai";
+import type { VertexModelConfig, Context, StreamOptions, AssistantMessage } from "../types.js";
 import { getAuthConfig, resolveLocation } from "../auth.js";
-import { sanitizeText, convertToGeminiMessages, calculateCost } from "../utils.js";
-import { createAssistantMessageEventStream, type AssistantMessageEventStream, type AssistantMessage } from "@mariozechner/pi-ai";
+import { sanitizeText, convertToGeminiMessages, convertToolsForGemini, retainThoughtSignature, calculateCost } from "../utils.js";
+import { createAssistantMessageEventStream, type AssistantMessageEventStream } from "@mariozechner/pi-ai";
+
+// Module-level counter for generating unique tool call IDs (matches pi-mono pattern)
+let toolCallCounter = 0;
+
+const THINKING_LEVEL_MAP: Record<string, ThinkingLevel> = {
+  minimal: ThinkingLevel.MINIMAL,
+  low: ThinkingLevel.LOW,
+  medium: ThinkingLevel.MEDIUM,
+  high: ThinkingLevel.HIGH,
+};
+
+function mapGeminiStopReason(reason: string): "stop" | "length" | "toolUse" | "error" {
+  switch (reason) {
+    case FinishReason.STOP:
+      return "stop";
+    case FinishReason.MAX_TOKENS:
+      return "length";
+    case FinishReason.SAFETY:
+    case FinishReason.RECITATION:
+    default:
+      return "error";
+  }
+}
 
 export function streamGemini(
   model: VertexModelConfig,
   context: Context,
-  options?: StreamOptions
+  options?: StreamOptions,
 ): AssistantMessageEventStream {
   const stream = createAssistantMessageEventStream();
-  
+
   (async () => {
     const output: AssistantMessage = {
       role: "assistant",
@@ -33,123 +62,203 @@ export function streamGemini(
       stopReason: "stop",
       timestamp: Date.now(),
     };
-    
+
     try {
       // Priority: config file > env var > model region > default
       const location = resolveLocation(model.region);
       const auth = getAuthConfig(location);
 
-      // Create client
+      // Create client with explicit API version (matches pi-mono)
       const client = new GoogleGenAI({
         vertexai: true,
         project: auth.projectId,
         location: auth.location,
+        apiVersion: "v1",
       });
-      
-      // Convert messages
-      const contents = convertToGeminiMessages(context.messages);
-      
-      // Build config
+
+      // Convert messages with model ID for proper thinking/tool handling
+      const contents = convertToGeminiMessages(context.messages, model.apiId);
+
+      // Build config — only set temperature when explicitly provided
       const config: any = {
         maxOutputTokens: options?.maxTokens || Math.floor(model.maxTokens / 2),
-        temperature: options?.temperature ?? 0.7,
+        ...(options?.temperature !== undefined && { temperature: options.temperature }),
       };
-      
+
       // Add system prompt if present
       if (context.systemPrompt) {
         config.systemInstruction = sanitizeText(context.systemPrompt);
       }
-      
-      // Add tools if present
+
+      // Add tools if present (using parametersJsonSchema for full JSON Schema support)
       if (context.tools && context.tools.length > 0) {
-        config.tools = [
-          {
-            functionDeclarations: context.tools.map((tool) => ({
-              name: tool.name,
-              description: tool.description,
-              parameters: tool.parameters,
-            })),
-          },
-        ];
+        config.tools = convertToolsForGemini(context.tools);
+      }
+
+      // Add thinking configuration (matches pi-mono's buildParams logic)
+      if (model.reasoning && options?.reasoning) {
+        const effort = options.reasoning === "xhigh" ? "high" : options.reasoning;
+        const isGemini3 = model.apiId.startsWith("gemini-3");
+
+        const thinkingConfig: any = { includeThoughts: true };
+
+        if (isGemini3) {
+          // Gemini 3 models use thinking levels (MINIMAL/LOW/MEDIUM/HIGH)
+          thinkingConfig.thinkingLevel = THINKING_LEVEL_MAP[effort];
+        } else {
+          // Gemini 2.5 models use thinking budgets (token counts)
+          const budgets: Record<string, number> = {
+            minimal: 128,
+            low: 2048,
+            medium: 8192,
+            high: model.apiId.includes("2.5-pro") ? 32768 : 24576,
+          };
+          thinkingConfig.thinkingBudget = budgets[effort] ?? 8192;
+        }
+
+        config.thinkingConfig = thinkingConfig;
+      }
+
+      // Pass abort signal to SDK for in-flight cancellation
+      if (options?.signal) {
+        if (options.signal.aborted) {
+          throw new Error("Request aborted");
+        }
+        config.abortSignal = options.signal;
       }
-      
+
       stream.push({ type: "start", partial: output });
-      
+
       // Start streaming
       const response = await client.models.generateContentStream({
         model: model.apiId,
         contents,
         config,
       });
-      
-      let textContent = "";
-      let textIndex = 0;
-      
+
+      // Track current content block for thinking/text transitions
+      let currentBlock: any = null;
+      let currentBlockType: "text" | "thinking" | null = null;
+
       for await (const chunk of response) {
-        if (options?.signal?.aborted) {
-          throw new Error("Request was aborted");
-        }
-        
-        // Update usage
-        if (chunk.usageMetadata) {
-          output.usage.input = chunk.usageMetadata.promptTokenCount || output.usage.input;
-          output.usage.output = chunk.usageMetadata.candidatesTokenCount || output.usage.output;
-          output.usage.totalTokens = chunk.usageMetadata.totalTokenCount || 
-            (output.usage.input + output.usage.output);
-          calculateCost(model.cost.input, model.cost.output, model.cost.cacheRead, model.cost.cacheWrite, output.usage);
-        }
-        
-        // Handle text
-        const text = chunk.text;
-        if (text) {
-          if (!textContent) {
-            // First text chunk
-            output.content.push({ type: "text", text: "" });
-            textIndex = output.content.length - 1;
-            stream.push({ type: "text_start", contentIndex: textIndex, partial: output });
-          }
-          textContent += text;
-          (output.content[textIndex] as any).text = textContent;
-          stream.push({ type: "text_delta", contentIndex: textIndex, delta: text, partial: output });
-        }
-        
-        // Handle function calls (tools)
-        if (chunk.functionCalls && chunk.functionCalls.length > 0) {
-          for (const call of chunk.functionCalls) {
-            output.content.push({
-              type: "toolCall",
-              id: call.id || `call_${Date.now()}`,
-              name: call.name,
-              arguments: call.args || {},
-            });
-            stream.push({
-              type: "toolcall_end",
-              contentIndex: output.content.length - 1,
-              toolCall: output.content[output.content.length - 1] as any,
-              partial: output,
-            });
+        const candidate = chunk.candidates?.[0];
+
+        // Process individual parts (handles thinking vs text detection)
+        if (candidate?.content?.parts) {
+          for (const part of candidate.content.parts) {
+            if (part.text !== undefined) {
+              const isThinking = part.thought === true;
+              const targetType = isThinking ? "thinking" : "text";
+
+              // Check if we need to transition to a new block
+              if (currentBlockType !== targetType) {
+                // End previous block
+                if (currentBlock && currentBlockType) {
+                  if (currentBlockType === "text") {
+                    stream.push({ type: "text_end", contentIndex: output.content.length - 1, content: currentBlock.text, partial: output });
+                  } else {
+                    stream.push({ type: "thinking_end", contentIndex: output.content.length - 1, content: currentBlock.thinking, partial: output });
+                  }
+                }
+
+                // Start new block
+                if (isThinking) {
+                  currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
+                  output.content.push(currentBlock);
+                  stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output });
+                } else {
+                  currentBlock = { type: "text", text: "", textSignature: undefined };
+                  output.content.push(currentBlock);
+                  stream.push({ type: "text_start", contentIndex: output.content.length - 1, partial: output });
+                }
+                currentBlockType = targetType;
+              }
+
+              // Accumulate content
+              if (currentBlockType === "thinking") {
+                currentBlock.thinking += part.text;
+                currentBlock.thinkingSignature = retainThoughtSignature(currentBlock.thinkingSignature, part.thoughtSignature);
+                stream.push({ type: "thinking_delta", contentIndex: output.content.length - 1, delta: part.text, partial: output });
+              } else {
+                currentBlock.text += part.text;
+                currentBlock.textSignature = retainThoughtSignature(currentBlock.textSignature, part.thoughtSignature);
+                stream.push({ type: "text_delta", contentIndex: output.content.length - 1, delta: part.text, partial: output });
+              }
+            }
+
+            if (part.functionCall) {
+              // End current text/thinking block before tool call
+              if (currentBlock && currentBlockType) {
+                if (currentBlockType === "text") {
+                  stream.push({ type: "text_end", contentIndex: output.content.length - 1, content: currentBlock.text, partial: output });
+                } else {
+                  stream.push({ type: "thinking_end", contentIndex: output.content.length - 1, content: currentBlock.thinking, partial: output });
+                }
+                currentBlock = null;
+                currentBlockType = null;
+              }
+
+              // Generate unique tool call ID with dedup (matches pi-mono pattern)
+              const providedId = part.functionCall.id;
+              const needsNewId =
+                !providedId || output.content.some((b: any) => b.type === "toolCall" && b.id === providedId);
+              const toolCallId = needsNewId
+                ? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
+                : providedId;
+
+              const toolCall = {
+                type: "toolCall" as const,
+                id: toolCallId,
+                name: part.functionCall.name || "",
+                arguments: (part.functionCall.args as Record<string, any>) ?? {},
+                ...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
+              };
+
+              output.content.push(toolCall);
+              const idx = output.content.length - 1;
+              stream.push({ type: "toolcall_start", contentIndex: idx, partial: output });
+              stream.push({ type: "toolcall_delta", contentIndex: idx, delta: JSON.stringify(toolCall.arguments), partial: output });
+              stream.push({ type: "toolcall_end", contentIndex: idx, toolCall, partial: output });
+            }
           }
         }
-        
+
         // Handle finish reason
-        if (chunk.candidates && chunk.candidates[0]?.finishReason) {
-          const reason = chunk.candidates[0].finishReason;
-          if (reason === "STOP") {
-            output.stopReason = "stop";
-          } else if (reason === "MAX_TOKENS") {
-            output.stopReason = "length";
-          } else if (reason === "SAFETY") {
-            output.stopReason = "error";
+        if (candidate?.finishReason) {
+          output.stopReason = mapGeminiStopReason(candidate.finishReason);
+          if (candidate.finishReason === FinishReason.SAFETY) {
             output.errorMessage = "Content blocked by safety filters";
           }
+          // Override to toolUse if any tool calls are present (matches pi-mono)
+          if (output.content.some((b: any) => b.type === "toolCall")) {
+            output.stopReason = "toolUse";
+          }
+        }
+
+        // Update usage — include thoughtsTokenCount in output (matches pi-mono)
+        if (chunk.usageMetadata) {
+          const meta = chunk.usageMetadata as any;
+          output.usage = {
+            input: meta.promptTokenCount || 0,
+            output: (meta.candidatesTokenCount || 0) + (meta.thoughtsTokenCount || 0),
+            cacheRead: meta.cachedContentTokenCount || 0,
+            cacheWrite: 0,
+            totalTokens: meta.totalTokenCount || 0,
+            cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+          };
+          calculateCost(model.cost.input, model.cost.output, model.cost.cacheRead, model.cost.cacheWrite, output.usage);
         }
       }
-      
-      // End text if we had any
-      if (textContent) {
-        stream.push({ type: "text_end", contentIndex: textIndex, content: textContent, partial: output });
+
+      // End final block
+      if (currentBlock && currentBlockType) {
+        if (currentBlockType === "text") {
+          stream.push({ type: "text_end", contentIndex: output.content.length - 1, content: currentBlock.text, partial: output });
+        } else {
+          stream.push({ type: "thinking_end", contentIndex: output.content.length - 1, content: currentBlock.thinking, partial: output });
+        }
       }
-      
+
       stream.push({ type: "done", reason: output.stopReason as any, message: output });
       stream.end();
     } catch (error) {
@@ -159,6 +268,6 @@ export function streamGemini(
       stream.end();
     }
   })();
-  
+
   return stream;
 }

package/streaming/maas.ts

@@ -1,8 +1,10 @@
 /**
  * MaaS streaming handler for Claude and all other models
  * Uses OpenAI-compatible Chat Completions endpoint
- * 
- * Delegates to pi-ai's built-in OpenAI streaming implementation
+ *
+ * Delegates to pi-ai's built-in OpenAI streaming implementation.
+ * Uses model.apiId directly in the request (no global fetch interceptor)
+ * and patches the model ID back to the friendly name in response events.
  */
 
 import type { VertexModelConfig, Context, StreamOptions } from "../types.js";
@@ -12,12 +14,11 @@ import { createAssistantMessageEventStream, type AssistantMessageEventStream, ty
 export function streamMaaS(
   model: VertexModelConfig,
   context: Context,
-  options?: StreamOptions
+  options?: StreamOptions,
 ): AssistantMessageEventStream {
   const stream = createAssistantMessageEventStream();
 
   (async () => {
-    const originalFetch = globalThis.fetch;
     try {
       // Priority: config file > env var > model region > default
       const location = resolveLocation(model.region);
@@ -26,12 +27,12 @@ export function streamMaaS(
 
       const baseUrl = buildBaseUrl(auth.projectId, auth.location);
       const endpoint = `${baseUrl}/endpoints/openapi`;
+
       // Create a model object compatible with pi-ai's OpenAI streaming.
-      // Note: baseUrl must point to the OpenAPI root; pi-ai appends /chat/completions.
-      // Use model.id (registered name like "glm-5") so pi can restore sessions correctly.
-      // The actual API model name (apiId like "zai-org/glm-5-maas") is injected via fetch interceptor below.
+      // Use model.apiId directly so the correct model name goes in the request body.
+      // The friendly model.id is patched back into response events below for session persistence.
       const modelForPi: Model<"openai-completions"> = {
-        id: model.id,
+        id: model.apiId,
         name: model.name,
         api: "openai-completions",
         provider: "vertex",
@@ -51,21 +52,6 @@ export function streamMaaS(
         },
       };
 
-      // Intercept fetch to replace model.id with the actual API model name (apiId)
-      // pi-ai's streaming uses model.id in the request body, but Vertex MaaS needs the full publisher-prefixed name
-      globalThis.fetch = async (input: any, init?: any) => {
-        if (init?.body && typeof init.body === "string") {
-          try {
-            const body = JSON.parse(init.body);
-            if (body.model === model.id) {
-              body.model = model.apiId;
-              init = { ...init, body: JSON.stringify(body) };
-            }
-          } catch {}
-        }
-        return originalFetch(input, init);
-      };
-
       // Delegate to pi-ai's built-in OpenAI streaming
       const innerStream = streamSimpleOpenAICompletions(
         modelForPi,
@@ -74,19 +60,26 @@ export function streamMaaS(
           ...options,
           apiKey: accessToken,
           maxTokens: options?.maxTokens || Math.floor(model.maxTokens / 2),
-          temperature: options?.temperature ?? 0.7,
-        }
+          temperature: options?.temperature,
+        },
       );
 
-      // Forward all events from inner stream to outer stream
+      // Forward all events, patching model ID back to the friendly name
+      // so pi-coding-agent can restore sessions correctly.
       for await (const event of innerStream) {
+        if ("partial" in event && event.partial) {
+          event.partial.model = model.id;
+        }
+        if ("message" in event && event.message) {
+          event.message.model = model.id;
+        }
+        if ("error" in event && event.error && typeof event.error === "object") {
+          (event.error as any).model = model.id;
+        }
         stream.push(event);
       }
-      globalThis.fetch = originalFetch;
       stream.end();
-
     } catch (error) {
-      globalThis.fetch = originalFetch;
       stream.push({
         type: "error",
         reason: options?.signal?.aborted ? "aborted" : "error",

package/types.ts

@@ -1,7 +1,31 @@
 /**
  * Type definitions for pi-vertex extension
+ *
+ * Core message/content types are re-exported from pi-ai to ensure pi-vertex
+ * handles the full message structure (thinking blocks, tool calls, tool results)
+ * that pi-coding-agent passes through the streamSimple callback.
  */
 
+// Re-export core types from pi-ai
+export type {
+  AssistantMessage,
+  AssistantMessageEvent,
+  AssistantMessageEventStream,
+  Context,
+  ImageContent,
+  Message,
+  StopReason,
+  TextContent,
+  ThinkingContent,
+  Tool,
+  ToolCall,
+  ToolResultMessage,
+  Usage,
+  UserMessage,
+} from "@mariozechner/pi-ai";
+
+// Vertex-specific types
+
 export type ModelInputType = "text" | "image";
 export type EndpointType = "gemini" | "maas";
 
@@ -33,44 +57,9 @@ export interface AuthConfig {
   credentials?: string;
 }
 
-export type MessageRole = "user" | "assistant" | "system";
-
-export interface TextContent {
-  type: "text";
-  text: string;
-}
-
-export interface ImageContent {
-  type: "image";
-  mimeType: string;
-  data: string;
-}
-
-export type MessageContent = TextContent | ImageContent;
-
-export interface Message {
-  role: MessageRole;
-  content: string | MessageContent[];
-}
-
-export interface Tool {
-  name: string;
-  description: string;
-  parameters: Record<string, unknown>;
-}
-
-export interface Context {
-  systemPrompt?: string;
-  messages: Message[];
-  tools?: Tool[];
-}
-
 export interface StreamOptions {
   maxTokens?: number;
   temperature?: number;
   reasoning?: "minimal" | "low" | "medium" | "high" | "xhigh";
   signal?: AbortSignal;
 }
-
-// Re-export types from pi-ai for convenience
-export type { AssistantMessage, AssistantMessageEvent, AssistantMessageEventStream } from "@mariozechner/pi-ai";