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

aethergraph/services/rag/facade.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 import hashlib
 import json
@@ -9,6 +10,7 @@ import time
 from typing import Any
 
 from aethergraph.contracts.services.llm import LLMClientProtocol
+from aethergraph.services.scope.scope import Scope
 
 from .chunker import TextSplitter
 from .utils.hybrid import topk_fuse
@@ -94,7 +96,13 @@ class RAGFacade:
         return os.path.join(self.root, make_fs_key(corpus_id))
 
     # ---------- ingestion ----------
-    async def add_corpus(self, corpus_id: str, meta: dict[str, Any] | None = None):
+    async def add_corpus(
+        self,
+        corpus_id: str,
+        meta: dict[str, Any] | None = None,
+        *,
+        scope_labels: dict[str, str] | None = None,
+    ):
         """Create a new corpus with optional metadata.
         Args:
             corpus_id: Unique identifier for the corpus.
@@ -104,18 +112,27 @@ class RAGFacade:
         os.makedirs(p, exist_ok=True)
         meta_path = os.path.join(p, "corpus.json")
         if not os.path.exists(meta_path):
+            full_meta = {
+                "corpus_id": corpus_id,
+                "fs_key": make_fs_key(corpus_id),
+                "created_at": _now_iso(),
+                "meta": meta or {},
+            }
+            if scope_labels:
+                full_meta.setdefault("meta", {})
+                full_meta["meta"]["scope"] = dict(scope_labels)
+
             with open(meta_path, "w", encoding="utf-8") as f:
-                json.dump(
-                    {
-                        "corpus_id": corpus_id,
-                        "fs_key": make_fs_key(corpus_id),  # for reference
-                        "created_at": _now_iso(),
-                        "meta": meta or {},
-                    },
-                    f,
-                )
+                json.dump(full_meta, f)
 
-    async def upsert_docs(self, corpus_id: str, docs: list[dict[str, Any]]) -> dict[str, Any]:
+    async def upsert_docs(
+        self,
+        corpus_id: str,
+        docs: list[dict[str, Any]],
+        *,
+        scope: Scope | None = None,
+        scope_id: str | None = None,  # e.g. memory_scope_id if tied to memory
+    ) -> dict[str, Any]:
         """Ingest and index a list of documents into the specified corpus.
         Args:
             corpus_id: The target corpus identifier.
