cognee-vector 0.1.3

Vector-store abstraction with brute-force, LanceDB, and pgvector adapters for cognee.
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

#![allow(
    clippy::unwrap_used,
    clippy::expect_used,
    reason = "test code — panics are acceptable failures"
)]
//! Span attribute integration tests for the PgVector adapter.
//!
//! Skipped silently when `cognee_test_utils::pg_test_url()` returns `None`
//! (i.e. `DB_PROVIDER` is not set to `postgres`). Mirrors the gating
//! pattern in `pgvector_integration.rs`.
#![cfg(feature = "pgvector")]

use cognee_test_utils::SpanCapture;
use cognee_vector::{PgVectorAdapter, VectorDB, VectorPoint};
use serial_test::serial;
use std::collections::HashMap;
use uuid::Uuid;

/// Returns a fresh adapter (with stale collections cleaned up) when a
/// Postgres URL is configured, or `None` to silently skip.
async fn make_adapter() -> Option<PgVectorAdapter> {
    let url = cognee_test_utils::pg_test_url()?;
    let adapter = PgVectorAdapter::new(&url, 4).await.ok()?;
    if let Ok(cols) = adapter.list_collections().await {
        for (dt, fname) in cols {
            let _ = adapter.delete_collection(&dt, &fname).await;
        }
    }
    Some(adapter)
}

#[tokio::test]
#[serial]
async fn upsert_emits_pgvector_span() {
    let Some(adapter) = make_adapter().await else {
        eprintln!("DB_PROVIDER!=postgres — skipping upsert_emits_pgvector_span");
        return;
    };
    let capture = SpanCapture::install();

    let points: Vec<VectorPoint> = (0..2)
        .map(|i| VectorPoint {
            id: Uuid::new_v4(),
            vector: vec![i as f32, 0.0, 0.0, 0.0],
            metadata: HashMap::new(),
        })
        .collect();
    adapter
        .index_points("DocumentChunk", "text", &points)
        .await
        .expect("upsert");

    let spans = capture.spans();
    let s = spans
        .iter()
        .find(|s| s.name == "cognee.db.vector.upsert")
        .expect("expected upsert span");
    assert_eq!(s.field_str("cognee.db.system").as_deref(), Some("pgvector"));
    assert_eq!(
        s.field_str("cognee.vector.collection").as_deref(),
        Some("DocumentChunk_text"),
    );
    assert_eq!(s.field_i64("cognee.db.row_count"), Some(2));
}

#[tokio::test]
#[serial]
async fn search_emits_pgvector_span() {
    let Some(adapter) = make_adapter().await else {
        eprintln!("DB_PROVIDER!=postgres — skipping search_emits_pgvector_span");
        return;
    };
    let capture = SpanCapture::install();

    let pid = Uuid::new_v4();
    adapter
        .index_points(
            "DocumentChunk",
            "text",
            &[VectorPoint {
                id: pid,
                vector: vec![0.1, 0.2, 0.3, 0.4],
                metadata: HashMap::new(),
            }],
        )
        .await
        .expect("seed");

    let results = adapter
        .search_similar("DocumentChunk", "text", &[0.1, 0.2, 0.3, 0.4], 5)
        .await
        .expect("search");
    assert!(!results.is_empty());

    let spans = capture.spans();
    let s = spans
        .iter()
        .find(|s| s.name == "cognee.db.vector.search")
        .expect("expected search span");
    assert_eq!(s.field_str("cognee.db.system").as_deref(), Some("pgvector"));
    assert_eq!(
        s.field_str("cognee.vector.collection").as_deref(),
        Some("DocumentChunk_text"),
    );
    assert!(s.field_i64("cognee.vector.result_count").unwrap_or(0) >= 1);
}

#[tokio::test]
#[serial]
async fn delete_emits_pgvector_span() {
    let Some(adapter) = make_adapter().await else {
        eprintln!("DB_PROVIDER!=postgres — skipping delete_emits_pgvector_span");
        return;
    };
    let capture = SpanCapture::install();

    let pid = Uuid::new_v4();
    adapter
        .index_points(
            "DocumentChunk",
            "text",
            &[VectorPoint {
                id: pid,
                vector: vec![0.1, 0.0, 0.0, 0.0],
                metadata: HashMap::new(),
            }],
        )
        .await
        .expect("seed");
    adapter
        .delete_points("DocumentChunk", "text", &[pid])
        .await
        .expect("delete");

    let spans = capture.spans();
    let s = spans
        .iter()
        .find(|s| s.name == "cognee.db.vector.delete")
        .expect("expected delete span");
    assert_eq!(s.field_str("cognee.db.system").as_deref(), Some("pgvector"));
}