gitdataai/libs/room/src/search.rs
ZhenYi f5e3da35b0 feat(room): store ordered streaming chunks + billing integration
- Save thinking_content as {"__chunks__": [{type, content}]} for replay
- Tool call sanitization — don't expose raw results to frontend
- Billing record_ai_usage integration
- Room service module refactoring into service/ directory
2026-04-26 13:10:42 +08:00

326 lines
12 KiB
Rust

use crate::error::RoomError;
use crate::service::RoomService;
use crate::types::RoomMessageSearchRequest;
use crate::ws_context::WsUserContext;
use chrono::Utc;
use models::rooms::{room_message, room_message_reaction};
use models::{DateTimeUtc, MessageId, RoomId, RoomThreadId, Seq, UserId};
use sea_orm::*;
use uuid::Uuid;
impl RoomService {
pub async fn room_message_search(
&self,
room_id: Uuid,
request: RoomMessageSearchRequest,
ctx: &WsUserContext,
) -> Result<super::MessageSearchResponse, RoomError> {
let user_id = ctx.user_id;
self.require_room_member(room_id, user_id).await?;
if request.q.trim().is_empty() {
return Ok(super::MessageSearchResponse {
messages: Vec::new(),
total: 0,
});
}
let limit = std::cmp::min(request.limit.unwrap_or(20), 100);
let offset = request.offset.unwrap_or(0);
// Build dynamic WHERE conditions
let mut conditions = vec![
"room = $1".to_string(),
"content_tsv @@ plainto_tsquery('simple', $2)".to_string(),
"revoked IS NULL".to_string(),
];
let mut param_index = 3;
let mut params: Vec<sea_orm::Value> = vec![room_id.into(), request.q.trim().into()];
// Add time range filter
if let Some(start_time) = request.start_time {
conditions.push(format!("send_at >= ${}", param_index));
params.push(start_time.into());
param_index += 1;
}
if let Some(end_time) = request.end_time {
conditions.push(format!("send_at <= ${}", param_index));
params.push(end_time.into());
param_index += 1;
}
// Add sender filter
if let Some(sender_id) = request.sender_id {
conditions.push(format!("sender_id = ${}", param_index));
params.push(sender_id.into());
param_index += 1;
}
// Add content type filter
if let Some(ref content_type) = request.content_type {
conditions.push(format!("content_type = ${}", param_index));
params.push(content_type.clone().into());
param_index += 1;
}
let where_clause = conditions.join(" AND ");
// PostgreSQL full-text search with highlighting via raw SQL.
// Uses ts_headline for result highlighting with <mark> tags.
let sql = format!(
r#"
SELECT id, seq, room, sender_type, sender_id, thread, in_reply_to,
content, content_type, edited_at, send_at, revoked, revoked_by,
ts_headline('simple', content, plainto_tsquery('simple', $2),
'StartSel=<mark>, StopSel=</mark>, MaxWords=50, MinWords=15') AS highlighted_content
FROM room_message
WHERE {}
ORDER BY send_at DESC
LIMIT ${} OFFSET ${}"#,
where_clause,
param_index,
param_index + 1
);
params.push(limit.into());
params.push(offset.into());
let stmt = Statement::from_sql_and_values(DbBackend::Postgres, &sql, params);
let rows = self.db.query_all_raw(stmt).await?;
// Parse results and build response with highlighted content
let mut results: Vec<super::RoomMessageResponse> = Vec::new();
for row in rows {
let sender_type_str = row.try_get::<String>("", "sender_type").unwrap_or_default();
let sender_type = match sender_type_str.as_str() {
"admin" => models::rooms::MessageSenderType::Admin,
"owner" => models::rooms::MessageSenderType::Owner,
"ai" => models::rooms::MessageSenderType::Ai,
"system" => models::rooms::MessageSenderType::System,
"tool" => models::rooms::MessageSenderType::Tool,
"guest" => models::rooms::MessageSenderType::Guest,
_ => models::rooms::MessageSenderType::Member,
};
let content_type_str = row.try_get::<String>("", "content_type").unwrap_or_default();
let content_type = match content_type_str.as_str() {
"image" => models::rooms::MessageContentType::Image,
"audio" => models::rooms::MessageContentType::Audio,
"video" => models::rooms::MessageContentType::Video,
"file" => models::rooms::MessageContentType::File,
_ => models::rooms::MessageContentType::Text,
};
let msg = room_message::Model {
id: row.try_get::<MessageId>("", "id").unwrap_or_default(),
seq: row.try_get::<Seq>("", "seq").unwrap_or_default(),
room: row.try_get::<RoomId>("", "room").unwrap_or_default(),
sender_type,
sender_id: row.try_get::<Option<UserId>>("", "sender_id").ok().flatten(),
model_id: row.try_get::<Option<Uuid>>("", "model_id").ok().flatten(),
thread: row.try_get::<Option<RoomThreadId>>("", "thread").ok().flatten(),
in_reply_to: row.try_get::<Option<MessageId>>("", "in_reply_to").ok().flatten(),
content: row.try_get::<String>("", "content").unwrap_or_default(),
content_type,
thinking_content: None,
edited_at: row.try_get::<Option<DateTimeUtc>>("", "edited_at").ok().flatten(),
send_at: row.try_get::<DateTimeUtc>("", "send_at").unwrap_or_default(),
revoked: row.try_get::<Option<DateTimeUtc>>("", "revoked").ok().flatten(),
revoked_by: row.try_get::<Option<UserId>>("", "revoked_by").ok().flatten(),
content_tsv: None,
};
let highlighted_content = row
.try_get::<String>("", "highlighted_content")
.unwrap_or_else(|_| msg.content.clone());
// Resolve display name for this message
let message_with_name = self.resolve_display_name(msg.clone(), room_id).await;
let mut msg_with_name = message_with_name;
msg_with_name.highlighted_content = Some(highlighted_content);
results.push(msg_with_name);
}
// COUNT query for total (without pagination)
let mut count_conditions = vec![
"room = $1".to_string(),
"content_tsv @@ plainto_tsquery('simple', $2)".to_string(),
"revoked IS NULL".to_string(),
];
let mut count_params: Vec<sea_orm::Value> = vec![room_id.into(), request.q.trim().into()];
let mut count_param_idx = 3;
if let Some(start_time) = request.start_time {
count_conditions.push(format!("send_at >= ${}", count_param_idx));
count_params.push(start_time.into());
count_param_idx += 1;
}
if let Some(end_time) = request.end_time {
count_conditions.push(format!("send_at <= ${}", count_param_idx));
count_params.push(end_time.into());
count_param_idx += 1;
}
if let Some(sender_id) = request.sender_id {
count_conditions.push(format!("sender_id = ${}", count_param_idx));
count_params.push(sender_id.into());
count_param_idx += 1;
}
if let Some(ref content_type) = request.content_type {
count_conditions.push(format!("content_type = ${}", count_param_idx));
count_params.push(content_type.clone().into());
}
let count_sql = format!(
"SELECT COUNT(*) AS count FROM room_message WHERE {}",
count_conditions.join(" AND ")
);
let count_stmt = Statement::from_sql_and_values(DbBackend::Postgres, &count_sql, count_params);
let count_row = self.db.query_one_raw(count_stmt).await?;
let total: i64 = count_row
.and_then(|r| r.try_get::<i64>("", "count").ok())
.unwrap_or(0);
Ok(super::MessageSearchResponse {
messages: results,
total,
})
}
pub async fn room_message_reaction_list(
&self,
room_id: Uuid,
message_id: Uuid,
ctx: &WsUserContext,
) -> Result<super::MessageReactionsResponse, RoomError> {
let user_id = ctx.user_id;
self.require_room_member(room_id, user_id).await?;
let _msg = room_message::Entity::find_by_id(message_id)
.one(&self.db)
.await?
.ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;
self.get_message_reactions(message_id, Some(user_id)).await
}
pub async fn room_message_reaction_toggle(
&self,
room_id: Uuid,
message_id: Uuid,
emoji: String,
ctx: &WsUserContext,
) -> Result<super::MessageReactionsResponse, RoomError> {
let user_id = ctx.user_id;
self.require_room_member(room_id, user_id).await?;
if emoji.is_empty() || emoji.len() > 50 {
return Err(RoomError::BadRequest("Invalid emoji format".to_string()));
}
if let Some(existing) = room_message_reaction::Entity::find()
.filter(room_message_reaction::Column::Room.eq(room_id))
.filter(room_message_reaction::Column::Message.eq(message_id))
.filter(room_message_reaction::Column::User.eq(user_id))
.filter(room_message_reaction::Column::Emoji.eq(&emoji))
.one(&self.db)
.await?
{
room_message_reaction::Entity::delete_by_id(existing.id)
.exec(&self.db)
.await?;
} else {
room_message_reaction::ActiveModel {
id: Set(Uuid::now_v7()),
room: Set(room_id),
message: Set(message_id),
user: Set(user_id),
emoji: Set(emoji),
created_at: Set(Utc::now()),
}
.insert(&self.db)
.await?;
}
self.get_message_reactions(message_id, Some(user_id)).await
}
pub async fn room_message_edit_history(
&self,
room_id: Uuid,
message_id: Uuid,
ctx: &WsUserContext,
) -> Result<super::MessageEditHistoryResponse, RoomError> {
let user_id = ctx.user_id;
self.require_room_member(room_id, user_id).await?;
let _msg = room_message::Entity::find_by_id(message_id)
.one(&self.db)
.await?
.ok_or_else(|| RoomError::NotFound("Message not found".to_string()))?;
let history = models::rooms::room_message_edit_history::Entity::find()
.filter(models::rooms::room_message_edit_history::Column::Message.eq(message_id))
.order_by_asc(models::rooms::room_message_edit_history::Column::EditedAt)
.all(&self.db)
.await?;
let total_edits = history.len() as i64;
let entries: Vec<super::MessageEditHistoryEntry> = history
.into_iter()
.map(|h| super::MessageEditHistoryEntry {
old_content: h.old_content,
new_content: h.new_content,
edited_at: h.edited_at,
})
.collect();
Ok(super::MessageEditHistoryResponse {
message_id,
history: entries,
total_edits,
})
}
pub async fn room_member_leave(
&self,
room_id: Uuid,
ctx: &WsUserContext,
) -> Result<(), RoomError> {
let user_id = ctx.user_id;
let member = self
.find_room_member(room_id, user_id)
.await?
.ok_or_else(|| RoomError::NotFound("You are not a member of this room".to_string()))?;
if member.role == models::rooms::RoomMemberRole::Owner {
return Err(RoomError::BadRequest(
"Owner cannot leave the room. Transfer ownership first.".to_string(),
));
}
models::rooms::room_member::Entity::delete_by_id((room_id, user_id))
.exec(&self.db)
.await?;
self.room_manager.unsubscribe(room_id, user_id).await;
let room = self.find_room_or_404(room_id).await?;
self.publish_room_event(
room.project,
super::RoomEventType::MemberRemoved,
Some(room_id),
None,
None,
None,
)
.await;
Ok(())
}
}