Skip to main content

typesec_integrations/did/
ollama.rs

1//! Ollama client that can send verified DID prompts.
2
3use std::sync::Arc;
4
5use serde_json::{Value, json};
6use typesec_core::{
7    Capability,
8    permissions::{AiCanInfer, CanReadSensitive},
9    resource::GenericResource,
10};
11
12use crate::http::{HttpClient, ReqwestHttpClient};
13
14use super::crypto::sha256_tagged;
15use super::document::DidResolver;
16use super::envelope::{DidEnvelope, DidReplyBinding};
17use super::error::DidError;
18use super::gateway::VerifiedDidPrompt;
19use super::identifier::Did;
20use super::keystore::DidKeyStore;
21
22/// Ollama client that can send verified DID prompts.
23pub struct DidOllamaClient {
24    base_url: String,
25    model: String,
26    http: Arc<dyn HttpClient>,
27}
28
29impl DidOllamaClient {
30    /// Create an Ollama client using reqwest.
31    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
32        Self::with_http(base_url, model, Arc::new(ReqwestHttpClient::new()))
33    }
34
35    /// Create an Ollama client with an injected HTTP client.
36    pub fn with_http(
37        base_url: impl Into<String>,
38        model: impl Into<String>,
39        http: Arc<dyn HttpClient>,
40    ) -> Self {
41        Self {
42            base_url: base_url.into().trim_end_matches('/').to_owned(),
43            model: model.into(),
44            http,
45        }
46    }
47
48    /// The `/api/chat` endpoint for this client's base URL.
49    fn chat_endpoint(&self) -> String {
50        format!("{}/api/chat", self.base_url)
51    }
52
53    /// A non-streaming single-user-turn chat request body.
54    fn chat_body(&self, content: &str) -> Value {
55        json!({
56            "model": self.model,
57            "stream": false,
58            "messages": [{
59                "role": "user",
60                "content": content
61            }]
62        })
63    }
64
65    /// POST a chat request and map a transport failure to [`DidError::Http`].
66    fn post_chat(&self, body: &Value) -> Result<Value, DidError> {
67        self.http
68            .post_json(&self.chat_endpoint(), &[], body)
69            .map_err(DidError::Http)
70    }
71
72    /// Reveal a verified prompt under typed authority and send it to Ollama.
73    pub fn chat_verified_prompt(
74        &self,
75        prompt: VerifiedDidPrompt,
76        _infer: &Capability<AiCanInfer, GenericResource>,
77        read: &Capability<CanReadSensitive, GenericResource>,
78    ) -> Result<Value, DidError> {
79        let plaintext = prompt.prompt.reveal(read)?;
80        self.post_chat(&self.chat_body(&plaintext))
81    }
82
83    /// Send a verified prompt to Ollama and bind the assistant reply to it.
84    pub fn chat_verified_prompt_bound(
85        &self,
86        prompt: VerifiedDidPrompt,
87        reply_from: Did,
88        resolver: &dyn DidResolver,
89        key_store: &dyn DidKeyStore,
90        _infer: &Capability<AiCanInfer, GenericResource>,
91        read: &Capability<CanReadSensitive, GenericResource>,
92    ) -> Result<DidEnvelope, DidError> {
93        let reply_to = prompt.subject.clone();
94        let binding = DidReplyBinding::for_prompt(&prompt);
95        let plaintext = prompt.prompt.reveal(read)?;
96        let response = self.post_chat(&self.chat_body(&plaintext))?;
97        let reply = ollama_reply_content(&response)?;
98        let reply_did = Did::key(sha256_tagged(
99            b"typesec-did-ollama-reply",
100            format!("{}\n{}", binding.prompt_ref.digest, reply).as_bytes(),
101        ));
102        DidEnvelope::reply(
103            reply_did, reply_from, reply_to, binding, reply, resolver, key_store,
104        )
105    }
106
107    /// Forward an already wrapped DID prompt to a DID-aware Ollama fork.
108    pub fn chat_wrapped_prompt(&self, envelope: &DidEnvelope) -> Result<Value, DidError> {
109        let body = json!({
110            "model": self.model,
111            "stream": false,
112            "did_envelope": envelope
113        });
114        self.post_chat(&body)
115    }
116}
117
118fn ollama_reply_content(response: &Value) -> Result<&str, DidError> {
119    response
120        .get("message")
121        .and_then(|message| message.get("content"))
122        .and_then(Value::as_str)
123        .ok_or(DidError::MissingOllamaReply)
124}