feat: Add CI, run pylint, and restructure files

Author: jimouris

.github/workflows/ci.yml

@@ -0,0 +1,49 @@
+name: CI
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+
+    - name: Create and activate virtual environment
+      run: |
+        uv venv
+        echo "$HOME/.venv/bin" >> $GITHUB_PATH
+
+    - name: Install dependencies
+      run: |
+        uv pip install -e ".[dev]"
+        uv pip install isort pylint
+
+    - name: Check code formatting with isort
+      run: |
+        uv run isort --check-only .
+
+    - name: Run pylint
+      run: |
+        uv run pylint src/nilrag test
+
+    - name: Run tests
+      run: |
+        uv run -m unittest test.rag

examples/1.init_schema_query.py

@@ -2,69 +2,70 @@
 Script to initialize schema and query to nilDB.
 """
 
-import os
-import json
-import sys
 import argparse
-import time
 import asyncio
-from nilrag.nildb_requests import NilDB, Node
+import json
+import time
 
+from nilrag.config import load_nil_db_config
 
 DEFAULT_CONFIG = "examples/nildb_config.json"
 
-async def main():
-    parser = argparse.ArgumentParser(description='Initialize schema and query for nilDB')
-    parser.add_argument('--config', type=str, default=DEFAULT_CONFIG,
-                      help=f'Path to nilDB config file (default: {DEFAULT_CONFIG})')
-    args = parser.parse_args()
 
-    # Load NilDB from JSON file if it exists
-    secret_key = None
-    if os.path.exists(args.config):
-        print(f"Loading NilDB configuration from {args.config}...")
-        with open(args.config, "r", encoding="utf-8") as f:
-            data = json.load(f)
-            secret_key = data["org_secret_key"]
-            nodes = []
-            for node_data in data["nodes"]:
-                nodes.append(
-                    Node(
-                        node_data["url"],
-                        node_data["node_id"],
-                        data["org_did"],
-                    )
-                )
-            nilDB = NilDB(nodes)
-    else:
-        print(f"Error: NilDB configuration file not found at {args.config}")
-        sys.exit(1)
+async def main():
+    """
+    Initialize schema and query for nilDB nodes.
 
-    jwts = nilDB.generate_jwt(secret_key, ttl=3600)
+    This script:
+    1. Loads the nilDB configuration from a JSON file
+    2. Generates JWT tokens for authentication
+    3. Creates a schema for storing embeddings and chunks
+    4. Creates a query for computing differences between embeddings
+    5. Updates the configuration file with the generated IDs and tokens
+    """
+    parser = argparse.ArgumentParser(
+        description="Initialize schema and query for nilDB"
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=DEFAULT_CONFIG,
+        help=f"Path to nilDB config file (default: {DEFAULT_CONFIG})",
+    )
+    args = parser.parse_args()
 
-    print(nilDB)
+    # Load NilDB configuration
+    nil_db, secret_key = load_nil_db_config(
+        args.config, require_secret_key=True
+    )
+    jwts = nil_db.generate_jwt(secret_key, ttl=3600)
+    print(nil_db)
     print()
 
     # Upload encrypted data to nilDB
     print("Initializing schema...")
     start_time = time.time()
-    schema_id = await nilDB.init_schema()
+    schema_id = await nil_db.init_schema()
     end_time = time.time()
     print(f"Schema initialized successfully in {end_time - start_time:.2f} seconds")
 
     print("Initializing query...")
     start_time = time.time()
-    diff_query_id = await nilDB.init_diff_query()
+    diff_query_id = await nil_db.init_diff_query()
     end_time = time.time()
     print(f"Query initialized successfully in {end_time - start_time:.2f} seconds")
 
+    # Update config file with new IDs and tokens
+    with open(args.config, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    for node_data, jwt in zip(data["nodes"], jwts):
+        node_data["schema_id"] = schema_id
+        node_data["diff_query_id"] = diff_query_id
+        node_data["bearer_token"] = jwt
     with open(args.config, "w", encoding="utf-8") as f:
-        for node_data, jwt in zip(data["nodes"], jwts):
-            node_data["schema_id"] = schema_id
-            node_data["diff_query_id"] = diff_query_id
-            node_data["bearer_token"] = jwt
         json.dump(data, f, indent=4)
     print("Updated nilDB configuration file with schema and query IDs.")
 
+
 if __name__ == "__main__":
     asyncio.run(main())

examples/2.data_owner_upload.py

@@ -2,89 +2,92 @@
 Script to upload data to nilDB using nilRAG.
 """
 
-import os
-import json
-import sys
 import argparse
-import time
 import asyncio
