cortexdb_client/
client.rs

1use tonic::codegen::InterceptedService;
2use tonic::metadata::{Ascii, MetadataValue};
3use tonic::service::Interceptor;
4use tonic::transport::{Channel, Endpoint};
5use tonic::{Request, Status};
6
7use crate::proto::admin_service_client::AdminServiceClient;
8use crate::proto::graph_rag_service_client::GraphRagServiceClient;
9use crate::proto::knowledge_graph_service_client::KnowledgeGraphServiceClient;
10use crate::proto::knowledge_service_client::KnowledgeServiceClient;
11use crate::proto::memory_service_client::MemoryServiceClient;
12use crate::proto::tools_service_client::ToolsServiceClient;
13
14/// Adds `authorization: Bearer <token>` to every request when a token is set.
15#[derive(Clone)]
16pub struct AuthInterceptor {
17    token: Option<MetadataValue<Ascii>>,
18}
19
20impl Interceptor for AuthInterceptor {
21    fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
22        if let Some(token) = &self.token {
23            request
24                .metadata_mut()
25                .insert("authorization", token.clone());
26        }
27        Ok(request)
28    }
29}
30
31type Svc = InterceptedService<Channel, AuthInterceptor>;
32
33/// One connection to a CortexDB sidecar, exposing all service clients.
34#[derive(Clone)]
35pub struct CortexClient {
36    channel: Channel,
37    interceptor: AuthInterceptor,
38}
39
40/// Builder for [`CortexClient`].
41pub struct CortexClientBuilder {
42    endpoint: String,
43    token: Option<String>,
44}
45
46impl CortexClientBuilder {
47    /// Authenticate every request with this bearer token.
48    pub fn token(mut self, token: impl Into<String>) -> Self {
49        self.token = Some(token.into());
50        self
51    }
52
53    /// Connect to the sidecar.
54    pub async fn connect(self) -> Result<CortexClient, Box<dyn std::error::Error + Send + Sync>> {
55        let channel = Endpoint::from_shared(self.endpoint)?.connect().await?;
56        let token = match self.token {
57            Some(t) => Some(format!("Bearer {t}").parse()?),
58            None => None,
59        };
60        Ok(CortexClient {
61            channel,
62            interceptor: AuthInterceptor { token },
63        })
64    }
65}
66
67impl CortexClient {
68    /// Connect without authentication.
69    pub async fn connect(
70        endpoint: impl Into<String>,
71    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
72        Self::builder(endpoint).connect().await
73    }
74
75    /// Start building a connection (set a token, then `connect`).
76    pub fn builder(endpoint: impl Into<String>) -> CortexClientBuilder {
77        CortexClientBuilder {
78            endpoint: endpoint.into(),
79            token: None,
80        }
81    }
82
83    fn svc(&self) -> Svc {
84        InterceptedService::new(self.channel.clone(), self.interceptor.clone())
85    }
86
87    /// AdminService: health + server info.
88    pub fn admin(&self) -> AdminServiceClient<Svc> {
89        AdminServiceClient::new(self.svc())
90    }
91
92    /// KnowledgeService: durable RAG knowledge save/search.
93    pub fn knowledge(&self) -> KnowledgeServiceClient<Svc> {
94        KnowledgeServiceClient::new(self.svc())
95    }
96
97    /// MemoryService: scoped agent memory.
98    pub fn memory(&self) -> MemoryServiceClient<Svc> {
99        MemoryServiceClient::new(self.svc())
100    }
101
102    /// KnowledgeGraphService: RDF triples, SPARQL, SHACL, inference, ontology.
103    pub fn graph(&self) -> KnowledgeGraphServiceClient<Svc> {
104        KnowledgeGraphServiceClient::new(self.svc())
105    }
106
107    /// GraphRagService: GraphRAG ingest/search + text search.
108    pub fn graphrag(&self) -> GraphRagServiceClient<Svc> {
109        GraphRagServiceClient::new(self.svc())
110    }
111
112    /// ToolsService: generic tool dispatch (same surface as MCP).
113    pub fn tools(&self) -> ToolsServiceClient<Svc> {
114        ToolsServiceClient::new(self.svc())
115    }
116}
cortexdb_client/client.rs

cortexdb_client/
client.rs
