From b9661e0f6f292c22b7b975bd2e1dadc866b8fe4e Mon Sep 17 00:00:00 2001
From: Grinion <Grinion@users.noreply.github.com>
Date: Thu, 5 Dec 2024 16:08:11 +0200
Subject: [PATCH] Add snapshot module to create an airdrop

---
 agent_server.py                                 |  2 ++
 injective_functions/factory.py                  |  2 ++
 injective_functions/function_schemas.json       |  9 +++++++++
 injective_functions/functions_schemas.json      |  9 +++++++++
 injective_functions/snapshot/__init__.py        | 17 +++++++++++++++++
 .../snapshot/snapshot_schema.json               | 13 +++++++++++++
 injective_functions/utils/function_helper.py    |  2 ++
 7 files changed, 54 insertions(+)
 create mode 100644 injective_functions/snapshot/__init__.py
 create mode 100644 injective_functions/snapshot/snapshot_schema.json

diff --git a/agent_server.py b/agent_server.py
index 8364a5f..d5fbbd7 100644
--- a/agent_server.py
+++ b/agent_server.py
@@ -5,6 +5,7 @@
 from datetime import datetime
 import argparse
 from injective_functions.factory import InjectiveClientFactory
+from injective_functions.snapshot import Snapshot
 from injective_functions.utils.function_helper import (
     FunctionSchemaLoader,
     FunctionExecutor,
@@ -46,6 +47,7 @@ def __init__(self):
             "./injective_functions/exchange/exchange_schema.json",
             "./injective_functions/staking/staking_schema.json",
             "./injective_functions/token_factory/token_factory_schema.json",
+            "./injective_functions/snapshot/snapshot_schema.json",
             "./injective_functions/utils/utils_schema.json",
         ]
         self.function_schemas = FunctionSchemaLoader.load_schemas(schema_paths)
diff --git a/injective_functions/factory.py b/injective_functions/factory.py
index 003e7d1..d5b608b 100644
--- a/injective_functions/factory.py
+++ b/injective_functions/factory.py
@@ -8,6 +8,7 @@
 from injective_functions.exchange.trader import InjectiveTrading
 from injective_functions.staking import InjectiveStaking
 from injective_functions.token_factory import InjectiveTokenFactory
+from injective_functions.snapshot import Snapshot
 
 
 class InjectiveClientFactory:
@@ -41,6 +42,7 @@ async def create_all(private_key: str, network_type: str = "mainnet") -> Dict:
             "trader": InjectiveTrading(chain_client),
             "staking": InjectiveStaking(chain_client),
             "token_factory": InjectiveTokenFactory(chain_client),
+            "snapshot": Snapshot(),
         }
         print(clients)
         return clients
diff --git a/injective_functions/function_schemas.json b/injective_functions/function_schemas.json
index 81b8721..bb99753 100644
--- a/injective_functions/function_schemas.json
+++ b/injective_functions/function_schemas.json
@@ -124,6 +124,15 @@
           },
           "required": ["to_address", "amount"]
         }
+      },
+      {
+        "name": "get_snapshot",
+        "description": "Get a snapshot to create an airdrop",
+        "parameters": {
+          "type": "object",
+          "properties": {},
+          "required": []
+        }
       }
     ]
   }
\ No newline at end of file
diff --git a/injective_functions/functions_schemas.json b/injective_functions/functions_schemas.json
index 891eabf..2510b82 100644
--- a/injective_functions/functions_schemas.json
+++ b/injective_functions/functions_schemas.json
@@ -906,6 +906,15 @@
           "uri_hash"
         ]
       }
+    },
+    {
+      "name": "get_snapshot",
+      "description": "Get a snapshot to create an airdrop",
+      "parameters": {
+        "type": "object",
+        "properties": {},
+        "required": []
+      }
     }
   ]
 }
\ No newline at end of file
diff --git a/injective_functions/snapshot/__init__.py b/injective_functions/snapshot/__init__.py
new file mode 100644
index 0000000..5dd0281
--- /dev/null
+++ b/injective_functions/snapshot/__init__.py
@@ -0,0 +1,17 @@
+import json
+from typing import Dict, List
+import requests
+
+"""This class handles all the functions related to the snapshot module"""
+
+
+class Snapshot():
+    def __init__(self) -> None:
+        # Initializes the network and the composer
+        super().__init__()
+
+    async def get_snapshot(self) -> List[Dict]:
+        data = requests.get(
+            "https://tokenstation.app/api/snapshots/last").json()
+
+        return data
diff --git a/injective_functions/snapshot/snapshot_schema.json b/injective_functions/snapshot/snapshot_schema.json
new file mode 100644
index 0000000..b30f1c5
--- /dev/null
+++ b/injective_functions/snapshot/snapshot_schema.json
@@ -0,0 +1,13 @@
+{
+  "functions": [
+    {
+      "name": "get_snapshot",
+      "description": "Get a snapshot to create an airdrop",
+      "parameters": {
+        "type": "object",
+        "properties": {},
+        "required": []
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/injective_functions/utils/function_helper.py b/injective_functions/utils/function_helper.py
index 74d947c..843f1bf 100644
--- a/injective_functions/utils/function_helper.py
+++ b/injective_functions/utils/function_helper.py
@@ -56,6 +56,8 @@ class InjectiveFunctionMapper:
         "mint": ("token_factory", "mint"),
         "burn": ("token_factory", "burn"),
         "set_denom_metadata": ("token_factory", "set_denom_metadata"),
+
+        "get_snapshot": ("snapshot", "get_snapshot"),
     }
 
     @classmethod