Tool use dsrs

Cargo.toml

@@ -21,6 +21,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tower-http = { version = "0.5.2", features = ["cors"] }
 thiserror = "1.0.63"
 async-trait = "0.1.81"
+anyhow = "1.0"
 jsonwebtoken = "9.3.0"
 jwt-compact = { version = "0.9.0-beta.1", features = ["es256k"] }
 diesel = { version = "=2.2.2", features = [
@@ -75,3 +76,7 @@ lazy_static = "1.4.0"
 subtle = "2.6.1"
 tiktoken-rs = "0.5"
 once_cell = "1.19"
+
+# DSPy for structured prompting and optimization
+# Using main branch since v0.6.0 tag doesn't exist yet
+dspy-rs = { git = "https://github.com/krypticmouse/DSRs", branch = "main" }

docs/nitro-deploy.md

@@ -830,6 +830,54 @@ A restart should not be needed but if you need to:
 sudo systemctl restart vsock-billing-proxy.service
 ```
 
+## Vsock Kagi Search proxy
+Create a vsock proxy service so that enclave program can talk to the Kagi Search API:
+
+First configure the endpoint into its allowlist:
+
+```sh
+sudo vim /etc/nitro_enclaves/vsock-proxy.yaml
+```
+
+Add this line:
+```
+- {address: kagi.com, port: 443}
+```
+
+Now create a service that spins this up automatically:
+
+```sh
+sudo vim /etc/systemd/system/vsock-kagi-proxy.service
+```
+
+```
+[Unit]
+Description=Vsock Kagi Search Proxy Service
+After=network.target
+
+[Service]
+User=root
+ExecStart=/usr/bin/vsock-proxy 8026 kagi.com 443
+Restart=always
+
+[Install]
+WantedBy=multi-user.target
+```
+
+Activate the service:
+
+```sh
+sudo systemctl daemon-reload
+sudo systemctl enable vsock-kagi-proxy.service
+sudo systemctl start vsock-kagi-proxy.service
+sudo systemctl status vsock-kagi-proxy.service
+```
+
+A restart should not be needed but if you need to:
+```sh
+sudo systemctl restart vsock-kagi-proxy.service
+```
+
 ## Vsock Tinfoil proxies
 Create vsock proxy services so that tinfoil-proxy can talk to Tinfoil services:
 
@@ -1319,6 +1367,26 @@ INSERT INTO enclave_secrets (key, value)
 VALUES ('billing_server_url', decode('your_base64_string', 'base64'));
 ```
 
+#### Kagi API Key
+
+After the DB is initialized, we need to store the Kagi Search API key encrypted to the enclave KMS key.
+
+```sh
+echo -n "KAGI_API_KEY" | base64 -w 0
+```
+
+Take that output and encrypt to the KMS key, from a machine that has encrypt access to the key:
+
+```sh
+aws kms encrypt --key-id "KEY_ARN" --plaintext "BASE64_KEY" --query CiphertextBlob --output text
+```
+
+Take that encrypted base64 and insert it into the `enclave_secrets` table with key as `kagi_api_key` and value as the base64.
+
+```sql
+INSERT INTO enclave_secrets (key, value)
+VALUES ('kagi_api_key', decode('your_base64_string', 'base64'));
+```
 
 ## Secrets Manager
 

entrypoint.sh

@@ -284,6 +284,10 @@ echo "127.0.0.21 doc-upload.model.tinfoil.sh" >> /etc/hosts
 echo "127.0.0.22 inference.tinfoil.sh" >> /etc/hosts
 log "Added Tinfoil proxy domains to /etc/hosts"
 
+# Add Kagi Search hostname to /etc/hosts
+echo "127.0.0.23 kagi.com" >> /etc/hosts
+log "Added Kagi Search domain to /etc/hosts"
+
 touch /app/libnsm.so
 log "Created /app/libnsm.so"
 
@@ -376,6 +380,10 @@ python3 /app/traffic_forwarder.py 127.0.0.21 443 3 8024 &
 log "Starting Tinfoil Inference traffic forwarder"
 python3 /app/traffic_forwarder.py 127.0.0.22 443 3 8025 &
 
+# Start the traffic forwarder for Kagi Search in the background
+log "Starting Kagi Search traffic forwarder"
+python3 /app/traffic_forwarder.py 127.0.0.23 443 3 8026 &
+
 # Wait for the forwarders to start
 log "Waiting for forwarders to start"
 sleep 5
@@ -539,6 +547,14 @@ else
     log "Tinfoil Inference connection failed"
 fi
 
+# Test the connection to Kagi Search
+log "Testing connection to Kagi Search:"
+if timeout 5 bash -c '</dev/tcp/127.0.0.23/443'; then
+    log "Kagi Search connection successful"
+else
+    log "Kagi Search connection failed"
+fi
+
 # Start the continuum-proxy if we're in AWS Nitro mode
 if [ "$APP_MODE" != "local" ]; then
     # Get Continuum Proxy API key from Secrets Manager

pcrDev.json

@@ -1,6 +1,6 @@
 {
   "HashAlgorithm": "Sha384 { ... }",
-  "PCR0": "a771d516c0373f30ae15b3296d188e376345d6d21294b6cc7c5a9d7c10e9989fcd1cf5ff45043644caac19e1644d768d",
+  "PCR0": "5d05893f150ce0237d4dbea01f83d1aed7ba8eeaa4239ce9ac17286ad4548cd5aab757aa20c15bd028703b6c3b8ff048",
   "PCR1": "f004075c672258b499f8e88d59701031a3b451f65c7de60c81d09da2b0799272675481ec390527594dd7069cb7de59d7",
-  "PCR2": "328f9ceecaa1bac78021274a5dd3f977f2c9ea4c96b1bf70d89b94f0958951da33a47881e2113ec27ae3072c05296f63"
+  "PCR2": "22485f598f10762f0a932469c6bbd9a1cd867da14703d7d5031e9f8f591c535ca80f9676706cea27768098297a9e2eb8"
 }

pcrDevHistory.json

@@ -348,5 +348,12 @@
     "PCR2": "328f9ceecaa1bac78021274a5dd3f977f2c9ea4c96b1bf70d89b94f0958951da33a47881e2113ec27ae3072c05296f63",
     "timestamp": 1760404411,
     "signature": "A/KDCzrSkDhl1P8RQK3DDTjaY4d/QS3sy9Lu1d9/Xt8DKLd/WuGdbNgm3e/FptvgUp01KW76BZ3m/CuGQ+AG+xBnRVIA9Wc8XZEApkMBgYH7EKkHvN03pvkSwgjl69tu"
+  },
+  {
+    "PCR0": "5d05893f150ce0237d4dbea01f83d1aed7ba8eeaa4239ce9ac17286ad4548cd5aab757aa20c15bd028703b6c3b8ff048",
+    "PCR1": "f004075c672258b499f8e88d59701031a3b451f65c7de60c81d09da2b0799272675481ec390527594dd7069cb7de59d7",
+    "PCR2": "22485f598f10762f0a932469c6bbd9a1cd867da14703d7d5031e9f8f591c535ca80f9676706cea27768098297a9e2eb8",
+    "timestamp": 1760667548,
+    "signature": "asao0r4Zq/fxkkZEi8ybubHU5IPzA70Fk2Um6KdcFgDFoV8Xyt7GUqSGFfozUHqdNL9HwR3116ihI9TQpXkFr7Niwuk52zonH9MQphZUYG8zWmp9WU2cidxq0LfFFTI9"
   }
 ]

