diff --git a/.devcontainer/requirements.txt b/.devcontainer/requirements.txt index d423946a..25ede636 100644 --- a/.devcontainer/requirements.txt +++ b/.devcontainer/requirements.txt @@ -3,6 +3,5 @@ ratio1 decentra-vision python-telegram-bot[rate-limiter] protobuf==5.28.3 -vectordb ngrok paramiko \ No newline at end of file diff --git a/.github/workflows/build_gpu.yml b/.github/workflows/build_gpu.yml new file mode 100644 index 00000000..ed48d77d --- /dev/null +++ b/.github/workflows/build_gpu.yml @@ -0,0 +1,99 @@ +name: Build GPU images + +on: + push: + branches: + - "develop" + - "main" + workflow_dispatch: + inputs: + networks: + description: "Networks to build (comma-separated: devnet,testnet,mainnet)" + required: false + default: "devnet,testnet" + +jobs: + build-and-push-gpu: + runs-on: ubuntu-latest + strategy: + matrix: + include: + - network: devnet + dockerfile: Dockerfile_devnet + branch: develop + - network: testnet + dockerfile: Dockerfile_testnet + branch: develop + - network: mainnet + dockerfile: Dockerfile_mainnet + branch: main + + # Only build if the branch matches, or if manually dispatched + if: >- + github.event_name == 'workflow_dispatch' || + (matrix.branch == 'develop' && github.ref == 'refs/heads/develop') || + (matrix.branch == 'main' && github.ref == 'refs/heads/main') + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Configure Git + run: | + git config --local user.email "action@github.com" + git config --local user.name "GitHub Action" + + - name: Retrieve edge node version + id: retrieve_version + run: | + echo "VERSION=$(cat ver.py | grep -o "'.*'")" >> $GITHUB_ENV + + - name: Check latest naeural_core version + id: check_core_latest_version + run: | + LATEST_VERSION=$(curl -s https://pypi.org/pypi/naeural-core/json | jq -r '.info.version') + echo "LATEST_NAEURAL_CORE_VERSION=$LATEST_VERSION" >> $GITHUB_ENV + + - name: Debug version + run: | + VERSION=${VERSION//\'/} + echo "GPU build for ${{ matrix.network }}" + echo "Edge node version: '$VERSION'" + echo "Latest naeural_core version on PyPI: '$LATEST_NAEURAL_CORE_VERSION'" + env: + VERSION: ${{ env.VERSION }} + LATEST_NAEURAL_CORE_VERSION: ${{ env.LATEST_NAEURAL_CORE_VERSION }} + + - name: Log in to Docker Hub + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + + - name: Set up Docker Buildx + id: buildx + uses: docker/setup-buildx-action@v3 + with: + version: "lab:latest" + driver: cloud + endpoint: "naeural/naeural-builder" + + - name: Cleanup space + run: | + echo "===========docker buildx du (before) ===================" + docker buildx du --builder "${{ steps.buildx.outputs.name }}" + echo "===========docker buildx prune -f ============" + docker buildx prune -f --verbose --builder "${{ steps.buildx.outputs.name }}" + echo "===========docker buildx du (after) =================" + docker buildx du --builder "${{ steps.buildx.outputs.name }}" + + - name: Build and push GPU ${{ matrix.network }} image + uses: docker/build-push-action@v6 + with: + builder: ${{ steps.buildx.outputs.name }} + context: . + file: ./${{ matrix.dockerfile }} + build-args: | + BASE_IMAGE=ratio1/base_edge_node_amd64_gpu:latest + push: true + tags: "ratio1/edge_node_gpu:${{ matrix.network }}" diff --git a/.gitignore b/.gitignore index aa2139b2..a6c4c30b 100644 --- a/.gitignore +++ b/.gitignore @@ -158,9 +158,6 @@ config_startup*.yaml config_startup*.yml inference/model_testing/_local_cache/_logs/MPTF.txt inference/model_testing/_local_cache/_logs/20211224_102325_MPTF_001_log.txt -vectordb -_vector_db_cache -_vector_db_cache_HNSWVectorDB db_cache plugins/libs/_cache/ diff --git a/Dockerfile_devnet b/Dockerfile_devnet index 69ed89f1..1d9ed735 100644 --- a/Dockerfile_devnet +++ b/Dockerfile_devnet @@ -1,45 +1,30 @@ -# -d for dind -FROM ratio1/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-dnctk -#FROM aidamian/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-d +# Base image: CPU by default, override with --build-arg BASE_IMAGE=ratio1/base_edge_node_amd64_gpu:latest for GPU +# The base image provides: Python 3.13, PyTorch, FFmpeg, Docker Engine (DIND), Node.js, uv, and ML/data stack +ARG BASE_IMAGE=ratio1/base_edge_node_amd64_cpu:latest +FROM ${BASE_IMAGE} -# Install IPFS -# The line below was creating issues due to flaky external repos -# RUN apt-get update && apt-get install -y wget && apt-get install -y tar -# Install tools needed for the next steps without hitting flaky external repos +# Install IPFS (Kubo) — needed for R1FS decentralized file system +ARG KUBO_VERSION=v0.35.0 RUN set -eux; \ - # disable NodeSource (if present) so apt won't read it - rm -f /etc/apt/sources.list.d/nodesource*.list /etc/apt/sources.list.d/nodesource*.sources || true; \ - apt-get update -o Acquire::Retries=3; \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - wget tar ca-certificates \ - ninja-build; \ - rm -rf /var/lib/apt/lists/* - -RUN wget https://dist.ipfs.tech/kubo/v0.35.0/kubo_v0.35.0_linux-amd64.tar.gz && \ - tar -xvzf kubo_v0.35.0_linux-amd64.tar.gz && \ - cd kubo && \ - bash install.sh -# End Install IPFS + wget -q https://dist.ipfs.tech/kubo/${KUBO_VERSION}/kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + tar -xzf kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + cd kubo && bash install.sh; \ + cd / && rm -rf kubo kubo_${KUBO_VERSION}_linux-amd64.tar.gz -# Install Cloudflared +# Install Cloudflared — tunnel engine for exposing webapp endpoints ARG CLOUDFLARED_VERSION=2025.7.0 - RUN set -eux; \ - # --- map architecture names to Cloudflare’s file names ------- - arch="$(uname -m)"; \ - case "$arch" in \ - x86_64) arch=amd64 ;; \ - aarch64 | arm64) arch=arm64 ;; \ - armv7l) arch=armv7 ;; \ + arch="$(uname -m)"; \ + case "$arch" in \ + x86_64) arch=amd64 ;; \ + aarch64 | arm64) arch=arm64 ;; \ + armv7l) arch=armv7 ;; \ *) echo "Unsupported arch: $arch" >&2; exit 1 ;; \ - esac; \ - # --- download the pinned release -------------------------------- + esac; \ curl -L "https://github.com/cloudflare/cloudflared/releases/download/${CLOUDFLARED_VERSION}/cloudflared-linux-${arch}" \ - -o /usr/local/bin/cloudflared; \ - chmod +x /usr/local/bin/cloudflared; \ - # --- (optional) show version so it appears in build logs -------- + -o /usr/local/bin/cloudflared; \ + chmod +x /usr/local/bin/cloudflared; \ cloudflared --version -# End Install Cloudflared COPY ./cmds /usr/local/bin/ RUN chmod +x /usr/local/bin/* @@ -84,9 +69,6 @@ ENV EE_DEBUG_R1FS=true # althouh this is not recommended as it should be in .env file # ENV EE_DEVICE=cuda:0 -# to avoid issues with hnswlib and building from source -ENV HNSWLIB_NO_NATIVE=1 - RUN pip install --no-cache-dir -r requirements.txt RUN pip install --no-cache-dir --no-deps naeural-core diff --git a/Dockerfile_mainnet b/Dockerfile_mainnet index 5a5fb190..7669faab 100644 --- a/Dockerfile_mainnet +++ b/Dockerfile_mainnet @@ -1,45 +1,30 @@ -# -d for dind -#FROM aidamian/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-d -FROM ratio1/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-dnctk - -# Install IPFS -# The line below was creating issues due to flaky external repos -# RUN apt-get update && apt-get install -y wget && apt-get install -y tar -# Install tools needed for the next steps without hitting flaky external repos +# Base image: CPU by default, override with --build-arg BASE_IMAGE=ratio1/base_edge_node_amd64_gpu:latest for GPU +# The base image provides: Python 3.13, PyTorch, FFmpeg, Docker Engine (DIND), Node.js, uv, and ML/data stack +ARG BASE_IMAGE=ratio1/base_edge_node_amd64_cpu:latest +FROM ${BASE_IMAGE} + +# Install IPFS (Kubo) — needed for R1FS decentralized file system +ARG KUBO_VERSION=v0.35.0 RUN set -eux; \ - # disable NodeSource (if present) so apt won't read it - rm -f /etc/apt/sources.list.d/nodesource*.list /etc/apt/sources.list.d/nodesource*.sources || true; \ - apt-get update -o Acquire::Retries=3; \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - wget tar ca-certificates \ - ninja-build; \ - rm -rf /var/lib/apt/lists/* - -RUN wget https://dist.ipfs.tech/kubo/v0.35.0/kubo_v0.35.0_linux-amd64.tar.gz && \ - tar -xvzf kubo_v0.35.0_linux-amd64.tar.gz && \ - cd kubo && \ - bash install.sh -# End Install IPFS - -# Install Cloudflared -ARG CLOUDFLARED_VERSION=2025.7.0 + wget -q https://dist.ipfs.tech/kubo/${KUBO_VERSION}/kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + tar -xzf kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + cd kubo && bash install.sh; \ + cd / && rm -rf kubo kubo_${KUBO_VERSION}_linux-amd64.tar.gz +# Install Cloudflared — tunnel engine for exposing webapp endpoints +ARG CLOUDFLARED_VERSION=2025.7.0 RUN set -eux; \ - # --- map architecture names to Cloudflare’s file names ------- - arch="$(uname -m)"; \ - case "$arch" in \ - x86_64) arch=amd64 ;; \ - aarch64 | arm64) arch=arm64 ;; \ - armv7l) arch=armv7 ;; \ + arch="$(uname -m)"; \ + case "$arch" in \ + x86_64) arch=amd64 ;; \ + aarch64 | arm64) arch=arm64 ;; \ + armv7l) arch=armv7 ;; \ *) echo "Unsupported arch: $arch" >&2; exit 1 ;; \ - esac; \ - # --- download the pinned release -------------------------------- + esac; \ curl -L "https://github.com/cloudflare/cloudflared/releases/download/${CLOUDFLARED_VERSION}/cloudflared-linux-${arch}" \ - -o /usr/local/bin/cloudflared; \ - chmod +x /usr/local/bin/cloudflared; \ - # --- (optional) show version so it appears in build logs -------- + -o /usr/local/bin/cloudflared; \ + chmod +x /usr/local/bin/cloudflared; \ cloudflared --version -# End Install Cloudflared COPY ./cmds /usr/local/bin/ diff --git a/Dockerfile_testnet b/Dockerfile_testnet index 370493d8..a6b02228 100644 --- a/Dockerfile_testnet +++ b/Dockerfile_testnet @@ -1,45 +1,30 @@ -# -d for dind -#FROM aidamian/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-d -FROM ratio1/base_edge_node:x86_64-py3.10.12-th2.3.1.cu121-tr4.43.3-dnctk +# Base image: CPU by default, override with --build-arg BASE_IMAGE=ratio1/base_edge_node_amd64_gpu:latest for GPU +# The base image provides: Python 3.13, PyTorch, FFmpeg, Docker Engine (DIND), Node.js, uv, and ML/data stack +ARG BASE_IMAGE=ratio1/base_edge_node_amd64_cpu:latest +FROM ${BASE_IMAGE} -# Install IPFS -# The line below was creating issues due to flaky external repos -# RUN apt-get update && apt-get install -y wget && apt-get install -y tar -# Install tools needed for the next steps without hitting flaky external repos +# Install IPFS (Kubo) — needed for R1FS decentralized file system +ARG KUBO_VERSION=v0.35.0 RUN set -eux; \ - # disable NodeSource (if present) so apt won't read it - rm -f /etc/apt/sources.list.d/nodesource*.list /etc/apt/sources.list.d/nodesource*.sources || true; \ - apt-get update -o Acquire::Retries=3; \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - wget tar ca-certificates \ - ninja-build; \ - rm -rf /var/lib/apt/lists/* - -RUN wget https://dist.ipfs.tech/kubo/v0.35.0/kubo_v0.35.0_linux-amd64.tar.gz && \ - tar -xvzf kubo_v0.35.0_linux-amd64.tar.gz && \ - cd kubo && \ - bash install.sh -# End Install IPFS + wget -q https://dist.ipfs.tech/kubo/${KUBO_VERSION}/kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + tar -xzf kubo_${KUBO_VERSION}_linux-amd64.tar.gz; \ + cd kubo && bash install.sh; \ + cd / && rm -rf kubo kubo_${KUBO_VERSION}_linux-amd64.tar.gz -# Install Cloudflared +# Install Cloudflared — tunnel engine for exposing webapp endpoints ARG CLOUDFLARED_VERSION=2025.7.0 - RUN set -eux; \ - # --- map architecture names to Cloudflare’s file names ------- - arch="$(uname -m)"; \ - case "$arch" in \ - x86_64) arch=amd64 ;; \ - aarch64 | arm64) arch=arm64 ;; \ - armv7l) arch=armv7 ;; \ + arch="$(uname -m)"; \ + case "$arch" in \ + x86_64) arch=amd64 ;; \ + aarch64 | arm64) arch=arm64 ;; \ + armv7l) arch=armv7 ;; \ *) echo "Unsupported arch: $arch" >&2; exit 1 ;; \ - esac; \ - # --- download the pinned release -------------------------------- + esac; \ curl -L "https://github.com/cloudflare/cloudflared/releases/download/${CLOUDFLARED_VERSION}/cloudflared-linux-${arch}" \ - -o /usr/local/bin/cloudflared; \ - chmod +x /usr/local/bin/cloudflared; \ - # --- (optional) show version so it appears in build logs -------- + -o /usr/local/bin/cloudflared; \ + chmod +x /usr/local/bin/cloudflared; \ cloudflared --version -# End Install Cloudflared COPY ./cmds /usr/local/bin/ RUN chmod +x /usr/local/bin/* @@ -84,9 +69,6 @@ ENV EE_DEBUG_R1FS=true # althouh this is not recommended as it should be in .env file # ENV EE_DEVICE=cuda:0 -# to avoid issues with hnswlib and building from source -ENV HNSWLIB_NO_NATIVE=1 - RUN pip install --no-cache-dir -r requirements.txt RUN pip install --no-cache-dir --no-deps naeural-core diff --git a/extensions/business/tutorials/test_faiss_vectordb.py b/extensions/business/tutorials/test_faiss_vectordb.py new file mode 100644 index 00000000..cb0e59b4 --- /dev/null +++ b/extensions/business/tutorials/test_faiss_vectordb.py @@ -0,0 +1,138 @@ +""" +Test plugin for FAISS vectordb replacement. + +Exposes FastAPI endpoints to test the FaissVectorDB adapter: + GET /status — check plugin is alive and show db stats + POST /add_docs — add documents to a context + POST /search — search a context with a query string + GET /list_contexts — list all contexts and their doc counts + POST /reset_context — delete and recreate a context +""" + +from naeural_core.business.default.web_app.fast_api_web_app import FastApiWebAppPlugin + +from extensions.utils.faiss_vectordb import FaissVectorDB + +__VER__ = '0.1.0.0' + +EMBEDDING_SIZE = 128 # small for testing, real uses 1024 + +_CONFIG = { + **FastApiWebAppPlugin.CONFIG, + + 'PORT': None, + + 'VALIDATION_RULES': { + **FastApiWebAppPlugin.CONFIG['VALIDATION_RULES'], + }, +} + + +class TestFaissVectordbPlugin(FastApiWebAppPlugin): + CONFIG = _CONFIG + + def __init__(self, **kwargs): + self._dbs = {} + super(TestFaissVectordbPlugin, self).__init__(**kwargs) + return + + + def on_init(self, **kwargs): + super(TestFaissVectordbPlugin, self).on_init(**kwargs) + self.P("TestFaissVectordb plugin initialized") + return + + def _get_db(self, context: str) -> FaissVectorDB: + if context not in self._dbs: + workspace = self.os_path.join( + self.get_data_folder(), 'faiss_test', context + ) + self._dbs[context] = FaissVectorDB( + workspace=workspace, + embedding_size=EMBEDDING_SIZE, + ) + self.P(f"Created new context: {context}") + return self._dbs[context] + + def _embed_texts(self, texts: list): + """Simple deterministic embedding for testing — hash-based.""" + embeddings = [] + for text in texts: + self.np.random.seed(hash(text) % (2**31)) + emb = self.np.random.randn(EMBEDDING_SIZE).astype(self.np.float32) + emb /= self.np.linalg.norm(emb) + embeddings.append(emb) + return self.np.array(embeddings, dtype=self.np.float32) + + @FastApiWebAppPlugin.endpoint + def status(self) -> dict: + """Health check and db stats.""" + contexts = {} + for name, db in self._dbs.items(): + contexts[name] = db.num_docs() + return { + "status": "ok", + "version": __VER__, + "embedding_size": EMBEDDING_SIZE, + "contexts": contexts, + } + + @FastApiWebAppPlugin.endpoint(method="post") + def add_docs(self, context: str = "default", documents: list = []) -> dict: + """Add documents (list of strings) to a context.""" + if not documents: + return {"error": "No documents provided"} + db = self._get_db(context) + embeddings = self._embed_texts(documents) + curr_size = db.num_docs() + docs = [ + {"text": doc, "embedding": emb, "idx": curr_size + i} + for i, (doc, emb) in enumerate(zip(documents, embeddings)) + ] + db.index(docs) + self.P(f"Indexed {len(docs)} docs in context '{context}', total={db.num_docs()}") + return { + "context": context, + "added": len(docs), + "total": db.num_docs(), + } + + @FastApiWebAppPlugin.endpoint(method="post") + def search(self, query: str, context: str = "default", k: int = 5) -> dict: + """Search a context with a query string.""" + if context not in self._dbs: + return {"error": f"Context '{context}' not found"} + db = self._get_db(context) + query_embedding = self._embed_texts([query])[0] + results = db.search(query_embedding, limit=k) + return { + "context": context, + "query": query, + "results": [ + {"text": r.text, "idx": r.idx, "score": round(r.score, 4)} + for r in results + ], + } + + @FastApiWebAppPlugin.endpoint + def list_contexts(self) -> dict: + """List all contexts and doc counts.""" + return { + "contexts": { + name: db.num_docs() for name, db in self._dbs.items() + } + } + + @FastApiWebAppPlugin.endpoint(method="post") + def reset_context(self, context: str = "default") -> dict: + """Delete and recreate a context.""" + if context in self._dbs: + self._dbs[context].close() + import shutil + workspace = self.os_path.join( + self.get_data_folder(), 'faiss_test', context + ) + shutil.rmtree(workspace, ignore_errors=True) + del self._dbs[context] + self.P(f"Reset context: {context}") + return {"status": "ok", "context": context} diff --git a/extensions/serving/base/base_doc_emb_serving.py b/extensions/serving/base/base_doc_emb_serving.py index cc1581a6..bb0422a3 100644 --- a/extensions/serving/base/base_doc_emb_serving.py +++ b/extensions/serving/base/base_doc_emb_serving.py @@ -7,9 +7,7 @@ from pypdf import PdfReader from docx import Document -from docarray import BaseDoc, DocList -from docarray.typing import NdArray -from vectordb import HNSWVectorDB +from extensions.utils.faiss_vectordb import FaissVectorDB """ @@ -77,15 +75,6 @@ class DocEmbCt: } -class NaeuralDoc(BaseDoc): - # TODO: find how the size of this can be configurable in case of different model. - # this should be done at initialization time only in order to avoid vectordb issues. - # TODO: encrypt the text for db and decrypt it when needed. - text: str = '' - embedding: NdArray[DOC_EMBEDDING_SIZE] - idx: int = -1 -# endclass - class DocSplitter: """ @@ -241,7 +230,7 @@ def __maybe_load_backup(self): embedding_size = saved_data.get('embedding_size', None) for context in contexts: if context not in self.__dbs: - self.__dbs[context] = HNSWVectorDB[NaeuralDoc](workspace=self.__db_cache_workspace(context)) + self.__dbs[context] = FaissVectorDB(workspace=self.__db_cache_workspace(context), embedding_size=DOC_EMBEDDING_SIZE) # endif sanity check in case of db already loaded # endfor each context # endif saved data available @@ -744,21 +733,19 @@ def __add_docs(self, docs, context: str = None): # endif context is None if context not in self.__dbs: self.P(f"Creating new context: {context}") - self.__dbs[context] = HNSWVectorDB[NaeuralDoc]( - workspace=self.__db_cache_workspace(context) - ) + self.__dbs[context] = FaissVectorDB(workspace=self.__db_cache_workspace(context), embedding_size=DOC_EMBEDDING_SIZE) self.__backup_contexts() # endif context not in dbs segments = self.__doc_splitter.split_documents(docs) segments_embeddings = self.embed_texts(segments) - curr_size = self.__dbs[context].num_docs()['num_docs'] + curr_size = self.__dbs[context].num_docs() lst_docs = [ - NaeuralDoc(text=segment, embedding=emb, idx=curr_size + i) + {"text": segment, "embedding": emb, "idx": curr_size + i} for i, (segment, emb) in enumerate(zip(segments, segments_embeddings)) ] # TODO: maybe check for duplicates self.P(f"Indexing {len(lst_docs)} documents in context '{context}'...") - self.__dbs[context].index(inputs=DocList[NaeuralDoc](lst_docs)) + self.__dbs[context].index(lst_docs) return def get_result_dict(self, request_id, docs=None, query=None, context_list=None, error_message=None, **kwargs): @@ -860,22 +847,13 @@ def _predict(self, processed_batch): # endif context not in dbs # Embed the query. query_embedding = self.embed_texts(query) - query_doc = NaeuralDoc(text=query, embedding=query_embedding, idx=-1) # Search for the closest documents. self.P(f"Searching for the closest {k} documents to the query in context '{context}'...") - search_results = self.__dbs[context].search( - inputs=DocList[NaeuralDoc]([query_doc]), limit=k - )[0] + search_results = self.__dbs[context].search(query_embedding, limit=k) self.P(f"Search results: {search_results}") - matches, scores = search_results.matches, search_results.scores - matches_with_scores = [ - (match, score) for match, score in zip(matches, scores) - ] - matches_ordered_by_idx = sorted(matches_with_scores, key=lambda x: x[0].idx) - # Sort the results by the index. - result_texts = [ - res[0].text for res in matches_ordered_by_idx - ] + # Sort the results by the document index. + results_ordered_by_idx = sorted(search_results, key=lambda r: r.idx) + result_texts = [r.text for r in results_ordered_by_idx] self.P(f"Result texts: {result_texts}") results.append( self.get_result_dict(request_id=req_id, docs=result_texts, query=query) diff --git a/extensions/utils/faiss_vectordb.py b/extensions/utils/faiss_vectordb.py new file mode 100644 index 00000000..33f6ddd7 --- /dev/null +++ b/extensions/utils/faiss_vectordb.py @@ -0,0 +1,132 @@ +""" +FAISS vector database adapter for Edge Node. + +Replaces jina-ai/vectordb (archived, broken on modern Python). +Uses IndexFlatIP for cosine similarity on L2-normalized vectors. +Auto-detects GPU and moves index there when available. + +Storage layout per context directory: + index.faiss — binary FAISS index + meta.json — sidecar with document text and idx +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import faiss +import numpy as np + + +@dataclass +class SearchResult: + text: str + idx: int + score: float + + +class FaissVectorDB: + """Thin FAISS wrapper matching the edge_node vectordb usage pattern.""" + + INDEX_FILE = "index.faiss" + META_FILE = "meta.json" + + def __init__(self, workspace: str | Path, embedding_size: int = 1024): + self.workspace = Path(workspace) + self.embedding_size = embedding_size + self._index = None + self._meta: list[dict] = [] + self._gpu_res = None + self._open() + + # -- lifecycle --------------------------------------------------------------- + + def _open(self) -> None: + self.workspace.mkdir(parents=True, exist_ok=True) + index_path = self.workspace / self.INDEX_FILE + meta_path = self.workspace / self.META_FILE + + if index_path.exists() and meta_path.exists(): + self._index = faiss.read_index(str(index_path)) + with open(meta_path) as f: + self._meta = json.load(f) + else: + self._index = faiss.IndexFlatIP(self.embedding_size) + self._meta = [] + + self._maybe_move_to_gpu() + + def _gpu_available(self) -> bool: + return hasattr(faiss, "get_num_gpus") and faiss.get_num_gpus() > 0 + + def _maybe_move_to_gpu(self) -> None: + if not self._gpu_available(): + return + self._gpu_res = faiss.StandardGpuResources() + self._index = faiss.index_cpu_to_gpu(self._gpu_res, 0, self._index) + + def _save(self) -> None: + if self._index is None: + return + + index_to_write = self._index + if self._gpu_res is not None: + index_to_write = faiss.index_gpu_to_cpu(self._index) + + faiss.write_index(index_to_write, str(self.workspace / self.INDEX_FILE)) + with open(self.workspace / self.META_FILE, "w") as f: + json.dump(self._meta, f) + + def close(self) -> None: + self._save() + self._index = None + self._meta = [] + self._gpu_res = None + + # -- operations -------------------------------------------------------------- + + def index(self, documents: list[dict]) -> None: + """Index a batch of documents. Each doc is a dict with text, embedding, idx.""" + embeddings = [] + for d in documents: + emb = d["embedding"] + if hasattr(emb, "numpy"): + emb = emb.numpy() + embeddings.append(emb) + + vectors = np.array(embeddings, dtype=np.float32) + if vectors.ndim == 1: + vectors = vectors.reshape(1, -1) + + faiss.normalize_L2(vectors) + self._index.add(vectors) + self._meta.extend({"text": d["text"], "idx": d["idx"]} for d in documents) + self._save() + + def search(self, query_embedding, limit: int = 10) -> list[SearchResult]: + """Search for nearest documents. Returns list of SearchResult.""" + if self._index is None or self._index.ntotal == 0: + return [] + + if hasattr(query_embedding, "numpy"): + query_embedding = query_embedding.numpy() + + query = np.array(query_embedding, dtype=np.float32) + if query.ndim == 1: + query = query.reshape(1, -1) + + faiss.normalize_L2(query) + scores, indices = self._index.search(query, min(limit, self._index.ntotal)) + + results = [] + for score, i in zip(scores[0], indices[0]): + if i < 0: + continue + meta = self._meta[i] + results.append(SearchResult(text=meta["text"], idx=meta["idx"], score=float(score))) + return results + + def num_docs(self) -> int: + return self._index.ntotal if self._index else 0 diff --git a/requirements.txt b/requirements.txt index 1e6c68f1..ac1ca384 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,7 +13,6 @@ decentra-vision python-telegram-bot[rate-limiter] protobuf==5.28.3 openai -vectordb ngrok openai diff --git a/ver.py b/ver.py index caa70f1a..b8f6c9af 100644 --- a/ver.py +++ b/ver.py @@ -1 +1 @@ -__VER__ = '2.10.150' +__VER__ = '2.10.151'