cognee_session/
session_store.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4
5use crate::error::SessionError;
6use crate::types::{SessionQAEntry, SessionTraceStep, UsedGraphElementIds};
7
8/// Partial update DTO for a QA entry.
9///
10/// - Outer `None` means "leave field unchanged".
11/// - `Some(None)` means "clear the field".
12/// - `Some(Some(value))` means "set the field to this value".
13///
14/// For non-optional fields (`question`, `answer`) the outer `Option` controls
15/// whether an update is applied.
16#[derive(Debug, Clone, Default)]
17pub struct SessionQAUpdate {
18    pub question: Option<String>,
19    pub answer: Option<String>,
20    pub context: Option<Option<String>>,
21    pub feedback_text: Option<Option<String>>,
22    pub feedback_score: Option<Option<i32>>,
23    pub used_graph_element_ids: Option<Option<UsedGraphElementIds>>,
24    pub memify_metadata: Option<Option<HashMap<String, bool>>>,
25}
26
27/// Abstraction over session Q&A storage backends (SQLite, Redis, filesystem, etc.).
28///
29/// Analogous to Python's `CacheDBInterface`. All backends implement this trait.
30#[async_trait]
31pub trait SessionStore: Send + Sync {
32    /// Store a Q&A entry in the session. Returns the generated `qa_id`.
33    async fn create_qa_entry(
34        &self,
35        session_id: &str,
36        user_id: Option<&str>,
37        question: &str,
38        answer: &str,
39        context: Option<&str>,
40    ) -> Result<String, SessionError>;
41
42    /// Retrieve the most recent `last_n` Q&A entries for a session, ordered oldest-first.
43    async fn get_latest_qa_entries(
44        &self,
45        session_id: &str,
46        user_id: Option<&str>,
47        last_n: usize,
48    ) -> Result<Vec<SessionQAEntry>, SessionError>;
49
50    /// Retrieve all Q&A entries for a session, ordered oldest-first.
51    async fn get_all_qa_entries(
52        &self,
53        session_id: &str,
54        user_id: Option<&str>,
55    ) -> Result<Vec<SessionQAEntry>, SessionError>;
56
57    /// Delete all entries for a session. Returns `true` if the session existed.
58    async fn delete_session(
59        &self,
60        session_id: &str,
61        user_id: Option<&str>,
62    ) -> Result<bool, SessionError>;
63
64    /// Delete a single Q&A entry by id. Returns `true` if found and deleted.
65    async fn delete_qa_entry(
66        &self,
67        session_id: &str,
68        user_id: Option<&str>,
69        qa_id: &str,
70    ) -> Result<bool, SessionError>;
71
72    /// Delete ALL session data across all users and sessions.
73    /// Equivalent to Python's `CacheDBInterface.prune()`.
74    async fn prune(&self) -> Result<(), SessionError>;
75
76    /// Update fields on a QA entry. Only non-`None` fields in `updates` are applied.
77    /// Returns `true` if the entry was found and updated.
78    async fn update_qa_entry(
79        &self,
80        session_id: &str,
81        user_id: Option<&str>,
82        qa_id: &str,
83        updates: SessionQAUpdate,
84    ) -> Result<bool, SessionError>;
85
86    /// Return the `qa_id` of the most-recent Q&A entry in the session, or
87    /// `None` when the session has no entries yet.
88    ///
89    /// Used by the search orchestrator to route conversationally-detected
90    /// feedback to the previous answer (mirrors Python `session_manager.py:462-469`).
91    ///
92    /// Default impl loads the latest entry via `get_latest_qa_entries(limit=1)`
93    /// and extracts its id. Backends may override with a more efficient query.
94    async fn latest_qa_id(
95        &self,
96        session_id: &str,
97        user_id: Option<&str>,
98    ) -> Result<Option<String>, SessionError> {
99        let entries = self.get_latest_qa_entries(session_id, user_id, 1).await?;
100        Ok(entries.into_iter().next().map(|e| e.id.to_string()))
101    }
102
103    /// Retrieve the graph knowledge snapshot for a session, or `None`.
104    async fn get_graph_context(
105        &self,
106        session_id: &str,
107        user_id: Option<&str>,
108    ) -> Result<Option<String>, SessionError>;
109
110    /// Store (or overwrite) the graph knowledge snapshot for a session.
111    async fn set_graph_context(
112        &self,
113        session_id: &str,
114        user_id: Option<&str>,
115        context: &str,
116    ) -> Result<(), SessionError>;
117
118    /// Append one agent-trace step to the session's trace list.
119    ///
120    /// Returns the persisted `trace_id` (caller-provided — `SessionManager`
121    /// generates it via UUID4 before invoking).
122    ///
123    /// Default impl returns `SessionError::StoreError` so backends that have
124    /// not been updated to support trace steps still compile. The fs / redis /
125    /// sea-orm backends override this with a real implementation.
126    async fn save_trace_step(
127        &self,
128        user_id: &str,
129        session_id: &str,
130        step: SessionTraceStep,
131    ) -> Result<String, SessionError> {
132        let _ = (user_id, session_id, step);
133        Err(SessionError::StoreError(
134            "save_trace_step not implemented for this backend".into(),
135        ))
136    }
137
138    /// Retrieve agent-trace steps for the given session (oldest-first).
139    ///
140    /// `SessionManager::get_agent_trace_session` performs `last_n` slicing on
141    /// top of the returned list. Default impl returns `SessionError::StoreError`.
142    async fn read_trace_steps(
143        &self,
144        user_id: &str,
145        session_id: &str,
146    ) -> Result<Vec<SessionTraceStep>, SessionError> {
147        let _ = (user_id, session_id);
148        Err(SessionError::StoreError(
149            "read_trace_steps not implemented for this backend".into(),
150        ))
151    }
152}
cognee_session/session_store.rs

cognee_session/
session_store.rs
