aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

aethergraph/storage/sessions/sqlite_session_store.py

@@ -0,0 +1,399 @@
+# aethergraph/storage/sqlite_session_store.py
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Sequence
+from datetime import datetime, timezone
+import json
+from pathlib import Path
+import sqlite3
+import threading
+from typing import Any
+import uuid
+
+from aethergraph.api.v1.schemas import Session
+from aethergraph.contracts.services.sessions import SessionStore
+from aethergraph.core.runtime.run_types import SessionKind
+
+
+def _dt_to_ts(dt: datetime | None) -> float | None:
+    if dt is None:
+        return None
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=timezone.utc)
+    return dt.timestamp()
+
+
+def _parse_dt(val: Any) -> datetime | None:
+    if val is None:
+        return None
+    if isinstance(val, datetime):
+        return val
+    if isinstance(val, str):
+        try:
+            return datetime.fromisoformat(val)
+        except Exception:
+            return None
+    if isinstance(val, int | float):
+        try:
+            return datetime.fromtimestamp(float(val), tz=timezone.utc)
+        except Exception:
+            return None
+    return None
+
+
+def _session_to_doc(sess: Session) -> dict[str, Any]:
+    # Support both Pydantic v1 (.dict) and v2 (.model_dump)
+    data = sess.model_dump() if hasattr(sess, "model_dump") else sess.dict()
+
+    # Normalize datetimes to ISO for JSON
+    for key in ("created_at", "updated_at", "last_artifact_at"):
+        if isinstance(data.get(key), datetime):
+            data[key] = data[key].isoformat()
+    return data
+
+
+def _doc_to_session(doc: dict[str, Any]) -> Session:
+    # Convert ISO/ts back to datetime
+    for key in ("created_at", "updated_at", "last_artifact_at"):
+        if key in doc:
+            parsed = _parse_dt(doc[key])
+            if parsed is not None:
+                doc[key] = parsed
+
+    # Normalize kind if stored as str
+    if "kind" in doc and isinstance(doc["kind"], str):
+        try:
+            doc["kind"] = SessionKind(doc["kind"])
+        except Exception:
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Unknown SessionKind stored in DB: {doc['kind']}")
+
+    return Session(**doc)
+
+
+class SQLiteSessionStoreSync:
+    """
+    SQLite-backed SessionStore.
+
+    - Stores full Session as JSON in `data_json`
+    - Promotes session_id, kind, user_id, org_id, created_at, updated_at,
+      artifact_count, last_artifact_at to columns for fast listing / stats.
+    """
+
+    def __init__(self, path: str):
+        path_obj = Path(path)
+        path_obj.parent.mkdir(parents=True, exist_ok=True)
+
+        self._db = sqlite3.connect(
+            str(path_obj),
+            check_same_thread=False,
+            isolation_level=None,
+        )
+        self._db.execute("PRAGMA journal_mode=WAL;")
+        self._db.execute("PRAGMA synchronous=NORMAL;")
+
+        # Base table
+        self._db.execute(
+            """
+            CREATE TABLE IF NOT EXISTS sessions (
+                session_id        TEXT PRIMARY KEY,
+                data_json         TEXT NOT NULL,
+                kind              TEXT NOT NULL,
+                user_id           TEXT,
+                org_id            TEXT,
+                created_at        REAL NOT NULL,
+                updated_at        REAL NOT NULL,
+                artifact_count    INTEGER NOT NULL DEFAULT 0,
+                last_artifact_at  REAL
+            )
+            """
+        )
+
+        # Indices
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_sessions_user_updated ON sessions(user_id, updated_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_sessions_org_updated ON sessions(org_id, updated_at DESC)"
+        )
+        self._db.execute(
+            "CREATE INDEX IF NOT EXISTS idx_sessions_kind_updated ON sessions(kind, updated_at DESC)"
+        )
+
+        self._lock = threading.RLock()
+
+    # -------- core helpers --------
+
+    def _upsert(self, sess: Session) -> Session:
+        doc = _session_to_doc(sess)
+        payload = json.dumps(doc, ensure_ascii=False)
+
+        created_ts = _dt_to_ts(sess.created_at)
+        updated_ts = _dt_to_ts(sess.updated_at)
+        last_art_ts = _dt_to_ts(sess.last_artifact_at)
+        artifact_count = sess.artifact_count or 0
+
+        kind_val = sess.kind.value if isinstance(sess.kind, SessionKind) else str(sess.kind)
+
+        with self._lock:
+            self._db.execute(
+                """
+                INSERT INTO sessions (
+                    session_id, data_json,
+                    kind, user_id, org_id,
+                    created_at, updated_at,
+                    artifact_count, last_artifact_at
+                )
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                ON CONFLICT(session_id) DO UPDATE SET
+                    data_json        = excluded.data_json,
+                    kind             = excluded.kind,
+                    user_id          = excluded.user_id,
+                    org_id           = excluded.org_id,
+                    created_at       = excluded.created_at,
+                    updated_at       = excluded.updated_at,
+                    artifact_count   = excluded.artifact_count,
+                    last_artifact_at = excluded.last_artifact_at
+                """,
+                (
+                    sess.session_id,
+                    payload,
+                    kind_val,
+                    sess.user_id,
+                    sess.org_id,
+                    created_ts,
+                    updated_ts,
+                    artifact_count,
+                    last_art_ts,
+                ),
+            )
+        return sess
+
+    # -------- SessionStore-style API (sync) --------
+
+    def create(
+        self,
+        *,
+        kind: SessionKind,
+        user_id: str | None = None,
+        org_id: str | None = None,
+        title: str | None = None,
+        source: str = "webui",
+        external_ref: str | None = None,
+    ) -> Session:
+        now = datetime.now(timezone.utc)
+        sess = Session(
+            session_id=str(uuid.uuid4()),
+            kind=kind,
+            title=title,
+            user_id=user_id,
+            org_id=org_id,
+            source=source,
+            external_ref=external_ref,
+            created_at=now,
+            updated_at=now,
+            artifact_count=0,
+            last_artifact_at=None,
+        )
+        return self._upsert(sess)
+
+    def get(self, session_id: str) -> Session | None:
+        with self._lock:
+            row = self._db.execute(
+                "SELECT data_json FROM sessions WHERE session_id = ?",
+                (session_id,),
+            ).fetchone()
+        if not row:
+            return None
+        doc = json.loads(row[0])
+        return _doc_to_session(doc)
+
+    def delete(self, session_id: str) -> None:
+        with self._lock:
+            self._db.execute(
+                "DELETE FROM sessions WHERE session_id = ?",
+                (session_id,),
+            )
+
+    def list_for_user(
+        self,
+        *,
+        user_id: str | None,
+        org_id: str | None = None,
+        kind: SessionKind | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> list[Session]:
+        where: list[str] = []
+        params: list[Any] = []
+
+        if user_id is not None:
+            where.append("user_id = ?")
+            params.append(user_id)
+        if org_id is not None:
+            where.append("org_id = ?")
+            params.append(org_id)
+        if kind is not None:
+            where.append("kind = ?")
+            params.append(kind.value if isinstance(kind, SessionKind) else str(kind))
+
+        sql = "SELECT data_json FROM sessions"
+        if where:
+            sql += " WHERE " + " AND ".join(where)
+        sql += " ORDER BY updated_at DESC"
+
+        sql += " LIMIT ? OFFSET ?"
+        params.extend([limit, offset])
+
+        with self._lock:
+            rows = self._db.execute(sql, params).fetchall()
+
+        return [_doc_to_session(json.loads(r[0])) for r in rows]
+
+    def touch(self, session_id: str, *, updated_at: datetime | None = None) -> None:
+        sess = self.get(session_id)
+        if not sess:
+            return
+        sess.updated_at = updated_at or datetime.now(timezone.utc)
+        self._upsert(sess)
+
+    def update(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        external_ref: str | None = None,
+    ) -> Session | None:
+        sess = self.get(session_id)
+        if not sess:
+            return None
+        if title is not None:
+            sess.title = title
+        if external_ref is not None:
+            sess.external_ref = external_ref
+        sess.updated_at = datetime.now(timezone.utc)
+        return self._upsert(sess)
+
+    def record_artifact(
+        self,
+        session_id: str,
+        *,
+        created_at: datetime | None = None,
+    ) -> None:
+        """
+        Optional API used by ArtifactFacade._record via getattr(..., 'record_artifact').
+        Updates artifact_count + last_artifact_at + updated_at.
+        """
+        sess = self.get(session_id)
+        if not sess:
+            return
+
+        ts = created_at or datetime.now(timezone.utc)
+
+        sess.artifact_count = (sess.artifact_count or 0) + 1
+        if sess.last_artifact_at is None or ts > sess.last_artifact_at:
+            sess.last_artifact_at = ts
+
+        # For UI, bump updated_at as well
+        if ts > sess.updated_at:
+            sess.updated_at = ts
+        else:
+            sess.updated_at = datetime.now(timezone.utc)
+
+        self._upsert(sess)
+
+
+class SQLiteSessionStore(SessionStore):
+    """
+    Async SessionStore implementation backed by SQLiteSessionStoreSync.
+    """
+
+    def __init__(self, path: str):
+        self._sync = SQLiteSessionStoreSync(path)
+
+    async def create(
+        self,
+        *,
+        kind: SessionKind,
+        user_id: str | None = None,
+        org_id: str | None = None,
+        title: str | None = None,
+        source: str = "webui",
+        external_ref: str | None = None,
+    ) -> Session:
+        # Delegate to sync create (which already constructs Session correctly)
+        return await asyncio.to_thread(
+            self._sync.create,
+            kind=kind,
+            user_id=user_id,
+            org_id=org_id,
+            title=title,
+            source=source,
+            external_ref=external_ref,
+        )
+
+    async def get(self, session_id: str) -> Session | None:
+        return await asyncio.to_thread(self._sync.get, session_id)
+
+    async def list_for_user(
+        self,
+        *,
+        user_id: str | None,
+        org_id: str | None = None,
+        kind: SessionKind | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> Sequence[Session]:
+        return await asyncio.to_thread(
+            self._sync.list_for_user,
+            user_id=user_id,
+            org_id=org_id,
+            kind=kind,
+            limit=limit,
+            offset=offset,
+        )
+
+    async def touch(
+        self,
+        session_id: str,
+        *,
+        updated_at: datetime | None = None,
+    ) -> None:
+        await asyncio.to_thread(self._sync.touch, session_id, updated_at=updated_at)
+
+    async def update(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        external_ref: str | None = None,
+    ) -> Session | None:
+        return await asyncio.to_thread(
+            self._sync.update,
+            session_id,
+            title=title,
+            external_ref=external_ref,
+        )
+
+    async def delete(self, session_id: str) -> None:
+        await asyncio.to_thread(self._sync.delete, session_id)
+
+    async def record_artifact(
+        self,
+        session_id: str,
+        *,
+        created_at: datetime | None = None,
+    ) -> None:
+        """
+        Optional method, called via getattr(..., 'record_artifact', None)
+        from ArtifactFacade._record.
+        """
+        await asyncio.to_thread(
+            self._sync.record_artifact,
+            session_id,
+            created_at=created_at,
+        )

aethergraph/storage/vector_index/chroma_index.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+from aethergraph.contracts.storage.vector_index import VectorIndex
+
+try:
+    import chromadb
+except Exception:
+    chromadb = None
+
+"""Adapter over a chromadb.Client.
+NOTE: this is just a stub implementation to show how vector indexes can be adapted.
+We can use the same VectorIndex interface for different backends with adaptation.
+
+e.g.
+- QdrantVectorIndex: corpus_id -> Qdrant collection, chunk_id -> point ID, vector + metadata stored per point.
+- PineconeVectorIndex: corpus_id -> Pinecone index, chunk_id -> vector ID, vector + metadata stored per vector.
+- WeaviateVectorIndex: corpus_id -> Weaviate class, chunk_id -> object ID, vector + metadata stored per object.
+"""
+
+
+class ChromaVectorIndex(VectorIndex):
+    """
+    Adapter over a chromadb.Client.
+    Each `corpus_id` is mapped to a Chroma collection:
+        collection_name = f"{collection_prefix}{corpus_id}"
+    """
+
+    def __init__(
+        self,
+        client: chromadb.ClientAPI,
+        *,
+        collection_prefix: str = "rag_",
+    ):
+        if chromadb is None:
+            raise RuntimeError("ChromaVectorIndex requires `chromadb` to be installed.")
+        self._client = client
+        self._prefix = collection_prefix
+
+    def _collection_name(self, corpus_id: str) -> str:
+        return f"{self._prefix}{corpus_id}"
+
+    def _collection(self, corpus_id: str):
+        name = self._collection_name(corpus_id)
+        return self._client.get_or_create_collection(name=name)
+
+    async def add(
+        self,
+        corpus_id: str,
+        chunk_ids: list[str],
+        vectors: list[list[float]],
+        metas: list[dict[str, Any]],
+    ) -> None:
+        if not chunk_ids:
+            return
+        col = self._collection(corpus_id)
+        # Chroma's API is sync, but fast; you can wrap in to_thread if needed.
+        col.upsert(
+            ids=chunk_ids,  # type: IDs
+            embeddings=vectors,  # type: Embeddings
+            metadatas=metas,  # type: Metadatas
+        )
+
+    async def delete(self, corpus_id: str, chunk_ids: list[str] | None = None) -> None:
+        name = self._collection_name(corpus_id)
+        if chunk_ids is None:
+            # Drop whole collection
+            try:
+                self._client.delete_collection(name=name)
+            except Exception as err:
+                raise RuntimeError(f"Failed to delete collection '{name}'") from err
+            return
+
+        try:
+            col = self._client.get_collection(name=name)
+        except Exception:
+            return
+        col.delete(ids=chunk_ids)
+
+    async def search(
+        self,
+        corpus_id: str,
+        query_vec: list[float],
+        k: int,
+    ) -> list[dict[str, Any]]:
+        try:
+            col = self._collection(corpus_id)
+        except Exception:
+            return []
+
+        def _search_sync() -> list[dict[str, Any]]:
+            # Chroma's client is sync; do this in a thread so we don't block the loop.
+            res = col.query(
+                query_embeddings=[query_vec],
+                n_results=k,
+            )
+
+            # Chroma returns lists-of-lists
+            ids = (res.get("ids") or [[]])[0]
+            dists = (res.get("distances") or [[]])[0]
+            metas = (res.get("metadatas") or [[]])[0]
+
+            out: list[dict[str, Any]] = []
+            for cid, dist, meta in zip(ids, dists, metas, strict=True):
+                # Chroma distance is "smaller is better".
+                # We convert to a "score" where larger is better.
+                # Any monotone transform is fine; here we use negative distance.
+                score = float(-dist)
+                out.append(
+                    {
+                        "chunk_id": cid,
+                        "score": score,
+                        "meta": meta or {},
+                    }
+                )
+            return out
+
+        return await asyncio.to_thread(_search_sync)
+
+    async def list_corpora(self) -> list[str]:
+        cols = self._client.list_collections()
+        out = []
+        for c in cols:
+            name = c.name
+            if name.startswith(self._prefix):
+                out.append(name[len(self._prefix) :])
+        return out
+
+    async def list_chunks(self, corpus_id: str) -> list[str]:
+        try:
+            col = self._collection(corpus_id)
+        except Exception:
+            return []
+        # Chroma `get()` can be expensive for huge collections; fine for small/mid-scale.
+        res = col.get()
+        return list(res.get("ids") or [])

aethergraph/storage/vector_index/faiss_index.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import asyncio
+from contextlib import suppress
+from pathlib import Path
+import pickle
+import threading
+from typing import Any
+
+import numpy as np
+
+from aethergraph.contracts.storage.vector_index import VectorIndex
+
+try:
+    import faiss  # type: ignore
+except Exception:
+    faiss = None
+
+
+class FAISSVectorIndex(VectorIndex):
+    """
+    Simple FAISS-backed index, one .index + .meta.pkl per corpus_id.
+    Uses cosine similarity via normalized vectors (IndexFlatIP).
+    """
+
+    def __init__(self, root: str, dim: int | None = None):
+        if faiss is None:
+            raise RuntimeError("FAISSVectorIndex requires `faiss` to be installed.")
+        self.root = Path(root)
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.dim = dim  # optional default dimension
+        self._lock = threading.RLock()
+
+    def _paths(self, corpus_id: str) -> tuple[Path, Path]:
+        base = self.root / corpus_id
+        return base.with_suffix(".index"), base.with_suffix(".meta.pkl")
+
+    def _load_sync(self, corpus_id: str):
+        idx_path, meta_path = self._paths(corpus_id)
+        if not (idx_path.exists() and meta_path.exists()):
+            return None, []
+        with self._lock:
+            index = faiss.read_index(str(idx_path))
+            with meta_path.open("rb") as f:
+                metas = pickle.load(f)
+        return index, metas
+
+    def _save_sync(self, corpus_id: str, index, metas: list[dict[str, Any]]) -> None:
+        idx_path, meta_path = self._paths(corpus_id)
+        self.root.mkdir(parents=True, exist_ok=True)
+        with self._lock:
+            faiss.write_index(index, str(idx_path))
+            with meta_path.open("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]],
+    ) -> None:
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+        if not chunk_ids:
+            return
+
+        vecs = np.asarray(vectors, dtype=np.float32)
+        # Normalize for cosine
+        norms = np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-9
+        vecs = vecs / norms
+
+        def _add_sync():
+            index, old_metas = self._load_sync(corpus_id)
+            d = vecs.shape[1]
+            if index is None:
+                # If dim was provided, optionally sanity-check
+                if self.dim is not None and self.dim != d:
+                    raise ValueError(f"FAISSVectorIndex: dim mismatch {self.dim} vs {d}")
+                index = faiss.IndexFlatIP(d)
+                old_metas = []
+                self.dim = d
+
+            index.add(vecs)
+            for cid, m in zip(chunk_ids, metas, strict=True):
+                old_metas.append({"chunk_id": cid, "meta": m})
+            self._save_sync(corpus_id, index, old_metas)
+
+        await asyncio.to_thread(_add_sync)
+
+    async def delete(self, corpus_id: str, chunk_ids: list[str] | None = None) -> None:
+        if chunk_ids is None:
+            # Delete entire corpus
+            idx_path, meta_path = self._paths(corpus_id)
+            for p in (idx_path, meta_path):
+                with suppress(Exception):
+                    p.unlink(missing_ok=True)
+            return
+
+        # Selective delete is tricky with plain FAISS; we’d need to rebuild.
+        # For now, keep explicit about limitations:
+        async def _delete_sync():
+            index, metas = self._load_sync(corpus_id)
+            if index is None:
+                return
+            keep_idxs = [i for i, m in enumerate(metas) if m["chunk_id"] not in set(chunk_ids)]
+            if not keep_idxs:
+                # Remove whole corpus
+                idx_path, meta_path = self._paths(corpus_id)
+                for p in (idx_path, meta_path):
+                    try:
+                        p.unlink(missing_ok=True)
+                    except Exception as e:
+                        import logging
+
+                        logger = logging.getLogger("aethergraph.storage.vector_index.faiss_index")
+                        logger.error(f"Failed to delete {p}: {e}")
+                return
+
+            # Rebuild index with kept vectors
+            # NOTE: we do NOT have original vectors here,
+            # so in a minimal implementation we simply raise.
+            raise NotImplementedError(
+                "FAISSVectorIndex: selective delete requires storing vectors; "
+                "either extend metadata to keep vectors, or rebuild from source."
+            )
+
+        await asyncio.to_thread(_delete_sync)
+
+    async def list_chunks(self, corpus_id: str) -> list[str]:
+        def _list_sync() -> list[str]:
+            _, metas = self._load_sync(corpus_id)
+            return [m["chunk_id"] for m in metas] if metas else []
+
+        return await asyncio.to_thread(_list_sync)
+
+    async def list_corpora(self) -> list[str]:
+        def _scan() -> list[str]:
+            out = []
+            for p in self.root.glob("*.meta.pkl"):
+                out.append(p.stem)  # strip .meta.pkl => corpus_id
+            return out
+
+        return await asyncio.to_thread(_scan)
+
+    async def search(
+        self,
+        corpus_id: str,
+        query_vec: list[float],
+        k: int,
+    ) -> list[dict[str, Any]]:
+        if faiss is None:
+            raise RuntimeError("FAISS not installed")
+
+        q = np.asarray([query_vec], dtype=np.float32)
+        q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-9)
+
+        def _search_sync() -> list[dict[str, Any]]:
+            index, metas = self._load_sync(corpus_id)
+            if index is None or not metas:
+                return []
+            D, I = index.search(q, k)  # noqa: E741
+            out: list[dict[str, Any]] = []
+            scores = D[0].tolist()
+            idxs = I[0].tolist()
+            for score, idx in zip(scores, idxs, strict=True):
+                if idx < 0 or idx >= len(metas):
+                    continue
+                m = metas[idx]
+                out.append(
+                    {
+                        "chunk_id": m["chunk_id"],
+                        "score": float(score),
+                        "meta": m["meta"],
+                    }
+                )
+            return out
+
+        return await asyncio.to_thread(_search_sync)