pythonclaw 0.2.0__py3-none-any.whl

pythonclaw/core/retrieval/reranker.py

@@ -0,0 +1,112 @@
+"""
+LLM-based re-ranker.
+
+Given a query and a list of candidate chunks (retrieved by sparse + dense),
+asks the LLM to sort them by relevance and returns the top-k.
+
+Prompt strategy
+---------------
+We ask the LLM to return a JSON array of 0-based indices sorted from most
+to least relevant.  This is compact, deterministic to parse, and works well
+with instruction-tuned models.
+
+The re-ranker is *optional*.  It adds one extra LLM call per retrieval but
+significantly improves precision, especially for ambiguous queries.
+
+Usage
+-----
+    from pythonclaw.core.retrieval.reranker import LLMReranker
+    reranker = LLMReranker(provider)
+    best = reranker.rerank(query="...", candidates=[...], top_k=3)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..llm.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+_RERANK_PROMPT = """\
+You are a relevance scoring assistant. Given a search query and a list of text passages, rank the passages by their relevance to the query.
+
+Query: {query}
+
+Passages:
+{passages}
+
+Return ONLY a valid JSON array of passage indices (0-based), ordered from most relevant to least relevant.
+Example: [2, 0, 3, 1]
+
+Your response (JSON array only):"""
+
+
+class LLMReranker:
+    """
+    Re-ranks retrieval candidates using a single LLM call.
+
+    Parameters
+    ----------
+    provider   : any LLMProvider instance.
+    max_chars  : truncate each candidate to this many characters in the prompt.
+    """
+
+    def __init__(self, provider: "LLMProvider", max_chars: int = 300) -> None:
+        self._provider = provider
+        self._max_chars = max_chars
+
+    def rerank(
+        self,
+        query: str,
+        candidates: list[dict],
+        top_k: int,
+    ) -> list[dict]:
+        """
+        Re-rank *candidates* for *query* and return the best *top_k*.
+
+        Falls back to the original order if the LLM response cannot be parsed.
+        """
+        if not candidates:
+            return []
+        if len(candidates) == 1:
+            return candidates[:top_k]
+
+        passages_text = "\n\n".join(
+            f"[{i}] {c['content'][: self._max_chars]}"
+            for i, c in enumerate(candidates)
+        )
+        prompt = _RERANK_PROMPT.format(query=query, passages=passages_text)
+
+        try:
+            response = self._provider.chat(
+                messages=[{"role": "user", "content": prompt}],
+                tools=None,
+                tool_choice=None,
+            )
+            raw = response.choices[0].message.content.strip()
+
+            # Extract first JSON array from the response
+            match = re.search(r"\[[\d,\s]+\]", raw)
+            if not match:
+                raise ValueError(f"No JSON array found in: {raw!r}")
+            indices: list[int] = json.loads(match.group())
+
+            reranked = [
+                candidates[i] for i in indices if 0 <= i < len(candidates)
+            ]
+            # Append any candidates the LLM missed (shouldn't happen, but be safe)
+            seen = set(indices)
+            for i, c in enumerate(candidates):
+                if i not in seen:
+                    reranked.append(c)
+
+            return reranked[:top_k]
+
+        except Exception as exc:
+            logger.warning("[LLMReranker] Reranking failed (%s), using original order.", exc)
+            return candidates[:top_k]

pythonclaw/core/retrieval/retriever.py

@@ -0,0 +1,166 @@
+"""
+HybridRetriever — the main retrieval class.
+
+Pipeline
+--------
+                corpus (list of chunk dicts)
+                        |
+            +-----------+-----------+
+            |                       |
+    BM25Retriever           EmbeddingRetriever
+     (sparse)                 (dense)
+            |                       |
+            +-----------+-----------+
+                     |
+          Reciprocal Rank Fusion
+                     |
+                     | (top fetch_k candidates)
+            LLMReranker  (optional)
+                     |
+                     | top_k final results
+
+Usage
+-----
+    from pythonclaw.core.retrieval import HybridRetriever, load_corpus_from_directory
+
+    retriever = HybridRetriever(provider=llm_provider)
+    retriever.fit(load_corpus_from_directory("context/knowledge"))
+    hits = retriever.retrieve("What is the refund policy?", top_k=5)
+    # hits = [{"source": "...", "content": "...", ...}, ...]
+
+Configuration
+-------------
+    use_sparse   : enable BM25 (default True)
+    use_dense    : enable embedding retriever (default True)
+    use_reranker : enable LLM re-ranking (default True, requires provider)
+    dense_model  : sentence-transformers model name
+    top_k        : number of results returned
+    fetch_k      : candidates fetched before fusion/reranking (>= top_k)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from .sparse import BM25Retriever
+from .dense import EmbeddingRetriever
+from .fusion import reciprocal_rank_fusion
+from .reranker import LLMReranker
+
+if TYPE_CHECKING:
+    from ..llm.base import LLMProvider
+
+logger = logging.getLogger(__name__)
+
+
+class HybridRetriever:
+    """
+    Combines sparse (BM25) + dense (embedding) retrieval with RRF fusion and
+    an optional LLM re-ranker.
+
+    Parameters
+    ----------
+    provider     : LLMProvider instance (required for use_reranker=True).
+    use_sparse   : include BM25 retrieval.
+    use_dense    : include embedding retrieval.
+    use_reranker : re-rank fused candidates with the LLM.
+    dense_model  : sentence-transformers model name.
+    """
+
+    def __init__(
+        self,
+        provider: "LLMProvider | None" = None,
+        use_sparse: bool = True,
+        use_dense: bool = True,
+        use_reranker: bool = True,
+        dense_model: str = "all-MiniLM-L6-v2",
+    ) -> None:
+        self._provider = provider
+        self.use_sparse = use_sparse
+        self.use_dense = use_dense
+        self.use_reranker = use_reranker and provider is not None
+
+        self._sparse = BM25Retriever() if use_sparse else None
+        self._dense = EmbeddingRetriever(dense_model) if use_dense else None
+        self._reranker = LLMReranker(provider) if self.use_reranker else None
+        self._corpus: list[dict] = []
+
+        if use_dense and self._dense:
+            logger.info("[HybridRetriever] Dense backend: %s", self._dense.backend_name)
+
+    # ── Indexing ──────────────────────────────────────────────────────────────
+
+    def fit(self, corpus: list[dict]) -> "HybridRetriever":
+        """
+        Index the corpus.  Each item must have a 'content' key.
+        Mutates corpus in-place by adding '_idx' for RRF deduplication.
+        """
+        for i, chunk in enumerate(corpus):
+            chunk["_idx"] = i
+        self._corpus = corpus
+
+        if self._sparse:
+            self._sparse.fit(corpus)
+        if self._dense:
+            self._dense.fit(corpus)
+
+        logger.info(
+            "[HybridRetriever] Indexed %d chunks (sparse=%s dense=%s reranker=%s)",
+            len(corpus), self.use_sparse, self.use_dense, self.use_reranker,
+        )
+        return self
+
+    # ── Retrieval ─────────────────────────────────────────────────────────────
+
+    def retrieve(self, query: str, top_k: int = 5) -> list[dict]:
+        """
+        Retrieve the *top_k* most relevant chunks for *query*.
+
+        Returns a list of chunk dicts (internal '_idx' field stripped).
+        """
+        if not self._corpus or not query.strip():
+            return []
+
+        # How many candidates to fetch before reranking
+        fetch_k = max(top_k * 3, top_k + 5)
+
+        ranked_lists: list[list[tuple[float, dict]]] = []
+
+        if self._sparse:
+            sparse_results = self._sparse.retrieve(query, top_k=fetch_k)
+            if sparse_results:
+                ranked_lists.append(sparse_results)
+
+        if self._dense:
+            dense_results = self._dense.retrieve(query, top_k=fetch_k)
+            if dense_results:
+                ranked_lists.append(dense_results)
+
+        if not ranked_lists:
+            return []
+
+        # Fusion
+        if len(ranked_lists) == 1:
+            fused = [(s, c) for s, c in ranked_lists[0]]
+        else:
+            fused = reciprocal_rank_fusion(ranked_lists)
+
+        candidates = [c for _, c in fused[: fetch_k if self._reranker else top_k]]
+
+        # Re-rank
+        if self._reranker and candidates:
+            candidates = self._reranker.rerank(query, candidates, top_k)
+        else:
+            candidates = candidates[:top_k]
+
+        # Strip internal index field before returning
+        return [{k: v for k, v in c.items() if k != "_idx"} for c in candidates]
+
+    # ── Convenience ──────────────────────────────────────────────────────────
+
+    def __len__(self) -> int:
+        return len(self._corpus)
+
+    def __bool__(self) -> bool:
+        return bool(self._corpus)

