src/storage/memory.rs - master - kleidinc/brain

src/storage/memory.rs
336 строк12 KB
Your Name
docs: complete rewrite of documentation suite
20 мар 2026, 14:16
20 мар 2026, 14:1664a640d
use anyhow::Result;
use arrow_array::{
    Array, FixedSizeListArray, Int64Array, RecordBatch, RecordBatchIterator, StringArray,
    types::Float32Type,
};
use arrow_schema::{DataType, Field, Schema};
use futures::StreamExt;
use lancedb::connection::connect;
use lancedb::query::{ExecutableQuery, QueryBase};
use serde::{Deserialize, Serialize};
use std::path::Path;
use std::sync::Arc;

/// Conversation memory schema for agentic behavior
/// Stores chat history, session context, and agent interactions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConversationMessage {
    pub id: String,
    pub session_id: String,
    pub agent_id: String,
    pub message_type: MessageType,
    pub content: String,
    pub timestamp: String,
    pub context_metadata: serde_json::Value, // Flexible JSON for agent-specific data
    pub embedding: Vec<f32>, // For semantic search of conversations
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum MessageType {
    Query,
    Response,
    ToolCall,
    ToolResult,
    System,
    AgentAction,
    UserContext,
}

impl std::fmt::Display for MessageType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MessageType::Query => write!(f, "query"),
            MessageType::Response => write!(f, "response"),
            MessageType::ToolCall => write!(f, "tool_call"),
            MessageType::ToolResult => write!(f, "tool_result"),
            MessageType::System => write!(f, "system"),
            MessageType::AgentAction => write!(f, "agent_action"),
            MessageType::UserContext => write!(f, "user_context"),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionMetadata {
    pub session_id: String,
    pub agent_id: String,
    pub client_type: String, // "claude-code", "cursor", "custom"
    pub project_context: String, // Current project/domain being worked on
    pub preferences: serde_json::Value, // Agent preferences and persona
    pub start_time: String,
    pub last_active: String,
    pub message_count: usize,
    pub status: SessionStatus,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SessionStatus {
    Active,
    Paused,
    Completed,
    Expired,
}

impl std::fmt::Display for SessionStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SessionStatus::Active => write!(f, "active"),
            SessionStatus::Paused => write!(f, "paused"),
            SessionStatus::Completed => write!(f, "completed"),
            SessionStatus::Expired => write!(f, "expired"),
        }
    }
}

/// Vector store extension for conversation memory
pub struct ConversationMemoryStore {
    messages_db: lancedb::connection::Connection,
    sessions_db: lancedb::connection::Connection,
    messages_table: String,
    sessions_table: String,
    embedding_manager: crate::embedding::EmbeddingManager,
    dimensions: usize,
}

impl ConversationMemoryStore {
    pub async fn new(
        data_dir: &Path,
        embedding_manager: crate::embedding::EmbeddingManager,
        dimensions: usize
    ) -> Result<Self> {
        let messages_db_path = data_dir.join("conversation_messages");
        let messages_db = connect(messages_db_path.to_str().unwrap()).execute().await?;

        let sessions_db_path = data_dir.join("conversation_sessions");
        let sessions_db = connect(sessions_db_path.to_str().unwrap()).execute().await?;

        let store = Self {
            messages_db,
            sessions_db,
            messages_table: "messages".to_string(),
            sessions_table: "sessions".to_string(),
            embedding_manager,
            dimensions,
        };

        store.ensure_message_table().await?;
        store.ensure_session_table().await?;

        Ok(store)
    }

