Iceberg/chat-api
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: log user message

Browse Source
This commit is contained in:
LucasDLTG
2026-05-13 12:39:02 +02:00
parent c976db7063
commit f01bc74059
7 changed files with 327 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.sqlx/query-426a96fdb055d248f1f73d52dd10014340c6174b1491c1da36a5ba63fa554a9f.json
Generated
+15
View File
@@ -0,0 +1,15 @@
{
"db_name": "PostgreSQL",
"query": "UPDATE chat.message SET tokens = $1 WHERE id = $2",
"describe": {
"columns": [],
"parameters": {
"Left": [
"Int4",
"Uuid"
]
},
"nullable": []
},
"hash": "426a96fdb055d248f1f73d52dd10014340c6174b1491c1da36a5ba63fa554a9f"
}
.sqlx/query-c98b93484092fc57098f4c6eac92af653bcf39aa1bb4d8add80ed069421ee859.json
Generated
+26
View File
@@ -0,0 +1,26 @@
{
"db_name": "PostgreSQL",
"query": "\n INSERT INTO chat.message (conversation_id, parent_id, role, content, tokens)\n VALUES ($1, $2, $3, $4, $5)\n RETURNING id\n ",
"describe": {
"columns": [
{
"ordinal": 0,
"name": "id",
"type_info": "Uuid"
}
],
"parameters": {
"Left": [
"Uuid",
"Uuid",
"Text",
"Text",
"Int4"
]
},
"nullable": [
false
]
},
"hash": "c98b93484092fc57098f4c6eac92af653bcf39aa1bb4d8add80ed069421ee859"
}
src/databases/postgres/chat.rs
+71
View File
@@ -97,3 +97,74 @@ pub async fn set_conversation_title(
tx.commit().await?;
Ok(())
}
#[derive(Debug, Clone, sqlx::Type)]
#[sqlx(type_name = "text")]
#[sqlx(rename_all = "lowercase")]
pub enum MessageRole {
User,
Assistant,
System,
}
pub async fn insert_message(
pool: &PgPool,
user_id: Uuid,
conversation_id: Uuid,
parent_id: Option<Uuid>,
role: MessageRole,
content: &str,
tokens: Option<i32>,
) -> Result<Uuid, DbError> {
let mut tx = pool.begin().await?;
let conn = tx.acquire().await?;
sqlx::query(r#"SELECT set_config('app.current_user_id', $1::text, true)"#)
.bind(user_id.to_string())
.execute(&mut *conn)
.await?;
let row = sqlx::query!(
r#"
INSERT INTO chat.message (conversation_id, parent_id, role, content, tokens)
VALUES ($1, $2, $3, $4, $5)
RETURNING id
"#,
conversation_id,
parent_id,
role as MessageRole,
content,
tokens
)
.fetch_one(&mut *conn)
.await?;
tx.commit().await?;
Ok(row.id)
}
pub async fn update_message_tokens(
pool: &PgPool,
user_id: Uuid,
message_id: Uuid,
tokens: i32,
) -> Result<(), DbError> {
let mut tx = pool.begin().await?;
let conn = tx.acquire().await?;
sqlx::query(r#"SELECT set_config('app.current_user_id', $1::text, true)"#)
.bind(user_id.to_string())
.execute(&mut *conn)
.await?;
sqlx::query!(
r#"UPDATE chat.message SET tokens = $1 WHERE id = $2"#,
tokens,
message_id
)
.execute(&mut *conn)
.await?;
tx.commit().await?;
Ok(())
}
src/dto/api.rs
+3 -1
View File
@@ -108,7 +108,7 @@ pub struct Choice {
pub finish_reason: FinishReason,
}
#[derive(Debug, Serialize, Deserialize, ToSchema)]
#[derive(Debug, Serialize, Deserialize, ToSchema, Clone, Copy)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
@@ -131,6 +131,7 @@ pub struct ChatRequest {
// Non standard Open AI
pub conversation_id: Option<Uuid>,
pub parent_id: Option<Uuid>,
}
#[derive(Debug, Deserialize, Serialize, ToSchema)]
@@ -171,6 +172,7 @@ pub struct ChatCompletionChunk {
pub id: String,
pub object: String,
pub choices: Vec<ChatChunkChoice>,
pub usage: Option<Usage>,
}
#[derive(Debug, Serialize, ToSchema)]
src/providers/ollama/client.rs
+15 -5
View File
@@ -345,7 +345,7 @@ impl OllamaProvider {
pub async fn chat_completions_stream(
&self,
body: &api::ChatRequest,
) -> Result<ReceiverStream<Result<Event, OllamaError>>, OllamaError> {
) -> Result<ReceiverStream<Result<api::ChatCompletionChunk, OllamaError>>, OllamaError> {
let url = format!("{}/api/chat", self.base_url);
let (messages, model) = self.extract_chat_params(body)?;
@@ -399,6 +399,18 @@ impl OllamaProvider {
Err(_) => continue,
};
let usage = if parsed.done {
Some(api::Usage {
prompt_tokens: parsed.prompt_eval_count.unwrap_or(0) as u32,
completion_tokens: parsed.eval_count.unwrap_or(0) as u32,
total_tokens: (parsed.prompt_eval_count.unwrap_or(0)
+ parsed.eval_count.unwrap_or(0))
as u32,
})
} else {
None
};
let event = api::ChatCompletionChunk {
id: stream_id.clone(),
object: "chat.completion.chunk".to_string(),
@@ -414,14 +426,12 @@ impl OllamaProvider {
None
},
}],
usage,
};
let event_data = serde_json::to_string(&event).unwrap_or_default();
let _ = tx.send(Ok(Event::default().data(event_data))).await;
let _ = tx.send(Ok(event)).await;
if parsed.done {
let _ = tx.send(Ok(Event::default().data("[DONE]"))).await;
break;
}
}
src/routes/v1/chat.rs
+193 -5
View File
@@ -7,11 +7,17 @@ use axum::{
extract::{Extension, State},
response::{
IntoResponse, Response,
sse::{KeepAlive, Sse},
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};
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;
@@ -126,6 +132,7 @@ pub async fn chat_completions(
) -> 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: {:?}",
@@ -173,16 +180,157 @@ pub async fn chat_completions(
.chat_completions_stream(&body)
.await
.map_err(into_http_response)?;
Ok(Sse::new(stream)
.keep_alive(KeepAlive::default())
.into_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())
}
}
@@ -207,3 +355,43 @@ async fn generate_conversation_title(ollama: &OllamaProvider, first_message: &st
.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
}
tests/ollama_provider.rs
+4
View File
@@ -184,6 +184,7 @@ async fn test_chat_completions_ok() {
content: "What is 2+2?".to_string(),
}],
conversation_id: None,
parent_id: None,
};
let res = provider.chat_completions(&req).await.unwrap();
@@ -219,6 +220,7 @@ async fn test_chat_completions_missing_messages() {
},
messages: vec![],
conversation_id: None,
parent_id: None,
};
let err = provider.chat_completions(&req).await.unwrap_err();
@@ -246,6 +248,7 @@ async fn test_chat_completions_no_user_message() {
content: "be helpful".to_string(),
}],
conversation_id: None,
parent_id: None,
};
let err = provider.chat_completions(&req).await.unwrap_err();
@@ -273,6 +276,7 @@ async fn test_chat_completions_model_not_found() {
content: "hi".to_string(),
}],
conversation_id: None,
parent_id: None,
};
let err = provider.chat_completions(&req).await.unwrap_err();