aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/services/rag/index/faiss_index.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import os
+import pickle
+from typing import Any
+
+import numpy as np
+
+try:
+    import faiss
+except Exception:
+    faiss = None
+
+from .base import VectorIndex
+
+"""A simple FAISS index per corpus (L2 on normalized vectors ~ cosine).
+Stores vectors as BLOBs along with metadata in a simple schema.
+"""
+
+
+class FAISSVectorIndex(VectorIndex):
+    """A simple FAISS index per corpus (L2 on normalized vectors ~ cosine)."""
+
+    def __init__(self, index_path: str, dim: int | None = None):
+        super().__init__(index_path)
+        self.dim = dim  # optional default; will infer on first add
+        os.makedirs(index_path, exist_ok=True)
+
+    def _paths(self, corpus_id: str):
+        base = os.path.join(self.index_path, corpus_id)
+        return base + ".index", base + ".meta.pkl"
+
+    def _load(self, corpus_id: str):
+        idx_path, meta_path = self._paths(corpus_id)
+        if not (os.path.exists(idx_path) and os.path.exists(meta_path)):
+            return None, []
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+        index = faiss.read_index(idx_path)
+        with open(meta_path, "rb") as f:
+            metas = pickle.load(f)
+        return index, metas
+
+    def _save(self, corpus_id: str, index, metas):
+        idx_path, meta_path = self._paths(corpus_id)
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+        faiss.write_index(index, idx_path)
+        with open(meta_path, "wb") as f:
+            pickle.dump(metas, f)
+
+    async def add(
+        self,
+        corpus_id: str,
+        chunk_ids: list[str],
+        vectors: list[list[float]],
+        metas: list[dict[str, Any]],
+    ):
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+        vecs = np.asarray(vectors, dtype=np.float32)
+        # normalize for cosine
+        norms = np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-9
+        vecs = vecs / norms
+        d = vecs.shape[1]
+        index, old_metas = self._load(corpus_id)
+        if index is None:
+            index = faiss.IndexFlatIP(d)  # cosine via normalized dot
+            old_metas = []
+        index.add(vecs)
+        for cid, m in zip(chunk_ids, metas, strict=True):
+            old_metas.append({"chunk_id": cid, "meta": m})
+        self._save(corpus_id, index, old_metas)
+
+    async def delete(self, corpus_id: str, chunk_ids: list[str] | None = None):
+        # Simple approach: rebuild if filtering; or delete entire corpus.
+        if not chunk_ids:
+            idx_path, meta_path = self._paths(corpus_id)
+            for p in (idx_path, meta_path):
+                if os.path.exists(p):
+                    os.remove(p)
+        else:
+            index, metas = self._load(corpus_id)
+            if index is None:
+                return
+            # Rebuild without those ids
+            keep = [i for i, m in enumerate(metas) if m["chunk_id"] not in set(chunk_ids)]
+            if not keep:
+                await self.delete(corpus_id, None)
+                return
+            # Need stored vectors to rebuild — this simple implementation does not persist them.
+            # In production, persist vectors or recompute from text.
+            raise NotImplementedError(
+                "Selective delete requires stored vectors; not implemented here."
+            )
+
+    async def list_chunks(self, corpus_id: str) -> list[str]:
+        _, metas = self._load(corpus_id)
+        return [m["chunk_id"] for m in metas] if metas else []
+
+    async def search(self, corpus_id: str, query_vec: list[float], k: int):
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+        index, metas = self._load(corpus_id)
+        if index is None:
+            return []
+        q = np.asarray([query_vec], dtype=np.float32)
+        q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-9)
+        D, I = index.search(q, k)  # noqa: E741
+        out = []
+        for score, idx in zip(D[0].tolist(), I[0].tolist(), strict=True):
+            if idx < 0 or idx >= len(metas):
+                continue
+            out.append(
+                {
+                    "chunk_id": metas[idx]["chunk_id"],
+                    "score": float(score),
+                    "meta": metas[idx]["meta"],
+                }
+            )
+        return out

