jarno/ws-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: storing chat messages and fetching latest messages

Browse Source
This commit is contained in:
Jarno
2026-02-13 21:34:35 +02:00
parent e63fd76d2f
commit 3ce2fa3841
4 changed files with 107 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
crates/daemon/src/chatpersistence.rs
View File

@@ -1,11 +1,13 @@
use anyhow::Result;
use directories::ProjectDirs;
use shared::ai::ChatMessage as CMessage;
use sqlx::sqlite::SqliteConnectOptions;
use sqlx::Row;
use sqlx::SqlitePool;
use tokio::fs;
use tonic::async_trait;
#[derive(Debug, sqlx::FromRow)]
pub struct ChatMessage {
pub id: i64,
pub text: String,
@@ -14,8 +16,16 @@ pub struct ChatMessage {
#[async_trait]
pub trait ChatRepository {
async fn save_message(&self, text: &str, is_user: &bool) -> Result<()>;
async fn get_all_messages(&self) -> Result<Vec<ChatMessage>>;
async fn save_message(&self, text: &str, is_user: &bool) -> Result<ChatMessage>;
async fn get_latest_messages(&self) -> Result<Vec<ChatMessage>>;
}
pub fn message_to_dto(msg: &ChatMessage) -> CMessage {
CMessage {
id: msg.id,
text: msg.text.clone(),
is_user: msg.is_user,
}
}
pub struct SqliteChatRepository {
@@ -40,14 +50,15 @@ impl SqliteChatRepository {
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS message (
"CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
text TEXT NOT NULL,
is_user BOOL NOT NULL
)",
)
.execute(&pool)
.await?;
.await
.inspect_err(|e| println!("sql error: {}", e))?;
Ok(Self { pool })
}
@@ -55,19 +66,35 @@ impl SqliteChatRepository {
#[async_trait]
impl ChatRepository for SqliteChatRepository {
async fn save_message(&self, text: &str, is_user: &bool) -> Result<()> {
sqlx::query("INSERT INTO messages (text, is_user) values (?, ?)")
.bind(text)
.bind(is_user)
.execute(&self.pool)
.await?;
Ok(())
async fn save_message(&self, text: &str, is_user: &bool) -> Result<ChatMessage> {
let result = sqlx::query_as::<_, ChatMessage>(
r#"
INSERT INTO messages (text, is_user)
VALUES (?, ?)
RETURNING id, text, is_user
"#,
)
.bind(text)
.bind(is_user)
.fetch_one(&self.pool)
.await
.inspect_err(|e| println!("sql error: {}", e))?;
Ok(result)
}
async fn get_all_messages(&self) -> Result<Vec<ChatMessage>> {
let rows = sqlx::query("SELECT id, text, is_user FROM messages ORDER BY id DESC LIMIT 10")
.fetch_all(&self.pool)
.await?;
async fn get_latest_messages(&self) -> Result<Vec<ChatMessage>> {
let rows = sqlx::query(
r#"
SELECT * FROM (
SELECT id, text, is_user
FROM messages
ORDER BY id DESC
LIMIT 10
) AS subquery ORDER BY id ASC"#,
)
.fetch_all(&self.pool)
.await
.inspect_err(|e| println!("sql error: {}", e))?;
let messages = rows
.into_iter()

57
crates/daemon/src/main.rs
View File

@@ -1,7 +1,7 @@
mod chatpersistence;
use std::cell::Cell;
use std::sync::atomic::AtomicI64;
use std::sync::Arc;
use genai::chat::{ChatMessage, ChatRequest};
use genai::Client;
@@ -10,18 +10,22 @@ use shared::ai::{
ChatHistoryRequest, ChatHistoryResponse, ChatMessage as CMessage, ChatRequest as CRequest,
ChatResponse as CResponse, PromptRequest, PromptResponse,
};
use tonic::{transport::Server, Request, Response, Status};
use tonic::{transport::Server, Code, Request, Response, Status};
use chatpersistence::SqliteChatRepository;
use crate::chatpersistence::{message_to_dto, ChatRepository};
pub struct DaemonServer {
message_counter: AtomicI64,
repo: Arc<dyn ChatRepository + Send + Sync>,
}
impl Default for DaemonServer {
fn default() -> Self {
impl DaemonServer {
pub fn new(repo: Arc<dyn ChatRepository + Send + Sync>) -> Self {
Self {
message_counter: AtomicI64::new(0),
repo,
}
}
}
@@ -46,36 +50,41 @@ impl AiDaemon for DaemonServer {
async fn chat(&self, request: Request<CRequest>) -> Result<Response<CResponse>, Status> {
let r = request.into_inner();
println!("<<<: {}", r.text());
let user_message = message_to_dto(
&self
.repo
.save_message(r.text(), &true)
.await
.map_err(|e| Status::new(Code::Internal, e.to_string()))?,
);
let response_text = format!("Pong: {}", r.text());
let ai_message = message_to_dto(
&self
.repo
.save_message(response_text.as_str(), &false)
.await
.map_err(|e| Status::new(Code::Internal, e.to_string()))?,
);
let response = CResponse {
chat_id: 1,
messages: vec![
CMessage {
id: self
.message_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
text: r.text().to_string(),
is_user: true,
},
CMessage {
id: self
.message_counter
.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
text: format!("Pong: {}", r.text()),
is_user: false,
},
],
messages: vec![user_message, ai_message],
};
return Ok(Response::new(response));
}
async fn chat_history(
&self,
request: Request<ChatHistoryRequest>,
_: Request<ChatHistoryRequest>,
) -> Result<Response<ChatHistoryResponse>, Status> {
let messages = self
.repo
.get_latest_messages()
.await
.map_err(|e| Status::new(Code::Internal, e.to_string()))?;
let response = ChatHistoryResponse {
chat_id: 1,
history: vec![],
history: messages.iter().map(|m| message_to_dto(m)).collect(),
};
Ok(Response::new(response))
}
@@ -101,7 +110,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let addr_s = "[::1]:50051";
let addr = addr_s.parse().unwrap();
let daemon = DaemonServer::default();
let daemon = DaemonServer::new(Arc::new(chat_repo));
let reflection_service = tonic_reflection::server::Builder::configure()
.register_encoded_file_descriptor_set(shared::ai::FILE_DESCRIPTOR_SET)
.build_v1()?;

2
crates/shared/proto/api.proto
View File

@@ -30,7 +30,7 @@ message ChatHistoryRequest {
message ChatHistoryResponse {
int64 chat_id = 1;
repeated ChatResponse history = 10;
repeated ChatMessage history = 10;
}
message PromptRequest {