agent-cli 0.70.5__py3-none-any.whl

agent_cli/core/openai_proxy.py

@@ -0,0 +1,139 @@
+"""Shared OpenAI-compatible forwarding helpers."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+
+    from fastapi import Request, Response
+
+LOGGER = logging.getLogger(__name__)
+
+
+@runtime_checkable
+class ChatRequestLike(Protocol):
+    """Minimal interface required to forward a chat request."""
+
+    stream: bool | None
+
+    def model_dump(self, *, exclude: set[str] | None = None) -> dict[str, Any]:
+        """Serialize request to a dict for forwarding."""
+
+
+async def proxy_request_to_upstream(
+    request: Request,
+    path: str,
+    upstream_base_url: str,
+    api_key: str | None = None,
+) -> Response:
+    """Forward a raw HTTP request to an upstream OpenAI-compatible provider."""
+    import httpx  # noqa: PLC0415
+    from fastapi import Response  # noqa: PLC0415
+
+    auth_header = request.headers.get("Authorization")
+    headers = {}
+    if auth_header:
+        headers["Authorization"] = auth_header
+    elif api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    if request.headers.get("Content-Type"):
+        headers["Content-Type"] = request.headers.get("Content-Type")
+
+    base = upstream_base_url.rstrip("/")
+    target_path = path
+
+    # Smart path joining to avoid /v1/v1/ if base already has it
+    if base.endswith("/v1") and (path == "v1" or path.startswith("v1/")):
+        target_path = path[2:].lstrip("/")
+
+    url = f"{base}/{target_path}"
+
+    try:
+        body = await request.body()
+        async with httpx.AsyncClient(timeout=60.0) as http:
+            req = http.build_request(
+                request.method,
+                url,
+                headers=headers,
+                content=body,
+                params=request.query_params,
+            )
+            resp = await http.send(req)
+
+            return Response(
+                content=resp.content,
+                status_code=resp.status_code,
+                media_type=resp.headers.get("Content-Type"),
+            )
+    except Exception:
+        LOGGER.warning("Proxy request failed to %s", url, exc_info=True)
+        return Response(status_code=502, content="Upstream Proxy Error")
+
+
+async def forward_chat_request(
+    request: ChatRequestLike,
+    openai_base_url: str,
+    api_key: str | None = None,
+    *,
+    exclude_fields: Iterable[str] = (),
+) -> Any:
+    """Forward a chat request to a backend LLM."""
+    import httpx  # noqa: PLC0415
+    from fastapi import HTTPException  # noqa: PLC0415
+    from fastapi.responses import StreamingResponse  # noqa: PLC0415
+
+    forward_payload = request.model_dump(exclude=set(exclude_fields))
+    headers = {"Authorization": f"Bearer {api_key}"} if api_key else None
+
+    if getattr(request, "stream", False):
+
+        async def generate() -> AsyncGenerator[str, None]:
+            try:
+                async with (
+                    httpx.AsyncClient(timeout=120.0) as client,
+                    client.stream(
+                        "POST",
+                        f"{openai_base_url.rstrip('/')}/chat/completions",
+                        json=forward_payload,
+                        headers=headers,
+                    ) as response,
+                ):
+                    if response.status_code != 200:  # noqa: PLR2004
+                        error_text = await response.aread()
+                        yield f"data: {json.dumps({'error': str(error_text)})}\n\n"
+                        return
+
+                    async for chunk in response.aiter_raw():
+                        if isinstance(chunk, bytes):
+                            yield chunk.decode("utf-8")
+                        else:
+                            yield chunk
+            except Exception as exc:
+                LOGGER.exception("Streaming error")
+                yield f"data: {json.dumps({'error': str(exc)})}\n\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+
+    async with httpx.AsyncClient(timeout=120.0) as client:
+        response = await client.post(
+            f"{openai_base_url.rstrip('/')}/chat/completions",
+            json=forward_payload,
+            headers=headers,
+        )
+        if response.status_code != 200:  # noqa: PLR2004
+            LOGGER.error(
+                "Upstream error %s: %s",
+                response.status_code,
+                response.text,
+            )
+            raise HTTPException(
+                status_code=response.status_code,
+                detail=f"Upstream error: {response.text}",
+            )
+
+        return response.json()

agent_cli/core/process.py

@@ -0,0 +1,195 @@
+"""Process management utilities for Agent CLI tools."""
+
+from __future__ import annotations
+
+import os
+import signal
+import sys
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+# Default location for PID files
+PID_DIR = Path.home() / ".cache" / "agent-cli"
+
+# Store the original process title before any modifications
+_original_proctitle: str | None = None
+
+
+def set_process_title(process_name: str) -> None:
+    """Set the process title and thread name for identification in ps/htop/btop.
+
+    Sets both:
+    - Process title: 'agent-cli-{name} ({original})' - identifiable prefix + original command
+    - Thread name: 'ag-{name}' (max 15 chars) - shown as program name in btop/htop
+
+    The original command line is captured on first call and reused on subsequent
+    calls to prevent nested titles like 'agent-cli-x (agent-cli-y (...))'.
+
+    Args:
+        process_name: The name of the process (e.g., 'transcribe', 'chat').
+
+    """
+    import setproctitle  # noqa: PLC0415
+
+    global _original_proctitle
+
+    # Capture the original command line only once, before any modification
+    if _original_proctitle is None:
+        _original_proctitle = setproctitle.getproctitle()
+
+    # Set the full process title: identifiable prefix + original command for debugging
+    setproctitle.setproctitle(f"agent-cli-{process_name} ({_original_proctitle})")
+
+    # Set the thread name (program name in htop/btop, limited to 15 chars on Linux)
+    # Use shorter prefix "ag-" to fit more of the command name
+    thread_name = f"ag-{process_name}"[:15]
+    setproctitle.setthreadtitle(thread_name)
+
+
+def _get_pid_file(process_name: str) -> Path:
+    """Get the path to the PID file for a given process name."""
+    PID_DIR.mkdir(parents=True, exist_ok=True)
+    return PID_DIR / f"{process_name}.pid"
+
+
+def _get_stop_file(process_name: str) -> Path:
+    """Get the path to the stop file for a given process name."""
+    PID_DIR.mkdir(parents=True, exist_ok=True)
+    return PID_DIR / f"{process_name}.stop"
+
+
+def check_stop_file(process_name: str) -> bool:
+    """Check if a stop file exists (used for cross-process signaling on Windows)."""
+    return _get_stop_file(process_name).exists()
+
+
+def clear_stop_file(process_name: str) -> None:
+    """Remove the stop file for the given process."""
+    stop_file = _get_stop_file(process_name)
+    if stop_file.exists():
+        stop_file.unlink()
+
+
+def _is_pid_running(pid: int) -> bool:
+    """Check if a process with the given PID is running."""
+    if sys.platform == "win32":
+        # On Windows, os.kill(pid, 0) would terminate the process!
+        import psutil  # noqa: PLC0415
+
+        return psutil.pid_exists(pid)
+    try:
+        os.kill(pid, 0)
+        return True
+    except (ProcessLookupError, PermissionError):
+        return False
+
+
+def _get_running_pid(process_name: str) -> int | None:
+    """Get PID if process is running, None otherwise. Cleans up stale files."""
+    pid_file = _get_pid_file(process_name)
+
+    if not pid_file.exists():
+        return None
+
+    try:
+        with pid_file.open() as f:
+            pid = int(f.read().strip())
+
+        # Check if process is actually running
+        if _is_pid_running(pid):
+            return pid
+
+    except (FileNotFoundError, ValueError):
+        pass
+
+    # Clean up stale/invalid PID file
+    if pid_file.exists():
+        pid_file.unlink()
+    return None
+
+
+def is_process_running(process_name: str) -> bool:
+    """Check if a process is currently running."""
+    return _get_running_pid(process_name) is not None
+
+
+def read_pid_file(process_name: str) -> int | None:
+    """Read PID from file if process is running."""
+    return _get_running_pid(process_name)
+
+
+def kill_process(process_name: str) -> bool:
+    """Kill a process by name.
+
+    Returns True if killed or cleaned up, False if not found.
+    On Windows, creates a stop file first to allow graceful shutdown.
+    """
+    pid_file = _get_pid_file(process_name)
+
+    # If no PID file exists at all, nothing to do
+    if not pid_file.exists():
+        return False
+
+    # Check if we have a running process
+    pid = _get_running_pid(process_name)
+
+    # If _get_running_pid returned None but file existed, it cleaned up a stale file
+    if pid is None:
+        return True
+
+    # On Windows, create stop file to signal graceful shutdown
+    if sys.platform == "win32":
+        _get_stop_file(process_name).touch()
+
+    # Send SIGINT for graceful shutdown
+    try:
+        os.kill(pid, signal.SIGINT)
+        # Wait for process to terminate
+        for _ in range(10):  # 1 second max
+            if not is_process_running(process_name):
+                break
+            time.sleep(0.1)
+    except (ProcessLookupError, PermissionError):
+        pass  # Process dead or no permission - we'll clean up regardless
+
+    # Clean up
+    if sys.platform == "win32":
+        clear_stop_file(process_name)
+    if pid_file.exists():
+        pid_file.unlink()
+
+    return True
+
+
+@contextmanager
+def pid_file_context(process_name: str) -> Generator[Path, None, None]:
+    """Context manager for PID file lifecycle.
+
+    Creates PID file on entry, cleans up on exit.
+    Exits with error if process already running.
+    """
+    if is_process_running(process_name):
+        existing_pid = _get_running_pid(process_name)
+        print(f"Process {process_name} is already running (PID: {existing_pid})")
+        sys.exit(1)
+
+    # Clear any stale stop file from previous run (Windows only)
+    if sys.platform == "win32":
+        clear_stop_file(process_name)
+
+    pid_file = _get_pid_file(process_name)
+    with pid_file.open("w") as f:
+        f.write(str(os.getpid()))
+
+    try:
+        yield pid_file
+    finally:
+        if pid_file.exists():
+            pid_file.unlink()
+        if sys.platform == "win32":
+            clear_stop_file(process_name)

agent_cli/core/reranker.py

@@ -0,0 +1,120 @@
+"""Shared ONNX Cross-Encoder for reranking (used by both RAG and Memory)."""
+
+from __future__ import annotations
+
+import logging
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _download_onnx_model(model_name: str, onnx_filename: str) -> str:
+    """Download the ONNX model, favoring the common `onnx/` folder layout."""
+    from huggingface_hub import hf_hub_download  # noqa: PLC0415
+
+    if "/" in onnx_filename:
+        return hf_hub_download(repo_id=model_name, filename=onnx_filename)
+
+    try:
+        return hf_hub_download(repo_id=model_name, filename=onnx_filename, subfolder="onnx")
+    except Exception as first_error:
+        LOGGER.debug(
+            "ONNX file not found under onnx/ for %s: %s. Falling back to repo root.",
+            model_name,
+            first_error,
+        )
+        try:
+            return hf_hub_download(repo_id=model_name, filename=onnx_filename)
+        except Exception as second_error:
+            LOGGER.exception(
+                "Failed to download ONNX model %s (filename=%s)",
+                model_name,
+                onnx_filename,
+                exc_info=second_error,
+            )
+            raise
+
+
+class OnnxCrossEncoder:
+    """A lightweight CrossEncoder using ONNX Runtime."""
+
+    def __init__(
+        self,
+        model_name: str = "Xenova/ms-marco-MiniLM-L-6-v2",
+        onnx_filename: str = "model.onnx",
+    ) -> None:
+        """Initialize the ONNX CrossEncoder."""
+        from onnxruntime import InferenceSession  # noqa: PLC0415
+        from transformers import AutoTokenizer  # noqa: PLC0415
+
+        self.model_name = model_name
+
+        # Download model if needed
+        LOGGER.info("Loading ONNX model: %s", model_name)
+        model_path = _download_onnx_model(model_name, onnx_filename)
+
+        self.session = InferenceSession(model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    def predict(
+        self,
+        pairs: list[tuple[str, str]],
+        batch_size: int = 32,
+    ) -> list[float]:
+        """Predict relevance scores for query-document pairs."""
+        import numpy as np  # noqa: PLC0415
+
+        if not pairs:
+            return []
+
+        all_scores = []
+
+        # Process in batches
+        for i in range(0, len(pairs), batch_size):
+            batch = pairs[i : i + batch_size]
+            queries = [q for q, d in batch]
+            docs = [d for q, d in batch]
+
+            # Tokenize
+            inputs = self.tokenizer(
+                queries,
+                docs,
+                padding=True,
+                truncation=True,
+                return_tensors="np",
+                max_length=512,
+            )
+
+            # ONNX Input
+            # Check what inputs the model expects. usually input_ids, attention_mask, token_type_ids
+            # specific models might not need token_type_ids
+            ort_inputs = {
+                "input_ids": inputs["input_ids"].astype(np.int64),
+                "attention_mask": inputs["attention_mask"].astype(np.int64),
+            }
+            if "token_type_ids" in inputs:
+                ort_inputs["token_type_ids"] = inputs["token_type_ids"].astype(np.int64)
+
+            # Run inference
+            logits = self.session.run(None, ort_inputs)[0]
+
+            # Extract scores (usually shape [batch, 1] or [batch])
+            batch_scores = logits.flatten() if logits.ndim > 1 else logits
+
+            all_scores.extend(batch_scores.tolist())
+
+        return all_scores
+
+
+def get_reranker_model(
+    model_name: str = "Xenova/ms-marco-MiniLM-L-6-v2",
+) -> OnnxCrossEncoder:
+    """Load the CrossEncoder model."""
+    return OnnxCrossEncoder(model_name)
+
+
+def predict_relevance(
+    model: OnnxCrossEncoder,
+    pairs: list[tuple[str, str]],
+) -> list[float]:
+    """Predict relevance scores for query-document pairs."""
+    return model.predict(pairs)

agent_cli/core/sse.py

@@ -0,0 +1,87 @@
+"""Shared SSE (Server-Sent Events) formatting helpers for OpenAI-compatible streaming."""
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any
+
+
+def format_chunk(
+    run_id: str,
+    model: str,
+    *,
+    content: str | None = None,
+    finish_reason: str | None = None,
+    extra: dict[str, Any] | None = None,
+) -> str:
+    """Format a single SSE chunk in OpenAI chat.completion.chunk format.
+
+    Args:
+        run_id: Unique identifier for this completion.
+        model: Model name to include in response.
+        content: Text content delta (None for finish chunk).
+        finish_reason: Reason for completion (e.g., "stop").
+        extra: Additional fields to include in the response.
+
+    Returns:
+        Formatted SSE data line.
+
+    """
+    data: dict[str, Any] = {
+        "id": f"chatcmpl-{run_id}",
+        "object": "chat.completion.chunk",
+        "created": int(time.time()),
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {"content": content} if content else {},
+                "finish_reason": finish_reason,
+            },
+        ],
+    }
+    if extra:
+        data.update(extra)
+    return f"data: {json.dumps(data)}\n\n"
+
+
+def format_done() -> str:
+    """Format the terminal [DONE] SSE message."""
+    return "data: [DONE]\n\n"
+
+
+def parse_chunk(line: str) -> dict[str, Any] | None:
+    """Parse an SSE data line into a dict.
+
+    Args:
+        line: Raw SSE line (e.g., "data: {...}").
+
+    Returns:
+        Parsed JSON dict, or None if not parseable or [DONE].
+
+    """
+    if not line.startswith("data:"):
+        return None
+    payload = line[5:].strip()
+    if payload == "[DONE]":
+        return None
+    try:
+        return json.loads(payload)
+    except json.JSONDecodeError:
+        return None
+
+
+def extract_content_from_chunk(chunk: dict[str, Any]) -> str:
+    """Extract text content from a parsed SSE chunk.
+
+    Args:
+        chunk: Parsed chunk dict from parse_chunk().
+
+    Returns:
+        Content string, or empty string if not found.
+
+    """
+    choices = chunk.get("choices") or [{}]
+    delta = choices[0].get("delta") or {}
+    return delta.get("content") or delta.get("text") or ""

agent_cli/core/transcription_logger.py

@@ -0,0 +1,70 @@
+"""Transcription logging utilities for automatic server-side logging."""
+
+from __future__ import annotations
+
+import json
+import logging
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+
+class TranscriptionLogger:
+    """Handles automatic logging of transcription results with timestamps."""
+
+    def __init__(self, log_file: Path | str | None = None) -> None:
+        """Initialize the transcription logger.
+
+        Args:
+            log_file: Path to the log file. If None, uses default location.
+
+        """
+        if log_file is None:
+            log_file = Path.home() / ".config" / "agent-cli" / "transcriptions.jsonl"
+        elif isinstance(log_file, str):
+            log_file = Path(log_file)
+
+        self.log_file = log_file
+
+        # Ensure the log directory exists
+        self.log_file.parent.mkdir(parents=True, exist_ok=True)
+
+    def log_transcription(
+        self,
+        *,
+        raw: str,
+        processed: str | None = None,
+    ) -> None:
+        """Log a transcription result.
+
+        Args:
+            raw: The raw transcript from ASR.
+            processed: The processed transcript from LLM.
+
+        """
+        log_entry: dict[str, Any] = {
+            "timestamp": datetime.now(UTC).isoformat(),
+            "raw": raw,
+            "processed": processed,
+        }
+
+        # Write to log file as JSON Lines format
+        try:
+            with self.log_file.open("a", encoding="utf-8") as f:
+                f.write(json.dumps(log_entry, ensure_ascii=False) + "\n")
+        except OSError:
+            # Use Python's logging module to log errors with the logger itself
+            logger = logging.getLogger(__name__)
+            logger.exception("Failed to write transcription log")
+
+
+# Default logger instance
+_default_logger: TranscriptionLogger | None = None
+
+
+def get_default_logger() -> TranscriptionLogger:
+    """Get the default transcription logger instance."""
+    global _default_logger
+    if _default_logger is None:
+        _default_logger = TranscriptionLogger()
+    return _default_logger