@@ -125,10 +142,16 @@ class RAGFacade:
             - File-based documents: {"path": "/path/to/doc.pdf", "labels": {...}}
             - Inline text documents: {"text": "Document content...", "title": "Doc Title", "labels": {...}}
         """
+
         if not self.embed:
             raise RuntimeError("RAGFacade: embed client not configured")
 
-        await self.add_corpus(corpus_id)
+        scope_labels: dict[str, str] = {}
+        if scope is not None:
+            scope_labels = scope.rag_labels(scope_id=scope_id)
+
+        await self.add_corpus(corpus_id, meta=None, scope_labels=scope_labels)
+
         cdir = self._cdir(corpus_id)
         docs_jl = os.path.join(cdir, "docs.jsonl")
         chunks_jl = os.path.join(cdir, "chunks.jsonl")
@@ -139,7 +162,8 @@ class RAGFacade:
         total_chunks = 0
 
         for d in docs:
-            labels = d.get("labels", {})
+            # Merge scope labels into provided labels
+            labels = {**scope_labels, **(d.get("labels", {}) or {})}
             title = d.get("title") or os.path.basename(d.get("path", "")) or "untitled"
             doc_id = _stable_id({"title": title, "labels": labels, "ts": _now_iso()})
             text = None
@@ -150,13 +174,13 @@ class RAGFacade:
                 uri = await self.artifacts.save_file(
                     path=d["path"],
                     kind="doc",
-                    run_id="rag",
-                    graph_id="rag",
-                    node_id="rag",
+                    run_id=scope.run_id if scope else "rag",
+                    graph_id=scope.graph_id if scope else "rag",
+                    node_id=scope.node_id if scope else "rag",
                     tool_name="rag.upsert",
                     tool_version="0.1.0",
                     labels=labels,
-                    cleanup=False,
+                    cleanup=False,  # keep source file as this is the original
                 )
                 path = d["path"].lower()
                 if path.endswith(".pdf"):
@@ -175,8 +199,29 @@ class RAGFacade:
             else:
                 # inline text doc — persist as artifact first
                 payload = d.get("text", "")
-                uri = await self.artifacts.save_text(payload=payload)  # store as temp artifact
-                doc_uri = uri.uri if hasattr(uri, "uri") else uri
+
+                # stage and save:
+                staged = await self.artifacts.plan_staging_path(".txt")
+                payload = d.get("text", "")
+
+                def _write_staged(path: str, content: str) -> None:
+                    with open(path, "w", encoding="utf-8") as f:
+                        f.write(content)
+
+                await asyncio.to_thread(_write_staged, staged, payload)
+
+                a = await self.artifacts.save_file(
+                    path=staged,
+                    kind="doc",
+                    run_id=scope.run_id if scope else "rag",
+                    graph_id=scope.graph_id if scope else "rag",
+                    node_id=scope.node_id if scope else "rag",
+                    tool_name="rag.upsert",
+                    tool_version="0.1.0",
+                    labels=labels,
+                )
+
+                doc_uri = a.uri if hasattr(a, "uri") else a
                 text = payload
 
             text = (text or "").strip()
@@ -185,7 +230,7 @@ class RAGFacade:
                     self.logger.warning(f"RAG: empty text for doc {title}")
                 continue
 
-            # write doc record
+            # write doc record with labels including scope
             with open(docs_jl, "a", encoding="utf-8") as f:
                 f.write(
                     json.dumps(
@@ -207,6 +252,7 @@ class RAGFacade:
             chunks = self.chunker.split(text)
             if not chunks:
                 continue
+
             # batch embed
             vecs = await self.embed.embed(chunks)
             for i, (chunk_text, vec) in enumerate(zip(chunks, vecs, strict=True)):
@@ -255,6 +301,40 @@ class RAGFacade:
                 out[obj["chunk_id"]] = obj
         return out
 
+    def _apply_filters(
+        self,
+        corpus_id: str,
+        hits: list[dict[str, Any]],
+        filters: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Apply filters to the search hits."""
+        if not filters:
+            return hits
+
+        # We need labels to test filters. They are in meta["labels"] for each chunk.
+        # hits come from index.search as [{"chunk_id", "score", "meta": {...}}, ...].
+        # It works as follows:
+        # 1. For each hit, we extract the labels from the meta information.
+        # 2. We then check if the labels match the desired filters.
+        out = []
+        for h in hits:
+            meta = h.get("meta", {}) or {}
+            labels = meta.get("labels", {}) or {}
+            ok = True
+            for k, want in filters.items():
+                val = labels.get(k)
+                if isinstance(want, list | tuple | set):
+                    if val not in want:
+                        ok = False
+                        break
+                else:
+                    if val != want:
+                        ok = False
+                        break
+            if ok:
+                out.append(h)
+        return out
+
     async def search(
         self,
         corpus_id: str,
@@ -277,7 +357,13 @@ class RAGFacade:
         # dense search via index then optional lexical fusion
         qvec = (await self.embed.embed([query]))[0]
         dense_hits = await self.index.search(corpus_id, qvec, max(24, k))
+
+        # apply filters before fusion
+        dense_hits = self._apply_filters(corpus_id, dense_hits, filters=filters)
+
         chunks_map = self._load_chunks_map(corpus_id)
+
+        # if only dense or no hits, return directly
         if mode == "dense" or not dense_hits:
             dense_hits = dense_hits[:k]
             return [
@@ -292,6 +378,7 @@ class RAGFacade:
                 for h in dense_hits
             ]
 
+        # hybrid fusion: i.e. dense + lexical
         fused = topk_fuse(
             query, dense_hits, {cid: rec.get("text", "") for cid, rec in chunks_map.items()}, k
         )
@@ -310,6 +397,38 @@ class RAGFacade:
             )
         return out
 
+    async def search_scoped(
+        self,
+        *,
+        curpus_id: str,
+        query: str,
+        scope: Scope | None = None,
+        scope_id: str | None = None,  # e.g. memory_scope_id if tied to memory, can be None
+        k: int = 8,
+        mode: str = "hybrid",
+    ) -> list[SearchHit]:
+        """
+        Convenience wrapper to search with scope-based filters.
+        Args:
+            curpus_id: Target corpus identifier.
+            query: The search query string.
+            scope: Scope object for filtering.
+            k: Number of top results to return.
+            mode: Search mode - "dense", "hybrid".
+        """
+        filters: dict[str, Any] | None = None
+        if scope is not None:
+            # build filters from scope labels
+            filters = scope.rag_filter(scope_id=scope_id)  # scope_id is optional
+
+        return await self.search(
+            curpus_id,
+            query,
+            k=k,
+            filters=filters,
+            mode=mode,
+        )
+
     async def retrieve(
         self, corpus_id: str, query: str, k: int = 6, rerank: bool = True
     ) -> list[SearchHit]:
@@ -320,6 +439,9 @@ class RAGFacade:
             k: Number of top results to return.
             rerank: Whether to rerank results using hybrid scoring.
         """
+        print(
+            f"🍏 RAGFacade.retrieve: corpus_id={corpus_id}, query={query}, k={k}, rerank={rerank}"
+        )
         # For now, rerank flag is ignored; fused hybrid already sorts reasonably.
         return await self.search(corpus_id, query, k=k, mode="hybrid")
 
@@ -332,6 +454,8 @@ class RAGFacade:
         style: str = "concise",
         with_citations: bool = True,
         k: int = 6,
+        scope: Scope | None = None,
+        scope_id: str | None = None,  # e.g. memory_scope_id if tied to memory, can be None
     ) -> dict[str, Any]:
         """Answer a question using retrieved context from the corpus.
         Args:
@@ -346,7 +470,19 @@ class RAGFacade:
             # use default LLM client
             llm = self.llm
 
-        hits = await self.retrieve(corpus_id, question, k=k, rerank=True)
+        filters: dict[str, Any] | None = None
+        if scope is not None:
+            # build filters from scope labels
+            filters = scope.rag_filter(scope_id=scope_id)  # scope_id is optional
+
+        hits = await self.search(
+            corpus_id,
+            question,
+            k=k,
+            filters=filters,
+            mode="hybrid",
+        )
+
         context = "\n\n".join([f"[{i + 1}] {h.text}" for i, h in enumerate(hits)])
         sys = "You answer strictly from the provided context. Cite chunk numbers like [1],[2]. If insufficient, say you don't know."
         if style == "detailed":
@@ -406,6 +542,26 @@ class RAGFacade:
         return out
 
     async def list_corpora(self) -> list[dict]:
+        """
+        List all available corpora managed by this RAGFacade.
+
+        This method scans the corpus root directory, loads metadata for each corpus,
+        and returns a list of corpus records with their logical IDs and metadata.
+
+        Examples:
+            Basic usage to enumerate corpora:
+            ```python
+            corpora = await context.rag().list_corpora()
+            for c in corpora:
+                print(c["corpus_id"], c["meta"].get("created_at"))
+            ```
+
+        Returns:
+            list[dict]: A list of dictionaries, each containing:
+
+                - "corpus_id": The logical identifier for the corpus.
+                - "meta": The metadata dictionary loaded from corpus.json (may be empty).
+        """
         out = []
         for d in sorted(os.listdir(self.root)):
             # cdir = self._cdir(d)
@@ -428,6 +584,32 @@ class RAGFacade:
     async def list_docs(
         self, corpus_id: str, limit: int = 200, after: str | None = None
     ) -> list[dict]:
+        """
+        List documents from a corpus in a paginated fashion.
+
+        This method reads documents from the `docs.jsonl` file associated with the given `corpus_id`,
+        returning up to `limit` documents after the specified `after` document ID.
+        It is typically accessed via `context.rag().list_docs(...)`.
+
+        Examples:
+            Basic usage to list the first 100 documents:
+            ```python
+            docs = await context.rag().list_docs("my-corpus", limit=100)
+            ```
+
+            Paginating after a specific document:
+            ```python
+            docs = await context.rag().list_docs("my-corpus", after="doc_123")
+            ```
+
+        Args:
+            corpus_id: The unique identifier for the corpus whose documents are to be listed.
+            limit: The maximum number of documents to return (default: 200).
+            after: If provided, only documents after this document ID will be returned.
+
+        Returns:
+            list[dict]: A list of document objects, each represented as a dictionary.
+        """
         cdir = self._cdir(corpus_id)
         docs_jl = os.path.join(cdir, "docs.jsonl")
         if not os.path.exists(docs_jl):
