add label embedding field

forta-network · Dec 11, 2023 · a621f0f · a621f0f
1 parent f16191b
commit a621f0f
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/cli/commands/run/server/agent.proto b/cli/commands/run/server/agent.proto
@@ -316,6 +316,7 @@ message AlertEvent {
       bool remove = 5;
       repeated string metadata = 6;
       string uniqueKey = 7;
+      repeated float embedding = 8;
     }
 
     // Unique string to identify this class of finding,

diff --git a/cli/commands/run/server/alert.proto b/cli/commands/run/server/alert.proto
@@ -91,6 +91,7 @@ message Label {
   string label = 7;
   repeated string metadata = 8;
   string uniqueKey = 9;
+  repeated float embedding = 10;
 }
 
 message Source {

diff --git a/python-sdk/src/forta_agent/label.py b/python-sdk/src/forta_agent/label.py
@@ -31,6 +31,7 @@ def __init__(self, dict):
         self.source = LabelSource(dict.get('source')) if dict.get(
             'source') is not None else None
         self.created_at = dict.get('createdAt', dict.get('created_at'))
+        self.embedding = dict.get('embedding')
 
     def toDict(self):
         d = dict(self.__dict__, **{

diff --git a/sdk/label.ts b/sdk/label.ts
@@ -38,6 +38,7 @@ type LabelInput = {
   createdAt?: string;
   id?: string;
   uniqueKey?: string;
+  embedding?: number[];
 };
 
 export class Label {
@@ -51,7 +52,8 @@ export class Label {
     readonly source?: LabelSource,
     readonly id?: string,
     readonly createdAt?: string,
-    readonly uniqueKey?: string
+    readonly uniqueKey?: string,
+    readonly embedding?: number[]
   ) {}
 
   static from(labelInput: LabelInput) {
@@ -68,7 +70,8 @@ export class Label {
     source,
     id,
     createdAt,
-    uniqueKey
+    uniqueKey,
+    embedding
   }: LabelInput) {
     if (typeof entityType == "string") {
       entityType = ENTITY_TYPE_STRING_TO_ENUM[entityType];
@@ -95,7 +98,8 @@ export class Label {
       source,
       id,
       createdAt,
-      uniqueKey
+      uniqueKey,
+      embedding
     );
   }
 }