aethergraph/services/rag/index/sqlite_index.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+from typing import Any
+
+import numpy as np
+
+from .base import VectorIndex
+
+"""A simple SQLite-based vector index per corpus (brute-force cosine similarity).
+Stores vectors as BLOBs along with metadata in a simple schema.
+"""
+
+
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS chunks (
+    corpus_id TEXT,
+    chunk_id  TEXT,
+    meta_json TEXT,
+    PRIMARY KEY (corpus_id, chunk_id)
+);
+CREATE TABLE IF NOT EXISTS embeddings (
+    corpus_id TEXT,
+    chunk_id  TEXT,
+    vec       BLOB,    -- np.float32 array bytes
+    norm      REAL,
+    PRIMARY KEY (corpus_id, chunk_id)
+);
+"""
+
+
+def _ensure_db(path: str):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    conn = sqlite3.connect(path)
+    try:
+        for stmt in SCHEMA.strip().split(";\n"):
+            s = stmt.strip()
+            if s:
+                conn.execute(s)
+        conn.commit()
+    finally:
+        conn.close()
+
+
+class SQLiteVectorIndex(VectorIndex):
+    def __init__(self, index_path: str):
+        super().__init__(index_path)
+        self.db_path = os.path.join(index_path, "index.sqlite")
+        _ensure_db(self.db_path)
+
+    def _connect(self):
+        return sqlite3.connect(self.db_path)
+
+    async def add(
+        self,
+        corpus_id: str,
+        chunk_ids: list[str],
+        vectors: list[list[float]],
+        metas: list[dict[str, Any]],
+    ):
+        conn = self._connect()
+        try:
+            cur = conn.cursor()
+            for cid, vec, meta in zip(chunk_ids, vectors, metas, strict=True):
+                v = np.asarray(vec, dtype=np.float32)
+                norm = float(np.linalg.norm(v) + 1e-9)
+                cur.execute(
+                    "REPLACE INTO chunks(corpus_id,chunk_id,meta_json) VALUES(?,?,?)",
+                    (corpus_id, cid, json.dumps(meta, ensure_ascii=False)),
+                )
+                cur.execute(
+                    "REPLACE INTO embeddings(corpus_id,chunk_id,vec,norm) VALUES(?,?,?,?)",
+                    (corpus_id, cid, v.tobytes(), norm),
+                )
+            conn.commit()
+        finally:
+            conn.close()
+
+    async def delete(self, corpus_id: str, chunk_ids: list[str] | None = None):
+        conn = self._connect()
+        try:
+            cur = conn.cursor()
+            if chunk_ids:
+                q = f"DELETE FROM chunks WHERE corpus_id=? AND chunk_id IN ({','.join(['?'] * len(chunk_ids))})"
+                cur.execute(q, [corpus_id, *chunk_ids])
+                q2 = f"DELETE FROM embeddings WHERE corpus_id=? AND chunk_id IN ({','.join(['?'] * len(chunk_ids))})"
+                cur.execute(q2, [corpus_id, *chunk_ids])
+            else:
+                cur.execute("DELETE FROM chunks WHERE corpus_id=?", (corpus_id,))
+                cur.execute("DELETE FROM embeddings WHERE corpus_id=?", (corpus_id,))
+            conn.commit()
+        finally:
+            conn.close()
+
+    async def list_chunks(self, corpus_id: str) -> list[str]:
+        conn = self._connect()
+        try:
+            cur = conn.cursor()
+            cur.execute("SELECT chunk_id FROM chunks WHERE corpus_id=?", (corpus_id,))
+            return [r[0] for r in cur.fetchall()]
+        finally:
+            conn.close()
+
+    async def search(
+        self, corpus_id: str, query_vec: list[float], k: int
+    ) -> list[dict[str, Any]]:  # Brute-force cosine similarity. Loads vectors for that corpus.
+        q = np.asarray(query_vec, dtype=np.float32)
+        qn = float(np.linalg.norm(q) + 1e-9)
+
+        conn = self._connect()
+        try:
+            cur = conn.cursor()
+            cur.execute(
+                "SELECT e.chunk_id, e.vec, e.norm, c.meta_json FROM embeddings e JOIN chunks c USING(corpus_id,chunk_id) WHERE e.corpus_id=?",
+                (corpus_id,),
+            )
+            rows = cur.fetchall()
+        finally:
+            conn.close()
+
+        scores = []
+        for chunk_id, vec_bytes, norm, meta_json in rows:
+            v = np.frombuffer(vec_bytes, dtype=np.float32)
+            score = float(np.dot(q, v) / (qn * norm))
+            scores.append((score, chunk_id, meta_json))
+
+        scores.sort(reverse=True, key=lambda x: x[0])
+        top = scores[:k]
+        out = []
+        for score, chunk_id, meta_json in top:
+            out.append({"chunk_id": chunk_id, "score": score, "meta": json.loads(meta_json)})
+        return out

aethergraph/services/rag/index_factory.py

@@ -0,0 +1,52 @@
+import logging
+from pathlib import Path
+
+from aethergraph.utils.optdeps import require
+
+logger = logging.getLogger("aethergraph.rag.index_factory")
+
+
+def _default_index_path(root: str, backend: str) -> str:
+    base = Path(root) / "rag_index"
+    if backend == "faiss":
+        return str(base / "faiss.index")
+    return str(base / "sqlite.index")
+
+
+def create_vector_index(
+    *, backend: str, index_path: str | None, dim: int | None, root: str = "./aethergraph_data/rag"
+):
+    """
+    Create a vector index instance. Supported backends: 'sqlite', 'faiss'.
+    Falls back to 'sqlite' if FAISS is unavailable.
+    """
+    backend = (backend or "sqlite").lower()
+    if backend not in {"sqlite", "faiss"}:
+        logger.warning(f"Unknown RAG backend {backend!r}; falling back to sqlite.")
+        backend = "sqlite"
+
+    if backend == "faiss":
+        # try FAISS, fallback to sqlite with a warning
+        try:
+            require("faiss", "faiss")  # faiss-cpu exposes module 'faiss'
+            from .index.faiss_index import FAISSVectorIndex
+
+            path = (
+                str(Path(index_path) / "faiss")
+                if index_path is not None
+                else _default_index_path(root, "faiss")
+            )
+            return FAISSVectorIndex(path, dim=dim)
+        except Exception as e:
+            logger.warning(f"FAISS backend unavailable ({e}); falling back to sqlite.")
+            backend = "sqlite"
+
+    # sqlite (default)
+    from .index.sqlite_index import SQLiteVectorIndex
+
+    path = (
+        str(Path(index_path) / "sqlite")
+        if index_path is not None
+        else _default_index_path(root, "sqlite")
+    )
+    return SQLiteVectorIndex(path)

aethergraph/services/rag/parsers/md.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+
+def extract_text(path: str) -> tuple[str, dict]:
+    with open(path, encoding="utf-8", errors="ignore") as f:
+        txt = f.read()
+    return txt, {}

aethergraph/services/rag/parsers/pdf.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from pypdf import PdfReader
+
+
+def extract_text(path: str) -> tuple[str, dict]:
+    reader = PdfReader(path)
+    texts = []
+    for page in reader.pages:
+        try:
+            texts.append(page.extract_text() or "")
+        except Exception:
+            continue
+    return "\n\n".join(texts), {"pages": len(reader.pages)}

aethergraph/services/rag/parsers/txt.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+
+def extract_text(path: str) -> tuple[str, dict]:
+    with open(path, encoding="utf-8", errors="ignore") as f:
+        txt = f.read()
+    return txt, {}

aethergraph/services/rag/utils/hybrid.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from collections import Counter
+import re
+from typing import Any
+
+
+def lexical_score(query: str, text: str) -> float:
+    # Extremely lightweight bag-of-words match score.
+    def ws(s: str) -> list[str]:
+        return re.findall(r"\w+", s.lower())
+
+    q = ws(query)
+    t = ws(text)
+    if not q or not t:
+        return 0.0
+    cq = Counter(q)
+    ct = Counter(t)
+    # normalized term overlap
+    overlap = sum(min(cq[w], ct.get(w, 0)) for w in cq)
+    return overlap / (sum(cq.values()) + 1e-9)
+
+
+def fuse_scores(dense_score: float, lexical: float, alpha: float = 0.8) -> float:
+    # Linear fusion; alpha favors dense similarity.
+    return alpha * dense_score + (1.0 - alpha) * lexical
+
+
+def topk_fuse(
+    query: str, dense_hits: list[dict[str, Any]], chunk_lookup: dict[str, str], k: int
+) -> list[dict[str, Any]]:
+    out = []
+    for h in dense_hits:
+        txt = chunk_lookup.get(h["chunk_id"], "")
+        lex = lexical_score(query, txt)
+        fused = fuse_scores(h.get("score", 0.0), lex)
+        out.append({**h, "score": fused})
+    out.sort(key=lambda x: x["score"], reverse=True)
+    return out[:k]

aethergraph/services/rag/utils/make_fs_key.py

@@ -0,0 +1,62 @@
+import base64
+import hashlib
+import re
+
+# Windows forbidden characters and device names
+_INVALID_CHARS_RE = re.compile(r'[<>:"/\\|?\*\x00-\x1F]')
+_RESERVED_WIN = {
+    "CON",
+    "PRN",
+    "AUX",
+    "NUL",
+    "COM1",
+    "COM2",
+    "COM3",
+    "COM4",
+    "COM5",
+    "COM6",
+    "COM7",
+    "COM8",
+    "COM9",
+    "LPT1",
+    "LPT2",
+    "LPT3",
+    "LPT4",
+    "LPT5",
+    "LPT6",
+    "LPT7",
+    "LPT8",
+    "LPT9",
+}
+
+
+def make_fs_key(cid: str, max_len: int = 128) -> str:
+    """
+    Convert any logical corpus_id (may include ':', Unicode, etc.)
+    into a portable filename segment: [a-zA-Z0-9._-] only, no trailing space/dot,
+    not a reserved device name.
+    """
+    # 1) Keep a short human-friendly prefix if present (e.g., "proj", "sess", "run")
+    if ":" in cid:
+        prefix, rest = cid.split(":", 1)
+    else:
+        prefix, rest = "cid", cid
+
+    # 2) Encode the rest to a compact, portable token (urlsafe base64 without padding)
+    #    This avoids collisions from simple character replacement.
+    token = base64.urlsafe_b64encode(rest.encode("utf-8")).decode("ascii").rstrip("=")
+
+    # 3) Build candidate and sanitize any stray chars just in case
+    key = f"{prefix}-{token}"
+    key = _INVALID_CHARS_RE.sub("_", key).rstrip(" .")
+
+    # 4) Avoid Windows reserved device names
+    if key.upper() in _RESERVED_WIN:
+        key = f"_{key}_"
+
+    # 5) Enforce a reasonable max length (append a short hash if truncated)
+    if len(key) > max_len:
+        h = hashlib.sha1(key.encode("utf-8")).hexdigest()[:8]
+        key = key[: max_len - 9] + "-" + h
+
+    return key

aethergraph/services/redactor/simple.py

@@ -0,0 +1,16 @@
+# services/redactor/simple.py
+# PII/secret scrubbing for logs/events/artifacts
+import re
+
+
+class RegexRedactor:
+    PATTERNS = [
+        (re.compile(r"sk-[A-Za-z0-9]{20,}"), "[REDACTED:APIKEY]"),
+        (re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"), "[REDACTED:EMAIL]"),
+        (re.compile(r"\b\d{16}\b"), "[REDACTED:NUM]"),
+    ]
+
+    def scrub(self, text: str) -> str:
+        for pat, repl in self.PATTERNS:
+            text = pat.sub(repl, text)
+        return text

aethergraph/services/registry/key_parsing.py

@@ -0,0 +1,44 @@
+from .registry_key import _REG_PREFIX, NS, Key
+
+
+def parse_ref(ref: str) -> Key:
+    """
+    Parse "<nspace>:<name>[@<version>]" or "<name>[@<version>]" (defaults to tool).
+    Also accepts "registry:<...>" prefix.
+
+    Examples:
+      "tool:my_tool@0.1.0"
+      "graph:my_graph"
+      "agent:router@latest"
+      "my_tool@0.1.0"           # -> tool
+      "registry:tool:my_tool@0.1.0"
+      "registry:my_tool@0.1.0"  # -> tool
+    """
+    if not ref:
+        raise ValueError("Empty ref")
+    m = _REG_PREFIX.match(ref)
+    s = m.group(1) if m else ref
+
+    # If a namespace is present, it looks like "ns:name..."
+    if ":" in s:
+        nspace, rest = s.split(":", 1)
+        if nspace not in NS:
+            # If the left side is not a namespace, treat whole thing as name (default to tool)
+            nspace, rest = "tool", s
+    else:
+        nspace, rest = "tool", s  # default namespace
+
+    if "@" in rest:
+        name, ver = rest.split("@", 1)
+        ver = ver or None
+    else:
+        name, ver = rest, None
+
+    # normalize @latest → None (caller treats None as "pick latest")
+    if ver and ver.lower() == "latest":
+        ver = None
+
+    if not name:
+        raise ValueError(f"Invalid ref (missing name): {ref}")
+
+    return Key(nspace=nspace, name=name, version=ver)

aethergraph/services/registry/registry_key.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+import re
+
+NS = {"tool", "graph", "graphfn", "agent"}
+
+# Simple ref regex to detect optional leading 'registry:'
+_REG_PREFIX = re.compile(r"^registry:(.+)$", re.I)
+
+
+@dataclass(frozen=True)
+class Key:
+    nspace: str
+    name: str
+    version: str | None = None  # None or "latest" means resolve latest
+
+    def canonical(self) -> str:
+        ver = self.version
+        # Normalize "latest" to omitted for display
+        return f"{self.nspace}:{self.name}" + (f"@{ver}" if ver and ver.lower() != "latest" else "")