svc-infra 0.1.595__py3-none-any.whl → 1.1.0__py3-none-any.whl

svc_infra/jobs/redis_queue.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import asdict
+from datetime import UTC, datetime
+from typing import Any, cast
+
+from redis import Redis
+
+from .queue import Job, JobQueue
+
+logger = logging.getLogger(__name__)
+
+# Lua script for atomic reserve: pop from ready, push to processing, set visibility timeout
+# Returns job_id if successful, nil if queue is empty
+_RESERVE_LUA = """
+local ready_key = KEYS[1]
+local processing_key = KEYS[2]
+local processing_vt_key = KEYS[3]
+local visible_at = ARGV[1]
+
+local job_id = redis.call('RPOPLPUSH', ready_key, processing_key)
+if job_id then
+    redis.call('ZADD', processing_vt_key, visible_at, job_id)
+end
+return job_id
+"""
+
+
+class RedisJobQueue(JobQueue):
+    """Redis-backed job queue with visibility timeout and delayed retries.
+
+    Keys (with optional prefix):
+      - {p}:ready (LIST)        ready job ids
+      - {p}:processing (LIST)   in-flight job ids
+      - {p}:processing_vt (ZSET) id -> visible_at (epoch seconds)
+      - {p}:delayed (ZSET)      id -> available_at (epoch seconds)
+      - {p}:seq (STRING)        INCR for job ids
+      - {p}:job:{id} (HASH)     job fields (json payload)
+      - {p}:dlq (LIST)          dead-letter job ids
+    """
+
+    def __init__(self, client: Redis, *, prefix: str = "jobs", visibility_timeout: int = 60):
+        self._r = client
+        self._p = prefix
+        self._vt = visibility_timeout
+        # Try to register Lua script for atomic reserve
+        # Falls back to non-atomic if Lua scripting isn't available (e.g., fakeredis in tests)
+        self._reserve_script = None
+        try:
+            self._reserve_script = client.register_script(_RESERVE_LUA)
+        except Exception as e:
+            logger.debug("Lua scripting not available, using non-atomic reserve: %s", e)
+
+    # Key helpers
+    def _k(self, name: str) -> str:
+        return f"{self._p}:{name}"
+
+    def _job_key(self, job_id: str) -> str:
+        return f"{self._p}:job:{job_id}"
+
+    # Core ops
+    def enqueue(self, name: str, payload: dict, *, delay_seconds: int = 0) -> Job:
+        now = datetime.now(UTC)
+        job_id = str(self._r.incr(self._k("seq")))
+        job = Job(id=job_id, name=name, payload=dict(payload))
+        # Persist job
+        data = asdict(job)
+        data["payload"] = json.dumps(data["payload"])  # store payload as JSON string
+        # available_at stored as ISO format
+        data["available_at"] = job.available_at.isoformat()
+        self._r.hset(
+            self._job_key(job_id),
+            mapping={k: str(v) for k, v in data.items() if v is not None},
+        )
+        if delay_seconds and delay_seconds > 0:
+            at = int(now.timestamp()) + int(delay_seconds)
+            self._r.zadd(self._k("delayed"), {job_id: at})
+        else:
+            # push to ready
+            self._r.lpush(self._k("ready"), job_id)
+        return job
+
+    def _move_due_delayed_to_ready(self) -> None:
+        now_ts = int(datetime.now(UTC).timestamp())
+        ids = cast("list[Any]", self._r.zrangebyscore(self._k("delayed"), "-inf", now_ts))
+        if not ids:
+            return
+        pipe = self._r.pipeline()
+        for jid in ids:
+            jid_s = jid.decode() if isinstance(jid, (bytes, bytearray)) else str(jid)
+            pipe.lpush(self._k("ready"), jid_s)
+            pipe.zrem(self._k("delayed"), jid_s)
+        pipe.execute()
+
+    def _requeue_timed_out_processing(self) -> None:
+        now_ts = int(datetime.now(UTC).timestamp())
+        ids = cast("list[Any]", self._r.zrangebyscore(self._k("processing_vt"), "-inf", now_ts))
+        if not ids:
+            return
+        pipe = self._r.pipeline()
+        for jid in ids:
+            jid_s = jid.decode() if isinstance(jid, (bytes, bytearray)) else str(jid)
+            pipe.lrem(self._k("processing"), 1, jid_s)
+            pipe.lpush(self._k("ready"), jid_s)
+            pipe.zrem(self._k("processing_vt"), jid_s)
+            # clear stale visibility timestamp so next reservation can set a fresh one
+            pipe.hdel(self._job_key(jid_s), "visible_at")
+        pipe.execute()
+
+    def reserve_next(self) -> Job | None:
+        # opportunistically move due delayed jobs
+        self._move_due_delayed_to_ready()
+        # move timed-out processing jobs back to ready before reserving
+        self._requeue_timed_out_processing()
+
+        # Calculate visibility timeout BEFORE reserve to prevent race condition
+        visible_at = int(datetime.now(UTC).timestamp()) + int(self._vt)
+
+        # Try atomic reserve using Lua script if available
+        # This prevents race condition where two workers could reserve the same job
+        if self._reserve_script is not None:
+            try:
+                jid = self._reserve_script(
+                    keys=[
+                        self._k("ready"),
+                        self._k("processing"),
+                        self._k("processing_vt"),
+                    ],
+                    args=[visible_at],
+                )
+            except Exception as e:
+                # Fall back to non-atomic if Lua fails at runtime
+                logger.warning("Lua script failed, using non-atomic reserve: %s", e)
+                jid = self._r.rpoplpush(self._k("ready"), self._k("processing"))
+                if jid:
+                    job_id_tmp = jid.decode() if isinstance(jid, (bytes, bytearray)) else str(jid)
+                    self._r.zadd(self._k("processing_vt"), {job_id_tmp: visible_at})
+        else:
+            # Non-atomic fallback (for fakeredis in tests, or older Redis versions)
+            jid = self._r.rpoplpush(self._k("ready"), self._k("processing"))
+            if jid:
+                job_id_tmp = jid.decode() if isinstance(jid, (bytes, bytearray)) else str(jid)
+                self._r.zadd(self._k("processing_vt"), {job_id_tmp: visible_at})
+
+        if not jid:
+            return None
+        job_id = jid.decode() if isinstance(jid, (bytes, bytearray)) else str(jid)
+        key = self._job_key(job_id)
+        data = cast("dict[Any, Any]", self._r.hgetall(key))
+        if not data:
+            # corrupted entry; ack and skip
+            self._r.lrem(self._k("processing"), 1, job_id)
+            self._r.zrem(self._k("processing_vt"), job_id)
+            return None
+
+        # Decode fields
+        def _get(field: str, default: str | None = None) -> str | None:
+            val = (
+                data.get(field.encode())
+                if isinstance(next(iter(data.keys())), bytes)
+                else data.get(field)
+            )
+            if val is None:
+                return default
+            return val.decode() if isinstance(val, (bytes, bytearray)) else str(val)
+
+        attempts = int(_get("attempts", "0") or "0") + 1
+        max_attempts = int(_get("max_attempts", "5") or "5")
+        backoff_seconds = int(_get("backoff_seconds", "60") or "60")
+        name = _get("name", "") or ""
+        payload_json = _get("payload", "{}") or "{}"
+        try:
+            payload = json.loads(payload_json)
+        except Exception:  # pragma: no cover
+            payload = {}
+        available_at_str = _get("available_at")
+        available_at = (
+            datetime.fromisoformat(available_at_str) if available_at_str else datetime.now(UTC)
+        )
+        # If exceeded max_attempts → DLQ and skip
+        if attempts > max_attempts:
+            self._r.lrem(self._k("processing"), 1, job_id)
+            self._r.zrem(self._k("processing_vt"), job_id)
+            self._r.lpush(self._k("dlq"), job_id)
+            return None
+        # Update attempts count in job hash (visibility timeout already set atomically in Lua script)
+        self._r.hset(key, mapping={"attempts": attempts, "visible_at": visible_at})
+        return Job(
+            id=job_id,
+            name=name,
+            payload=payload,
+            available_at=available_at,
+            attempts=attempts,
+            max_attempts=max_attempts,
+            backoff_seconds=backoff_seconds,
+        )
+
+    def ack(self, job_id: str) -> None:
+        self._r.lrem(self._k("processing"), 1, job_id)
+        self._r.zrem(self._k("processing_vt"), job_id)
+        self._r.delete(self._job_key(job_id))
+
+    def fail(self, job_id: str, *, error: str | None = None) -> None:
+        key = self._job_key(job_id)
+        data = cast("dict[Any, Any]", self._r.hgetall(key))
+        if not data:
+            # nothing to do
+            self._r.lrem(self._k("processing"), 1, job_id)
+            return
+
+        def _get(field: str, default: str | None = None) -> str | None:
+            val = (
+                data.get(field.encode())
+                if isinstance(next(iter(data.keys())), bytes)
+                else data.get(field)
+            )
+            if val is None:
+                return default
+            return val.decode() if isinstance(val, (bytes, bytearray)) else str(val)
+
+        attempts = int(_get("attempts", "0") or "0")
+        max_attempts = int(_get("max_attempts", "5") or "5")
+        backoff_seconds = int(_get("backoff_seconds", "60") or "60")
+        now_ts = int(datetime.now(UTC).timestamp())
+        # DLQ if at or beyond max_attempts
+        if attempts >= max_attempts:
+            self._r.lrem(self._k("processing"), 1, job_id)
+            self._r.zrem(self._k("processing_vt"), job_id)
+            self._r.lpush(self._k("dlq"), job_id)
+            return
+        delay = backoff_seconds * max(1, attempts)
+        available_at_ts = now_ts + delay
+        mapping: dict[str, str] = {
+            "last_error": error or "",
+            "available_at": datetime.fromtimestamp(available_at_ts, tz=UTC).isoformat(),
+        }
+        self._r.hset(key, mapping=mapping)
+        self._r.lrem(self._k("processing"), 1, job_id)
+        self._r.zrem(self._k("processing_vt"), job_id)
+        self._r.zadd(self._k("delayed"), {job_id: available_at_ts})

