gitdataai/libs/room/src/ai.rs

use crate::error::RoomError;
use crate::service::RoomService;
use crate::ws_context::WsUserContext;
use chrono::Utc;
use models::agents::model as ai_model;
use models::rooms::room_ai;
use sea_orm::*;
use uuid::Uuid;

impl RoomService {
    pub async fn room_ai_list(
        &self,
        room_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<Vec<super::RoomAiResponse>, RoomError> {
        let user_id = ctx.user_id;
        self.require_room_access(room_id, user_id).await?;

        let configs = room_ai::Entity::find()
            .filter(room_ai::Column::Room.eq(room_id))
            .all(&self.db)
            .await?;

        if configs.is_empty() {
            return Ok(Vec::new());
        }

        // Batch-fetch all referenced models to avoid N+1 queries
        let model_ids: Vec<Uuid> = configs.iter().map(|c| c.model).collect();
        let models = ai_model::Entity::find()
            .filter(ai_model::Column::Id.is_in(model_ids))
            .all(&self.db)
            .await?;

        let model_names: std::collections::HashMap<Uuid, String> = models
            .into_iter()
            .map(|m| (m.id, m.name))
            .collect();

        let mut responses = Vec::with_capacity(configs.len());
        for config in configs {
            // Skip entries where model_name cannot be resolved — never expose UID
            let Some(model_name) = model_names.get(&config.model).cloned() else {
                tracing::warn!(
                    "room_ai_list: skipping config with unknown model_id={}",
                    config.model
                );
                continue;
            };
            let mut resp = super::RoomAiResponse::from(config);
            resp.model_name = Some(model_name);
            responses.push(resp);
        }

        Ok(responses)
    }

    pub async fn room_ai_upsert(
        &self,
        room_id: Uuid,
        request: super::RoomAiUpsertRequest,
        ctx: &WsUserContext,
    ) -> Result<super::RoomAiResponse, RoomError> {
        let user_id = ctx.user_id;
        self.require_room_admin(room_id, user_id).await?;

        let now = Utc::now();
        let existing = room_ai::Entity::find_by_id((room_id, request.model))
            .one(&self.db)
            .await?;

        let saved = if let Some(existing) = existing {
            let mut active: room_ai::ActiveModel = existing.into();
            if request.version.is_some() {
                active.version = Set(request.version);
            }
            if request.history_limit.is_some() {
                active.history_limit = Set(request.history_limit);
            }
            if request.system_prompt.is_some() {
                active.system_prompt = Set(request.system_prompt);
            }
            if request.temperature.is_some() {
                active.temperature = Set(request.temperature);
            }
            if request.max_tokens.is_some() {
                active.max_tokens = Set(request.max_tokens);
            }
            if request.use_exact.is_some() {
                active.use_exact = Set(request.use_exact.unwrap_or(true));
            }
            if request.think.is_some() {
                active.think = Set(request.think.unwrap_or(false));
            }
            if request.stream.is_some() {
                active.stream = Set(request.stream.unwrap_or(false));
            }
            if request.min_score.is_some() {
                active.min_score = Set(request.min_score);
            }
            if request.agent_type.is_some() {
                active.agent_type = Set(request.agent_type);
            }
            active.updated_at = Set(now);
            active.update(&self.db).await?
        } else {
            room_ai::ActiveModel {
                room: Set(room_id),
                model: Set(request.model),
                version: Set(request.version),
                call_count: Set(0),
                last_call_at: Set(None),
                history_limit: Set(request.history_limit),
                system_prompt: Set(request.system_prompt),
                temperature: Set(request.temperature),
                max_tokens: Set(request.max_tokens),
                use_exact: Set(request.use_exact.unwrap_or(true)),
                think: Set(request.think.unwrap_or(false)),
                stream: Set(request.stream.unwrap_or(false)),
                min_score: Set(request.min_score),
                agent_type: Set(request.agent_type),
                created_at: Set(now),
                updated_at: Set(now),
            }
            .insert(&self.db)
            .await?
        };

        let model_name = ai_model::Entity::find_by_id(saved.model)
            .one(&self.db)
            .await
            .ok()
            .flatten()
            .map(|m| m.name);
        let mut resp = super::RoomAiResponse::from(saved);
        resp.model_name = model_name;

        if let Ok(room) = self.find_room_or_404(room_id).await {
            self.publish_room_event(
                room.project,
                super::RoomEventType::RoomAiUpdated,
                Some(room_id),
                None,
                None,
                None,
            )
            .await;
        }

        Ok(resp)
    }

    pub async fn room_ai_delete(
        &self,
        room_id: Uuid,
        model_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<(), RoomError> {
        let user_id = ctx.user_id;
        self.require_room_admin(room_id, user_id).await?;

        room_ai::Entity::delete_by_id((room_id, model_id))
            .exec(&self.db)
            .await?;

        if let Ok(room) = self.find_room_or_404(room_id).await {
            self.publish_room_event(
                room.project,
                super::RoomEventType::RoomAiUpdated,
                Some(room_id),
                None,
                None,
                None,
            )
            .await;
        }

        Ok(())
    }

    pub async fn room_ai_stop(
        &self,
        room_id: Uuid,
        ctx: &WsUserContext,
    ) -> Result<(), RoomError> {
        let user_id = ctx.user_id;
        self.require_room_access(room_id, user_id).await?;
        tracing::info!(%room_id, %user_id, "AI stream stop requested");
        Ok(())
    }
}