claude-smart 0.2.41 → 0.2.43

package/plugin/vendor/reflexio/reflexio/server/services/braintrust/service.py

@@ -0,0 +1,281 @@
+"""Orchestrates the Braintrust connector: connect, select-projects, sync, status.
+
+The service is stateless across requests; it loads the connection from
+storage on each call. API keys are encrypted in storage and decrypted
+only when needed for a Braintrust HTTP call.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from reflexio.models.api_schema.braintrust_schema import (
+    BraintrustConnection,
+    BraintrustProjectSummary,
+    BraintrustStatusResponse,
+    BraintrustWorkspaceSummary,
+    ConnectBraintrustRequest,
+    ConnectBraintrustResponse,
+    ImportedScore,
+    SelectProjectsRequest,
+    SelectProjectsResponse,
+    SyncBraintrustResponse,
+)
+from reflexio.server.services.braintrust._encryption import decrypt, encrypt
+from reflexio.server.services.braintrust.client import (
+    DEFAULT_BASE_URL,
+    BraintrustAuthError,
+    BraintrustClient,
+    BraintrustHTTPError,
+)
+
+
+def _default_client_factory(api_key: str) -> BraintrustClient:
+    """Default factory that honors `BRAINTRUST_BASE_URL` env override.
+
+    Set the env var to point at a staging or mock Braintrust instance
+    (useful for dev / cron testing without hitting the real API).
+    """
+    import os
+
+    base_url = os.environ.get("BRAINTRUST_BASE_URL", "").strip() or DEFAULT_BASE_URL
+    return BraintrustClient(api_key, base_url=base_url)
+
+
+logger = logging.getLogger(__name__)
+
+
+_METADATA_SESSION_KEY = "reflexio_session_id"
+
+
+@dataclass
+class BraintrustConnectorService:
+    """Top-level orchestrator for the Braintrust connector.
+
+    Args:
+        storage: BaseStorage-like; uses save/get/delete_braintrust_connection
+            and save_imported_scores. Default no-op storage methods keep
+            this workable until per-backend implementations land.
+        org_id (str): The Reflexio org for which the connector operates.
+        client_factory (Callable[[str], BraintrustClient]): Factory that
+            takes an API key and returns a BraintrustClient. Injectable so
+            tests can pass mocks.
+    """
+
+    storage: object
+    org_id: str
+    client_factory: Callable[[str], BraintrustClient] = field(
+        default=_default_client_factory
+    )
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def connect(self, request: ConnectBraintrustRequest) -> ConnectBraintrustResponse:
+        """Step 1: validate the API key and return the workspace/project tree.
+
+        Persists nothing — the caller chooses which projects to sync and
+        then calls `select_projects` with the same key.
+        """
+        client = self.client_factory(request.api_key)
+        try:
+            if not client.validate_key():
+                return ConnectBraintrustResponse(
+                    success=False, msg="Braintrust rejected the API key."
+                )
+            workspaces = self._fetch_workspace_tree(client)
+        except BraintrustAuthError:
+            return ConnectBraintrustResponse(
+                success=False, msg="Braintrust rejected the API key."
+            )
+        except BraintrustHTTPError as e:
+            return ConnectBraintrustResponse(
+                success=False, msg=f"Braintrust HTTP error: {e}"
+            )
+        finally:
+            client.close()
+
+        return ConnectBraintrustResponse(success=True, workspaces=workspaces, msg="")
+
+    def select_projects(self, request: SelectProjectsRequest) -> SelectProjectsResponse:
+        """Step 2: persist the connection with the selected projects.
+
+        Stores the API key encrypted. Overwrites any existing connection
+        for the org.
+        """
+        connection = BraintrustConnection(
+            org_id=self.org_id,
+            api_key_enc=encrypt(request.api_key),
+            workspace_id=request.workspace_id,
+            workspace_name=request.workspace_name,
+            project_ids=request.project_ids,
+            last_sync_ts=None,
+            last_error=None,
+        )
+        self._save_connection(connection)
+        return SelectProjectsResponse(success=True, msg="Connected.")
+
+    def status(self) -> BraintrustStatusResponse:
+        """Return whether the org is connected + sync state. Never echoes the key."""
+        connection = self._load_connection()
+        if connection is None:
+            return BraintrustStatusResponse(connected=False)
+        return BraintrustStatusResponse(
+            connected=True,
+            workspace_id=connection.workspace_id,
+            workspace_name=connection.workspace_name,
+            project_count=len(connection.project_ids),
+            last_sync_ts=connection.last_sync_ts,
+            last_error=connection.last_error,
+        )
+
+    def disconnect(self) -> None:
+        """Delete the persisted connection."""
+        self.storage.delete_braintrust_connection(self.org_id)  # type: ignore[attr-defined]
+
+    def sync_once(self, backfill_days: int = 90) -> SyncBraintrustResponse:
+        """One-shot sync: walk projects → experiments → spans → write scores.
+
+        Updates `last_sync_ts` and `last_error` on the persisted connection.
+        Idempotent at the storage layer (per-backend overrides should
+        upsert by (source, source_run_id, scorer_name)).
+        """
+        connection = self._load_connection()
+        if connection is None:
+            return SyncBraintrustResponse(
+                success=False, msg="Not connected to Braintrust."
+            )
+
+        try:
+            api_key = decrypt(connection.api_key_enc)
+        except Exception as e:  # noqa: BLE001
+            return SyncBraintrustResponse(
+                success=False, msg=f"Failed to decrypt API key: {e}"
+            )
+
+        since_ts = connection.last_sync_ts
+        if since_ts is None:
+            since_ts = max(0, int(time.time()) - backfill_days * 24 * 60 * 60)
+
+        all_scores: list[ImportedScore] = []
+        client = self.client_factory(api_key)
+        try:
+            for project_id in connection.project_ids:
+                experiments = client.list_experiments(project_id, since_ts=since_ts)
+                for exp in experiments:
+                    spans = client.list_spans(exp["id"])
+                    all_scores.extend(_scores_from_spans(spans, self.org_id))
+        except BraintrustAuthError:
+            self._persist_sync_outcome(connection, error="API key invalid.")
+            return SyncBraintrustResponse(
+                success=False, msg="API key invalid; halting sync."
+            )
+        except BraintrustHTTPError as e:
+            self._persist_sync_outcome(connection, error=str(e))
+            return SyncBraintrustResponse(success=False, msg=str(e))
+        finally:
+            client.close()
+
+        self.storage.save_imported_scores(all_scores)  # type: ignore[attr-defined]
+        self._persist_sync_outcome(connection, error=None)
+        return SyncBraintrustResponse(
+            success=True, scored_count=len(all_scores), msg=""
+        )
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    def _fetch_workspace_tree(
+        self, client: BraintrustClient
+    ) -> list[BraintrustWorkspaceSummary]:
+        """Pull workspaces + their projects in one shot for the connect UX."""
+        result: list[BraintrustWorkspaceSummary] = []
+        for ws in client.list_organizations():
+            workspace_id = str(ws.get("id", ""))
+            workspace_name = str(ws.get("name", ""))
+            if not workspace_id:
+                continue
+            projects_raw = client.list_projects(workspace_id)
+            projects = [
+                BraintrustProjectSummary(
+                    project_id=str(p.get("id", "")),
+                    project_name=str(p.get("name", "")),
+                )
+                for p in projects_raw
+                if p.get("id")
+            ]
+            result.append(
+                BraintrustWorkspaceSummary(
+                    workspace_id=workspace_id,
+                    workspace_name=workspace_name,
+                    projects=projects,
+                )
+            )
+        return result
+
+    def _load_connection(self) -> BraintrustConnection | None:
+        return self.storage.get_braintrust_connection(self.org_id)  # type: ignore[attr-defined]
+
+    def _save_connection(self, connection: BraintrustConnection) -> None:
+        self.storage.save_braintrust_connection(connection)  # type: ignore[attr-defined]
+
+    def _persist_sync_outcome(
+        self, connection: BraintrustConnection, *, error: str | None
+    ) -> None:
+        update: dict[str, int | str | None] = {"last_error": error}
+        if error is None:
+            update["last_sync_ts"] = int(time.time())
+        updated = connection.model_copy(update=update)
+        self._save_connection(updated)
+
+
+def _scores_from_spans(
+    spans: list[dict],
+    org_id: str,
+) -> list[ImportedScore]:
+    """Extract ImportedScore rows from a list of Braintrust spans.
+
+    A single span may carry multiple scorers — one row per scorer.
+    `session_id` is pulled from `span.metadata.reflexio_session_id` when
+    present (the progressive-matching upgrade path).
+    """
+    out: list[ImportedScore] = []
+    for span in spans:
+        span_id = str(span.get("id", ""))
+        if not span_id:
+            continue
+        ts_raw = span.get("created_at", 0)
+        try:
+            ts = int(ts_raw)
+        except (TypeError, ValueError):
+            ts = 0
+        metadata = span.get("metadata") or {}
+        session_id: str | None = None
+        if isinstance(metadata, dict):
+            v = metadata.get(_METADATA_SESSION_KEY)
+            if isinstance(v, str) and v:
+                session_id = v
+        scores = span.get("scores") or {}
+        if not isinstance(scores, dict):
+            continue
+        for scorer_name, value in scores.items():
+            try:
+                fvalue = float(value)
+            except (TypeError, ValueError):
+                continue
+            out.append(
+                ImportedScore(
+                    org_id=org_id,
+                    source_run_id=span_id,
+                    session_id=session_id,
+                    scorer_name=str(scorer_name),
+                    value=fvalue,
+                    ts=ts,
+                )
+            )
+    return out

package/plugin/vendor/reflexio/reflexio/server/services/configurator/base_configurator.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+from pydantic import BaseModel
+
+from reflexio.models.config_schema import Config, StorageConfig, StorageConfigTest
+from reflexio.server.services.configurator.config_storage import ConfigStorage
+from reflexio.server.services.storage.error import StorageError
+from reflexio.server.services.storage.storage_base import BaseStorage
+
+logger = logging.getLogger(__name__)
+
+_CONFIG_NAME_ALIASES = {
+    "batch_size": "window_size",
+    "batch_interval": "stride_size",
+    "extraction_window_size": "window_size",
+    "extraction_window_stride": "stride_size",
+}
+
+
+class BaseConfigurator(ABC):
+    """Abstract base for organization configurators.
+
+    Subclasses must:
+      - Set ``_STORAGE_FACTORIES`` and ``_STORAGE_READINESS_CHECKS`` class variables.
+      - Implement ``_select_config_storage`` to choose the config persistence backend.
+
+    All shared configuration and storage-dispatch logic lives here so that OS
+    and enterprise implementations only override what differs.
+    """
+
+    _STORAGE_FACTORIES: ClassVar[dict[type[StorageConfig], Callable[..., BaseStorage]]]
+    _STORAGE_READINESS_CHECKS: ClassVar[
+        dict[type[StorageConfig], Callable[[Any], bool]]
+    ]
+
+    def __init__(
+        self,
+        org_id: str,
+        base_dir: str | None = None,
+        config: Config | None = None,
+    ) -> None:
+        self.org_id = org_id
+        self.base_dir = base_dir
+
+        self.config_storage = self._select_config_storage(org_id, base_dir)
+        if not config:
+            self.config = self.config_storage.load_config()
+        else:
+            self.config = config
+
+        if not self.config:
+            raise ValueError(f"Failed to load configuration for organization {org_id}")
+
+    @abstractmethod
+    def _select_config_storage(
+        self, org_id: str, base_dir: str | None
+    ) -> ConfigStorage:
+        """Select and return the appropriate ConfigStorage backend.
+
+        This is the primary extension point: OS uses LocalJson only, while
+        enterprise uses a multi-tier priority system.
+        """
+
+    # ==========================
+    # Configuration
+    # ==========================
+
+    def get_config(self) -> Config:
+        return self.config
+
+    def get_config_for_response(self) -> dict[str, Any]:
+        """Return config serialized for API responses."""
+        return self.config.model_dump(mode="json")
+
+    def normalize_config_payload(self, config: dict[str, Any]) -> dict[str, Any]:
+        """Normalize raw API config payloads before Pydantic validation."""
+        return config
+
+    def get_agent_context(self) -> str:
+        context = self.get_config().agent_context_prompt
+        if not context:
+            return ""
+        return context.strip()
+
+    def set_config(self, config: Config) -> None:
+        self.config_storage.save_config(config=config)
+        self.config = config
+
+    def set_config_by_name(
+        self,
+        config_name: str,
+        config_value: str | int | float | bool | list | dict | BaseModel | None,
+    ) -> None:
+        original_config_name = config_name
+        config_name = _CONFIG_NAME_ALIASES.get(config_name, config_name)
+        if config_name not in type(self.config).model_fields:
+            raise ValueError(f"Invalid config name: {config_name}")
+
+        # Aliased plural fields (e.g. ``profile_extractor_configs``) map to
+        # singular Optional fields on the model. Collapse the list to its
+        # first element, or ``None`` if empty, before assignment.
+        resolved_value: str | int | float | bool | list | dict | BaseModel | None = (
+            config_value
+        )
+        if (
+            original_config_name in _CONFIG_NAME_ALIASES
+            and config_name
+            in {"profile_extractor_config", "user_playbook_extractor_config"}
+            and isinstance(config_value, list)
+        ):
+            resolved_value = config_value[0] if config_value else None
+
+        setattr(self.config, config_name, resolved_value)
+        self.set_config(config=self.config)
+
+    # ==========================
+    # Storage
+    # ==========================
+
+    def get_current_storage_configuration(self) -> StorageConfig:
+        """Return the currently configured storage config."""
+        return self.get_config().storage_config
+
+    def create_storage(self, storage_config: StorageConfig) -> BaseStorage | None:
+        """Create a storage from the ``_STORAGE_FACTORIES`` registry.
+
+        Returns None when *storage_config* is None (enterprise cloud mode
+        where storage is not yet configured).
+        """
+        if storage_config is None:
+            return None
+        factory = self._STORAGE_FACTORIES.get(type(storage_config))
+        if factory is None:
+            raise ValueError(
+                f"No storage factory registered for {type(storage_config).__name__}"
+            )
+        return factory(self, storage_config)
+
+    def is_storage_configured(self) -> bool:
+        """Check whether a valid, non-failed storage option is configured."""
+        if not self.is_storage_config_ready_to_test(
+            storage_config=self.get_current_storage_configuration(),
+        ):
+            return False
+        return self.get_config().storage_config_test != StorageConfigTest.FAILED
+
+    def is_storage_config_ready_to_test(self, storage_config: StorageConfig) -> bool:
+        """Check whether *storage_config* is fully filled in and ready for a test connection."""
+        check = self._STORAGE_READINESS_CHECKS.get(type(storage_config))
+        return check(storage_config) if check else False
+
+    def test_and_init_storage_config(
+        self, storage_config: StorageConfig
+    ) -> tuple[bool, str]:
+        """Test whether *storage_config* is valid and initialise the storage.
+
+        Returns:
+            tuple[bool, str]: (success, message)
+        """
+        if not self.is_storage_config_ready_to_test(storage_config=storage_config):
+            return False, "Storage configuration is not ready to test"
+
+        try:
+            storage = self.create_storage(storage_config=storage_config)
+            if storage is not None:
+                storage.migrate()
+                return True, "Storage initialized successfully"
+            return False, "Failed to create storage"
+        except StorageError as e:
+            logger.error("Storage initialization failed: %s", e.message)
+            return False, e.message
+        except Exception as e:
+            logger.error("Storage initialization failed: %s", e)
+            return False, str(e)

package/plugin/vendor/reflexio/reflexio/server/services/configurator/config_storage.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from reflexio.models.config_schema import Config
+
+
+class ConfigStorage(ABC):
+    """
+    Abstract base class for configuration storage operations.
+    Defines the interface for saving/loading configuration to/from persistent storage.
+    """
+
+    def __init__(self, org_id: str):
+        self.org_id: str = org_id
+
+    @abstractmethod
+    def get_default_config(self) -> Config:
+        """
+        Returns a default configuration that is uninitialized.
+
+        Returns:
+            Config: Default configuration object
+        """
+
+    @abstractmethod
+    def load_config(self) -> Config:
+        """
+        Loads the current configuration of the organization. If the organization does
+        not exist, or if no config exists, or if the current saved config is no longer valid,
+        this routine creates a default one but will not update the persistent storage.
+
+        Returns:
+            Config: Loaded configuration object
+        """
+
+    @abstractmethod
+    def save_config(self, config: Config) -> None:
+        """
+        Saves the configuration to the persistent storage.
+
+        Args:
+            config (Config): Configuration object to save
+        """
+
+    def get_version(self) -> tuple[str, Any] | None:
+        """Cheap probe of the persisted config version, used by the per-org cache.
+
+        Backends that can cheaply detect out-of-band mutations should
+        override this. The Reflexio cache stamps the returned tuple at
+        load time and re-probes on every cache hit; if the value
+        changes, the cached instance is evicted and rebuilt with fresh
+        configuration.
+
+        Returns:
+            tuple[str, Any] | None: A ``(kind, value)`` tuple where
+            ``kind`` is a short string identifying the probe type
+            (e.g. ``"file"`` for mtime, ``"db"`` for a row version).
+            Returns ``None`` when probing is unsupported or fails — the
+            cache treats ``None`` as "still fresh", deferring to the
+            TTL safety net.
+        """
+        return None

package/plugin/vendor/reflexio/reflexio/server/services/configurator/configurator.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from reflexio.models.config_schema import (
+    StorageConfig,
+    StorageConfigSQLite,
+)
+from reflexio.server.services.configurator.base_configurator import BaseConfigurator
+from reflexio.server.services.configurator.config_storage import ConfigStorage
+from reflexio.server.services.configurator.local_file_config_storage import (
+    LocalFileConfigStorage,
+)
+from reflexio.server.services.storage.sqlite_storage import SQLiteStorage
+from reflexio.server.services.storage.storage_base import BaseStorage
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Storage factory functions — one per StorageConfig type
+# ---------------------------------------------------------------------------
+
+
+def _create_sqlite_storage(
+    configurator: BaseConfigurator, config: StorageConfigSQLite
+) -> BaseStorage:
+    logger.info("Using SQLite storage for org %s", configurator.org_id)
+    full_config = configurator.get_config()
+    api_key_config = full_config.api_key_config if full_config else None
+    llm_config = full_config.llm_config if full_config else None
+    enable_document_expansion = (
+        full_config.enable_document_expansion if full_config else False
+    )
+    # When base_dir is explicitly provided (e.g. tests with temp dirs)
+    # and no db_path is configured, use base_dir for the SQLite DB
+    # so the storage is isolated from the shared default database.
+    db_path = config.db_path
+    if db_path is None and configurator.base_dir:
+        db_path = str(Path(configurator.base_dir) / "reflexio.db")
+    return SQLiteStorage(
+        org_id=configurator.org_id,
+        db_path=db_path,
+        api_key_config=api_key_config,
+        llm_config=llm_config,
+        enable_document_expansion=enable_document_expansion,
+    )
+
+
+class DefaultConfigurator(BaseConfigurator):
+    """OS configurator with LocalJson config storage and SQLite data storage."""
+
+    _STORAGE_FACTORIES: dict[type[StorageConfig], Callable[..., BaseStorage]] = {
+        StorageConfigSQLite: _create_sqlite_storage,
+    }
+
+    _STORAGE_READINESS_CHECKS: dict[type[StorageConfig], Callable[[Any], bool]] = {
+        StorageConfigSQLite: lambda _: True,  # db_path defaults via env var if None
+    }
+
+    def _select_config_storage(
+        self, org_id: str, base_dir: str | None
+    ) -> ConfigStorage:
+        if base_dir:
+            return LocalFileConfigStorage(org_id=org_id, base_dir=base_dir)
+        return LocalFileConfigStorage(org_id=org_id)
+
+
+# ---------------------------------------------------------------------------
+# Configurator class registry — allows enterprise to swap in its own class
+# ---------------------------------------------------------------------------
+
+_configurator_class: type[BaseConfigurator] = DefaultConfigurator
+
+
+def set_configurator_class(cls: type[BaseConfigurator]) -> None:
+    """Register a configurator class to be used by RequestContext."""
+    global _configurator_class  # noqa: PLW0603
+    _configurator_class = cls
+
+
+def get_configurator_class() -> type[BaseConfigurator]:
+    """Return the currently registered configurator class."""
+    return _configurator_class