Skip to main content

gaise_api/
lib.rs

1use std::sync::Arc;
2use axum::{
3    extract::State,
4    http::StatusCode,
5    response::{sse::{Event, Sse}, IntoResponse},
6    routing::post,
7    Json, Router,
8};
9use futures_util::{StreamExt};
10use gaise_core::{
11    contracts::{GaiseEmbeddingsRequest, GaiseInstructRequest},
12    GaiseClient,
13};
14use gaise_client::{GaiseClientService};
15use tracing::error;
16
17pub struct AppState {
18    pub client_service: GaiseClientService,
19}
20
21pub fn create_app(state: Arc<AppState>) -> Router {
22    Router::new()
23        .route("/v1/instruct", post(handle_instruct))
24        .route("/v1/instruct/stream", post(handle_instruct_stream))
25        .route("/v1/embeddings", post(handle_embeddings))
26        .with_state(state)
27}
28
29async fn handle_instruct(
30    State(state): State<Arc<AppState>>,
31    Json(request): Json<GaiseInstructRequest>,
32) -> impl IntoResponse {
33    match state.client_service.instruct(&request).await {
34        Ok(response) => Json(response).into_response(),
35        Err(e) => {
36            error!("Instruct error: {}", e);
37            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
38        }
39    }
40}
41
42async fn handle_instruct_stream(
43    State(state): State<Arc<AppState>>,
44    Json(request): Json<GaiseInstructRequest>,
45) -> impl IntoResponse {
46    match state.client_service.instruct_stream(&request).await {
47        Ok(stream) => {
48            let sse_stream = stream.map(|item| {
49                match item {
50                    Ok(chunk) => {
51                        Event::default().json_data(chunk)
52                    }
53                    Err(e) => {
54                        Ok(Event::default().event("error").data(e.to_string()))
55                    }
56                }
57            });
58            Sse::new(sse_stream).into_response()
59        }
60        Err(e) => {
61            error!("Instruct stream error: {}", e);
62            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
63        }
64    }
65}
66
67async fn handle_embeddings(
68    State(state): State<Arc<AppState>>,
69    Json(request): Json<GaiseEmbeddingsRequest>,
70) -> impl IntoResponse {
71    match state.client_service.embeddings(&request).await {
72        Ok(response) => Json(response).into_response(),
73        Err(e) => {
74            error!("Embeddings error: {}", e);
75            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
76        }
77    }
78}