longloom 1.0.0__py3-none-any.whl

longloom/__init__.py

@@ -0,0 +1,87 @@
+"""longloom — long-context map/merge/repair for any LLM and any storage.
+
+Process documents far larger than a model's context window using a LangGraph
+map/merge/repair workflow, with dependency injection for both the LLM and the
+storage backend::
+
+    import os
+    from longloom import MapReduceProcessor, SQLiteContextStore, init_llm
+
+    processor = MapReduceProcessor(
+        llm=init_llm("nvidia", "nvidia/nemotron-3-ultra-550b-a55b",
+                     api_key=os.environ["NVIDIA_API_KEY"]),
+        store=SQLiteContextStore("ctx.db"),
+    )
+    result = processor.process(long_document, task="Summarize")
+"""
+
+from .chunks import ContextChunk, content_hash, split_text
+from .factory import SUPPORTED_SCHEMES, create_store
+from .stores.naming import FieldMap
+from .graph import MapReduceState, build_graph
+from .llm import LLMAdapter, LLMLike, init_llm
+from .processor import MapReduceProcessor
+from .prompts import DEFAULT_PROMPTS, PromptTemplates, render_prompt
+from .stores import (
+    MYSQL_DIALECT,
+    POSTGRES_DIALECT,
+    SQLITE_DIALECT,
+    BaseContextStore,
+    InMemoryContextStore,
+    JSONFileContextStore,
+    MongoContextStore,
+    SQLiteContextStore,
+    SqlContextStore,
+    SupabaseContextStore,
+)
+from .tokens import EstimateCounter, TiktokenCounter, TokenCounter, default_counter
+
+__version__ = "1.0.0"
+
+__all__ = [
+    # processor / graph
+    "MapReduceProcessor",
+    "MapReduceState",
+    "build_graph",
+    # llm
+    "LLMAdapter",
+    "LLMLike",
+    "init_llm",
+    "NVIDIAChatModel",
+    # stores
+    "BaseContextStore",
+    "InMemoryContextStore",
+    "SQLiteContextStore",
+    "JSONFileContextStore",
+    "SqlContextStore",
+    "POSTGRES_DIALECT",
+    "MYSQL_DIALECT",
+    "SQLITE_DIALECT",
+    "MongoContextStore",
+    "SupabaseContextStore",
+    "create_store",
+    "SUPPORTED_SCHEMES",
+    "FieldMap",
+    # tokens
+    "TokenCounter",
+    "TiktokenCounter",
+    "EstimateCounter",
+    "default_counter",
+    # chunks
+    "ContextChunk",
+    "split_text",
+    "content_hash",
+    # prompts
+    "PromptTemplates",
+    "DEFAULT_PROMPTS",
+    "render_prompt",
+]
+
+
+def __getattr__(name: str):
+    # Optional import: NVIDIAChatModel only pulls `requests` when used.
+    if name == "NVIDIAChatModel":
+        from .nvidia import NVIDIAChatModel
+
+        return NVIDIAChatModel
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

longloom/chunks.py

