cognate_rag/
lib.rs

1//! Cognate RAG — Retrieval-Augmented Generation pipeline.
2//!
3//! # Overview
4//!
5//! 1. Choose an [`EmbeddingProvider`] (e.g. `OpenAiProvider` from `cognate-providers`).
6//! 2. Choose a [`VectorStore`] backend (e.g. [`MemoryVectorStore`] for prototyping).
7//! 3. Wrap them in [`RagPipeline`] to get [`ingest`](RagPipeline::ingest) and
8//!    [`retrieve`](RagPipeline::retrieve).
9//!
10//! # Example
11//!
12//! ```rust,no_run
13//! use cognate_rag::{RagPipeline, MemoryVectorStore};
14//! use cognate_core::EmbeddingProvider;
15//!
16//! async fn run(embedder: impl EmbeddingProvider) {
17//!     let store = MemoryVectorStore::new();
18//!     let pipeline = RagPipeline::new(embedder, store);
19//!
20//!     pipeline
21//!         .ingest(
22//!             vec!["Rust is fast".to_string(), "Rust is safe".to_string()],
23//!             vec![serde_json::json!({"source": "doc1"}), serde_json::json!({"source": "doc2"})],
24//!         )
25//!         .await
26//!         .unwrap();
27//!
28//!     let results = pipeline.retrieve("fast systems language", 1).await.unwrap();
29//!     println!("{}", results[0].content);
30//! }
31//! ```
32#![warn(missing_docs)]
33
34pub mod memory;
35
36pub use memory::MemoryVectorStore;
37
38use async_trait::async_trait;
39use cognate_core::{EmbeddingProvider, Error};
40use serde::{Deserialize, Serialize};
41
42// ─── Core types ────────────────────────────────────────────────────────────
43
44/// A dense embedding vector.
45pub type Vector = Vec<f32>;
46
47/// A document stored in a vector store.
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct Document {
50    /// Unique identifier for this document.
51    pub id: String,
52    /// The text content of the document.
53    pub content: String,
54    /// Arbitrary key-value metadata associated with this document.
55    pub metadata: serde_json::Value,
56    /// The embedding vector, if one has been computed.
57    pub embedding: Option<Vector>,
58}
59
60// ─── VectorStore trait ─────────────────────────────────────────────────────
61
62/// A persistent or in-memory store of embedded documents.
63///
64/// Implement this trait to add support for a new vector database backend
65/// (pgvector, Qdrant, Pinecone, etc.).
66#[async_trait]
67pub trait VectorStore: Send + Sync {
68    /// Persist a batch of documents (with their embeddings) to the store.
69    async fn add_documents(
70        &self,
71        docs: Vec<Document>,
72    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
73
74    /// Return the `limit` documents whose embeddings are most similar to
75    /// `query_vector` (cosine similarity, descending).
76    async fn search(
77        &self,
78        query_vector: Vector,
79        limit: usize,
80    ) -> Result<Vec<Document>, Box<dyn std::error::Error + Send + Sync>>;
81}
82
83// ─── RagPipeline ───────────────────────────────────────────────────────────
84
85/// A high-level RAG pipeline combining an embedding provider with a vector store.
86///
87/// # Type parameters
88///
89/// * `E` — any [`EmbeddingProvider`] (e.g. `OpenAiProvider`).
90/// * `V` — any [`VectorStore`] (e.g. [`MemoryVectorStore`]).
91pub struct RagPipeline<E, V> {
92    embedder: E,
93    store: V,
94}
95
96impl<E: EmbeddingProvider, V: VectorStore> RagPipeline<E, V> {
97    /// Create a new pipeline from an embedder and a vector store.
98    pub fn new(embedder: E, store: V) -> Self {
99        Self { embedder, store }
100    }
101
102    /// Embed `texts` and store them alongside `metadata` in the vector store.
103    ///
104    /// `texts` and `metadata` must have the same length.
105    ///
106    /// # Errors
107    ///
108    /// Returns [`Error::VectorStore`] if embedding or storage fails.
109    pub async fn ingest(
110        &self,
111        texts: Vec<String>,
112        metadata: Vec<serde_json::Value>,
113    ) -> cognate_core::Result<()> {
114        if texts.len() != metadata.len() {
115            return Err(Error::InvalidRequest(
116                "texts and metadata must have the same length".to_string(),
117            ));
118        }
119
120        let embeddings = self
121            .embedder
122            .embed(texts.clone())
123            .await
124            .map_err(|e| Error::VectorStore(e.to_string()))?;
125
126        let docs: Vec<Document> = texts
127            .into_iter()
128            .zip(embeddings)
129            .zip(metadata)
130            .enumerate()
131            .map(|(i, ((content, emb), meta))| Document {
132                id: i.to_string(),
133                content,
134                metadata: meta,
135                embedding: Some(emb),
136            })
137            .collect();
138
139        self.store
140            .add_documents(docs)
141            .await
142            .map_err(|e| Error::VectorStore(e.to_string()))
143    }
144
145    /// Embed `query` and return the `limit` most similar documents.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`Error::VectorStore`] if embedding or search fails.
150    pub async fn retrieve(
151        &self,
152        query: &str,
153        limit: usize,
154    ) -> cognate_core::Result<Vec<Document>> {
155        let mut embeddings = self
156            .embedder
157            .embed(vec![query.to_string()])
158            .await
159            .map_err(|e| Error::VectorStore(e.to_string()))?;
160
161        let query_vec = if embeddings.is_empty() {
162            return Err(Error::VectorStore(
163                "embedding provider returned no vectors".to_string(),
164            ));
165        } else {
166            embeddings.remove(0)
167        };
168
169        self.store
170            .search(query_vec, limit)
171            .await
172            .map_err(|e| Error::VectorStore(e.to_string()))
173    }
174}
cognate_rag/lib.rs

cognate_rag/
lib.rs
