1
1
package ee .carlrobert .codegpt .completions ;
2
2
3
3
import com .intellij .openapi .diagnostic .Logger ;
4
- import ee .carlrobert .codegpt .completions .you .YouUserManager ;
5
4
import ee .carlrobert .codegpt .conversations .Conversation ;
6
5
import ee .carlrobert .codegpt .conversations .message .Message ;
7
- import ee .carlrobert .codegpt .settings .service .ServiceType ;
8
- import ee .carlrobert .codegpt .settings .state .AzureSettingsState ;
9
- import ee .carlrobert .codegpt .settings .state .OpenAISettingsState ;
10
6
import ee .carlrobert .codegpt .settings .state .SettingsState ;
11
7
import ee .carlrobert .codegpt .telemetry .TelemetryAction ;
12
8
import ee .carlrobert .llm .client .openai .completion .ErrorDetails ;
13
9
import ee .carlrobert .llm .client .you .completion .YouCompletionEventListener ;
14
10
import ee .carlrobert .llm .client .you .completion .YouSerpResult ;
15
11
import ee .carlrobert .llm .completion .CompletionEventListener ;
16
12
import java .util .List ;
17
- import java .util .function .BiConsumer ;
18
- import java .util .function .Consumer ;
19
13
import javax .swing .SwingWorker ;
20
14
import okhttp3 .sse .EventSource ;
21
15
import org .jetbrains .annotations .NotNull ;
22
- import org .jetbrains .annotations .Nullable ;
23
16
24
17
public class CompletionRequestHandler {
25
18
26
19
private static final Logger LOG = Logger .getInstance (CompletionRequestHandler .class );
27
20
28
21
private final StringBuilder messageBuilder = new StringBuilder ();
22
+ private final boolean useContextualSearch ;
23
+ private final ToolWindowCompletionEventListener toolWindowCompletionEventListener ;
29
24
private SwingWorker <Void , String > swingWorker ;
30
25
private EventSource eventSource ;
31
- private @ Nullable Consumer <String > messageListener ;
32
- private @ Nullable BiConsumer <ErrorDetails , Throwable > errorListener ;
33
- private @ Nullable Consumer <String > completedListener ;
34
- private @ Nullable Consumer <List <YouSerpResult >> serpResultsListener ;
35
- private @ Nullable Runnable tokensExceededListener ;
36
- private boolean useContextualSearch ;
37
-
38
- public CompletionRequestHandler withContextualSearch (boolean useContextualSearch ) {
39
- this .useContextualSearch = useContextualSearch ;
40
- return this ;
41
- }
42
-
43
- public void addMessageListener (Consumer <String > messageListener ) {
44
- this .messageListener = messageListener ;
45
- }
46
-
47
- public void addErrorListener (BiConsumer <ErrorDetails , Throwable > errorListener ) {
48
- this .errorListener = errorListener ;
49
- }
50
-
51
- public void addRequestCompletedListener (Consumer <String > completedListener ) {
52
- this .completedListener = completedListener ;
53
- }
54
-
55
- public void addTokensExceededListener (Runnable tokensExceededListener ) {
56
- this .tokensExceededListener = tokensExceededListener ;
57
- }
58
26
59
- public void addSerpResultsListener (Consumer <List <YouSerpResult >> serpResultsListener ) {
60
- this .serpResultsListener = serpResultsListener ;
27
+ public CompletionRequestHandler (
28
+ boolean useContextualSearch ,
29
+ ToolWindowCompletionEventListener toolWindowCompletionEventListener ) {
30
+ this .useContextualSearch = useContextualSearch ;
31
+ this .toolWindowCompletionEventListener = toolWindowCompletionEventListener ;
61
32
}
62
33
63
34
public void call (Conversation conversation , Message message , boolean isRetry ) {
@@ -75,60 +46,13 @@ public void cancel() {
75
46
private EventSource startCall (
76
47
@ NotNull Conversation conversation ,
77
48
@ NotNull Message message ,
78
- boolean isRetry ,
49
+ boolean retry ,
79
50
CompletionEventListener eventListener ) {
80
- var settings = SettingsState .getInstance ();
81
- var requestProvider = new CompletionRequestProvider (conversation );
82
-
83
51
try {
84
- if (settings .getSelectedService () == ServiceType .LLAMA_CPP ) {
85
- return CompletionClientProvider .getLlamaClient ()
86
- .getChatCompletion (requestProvider .buildLlamaCompletionRequest (message ), eventListener );
87
- }
88
-
89
- if (settings .getSelectedService () == ServiceType .YOU ) {
90
- var sessionId = "" ;
91
- var accessToken = "" ;
92
- var youUserManager = YouUserManager .getInstance ();
93
- if (youUserManager .isAuthenticated ()) {
94
- var authenticationResponse =
95
- youUserManager .getAuthenticationResponse ().getData ();
96
- sessionId = authenticationResponse .getSession ().getSessionId ();
97
- accessToken = authenticationResponse .getSessionJwt ();
98
- }
99
- var request = requestProvider .buildYouCompletionRequest (message );
100
- LOG .info ("Initiating completion request using model: " +
101
- (request .isUseGPT4Model () ? "GPT-4" : "YouBot" ));
102
-
103
- return CompletionClientProvider .getYouClient (sessionId , accessToken )
104
- .getChatCompletion (request , eventListener );
105
- }
106
-
107
- if (settings .getSelectedService () == ServiceType .AZURE ) {
108
- var azureSettings = AzureSettingsState .getInstance ();
109
- return CompletionClientProvider .getAzureClient ().getChatCompletion (
110
- requestProvider .buildOpenAIChatCompletionRequest (
111
- azureSettings .getModel (),
112
- message ,
113
- isRetry ,
114
- useContextualSearch ,
115
- azureSettings .isUsingCustomPath () ? azureSettings .getPath () : null ),
116
- eventListener );
117
- }
118
-
119
- var openAISettings = OpenAISettingsState .getInstance ();
120
- return CompletionClientProvider .getOpenAIClient ().getChatCompletion (
121
- requestProvider .buildOpenAIChatCompletionRequest (
122
- openAISettings .getModel (),
123
- message ,
124
- isRetry ,
125
- useContextualSearch ,
126
- openAISettings .isUsingCustomPath () ? openAISettings .getPath () : null ),
127
- eventListener );
52
+ return CompletionRequestService .getInstance ()
53
+ .getChatCompletionAsync (conversation , message , retry , useContextualSearch , eventListener );
128
54
} catch (Throwable t ) {
129
- if (errorListener != null ) {
130
- errorListener .accept (new ErrorDetails ("Something went wrong" ), t );
131
- }
55
+ toolWindowCompletionEventListener .handleError (new ErrorDetails ("Something went wrong" ), t );
132
56
throw t ;
133
57
}
134
58
}
@@ -152,13 +76,9 @@ protected Void doInBackground() {
152
76
conversation ,
153
77
message ,
154
78
isRetry ,
155
- settings .getSelectedService () == ServiceType .YOU ?
156
- new YouRequestCompletionEventListener () :
157
- new BaseCompletionEventListener ());
79
+ new YouRequestCompletionEventListener ());
158
80
} catch (TotalUsageExceededException e ) {
159
- if (tokensExceededListener != null ) {
160
- tokensExceededListener .run ();
161
- }
81
+ toolWindowCompletionEventListener .handleTokensExceeded (conversation , message );
162
82
} finally {
163
83
sendInfo (settings );
164
84
}
@@ -169,13 +89,16 @@ protected void process(List<String> chunks) {
169
89
message .setResponse (messageBuilder .toString ());
170
90
for (String text : chunks ) {
171
91
messageBuilder .append (text );
172
- if (messageListener != null ) {
173
- messageListener .accept (text );
174
- }
92
+ toolWindowCompletionEventListener .handleMessage (text );
175
93
}
176
94
}
177
95
178
- class BaseCompletionEventListener implements CompletionEventListener {
96
+ class YouRequestCompletionEventListener implements YouCompletionEventListener {
97
+
98
+ @ Override
99
+ public void onSerpResults (List <YouSerpResult > results ) {
100
+ toolWindowCompletionEventListener .handleSerpResults (results , message );
101
+ }
179
102
180
103
@ Override
181
104
public void onMessage (String message ) {
@@ -184,34 +107,20 @@ public void onMessage(String message) {
184
107
185
108
@ Override
186
109
public void onComplete (StringBuilder messageBuilder ) {
187
- if (completedListener != null ) {
188
- completedListener .accept (messageBuilder .toString ());
189
- }
110
+ toolWindowCompletionEventListener .handleCompleted (messageBuilder .toString (), message ,
111
+ conversation , isRetry );
190
112
}
191
113
192
114
@ Override
193
115
public void onError (ErrorDetails error , Throwable ex ) {
194
116
try {
195
- if (errorListener != null ) {
196
- errorListener .accept (error , ex );
197
- }
117
+ toolWindowCompletionEventListener .handleError (error , ex );
198
118
} finally {
199
119
sendError (error , ex );
200
120
}
201
121
}
202
122
}
203
123
204
- class YouRequestCompletionEventListener extends BaseCompletionEventListener
205
- implements YouCompletionEventListener {
206
-
207
- @ Override
208
- public void onSerpResults (List <YouSerpResult > results ) {
209
- if (serpResultsListener != null ) {
210
- serpResultsListener .accept (results );
211
- }
212
- }
213
- }
214
-
215
124
private void sendInfo (SettingsState settings ) {
216
125
TelemetryAction .COMPLETION .createActionMessage ()
217
126
.property ("conversationId" , conversation .getId ().toString ())
0 commit comments