banking/vectorstore.py

from typing import Dict, List

import faiss
import numpy as np


class VectorStore:
    def __init__(self, dim: int = 64):
        self.dim = dim
        self.stores: Dict[str, Dict] = {}

    def _embed(self, text: str) -> np.ndarray:
        vec = np.zeros((self.dim,), dtype="float32")
        encoded = text.encode("utf-8")
        for idx, byte in enumerate(encoded):
            vec[idx % self.dim] += byte / 255.0
        norm = np.linalg.norm(vec)
        if norm > 0:
            vec /= norm
        return vec.reshape(1, -1)

    def rebuild_user(self, username: str, user_doc: Dict) -> None:
        entries: List[str] = []
        for acct in user_doc.get("accounts", []):
            entries.append(
                f"Account {acct['name']} ({acct['id']}) balance {acct['currency']} {acct['balance']:.2f}"
            )
        for tx in user_doc.get("transactions", []):
            entries.append(
                f"Transaction {tx['id']} on {tx['account_id']} amount {tx['currency']} {tx['amount']:.2f} for {tx['description']}"
            )
        for ben in user_doc.get("beneficiaries", []):
            entries.append(
                f"Beneficiary {ben['name']} at {ben['bank']} account {ben['account_number']} notes {ben.get('notes','')}"
            )

        index = faiss.IndexFlatL2(self.dim)
        if entries:
            vectors = np.vstack([self._embed(item) for item in entries])
            index.add(vectors)
        self.stores[username] = {"index": index, "entries": entries}

    def bootstrap(self, users: List[Dict]) -> None:
        for user in users:
            self.rebuild_user(user["username"], user)

    def search(self, username: str, query: str, top_k: int = 5) -> List[str]:
        store = self.stores.get(username)
        if not store or store["index"].ntotal == 0:
            return []
        vector = self._embed(query)
        k = min(top_k, store["index"].ntotal)
        _, neighbors = store["index"].search(vector, k)
        return [store["entries"][idx] for idx in neighbors[0] if idx != -1]