e71846aa Andy McCall 1 month ago ↶
1
2
3
4
5
6
7
8
9
10
11
12
use crate :: db:: { now_ms, Pool} ;
13
use crate :: Shared;
14
use rand:: RngCore;
15
use ring:: aead;
16
use sqlx:: FromRow;
17
use std:: path:: Path;
18
use std:: time:: Duration;
19
20
21
pub const APIS : [ & str ; 2 ] = [ " chat-completions" , " anthropic-messages" ] ;
22
23
24
pub const DEFAULT_BUDGET : i64 = 120_000 ;
5360cea3 Andy McCall 21 days ago ↶
25
26
27
pub const CALL_TIMEOUT : Duration = Duration:: from_secs( 90 ) ;
28
29
30
31
pub const BACKGROUND_TIMEOUT : Duration = Duration:: from_secs( 300 ) ;
32
33
34
pub const MAX_ANSWER_TOKENS : u32 = 16_000 ;
e71846aa Andy McCall 1 month ago ↶
35
36
const MAX_ANSWER : usize = 64 * 1024 ;
5436f0dd Andy McCall 19 days ago ↶
37
38
39
const STREAM_LIMIT : Duration = Duration:: from_secs( 15 * 60 ) ;
e71846aa Andy McCall 1 month ago ↶
40
2a83f530 Andy McCall 1 month ago ↶
41
pub use crate :: levels:: Scope;
e71846aa Andy McCall 1 month ago ↶
42
43
# [ derive ( Clone, Debug, FromRow ) ]
44
pub struct Assistant {
45
pub id : i64 ,
46
pub scope : String,
47
pub scope_id : i64 ,
48
pub name : String,
49
pub endpoint : String,
50
51
pub api : String,
52
pub model : String,
53
54
pub headers : String,
55
56
pub token : Vec< u8 > ,
57
pub enabled : bool ,
58
pub budget : i64 ,
59
pub created_ms : i64 ,
60
pub updated_ms : i64 ,
61
}
62
63
impl Assistant {
64
pub fn scope ( & self ) -> Scope {
65
match self .scope.as_str ( ) {
66
" org" => Scope:: Org( self .scope_id) ,
67
" repo" => Scope:: Repo( self .scope_id) ,
68
_ => Scope:: Site,
69
}
70
}
71
72
pub fn has_token ( & self ) -> bool {
73
! self .token.is_empty ( )
74
}
75
76
77
pub fn header_list ( & self ) -> Vec< ( String, String) > {
78
parse_headers ( & self .headers) .unwrap_or_default ( )
79
}
80
}
81
82
83
# [ derive ( Clone, Debug, Default ) ]
84
pub struct Fields {
85
pub name : String,
86
pub endpoint : String,
87
pub api : String,
88
pub model : String,
89
pub headers : String,
90
pub enabled : bool ,
91
pub budget : i64 ,
92
}
93
94
95
pub fn parse_headers ( text : & str ) -> Result< Vec< ( String, String) > , String> {
96
let mut out = vec! [ ] ;
97
for line in text.lines ( ) .map ( str :: trim) .filter ( | l | ! l.is_empty ( ) ) {
98
let Some ( ( name, value) ) = line.split_once ( ' :' ) else { return Err ( line.to_string ( ) ) } ;
99
let ( name, value) = ( name.trim ( ) , value.trim ( ) ) ;
100
let name_ok = ! name.is_empty ( ) & & name.chars ( ) .all ( | c | c.is_ascii_alphanumeric ( ) | | c = = ' -' | | c = = ' _' ) ;
101
if ! name_ok | | value.chars ( ) .any ( | c | c.is_control ( ) ) {
102
return Err ( line.to_string ( ) ) ;
103
}
104
out.push ( ( name.to_string ( ) , value.to_string ( ) ) ) ;
105
}
106
Ok ( out)
107
}
108
109
110
111
112
# [ derive ( Clone ) ]
113
pub struct SecretKey ( [u8 ; 32]) ;
114
115
impl SecretKey {
116
pub fn load_or_create ( keys : & Path) -> anyhow:: Result< SecretKey> {
4ac01a73 Andy McCall 1 month ago ↶
117
SecretKey:: load_or_create_named( keys, " assistants.key" )
118
}
119
120
121
122
pub fn load_or_create_named ( keys : & Path, file : & str ) -> anyhow:: Result< SecretKey> {
e71846aa Andy McCall 1 month ago ↶
123
std:: fs:: create_dir_all( keys) ? ;
4ac01a73 Andy McCall 1 month ago ↶
124
let path = keys.join ( file) ;
e71846aa Andy McCall 1 month ago ↶
125
if let Ok ( bytes) = std:: fs:: read( & path) {
126
if bytes.len ( ) = = 32 {
127
return Ok ( SecretKey( bytes.try_into ( ) .unwrap ( ) ) ) ;
128
}
129
}
130
let mut k = [ 0 u8 ; 32 ] ;
131
rand:: rngs:: OsRng.fill_bytes ( & mut k) ;
132
std:: fs:: write( & path, k) ? ;
133
# [ cfg ( unix ) ]
134
{
135
use std:: os:: unix:: fs:: PermissionsExt;
136
std:: fs:: set_permissions( & path, std:: fs:: Permissions:: from_mode( 0o600 ) ) ? ;
137
}
138
Ok ( SecretKey( k) )
139
}
140
141
fn key ( & self ) -> aead:: LessSafeKey {
142
aead:: LessSafeKey:: new( aead:: UnboundKey:: new( & aead:: CHACHA20_POLY1305 , & self .0 ) .expect ( " key length" ) )
143
}
144
145
146
147
pub fn seal ( & self , plain : & str ) -> Vec< u8 > {
148
if plain.is_empty ( ) {
149
return vec! [ ] ;
150
}
151
let mut nonce = [ 0 u8 ; 12 ] ;
152
rand:: rngs:: OsRng.fill_bytes ( & mut nonce) ;
153
let mut out = plain.as_bytes ( ) .to_vec ( ) ;
154
self .key ( ) .seal_in_place_append_tag ( aead:: Nonce:: assume_unique_for_key( nonce) , aead:: Aad:: empty( ) , & mut out) .expect ( " seal" ) ;
155
let mut sealed = nonce.to_vec ( ) ;
156
sealed.extend ( out) ;
157
sealed
158
}
159
160
161
162
pub fn open ( & self , sealed : & [u8 ]) -> Option< String> {
163
if sealed.is_empty ( ) {
164
return Some ( String :: new( ) ) ;
165
}
166
let ( nonce, rest) = sealed.split_at_checked ( 12 ) ? ;
167
let mut buf = rest.to_vec ( ) ;
168
let nonce = aead:: Nonce:: try_assume_unique_for_key( nonce) .ok ( ) ? ;
169
let plain = self .key ( ) .open_in_place ( nonce, aead:: Aad:: empty( ) , & mut buf) .ok ( ) ? ;
170
String :: from_utf8( plain.to_vec ( ) ) .ok ( )
171
}
172
}
173
174
175
176
pub async fn list ( pool : & Pool, scope : Scope) -> sqlx:: Result< Vec< Assistant> > {
5d90e0e0 Andy McCall 20 days ago
177
sqlx:: query_as( " SELECT * FROM assistants WHERE scope = $1 AND scope_id = $2 ORDER BY name" )
e71846aa Andy McCall 1 month ago ↶
178
.bind ( scope.kind ( ) )
179
.bind ( scope.id ( ) )
180
.fetch_all ( pool)
181
.await
182
}
183
184
pub async fn get ( pool : & Pool, id : i64 ) -> sqlx:: Result< Option< Assistant> > {
5d90e0e0 Andy McCall 20 days ago
185
sqlx:: query_as( " SELECT * FROM assistants WHERE id = $1" ) .bind ( id) .fetch_optional ( pool) .await
e71846aa Andy McCall 1 month ago ↶
186
}
187
188
189
pub async fn create ( pool : & Pool, scope : Scope, f : & Fields, token : & [u8 ]) -> sqlx:: Result< Option< Assistant> > {
190
let now = now_ms ( ) ;
191
let r = sqlx:: query_as(
192
" INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \
5d90e0e0 Andy McCall 20 days ago
193
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT DO NOTHING RETURNING *" ,
e71846aa Andy McCall 1 month ago ↶
194
)
195
.bind ( scope.kind ( ) )
196
.bind ( scope.id ( ) )
197
.bind ( & f.name)
198
.bind ( & f.endpoint)
199
.bind ( & f.api)
200
.bind ( & f.model)
201
.bind ( & f.headers)
202
.bind ( token)
203
.bind ( f.enabled)
204
.bind ( f.budget)
205
.bind ( now)
206
.bind ( now)
207
.fetch_optional ( pool)
208
.await? ;
209
Ok ( r)
210
}
211
212
213
214
pub async fn update ( pool : & Pool, id : i64 , f : & Fields, token : Option< & [ u8 ] > ) -> sqlx:: Result< bool > {
215
let r = match token {
216
Some ( t) => sqlx:: query(
5d90e0e0 Andy McCall 20 days ago
217
" UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, token = $8, updated_ms = $9 \
218
WHERE id = $10 AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = $11 AND o.id <> $12)" ,
e71846aa Andy McCall 1 month ago ↶
219
)
220
.bind ( & f.name)
221
.bind ( & f.endpoint)
222
.bind ( & f.api)
223
.bind ( & f.model)
224
.bind ( & f.headers)
225
.bind ( f.enabled)
226
.bind ( f.budget)
227
.bind ( t)
228
.bind ( now_ms ( ) )
229
.bind ( id)
230
.bind ( & f.name)
231
.bind ( id)
232
.execute ( pool)
233
.await? ,
234
None => sqlx:: query(
5d90e0e0 Andy McCall 20 days ago
235
" UPDATE assistants SET name = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, updated_ms = $8 \
236
WHERE id = $9 AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = $10 AND o.id <> $11)" ,
e71846aa Andy McCall 1 month ago ↶
237
)
238
.bind ( & f.name)
239
.bind ( & f.endpoint)
240
.bind ( & f.api)
241
.bind ( & f.model)
242
.bind ( & f.headers)
243
.bind ( f.enabled)
244
.bind ( f.budget)
245
.bind ( now_ms ( ) )
246
.bind ( id)
247
.bind ( & f.name)
248
.bind ( id)
249
.execute ( pool)
250
.await? ,
251
} ;
252
Ok ( r.rows_affected ( ) = = 1 )
253
}
254
ed89fae5 Andy McCall 1 month ago ↶
255
256
257
258
pub async fn rename_in_history ( pool : & Pool, scope : Scope, owner_hint : & str , old : & str , new : & str ) -> sqlx:: Result< ( ) > {
259
if old = = new | | old.is_empty ( ) | | new.is_empty ( ) {
260
return Ok ( ( ) ) ;
261
}
262
let like = format! ( " %{old} %" ) ;
263
match scope {
264
Scope:: Site => {
ffb67a81 Andy McCall 22 days ago
265
for table in [ " proposal_reviews" , " proposal_comments" ] {
5d90e0e0 Andy McCall 20 days ago
266
sqlx:: query( & format! ( " UPDATE {table} SET via = $1 WHERE via = $2" ) ) .bind ( new) .bind ( old) .execute ( pool) .await? ;
ffb67a81 Andy McCall 22 days ago
267
}
5d90e0e0 Andy McCall 20 days ago
268
sqlx:: query( " UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3" )
ed89fae5 Andy McCall 1 month ago ↶
269
.bind ( old)
270
.bind ( new)
271
.bind ( & like)
272
.execute ( pool)
273
.await? ;
274
}
275
Scope:: Org( _ ) => {
5d90e0e0 Andy McCall 20 days ago
276
let in_org = " proposal_id IN (SELECT p.id FROM proposals p JOIN repositories r ON r.id = p.repo_id WHERE r.owner = $1)" ;
ffb67a81 Andy McCall 22 days ago
277
for table in [ " proposal_reviews" , " proposal_comments" ] {
5d90e0e0 Andy McCall 20 days ago
278
sqlx:: query( & format! ( " UPDATE {table} SET via = $1 WHERE via = $2 AND {in_org} " ) )
ffb67a81 Andy McCall 22 days ago
279
.bind ( new)
280
.bind ( old)
281
.bind ( owner_hint)
282
.execute ( pool)
283
.await? ;
284
}
5d90e0e0 Andy McCall 20 days ago
285
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_org} " ) )
ed89fae5 Andy McCall 1 month ago ↶
286
.bind ( old)
287
.bind ( new)
288
.bind ( & like)
289
.bind ( owner_hint)
290
.execute ( pool)
291
.await? ;
292
}
293
Scope:: Repo( id) => {
5d90e0e0 Andy McCall 20 days ago
294
let in_repo = " proposal_id IN (SELECT id FROM proposals WHERE repo_id = $1)" ;
ffb67a81 Andy McCall 22 days ago
295
for table in [ " proposal_reviews" , " proposal_comments" ] {
5d90e0e0 Andy McCall 20 days ago
296
sqlx:: query( & format! ( " UPDATE {table} SET via = $1 WHERE via = $2 AND {in_repo} " ) )
ffb67a81 Andy McCall 22 days ago
297
.bind ( new)
298
.bind ( old)
299
.bind ( id)
300
.execute ( pool)
301
.await? ;
302
}
5d90e0e0 Andy McCall 20 days ago
303
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_repo} " ) )
ed89fae5 Andy McCall 1 month ago ↶
304
.bind ( old)
305
.bind ( new)
306
.bind ( & like)
307
.bind ( id)
308
.execute ( pool)
309
.await? ;
310
}
311
}
312
Ok ( ( ) )
313
}
314
e71846aa Andy McCall 1 month ago ↶
315
pub async fn set_enabled ( pool : & Pool, id : i64 , enabled : bool ) -> sqlx:: Result< ( ) > {
5d90e0e0 Andy McCall 20 days ago
316
sqlx:: query( " UPDATE assistants SET enabled = $1, updated_ms = $2 WHERE id = $3" )
e71846aa Andy McCall 1 month ago ↶
317
.bind ( enabled)
318
.bind ( now_ms ( ) )
319
.bind ( id)
320
.execute ( pool)
321
.await? ;
322
Ok ( ( ) )
323
}
324
325
pub async fn delete ( pool : & Pool, id : i64 ) -> sqlx:: Result< ( ) > {
5d90e0e0 Andy McCall 20 days ago
326
sqlx:: query( " DELETE FROM assistants WHERE id = $1" ) .bind ( id) .execute ( pool) .await? ;
e71846aa Andy McCall 1 month ago ↶
327
Ok ( ( ) )
328
}
329
330
331
332
pub async fn delete_scope ( pool : & Pool, scope : Scope) -> sqlx:: Result< ( ) > {
5d90e0e0 Andy McCall 20 days ago
333
sqlx:: query( " DELETE FROM assistants WHERE scope = $1 AND scope_id = $2" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
334
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2" )
335
.bind ( scope.kind ( ) )
336
.bind ( scope.id ( ) )
337
.execute ( pool)
338
.await? ;
e71846aa Andy McCall 1 month ago ↶
339
Ok ( ( ) )
340
}
341
342
343
pub async fn prompt ( pool : & Pool, scope : Scope) -> sqlx:: Result< String> {
5d90e0e0 Andy McCall 20 days ago
344
Ok ( sqlx:: query_scalar( " SELECT prompt FROM assistant_prompts WHERE scope = $1 AND scope_id = $2" )
e71846aa Andy McCall 1 month ago ↶
345
.bind ( scope.kind ( ) )
346
.bind ( scope.id ( ) )
347
.fetch_optional ( pool)
348
.await?
349
.unwrap_or_default ( ) )
350
}
351
352
pub async fn set_prompt ( pool : & Pool, scope : Scope, text : & str ) -> sqlx:: Result< ( ) > {
353
if text.trim ( ) .is_empty ( ) {
5d90e0e0 Andy McCall 20 days ago
354
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2" )
e71846aa Andy McCall 1 month ago ↶
355
.bind ( scope.kind ( ) )
356
.bind ( scope.id ( ) )
357
.execute ( pool)
358
.await? ;
359
} else {
360
sqlx:: query(
5d90e0e0 Andy McCall 20 days ago
361
" INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES ($1, $2, $3, $4) \
e71846aa Andy McCall 1 month ago ↶
362
ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms" ,
363
)
364
.bind ( scope.kind ( ) )
365
.bind ( scope.id ( ) )
366
.bind ( text.trim ( ) )
367
.bind ( now_ms ( ) )
368
.execute ( pool)
369
.await? ;
370
}
371
Ok ( ( ) )
372
}
373
374
375
376
377
378
# [ derive ( Clone, Debug ) ]
379
pub struct Available {
380
pub assistant : Assistant,
381
382
pub from : String,
383
384
pub off_above : bool ,
385
}
386
387
impl Available {
388
pub fn on ( & self ) -> bool {
389
self .assistant.enabled & & ! self .off_above
390
}
391
}
392
393
394
395
396
pub async fn available ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< Available> > {
2a83f530 Andy McCall 1 month ago ↶
397
let mut levels = crate :: levels:: above( pool, scope, owner) .await? ;
e71846aa Andy McCall 1 month ago ↶
398
levels.push ( ( scope, String :: new( ) ) ) ;
399
let mut out: Vec< Available> = vec! [ ] ;
400
let mut off: Vec< String> = vec! [ ] ;
401
for ( level, from) in levels {
402
for a in list ( pool, level) .await? {
403
let off_above = off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) ;
404
if ! a.enabled & & ! off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) {
405
off.push ( a.name.clone ( ) ) ;
406
}
407
out.retain ( | o | ! o.assistant.name.eq_ignore_ascii_case ( & a.name) ) ;
408
out.push ( Available { assistant: a, from: from.clone ( ) , off_above } ) ;
409
}
410
}
411
Ok ( out)
412
}
413
414
415
pub async fn enabled_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< Vec< Assistant> > {
416
Ok ( available ( pool, Scope:: Repo( repo.id) , & repo.owner) .await? .into_iter ( ) .filter ( Available:: on) .map ( | a | a.assistant ) .collect ( ) )
417
}
418
419
420
421
pub async fn inherited_prompts ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< ( String, String) > > {
422
let mut out = vec! [ ] ;
2a83f530 Andy McCall 1 month ago ↶
423
for ( level, from) in crate :: levels:: above( pool, scope, owner) .await? {
e71846aa Andy McCall 1 month ago ↶
424
let text = prompt ( pool, level) .await? ;
425
if ! text.is_empty ( ) {
426
out.push ( ( from, text) ) ;
427
}
428
}
429
Ok ( out)
430
}
431
432
433
434
pub async fn prompt_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< String> {
435
let scope = Scope:: Repo( repo.id) ;
436
let mut parts = inherited_prompts ( pool, scope, & repo.owner) .await? ;
437
let own = prompt ( pool, scope) .await? ;
438
if ! own.is_empty ( ) {
439
parts.push ( ( " this repository" .to_string ( ) , own) ) ;
440
}
441
Ok ( parts.into_iter ( ) .map ( | ( from , text ) | format! ( " ## Instructions from {from} \n \n {text} \n " ) ) .collect:: < Vec< _ > > ( ) .join ( " \n " ) )
442
}
443
444
445
446
447
5436f0dd Andy McCall 19 days ago ↶
448
449
450
451
452
5360cea3 Andy McCall 21 days ago ↶
453
pub async fn ask ( app : & Shared, a : & Assistant, system : & str , user : & str , wait : Duration) -> Result< String, String> {
5436f0dd Andy McCall 19 days ago ↶
454
ask_streaming ( app, a, system, user, wait, | _| { } ) .await
455
}
456
457
458
pub async fn ask_streaming (
459
app : & Shared,
460
a : & Assistant,
461
system : & str ,
462
user : & str ,
463
wait : Duration,
464
mut part : impl FnMut( &str ) ,
465
) -> Result< String, String> {
e71846aa Andy McCall 1 month ago ↶
466
let Some ( token) = app.assistant_key.open ( & a.token) else {
467
return Err ( " its token cannot be read: set it again" .into ( ) ) ;
468
} ;
5436f0dd Andy McCall 19 days ago ↶
469
let client = crate :: net:: client( & a.endpoint, & app.config.outbound.allow, STREAM_LIMIT .max ( wait) ) .await? ;
e71846aa Andy McCall 1 month ago ↶
470
let mut req = client.post ( & a.endpoint) .header ( " content-type" , " application/json" ) ;
471
let body = match a.api.as_str ( ) {
472
" anthropic-messages" => {
473
if ! token.is_empty ( ) {
474
req = req.header ( " x-api-key" , & token) ;
475
}
476
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
477
serde_json:: json! ( {
478
" model" : a.model,
5360cea3 Andy McCall 21 days ago ↶
479
" max_tokens" : MAX_ANSWER_TOKENS ,
e71846aa Andy McCall 1 month ago ↶
480
" system" : system,
481
" messages" : [ { " role" : " user" , " content" : user } ] ,
5436f0dd Andy McCall 19 days ago ↶
482
" stream" : true ,
e71846aa Andy McCall 1 month ago ↶
483
} )
484
}
485
_ => {
486
if ! token.is_empty ( ) {
487
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
488
}
5360cea3 Andy McCall 21 days ago ↶
489
490
491
e71846aa Andy McCall 1 month ago ↶
492
serde_json:: json! ( {
493
" model" : a.model,
494
" messages" : [
495
{ " role" : " system" , " content" : system } ,
496
{ " role" : " user" , " content" : user } ,
497
] ,
5436f0dd Andy McCall 19 days ago ↶
498
" stream" : true ,
e71846aa Andy McCall 1 month ago ↶
499
} )
500
}
501
} ;
502
for ( name, value) in a.header_list ( ) {
503
req = req.header ( name, value) ;
504
}
5436f0dd Andy McCall 19 days ago ↶
505
let quiet = | | format! ( " no answer within {} seconds" , wait.as_secs ( ) ) ;
506
let r = tokio:: time:: timeout( wait, req.json ( & body) .send ( ) ) .await.map_err ( | _| quiet ( ) ) ? .map_err ( | e | {
e71846aa Andy McCall 1 month ago ↶
507
if e.is_timeout ( ) {
5436f0dd Andy McCall 19 days ago ↶
508
quiet ( )
e71846aa Andy McCall 1 month ago ↶
509
} else {
510
scrub ( & e.to_string ( ) , & token)
511
}
512
} ) ? ;
513
let status = r.status ( ) ;
5436f0dd Andy McCall 19 days ago ↶
514
let streamed =
515
r.headers ( ) .get ( reqwest:: header:: CONTENT_TYPE ) .and_then ( | v | v.to_str ( ) .ok ( ) ) .is_some_and ( | t | t.starts_with ( " text/event-stream" ) ) ;
516
if status.is_success ( ) & & streamed {
517
let ( answer, stop) = read_stream ( & a.api, r, wait, & mut part) .await.map_err ( | e | scrub ( & e, & token) ) ? ;
518
return finish ( answer, stop) ;
519
}
520
let text: String =
521
tokio:: time:: timeout( wait, r.text ( ) ) .await.map_err ( | _| quiet ( ) ) ? .unwrap_or_default ( ) .chars ( ) .take ( MAX_ANSWER ) .collect ( ) ;
e71846aa Andy McCall 1 month ago ↶
522
if ! status.is_success ( ) {
523
let why = serde_json:: from_str:: < serde_json:: Value> ( & text)
524
.ok ( )
525
.and_then ( | v | v[ " error" ] [ " message" ] .as_str ( ) .or ( v[ " error" ] .as_str ( ) ) .or ( v[ " message" ] .as_str ( ) ) .map ( str :: to_string) )
526
.unwrap_or_else ( | | text.lines ( ) .next ( ) .unwrap_or ( " " ) .chars ( ) .take ( 160 ) .collect ( ) ) ;
527
return Err ( scrub ( & format! ( " {} {} " , status.as_u16 ( ) , why.trim ( ) ) , & token) ) ;
528
}
529
let v: serde_json:: Value = serde_json:: from_str( & text) .map_err ( | _| " the answer was not JSON" .to_string ( ) ) ? ;
5360cea3 Andy McCall 21 days ago ↶
530
let ( answer, stop) = answer_of ( & a.api, & v) ;
5436f0dd Andy McCall 19 days ago ↶
531
if ! answer.is_empty ( ) {
532
part ( & answer) ;
533
}
5360cea3 Andy McCall 21 days ago ↶
534
finish ( answer, stop)
5436f0dd Andy McCall 19 days ago ↶
535
}
536
537
538
539
async fn read_stream ( api : & str , mut r : reqwest:: Response, quiet : Duration, part : & mut impl FnMut( &str ) ) -> Result< ( String, Stop) , String> {
540
let mut events = Events:: new( api) ;
541
loop {
542
let chunk = match tokio:: time:: timeout( quiet, r.chunk ( ) ) .await {
543
Err ( _ ) if events.text.is_empty ( ) => return Err ( format! ( " no answer within {} seconds" , quiet.as_secs ( ) ) ) ,
544
545
Err ( _ ) => return Ok ( ( events.text, Stop:: Quiet) ) ,
546
Ok ( Err ( e) ) if events.text.is_empty ( ) => return Err ( e.to_string ( ) ) ,
547
Ok ( Err ( _ ) ) => return Ok ( ( events.text, Stop:: Quiet) ) ,
548
Ok ( Ok ( None ) ) => return Ok ( ( events.text, events.stop.unwrap_or ( Stop:: Done) ) ) ,
549
Ok ( Ok ( Some ( c) ) ) => c,
550
} ;
551
let before = events.text.len ( ) ;
552
events.feed ( & chunk) ? ;
553
if events.text.len ( ) > before {
554
part ( & events.text[ before.. ] ) ;
555
}
556
if events.done | | events.text.len ( ) > = MAX_ANSWER {
557
return Ok ( ( events.text, events.stop.unwrap_or ( Stop:: Done) ) ) ;
558
}
559
}
560
}
561
562
563
564
565
struct Events {
566
anthropic : bool ,
567
buf : Vec< u8 > ,
568
text : String,
569
stop : Option< Stop> ,
570
done : bool ,
571
}
572
573
impl Events {
574
fn new ( api : & str ) -> Events {
575
Events { anthropic: api = = " anthropic-messages" , buf: Vec :: new( ) , text: String :: new( ) , stop: None , done: false }
576
}
577
578
fn feed ( & mut self , bytes : & [u8 ]) -> Result< ( ) , String> {
579
self .buf.extend_from_slice ( bytes) ;
580
while let Some ( end) = self .buf.iter ( ) .position ( | b | * b = = b ' \n ' ) {
581
let line: Vec< u8 > = self .buf.drain ( .. = end) .collect ( ) ;
582
let line = String :: from_utf8_lossy( & line) ;
583
let Some ( data) = line.trim_end_matches ( [ ' \r ' , ' \n ' ] ) .strip_prefix ( " data:" ) else {
584
continue ;
585
} ;
586
let data = data.trim_start ( ) ;
587
if data = = " [DONE]" {
588
self .done = true ;
589
continue ;
590
}
591
let Ok ( v) = serde_json:: from_str:: < serde_json:: Value> ( data) else {
592
continue ;
593
} ;
594
self .event ( & v) ? ;
595
}
596
Ok ( ( ) )
597
}
598
599
fn event ( & mut self , v : & serde_json:: Value) -> Result< ( ) , String> {
600
if self .anthropic {
601
match v[ " type" ] .as_str ( ) {
602
Some ( " content_block_delta" ) if v[ " delta" ] [ " type" ] = = " text_delta" => {
603
self .text.push_str ( v[ " delta" ] [ " text" ] .as_str ( ) .unwrap_or ( " " ) ) ;
604
}
605
Some ( " message_delta" ) => {
606
self .stop = match v[ " delta" ] [ " stop_reason" ] .as_str ( ) {
607
Some ( " max_tokens" ) => Some ( Stop:: Length) ,
608
Some ( " refusal" ) => Some ( Stop:: Declined( v[ " delta" ] [ " stop_details" ] [ " category" ] .as_str ( ) .map ( str :: to_string) ) ) ,
609
Some ( _ ) => Some ( Stop:: Done) ,
610
None => self .stop.take ( ) ,
611
} ;
612
}
613
Some ( " message_stop" ) => self .done = true ,
614
Some ( " error" ) => {
615
return Err ( v[ " error" ] [ " message" ] .as_str ( ) .unwrap_or ( " the provider reported an error" ) .to_string ( ) ) ;
616
}
617
_ => { }
618
}
619
} else {
620
if let Some ( e) = v[ " error" ] [ " message" ] .as_str ( ) {
621
return Err ( e.to_string ( ) ) ;
622
}
623
let choice = & v[ " choices" ] [ 0 ] ;
624
if let Some ( t) = choice[ " delta" ] [ " content" ] .as_str ( ) {
625
self .text.push_str ( t) ;
626
}
627
match choice[ " finish_reason" ] .as_str ( ) {
628
Some ( " length" ) => self .stop = Some ( Stop:: Length) ,
629
Some ( " content_filter" ) => self .stop = Some ( Stop:: Declined( None ) ) ,
630
Some ( _ ) => self .stop = Some ( Stop:: Done) ,
631
None => { }
632
}
633
}
634
Ok ( ( ) )
635
}
5360cea3 Andy McCall 21 days ago ↶
636
}
637
638
639
# [ derive ( Debug, PartialEq, Eq ) ]
640
enum Stop {
641
642
Done,
643
644
Length,
645
646
Declined( Option< String> ) ,
5436f0dd Andy McCall 19 days ago ↶
647
648
Quiet,
5360cea3 Andy McCall 21 days ago ↶
649
}
650
651
652
653
fn answer_of ( api : & str , v : & serde_json:: Value) -> ( String, Stop) {
654
match api {
655
" anthropic-messages" => {
656
let text = v[ " content" ]
657
.as_array ( )
658
.map ( | parts | {
659
parts
660
.iter ( )
661
.filter ( | p | p[ " type" ] .as_str ( ) .is_none_or ( | t | t = = " text" ) )
662
.filter_map ( | p | p[ " text" ] .as_str ( ) )
663
.collect:: < Vec< _ > > ( )
664
.join ( " " )
665
} )
666
.unwrap_or_default ( ) ;
667
let stop = match v[ " stop_reason" ] .as_str ( ) {
668
Some ( " max_tokens" ) => Stop:: Length,
669
Some ( " refusal" ) => Stop:: Declined( v[ " stop_details" ] [ " category" ] .as_str ( ) .map ( str :: to_string) ) ,
670
_ => Stop:: Done,
671
} ;
672
( text, stop)
673
}
674
_ => {
675
let choice = & v[ " choices" ] [ 0 ] ;
676
let text = match & choice[ " message" ] [ " content" ] {
677
serde_json:: Value:: String( s) => s.clone ( ) ,
678
serde_json:: Value:: Array( parts) => parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) ,
679
_ => String :: new( ) ,
680
} ;
681
let stop = match choice[ " finish_reason" ] .as_str ( ) {
682
Some ( " length" ) => Stop:: Length,
683
Some ( " content_filter" ) => Stop:: Declined( None ) ,
684
_ => Stop:: Done,
685
} ;
686
( text, stop)
687
}
e71846aa Andy McCall 1 month ago ↶
688
}
689
}
690
5360cea3 Andy McCall 21 days ago ↶
691
692
693
fn finish ( answer : String, stop : Stop) -> Result< String, String> {
694
match stop {
695
Stop:: Declined( category) => Err ( match category {
696
Some ( c) => format! ( " it declined to answer ({c} )" ) ,
697
None => " it declined to answer" .into ( ) ,
698
} ) ,
699
Stop:: Length if answer.trim ( ) .is_empty ( ) => {
700
Err ( " it ran out of room before writing an answer (it spent its length limit thinking): try a smaller proposal" .into ( ) )
701
}
702
Stop:: Length => Ok ( format! ( " {} \n \n *(Cut short: the answer reached its length limit.)*" , answer.trim_end ( ) ) ) ,
5436f0dd Andy McCall 19 days ago ↶
703
Stop:: Quiet if answer.trim ( ) .is_empty ( ) => Err ( " it went quiet before writing an answer" .into ( ) ) ,
704
Stop:: Quiet => Ok ( format! ( " {} \n \n *(Cut short: the assistant went quiet part of the way through.)*" , answer.trim_end ( ) ) ) ,
5360cea3 Andy McCall 21 days ago ↶
705
Stop:: Done if answer.trim ( ) .is_empty ( ) => Err ( " the answer carried no text" .into ( ) ) ,
706
Stop:: Done => Ok ( answer) ,
707
}
708
}
709
155cab14 Andy McCall 1 month ago ↶
710
711
712
pub const PRESETS : & [ ( & str , & str , & str , & str , & str ) ] = & [
713
( " anthropic" , " Anthropic (Claude)" , " Claude" , " https://api.anthropic.com/v1/messages" , " anthropic-messages" ) ,
714
( " openai" , " OpenAI (ChatGPT)" , " ChatGPT" , " https://api.openai.com/v1/chat/completions" , " chat-completions" ) ,
715
( " gemini" , " Google Gemini" , " Gemini" , " https://generativelanguage.googleapis.com/v1beta/openai/chat/completions" , " chat-completions" ) ,
716
( " deepseek" , " DeepSeek" , " DeepSeek" , " https://api.deepseek.com/chat/completions" , " chat-completions" ) ,
717
( " mistral" , " Mistral" , " Mistral" , " https://api.mistral.ai/v1/chat/completions" , " chat-completions" ) ,
718
( " perplexity" , " Perplexity" , " Perplexity" , " https://api.perplexity.ai/chat/completions" , " chat-completions" ) ,
719
( " ollama" , " Ollama (on your network)" , " Ollama" , " http://localhost:11434/v1/chat/completions" , " chat-completions" ) ,
720
] ;
721
722
723
724
pub fn models_url ( endpoint : & str ) -> Option< String> {
725
let e = endpoint.trim_end_matches ( ' /' ) ;
726
e.strip_suffix ( " /chat/completions" ) .or_else ( | | e.strip_suffix ( " /messages" ) ) .map ( | base | format! ( " {base} /models" ) )
727
}
728
729
730
731
9a777fec Andy McCall 21 days ago ↶
732
pub async fn list_models ( endpoint : & str , api : & str , token : & str , allow : & [String]) -> Result< Vec< String> , String> {
155cab14 Andy McCall 1 month ago ↶
733
let url = models_url ( endpoint) .ok_or ( " this endpoint has no model list next to it" ) ? ;
9a777fec Andy McCall 21 days ago ↶
734
let client = crate :: net:: client( & url, allow, std:: time:: Duration:: from_secs( 15 ) ) .await? ;
155cab14 Andy McCall 1 month ago ↶
735
let mut req = client.get ( & url) ;
736
if api = = " anthropic-messages" {
737
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
738
if ! token.is_empty ( ) {
739
req = req.header ( " x-api-key" , token) ;
740
}
741
} else if ! token.is_empty ( ) {
742
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
743
}
744
let r = req.send ( ) .await.map_err ( | e | scrub ( & e.to_string ( ) , token) ) ? ;
745
let status = r.status ( ) ;
746
if ! status.is_success ( ) {
747
return Err ( format! ( " the model list answered {} " , status.as_u16 ( ) ) ) ;
748
}
749
let v: serde_json:: Value = r.json ( ) .await.map_err ( | _| " the model list was not JSON" .to_string ( ) ) ? ;
750
let skip = [ " embed" , " tts" , " whisper" , " dall-e" , " moderation" , " audio" , " image" , " realtime" , " transcribe" , " search-preview" ] ;
751
let mut out: Vec< String> = v[ " data" ]
752
.as_array ( )
753
.or ( v[ " models" ] .as_array ( ) )
754
.ok_or ( " the model list had no data" ) ?
755
.iter ( )
756
.filter_map ( | m | m[ " id" ] .as_str ( ) .or ( m[ " name" ] .as_str ( ) ) )
757
.map ( | id | id.strip_prefix ( " models/" ) .unwrap_or ( id) .to_string ( ) )
758
.filter ( | id | ! skip.iter ( ) .any ( | s | id.to_lowercase ( ) .contains ( s) ) )
759
.collect ( ) ;
760
out.sort ( ) ;
761
out.dedup ( ) ;
762
if out.is_empty ( ) {
763
return Err ( " the model list was empty" .into ( ) ) ;
764
}
765
Ok ( out)
766
}
767
e71846aa Andy McCall 1 month ago ↶
768
769
pub async fn test ( app : & Shared, a : & Assistant) -> Result< String, String> {
5360cea3 Andy McCall 21 days ago ↶
770
let answer =
771
ask ( app, a, " You are being connected to an OxygenHub site. Answer briefly." , " Reply with the single word OK." , CALL_TIMEOUT )
772
.await? ;
e71846aa Andy McCall 1 month ago ↶
773
Ok ( answer.trim ( ) .chars ( ) .take ( 80 ) .collect ( ) )
774
}
775
776
777
fn scrub ( text : & str , token : & str ) -> String {
778
if token.is_empty ( ) {
779
text.to_string ( )
780
} else {
781
text.replace ( token, " ***" )
782
}
783
}
784
785
# [ cfg ( test ) ]
786
mod tests {
787
use super :: * ;
5436f0dd Andy McCall 19 days ago ↶
788
c43e4f63 Andy McCall 11 days ago ↶
789
790
791
# [ test ]
792
fn a_draft_s_title_line_is_split_off ( ) {
793
assert_eq! (
794
split_title ( " Title: feat(flow): hand jobs back\n \n **Summary.** Runners leave." ) ,
795
( Some ( " feat(flow): hand jobs back" .to_string ( ) ) , " **Summary.** Runners leave." .to_string ( ) )
796
) ;
797
assert_eq! ( split_title ( " title: \" Quoted\" \n Body" ) , ( Some ( " Quoted" .to_string ( ) ) , " Body" .to_string ( ) ) ) ;
798
assert_eq! ( split_title ( " **Summary.** No title here." ) , ( None , " **Summary.** No title here." .to_string ( ) ) ) ;
799
assert_eq! ( split_title ( " Title:\n Body" ) , ( None , " Body" .to_string ( ) ) , " an empty title is none" ) ;
800
}
801
5436f0dd Andy McCall 19 days ago ↶
802
803
804
# [ test ]
805
fn anthropic_events_give_the_text_whole ( ) {
806
let sse = " event: message_start\n data: {\" type\" :\" message_start\" }\n \n \
807
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" thinking_delta\" ,\" thinking\" :\" hmm\" }}\n \n \
808
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" text_delta\" ,\" text\" :\" Adds a \" }}\n \n \
809
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" text_delta\" ,\" text\" :\" line — café.\" }}\n \n \
810
data: {\" type\" :\" message_delta\" ,\" delta\" :{\" stop_reason\" :\" end_turn\" }}\n \n \
811
data: {\" type\" :\" message_stop\" }\n \n " ;
812
for size in [ 1 , 3 , 7 , 64 , sse.len ( ) ] {
813
let mut e = Events:: new( " anthropic-messages" ) ;
814
for chunk in sse.as_bytes ( ) .chunks ( size) {
815
e.feed ( chunk) .unwrap ( ) ;
816
}
817
assert_eq! ( ( e.text.as_str ( ) , e.stop.as_ref ( ) , e.done) , ( " Adds a line — café." , Some ( & Stop:: Done) , true ) , " chunks of {size}" ) ;
818
}
819
let mut cut = Events:: new( " anthropic-messages" ) ;
820
cut.feed ( b " data: {\" type\" :\" message_delta\" ,\" delta\" :{\" stop_reason\" :\" max_tokens\" }}\n " ) .unwrap ( ) ;
821
assert_eq! ( cut.stop, Some ( Stop:: Length) ) ;
822
let mut bad = Events:: new( " anthropic-messages" ) ;
823
let e = bad.feed ( b " data: {\" type\" :\" error\" ,\" error\" :{\" type\" :\" overloaded_error\" ,\" message\" :\" Overloaded\" }}\n " ) .unwrap_err ( ) ;
824
assert_eq! ( e, " Overloaded" ) ;
825
}
826
827
# [ test ]
828
fn chat_completion_events_give_the_text ( ) {
829
let sse = " data: {\" choices\" :[{\" delta\" :{\" role\" :\" assistant\" }}]}\n \n \
830
data: {\" choices\" :[{\" delta\" :{\" content\" :\" Hello\" }}]}\n \n \
831
data: {\" choices\" :[{\" delta\" :{\" content\" :\" , world\" },\" finish_reason\" :\" length\" }]}\n \n \
832
data: [DONE]\n \n " ;
833
let mut e = Events:: new( " chat-completions" ) ;
834
for chunk in sse.as_bytes ( ) .chunks ( 5 ) {
835
e.feed ( chunk) .unwrap ( ) ;
836
}
837
assert_eq! ( ( e.text.as_str ( ) , e.stop.as_ref ( ) , e.done) , ( " Hello, world" , Some ( & Stop:: Length) , true ) ) ;
838
}
839
840
841
842
843
# [ tokio ::test ]
844
async fn a_slow_steady_stream_is_read_and_a_silent_one_cut ( ) {
845
async fn serve ( parts : usize , gap_ms : u64 , then_silent : bool ) -> String {
846
let router = axum:: Router:: new( ) .route (
847
" /v1/messages" ,
848
axum:: routing:: post( move | | async move {
849
let ( tx, rx) = tokio:: sync:: mpsc:: channel:: < Result< String, std:: io:: Error> > ( 4 ) ;
850
tokio:: spawn( async move {
851
for i in 0 .. parts {
852
tokio:: time:: sleep( Duration:: from_millis( gap_ms) ) .await;
853
let ev = format! (
854
" data: {{ \" type\" :\" content_block_delta\" ,\" delta\" :{{ \" type\" :\" text_delta\" ,\" text\" :\" {i} \" }} }} \n \n "
855
) ;
856
if tx.send ( Ok ( ev) ) .await.is_err ( ) {
857
return ;
858
}
859
}
860
if then_silent {
861
tokio:: time:: sleep( Duration:: from_secs( 5 ) ) .await;
862
}
863
let _ = tx.send ( Ok ( " data: {\" type\" :\" message_stop\" }\n \n " .into ( ) ) ) .await;
864
} ) ;
865
let body = axum:: body:: Body:: from_stream( tokio_stream:: wrappers:: ReceiverStream:: new( rx) ) ;
866
( [ ( axum:: http:: header:: CONTENT_TYPE , " text/event-stream" ) ] , body)
867
} ) ,
868
) ;
869
let listener = tokio:: net:: TcpListener:: bind( " 127.0.0.1:0" ) .await.unwrap ( ) ;
870
let port = listener.local_addr ( ) .unwrap ( ) .port ( ) ;
871
tokio:: spawn( async move { axum:: serve( listener, router) .await.unwrap ( ) } ) ;
872
format! ( " http://127.0.0.1:{port} /v1/messages" )
873
}
874
let quiet = Duration:: from_millis( 400 ) ;
875
876
let url = serve ( 6 , 200 , false ) .await;
877
let r = reqwest:: Client:: new( ) .post ( & url) .send ( ) .await.unwrap ( ) ;
878
let mut parts = Vec :: new( ) ;
879
let ( text, stop) = read_stream ( " anthropic-messages" , r, quiet, & mut | p: & str | parts.push ( p.to_string ( ) ) ) .await.unwrap ( ) ;
880
assert_eq! ( ( text.as_str ( ) , stop) , ( " 0 1 2 3 4 5 " , Stop:: Done) ) ;
881
assert_eq! ( parts.len ( ) , 6 , " each part handed out as it came" ) ;
882
883
let url = serve ( 2 , 50 , true ) .await;
884
let r = reqwest:: Client:: new( ) .post ( & url) .send ( ) .await.unwrap ( ) ;
885
let ( text, stop) = read_stream ( " anthropic-messages" , r, quiet, & mut | _ : & str | { } ) .await.unwrap ( ) ;
886
assert_eq! ( ( text.as_str ( ) , stop) , ( " 0 1 " , Stop:: Quiet) ) ;
887
assert! ( finish ( text, Stop:: Quiet) .unwrap ( ) .contains ( " went quiet" ) ) ;
888
}
e71846aa Andy McCall 1 month ago ↶
889
890
# [ test ]
891
fn tokens_seal_and_open_only_under_their_key ( ) {
892
let dir = tempfile:: tempdir( ) .unwrap ( ) ;
893
let k = SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) ;
894
let sealed = k.seal ( " sk-very-secret" ) ;
895
assert_ne! ( sealed, b " sk-very-secret" ) ;
896
assert! ( ! String :: from_utf8_lossy( & sealed) .contains ( " secret" ) ) ;
897
assert_eq! ( k.open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) ) ;
898
assert_eq! ( SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) .open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) , " the key is kept" ) ;
899
let other = SecretKey:: load_or_create( tempfile:: tempdir( ) .unwrap ( ) .path ( ) ) .unwrap ( ) ;
900
assert_eq! ( other.open ( & sealed) , None ) ;
901
assert_eq! ( k.open ( & [ 1 , 2 , 3 ] ) , None ) ;
902
assert_eq! ( k.seal ( " " ) , Vec : : < u8 > :: new( ) ) ;
903
assert_eq! ( k.open ( & [ ] ) .as_deref ( ) , Some ( " " ) ) ;
5360cea3 Andy McCall 21 days ago ↶
904
}
905
906
# [ test ]
907
fn answers_say_why_they_stopped ( ) {
908
let a = | v : serde_json:: Value| {
909
let ( text, stop) = answer_of ( " anthropic-messages" , & v) ;
910
finish ( text, stop)
911
} ;
912
913
let thought = serde_json:: json! ( { " stop_reason" : " max_tokens" , " content" : [ { " type" : " thinking" , " thinking" : " " } ] } ) ;
914
assert! ( a ( thought) .unwrap_err ( ) .contains ( " ran out of room" ) ) ;
915
let cut = serde_json:: json! ( { " stop_reason" : " max_tokens" , " content" : [ { " type" : " thinking" , " thinking" : " " } , { " type" : " text" , " text" : " Half" } ] } ) ;
916
assert! ( a ( cut) .unwrap ( ) .starts_with ( " Half\n \n *(Cut short" ) ) ;
917
let declined = serde_json:: json! ( { " stop_reason" : " refusal" , " stop_details" : { " category" : " cyber" } , " content" : [ ] } ) ;
918
assert_eq! ( a ( declined) .unwrap_err ( ) , " it declined to answer (cyber)" ) ;
919
let ok = serde_json:: json! ( { " stop_reason" : " end_turn" , " content" : [ { " type" : " text" , " text" : " Fine." } ] } ) ;
920
assert_eq! ( a ( ok) .unwrap ( ) , " Fine." ) ;
921
let chat = | v : serde_json:: Value| {
922
let ( text, stop) = answer_of ( " chat-completions" , & v) ;
923
finish ( text, stop)
924
} ;
925
assert! ( chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " length" , " message" : { " content" : " " } } ] } ) )
926
.unwrap_err ( )
927
.contains ( " ran out of room" ) ) ;
928
assert_eq! (
929
chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " content_filter" , " message" : { " content" : " " } } ] } ) ) .unwrap_err ( ) ,
930
" it declined to answer"
931
) ;
932
assert_eq! (
933
chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " stop" , " message" : { " content" : " Yes." } } ] } ) ) .unwrap ( ) ,
934
" Yes."
935
) ;
e71846aa Andy McCall 1 month ago ↶
936
}
937
938
# [ test ]
c08cec37 Andy McCall 19 days ago ↶
939
fn a_refusal_for_length_is_told_from_others ( ) {
940
for e in [
941
" 400 This model's maximum context length is 8192 tokens. Please reduce the length of the messages." ,
942
" 400 prompt is too long: 210000 tokens > 200000 maximum" ,
943
" 413 Request too large" ,
944
" 400 input length exceeds the context window" ,
945
] {
946
assert! ( super :: too_long( e) , " {e}" ) ;
947
}
948
for e in [ " 401 invalid api key" , " 500 too long a queue" , " no answer within 90 seconds" , " 400 model not found" ] {
949
assert! ( ! super :: too_long( e) , " {e}" ) ;
950
}
951
}
952
953
# [ test ]
e71846aa Andy McCall 1 month ago ↶
954
fn headers_parse_or_name_the_bad_line ( ) {
955
assert_eq! (
956
parse_headers ( " X-Org: acme\n \n OpenAI-Beta : v1 " ) .unwrap ( ) ,
957
vec! [ ( " X-Org" .into ( ) , " acme" .into ( ) ) , ( " OpenAI-Beta" .into ( ) , " v1" .into ( ) ) ]
958
) ;
959
assert_eq! ( parse_headers ( " no colon" ) .unwrap_err ( ) , " no colon" ) ;
960
assert_eq! ( parse_headers ( " bad name: x" ) .unwrap_err ( ) , " bad name: x" ) ;
961
}
962
}
7d2eb590 Andy McCall 1 month ago ↶
963
964
965
966
967
pub const CALLS_PER_HOUR : usize = 30 ;
968
969
970
# [ derive ( Default ) ]
971
pub struct Gate {
972
busy : std:: collections:: HashSet< ( i64 , String) > ,
973
recent : std:: collections:: HashMap< i64 , Vec< i64 > > ,
974
}
975
976
impl Gate {
977
978
fn take ( & mut self , repo_id : i64 , proposal_id : i64 , name : & str ) -> Result< ( ) , String> {
979
let key = ( proposal_id, name.to_lowercase ( ) ) ;
980
if self .busy.contains ( & key) {
981
return Err ( format! ( " {name} is already answering this proposal" ) ) ;
982
}
983
let now = now_ms ( ) ;
984
let calls = self .recent.entry ( repo_id) .or_default ( ) ;
985
calls.retain ( | t | now - * t < 3_600_000 ) ;
986
if calls.len ( ) > = CALLS_PER_HOUR {
987
return Err ( format! ( " this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later" ) ) ;
988
}
989
calls.push ( now) ;
990
self .busy.insert ( key) ;
991
Ok ( ( ) )
992
}
993
994
fn release ( & mut self , proposal_id : i64 , name : & str ) {
995
self .busy.remove ( & ( proposal_id, name.to_lowercase ( ) ) ) ;
996
}
b5a0a69b Andy McCall 20 days ago ↶
997
998
999
pub fn answering ( & self , proposal_id : i64 ) -> Vec< String> {
1000
let mut v: Vec< String> = self .busy.iter ( ) .filter ( | ( p , _) | * p = = proposal_id) .map ( | ( _, n ) | n.clone ( ) ) .collect ( ) ;
1001
v.sort ( ) ;
1002
v
1003
}
7d2eb590 Andy McCall 1 month ago ↶
1004
}
1005
1006
1007
1008
pub fn review_context (
1009
row : & crate:: db:: repos:: RepoRow,
1010
p : & crate:: db:: proposals:: ProposalRow,
1011
loaded : & crate:: web:: proposals:: Loaded,
1012
budget : usize ,
1013
) -> String {
1014
let mut head =
1015
format! ( " # Proposal #{} : {} \n \n Repository {} /{} · into {} · by {} \n \n " , p.number, p.title, row.owner, row.name, p.target, p.author) ;
1016
if ! p.body.trim ( ) .is_empty ( ) {
1017
head.push_str ( " ## Description\n \n " ) ;
1018
head.push_str ( p.body.trim ( ) ) ;
1019
head.push_str ( " \n \n " ) ;
1020
}
1021
context_text ( head, & loaded.changes, & loaded.diffs, budget)
1022
}
1023
1024
1025
pub fn draft_context (
1026
row : & crate:: db:: repos:: RepoRow,
1027
title : & str ,
1028
target : & str ,
1029
changes : & [crate:: repos:: view:: ChangeView],
1030
diffs : & [oxy_core:: repo:: FileDiff],
1031
budget : usize ,
1032
) -> String {
1033
let head = format! (
1034
" # Proposal: {} \n \n Repository {} /{} · into {} \n \n " ,
1035
if title.trim ( ) .is_empty ( ) { " (untitled)" } else { title.trim ( ) } ,
1036
row.owner,
1037
row.name,
1038
target
1039
) ;
1040
context_text ( head, changes, diffs, budget)
1041
}
1042
1043
1044
1045
fn context_text ( mut text : String, changes : & [crate:: repos:: view:: ChangeView], diffs : & [oxy_core:: repo:: FileDiff], budget : usize ) -> String {
1046
text.push_str ( & format! ( " ## Changes ({} )\n \n " , changes.len ( ) ) ) ;
1047
for c in changes {
1048
text.push_str ( & format! ( " - {} {} \n " , c.short, c.title) ) ;
1049
let body = c.body.trim ( ) ;
1050
if ! body.is_empty ( ) {
1051
for line in body.lines ( ) .take ( 20 ) {
1052
text.push_str ( " " ) ;
1053
text.push_str ( line) ;
1054
text.push ( ' \n ' ) ;
1055
}
1056
}
1057
}
c08cec37 Andy McCall 19 days ago ↶
1058
1059
1060
text.push_str ( & format! ( " \n ## Files ({} )\n \n " , diffs.len ( ) ) ) ;
1061
let mut parts: Vec< ( String, String) > = Vec :: with_capacity( diffs.len ( ) ) ;
7d2eb590 Andy McCall 1 month ago ↶
1062
for f in diffs {
c08cec37 Andy McCall 19 days ago ↶
1063
let ( added, removed) = f.hunks.iter ( ) .fold ( ( 0 , 0 ) , | ( a , r ) , h | ( a + h.added.len ( ) , r + h.removed.len ( ) ) ) ;
1064
text.push_str ( & format! ( " - {} ({} , +{added} -{removed} )\n " , f.path, f.kind) ) ;
1065
let mut diff = format! ( " ### {} ({} )\n " , f.path, f.kind) ;
7d2eb590 Andy McCall 1 month ago ↶
1066
if let Some ( from) = & f.from {
1067
diff.push_str ( & format! ( " renamed from {from} \n " ) ) ;
1068
}
1069
let binary = f.hunks.iter ( ) .any ( | h | h.added.iter ( ) .chain ( h.removed.iter ( ) .map ( | ( _, l ) | l ) ) .any ( | l | l.contains ( & 0 ) ) ) ;
1070
if binary {
1071
diff.push_str ( " (binary)\n \n " ) ;
c08cec37 Andy McCall 19 days ago ↶
1072
parts.push ( ( f.path.clone ( ) , diff) ) ;
7d2eb590 Andy McCall 1 month ago ↶
1073
continue ;
1074
}
1075
diff.push_str ( " ```\n " ) ;
1076
for h in & f.hunks {
1077
diff.push_str ( & format! ( " @@ after line {} @@\n " , h.base_line) ) ;
1078
for ( _ , l) in & h.removed {
1079
diff.push ( ' -' ) ;
1080
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
1081
diff.push ( ' \n ' ) ;
1082
}
1083
for l in & h.added {
1084
diff.push ( ' +' ) ;
1085
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
1086
diff.push ( ' \n ' ) ;
1087
}
1088
}
1089
diff.push_str ( " ```\n \n " ) ;
c08cec37 Andy McCall 19 days ago ↶
1090
parts.push ( ( f.path.clone ( ) , diff) ) ;
7d2eb590 Andy McCall 1 month ago ↶
1091
}
c08cec37 Andy McCall 19 days ago ↶
1092
text.push_str ( " \n ## Diff\n \n " ) ;
1093
1094
1095
1096
let head = text.chars ( ) .count ( ) ;
1097
let mut room = budget.saturating_sub ( head + 200 ) ;
1098
let mut left_out = Vec :: new( ) ;
1099
for ( path, diff) in parts {
1100
let n = diff.chars ( ) .count ( ) ;
1101
if n < = room {
1102
text.push_str ( & diff) ;
1103
room - = n;
1104
} else if room > = 1_500 {
1105
let keep = room - 80 ;
1106
text.extend ( diff.chars ( ) .take ( keep) ) ;
1107
text.push_str ( " \n (the rest of this file was cut to fit)\n ```\n \n " ) ;
1108
left_out.push ( format! ( " {path} (cut)" ) ) ;
1109
room = 0 ;
1110
} else {
1111
left_out.push ( path) ;
1112
}
1113
}
1114
if ! left_out.is_empty ( ) {
1115
let names: String = left_out.join ( " , " ) .chars ( ) .take ( room.max ( 400 ) ) .collect ( ) ;
1116
text.push_str ( & format! ( " (left out to fit the context budget, {} file(s): {names} )\n " , left_out.len ( ) ) ) ;
7d2eb590 Andy McCall 1 month ago ↶
1117
}
1118
text
1119
}
1120
1121
1122
1123
pub fn review_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
1124
let mut s = format! (
1125
" You are {name} , an AI assistant reviewing a proposal on {site} , an OxygenHub site. \
1126
A person asked for your review; it is recorded as a review comment made via you, under their name. \
1127
You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \
1128
then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \
1129
keep praise out unless something is notably well done; say so plainly when you find nothing wrong." ,
1130
name = a.name
1131
) ;
1132
if ! prompts.trim ( ) .is_empty ( ) {
1133
s.push_str ( " \n \n " ) ;
1134
s.push_str ( prompts.trim ( ) ) ;
1135
}
1136
s
1137
}
1138
1139
1140
1141
1142
1143
pub async fn request_review (
1144
app : & Shared,
1145
row : & crate:: db:: repos:: RepoRow,
1146
p : & crate:: db:: proposals:: ProposalRow,
1147
by : & crate:: db:: accounts:: Account,
1148
a : Assistant,
1149
) -> Result< ( ) , String> {
1150
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
1151
let pool = & app.db;
1152
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} to review" , a.name) , " " ) .await.ok ( ) ;
33d98147 Andy McCall 16 days ago
1153
app.bus.live ( & crate :: events:: topic:: proposal( row.id, p.number) , " assistant" ) ;
7d2eb590 Andy McCall 1 month ago ↶
1154
crate :: db:: accounts:: audit(
1155
pool,
1156
Some ( by) ,
1157
" assistant.asked" ,
1158
& format! ( " {} /{} " , row.owner, row.name) ,
1159
& format! ( " {} #{} " , a.name, p.number) ,
1160
)
1161
.await
1162
.ok ( ) ;
1163
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
1164
tokio:: spawn( async move {
1165
let outcome = answer_review ( & app, & row, & p, & a) .await;
1166
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
1167
match outcome {
1168
Ok ( body) => {
1169
if crate :: db:: proposals:: add_review_via( & app.db, p.id, by.id, & by.username, " comment" , & body, & a.name) .await.is_ok ( ) {
1170
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " review.comment via {} " , a.name) , " " ) .await.ok ( ) ;
1171
crate :: web:: proposals:: reviewed( & app, & row, & p, & by, " comment" ) .await;
1172
crate :: hooks:: emit(
1173
& app,
1174
& row,
1175
" review" ,
1176
& by.username,
1177
serde_json:: json! ( { " number" : p.number, " verdict" : " comment" , " via" : a.name } ) ,
1178
)
1179
.await;
1180
}
1181
}
1182
Err ( why) => {
1183
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
33d98147 Andy McCall 16 days ago
1184
app.bus.live ( & crate :: events:: topic:: proposal( row.id, p.number) , " assistant" ) ;
7d2eb590 Andy McCall 1 month ago ↶
1185
}
1186
}
1187
} ) ;
1188
Ok ( ( ) )
1189
}
1190
1191
async fn answer_review (
1192
app : & Shared,
1193
row : & crate:: db:: repos:: RepoRow,
1194
p : & crate:: db:: proposals:: ProposalRow,
1195
a : & Assistant,
1196
) -> Result< String, String> {
1197
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
1198
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
1199
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1200
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
1201
let system = review_system_prompt ( a, & app.config.site.name, & prompts) ;
119fd17d Andy McCall 19 days ago ↶
1202
ask_within ( app, a, & system, budget, BACKGROUND_TIMEOUT , | room | review_context ( row, p, & loaded, room) , | _| { } ) .await
c08cec37 Andy McCall 19 days ago ↶
1203
}
1204
1205
1206
1207
1208
1209
async fn ask_within (
1210
app : & Shared,
1211
a : & Assistant,
1212
system : & str ,
1213
budget : usize ,
1214
wait : Duration,
1215
context : impl Fn( usize ) -> String,
119fd17d Andy McCall 19 days ago ↶
1216
mut part : impl FnMut( &str ) ,
c08cec37 Andy McCall 19 days ago ↶
1217
) -> Result< String, String> {
1218
let room = | b : usize | b.saturating_sub ( system.chars ( ) .count ( ) ) ;
119fd17d Andy McCall 19 days ago ↶
1219
1220
1221
match ask_streaming ( app, a, system, & context ( room ( budget) ) , wait, & mut part) .await {
c08cec37 Andy McCall 19 days ago ↶
1222
Err ( e) if too_long ( & e) => {
1223
let smaller = ( budget / 3 ) .max ( 1_000 ) ;
1224
tracing:: info! ( " assistant {}: too long at {budget} characters ({e}); asking again at {smaller}" , a.name) ;
119fd17d Andy McCall 19 days ago ↶
1225
ask_streaming ( app, a, system, & context ( room ( smaller) ) , wait, & mut part) .await.map_err ( | e2 | {
c08cec37 Andy McCall 19 days ago ↶
1226
format! (
1227
" the proposal was too large for {} 's model at {budget} characters and again at {smaller} ({e2} ): lower this assistant's context budget" ,
1228
a.name
1229
)
1230
} )
1231
}
1232
Err ( e) if e.starts_with ( " no answer within" ) => Err ( format! (
1233
" {e} : a large proposal on a slow or local model can take longer; lower this assistant's context budget, or use a faster model"
1234
) ) ,
1235
other => other,
1236
}
1237
}
1238
1239
1240
fn too_long ( e : & str ) -> bool {
1241
let l = e.to_lowercase ( ) ;
1242
( l.starts_with ( " 400" ) | | l.starts_with ( " 413" ) | | l.starts_with ( " 422" ) )
1243
& & [
1244
" too long" ,
1245
" too large" ,
1246
" context length" ,
1247
" context window" ,
1248
" maximum context" ,
1249
" too many tokens" ,
1250
" exceeds" ,
1251
" reduce the length" ,
1252
" input length" ,
1253
]
1254
.iter ( )
1255
.any ( | k | l.contains ( k) )
ffb67a81 Andy McCall 22 days ago
1256
}
1257
1258
1259
1260
1261
pub fn summary_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
1262
let mut s = format! (
1263
" You are {name} , an AI assistant working on {site} , an OxygenHub site, with a next-generation version control system \
1264
called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \
1265
Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \
1266
under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \
1267
should look at closely; keep to what the description and the diff show." ,
1268
name = a.name
1269
) ;
1270
if ! prompts.trim ( ) .is_empty ( ) {
1271
s.push_str ( " \n \n " ) ;
1272
s.push_str ( prompts.trim ( ) ) ;
1273
}
1274
s
1275
}
1276
1277
1278
1279
pub async fn request_summary (
1280
app : & Shared,
1281
row : & crate:: db:: repos:: RepoRow,
1282
p : & crate:: db:: proposals:: ProposalRow,
1283
by : & crate:: db:: accounts:: Account,
1284
a : Assistant,
1285
) -> Result< ( ) , String> {
1286
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
1287
let pool = & app.db;
1288
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} for a summary" , a.name) , " " ) .await.ok ( ) ;
33d98147 Andy McCall 16 days ago
1289
app.bus.live ( & crate :: events:: topic:: proposal( row.id, p.number) , " assistant" ) ;
ffb67a81 Andy McCall 22 days ago
1290
crate :: db:: accounts:: audit(
1291
pool,
1292
Some ( by) ,
1293
" assistant.asked" ,
1294
& format! ( " {} /{} " , row.owner, row.name) ,
1295
& format! ( " {} #{} summary" , a.name, p.number) ,
1296
)
1297
.await
1298
.ok ( ) ;
1299
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
1300
tokio:: spawn( async move {
1301
let outcome = answer_summary ( & app, & row, & p, & a) .await;
1302
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
1303
match outcome {
1304
Ok ( body) => {
1305
crate :: db:: proposals:: add_comment_via( & app.db, p.id, by.id, & by.username, & body, & a.name) .await.ok ( ) ;
33d98147 Andy McCall 16 days ago
1306
app.bus.live ( & crate :: events:: topic:: proposal( row.id, p.number) , " comment" ) ;
ffb67a81 Andy McCall 22 days ago
1307
}
1308
Err ( why) => {
1309
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
33d98147 Andy McCall 16 days ago
1310
app.bus.live ( & crate :: events:: topic:: proposal( row.id, p.number) , " assistant" ) ;
ffb67a81 Andy McCall 22 days ago
1311
}
1312
}
1313
} ) ;
1314
Ok ( ( ) )
1315
}
1316
1317
async fn answer_summary (
1318
app : & Shared,
1319
row : & crate:: db:: repos:: RepoRow,
1320
p : & crate:: db:: proposals:: ProposalRow,
1321
a : & Assistant,
1322
) -> Result< String, String> {
1323
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
1324
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
1325
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1326
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
1327
let system = summary_system_prompt ( a, & app.config.site.name, & prompts) ;
119fd17d Andy McCall 19 days ago ↶
1328
ask_within ( app, a, & system, budget, BACKGROUND_TIMEOUT , | room | review_context ( row, p, & loaded, room) , | _| { } ) .await
7d2eb590 Andy McCall 1 month ago ↶
1329
}
1330
1331
1332
c43e4f63 Andy McCall 11 days ago ↶
1333
1334
1335
pub fn draft_system_prompt ( a : & Assistant, site : & str , prompts : & str , with_title : bool ) -> String {
1336
let shape = if with_title {
1337
" Begin with one line `Title: ` and a title of at most 72 characters, in the style of the changes' own titles \
1338
(Conventional Commits if they use them), then a blank line, then the description in Markdown: \
1339
no greeting, no sign-off."
1340
} else {
1341
" Write Markdown for the description field only: no title line, no greeting, no sign-off."
1342
} ;
7d2eb590 Andy McCall 1 month ago ↶
1343
let mut s = format! (
1344
" You are {name} , an AI assistant drafting the description of a proposal on {site} , an OxygenHub site, \
1345
for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \
c43e4f63 Andy McCall 11 days ago ↶
1346
{shape} \
7d2eb590 Andy McCall 1 month ago ↶
1347
Say what the changes do and why, what a reviewer should look at, and how to verify them; \
1348
keep to what the diff shows, and keep it short." ,
1349
name = a.name
1350
) ;
1351
if ! prompts.trim ( ) .is_empty ( ) {
1352
s.push_str ( " \n \n " ) ;
1353
s.push_str ( prompts.trim ( ) ) ;
1354
}
1355
s
1356
}
1357
1358
1359
pub struct Draft < 'a > {
1360
1361
1362
pub slot : i64 ,
1363
pub title : & 'a str ,
1364
pub target : & 'a str ,
1365
pub changes : & 'a [crate:: repos:: view:: ChangeView],
1366
pub diffs : & 'a [oxy_core:: repo:: FileDiff],
c43e4f63 Andy McCall 11 days ago ↶
1367
1368
pub with_title : bool ,
1369
}
1370
1371
1372
1373
pub fn split_title ( text : & str ) -> ( Option< String> , String) {
1374
let text = text.trim_start ( ) ;
1375
let ( first, rest) = text.split_once ( ' \n ' ) .unwrap_or ( ( text, " " ) ) ;
1376
let Some ( t) = first.trim ( ) .strip_prefix ( " Title:" ) .or_else ( | | first.trim ( ) .strip_prefix ( " title:" ) ) else {
1377
return ( None , text.trim ( ) .to_string ( ) ) ;
1378
} ;
1379
let t: String = t.trim ( ) .trim_matches ( | c | c = = ' "' | | c = = ' `' | | c = = ' *' ) .trim ( ) .chars ( ) .take ( 120 ) .collect ( ) ;
1380
( ( ! t.is_empty ( ) ) .then_some ( t) , rest.trim ( ) .to_string ( ) )
7d2eb590 Andy McCall 1 month ago ↶
1381
}
1382
1383
1384
1385
pub async fn draft_description (
119fd17d Andy McCall 19 days ago ↶
1386
app : & Shared,
1387
row : & crate:: db:: repos:: RepoRow,
1388
by : & crate:: db:: accounts:: Account,
1389
a : & Assistant,
1390
d : Draft< '_ > ,
1391
) -> Result< String, String> {
1392
draft_description_streaming ( app, row, by, a, d, | _| { } ) .await
1393
}
1394
1395
1396
1397
pub async fn draft_description_streaming (
7d2eb590 Andy McCall 1 month ago ↶
1398
app : & Shared,
1399
row : & crate:: db:: repos:: RepoRow,
1400
by : & crate:: db:: accounts:: Account,
1401
a : & Assistant,
1402
d : Draft< '_ > ,
119fd17d Andy McCall 19 days ago ↶
1403
part : impl FnMut( &str ) ,
7d2eb590 Andy McCall 1 month ago ↶
1404
) -> Result< String, String> {
c43e4f63 Andy McCall 11 days ago ↶
1405
let Draft { slot, title, target, changes, diffs, with_title } = d;
7d2eb590 Andy McCall 1 month ago ↶
1406
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, slot, & a.name) ? ;
1407
crate :: db:: accounts:: audit( & app.db, Some ( by) , " assistant.drafted" , & format! ( " {} /{} " , row.owner, row.name) , & a.name) .await.ok ( ) ;
1408
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1409
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ;
1410
let outcome = match prompts {
1411
Ok ( prompts) => {
c43e4f63 Andy McCall 11 days ago ↶
1412
let system = draft_system_prompt ( a, & app.config.site.name, & prompts, with_title) ;
119fd17d Andy McCall 19 days ago ↶
1413
ask_within ( app, a, & system, budget, CALL_TIMEOUT , | room | draft_context ( row, title, target, changes, diffs, room) , part) .await
7d2eb590 Andy McCall 1 month ago ↶
1414
}
1415
Err ( e) => Err ( e) ,
1416
} ;
1417
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( slot, & a.name) ;
1418
outcome.map ( | t | t.trim ( ) .to_string ( ) )
1419
}