@agentikos/omega-os 0.1.0 → 0.2.0

package/omega/Agentik_Engine/omega_engine/rag/corrective.py

@@ -0,0 +1,119 @@
+"""Corrective RAG (CRAG) — the envelope every strategy wears.
+
+For each retrieval, the provider (role "rag-grader") scores every doc 0-100.
+If the average is below `threshold`, the provider (role "rag-agent") produces
+a refined query and we re-retrieve. Bounded retries — never an infinite loop.
+
+This is the user's explicit design: corrective wraps the chosen strategy.
+The router resolves WHICH retriever runs; the corrective layer ensures the
+RESULT meets a quality floor.
+"""
+from __future__ import annotations
+
+from omega_engine.provider import AgentProvider, AgentRequest
+from omega_engine.rag.base import Document, RetrievalResult, Retriever
+
+
+class CorrectiveRetriever:
+    """Grade → if low, refine query and retry. Bounded by `max_retries`."""
+
+    strategy = "corrective"
+
+    def __init__(
+        self,
+        inner: Retriever,
+        provider: AgentProvider,
+        *,
+        threshold: float = 70.0,
+        max_retries: int = 2,
+    ) -> None:
+        self.inner = inner
+        self.provider = provider
+        self.threshold = threshold
+        self.max_retries = max_retries
+
+    def retrieve(self, query: str, k: int = 5) -> RetrievalResult:
+        current = query
+        last: RetrievalResult | None = None
+
+        for attempt in range(self.max_retries + 1):
+            result = self.inner.retrieve(current, k=k)
+            last = result
+            if not result.documents:
+                break
+
+            grade_avg = self._grade(current, result.documents)
+            # Stamp the grade onto every document so the caller can see it.
+            for d in result.documents:
+                d.metadata["grader_avg"] = grade_avg
+            result.score = max(result.score, grade_avg / 100.0)
+
+            if grade_avg >= self.threshold or attempt == self.max_retries:
+                # Either good enough OR we've used up our retries.
+                inner_strategy = result.strategy
+                return RetrievalResult(
+                    query=query,
+                    documents=result.documents,
+                    score=result.score,
+                    strategy=f"corrective+{inner_strategy}",
+                )
+
+            # Below threshold AND retries remain — ask provider for a refined
+            # query and try again.
+            current = self._refine(query, current, result.documents)
+
+        # Defensive fallback — exhausted retries but never returned (e.g.
+        # zero docs from inner). Surface what we have, untouched.
+        if last is None:
+            return RetrievalResult(query=query, documents=[], score=0.0,
+                                   strategy="corrective+empty")
+        return RetrievalResult(
+            query=query,
+            documents=last.documents,
+            score=last.score,
+            strategy=f"corrective+{last.strategy}",
+        )
+
+    # ----- internals -----
+
+    def _grade(self, query: str, docs: list[Document]) -> float:
+        """Average grader score in 0-100. Bounded; missing scores → 0."""
+        result = self.provider.run(AgentRequest(
+            role="rag-grader",
+            prompt=f"Score relevance 0-100 for query: {query}",
+            context={
+                "query": query,
+                "documents": [
+                    {"id": d.id, "text": d.text[:300]} for d in docs
+                ],
+            },
+        ))
+        raw = (result.artifacts or {}).get("scores", [])
+        scores = [float(s) for s in raw if isinstance(s, (int, float))]
+        if not scores:
+            return 0.0
+        avg = sum(scores) / len(scores)
+        return max(0.0, min(100.0, avg))
+
+    def _refine(self, original: str, current: str, docs: list[Document]) -> str:
+        """Ask the provider for a refined query. Fallback: append context."""
+        result = self.provider.run(AgentRequest(
+            role="rag-agent",
+            prompt=(
+                f"Original query: {original}\n"
+                f"Current query: {current}\n"
+                "These docs scored below the relevance threshold. "
+                "Suggest a better query."
+            ),
+            context={
+                "original_query": original,
+                "current_query": current,
+                "snippets": [d.text[:200] for d in docs[:3]],
+            },
+        ))
+        next_q = (result.artifacts or {}).get("next_query")
+        if next_q and next_q != current:
+            return str(next_q)
+        # Provider had nothing useful — perturb the query so the inner
+        # retriever sees something different.
+        return f"{current} context"

