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 ;
25
26
pub const CALL_TIMEOUT : Duration = Duration:: from_secs( 60 ) ;
27
28
const MAX_ANSWER : usize = 64 * 1024 ;
29
2a83f530 Andy McCall 1 month ago ↶
30
pub use crate :: levels:: Scope;
e71846aa Andy McCall 1 month ago ↶
31
32
# [ derive ( Clone, Debug, FromRow ) ]
33
pub struct Assistant {
34
pub id : i64 ,
35
pub scope : String,
36
pub scope_id : i64 ,
37
pub name : String,
38
pub endpoint : String,
39
40
pub api : String,
41
pub model : String,
42
43
pub headers : String,
44
45
pub token : Vec< u8 > ,
46
pub enabled : bool ,
47
pub budget : i64 ,
48
pub created_ms : i64 ,
49
pub updated_ms : i64 ,
50
}
51
52
impl Assistant {
53
pub fn scope ( & self ) -> Scope {
54
match self .scope.as_str ( ) {
55
" org" => Scope:: Org( self .scope_id) ,
56
" repo" => Scope:: Repo( self .scope_id) ,
57
_ => Scope:: Site,
58
}
59
}
60
61
pub fn has_token ( & self ) -> bool {
62
! self .token.is_empty ( )
63
}
64
65
66
pub fn header_list ( & self ) -> Vec< ( String, String) > {
67
parse_headers ( & self .headers) .unwrap_or_default ( )
68
}
69
}
70
71
72
# [ derive ( Clone, Debug, Default ) ]
73
pub struct Fields {
74
pub name : String,
75
pub endpoint : String,
76
pub api : String,
77
pub model : String,
78
pub headers : String,
79
pub enabled : bool ,
80
pub budget : i64 ,
81
}
82
83
84
pub fn parse_headers ( text : & str ) -> Result< Vec< ( String, String) > , String> {
85
let mut out = vec! [ ] ;
86
for line in text.lines ( ) .map ( str :: trim) .filter ( | l | ! l.is_empty ( ) ) {
87
let Some ( ( name, value) ) = line.split_once ( ' :' ) else { return Err ( line.to_string ( ) ) } ;
88
let ( name, value) = ( name.trim ( ) , value.trim ( ) ) ;
89
let name_ok = ! name.is_empty ( ) & & name.chars ( ) .all ( | c | c.is_ascii_alphanumeric ( ) | | c = = ' -' | | c = = ' _' ) ;
90
if ! name_ok | | value.chars ( ) .any ( | c | c.is_control ( ) ) {
91
return Err ( line.to_string ( ) ) ;
92
}
93
out.push ( ( name.to_string ( ) , value.to_string ( ) ) ) ;
94
}
95
Ok ( out)
96
}
97
98
99
100
101
# [ derive ( Clone ) ]
102
pub struct SecretKey ( [u8 ; 32]) ;
103
104
impl SecretKey {
105
pub fn load_or_create ( keys : & Path) -> anyhow:: Result< SecretKey> {
4ac01a73 Andy McCall 1 month ago ↶
106
SecretKey:: load_or_create_named( keys, " assistants.key" )
107
}
108
109
110
111
pub fn load_or_create_named ( keys : & Path, file : & str ) -> anyhow:: Result< SecretKey> {
e71846aa Andy McCall 1 month ago ↶
112
std:: fs:: create_dir_all( keys) ? ;
4ac01a73 Andy McCall 1 month ago ↶
113
let path = keys.join ( file) ;
e71846aa Andy McCall 1 month ago ↶
114
if let Ok ( bytes) = std:: fs:: read( & path) {
115
if bytes.len ( ) = = 32 {
116
return Ok ( SecretKey( bytes.try_into ( ) .unwrap ( ) ) ) ;
117
}
118
}
119
let mut k = [ 0 u8 ; 32 ] ;
120
rand:: rngs:: OsRng.fill_bytes ( & mut k) ;
121
std:: fs:: write( & path, k) ? ;
122
# [ cfg ( unix ) ]
123
{
124
use std:: os:: unix:: fs:: PermissionsExt;
125
std:: fs:: set_permissions( & path, std:: fs:: Permissions:: from_mode( 0o600 ) ) ? ;
126
}
127
Ok ( SecretKey( k) )
128
}
129
130
fn key ( & self ) -> aead:: LessSafeKey {
131
aead:: LessSafeKey:: new( aead:: UnboundKey:: new( & aead:: CHACHA20_POLY1305 , & self .0 ) .expect ( " key length" ) )
132
}
133
134
135
136
pub fn seal ( & self , plain : & str ) -> Vec< u8 > {
137
if plain.is_empty ( ) {
138
return vec! [ ] ;
139
}
140
let mut nonce = [ 0 u8 ; 12 ] ;
141
rand:: rngs:: OsRng.fill_bytes ( & mut nonce) ;
142
let mut out = plain.as_bytes ( ) .to_vec ( ) ;
143
self .key ( ) .seal_in_place_append_tag ( aead:: Nonce:: assume_unique_for_key( nonce) , aead:: Aad:: empty( ) , & mut out) .expect ( " seal" ) ;
144
let mut sealed = nonce.to_vec ( ) ;
145
sealed.extend ( out) ;
146
sealed
147
}
148
149
150
151
pub fn open ( & self , sealed : & [u8 ]) -> Option< String> {
152
if sealed.is_empty ( ) {
153
return Some ( String :: new( ) ) ;
154
}
155
let ( nonce, rest) = sealed.split_at_checked ( 12 ) ? ;
156
let mut buf = rest.to_vec ( ) ;
157
let nonce = aead:: Nonce:: try_assume_unique_for_key( nonce) .ok ( ) ? ;
158
let plain = self .key ( ) .open_in_place ( nonce, aead:: Aad:: empty( ) , & mut buf) .ok ( ) ? ;
159
String :: from_utf8( plain.to_vec ( ) ) .ok ( )
160
}
161
}
162
163
164
165
pub async fn list ( pool : & Pool, scope : Scope) -> sqlx:: Result< Vec< Assistant> > {
166
sqlx:: query_as( " SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name" )
167
.bind ( scope.kind ( ) )
168
.bind ( scope.id ( ) )
169
.fetch_all ( pool)
170
.await
171
}
172
173
pub async fn get ( pool : & Pool, id : i64 ) -> sqlx:: Result< Option< Assistant> > {
174
sqlx:: query_as( " SELECT * FROM assistants WHERE id = ?" ) .bind ( id) .fetch_optional ( pool) .await
175
}
176
177
178
pub async fn create ( pool : & Pool, scope : Scope, f : & Fields, token : & [u8 ]) -> sqlx:: Result< Option< Assistant> > {
179
let now = now_ms ( ) ;
180
let r = sqlx:: query_as(
181
" INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \
182
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *" ,
183
)
184
.bind ( scope.kind ( ) )
185
.bind ( scope.id ( ) )
186
.bind ( & f.name)
187
.bind ( & f.endpoint)
188
.bind ( & f.api)
189
.bind ( & f.model)
190
.bind ( & f.headers)
191
.bind ( token)
192
.bind ( f.enabled)
193
.bind ( f.budget)
194
.bind ( now)
195
.bind ( now)
196
.fetch_optional ( pool)
197
.await? ;
198
Ok ( r)
199
}
200
201
202
203
pub async fn update ( pool : & Pool, id : i64 , f : & Fields, token : Option< & [ u8 ] > ) -> sqlx:: Result< bool > {
204
let r = match token {
205
Some ( t) => sqlx:: query(
206
" UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \
207
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 <> ?)" ,
208
)
209
.bind ( & f.name)
210
.bind ( & f.endpoint)
211
.bind ( & f.api)
212
.bind ( & f.model)
213
.bind ( & f.headers)
214
.bind ( f.enabled)
215
.bind ( f.budget)
216
.bind ( t)
217
.bind ( now_ms ( ) )
218
.bind ( id)
219
.bind ( & f.name)
220
.bind ( id)
221
.execute ( pool)
222
.await? ,
223
None => sqlx:: query(
224
" UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \
225
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 <> ?)" ,
226
)
227
.bind ( & f.name)
228
.bind ( & f.endpoint)
229
.bind ( & f.api)
230
.bind ( & f.model)
231
.bind ( & f.headers)
232
.bind ( f.enabled)
233
.bind ( f.budget)
234
.bind ( now_ms ( ) )
235
.bind ( id)
236
.bind ( & f.name)
237
.bind ( id)
238
.execute ( pool)
239
.await? ,
240
} ;
241
Ok ( r.rows_affected ( ) = = 1 )
242
}
243
ed89fae5 Andy McCall 1 month ago ↶
244
245
246
247
pub async fn rename_in_history ( pool : & Pool, scope : Scope, owner_hint : & str , old : & str , new : & str ) -> sqlx:: Result< ( ) > {
248
if old = = new | | old.is_empty ( ) | | new.is_empty ( ) {
249
return Ok ( ( ) ) ;
250
}
251
let like = format! ( " %{old} %" ) ;
252
match scope {
253
Scope:: Site => {
254
sqlx:: query( " UPDATE proposal_reviews SET via = ? WHERE via = ?" ) .bind ( new) .bind ( old) .execute ( pool) .await? ;
255
sqlx:: query( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ?" )
256
.bind ( old)
257
.bind ( new)
258
.bind ( & like)
259
.execute ( pool)
260
.await? ;
261
}
262
Scope:: Org( _ ) => {
263
let in_org = " proposal_id IN (SELECT p.id FROM proposals p JOIN repositories r ON r.id = p.repo_id WHERE r.owner = ?)" ;
264
sqlx:: query( & format! ( " UPDATE proposal_reviews SET via = ? WHERE via = ? AND {in_org} " ) )
265
.bind ( new)
266
.bind ( old)
267
.bind ( owner_hint)
268
.execute ( pool)
269
.await? ;
270
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_org} " ) )
271
.bind ( old)
272
.bind ( new)
273
.bind ( & like)
274
.bind ( owner_hint)
275
.execute ( pool)
276
.await? ;
277
}
278
Scope:: Repo( id) => {
279
let in_repo = " proposal_id IN (SELECT id FROM proposals WHERE repo_id = ?)" ;
280
sqlx:: query( & format! ( " UPDATE proposal_reviews SET via = ? WHERE via = ? AND {in_repo} " ) )
281
.bind ( new)
282
.bind ( old)
283
.bind ( id)
284
.execute ( pool)
285
.await? ;
286
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_repo} " ) )
287
.bind ( old)
288
.bind ( new)
289
.bind ( & like)
290
.bind ( id)
291
.execute ( pool)
292
.await? ;
293
}
294
}
295
Ok ( ( ) )
296
}
297
e71846aa Andy McCall 1 month ago ↶
298
pub async fn set_enabled ( pool : & Pool, id : i64 , enabled : bool ) -> sqlx:: Result< ( ) > {
299
sqlx:: query( " UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?" )
300
.bind ( enabled)
301
.bind ( now_ms ( ) )
302
.bind ( id)
303
.execute ( pool)
304
.await? ;
305
Ok ( ( ) )
306
}
307
308
pub async fn delete ( pool : & Pool, id : i64 ) -> sqlx:: Result< ( ) > {
309
sqlx:: query( " DELETE FROM assistants WHERE id = ?" ) .bind ( id) .execute ( pool) .await? ;
310
Ok ( ( ) )
311
}
312
313
314
315
pub async fn delete_scope ( pool : & Pool, scope : Scope) -> sqlx:: Result< ( ) > {
316
sqlx:: query( " DELETE FROM assistants WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
317
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
318
Ok ( ( ) )
319
}
320
321
322
pub async fn prompt ( pool : & Pool, scope : Scope) -> sqlx:: Result< String> {
323
Ok ( sqlx:: query_scalar( " SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
324
.bind ( scope.kind ( ) )
325
.bind ( scope.id ( ) )
326
.fetch_optional ( pool)
327
.await?
328
.unwrap_or_default ( ) )
329
}
330
331
pub async fn set_prompt ( pool : & Pool, scope : Scope, text : & str ) -> sqlx:: Result< ( ) > {
332
if text.trim ( ) .is_empty ( ) {
333
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
334
.bind ( scope.kind ( ) )
335
.bind ( scope.id ( ) )
336
.execute ( pool)
337
.await? ;
338
} else {
339
sqlx:: query(
340
" INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
341
ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms" ,
342
)
343
.bind ( scope.kind ( ) )
344
.bind ( scope.id ( ) )
345
.bind ( text.trim ( ) )
346
.bind ( now_ms ( ) )
347
.execute ( pool)
348
.await? ;
349
}
350
Ok ( ( ) )
351
}
352
353
354
355
356
357
# [ derive ( Clone, Debug ) ]
358
pub struct Available {
359
pub assistant : Assistant,
360
361
pub from : String,
362
363
pub off_above : bool ,
364
}
365
366
impl Available {
367
pub fn on ( & self ) -> bool {
368
self .assistant.enabled & & ! self .off_above
369
}
370
}
371
372
373
374
375
pub async fn available ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< Available> > {
2a83f530 Andy McCall 1 month ago ↶
376
let mut levels = crate :: levels:: above( pool, scope, owner) .await? ;
e71846aa Andy McCall 1 month ago ↶
377
levels.push ( ( scope, String :: new( ) ) ) ;
378
let mut out: Vec< Available> = vec! [ ] ;
379
let mut off: Vec< String> = vec! [ ] ;
380
for ( level, from) in levels {
381
for a in list ( pool, level) .await? {
382
let off_above = off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) ;
383
if ! a.enabled & & ! off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) {
384
off.push ( a.name.clone ( ) ) ;
385
}
386
out.retain ( | o | ! o.assistant.name.eq_ignore_ascii_case ( & a.name) ) ;
387
out.push ( Available { assistant: a, from: from.clone ( ) , off_above } ) ;
388
}
389
}
390
Ok ( out)
391
}
392
393
394
pub async fn enabled_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< Vec< Assistant> > {
395
Ok ( available ( pool, Scope:: Repo( repo.id) , & repo.owner) .await? .into_iter ( ) .filter ( Available:: on) .map ( | a | a.assistant ) .collect ( ) )
396
}
397
398
399
400
pub async fn inherited_prompts ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< ( String, String) > > {
401
let mut out = vec! [ ] ;
2a83f530 Andy McCall 1 month ago ↶
402
for ( level, from) in crate :: levels:: above( pool, scope, owner) .await? {
e71846aa Andy McCall 1 month ago ↶
403
let text = prompt ( pool, level) .await? ;
404
if ! text.is_empty ( ) {
405
out.push ( ( from, text) ) ;
406
}
407
}
408
Ok ( out)
409
}
410
411
412
413
pub async fn prompt_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< String> {
414
let scope = Scope:: Repo( repo.id) ;
415
let mut parts = inherited_prompts ( pool, scope, & repo.owner) .await? ;
416
let own = prompt ( pool, scope) .await? ;
417
if ! own.is_empty ( ) {
418
parts.push ( ( " this repository" .to_string ( ) , own) ) ;
419
}
420
Ok ( parts.into_iter ( ) .map ( | ( from , text ) | format! ( " ## Instructions from {from} \n \n {text} \n " ) ) .collect:: < Vec< _ > > ( ) .join ( " \n " ) )
421
}
422
423
424
425
426
427
428
pub async fn ask ( app : & Shared, a : & Assistant, system : & str , user : & str ) -> Result< String, String> {
429
let Some ( token) = app.assistant_key.open ( & a.token) else {
430
return Err ( " its token cannot be read: set it again" .into ( ) ) ;
431
} ;
432
let client = reqwest:: Client:: builder( )
433
.timeout ( CALL_TIMEOUT )
434
.user_agent ( format! ( " oxygenhub/{} " , env! ( " CARGO_PKG_VERSION" ) ) )
435
.build ( )
436
.map_err ( | e | e.to_string ( ) ) ? ;
437
let mut req = client.post ( & a.endpoint) .header ( " content-type" , " application/json" ) ;
438
let body = match a.api.as_str ( ) {
439
" anthropic-messages" => {
440
if ! token.is_empty ( ) {
441
req = req.header ( " x-api-key" , & token) ;
442
}
443
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
444
serde_json:: json! ( {
445
" model" : a.model,
446
" max_tokens" : 4096 ,
447
" system" : system,
448
" messages" : [ { " role" : " user" , " content" : user } ] ,
449
} )
450
}
451
_ => {
452
if ! token.is_empty ( ) {
453
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
454
}
455
serde_json:: json! ( {
456
" model" : a.model,
457
" messages" : [
458
{ " role" : " system" , " content" : system } ,
459
{ " role" : " user" , " content" : user } ,
460
] ,
461
} )
462
}
463
} ;
464
for ( name, value) in a.header_list ( ) {
465
req = req.header ( name, value) ;
466
}
467
let r = req.json ( & body) .send ( ) .await.map_err ( | e | {
468
if e.is_timeout ( ) {
469
" no answer within 60 seconds" .to_string ( )
470
} else {
471
scrub ( & e.to_string ( ) , & token)
472
}
473
} ) ? ;
474
let status = r.status ( ) ;
475
let text: String = r.text ( ) .await.unwrap_or_default ( ) .chars ( ) .take ( MAX_ANSWER ) .collect ( ) ;
476
if ! status.is_success ( ) {
477
let why = serde_json:: from_str:: < serde_json:: Value> ( & text)
478
.ok ( )
479
.and_then ( | v | v[ " error" ] [ " message" ] .as_str ( ) .or ( v[ " error" ] .as_str ( ) ) .or ( v[ " message" ] .as_str ( ) ) .map ( str :: to_string) )
480
.unwrap_or_else ( | | text.lines ( ) .next ( ) .unwrap_or ( " " ) .chars ( ) .take ( 160 ) .collect ( ) ) ;
481
return Err ( scrub ( & format! ( " {} {} " , status.as_u16 ( ) , why.trim ( ) ) , & token) ) ;
482
}
483
let v: serde_json:: Value = serde_json:: from_str( & text) .map_err ( | _| " the answer was not JSON" .to_string ( ) ) ? ;
484
let answer = match a.api.as_str ( ) {
485
" anthropic-messages" => v[ " content" ]
486
.as_array ( )
487
.map ( | parts | parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) )
488
.unwrap_or_default ( ) ,
489
_ => match & v[ " choices" ] [ 0 ] [ " message" ] [ " content" ] {
490
serde_json:: Value:: String( s) => s.clone ( ) ,
491
serde_json:: Value:: Array( parts) => parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) ,
492
_ => String :: new( ) ,
493
} ,
494
} ;
495
if answer.trim ( ) .is_empty ( ) {
496
return Err ( " the answer carried no text" .into ( ) ) ;
497
}
498
Ok ( answer)
499
}
500
155cab14 Andy McCall 1 month ago ↶
501
502
503
pub const PRESETS : & [ ( & str , & str , & str , & str , & str ) ] = & [
504
( " anthropic" , " Anthropic (Claude)" , " Claude" , " https://api.anthropic.com/v1/messages" , " anthropic-messages" ) ,
505
( " openai" , " OpenAI (ChatGPT)" , " ChatGPT" , " https://api.openai.com/v1/chat/completions" , " chat-completions" ) ,
506
( " gemini" , " Google Gemini" , " Gemini" , " https://generativelanguage.googleapis.com/v1beta/openai/chat/completions" , " chat-completions" ) ,
507
( " deepseek" , " DeepSeek" , " DeepSeek" , " https://api.deepseek.com/chat/completions" , " chat-completions" ) ,
508
( " mistral" , " Mistral" , " Mistral" , " https://api.mistral.ai/v1/chat/completions" , " chat-completions" ) ,
509
( " perplexity" , " Perplexity" , " Perplexity" , " https://api.perplexity.ai/chat/completions" , " chat-completions" ) ,
510
( " ollama" , " Ollama (on your network)" , " Ollama" , " http://localhost:11434/v1/chat/completions" , " chat-completions" ) ,
511
] ;
512
513
514
515
pub fn models_url ( endpoint : & str ) -> Option< String> {
516
let e = endpoint.trim_end_matches ( ' /' ) ;
517
e.strip_suffix ( " /chat/completions" ) .or_else ( | | e.strip_suffix ( " /messages" ) ) .map ( | base | format! ( " {base} /models" ) )
518
}
519
520
521
522
523
pub async fn list_models ( endpoint : & str , api : & str , token : & str ) -> Result< Vec< String> , String> {
524
let url = models_url ( endpoint) .ok_or ( " this endpoint has no model list next to it" ) ? ;
525
let client = reqwest:: Client:: builder( )
526
.timeout ( std:: time:: Duration:: from_secs( 15 ) )
527
.user_agent ( format! ( " oxygenhub/{} " , env! ( " CARGO_PKG_VERSION" ) ) )
528
.build ( )
529
.map_err ( | e | e.to_string ( ) ) ? ;
530
let mut req = client.get ( & url) ;
531
if api = = " anthropic-messages" {
532
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
533
if ! token.is_empty ( ) {
534
req = req.header ( " x-api-key" , token) ;
535
}
536
} else if ! token.is_empty ( ) {
537
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
538
}
539
let r = req.send ( ) .await.map_err ( | e | scrub ( & e.to_string ( ) , token) ) ? ;
540
let status = r.status ( ) ;
541
if ! status.is_success ( ) {
542
return Err ( format! ( " the model list answered {} " , status.as_u16 ( ) ) ) ;
543
}
544
let v: serde_json:: Value = r.json ( ) .await.map_err ( | _| " the model list was not JSON" .to_string ( ) ) ? ;
545
let skip = [ " embed" , " tts" , " whisper" , " dall-e" , " moderation" , " audio" , " image" , " realtime" , " transcribe" , " search-preview" ] ;
546
let mut out: Vec< String> = v[ " data" ]
547
.as_array ( )
548
.or ( v[ " models" ] .as_array ( ) )
549
.ok_or ( " the model list had no data" ) ?
550
.iter ( )
551
.filter_map ( | m | m[ " id" ] .as_str ( ) .or ( m[ " name" ] .as_str ( ) ) )
552
.map ( | id | id.strip_prefix ( " models/" ) .unwrap_or ( id) .to_string ( ) )
553
.filter ( | id | ! skip.iter ( ) .any ( | s | id.to_lowercase ( ) .contains ( s) ) )
554
.collect ( ) ;
555
out.sort ( ) ;
556
out.dedup ( ) ;
557
if out.is_empty ( ) {
558
return Err ( " the model list was empty" .into ( ) ) ;
559
}
560
Ok ( out)
561
}
562
e71846aa Andy McCall 1 month ago ↶
563
564
pub async fn test ( app : & Shared, a : & Assistant) -> Result< String, String> {
565
let answer = ask ( app, a, " You are being connected to an OxygenHub site. Answer briefly." , " Reply with the single word OK." ) .await? ;
566
Ok ( answer.trim ( ) .chars ( ) .take ( 80 ) .collect ( ) )
567
}
568
569
570
fn scrub ( text : & str , token : & str ) -> String {
571
if token.is_empty ( ) {
572
text.to_string ( )
573
} else {
574
text.replace ( token, " ***" )
575
}
576
}
577
578
# [ cfg ( test ) ]
579
mod tests {
580
use super :: * ;
581
582
# [ test ]
583
fn tokens_seal_and_open_only_under_their_key ( ) {
584
let dir = tempfile:: tempdir( ) .unwrap ( ) ;
585
let k = SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) ;
586
let sealed = k.seal ( " sk-very-secret" ) ;
587
assert_ne! ( sealed, b " sk-very-secret" ) ;
588
assert! ( ! String :: from_utf8_lossy( & sealed) .contains ( " secret" ) ) ;
589
assert_eq! ( k.open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) ) ;
590
assert_eq! ( SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) .open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) , " the key is kept" ) ;
591
let other = SecretKey:: load_or_create( tempfile:: tempdir( ) .unwrap ( ) .path ( ) ) .unwrap ( ) ;
592
assert_eq! ( other.open ( & sealed) , None ) ;
593
assert_eq! ( k.open ( & [ 1 , 2 , 3 ] ) , None ) ;
594
assert_eq! ( k.seal ( " " ) , Vec : : < u8 > :: new( ) ) ;
595
assert_eq! ( k.open ( & [ ] ) .as_deref ( ) , Some ( " " ) ) ;
596
}
597
598
# [ test ]
599
fn headers_parse_or_name_the_bad_line ( ) {
600
assert_eq! (
601
parse_headers ( " X-Org: acme\n \n OpenAI-Beta : v1 " ) .unwrap ( ) ,
602
vec! [ ( " X-Org" .into ( ) , " acme" .into ( ) ) , ( " OpenAI-Beta" .into ( ) , " v1" .into ( ) ) ]
603
) ;
604
assert_eq! ( parse_headers ( " no colon" ) .unwrap_err ( ) , " no colon" ) ;
605
assert_eq! ( parse_headers ( " bad name: x" ) .unwrap_err ( ) , " bad name: x" ) ;
606
}
607
}
7d2eb590 Andy McCall 1 month ago ↶
608
609
610
611
612
pub const CALLS_PER_HOUR : usize = 30 ;
613
614
615
# [ derive ( Default ) ]
616
pub struct Gate {
617
busy : std:: collections:: HashSet< ( i64 , String) > ,
618
recent : std:: collections:: HashMap< i64 , Vec< i64 > > ,
619
}
620
621
impl Gate {
622
623
fn take ( & mut self , repo_id : i64 , proposal_id : i64 , name : & str ) -> Result< ( ) , String> {
624
let key = ( proposal_id, name.to_lowercase ( ) ) ;
625
if self .busy.contains ( & key) {
626
return Err ( format! ( " {name} is already answering this proposal" ) ) ;
627
}
628
let now = now_ms ( ) ;
629
let calls = self .recent.entry ( repo_id) .or_default ( ) ;
630
calls.retain ( | t | now - * t < 3_600_000 ) ;
631
if calls.len ( ) > = CALLS_PER_HOUR {
632
return Err ( format! ( " this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later" ) ) ;
633
}
634
calls.push ( now) ;
635
self .busy.insert ( key) ;
636
Ok ( ( ) )
637
}
638
639
fn release ( & mut self , proposal_id : i64 , name : & str ) {
640
self .busy.remove ( & ( proposal_id, name.to_lowercase ( ) ) ) ;
641
}
642
}
643
644
645
646
pub fn review_context (
647
row : & crate:: db:: repos:: RepoRow,
648
p : & crate:: db:: proposals:: ProposalRow,
649
loaded : & crate:: web:: proposals:: Loaded,
650
budget : usize ,
651
) -> String {
652
let mut head =
653
format! ( " # Proposal #{} : {} \n \n Repository {} /{} · into {} · by {} \n \n " , p.number, p.title, row.owner, row.name, p.target, p.author) ;
654
if ! p.body.trim ( ) .is_empty ( ) {
655
head.push_str ( " ## Description\n \n " ) ;
656
head.push_str ( p.body.trim ( ) ) ;
657
head.push_str ( " \n \n " ) ;
658
}
659
context_text ( head, & loaded.changes, & loaded.diffs, budget)
660
}
661
662
663
pub fn draft_context (
664
row : & crate:: db:: repos:: RepoRow,
665
title : & str ,
666
target : & str ,
667
changes : & [crate:: repos:: view:: ChangeView],
668
diffs : & [oxy_core:: repo:: FileDiff],
669
budget : usize ,
670
) -> String {
671
let head = format! (
672
" # Proposal: {} \n \n Repository {} /{} · into {} \n \n " ,
673
if title.trim ( ) .is_empty ( ) { " (untitled)" } else { title.trim ( ) } ,
674
row.owner,
675
row.name,
676
target
677
) ;
678
context_text ( head, changes, diffs, budget)
679
}
680
681
682
683
fn context_text ( mut text : String, changes : & [crate:: repos:: view:: ChangeView], diffs : & [oxy_core:: repo:: FileDiff], budget : usize ) -> String {
684
text.push_str ( & format! ( " ## Changes ({} )\n \n " , changes.len ( ) ) ) ;
685
for c in changes {
686
text.push_str ( & format! ( " - {} {} \n " , c.short, c.title) ) ;
687
let body = c.body.trim ( ) ;
688
if ! body.is_empty ( ) {
689
for line in body.lines ( ) .take ( 20 ) {
690
text.push_str ( " " ) ;
691
text.push_str ( line) ;
692
text.push ( ' \n ' ) ;
693
}
694
}
695
}
696
text.push_str ( " \n ## Diff\n \n " ) ;
697
let head = text.chars ( ) .count ( ) ;
698
let mut diff = String :: new( ) ;
699
for f in diffs {
700
diff.push_str ( & format! ( " ### {} ({} )\n " , f.path, f.kind) ) ;
701
if let Some ( from) = & f.from {
702
diff.push_str ( & format! ( " renamed from {from} \n " ) ) ;
703
}
704
let binary = f.hunks.iter ( ) .any ( | h | h.added.iter ( ) .chain ( h.removed.iter ( ) .map ( | ( _, l ) | l ) ) .any ( | l | l.contains ( & 0 ) ) ) ;
705
if binary {
706
diff.push_str ( " (binary)\n \n " ) ;
707
continue ;
708
}
709
diff.push_str ( " ```\n " ) ;
710
for h in & f.hunks {
711
diff.push_str ( & format! ( " @@ after line {} @@\n " , h.base_line) ) ;
712
for ( _ , l) in & h.removed {
713
diff.push ( ' -' ) ;
714
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
715
diff.push ( ' \n ' ) ;
716
}
717
for l in & h.added {
718
diff.push ( ' +' ) ;
719
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
720
diff.push ( ' \n ' ) ;
721
}
722
}
723
diff.push_str ( " ```\n \n " ) ;
724
}
725
let room = budget.saturating_sub ( head + 80 ) ;
726
if diff.chars ( ) .count ( ) > room {
727
let cut: String = diff.chars ( ) .take ( room) .collect ( ) ;
728
text.push_str ( & cut) ;
729
text.push_str ( " \n \n (the diff was cut here to fit the context budget)\n " ) ;
730
} else {
731
text.push_str ( & diff) ;
732
}
733
text
734
}
735
736
737
738
pub fn review_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
739
let mut s = format! (
740
" You are {name} , an AI assistant reviewing a proposal on {site} , an OxygenHub site. \
741
A person asked for your review; it is recorded as a review comment made via you, under their name. \
742
You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \
743
then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \
744
keep praise out unless something is notably well done; say so plainly when you find nothing wrong." ,
745
name = a.name
746
) ;
747
if ! prompts.trim ( ) .is_empty ( ) {
748
s.push_str ( " \n \n " ) ;
749
s.push_str ( prompts.trim ( ) ) ;
750
}
751
s
752
}
753
754
755
756
757
758
pub async fn request_review (
759
app : & Shared,
760
row : & crate:: db:: repos:: RepoRow,
761
p : & crate:: db:: proposals:: ProposalRow,
762
by : & crate:: db:: accounts:: Account,
763
a : Assistant,
764
) -> Result< ( ) , String> {
765
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
766
let pool = & app.db;
767
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} to review" , a.name) , " " ) .await.ok ( ) ;
768
crate :: db:: accounts:: audit(
769
pool,
770
Some ( by) ,
771
" assistant.asked" ,
772
& format! ( " {} /{} " , row.owner, row.name) ,
773
& format! ( " {} #{} " , a.name, p.number) ,
774
)
775
.await
776
.ok ( ) ;
777
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
778
tokio:: spawn( async move {
779
let outcome = answer_review ( & app, & row, & p, & a) .await;
780
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
781
match outcome {
782
Ok ( body) => {
783
if crate :: db:: proposals:: add_review_via( & app.db, p.id, by.id, & by.username, " comment" , & body, & a.name) .await.is_ok ( ) {
784
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " review.comment via {} " , a.name) , " " ) .await.ok ( ) ;
785
crate :: web:: proposals:: reviewed( & app, & row, & p, & by, " comment" ) .await;
786
crate :: hooks:: emit(
787
& app,
788
& row,
789
" review" ,
790
& by.username,
791
serde_json:: json! ( { " number" : p.number, " verdict" : " comment" , " via" : a.name } ) ,
792
)
793
.await;
794
}
795
}
796
Err ( why) => {
797
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
798
}
799
}
800
} ) ;
801
Ok ( ( ) )
802
}
803
804
async fn answer_review (
805
app : & Shared,
806
row : & crate:: db:: repos:: RepoRow,
807
p : & crate:: db:: proposals:: ProposalRow,
808
a : & Assistant,
809
) -> Result< String, String> {
810
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
811
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
812
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
813
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
814
let system = review_system_prompt ( a, & app.config.site.name, & prompts) ;
815
let user = review_context ( row, p, & loaded, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
816
ask ( app, a, & system, & user) .await
817
}
818
819
820
821
822
pub fn draft_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
823
let mut s = format! (
824
" You are {name} , an AI assistant drafting the description of a proposal on {site} , an OxygenHub site, \
825
for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \
826
Write Markdown for the description field only: no title line, no greeting, no sign-off. \
827
Say what the changes do and why, what a reviewer should look at, and how to verify them; \
828
keep to what the diff shows, and keep it short." ,
829
name = a.name
830
) ;
831
if ! prompts.trim ( ) .is_empty ( ) {
832
s.push_str ( " \n \n " ) ;
833
s.push_str ( prompts.trim ( ) ) ;
834
}
835
s
836
}
837
838
839
pub struct Draft < 'a > {
840
841
842
pub slot : i64 ,
843
pub title : & 'a str ,
844
pub target : & 'a str ,
845
pub changes : & 'a [crate:: repos:: view:: ChangeView],
846
pub diffs : & 'a [oxy_core:: repo:: FileDiff],
847
}
848
849
850
851
pub async fn draft_description (
852
app : & Shared,
853
row : & crate:: db:: repos:: RepoRow,
854
by : & crate:: db:: accounts:: Account,
855
a : & Assistant,
856
d : Draft< '_ > ,
857
) -> Result< String, String> {
858
let Draft { slot, title, target, changes, diffs } = d;
859
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, slot, & a.name) ? ;
860
crate :: db:: accounts:: audit( & app.db, Some ( by) , " assistant.drafted" , & format! ( " {} /{} " , row.owner, row.name) , & a.name) .await.ok ( ) ;
861
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
862
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ;
863
let outcome = match prompts {
864
Ok ( prompts) => {
865
let system = draft_system_prompt ( a, & app.config.site.name, & prompts) ;
866
let user = draft_context ( row, title, target, changes, diffs, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
867
ask ( app, a, & system, & user) .await
868
}
869
Err ( e) => Err ( e) ,
870
} ;
871
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( slot, & a.name) ;
872
outcome.map ( | t | t.trim ( ) .to_string ( ) )
873
}