Zvec Linux Demo

zvec-linux-demo.py

from __future__ import annotations

import shutil
import time
from pathlib import Path

import zvec
from zvec import (
    CollectionOption,
    DataType,
    Doc,
    FieldSchema,
    HnswIndexParam,
    InvertIndexParam,
    LogLevel,
    LogType,
    VectorQuery,
    VectorSchema,
)


DEMO_DIR = Path(__file__).resolve().parent / ".zvec_linux_terminal_demo"
COLLECTION_PATH = DEMO_DIR / "linux_demo_collection"
VECTOR_DIM = 8
WIDTH = 108


def style(text: str, code: str) -> str:
    return f"\033[{code}m{text}\033[0m"


def banner(title: str) -> None:
    line = "=" * WIDTH
    print(style(f"\n{line}", "36"))
    print(style(f"{title:^{WIDTH}}", "1;36"))
    print(style(f"{line}", "36"))


def ok(message: str) -> None:
    print(style(f"[ok]   {message}", "32"))


def info(message: str) -> None:
    print(style(f"[info] {message}", "37"))


def timed(label: str, fn):
    started = time.perf_counter()
    result = fn()
    elapsed_ms = (time.perf_counter() - started) * 1000
    return result, elapsed_ms


def short_vector(values: list[float]) -> str:
    return "[" + ", ".join(f"{v:.2f}" for v in values[:4]) + ", ...]"


def build_schema() -> zvec.CollectionSchema:
    return zvec.CollectionSchema(
        name="linux_terminal_demo",
        fields=[
            FieldSchema(
                "title",
                DataType.STRING,
                nullable=False,
                index_param=InvertIndexParam(),
            ),
            FieldSchema(
                "platform",
                DataType.STRING,
                nullable=False,
                index_param=InvertIndexParam(),
            ),
            FieldSchema(
                "kind",
                DataType.STRING,
                nullable=False,
                index_param=InvertIndexParam(),
            ),
        ],
        vectors=[
            VectorSchema(
                "embedding",
                DataType.VECTOR_FP32,
                dimension=VECTOR_DIM,
                index_param=HnswIndexParam(),
            )
        ],
    )


def build_docs() -> list[Doc]:
    return [
        Doc(
            id="linux-1",
            fields={
                "title": "Ubuntu terminal build log screenshot",
                "platform": "linux",
                "kind": "image",
            },
            vectors={"embedding": [0.98, 0.91, 0.84, 0.20, 0.15, 0.11, 0.08, 0.05]},
        ),
        Doc(
            id="linux-2",
            fields={
                "title": "Systemd service debugging note",
                "platform": "linux",
                "kind": "note",
            },
            vectors={"embedding": [0.95, 0.89, 0.81, 0.22, 0.18, 0.10, 0.07, 0.04]},
        ),
        Doc(
            id="linux-3",
            fields={
                "title": "Docker bridge network diagram",
                "platform": "linux",
                "kind": "image",
            },
            vectors={"embedding": [0.14, 0.20, 0.26, 0.95, 0.90, 0.82, 0.18, 0.10]},
        ),
        Doc(
            id="mac-1",
            fields={
                "title": "macOS console crash screenshot",
                "platform": "macos",
                "kind": "image",
            },
            vectors={"embedding": [0.80, 0.76, 0.72, 0.18, 0.16, 0.14, 0.10, 0.08]},
        ),
        Doc(
            id="win-1",
            fields={
                "title": "Windows PowerShell troubleshooting note",
                "platform": "windows",
                "kind": "note",
            },
            vectors={"embedding": [0.76, 0.73, 0.68, 0.16, 0.12, 0.10, 0.11, 0.09]},
        ),
    ]


def print_result_table(results: list[Doc]) -> None:
    print(style("\nSearch results", "1;33"))
    print(style("-" * WIDTH, "33"))
    print(f"{'rank':<6}{'id':<12}{'score':<12}{'platform':<12}{'kind':<10}title")
    print(style("-" * WIDTH, "33"))
    for idx, doc in enumerate(results, start=1):
        score = f"{doc.score:.4f}" if doc.score is not None else "-"
        title = doc.field("title") or ""
        platform = doc.field("platform") or ""
        kind = doc.field("kind") or ""
        print(
            f"{idx:<6}{doc.id:<12}{score:<12}{platform:<12}{kind:<10}{title}"
        )


def print_metrics(metrics: dict[str, float], doc_count: int) -> None:
    print(style("\nLifecycle", "1;35"))
    print(style("-" * WIDTH, "35"))
    print(
        f"{'create':<18}{'insert':<18}{'query':<18}{'destroy':<18}{'docs':<18}{'vector_dim':<18}"
    )
    print(style("-" * WIDTH, "35"))
    print(
        f"{metrics['create']:.2f} ms{'':<10}"
        f"{metrics['insert']:.2f} ms{'':<10}"
        f"{metrics['query']:.2f} ms{'':<10}"
        f"{metrics['destroy']:.2f} ms{'':<10}"
        f"{doc_count:<18}"
        f"{VECTOR_DIM:<18}"
    )


def print_catalog(docs: list[Doc]) -> None:
    summary = ", ".join(
        f"{doc.id}:{doc.field('platform')}/{doc.field('kind')}" for doc in docs
    )
    print(style("Corpus", "1;37"))
    print(summary)


def main() -> int:
    banner("Zvec Linux Terminal Demo")
    info(f"collection path: {COLLECTION_PATH}")

    if DEMO_DIR.exists():
        shutil.rmtree(DEMO_DIR)

    DEMO_DIR.mkdir(parents=True, exist_ok=True)

    zvec.init(log_type=LogType.CONSOLE, log_level=LogLevel.ERROR)
    ok("zvec initialized for a compact Linux terminal screenshot")

    collection, create_ms = timed(
        "create_and_open",
        lambda: zvec.create_and_open(
            path=str(COLLECTION_PATH),
            schema=build_schema(),
            option=CollectionOption(read_only=False, enable_mmap=True),
        ),
    )

    docs = build_docs()
    print_catalog(docs)
    statuses, insert_ms = timed("insert", lambda: collection.insert(docs))
    if not all(status.ok() for status in statuses):
        print(style("[error] insert failed", "31"))
        return 1

    query_vector = [0.97, 0.90, 0.83, 0.19, 0.16, 0.10, 0.09, 0.05]
    info("scenario: a Linux troubleshooting note automatically recalls a terminal screenshot and related notes")
    info(f"query vector: {short_vector(query_vector)}")
    results, query_ms = timed(
        "query",
        lambda: collection.query(
            vectors=VectorQuery(field_name="embedding", vector=query_vector),
            topk=3,
            output_fields=["title", "platform", "kind"],
        ),
    )
    print_result_table(results)

    _, destroy_ms = timed("destroy", collection.destroy)
    if DEMO_DIR.exists():
        shutil.rmtree(DEMO_DIR, ignore_errors=True)
    print_metrics(
        {
            "create": create_ms,
            "insert": insert_ms,
            "query": query_ms,
            "destroy": destroy_ms,
        },
        len(docs),
    )
    ok("demo artifacts removed")

    banner("Demo Complete")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print(style("\n[warn] interrupted by user", "33"))
        raise