pythonclaw/core/retrieval/sparse.py

@@ -0,0 +1,69 @@
+"""
+Sparse retriever — BM25Okapi.
+
+Falls back to a simple TF-weighted word-overlap scorer when `rank_bm25`
+is not installed.  Install rank-bm25 for best quality:
+
+    pip install rank-bm25
+"""
+
+from __future__ import annotations
+
+import re
+
+try:
+    from rank_bm25 import BM25Okapi
+    _HAS_BM25 = True
+except ImportError:
+    _HAS_BM25 = False
+
+
+def _tokenize(text: str) -> list[str]:
+    return re.findall(r"\w+", text.lower())
+
+
+class BM25Retriever:
+    """
+    Wraps BM25Okapi (or a simple fallback) for sparse retrieval.
+
+    Usage
+    -----
+        r = BM25Retriever()
+        r.fit(corpus)                         # corpus = list of {"content": ..., ...}
+        results = r.retrieve("my query", 10)  # -> [(score, chunk_dict), ...]
+    """
+
+    def __init__(self) -> None:
+        self._corpus: list[dict] = []
+        self._bm25: object | None = None
+        self._tokenized: list[list[str]] = []
+
+    def fit(self, corpus: list[dict]) -> None:
+        self._corpus = corpus
+        self._tokenized = [_tokenize(c["content"]) for c in corpus]
+        if _HAS_BM25 and corpus:
+            self._bm25 = BM25Okapi(self._tokenized)
+
+    def retrieve(self, query: str, top_k: int) -> list[tuple[float, dict]]:
+        if not self._corpus:
+            return []
+
+        tokens = _tokenize(query)
+
+        if _HAS_BM25 and self._bm25:
+            raw_scores = self._bm25.get_scores(tokens)
+            pairs = [(float(s), c) for s, c in zip(raw_scores, self._corpus) if s > 0]
+        else:
+            # Fallback: term-frequency word-overlap
+            pairs = []
+            query_set = set(tokens)
+            for chunk in self._corpus:
+                chunk_tokens = _tokenize(chunk["content"])
+                if not chunk_tokens:
+                    continue
+                tf = sum(1 for t in chunk_tokens if t in query_set)
+                if tf > 0:
+                    pairs.append((float(tf) / len(chunk_tokens), chunk))
+
+        pairs.sort(key=lambda x: x[0], reverse=True)
+        return pairs[:top_k]

