@auto-engineer/ai-gateway 0.1.1

package/src/index.ts

@@ -0,0 +1,499 @@
+import { generateText, streamText, generateObject, streamObject } from 'ai';
+import { openai } from '@ai-sdk/openai';
+import { anthropic } from '@ai-sdk/anthropic';
+import { google } from '@ai-sdk/google';
+import { xai } from '@ai-sdk/xai';
+import { configureAIProvider } from './config';
+import { z } from 'zod';
+import { getRegisteredToolsForAI } from './mcp-server';
+import { startServer } from './mcp-server';
+
+interface AIToolValidationError extends Error {
+  cause?: {
+    issues?: unknown[];
+    [key: string]: unknown;
+  };
+  issues?: unknown[];
+  errors?: unknown[];
+  zodIssues?: unknown[];
+  validationDetails?: {
+    cause?: {
+      issues?: unknown[];
+    };
+    issues?: unknown[];
+    errors?: unknown[];
+  };
+  schemaDescription?: string;
+  [key: string]: unknown;
+}
+
+export enum AIProvider {
+  OpenAI = 'openai',
+  Anthropic = 'anthropic',
+  Google = 'google',
+  XAI = 'xai',
+}
+
+export interface AIOptions {
+  model?: string;
+  temperature?: number;
+  maxTokens?: number;
+  streamCallback?: (token: string) => void;
+  includeTools?: boolean;
+}
+
+export interface StructuredAIOptions<T> extends Omit<AIOptions, 'streamCallback'> {
+  schema: z.ZodSchema<T>;
+  schemaName?: string;
+  schemaDescription?: string;
+}
+
+export interface StreamStructuredAIOptions<T> extends StructuredAIOptions<T> {
+  onPartialObject?: (partialObject: unknown) => void;
+}
+
+const defaultOptions: AIOptions = {
+  temperature: 0.7,
+  maxTokens: 1000,
+};
+
+function getDefaultModel(provider: AIProvider): string {
+  switch (provider) {
+    case AIProvider.OpenAI:
+      return 'gpt-4o-mini';
+    case AIProvider.Anthropic:
+      return 'claude-sonnet-4-20250514';
+    case AIProvider.Google:
+      return 'gemini-2.5-pro';
+    case AIProvider.XAI:
+      return 'grok-3';
+    default:
+      // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
+      throw new Error(`Unknown provider: ${provider}`);
+  }
+}
+
+function getModel(provider: AIProvider, model?: string) {
+  const modelName = model ?? getDefaultModel(provider);
+
+  switch (provider) {
+    case AIProvider.OpenAI:
+      return openai(modelName);
+    case AIProvider.Anthropic:
+      return anthropic(modelName);
+    case AIProvider.Google:
+      return google(modelName);
+    case AIProvider.XAI:
+      return xai(modelName);
+    default:
+      throw new Error(`Unknown provider: ${provider as string}`);
+  }
+}
+
+export async function generateTextWithAI(
+  prompt: string,
+  provider: AIProvider,
+  options: AIOptions = {},
+): Promise<string> {
+  const finalOptions = { ...defaultOptions, ...options };
+  const model = finalOptions.model ?? getDefaultModel(provider);
+  const modelInstance = getModel(provider, model);
+
+  if (finalOptions.includeTools === true) {
+    await startServer();
+    const result = await generateTextWithToolsAI(prompt, provider, options);
+
+    return result.text;
+  }
+
+  const result = await generateText({
+    model: modelInstance,
+    prompt,
+    temperature: finalOptions.temperature,
+    maxTokens: finalOptions.maxTokens,
+  });
+
+  return result.text;
+}
+
+export async function* streamTextWithAI(
+  prompt: string,
+  provider: AIProvider,
+  options: AIOptions = {},
+): AsyncGenerator<string> {
+  const finalOptions = { ...defaultOptions, ...options };
+  const model = getModel(provider, finalOptions.model);
+
+  const stream = await streamText({
+    model,
+    prompt,
+    temperature: finalOptions.temperature,
+    maxTokens: finalOptions.maxTokens,
+  });
+
+  for await (const chunk of stream.textStream) {
+    yield chunk;
+  }
+}
+
+/**
+ * Generates text using streaming and collects the full response.
+ * Optionally calls a stream callback for each token if provided.
+ * Always returns the complete collected response.
+ */
+export async function generateTextStreamingWithAI(
+  prompt: string,
+  provider: AIProvider,
+  options: AIOptions = {},
+): Promise<string> {
+  const finalOptions = { ...defaultOptions, ...options };
+  let collectedResult = '';
+
+  const stream = streamTextWithAI(prompt, provider, finalOptions);
+
+  for await (const token of stream) {
+    // Collect all tokens for the final result
+    collectedResult += token;
+
+    // Call the stream callback if provided
+    if (finalOptions.streamCallback) {
+      finalOptions.streamCallback(token);
+    }
+  }
+
+  return collectedResult;
+}
+
+export async function generateTextWithToolsAI(
+  prompt: string,
+  provider: AIProvider,
+  options: AIOptions = {},
+): Promise<{ text: string; toolCalls?: unknown[] }> {
+  const finalOptions = { ...defaultOptions, ...options };
+  const model = finalOptions.model ?? getDefaultModel(provider);
+  const modelInstance = getModel(provider, model);
+
+  const registeredTools = finalOptions.includeTools === true ? getRegisteredToolsForAI() : {};
+  console.log('registeredTools', registeredTools);
+  const hasTools = Object.keys(registeredTools).length > 0;
+  console.log('hasTools', hasTools);
+
+  // Build conversation messages
+  const messages: Array<{ role: 'user' | 'assistant'; content: string }> = [{ role: 'user', content: prompt }];
+
+  let finalResult = '';
+  const allToolCalls: unknown[] = [];
+  let attempts = 0;
+  const maxAttempts = 5;
+
+  while (attempts < maxAttempts) {
+    attempts++;
+
+    const opts = {
+      model: modelInstance,
+      messages,
+      temperature: finalOptions.temperature,
+      maxTokens: finalOptions.maxTokens,
+      ...(hasTools && {
+        tools: registeredTools,
+        toolChoice: 'auto' as const,
+      }),
+    };
+    console.log('opts', opts);
+    const result = await generateText(opts);
+    console.log('result', JSON.stringify(result, null, 2));
+
+    // Add assistant message to conversation
+    if (result.text) {
+      messages.push({ role: 'assistant', content: result.text });
+      finalResult = result.text;
+    }
+
+    // If there are tool calls, execute them and continue conversation
+    if (result.toolCalls !== undefined && result.toolCalls.length > 0) {
+      allToolCalls.push(...result.toolCalls);
+
+      // Execute tools and create a simple follow-up prompt
+      const toolResults = await executeToolCalls(result.toolCalls, registeredTools);
+
+      // Add the tool results as a user message and request a final response
+      messages.push({
+        role: 'user',
+        content: `${toolResults}Based on this product catalog data, please provide specific product recommendations for a soccer-loving daughter. Include product names, prices, and reasons why each item would be suitable.`,
+      });
+
+      // Continue the conversation to get AI's response to tool results
+      continue;
+    }
+
+    // If no tool calls, we're done
+    break;
+  }
+
+  return {
+    text: finalResult,
+    toolCalls: allToolCalls,
+  };
+}
+
+async function executeToolCalls(
+  toolCalls: unknown[],
+  registeredTools: Record<
+    string,
+    { execute?: (args: Record<string, unknown>) => Promise<string>; description?: string }
+  >,
+): Promise<string> {
+  let toolResults = '';
+
+  for (const toolCall of toolCalls) {
+    try {
+      const toolCallObj = toolCall as { toolName: string; args: Record<string, unknown> };
+      const tool = registeredTools[toolCallObj.toolName];
+      if (tool?.execute) {
+        const toolResult = await tool.execute(toolCallObj.args);
+        toolResults += `Tool ${toolCallObj.toolName} returned: ${String(toolResult)}\n\n`;
+      } else {
+        toolResults += `Error: Tool ${toolCallObj.toolName} not found or missing execute function\n\n`;
+      }
+    } catch (error) {
+      const toolCallObj = toolCall as { toolName: string };
+      console.error(`Tool execution error for ${toolCallObj.toolName}:`, error);
+      toolResults += `Error executing tool ${toolCallObj.toolName}: ${String(error)}\n\n`;
+    }
+  }
+
+  return toolResults;
+}
+
+export async function generateTextWithImageAI(
+  text: string,
+  imageBase64: string,
+  provider: AIProvider,
+  options: AIOptions = {},
+): Promise<string> {
+  const finalOptions = { ...defaultOptions, ...options };
+  const model = finalOptions.model ?? getDefaultModel(provider);
+  const modelInstance = getModel(provider, model);
+
+  if (provider !== AIProvider.OpenAI && provider !== AIProvider.XAI) {
+    throw new Error(`Provider ${provider} does not support image inputs`);
+  }
+
+  const result = await generateText({
+    model: modelInstance,
+    messages: [
+      {
+        role: 'user',
+        content: [
+          { type: 'text', text },
+          { type: 'image', image: imageBase64 },
+        ],
+      },
+    ],
+    temperature: finalOptions.temperature,
+    maxTokens: finalOptions.maxTokens,
+  });
+
+  return result.text;
+}
+
+export function getAvailableProviders(): AIProvider[] {
+  const config = configureAIProvider();
+  const providers: AIProvider[] = [];
+  if (config.openai != null) providers.push(AIProvider.OpenAI);
+  if (config.anthropic != null) providers.push(AIProvider.Anthropic);
+  if (config.google != null) providers.push(AIProvider.Google);
+  if (config.xai != null) providers.push(AIProvider.XAI);
+  return providers;
+}
+
+function getEnhancedPrompt(prompt: string, lastError: AIToolValidationError): string {
+  const errorDetails = lastError.zodIssues
+    ? JSON.stringify(lastError.zodIssues, null, 2)
+    : lastError.validationDetails?.cause?.issues
+      ? JSON.stringify(lastError.validationDetails.cause.issues, null, 2)
+      : lastError.message;
+
+  return `${prompt}\\n\\n⚠️ IMPORTANT: Your previous response failed validation with the following errors:\\n${errorDetails}\\n\\nPlease fix these errors and ensure your response EXACTLY matches the required schema structure.`;
+}
+
+function isSchemaError(error: Error): boolean {
+  return (
+    error.message.includes('response did not match schema') ||
+    error.message.includes('TypeValidationError') ||
+    error.name === 'AI_TypeValidationError'
+  );
+}
+
+function enhanceValidationError(error: AIToolValidationError): AIToolValidationError {
+  const enhancedError = new Error(error.message) as AIToolValidationError;
+  Object.assign(enhancedError, error);
+
+  if (error.cause !== undefined || error.issues !== undefined || error.errors !== undefined) {
+    enhancedError.validationDetails = {
+      cause: error.cause,
+      issues: error.issues,
+      errors: error.errors,
+    };
+  }
+
+  if (
+    error.message.includes('response did not match schema') &&
+    'cause' in error &&
+    typeof error.cause === 'object' &&
+    error.cause !== null &&
+    'issues' in error.cause
+  ) {
+    enhancedError.zodIssues = error.cause.issues as unknown[];
+  }
+  return enhancedError;
+}
+
+function handleFailedRequest(
+  lastError: AIToolValidationError | undefined,
+  maxRetries: number,
+  attempt: number,
+): { shouldRetry: boolean; enhancedError?: AIToolValidationError } {
+  if (!lastError || !isSchemaError(lastError) || attempt >= maxRetries - 1) {
+    return { shouldRetry: false, enhancedError: lastError };
+  }
+
+  console.log(`Schema validation failed on attempt ${attempt + 1}/${maxRetries}, retrying...`);
+  const enhancedError = enhanceValidationError(lastError);
+
+  return { shouldRetry: true, enhancedError };
+}
+
+export async function generateStructuredDataWithAI<T>(
+  prompt: string,
+  provider: AIProvider,
+  options: StructuredAIOptions<T>,
+): Promise<T> {
+  const maxSchemaRetries = 3;
+  let lastError: AIToolValidationError | undefined;
+
+  if (options.includeTools === true) await startServer();
+  const registeredTools = options.includeTools === true ? getRegisteredToolsForAI() : {};
+  console.log('registeredTools', registeredTools);
+  const hasTools = Object.keys(registeredTools).length > 0;
+  console.log('hasTools', hasTools);
+
+  for (let attempt = 0; attempt < maxSchemaRetries; attempt++) {
+    try {
+      const model = getModel(provider, options.model);
+
+      const enhancedPrompt = attempt > 0 && lastError ? getEnhancedPrompt(prompt, lastError) : prompt;
+      console.log('enhancedPrompt', enhancedPrompt);
+
+      const opts = {
+        model,
+        prompt: enhancedPrompt,
+        schema: options.schema,
+        schemaName: options.schemaName,
+        schemaDescription: options.schemaDescription,
+        temperature: options.temperature,
+        maxTokens: options.maxTokens,
+        ...(hasTools && {
+          tools: registeredTools,
+          toolChoice: 'auto' as const,
+        }),
+      };
+      console.log('opts', opts);
+      const result = await generateObject(opts);
+      console.log('result', JSON.stringify(result, null, 2));
+      return result.object;
+    } catch (error: unknown) {
+      lastError =
+        error instanceof Error
+          ? (error as AIToolValidationError)
+          : (new Error('An unknown error occurred') as AIToolValidationError);
+
+      const { shouldRetry, enhancedError } = handleFailedRequest(lastError, maxSchemaRetries, attempt);
+      lastError = enhancedError;
+
+      if (!shouldRetry) {
+        throw lastError;
+      }
+    }
+  }
+
+  // If we exhausted all retries, throw the last error
+  throw lastError;
+}
+
+export async function streamStructuredDataWithAI<T>(
+  prompt: string,
+  provider: AIProvider,
+  options: StreamStructuredAIOptions<T>,
+): Promise<T> {
+  const maxSchemaRetries = 3;
+  let lastError: AIToolValidationError | undefined;
+
+  for (let attempt = 0; attempt < maxSchemaRetries; attempt++) {
+    try {
+      const model = getModel(provider, options.model);
+
+      const enhancedPrompt = attempt > 0 && lastError ? getEnhancedPrompt(prompt, lastError) : prompt;
+
+      const result = streamObject({
+        model,
+        prompt: enhancedPrompt,
+        schema: options.schema,
+        schemaName: options.schemaName,
+        schemaDescription: options.schemaDescription,
+        temperature: options.temperature,
+        maxTokens: options.maxTokens,
+      });
+
+      // Stream partial objects if callback provided
+      if (options.onPartialObject) {
+        void (async () => {
+          try {
+            for await (const partialObject of result.partialObjectStream) {
+              options.onPartialObject?.(partialObject);
+            }
+          } catch (streamError) {
+            console.error('Error in partial object stream:', streamError);
+          }
+        })();
+      }
+
+      // Return the final complete object
+      return await result.object;
+    } catch (error: unknown) {
+      lastError =
+        error instanceof Error
+          ? (error as AIToolValidationError)
+          : (new Error('An unknown error occurred') as AIToolValidationError);
+
+      const { shouldRetry, enhancedError } = handleFailedRequest(lastError, maxSchemaRetries, attempt);
+      lastError = enhancedError;
+
+      if (!shouldRetry) {
+        throw lastError;
+      }
+    }
+  }
+
+  // If we exhausted all retries, throw the last error
+  throw lastError;
+}
+
+// Export zod for convenience
+export { z };
+
+// Export MCP server singleton functions
+export {
+  server as mcpServer,
+  registerTool,
+  registerTools,
+  startServer,
+  isServerStarted,
+  getRegisteredTools,
+  getRegisteredToolsForAI,
+  getToolHandler,
+  getSchemaByName,
+  executeRegisteredTool,
+  type ToolHandler,
+  type RegisteredTool,
+} from './mcp-server';

