wave-agent-sdk 0.0.6 → 0.0.8

package/src/services/aiService.ts

@@ -2,14 +2,30 @@ import OpenAI from "openai";
 import { ChatCompletionMessageToolCall } from "openai/resources";
 import {
   ChatCompletionCreateParamsNonStreaming,
+  ChatCompletionCreateParamsStreaming,
   ChatCompletionMessageParam,
   ChatCompletionFunctionTool,
+  ChatCompletionChunk,
 } from "openai/resources.js";
+import { logger } from "../utils/globalLogger.js";
 import type { GatewayConfig, ModelConfig } from "../types/index.js";
+import {
+  transformMessagesForClaudeCache,
+  addCacheControlToLastTool,
+  isClaudeModel,
+  extendUsageWithCacheMetrics,
+  type ClaudeUsage,
+} from "../utils/cacheControlUtils.js";
+
 import * as os from "os";
 import * as fs from "fs";
 import * as path from "path";
 
+/**
+ * Use parametersChunk as compact param for better performance
+ * Instead of parsing JSON, we use the raw chunk for efficient streaming
+ */
+
 /**
  * Check if a directory is a git repository
  * @param dirPath Directory path to check
@@ -79,16 +95,31 @@ export interface CallAgentOptions {
   tools?: ChatCompletionFunctionTool[]; // Tool configuration
   model?: string; // Custom model
   systemPrompt?: string; // Custom system prompt
+
+  // NEW: Streaming callbacks
+  onContentUpdate?: (content: string) => void;
+  onToolUpdate?: (toolCall: {
+    id: string;
+    name: string;
+    parameters: string;
+    parametersChunk?: string;
+    stage?: "start" | "streaming" | "running" | "end";
+  }) => void;
 }
 
 export interface CallAgentResult {
   content?: string;
   tool_calls?: ChatCompletionMessageToolCall[];
-  usage?: {
-    prompt_tokens: number;
-    completion_tokens: number;
-    total_tokens: number;
-  };
+  usage?: ClaudeUsage;
+  finish_reason?:
+    | "stop"
+    | "length"
+    | "tool_calls"
+    | "content_filter"
+    | "function_call"
+    | null;
+  response_headers?: Record<string, string>;
+  metadata?: Record<string, unknown>;
 }
 
 export async function callAgent(
@@ -104,6 +135,8 @@ export async function callAgent(
     tools,
     model,
     systemPrompt,
+    onContentUpdate,
+    onToolUpdate,
   } = options;
 
   try {
@@ -150,65 +183,371 @@ Today's date: ${new Date().toISOString().split("T")[0]}
     };
 
     // ChatCompletionMessageParam[] is already in OpenAI format, add system prompt to the beginning
-    const openaiMessages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] =
-      [systemMessage, ...messages];
+    let openaiMessages: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = [
+      systemMessage,
+      ...messages,
+    ];
+
+    // Apply cache control for Claude models
+    const currentModel = model || modelConfig.agentModel;
+
+    if (isClaudeModel(currentModel)) {
+      openaiMessages = transformMessagesForClaudeCache(
+        openaiMessages,
+        currentModel,
+      );
+    }
 
     // Get model configuration - use injected modelConfig with optional override
     const openaiModelConfig = getModelConfig(model || modelConfig.agentModel, {
       temperature: 0,
-      max_completion_tokens: 32768,
     });
 
+    // Determine if streaming is needed
+    const isStreaming = !!(onContentUpdate || onToolUpdate);
+
     // Prepare API call parameters
-    const createParams: ChatCompletionCreateParamsNonStreaming = {
+    const createParams = {
       ...openaiModelConfig,
       messages: openaiMessages,
-    };
+      stream: isStreaming,
+    } as
+      | ChatCompletionCreateParamsNonStreaming
+      | ChatCompletionCreateParamsStreaming;
 
     // Only add tools if they exist
     if (tools && tools.length > 0) {
-      createParams.tools = tools;
+      // Apply cache control to tools for Claude models
+      if (isClaudeModel(currentModel)) {
+        createParams.tools = addCacheControlToLastTool(tools);
+      } else {
+        createParams.tools = tools;
+      }
     }
 
-    // Call OpenAI API (non-streaming)
-    const response = await openai.chat.completions.create(createParams, {
-      signal: abortSignal,
-    });
+    if (isStreaming) {
+      // Handle streaming response
+      const { data: stream, response } = await openai.chat.completions
+        .create(createParams as ChatCompletionCreateParamsStreaming, {
+          signal: abortSignal,
+        })
+        .withResponse();
+
+      // Extract response headers
+      const responseHeaders: Record<string, string> = {};
+      response.headers.forEach((value, key) => {
+        responseHeaders[key] = value;
+      });
+
+      return await processStreamingResponse(
+        stream,
+        onContentUpdate,
+        onToolUpdate,
+        abortSignal,
+        responseHeaders,
+        currentModel,
+      );
+    } else {
+      // Handle non-streaming response
+      const { data: response, response: rawResponse } =
+        await openai.chat.completions
+          .create(createParams as ChatCompletionCreateParamsNonStreaming, {
+            signal: abortSignal,
+          })
+          .withResponse();
+
+      // Extract response headers
+      const responseHeaders: Record<string, string> = {};
+      rawResponse.headers.forEach((value, key) => {
+        responseHeaders[key] = value;
+      });
+
+      const finalMessage = response.choices[0]?.message;
+      const finishReason = response.choices[0]?.finish_reason || null;
+
+      let totalUsage = response.usage
+        ? {
+            prompt_tokens: response.usage.prompt_tokens,
+            completion_tokens: response.usage.completion_tokens,
+            total_tokens: response.usage.total_tokens,
+          }
+        : undefined;
+
+      // Extend usage with cache metrics for Claude models
+      if (totalUsage && isClaudeModel(currentModel) && response.usage) {
+        totalUsage = extendUsageWithCacheMetrics(
+          totalUsage,
+          response.usage as Partial<ClaudeUsage>,
+        );
+      }
 
-    const finalMessage = response.choices[0]?.message;
-    const totalUsage = response.usage
-      ? {
-          prompt_tokens: response.usage.prompt_tokens,
-          completion_tokens: response.usage.completion_tokens,
-          total_tokens: response.usage.total_tokens,
+      const result: CallAgentResult = {};
+
+      if (finalMessage) {
+        const {
+          content: finalContent,
+          tool_calls: finalToolCalls,
+          ...otherFields
+        } = finalMessage;
+
+        if (typeof finalContent === "string" && finalContent.length > 0) {
+          result.content = finalContent;
         }
-      : undefined;
 
-    const result: CallAgentResult = {};
+        if (Array.isArray(finalToolCalls) && finalToolCalls.length > 0) {
+          result.tool_calls = finalToolCalls;
+        }
 
-    // Return content
-    if (finalMessage?.content) {
-      result.content = finalMessage.content;
-    }
+        if (Object.keys(otherFields).length > 0) {
+          const metadata: Record<string, unknown> = {};
+          for (const [key, value] of Object.entries(otherFields)) {
+            if (value !== undefined) {
+              metadata[key] = value;
+            }
+          }
+          if (Object.keys(metadata).length > 0) {
+            result.metadata = metadata;
+          }
+        }
+      }
 
-    // Return tool call
-    if (finalMessage?.tool_calls && finalMessage.tool_calls.length > 0) {
-      result.tool_calls = finalMessage.tool_calls;
-    }
+      if (totalUsage) {
+        result.usage = totalUsage;
+      }
 
-    // Return token usage information
-    if (totalUsage) {
-      result.usage = totalUsage;
-    }
+      if (finishReason) {
+        result.finish_reason = finishReason;
+      }
+
+      if (Object.keys(responseHeaders).length > 0) {
+        result.response_headers = responseHeaders;
+      }
 
-    return result;
+      return result;
+    }
   } catch (error) {
     if ((error as Error).name === "AbortError") {
       throw new Error("Request was aborted");
     }
-    // // logger.error("Failed to call OpenAI:", error);
+    logger.error("Failed to call OpenAI:", error);
+    throw error;
+  }
+}
+
+/**
+ * Process streaming response from OpenAI API
+ * @param stream Async iterator of chat completion chunks
+ * @param onContentUpdate Callback for content updates
+ * @param onToolUpdate Callback for tool updates
+ * @param abortSignal Optional abort signal
+ * @param responseHeaders Response headers from the initial request
+ * @param modelName Model name for cache control processing
+ * @returns Final result with accumulated content and tool calls
+ */
+async function processStreamingResponse(
+  stream: AsyncIterable<ChatCompletionChunk>,
+  onContentUpdate?: (content: string) => void,
+  onToolUpdate?: (toolCall: {
+    id: string;
+    name: string;
+    parameters: string;
+    parametersChunk?: string;
+    stage?: "start" | "streaming" | "running" | "end";
+  }) => void,
+  abortSignal?: AbortSignal,
+  responseHeaders?: Record<string, string>,
+  modelName?: string,
+): Promise<CallAgentResult> {
+  let accumulatedContent = "";
+  const toolCalls: {
+    id: string;
+    type: "function";
+    function: {
+      name: string;
+      arguments: string;
+    };
+  }[] = [];
+  const additionalDeltaFields: Record<string, unknown> = {};
+  let usage: CallAgentResult["usage"] = undefined;
+  let finishReason: CallAgentResult["finish_reason"] = null;
+
+  try {
+    for await (const chunk of stream) {
+      // Check for abort signal
+      if (abortSignal?.aborted) {
+        throw new Error("Request was aborted");
+      }
+
+      // Check for usage information in any chunk
+      if (chunk.usage) {
+        let chunkUsage = {
+          prompt_tokens: chunk.usage.prompt_tokens,
+          completion_tokens: chunk.usage.completion_tokens,
+          total_tokens: chunk.usage.total_tokens,
+        };
+
+        // Extend usage with cache metrics for Claude models
+        if (modelName && isClaudeModel(modelName)) {
+          chunkUsage = extendUsageWithCacheMetrics(
+            chunkUsage,
+            chunk.usage as Partial<ClaudeUsage>,
+          );
+        }
+
+        usage = chunkUsage;
+      }
+
+      // Check for finish_reason in the choice
+      const choice = chunk.choices?.[0];
+      if (choice?.finish_reason) {
+        finishReason = choice.finish_reason;
+      }
+
+      const delta = choice?.delta;
+      if (!delta) {
+        continue;
+      }
+
+      const {
+        content,
+        tool_calls: toolCallUpdates,
+        ...deltaMetadata
+      } = delta as unknown as {
+        content?: string;
+        tool_calls?: ChatCompletionChunk.Choice.Delta.ToolCall[];
+        [key: string]: unknown;
+      };
+
+      if (Object.keys(deltaMetadata).length > 0) {
+        Object.assign(additionalDeltaFields, deltaMetadata);
+      }
+
+      if (typeof content === "string" && content.length > 0) {
+        // Note: OpenAI API already handles UTF-8 character boundaries correctly in streaming,
+        // ensuring that delta.content always contains complete UTF-8 strings
+        accumulatedContent += content;
+        if (onContentUpdate) {
+          onContentUpdate(accumulatedContent);
+        }
+      }
+
+      if (Array.isArray(toolCallUpdates)) {
+        for (const rawToolCall of toolCallUpdates) {
+          const toolCallDelta =
+            rawToolCall as ChatCompletionChunk.Choice.Delta.ToolCall;
+
+          if (!toolCallDelta.function) {
+            continue;
+          }
+
+          const functionDelta = toolCallDelta.function;
+
+          let existingCall;
+          let isNew = false;
+
+          if (toolCallDelta.id) {
+            existingCall = toolCalls.find((t) => t.id === toolCallDelta.id);
+            if (!existingCall) {
+              existingCall = {
+                id: toolCallDelta.id,
+                type: "function" as const,
+                function: {
+                  name: functionDelta.name || "",
+                  arguments: "",
+                },
+              };
+              toolCalls.push(existingCall);
+              isNew = true;
+            }
+          } else {
+            existingCall = toolCalls[toolCalls.length - 1];
+          }
+
+          if (!existingCall) {
+            continue;
+          }
+
+          if (functionDelta.name) {
+            existingCall.function.name = functionDelta.name;
+          }
+
+          // Emit start stage when a new tool call is created and we have the tool name
+          if (onToolUpdate && isNew && existingCall.function.name) {
+            onToolUpdate({
+              id: existingCall.id,
+              name: existingCall.function.name,
+              parameters: "", // Empty parameters for start stage
+              parametersChunk: "", // Empty chunk for start stage
+              stage: "start", // New tool call triggers start stage
+            });
+            isNew = false; // Prevent duplicate start emissions
+          }
+
+          if (functionDelta.arguments) {
+            existingCall.function.arguments += functionDelta.arguments;
+          }
+
+          // Emit streaming updates for all chunks with actual content (including first chunk)
+          if (
+            onToolUpdate &&
+            existingCall.function.name &&
+            functionDelta.arguments &&
+            functionDelta.arguments.length > 0 // Only emit streaming for chunks with actual content
+          ) {
+            onToolUpdate({
+              id: existingCall.id,
+              name: existingCall.function.name,
+              parameters: existingCall.function.arguments,
+              parametersChunk: functionDelta.arguments,
+              stage: "streaming",
+            });
+          }
+        }
+      }
+    }
+  } catch (error) {
+    if ((error as Error).message === "Request was aborted") {
+      throw error;
+    }
     throw error;
   }
+
+  // Prepare final result
+  const result: CallAgentResult = {};
+
+  if (accumulatedContent) {
+    result.content = accumulatedContent;
+  }
+
+  if (toolCalls.length > 0) {
+    result.tool_calls = toolCalls;
+  }
+
+  if (usage) {
+    result.usage = usage;
+  }
+
+  if (finishReason) {
+    result.finish_reason = finishReason;
+  }
+
+  if (responseHeaders && Object.keys(responseHeaders).length > 0) {
+    result.response_headers = responseHeaders;
+  }
+
+  if (Object.keys(additionalDeltaFields).length > 0) {
+    result.metadata = {};
+    for (const [key, value] of Object.entries(additionalDeltaFields)) {
+      if (value !== undefined) {
+        result.metadata[key] = value;
+      }
+    }
+    if (Object.keys(result.metadata).length === 0) {
+      delete result.metadata;
+    }
+  }
+
+  return result;
 }
 
 export interface CompressMessagesOptions {
@@ -244,7 +583,7 @@ export async function compressMessages(
   // Get model configuration - use injected fast model
   const openaiModelConfig = getModelConfig(modelConfig.fastModel, {
     temperature: 0.1,
-    max_tokens: 1500,
+    max_tokens: 2048,
   });
 
   try {
@@ -254,54 +593,106 @@ export async function compressMessages(
         messages: [
           {
             role: "system",
-            content: `You are an expert conversation history compression specialist. Your task is to create comprehensive yet concise summaries that preserve critical development context.
-
-## Primary Request and Intent
-Compress conversation history while maintaining all essential technical and procedural information.
-
-## Key Technical Concepts
-- Code modifications and file operations
-- Tool executions and their results  
-- Error handling and debugging processes
-- User requirements and assistant solutions
-- Technical discussions and decisions
-
-## Compression Strategy
-1. **Preserve Critical Information**:
-   - All file paths, function names, and code examples
-   - Tool execution results and outcomes
-   - Error messages and resolution steps
-   - User requirements and implementation approaches
-   - Technical decisions and their reasoning
-
-2. **Structure Organization**:
-   - Group related actions and discussions
-   - Maintain chronological flow for complex operations
-   - Separate different technical topics clearly
-   
-3. **Context Preservation**:
-   - Keep enough detail for future reference
-   - Maintain relationships between requests and solutions
-   - Preserve debugging context and error resolution paths
-
-## Output Requirements:
-- Use third-person narrative format
-- Target 300-800 words (scale based on complexity)
-- Maintain the original conversation language
-- Structure with clear sections for multi-topic conversations
-- Focus on actionable information and outcomes
-
-## Format Template:
-For technical conversations, structure as:
-- **User Requests**: Key requirements and goals
-- **Technical Implementation**: Code changes, file operations, tool usage
-- **Problem Resolution**: Errors encountered and solutions applied
-- **Outcomes**: Final results and current state`,
+            content: `Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
+This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
+
+Before providing your final summary, wrap your analysis in <analysis> tags to organize your thoughts and ensure you've covered all necessary points. In your analysis process:
+
+1. Chronologically analyze each message and section of the conversation. For each section thoroughly identify:
+   - The user's explicit requests and intents
+   - Your approach to addressing the user's requests
+   - Key decisions, technical concepts and code patterns
+   - Specific details like:
+     - file names
+     - full code snippets
+     - function signatures
+     - file edits
+  - Errors that you ran into and how you fixed them
+  - Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
+2. Double-check for technical accuracy and completeness, addressing each required element thoroughly.
+
+Your summary should include the following sections:
+
+1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
+2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
+3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
+4. Errors and fixes: List all errors that you ran into, and how you fixed them. Pay special attention to specific user feedback that you received, especially if the user told you to do something differently.
+5. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
+6. All user messages: List ALL user messages that are not tool results. These are critical for understanding the users' feedback and changing intent.
+6. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
+7. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
+8. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's most recent explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests or really old requests that were already completed without confirming with the user first.
+                       If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
+
+Here's an example of how your output should be structured:
+
+<example>
+<analysis>
+[Your thought process, ensuring all points are covered thoroughly and accurately]
+</analysis>
+
+<summary>
+1. Primary Request and Intent:
+   [Detailed description]
+
+2. Key Technical Concepts:
+   - [Concept 1]
+   - [Concept 2]
+   - [...]
+
+3. Files and Code Sections:
+   - [File Name 1]
+      - [Summary of why this file is important]
+      - [Summary of the changes made to this file, if any]
+      - [Important Code Snippet]
+   - [File Name 2]
+      - [Important Code Snippet]
+   - [...]
+
+4. Errors and fixes:
+    - [Detailed description of error 1]:
+      - [How you fixed the error]
+      - [User feedback on the error if any]
+    - [...]
+
+5. Problem Solving:
+   [Description of solved problems and ongoing troubleshooting]
+
+6. All user messages: 
+    - [Detailed non tool use user message]
+    - [...]
+
+7. Pending Tasks:
+   - [Task 1]
+   - [Task 2]
+   - [...]
+
+8. Current Work:
+   [Precise description of current work]
+
+9. Optional Next Step:
+   [Optional Next step to take]
+
+</summary>
+</example>
+
+Please provide your summary based on the conversation so far, following this structure and ensuring precision and thoroughness in your response. 
+
+There may be additional summarization instructions provided in the included context. If so, remember to follow these instructions when creating the above summary. Examples of instructions include:
+<example>
+## Compact Instructions
+When summarizing the conversation focus on typescript code changes and also remember the mistakes you made and how you fixed them.
+</example>
+
+<example>
+# Summary instructions
+When you are using compact - please focus on test output and code changes. Include file reads verbatim.
+</example>`,
           },
           ...messages,
           {
             role: "user",
-            content: `Please compress this conversation following the structured approach. Focus on preserving all technical details, file operations, and problem-solving context while creating a concise summary.`,
+            content: `Please create a detailed summary of the conversation so far.`,
           },
         ],
       },
@@ -329,7 +720,7 @@ For technical conversations, structure as:
     if ((error as Error).name === "AbortError") {
       throw new Error("Compression request was aborted");
     }
-    // // logger.error("Failed to compress messages:", error);
+    logger.error("Failed to compress messages:", error);
     return {
       content: "Failed to compress conversation history",
       usage: undefined,