398 lines
13 KiB
Rust
398 lines
13 KiB
Rust
use crate::{
|
|
databases::postgres::chat::ConversationState, dto::api::BaseLLMRequest,
|
|
middlewares::auth::middleware::Auth, providers::ollama::client::OllamaProvider,
|
|
};
|
|
use axum::{
|
|
Json,
|
|
extract::{Extension, State},
|
|
response::{
|
|
IntoResponse, Response,
|
|
sse::{Event, KeepAlive, Sse},
|
|
},
|
|
};
|
|
use sqlx::PgPool;
|
|
use tokio_stream::StreamExt;
|
|
use uuid::Uuid;
|
|
|
|
use crate::databases::postgres::chat::{
|
|
get_or_create_conversation, set_conversation_title, update_message_tokens,
|
|
};
|
|
use crate::databases::postgres::{chat, errors};
|
|
use crate::dto::api;
|
|
use crate::dto::api::CompletionRequest;
|
|
use crate::providers::ollama::errors::into_http_response;
|
|
use crate::state::app_state::AppState;
|
|
|
|
#[utoipa::path(
|
|
post,
|
|
path = "/completions",
|
|
tag = "chat",
|
|
request_body(
|
|
content = api::CompletionRequest,
|
|
description = "Text completion request",
|
|
content_type = "application/json"
|
|
),
|
|
responses(
|
|
(
|
|
status = 200,
|
|
description = "Text completion response. If stream=true, response is SSE stream of chunks ending in [DONE].",
|
|
body = api::CompletionResponse,
|
|
content_type = "application/json"
|
|
),
|
|
(
|
|
status = 400,
|
|
description = "Invalid request: missing prompt, model, or invalid format",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "prompt is required and cannot be empty" })
|
|
),
|
|
(
|
|
status = 422,
|
|
description = "Model not found or not available locally",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "model 'llama3' is not available — run `ollama pull llama3` first" })
|
|
),
|
|
(
|
|
status = 500,
|
|
description = "Internal server error (Ollama or network failure)",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "connection refused" })
|
|
)
|
|
)
|
|
)]
|
|
pub async fn completions(
|
|
State(state): State<AppState>,
|
|
Json(body): Json<api::CompletionRequest>,
|
|
) -> Result<Response, (axum::http::StatusCode, Json<api::ErrorResponse>)> {
|
|
tracing::debug!("Received /completion with body {:?}", body);
|
|
if body.base.stream {
|
|
let stream = state.ollama.completions_stream(&body).await.map_err(|e| {
|
|
let (code, msg) = into_http_response(e);
|
|
(code, Json(api::ErrorResponse::new(msg)))
|
|
})?;
|
|
|
|
Ok(Sse::new(stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response())
|
|
} else {
|
|
let response = state.ollama.completions(&body).await.map_err(|e| {
|
|
let (code, msg) = into_http_response(e);
|
|
(code, Json(api::ErrorResponse::new(msg)))
|
|
})?;
|
|
|
|
Ok(Json(response).into_response())
|
|
}
|
|
}
|
|
|
|
#[utoipa::path(
|
|
post,
|
|
path = "/chat/completions",
|
|
tag = "chat",
|
|
request_body(
|
|
content = api::ChatRequest,
|
|
description = "Chat completion request with message history",
|
|
content_type = "application/json"
|
|
),
|
|
responses(
|
|
(
|
|
status = 200,
|
|
description = "Chat completion response. If stream=false returns JSON. If stream=true returns SSE stream of chunks ending with [DONE].",
|
|
body = api::ChatCompletionResponse,
|
|
content_type = "application/json"
|
|
),
|
|
(
|
|
status = 400,
|
|
description = "Invalid request",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "messages array with at least one user message is required" })
|
|
),
|
|
(
|
|
status = 401,
|
|
description = "Unauthorized",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "missing or invalid token" })
|
|
),
|
|
(
|
|
status = 422,
|
|
description = "Model not found or unavailable",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "model 'llama3' is not available — run `ollama pull llama3` first" })
|
|
),
|
|
(
|
|
status = 500,
|
|
description = "Internal server error",
|
|
body = api::ErrorResponse,
|
|
example = json!({ "error": "connection refused" })
|
|
)
|
|
)
|
|
)]
|
|
pub async fn chat_completions(
|
|
State(state): State<AppState>,
|
|
Extension(auth): Extension<Auth>,
|
|
Json(body): Json<api::ChatRequest>,
|
|
) -> Result<Response, (axum::http::StatusCode, String)> {
|
|
tracing::debug!("Received /chat/completion with body {:?}", body);
|
|
|
|
// Get current conversation or create it
|
|
let conversation_id = if matches!(&auth, Auth::Jwt(_)) {
|
|
tracing::debug!(
|
|
"Is conversation_id existing: {:?}",
|
|
body.conversation_id.is_some()
|
|
);
|
|
|
|
let conversation_state =
|
|
get_or_create_conversation(&state.postgres, body.conversation_id, auth.user_id())
|
|
.await
|
|
.map_err(crate::databases::postgres::errors::into_http_response)?;
|
|
|
|
let id = match &conversation_state {
|
|
ConversationState::Existing(uuid) => *uuid,
|
|
ConversationState::Created(uuid) => {
|
|
let conversation_id = *uuid;
|
|
let title =
|
|
generate_conversation_title(&state.ollama, body.messages[0].content.as_str())
|
|
.await;
|
|
|
|
if let Err(e) =
|
|
set_conversation_title(&state.postgres, conversation_id, auth.user_id(), &title)
|
|
.await
|
|
{
|
|
tracing::warn!(
|
|
conversation_id = %conversation_id,
|
|
error = %e,
|
|
"Failed to set conversation title, continuing with no title"
|
|
);
|
|
}
|
|
|
|
tracing::debug!(conversation_id = %conversation_id, %title, "generated conversation title");
|
|
conversation_id
|
|
}
|
|
};
|
|
|
|
tracing::debug!("Using conversation_id: {:?}", id);
|
|
Some(id)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
if body.base.stream {
|
|
let stream = state
|
|
.ollama
|
|
.chat_completions_stream(&body)
|
|
.await
|
|
.map_err(into_http_response)?;
|
|
|
|
if let Some(conversation_id) = conversation_id {
|
|
let user_msg_id = log_user_message(
|
|
&state.postgres,
|
|
auth.user_id(),
|
|
conversation_id,
|
|
body.parent_id,
|
|
body.messages
|
|
.last()
|
|
.map(|m| m.content.as_str())
|
|
.unwrap_or(""),
|
|
None, // not known at this stage
|
|
)
|
|
.await
|
|
.map_err(errors::into_http_response)?;
|
|
|
|
let start_event = serde_json::json!({
|
|
"type": "start",
|
|
"model": body.base.model,
|
|
"conversation_id": conversation_id,
|
|
"created": chrono::Utc::now().timestamp() as u64,
|
|
});
|
|
|
|
let start_stream = futures::stream::once(async move {
|
|
Ok::<Event, crate::providers::ollama::errors::OllamaError>(
|
|
Event::default()
|
|
.event("metadata")
|
|
.data(start_event.to_string()),
|
|
)
|
|
});
|
|
|
|
let mut accumulated = String::new();
|
|
|
|
let wrapped_stream = stream.map(move |item| match item {
|
|
Ok(chunk) => {
|
|
let is_done = chunk.choices[0].finish_reason == Some(api::FinishReason::Stop);
|
|
|
|
if let Some(content) = chunk.choices[0].delta.content.as_ref() {
|
|
accumulated.push_str(content);
|
|
}
|
|
|
|
if is_done {
|
|
let content = accumulated.clone();
|
|
let pool = state.postgres.clone();
|
|
let user_id = auth.user_id();
|
|
let prompt_tokens = chunk.usage.as_ref().map(|u| u.prompt_tokens as i32);
|
|
let completion_tokens =
|
|
chunk.usage.as_ref().map(|u| u.completion_tokens as i32);
|
|
|
|
let end_event = serde_json::json!({
|
|
"type": "end",
|
|
"usage": {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": chunk.usage.as_ref().map(|u| u.total_tokens),
|
|
}
|
|
});
|
|
|
|
tokio::spawn(async move {
|
|
if let Some(prompt_tokens) = prompt_tokens {
|
|
let _ = update_message_tokens(
|
|
&pool,
|
|
user_id,
|
|
user_msg_id,
|
|
prompt_tokens,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let _ = log_assistant_message(
|
|
&pool,
|
|
user_id,
|
|
conversation_id,
|
|
user_msg_id,
|
|
&content,
|
|
completion_tokens,
|
|
)
|
|
.await;
|
|
});
|
|
|
|
return Ok(Event::default()
|
|
.event("metadata")
|
|
.data(end_event.to_string()));
|
|
}
|
|
|
|
let data = serde_json::to_string(&chunk).unwrap_or_default();
|
|
Ok(Event::default().data(data))
|
|
}
|
|
Err(e) => Err(e),
|
|
});
|
|
|
|
let full_stream = start_stream.chain(wrapped_stream);
|
|
|
|
Ok(Sse::new(full_stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response())
|
|
} else {
|
|
let plain_stream = stream.map(
|
|
|item| -> Result<Event, crate::providers::ollama::errors::OllamaError> {
|
|
match item {
|
|
Ok(chunk) => {
|
|
let data = serde_json::to_string(&chunk).unwrap_or_default();
|
|
Ok(Event::default().data(data))
|
|
}
|
|
Err(e) => Err(e),
|
|
}
|
|
},
|
|
);
|
|
|
|
Ok(Sse::new(plain_stream)
|
|
.keep_alive(KeepAlive::default())
|
|
.into_response())
|
|
}
|
|
} else {
|
|
let mut response = state
|
|
.ollama
|
|
.chat_completions(&body)
|
|
.await
|
|
.map_err(into_http_response)?;
|
|
|
|
response.conversation_id = conversation_id;
|
|
|
|
if let Some(conversation_id) = conversation_id {
|
|
// log user message, get its id as parent for assistant
|
|
let user_msg_id = log_user_message(
|
|
&state.postgres,
|
|
auth.user_id(),
|
|
conversation_id,
|
|
body.parent_id,
|
|
body.messages
|
|
.last()
|
|
.map(|m| m.content.as_str())
|
|
.unwrap_or(""),
|
|
response.usage.map(|u| u.prompt_tokens as i32),
|
|
)
|
|
.await
|
|
.map_err(errors::into_http_response)?;
|
|
|
|
// log assistant response
|
|
log_assistant_message(
|
|
&state.postgres,
|
|
auth.user_id(),
|
|
conversation_id,
|
|
user_msg_id,
|
|
&response.choices[0].message.content,
|
|
response.usage.map(|u| u.completion_tokens as i32),
|
|
)
|
|
.await
|
|
.map_err(errors::into_http_response)?;
|
|
}
|
|
|
|
Ok(Json(response).into_response())
|
|
}
|
|
}
|
|
|
|
async fn generate_conversation_title(ollama: &OllamaProvider, first_message: &str) -> String {
|
|
let request = CompletionRequest {
|
|
base: BaseLLMRequest {
|
|
model: "llama3:latest".to_string(),
|
|
..Default::default()
|
|
},
|
|
prompt: format!(
|
|
"Generate a short title (max 6 words) for the following chat conversation: {}",
|
|
first_message
|
|
),
|
|
};
|
|
|
|
ollama
|
|
.completions(&request)
|
|
.await
|
|
.ok()
|
|
.and_then(|r| r.choices.first().map(|c| c.text.trim().to_string()))
|
|
.filter(|t| !t.is_empty())
|
|
.unwrap_or_else(|| "New Conversation".to_string()) // ← default on any error
|
|
}
|
|
|
|
async fn log_user_message(
|
|
pool: &PgPool,
|
|
user_id: Uuid,
|
|
conversation_id: Uuid,
|
|
parent_id: Option<Uuid>,
|
|
content: &str,
|
|
tokens: Option<i32>,
|
|
) -> Result<Uuid, errors::DbError> {
|
|
chat::insert_message(
|
|
pool,
|
|
user_id,
|
|
conversation_id,
|
|
parent_id,
|
|
chat::MessageRole::User,
|
|
content,
|
|
tokens,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn log_assistant_message(
|
|
pool: &PgPool,
|
|
user_id: Uuid,
|
|
conversation_id: Uuid,
|
|
parent_id: Uuid,
|
|
content: &str,
|
|
tokens: Option<i32>,
|
|
) -> Result<Uuid, errors::DbError> {
|
|
chat::insert_message(
|
|
pool,
|
|
user_id,
|
|
conversation_id,
|
|
Some(parent_id),
|
|
chat::MessageRole::Assistant,
|
|
content,
|
|
tokens,
|
|
)
|
|
.await
|
|
}
|