cognee-lib 0.1.3

Cognee — an AI-memory pipeline that turns raw data into queryable knowledge graphs (umbrella crate).
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

//! Pipeline context trait for shared component access.

use std::sync::Arc;

use async_trait::async_trait;

use cognee_database::DatabaseConnection;
use cognee_embedding::EmbeddingEngine;
use cognee_graph::GraphDBTrait;
use cognee_llm::Llm;
use cognee_storage::StorageTrait;
use cognee_vector::VectorDB;

use crate::error::ComponentError;

/// Trait providing access to shared pipeline components.
///
/// Implementations lazily initialize and cache expensive components
/// (e.g., embedding models, database connections) so they can be
/// reused across multiple pipeline invocations.
#[async_trait]
pub trait PipelineContext: Send + Sync {
    async fn storage(&self) -> Result<Arc<dyn StorageTrait>, ComponentError>;
    async fn database(&self) -> Result<Arc<DatabaseConnection>, ComponentError>;
    async fn graph_db(&self) -> Result<Arc<dyn GraphDBTrait>, ComponentError>;
    async fn vector_db(&self) -> Result<Arc<dyn VectorDB>, ComponentError>;
    async fn embedding_engine(&self) -> Result<Arc<dyn EmbeddingEngine>, ComponentError>;
    async fn llm(&self) -> Result<Arc<dyn Llm>, ComponentError>;
}