source-kb 0.2.2__py3-none-any.whl

engine/runner.py

@@ -0,0 +1,335 @@
+"""Batch execution runner — concurrent LLM task execution with retry and heartbeat.
+
+Manages the lifecycle of sub-agent tasks: dispatch, monitor, retry, validate.
+
+Usage:
+    from engine.runner import run_batch, SubagentTask, SubagentResult
+
+    tasks = [SubagentTask(task_id="mod__business", prompt="...", output_path=...)]
+    results = run_batch(tasks, strategy, max_concurrent=5)
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+from core.interfaces import LlmRequest, LlmResponse, LlmStrategy
+from core.monitor.progress import write_progress
+
+logger = logging.getLogger(__name__)
+
+MAX_RETRIES = 2
+MIN_DOC_SIZE_BYTES = 500
+
+
+class BatchAbortError(Exception):
+    """Raised when circuit breaker triggers batch abort."""
+
+    def __init__(self, reason: str, diagnosis: str, failed_count: int, total_count: int):
+        self.reason = reason
+        self.diagnosis = diagnosis
+        self.failed_count = failed_count
+        self.total_count = total_count
+        super().__init__(f"Batch aborted: {reason} ({failed_count}/{total_count} failed)")
+
+
+@dataclass
+class CircuitBreakerConfig:
+    """Configuration for batch circuit breaker."""
+    max_consecutive_failures: int = 3
+    failure_rate_threshold: float = 0.5
+    timeout_error_keywords: tuple[str, ...] = ("timeout", "Timeout", "timed out", "PoolTimeout", "ConnectTimeout")
+    garbage_error_keywords: tuple[str, ...] = ("tool-call artifacts", "web_search", "read_file")
+
+    @classmethod
+    def from_config(cls, config: dict) -> CircuitBreakerConfig:
+        limits = config.get("limits", {})
+        return cls(
+            max_consecutive_failures=limits.get("max_consecutive_failures", 3),
+            failure_rate_threshold=limits.get("failure_rate_threshold", 0.5),
+        )
+
+
+class _CircuitBreaker:
+    """Tracks failures within a batch and triggers abort when thresholds are exceeded."""
+
+    def __init__(self, total_tasks: int, config: CircuitBreakerConfig):
+        self._total = total_tasks
+        self._config = config
+        self._lock = threading.Lock()
+        self._consecutive_failures = 0
+        self._total_failures = 0
+        self._total_completed = 0
+        self._timeout_count = 0
+        self._garbage_count = 0
+        self._tripped = False
+        self._trip_reason = ""
+        self._trip_diagnosis = ""
+
+    def record_success(self) -> None:
+        with self._lock:
+            self._consecutive_failures = 0
+            self._total_completed += 1
+
+    def record_failure(self, error: str) -> None:
+        with self._lock:
+            self._consecutive_failures += 1
+            self._total_failures += 1
+            self._total_completed += 1
+
+            if any(kw in error for kw in self._config.timeout_error_keywords):
+                self._timeout_count += 1
+            if any(kw in error for kw in self._config.garbage_error_keywords):
+                self._garbage_count += 1
+
+            self._check_trip(error)
+
+    def _check_trip(self, last_error: str) -> None:
+        if self._consecutive_failures >= self._config.max_consecutive_failures:
+            self._tripped = True
+            self._trip_reason = f"{self._consecutive_failures} consecutive failures"
+            self._trip_diagnosis = self._diagnose()
+        elif (self._total_completed >= 3 and
+              self._total_failures / self._total_completed > self._config.failure_rate_threshold):
+            self._tripped = True
+            self._trip_reason = f"Failure rate too high ({self._total_failures}/{self._total_completed})"
+            self._trip_diagnosis = self._diagnose()
+
+    def _diagnose(self) -> str:
+        if self._timeout_count >= 2:
+            return "Network timeout. Check: 1) Is the LLM API reachable? 2) Is the network stable? 3) Consider increasing subagent_timeout"
+        if self._garbage_count >= 2:
+            return "Model output error (tool-call artifacts). Check: 1) Does the model support plain text generation? 2) Is the prompt too large causing model confusion? 3) Try a different model"
+        if self._total_failures == self._total_completed:
+            return "All tasks failed. Check: 1) Is the API connection working? 2) Is the API key valid? 3) Is the model name correct?"
+        return "Most tasks failed. Check prompt content and size in .meta/prompts/ to confirm it does not exceed the model context window"
+
+    @property
+    def is_tripped(self) -> bool:
+        with self._lock:
+            return self._tripped
+
+    def get_abort_error(self) -> BatchAbortError | None:
+        with self._lock:
+            if not self._tripped:
+                return None
+            return BatchAbortError(
+                reason=self._trip_reason,
+                diagnosis=self._trip_diagnosis,
+                failed_count=self._total_failures,
+                total_count=self._total,
+            )
+
+
+@dataclass
+class SubagentTask:
+    """A single document generation task."""
+    task_id: str
+    prompt: str
+    output_path: Path
+    doc_type: str
+    system_prompt: str = (
+        "You are a technical documentation expert. "
+        "Generate a knowledge base document based ONLY on the source code provided in the user prompt. "
+        "You have NO tools available — do NOT output <web_search>, <read_file>, or any XML tool calls. "
+        "Do NOT search the web. Do NOT attempt to read files. All source code is already inlined in the prompt. "
+        "Output ONLY the final markdown document content, nothing else."
+    )
+    timeout: int = 900
+    model: str | None = None
+
+
+@dataclass
+class SubagentResult:
+    """Result of a sub-agent task execution."""
+    task_id: str
+    status: Literal["done", "failed", "empty", "delegated"]
+    output_path: Path | None = None
+    content: str = ""
+    elapsed: float = 0.0
+    error: str = ""
+    usage: dict[str, int] = field(default_factory=dict)
+    prompt_file: str = ""
+
+
+def run_single(task: SubagentTask, strategy: LlmStrategy, *, save_prompts: bool = False) -> SubagentResult:
+    """Execute a single task via the LLM strategy."""
+    module_dir = task.output_path.parent
+    write_progress(module_dir, task.doc_type, "START")
+
+    if save_prompts:
+        from core.monitor.prompt_store import save_prompt
+        save_prompt(module_dir, task.task_id, task.system_prompt, task.prompt, model=task.model or "")
+
+    t0 = time.time()
+    try:
+        resp: LlmResponse = strategy.call(LlmRequest(
+            system=task.system_prompt,
+            user=task.prompt,
+            model=task.model,
+            max_tokens=16384,
+            temperature=0.3,
+        ))
+        elapsed = time.time() - t0
+
+        if resp.status == "failed":
+            write_progress(module_dir, task.doc_type, f"FAILED: {resp.error}")
+            return SubagentResult(
+                task_id=task.task_id, status="failed", elapsed=elapsed,
+                error=resp.error or "LLM returned failed status", usage=resp.usage,
+            )
+
+        if resp.status == "delegated":
+            write_progress(module_dir, task.doc_type, "DELEGATED")
+            return SubagentResult(
+                task_id=task.task_id, status="delegated",
+                output_path=task.output_path, elapsed=elapsed,
+                usage=resp.usage, prompt_file=resp.prompt_file,
+            )
+
+        content = resp.content.strip()
+        content = _strip_markdown_fence(content)
+
+        if _is_garbage_output(content):
+            write_progress(module_dir, task.doc_type, "GARBAGE")
+            return SubagentResult(
+                task_id=task.task_id, status="empty",
+                content="", elapsed=elapsed,
+                error="Output contains tool-call artifacts (web_search/read_file)", usage=resp.usage,
+            )
+
+        if not content or len(content.encode("utf-8")) < MIN_DOC_SIZE_BYTES:
+            write_progress(module_dir, task.doc_type, "EMPTY")
+            return SubagentResult(
+                task_id=task.task_id, status="empty",
+                content=content, elapsed=elapsed,
+                error=f"Output too small ({len(content)} chars)", usage=resp.usage,
+            )
+
+        # Write output
+        task.output_path.parent.mkdir(parents=True, exist_ok=True)
+        task.output_path.write_text(content, encoding="utf-8")
+
+        # Verify write
+        if not task.output_path.exists() or task.output_path.stat().st_size < MIN_DOC_SIZE_BYTES:
+            write_progress(module_dir, task.doc_type, "WRITE_FAILED")
+            return SubagentResult(
+                task_id=task.task_id, status="failed", elapsed=elapsed,
+                error="Write verification failed", usage=resp.usage,
+            )
+
+        write_progress(module_dir, task.doc_type, f"DONE size={task.output_path.stat().st_size}")
+        return SubagentResult(
+            task_id=task.task_id, status="done",
+            output_path=task.output_path, content=content,
+            elapsed=elapsed, usage=resp.usage,
+        )
+
+    except Exception as e:
+        elapsed = time.time() - t0
+        write_progress(module_dir, task.doc_type, f"ERROR: {e}")
+        return SubagentResult(task_id=task.task_id, status="failed", elapsed=elapsed, error=str(e))
+
+
+_GARBAGE_PATTERNS_DEFAULT = ("<web_search>", "<read_file>", "<search_results>", "Search results for")
+
+_garbage_patterns_override: tuple[str, ...] | None = None
+
+
+def configure_garbage_patterns(patterns: list[str] | None) -> None:
+    """Set garbage patterns from preset config. Call once at pipeline startup."""
+    global _garbage_patterns_override
+    if patterns:
+        _garbage_patterns_override = tuple(patterns)
+
+
+def _strip_markdown_fence(content: str) -> str:
+    """Strip outer markdown code fence if LLM wrapped the output."""
+    if content.startswith("```markdown") and content.endswith("```"):
+        return content[len("```markdown"):-(len("```"))].strip()
+    if content.startswith("```md") and content.endswith("```"):
+        return content[len("```md"):-(len("```"))].strip()
+    if content.startswith("```") and content.endswith("```"):
+        first_nl = content.index("\n") if "\n" in content else 3
+        return content[first_nl + 1:-(len("```"))].strip()
+    return content
+
+
+def _is_garbage_output(content: str) -> bool:
+    """Detect LLM outputs that contain tool-call artifacts instead of documentation."""
+    patterns = _garbage_patterns_override or _GARBAGE_PATTERNS_DEFAULT
+    first_500 = content[:500]
+    return any(p in first_500 for p in patterns)
+
+
+def _run_with_retry(task: SubagentTask, strategy: LlmStrategy, *, save_prompts: bool = False) -> SubagentResult:
+    """Run a task with retry on failure/empty."""
+    for attempt in range(1, MAX_RETRIES + 1):
+        result = run_single(task, strategy, save_prompts=save_prompts and attempt == 1)
+        if result.status in ("done", "delegated"):
+            return result
+        if attempt < MAX_RETRIES:
+            logger.warning("[%s] attempt %d/%d %s — retrying", task.task_id, attempt, MAX_RETRIES, result.status)
+            time.sleep(2 ** attempt)
+    return result
+
+
+def run_batch(
+    tasks: list[SubagentTask],
+    strategy: LlmStrategy,
+    max_concurrent: int = 5,
+    *,
+    save_prompts: bool = False,
+    breaker_config: CircuitBreakerConfig | None = None,
+) -> list[SubagentResult]:
+    """Execute a batch of tasks with concurrency control, retry, and circuit breaker.
+
+    Raises BatchAbortError if the circuit breaker trips (consecutive failures
+    or failure rate exceeds threshold).
+
+    Returns results in the same order as input tasks.
+    """
+    if not tasks:
+        return []
+
+    results: dict[str, SubagentResult] = {}
+    effective = min(max_concurrent, len(tasks))
+    breaker = _CircuitBreaker(len(tasks), breaker_config or CircuitBreakerConfig())
+
+    logger.info("Starting batch: %d tasks, concurrency=%d", len(tasks), effective)
+
+    with ThreadPoolExecutor(max_workers=effective) as pool:
+        futures = {pool.submit(_run_with_retry, t, strategy, save_prompts=save_prompts): t.task_id for t in tasks}
+        for future in as_completed(futures):
+            tid = futures[future]
+            try:
+                result = future.result()
+                results[tid] = result
+                if result.status in ("done", "delegated"):
+                    breaker.record_success()
+                else:
+                    breaker.record_failure(result.error)
+            except Exception as e:
+                results[tid] = SubagentResult(task_id=tid, status="failed", error=str(e))
+                breaker.record_failure(str(e))
+
+            if breaker.is_tripped:
+                abort_err = breaker.get_abort_error()
+                # Cancel remaining futures
+                for f in futures:
+                    f.cancel()
+                logger.error("Circuit breaker tripped: %s", abort_err)
+                raise abort_err
+
+    ordered = [results[t.task_id] for t in tasks]
+    done = sum(1 for r in ordered if r.status == "done")
+    failed = sum(1 for r in ordered if r.status in ("failed", "empty"))
+    logger.info("Batch complete: %d done, %d failed", done, failed)
+    return ordered
+

engine/strategies/__init__.py

@@ -0,0 +1,86 @@
+"""LLM strategy selection based on configuration.
+
+Selects the appropriate LlmStrategy implementation based on agent.model config:
+  - "delegated" → DelegatedLlmStrategy
+  - "dry-run"   → DryRunLlmStrategy
+  - anything else → ApiLlmStrategy (OpenAI-compatible)
+
+Usage:
+    from engine.strategies import create_strategy, ConfigProxy
+    strategy = create_strategy(ConfigProxy(raw_config))
+"""
+
+from __future__ import annotations
+
+import os
+
+from core.interfaces import LlmStrategy
+
+
+class ConfigProxy:
+    """Adapts a raw config dict to the interface expected by create_strategy().
+
+    Resolves values from environment variables first, then config dict.
+    """
+
+    def __init__(self, config: dict, *, timeout_override: int | None = None):
+        self._raw = config
+        self._timeout_override = timeout_override
+
+    @property
+    def agent_backend(self):
+        model = os.getenv("LLM_MODEL", "") or self._raw.get("agent", {}).get("model", "")
+        if model == "delegated": return "delegated"
+        if model == "dry-run": return "dry-run"
+        return "api"
+
+    @property
+    def agent_base_url(self):
+        return os.getenv("LLM_BASE_URL", "") or self._raw.get("agent", {}).get("base_url", "")
+
+    @property
+    def agent_api_key(self):
+        return os.getenv("LLM_API_KEY", "") or self._raw.get("agent", {}).get("api_key", "")
+
+    @property
+    def agent_model(self):
+        return os.getenv("LLM_MODEL", "") or self._raw.get("agent", {}).get("model", "")
+
+    @property
+    def agent_timeout(self):
+        if self._timeout_override is not None:
+            return self._timeout_override
+        env = os.getenv("KB_AGENT_TIMEOUT", "")
+        if env:
+            try:
+                return int(env)
+            except ValueError:
+                pass
+        return self._raw.get("agent", {}).get("subagent_timeout", 900)
+
+
+def create_strategy(config) -> LlmStrategy:
+    """Create the appropriate LLM strategy from config.
+
+    Args:
+        config: ConfigProxy instance or any object with agent_* properties
+
+    Returns:
+        LlmStrategy implementation
+    """
+    backend = config.agent_backend
+
+    if backend == "delegated":
+        from engine.strategies.delegated import DelegatedLlmStrategy
+        return DelegatedLlmStrategy()
+    elif backend == "dry-run":
+        from engine.strategies.dryrun import DryRunLlmStrategy
+        return DryRunLlmStrategy()
+    else:
+        from engine.strategies.api import ApiLlmStrategy
+        return ApiLlmStrategy(
+            base_url=config.agent_base_url,
+            api_key=config.agent_api_key,
+            model=config.agent_model,
+            timeout=config.agent_timeout,
+        )

engine/strategies/api.py

@@ -0,0 +1,128 @@
+"""OpenAI-compatible API LLM strategy.
+
+Calls any LLM provider supporting /v1/chat/completions endpoint.
+Includes retry with exponential backoff, rate-limit handling, and timeout.
+
+Supports: Anthropic, OpenAI, DeepSeek, Moonshot, Ollama, vLLM, LocalAI, etc.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import requests
+
+from core.interfaces import LlmRequest, LlmResponse, LlmStrategy
+
+logger = logging.getLogger(__name__)
+
+MAX_RETRIES = 3
+
+
+class ApiLlmStrategy(LlmStrategy):
+    """LLM execution via OpenAI-compatible HTTP API."""
+
+    def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 900):
+        if not base_url:
+            raise RuntimeError(
+                "LLM API endpoint not configured. Set LLM_BASE_URL env var.\n"
+                "Examples: https://api.anthropic.com, https://api.openai.com, http://localhost:11434"
+            )
+        self._base_url = base_url.rstrip("/")
+        self._api_key = api_key
+        self._model = model
+        self._timeout = timeout
+        self._session = requests.Session()
+        self._session.headers.update({"Content-Type": "application/json"})
+        if api_key:
+            self._session.headers["Authorization"] = f"Bearer {api_key}"
+
+    def call(self, request: LlmRequest) -> LlmResponse:
+        model = request.model or self._model
+
+        body = {
+            "model": model,
+            "max_tokens": request.max_tokens,
+            "temperature": request.temperature,
+            "messages": [
+                {"role": "system", "content": request.system},
+                {"role": "user", "content": request.user},
+            ],
+        }
+
+        url = f"{self._base_url}/v1/chat/completions"
+
+        for attempt in range(1, MAX_RETRIES + 1):
+            t0 = time.time()
+            try:
+                resp = self._session.post(url, json=body, timeout=self._timeout)
+            except requests.exceptions.Timeout:
+                if attempt >= MAX_RETRIES:
+                    return LlmResponse(content="", status="failed", error=f"Timeout after {MAX_RETRIES} retries")
+                time.sleep(2 ** attempt)
+                continue
+            except requests.exceptions.ConnectionError as e:
+                if attempt >= MAX_RETRIES:
+                    return LlmResponse(content="", status="failed", error=f"Connection failed: {e}")
+                time.sleep(2 ** attempt)
+                continue
+
+            elapsed = time.time() - t0
+
+            if resp.status_code == 429:
+                retry_after = int(resp.headers.get("Retry-After", 2 ** attempt))
+                if attempt >= MAX_RETRIES:
+                    return LlmResponse(content="", status="failed", error="Rate limited")
+                time.sleep(retry_after)
+                continue
+
+            if resp.status_code >= 500:
+                if attempt >= MAX_RETRIES:
+                    return LlmResponse(content="", status="failed", error=f"Server error {resp.status_code}")
+                time.sleep(2 ** attempt)
+                continue
+
+            if resp.status_code >= 400:
+                return LlmResponse(content="", status="failed",
+                                   error=f"HTTP {resp.status_code}: {resp.text[:300]}")
+
+            try:
+                data = resp.json()
+            except ValueError:
+                return LlmResponse(content="", status="failed", error="Non-JSON response")
+
+            choices = data.get("choices", [])
+            if not choices:
+                return LlmResponse(content="", status="failed", error="Empty choices")
+
+            usage = data.get("usage", {})
+            return LlmResponse(
+                content=choices[0]["message"]["content"],
+                status="done",
+                usage={"prompt_tokens": usage.get("prompt_tokens", 0),
+                       "completion_tokens": usage.get("completion_tokens", 0)},
+                elapsed=elapsed,
+            )
+
+        return LlmResponse(content="", status="failed", error="All retries exhausted")
+
+    def call_batch(self, requests_list: list[LlmRequest], max_concurrent: int = 5) -> list[LlmResponse]:
+        """Execute multiple requests with ThreadPoolExecutor."""
+        if not requests_list:
+            return []
+
+        results: dict[int, LlmResponse] = {}
+        effective = min(max_concurrent, len(requests_list))
+
+        with ThreadPoolExecutor(max_workers=effective) as pool:
+            futures = {pool.submit(self.call, req): i for i, req in enumerate(requests_list)}
+            for future in as_completed(futures):
+                idx = futures[future]
+                try:
+                    results[idx] = future.result()
+                except Exception as e:
+                    results[idx] = LlmResponse(content="", status="failed", error=str(e))
+
+        return [results[i] for i in range(len(requests_list))]

