nillion tests

Author: Memo Ugurbil

plugins/nillion/README.md

@@ -14,4 +14,7 @@ pip install -e plugins/nillion
 Set up your environment variables in a `.env` file:
 
 ```env
+NILLION_SECRET_KEY=your_nillion_secret_key
+NILLION_ORG_ID=your_nillion_org_id
+OPENAI_API_KEY=your_openai_api_key
 ```

plugins/nillion/nillion_game_sdk/nillion_plugin.py

@@ -1,18 +1,24 @@
 import jwt
-import requests
+import json
 import nilql
 import os
+import re
+import requests
 import time
 import uuid
 
-from jsonschema import Draft7Validator, validators
+from collections import deque, defaultdict
 from ecdsa import SECP256k1, SigningKey
+from jsonschema import Draft7Validator, validators
 from langchain_openai import ChatOpenAI
 from typing import Any, Generator
 
-# from collections import deque, defaultdict
-# import json
-# import re
+from .schemas import (
+    NillionCreateSchemaInput,
+    NillionLookupSchemaInput,
+    NillionDataUploadInput,
+    NillionDataDownloadInput,
+)
 
 class NillionPlugin:
     """
@@ -25,11 +31,6 @@ def __init__(self) -> None:
         self.name: str = "Nillion Plugin"
         self.secret_key = os.getenv("NILLION_SECRET_KEY")
         self.org_did = os.getenv("NILLION_ORG_ID")
-        self.key = None
-        self.nodes = None
-        
-    def initialize(self):
-        """Initialize the plugin"""
         if not self.secret_key:
             raise ValueError("NILLION_SECRET_KEY is not configured.")
         if not self.org_did:
@@ -39,7 +40,7 @@ def initialize(self):
 
         """Initialize config with JWTs signed with ES256K for multiple node_ids; Add cluster key."""
         response = requests.post(
-            "https://sv-sda-registration.replit.app/api/config",
+            "https://secret-vault-registration.replit.app/api/config",
             headers={
                 "Content-Type": "application/json",
             },
@@ -183,7 +184,6 @@ def create_schema(self, args: dict[str, Any]) -> tuple:
             tuple[str, dict]: The schema_uuid and the corresponding schema definition
         """
         try:
-
             validated_args = NillionCreateSchemaInput(**args)
             print(f"fn:create_schema [{validated_args.schema_description}]")
 
@@ -253,7 +253,6 @@ def create_schema(self, args: dict[str, Any]) -> tuple:
             response = llm.invoke(schema_prompt)
 
             schema = json.loads(str(response.content))
-
             schema["_id"] = str(uuid.uuid4())
             schema["owner"] = self.org_did
 
@@ -263,13 +262,14 @@ def create_schema(self, args: dict[str, Any]) -> tuple:
             print(f'fn:create_schema [{schema["_id"]}]')
             return schema["_id"], schema
         except Exception as e:
-            print(f"Error creating schema: {str(e)}")
+            print(f"Error creating schema: {e!r}")
             return None, None
 
     def data_upload(self, args: dict[str, Any]) -> list[str]:
-        """Create a schema in your privacy preserving database, called the Nillion SecretVault
-        (or nildb), based on a natural language description. Do not use this tool for any other
-        purpose.
+        """Upload specified data into your privacy preserving database, called the Nillion SecretVault
+        (or nildb), using the specified schema UUID. The data must exactly fit the requirements of
+        the desired schema itself.
+        Success will return  a list of created record UUIDs, failure is an empty list.
         Args:
             args (dict[str, Any]): Arguments containing a UUID and the data to upload.
         Returns:
@@ -322,9 +322,11 @@ def data_upload(self, args: dict[str, Any]) -> list[str]:
             return []
 
     def data_download(self, args: dict[str, Any]) -> list[dict]:
-        """Create a schema in your privacy preserving database, called the Nillion SecretVault
-        (or nildb), based on a natural language description. Do not use this tool for any other
-        purpose.
+        """Download all the data from your privacy preserving database, called the Nillion SecretVault
+        (or nildb), using the specified schema UUID. You must know the schema UUID for the remote schema
+        that you require. If you do not have the schema UUID you must use the lookup_schema action of the
+        NillionActionProvider.
+        Success will return true, whereas a failure response will return false.
         Args:
             args (dict[str, Any]): Arguments containing a target schema UUID
         Returns:

plugins/nillion/nillion_game_sdk/schemas.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, Field
+
+
+class NillionCreateSchemaInput(BaseModel):
+    """Input argument schema for schema create action."""
+
+    schema_description: str = Field(
+        ..., description="a complete description of the desired nildb schema"
+    )
+
+
+class NillionDataDownloadInput(BaseModel):
+    """Input argument schema for data download action."""
+
+    schema_uuid: str = Field(
+        description="the UUID4 obtained from the nildb_schema_lookup_tool"
+    )
+
+
+class NillionDataUploadInput(BaseModel):
+    """Input argument schema for data upload action."""
+
+    schema_uuid: str = Field(
+        description="the UUID obtained from the nillion_lookup_schema tool."
+    )
+    data_to_store: list = Field(
+        description="data to store in the database that validates against desired schema"
+    )
+
+
+class NillionLookupSchemaInput(BaseModel):
+    """Input argument schema for lookup schema action."""
+
+    schema_description: str = Field(
+        ..., description="a complete description of the desired nildb schema"
+    )

plugins/nillion/nillion_test.py

@@ -0,0 +1,62 @@
+from dotenv import load_dotenv
+from nillion_game_sdk.nillion_plugin import NillionPlugin
+
+load_dotenv()
+
+plugin = NillionPlugin()
+
+# Test Create Schema
+prompt = {
+    "schema_description": "Create me a schema to track how much I spend on my boat collection"
+}
+
+schema_id, schema = plugin.create_schema(prompt)
+
+print(schema_id, schema)
+
+# Test Lookup Schema
+prompt = {
+    "schema_description": "Find me my schema that tracks how much I spend on my boat collection"
+}
+
+my_uuid, my_schema = plugin.lookup_schema(prompt)
+
+print(my_uuid, my_schema)
+
+# Test Data Upload
+prompt = {
+    "schema_uuid": "e868cc72-42f1-488d-ab3a-c8ffdbfba9b1",
+    "data_to_store": [
+        {
+            "_id": "B0AB94E9-CC60-4DC0-9D15-8AD41B09B6C6",
+            "boat_name": "Aquatic Escape",
+            "purchase_date": "2024-10-12T21:32:49.208Z",
+            "purchase_price": {
+                "$share": "$342,077"
+            },
+            "maintenance_costs": {
+                "$share": "$394,090"
+            },
+            "insurance_cost": {
+                "$share": "$200,177"
+            },
+            "storage_fee": {
+                "$share": "$254,353"
+            }
+        },
+    ]
+}
+
+record_uuids = plugin.data_upload(prompt)
+
+print(f"{record_uuids=}")
+
+
+# Test Data Download
+prompt = {
+    "schema_uuid": "e868cc72-42f1-488d-ab3a-c8ffdbfba9b1"
+}
+
+data = plugin.data_download(prompt)
+
+print(f"{data=}")

plugins/nillion/pyproject.toml

@@ -20,7 +20,7 @@ classifiers = [
     "Intended Audience :: Developers",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
-dependencies = ["game-sdk>=0.1.1", "requests>=2.31.0"]
+dependencies = ["game-sdk>=0.1.1", "requests>=2.31.0", "py-dotenv>=0.1"]
 
 [project.urls]
 "Homepage" = "https://github.com/game-by-virtuals/game-python/plugins/nillion"