about summary refs log tree commit diff
path: root/R2R/r2r/base/abstractions/vector.py
diff options
context:
space:
mode:
Diffstat (limited to 'R2R/r2r/base/abstractions/vector.py')
-rwxr-xr-xR2R/r2r/base/abstractions/vector.py66
1 files changed, 66 insertions, 0 deletions
diff --git a/R2R/r2r/base/abstractions/vector.py b/R2R/r2r/base/abstractions/vector.py
new file mode 100755
index 00000000..445f3302
--- /dev/null
+++ b/R2R/r2r/base/abstractions/vector.py
@@ -0,0 +1,66 @@
+"""Abstraction for a vector that can be stored in the system."""
+
+from enum import Enum
+from typing import Any
+from uuid import UUID
+
+
+class VectorType(Enum):
+    FIXED = "FIXED"
+
+
+class Vector:
+    """A vector with the option to fix the number of elements."""
+
+    def __init__(
+        self,
+        data: list[float],
+        type: VectorType = VectorType.FIXED,
+        length: int = -1,
+    ):
+        self.data = data
+        self.type = type
+        self.length = length
+
+        if (
+            self.type == VectorType.FIXED
+            and length > 0
+            and len(data) != length
+        ):
+            raise ValueError(f"Vector must be exactly {length} elements long.")
+
+    def __repr__(self) -> str:
+        return (
+            f"Vector(data={self.data}, type={self.type}, length={self.length})"
+        )
+
+
+class VectorEntry:
+    """A vector entry that can be stored directly in supported vector databases."""
+
+    def __init__(self, id: UUID, vector: Vector, metadata: dict[str, Any]):
+        """Create a new VectorEntry object."""
+        self.vector = vector
+        self.id = id
+        self.metadata = metadata
+
+    def to_serializable(self) -> str:
+        """Return a serializable representation of the VectorEntry."""
+        metadata = self.metadata
+
+        for key in metadata:
+            if isinstance(metadata[key], UUID):
+                metadata[key] = str(metadata[key])
+        return {
+            "id": str(self.id),
+            "vector": self.vector.data,
+            "metadata": metadata,
+        }
+
+    def __str__(self) -> str:
+        """Return a string representation of the VectorEntry."""
+        return f"VectorEntry(id={self.id}, vector={self.vector}, metadata={self.metadata})"
+
+    def __repr__(self) -> str:
+        """Return an unambiguous string representation of the VectorEntry."""
+        return f"VectorEntry(id={self.id}, vector={self.vector}, metadata={self.metadata})"