refactor(llm-api): implement clean architecture following scraper pattern
Split monolithic 1012-line main.rs into layered hexagonal architecture: - Domain: entity types and LlmError enum - Application: prompt building, sampler construction, tool call parsing - Infrastructure: LlamaEngine wrapping llama-cpp-2 with isolated unsafe transmute - Presentation: Axum handlers, middleware (auth), error chain, router - Config: type-safe AppConfig with LazyLock - Bootstrap: Application struct with build() + run() Resolves build_sampler/build_sampler_params duplication. Adds simple web chat UI at GET /. Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
asepharyana
co-authored by
Claude Code
Claude Code
parent
dfd6fa66a7
commit
e351d74fa4
@@ -0,0 +1,310 @@
|
||||
//! Chat completions endpoint — streaming and non-streaming.
|
||||
|
||||
use std::convert::Infallible;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::State;
|
||||
use axum::response::sse::{Event, KeepAlive, Sse};
|
||||
use axum::{Json, response::{IntoResponse, Response}};
|
||||
use chrono::Utc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use tracing::info;
|
||||
|
||||
use crate::application::chat;
|
||||
use crate::domain::entity::{
|
||||
ChatRequest, ChatResponse, Choice, ResponseMessage, SseChunk, SseChoice, SseDelta, Usage,
|
||||
};
|
||||
use crate::infrastructure::llama::SendSampler;
|
||||
use crate::presentation::error::AppError;
|
||||
use crate::presentation::state::AppState;
|
||||
|
||||
/// POST /v1/chat/completions
|
||||
pub async fn chat_completions(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<ChatRequest>,
|
||||
) -> Result<Response, AppError> {
|
||||
let max_tokens = req.max_tokens.unwrap_or(256).min(1024);
|
||||
let stop = req.stop.clone().unwrap_or_default();
|
||||
let prompt = chat::build_prompt(&req.messages, &req.tools);
|
||||
|
||||
// Tokenize
|
||||
let input_tokens = state
|
||||
.engine
|
||||
.tokenize(&prompt)
|
||||
.map_err(|e| AppError::LlmError(e.to_string()))?;
|
||||
|
||||
let prompt_tokens = input_tokens.len() as u32;
|
||||
info!(
|
||||
" Chat: {} prompt tokens, max_tokens={}, tools={}",
|
||||
prompt_tokens,
|
||||
max_tokens,
|
||||
req.tools.as_ref().is_some_and(|t| !t.is_empty())
|
||||
);
|
||||
|
||||
let response = if req.stream.unwrap_or(false) {
|
||||
handle_streaming(state.clone(), req, max_tokens, stop, input_tokens).await?
|
||||
} else {
|
||||
handle_non_streaming(state.clone(), req, max_tokens, stop, input_tokens).await?
|
||||
};
|
||||
Ok(response.into_response())
|
||||
}
|
||||
|
||||
// ── Non-streaming path ──
|
||||
|
||||
async fn handle_non_streaming(
|
||||
state: Arc<AppState>,
|
||||
req: ChatRequest,
|
||||
max_tokens: u32,
|
||||
stop: Vec<String>,
|
||||
input_tokens: Vec<llama_cpp_2::token::LlamaToken>,
|
||||
) -> Result<Response, AppError> {
|
||||
let chat_id = format!("chatcmpl-{}", uuid::Uuid::new_v4());
|
||||
let created = Utc::now().timestamp();
|
||||
let prompt_tokens = input_tokens.len() as u32;
|
||||
let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty());
|
||||
let params = chat::SamplerParams::from_request(&req);
|
||||
let mut sampler = SendSampler(chat::build_sampler(¶ms));
|
||||
|
||||
let (output_tokens, raw_text) = state
|
||||
.engine
|
||||
.generate(&input_tokens, &mut sampler, max_tokens, &stop)
|
||||
.await?;
|
||||
|
||||
let (output_text, tool_calls) = chat::parse_tool_calls(&chat::clean_text(&raw_text));
|
||||
|
||||
let completion_tokens = output_tokens.len() as u32;
|
||||
info!(" {} generated tokens", completion_tokens);
|
||||
|
||||
let finish_reason = if has_tools && !tool_calls.is_empty() {
|
||||
"tool_calls"
|
||||
} else if completion_tokens < max_tokens {
|
||||
"stop"
|
||||
} else {
|
||||
"length"
|
||||
};
|
||||
|
||||
Ok(Json(ChatResponse {
|
||||
id: chat_id,
|
||||
object: "chat.completion".into(),
|
||||
created,
|
||||
model: req.model,
|
||||
choices: vec![Choice {
|
||||
index: 0,
|
||||
message: ResponseMessage {
|
||||
role: "assistant".into(),
|
||||
content: Some(output_text),
|
||||
tool_calls: if tool_calls.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_calls)
|
||||
},
|
||||
},
|
||||
finish_reason: finish_reason.into(),
|
||||
}],
|
||||
usage: Usage {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens: prompt_tokens + completion_tokens,
|
||||
},
|
||||
})
|
||||
.into_response())
|
||||
}
|
||||
|
||||
// ── Streaming path ──
|
||||
|
||||
async fn handle_streaming(
|
||||
state: Arc<AppState>,
|
||||
req: ChatRequest,
|
||||
max_tokens: u32,
|
||||
stop: Vec<String>,
|
||||
input_tokens: Vec<llama_cpp_2::token::LlamaToken>,
|
||||
) -> Result<Response, AppError> {
|
||||
let chat_id = format!("chatcmpl-{}", uuid::Uuid::new_v4());
|
||||
let created = Utc::now().timestamp();
|
||||
let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty());
|
||||
let model_name = req.model.clone();
|
||||
let params = chat::SamplerParams::from_request(&req);
|
||||
let (tx, rx) = mpsc::channel::<Result<Event, Infallible>>(64);
|
||||
|
||||
tokio::spawn(async move {
|
||||
// Role chunk
|
||||
let role_chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: Some("assistant".into()),
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: None,
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
if tx.send(Ok(Event::default().data(role_chunk))).await.is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
// Build sampler
|
||||
let mut sampler = SendSampler(chat::build_sampler(¶ms));
|
||||
|
||||
// Lock context
|
||||
let mut inner = state.engine.ctx().lock().await;
|
||||
inner.clear();
|
||||
if let Err(e) = inner.prefill(&input_tokens) {
|
||||
info!(" Prefill error: {e}");
|
||||
return;
|
||||
}
|
||||
|
||||
let mut count = 0u32;
|
||||
let mut text_buf = String::new();
|
||||
let mut current = inner.sample(&mut sampler);
|
||||
|
||||
loop {
|
||||
if count >= max_tokens {
|
||||
let chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: None,
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: Some("length".into()),
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||
break;
|
||||
}
|
||||
|
||||
if state.engine.is_eog(current) {
|
||||
let reason = if has_tools && text_buf.contains("<tool_call>") {
|
||||
"tool_calls"
|
||||
} else {
|
||||
"stop"
|
||||
};
|
||||
let chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: None,
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: Some(reason.into()),
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||
break;
|
||||
}
|
||||
|
||||
let piece = state.engine.decode_token(current);
|
||||
let content = chat::clean_text(&piece);
|
||||
|
||||
if !content.is_empty() {
|
||||
let chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: None,
|
||||
content: Some(content.clone()),
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: None,
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
if tx.send(Ok(Event::default().data(chunk))).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
text_buf.push_str(&piece);
|
||||
|
||||
// Check stop sequences
|
||||
let mut stop_now = false;
|
||||
for s in &stop {
|
||||
if text_buf.contains(s) {
|
||||
stop_now = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if stop_now {
|
||||
let chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: None,
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: Some("stop".into()),
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||
break;
|
||||
}
|
||||
|
||||
// Check tool call completeness
|
||||
if has_tools && text_buf.contains("<tool_call>") {
|
||||
let open = text_buf.matches("<tool_call>").count();
|
||||
let close = text_buf.matches("</tool_call>").count();
|
||||
if close >= open {
|
||||
let chunk = serde_json::to_string(&SseChunk {
|
||||
id: chat_id.clone(),
|
||||
object: "chat.completion.chunk".into(),
|
||||
created,
|
||||
model: model_name.clone(),
|
||||
choices: vec![SseChoice {
|
||||
index: 0,
|
||||
delta: SseDelta {
|
||||
role: None,
|
||||
content: None,
|
||||
tool_calls: None,
|
||||
},
|
||||
finish_reason: Some("tool_calls".into()),
|
||||
}],
|
||||
})
|
||||
.unwrap();
|
||||
let _ = tx.send(Ok(Event::default().data(chunk))).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let pos = input_tokens.len() as i32 + count as i32;
|
||||
if let Err(e) = inner.decode(current, pos) {
|
||||
info!(" Decode error: {e}");
|
||||
break;
|
||||
}
|
||||
count += 1;
|
||||
current = inner.sample(&mut sampler);
|
||||
}
|
||||
});
|
||||
|
||||
let stream = ReceiverStream::new(rx);
|
||||
let sse = Sse::new(stream).keep_alive(KeepAlive::default());
|
||||
Ok(sse.into_response())
|
||||
}
|
||||
Reference in New Issue
Block a user