typesec-integrations 0.10.0

OAuth, OIDC, WorkOS, Arcade, Pydantic AI, and DID integrations for typesec
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Ollama client that can send verified DID prompts.

use std::sync::Arc;

use serde_json::{Value, json};
use typesec_core::{
    Capability,
    permissions::{AiCanInfer, CanReadSensitive},
    resource::GenericResource,
};

use crate::http::{HttpClient, ReqwestHttpClient};

use super::crypto::sha256_tagged;
use super::document::DidResolver;
use super::envelope::{DidEnvelope, DidReplyBinding};
use super::error::DidError;
use super::gateway::VerifiedDidPrompt;
use super::identifier::Did;
use super::keystore::DidKeyStore;

/// Ollama client that can send verified DID prompts.
pub struct DidOllamaClient {
    base_url: String,
    model: String,
    http: Arc<dyn HttpClient>,
}

impl DidOllamaClient {
    /// Create an Ollama client using reqwest.
    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
        Self::with_http(base_url, model, Arc::new(ReqwestHttpClient::new()))
    }

    /// Create an Ollama client with an injected HTTP client.
    pub fn with_http(
        base_url: impl Into<String>,
        model: impl Into<String>,
        http: Arc<dyn HttpClient>,
    ) -> Self {
        Self {
            base_url: base_url.into().trim_end_matches('/').to_owned(),
            model: model.into(),
            http,
        }
    }

    /// Reveal a verified prompt under typed authority and send it to Ollama.
    pub fn chat_verified_prompt(
        &self,
        prompt: VerifiedDidPrompt,
        _infer: &Capability<AiCanInfer, GenericResource>,
        read: &Capability<CanReadSensitive, GenericResource>,
    ) -> Result<Value, DidError> {
        let plaintext = prompt.prompt.reveal(read)?;
        let body = json!({
            "model": self.model,
            "stream": false,
            "messages": [{
                "role": "user",
                "content": plaintext
            }]
        });
        self.http
            .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
            .map_err(DidError::Http)
    }

    /// Send a verified prompt to Ollama and bind the assistant reply to it.
    pub fn chat_verified_prompt_bound(
        &self,
        prompt: VerifiedDidPrompt,
        reply_from: Did,
        resolver: &dyn DidResolver,
        key_store: &dyn DidKeyStore,
        _infer: &Capability<AiCanInfer, GenericResource>,
        read: &Capability<CanReadSensitive, GenericResource>,
    ) -> Result<DidEnvelope, DidError> {
        let reply_to = prompt.subject.clone();
        let binding = DidReplyBinding::for_prompt(&prompt);
        let plaintext = prompt.prompt.reveal(read)?;
        let body = json!({
            "model": self.model,
            "stream": false,
            "messages": [{
                "role": "user",
                "content": plaintext
            }]
        });
        let response = self
            .http
            .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
            .map_err(DidError::Http)?;
        let reply = ollama_reply_content(&response)?;
        let reply_did = Did::key(sha256_tagged(
            b"typesec-did-ollama-reply",
            format!("{}\n{}", binding.prompt_ref.digest, reply).as_bytes(),
        ));
        DidEnvelope::reply(
            reply_did, reply_from, reply_to, binding, reply, resolver, key_store,
        )
    }

    /// Forward an already wrapped DID prompt to a DID-aware Ollama fork.
    pub fn chat_wrapped_prompt(&self, envelope: &DidEnvelope) -> Result<Value, DidError> {
        let body = json!({
            "model": self.model,
            "stream": false,
            "did_envelope": envelope
        });
        self.http
            .post_json(&format!("{}/api/chat", self.base_url), &[], &body)
            .map_err(DidError::Http)
    }
}

fn ollama_reply_content(response: &Value) -> Result<&str, DidError> {
    response
        .get("message")
        .and_then(|message| message.get("content"))
        .and_then(Value::as_str)
        .ok_or(DidError::MissingOllamaReply)
}