feat: ✨ Working PoC of the Dungeon Masters Vault

src/embedding.py

@@ -0,0 +1,45 @@
+import requests
+from langchain_core.embeddings import Embeddings
+
+
+class LocalLMEmbeddings(Embeddings):
+    def __init__(
+        self, model: str, base_url: str = "http://192.168.0.49:1234", batch_size: int = 32
+    ):
+        self.url = f"{base_url}/v1/embeddings"
+        self.model = model
+        self.batch_size = batch_size
+
+    def _post_request(self, input_texts: list[str]) -> list[list[float]]:
+        """Handles the actual HTTP POST to the local server."""
+        payload = {"model": self.model, "input": input_texts}
+        try:
+            response = requests.post(
+                self.url, json=payload, timeout=120
+            )  # Longer timeout for batches
+            response.raise_for_status()
+            data = response.json()
+            return [item["embedding"] for item in data["data"]]
+        except Exception as e:
+            print(f"❌ Batch request failed: {e}")
+            # Returning empty lists to maintain index integrity if needed,
+            # or you could raise the error to stop the pipeline.
+            return [[] for _ in input_texts]
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Splits 500+ chunks into batches of 32 and processes them."""
+        all_embeddings = []
+
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+            print(f"🚀 Processing batch {(i // self.batch_size) + 1} (Size: {len(batch)})...")
+
+            batch_vectors = self._post_request(batch)
+            all_embeddings.extend(batch_vectors)
+
+        return all_embeddings
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embeds the single search query."""
+        result = self._post_request([text])
+        return result[0] if result else []