|
1 | | -use std::net::Ipv4Addr; |
2 | | -use std::path::Path; |
| 1 | +#[cfg(target_arch = "aarch64")] |
| 2 | +mod inner { |
| 3 | + use std::net::Ipv4Addr; |
| 4 | + use std::path::Path; |
| 5 | + |
| 6 | + use axum::http::StatusCode; |
| 7 | + use axum::{Router, error_handling::HandleError}; |
| 8 | + use hypr_llm_cactus::{CompleteService, ModelManagerBuilder}; |
| 9 | + use tokio::net::TcpListener; |
| 10 | + use tower_http::cors::CorsLayer; |
| 11 | + |
| 12 | + use crate::{Error, ModelSelection}; |
| 13 | + |
| 14 | + pub struct LlmServer { |
| 15 | + base_url: String, |
| 16 | + shutdown_tx: tokio::sync::watch::Sender<()>, |
| 17 | + task: tokio::task::JoinHandle<()>, |
| 18 | + } |
3 | 19 |
|
4 | | -use axum::http::StatusCode; |
5 | | -use axum::{Router, error_handling::HandleError}; |
6 | | -use hypr_llm_cactus::{CompleteService, ModelManagerBuilder}; |
7 | | -use tokio::net::TcpListener; |
8 | | -use tower_http::cors::CorsLayer; |
| 20 | + impl LlmServer { |
| 21 | + pub async fn start(selection: &ModelSelection, models_dir: &Path) -> Result<Self, Error> { |
| 22 | + let file_path = selection.file_path(models_dir); |
| 23 | + let name = selection.display_name(); |
9 | 24 |
|
10 | | -use crate::{Error, ModelSelection}; |
| 25 | + if !file_path.exists() { |
| 26 | + return Err(Error::ModelNotDownloaded); |
| 27 | + } |
11 | 28 |
|
12 | | -pub struct LlmServer { |
13 | | - base_url: String, |
14 | | - shutdown_tx: tokio::sync::watch::Sender<()>, |
15 | | - task: tokio::task::JoinHandle<()>, |
16 | | -} |
| 29 | + let manager = ModelManagerBuilder::default() |
| 30 | + .register(name.clone(), file_path) |
| 31 | + .default_model(name) |
| 32 | + .build(); |
17 | 33 |
|
18 | | -impl LlmServer { |
19 | | - pub async fn start(selection: &ModelSelection, models_dir: &Path) -> Result<Self, Error> { |
20 | | - let file_path = selection.file_path(models_dir); |
21 | | - let name = selection.display_name(); |
| 34 | + let service = HandleError::new(CompleteService::new(manager), handle_error); |
22 | 35 |
|
23 | | - if !file_path.exists() { |
24 | | - return Err(Error::ModelNotDownloaded); |
25 | | - } |
| 36 | + let router = Router::new() |
| 37 | + .route_service("/v1/chat/completions", service) |
| 38 | + .layer(CorsLayer::permissive()); |
26 | 39 |
|
27 | | - let manager = ModelManagerBuilder::default() |
28 | | - .register(name.clone(), file_path) |
29 | | - .default_model(name) |
30 | | - .build(); |
| 40 | + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0u16)).await?; |
| 41 | + let addr = listener.local_addr()?; |
| 42 | + let base_url = format!("http://{}/v1", addr); |
31 | 43 |
|
32 | | - let service = HandleError::new(CompleteService::new(manager), handle_error); |
| 44 | + let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(()); |
33 | 45 |
|
34 | | - let router = Router::new() |
35 | | - .route_service("/v1/chat/completions", service) |
36 | | - .layer(CorsLayer::permissive()); |
| 46 | + let task = tokio::spawn(async move { |
| 47 | + axum::serve(listener, router) |
| 48 | + .with_graceful_shutdown(async move { |
| 49 | + let _ = shutdown_rx.changed().await; |
| 50 | + }) |
| 51 | + .await |
| 52 | + .ok(); |
| 53 | + }); |
37 | 54 |
|
38 | | - let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0u16)).await?; |
39 | | - let addr = listener.local_addr()?; |
40 | | - let base_url = format!("http://{}/v1", addr); |
| 55 | + tracing::info!(url = %base_url, "local LLM server started"); |
41 | 56 |
|
42 | | - let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(()); |
| 57 | + Ok(Self { |
| 58 | + base_url, |
| 59 | + shutdown_tx, |
| 60 | + task, |
| 61 | + }) |
| 62 | + } |
43 | 63 |
|
44 | | - let task = tokio::spawn(async move { |
45 | | - axum::serve(listener, router) |
46 | | - .with_graceful_shutdown(async move { |
47 | | - let _ = shutdown_rx.changed().await; |
48 | | - }) |
49 | | - .await |
50 | | - .ok(); |
51 | | - }); |
| 64 | + pub fn url(&self) -> &str { |
| 65 | + &self.base_url |
| 66 | + } |
52 | 67 |
|
53 | | - tracing::info!(url = %base_url, "local LLM server started"); |
| 68 | + pub async fn stop(self) { |
| 69 | + let _ = self.shutdown_tx.send(()); |
| 70 | + let _ = self.task.await; |
| 71 | + tracing::info!("local LLM server stopped"); |
| 72 | + } |
| 73 | + } |
54 | 74 |
|
55 | | - Ok(Self { |
56 | | - base_url, |
57 | | - shutdown_tx, |
58 | | - task, |
59 | | - }) |
| 75 | + async fn handle_error(err: hypr_llm_cactus::Error) -> (StatusCode, String) { |
| 76 | + (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) |
60 | 77 | } |
| 78 | +} |
61 | 79 |
|
62 | | - pub fn url(&self) -> &str { |
63 | | - &self.base_url |
| 80 | +#[cfg(not(target_arch = "aarch64"))] |
| 81 | +mod inner { |
| 82 | + use std::path::Path; |
| 83 | + |
| 84 | + use crate::{Error, ModelSelection}; |
| 85 | + |
| 86 | + pub struct LlmServer { |
| 87 | + _private: (), |
64 | 88 | } |
65 | 89 |
|
66 | | - pub async fn stop(self) { |
67 | | - let _ = self.shutdown_tx.send(()); |
68 | | - let _ = self.task.await; |
69 | | - tracing::info!("local LLM server stopped"); |
| 90 | + impl LlmServer { |
| 91 | + pub async fn start(_selection: &ModelSelection, _models_dir: &Path) -> Result<Self, Error> { |
| 92 | + Err(Error::Other( |
| 93 | + "Local LLM is not supported on this platform".to_string(), |
| 94 | + )) |
| 95 | + } |
| 96 | + |
| 97 | + pub fn url(&self) -> &str { |
| 98 | + unreachable!() |
| 99 | + } |
| 100 | + |
| 101 | + pub async fn stop(self) {} |
70 | 102 | } |
71 | 103 | } |
72 | 104 |
|
73 | | -async fn handle_error(err: hypr_llm_cactus::Error) -> (StatusCode, String) { |
74 | | - (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()) |
75 | | -} |
| 105 | +pub use inner::LlmServer; |
0 commit comments