pythonclaw/core/session_store.py

@@ -0,0 +1,269 @@
+"""
+Markdown-backed session store for pythonclaw.
+
+Each session gets its own Markdown file with timestamped messages::
+
+    context/sessions/telegram_1285451567.md
+
+File format
+-----------
+Human-readable Markdown with embedded metadata in HTML comments for reliable
+round-trip parsing.  Each message block::
+
+    <!-- msg:{"role":"user","ts":"2026-02-23T15:18:58"} -->
+    ### 2026-02-23 15:18:58 — User
+
+    Hello, how are you?
+
+    ---
+
+Tool calls are stored as JSON code blocks inside the message.  System
+injection messages (skill loads, compaction summaries) are also recorded.
+
+Truncation
+----------
+When a session file grows beyond *max_messages*, older messages are dropped
+(keeping only the most recent ones by timestamp).  The system prompt
+(messages[0]) is never saved — it is always rebuilt fresh on restore.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_STORE_DIR = os.path.join("context", "sessions")
+DEFAULT_MAX_MESSAGES = 200
+
+_META_PATTERN = re.compile(r"<!-- msg:(.*?) -->")
+_ROLE_LABELS = {
+    "user": "User",
+    "assistant": "Assistant",
+    "system": "System",
+    "tool": "Tool",
+}
+
+
+class SessionStore:
+    """Reads and writes per-session message history as Markdown files."""
+
+    def __init__(
+        self,
+        base_dir: str = DEFAULT_STORE_DIR,
+        max_messages: int = DEFAULT_MAX_MESSAGES,
+    ) -> None:
+        self.base_dir = base_dir
+        self.max_messages = max_messages
+        os.makedirs(base_dir, exist_ok=True)
+
+    # ── File path ─────────────────────────────────────────────────────────────
+
+    def _path(self, session_id: str) -> str:
+        """Convert a session_id like 'telegram:123' to a safe filename."""
+        safe = re.sub(r"[^\w\-]", "_", session_id)
+        return os.path.join(self.base_dir, f"{safe}.md")
+
+    # ── Serialisation ─────────────────────────────────────────────────────────
+
+    @staticmethod
+    def _msg_to_markdown(msg: dict) -> str:
+        """Convert a single message dict to a Markdown block."""
+        role = msg.get("role", "unknown")
+        content = msg.get("content", "") or ""
+        ts = msg.get("_ts") or datetime.now().isoformat(timespec="seconds")
+
+        # Build metadata for round-trip parsing
+        meta: dict = {"role": role, "ts": ts}
+        if msg.get("tool_call_id"):
+            meta["tool_call_id"] = msg["tool_call_id"]
+
+        meta_json = json.dumps(meta, ensure_ascii=False)
+        label = _ROLE_LABELS.get(role, role.title())
+
+        # Format timestamp for display
+        try:
+            dt = datetime.fromisoformat(ts)
+            display_ts = dt.strftime("%Y-%m-%d %H:%M:%S")
+        except (ValueError, TypeError):
+            display_ts = ts
+
+        lines = [f"<!-- msg:{meta_json} -->"]
+        lines.append(f"### {display_ts} — {label}")
+        lines.append("")
+
+        if content:
+            lines.append(content)
+            lines.append("")
+
+        # Embed tool_calls as JSON
+        tool_calls = msg.get("tool_calls")
+        if tool_calls:
+            lines.append("<details><summary>Tool Calls</summary>")
+            lines.append("")
+            lines.append("```json")
+            lines.append(json.dumps(tool_calls, ensure_ascii=False, indent=2))
+            lines.append("```")
+            lines.append("")
+            lines.append("</details>")
+            lines.append("")
+
+        lines.append("---")
+        lines.append("")
+        return "\n".join(lines)
+
+    @staticmethod
+    def _parse_markdown(text: str) -> list[dict]:
+        """Parse a session Markdown file back into message dicts."""
+        messages: list[dict] = []
+
+        # Split into blocks by the HTML comment markers
+        blocks = _META_PATTERN.split(text)
+        # blocks = [preamble, meta1, content1, meta2, content2, ...]
+
+        i = 1  # skip preamble (title / header)
+        while i < len(blocks) - 1:
+            meta_str = blocks[i].strip()
+            body = blocks[i + 1].strip()
+            i += 2
+
+            try:
+                meta = json.loads(meta_str)
+            except json.JSONDecodeError:
+                continue
+
+            role = meta.get("role", "unknown")
+            msg: dict = {"role": role}
+
+            if meta.get("tool_call_id"):
+                msg["tool_call_id"] = meta["tool_call_id"]
+            if meta.get("ts"):
+                msg["_ts"] = meta["ts"]
+
+            # Extract content: everything between the header line and
+            # optional <details> / --- markers
+            content_lines = []
+            tool_calls_json = None
+            in_details = False
+            in_json_block = False
+            json_lines: list[str] = []
+
+            for line in body.split("\n"):
+                stripped = line.strip()
+
+                # Skip the ### header line and trailing ---
+                if stripped.startswith("### ") or stripped == "---":
+                    continue
+
+                if stripped == "<details><summary>Tool Calls</summary>":
+                    in_details = True
+                    continue
+                if stripped == "</details>":
+                    in_details = False
+                    continue
+
+                if in_details:
+                    if stripped == "```json":
+                        in_json_block = True
+                        continue
+                    if stripped == "```" and in_json_block:
+                        in_json_block = False
+                        try:
+                            tool_calls_json = json.loads("\n".join(json_lines))
+                        except json.JSONDecodeError:
+                            pass
+                        json_lines = []
+                        continue
+                    if in_json_block:
+                        json_lines.append(line)
+                        continue
+                    continue
+
+                content_lines.append(line)
+
+            content = "\n".join(content_lines).strip()
+            if content:
+                msg["content"] = content
+            else:
+                msg["content"] = ""
+
+            if tool_calls_json:
+                msg["tool_calls"] = tool_calls_json
+
+            messages.append(msg)
+
+        return messages
+
+    # ── Core API ──────────────────────────────────────────────────────────────
+
+    def save(self, session_id: str, messages: list[dict]) -> None:
+        """
+        Persist messages[1:] to Markdown.
+        messages[0] is the initial system prompt — always rebuilt fresh.
+        """
+        to_save = messages[1:] if len(messages) > 1 else []
+
+        # Add timestamps to messages that don't have one
+        for msg in to_save:
+            if "_ts" not in msg:
+                msg["_ts"] = datetime.now().isoformat(timespec="seconds")
+
+        # Truncate by time — keep only the most recent max_messages
+        if len(to_save) > self.max_messages:
+            to_save = to_save[-self.max_messages:]
+
+        path = self._path(session_id)
+        try:
+            lines = [f"# Session: {session_id}\n\n"]
+            for msg in to_save:
+                lines.append(self._msg_to_markdown(msg))
+            with open(path, "w", encoding="utf-8") as f:
+                f.write("\n".join(lines))
+        except OSError as exc:
+            logger.error("[SessionStore] Failed to save session '%s': %s", session_id, exc)
+
+    def load(self, session_id: str) -> list[dict]:
+        """
+        Return saved messages (messages[1:] from a previous run).
+        Applies time-based truncation: only the most recent max_messages
+        are returned.
+        """
+        path = self._path(session_id)
+        if not os.path.exists(path):
+            return []
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                text = f.read()
+        except OSError as exc:
+            logger.error("[SessionStore] Failed to load session '%s': %s", session_id, exc)
+            return []
+
+        messages = self._parse_markdown(text)
+
+        # Time-based truncation: keep only the most recent messages
+        if len(messages) > self.max_messages:
+            messages = messages[-self.max_messages:]
+
+        return messages
+
+    def delete(self, session_id: str) -> None:
+        """Remove the Markdown file for session_id."""
+        path = self._path(session_id)
+        if os.path.exists(path):
+            try:
+                os.remove(path)
+                logger.info("[SessionStore] Deleted session '%s'", session_id)
+            except OSError as exc:
+                logger.error("[SessionStore] Failed to delete '%s': %s", session_id, exc)
+
+    def list_session_ids(self) -> list[str]:
+        """Return all session IDs that have a persisted Markdown file."""
+        ids = []
+        for fname in os.listdir(self.base_dir):
+            if fname.endswith(".md"):
+                ids.append(fname[: -len(".md")])
+        return ids