Skip to main content

cortex_memory_client/
lib.rs

1//! Rust client for the Cortex graph memory engine.
2//!
3//! Thin wrapper over the tonic-generated gRPC client with ergonomic convenience methods.
4//!
5//! # Example
6//! ```rust,no_run
7//! use cortex_client::CortexClient;
8//! use cortex_proto::cortex::v1::CreateNodeRequest;
9//!
10//! #[tokio::main]
11//! async fn main() -> anyhow::Result<()> {
12//!     let mut client = CortexClient::connect("http://localhost:9090").await?;
13//!
14//!     let node = client.create_node(CreateNodeRequest {
15//!         kind: "decision".into(),
16//!         title: "Use Rust for performance-critical paths".into(),
17//!         body: "Go for I/O-bound, Rust for CPU-bound.".into(),
18//!         importance: 0.8,
19//!         ..Default::default()
20//!     }).await?;
21//!
22//!     let results = client.search("language choices", 5).await?;
23//!     let briefing = client.briefing("kai").await?;
24//!
25//!     println!("Node: {}", node.id);
26//!     println!("Briefing:\n{}", briefing);
27//!     Ok(())
28//! }
29//! ```
30use cortex_proto::cortex::v1::{
31    cortex_service_client::CortexServiceClient, BriefingRequest, CreateEdgeRequest,
32    CreateNodeRequest, GetNodeRequest, HybridResultEntry, HybridSearchRequest, NodeResponse,
33    SearchResponse, SimilaritySearchRequest, StatsRequest, StatsResponse, SubgraphResponse,
34    TraverseRequest,
35};
36use tonic::transport::Channel;
37
38/// Re-export generated proto types for callers that need raw access.
39pub use cortex_proto::cortex::v1 as proto;
40
41/// A connected Cortex client.
42///
43/// Wraps the tonic gRPC client with ergonomic methods for common operations.
44/// For full proto access use the [`proto`] re-export and call [`CortexClient::inner`].
45pub struct CortexClient {
46    inner: CortexServiceClient<Channel>,
47}
48
49impl CortexClient {
50    /// Connect to a running Cortex server.
51    ///
52    /// `addr` should be a full URI, e.g. `"http://localhost:9090"`.
53    pub async fn connect(addr: impl Into<String>) -> anyhow::Result<Self> {
54        let channel = Channel::from_shared(addr.into())?.connect().await?;
55        Ok(Self {
56            inner: CortexServiceClient::new(channel),
57        })
58    }
59
60    /// Expose the raw gRPC client for full proto access.
61    pub fn inner(&mut self) -> &mut CortexServiceClient<Channel> {
62        &mut self.inner
63    }
64
65    /// Create a node. Returns the stored [`NodeResponse`].
66    pub async fn create_node(&mut self, req: CreateNodeRequest) -> anyhow::Result<NodeResponse> {
67        let resp = self.inner.create_node(req).await?;
68        Ok(resp.into_inner())
69    }
70
71    /// Get a node by ID. Returns `None` if not found.
72    pub async fn get_node(&mut self, id: &str) -> anyhow::Result<Option<NodeResponse>> {
73        match self.inner.get_node(GetNodeRequest { id: id.into() }).await {
74            Ok(resp) => Ok(Some(resp.into_inner())),
75            Err(status) if status.code() == tonic::Code::NotFound => Ok(None),
76            Err(e) => Err(e.into()),
77        }
78    }
79
80    /// Semantic similarity search. Returns scored result entries.
81    pub async fn search(&mut self, query: &str, limit: u32) -> anyhow::Result<SearchResponse> {
82        let resp = self
83            .inner
84            .similarity_search(SimilaritySearchRequest {
85                query: query.into(),
86                limit,
87                ..Default::default()
88            })
89            .await?;
90        Ok(resp.into_inner())
91    }
92
93    /// Hybrid search combining vector similarity with graph proximity.
94    ///
95    /// `anchor_ids` are node IDs that anchor the graph proximity component.
96    /// Pass an empty `Vec` for pure hybrid mode with no anchors.
97    pub async fn search_hybrid(
98        &mut self,
99        query: &str,
100        anchor_ids: Vec<String>,
101        limit: u32,
102    ) -> anyhow::Result<Vec<HybridResultEntry>> {
103        let resp = self
104            .inner
105            .hybrid_search(HybridSearchRequest {
106                query: query.into(),
107                anchor_ids,
108                limit,
109                ..Default::default()
110            })
111            .await?;
112        Ok(resp.into_inner().results)
113    }
114
115    /// Generate a rendered context briefing for an agent. Returns markdown text.
116    pub async fn briefing(&mut self, agent_id: &str) -> anyhow::Result<String> {
117        let resp = self
118            .inner
119            .get_briefing(BriefingRequest {
120                agent_id: agent_id.into(),
121                ..Default::default()
122            })
123            .await?;
124        Ok(resp.into_inner().rendered)
125    }
126
127    /// Graph traversal starting from `node_id` up to `depth` hops.
128    pub async fn traverse(
129        &mut self,
130        node_id: &str,
131        depth: u32,
132    ) -> anyhow::Result<SubgraphResponse> {
133        let resp = self
134            .inner
135            .traverse(TraverseRequest {
136                start_ids: vec![node_id.into()],
137                max_depth: depth,
138                ..Default::default()
139            })
140            .await?;
141        Ok(resp.into_inner())
142    }
143
144    /// Create an edge between two nodes. Returns the edge ID.
145    pub async fn create_edge(
146        &mut self,
147        from_id: &str,
148        to_id: &str,
149        relation: &str,
150    ) -> anyhow::Result<String> {
151        let resp = self
152            .inner
153            .create_edge(CreateEdgeRequest {
154                from_id: from_id.into(),
155                to_id: to_id.into(),
156                relation: relation.into(),
157                weight: 1.0,
158                metadata: Default::default(),
159            })
160            .await?;
161        Ok(resp.into_inner().id)
162    }
163
164    /// Get graph statistics.
165    pub async fn stats(&mut self) -> anyhow::Result<StatsResponse> {
166        let resp = self.inner.stats(StatsRequest {}).await?;
167        Ok(resp.into_inner())
168    }
169}