package/src/mcp-server.ts

@@ -0,0 +1,203 @@
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
+import { z } from 'zod';
+
+interface ToolResult {
+  content: Array<{
+    type: 'text';
+    text: string;
+  }>;
+  isError?: boolean;
+  [key: string]: unknown;
+}
+
+interface ToolDescription {
+  title: string;
+  description: string;
+  inputSchema: Record<string, z.ZodSchema>;
+  schema?: z.ZodSchema;
+  schemaName?: string;
+  schemaDescription?: string;
+}
+
+type ToolHandler<T extends Record<string, unknown> = Record<string, unknown>> = (params: T) => Promise<ToolResult>;
+
+interface RegisteredTool {
+  name: string;
+  description: ToolDescription;
+  handler: ToolHandler<Record<string, unknown>>;
+  aiSdkTool: {
+    parameters: z.ZodSchema;
+    description: string;
+    execute?: (args: Record<string, unknown>) => Promise<string>;
+  };
+}
+
+const server = new McpServer({
+  name: 'frontend-implementation',
+  version: '0.1.0',
+});
+
+const transport = new StdioServerTransport();
+
+let isStarted = false;
+const toolRegistry = new Map<string, RegisteredTool>();
+
+async function cleanup() {
+  console.log('Cleaning up...');
+  await transport.close();
+  process.exit(0);
+}
+
+process.on('SIGTERM', () => {
+  void cleanup();
+});
+process.on('SIGINT', () => {
+  void cleanup();
+});
+
+export { server };
+
+function createMcpHandler<T extends Record<string, unknown>>(
+  handler: ToolHandler<T>,
+): (args: Record<string, unknown>, extra: unknown) => Promise<ToolResult> {
+  return (args: Record<string, unknown>, _extra: unknown) => {
+    return handler(args as T);
+  };
+}
+
+function createAiSdkTool(description: ToolDescription, handler: ToolHandler) {
+  const parameterSchema =
+    Object.keys(description.inputSchema).length > 0 ? z.object(description.inputSchema) : z.object({}).passthrough();
+
+  return {
+    parameters: parameterSchema,
+    description: description.description,
+    execute: async (args: Record<string, unknown>) => {
+      const result = await handler(args);
+      const textOutput = result.content[0]?.text || '';
+
+      // If a schema is provided, parse and validate the JSON output
+      if (description.schema) {
+        try {
+          const parsed = JSON.parse(textOutput) as unknown;
+          description.schema.parse(parsed);
+          return textOutput; // Return original text output for consistency
+        } catch (parseError) {
+          console.warn(`Tool failed to parse/validate JSON output:`, parseError);
+          return textOutput; // Fallback to raw text
+        }
+      }
+
+      return textOutput;
+    },
+    ...(description.schema && { schema: description.schema }),
+    ...(description.schemaName != null && { schemaName: description.schemaName }),
+    ...(description.schemaDescription != null && { schemaDescription: description.schemaDescription }),
+  };
+}
+
+export function registerTool<T extends Record<string, unknown> = Record<string, unknown>>(
+  name: string,
+  description: ToolDescription,
+  handler: ToolHandler<T>,
+) {
+  console.log(`🔧 Registering MCP tool: ${name}`);
+
+  if (isStarted) {
+    throw new Error('Cannot register tools after server has started');
+  }
+
+  const mcpHandler = createMcpHandler(handler);
+  const aiSdkTool = createAiSdkTool(description, handler as ToolHandler<Record<string, unknown>>);
+
+  const registeredTool: RegisteredTool = {
+    name,
+    description,
+    handler: handler as ToolHandler<Record<string, unknown>>,
+    aiSdkTool,
+  };
+
+  toolRegistry.set(name, registeredTool);
+  server.registerTool(name, description, mcpHandler);
+  console.log(`✅ Tool ${name} registered successfully. Total tools: ${toolRegistry.size}`);
+}
+
+export function registerTools<T extends Record<string, unknown> = Record<string, unknown>>(
+  tools: Array<{
+    name: string;
+    description: ToolDescription;
+    handler: ToolHandler<T>;
+  }>,
+) {
+  if (isStarted) {
+    throw new Error('Cannot register tools after server has started');
+  }
+  tools.forEach((tool) => {
+    registerTool(tool.name, tool.description, tool.handler);
+  });
+}
+
+export async function startServer() {
+  if (isStarted) return;
+
+  await server.connect(transport);
+  isStarted = true;
+}
+
+export function isServerStarted() {
+  return isStarted;
+}
+
+export function getRegisteredTools(): RegisteredTool[] {
+  return Array.from(toolRegistry.values());
+}
+
+export function getRegisteredToolsForAI(): Record<
+  string,
+  {
+    parameters: z.ZodSchema;
+    description: string;
+    execute?: (args: Record<string, unknown>) => Promise<string>;
+  }
+> {
+  console.log(`📋 Getting registered tools for AI. Registry size: ${toolRegistry.size}`);
+  const tools: Record<
+    string,
+    {
+      parameters: z.ZodSchema;
+      description: string;
+      execute?: (args: Record<string, unknown>) => Promise<string>;
+    }
+  > = {};
+  for (const tool of toolRegistry.values()) {
+    tools[tool.name] = tool.aiSdkTool;
+    console.log(`  - Tool: ${tool.name}`);
+  }
+  console.log(`📊 Returning ${Object.keys(tools).length} tools for AI`);
+  return tools;
+}
+
+export function getToolHandler(name: string): ToolHandler | undefined {
+  return toolRegistry.get(name)?.handler;
+}
+
+export async function executeRegisteredTool(name: string, params: Record<string, unknown>): Promise<ToolResult> {
+  const tool = toolRegistry.get(name);
+  if (!tool) {
+    throw new Error(`Tool '${name}' not found`);
+  }
+  return await tool.handler(params);
+}
+
+export function getSchemaByName(schemaName: string): z.ZodSchema | undefined {
+  console.log();
+  for (const tool of toolRegistry.values()) {
+    if (tool.description?.schemaName === schemaName) {
+      return tool.description.schema;
+    }
+  }
+  return undefined;
+}
+
+export type { ToolResult, ToolDescription, ToolHandler, RegisteredTool };

package/tsconfig.json

@@ -0,0 +1,10 @@
+{
+  "extends": "../../tsconfig.base.json",
+  "compilerOptions": {
+    "composite": true,
+    "outDir": "./dist",
+    "rootDir": "./src"
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "dist"]
+}