|
| 1 | +//! [`CogneeGraph`] — a read-only [`MemoryGraph`] over Cognee's derived |
| 2 | +//! knowledge graph. |
| 3 | +//! |
| 4 | +//! Cognee's graph is **built by its `cognify` pipeline** over ingested |
| 5 | +//! documents, not a generic key/value store with hand-editable relations: |
| 6 | +//! there is no endpoint to write an arbitrary KV record, and no endpoint to |
| 7 | +//! insert a graph edge directly. So this implements exactly the one method |
| 8 | +//! that has a genuine Cognee counterpart — |
| 9 | +//! `relations`, backed by `GET /api/v1/datasets/{dataset_id}/graph` — and |
| 10 | +//! returns [`MemoryError::Other`] for every method that has none (`kv_get`, |
| 11 | +//! `kv_put`, `kv_delete`, `kv_list`, `put_relation`), rather than faking |
| 12 | +//! empty success. |
| 13 | +
|
| 14 | +use anyhow::anyhow; |
| 15 | +use async_trait::async_trait; |
| 16 | +use reqwest::Method; |
| 17 | +use serde_json::Value; |
| 18 | +use tinymemory_api::error::MemoryError; |
| 19 | +use tinymemory_api::provider::MemoryGraph; |
| 20 | +use tinymemory_api::types::{GraphRelationRecord, MemoryKvRecord}; |
| 21 | + |
| 22 | +use crate::common::{stable_id, Attempts, HttpClient}; |
| 23 | + |
| 24 | +/// Read-only relation queries over one Cognee dataset's knowledge graph. |
| 25 | +#[derive(Debug)] |
| 26 | +pub struct CogneeGraph { |
| 27 | + client: HttpClient, |
| 28 | +} |
| 29 | + |
| 30 | +impl CogneeGraph { |
| 31 | + /// Connect to the same self-hosted Cognee server a [`crate::CogneeMemory`] |
| 32 | + /// targets (`::new`/`::self_hosted`). |
| 33 | + /// |
| 34 | + /// # Errors |
| 35 | + /// |
| 36 | + /// Returns an error when `endpoint` is not an HTTP(S) URL. |
| 37 | + pub fn new(endpoint: &str, access_token: Option<&str>) -> anyhow::Result<Self> { |
| 38 | + Ok(Self { |
| 39 | + client: HttpClient::bearer(endpoint, access_token)?, |
| 40 | + }) |
| 41 | + } |
| 42 | + |
| 43 | + /// Connect to a Cognee Cloud tenant using `X-Api-Key` authentication. |
| 44 | + /// |
| 45 | + /// # Errors |
| 46 | + /// |
| 47 | + /// Returns an error when `endpoint` is invalid or `api_key` is blank. |
| 48 | + pub fn api(endpoint: &str, api_key: &str) -> anyhow::Result<Self> { |
| 49 | + anyhow::ensure!( |
| 50 | + !api_key.trim().is_empty(), |
| 51 | + "cognee API key must not be empty" |
| 52 | + ); |
| 53 | + Ok(Self { |
| 54 | + client: HttpClient::api_key(endpoint, Some(api_key))?, |
| 55 | + }) |
| 56 | + } |
| 57 | + |
| 58 | + /// Matches [`crate::cognee`]'s private `CogneeDialect::dataset_name` |
| 59 | + /// exactly, so both halves resolve one TinyMemory namespace to the same |
| 60 | + /// Cognee dataset. |
| 61 | + fn dataset_name(namespace: &str) -> String { |
| 62 | + format!("tinymemory__{}", stable_id("dataset", namespace)) |
| 63 | + } |
| 64 | + |
| 65 | + async fn find_dataset_id(&self, namespace: &str) -> anyhow::Result<Option<String>> { |
| 66 | + let name = Self::dataset_name(namespace); |
| 67 | + let response: Value = self |
| 68 | + .client |
| 69 | + .json( |
| 70 | + Method::GET, |
| 71 | + "api/v1/datasets/", |
| 72 | + None, |
| 73 | + Attempts::RetryTransient, |
| 74 | + ) |
| 75 | + .await?; |
| 76 | + Ok(response |
| 77 | + .as_array() |
| 78 | + .into_iter() |
| 79 | + .flatten() |
| 80 | + .find(|value| value.get("name").and_then(Value::as_str) == Some(name.as_str())) |
| 81 | + .and_then(|value| value.get("id").and_then(Value::as_str)) |
| 82 | + .map(str::to_owned)) |
| 83 | + } |
| 84 | +} |
| 85 | + |
| 86 | +const NO_KV_STORE: &str = "cognee has no generic key/value store to read or write"; |
| 87 | +const NO_WRITABLE_GRAPH: &str = |
| 88 | + "cognee's graph is derived by the cognify pipeline over ingested documents and cannot be edited directly"; |
| 89 | + |
| 90 | +#[async_trait] |
| 91 | +impl MemoryGraph for CogneeGraph { |
| 92 | + async fn kv_get( |
| 93 | + &self, |
| 94 | + _namespace: Option<&str>, |
| 95 | + _key: &str, |
| 96 | + ) -> Result<Option<MemoryKvRecord>, MemoryError> { |
| 97 | + Err(MemoryError::Other(anyhow!(NO_KV_STORE))) |
| 98 | + } |
| 99 | + |
| 100 | + async fn kv_put( |
| 101 | + &self, |
| 102 | + _namespace: Option<&str>, |
| 103 | + _key: &str, |
| 104 | + _value: serde_json::Value, |
| 105 | + ) -> Result<(), MemoryError> { |
| 106 | + Err(MemoryError::Other(anyhow!(NO_KV_STORE))) |
| 107 | + } |
| 108 | + |
| 109 | + async fn kv_delete(&self, _namespace: Option<&str>, _key: &str) -> Result<bool, MemoryError> { |
| 110 | + Err(MemoryError::Other(anyhow!(NO_KV_STORE))) |
| 111 | + } |
| 112 | + |
| 113 | + async fn kv_list( |
| 114 | + &self, |
| 115 | + _namespace: Option<&str>, |
| 116 | + _prefix: Option<&str>, |
| 117 | + _limit: usize, |
| 118 | + ) -> Result<Vec<MemoryKvRecord>, MemoryError> { |
| 119 | + Err(MemoryError::Other(anyhow!(NO_KV_STORE))) |
| 120 | + } |
| 121 | + |
| 122 | + /// Reads the dataset's derived graph and reshapes it into |
| 123 | + /// `(subject, predicate, object)` triples. |
| 124 | + /// |
| 125 | + /// Cognee's graph endpoint takes only a dataset id, not a subject or |
| 126 | + /// predicate filter, so this fetches the whole dataset graph and filters |
| 127 | + /// client-side. `namespace: None` ("the global, namespace-less slice") has |
| 128 | + /// no Cognee counterpart — every dataset is namespace-scoped — so it is |
| 129 | + /// rejected as invalid input rather than silently returning nothing. |
| 130 | + async fn relations( |
| 131 | + &self, |
| 132 | + namespace: Option<&str>, |
| 133 | + subject: Option<&str>, |
| 134 | + predicate: Option<&str>, |
| 135 | + limit: usize, |
| 136 | + ) -> Result<Vec<GraphRelationRecord>, MemoryError> { |
| 137 | + let namespace = namespace.ok_or_else(|| { |
| 138 | + MemoryError::Invalid( |
| 139 | + "cognee requires a namespace to resolve a dataset graph".to_string(), |
| 140 | + ) |
| 141 | + })?; |
| 142 | + let Some(dataset_id) = self.find_dataset_id(namespace).await? else { |
| 143 | + return Ok(Vec::new()); |
| 144 | + }; |
| 145 | + let graph: Value = self |
| 146 | + .client |
| 147 | + .json( |
| 148 | + Method::GET, |
| 149 | + &format!("api/v1/datasets/{dataset_id}/graph"), |
| 150 | + None, |
| 151 | + Attempts::RetryTransient, |
| 152 | + ) |
| 153 | + .await?; |
| 154 | + let nodes = graph.get("nodes").and_then(Value::as_array); |
| 155 | + let labels: std::collections::HashMap<&str, &str> = nodes |
| 156 | + .into_iter() |
| 157 | + .flatten() |
| 158 | + .filter_map(|node| { |
| 159 | + Some(( |
| 160 | + node.get("id")?.as_str()?, |
| 161 | + node.get("label")?.as_str().unwrap_or_default(), |
| 162 | + )) |
| 163 | + }) |
| 164 | + .collect(); |
| 165 | + |
| 166 | + let edges = graph.get("edges").and_then(Value::as_array); |
| 167 | + let relations = edges |
| 168 | + .into_iter() |
| 169 | + .flatten() |
| 170 | + .filter_map(|edge| { |
| 171 | + let source = edge.get("source")?.as_str()?; |
| 172 | + let target = edge.get("target")?.as_str()?; |
| 173 | + let label = edge.get("label")?.as_str().unwrap_or_default(); |
| 174 | + Some(GraphRelationRecord { |
| 175 | + namespace: Some(namespace.to_string()), |
| 176 | + subject: labels.get(source).copied().unwrap_or(source).to_string(), |
| 177 | + predicate: label.to_string(), |
| 178 | + object: labels.get(target).copied().unwrap_or(target).to_string(), |
| 179 | + attrs: Value::Null, |
| 180 | + updated_at: 0.0, |
| 181 | + evidence_count: 1, |
| 182 | + order_index: None, |
| 183 | + document_ids: Vec::new(), |
| 184 | + chunk_ids: Vec::new(), |
| 185 | + }) |
| 186 | + }) |
| 187 | + .filter(|relation| subject.is_none_or(|s| relation.subject == s)) |
| 188 | + .filter(|relation| predicate.is_none_or(|p| relation.predicate == p)) |
| 189 | + .take(limit) |
| 190 | + .collect(); |
| 191 | + Ok(relations) |
| 192 | + } |
| 193 | + |
| 194 | + async fn put_relation(&self, _relation: GraphRelationRecord) -> Result<(), MemoryError> { |
| 195 | + Err(MemoryError::Other(anyhow!(NO_WRITABLE_GRAPH))) |
| 196 | + } |
| 197 | +} |
0 commit comments