+import time
+
 import nilql
-from nilrag.util import (
-    create_chunks,
-    encrypt_float_list,
-    generate_embeddings_huggingface,
-    load_file,
-)
-from nilrag.nildb_requests import NilDB, Node
 
+from nilrag.config import load_nil_db_config
+from nilrag.util import (create_chunks, encrypt_float_list,
+                         generate_embeddings_huggingface, load_file)
 
 DEFAULT_CONFIG = "examples/nildb_config.json"
-DEFAULT_FILE_PATH = 'examples/data/20-fake.txt'
+DEFAULT_FILE_PATH = "examples/data/20-fake.txt"
+
 
 async def main():
-    parser = argparse.ArgumentParser(description='Upload data to nilDB using nilRAG')
-    parser.add_argument('--config', type=str, default=DEFAULT_CONFIG,
-                      help=f'Path to nilDB config file (default: {DEFAULT_CONFIG})')
-    parser.add_argument('--file', type=str, default=DEFAULT_FILE_PATH,
-                      help=f'Path to data file to upload (default: {DEFAULT_FILE_PATH})')
-    args = parser.parse_args()
+    """
+    Upload data to nilDB using nilRAG.
 
-    # Load NilDB from JSON file if it exists
-    if os.path.exists(args.config):
-        print(f"Loading NilDB configuration from {args.config}...")
-        with open(args.config, "r", encoding="utf-8") as f:
-            data = json.load(f)
-            nodes = []
-            for node_data in data["nodes"]:
-                nodes.append(
-                    Node(
-                        node_data["url"],
-                        node_data["node_id"],
-                        data["org_did"],
-                        node_data["bearer_token"],
-                        node_data.get("schema_id"),
-                    )
-                )
-            nilDB = NilDB(nodes)
-    else:
-        print(f"Error: NilDB configuration file not found at {args.config}")
-        sys.exit(1)
+    This script:
+    1. Loads the nilDB configuration
+    2. Initializes encryption keys for different modes
+    3. Processes the input file into chunks and embeddings
+    4. Encrypts the data using nilQL
+    5. Uploads the encrypted data to nilDB nodes
+    """
+    parser = argparse.ArgumentParser(description="Upload data to nilDB using nilRAG")
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=DEFAULT_CONFIG,
+        help=f"Path to nilDB config file (default: {DEFAULT_CONFIG})",
+    )
+    parser.add_argument(
+        "--file",
+        type=str,
+        default=DEFAULT_FILE_PATH,
+        help=f"Path to data file to upload (default: {DEFAULT_FILE_PATH})",
+    )
+    args = parser.parse_args()
 
-    print(nilDB)
+    # Load NilDB configuration
+    nil_db, _ = load_nil_db_config(
+        args.config,
+        require_bearer_token=True,
+        require_schema_id=True,
+    )
+    print(nil_db)
     print()
 
     # Initialize secret keys for different modes of operation
-    num_nodes = len(nilDB.nodes)
-    additive_key = nilql.ClusterKey.generate({'nodes': [{}] * num_nodes}, {'sum': True})
-    xor_key = nilql.ClusterKey.generate({'nodes': [{}] * num_nodes}, {'store': True})
+    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}
+    )
 
     # Load and process input file
     paragraphs = load_file(args.file)
     chunks = create_chunks(paragraphs, chunk_size=50, overlap=10)
 
     # Generate embeddings and chunks
-    print('Generating embeddings and chunks...')
+    print("Generating embeddings and chunks...")
     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 and chunks generated in {end_time - start_time:.2f} seconds!")
 
     # Encrypt chunks and embeddings
-    print('Encrypting data...')
+    print("Encrypting data...")
     start_time = time.time()
     chunks_shares = [nilql.encrypt(xor_key, chunk) for chunk in chunks]
     embeddings_shares = [
         encrypt_float_list(additive_key, embedding) for embedding in embeddings
     ]
     end_time = time.time()
-    print(f'Data encrypted in {end_time - start_time:.2f} seconds')
+    print(f"Data encrypted in {end_time - start_time:.2f} seconds")
 
     # Upload encrypted data to nilDB
-    print('Uploading data...')
+    print("Uploading data...")
     start_time = time.time()
-    await nilDB.upload_data(embeddings_shares, chunks_shares)
+    await nil_db.upload_data(embeddings_shares, chunks_shares)
     end_time = time.time()
-    print(f'Data uploaded in {end_time - start_time:.2f} seconds')
+    print(f"Data uploaded in {end_time - start_time:.2f} seconds")
+
 
 if __name__ == "__main__":
     asyncio.run(main())