Refactoring the structure a bit
This commit is contained in:
Dave
@@ -1,354 +1,17 @@
|
||||
mod commands;
|
||||
mod http;
|
||||
mod io;
|
||||
mod llm;
|
||||
mod state;
|
||||
mod store;
|
||||
|
||||
use crate::commands::{chat, fs};
|
||||
use crate::llm::types::Message;
|
||||
use crate::http::build_routes;
|
||||
use crate::http::context::AppContext;
|
||||
use crate::state::SessionState;
|
||||
use crate::store::JsonFileStore;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use poem::web::websocket::{Message as WsMessage, WebSocket};
|
||||
use poem::{
|
||||
EndpointExt, Response, Route, Server, get, handler,
|
||||
http::{StatusCode, header},
|
||||
listener::TcpListener,
|
||||
web::{Data, Path},
|
||||
};
|
||||
use poem_openapi::{Object, OpenApi, OpenApiService, param::Query, payload::Json};
|
||||
use rust_embed::RustEmbed;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use poem::Server;
|
||||
use poem::listener::TcpListener;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppContext {
|
||||
state: Arc<SessionState>,
|
||||
store: Arc<JsonFileStore>,
|
||||
}
|
||||
|
||||
#[derive(RustEmbed)]
|
||||
#[folder = "../frontend/dist"]
|
||||
struct EmbeddedAssets;
|
||||
|
||||
type OpenApiResult<T> = poem::Result<T>;
|
||||
|
||||
fn bad_request(message: String) -> poem::Error {
|
||||
poem::Error::from_string(message, StatusCode::BAD_REQUEST)
|
||||
}
|
||||
|
||||
#[handler]
|
||||
fn health() -> &'static str {
|
||||
"ok"
|
||||
}
|
||||
|
||||
fn serve_embedded(path: &str) -> Response {
|
||||
let normalized = if path.is_empty() {
|
||||
"index.html"
|
||||
} else {
|
||||
path.trim_start_matches('/')
|
||||
};
|
||||
let is_asset_request = normalized.starts_with("assets/");
|
||||
let asset = if is_asset_request {
|
||||
EmbeddedAssets::get(normalized)
|
||||
} else {
|
||||
EmbeddedAssets::get(normalized).or_else(|| {
|
||||
if normalized == "index.html" {
|
||||
None
|
||||
} else {
|
||||
EmbeddedAssets::get("index.html")
|
||||
}
|
||||
})
|
||||
};
|
||||
|
||||
match asset {
|
||||
Some(content) => {
|
||||
let body = content.data.into_owned();
|
||||
let mime = mime_guess::from_path(normalized)
|
||||
.first_or_octet_stream()
|
||||
.to_string();
|
||||
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(header::CONTENT_TYPE, mime)
|
||||
.body(body)
|
||||
}
|
||||
None => Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.body("Not Found"),
|
||||
}
|
||||
}
|
||||
|
||||
#[handler]
|
||||
fn embedded_asset(Path(path): Path<String>) -> Response {
|
||||
let asset_path = format!("assets/{path}");
|
||||
serve_embedded(&asset_path)
|
||||
}
|
||||
|
||||
#[handler]
|
||||
fn embedded_file(Path(path): Path<String>) -> Response {
|
||||
serve_embedded(&path)
|
||||
}
|
||||
|
||||
#[handler]
|
||||
fn embedded_index() -> Response {
|
||||
serve_embedded("index.html")
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct PathPayload {
|
||||
path: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct ModelPayload {
|
||||
model: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct ApiKeyPayload {
|
||||
api_key: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct FilePathPayload {
|
||||
path: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct WriteFilePayload {
|
||||
path: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct SearchPayload {
|
||||
query: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Object)]
|
||||
struct ExecShellPayload {
|
||||
command: String,
|
||||
args: Vec<String>,
|
||||
}
|
||||
struct Api {
|
||||
ctx: Arc<AppContext>,
|
||||
}
|
||||
|
||||
#[OpenApi]
|
||||
impl Api {
|
||||
#[oai(path = "/project", method = "get")]
|
||||
async fn get_current_project(&self) -> OpenApiResult<Json<Option<String>>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let result =
|
||||
fs::get_current_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
#[oai(path = "/project", method = "post")]
|
||||
async fn open_project(&self, payload: Json<PathPayload>) -> OpenApiResult<Json<String>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let confirmed = fs::open_project(payload.0.path, &ctx.state, ctx.store.as_ref())
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(confirmed))
|
||||
}
|
||||
|
||||
#[oai(path = "/project", method = "delete")]
|
||||
async fn close_project(&self) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
fs::close_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?;
|
||||
Ok(Json(true))
|
||||
}
|
||||
|
||||
#[oai(path = "/model", method = "get")]
|
||||
async fn get_model_preference(&self) -> OpenApiResult<Json<Option<String>>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let result = fs::get_model_preference(ctx.store.as_ref()).map_err(bad_request)?;
|
||||
Ok(Json(result))
|
||||
}
|
||||
|
||||
#[oai(path = "/model", method = "post")]
|
||||
async fn set_model_preference(&self, payload: Json<ModelPayload>) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
fs::set_model_preference(payload.0.model, ctx.store.as_ref()).map_err(bad_request)?;
|
||||
Ok(Json(true))
|
||||
}
|
||||
|
||||
#[oai(path = "/ollama/models", method = "get")]
|
||||
async fn get_ollama_models(
|
||||
&self,
|
||||
base_url: Query<Option<String>>,
|
||||
) -> OpenApiResult<Json<Vec<String>>> {
|
||||
let models = chat::get_ollama_models(base_url.0)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(models))
|
||||
}
|
||||
|
||||
#[oai(path = "/anthropic/key/exists", method = "get")]
|
||||
async fn get_anthropic_api_key_exists(&self) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let exists = chat::get_anthropic_api_key_exists(ctx.store.as_ref()).map_err(bad_request)?;
|
||||
Ok(Json(exists))
|
||||
}
|
||||
|
||||
#[oai(path = "/anthropic/key", method = "post")]
|
||||
async fn set_anthropic_api_key(
|
||||
&self,
|
||||
payload: Json<ApiKeyPayload>,
|
||||
) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
chat::set_anthropic_api_key(ctx.store.as_ref(), payload.0.api_key).map_err(bad_request)?;
|
||||
Ok(Json(true))
|
||||
}
|
||||
|
||||
#[oai(path = "/fs/read", method = "post")]
|
||||
async fn read_file(&self, payload: Json<FilePathPayload>) -> OpenApiResult<Json<String>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let content = fs::read_file(payload.0.path, &ctx.state)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(content))
|
||||
}
|
||||
|
||||
#[oai(path = "/fs/write", method = "post")]
|
||||
async fn write_file(&self, payload: Json<WriteFilePayload>) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
fs::write_file(payload.0.path, payload.0.content, &ctx.state)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(true))
|
||||
}
|
||||
|
||||
#[oai(path = "/fs/list", method = "post")]
|
||||
async fn list_directory(
|
||||
&self,
|
||||
payload: Json<FilePathPayload>,
|
||||
) -> OpenApiResult<Json<Vec<fs::FileEntry>>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let entries = fs::list_directory(payload.0.path, &ctx.state)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(entries))
|
||||
}
|
||||
|
||||
#[oai(path = "/fs/search", method = "post")]
|
||||
async fn search_files(
|
||||
&self,
|
||||
payload: Json<SearchPayload>,
|
||||
) -> OpenApiResult<Json<Vec<crate::commands::search::SearchResult>>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let results = crate::commands::search::search_files(payload.0.query, &ctx.state)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(results))
|
||||
}
|
||||
|
||||
#[oai(path = "/shell/exec", method = "post")]
|
||||
async fn exec_shell(
|
||||
&self,
|
||||
payload: Json<ExecShellPayload>,
|
||||
) -> OpenApiResult<Json<crate::commands::shell::CommandOutput>> {
|
||||
let ctx = self.ctx.clone();
|
||||
let output =
|
||||
crate::commands::shell::exec_shell(payload.0.command, payload.0.args, &ctx.state)
|
||||
.await
|
||||
.map_err(bad_request)?;
|
||||
Ok(Json(output))
|
||||
}
|
||||
|
||||
#[oai(path = "/chat/cancel", method = "post")]
|
||||
async fn cancel_chat(&self) -> OpenApiResult<Json<bool>> {
|
||||
let ctx = self.ctx.clone();
|
||||
chat::cancel_chat(&ctx.state).map_err(bad_request)?;
|
||||
Ok(Json(true))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
enum WsRequest {
|
||||
Chat {
|
||||
messages: Vec<Message>,
|
||||
config: chat::ProviderConfig,
|
||||
},
|
||||
Cancel,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
enum WsResponse {
|
||||
Token { content: String },
|
||||
Update { messages: Vec<Message> },
|
||||
Error { message: String },
|
||||
}
|
||||
|
||||
#[handler]
|
||||
async fn ws_handler(ws: WebSocket, ctx: Data<&AppContext>) -> impl poem::IntoResponse {
|
||||
let ctx = ctx.0.clone();
|
||||
ws.on_upgrade(move |socket| async move {
|
||||
let (mut sink, mut stream) = socket.split();
|
||||
let (tx, mut rx) = mpsc::unbounded_channel::<WsResponse>();
|
||||
|
||||
let forward = tokio::spawn(async move {
|
||||
while let Some(msg) = rx.recv().await {
|
||||
if let Ok(text) = serde_json::to_string(&msg)
|
||||
&& sink.send(WsMessage::Text(text)).await.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
while let Some(Ok(msg)) = stream.next().await {
|
||||
if let WsMessage::Text(text) = msg {
|
||||
let parsed: Result<WsRequest, _> = serde_json::from_str(&text);
|
||||
match parsed {
|
||||
Ok(WsRequest::Chat { messages, config }) => {
|
||||
let tx_updates = tx.clone();
|
||||
let tx_tokens = tx.clone();
|
||||
let ctx_clone = ctx.clone();
|
||||
|
||||
let result = chat::chat(
|
||||
messages,
|
||||
config,
|
||||
&ctx_clone.state,
|
||||
ctx_clone.store.as_ref(),
|
||||
|history| {
|
||||
let _ = tx_updates.send(WsResponse::Update {
|
||||
messages: history.to_vec(),
|
||||
});
|
||||
},
|
||||
|token| {
|
||||
let _ = tx_tokens.send(WsResponse::Token {
|
||||
content: token.to_string(),
|
||||
});
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Err(err) = result {
|
||||
let _ = tx.send(WsResponse::Error { message: err });
|
||||
}
|
||||
}
|
||||
Ok(WsRequest::Cancel) => {
|
||||
let _ = chat::cancel_chat(&ctx.state);
|
||||
}
|
||||
Err(err) => {
|
||||
let _ = tx.send(WsResponse::Error {
|
||||
message: format!("Invalid request: {err}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
drop(tx);
|
||||
let _ = forward.await;
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), std::io::Error> {
|
||||
@@ -361,34 +24,8 @@ async fn main() -> Result<(), std::io::Error> {
|
||||
state: app_state,
|
||||
store,
|
||||
};
|
||||
let ctx_arc = Arc::new(ctx.clone());
|
||||
|
||||
let api_service = OpenApiService::new(
|
||||
Api {
|
||||
ctx: ctx_arc.clone(),
|
||||
},
|
||||
"Living Spec API",
|
||||
"1.0",
|
||||
)
|
||||
.server("http://127.0.0.1:3001/api");
|
||||
let docs_service = OpenApiService::new(
|
||||
Api {
|
||||
ctx: ctx_arc.clone(),
|
||||
},
|
||||
"Living Spec API",
|
||||
"1.0",
|
||||
)
|
||||
.server("http://127.0.0.1:3001/api");
|
||||
|
||||
let app = Route::new()
|
||||
.nest("/api", api_service)
|
||||
.nest("/docs", docs_service.swagger_ui())
|
||||
.at("/ws", get(ws_handler))
|
||||
.at("/health", get(health))
|
||||
.at("/assets/*path", get(embedded_asset))
|
||||
.at("/", get(embedded_index))
|
||||
.at("/*path", get(embedded_file))
|
||||
.data(ctx);
|
||||
let app = build_routes(ctx);
|
||||
|
||||
Server::new(TcpListener::bind("127.0.0.1:3001"))
|
||||
.run(app)
|
||||
|
||||
Reference in New Issue
Block a user