@@ -0,0 +1,109 @@
+"""Chunk data model and token-aware text splitting.
+
+Backends and counters are pluggable: :func:`split_text` works with any
+callable token counter (tiktoken-based or estimator) and never depends on
+a specific tokenizer's encode/decode round-trip.
+"""
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, List
+
+__all__ = ["ContextChunk", "content_hash", "utc_now_iso", "split_text"]
+
+
+def content_hash(text: str, length: int = 16) -> str:
+    """Short SHA-256 hash of text, used for content addressing."""
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:length]
+
+
+def utc_now_iso() -> str:
+    """UTC timestamp in ISO-8601 format."""
+    return datetime.now(timezone.utc).isoformat()
+
+
+@dataclass
+class ContextChunk:
+    """A single chunk of a larger document."""
+
+    id: str
+    content: str
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    token_count: int = 0
+    created_at: str = ""
+    chunk_index: int = 0
+    total_chunks: int = 1
+
+
+def split_text(
+    text: str,
+    max_tokens: int,
+    counter: Callable[[str], int],
+) -> List[str]:
+    """Split *text* into pieces of at most *max_tokens* (counter-defined tokens).
+
+    Strategy — LINE-AWARE, so facts on a single line are never cut in half:
+
+      1. Greedily pack whole lines into chunks while the chunk fits
+         (``counter()`` measures each line; works with any counter).
+      2. A single line longer than *max_tokens* is hard-split by token count
+         (binary search with the counter) — unavoidable last resort.
+
+    Returns a list of non-empty strings whose concatenation reconstructs the
+    original text.
+    """
+    if max_tokens < 1:
+        raise ValueError("max_tokens must be >= 1")
+    if not text:
+        return []
+
+    def _hard_split(segment: str) -> List[str]:
+        """Token-count split for a single line that exceeds max_tokens."""
+        if chr(10) in segment:  # defensive; caller passes single lines
+            return [segment]
+        out: List[str] = []
+        remaining = segment
+        while remaining:
+            if counter(remaining) <= max_tokens:
+                out.append(remaining)
+                break
+            guess = max(1, int(len(remaining) * max_tokens / max(counter(remaining), 1)))
+            steps = 0
+            while guess > 1 and counter(remaining[:guess]) > max_tokens and steps < 64:
+                guess = max(1, int(guess * 0.9))
+                steps += 1
+            while guess < len(remaining) and counter(remaining[: guess + 1]) <= max_tokens:
+                guess += 1
+            if guess <= 0:
+                guess = 1
+            # Snap to a whitespace boundary so identifier tokens (IDs, numbers,
+            # handles — anything without spaces) are never cut in half.
+            ws = remaining.rfind(" ", 0, guess)
+            if ws > guess * 0.5:
+                guess = ws + 1
+            out.append(remaining[:guess])
+            remaining = remaining[guess:]
+        return out
+
+    chunks: List[str] = []
+    current: List[str] = []
+    current_tokens = 0
+    for line in text.splitlines(keepends=True):
+        line_tokens = counter(line)
+        if line_tokens > max_tokens:
+            if current:
+                chunks.append("".join(current))
+                current, current_tokens = [], 0
+            chunks.extend(_hard_split(line.rstrip("\n")) if line.strip() else [line])
+            continue
+        if current_tokens + line_tokens > max_tokens and current:
+            chunks.append("".join(current))
+            current, current_tokens = [], 0
+        current.append(line)
+        current_tokens += line_tokens
+    if current:
+        chunks.append("".join(current))
+
+    return [c for c in chunks if c.strip()]

longloom/factory.py

