raglite-chromadb 1.0.1__py3-none-any.whl

raglite/__init__.py

@@ -0,0 +1 @@
+__all__ = []

raglite/chroma_rest.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import json
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+
+
+@dataclass
+class ChromaLoc:
+    base_url: str = "http://127.0.0.1:8100"
+    tenant: str = "default_tenant"
+    database: str = "default_database"
+
+    def collections_url(self) -> str:
+        return f"{self.base_url}/api/v2/tenants/{self.tenant}/databases/{self.database}/collections"
+
+
+def _req_json(method: str, url: str, body: dict | None = None, timeout: int = 120) -> dict | list:
+    data = None if body is None else json.dumps(body).encode("utf-8")
+    req = urllib.request.Request(url, data=data, method=method)
+    req.add_header("Content-Type", "application/json")
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            raw = resp.read().decode("utf-8")
+            if not raw:
+                return {}
+            return json.loads(raw)
+    except urllib.error.HTTPError as e:  # type: ignore[attr-defined]
+        detail = ""
+        try:
+            detail = e.read().decode("utf-8", errors="ignore")
+        except Exception:
+            detail = ""
+        raise RuntimeError(f"Chroma HTTP {e.code} calling {url}: {detail[:500]}")
+
+
+def list_collections(loc: ChromaLoc) -> list[dict]:
+    res = _req_json("GET", loc.collections_url())
+    if not isinstance(res, list):
+        raise RuntimeError(f"Expected list from Chroma list_collections, got {type(res).__name__}")
+    return res
+
+
+def get_or_create_collection(loc: ChromaLoc, name: str, *, space: str = "cosine") -> dict:
+    cols = list_collections(loc)
+    for c in cols:
+        if c.get("name") == name:
+            return c
+
+    created = _req_json(
+        "POST",
+        loc.collections_url(),
+        {"name": name, "metadata": {"hnsw:space": space}},
+    )
+    if not isinstance(created, dict):
+        raise RuntimeError(f"Expected dict from Chroma create_collection, got {type(created).__name__}")
+    return created
+
+
+def add(
+    loc: ChromaLoc,
+    collection_id: str,
+    *,
+    ids: list[str],
+    documents: list[str],
+    embeddings: list[list[float]],
+    metadatas: list[dict] | None = None,
+) -> None:
+    url = f"{loc.collections_url()}/{collection_id}/add"
+    body: dict = {"ids": ids, "documents": documents, "embeddings": embeddings}
+    if metadatas is not None:
+        body["metadatas"] = metadatas
+    _req_json("POST", url, body, timeout=600)
+
+
+def upsert(
+    loc: ChromaLoc,
+    collection_id: str,
+    *,
+    ids: list[str],
+    documents: list[str],
+    embeddings: list[list[float]],
+    metadatas: list[dict] | None = None,
+) -> None:
+    """Upsert records into a collection.
+
+    Chroma's /upsert updates existing ids and inserts new ones.
+    """
+    url = f"{loc.collections_url()}/{collection_id}/upsert"
+    body: dict = {"ids": ids, "documents": documents, "embeddings": embeddings}
+    if metadatas is not None:
+        body["metadatas"] = metadatas
+    _req_json("POST", url, body, timeout=600)
+
+
+def query(
+    loc: ChromaLoc,
+    collection_id: str,
+    *,
+    query_embeddings: list[list[float]],
+    n_results: int = 10,
+    include: list[str] | None = None,
+) -> dict:
+    url = f"{loc.collections_url()}/{collection_id}/query"
+    body: dict = {"query_embeddings": query_embeddings, "n_results": n_results}
+    if include is not None:
+        body["include"] = include
+    res = _req_json("POST", url, body, timeout=600)
+    assert isinstance(res, dict)
+    return res

raglite/extract.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+from bs4 import BeautifulSoup
+from pypdf import PdfReader
+
+
+FileKind = Literal["pdf", "html", "htm", "txt"]
+
+
+@dataclass
+class ExtractResult:
+    kind: FileKind
+    text: str
+
+
+def _clean_text(s: str) -> str:
+    s = s.replace("\r\n", "\n").replace("\r", "\n")
+    s = re.sub(r"[ \t]+", " ", s)
+    s = re.sub(r"\n{3,}", "\n\n", s)
+    return s.strip()
+
+
+def extract_pdf(path: Path) -> ExtractResult:
+    reader = PdfReader(str(path))
+    parts = []
+    for i, page in enumerate(reader.pages):
+        try:
+            txt = page.extract_text() or ""
+        except Exception:
+            txt = ""
+        if txt.strip():
+            parts.append(f"\n\n--- Page {i+1} ---\n\n{txt}")
+    return ExtractResult(kind="pdf", text=_clean_text("\n".join(parts)))
+
+
+def extract_html(path: Path) -> ExtractResult:
+    html = path.read_text(encoding="utf-8", errors="ignore")
+    soup = BeautifulSoup(html, "lxml")
+
+    # Remove scripts/styles
+    for tag in soup(["script", "style", "noscript"]):
+        tag.decompose()
+
+    text = soup.get_text("\n")
+    return ExtractResult(kind="html", text=_clean_text(text))
+
+
+def extract_txt(path: Path) -> ExtractResult:
+    return ExtractResult(kind="txt", text=_clean_text(path.read_text(encoding="utf-8", errors="ignore")))
+
+
+def extract_file(path: Path) -> ExtractResult:
+    suffix = path.suffix.lower().lstrip(".")
+    if suffix == "pdf":
+        return extract_pdf(path)
+    if suffix in ("html", "htm"):
+        return extract_html(path)
+    return extract_txt(path)

raglite/prompts.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class PromptPair:
+    execution_notes_prompt: str
+    tool_summary_prompt: str
+
+
+def build_prompts(*, token_cap_hint: str = "~1200 tokens max") -> PromptPair:
+    # These prompts are designed to be copy/pasted into Cosmo.
+    execution_notes = f"""You are an expert at converting documentation into EXECUTION-RELEVANT notes for an AI agent that can run tools (CLI commands, HTTP calls, scripts, and functions).
+
+OUTPUT FORMAT (Markdown):
+- Title
+- What this tool/service is
+- When to use
+- Inputs (required/optional)
+- Outputs
+- Preconditions / assumptions
+- Step-by-step 'golden path' (numbered)
+- Verification checks (how to confirm success)
+- Common errors + fixes
+- Safety/rollback notes (what not to do / how to undo)
+
+RULES:
+- Be concise and operational; no marketing.
+- Prefer concrete commands, flags, endpoints, and example payloads.
+- If the doc is long, extract only what is needed to execute.
+- Keep the final output within {token_cap_hint}.
+
+SOURCE DOCUMENT (extracted text) is below. Use it as the only source of truth.
+---
+"""
+
+    tool_summary = f"""You are an expert at writing ultra-condensed TOOL INDEX summaries for an agent tool library.
+
+Write a short Markdown file with:
+- Tool name
+- 1-sentence purpose
+- Capabilities (3-7 bullets)
+- Required environment/dependencies
+- Primary entrypoints (commands/endpoints)
+- Key limitations / footguns (1-3 bullets)
+
+RULES:
+- No fluff. Assume the reader is an executor agent.
+- Keep within ~250-400 tokens.
+
+SOURCE DOCUMENT (extracted text) is below. Use it as the only source of truth.
+---
+"""
+
+    return PromptPair(execution_notes_prompt=execution_notes, tool_summary_prompt=tool_summary)