@ssweens/pi-vertex 1.0.0

package/streaming/maas.ts

@@ -0,0 +1,97 @@
+/**
+ * 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
+ */
+
+import type { VertexModelConfig, Context, StreamOptions } from "../types.js";
+import { getAuthConfig, buildBaseUrl, getAccessToken, resolveLocation } from "../auth.js";
+import { createAssistantMessageEventStream, type AssistantMessageEventStream, type Model, streamSimpleOpenAICompletions } from "@mariozechner/pi-ai";
+
+export function streamMaaS(
+  model: VertexModelConfig,
+  context: Context,
+  options?: StreamOptions
+): AssistantMessageEventStream {
+  const stream = createAssistantMessageEventStream();
+
+  (async () => {
+    try {
+      // Priority: config file > env var > model region > default
+      const location = resolveLocation(model.region);
+      const auth = getAuthConfig(location);
+      const accessToken = await getAccessToken();
+
+      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.
+      const modelForPi: Model<"openai-completions"> = {
+        id: model.apiId, // Use the full API ID with publisher prefix
+        name: model.name,
+        api: "openai-completions",
+        provider: "vertex",
+        baseUrl: endpoint,
+        reasoning: model.reasoning,
+        input: model.input,
+        cost: model.cost,
+        contextWindow: model.contextWindow,
+        maxTokens: model.maxTokens,
+        headers: {},
+        compat: {
+          supportsStore: false,
+          supportsDeveloperRole: false,
+          supportsReasoningEffort: false,
+          maxTokensField: "max_tokens",
+          thinkingFormat: model.publisher === "qwen" ? "qwen" : "openai",
+        },
+      };
+
+      // Delegate to pi-ai's built-in OpenAI streaming
+      const innerStream = streamSimpleOpenAICompletions(
+        modelForPi,
+        context as any,
+        {
+          ...options,
+          apiKey: accessToken,
+          maxTokens: options?.maxTokens || Math.floor(model.maxTokens / 2),
+          temperature: options?.temperature ?? 0.7,
+        }
+      );
+
+      // Forward all events from inner stream to outer stream
+      for await (const event of innerStream) {
+        stream.push(event);
+      }
+      stream.end();
+
+    } catch (error) {
+      stream.push({
+        type: "error",
+        reason: options?.signal?.aborted ? "aborted" : "error",
+        error: {
+          role: "assistant",
+          content: [],
+          api: "openai-completions",
+          provider: "vertex",
+          model: model.id,
+          usage: {
+            input: 0,
+            output: 0,
+            cacheRead: 0,
+            cacheWrite: 0,
+            totalTokens: 0,
+            cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+          },
+          stopReason: options?.signal?.aborted ? "aborted" : "error",
+          errorMessage: error instanceof Error ? error.message : String(error),
+          timestamp: Date.now(),
+        },
+      });
+      stream.end();
+    }
+  })();
+
+  return stream;
+}

package/types.ts

