adk_server/rest/controllers/
artifacts.rs

1use crate::ServerConfig;
2use adk_artifact::{ListRequest, LoadRequest};
3use axum::{
4    body::Body,
5    extract::{Path, State},
6    http::{header, StatusCode},
7    response::IntoResponse,
8    Json,
9};
10
11#[derive(Clone)]
12pub struct ArtifactsController {
13    config: ServerConfig,
14}
15
16impl ArtifactsController {
17    pub fn new(config: ServerConfig) -> Self {
18        Self { config }
19    }
20}
21
22pub async fn list_artifacts(
23    State(controller): State<ArtifactsController>,
24    Path((app_name, user_id, session_id)): Path<(String, String, String)>,
25) -> Result<Json<Vec<String>>, StatusCode> {
26    if let Some(service) = &controller.config.artifact_service {
27        let resp = service
28            .list(ListRequest { app_name, user_id, session_id })
29            .await
30            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
31        Ok(Json(resp.file_names))
32    } else {
33        Ok(Json(vec![]))
34    }
35}
36
37pub async fn get_artifact(
38    State(controller): State<ArtifactsController>,
39    Path((app_name, user_id, session_id, artifact_name)): Path<(String, String, String, String)>,
40) -> Result<impl IntoResponse, StatusCode> {
41    if let Some(service) = &controller.config.artifact_service {
42        let resp = service
43            .load(LoadRequest {
44                app_name,
45                user_id,
46                session_id,
47                file_name: artifact_name.clone(),
48                version: None,
49            })
50            .await
51            .map_err(|_| StatusCode::NOT_FOUND)?;
52
53        let mime = mime_guess::from_path(&artifact_name).first_or_octet_stream();
54        let mime_header = header::HeaderValue::from_str(mime.as_ref())
55            .unwrap_or_else(|_| header::HeaderValue::from_static("application/octet-stream"));
56
57        match resp.part {
58            adk_core::Part::InlineData { data, .. } => {
59                Ok(([(header::CONTENT_TYPE, mime_header)], Body::from(data)))
60            }
61            adk_core::Part::Text { text } => {
62                Ok(([(header::CONTENT_TYPE, mime_header)], Body::from(text)))
63            }
64            _ => Err(StatusCode::INTERNAL_SERVER_ERROR),
65        }
66    } else {
67        Err(StatusCode::NOT_FOUND)
68    }
69}
adk_server/rest/controllers/artifacts.rs

adk_server/rest/controllers/
artifacts.rs
