dulus 0.2.0__py3-none-any.whl

batch_api.py

@@ -0,0 +1,307 @@
+"""
+Dulus Batch API — provider-agnostic OpenAI-compatible batch processing.
+
+Works with any provider that supports the OpenAI Batch API format:
+  - OpenAI (api.openai.com)
+  - Kimi/Moonshot (api.moonshot.ai)
+  - Any OpenAI-compatible endpoint
+
+Usage:
+    mgr = BatchManager(api_key="sk-...", base_url="https://api.openai.com")
+    jsonl = mgr.prepare_jsonl(["prompt1", "prompt2"], model="gpt-4o-mini")
+    file_id = mgr.upload_file(jsonl)
+    batch_id = mgr.create_batch(file_id)
+"""
+
+import json
+import urllib.request
+import os
+import time
+from typing import Optional, List, Dict, Any
+
+# ── Defaults ─────────────────────────────────────────────────────────────────
+
+OPENAI_BASE_URL = "https://api.openai.com"
+KIMI_BASE_URL   = "https://api.moonshot.ai"
+
+BATCH_SYSTEM_PROMPT = (
+    "You are Dulus, an AI assistant. You are processing a batch request — "
+    "respond directly to each task. Be concise, precise, and complete. "
+    "Output in the same language the user writes in. "
+    "No tool calls available — just answer with text."
+)
+
+
+# ── BatchManager ─────────────────────────────────────────────────────────────
+
+class BatchManager:
+    """Provider-agnostic manager for the OpenAI-compatible Batch API."""
+
+    def __init__(self, api_key: str, base_url: str = OPENAI_BASE_URL):
+        self.api_key  = api_key
+        self.base_url = base_url.rstrip("/")
+
+    def _headers(self, content_type: str = "application/json") -> dict:
+        return {
+            "Content-Type": content_type,
+            "Authorization": f"Bearer {self.api_key}",
+        }
+
+    # ── JSONL preparation ────────────────────────────────────────────────
+
+    def prepare_jsonl(
+        self,
+        prompts: List[str],
+        model: str = "gpt-4o-mini",
+        system_prompt: str = None,
+        endpoint: str = "/v1/chat/completions",
+    ) -> str:
+        """Convert a list of prompts into JSONL content for the Batch API.
+
+        Args:
+            prompts:       List of user prompts.
+            model:         Model name (provider-specific).
+            system_prompt: Defaults to BATCH_SYSTEM_PROMPT. Pass "" to omit.
+            endpoint:      API endpoint for each request.
+        """
+        if system_prompt is None:
+            system_prompt = BATCH_SYSTEM_PROMPT
+
+        lines = []
+        ts = int(time.time())
+        for i, prompt in enumerate(prompts):
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
+            messages.append({"role": "user", "content": prompt})
+
+            request = {
+                "custom_id": f"req_{ts}_{i}",
+                "method": "POST",
+                "url": endpoint,
+                "body": {
+                    "model": model,
+                    "messages": messages,
+                },
+            }
+            lines.append(json.dumps(request, ensure_ascii=False))
+        return "\n".join(lines)
+
+    # ── File upload (multipart/form-data) ────────────────────────────────
+
+    def upload_file(self, jsonl_content: str, filename: str = "batch_input.jsonl") -> str:
+        """Upload JSONL content and return the file_id."""
+        url = f"{self.base_url}/v1/files"
+        boundary = f"----DulusBatch{int(time.time())}"
+
+        parts = []
+        # purpose field
+        parts.append(f"--{boundary}\r\n"
+                      f'Content-Disposition: form-data; name="purpose"\r\n\r\n'
+                      f"batch")
+        # file field
+        parts.append(f"--{boundary}\r\n"
+                      f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'
+                      f"Content-Type: application/octet-stream\r\n\r\n"
+                      f"{jsonl_content}")
+        parts.append(f"--{boundary}--\r\n")
+
+        full_body = "\r\n".join(parts).encode("utf-8")
+
+        req = urllib.request.Request(
+            url,
+            data=full_body,
+            headers={
+                "Content-Type": f"multipart/form-data; boundary={boundary}",
+                "Authorization": f"Bearer {self.api_key}",
+            },
+            method="POST",
+        )
+
+        with urllib.request.urlopen(req) as resp:
+            return json.loads(resp.read().decode("utf-8"))["id"]
+
+    # ── Batch lifecycle ──────────────────────────────────────────────────
+
+    def create_batch(
+        self,
+        file_id: str,
+        endpoint: str = "/v1/chat/completions",
+        completion_window: str = "24h",
+    ) -> str:
+        """Create a batch from an uploaded file. Returns batch_id."""
+        url = f"{self.base_url}/v1/batches"
+        payload = {
+            "input_file_id": file_id,
+            "endpoint": endpoint,
+            "completion_window": completion_window,
+        }
+        req = urllib.request.Request(
+            url,
+            data=json.dumps(payload).encode("utf-8"),
+            headers=self._headers(),
+            method="POST",
+        )
+        with urllib.request.urlopen(req) as resp:
+            return json.loads(resp.read().decode("utf-8"))["id"]
+
+    def retrieve_batch(self, batch_id: str) -> Dict[str, Any]:
+        """Get batch status/info."""
+        url = f"{self.base_url}/v1/batches/{batch_id}"
+        req = urllib.request.Request(url, headers=self._headers(), method="GET")
+        with urllib.request.urlopen(req) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+
+    def cancel_batch(self, batch_id: str) -> Dict[str, Any]:
+        """Cancel a running batch."""
+        url = f"{self.base_url}/v1/batches/{batch_id}/cancel"
+        req = urllib.request.Request(url, headers=self._headers(), method="POST")
+        with urllib.request.urlopen(req) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+
+    def get_file_content(self, file_id: str) -> str:
+        """Download file content (e.g. batch results)."""
+        url = f"{self.base_url}/v1/files/{file_id}/content"
+        req = urllib.request.Request(url, headers=self._headers(), method="GET")
+        with urllib.request.urlopen(req) as resp:
+            return resp.read().decode("utf-8")
+
+
+# ── Backward compat alias ────────────────────────────────────────────────────
+KimiBatchManager = BatchManager  # old name still works
+
+
+# ── Local job persistence ────────────────────────────────────────────────────
+
+_JOBS_DIR = os.path.join(os.path.expanduser("~"), ".dulus", "jobs")
+
+
+def save_batch_job(batch_id: str, description: str = "", file_id: str = "",
+                   provider: str = "unknown") -> str:
+    """Save a batch job record locally in ~/.dulus/jobs/."""
+    os.makedirs(_JOBS_DIR, exist_ok=True)
+    job_file = os.path.join(_JOBS_DIR, f"{batch_id}.json")
+
+    job_data = {
+        "job_id": batch_id,
+        "id": batch_id,
+        "tool_name": "batch",
+        "provider": provider,
+        "params": {"description": description, "file_id": file_id},
+        "status": "created",
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
+        "description": description,
+        "file_id": file_id,
+        "batch_id": batch_id,
+    }
+
+    with open(job_file, "w", encoding="utf-8") as f:
+        json.dump(job_data, f, indent=2)
+    return job_file
+
+
+def list_batch_jobs(include_pollers: bool = True, **_kw) -> List[Dict]:
+    """List saved batch jobs from ~/.dulus/jobs/."""
+    if not os.path.exists(_JOBS_DIR):
+        return []
+
+    batch_map: Dict[str, Dict] = {}
+    poller_jobs: List[Dict] = []
+    # Accept both old "kimi_batch" and new "batch" tool_name
+    _batch_names  = {"kimi_batch", "batch"}
+    _poller_names = {"kimi_batch_poll", "batch_poll"}
+
+    for fname in os.listdir(_JOBS_DIR):
+        if not fname.endswith(".json"):
+            continue
+        try:
+            with open(os.path.join(_JOBS_DIR, fname), "r", encoding="utf-8") as f:
+                job = json.load(f)
+
+            tn = job.get("tool_name", "")
+            if tn in _batch_names:
+                bid = job.get("batch_id") or job.get("id")
+                if bid:
+                    batch_map[bid] = job
+
+            elif include_pollers and tn in _poller_names:
+                poller_jobs.append(job)
+                br = job.get("batch_result", {})
+                if br:
+                    bid = br.get("id")
+                    if bid and bid in batch_map:
+                        batch_map[bid]["status"]         = br.get("status", "unknown")
+                        batch_map[bid]["request_counts"]  = br.get("request_counts", {})
+                        batch_map[bid]["output_file_id"]  = br.get("output_file_id")
+                        batch_map[bid]["completed_at"]    = br.get("completed_at")
+                        batch_map[bid]["_poller_job_id"]  = job.get("job_id")
+        except Exception:
+            continue
+
+    # Pollers for batches not yet in map → synthetic entry
+    for poller in poller_jobs:
+        br  = poller.get("batch_result", {})
+        bid = br.get("id")
+        if bid and bid not in batch_map:
+            batch_map[bid] = {
+                "job_id": bid, "id": bid,
+                "tool_name": "batch",
+                "status": br.get("status", "unknown"),
+                "created_at": poller.get("created_at"),
+                "description": f"(from poller {poller.get('job_id', '?')[:8]}...)",
+                "batch_id": bid,
+                "request_counts": br.get("request_counts", {}),
+                "output_file_id": br.get("output_file_id"),
+                "completed_at": br.get("completed_at"),
+                "_from_poller": True,
+                "_poller_job_id": poller.get("job_id"),
+            }
+
+    jobs = list(batch_map.values())
+    jobs.sort(key=lambda x: x.get("created_at", ""), reverse=True)
+    return jobs
+
+
+def update_batch_job_status(batch_id: str, status_info: Dict[str, Any]) -> bool:
+    """Update a batch job's status in its local file."""
+    job_file = os.path.join(_JOBS_DIR, f"{batch_id}.json")
+    if not os.path.exists(job_file):
+        return False
+    try:
+        with open(job_file, "r", encoding="utf-8") as f:
+            job = json.load(f)
+        for key in ("status", "request_counts", "output_file_id", "completed_at"):
+            if key in status_info:
+                job[key] = status_info[key]
+        with open(job_file, "w", encoding="utf-8") as f:
+            json.dump(job, f, indent=2)
+        return True
+    except Exception:
+        return False
+
+
+def get_batch_job_by_id(batch_id: str) -> Optional[Dict]:
+    """Get a batch job by ID (checks both batch and poller files)."""
+    # Direct file
+    job_file = os.path.join(_JOBS_DIR, f"{batch_id}.json")
+    if os.path.exists(job_file):
+        try:
+            with open(job_file, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            pass
+
+    # Scan pollers
+    if os.path.exists(_JOBS_DIR):
+        for fname in os.listdir(_JOBS_DIR):
+            if not fname.endswith(".json"):
+                continue
+            try:
+                with open(os.path.join(_JOBS_DIR, fname), "r", encoding="utf-8") as f:
+                    job = json.load(f)
+                if job.get("tool_name") in ("kimi_batch_poll", "batch_poll"):
+                    if job.get("params", {}).get("batch_id") == batch_id:
+                        return job
+            except Exception:
+                continue
+    return None

checkpoint/__init__.py

@@ -0,0 +1,27 @@
+"""Checkpoint system: automatic file snapshots with rewind support."""
+from .types import FileBackup, Snapshot, MAX_SNAPSHOTS
+from .store import (
+    track_file_edit,
+    make_snapshot,
+    list_snapshots,
+    get_snapshot,
+    rewind_files,
+    files_changed_since,
+    delete_session_checkpoints,
+    cleanup_old_sessions,
+    reset_file_versions,
+)
+from .hooks import (
+    set_session,
+    get_tracked_edits,
+    reset_tracked,
+    install_hooks,
+)
+
+__all__ = [
+    "FileBackup", "Snapshot", "MAX_SNAPSHOTS",
+    "track_file_edit", "make_snapshot", "list_snapshots", "get_snapshot",
+    "rewind_files", "files_changed_since",
+    "delete_session_checkpoints", "cleanup_old_sessions", "reset_file_versions",
+    "set_session", "get_tracked_edits", "reset_tracked", "install_hooks",
+]

checkpoint/hooks.py

@@ -0,0 +1,90 @@
+"""Checkpoint hooks: intercept Write/Edit/NotebookEdit to back up files before modification.
+
+Import this module after tools are registered to install the hooks.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+from . import store
+
+# ── Module state ────────────────────────────────────────────────────────────
+
+_current_session_id: str | None = None
+_tracked_edits: dict[str, str | None] = {}   # file_path → backup_filename
+
+
+def set_session(session_id: str) -> None:
+    global _current_session_id
+    _current_session_id = session_id
+
+
+def get_tracked_edits() -> dict[str, str | None]:
+    """Return the current interval's tracked edits (for make_snapshot)."""
+    return dict(_tracked_edits)
+
+
+def reset_tracked() -> None:
+    """Clear tracked edits after a snapshot is created."""
+    _tracked_edits.clear()
+
+
+# ── Backup logic ────────────────────────────────────────────────────────────
+
+def _backup_before_write(file_path: str) -> None:
+    """Back up a file before it is modified (first-write-wins per snapshot interval)."""
+    if _current_session_id is None:
+        return
+    if file_path in _tracked_edits:
+        return  # already backed up this interval
+
+    backup_name = store.track_file_edit(_current_session_id, file_path)
+    _tracked_edits[file_path] = backup_name
+
+
+# ── Hook installation ───────────────────────────────────────────────────────
+
+_hooks_installed = False
+
+
+def install_hooks() -> None:
+    """Wrap Write/Edit/NotebookEdit tool functions to call backup before execution."""
+    global _hooks_installed
+    if _hooks_installed:
+        return
+    _hooks_installed = True
+
+    from tool_registry import get_tool
+
+    # Hook Write
+    write_tool = get_tool("Write")
+    if write_tool:
+        original_write = write_tool.func
+        def hooked_write(params, config):
+            fp = params.get("file_path", "")
+            if fp:
+                _backup_before_write(fp)
+            return original_write(params, config)
+        write_tool.func = hooked_write
+
+    # Hook Edit
+    edit_tool = get_tool("Edit")
+    if edit_tool:
+        original_edit = edit_tool.func
+        def hooked_edit(params, config):
+            fp = params.get("file_path", "")
+            if fp:
+                _backup_before_write(fp)
+            return original_edit(params, config)
+        edit_tool.func = hooked_edit
+
+    # Hook NotebookEdit
+    nb_tool = get_tool("NotebookEdit")
+    if nb_tool:
+        original_nb = nb_tool.func
+        def hooked_nb(params, config):
+            fp = params.get("notebook_path", "")
+            if fp:
+                _backup_before_write(fp)
+            return original_nb(params, config)
+        nb_tool.func = hooked_nb