blob: 57e0980cfda8c61343f8f30f00db1fe734859bba (
plain)
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
|
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { createHash } from 'crypto';
import { GeminiEventType, ServerGeminiStreamEvent } from '../core/turn.js';
const TOOL_CALL_LOOP_THRESHOLD = 5;
const CONTENT_LOOP_THRESHOLD = 10;
const SENTENCE_ENDING_PUNCTUATION_REGEX = /[.!?]+(?=\s|$)/;
/**
* Service for detecting and preventing infinite loops in AI responses.
* Monitors tool call repetitions and content sentence repetitions.
*/
export class LoopDetectionService {
// Tool call tracking
private lastToolCallKey: string | null = null;
private toolCallRepetitionCount: number = 0;
// Content streaming tracking
private lastRepeatedSentence: string = '';
private sentenceRepetitionCount: number = 0;
private partialContent: string = '';
private getToolCallKey(toolCall: { name: string; args: object }): string {
const argsString = JSON.stringify(toolCall.args);
const keyString = `${toolCall.name}:${argsString}`;
return createHash('sha256').update(keyString).digest('hex');
}
/**
* Processes a stream event and checks for loop conditions.
* @param event - The stream event to process
* @returns true if a loop is detected, false otherwise
*/
addAndCheck(event: ServerGeminiStreamEvent): boolean {
switch (event.type) {
case GeminiEventType.ToolCallRequest:
// content chanting only happens in one single stream, reset if there
// is a tool call in between
this.resetSentenceCount();
return this.checkToolCallLoop(event.value);
case GeminiEventType.Content:
return this.checkContentLoop(event.value);
default:
this.reset();
return false;
}
}
private checkToolCallLoop(toolCall: { name: string; args: object }): boolean {
const key = this.getToolCallKey(toolCall);
if (this.lastToolCallKey === key) {
this.toolCallRepetitionCount++;
} else {
this.lastToolCallKey = key;
this.toolCallRepetitionCount = 1;
}
return this.toolCallRepetitionCount >= TOOL_CALL_LOOP_THRESHOLD;
}
private checkContentLoop(content: string): boolean {
this.partialContent += content;
if (!SENTENCE_ENDING_PUNCTUATION_REGEX.test(this.partialContent)) {
return false;
}
const completeSentences =
this.partialContent.match(/[^.!?]+[.!?]+(?=\s|$)/g) || [];
if (completeSentences.length === 0) {
return false;
}
const lastSentence = completeSentences[completeSentences.length - 1];
const lastCompleteIndex = this.partialContent.lastIndexOf(lastSentence);
const endOfLastSentence = lastCompleteIndex + lastSentence.length;
this.partialContent = this.partialContent.slice(endOfLastSentence);
for (const sentence of completeSentences) {
const trimmedSentence = sentence.trim();
if (trimmedSentence === '') {
continue;
}
if (this.lastRepeatedSentence === trimmedSentence) {
this.sentenceRepetitionCount++;
} else {
this.lastRepeatedSentence = trimmedSentence;
this.sentenceRepetitionCount = 1;
}
if (this.sentenceRepetitionCount >= CONTENT_LOOP_THRESHOLD) {
return true;
}
}
return false;
}
/**
* Resets all loop detection state.
*/
reset(): void {
this.resetToolCallCount();
this.resetSentenceCount();
}
private resetToolCallCount(): void {
this.lastToolCallKey = null;
this.toolCallRepetitionCount = 0;
}
private resetSentenceCount(): void {
this.lastRepeatedSentence = '';
this.sentenceRepetitionCount = 0;
this.partialContent = '';
}
}
|
