imphnen-backend-service/imphnen-dimentorin/src/ai_agent/application/rag_service.rs

use std::sync::Arc;

use chrono::Utc;
use sea_orm::DatabaseConnection;
use uuid::Uuid;

use super::super::domain::{
	ChatRequest, ChatResponse, RagDocument, RagRepository, RagService, RagSource,
};
use crate::ai_agent::infrastructure::llm_provider::{chat_completion, embed_text};
use crate::materials::domain::MaterialRepository;
use imphnen_entities::seaorm::common::materials::{
	Column as MaterialColumn, Entity as MaterialsEntity,
};
use imphnen_utils::AppError;
use sea_orm::{ColumnTrait, EntityTrait, QueryFilter};

const CHUNK_CHARS: usize = 700;
const CHUNK_OVERLAP: usize = 80;
const SEARCH_LIMIT: u64 = 4;
const MAX_ANSWER_TOKENS: u32 = 600;

/// Stable u64 point id for a material chunk (uuid bytes xor chunk index).
fn qdrant_point_id(material_id: Uuid, chunk_index: usize) -> u64 {
	let bytes = material_id.as_bytes();
	let mut val: u64 = 0;
	for (i, b) in bytes.iter().enumerate() {
		val ^= (*b as u64) << ((i % 8) * 8);
	}
	val ^ (chunk_index as u64)
}

pub struct RagServiceImpl {
	repo: Arc<dyn RagRepository>,
	db: DatabaseConnection,
}

impl RagServiceImpl {
	pub fn new(repo: Arc<dyn RagRepository>, db: DatabaseConnection) -> Self {
		Self { repo, db }
	}

	fn split_chunks(title: &str, content: &str) -> Vec<String> {
		let mut chunks = Vec::new();
		let text = format!("{}\n{}", title, content);
		let bytes = text.as_bytes();
		let mut start = 0usize;
		while start < bytes.len() {
			let end = (start + CHUNK_CHARS).min(bytes.len());
			// don't split mid-utf8 char
			let mut cut = end;
			while cut > start && !bytes[cut - 1].is_ascii() && cut < end + 3 {
				cut = end;
				break;
			}
			let chunk = &text[start..cut];
			if !chunk.trim().is_empty() {
				chunks.push(chunk.to_string());
			}
			if end >= bytes.len() {
				break;
			}
			start = end.saturating_sub(CHUNK_OVERLAP);
		}
		chunks
	}

	async fn index_entity(&self, material_id: Uuid) -> Result<u64, AppError> {
		let material = MaterialsEntity::find_by_id(material_id)
			.one(&self.db)
			.await
			.map_err(|e| AppError::InternalServerError(e.to_string()))?
			.ok_or_else(|| AppError::NotFoundError("Material not found".into()))?;
		if !material.is_published {
			return Ok(0);
		}
		// re-index: clear old chunks first
		self.repo.delete_material(material_id).await?;

		let chunks = Self::split_chunks(&material.title, &material.content);
		let mut indexed = 0u64;
		let now = Utc::now();
		for (i, chunk) in chunks.iter().enumerate() {
			let embedding = embed_text(chunk).await?;
			let doc = RagDocument {
				material_id,
				title: material.title.clone(),
				chunk: chunk.clone(),
				indexed_at: now,
			};
			let point_id = qdrant_point_id(material_id, i);
			self.repo.upsert_document(point_id, &doc, embedding).await?;
			indexed += 1;
		}
		Ok(indexed)
	}
}

#[async_trait::async_trait]
impl RagService for RagServiceImpl {
	async fn index_material(&self, material_id: Uuid) -> Result<u64, AppError> {
		self.index_entity(material_id).await
	}

	async fn reindex_all(&self) -> Result<u64, AppError> {
		let materials = MaterialsEntity::find()
			.filter(MaterialColumn::IsPublished.eq(true))
			.all(&self.db)
			.await
			.map_err(|e| AppError::InternalServerError(e.to_string()))?;
		let mut total = 0u64;
		for m in materials {
			total += self.index_entity(m.id).await?;
		}
		Ok(total)
	}

	async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, AppError> {
		let question = request.question.trim().to_string();
		if question.is_empty() {
			return Err(AppError::BadRequestError("question tidak boleh kosong".into()));
		}

		let query_emb = embed_text(&question).await?;
		let hits = self
			.repo
			.search(query_emb, SEARCH_LIMIT, request.material_id)
			.await?;

		if hits.is_empty() {
			// No indexed context: answer without RAG context but stay honest.
			let answer = chat_completion(
				"Kamu adalah asisten AI Dimentorin. Jawab pertanyaan singkat dan jelas. Jika tidak tahu, akui tidak tahu.",
				&question,
				MAX_ANSWER_TOKENS,
			)
			.await?;
			return Ok(ChatResponse {
				answer,
				sources: vec![],
			});
		}

		let context: Vec<String> = hits
			.iter()
			.map(|(doc, _)| format!("[{}]\n{}", doc.title, doc.chunk))
			.collect();
		let context_block = context.join("\n\n---\n\n");
		let system = format!(
			"Kamu adalah asisten AI Dimentorin yang menjawab berdasarkan materi mentoring berikut.\n\
			 Jawab dalam bahasa Indonesia, singkat, jelas, dan berfokus pada konteks yang diberikan.\n\
			 Jika pertanyaan di luar materi, katakan bahwa hal itu di luar materi yang tersedia.\n\n\
			 === MATERI ===\n{}",
			context_block
		);
		let answer = chat_completion(&system, &question, MAX_ANSWER_TOKENS).await?;

		let sources = hits
			.into_iter()
			.map(|(doc, score)| RagSource {
				material_id: doc.material_id,
				title: doc.title,
				score,
				snippet: doc.chunk.chars().take(180).collect(),
			})
			.collect();

		Ok(ChatResponse { answer, sources })
	}
}