misc fixes

.env.example

@@ -0,0 +1,6 @@
+# nilDB organization DID and secret key
+NILDB_ORG_DID=YOUR_ORG_DID
+NILDB_ORG_SECRET_KEY=YOUR_ORG_SECRET_KEY
+
+# nilAI API token
+NILAI_API_KEY=Nillion2025

README.md

@@ -105,16 +105,16 @@ pip install nilrag
 ## Data owner
 
 ### 1. Initialization
-First, copy the sample nilDB config file:
+First, copy the example nilDB `.env` and nilDB config:
+
 ```shell
 cp ./examples/nildb_config.sample.json ./examples/nildb_config.json
+cp .env.example .env
 ```
+
 Next, register a new organization in Nillion's [SecretVault Registration
 Portal](https://sv-sda-registration.replit.app/) and fill in the details in your
-nilDB config file `./examples/nildb_config.json`.
-
-You can safely ignore `bearer_token`, `schema_id`, and `diff_query_id` as we'll
-fill these out later.
+`.env` file.
 
 You are all set to create your first schema and query for RAG. At the minimum,
 they should look like:
@@ -132,9 +132,9 @@ uv run examples/1.init_schema_query.py
 # Or specify a custom config file
 uv run examples/1.init_schema_query.py --config /path/to/your/config.json
 ```
-This, will fill out `bearer_token`, `schema_id`, and `diff_query_id` in your
-config file. Verify that it has been populated successfully.
 
+This will fill out `bearer_token`, `schema_id`, and `diff_query_id` in your
+config file. Verify that it has been populated successfully.
 
 ### 2. Uploading Documents
 After initialization, the data owner can upload their documents to the nilDB
@@ -170,3 +170,14 @@ uv run examples/3.client_query.py --config /path/to/config.json --prompt "Your c
 # Run a specific test file
 uv run -m unittest test.rag
 ```
+
+## Note
+
+When querying nilRAG on data which is public knowledge, it is not always evident if the answer
+you get back comes from the RAG itself or from the model's general knowledge. For example, if
+you query the nilRAG which uses the provided `examples/data/20-fake.txt` dataset about a famous
+person, it will reply with details of the famous person even though that data is not present in
+the `examples/data/20-fake.txt` dataset. This is easily fixed by using a prompt which includes an
+instruction like:
+
+> If you don't know, just say you don't know.

examples/1.init_schema_query.py

@@ -35,13 +35,22 @@ async def main():
     args = parser.parse_args()
 
     # Load NilDB configuration
+    start_time = time.time()
     nil_db, secret_key = load_nil_db_config(
         args.config, require_secret_key=True
     )
-    jwts = nil_db.generate_jwt(secret_key, ttl=3600)
+    end_time = time.time()
+    print(f"Config loaded successfully in {end_time - start_time:.2f} seconds")
     print(nil_db)
     print()
 
+    # Generate tokens
+    print("Generating JWTs...")
+    start_time = time.time()
+    jwts = nil_db.generate_jwt(secret_key, ttl=3600)
+    end_time = time.time()
+    print(f"JWTs generated successfully in {end_time - start_time:.2f} seconds")
+
     # Upload encrypted data to nilDB
     print("Initializing schema...")
     start_time = time.time()

examples/2.data_owner_upload.py

@@ -43,33 +43,44 @@ async def main():
     args = parser.parse_args()
 
     # Load NilDB configuration
+    start_time = time.time()
     nil_db, _ = load_nil_db_config(
         args.config,
         require_bearer_token=True,
         require_schema_id=True,
     )
+    end_time = time.time()
+    print(f"Config loaded successfully in {end_time - start_time:.2f} seconds")
     print(nil_db)
     print()
 
     # Initialize secret keys for different modes of operation
+    print("Generating cluster keys...")
+    start_time = time.time()
     num_nodes = len(nil_db.nodes)
     additive_key = nilql.ClusterKey.generate(
         {"nodes": [{}] * num_nodes}, {"sum": True}
     )
     xor_key = nilql.ClusterKey.generate(
         {"nodes": [{}] * num_nodes}, {"store": True}
     )
+    end_time = time.time()
+    print(f"Cluster keys generated successfully in {end_time - start_time:.2f} seconds")
 
     # Load and process input file
+    print("Loading and chunking input file...")
+    start_time = time.time()
     paragraphs = load_file(args.file)
     chunks = create_chunks(paragraphs, chunk_size=50, overlap=10)
+    end_time = time.time()
+    print(f"Input file chunked successfully in {end_time - start_time:.2f} seconds")
 
     # Generate embeddings and chunks
-    print("Generating embeddings and chunks...")
+    print("Generating embeddings...")
     start_time = time.time()
     embeddings = generate_embeddings_huggingface(chunks)
     end_time = time.time()
-    print(f"Embeddings and chunks generated in {end_time - start_time:.2f} seconds!")
+    print(f"Embeddings generated in {end_time - start_time:.2f} seconds!")
 
     # Encrypt chunks and embeddings
     print("Encrypting data...")

examples/3.client_query.py

@@ -6,6 +6,7 @@
 import asyncio
 import json
 import time
+import os
 
 from nilrag.config import load_nil_db_config
 from nilrag.nildb_requests import ChatCompletionConfig
@@ -40,20 +41,29 @@ async def main():
     args = parser.parse_args()
 
     # Load NilDB configuration
+    start_time = time.time()
     nil_db, _ = load_nil_db_config(
         args.config,
         require_bearer_token=True,
         require_schema_id=True,
         require_diff_query_id=True,
     )
+    end_time = time.time()
+    print(f"Config loaded successfully in {end_time - start_time:.2f} seconds")
     print(nil_db)
     print()
 
-    print("Query nilAI with nilRAG...")
+    print("Using prompt:", args.prompt)
+    token = os.getenv("NILAI_API_KEY")
+
+    if not token:
+        raise ValueError("No NILAI_API_KEY supplied")
+
+    print("Querying nilAI with nilRAG...")
     start_time = time.time()
     config = ChatCompletionConfig(
         nilai_url="https://nilai-a779.nillion.network",
-        token="Nillion2025",
+        token=token,
         messages=[{"role": "user", "content": args.prompt}],
         model="meta-llama/Llama-3.1-8B-Instruct",
         temperature=0.2,

examples/nildb_config.sample.json

@@ -1,27 +1,25 @@
 {
-    "nodes": [
-        {
-            "url": "https://nildb-nx8v.nillion.network/api/v1",
-            "node_id": "did:nil:testnet:nillion1qfrl8nje3nvwh6cryj63mz2y6gsdptvn07nx8v",
-            "bearer_token": "",
-            "schema_id": "",
-            "diff_query_id": ""
-        },
-        {
-            "url": "https://nildb-p3mx.nillion.network/api/v1",
-            "node_id": "did:nil:testnet:nillion1uak7fgsp69kzfhdd6lfqv69fnzh3lprg2mp3mx",
-            "bearer_token": "",
-            "schema_id": "",
-            "diff_query_id": ""
-        },
-        {
-            "url": "https://nildb-rugk.nillion.network/api/v1",
-            "node_id": "did:nil:testnet:nillion1kfremrp2mryxrynx66etjl8s7wazxc3rssrugk",
-            "bearer_token": "",
-            "schema_id": "",
-            "diff_query_id": ""
-        }
-    ],
-    "org_secret_key": "Add you SECRET_KEY here",
-    "org_did": "Add your ORG DID here"
-}
+  "nodes": [
+    {
+      "url": "https://nildb-nx8v.nillion.network/api/v1",
+      "node_id": "did:nil:testnet:nillion1qfrl8nje3nvwh6cryj63mz2y6gsdptvn07nx8v",
+      "bearer_token": "",
+      "schema_id": "",
+      "diff_query_id": ""
+    },
+    {
+      "url": "https://nildb-p3mx.nillion.network/api/v1",
+      "node_id": "did:nil:testnet:nillion1uak7fgsp69kzfhdd6lfqv69fnzh3lprg2mp3mx",
+      "bearer_token": "",
+      "schema_id": "",
+      "diff_query_id": ""
+    },
+    {
+      "url": "https://nildb-rugk.nillion.network/api/v1",
+      "node_id": "did:nil:testnet:nillion1kfremrp2mryxrynx66etjl8s7wazxc3rssrugk",
+      "bearer_token": "",
+      "schema_id": "",
+      "diff_query_id": ""
+    }
+  ]
+}

pyproject.toml

@@ -11,6 +11,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "aiohttp>=3.11.14",
+    "dotenv>=0.9.9",
     "ecdsa>=0.19.0",
     "faker>=37.0.0",
     "nilql==0.0.0a12",

src/nilrag/config.py

@@ -6,8 +6,10 @@
 import os
 from typing import Optional, Tuple
 
-from nilrag.nildb_requests import NilDB, Node
+from dotenv import load_dotenv
+_ = load_dotenv()
 
+from nilrag.nildb_requests import NilDB, Node
 
 def load_nil_db_config(
     config_path: str,
@@ -44,11 +46,15 @@ def load_nil_db_config(
         raise ValueError(f'Error: Invalid JSON in configuration file {config_path}') from exc
 
     # Get secret key if required
-    secret_key = None
-    if require_secret_key:
-        if "org_secret_key" not in data:
-            raise ValueError("Error: org_secret_key not found in configuration")
-        secret_key = data["org_secret_key"]
+    secret_key = os.getenv("NILDB_ORG_SECRET_KEY")
+    if require_secret_key and not secret_key:
+        raise ValueError("Error: org_secret_key not found in NILDB_ORG_SECRET_KEY envvar")
+
+    # Get org DID
+    org_did = os.getenv("NILDB_ORG_DID")
+
+    if not org_did:
+        raise ValueError("Error: org_did not found in NILDB_ORG_DID envvar")
 
     # Create nodes
     nodes = []
@@ -65,7 +71,7 @@ def load_nil_db_config(
         node = Node(
             url=node_data["url"],
             node_id=node_data.get("node_id"),
-            org=data.get("org_did"),
+            org=org_did,
             bearer_token=node_data.get("bearer_token"),
             schema_id=node_data.get("schema_id"),
             diff_query_id=node_data.get("diff_query_id"),