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 => {
ffb67a81 Andy McCall 22 days ago
254
for table in [ " proposal_reviews" , " proposal_comments" ] {
255
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ?" ) ) .bind ( new) .bind ( old) .execute ( pool) .await? ;
256
}
ed89fae5 Andy McCall 1 month ago ↶
257
sqlx:: query( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ?" )
258
.bind ( old)
259
.bind ( new)
260
.bind ( & like)
261
.execute ( pool)
262
.await? ;
263
}
264
Scope:: Org( _ ) => {
265
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 22 days ago
266
for table in [ " proposal_reviews" , " proposal_comments" ] {
267
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ? AND {in_org} " ) )
268
.bind ( new)
269
.bind ( old)
270
.bind ( owner_hint)
271
.execute ( pool)
272
.await? ;
273
}
ed89fae5 Andy McCall 1 month ago ↶
274
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_org} " ) )
275
.bind ( old)
276
.bind ( new)
277
.bind ( & like)
278
.bind ( owner_hint)
279
.execute ( pool)
280
.await? ;
281
}
282
Scope:: Repo( id) => {
283
let in_repo = " proposal_id IN (SELECT id FROM proposals WHERE repo_id = ?)" ;
ffb67a81 Andy McCall 22 days ago
284
for table in [ " proposal_reviews" , " proposal_comments" ] {
285
sqlx:: query( & format! ( " UPDATE {table} SET via = ? WHERE via = ? AND {in_repo} " ) )
286
.bind ( new)
287
.bind ( old)
288
.bind ( id)
289
.execute ( pool)
290
.await? ;
291
}
ed89fae5 Andy McCall 1 month ago ↶
292
sqlx:: query( & format! ( " UPDATE proposal_events SET event = replace(event, ?, ?) WHERE event LIKE ? AND {in_repo} " ) )
293
.bind ( old)
294
.bind ( new)
295
.bind ( & like)
296
.bind ( id)
297
.execute ( pool)
298
.await? ;
299
}
300
}
301
Ok ( ( ) )
302
}
303
e71846aa Andy McCall 1 month ago ↶
304
pub async fn set_enabled ( pool : & Pool, id : i64 , enabled : bool ) -> sqlx:: Result< ( ) > {
305
sqlx:: query( " UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?" )
306
.bind ( enabled)
307
.bind ( now_ms ( ) )
308
.bind ( id)
309
.execute ( pool)
310
.await? ;
311
Ok ( ( ) )
312
}
313
314
pub async fn delete ( pool : & Pool, id : i64 ) -> sqlx:: Result< ( ) > {
315
sqlx:: query( " DELETE FROM assistants WHERE id = ?" ) .bind ( id) .execute ( pool) .await? ;
316
Ok ( ( ) )
317
}
318
319
320
321
pub async fn delete_scope ( pool : & Pool, scope : Scope) -> sqlx:: Result< ( ) > {
322
sqlx:: query( " DELETE FROM assistants WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
323
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" ) .bind ( scope.kind ( ) ) .bind ( scope.id ( ) ) .execute ( pool) .await? ;
324
Ok ( ( ) )
325
}
326
327
328
pub async fn prompt ( pool : & Pool, scope : Scope) -> sqlx:: Result< String> {
329
Ok ( sqlx:: query_scalar( " SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
330
.bind ( scope.kind ( ) )
331
.bind ( scope.id ( ) )
332
.fetch_optional ( pool)
333
.await?
334
.unwrap_or_default ( ) )
335
}
336
337
pub async fn set_prompt ( pool : & Pool, scope : Scope, text : & str ) -> sqlx:: Result< ( ) > {
338
if text.trim ( ) .is_empty ( ) {
339
sqlx:: query( " DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?" )
340
.bind ( scope.kind ( ) )
341
.bind ( scope.id ( ) )
342
.execute ( pool)
343
.await? ;
344
} else {
345
sqlx:: query(
346
" INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \
347
ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms" ,
348
)
349
.bind ( scope.kind ( ) )
350
.bind ( scope.id ( ) )
351
.bind ( text.trim ( ) )
352
.bind ( now_ms ( ) )
353
.execute ( pool)
354
.await? ;
355
}
356
Ok ( ( ) )
357
}
358
359
360
361
362
363
# [ derive ( Clone, Debug ) ]
364
pub struct Available {
365
pub assistant : Assistant,
366
367
pub from : String,
368
369
pub off_above : bool ,
370
}
371
372
impl Available {
373
pub fn on ( & self ) -> bool {
374
self .assistant.enabled & & ! self .off_above
375
}
376
}
377
378
379
380
381
pub async fn available ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< Available> > {
2a83f530 Andy McCall 1 month ago ↶
382
let mut levels = crate :: levels:: above( pool, scope, owner) .await? ;
e71846aa Andy McCall 1 month ago ↶
383
levels.push ( ( scope, String :: new( ) ) ) ;
384
let mut out: Vec< Available> = vec! [ ] ;
385
let mut off: Vec< String> = vec! [ ] ;
386
for ( level, from) in levels {
387
for a in list ( pool, level) .await? {
388
let off_above = off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) ;
389
if ! a.enabled & & ! off.iter ( ) .any ( | n | n.eq_ignore_ascii_case ( & a.name) ) {
390
off.push ( a.name.clone ( ) ) ;
391
}
392
out.retain ( | o | ! o.assistant.name.eq_ignore_ascii_case ( & a.name) ) ;
393
out.push ( Available { assistant: a, from: from.clone ( ) , off_above } ) ;
394
}
395
}
396
Ok ( out)
397
}
398
399
400
pub async fn enabled_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< Vec< Assistant> > {
401
Ok ( available ( pool, Scope:: Repo( repo.id) , & repo.owner) .await? .into_iter ( ) .filter ( Available:: on) .map ( | a | a.assistant ) .collect ( ) )
402
}
403
404
405
406
pub async fn inherited_prompts ( pool : & Pool, scope : Scope, owner : & str ) -> sqlx:: Result< Vec< ( String, String) > > {
407
let mut out = vec! [ ] ;
2a83f530 Andy McCall 1 month ago ↶
408
for ( level, from) in crate :: levels:: above( pool, scope, owner) .await? {
e71846aa Andy McCall 1 month ago ↶
409
let text = prompt ( pool, level) .await? ;
410
if ! text.is_empty ( ) {
411
out.push ( ( from, text) ) ;
412
}
413
}
414
Ok ( out)
415
}
416
417
418
419
pub async fn prompt_for_repo ( pool : & Pool, repo : & crate:: db:: repos:: RepoRow) -> sqlx:: Result< String> {
420
let scope = Scope:: Repo( repo.id) ;
421
let mut parts = inherited_prompts ( pool, scope, & repo.owner) .await? ;
422
let own = prompt ( pool, scope) .await? ;
423
if ! own.is_empty ( ) {
424
parts.push ( ( " this repository" .to_string ( ) , own) ) ;
425
}
426
Ok ( parts.into_iter ( ) .map ( | ( from , text ) | format! ( " ## Instructions from {from} \n \n {text} \n " ) ) .collect:: < Vec< _ > > ( ) .join ( " \n " ) )
427
}
428
429
430
431
432
433
434
pub async fn ask ( app : & Shared, a : & Assistant, system : & str , user : & str ) -> Result< String, String> {
435
let Some ( token) = app.assistant_key.open ( & a.token) else {
436
return Err ( " its token cannot be read: set it again" .into ( ) ) ;
437
} ;
438
let client = reqwest:: Client:: builder( )
439
.timeout ( CALL_TIMEOUT )
440
.user_agent ( format! ( " oxygenhub/{} " , env! ( " CARGO_PKG_VERSION" ) ) )
441
.build ( )
442
.map_err ( | e | e.to_string ( ) ) ? ;
443
let mut req = client.post ( & a.endpoint) .header ( " content-type" , " application/json" ) ;
444
let body = match a.api.as_str ( ) {
445
" anthropic-messages" => {
446
if ! token.is_empty ( ) {
447
req = req.header ( " x-api-key" , & token) ;
448
}
449
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
450
serde_json:: json! ( {
451
" model" : a.model,
452
" max_tokens" : 4096 ,
453
" system" : system,
454
" messages" : [ { " role" : " user" , " content" : user } ] ,
455
} )
456
}
457
_ => {
458
if ! token.is_empty ( ) {
459
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
460
}
461
serde_json:: json! ( {
462
" model" : a.model,
463
" messages" : [
464
{ " role" : " system" , " content" : system } ,
465
{ " role" : " user" , " content" : user } ,
466
] ,
467
} )
468
}
469
} ;
470
for ( name, value) in a.header_list ( ) {
471
req = req.header ( name, value) ;
472
}
473
let r = req.json ( & body) .send ( ) .await.map_err ( | e | {
474
if e.is_timeout ( ) {
475
" no answer within 60 seconds" .to_string ( )
476
} else {
477
scrub ( & e.to_string ( ) , & token)
478
}
479
} ) ? ;
480
let status = r.status ( ) ;
481
let text: String = r.text ( ) .await.unwrap_or_default ( ) .chars ( ) .take ( MAX_ANSWER ) .collect ( ) ;
482
if ! status.is_success ( ) {
483
let why = serde_json:: from_str:: < serde_json:: Value> ( & text)
484
.ok ( )
485
.and_then ( | v | v[ " error" ] [ " message" ] .as_str ( ) .or ( v[ " error" ] .as_str ( ) ) .or ( v[ " message" ] .as_str ( ) ) .map ( str :: to_string) )
486
.unwrap_or_else ( | | text.lines ( ) .next ( ) .unwrap_or ( " " ) .chars ( ) .take ( 160 ) .collect ( ) ) ;
487
return Err ( scrub ( & format! ( " {} {} " , status.as_u16 ( ) , why.trim ( ) ) , & token) ) ;
488
}
489
let v: serde_json:: Value = serde_json:: from_str( & text) .map_err ( | _| " the answer was not JSON" .to_string ( ) ) ? ;
490
let answer = match a.api.as_str ( ) {
491
" anthropic-messages" => v[ " content" ]
492
.as_array ( )
493
.map ( | parts | parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) )
494
.unwrap_or_default ( ) ,
495
_ => match & v[ " choices" ] [ 0 ] [ " message" ] [ " content" ] {
496
serde_json:: Value:: String( s) => s.clone ( ) ,
497
serde_json:: Value:: Array( parts) => parts.iter ( ) .filter_map ( | p | p[ " text" ] .as_str ( ) ) .collect:: < Vec< _ > > ( ) .join ( " " ) ,
498
_ => String :: new( ) ,
499
} ,
500
} ;
501
if answer.trim ( ) .is_empty ( ) {
502
return Err ( " the answer carried no text" .into ( ) ) ;
503
}
504
Ok ( answer)
505
}
506
155cab14 Andy McCall 1 month ago ↶
507
508
509
pub const PRESETS : & [ ( & str , & str , & str , & str , & str ) ] = & [
510
( " anthropic" , " Anthropic (Claude)" , " Claude" , " https://api.anthropic.com/v1/messages" , " anthropic-messages" ) ,
511
( " openai" , " OpenAI (ChatGPT)" , " ChatGPT" , " https://api.openai.com/v1/chat/completions" , " chat-completions" ) ,
512
( " gemini" , " Google Gemini" , " Gemini" , " https://generativelanguage.googleapis.com/v1beta/openai/chat/completions" , " chat-completions" ) ,
513
( " deepseek" , " DeepSeek" , " DeepSeek" , " https://api.deepseek.com/chat/completions" , " chat-completions" ) ,
514
( " mistral" , " Mistral" , " Mistral" , " https://api.mistral.ai/v1/chat/completions" , " chat-completions" ) ,
515
( " perplexity" , " Perplexity" , " Perplexity" , " https://api.perplexity.ai/chat/completions" , " chat-completions" ) ,
516
( " ollama" , " Ollama (on your network)" , " Ollama" , " http://localhost:11434/v1/chat/completions" , " chat-completions" ) ,
517
] ;
518
519
520
521
pub fn models_url ( endpoint : & str ) -> Option< String> {
522
let e = endpoint.trim_end_matches ( ' /' ) ;
523
e.strip_suffix ( " /chat/completions" ) .or_else ( | | e.strip_suffix ( " /messages" ) ) .map ( | base | format! ( " {base} /models" ) )
524
}
525
526
527
528
529
pub async fn list_models ( endpoint : & str , api : & str , token : & str ) -> Result< Vec< String> , String> {
530
let url = models_url ( endpoint) .ok_or ( " this endpoint has no model list next to it" ) ? ;
531
let client = reqwest:: Client:: builder( )
532
.timeout ( std:: time:: Duration:: from_secs( 15 ) )
533
.user_agent ( format! ( " oxygenhub/{} " , env! ( " CARGO_PKG_VERSION" ) ) )
534
.build ( )
535
.map_err ( | e | e.to_string ( ) ) ? ;
536
let mut req = client.get ( & url) ;
537
if api = = " anthropic-messages" {
538
req = req.header ( " anthropic-version" , " 2023-06-01" ) ;
539
if ! token.is_empty ( ) {
540
req = req.header ( " x-api-key" , token) ;
541
}
542
} else if ! token.is_empty ( ) {
543
req = req.header ( " authorization" , format! ( " Bearer {token} " ) ) ;
544
}
545
let r = req.send ( ) .await.map_err ( | e | scrub ( & e.to_string ( ) , token) ) ? ;
546
let status = r.status ( ) ;
547
if ! status.is_success ( ) {
548
return Err ( format! ( " the model list answered {} " , status.as_u16 ( ) ) ) ;
549
}
550
let v: serde_json:: Value = r.json ( ) .await.map_err ( | _| " the model list was not JSON" .to_string ( ) ) ? ;
551
let skip = [ " embed" , " tts" , " whisper" , " dall-e" , " moderation" , " audio" , " image" , " realtime" , " transcribe" , " search-preview" ] ;
552
let mut out: Vec< String> = v[ " data" ]
553
.as_array ( )
554
.or ( v[ " models" ] .as_array ( ) )
555
.ok_or ( " the model list had no data" ) ?
556
.iter ( )
557
.filter_map ( | m | m[ " id" ] .as_str ( ) .or ( m[ " name" ] .as_str ( ) ) )
558
.map ( | id | id.strip_prefix ( " models/" ) .unwrap_or ( id) .to_string ( ) )
559
.filter ( | id | ! skip.iter ( ) .any ( | s | id.to_lowercase ( ) .contains ( s) ) )
560
.collect ( ) ;
561
out.sort ( ) ;
562
out.dedup ( ) ;
563
if out.is_empty ( ) {
564
return Err ( " the model list was empty" .into ( ) ) ;
565
}
566
Ok ( out)
567
}
568
e71846aa Andy McCall 1 month ago ↶
569
570
pub async fn test ( app : & Shared, a : & Assistant) -> Result< String, String> {
571
let answer = ask ( app, a, " You are being connected to an OxygenHub site. Answer briefly." , " Reply with the single word OK." ) .await? ;
572
Ok ( answer.trim ( ) .chars ( ) .take ( 80 ) .collect ( ) )
573
}
574
575
576
fn scrub ( text : & str , token : & str ) -> String {
577
if token.is_empty ( ) {
578
text.to_string ( )
579
} else {
580
text.replace ( token, " ***" )
581
}
582
}
583
584
# [ cfg ( test ) ]
585
mod tests {
586
use super :: * ;
587
588
# [ test ]
589
fn tokens_seal_and_open_only_under_their_key ( ) {
590
let dir = tempfile:: tempdir( ) .unwrap ( ) ;
591
let k = SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) ;
592
let sealed = k.seal ( " sk-very-secret" ) ;
593
assert_ne! ( sealed, b " sk-very-secret" ) ;
594
assert! ( ! String :: from_utf8_lossy( & sealed) .contains ( " secret" ) ) ;
595
assert_eq! ( k.open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) ) ;
596
assert_eq! ( SecretKey:: load_or_create( dir.path ( ) ) .unwrap ( ) .open ( & sealed) .as_deref ( ) , Some ( " sk-very-secret" ) , " the key is kept" ) ;
597
let other = SecretKey:: load_or_create( tempfile:: tempdir( ) .unwrap ( ) .path ( ) ) .unwrap ( ) ;
598
assert_eq! ( other.open ( & sealed) , None ) ;
599
assert_eq! ( k.open ( & [ 1 , 2 , 3 ] ) , None ) ;
600
assert_eq! ( k.seal ( " " ) , Vec : : < u8 > :: new( ) ) ;
601
assert_eq! ( k.open ( & [ ] ) .as_deref ( ) , Some ( " " ) ) ;
602
}
603
604
# [ test ]
605
fn headers_parse_or_name_the_bad_line ( ) {
606
assert_eq! (
607
parse_headers ( " X-Org: acme\n \n OpenAI-Beta : v1 " ) .unwrap ( ) ,
608
vec! [ ( " X-Org" .into ( ) , " acme" .into ( ) ) , ( " OpenAI-Beta" .into ( ) , " v1" .into ( ) ) ]
609
) ;
610
assert_eq! ( parse_headers ( " no colon" ) .unwrap_err ( ) , " no colon" ) ;
611
assert_eq! ( parse_headers ( " bad name: x" ) .unwrap_err ( ) , " bad name: x" ) ;
612
}
613
}
7d2eb590 Andy McCall 1 month ago ↶
614
615
616
617
618
pub const CALLS_PER_HOUR : usize = 30 ;
619
620
621
# [ derive ( Default ) ]
622
pub struct Gate {
623
busy : std:: collections:: HashSet< ( i64 , String) > ,
624
recent : std:: collections:: HashMap< i64 , Vec< i64 > > ,
625
}
626
627
impl Gate {
628
629
fn take ( & mut self , repo_id : i64 , proposal_id : i64 , name : & str ) -> Result< ( ) , String> {
630
let key = ( proposal_id, name.to_lowercase ( ) ) ;
631
if self .busy.contains ( & key) {
632
return Err ( format! ( " {name} is already answering this proposal" ) ) ;
633
}
634
let now = now_ms ( ) ;
635
let calls = self .recent.entry ( repo_id) .or_default ( ) ;
636
calls.retain ( | t | now - * t < 3_600_000 ) ;
637
if calls.len ( ) > = CALLS_PER_HOUR {
638
return Err ( format! ( " this repository has asked assistants {CALLS_PER_HOUR} times in the last hour; try later" ) ) ;
639
}
640
calls.push ( now) ;
641
self .busy.insert ( key) ;
642
Ok ( ( ) )
643
}
644
645
fn release ( & mut self , proposal_id : i64 , name : & str ) {
646
self .busy.remove ( & ( proposal_id, name.to_lowercase ( ) ) ) ;
647
}
648
}
649
650
651
652
pub fn review_context (
653
row : & crate:: db:: repos:: RepoRow,
654
p : & crate:: db:: proposals:: ProposalRow,
655
loaded : & crate:: web:: proposals:: Loaded,
656
budget : usize ,
657
) -> String {
658
let mut head =
659
format! ( " # Proposal #{} : {} \n \n Repository {} /{} · into {} · by {} \n \n " , p.number, p.title, row.owner, row.name, p.target, p.author) ;
660
if ! p.body.trim ( ) .is_empty ( ) {
661
head.push_str ( " ## Description\n \n " ) ;
662
head.push_str ( p.body.trim ( ) ) ;
663
head.push_str ( " \n \n " ) ;
664
}
665
context_text ( head, & loaded.changes, & loaded.diffs, budget)
666
}
667
668
669
pub fn draft_context (
670
row : & crate:: db:: repos:: RepoRow,
671
title : & str ,
672
target : & str ,
673
changes : & [crate:: repos:: view:: ChangeView],
674
diffs : & [oxy_core:: repo:: FileDiff],
675
budget : usize ,
676
) -> String {
677
let head = format! (
678
" # Proposal: {} \n \n Repository {} /{} · into {} \n \n " ,
679
if title.trim ( ) .is_empty ( ) { " (untitled)" } else { title.trim ( ) } ,
680
row.owner,
681
row.name,
682
target
683
) ;
684
context_text ( head, changes, diffs, budget)
685
}
686
687
688
689
fn context_text ( mut text : String, changes : & [crate:: repos:: view:: ChangeView], diffs : & [oxy_core:: repo:: FileDiff], budget : usize ) -> String {
690
text.push_str ( & format! ( " ## Changes ({} )\n \n " , changes.len ( ) ) ) ;
691
for c in changes {
692
text.push_str ( & format! ( " - {} {} \n " , c.short, c.title) ) ;
693
let body = c.body.trim ( ) ;
694
if ! body.is_empty ( ) {
695
for line in body.lines ( ) .take ( 20 ) {
696
text.push_str ( " " ) ;
697
text.push_str ( line) ;
698
text.push ( ' \n ' ) ;
699
}
700
}
701
}
702
text.push_str ( " \n ## Diff\n \n " ) ;
703
let head = text.chars ( ) .count ( ) ;
704
let mut diff = String :: new( ) ;
705
for f in diffs {
706
diff.push_str ( & format! ( " ### {} ({} )\n " , f.path, f.kind) ) ;
707
if let Some ( from) = & f.from {
708
diff.push_str ( & format! ( " renamed from {from} \n " ) ) ;
709
}
710
let binary = f.hunks.iter ( ) .any ( | h | h.added.iter ( ) .chain ( h.removed.iter ( ) .map ( | ( _, l ) | l ) ) .any ( | l | l.contains ( & 0 ) ) ) ;
711
if binary {
712
diff.push_str ( " (binary)\n \n " ) ;
713
continue ;
714
}
715
diff.push_str ( " ```\n " ) ;
716
for h in & f.hunks {
717
diff.push_str ( & format! ( " @@ after line {} @@\n " , h.base_line) ) ;
718
for ( _ , l) in & h.removed {
719
diff.push ( ' -' ) ;
720
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
721
diff.push ( ' \n ' ) ;
722
}
723
for l in & h.added {
724
diff.push ( ' +' ) ;
725
diff.push_str ( String :: from_utf8_lossy( l) .trim_end_matches ( ' \n ' ) ) ;
726
diff.push ( ' \n ' ) ;
727
}
728
}
729
diff.push_str ( " ```\n \n " ) ;
730
}
731
let room = budget.saturating_sub ( head + 80 ) ;
732
if diff.chars ( ) .count ( ) > room {
733
let cut: String = diff.chars ( ) .take ( room) .collect ( ) ;
734
text.push_str ( & cut) ;
735
text.push_str ( " \n \n (the diff was cut here to fit the context budget)\n " ) ;
736
} else {
737
text.push_str ( & diff) ;
738
}
739
text
740
}
741
742
743
744
pub fn review_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
745
let mut s = format! (
746
" You are {name} , an AI assistant reviewing a proposal on {site} , an OxygenHub site. \
747
A person asked for your review; it is recorded as a review comment made via you, under their name. \
748
You cannot approve: only people approve. Answer in Markdown: start with a one-paragraph summary, \
749
then findings (correctness first, then clarity), each naming the file and quoting the lines it is about; \
750
keep praise out unless something is notably well done; say so plainly when you find nothing wrong." ,
751
name = a.name
752
) ;
753
if ! prompts.trim ( ) .is_empty ( ) {
754
s.push_str ( " \n \n " ) ;
755
s.push_str ( prompts.trim ( ) ) ;
756
}
757
s
758
}
759
760
761
762
763
764
pub async fn request_review (
765
app : & Shared,
766
row : & crate:: db:: repos:: RepoRow,
767
p : & crate:: db:: proposals:: ProposalRow,
768
by : & crate:: db:: accounts:: Account,
769
a : Assistant,
770
) -> Result< ( ) , String> {
771
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
772
let pool = & app.db;
773
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} to review" , a.name) , " " ) .await.ok ( ) ;
774
crate :: db:: accounts:: audit(
775
pool,
776
Some ( by) ,
777
" assistant.asked" ,
778
& format! ( " {} /{} " , row.owner, row.name) ,
779
& format! ( " {} #{} " , a.name, p.number) ,
780
)
781
.await
782
.ok ( ) ;
783
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
784
tokio:: spawn( async move {
785
let outcome = answer_review ( & app, & row, & p, & a) .await;
786
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
787
match outcome {
788
Ok ( body) => {
789
if crate :: db:: proposals:: add_review_via( & app.db, p.id, by.id, & by.username, " comment" , & body, & a.name) .await.is_ok ( ) {
790
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " review.comment via {} " , a.name) , " " ) .await.ok ( ) ;
791
crate :: web:: proposals:: reviewed( & app, & row, & p, & by, " comment" ) .await;
792
crate :: hooks:: emit(
793
& app,
794
& row,
795
" review" ,
796
& by.username,
797
serde_json:: json! ( { " number" : p.number, " verdict" : " comment" , " via" : a.name } ) ,
798
)
799
.await;
800
}
801
}
802
Err ( why) => {
803
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
804
}
805
}
806
} ) ;
807
Ok ( ( ) )
808
}
809
810
async fn answer_review (
811
app : & Shared,
812
row : & crate:: db:: repos:: RepoRow,
813
p : & crate:: db:: proposals:: ProposalRow,
814
a : & Assistant,
815
) -> Result< String, String> {
816
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
817
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
818
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
819
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
820
let system = review_system_prompt ( a, & app.config.site.name, & prompts) ;
ffb67a81 Andy McCall 22 days ago
821
let user = review_context ( row, p, & loaded, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
822
ask ( app, a, & system, & user) .await
823
}
824
825
826
827
828
pub fn summary_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
829
let mut s = format! (
830
" You are {name} , an AI assistant working on {site} , an OxygenHub site, with a next-generation version control system \
831
called Oxygen (its command is oxy). A proposal is Oxygen's equivalent of a merge request or pull request. \
832
Please summarise the proposal and comment on it. A person asked for this; it is posted as a comment made via you, \
833
under their name. Answer in Markdown: a short summary of what the proposal changes and why, then anything a reviewer \
834
should look at closely; keep to what the description and the diff show." ,
835
name = a.name
836
) ;
837
if ! prompts.trim ( ) .is_empty ( ) {
838
s.push_str ( " \n \n " ) ;
839
s.push_str ( prompts.trim ( ) ) ;
840
}
841
s
842
}
843
844
845
846
pub async fn request_summary (
847
app : & Shared,
848
row : & crate:: db:: repos:: RepoRow,
849
p : & crate:: db:: proposals:: ProposalRow,
850
by : & crate:: db:: accounts:: Account,
851
a : Assistant,
852
) -> Result< ( ) , String> {
853
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, p.id, & a.name) ? ;
854
let pool = & app.db;
855
crate :: db:: proposals:: add_event( pool, p.id, & by.username, & format! ( " asked {} for a summary" , a.name) , " " ) .await.ok ( ) ;
856
crate :: db:: accounts:: audit(
857
pool,
858
Some ( by) ,
859
" assistant.asked" ,
860
& format! ( " {} /{} " , row.owner, row.name) ,
861
& format! ( " {} #{} summary" , a.name, p.number) ,
862
)
863
.await
864
.ok ( ) ;
865
let ( app, row, p, by) = ( app.clone ( ) , row.clone ( ) , p.clone ( ) , by.clone ( ) ) ;
866
tokio:: spawn( async move {
867
let outcome = answer_summary ( & app, & row, & p, & a) .await;
868
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( p.id, & a.name) ;
869
match outcome {
870
Ok ( body) => {
871
crate :: db:: proposals:: add_comment_via( & app.db, p.id, by.id, & by.username, & body, & a.name) .await.ok ( ) ;
872
}
873
Err ( why) => {
874
crate :: db:: proposals:: add_event( & app.db, p.id, & by.username, & format! ( " {} did not answer" , a.name) , & why) .await.ok ( ) ;
875
}
876
}
877
} ) ;
878
Ok ( ( ) )
879
}
880
881
async fn answer_summary (
882
app : & Shared,
883
row : & crate:: db:: repos:: RepoRow,
884
p : & crate:: db:: proposals:: ProposalRow,
885
a : & Assistant,
886
) -> Result< String, String> {
887
let repo = crate :: repos:: names:: RepoName:: new( & row.owner, & row.name) .ok_or ( " bad repository name" ) ? ;
888
let loaded = crate :: web:: proposals:: load( app, & repo, p.clone ( ) ) .await.map_err ( | e | format! ( " the proposal could not be loaded: {e} " ) ) ? ;
889
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
890
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ? ;
891
let system = summary_system_prompt ( a, & app.config.site.name, & prompts) ;
7d2eb590 Andy McCall 1 month ago ↶
892
let user = review_context ( row, p, & loaded, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
893
ask ( app, a, & system, & user) .await
894
}
895
896
897
898
899
pub fn draft_system_prompt ( a : & Assistant, site : & str , prompts : & str ) -> String {
900
let mut s = format! (
901
" You are {name} , an AI assistant drafting the description of a proposal on {site} , an OxygenHub site, \
902
for the person who will open or edit it; they read and change your draft before saving, and the draft is theirs. \
903
Write Markdown for the description field only: no title line, no greeting, no sign-off. \
904
Say what the changes do and why, what a reviewer should look at, and how to verify them; \
905
keep to what the diff shows, and keep it short." ,
906
name = a.name
907
) ;
908
if ! prompts.trim ( ) .is_empty ( ) {
909
s.push_str ( " \n \n " ) ;
910
s.push_str ( prompts.trim ( ) ) ;
911
}
912
s
913
}
914
915
916
pub struct Draft < 'a > {
917
918
919
pub slot : i64 ,
920
pub title : & 'a str ,
921
pub target : & 'a str ,
922
pub changes : & 'a [crate:: repos:: view:: ChangeView],
923
pub diffs : & 'a [oxy_core:: repo:: FileDiff],
924
}
925
926
927
928
pub async fn draft_description (
929
app : & Shared,
930
row : & crate:: db:: repos:: RepoRow,
931
by : & crate:: db:: accounts:: Account,
932
a : & Assistant,
933
d : Draft< '_ > ,
934
) -> Result< String, String> {
935
let Draft { slot, title, target, changes, diffs } = d;
936
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .take ( row.id, slot, & a.name) ? ;
937
crate :: db:: accounts:: audit( & app.db, Some ( by) , " assistant.drafted" , & format! ( " {} /{} " , row.owner, row.name) , & a.name) .await.ok ( ) ;
938
let budget = usize :: try_from( a.budget) .unwrap_or ( DEFAULT_BUDGET as usize ) .max ( 1_000 ) ;
939
let prompts = prompt_for_repo ( & app.db, row) .await.map_err ( | e | e.to_string ( ) ) ;
940
let outcome = match prompts {
941
Ok ( prompts) => {
942
let system = draft_system_prompt ( a, & app.config.site.name, & prompts) ;
943
let user = draft_context ( row, title, target, changes, diffs, budget.saturating_sub ( system.chars ( ) .count ( ) ) ) ;
944
ask ( app, a, & system, & user) .await
945
}
946
Err ( e) => Err ( e) ,
947
} ;
948
app.assistant_gate.lock ( ) .unwrap_or_else ( | e | e.into_inner ( ) ) .release ( slot, & a.name) ;
949
outcome.map ( | t | t.trim ( ) .to_string ( ) )
950
}