code-muse 0.0.1__py3-none-any.whl

code_muse/round_robin_model.py

@@ -0,0 +1,151 @@
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.models import (
+    Model,
+    ModelMessage,
+    ModelRequestParameters,
+    ModelResponse,
+    ModelSettings,
+    StreamedResponse,
+)
+
+try:
+    from opentelemetry.context import get_current_span
+except ImportError:
+    # If opentelemetry is not installed, provide a dummy implementation
+    def get_current_span():
+        class DummySpan:
+            def is_recording(self):
+                return False
+
+            def set_attributes(self, attributes):
+                pass
+
+        return DummySpan()
+
+
+@dataclass(init=False)
+class RoundRobinModel(Model):
+    """A model that cycles through multiple models in a round-robin fashion.
+
+    This model distributes requests across multiple candidate models to help
+    overcome rate limits or distribute load.
+    """
+
+    models: list[Model]
+    _current_index: int = field(default=0, repr=False)
+    _model_name: str = field(repr=False)
+    _rotate_every: int = field(default=1, repr=False)
+    _request_count: int = field(default=0, repr=False)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
+
+    def __init__(
+        self,
+        *models: Model,
+        rotate_every: int = 1,
+        settings: ModelSettings | None = None,
+    ):
+        """Initialize a round-robin model instance.
+
+        Args:
+            models: The model instances to cycle through.
+            rotate_every: Number of requests before rotating to the next model (default: 1).
+            settings: Model settings that will be used as defaults for this model.
+        """
+        super().__init__(settings=settings)
+        if not models:
+            raise ValueError("At least one model must be provided")
+        if rotate_every < 1:
+            raise ValueError("rotate_every must be at least 1")
+        self.models = list(models)
+        self._current_index = 0
+        self._request_count = 0
+        self._rotate_every = rotate_every
+        self._lock = asyncio.Lock()
+
+    @property
+    def model_name(self) -> str:
+        """The model name showing this is a round-robin model with its candidates."""
+        base_name = f"round_robin:{','.join(model.model_name for model in self.models)}"
+        if self._rotate_every != 1:
+            return f"{base_name}:rotate_every={self._rotate_every}"
+        return base_name
+
+    @property
+    def system(self) -> str:
+        """System prompt from the current model."""
+        return self.models[self._current_index].system
+
+    @property
+    def base_url(self) -> str | None:
+        """Base URL from the current model."""
+        return self.models[self._current_index].base_url
+
+    async def _get_next_model(self) -> Model:
+        """Get the next model in the round-robin sequence and update the index."""
+        async with self._lock:
+            model = self.models[self._current_index]
+            self._request_count += 1
+            if self._request_count >= self._rotate_every:
+                self._current_index = (self._current_index + 1) % len(self.models)
+                self._request_count = 0
+            return model
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        """Make a request using the next model in the round-robin sequence."""
+        current_model = await self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        try:
+            response = await current_model.request(
+                messages, merged_settings, prepared_params
+            )
+            self._set_span_attributes(current_model)
+            return response
+        except Exception:
+            # Unlike FallbackModel, we don't try other models here
+            # The round-robin strategy is about distribution, not failover
+            raise
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        """Make a streaming request using the next model in the round-robin sequence."""
+        current_model = await self._get_next_model()
+        # Use prepare_request to merge settings and customize parameters
+        merged_settings, prepared_params = current_model.prepare_request(
+            model_settings, model_request_parameters
+        )
+
+        async with current_model.request_stream(
+            messages, merged_settings, prepared_params, run_context
+        ) as response:
+            self._set_span_attributes(current_model)
+            yield response
+
+    def _set_span_attributes(self, model: Model):
+        """Set span attributes for observability."""
+        with suppress(Exception):
+            span = get_current_span()
+            if span.is_recording():
+                attributes = getattr(span, "attributes", {})
+                if attributes.get("gen_ai.request.model") == self.model_name:
+                    span.set_attributes(model.model_attributes(model))

code_muse/secret_storage.py

@@ -0,0 +1,74 @@
+"""Atomic private file storage helpers for sensitive data.
+
+Implements fail-closed helpers that create files with restrictive permissions
+from the start, fsync before rename, and warn when existing files are too
+broadly permissioned.
+"""
+
+import json
+import logging
+import os
+import stat
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_private_dir(path: Path) -> Path:
+    """Create *path* (and parents) with mode ``0o700``.
+
+    If the directory already exists its mode is corrected if necessary.
+    """
+    path.mkdir(parents=True, exist_ok=True)
+    os.chmod(path, 0o700)
+    return path
+
+
+def atomic_write_private_bytes(path: Path, data: bytes) -> None:
+    """Atomically write *data* to *path* with mode ``0o600``.
+
+    Uses a temporary sibling file created with ``O_CREAT | O_EXCL`` so there
+    is no race window where another observer can open the file before the
+    mode is set.  After writing, ``fsync`` flushes data to disk and
+    ``os.replace`` atomically moves the temp file into place.
+    """
+    tmp_path = Path(str(path) + ".tmp")
+    fd: int | None = None
+    try:
+        fd = os.open(
+            str(tmp_path),
+            os.O_CREAT | os.O_EXCL | os.O_WRONLY,
+            0o600,
+        )
+        os.write(fd, data)
+        os.fsync(fd)
+    finally:
+        if fd is not None:
+            os.close(fd)
+    os.replace(str(tmp_path), path)
+
+
+def atomic_write_private_json(path: Path, data: dict[str, Any]) -> None:
+    """Atomically write a JSON object to *path* with mode ``0o600``."""
+    payload = json.dumps(data, indent=2).encode("utf-8")
+    atomic_write_private_bytes(path, payload)
+
+
+def warn_or_fix_private_file_mode(path: Path, expected_mode: int = 0o600) -> None:
+    """Check *path* permissions and fix them if they are too broad.
+
+    Emits a warning when the mode differs from *expected_mode* so operators
+    know the file was created by an older version or another process.
+    """
+    if not path.exists():
+        return
+    current_mode = stat.S_IMODE(path.stat().st_mode)
+    if current_mode != expected_mode:
+        logger.warning(
+            "File %s has mode %04o; fixing to %04o",
+            path,
+            current_mode,
+            expected_mode,
+        )
+        path.chmod(expected_mode)