svc_infra/jobs/runner.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import Awaitable, Callable
+
+from .queue import JobQueue
+
+ProcessFunc = Callable[[object], Awaitable[None]]
+
+
+class WorkerRunner:
+    """Cooperative worker loop with graceful stop.
+
+    - start(): begin polling the queue and processing jobs
+    - stop(grace_seconds): signal stop, wait up to grace for current job to finish
+    """
+
+    def __init__(self, queue: JobQueue, handler: ProcessFunc, *, poll_interval: float = 0.25):
+        self._queue = queue
+        self._handler = handler
+        self._poll_interval = poll_interval
+        self._task: asyncio.Task | None = None
+        self._stopping = asyncio.Event()
+        self._inflight: asyncio.Task | None = None
+
+    async def _loop(self) -> None:
+        try:
+            while not self._stopping.is_set():
+                job = self._queue.reserve_next()
+                if not job:
+                    await asyncio.sleep(self._poll_interval)
+                    continue
+
+                # Process one job; track in-flight task for stop()
+                async def _run():
+                    try:
+                        await self._handler(job)
+                    except Exception as exc:  # pragma: no cover
+                        self._queue.fail(job.id, error=str(exc))
+                        return
+                    self._queue.ack(job.id)
+
+                self._inflight = asyncio.create_task(_run())
+                try:
+                    await self._inflight
+                finally:
+                    self._inflight = None
+        finally:
+            # exiting loop
+            pass
+
+    def start(self) -> asyncio.Task:
+        if self._task is None or self._task.done():
+            self._task = asyncio.create_task(self._loop())
+        return self._task
+
+    async def stop(self, *, grace_seconds: float = 10.0) -> None:
+        self._stopping.set()
+        # Wait for in-flight job to complete, up to grace
+        if self._inflight is not None and not self._inflight.done():
+            try:
+                await asyncio.wait_for(self._inflight, timeout=grace_seconds)
+            except TimeoutError:
+                # Give up; job will be retried if your queue supports visibility timeouts
+                pass
+        # Finally, wait for loop to exit (should be quick since stopping is set)
+        if self._task is not None:
+            try:
+                await asyncio.wait_for(self._task, timeout=max(0.1, self._poll_interval + 0.1))
+            except TimeoutError:
+                # Cancel as a last resort
+                self._task.cancel()
+                with contextlib.suppress(Exception):
+                    await self._task