pcrProd.json

@@ -1,6 +1,6 @@
 {
   "HashAlgorithm": "Sha384 { ... }",
-  "PCR0": "858fe55e3736da573ab19719b92df9793b106451fcb1750b35e46b788f58ce294a1048795723f57628ec2347ebd0c3a9",
+  "PCR0": "c2d7a330881cd2f394d85e3a04def8714d611138560785806f116ea52e8748fcf6f2453bff89977faa9b1c20af6e8778",
   "PCR1": "f004075c672258b499f8e88d59701031a3b451f65c7de60c81d09da2b0799272675481ec390527594dd7069cb7de59d7",
-  "PCR2": "82c53aa649f676766a839e93699ecca83f1207dd9ab59590dfba33149cd274dc59afe10ea1af0feaa5aec1f4d840efd3"
+  "PCR2": "97845cf3189f0da75d5bf893c91461d3cc8f19c2f8955c4e097f7b69ffe4c18dfc8d55590d503cb73a4dcd64ba225c90"
 }

pcrProdHistory.json

@@ -348,5 +348,12 @@
     "PCR2": "82c53aa649f676766a839e93699ecca83f1207dd9ab59590dfba33149cd274dc59afe10ea1af0feaa5aec1f4d840efd3",
     "timestamp": 1760404429,
     "signature": "/obyBdTdhNbUs59Kz9iOfffAt0Fdmoy+yAwIw9mPpDsrsMcIjj0b26gyq3ynhxFL+APY1e8jw5iOAVCHWYyh1PLJoye5Tg2f4Q8j+eoeQ22Th4Q/BW/qdqU4933SdVtr"
+  },
+  {
+    "PCR0": "c2d7a330881cd2f394d85e3a04def8714d611138560785806f116ea52e8748fcf6f2453bff89977faa9b1c20af6e8778",
+    "PCR1": "f004075c672258b499f8e88d59701031a3b451f65c7de60c81d09da2b0799272675481ec390527594dd7069cb7de59d7",
+    "PCR2": "97845cf3189f0da75d5bf893c91461d3cc8f19c2f8955c4e097f7b69ffe4c18dfc8d55590d503cb73a4dcd64ba225c90",
+    "timestamp": 1760667566,
+    "signature": "JQ6m7fjhIcy0CAWqzuIvoHYZkyojDHjdo/toYALN8plqoDBRS/qUqREWLe4oUVuJmqa969DvTxw7Y88DuW+Zm9Ibw/wI092UXStjAnwhQqI63B3MFQxYPPC1gAACW9jw"
   }
 ]

src/db.rs

@@ -572,6 +572,7 @@ pub trait DBConnection {
 
     // Tool calls / outputs
     fn create_tool_call(&self, new_call: NewToolCall) -> Result<ToolCall, DBError>;
+    fn get_tool_call_by_uuid(&self, uuid: Uuid, user_id: Uuid) -> Result<ToolCall, DBError>;
     fn create_tool_output(&self, new_output: NewToolOutput) -> Result<ToolOutput, DBError>;
 
     // Context reconstruction
@@ -2340,6 +2341,12 @@ impl DBConnection for PostgresConnection {
         new_call.insert(conn).map_err(DBError::from)
     }
 
+    fn get_tool_call_by_uuid(&self, uuid: Uuid, user_id: Uuid) -> Result<ToolCall, DBError> {
+        debug!("Getting tool call by UUID: {} for user: {}", uuid, user_id);
+        let conn = &mut self.db.get().map_err(|_| DBError::ConnectionError)?;
+        ToolCall::get_by_uuid(conn, uuid, user_id).map_err(DBError::from)
+    }
+
     fn create_tool_output(&self, new_output: NewToolOutput) -> Result<ToolOutput, DBError> {
         debug!("Creating new tool output");
         let conn = &mut self.db.get().map_err(|_| DBError::ConnectionError)?;