@@ -450,7 +632,32 @@ class RAGFacade:
 
     async def delete_docs(self, corpus_id: str, doc_ids: list[str]) -> dict:
         """
-        Removes docs from docs.jsonl and any chunks in chunks.jsonl; asks the index to drop vectors if supported.
+        Remove one or more documents and their associated chunks from a corpus.
+
+        This method deletes all records for the specified `doc_ids` from both the `docs.jsonl`
+        and `chunks.jsonl` files within the given corpus. It also instructs the vector index
+        backend to remove any vectors associated with the deleted chunks, if supported.
+
+        Examples:
+            Basic usage to delete a single document:
+            ```python
+            await context.rag().delete_docs("my-corpus", ["doc_123"])
+            ```
+
+            Deleting multiple documents at once:
+            ```python
+            await context.rag().delete_docs("my-corpus", ["doc_1", "doc_2", "doc_3"])
+            ```
+
+        Args:
+            corpus_id: The unique identifier for the corpus from which documents will be removed.
+            doc_ids: A list of document IDs to delete. All chunks belonging to these documents
+                will also be removed.
+
+        Returns:
+            dict: A dictionary containing:
+                - "removed_docs": The number of documents removed.
+                - "removed_chunks": The number of chunks removed from the index and storage.
         """
         cdir = self._cdir(corpus_id)
         docs_jl = os.path.join(cdir, "docs.jsonl")
@@ -495,7 +702,32 @@ class RAGFacade:
         self, corpus_id: str, *, doc_ids: list[str] | None = None, batch: int = 64
     ) -> dict:
         """
-        Re-embeds selected docs (or all) and re-adds vectors. Uses the configured embed client or a model override if your client supports it.
+        Re-embed vectors for selected documents (or all) in a corpus.
+
+        This method re-computes embeddings for all chunks belonging to the specified `doc_ids`
+        (or for all documents if `doc_ids` is None) and updates the vector index accordingly.
+        It uses the currently configured embedding client and can be accessed via `context.rag().reembed(...)`.
+
+        Examples:
+            Re-embed all documents in a corpus:
+            ```python
+            await context.rag().reembed("my-corpus")
+            ```
+
+            Re-embed only specific documents:
+            ```python
+            await context.rag().reembed("my-corpus", doc_ids=["doc_123", "doc_456"])
+            ```
+
+        Args:
+            corpus_id: The unique identifier for the corpus whose vectors will be re-embedded.
+            doc_ids: Optional list of document IDs to re-embed. If None, all documents are processed.
+            batch: The number of chunks to embed per batch (default: 64).
+
+        Returns:
+            dict: A dictionary containing:
+                - "reembedded": The number of chunks re-embedded.
+                - "model": The embedding model used (if available).
         """
         cdir = self._cdir(corpus_id)
         chunks_jl = os.path.join(cdir, "chunks.jsonl")
@@ -526,6 +758,30 @@ class RAGFacade:
         return {"reembedded": added, "model": getattr(embed, "embed_model", None)}
 
     async def stats(self, corpus_id: str) -> dict:
+        """
+        Retrieve summary statistics for a given corpus.
+
+        This method counts the number of documents and chunks in the specified corpus,
+        and loads the associated corpus metadata. It is typically accessed via
+        `context.rag().stats(...)`.
+
+        Examples:
+            Basic usage to get corpus statistics:
+            ```python
+            stats = await context.rag().stats("my-corpus")
+            print(stats["docs"], stats["chunks"])
+            ```
+
+        Args:
+            corpus_id: The unique identifier for the corpus whose statistics are to be retrieved.
+
+        Returns:
+            dict: A dictionary containing:
+                - "corpus_id": The logical identifier for the corpus.
+                - "docs": The number of documents in the corpus.
+                - "chunks": The number of text chunks in the corpus.
+                - "meta": The metadata dictionary loaded from corpus.json (may be empty).
+        """
         cdir = self._cdir(corpus_id)
         docs_jl = os.path.join(cdir, "docs.jsonl")
         chunks_jl = os.path.join(cdir, "chunks.jsonl")

aethergraph/services/rag/index_factory.py

@@ -29,7 +29,7 @@ def create_vector_index(
         # try FAISS, fallback to sqlite with a warning
         try:
             require("faiss", "faiss")  # faiss-cpu exposes module 'faiss'
-            from .index.faiss_index import FAISSVectorIndex
+            from aethergraph.storage.vector_index.faiss_index import FAISSVectorIndex
 
             path = (
                 str(Path(index_path) / "faiss")
@@ -42,7 +42,7 @@ def create_vector_index(
             backend = "sqlite"
 
     # sqlite (default)
-    from .index.sqlite_index import SQLiteVectorIndex
+    from aethergraph.storage.vector_index.sqlite_index import SQLiteVectorIndex
 
     path = (
         str(Path(index_path) / "sqlite")