@@ -0,0 +1,76 @@
+/**
+ * Type definitions for pi-vertex extension
+ */
+
+export type ModelInputType = "text" | "image";
+export type EndpointType = "gemini" | "maas";
+
+export interface ModelCost {
+  input: number;
+  output: number;
+  cacheRead: number;
+  cacheWrite: number;
+}
+
+export interface VertexModelConfig {
+  id: string;
+  name: string;
+  apiId: string;
+  publisher: string;
+  endpointType: EndpointType;
+  contextWindow: number;
+  maxTokens: number;
+  input: ModelInputType[];
+  reasoning: boolean;
+  tools: boolean;
+  cost: ModelCost;
+  region: string;
+}
+
+export interface AuthConfig {
+  projectId: string;
+  location: string;
+  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";

package/utils.ts

@@ -0,0 +1,194 @@
+/**
+ * Utility functions for pi-vertex extension
+ */
+
+import type { Message, MessageContent, TextContent, ToolCall, AssistantMessage } from "./types.js";
+
+/**
+ * Sanitize text by removing invalid surrogate pairs
+ */
+export function sanitizeText(text: string): string {
+  return text.replace(/[\uD800-\uDFFF]/g, "\uFFFD");
+}
+
+/**
+ * Convert messages to Gemini format
+ */
+export function convertToGeminiMessages(messages: Message[]): any[] {
+  const result: any[] = [];
+  
+  for (const msg of messages) {
+    if (msg.role === "user") {
+      if (typeof msg.content === "string") {
+        if (msg.content.trim()) {
+          result.push({
+            role: "user",
+            parts: [{ text: sanitizeText(msg.content) }],
+          });
+        }
+      } else {
+        const parts = msg.content.map((item) => {
+          if (item.type === "text") {
+            return { text: sanitizeText(item.text) };
+          } else {
+            return {
+              inlineData: {
+                mimeType: item.mimeType,
+                data: item.data,
+              },
+            };
+          }
+        });
+        result.push({ role: "user", parts });
+      }
+    } else if (msg.role === "assistant") {
+      // Gemini doesn't have a separate assistant role in the same way
+      // We'll handle this in the conversation history
+      if (typeof msg.content === "string") {
+        if (msg.content.trim()) {
+          result.push({
+            role: "model",
+            parts: [{ text: sanitizeText(msg.content) }],
+          });
+        }
+      }
+    }
+  }
+  
+  return result;
+}
+
+/**
+ * Convert messages to OpenAI-compatible format (for Claude and MaaS)
+ */
+export function convertToOpenAIMessages(messages: Message[]): any[] {
+  const result: any[] = [];
+  
+  for (const msg of messages) {
+    if (msg.role === "user") {
+      if (typeof msg.content === "string") {
+        if (msg.content.trim()) {
+          result.push({
+            role: "user",
+            content: sanitizeText(msg.content),
+          });
+        }
+      } else {
+        const content = msg.content.map((item) => {
+          if (item.type === "text") {
+            return { type: "text", text: sanitizeText(item.text) };
+          } else {
+            return {
+              type: "image_url",
+              image_url: {
+                url: `data:${item.mimeType};base64,${item.data}`,
+              },
+            };
+          }
+        });
+        result.push({ role: "user", content });
+      }
+    } else if (msg.role === "assistant") {
+      if (typeof msg.content === "string") {
+        if (msg.content.trim()) {
+          result.push({
+            role: "assistant",
+            content: sanitizeText(msg.content),
+          });
+        }
+      }
+    } else if (msg.role === "system") {
+      // System messages handled separately
+    }
+  }
+  
+  return result;
+}
+
+/**
+ * Convert tools to OpenAI format
+ */
+export function convertTools(tools: any[]): any[] {
+  return tools.map((tool) => ({
+    type: "function",
+    function: {
+      name: tool.name,
+      description: tool.description,
+      parameters: tool.parameters,
+    },
+  }));
+}
+
+/**
+ * Parse SSE (Server-Sent Events) stream
+ */
+export async function* parseSSEStream(response: Response): AsyncGenerator<string> {
+  const reader = response.body?.getReader();
+  if (!reader) {
+    throw new Error("No response body");
+  }
+
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (trimmed.startsWith("data: ")) {
+          const data = trimmed.slice(6);
+          if (data === "[DONE]") return;
+          yield data;
+        }
+      }
+    }
+
+    // Process remaining buffer
+    const trimmed = buffer.trim();
+    if (trimmed.startsWith("data: ")) {
+      const data = trimmed.slice(6);
+      if (data !== "[DONE]") {
+        yield data;
+      }
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+/**
+ * Map stop reason to standard format
+ */
+export function mapStopReason(reason: string): "stop" | "length" | "toolUse" | "error" {
+  switch (reason) {
+    case "stop":
+    case "end_turn":
+      return "stop";
+    case "length":
+    case "max_tokens":
+      return "length";
+    case "tool_calls":
+    case "tool_use":
+      return "toolUse";
+    default:
+      return "error";
+  }
+}
+
+/**
+ * Calculate cost based on usage and model cost config
+ */
+export function calculateCost(inputCost: number, outputCost: number, cacheReadCost: number, cacheWriteCost: number, usage: AssistantMessage["usage"]): void {
+  usage.cost.input = (inputCost / 1000000) * usage.input;
+  usage.cost.output = (outputCost / 1000000) * usage.output;
+  usage.cost.cacheRead = (cacheReadCost / 1000000) * usage.cacheRead;
+  usage.cost.cacheWrite = (cacheWriteCost / 1000000) * usage.cacheWrite;
+  usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite;
+}