svc_infra/jobs/scheduler.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+
+CronFunc = Callable[[], Awaitable[None]]
+
+
+@dataclass
+class ScheduledTask:
+    name: str
+    interval_seconds: int
+    func: CronFunc
+    next_run_at: datetime
+
+
+class InMemoryScheduler:
+    """Interval-based scheduler for simple periodic tasks (tests/local).
+
+    Not a full cron parser. Tracks next_run_at per task.
+    """
+
+    def __init__(self, tick_interval: float = 60.0):
+        self._tasks: dict[str, ScheduledTask] = {}
+        self._tick_interval = tick_interval
+
+    def add_task(self, name: str, interval_seconds: int, func: CronFunc) -> None:
+        now = datetime.now(UTC)
+        self._tasks[name] = ScheduledTask(
+            name=name,
+            interval_seconds=interval_seconds,
+            func=func,
+            next_run_at=now + timedelta(seconds=interval_seconds),
+        )
+
+    async def tick(self) -> None:
+        now = datetime.now(UTC)
+        for task in self._tasks.values():
+            if task.next_run_at <= now:
+                await task.func()
+                task.next_run_at = now + timedelta(seconds=task.interval_seconds)
+
+    async def run(self) -> None:
+        """Run the scheduler loop indefinitely.
+
+        Calls tick() at regular intervals to check and execute due tasks.
+        This method runs forever until cancelled.
+        """
+        while True:
+            await self.tick()
+            await asyncio.sleep(self._tick_interval)

