From 0e081a24819c6730c959f59d2c6a369ce349de26 Mon Sep 17 00:00:00 2001
From: "Eric O. Korman" <eric.korman@gmail.com>
Date: Sat, 29 Jun 2024 00:44:30 -0500
Subject: [PATCH] add collection.py and engine.py

---
 affine/collection.py | 43 +++++++++++++++++++++++++++++++++++++++++++
 affine/engine.py     | 24 ++++++++++++++++++++++++
 2 files changed, 67 insertions(+)
 create mode 100644 affine/collection.py
 create mode 100644 affine/engine.py

diff --git a/affine/collection.py b/affine/collection.py
new file mode 100644
index 0000000..6866967
--- /dev/null
+++ b/affine/collection.py
@@ -0,0 +1,43 @@
+from typing import get_args
+
+import numpy as np
+from affine.engine import Engine, Operation
+from pydantic import BaseModel
+
+
+def get_attribute_name_and_operation_from_key(
+    k: str,
+) -> tuple[str, Operation]:
+    s = k.split("__")
+    if len(s) == 1:
+        s.append("eq")
+    if s[1] not in get_args(Operation):
+        raise ValueError(
+            f"Operation {s[1]} not supported. Supported operations are {get_args(Operation)}"
+        )
+    return tuple(s)
+
+
+class Collection(
+    BaseModel
+):  # think of Collection as a namespace? as a table? engine should be the db
+    engine: Engine
+
+    def query(self, **kwargs):
+        pass
+
+
+class TopK(BaseModel):
+    vector: np.ndarray
+    k: int
+
+
+# Example
+class Person(Collection):
+    age: int
+    face_embedding: np.ndarray
+
+
+Person.query(age__gte=18, face_embedding=TopK(vector=np.array([1, 2, 3]), k=3))
+
+# use a global connection (like mongoengine and others?)
diff --git a/affine/engine.py b/affine/engine.py
new file mode 100644
index 0000000..e7201c5
--- /dev/null
+++ b/affine/engine.py
@@ -0,0 +1,24 @@
+from typing import Any, Literal
+
+import numpy as np
+from pydantic import BaseModel
+
+Operation = Literal["eq", "lte", "gte"]
+
+
+class Filter(BaseModel):
+    field: str
+    operation: Operation
+    value: Any
+
+
+class NumPyEngine:
+    def __init__(
+        self, dim: int, name: str = None
+    ):  # maybe add option to the init for ANN algo
+        self.data = np.empty((0, dim))
+
+    def query(
+        self, vector: np.ndarray, top_k: int, filters: list[Filter] = None
+    ):
+        pass