engine/strategies/delegated.py

@@ -0,0 +1,50 @@
+"""Delegated LLM strategy — writes prompt JSON for external agent execution.
+
+Instead of calling an API, generates prompt files that an external agent
+(or human) can execute. Used with `--resume` to continue the pipeline
+after prompts are fulfilled.
+
+Usage:
+    strategy = DelegatedLlmStrategy()
+    response = strategy.call(request)  # response.status == "delegated"
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+
+from core.interfaces import LlmRequest, LlmResponse, LlmStrategy
+
+_counter = itertools.count(1)
+
+
+class DelegatedLlmStrategy(LlmStrategy):
+    """Write prompt JSON files for external execution."""
+
+    def call(self, request: LlmRequest) -> LlmResponse:
+        seq = next(_counter)
+        prompt_data = {
+            "id": f"prompt_{seq:04d}",
+            "system": request.system,
+            "user": request.user,
+            "model": request.model or "",
+            "max_tokens": request.max_tokens,
+            "temperature": request.temperature,
+        }
+
+        return LlmResponse(
+            content="",
+            status="delegated",
+            prompt_file=json.dumps(prompt_data, ensure_ascii=False),
+        )
+
+    def call_batch(self, requests_list: list[LlmRequest], max_concurrent: int = 5) -> list[LlmResponse]:
+        """Delegated mode: generate all prompts sequentially (no concurrency needed)."""
+        return [self.call(r) for r in requests_list]
+
+
+def reset_counter() -> None:
+    """Reset the sequence counter (for testing)."""
+    global _counter
+    _counter = itertools.count(1)

engine/strategies/dryrun.py

@@ -0,0 +1,25 @@
+"""Dry-run LLM strategy — returns placeholder responses without network calls.
+
+Used for testing pipeline logic, validating prompt construction, and
+estimating token usage without incurring API costs.
+"""
+
+from __future__ import annotations
+
+from core.interfaces import LlmRequest, LlmResponse, LlmStrategy
+
+
+class DryRunLlmStrategy(LlmStrategy):
+    """Placeholder strategy that returns synthetic responses."""
+
+    def call(self, request: LlmRequest) -> LlmResponse:
+        model = request.model or "dry-run"
+        return LlmResponse(
+            content=(
+                f"[dry-run] model={model} "
+                f"system_len={len(request.system)} "
+                f"user_len={len(request.user)}"
+            ),
+            status="dry-run",
+            usage={"prompt_tokens": 0, "completion_tokens": 0},
+        )