@@ -0,0 +1,245 @@
+"""``create_store(url)`` — build any supported backend from a connection string.
+
+Supported URL schemes:
+
+==========================  ==============================================
+``memory://``               :class:`InMemoryContextStore`
+``json://<dir>``            :class:`JSONFileContextStore`
+``sqlite:///<path>``        :class:`SQLiteContextStore` (``sqlite://`` = in-memory)
+``postgresql://<dsn>``      :class:`SqlContextStore` (PostgreSQL dialect)
+``mysql://<dsn>``           :class:`SqlContextStore` (MySQL/MariaDB dialect)
+``mongodb://<uri>``         :class:`MongoContextStore` (path = database name)
+``https://<proj>.supabase.co``  :class:`SupabaseContextStore` (key via SUPABASE_KEY)
+==========================  ==============================================
+
+SQL URLs prefer SQLAlchemy when installed (pooled, works with every driver);
+otherwise known drivers (psycopg2/psycopg, pymysql/MySQLdb) are tried directly.
+"""
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, Callable, Optional
+from urllib.parse import unquote, urlparse
+
+from .stores.base import BaseContextStore
+from .stores.memory import InMemoryContextStore
+
+__all__ = ["create_store", "SUPPORTED_SCHEMES"]
+
+SUPPORTED_SCHEMES = (
+    "memory", "json", "sqlite", "postgres", "postgresql", "mysql", "mariadb",
+    "mongodb", "mongodb+srv", "supabase", "https",
+)
+
+
+def create_store(
+    url: str,
+    table: str = "context_chunks",
+    columns: Optional[Dict[str, str]] = None,
+    namespace: Optional[str] = None,
+    **kwargs,
+) -> BaseContextStore:
+    """Create a storage backend from a connection string (see module docstring).
+
+    ``table``, ``columns`` (logical → physical names) and ``namespace`` (a
+    discriminator letting several apps share one table/collection) embed
+    LongLoom into existing databases — supported by the SQL, MongoDB and
+    Supabase backends; ``memory`` and ``json`` ignore them.
+    """
+    if not url:
+        raise ValueError("url must be a non-empty connection string")
+    import re as _re
+
+    if table and not _re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", table):
+        raise ValueError(f"invalid table/collection name: {table!r}")
+    parsed = urlparse(url)
+    scheme = (parsed.scheme or "").lower()
+
+    if scheme == "memory":
+        return InMemoryContextStore(**kwargs)
+
+    if scheme == "json":
+        root = parsed.netloc + parsed.path if parsed.netloc else parsed.path
+        from .stores.jsonfile import JSONFileContextStore
+
+        return JSONFileContextStore(root_dir=unquote(root) or "longloom_data", **kwargs)
+
+    if scheme == "sqlite":
+        path = parsed.path or ""
+        if parsed.netloc and parsed.netloc != ":memory:":
+            path = parsed.netloc + path
+        db_path = unquote(path) if path and path != ":memory:" else ":memory:"
+        from .stores.sqlite import SQLiteContextStore
+
+        return SQLiteContextStore(db_path=db_path, table=table, columns=columns,
+                                  namespace=namespace, **kwargs)
+
+    if scheme in ("postgres", "postgresql"):
+        from .stores.sql import POSTGRES_DIALECT, SqlContextStore
+
+        return SqlContextStore(
+            connect=_sql_connect_factory(url, "postgres"),
+            dialect=POSTGRES_DIALECT,
+            table=table, columns=columns, namespace=namespace,
+            **kwargs,
+        )
+
+    if scheme in ("mysql", "mariadb"):
+        from .stores.sql import MYSQL_DIALECT, SqlContextStore
+
+        return SqlContextStore(
+            connect=_sql_connect_factory(url, "mysql"),
+            dialect=MYSQL_DIALECT,
+            table=table, columns=columns, namespace=namespace,
+            **kwargs,
+        )
+
+    if scheme in ("mongodb", "mongodb+srv"):
+        from .stores.mongo import MongoContextStore
+
+        db_name = unquote(parsed.path.lstrip("/")) or kwargs.pop("db_name", "longloom")
+        return MongoContextStore(client_or_uri=url, db_name=db_name,
+                                 collection=kwargs.pop("collection", table),
+                                 columns=columns, namespace=namespace, **kwargs)
+
+    if scheme == "supabase":
+        from .stores.supabase import SupabaseContextStore
+
+        host = parsed.netloc
+        # supabase://x.supabase.co → https://x.supabase.co (cloud);
+        # supabase://127.0.0.1:8905 → http://… (local PostgREST stub).
+        project_url = f"https://{host}" if host.endswith(".supabase.co") else f"http://{host}"
+        return SupabaseContextStore(url=project_url, table=table, columns=columns,
+                                    namespace=namespace, **kwargs)
+
+    if scheme == "https" and re.search(r"\.supabase\.", parsed.netloc or ""):
+        from .stores.supabase import SupabaseContextStore
+
+        return SupabaseContextStore(url=url, table=table, **kwargs)
+
+    raise ValueError(
+        f"Unsupported store URL scheme: {scheme!r}. Supported: memory://, json://, "
+        "sqlite://, postgresql://, mysql://, mongodb://, supabase://, https://<proj>.supabase.co"
+    )
+
+
+# ----------------------------------------------------------------------
+# SQL connection helpers
+# ----------------------------------------------------------------------
+
+def _sql_connect_factory(url: str, kind: str) -> Callable[[], Any]:
+    """Return a zero-arg callable yielding a DB-API-style connection."""
+
+    # Preferred: SQLAlchemy — uniform support for every driver and pooling.
+    try:
+        import sqlalchemy  # noqa: F401
+        from sqlalchemy import create_engine
+
+        norm = url
+        if norm.startswith("postgres://"):
+            norm = "postgresql://" + norm[len("postgres://"):]
+        if kind == "mysql" and not norm.startswith("mysql+"):
+            try:
+                import MySQLdb  # type: ignore  # noqa: F401
+            except ImportError:
+                try:
+                    import pymysql  # type: ignore  # noqa: F401
+
+                    norm = "mysql+pymysql://" + norm.split("://", 1)[1]
+                except ImportError:
+                    pass  # let SQLAlchemy report the missing driver
+
+        engine = create_engine(norm, pool_pre_ping=True)
+
+        class _SaConn:
+            """DB-API-style adapter over a SQLAlchemy Connection.
+
+            SQLAlchemy 2.0 ``Connection.execute`` is statement-oriented, not
+            cursor-oriented (returns the statement result only via
+            ``exec_driver_sql``), so expose exactly the interface
+            SqlContextStore needs.
+            """
+
+            def __init__(self, sa_conn: Any):
+                self._c = sa_conn
+
+            def execute(self, sql: str, params: Optional[list] = None) -> Any:
+                # exec_driver_sql requires a tuple for positional params
+                # (a bare list fails SQLAlchemy's param validation).
+                return self._c.exec_driver_sql(sql, tuple(params or []))
+
+            def commit(self) -> None:
+                self._c.commit()
+
+            def rollback(self) -> None:
+                self._c.rollback()
+
+            def close(self) -> None:
+                self._c.close()
+
+        def connect_sa() -> Any:
+            return _SaConn(engine.connect())
+
+        # Validate eagerly so misconfigurations surface at construction time.
+        connect_sa().execute("SELECT 1")
+        return connect_sa
+    except ImportError:
+        pass
+
+    # Fallback: known drivers directly.
+    if kind == "postgres":
+        def connect_pg() -> Any:
+            try:
+                import psycopg2
+
+                return _as_executable(psycopg2.connect(url))
+            except ImportError:
+                pass
+            try:
+                import psycopg
+
+                return _as_executable(psycopg.connect(url))
+            except ImportError as exc:
+                raise ImportError(
+                    "No PostgreSQL driver found. Install one of: "
+                    "psycopg2-binary, psycopg, or sqlalchemy. "
+                    "pip install 'longloom[postgres]'"
+                ) from exc
+
+        connect_pg()  # validate eagerly
+        return connect_pg
+
+    def connect_mysql() -> Any:
+        parsed = urlparse(url)
+        conn_kwargs = {
+            "host": parsed.hostname or "127.0.0.1",
+            "port": parsed.port or 3306,
+            "user": unquote(parsed.username or "root"),
+            "password": unquote(parsed.password or ""),
+            "database": unquote(parsed.path.lstrip("/")) or None,
+        }
+        conn_kwargs = {k: v for k, v in conn_kwargs.items() if v is not None}
+        try:
+            import pymysql
+
+            return _as_executable(pymysql.connect(**conn_kwargs))
+        except ImportError:
+            pass
+        try:
+            import MySQLdb
+
+            return _as_executable(MySQLdb.connect(**conn_kwargs))
+        except ImportError as exc:
+            raise ImportError(
+                "No MySQL driver found. Install pymysql or mysqlclient, "
+                "or sqlalchemy. pip install 'longloom[mysql]'"
+            ) from exc
+
+    connect_mysql()  # validate eagerly
+    return connect_mysql
+
+
+def _as_executable(conn: Any) -> Any:
+    from .stores.sql import _as_executable as _adapt
+    return _adapt(conn)

