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 22 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 20 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> > {
177
sqlx:: query_as( " SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name" )
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> > {
185
sqlx:: query_as( " SELECT * FROM assistants WHERE id = ?" ) .bind ( id) .fetch_optional ( pool) .await
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) \
193
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *" ,
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(
217
" UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \
218
WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)" ,
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(
235
" UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \
236
WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)" ,
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 23 days ago
265
for table in [ " proposal_reviews" , " proposal_comments" ] {
266
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ?" ) ) .bind ( new) .bind ( old) .execute ( pool) .await? ;
267
}
ed89fae5 Andy McCall 1 month ago ↶
268
sqlx:: query( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ?" )
269
.bind ( old)
270
.bind ( new)
271
.bind ( & like)
272
.execute ( pool)
273
.await? ;
274
}
275
Scope:: Org( _ ) => {
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 = ?)" ;
ffb67a81 Andy McCall 23 days ago
277
for table in [ " proposal_reviews" , " proposal_comments" ] {
278
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ? AND {in_org} " ) )
279
.bind ( new)
280
.bind ( old)
281
.bind ( owner_hint)
282
.execute ( pool)
283
.await? ;
284
}
ed89fae5 Andy McCall 1 month ago ↶
285
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_org} " ) )
286
.bind ( old)
287
.bind ( new)
288
.bind ( & like)
289
.bind ( owner_hint)
290
.execute ( pool)
291
.await? ;
292
}
293
Scope:: Repo( id) => {
294
let in_repo = " proposal_id IN (SELECT id FROM proposals WHERE repo_id = ?)" ;
ffb67a81 Andy McCall 23 days ago
295
for table in [ " proposal_reviews" , " proposal_comments" ] {
296
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ? AND {in_repo} " ) )
297
.bind ( new)
298
.bind ( old)
299
.bind ( id)
300
.execute ( pool)
301
.await? ;
302
}
ed89fae5 Andy McCall 1 month ago ↶
303
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_repo} " ) )
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< ( ) > {
316
sqlx:: query( " UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?" )
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< ( ) > {
326
sqlx:: query( " DELETE FROM assistants WHERE id = ?" ) .bind ( id) .execute ( pool) .await? ;
327
Ok ( ( ) )
328
}
329
330
331
332
pub async fn delete_scope ( pool : & Pool, scope : Scope) -> sqlx:: Result< ( ) > {
333
sqlx:: query( " DELETE FROM assistants WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
334
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
335
Ok ( ( ) )
336
}
337
338
339
pub async fn prompt ( pool : & Pool, scope : Scope) -> sqlx:: Result< String> {
340
Ok ( sqlx:: query_scalar( " SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
341
.bind ( scope.kind ( ) )
342
.bind ( scope.id ( ) )
343
.fetch_optional ( pool)
344
.await?
345
.unwrap_or_default ( ) )
346
}
347
348
pub async fn set_prompt ( pool : & Pool, scope : Scope, text : & str ) -> sqlx:: Result< ( ) > {
349
if text.trim ( ) .is_empty ( ) {
350
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
351
.bind ( scope.kind ( ) )
352
.bind ( scope.id ( ) )
353
.execute ( pool)
354
.await? ;
355
} else {
356
sqlx:: query(
357
" INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
358
ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms" ,
359
)
360
.bind ( scope.kind ( ) )
361
.bind ( scope.id ( ) )
362
.bind ( text.trim ( ) )
363
.bind ( now_ms ( ) )
364
.execute ( pool)
365
.await? ;
366
}
367
Ok ( ( ) )
368
}
369
370
371
372
373
374
# [ derive ( Clone, Debug ) ]
375
pub struct Available {
376
pub assistant : Assistant,
377
378
pub from : String,
379
380
pub off_above : bool ,
381
}
382
383
impl Available {
384
pub fn on ( & self ) -> bool {
385
self .assistant.enabled & & ! self .off_above
386
}
387
}
388
389
390
391
392
pub async fn available ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< Available> > {
2a83f530 Andy McCall 1 month ago ↶
393
let mut levels = crate :: levels:: above( pool, scope, owner) .await? ;
e71846aa Andy McCall 1 month ago ↶
394
levels.push ( ( scope, String :: new( ) ) ) ;
395
let mut out: Vec< Available> = vec! [ ] ;
396
let mut off: Vec< String> = vec! [ ] ;
397
for ( level, from) in levels {
398
for a in list ( pool, level) .await? {
399
let off_above = off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) ;
400
if ! a.enabled & & ! off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) {
401
off.push ( a.name.clone ( ) ) ;
402
}
403
out.retain ( | o | ! o.assistant.name.eq_ignore_ascii_case ( & a.name) ) ;
404
out.push ( Available { assistant: a, from: from.clone ( ) , off_above } ) ;
405
}
406
}
407
Ok ( out)
408
}
409
410
411
pub async fn enabled_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< Vec< Assistant> > {
412
Ok ( available ( pool, Scope:: Repo( repo.id) , & repo.owner) .await? .into_iter ( ) .filter ( Available:: on) .map ( | a | a.assistant ) .collect ( ) )
413
}
414
415
416
417
pub async fn inherited_prompts ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< ( String, String) > > {
418
let mut out = vec! [ ] ;
2a83f530 Andy McCall 1 month ago ↶
419
for ( level, from) in crate :: levels:: above( pool, scope, owner) .await? {
e71846aa Andy McCall 1 month ago ↶
420
let text = prompt ( pool, level) .await? ;
421
if ! text.is_empty ( ) {
422
out.push ( ( from, text) ) ;
423
}
424
}
425
Ok ( out)
426
}
427
428
429
430
pub async fn prompt_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< String> {
431
let scope = Scope:: Repo( repo.id) ;
432
let mut parts = inherited_prompts ( pool, scope, & repo.owner) .await? ;
433
let own = prompt ( pool, scope) .await? ;
434
if ! own.is_empty ( ) {
435
parts.push ( ( " this repository" .to_string ( ) , own) ) ;
436
}
437
Ok ( parts.into_iter ( ) .map ( | ( from , text ) | format! ( " ## Instructions from {from} \n \n {text} \n " ) ) .collect:: < Vec< _ > > ( ) .join ( " \n " ) )
438
}
439
440
441
442
443
5436f0dd Andy McCall 20 days ago ↶
444
445
446
447
448
5360cea3 Andy McCall 22 days ago ↶
449
pub async fn ask ( app : & Shared, a : & Assistant, system : & str , user : & str , wait : Duration) -> Result< String, String> {
5436f0dd Andy McCall 20 days ago ↶
450
ask_streaming ( app, a, system, user, wait, | _| { } ) .await
451
}
452
453
454
pub async fn ask_streaming (
455
app : & Shared,
456
a : & Assistant,
457
system : & str ,
458
user : & str ,
459
wait : Duration,
460
mut part : impl FnMut( &str ) ,
461
) -> Result< String, String> {
e71846aa Andy McCall 1 month ago ↶
462
let Some ( token) = app.assistant_key.open ( & a.token) else {
463
return Err ( " its token cannot be read: set it again" .into ( ) ) ;
464
} ;
5436f0dd Andy McCall 20 days ago ↶
465
let client = crate :: net:: client( & a.endpoint, & app.config.outbound.allow, STREAM_LIMIT .max ( wait) ) .await? ;
e71846aa Andy McCall 1 month ago ↶
466
let mut req = client.post ( & a.endpoint) .header ( " content-type" , " application/json" ) ;
467
let body = match a.api.as_str ( ) {
468
" anthropic-messages" => {
469
if ! token.is_empty ( ) {
470
req = req.header ( " x-api-key" , & token) ;
471
}
472
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
473
serde_json:: json! ( {
474
" model" : a.model,
5360cea3 Andy McCall 22 days ago ↶
475
" max_tokens" : MAX_ANSWER_TOKENS ,
e71846aa Andy McCall 1 month ago ↶
476
" system" : system,
477
" messages" : [ { " role" : " user" , " content" : user } ] ,
5436f0dd Andy McCall 20 days ago ↶
478
" stream" : true ,
e71846aa Andy McCall 1 month ago ↶
479
} )
480
}
481
_ => {
482
if ! token.is_empty ( ) {
483
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
484
}
5360cea3 Andy McCall 22 days ago ↶
485
486
487
e71846aa Andy McCall 1 month ago ↶
488
serde_json:: json! ( {
489
" model" : a.model,
490
" messages" : [
491
{ " role" : " system" , " content" : system } ,
492
{ " role" : " user" , " content" : user } ,
493
] ,
5436f0dd Andy McCall 20 days ago ↶
494
" stream" : true ,
e71846aa Andy McCall 1 month ago ↶
495
} )
496
}
497
} ;
498
for ( name, value) in a.header_list ( ) {
499
req = req.header ( name, value) ;
500
}
5436f0dd Andy McCall 20 days ago ↶
501
let quiet = | | format! ( " no answer within {} seconds" , wait.as_secs ( ) ) ;
502
let r = tokio:: time:: timeout( wait, req.json ( & body) .send ( ) ) .await.map_err ( | _| quiet ( ) ) ? .map_err ( | e | {
e71846aa Andy McCall 1 month ago ↶
503
if e.is_timeout ( ) {
5436f0dd Andy McCall 20 days ago ↶
504
quiet ( )
e71846aa Andy McCall 1 month ago ↶
505
} else {
506
scrub ( & e.to_string ( ) , & token)
507
}
508
} ) ? ;
509
let status = r.status ( ) ;
5436f0dd Andy McCall 20 days ago ↶
510
let streamed =
511
r.headers ( ) .get ( reqwest:: header:: CONTENT_TYPE ) .and_then ( | v | v.to_str ( ) .ok ( ) ) .is_some_and ( | t | t.starts_with ( " text/event-stream" ) ) ;
512
if status.is_success ( ) & & streamed {
513
let ( answer, stop) = read_stream ( & a.api, r, wait, & mut part) .await.map_err ( | e | scrub ( & e, & token) ) ? ;
514
return finish ( answer, stop) ;
515
}
516
let text: String =
517
tokio:: time:: timeout( wait, r.text ( ) ) .await.map_err ( | _| quiet ( ) ) ? .unwrap_or_default ( ) .chars ( ) .take ( MAX_ANSWER ) .collect ( ) ;
e71846aa Andy McCall 1 month ago ↶
518
if ! status.is_success ( ) {
519
let why = serde_json:: from_str:: < serde_json:: Value> ( & text)
520
.ok ( )
521
.and_then ( | v | v[ " error" ] [ " message" ] .as_str ( ) .or ( v[ " error" ] .as_str ( ) ) .or ( v[ " message" ] .as_str ( ) ) .map ( str :: to_string) )
522
.unwrap_or_else ( | | text.lines ( ) .next ( ) .unwrap_or ( " " ) .chars ( ) .take ( 160 ) .collect ( ) ) ;
523
return Err ( scrub ( & format! ( " {} {} " , status.as_u16 ( ) , why.trim ( ) ) , & token) ) ;
524
}
525
let v: serde_json:: Value = serde_json:: from_str( & text) .map_err ( | _| " the answer was not JSON" .to_string ( ) ) ? ;
5360cea3 Andy McCall 22 days ago ↶
526
let ( answer, stop) = answer_of ( & a.api, & v) ;
5436f0dd Andy McCall 20 days ago ↶
527
if ! answer.is_empty ( ) {
528
part ( & answer) ;
529
}
5360cea3 Andy McCall 22 days ago ↶
530
finish ( answer, stop)
5436f0dd Andy McCall 20 days ago ↶
531
}
532
533
534
535
async fn read_stream ( api : & str , mut r : reqwest:: Response, quiet : Duration, part : & mut impl FnMut( &str ) ) -> Result< ( String, Stop) , String> {
536
let mut events = Events:: new( api) ;
537
loop {
538
let chunk = match tokio:: time:: timeout( quiet, r.chunk ( ) ) .await {
539
Err ( _ ) if events.text.is_empty ( ) => return Err ( format! ( " no answer within {} seconds" , quiet.as_secs ( ) ) ) ,
540
541
Err ( _ ) => return Ok ( ( events.text, Stop:: Quiet) ) ,
542
Ok ( Err ( e) ) if events.text.is_empty ( ) => return Err ( e.to_string ( ) ) ,
543
Ok ( Err ( _ ) ) => return Ok ( ( events.text, Stop:: Quiet) ) ,
544
Ok ( Ok ( None ) ) => return Ok ( ( events.text, events.stop.unwrap_or ( Stop:: Done) ) ) ,
545
Ok ( Ok ( Some ( c) ) ) => c,
546
} ;
547
let before = events.text.len ( ) ;
548
events.feed ( & chunk) ? ;
549
if events.text.len ( ) > before {
550
part ( & events.text[ before.. ] ) ;
551
}
552
if events.done | | events.text.len ( ) > = MAX_ANSWER {
553
return Ok ( ( events.text, events.stop.unwrap_or ( Stop:: Done) ) ) ;
554
}
555
}
556
}
557
558
559
560
561
struct Events {
562
anthropic : bool ,
563
buf : Vec< u8 > ,
564
text : String,
565
stop : Option< Stop> ,
566
done : bool ,
567
}
568
569
impl Events {
570
fn new ( api : & str ) -> Events {
571
Events { anthropic: api = = " anthropic-messages" , buf: Vec :: new( ) , text: String :: new( ) , stop: None , done: false }
572
}
573
574
fn feed ( & mut self , bytes : & [u8 ]) -> Result< ( ) , String> {
575
self .buf.extend_from_slice ( bytes) ;
576
while let Some ( end) = self .buf.iter ( ) .position ( | b | * b = = b ' \n ' ) {
577
let line: Vec< u8 > = self .buf.drain ( .. = end) .collect ( ) ;
578
let line = String :: from_utf8_lossy( & line) ;
579
let Some ( data) = line.trim_end_matches ( [ ' \r ' , ' \n ' ] ) .strip_prefix ( " data:" ) else {
580
continue ;
581
} ;
582
let data = data.trim_start ( ) ;
583
if data = = " [DONE]" {
584
self .done = true ;
585
continue ;
586
}
587
let Ok ( v) = serde_json:: from_str:: < serde_json:: Value> ( data) else {
588
continue ;
589
} ;
590
self .event ( & v) ? ;
591
}
592
Ok ( ( ) )
593
}
594
595
fn event ( & mut self , v : & serde_json:: Value) -> Result< ( ) , String> {
596
if self .anthropic {
597
match v[ " type" ] .as_str ( ) {
598
Some ( " content_block_delta" ) if v[ " delta" ] [ " type" ] = = " text_delta" => {
599
self .text.push_str ( v[ " delta" ] [ " text" ] .as_str ( ) .unwrap_or ( " " ) ) ;
600
}
601
Some ( " message_delta" ) => {
602
self .stop = match v[ " delta" ] [ " stop_reason" ] .as_str ( ) {
603
Some ( " max_tokens" ) => Some ( Stop:: Length) ,
604
Some ( " refusal" ) => Some ( Stop:: Declined( v[ " delta" ] [ " stop_details" ] [ " category" ] .as_str ( ) .map ( str :: to_string) ) ) ,
605
Some ( _ ) => Some ( Stop:: Done) ,
606
None => self .stop.take ( ) ,
607
} ;
608
}
609
Some ( " message_stop" ) => self .done = true ,
610
Some ( " error" ) => {
611
return Err ( v[ " error" ] [ " message" ] .as_str ( ) .unwrap_or ( " the provider reported an error" ) .to_string ( ) ) ;
612
}
613
_ => { }
614
}
615
} else {
616
if let Some ( e) = v[ " error" ] [ " message" ] .as_str ( ) {
617
return Err ( e.to_string ( ) ) ;
618
}
619
let choice = & v[ " choices" ] [ 0 ] ;
620
if let Some ( t) = choice[ " delta" ] [ " content" ] .as_str ( ) {
621
self .text.push_str ( t) ;
622
}
623
match choice[ " finish_reason" ] .as_str ( ) {
624
Some ( " length" ) => self .stop = Some ( Stop:: Length) ,
625
Some ( " content_filter" ) => self .stop = Some ( Stop:: Declined( None ) ) ,
626
Some ( _ ) => self .stop = Some ( Stop:: Done) ,
627
None => { }
628
}
629
}
630
Ok ( ( ) )
631
}
5360cea3 Andy McCall 22 days ago ↶
632
}
633
634
635
# [ derive ( Debug, PartialEq, Eq ) ]
636
enum Stop {
637
638
Done,
639
640
Length,
641
642
Declined( Option< String> ) ,
5436f0dd Andy McCall 20 days ago ↶
643
644
Quiet,
5360cea3 Andy McCall 22 days ago ↶
645
}
646
647
648
649
fn answer_of ( api : & str , v : & serde_json:: Value) -> ( String, Stop) {
650
match api {
651
" anthropic-messages" => {
652
let text = v[ " content" ]
653
.as_array ( )
654
.map ( | parts | {
655
parts
656
.iter ( )
657
.filter ( | p | p[ " type" ] .as_str ( ) .is_none_or ( | t | t = = " text" ) )
658
.filter_map ( | p | p[ " text" ] .as_str ( ) )
659
.collect:: < Vec< _ > > ( )
660
.join ( " " )
661
} )
662
.unwrap_or_default ( ) ;
663
let stop = match v[ " stop_reason" ] .as_str ( ) {
664
Some ( " max_tokens" ) => Stop:: Length,
665
Some ( " refusal" ) => Stop:: Declined( v[ " stop_details" ] [ " category" ] .as_str ( ) .map ( str :: to_string) ) ,
666
_ => Stop:: Done,
667
} ;
668
( text, stop)
669
}
670
_ => {
671
let choice = & v[ " choices" ] [ 0 ] ;
672
let text = match & choice[ " message" ] [ " content" ] {
673
serde_json:: Value:: String( s) => s.clone ( ) ,
674
serde_json:: Value:: Array( parts) => parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) ,
675
_ => String :: new( ) ,
676
} ;
677
let stop = match choice[ " finish_reason" ] .as_str ( ) {
678
Some ( " length" ) => Stop:: Length,
679
Some ( " content_filter" ) => Stop:: Declined( None ) ,
680
_ => Stop:: Done,
681
} ;
682
( text, stop)
683
}
e71846aa Andy McCall 1 month ago ↶
684
}
685
}
686
5360cea3 Andy McCall 22 days ago ↶
687
688
689
fn finish ( answer : String, stop : Stop) -> Result< String, String> {
690
match stop {
691
Stop:: Declined( category) => Err ( match category {
692
Some ( c) => format! ( " it declined to answer ({c} )" ) ,
693
None => " it declined to answer" .into ( ) ,
694
} ) ,
695
Stop:: Length if answer.trim ( ) .is_empty ( ) => {
696
Err ( " it ran out of room before writing an answer (it spent its length limit thinking): try a smaller proposal" .into ( ) )
697
}
698
Stop:: Length => Ok ( format! ( " {} \n \n *(Cut short: the answer reached its length limit.)*" , answer.trim_end ( ) ) ) ,
5436f0dd Andy McCall 20 days ago ↶
699
Stop:: Quiet if answer.trim ( ) .is_empty ( ) => Err ( " it went quiet before writing an answer" .into ( ) ) ,
700
Stop:: Quiet => Ok ( format! ( " {} \n \n *(Cut short: the assistant went quiet part of the way through.)*" , answer.trim_end ( ) ) ) ,
5360cea3 Andy McCall 22 days ago ↶
701
Stop:: Done if answer.trim ( ) .is_empty ( ) => Err ( " the answer carried no text" .into ( ) ) ,
702
Stop:: Done => Ok ( answer) ,
703
}
704
}
705
155cab14 Andy McCall 1 month ago ↶
706
707
708
pub const PRESETS : & [ ( & str , & str , & str , & str , & str ) ] = & [
709
( " anthropic" , " Anthropic (Claude)" , " Claude" , " https://api.anthropic.com/v1/messages" , " anthropic-messages" ) ,
710
( " openai" , " OpenAI (ChatGPT)" , " ChatGPT" , " https://api.openai.com/v1/chat/completions" , " chat-completions" ) ,
711
( " gemini" , " Google Gemini" , " Gemini" , " https://generativelanguage.googleapis.com/v1beta/openai/chat/completions" , " chat-completions" ) ,
712
( " deepseek" , " DeepSeek" , " DeepSeek" , " https://api.deepseek.com/chat/completions" , " chat-completions" ) ,
713
( " mistral" , " Mistral" , " Mistral" , " https://api.mistral.ai/v1/chat/completions" , " chat-completions" ) ,
714
( " perplexity" , " Perplexity" , " Perplexity" , " https://api.perplexity.ai/chat/completions" , " chat-completions" ) ,
715
( " ollama" , " Ollama (on your network)" , " Ollama" , " http://localhost:11434/v1/chat/completions" , " chat-completions" ) ,
716
] ;
717
718
719
720
pub fn models_url ( endpoint : & str ) -> Option< String> {
721
let e = endpoint.trim_end_matches ( ' /' ) ;
722
e.strip_suffix ( " /chat/completions" ) .or_else ( | | e.strip_suffix ( " /messages" ) ) .map ( | base | format! ( " {base} /models" ) )
723
}
724
725
726
727
9a777fec Andy McCall 22 days ago ↶
728
pub async fn list_models ( endpoint : & str , api : & str , token : & str , allow : & [String]) -> Result< Vec< String> , String> {
155cab14 Andy McCall 1 month ago ↶
729
let url = models_url ( endpoint) .ok_or ( " this endpoint has no model list next to it" ) ? ;
9a777fec Andy McCall 22 days ago ↶
730
let client = crate :: net:: client( & url, allow, std:: time:: Duration:: from_secs( 15 ) ) .await? ;
155cab14 Andy McCall 1 month ago ↶
731
let mut req = client.get ( & url) ;
732
if api = = " anthropic-messages" {
733
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
734
if ! token.is_empty ( ) {
735
req = req.header ( " x-api-key" , token) ;
736
}
737
} else if ! token.is_empty ( ) {
738
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
739
}
740
let r = req.send ( ) .await.map_err ( | e | scrub ( & e.to_string ( ) , token) ) ? ;
741
let status = r.status ( ) ;
742
if ! status.is_success ( ) {
743
return Err ( format! ( " the model list answered {} " , status.as_u16 ( ) ) ) ;
744
}
745
let v: serde_json:: Value = r.json ( ) .await.map_err ( | _| " the model list was not JSON" .to_string ( ) ) ? ;
746
let skip = [ " embed" , " tts" , " whisper" , " dall-e" , " moderation" , " audio" , " image" , " realtime" , " transcribe" , " search-preview" ] ;
747
let mut out: Vec< String> = v[ " data" ]
748
.as_array ( )
749
.or ( v[ " models" ] .as_array ( ) )
750
.ok_or ( " the model list had no data" ) ?
751
.iter ( )
752
.filter_map ( | m | m[ " id" ] .as_str ( ) .or ( m[ " name" ] .as_str ( ) ) )
753
.map ( | id | id.strip_prefix ( " models/" ) .unwrap_or ( id) .to_string ( ) )
754
.filter ( | id | ! skip.iter ( ) .any ( | s | id.to_lowercase ( ) .contains ( s) ) )
755
.collect ( ) ;
756
out.sort ( ) ;
757
out.dedup ( ) ;
758
if out.is_empty ( ) {
759
return Err ( " the model list was empty" .into ( ) ) ;
760
}
761
Ok ( out)
762
}
763
e71846aa Andy McCall 1 month ago ↶
764
765
pub async fn test ( app : & Shared, a : & Assistant) -> Result< String, String> {
5360cea3 Andy McCall 22 days ago ↶
766
let answer =
767
ask ( app, a, " You are being connected to an OxygenHub site. Answer briefly." , " Reply with the single word OK." , CALL_TIMEOUT )
768
.await? ;
e71846aa Andy McCall 1 month ago ↶
769
Ok ( answer.trim ( ) .chars ( ) .take ( 80 ) .collect ( ) )
770
}
771
772
773
fn scrub ( text : & str , token : & str ) -> String {
774
if token.is_empty ( ) {
775
text.to_string ( )
776
} else {
777
text.replace ( token, " ***" )
778
}
779
}
780
781
# [ cfg ( test ) ]
782
mod tests {
783
use super :: * ;
5436f0dd Andy McCall 20 days ago ↶
784
785
786
787
# [ test ]
788
fn anthropic_events_give_the_text_whole ( ) {
789
let sse = " event: message_start\n data: {\" type\" :\" message_start\" }\n \n \
790
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" thinking_delta\" ,\" thinking\" :\" hmm\" }}\n \n \
791
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" text_delta\" ,\" text\" :\" Adds a \" }}\n \n \
792
data: {\" type\" :\" content_block_delta\" ,\" delta\" :{\" type\" :\" text_delta\" ,\" text\" :\" line — café.\" }}\n \n \
793
data: {\" type\" :\" message_delta\" ,\" delta\" :{\" stop_reason\" :\" end_turn\" }}\n \n \
794
data: {\" type\" :\" message_stop\" }\n \n " ;
795
for size in [ 1 , 3 , 7 , 64 , sse.len ( ) ] {
796
let mut e = Events:: new( " anthropic-messages" ) ;
797
for chunk in sse.as_bytes ( ) .chunks ( size) {
798
e.feed ( chunk) .unwrap ( ) ;
799
}
800
assert_eq! ( ( e.text.as_str ( ) , e.stop.as_ref ( ) , e.done) , ( " Adds a line — café." , Some ( & Stop:: Done) , true ) , " chunks of {size}" ) ;
801
}
802
let mut cut = Events:: new( " anthropic-messages" ) ;
803
cut.feed ( b " data: {\" type\" :\" message_delta\" ,\" delta\" :{\" stop_reason\" :\" max_tokens\" }}\n " ) .unwrap ( ) ;
804
assert_eq! ( cut.stop, Some ( Stop:: Length) ) ;
805
let mut bad = Events:: new( " anthropic-messages" ) ;
806
let e = bad.feed ( b " data: {\" type\" :\" error\" ,\" error\" :{\" type\" :\" overloaded_error\" ,\" message\" :\" Overloaded\" }}\n " ) .unwrap_err ( ) ;
807
assert_eq! ( e, " Overloaded" ) ;
808
}
809
810
# [ test ]
811
fn chat_completion_events_give_the_text ( ) {
812
let sse = " data: {\" choices\" :[{\" delta\" :{\" role\" :\" assistant\" }}]}\n \n \
813
data: {\" choices\" :[{\" delta\" :{\" content\" :\" Hello\" }}]}\n \n \
814
data: {\" choices\" :[{\" delta\" :{\" content\" :\" , world\" },\" finish_reason\" :\" length\" }]}\n \n \
815
data: [DONE]\n \n " ;
816
let mut e = Events:: new( " chat-completions" ) ;
817
for chunk in sse.as_bytes ( ) .chunks ( 5 ) {
818
e.feed ( chunk) .unwrap ( ) ;
819
}
820
assert_eq! ( ( e.text.as_str ( ) , e.stop.as_ref ( ) , e.done) , ( " Hello, world" , Some ( & Stop:: Length) , true ) ) ;
821
}
822
823
824
825
826
# [ tokio ::test ]
827
async fn a_slow_steady_stream_is_read_and_a_silent_one_cut ( ) {
828
async fn serve ( parts : usize , gap_ms : u64 , then_silent : bool ) -> String {
829
let router = axum:: Router:: new( ) .route (
830
" /v1/messages" ,
831
axum:: routing:: post( move | | async move {
832
let ( tx, rx) = tokio:: sync:: mpsc:: channel:: < Result< String, std:: io:: Error> > ( 4 ) ;
833
tokio:: spawn( async move {
834
for i in 0 .. parts {
835
tokio:: time:: sleep( Duration:: from_millis( gap_ms) ) .await;
836
let ev = format! (
837
" data: {{ \" type\" :\" content_block_delta\" ,\" delta\" :{{ \" type\" :\" text_delta\" ,\" text\" :\" {i} \" }} }} \n \n "
838
) ;
839
if tx.send ( Ok ( ev) ) .await.is_err ( ) {
840
return ;
841
}
842
}
843
if then_silent {
844
tokio:: time:: sleep( Duration:: from_secs( 5 ) ) .await;
845
}
846
let _ = tx.send ( Ok ( " data: {\" type\" :\" message_stop\" }\n \n " .into ( ) ) ) .await;
847
} ) ;
848
let body = axum:: body:: Body:: from_stream( tokio_stream:: wrappers:: ReceiverStream:: new( rx) ) ;
849
( [ ( axum:: http:: header:: CONTENT_TYPE , " text/event-stream" ) ] , body)
850
} ) ,
851
) ;
852
let listener = tokio:: net:: TcpListener:: bind( " 127.0.0.1:0" ) .await.unwrap ( ) ;
853
let port = listener.local_addr ( ) .unwrap ( ) .port ( ) ;
854
tokio:: spawn( async move { axum:: serve( listener, router) .await.unwrap ( ) } ) ;
855
format! ( " http://127.0.0.1:{port} /v1/messages" )
856
}
857
let quiet = Duration:: from_millis( 400 ) ;
858
859
let url = serve ( 6 , 200 , false ) .await;
860
let r = reqwest:: Client:: new( ) .post ( & url) .send ( ) .await.unwrap ( ) ;
861
let mut parts = Vec :: new( ) ;
862
let ( text, stop) = read_stream ( " anthropic-messages" , r, quiet, & mut | p: & str | parts.push ( p.to_string ( ) ) ) .await.unwrap ( ) ;
863
assert_eq! ( ( text.as_str ( ) , stop) , ( " 0 1 2 3 4 5 " , Stop:: Done) ) ;
864
assert_eq! ( parts.len ( ) , 6 , " each part handed out as it came" ) ;
865
866
let url = serve ( 2 , 50 , true ) .await;
867
let r = reqwest:: Client:: new( ) .post ( & url) .send ( ) .await.unwrap ( ) ;
868
let ( text, stop) = read_stream ( " anthropic-messages" , r, quiet, & mut | _ : & str | { } ) .await.unwrap ( ) ;
869
assert_eq! ( ( text.as_str ( ) , stop) , ( " 0 1 " , Stop:: Quiet) ) ;
870
assert! ( finish ( text, Stop:: Quiet) .unwrap ( ) .contains ( " went quiet" ) ) ;
871
}
e71846aa Andy McCall 1 month ago ↶
872
873
# [ test ]
874
fn tokens_seal_and_open_only_under_their_key ( ) {
875
let dir = tempfile:: tempdir( ) .unwrap ( ) ;
876
let k = SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) ;
877
let sealed = k.seal ( " sk-very-secret" ) ;
878
assert_ne! ( sealed, b " sk-very-secret" ) ;
879
assert! ( ! String :: from_utf8_lossy( & sealed) .contains ( " secret" ) ) ;
880
assert_eq! ( k.open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) ) ;
881
assert_eq! ( SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) .open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) , " the key is kept" ) ;
882
let other = SecretKey:: load_or_create( tempfile:: tempdir( ) .unwrap ( ) .path ( ) ) .unwrap ( ) ;
883
assert_eq! ( other.open ( & sealed) , None ) ;
884
assert_eq! ( k.open ( & [ 1 , 2 , 3 ] ) , None ) ;
885
assert_eq! ( k.seal ( " " ) , Vec : : < u8 > :: new( ) ) ;
886
assert_eq! ( k.open ( & [ ] ) .as_deref ( ) , Some ( " " ) ) ;
5360cea3 Andy McCall 22 days ago ↶
887
}
888
889
# [ test ]
890
fn answers_say_why_they_stopped ( ) {
891
let a = | v : serde_json:: Value| {
892
let ( text, stop) = answer_of ( " anthropic-messages" , & v) ;
893
finish ( text, stop)
894
} ;
895
896
let thought = serde_json:: json! ( { " stop_reason" : " max_tokens" , " content" : [ { " type" : " thinking" , " thinking" : " " } ] } ) ;
897
assert! ( a ( thought) .unwrap_err ( ) .contains ( " ran out of room" ) ) ;
898
let cut = serde_json:: json! ( { " stop_reason" : " max_tokens" , " content" : [ { " type" : " thinking" , " thinking" : " " } , { " type" : " text" , " text" : " Half" } ] } ) ;
899
assert! ( a ( cut) .unwrap ( ) .starts_with ( " Half\n \n *(Cut short" ) ) ;
900
let declined = serde_json:: json! ( { " stop_reason" : " refusal" , " stop_details" : { " category" : " cyber" } , " content" : [ ] } ) ;
901
assert_eq! ( a ( declined) .unwrap_err ( ) , " it declined to answer (cyber)" ) ;
902
let ok = serde_json:: json! ( { " stop_reason" : " end_turn" , " content" : [ { " type" : " text" , " text" : " Fine." } ] } ) ;
903
assert_eq! ( a ( ok) .unwrap ( ) , " Fine." ) ;
904
let chat = | v : serde_json:: Value| {
905
let ( text, stop) = answer_of ( " chat-completions" , & v) ;
906
finish ( text, stop)
907
} ;
908
assert! ( chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " length" , " message" : { " content" : " " } } ] } ) )
909
.unwrap_err ( )
910
.contains ( " ran out of room" ) ) ;
911
assert_eq! (
912
chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " content_filter" , " message" : { " content" : " " } } ] } ) ) .unwrap_err ( ) ,
913
" it declined to answer"
914
) ;
915
assert_eq! (
916
chat ( serde_json:: json! ( { " choices" : [ { " finish_reason" : " stop" , " message" : { " content" : " Yes." } } ] } ) ) .unwrap ( ) ,
917
" Yes."
918
) ;
e71846aa Andy McCall 1 month ago ↶
919
}
920
921
# [ test ]
922
fn headers_parse_or_name_the_bad_line ( ) {
923
assert_eq! (
924
parse_headers ( " X-Org: acme\n \n OpenAI-Beta : v1 " ) .unwrap ( ) ,
925
vec! [ ( " X-Org" .into ( ) , " acme" .into ( ) ) , ( " OpenAI-Beta" .into ( ) , " v1" .into ( ) ) ]
926
) ;
927
assert_eq! ( parse_headers ( " no colon" ) .unwrap_err ( ) , " no colon" ) ;
928
assert_eq! ( parse_headers ( " bad name: x" ) .unwrap_err ( ) , " bad name: x" ) ;
929
}
930
}
7d2eb590 Andy McCall 1 month ago ↶
931
932
933
934
935
pub const CALLS_PER_HOUR : usize = 30 ;
936
937
938
# [ derive ( Default ) ]
939
pub struct Gate {
940
busy : std:: collections:: HashSet< ( i64 , String) > ,
941
recent : std:: collections:: HashMap< i64 , Vec< i64 > > ,
942
}
943
944
impl Gate {
945
946
fn take ( & mut self , repo_id : i64 , proposal_id : i64 , name : & str ) -> Result< ( ) , String> {
947
let key = ( proposal_id, name.to_lowercase ( ) ) ;
948
if self .busy.contains ( & key) {
949
return Err ( format! ( " {name} is already answering this proposal" ) ) ;
950
}
951
let now = now_ms ( ) ;
952
let calls = self .recent.entry ( repo_id) .or_default ( ) ;
953
calls.retain ( | t | now - * t < 3_600_000 ) ;
954
if calls.len ( ) > = CALLS_PER_HOUR {
955
return Err ( format! ( " this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later" ) ) ;
956
}
957
calls.push ( now) ;
958
self .busy.insert ( key) ;
959
Ok ( ( ) )
960
}
961
962
fn release ( & mut self , proposal_id : i64 , name : & str ) {
963
self .busy.remove ( & ( proposal_id, name.to_lowercase ( ) ) ) ;
964
}
965
}
966
967
968
969
pub fn review_context (
970
row : & crate:: db:: repos:: RepoRow,
971
p : & crate:: db:: proposals:: ProposalRow,
972
loaded : & crate:: web:: proposals:: Loaded,
973
budget : usize ,
974
) -> String {
975
let mut head =
976
format! ( " # Proposal #{} : {} \n \n Repository {} /{} · into {} · by {} \n \n " , p.number, p.title, row.owner, row.name, p.target, p.author) ;
977
if ! p.body.trim ( ) .is_empty ( ) {
978
head.push_str ( " ## Description\n \n " ) ;
979
head.push_str ( p.body.trim ( ) ) ;
980
head.push_str ( " \n \n " ) ;
981
}
982
context_text ( head, & loaded.changes, & loaded.diffs, budget)
983
}
984
985
986
pub fn draft_context (
987
row : & crate:: db:: repos:: RepoRow,
988
title : & str ,
989
target : & str ,
990
changes : & [crate:: repos:: view:: ChangeView],
991
diffs : & [oxy_core:: repo:: FileDiff],
992
budget : usize ,
993
) -> String {
994
let head = format! (
995
" # Proposal: {} \n \n Repository {} /{} · into {} \n \n " ,
996
if title.trim ( ) .is_empty ( ) { " (untitled)" } else { title.trim ( ) } ,
997
row.owner,
998
row.name,
999
target
1000
) ;
1001
context_text ( head, changes, diffs, budget)
1002
}
1003
1004
1005
1006
fn context_text ( mut text : String, changes : & [crate:: repos:: view:: ChangeView], diffs : & [oxy_core:: repo:: FileDiff], budget : usize ) -> String {
1007
text.push_str ( & format! ( " ## Changes ({} )\n \n " , changes.len ( ) ) ) ;
1008
for c in changes {
1009
text.push_str ( & format! ( " - {} {} \n " , c.short, c.title) ) ;
1010
let body = c.body.trim ( ) ;
1011
if ! body.is_empty ( ) {
1012
for line in body.lines ( ) .take ( 20 ) {
1013
text.push_str ( " " ) ;
1014
text.push_str ( line) ;
1015
text.push ( ' \n ' ) ;
1016
}
1017
}
1018
}
1019
text.push_str ( " \n ## Diff\n \n " ) ;
1020
let head = text.chars ( ) .count ( ) ;
1021
let mut diff = String :: new( ) ;
1022
for f in diffs {
1023
diff.push_str ( & format! ( " ### {} ({} )\n " , f.path, f.kind) ) ;
1024
if let Some ( from) = & f.from {
1025
diff.push_str ( & format! ( " renamed from {from} \n " ) ) ;
1026
}
1027
let binary = f.hunks.iter ( ) .any ( | h | h.added.iter ( ) .chain ( h.removed.iter ( ) .map ( | ( _, l ) | l ) ) .any ( | l | l.contains ( & 0 ) ) ) ;
1028
if binary {
1029
diff.push_str ( " (binary)\n \n " ) ;
1030
continue ;
1031
}
1032
diff.push_str ( " ```\n " ) ;
1033
for h in & f.hunks {
1034
diff.push_str ( & format! ( " @@ after line {} @@\n " , h.base_line) ) ;
1035
for ( _ , l) in & h.removed {
1036
diff.push ( ' -' ) ;
1037
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
1038
diff.push ( ' \n ' ) ;
1039
}
1040
for l in & h.added {
1041
diff.push ( ' +' ) ;
1042
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
1043
diff.push ( ' \n ' ) ;
1044
}
1045
}
1046
diff.push_str ( " ```\n \n " ) ;
1047
}
1048
let room = budget.saturating_sub ( head + 80 ) ;
1049
if diff.chars ( ) .count ( ) > room {
1050
let cut: String = diff.chars ( ) .take ( room) .collect ( ) ;
1051
text.push_str ( & cut) ;
1052
text.push_str ( " \n \n (the diff was cut here to fit the context budget)\n " ) ;
1053
} else {
1054
text.push_str ( & diff) ;
1055
}
1056
text
1057
}
1058
1059
1060
1061
pub fn review_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
1062
let mut s = format! (
1063
" You are {name} , an AI assistant reviewing a proposal on {site} , an OxygenHub site. \
1064
A person asked for your review; it is recorded as a review comment made via you, under their name. \
1065
You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \
1066
then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \
1067
keep praise out unless something is notably well done; say so plainly when you find nothing wrong." ,
1068
name = a.name
1069
) ;
1070
if ! prompts.trim ( ) .is_empty ( ) {
1071
s.push_str ( " \n \n " ) ;
1072
s.push_str ( prompts.trim ( ) ) ;
1073
}
1074
s
1075
}
1076
1077
1078
1079
1080
1081
pub async fn request_review (
1082
app : & Shared,
1083
row : & crate:: db:: repos:: RepoRow,
1084
p : & crate:: db:: proposals:: ProposalRow,
1085
by : & crate:: db:: accounts:: Account,
1086
a : Assistant,
1087
) -> Result< ( ) , String> {
1088
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
1089
let pool = & app.db;
1090
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} to review" , a.name) , " " ) .await.ok ( ) ;
1091
crate :: db:: accounts:: audit(
1092
pool,
1093
Some ( by) ,
1094
" assistant.asked" ,
1095
& format! ( " {} /{} " , row.owner, row.name) ,
1096
& format! ( " {} #{} " , a.name, p.number) ,
1097
)
1098
.await
1099
.ok ( ) ;
1100
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
1101
tokio:: spawn( async move {
1102
let outcome = answer_review ( & app, & row, & p, & a) .await;
1103
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
1104
match outcome {
1105
Ok ( body) => {
1106
if crate :: db:: proposals:: add_review_via( & app.db, p.id, by.id, & by.username, " comment" , & body, & a.name) .await.is_ok ( ) {
1107
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " review.comment via {} " , a.name) , " " ) .await.ok ( ) ;
1108
crate :: web:: proposals:: reviewed( & app, & row, & p, & by, " comment" ) .await;
1109
crate :: hooks:: emit(
1110
& app,
1111
& row,
1112
" review" ,
1113
& by.username,
1114
serde_json:: json! ( { " number" : p.number, " verdict" : " comment" , " via" : a.name } ) ,
1115
)
1116
.await;
1117
}
1118
}
1119
Err ( why) => {
1120
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
1121
}
1122
}
1123
} ) ;
1124
Ok ( ( ) )
1125
}
1126
1127
async fn answer_review (
1128
app : & Shared,
1129
row : & crate:: db:: repos:: RepoRow,
1130
p : & crate:: db:: proposals:: ProposalRow,
1131
a : & Assistant,
1132
) -> Result< String, String> {
1133
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
1134
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
1135
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1136
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
1137
let system = review_system_prompt ( a, & app.config.site.name, & prompts) ;
ffb67a81 Andy McCall 23 days ago
1138
let user = review_context ( row, p, & loaded, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
5360cea3 Andy McCall 22 days ago ↶
1139
ask ( app, a, & system, & user, BACKGROUND_TIMEOUT ) .await
ffb67a81 Andy McCall 23 days ago
1140
}
1141
1142
1143
1144
1145
pub fn summary_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
1146
let mut s = format! (
1147
" You are {name} , an AI assistant working on {site} , an OxygenHub site, with a next-generation version control system \
1148
called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \
1149
Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \
1150
under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \
1151
should look at closely; keep to what the description and the diff show." ,
1152
name = a.name
1153
) ;
1154
if ! prompts.trim ( ) .is_empty ( ) {
1155
s.push_str ( " \n \n " ) ;
1156
s.push_str ( prompts.trim ( ) ) ;
1157
}
1158
s
1159
}
1160
1161
1162
1163
pub async fn request_summary (
1164
app : & Shared,
1165
row : & crate:: db:: repos:: RepoRow,
1166
p : & crate:: db:: proposals:: ProposalRow,
1167
by : & crate:: db:: accounts:: Account,
1168
a : Assistant,
1169
) -> Result< ( ) , String> {
1170
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
1171
let pool = & app.db;
1172
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} for a summary" , a.name) , " " ) .await.ok ( ) ;
1173
crate :: db:: accounts:: audit(
1174
pool,
1175
Some ( by) ,
1176
" assistant.asked" ,
1177
& format! ( " {} /{} " , row.owner, row.name) ,
1178
& format! ( " {} #{} summary" , a.name, p.number) ,
1179
)
1180
.await
1181
.ok ( ) ;
1182
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
1183
tokio:: spawn( async move {
1184
let outcome = answer_summary ( & app, & row, & p, & a) .await;
1185
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
1186
match outcome {
1187
Ok ( body) => {
1188
crate :: db:: proposals:: add_comment_via( & app.db, p.id, by.id, & by.username, & body, & a.name) .await.ok ( ) ;
1189
}
1190
Err ( why) => {
1191
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
1192
}
1193
}
1194
} ) ;
1195
Ok ( ( ) )
1196
}
1197
1198
async fn answer_summary (
1199
app : & Shared,
1200
row : & crate:: db:: repos:: RepoRow,
1201
p : & crate:: db:: proposals:: ProposalRow,
1202
a : & Assistant,
1203
) -> Result< String, String> {
1204
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
1205
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
1206
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1207
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
1208
let system = summary_system_prompt ( a, & app.config.site.name, & prompts) ;
7d2eb590 Andy McCall 1 month ago ↶
1209
let user = review_context ( row, p, & loaded, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
5360cea3 Andy McCall 22 days ago ↶
1210
ask ( app, a, & system, & user, BACKGROUND_TIMEOUT ) .await
7d2eb590 Andy McCall 1 month ago ↶
1211
}
1212
1213
1214
1215
1216
pub fn draft_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
1217
let mut s = format! (
1218
" You are {name} , an AI assistant drafting the description of a proposal on {site} , an OxygenHub site, \
1219
for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \
1220
Write Markdown for the description field only: no title line, no greeting, no sign-off. \
1221
Say what the changes do and why, what a reviewer should look at, and how to verify them; \
1222
keep to what the diff shows, and keep it short." ,
1223
name = a.name
1224
) ;
1225
if ! prompts.trim ( ) .is_empty ( ) {
1226
s.push_str ( " \n \n " ) ;
1227
s.push_str ( prompts.trim ( ) ) ;
1228
}
1229
s
1230
}
1231
1232
1233
pub struct Draft < 'a > {
1234
1235
1236
pub slot : i64 ,
1237
pub title : & 'a str ,
1238
pub target : & 'a str ,
1239
pub changes : & 'a [crate:: repos:: view:: ChangeView],
1240
pub diffs : & 'a [oxy_core:: repo:: FileDiff],
1241
}
1242
1243
1244
1245
pub async fn draft_description (
1246
app : & Shared,
1247
row : & crate:: db:: repos:: RepoRow,
1248
by : & crate:: db:: accounts:: Account,
1249
a : & Assistant,
1250
d : Draft< '_ > ,
1251
) -> Result< String, String> {
1252
let Draft { slot, title, target, changes, diffs } = d;
1253
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, slot, & a.name) ? ;
1254
crate :: db:: accounts:: audit( & app.db, Some ( by) , " assistant.drafted" , & format! ( " {} /{} " , row.owner, row.name) , & a.name) .await.ok ( ) ;
1255
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
1256
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ;
1257
let outcome = match prompts {
1258
Ok ( prompts) => {
1259
let system = draft_system_prompt ( a, & app.config.site.name, & prompts) ;
1260
let user = draft_context ( row, title, target, changes, diffs, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
5360cea3 Andy McCall 22 days ago ↶
1261
ask ( app, a, & system, & user, CALL_TIMEOUT ) .await
7d2eb590 Andy McCall 1 month ago ↶
1262
}
1263
Err ( e) => Err ( e) ,
1264
} ;
1265
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( slot, & a.name) ;
1266
outcome.map ( | t | t.trim ( ) .to_string ( ) )
1267
}