typesec_integrations/did/
ollama.rs1use 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
22pub struct DidOllamaClient {
24 base_url: String,
25 model: String,
26 http: Arc<dyn HttpClient>,
27}
28
29impl DidOllamaClient {
30 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 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 pub fn chat_verified_prompt(
50 &self,
51 prompt: VerifiedDidPrompt,
52 _infer: &Capability<AiCanInfer, GenericResource>,
53 read: &Capability<CanReadSensitive, GenericResource>,
54 ) -> Result<Value, DidError> {
55 let plaintext = prompt.prompt.reveal(read)?;
56 let body = json!({
57 "model": self.model,
58 "stream": false,
59 "messages": [{
60 "role": "user",
61 "content": plaintext
62 }]
63 });
64 self.http
65 .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
66 .map_err(DidError::Http)
67 }
68
69 pub fn chat_verified_prompt_bound(
71 &self,
72 prompt: VerifiedDidPrompt,
73 reply_from: Did,
74 resolver: &dyn DidResolver,
75 key_store: &dyn DidKeyStore,
76 _infer: &Capability<AiCanInfer, GenericResource>,
77 read: &Capability<CanReadSensitive, GenericResource>,
78 ) -> Result<DidEnvelope, DidError> {
79 let reply_to = prompt.subject.clone();
80 let binding = DidReplyBinding::for_prompt(&prompt);
81 let plaintext = prompt.prompt.reveal(read)?;
82 let body = json!({
83 "model": self.model,
84 "stream": false,
85 "messages": [{
86 "role": "user",
87 "content": plaintext
88 }]
89 });
90 let response = self
91 .http
92 .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
93 .map_err(DidError::Http)?;
94 let reply = ollama_reply_content(&response)?;
95 let reply_did = Did::key(sha256_tagged(
96 b"typesec-did-ollama-reply",
97 format!("{}\n{}", binding.prompt_ref.digest, reply).as_bytes(),
98 ));
99 DidEnvelope::reply(
100 reply_did, reply_from, reply_to, binding, reply, resolver, key_store,
101 )
102 }
103
104 pub fn chat_wrapped_prompt(&self, envelope: &DidEnvelope) -> Result<Value, DidError> {
106 let body = json!({
107 "model": self.model,
108 "stream": false,
109 "did_envelope": envelope
110 });
111 self.http
112 .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
113 .map_err(DidError::Http)
114 }
115}
116
117fn ollama_reply_content(response: &Value) -> Result<&str, DidError> {
118 response
119 .get("message")
120 .and_then(|message| message.get("content"))
121 .and_then(Value::as_str)
122 .ok_or(DidError::MissingOllamaReply)
123}