rain_engine_core/
retrieval.rs

1use crate::{GoalRecord, ObservationRecord, RelationshipEdge, SessionRecord, TaskRecord};
2use async_trait::async_trait;
3use serde::{Deserialize, Serialize};
4use thiserror::Error;
5
6#[derive(Debug, Error, Clone, PartialEq, Eq)]
7#[error("{message}")]
8pub struct RetrievalError {
9    pub message: String,
10}
11
12impl RetrievalError {
13    pub fn new(message: impl Into<String>) -> Self {
14        Self {
15            message: message.into(),
16        }
17    }
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
21pub enum RetrievedItemKind {
22    Observation,
23    Task,
24    Goal,
25    SessionRecord,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
29pub struct RetrievedItem {
30    pub kind: RetrievedItemKind,
31    pub key: String,
32    pub score: f32,
33    pub snippet: String,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
37pub struct WorkingSet {
38    pub observations: Vec<ObservationRecord>,
39    pub tasks: Vec<TaskRecord>,
40    pub goals: Vec<GoalRecord>,
41}
42
43#[async_trait]
44pub trait RetrievalStore: Send + Sync {
45    async fn exact_replay(
46        &self,
47        session_id: &str,
48        limit: usize,
49    ) -> Result<Vec<SessionRecord>, RetrievalError>;
50
51    async fn semantic_search(
52        &self,
53        session_id: &str,
54        query: &str,
55        limit: usize,
56    ) -> Result<Vec<RetrievedItem>, RetrievalError>;
57
58    async fn graph_neighbors(
59        &self,
60        session_id: &str,
61        resource_id: &str,
62        max_hops: usize,
63    ) -> Result<Vec<RelationshipEdge>, RetrievalError>;
64
65    async fn recent_working_set(
66        &self,
67        session_id: &str,
68        limit: usize,
69    ) -> Result<WorkingSet, RetrievalError>;
70}
rain_engine_core/retrieval.rs

rain_engine_core/
retrieval.rs