    fn messages_schema(&self) -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("id", DataType::Utf8, false),
            Field::new("session_id", DataType::Utf8, false),
            Field::new("agent_id", DataType::Utf8, false),
            Field::new("message_type", DataType::Utf8, false),
            Field::new("content", DataType::LargeUtf8, false),
            Field::new("timestamp", DataType::Utf8, false),
            Field::new("context_metadata", DataType::LargeUtf8, false),
            Field::new(
                "embedding",
                DataType::FixedSizeList(
                    Arc::new(Field::new("item", DataType::Float32, true)),
                    self.dimensions as i32,
                ),
                false,
            ),
        ]))
    }

    fn sessions_schema(&self) -> Arc<Schema> {
        Arc::new(Schema::new(vec![
            Field::new("session_id", DataType::Utf8, false),
            Field::new("agent_id", DataType::Utf8, false),
            Field::new("client_type", DataType::Utf8, false),
            Field::new("project_context", DataType::Utf8, false),
            Field::new("preferences", DataType::LargeUtf8, false),
            Field::new("start_time", DataType::Utf8, false),
            Field::new("last_active", DataType::Utf8, false),
            Field::new("message_count", DataType::Int64, false),
            Field::new("status", DataType::Utf8, false),
        ]))
    }

    async fn ensure_message_table(&self) -> Result<()> {
        let tables = self.messages_db.table_names().execute().await?;
        if !tables.contains(&self.messages_table) {
            let schema = self.messages_schema();
            let batch = RecordBatch::new_empty(schema.clone());
            let batches: Vec<Result<RecordBatch, arrow_schema::ArrowError>> = vec![Ok(batch)];
            let reader = RecordBatchIterator::new(batches.into_iter(), schema);

            self.messages_db
                .create_table(&self.messages_table, reader)
                .execute()
                .await?;

            tracing::info!("Created conversation messages table");
        }
        Ok(())
    }

    async fn ensure_session_table(&self) -> Result<()> {
        let tables = self.sessions_db.table_names().execute().await?;
        if !tables.contains(&self.sessions_table) {
            let schema = self.sessions_schema();
            let batch = RecordBatch::new_empty(schema.clone());
            let batches: Vec<Result<RecordBatch, arrow_schema::ArrowError>> = vec![Ok(batch)];
            let reader = RecordBatchIterator::new(batches.into_iter(), schema);

            self.sessions_db
                .create_table(&self.sessions_table, reader)
                .execute()
                .await?;

            tracing::info!("Created conversation sessions table");
        }
        Ok(())
    }

    /// Store a conversation message with semantic embedding
    pub async fn store_message(
        &self,
        session_id: &str,
        agent_id: &str,
        message_type: MessageType,
        content: &str,
        context_metadata: serde_json::Value,
    ) -> Result<String> {
        let id = uuid::Uuid::new_v4().to_string();
        let timestamp = chrono::Utc::now().to_rfc3339();

        // Generate embedding for semantic search
        let embedding = self.embedding_manager.embed_single(content)?;

        let batch = RecordBatch::try_new(
            self.messages_schema(),
            vec![
                Arc::new(StringArray::from(vec![id.as_str()])),
                Arc::new(StringArray::from(vec![session_id])),
                Arc::new(StringArray::from(vec![agent_id])),
                Arc::new(StringArray::from(vec![message_type.to_string()])),
                Arc::new(arrow_array::LargeStringArray::from(vec![content])),
                Arc::new(StringArray::from(vec![timestamp.as_str()])),
                Arc::new(arrow_array::LargeStringArray::from(vec![context_metadata.to_string()])),
                Arc::new(FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
                    vec![Some(embedding.iter().map(|&v| Some(v)).collect::<Vec<_>>())],
                    self.dimensions as i32,
                )),
            ],
        )?;

        let table = self.messages_db.open_table(&self.messages_table).execute().await?;
        let batches: Vec<Result<RecordBatch, arrow_schema::ArrowError>> = vec![Ok(batch)];
        let reader = RecordBatchIterator::new(batches.into_iter(), self.messages_schema());
        table.add(reader).execute().await?;

        Ok(id)
    }

    /// Create or update a session
    pub async fn upsert_session(&self, metadata: &SessionMetadata) -> Result<()> {
        // Delete existing session if it exists
        let predicate = format!("session_id = '{}'", metadata.session_id.escape_default());
        if let Ok(table) = self.sessions_db.open_table(&self.sessions_table).execute().await {
            let _ = table.delete(&predicate).await;
        }

        let batch = RecordBatch::try_new(
            self.sessions_schema(),
            vec![
                Arc::new(StringArray::from(vec![metadata.session_id.as_str()])),
                Arc::new(StringArray::from(vec![metadata.agent_id.as_str()])),
                Arc::new(StringArray::from(vec![metadata.client_type.as_str()])),
                Arc::new(StringArray::from(vec![metadata.project_context.as_str()])),
                Arc::new(arrow_array::LargeStringArray::from(vec![metadata.preferences.to_string()])),
                Arc::new(StringArray::from(vec![metadata.start_time.as_str()])),
                Arc::new(StringArray::from(vec![metadata.last_active.as_str()])),
                Arc::new(Int64Array::from(vec![metadata.message_count as i64])),
                Arc::new(StringArray::from(vec![metadata.status.to_string()])),
            ],
        )?;

        let table = self.sessions_db.open_table(&self.sessions_table).execute().await?;
        let batches: Vec<Result<RecordBatch, arrow_schema::ArrowError>> = vec![Ok(batch)];
        let reader = RecordBatchIterator::new(batches.into_iter(), self.sessions_schema());
        table.add(reader).execute().await?;

        Ok(())
    }

    /// Retrieve conversation context for a session
    pub async fn get_session_context(
        &self,
        session_id: &str,
        max_messages: usize
    ) -> Result<Vec<ConversationMessage>> {
        let table = self.messages_db.open_table(&self.messages_table).execute().await?;

        let mut stream = table
            .query()
            .filter(format!("session_id = '{}'", session_id.escape_default()))
            .limit(max_messages)
            .execute()
            .await?;

        let mut messages = Vec::new();

        while let Some(batch) = stream.next().await {
            let batch = batch?;
            // Parse batch into ConversationMessage structs
            // (implementation details for brevity - full implementation would handle all fields)
            messages.push(self.parse_message_from_batch(&batch, 0));
        }

        // Sort by timestamp, most recent first
        messages.sort_by(|a, b| b.timestamp.cmp(&a.timestamp));

        Ok(messages)
    }

    /// Semantic search across conversation history
    pub async fn search_conversations(
        &self,
        query_embedding: &[f32],
        session_filter: Option<&str>,
        limit: usize,
    ) -> Result<Vec<ConversationMessage>> {
        let table = self.messages_db.open_table(&self.messages_table).execute().await?;

        let mut query = table
            .query()
            .nearest_to(query_embedding.to_vec())?
            .limit(limit);

        if let Some(session_id) = session_filter {
            query = query.filter(format!("session_id = '{}'", session_id.escape_default()));
        }

        let mut stream = query.execute().await?;
        let mut results = Vec::new();

        while let Some(batch) = stream.next().await {
            let batch = batch?;
            results.push(self.parse_message_from_batch(&batch, 0));
        }

        Ok(results)
    }

    fn parse_message_from_batch(&self, batch: &RecordBatch, row_idx: usize) -> ConversationMessage {
        // Implementation to parse arrow array data into ConversationMessage
        // This is a placeholder - full implementation would extract all fields
        ConversationMessage {
            id: "placeholder".to_string(),
            session_id: "placeholder".to_string(),
            agent_id: "placeholder".to_string(),
            message_type: MessageType::Query,
            content: "placeholder".to_string(),
            timestamp: chrono::Utc::now().to_rfc3339(),
            context_metadata: serde_json::Value::Null,
            embedding: vec![],
        }
    }
}