little-librarian/src/api.rs

81 lines
2 KiB
Rust

//! HTTP API server for semantic search queries.
use std::{
io,
net::{AddrParseError, SocketAddr},
str::FromStr,
};
use snafu::{ResultExt, Snafu};
use state::AppState;
use tokio::net::TcpListener;
use utoipa::OpenApi;
use utoipa_axum::router::OpenApiRouter;
use utoipa_swagger_ui::SwaggerUi;
use crate::{storage::Postgres, text_encoder::TextEncoder, tokenize::Tokenizer};
pub mod code;
pub mod error;
pub mod query;
pub mod routes;
pub mod state;
const TAG: &str = "little-librarian";
/// OpenAPI documentation configuration.
#[derive(OpenApi)]
#[openapi(
tags(
(name = TAG, description = "Calibre Semantic Search API")
)
)]
struct ApiDoc;
/// Errors that occur when starting the HTTP server.
#[derive(Debug, Snafu)]
pub enum ServeError {
#[snafu(display("Failed to parse address into <ip>:<port>."))]
AddressParse { source: AddrParseError },
#[snafu(display("Failed to bind to {address}."))]
Bind {
source: io::Error,
address: SocketAddr,
},
#[snafu(display("Failed to run http server."))]
Serve { source: io::Error },
}
/// Start the HTTP API server with the given configuration.
pub async fn serve(
address: &str,
db: Postgres,
tokenizer: Tokenizer,
embedder: TextEncoder,
reranker: TextEncoder,
chunk_size: usize,
) -> Result<(), ServeError> {
let state = AppState {
db,
tokenizer,
embedder,
reranker,
chunk_size,
};
let (router, api) = OpenApiRouter::with_openapi(ApiDoc::openapi())
.nest("/api/v1/", routes::router(state))
.split_for_parts();
let router =
router.merge(SwaggerUi::new("/swagger-ui").url("/api-docs/openapi.json", api.clone()));
let address = SocketAddr::from_str(address).context(AddressParseSnafu)?;
let listener = TcpListener::bind(&address)
.await
.context(BindSnafu { address })?;
axum::serve(listener, router.into_make_service())
.await
.context(ServeSnafu)
}