longloom/graph.py

@@ -0,0 +1,220 @@
+"""The LangGraph map → merge → repair workflow, with all dependencies injected.
+
+Every node is a closure over an :class:`LLMAdapter` and a
+:class:`BaseContextStore` supplied by the caller — the graph itself has zero
+knowledge of any vendor or storage engine.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, TypedDict
+
+from langgraph.graph import END, StateGraph
+
+from .chunks import ContextChunk
+from .llm import LLMAdapter
+from .prompts import DEFAULT_PROMPTS, PromptTemplates, render_prompt
+from .stores.base import BaseContextStore
+from .tokens import TokenCounter, default_counter
+
+__all__ = ["MapReduceState", "build_graph"]
+
+
+class MapReduceState(TypedDict, total=False):
+    # Input
+    input_content: str
+    session_id: str
+    task: str
+    task_focus: str
+    # Config
+    max_chunk_tokens: int
+    max_context_tokens: int
+    max_repair_attempts: int
+    # Map phase
+    chunks: List[ContextChunk]
+    map_results: List[str]
+    # Merge phase
+    merged_result: str
+    merge_token_count: int
+    # Repair phase
+    needs_repair: bool
+    repair_attempts: int
+    final_result: str
+    # Metadata
+    total_input_tokens: int
+    processing_log: List[str]
+
+
+def build_graph(
+    llm: LLMAdapter,
+    store: BaseContextStore,
+    token_counter: Optional[TokenCounter] = None,
+    prompts: Optional[PromptTemplates] = None,
+) -> Any:
+    """Compile the map/merge/repair graph with injected dependencies."""
+    prompts = prompts or DEFAULT_PROMPTS
+    counter: TokenCounter = token_counter or store.counter
+
+    def _combine(parts: List[str], state: MapReduceState, label: str) -> str:
+        """One combine call over a group of results (fits the merge budget)."""
+        chunk_results_text = "\n\n---\n\n".join(
+            f"Part {i + 1}:\n{part}" for i, part in enumerate(parts)
+        )
+        prompt = render_prompt(
+            prompts.merge,
+            {
+                "task": state["task"],
+                "task_focus": state.get("task_focus", "key information and insights"),
+                "chunk_results": chunk_results_text,
+            },
+        )
+        return llm.complete(prompt)
+
+    # ------------------------------------------------------------------
+    # Nodes (closures over llm/store/prompts/counter)
+    # ------------------------------------------------------------------
+
+    def chunk_input(state: MapReduceState) -> Dict[str, Any]:
+        chunks = store.store_context(
+            content=state["input_content"],
+            session_id=state["session_id"],
+            metadata={"task": state["task"]},
+            max_chunk_tokens=state.get("max_chunk_tokens", 3000),
+        )
+        total_tokens = sum(c.token_count for c in chunks)
+        return {
+            "chunks": chunks,
+            "total_input_tokens": total_tokens,
+            "repair_attempts": 0,
+            "processing_log": [
+                f"Chunked input into {len(chunks)} chunks ({total_tokens} total tokens)"
+            ],
+        }
+
+    def map_chunks(state: MapReduceState) -> Dict[str, Any]:
+        chunks = state["chunks"]
+        total = len(chunks)
+        log = list(state.get("processing_log", []))
+        results: List[str] = []
+        for chunk in chunks:
+            prompt = render_prompt(
+                prompts.map,
+                {
+                    "task": state["task"],
+                    "task_focus": state.get("task_focus", "key information and insights"),
+                    "chunk_index": chunk.chunk_index + 1,
+                    "total_chunks": total,
+                    "chunk_content": chunk.content,
+                },
+            )
+            results.append(llm.complete(prompt))
+            log.append(f"Mapped chunk {chunk.chunk_index + 1}/{total}: {len(results[-1])} chars")
+        return {"map_results": results, "processing_log": log}
+
+    def merge_results(state: MapReduceState) -> Dict[str, Any]:
+        results = state["map_results"]
+        log = list(state.get("processing_log", []))
+        max_context = state.get("max_context_tokens", 8192)
+
+        # Merge-input budget: keep every combine call comfortably inside the
+        # model window (60% head-room for instructions + output).
+        budget = max(1000, int(max_context * 0.6))
+
+        # HIERARCHICAL MERGE: with verbatim/exhaustive map phases the combined
+        # map results can exceed any model's input window. Merge in groups
+        # that fit the budget, recursively, until one result remains.
+        current = list(results)
+        passes = 0
+        while True:
+            total = sum(counter(r) for r in current) + 30 * len(current)
+            if total <= budget or len(current) <= 1:
+                break
+            avg = max(1, total // len(current))
+            # group size from CURRENT average result size: k * avg <= budget
+            k = max(2, budget // avg)
+            if k >= len(current):
+                break  # cannot split further; final merge below takes the risk
+            nxt: List[str] = []
+            for i in range(0, len(current), k):
+                group = current[i:i + k]
+                nxt.append(_combine(group, state, f"pass{passes + 1}-g{i // k + 1}"))
+            current = nxt
+            passes += 1
+            log.append(f"Hierarchical merge pass {passes}: "
+                       f"{len(results)} map results -> {len(current)} partials")
+
+        if len(current) == 1:
+            merged = current[0]
+        else:
+            merged = _combine(current, state, "final")
+
+        token_count = counter(merged)
+        log_msg = (f"Merged {len(results)} results into {token_count} tokens"
+                   + (f" ({passes} hierarchical pass(es))" if passes else ""))
+        log.append(log_msg)
+        return {
+            "merged_result": merged,
+            "merge_token_count": token_count,
+            "needs_repair": token_count > max_context,
+            "processing_log": log,
+        }
+
+    def repair_result(state: MapReduceState) -> Dict[str, Any]:
+        current_tokens = state["merge_token_count"]
+        max_context = state.get("max_context_tokens", 8192)
+        attempts = state.get("repair_attempts", 0)
+        prompt = render_prompt(
+            prompts.repair,
+            {
+                "task": state["task"],
+                "task_focus": state.get("task_focus", "key information and insights"),
+                "current_tokens": current_tokens,
+                "max_tokens": max_context,
+                "merged_result": state["merged_result"],
+            },
+        )
+        repaired = llm.complete(prompt)
+        repaired_tokens = counter(repaired)
+        attempts += 1
+        log = list(state.get("processing_log", []))
+        log.append(f"Repair attempt {attempts}: {current_tokens} -> {repaired_tokens} tokens")
+
+        still_too_long = repaired_tokens > max_context
+        can_retry = attempts < state.get("max_repair_attempts", 3)
+        update: Dict[str, Any] = {
+            "merged_result": repaired,
+            "merge_token_count": repaired_tokens,
+            "needs_repair": still_too_long and can_retry,
+            "repair_attempts": attempts,
+            "processing_log": log,
+        }
+        if not (still_too_long and can_retry):
+            update["final_result"] = repaired
+        return update
+
+    def finalize(state: MapReduceState) -> Dict[str, Any]:
+        final = state.get("final_result") or state.get("merged_result", "")
+        log = list(state.get("processing_log", [])) + ["Processing complete"]
+        return {"final_result": final, "processing_log": log}
+
+    # ------------------------------------------------------------------
+    # Wiring
+    # ------------------------------------------------------------------
+
+    def _route(state: MapReduceState) -> str:
+        return "repair_result" if state.get("needs_repair", False) else "finalize"
+
+    workflow = StateGraph(MapReduceState)
+    workflow.add_node("chunk_input", chunk_input)
+    workflow.add_node("map_chunks", map_chunks)
+    workflow.add_node("merge_results", merge_results)
+    workflow.add_node("repair_result", repair_result)
+    workflow.add_node("finalize", finalize)
+
+    workflow.set_entry_point("chunk_input")
+    workflow.add_edge("chunk_input", "map_chunks")
+    workflow.add_edge("map_chunks", "merge_results")
+    workflow.add_conditional_edges("merge_results", _route)
+    workflow.add_conditional_edges("repair_result", _route)
+    workflow.add_edge("finalize", END)
+
+    return workflow.compile()