refactor and add tests, v0.2.0

src/embeddingbuddy/models/schemas.py

@@ -0,0 +1,58 @@
+from typing import List, Optional, Any, Dict
+from dataclasses import dataclass
+import numpy as np
+
+
+@dataclass
+class Document:
+    id: str
+    text: str
+    embedding: List[float]
+    category: Optional[str] = None
+    subcategory: Optional[str] = None
+    tags: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.tags is None:
+            self.tags = []
+        if self.category is None:
+            self.category = "Unknown"
+        if self.subcategory is None:
+            self.subcategory = "Unknown"
+
+
+@dataclass
+class ProcessedData:
+    documents: List[Document]
+    embeddings: np.ndarray
+    error: Optional[str] = None
+
+    def __post_init__(self):
+        if self.embeddings is not None and not isinstance(self.embeddings, np.ndarray):
+            self.embeddings = np.array(self.embeddings)
+
+
+@dataclass
+class ReducedData:
+    reduced_embeddings: np.ndarray
+    variance_explained: Optional[np.ndarray] = None
+    method: str = "unknown"
+    n_components: int = 2
+
+    def __post_init__(self):
+        if not isinstance(self.reduced_embeddings, np.ndarray):
+            self.reduced_embeddings = np.array(self.reduced_embeddings)
+
+
+@dataclass
+class PlotData:
+    documents: List[Document]
+    coordinates: np.ndarray
+    prompts: Optional[List[Document]] = None
+    prompt_coordinates: Optional[np.ndarray] = None
+    
+    def __post_init__(self):
+        if not isinstance(self.coordinates, np.ndarray):
+            self.coordinates = np.array(self.coordinates)
+        if self.prompt_coordinates is not None and not isinstance(self.prompt_coordinates, np.ndarray):
+            self.prompt_coordinates = np.array(self.prompt_coordinates)