code_muse/security/__init__.py

@@ -0,0 +1 @@
+"""Security helpers for Muse."""

code_muse/security/redaction.pyx

@@ -0,0 +1,135 @@
+# cython: language_level=3
+"""Redaction helpers to prevent secret leakage in logs and output."""
+
+import json
+import re
+import urllib.parse
+from typing import Any
+
+# Keys whose values should always be redacted in structured data.
+SENSITIVE_KEYS: frozenset = frozenset(
+    {
+        "access_token",
+        "refresh_token",
+        "id_token",
+        "api_key",
+        "code",
+        "code_verifier",
+        "code_challenge",
+        "client_secret",
+        "client_id",
+        "password",
+        "token",
+        "secret",
+        "authorization",
+        "bearer",
+        "apikey",
+        "auth_token",
+        "session_token",
+        "csrf_token",
+    }
+)
+
+REDACTED = "<redacted>"
+
+# Pre-compiled regex patterns for performance.
+cdef object _BEARER_HEADER_RE = re.compile(
+    r"(?i)(authorization\s*:\s*bearer\s+)\S+"
+)
+cdef object _BEARER_STANDALONE_RE = re.compile(r"(?i)(bearer\s+)\S+")
+cdef object _ENV_ASSIGNMENT_RE = re.compile(
+    r"(?i)([A-Z_]*(?:API_KEY|SECRET|TOKEN|PASSWORD|AUTH|CREDENTIALS)[A-Z_]*=)(.+?)(?=[\s&]+[A-Z_][A-Z0-9_]+=|$)"
+)
+
+
+cdef bint _is_sensitive_key(str key):
+    return key.lower() in SENSITIVE_KEYS
+
+
+cpdef str _redact_url_query_params(str value):
+    cdef object parsed
+    cdef dict qs
+    cdef bint redacted = False
+    cdef str k
+    try:
+        parsed = urllib.parse.urlparse(value)
+        if not parsed.query:
+            return value
+        qs = urllib.parse.parse_qs(parsed.query)
+        for k in list(qs):
+            if _is_sensitive_key(k):
+                qs[k] = [REDACTED]
+                redacted = True
+        if redacted:
+            new_query = urllib.parse.urlencode(qs, doseq=True)
+            return urllib.parse.urlunparse(parsed._replace(query=new_query))
+    except Exception:
+        pass
+    return value
+
+
+cpdef str _redact_bearer_tokens(str value):
+    value = _BEARER_HEADER_RE.sub(r"\1" + REDACTED, value)
+    value = _BEARER_STANDALONE_RE.sub(r"\1" + REDACTED, value)
+    return value
+
+
+cpdef str _redact_env_assignments(str value):
+    return _ENV_ASSIGNMENT_RE.sub(r"\1" + REDACTED, value)
+
+
+cpdef str _redact_json_string(str value):
+    cdef str stripped = value.strip()
+    if not stripped.startswith(("{", "[")):
+        return value
+    try:
+        parsed = json.loads(stripped)
+        redacted = redact_secrets(parsed)
+        return json.dumps(redacted, separators=(",", ":"))
+    except Exception:
+        return value
+
+
+def redact_secrets(value: Any, _parent_key: str = "") -> Any:
+    """Recursively redact secrets from strings, dicts, and lists.
+
+    Covers:
+    - URL query parameters containing sensitive keys
+    - JSON/dict keys matching known sensitive names
+    - ``Authorization: Bearer ...`` headers
+    - OAuth response bodies (via recursive dict redaction)
+    - Environment-variable-style assignments of sensitive names
+
+    Returns deterministic ``<redacted>`` output; never logs token length.
+    """
+    cdef str s
+    cdef dict d
+    cdef list lst
+    cdef object k, v
+
+    if isinstance(value, bytes):
+        try:
+            return redact_secrets(value.decode("utf-8"))
+        except UnicodeDecodeError:
+            return REDACTED.encode("utf-8")
+    if isinstance(value, str):
+        s = value
+        s = _redact_bearer_tokens(s)
+        s = _redact_url_query_params(s)
+        s = _redact_env_assignments(s)
+        s = _redact_json_string(s)
+        return s
+    if isinstance(value, dict):
+        d = {}
+        for k, v in value.items():
+            if _is_sensitive_key(k):
+                d[k] = REDACTED
+            else:
+                d[k] = redact_secrets(v, k)
+        return d
+    if isinstance(value, list):
+        lst = []
+        for v in value:
+            lst.append(redact_secrets(v, _parent_key))
+        return lst
+    return value