otherone-storage 0.1.2

// 作用：MongoDB 数据库存储模块
use futures::TryStreamExt;
use mongodb::bson::{doc, Document};
use mongodb::{Client, Collection};
use uuid::Uuid;

use crate::error::StorageError;
use crate::types::{Entry, Session, SessionData};

pub struct MongoClient { client: Client, db_name: String }

impl MongoClient {
    pub async fn new(connection_string: &str, db_name: &str) -> Result<Self, StorageError> {
        let client = Client::with_uri_str(connection_string).await
            .map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        Ok(MongoClient { client, db_name: db_name.to_string() })
    }
    fn sc(&self) -> Collection<Document> { self.client.database(&self.db_name).collection("otherone_sessions") }
    fn ec(&self) -> Collection<Document> { self.client.database(&self.db_name).collection("otherone_entries") }

    pub async fn create_session(&self) -> Result<String, StorageError> {
        let sid = Uuid::new_v4().to_string();
        self.sc().insert_one(doc!{"session_id":&sid,"status":0i32,"create_at":chrono::Utc::now().to_rfc3339()}).await
            .map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        Ok(sid)
    }

    pub async fn write_entry(&self, session_id: &str, role: &str, content: &str, _tools: Option<&serde_json::Value>, token_consumption: Option<u32>) -> Result<(), StorageError> {
        let eid = Uuid::new_v4().to_string();
        self.ec().insert_one(doc!{
            "entry_id":&eid,"session_id":session_id,"role":role,"content":content,
            "token_consumption":token_consumption.map(|t| t as i32).unwrap_or(0),
            "status":0i32,"create_at":chrono::Utc::now().to_rfc3339(),"is_compaction":1i32
        }).await.map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        Ok(())
    }

    pub async fn get_all_sessions(&self) -> Result<Vec<Session>, StorageError> {
        let mut cursor = self.sc().find(doc!{"status":0i32}).await
            .map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        let mut sessions = vec![];
        while let Some(d) = cursor.try_next().await.map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))? {
            sessions.push(Session {
                session_id: d.get_str("session_id").unwrap_or("").into(),
                status: d.get_i32("status").unwrap_or(0) as i16,
                create_at: d.get_str("create_at").unwrap_or("").into(),
            });
        }
        Ok(sessions)
    }

    pub async fn read_session(&self, session_id: &str) -> Result<SessionData, StorageError> {
        let s_opt = self.sc().find_one(doc!{"session_id":session_id,"status":0i32}).await
            .map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        let session = match s_opt {
            None => return Ok(SessionData { session: None, entries: vec![], compacted_entries: vec![] }),
            Some(d) => Session { session_id: d.get_str("session_id").unwrap_or("").into(), status: d.get_i32("status").unwrap_or(0) as i16, create_at: d.get_str("create_at").unwrap_or("").into() },
        };
        let mut cursor = self.ec().find(doc!{"session_id":session_id,"status":0i32}).await
            .map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))?;
        let mut entries = vec![];
        while let Some(d) = cursor.try_next().await.map_err(|e| StorageError::ConfigError(format!("MongoDB: {}", e)))? {
            entries.push(Entry {
                entry_id: d.get_str("entry_id").unwrap_or("").into(), session_id: d.get_str("session_id").unwrap_or("").into(),
                content: d.get_str("content").unwrap_or("").into(), role: d.get_str("role").unwrap_or("").into(),
                token_consumption: d.get_i32("token_consumption").ok().map(|v| v as u32),
                status: d.get_i32("status").unwrap_or(0) as i16, tools: None,
                create_at: d.get_str("create_at").unwrap_or("").into(),
                is_compaction: d.get_i32("is_compaction").unwrap_or(1) as i16,
            });
        }
        Ok(SessionData { session: Some(session), entries, compacted_entries: vec![] })
    }
}