svc_infra/jobs/worker.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import Awaitable, Callable
+
+from .queue import Job, JobQueue
+
+ProcessFunc = Callable[[Job], Awaitable[None]]
+
+
+def _get_job_timeout_seconds() -> float | None:
+    raw = os.getenv("JOB_DEFAULT_TIMEOUT_SECONDS")
+    if not raw:
+        return None
+    try:
+        return float(raw)
+    except ValueError:
+        return None
+
+
+async def process_one(queue: JobQueue, handler: ProcessFunc) -> bool:
+    """Reserve a job, process with handler, ack on success or fail with backoff.
+
+    Returns True if a job was processed (success or fail), False if no job was available.
+    """
+    job = queue.reserve_next()
+    if not job:
+        return False
+    try:
+        timeout = _get_job_timeout_seconds()
+        if timeout and timeout > 0:
+            await asyncio.wait_for(handler(job), timeout=timeout)
+        else:
+            await handler(job)
+    except Exception as exc:  # pragma: no cover - exercise in tests by raising
+        queue.fail(job.id, error=str(exc))
+        return True
+    queue.ack(job.id)
+    return True

svc_infra/loaders/__init__.py

@@ -0,0 +1,186 @@
+"""
+Content loaders for fetching from remote and local sources.
+
+This module provides async-first loaders for GitHub, URLs, and other sources.
+All loaders return a consistent `LoadedContent` format that is compatible
+with ai-infra's Retriever.add_text() method.
+
+Quick Start:
+    >>> from svc_infra.loaders import GitHubLoader, URLLoader
+    >>>
+    >>> # Load from GitHub
+    >>> loader = GitHubLoader("nfraxlab/svc-infra", path="docs")
+    >>> contents = await loader.load()
+    >>>
+    >>> # Load from URL
+    >>> loader = URLLoader("https://example.com/guide.md")
+    >>> contents = await loader.load()
+    >>>
+    >>> # Sync usage (for scripts/notebooks)
+    >>> contents = loader.load_sync()
+
+With ai-infra Retriever:
+    >>> from ai_infra import Retriever
+    >>> from svc_infra.loaders import GitHubLoader
+    >>>
+    >>> retriever = Retriever()
+    >>> loader = GitHubLoader("nfraxlab/svc-infra", path="docs")
+    >>>
+    >>> for content in await loader.load():
+    ...     retriever.add_text(content.content, metadata=content.metadata)
+
+Convenience Functions:
+    >>> from svc_infra.loaders import load_github, load_url
+    >>>
+    >>> # One-liner loading
+    >>> contents = await load_github("nfraxlab/svc-infra", path="docs")
+    >>> contents = await load_url("https://example.com/guide.md")
+
+Available Loaders:
+    - GitHubLoader: Load files from GitHub repositories
+    - URLLoader: Load content from URLs (with HTML text extraction)
+
+Future Loaders (planned):
+    - S3Loader: Load files from S3-compatible storage
+    - NotionLoader: Load pages from Notion
+    - ConfluenceLoader: Load pages from Confluence
+"""
+
+from .base import BaseLoader
+from .github import GitHubLoader
+from .models import LoadedContent, LoadedDocument, to_loaded_documents
+from .url import URLLoader
+
+
+async def load_github(
+    repo: str,
+    path: str = "",
+    branch: str = "main",
+    pattern: str = "*.md",
+    **kwargs,
+) -> list[LoadedContent]:
+    """Convenience function to load content from GitHub.
+
+    This is a shortcut for creating a GitHubLoader and calling load().
+
+    Args:
+        repo: Repository in "owner/repo" format
+        path: Path within repo (empty for root)
+        branch: Branch name (default: "main")
+        pattern: Glob pattern for files (default: "*.md")
+        **kwargs: Additional arguments passed to GitHubLoader
+
+    Returns:
+        List of LoadedContent objects.
+
+    Example:
+        >>> contents = await load_github("nfraxlab/svc-infra", path="docs")
+        >>> for c in contents:
+        ...     print(f"{c.source}: {len(c.content)} chars")
+    """
+    loader = GitHubLoader(repo, path=path, branch=branch, pattern=pattern, **kwargs)
+    return await loader.load()
+
+
+async def load_url(
+    urls: str | list[str],
+    **kwargs,
+) -> list[LoadedContent]:
+    """Convenience function to load content from URL(s).
+
+    This is a shortcut for creating a URLLoader and calling load().
+
+    Args:
+        urls: Single URL or list of URLs to load
+        **kwargs: Additional arguments passed to URLLoader
+
+    Returns:
+        List of LoadedContent objects.
+
+    Example:
+        >>> # Single URL
+        >>> contents = await load_url("https://example.com/guide.md")
+        >>>
+        >>> # Multiple URLs
+        >>> contents = await load_url([
+        ...     "https://example.com/page1",
+        ...     "https://example.com/page2",
+        ... ])
+    """
+    loader = URLLoader(urls, **kwargs)
+    return await loader.load()
+
+
+def load_github_sync(
+    repo: str,
+    path: str = "",
+    branch: str = "main",
+    pattern: str = "*.md",
+    **kwargs,
+) -> list[LoadedContent]:
+    """Synchronous convenience function to load content from GitHub.
+
+    This is a shortcut for creating a GitHubLoader and calling load_sync().
+    Use this in scripts, notebooks, or non-async contexts.
+
+    Args:
+        repo: Repository in "owner/repo" format
+        path: Path within repo (empty for root)
+        branch: Branch name (default: "main")
+        pattern: Glob pattern for files (default: "*.md")
+        **kwargs: Additional arguments passed to GitHubLoader
+
+    Returns:
+        List of LoadedContent objects.
+
+    Example:
+        >>> # In a script or notebook (no await needed)
+        >>> contents = load_github_sync("nfraxlab/svc-infra", path="docs")
+        >>> for c in contents:
+        ...     print(f"{c.source}: {len(c.content)} chars")
+    """
+    loader = GitHubLoader(repo, path=path, branch=branch, pattern=pattern, **kwargs)
+    return loader.load_sync()
+
+
+def load_url_sync(
+    urls: str | list[str],
+    **kwargs,
+) -> list[LoadedContent]:
+    """Synchronous convenience function to load content from URL(s).
+
+    This is a shortcut for creating a URLLoader and calling load_sync().
+    Use this in scripts, notebooks, or non-async contexts.
+
+    Args:
+        urls: Single URL or list of URLs to load
+        **kwargs: Additional arguments passed to URLLoader
+
+    Returns:
+        List of LoadedContent objects.
+
+    Example:
+        >>> # In a script or notebook (no await needed)
+        >>> contents = load_url_sync("https://example.com/guide.md")
+    """
+    loader = URLLoader(urls, **kwargs)
+    return loader.load_sync()
+
+
+__all__ = [
+    # Base classes
+    "BaseLoader",
+    "LoadedContent",
+    # Compatibility
+    "LoadedDocument",
+    "to_loaded_documents",
+    # Loaders
+    "GitHubLoader",
+    "URLLoader",
+    # Async convenience functions
+    "load_github",
+    "load_url",
+    # Sync convenience functions
+    "load_github_sync",
+    "load_url_sync",
+]