File size: 4,201 Bytes
7a1ad33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import type { Part } from '@google/genai';
import {
  estimateTokenCountSync,
  ASCII_TOKENS_PER_CHAR,
  NON_ASCII_TOKENS_PER_CHAR,
} from '../utils/tokenCalculation.js';

export const MIN_TARGET_TOKENS = 10;
export const MIN_CHARS_FOR_TRUNCATION = 100;
export const TEXT_TRUNCATION_PREFIX =
  '[Message Normalized: Exceeded size limit]';
export const TOOL_TRUNCATION_PREFIX =
  '[Message Normalized: Tool output exceeded size limit]';

/**
 * Estimates the character limit for a target token count, accounting for ASCII vs Non-ASCII.
 * Uses a weighted average based on the provided text to decide how many characters
 * fit into the target token budget.
 */
export function estimateCharsFromTokens(
  text: string,
  targetTokens: number,
): number {
  if (text.length === 0) return 0;

  // Count ASCII vs Non-ASCII in a sample of the text.
  let asciiCount = 0;
  const sampleLen = Math.min(text.length, 1000);
  for (let i = 0; i < sampleLen; i++) {
    if (text.charCodeAt(i) <= 127) {
      asciiCount++;
    }
  }

  const asciiRatio = asciiCount / sampleLen;
  // Weighted tokens per character:
  const avgTokensPerChar =
    asciiRatio * ASCII_TOKENS_PER_CHAR +
    (1 - asciiRatio) * NON_ASCII_TOKENS_PER_CHAR;

  // Characters = Tokens / (Tokens per Character)
  return Math.floor(targetTokens / avgTokensPerChar);
}

/**
 * Truncates a string to a target length, keeping a proportional amount of the head and tail,
 * and prepending a prefix.
 */
export function truncateProportionally(
  str: string,
  targetChars: number,
  prefix: string,
  headRatio: number = 0.2,
): string {
  if (str.length <= targetChars) return str;

  const ellipsis = '\n...\n';
  const overhead = prefix.length + ellipsis.length + 1; // +1 for the newline after prefix
  const availableChars = Math.max(0, targetChars - overhead);

  if (availableChars <= 0) {
    return prefix; // Safe fallback if target is extremely small
  }

  const headChars = Math.floor(availableChars * headRatio);
  const tailChars = availableChars - headChars;

  return `${prefix}\n${str.substring(0, headChars)}${ellipsis}${str.substring(str.length - tailChars)}`;
}

/**
 * Safely normalizes a function response by truncating large string values
 * within the response object while maintaining its JSON structure.
 */
export function normalizeFunctionResponse(
  part: Part,
  ratio: number,
  headRatio: number = 0.2,
  savedPath?: string,
  intentSummary?: string,
): Part {
  const fr = part.functionResponse;
  if (!fr || !fr.response) return part;

  const responseObj = fr.response;
  if (typeof responseObj !== 'object' || responseObj === null) return part;

  let hasChanges = false;
  const newResponse: Record<string, unknown> = {};

  // For function responses, we truncate individual string values that are large.
  // This preserves the schema keys (stdout, stderr, etc).
  for (const [key, value] of Object.entries(responseObj)) {
    if (typeof value === 'string' && value.length > MIN_CHARS_FOR_TRUNCATION) {
      const valueTokens = estimateTokenCountSync([{ text: value }]);
      const targetValueTokens = Math.max(
        MIN_TARGET_TOKENS,
        Math.floor(valueTokens * ratio),
      );
      const targetChars = estimateCharsFromTokens(value, targetValueTokens);

      if (value.length > targetChars) {
        let truncated = truncateProportionally(
          value,
          targetChars,
          TOOL_TRUNCATION_PREFIX,
          headRatio,
        );
        if (savedPath) {
          truncated += `\n\nFull output saved to: ${savedPath}`;
        }
        if (intentSummary) {
          truncated += intentSummary;
        }
        newResponse[key] = truncated;
        hasChanges = true;
      } else {
        newResponse[key] = value;
      }
    } else {
      newResponse[key] = value;
    }
  }

  if (!hasChanges) return part;

  return {
    functionResponse: {
      // This spread should be safe as we mostly care about the function
      // response properties.
      // eslint-disable-next-line @typescript-eslint/no-misused-spread
      ...fr,
      response: newResponse,
    },
  };
}