package/omega/Agentik_Engine/omega_engine/rag/graph.py

@@ -0,0 +1,169 @@
+"""Graph retriever — entity-relation graph with depth-limited expansion.
+
+Plain Python: dict-of-sets adjacency, typed edges. Persists to JSON. The
+retrieval contract picks seed nodes from query tokens, expands `depth` hops,
+and returns each visited node as a Document.
+
+This is the "second strategy" in the multi-RAG router: when a query is
+relational ("who works on X", "what depends on Y"), the graph wins over
+hybrid.
+"""
+from __future__ import annotations
+
+import json
+import re
+from collections import defaultdict, deque
+from pathlib import Path
+from typing import Any, Iterable
+
+from omega_engine.rag.base import Document, RetrievalResult
+
+
+_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
+
+
+def _tokens(text: str) -> list[str]:
+    return [t.lower() for t in _TOKEN_RE.findall(text)]
+
+
+class GraphRetriever:
+    """In-memory entity-relation graph with JSON persistence.
+
+    Edge tuples are `(neighbour, edge_type)`. The graph is undirected for
+    expansion — adding `(a, b, "depends_on")` makes `b` a neighbour of `a`
+    AND `a` a neighbour of `b`. The edge type travels with the neighbour so
+    the caller can filter on it.
+    """
+
+    strategy = "graph"
+
+    def __init__(self, json_path: str | Path | None = None) -> None:
+        self.json_path = str(json_path) if json_path is not None else None
+        self._adj: dict[str, set[tuple[str, str]]] = defaultdict(set)
+        self._node_text: dict[str, str] = {}      # optional node text body
+        self._node_meta: dict[str, dict[str, Any]] = {}
+        if self.json_path and Path(self.json_path).exists():
+            self._load()
+
+    # ----- mutation -----
+
+    def add_node(
+        self,
+        node: str,
+        text: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        self._adj.setdefault(node, set())
+        if text is not None:
+            self._node_text[node] = text
+        if metadata is not None:
+            self._node_meta[node] = dict(metadata)
+
+    def add_edge(self, a: str, b: str, edge_type: str = "rel") -> None:
+        """Add an undirected typed edge between two nodes. Idempotent."""
+        self.add_node(a)
+        self.add_node(b)
+        self._adj[a].add((b, edge_type))
+        self._adj[b].add((a, edge_type))
+        if self.json_path:
+            self._save()
+
+    def neighbors(self, node: str, depth: int = 1) -> list[str]:
+        """BFS up to `depth` hops. Returns nodes in discovery order, excluding
+        the seed itself."""
+        if depth <= 0 or node not in self._adj:
+            return []
+        seen: set[str] = {node}
+        out: list[str] = []
+        q: deque[tuple[str, int]] = deque([(node, 0)])
+        while q:
+            cur, d = q.popleft()
+            if d == depth:
+                continue
+            for nbr, _etype in self._adj[cur]:
+                if nbr in seen:
+                    continue
+                seen.add(nbr)
+                out.append(nbr)
+                q.append((nbr, d + 1))
+        return out
+
+    # ----- retrieval -----
+
+    def retrieve(self, query: str, k: int = 5, depth: int = 2) -> RetrievalResult:
+        """Pick seeds whose name overlaps a query token; expand `depth` hops.
+
+        Score = closeness (closer hops rank higher) + seed-name overlap.
+        """
+        if not self._adj:
+            return RetrievalResult(query=query, documents=[], score=0.0,
+                                   strategy=self.strategy)
+
+        q_tokens = set(_tokens(query))
+        seeds: list[str] = []
+        for node in self._adj:
+            if any(t in node.lower() for t in q_tokens):
+                seeds.append(node)
+        if not seeds:
+            # No direct hit — fall back to seeding from every node, capped.
+            # Lets the graph still surface SOMETHING for ambient queries.
+            seeds = list(self._adj.keys())[:3]
+
+        scored: dict[str, float] = {}
+        for seed in seeds:
+            scored.setdefault(seed, max(scored.get(seed, 0.0), 1.0))
+            # BFS, score = 1 / (1 + hops)
+            seen: set[str] = {seed}
+            q: deque[tuple[str, int]] = deque([(seed, 0)])
+            while q:
+                cur, d = q.popleft()
+                if d == depth:
+                    continue
+                for nbr, _etype in self._adj[cur]:
+                    if nbr in seen:
+                        continue
+                    seen.add(nbr)
+                    s = 1.0 / (1.0 + d + 1)
+                    scored[nbr] = max(scored.get(nbr, 0.0), s)
+                    q.append((nbr, d + 1))
+
+        ranked = sorted(scored.items(), key=lambda p: p[1], reverse=True)[:k]
+        docs: list[Document] = []
+        for node, s in ranked:
+            text = self._node_text.get(node, node)
+            meta = dict(self._node_meta.get(node, {}))
+            meta["score"] = s
+            meta["edges"] = sorted(
+                (nbr, etype) for nbr, etype in self._adj[node]
+            )
+            docs.append(Document(id=node, text=text, metadata=meta))
+        agg = sum(s for _, s in ranked) / len(ranked) if ranked else 0.0
+        return RetrievalResult(query=query, documents=docs, score=agg,
+                               strategy=self.strategy)
+
+    # ----- persistence -----
+
+    def _save(self) -> None:
+        assert self.json_path is not None
+        Path(self.json_path).parent.mkdir(parents=True, exist_ok=True)
+        # Sets aren't JSON-serialisable — flatten to lists.
+        data = {
+            "adj": {
+                k: sorted([list(t) for t in v]) for k, v in self._adj.items()
+            },
+            "node_text": self._node_text,
+            "node_meta": self._node_meta,
+        }
+        Path(self.json_path).write_text(json.dumps(data, indent=2))
+
+    def _load(self) -> None:
+        assert self.json_path is not None
+        data = json.loads(Path(self.json_path).read_text())
+        self._adj = defaultdict(set, {
+            k: {tuple(t) for t in v} for k, v in data.get("adj", {}).items()
+        })
+        self._node_text = data.get("node_text", {})
+        self._node_meta = data.get("node_meta", {})
+
+    def nodes(self) -> Iterable[str]:
+        return self._adj.keys()

package/omega/Agentik_Engine/omega_engine/rag/hybrid.py

@@ -0,0 +1,205 @@
+"""Hybrid retriever — sparse FTS5 (BM25) blended with dense cosine.
+
+SQLite FTS5 gives a real BM25 score via its `bm25()` ranking function.
+For the dense leg, we use a *hashing trick* — every token is hashed into one
+of `dim` buckets, weighted by sublinear TF, then ℓ2-normalised. Cosine
+similarity is the dot product of two normalised vectors. Pure stdlib.
+
+The combined score is `alpha * dense + (1 - alpha) * sparse`, both legs
+normalised to [0, 1] across the candidate pool.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import re
+import sqlite3
+from collections import Counter
+from pathlib import Path
+
+from omega_engine.rag.base import Document, RetrievalResult
+
+
+_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
+
+
+def _tokens(text: str) -> list[str]:
+    return [t.lower() for t in _TOKEN_RE.findall(text)]
+
+
+def _hash_to_bucket(token: str, dim: int) -> int:
+    # md5 is fast and stable — we only need a uniform-ish bucket index, not
+    # cryptographic security.
+    h = hashlib.md5(token.encode("utf-8")).digest()
+    return int.from_bytes(h[:4], "big") % dim
+
+
+def _vectorise(text: str, dim: int) -> list[float]:
+    """Hashed-feature vector with sublinear TF, ℓ2-normalised."""
+    counts: Counter[int] = Counter()
+    for tok in _tokens(text):
+        counts[_hash_to_bucket(tok, dim)] += 1
+    vec = [0.0] * dim
+    for bucket, c in counts.items():
+        # 1 + log(tf) damps very common tokens — the classic tf-idf weighting.
+        vec[bucket] = 1.0 + math.log(c)
+    norm = math.sqrt(sum(v * v for v in vec))
+    if norm == 0.0:
+        return vec
+    return [v / norm for v in vec]
+
+
+def _cosine(a: list[float], b: list[float]) -> float:
+    # Both vectors are already ℓ2-normalised → cosine == dot product.
+    return sum(x * y for x, y in zip(a, b))
+
+
+def _normalise(values: list[float]) -> list[float]:
+    """Min-max scale a list into [0, 1]. Empty / flat → all zeros."""
+    if not values:
+        return values
+    lo, hi = min(values), max(values)
+    if hi - lo < 1e-12:
+        return [0.0 for _ in values]
+    return [(v - lo) / (hi - lo) for v in values]
+
+
+_SCHEMA = """
+CREATE VIRTUAL TABLE IF NOT EXISTS docs_fts USING fts5(
+    id UNINDEXED, text, tokenize = 'unicode61'
+);
+CREATE TABLE IF NOT EXISTS docs_meta (
+    id   TEXT PRIMARY KEY,
+    text TEXT NOT NULL,
+    meta TEXT NOT NULL DEFAULT '{}',
+    vec  TEXT NOT NULL
+);
+"""
+
+
+class HybridRetriever:
+    """Sparse FTS5 + dense hashed-cosine, combined per `alpha`.
+
+    `alpha = 1` → pure dense.  `alpha = 0` → pure sparse.  Default 0.5 blends.
+    """
+
+    strategy = "hybrid"
+
+    def __init__(
+        self,
+        db_path: str | Path,
+        *,
+        dim: int = 256,
+        alpha: float = 0.5,
+    ) -> None:
+        self.db_path = str(db_path)
+        self.dim = dim
+        self.alpha = alpha
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+        self._conn = sqlite3.connect(self.db_path, isolation_level=None)
+        self._conn.row_factory = sqlite3.Row
+        self._conn.execute("PRAGMA journal_mode=WAL;")
+        self._conn.execute("PRAGMA synchronous=NORMAL;")
+        self._conn.executescript(_SCHEMA)
+
+    # ----- ingest -----
+
+    def index(self, documents: list[Document]) -> int:
+        """Index a batch. Returns the number of docs written. Idempotent on id."""
+        n = 0
+        with self._conn:
+            for d in documents:
+                vec = _vectorise(d.text, self.dim)
+                self._conn.execute(
+                    "INSERT OR REPLACE INTO docs_meta (id, text, meta, vec) "
+                    "VALUES (?, ?, ?, ?)",
+                    (d.id, d.text, json.dumps(d.metadata), json.dumps(vec)),
+                )
+                # FTS5 has no native upsert — delete-then-insert is the idiom.
+                self._conn.execute("DELETE FROM docs_fts WHERE id = ?", (d.id,))
+                self._conn.execute(
+                    "INSERT INTO docs_fts (id, text) VALUES (?, ?)", (d.id, d.text)
+                )
+                n += 1
+        return n
+
+    def count(self) -> int:
+        cur = self._conn.execute("SELECT COUNT(*) AS c FROM docs_meta")
+        return int(cur.fetchone()["c"])
+
+    # ----- retrieval -----
+
+    def retrieve(self, query: str, k: int = 5) -> RetrievalResult:
+        if k <= 0 or self.count() == 0:
+            return RetrievalResult(query=query, documents=[], score=0.0,
+                                   strategy=self.strategy)
+
+        # --- sparse leg: FTS5 + bm25. bm25() returns a NEGATIVE score (lower
+        # is better), so we negate it before normalising.
+        sparse_rows: dict[str, float] = {}
+        fts_query = _fts5_safe(query)
+        if fts_query:
+            try:
+                cur = self._conn.execute(
+                    "SELECT id, bm25(docs_fts) AS score FROM docs_fts "
+                    "WHERE docs_fts MATCH ? ORDER BY score LIMIT ?",
+                    (fts_query, k * 4),
+                )
+                for row in cur.fetchall():
+                    sparse_rows[row["id"]] = -float(row["score"])
+            except sqlite3.OperationalError:
+                # malformed FTS5 query — fall back to dense-only
+                sparse_rows = {}
+
+        # --- dense leg: cosine over every indexed doc (small corpora; for
+        # larger corpora a clustered ANN goes here, contract identical).
+        qvec = _vectorise(query, self.dim)
+        dense_rows: dict[str, float] = {}
+        cur = self._conn.execute("SELECT id, vec FROM docs_meta")
+        for row in cur.fetchall():
+            v = json.loads(row["vec"])
+            dense_rows[row["id"]] = _cosine(qvec, v)
+
+        # --- combine: union of candidates, min-max each leg, alpha blend.
+        candidates = set(sparse_rows) | set(dense_rows)
+        if not candidates:
+            return RetrievalResult(query=query, documents=[], score=0.0,
+                                   strategy=self.strategy)
+
+        ids = list(candidates)
+        sparse_vec = [sparse_rows.get(i, 0.0) for i in ids]
+        dense_vec = [dense_rows.get(i, 0.0) for i in ids]
+        sparse_n = _normalise(sparse_vec)
+        dense_n = _normalise(dense_vec)
+        combined = [
+            (self.alpha * d) + ((1 - self.alpha) * s)
+            for d, s in zip(dense_n, sparse_n)
+        ]
+
+        ranked = sorted(zip(ids, combined), key=lambda p: p[1], reverse=True)[:k]
+        docs: list[Document] = []
+        for doc_id, score in ranked:
+            cur = self._conn.execute(
+                "SELECT text, meta FROM docs_meta WHERE id = ?", (doc_id,)
+            )
+            row = cur.fetchone()
+            if not row:
+                continue
+            meta = json.loads(row["meta"])
+            meta["score"] = score
+            docs.append(Document(id=doc_id, text=row["text"], metadata=meta))
+
+        agg = sum(s for _, s in ranked) / len(ranked) if ranked else 0.0
+        return RetrievalResult(query=query, documents=docs, score=agg,
+                               strategy=self.strategy)
+
+    def close(self) -> None:
+        self._conn.close()
+
+
+def _fts5_safe(query: str) -> str:
+    """Coerce a freeform query into something FTS5 can parse — strip syntax
+    chars, OR the surviving tokens together. Returns "" if nothing remains."""
+    toks = _tokens(query)
+    return " OR ".join(toks) if toks else ""

package/omega/Agentik_Engine/omega_engine/rag/multimodal.py

@@ -0,0 +1,136 @@
+"""Multimodal retriever — PDFs and images on top of a text retriever.
+
+v1 keeps the surface honest:
+  * PDFs are converted to text via `pdftotext` if the binary is on PATH,
+    otherwise the doc is registered with its path + an empty body (no fake
+    text, no hallucinated extraction).
+  * Images get a caption from the provider (role "rag-caption") and are
+    indexed by that caption + their path. If the provider returns nothing,
+    the image still indexes by filename.
+
+The underlying text index is any text-mode Retriever (typically
+HybridRetriever) — so a query "the architecture diagram explaining the
+event log" can surface an image whose caption matches.
+"""
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+
+from omega_engine.provider import AgentProvider, AgentRequest
+from omega_engine.rag.base import Document, RetrievalResult, Retriever
+
+
+def _pdftotext_available() -> bool:
+    return shutil.which("pdftotext") is not None
+
+
+def _extract_pdf_text(path: Path) -> str:
+    """Best-effort PDF → text. Returns "" when pdftotext is missing or fails.
+    NEVER fakes output (Karpathy: think before coding — no hallucination)."""
+    if not _pdftotext_available():
+        return ""
+    try:
+        out = subprocess.run(
+            ["pdftotext", "-q", str(path), "-"],
+            capture_output=True, text=True, timeout=15, check=False,
+        )
+        return out.stdout if out.returncode == 0 else ""
+    except (subprocess.TimeoutExpired, OSError):
+        return ""
+
+
+class MultimodalRetriever:
+    """Wraps a text retriever; understands `add_pdf` and `add_image`."""
+
+    strategy = "multimodal"
+
+    # `Retriever` (Protocol) needs an `index` method when we use HybridRetriever.
+    # We type-hint loosely so any object with `retrieve` works AND any
+    # object that *also* has `index` accepts new docs.
+    def __init__(
+        self,
+        inner: Retriever,
+        provider: AgentProvider | None = None,
+    ) -> None:
+        self.inner = inner
+        self.provider = provider
+        self._registered: dict[str, dict] = {}   # id → {kind, path, caption}
+
+    # ----- ingest -----
+
+    def add_pdf(self, doc_id: str, path: str | Path) -> Document:
+        """Extract text via pdftotext if available; index whatever survives."""
+        p = Path(path)
+        text = _extract_pdf_text(p)
+        meta = {
+            "modality": "pdf",
+            "source": str(p),
+            "extracted": bool(text),
+        }
+        doc = Document(id=doc_id, text=text or f"[pdf:{p.name}]", metadata=meta)
+        self._register(doc)
+        return doc
+
+    def add_image(
+        self,
+        doc_id: str,
+        path: str | Path,
+        *,
+        caption: str | None = None,
+    ) -> Document:
+        """Index an image by its caption. Caption is provider-generated
+        unless one is passed in. With no provider and no caption, we fall
+        back to the filename so the image is still discoverable."""
+        p = Path(path)
+        final_caption = caption
+        if final_caption is None and self.provider is not None:
+            result = self.provider.run(AgentRequest(
+                role="rag-caption",
+                prompt=f"Caption the image at {p}",
+                context={"path": str(p)},
+            ))
+            final_caption = (result.artifacts or {}).get("caption")
+        if not final_caption:
+            final_caption = p.stem.replace("_", " ").replace("-", " ")
+        meta = {
+            "modality": "image",
+            "source": str(p),
+            "caption": final_caption,
+        }
+        doc = Document(id=doc_id, text=final_caption, metadata=meta)
+        self._register(doc)
+        return doc
+
+    def add_text(self, doc_id: str, text: str, **meta) -> Document:
+        """Plain text passthrough — keeps the retriever uniformly usable."""
+        m = {"modality": "text", **meta}
+        doc = Document(id=doc_id, text=text, metadata=m)
+        self._register(doc)
+        return doc
+
+    def _register(self, doc: Document) -> None:
+        self._registered[doc.id] = doc.metadata
+        index = getattr(self.inner, "index", None)
+        if callable(index):
+            index([doc])
+        # If the inner retriever has no `index`, the caller is responsible
+        # for populating it elsewhere — we still track the modality metadata.
+
+    # ----- retrieval -----
+
+    def retrieve(self, query: str, k: int = 5) -> RetrievalResult:
+        result = self.inner.retrieve(query, k=k)
+        # Inject modality metadata for any doc we registered.
+        for d in result.documents:
+            if d.id in self._registered:
+                d.metadata.setdefault(
+                    "modality", self._registered[d.id].get("modality", "text"),
+                )
+        return RetrievalResult(
+            query=query,
+            documents=result.documents,
+            score=result.score,
+            strategy=self.strategy,
+        )