81 lines
2 KiB
Rust
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)
|
|
}
|