cortex_memory_client/
lib.rs1use 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
38pub use cortex_proto::cortex::v1 as proto;
40
41pub struct CortexClient {
46 inner: CortexServiceClient<Channel>,
47}
48
49impl CortexClient {
50 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 pub fn inner(&mut self) -> &mut CortexServiceClient<Channel> {
62 &mut self.inner
63 }
64
65 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 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 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 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 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 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 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 pub async fn stats(&mut self) -> anyhow::Result<StatsResponse> {
166 let resp = self.inner.stats(StatsRequest {}).await?;
167 Ok(resp.into_inner())
168 }
169}