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}