From ef43df8a1045fdef6ab5b0c1d5299d292df6d857 Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 15:56:31 -0700 Subject: [PATCH 1/6] Realign SDK with whiteboxxai.com branding for v1.0.0 Replaces the stale whiteboxai/whiteboxai.io (single "x") branding and src/ layout with the fixed, canonical whiteboxxai/whiteboxxai.com content from the whitebox-xai-azure monorepo's sdk/ directory: - Flat whiteboxxai/ package layout (was src/whiteboxai/); dropped the legacy setup.py in favor of pyproject.toml as the single build config. - Fixed an import-time crash affecting anyone without PyTorch/TensorFlow installed, corrected ModelMonitor's resource API calls, added structured exception metadata, and added buffering/context-manager/ drift-report/alert convenience methods (see CHANGELOG.md [1.0.0] for the full list). - Preserved and rebranded (rather than discarded) this repo's existing release scaffolding: CHANGELOG.md (added [1.0.0] and a reconstructed [0.2.1] entry), the mkdocs documentation site, and the tests/ unit+integration structure -- now populated with the repaired test suite (168 passed, 2 skipped locally, one for a missing sklearn install, one for a not-yet-implemented metrics resource). - Updated test.yml for the new layout and .env.example/MANIFEST.in for accuracy. publish.yml needs no changes (package-name-agnostic, OIDC trusted publishing). Co-Authored-By: Claude Sonnet 5 --- .env.example | 26 +- .github/workflows/test.yml | 10 +- CHANGELOG.md | 77 +- LICENSE | 2 +- MANIFEST.in | 3 +- README.md | 324 ++--- __init__.py | 1 - docs/HUGGINGFACE_INTEGRATION.md | 32 +- docs/LANGCHAIN_INTEGRATION.md | 32 +- docs/PRODUCTION_DEPLOYMENT.md | 136 +- docs/PYTORCH_INTEGRATION.md | 44 +- docs/SDK_MIGRATION_GUIDE.md | 1214 +++++++++++++++++ docs/SKLEARN_INTEGRATION.md | 30 +- docs/TENSORFLOW_INTEGRATION.md | 60 +- docs/api-reference.md | 48 +- docs/getting-started.md | 28 +- docs/index.md | 36 +- docs/integrations.md | 42 +- docs/offline-mode.md | 16 +- examples/async_monitoring.py | 12 +- examples/basic_monitoring.py | 4 +- examples/boosting_example.py | 130 +- examples/decorator_monitoring.py | 11 +- examples/langchain_example.py | 91 +- examples/offline_mode_example.py | 66 +- examples/pytorch_integration.py | 9 +- examples/sklearn_integration.py | 14 +- examples/tensorflow_example.py | 37 +- examples/transformers_example.py | 53 +- mkdocs.yml | 8 +- pyproject.toml | 126 +- requirements-dev.txt | 26 - setup.py | 100 -- src/whiteboxai/__init__.py | 24 - src/whiteboxai/__version__.py | 6 - src/whiteboxai/exceptions.py | 57 - src/whiteboxai/models/__init__.py | 1 - src/whiteboxai/resources.py | 460 ------- tests/conftest.py | 6 +- tests/integration/test_crewai_integration.py | 455 ++++++ .../integration/test_langchain_integration.py | 452 ++++++ tests/integration/test_sklearn.py | 2 +- tests/unit/test_client.py | 241 +++- tests/unit/test_config_exceptions.py | 301 ++++ tests/unit/test_decorators.py | 297 ++++ tests/unit/test_monitor.py | 335 ++++- tests/unit/test_utils.py | 277 ++++ whiteboxxai/__init__.py | 31 + {src/whiteboxai => whiteboxxai}/cache.py | 4 +- {src/whiteboxai => whiteboxxai}/client.py | 79 +- {src/whiteboxai => whiteboxxai}/config.py | 24 +- {src/whiteboxai => whiteboxxai}/decorators.py | 16 +- whiteboxxai/exceptions.py | 109 ++ {src/whiteboxai => whiteboxxai}/git_utils.py | 48 +- .../integrations/__init__.py | 59 +- .../integrations/boosting.py | 68 +- .../integrations/crewai_monitor.py | 115 +- .../integrations/langchain.py | 40 +- .../integrations/langchain_agents.py | 78 +- .../integrations/pytorch.py | 44 +- .../integrations/sklearn.py | 19 +- .../integrations/tensorflow.py | 69 +- .../integrations/transformers.py | 33 +- {src/whiteboxai => whiteboxxai}/monitor.py | 188 ++- {src/whiteboxai => whiteboxxai}/offline.py | 50 +- {src/whiteboxai => whiteboxxai}/privacy.py | 10 +- whiteboxxai/resources.py | 872 ++++++++++++ {src/whiteboxai => whiteboxxai}/utils.py | 80 +- 68 files changed, 6073 insertions(+), 1725 deletions(-) delete mode 100644 __init__.py create mode 100644 docs/SDK_MIGRATION_GUIDE.md delete mode 100644 requirements-dev.txt delete mode 100644 setup.py delete mode 100644 src/whiteboxai/__init__.py delete mode 100644 src/whiteboxai/__version__.py delete mode 100644 src/whiteboxai/exceptions.py delete mode 100644 src/whiteboxai/models/__init__.py delete mode 100644 src/whiteboxai/resources.py create mode 100644 tests/integration/test_crewai_integration.py create mode 100644 tests/integration/test_langchain_integration.py create mode 100644 tests/unit/test_config_exceptions.py create mode 100644 tests/unit/test_decorators.py create mode 100644 tests/unit/test_utils.py create mode 100644 whiteboxxai/__init__.py rename {src/whiteboxai => whiteboxxai}/cache.py (92%) rename {src/whiteboxai => whiteboxxai}/client.py (83%) rename {src/whiteboxai => whiteboxxai}/config.py (77%) rename {src/whiteboxai => whiteboxxai}/decorators.py (95%) create mode 100644 whiteboxxai/exceptions.py rename {src/whiteboxai => whiteboxxai}/git_utils.py (85%) rename {src/whiteboxai => whiteboxxai}/integrations/__init__.py (56%) rename {src/whiteboxai => whiteboxxai}/integrations/boosting.py (91%) rename {src/whiteboxai => whiteboxxai}/integrations/crewai_monitor.py (79%) rename {src/whiteboxai => whiteboxxai}/integrations/langchain.py (94%) rename {src/whiteboxai => whiteboxxai}/integrations/langchain_agents.py (88%) rename {src/whiteboxai => whiteboxxai}/integrations/pytorch.py (85%) rename {src/whiteboxai => whiteboxxai}/integrations/sklearn.py (92%) rename {src/whiteboxai => whiteboxxai}/integrations/tensorflow.py (87%) rename {src/whiteboxai => whiteboxxai}/integrations/transformers.py (94%) rename {src/whiteboxai => whiteboxxai}/monitor.py (54%) rename {src/whiteboxai => whiteboxxai}/offline.py (91%) rename {src/whiteboxai => whiteboxxai}/privacy.py (95%) create mode 100644 whiteboxxai/resources.py rename {src/whiteboxai => whiteboxxai}/utils.py (74%) diff --git a/.env.example b/.env.example index 67a8cc2..f256f64 100644 --- a/.env.example +++ b/.env.example @@ -1,20 +1,10 @@ -# WhiteBoxAI Python SDK - Environment Variables Template +# WhiteBoxXAI Python SDK - Environment Variables Template +# +# Only the two variables below are actually read from the environment by +# the SDK (see whiteboxxai/config.py). Everything else (offline mode, +# caching, privacy filters, sampling) is configured via WhiteBoxXAI() +# constructor keyword arguments instead — see README.md. # API Configuration -EXPLAINAI_API_KEY=your-api-key-here -EXPLAINAI_BASE_URL=https://api.whiteboxai.io - -# Offline Mode (Optional) -EXPLAINAI_ENABLE_OFFLINE=true -EXPLAINAI_OFFLINE_DIR=./whiteboxai_offline -EXPLAINAI_OFFLINE_AUTO_SYNC=true - -# Privacy & Security (Optional) -EXPLAINAI_ENABLE_PRIVACY_FILTERS=true - -# Caching (Optional) -EXPLAINAI_ENABLE_CACHING=true -EXPLAINAI_CACHE_TTL=3600 - -# Sampling (Optional) -EXPLAINAI_SAMPLING_RATE=1.0 +WHITEBOXXAI_API_KEY=your-api-key-here +WHITEBOXXAI_BASE_URL=https://api.whiteboxxai.com diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a1b55aa..8c345d3 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -40,15 +40,15 @@ jobs: - name: Run flake8 linter run: | - flake8 src/ tests/ examples/ --count --statistics --max-line-length=120 + flake8 whiteboxxai/ tests/ examples/ --count --statistics --max-line-length=120 - name: Run pylint run: | - pylint src/ --exit-zero --score=yes + pylint whiteboxxai/ --exit-zero --score=yes - name: Run security check (Bandit) run: | - bandit -r src/ -f json -o bandit-report.json || true + bandit -r whiteboxxai/ -f json -o bandit-report.json || true - name: Upload Bandit report uses: actions/upload-artifact@v4 @@ -62,7 +62,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.9', '3.10', '3.11'] + python-version: ['3.9', '3.10', '3.11', '3.12'] steps: - name: Checkout code @@ -81,7 +81,7 @@ jobs: - name: Run unit tests run: | - pytest tests/unit -v --cov=src/whiteboxai --cov-report=xml + pytest tests/unit -v --cov=whiteboxxai --cov-report=xml - name: Upload coverage to Codecov uses: codecov/codecov-action@v3 diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d9d991..2c0b495 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,87 @@ # Changelog -All notable changes to the WhiteBoxAI Python SDK will be documented in this file. +All notable changes to the WhiteBoxXAI Python SDK will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +## [1.0.0] - 2026-07-14 + +First stable, production release. This release realigns the SDK with the +`whiteboxxai.com` product branding and folds in the fixes/features +developed in the `whitebox-xai-azure` monorepo's `sdk/` directory. + +### Changed +- **BREAKING**: Rebranded from `whiteboxai`/`whiteboxai-sdk` (single "x", + `whiteboxai.io`/`whitebox.agentaflow.com`) to `whiteboxxai`/`whitebox-xai-sdk` + (double "x", `whiteboxxai.com`). Update imports: `from whiteboxai import ...` + → `from whiteboxxai import ...`; update the install command to + `pip install whitebox-xai-sdk`; update the `EXPLAINAI_*` env vars to + `WHITEBOXXAI_*` (`WHITEBOXXAI_API_KEY`, `WHITEBOXXAI_BASE_URL`). +- **BREAKING**: Switched from a `src/whiteboxai/` layout to a flat + `whiteboxxai/` layout at the repository root, matching the canonical + source in `whitebox-xai-azure/sdk/`. Removed the legacy `setup.py`; + `pyproject.toml` is now the single build configuration. +- Version is now read from installed package metadata + (`importlib.metadata.version("whitebox-xai-sdk")`) instead of being + hardcoded in three separate places. +- `Development Status` classifier bumped from Beta to Production/Stable. +- `ModelMonitor.log_prediction()`/`alog_prediction()`/`log_batch()` now call + the predictions API with the correct `input_data`/`output_data` field + names (previously sent `inputs`/`outputs`, which the real API doesn't + accept). + +### Added +- `ModelMonitor` local buffering (`buffer_size=`) with `flush()`/`aflush()` + and context-manager support (`with ModelMonitor(...) as monitor:` + flushes on exit). +- `ModelMonitor.get_prediction_count()`, `get_drift_reports()`, + `create_alert_rule()`, and `get_active_alerts()` convenience methods. +- `log_prediction(explain=True)` now actually triggers explanation + generation via the explanations resource, instead of silently + accepting and dropping the flag. +- SDK exception classes (`APIError`, `AuthenticationError`, `RateLimitError`, + `ValidationError`, `NotFoundError`) now carry structured metadata + (`status_code`, `response`, `request_id`, `retry_after`, `fields`, etc.) + instead of being bare, attribute-less `Exception` subclasses. +- `Config` now validates `timeout`/`max_retries` are positive. +- MCP (Model Context Protocol) usage documented in `README.md` for + non-Python integrations, pointing at the companion `whiteboxxai-mcp` + package. +- `LICENSE` file (MIT), matching `pyproject.toml`'s declared license. + +### Fixed +- Fixed an import-time crash: `whiteboxxai.integrations` (and therefore the + whole SDK) failed to import for anyone without PyTorch or TensorFlow + installed, due to bad optional-dependency fallback stubs in + `integrations/pytorch.py` and `integrations/tensorflow.py`, and an + unguarded/unbound name in `integrations/langchain.py`. +- `README.md` code samples fixed throughout: wrong import path + (`whiteboxai` → `whiteboxxai`), wrong class name (`WhiteBoxAI` → + `WhiteBoxXAI`), and wrong install command. +- Corrected `WhiteBoxXAI_*`-cased env var references (wrong case, and in + one place the wrong variable name entirely) to the real + `WHITEBOXXAI_API_KEY`/`WHITEBOXXAI_BASE_URL` across guides and examples. +- `docs/` guides (`api-reference.md`, `getting-started.md`, `index.md`, + `integrations.md`, `offline-mode.md`) rebranded from the stale + `whiteboxai.io`/`agentaflow.com`/`EXPLAINAI_*` naming. + +## [0.2.1] - 2026-05-03 + +No changelog entry was recorded for this release at the time. Reconstructed +from the git history between the `0.2.0` and `0.2.1` tags: + +### Added +- `AgentWorkflows` resource and client binding for multi-agent workflow + tracking. + +### Fixed +- e2e pytest exit-code handling. +- Import normalization and formatting cleanup across integrations and + examples. + ## [0.2.0] - 2026-02-10 ### Added diff --git a/LICENSE b/LICENSE index c2063f0..2ac025c 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2026 AgentaFlow +Copyright (c) 2026 WhiteBoxXAI Team Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/MANIFEST.in b/MANIFEST.in index 52c34f3..f1e9f9c 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -5,7 +5,6 @@ include CHANGELOG.md # Include configuration files include pyproject.toml -include setup.py include pytest.ini include .gitignore @@ -22,7 +21,7 @@ recursive-include examples *.py recursive-include tests *.py # Include source package data -recursive-include src *.py +recursive-include whiteboxxai *.py # Exclude compiled files global-exclude *.py[cod] diff --git a/README.md b/README.md index e1d161c..5df72c3 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,27 @@ -# WhiteBoxAI Python SDK +# WhiteBoxXAI Python SDK -Official Python SDK for integrating WhiteBoxAI monitoring into your ML applications. +Official Python SDK for integrating WhiteBoxXAI monitoring into your ML applications. ## Features - 🚀 **Easy Integration** - Monitor models with just a few lines of code -- 📊 **Framework Support** - Native integrations for Scikit-learn, PyTorch, TensorFlow, XGBoost, LightGBM, Hugging Face Transformers, and LangChain -- 🤖 **Multi-Agent Monitoring** - First-class support for CrewAI and LangGraph multi-agent workflows +- 📊 **Framework Support** - Native integrations for Scikit-learn, PyTorch, TensorFlow, XGBoost, and more - 🎯 **Decorator-based Monitoring** - Zero-code-change monitoring with decorators - ⚡ **Async/Sync Interfaces** - Support for both synchronous and asynchronous workflows - 🔒 **Privacy-First** - Built-in PII detection and data masking - 💾 **Local Caching** - TTL-based caching to reduce API calls - 📈 **Drift Detection** - Automatic model and data drift monitoring -- 🗂️ **Git Context** - Automatic capture of git branch, commit, and author metadata - 🎨 **Flexible Configuration** - Extensive configuration options and feature flags ## Installation ```bash -pip install whiteboxai-sdk +pip install whitebox-xai-sdk # With specific framework support -pip install whiteboxai-sdk[sklearn] -pip install whiteboxai-sdk[pytorch] -pip install whiteboxai-sdk[all] # All integrations +pip install whitebox-xai-sdk[sklearn] +pip install whitebox-xai-sdk[pytorch] +pip install whitebox-xai-sdk[all] # All integrations ``` ## Quick Start @@ -31,10 +29,10 @@ pip install whiteboxai-sdk[all] # All integrations ### Basic Usage ```python -from whiteboxai import WhiteBoxAI, ModelMonitor +from whiteboxxai import WhiteBoxXAI, ModelMonitor # Initialize client -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") # Create monitor monitor = ModelMonitor(client) @@ -57,15 +55,15 @@ monitor.log_prediction( ```python from sklearn.ensemble import RandomForestClassifier -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.sklearn import SklearnMonitor # Train model model = RandomForestClassifier() model.fit(X_train, y_train) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = SklearnMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -81,8 +79,8 @@ predictions = monitored_model.predict(X_test) ```python import torch import torch.nn as nn -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.pytorch import TorchMonitor # Define model model = nn.Sequential( @@ -92,7 +90,7 @@ model = nn.Sequential( ) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = TorchMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -108,8 +106,8 @@ with torch.no_grad(): ```python from tensorflow import keras -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.tensorflow import KerasMonitor, WhiteBoxAICallback +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.tensorflow import KerasMonitor, WhiteBoxXAICallback # Build model model = keras.Sequential([ @@ -119,12 +117,12 @@ model = keras.Sequential([ model.compile(optimizer='adam', loss='mse') # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = KerasMonitor(client, model=model, model_name="keras_model") monitor.register_from_model(model_type="regression") # Train with monitoring callback -callback = WhiteBoxAICallback(monitor, log_frequency=1) +callback = WhiteBoxXAICallback(monitor, log_frequency=1) model.fit(X_train, y_train, validation_split=0.2, callbacks=[callback], @@ -138,14 +136,14 @@ predictions = monitor.predict(X_test, log=True) ```python from transformers import pipeline -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.transformers import TransformersMonitor, wrap_transformers_pipeline +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.transformers import TransformersMonitor, wrap_transformers_pipeline # Load model classifier = pipeline("sentiment-analysis") # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = TransformersMonitor( client=client, pipeline=classifier, @@ -169,11 +167,11 @@ result = wrapped("Great service!") # Automatically logged from langchain.chains import LLMChain from langchain.llms import OpenAI from langchain.prompts import PromptTemplate -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.langchain import LangChainMonitor, wrap_langchain_chain +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.langchain import LangChainMonitor, wrap_langchain_chain # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = LangChainMonitor( client=client, application_name="qa_bot", @@ -198,45 +196,20 @@ wrapped_chain = wrap_langchain_chain(chain, monitor) result = wrapped_chain.run(question="What is AI?") # Automatically logged ``` -### Multi-Agent Monitoring (CrewAI / LangGraph) - -```python -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.crewai_monitor import CrewAIMonitor - -client = WhiteBoxAI(api_key="your-api-key") -monitor = CrewAIMonitor(client=client) - -# Start workflow monitoring -workflow_id = monitor.start_monitoring( - workflow_name="research_pipeline", - description="Multi-agent research and writing workflow" -) - -# Register and track agents -monitor.register_agent(researcher, role="Research Analyst") -monitor.register_agent(writer, role="Content Writer") - -# Log task completions and agent interactions -monitor.log_task_completion(task, status="completed", output_data=result) -monitor.log_interaction(from_agent=researcher, to_agent=writer, interaction_type="handoff") - -# Complete and retrieve analytics -summary = monitor.complete_monitoring(status="completed") -``` - -For LangGraph multi-agent workflows, use `LangGraphMultiAgentMonitor` and `MultiAgentCallbackHandler` -from `whiteboxai.integrations.langchain_agents`. +For multi-agent LangChain/LangGraph workflows and CrewAI, see +`whiteboxxai.integrations.langchain_agents` (`MultiAgentCallbackHandler`, +`LangGraphMultiAgentMonitor`, `monitor_langchain_agent`) and +`whiteboxxai.integrations.crewai_monitor` (`CrewAIMonitor`, `monitor_crew`). ### XGBoost/LightGBM Monitoring ```python import xgboost as xgb import lightgbm as lgb -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.boosting import XGBoostMonitor, LightGBMMonitor, wrap_xgboost_model +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.boosting import XGBoostMonitor, LightGBMMonitor, wrap_xgboost_model -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") # XGBoost monitoring xgb_monitor = XGBoostMonitor( @@ -275,9 +248,9 @@ predictions = lgb_monitor.predict(model, X_test, y_test) ### Decorator-based Monitoring ```python -from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_model +from whiteboxxai import WhiteBoxXAI, ModelMonitor, monitor_model -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = ModelMonitor(client, model_id=123) @monitor_model(monitor, input_keys=["features"], explain=True) @@ -293,10 +266,10 @@ result = predict(features=[1.0, 2.0, 3.0]) ```python import asyncio -from whiteboxai import WhiteBoxAI, ModelMonitor +from whiteboxxai import WhiteBoxXAI, ModelMonitor async def main(): - async with WhiteBoxAI(api_key="your-api-key") as client: + async with WhiteBoxXAI(api_key="your-api-key") as client: monitor = ModelMonitor(client) # Register model @@ -316,25 +289,24 @@ asyncio.run(main()) ## Advanced Features -### Git Context +### Local Buffering & Batch Flushing -Capture git metadata automatically alongside model monitoring: +For high-throughput logging, buffer predictions locally and flush them as a +batch instead of sending every prediction immediately: ```python -from whiteboxai import WhiteBoxAI, GitContext, detect_git_context +from whiteboxxai import WhiteBoxXAI, ModelMonitor -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") -# Auto-detect from current git repo -git_ctx = detect_git_context() -print(git_ctx.branch, git_ctx.commit_hash, git_ctx.author) +# Predictions are buffered locally; a batch is sent once 100 accumulate, +# or when the buffer is flushed (explicitly, or on context-manager exit). +with ModelMonitor(client, model_id=123, buffer_size=100) as monitor: + for features, output in predictions: + monitor.log_prediction(inputs=features, output=output) + # Any remaining buffered predictions are flushed automatically here. -# Or build manually -git_ctx = GitContext( - branch="main", - commit_hash="abc1234", - author="dev@example.com" -) +print(monitor.get_prediction_count()) ``` ### Offline Mode @@ -342,13 +314,13 @@ git_ctx = GitContext( Enable robust operation with unreliable network connectivity. Operations are queued locally and synced automatically. ```python -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI # Enable offline mode with auto-sync -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_offline=True, - offline_dir="./whiteboxai_offline", + offline_dir="./whiteboxxai_offline", offline_auto_sync=True, offline_sync_interval=60 # Sync every 60 seconds ) @@ -375,7 +347,7 @@ client.cleanup_offline_queue(older_than_days=7) **Configuration:** ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_offline=True, offline_dir="./offline_queue", # Storage directory @@ -385,15 +357,13 @@ client = WhiteBoxAI( ) ``` -See [Offline Mode Guide](docs/offline-mode.md) for complete documentation. - ### Privacy Filters ```python -from whiteboxai import WhiteBoxAI -from whiteboxai.privacy import mask_data +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.privacy import mask_data -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_privacy_filters=True ) @@ -413,7 +383,7 @@ masked = mask_data(data) ### Local Caching ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_caching=True, cache_ttl=3600, @@ -444,6 +414,9 @@ monitor.set_baseline(baseline) # Detect drift current_data = np.random.randn(100, 10) drift_report = monitor.detect_drift(current_data) + +# Retrieve previously persisted drift reports +reports = monitor.get_drift_reports(limit=10) ``` ## Configuration @@ -451,20 +424,20 @@ drift_report = monitor.detect_drift(current_data) The SDK can be configured via constructor parameters or environment variables: ```python -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI -client = WhiteBoxAI( - api_key="your-api-key", # or EXPLAINAI_API_KEY env var - base_url="https://api.whiteboxai.io", # or EXPLAINAI_BASE_URL env var +client = WhiteBoxXAI( + api_key="your-api-key", # or WHITEBOXXAI_API_KEY env var + base_url="https://api.whiteboxxai.com", # Custom API endpoint, or WHITEBOXXAI_BASE_URL env var timeout=30, # Request timeout (seconds) max_retries=3, # Retry attempts # Offline mode enable_offline=True, # Enable offline queueing - offline_dir="./whiteboxai_offline", # Queue storage directory - offline_max_queue_size=10000, # Max queued operations - offline_auto_sync=True, # Auto-sync in background - offline_sync_interval=60, # Sync interval (seconds) + offline_dir="./whiteboxxai_offline", # Queue storage directory + offline_max_queue_size=10000, # Max queued operations + offline_auto_sync=True, # Auto-sync in background + offline_sync_interval=60, # Sync interval (seconds) # Other features enable_caching=True, # Enable local caching @@ -474,19 +447,67 @@ client = WhiteBoxAI( ) ``` +### Authentication + +The `api_key` parameter (or `WHITEBOXXAI_API_KEY` env var) is sent as a +bearer token: `Authorization: Bearer `. As of this SDK's current +release, the WhiteBoxXAI backend validates this token as a standard JWT +obtained via account login (`/api/v1/auth/login`), not a separate, +dedicated API-key entity — there is no self-service API key issuance/ +management endpoint yet. In practice: log in through your WhiteBoxXAI +account (dashboard or the `/auth/login` endpoint) and use the returned +token as `api_key`. This also matches how the companion MCP server +authenticates (see below) and will be updated here once dedicated, +long-lived API keys ship on the backend. + +## Using WhiteBoxXAI from Other Languages (MCP) + +This SDK is the primary, first-class integration path and is Python-only. +For other languages, or for agentic clients (Claude Desktop, Claude Code, +LangChain, custom agent harnesses), use the companion +[Model Context Protocol](https://modelcontextprotocol.io/) server, +`whiteboxxai-mcp`, instead of calling the API directly: + +```bash +pip install whiteboxxai-mcp + +# Configure a service-account credential, then run the stdio server +export WHITEBOXXAI_MCP_API_BASE_URL="https://api.whiteboxxai.com" +export WHITEBOXXAI_MCP_EMAIL="mcp-service@yourorg.com" +export WHITEBOXXAI_MCP_PASSWORD="..." +whiteboxxai-mcp +``` + +Point any MCP-compatible client at this command (e.g. Claude Desktop's +`mcpServers` config, or Claude Code's `.mcp.json`). MCP is a language- and +client-agnostic protocol, so this is the recommended path for non-Python +integrations. + +**Current limitations** (as of this SDK's release): only models, +predictions, drift, and bias/fairness tools are available; explanations, +LLM/RAG observability, safety, alerts, and multi-agent workflow tools are +tracked as follow-on milestones, and `whitebox_explanations_generate` +currently returns placeholder feature-importance values rather than real +SHAP/LIME output. Auth uses the same username/password (or short-lived +JWT) mechanism described above — there is no long-lived API-key/ +service-account system yet. + +See the `whiteboxxai-mcp` package's own README for full setup and its +tool reference. + ## API Reference -### WhiteBoxAI Client +### WhiteBoxXAI Client Main client for API interaction. -**Resources:** -- `models` - Register and manage models -- `predictions` - Log predictions and batches -- `explanations` - Generate and retrieve explanations -- `drift` - Drift detection and reports -- `alerts` - Alert management -- `agent_workflows` - Multi-agent workflow tracking +**Methods:** +- `models` - Models resource +- `predictions` - Predictions resource +- `explanations` - Explanations resource +- `drift` - Drift detection resource +- `fairness` - Bias/fairness auditing resource +- `alerts` - Alerts resource ### ModelMonitor @@ -494,10 +515,16 @@ Simplified monitoring interface. **Methods:** - `register_model()` - Register a new model -- `log_prediction()` - Log a single prediction +- `log_prediction()` - Log a single prediction (or buffer it, if `buffer_size` is set) - `log_batch()` - Log multiple predictions +- `flush()` - Send any buffered predictions immediately +- `get_prediction_count()` - Number of predictions logged by this monitor instance - `set_baseline()` - Set baseline data for drift detection - `detect_drift()` - Detect model drift +- `get_drift_reports()` - Retrieve previously persisted drift reports +- `create_alert_rule()` - Create a threshold-based alert rule for this model +- `get_active_alerts()` - List alert rules for this model +- Can be used as a context manager (`with ModelMonitor(...) as monitor:`); buffered predictions are flushed on exit. ### Decorators @@ -506,95 +533,36 @@ Simplified monitoring interface. ### Framework Integrations -- `whiteboxai.integrations.sklearn` - Scikit-learn integration -- `whiteboxai.integrations.pytorch` - PyTorch integration -- `whiteboxai.integrations.tensorflow` - TensorFlow/Keras integration -- `whiteboxai.integrations.transformers` - Hugging Face Transformers integration -- `whiteboxai.integrations.langchain` - LangChain chains and agents -- `whiteboxai.integrations.langchain_agents` - LangGraph multi-agent monitoring -- `whiteboxai.integrations.crewai_monitor` - CrewAI workflow monitoring -- `whiteboxai.integrations.boosting` - XGBoost and LightGBM integration +- `whiteboxxai.integrations.sklearn` - Scikit-learn integration +- `whiteboxxai.integrations.pytorch` - PyTorch integration +- `whiteboxxai.integrations.tensorflow` - TensorFlow/Keras integration +- `whiteboxxai.integrations.transformers` - Hugging Face Transformers integration +- `whiteboxxai.integrations.langchain` - LangChain chains/agents integration +- `whiteboxxai.integrations.langchain_agents` - LangChain/LangGraph multi-agent integration +- `whiteboxxai.integrations.crewai_monitor` - CrewAI multi-agent integration +- `whiteboxxai.integrations.boosting` - XGBoost/LightGBM integration ## Examples See the `examples/` directory for more examples: - `basic_monitoring.py` - Basic monitoring example -- `async_monitoring.py` - Async API usage -- `decorator_monitoring.py` - Zero-code-change decorator usage - `sklearn_integration.py` - Scikit-learn integration - `pytorch_integration.py` - PyTorch integration - `tensorflow_example.py` - TensorFlow/Keras integration - `transformers_example.py` - Hugging Face Transformers integration -- `langchain_example.py` - LangChain chains and agents +- `langchain_example.py` - LangChain integration - `boosting_example.py` - XGBoost/LightGBM integration +- `decorator_monitoring.py` - Decorator-based monitoring +- `async_monitoring.py` - Async API usage - `offline_mode_example.py` - Offline mode with queue management ## Support -- Documentation: https://whitebox.agentaflow.com +- Documentation: https://docs.whiteboxxai.com - Issues: https://github.com/AgentaFlow/whitebox-python-sdk/issues +- Email: support@whiteboxxai.com -``` -whiteboxai-python-sdk/ -├── src/ -│ └── whiteboxai/ -│ ├── __init__.py -│ ├── __version__.py -│ ├── client.py -│ ├── config.py -│ ├── monitor.py -│ ├── decorators.py -│ ├── privacy.py -│ ├── offline.py -│ ├── cache.py -│ ├── git_utils.py -│ ├── resources.py -│ ├── utils.py -│ ├── exceptions.py -│ ├── integrations/ -│ │ ├── sklearn.py -│ │ ├── pytorch.py -│ │ ├── tensorflow.py -│ │ ├── transformers.py -│ │ ├── langchain.py -│ │ ├── langchain_agents.py -│ │ ├── crewai_monitor.py -│ │ └── boosting.py -│ └── models/ -├── tests/ -│ ├── unit/ -│ ├── integration/ -│ └── e2e/ -├── examples/ -│ ├── basic_monitoring.py -│ ├── async_monitoring.py -│ ├── decorator_monitoring.py -│ ├── sklearn_integration.py -│ ├── pytorch_integration.py -│ ├── tensorflow_example.py -│ ├── transformers_example.py -│ ├── langchain_example.py -│ ├── boosting_example.py -│ └── offline_mode_example.py -├── docs/ -│ ├── getting-started.md -│ ├── integrations.md -│ ├── offline-mode.md -│ ├── api-reference.md -│ ├── SKLEARN_INTEGRATION.md -│ ├── PYTORCH_INTEGRATION.md -│ ├── TENSORFLOW_INTEGRATION.md -│ ├── HUGGINGFACE_INTEGRATION.md -│ ├── LANGCHAIN_INTEGRATION.md -│ └── PRODUCTION_DEPLOYMENT.md -├── pyproject.toml -├── setup.py -├── README.md -├── CHANGELOG.md -├── LICENSE -└── .github/ - └── workflows/ - ├── test.yml - └── publish.yml -``` +## License + +MIT License - see LICENSE file for details diff --git a/__init__.py b/__init__.py deleted file mode 100644 index b2b2619..0000000 --- a/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# WhiteBoxAI SDK diff --git a/docs/HUGGINGFACE_INTEGRATION.md b/docs/HUGGINGFACE_INTEGRATION.md index 5b86a2d..c76e9f9 100644 --- a/docs/HUGGINGFACE_INTEGRATION.md +++ b/docs/HUGGINGFACE_INTEGRATION.md @@ -1,6 +1,6 @@ # Hugging Face Transformers Integration Guide -This guide demonstrates how to use WhiteBoxAI to monitor Hugging Face Transformers models for NLP tasks including text classification, named entity recognition, question answering, and text generation. +This guide demonstrates how to use WhiteBoxXAI to monitor Hugging Face Transformers models for NLP tasks including text classification, named entity recognition, question answering, and text generation. ## Table of Contents @@ -14,27 +14,27 @@ This guide demonstrates how to use WhiteBoxAI to monitor Hugging Face Transforme ## Installation -Install WhiteBoxAI SDK with Transformers support: +Install WhiteBoxXAI SDK with Transformers support: ```bash -pip install whiteboxai[transformers] +pip install whiteboxxai[transformers] ``` Or install dependencies separately: ```bash -pip install whiteboxai transformers torch +pip install whiteboxxai transformers torch ``` ## Quick Start ```python from transformers import pipeline -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.transformers import TransformersMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.transformers import TransformersMonitor # Initialize client -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") # Load Hugging Face pipeline classifier = pipeline("sentiment-analysis") @@ -59,13 +59,13 @@ print(result) # [{'label': 'POSITIVE', 'score': 0.9998}] ## Supported Tasks -WhiteBoxAI supports all major Hugging Face Transformers tasks: +WhiteBoxXAI supports all major Hugging Face Transformers tasks: ### Text Classification ```python from transformers import pipeline -from whiteboxai.integrations.transformers import TransformersMonitor +from whiteboxxai.integrations.transformers import TransformersMonitor # Sentiment analysis classifier = pipeline("sentiment-analysis") @@ -182,7 +182,7 @@ summary = monitor.predict(article, max_length=130, min_length=30, log=True) ### Method 1: Using TransformersMonitor ```python -from whiteboxai.integrations.transformers import TransformersMonitor +from whiteboxxai.integrations.transformers import TransformersMonitor # Create monitor monitor = TransformersMonitor( @@ -208,7 +208,7 @@ results = monitor.predict(["Text 1", "Text 2"], log=True) ### Method 2: Using Pipeline Wrapper ```python -from whiteboxai.integrations.transformers import wrap_transformers_pipeline +from whiteboxxai.integrations.transformers import wrap_transformers_pipeline # Wrap pipeline for auto-logging wrapped_classifier = wrap_transformers_pipeline(classifier, monitor) @@ -220,7 +220,7 @@ result = wrapped_classifier("Input text") ### Method 3: Using TransformersPipelineWrapper ```python -from whiteboxai.integrations.transformers import TransformersPipelineWrapper +from whiteboxxai.integrations.transformers import TransformersPipelineWrapper # Wrap pipeline (with auto-registration) wrapper = TransformersPipelineWrapper( @@ -528,7 +528,7 @@ See complete examples in: Main class for monitoring Transformers models. **Methods**: -- `register_from_model()` - Register model with WhiteBoxAI +- `register_from_model()` - Register model with WhiteBoxXAI - `predict()` - Make predictions with logging - `set_baseline()` - Set baseline data for drift detection - `log_generation_metrics()` - Log text generation metrics @@ -544,6 +544,6 @@ Function to wrap existing pipelines. ## Support For issues or questions: -- GitHub Issues: https://github.com/whiteboxai/whiteboxai -- Documentation: https://docs.whiteboxai.com -- Email: support@whiteboxai.com +- GitHub Issues: https://github.com/whiteboxxai/whiteboxxai +- Documentation: https://docs.whiteboxxai.com +- Email: support@whiteboxxai.com diff --git a/docs/LANGCHAIN_INTEGRATION.md b/docs/LANGCHAIN_INTEGRATION.md index 17711e6..6e77b7a 100644 --- a/docs/LANGCHAIN_INTEGRATION.md +++ b/docs/LANGCHAIN_INTEGRATION.md @@ -1,6 +1,6 @@ # LangChain Integration Guide -This guide demonstrates how to use WhiteBoxAI to monitor LangChain applications including chains, agents, and RAG pipelines. +This guide demonstrates how to use WhiteBoxXAI to monitor LangChain applications including chains, agents, and RAG pipelines. ## Table of Contents @@ -14,10 +14,10 @@ This guide demonstrates how to use WhiteBoxAI to monitor LangChain applications ## Installation -Install WhiteBoxAI SDK with LangChain support: +Install WhiteBoxXAI SDK with LangChain support: ```bash -pip install whiteboxai langchain +pip install whiteboxxai langchain ``` For specific LLM providers, install additional dependencies: @@ -39,11 +39,11 @@ pip install transformers from langchain.chains import LLMChain from langchain.llms import OpenAI from langchain.prompts import PromptTemplate -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.langchain import LangChainMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.langchain import LangChainMonitor # Initialize client -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") # Create monitor monitor = LangChainMonitor( @@ -85,9 +85,9 @@ Main class for monitoring LangChain applications. Tracks: - Tool usage - RAG retrievals -### WhiteBoxAICallbackHandler +### WhiteBoxXAICallbackHandler -LangChain callback handler that automatically logs events to WhiteBoxAI. +LangChain callback handler that automatically logs events to WhiteBoxXAI. ### Supported Components @@ -102,7 +102,7 @@ LangChain callback handler that automatically logs events to WhiteBoxAI. ### Method 1: Callback Handler (Recommended) ```python -from whiteboxai.integrations.langchain import LangChainMonitor +from whiteboxxai.integrations.langchain import LangChainMonitor # Create monitor monitor = LangChainMonitor(client, application_name="my_app") @@ -119,7 +119,7 @@ agent.run(input="...", callbacks=[callback]) ### Method 2: Wrap Chain ```python -from whiteboxai.integrations.langchain import wrap_langchain_chain +from whiteboxxai.integrations.langchain import wrap_langchain_chain # Wrap chain for automatic logging wrapped_chain = wrap_langchain_chain(chain, monitor) @@ -316,7 +316,7 @@ monitor = LangChainMonitor( callback = monitor.create_callback_handler() result = chain.run(input="...", callbacks=[callback]) -# Access metrics via WhiteBoxAI dashboard +# Access metrics via WhiteBoxXAI dashboard ``` ### Custom Metadata @@ -504,7 +504,7 @@ See complete examples in: Main class for monitoring LangChain applications. **Methods**: -- `register_application()` - Register application with WhiteBoxAI +- `register_application()` - Register application with WhiteBoxXAI - `create_callback_handler()` - Create callback handler - `log_chain_execution()` - Log chain execution - `log_agent_execution()` - Log agent run @@ -512,7 +512,7 @@ Main class for monitoring LangChain applications. - `log_tool_call()` - Log tool usage - `log_rag_retrieval()` - Log RAG retrieval -### WhiteBoxAICallbackHandler +### WhiteBoxXAICallbackHandler LangChain callback handler for automatic logging. @@ -529,6 +529,6 @@ Function to wrap chains for automatic logging. ## Support For issues or questions: -- GitHub Issues: https://github.com/whiteboxai/whiteboxai -- Documentation: https://docs.whiteboxai.com -- Email: support@whiteboxai.com +- GitHub Issues: https://github.com/whiteboxxai/whiteboxxai +- Documentation: https://docs.whiteboxxai.com +- Email: support@whiteboxxai.com diff --git a/docs/PRODUCTION_DEPLOYMENT.md b/docs/PRODUCTION_DEPLOYMENT.md index 821709a..2e8ea84 100644 --- a/docs/PRODUCTION_DEPLOYMENT.md +++ b/docs/PRODUCTION_DEPLOYMENT.md @@ -1,6 +1,6 @@ # Production Deployment Guide -This guide covers best practices for deploying WhiteBoxAI monitoring in production environments. +This guide covers best practices for deploying WhiteBoxXAI monitoring in production environments. ## Table of Contents @@ -25,7 +25,7 @@ This guide covers best practices for deploying WhiteBoxAI monitoring in producti │ Application Layer │ │ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ │ │ ML Service 1 │ │ ML Service 2 │ │ ML Service N │ │ -│ │ + WhiteBoxAI │ │ + WhiteBoxAI │ │ + WhiteBoxAI │ │ +│ │ + WhiteBoxXAI │ │ + WhiteBoxXAI │ │ + WhiteBoxXAI │ │ │ │ SDK │ │ SDK │ │ SDK │ │ │ └──────┬───────┘ └──────┬───────┘ └──────┬───────┘ │ └─────────┼──────────────────┼──────────────────┼─────────────┘ @@ -37,7 +37,7 @@ This guide covers best practices for deploying WhiteBoxAI monitoring in producti └────────┬────────┘ │ ┌────────▼────────┐ - │ WhiteBoxAI API │ + │ WhiteBoxXAI API │ │ (FastAPI) │ └────────┬────────┘ │ @@ -52,7 +52,7 @@ This guide covers best practices for deploying WhiteBoxAI monitoring in producti ### Components 1. **SDK Integration**: Lightweight SDK embedded in ML services -2. **API Gateway**: Load-balanced WhiteBoxAI API endpoints +2. **API Gateway**: Load-balanced WhiteBoxXAI API endpoints 3. **Data Layer**: PostgreSQL + TimescaleDB, Redis cache, object storage 4. **Worker Layer**: Background workers for explanations and reports @@ -64,42 +64,42 @@ This guide covers best practices for deploying WhiteBoxAI monitoring in producti ```bash # API Configuration -WHITEBOXAI_API_KEY=prod-api-key-here -WHITEBOXAI_BASE_URL=https://api.whiteboxai.yourcompany.com -WHITEBOXAI_ENVIRONMENT=production +WHITEBOXXAI_API_KEY=prod-api-key-here +WHITEBOXXAI_BASE_URL=https://api.whiteboxxai.yourcompany.com +WHITEBOXXAI_ENVIRONMENT=production # Performance Tuning -WHITEBOXAI_BATCH_SIZE=100 -WHITEBOXAI_ASYNC_LOGGING=true -WHITEBOXAI_SAMPLING_RATE=0.1 -WHITEBOXAI_CACHE_ENABLED=true -WHITEBOXAI_CACHE_TTL=3600 +WHITEBOXXAI_BATCH_SIZE=100 +WHITEBOXXAI_ASYNC_LOGGING=true +WHITEBOXXAI_SAMPLING_RATE=0.1 +WHITEBOXXAI_CACHE_ENABLED=true +WHITEBOXXAI_CACHE_TTL=3600 # Reliability -WHITEBOXAI_MAX_RETRIES=3 -WHITEBOXAI_RETRY_BACKOFF=exponential -WHITEBOXAI_TIMEOUT=30 -WHITEBOXAI_CIRCUIT_BREAKER_ENABLED=true +WHITEBOXXAI_MAX_RETRIES=3 +WHITEBOXXAI_RETRY_BACKOFF=exponential +WHITEBOXXAI_TIMEOUT=30 +WHITEBOXXAI_CIRCUIT_BREAKER_ENABLED=true # Security -WHITEBOXAI_TLS_VERIFY=true -WHITEBOXAI_PII_DETECTION=true -WHITEBOXAI_DATA_MASKING=true +WHITEBOXXAI_TLS_VERIFY=true +WHITEBOXXAI_PII_DETECTION=true +WHITEBOXXAI_DATA_MASKING=true # Monitoring -WHITEBOXAI_ENABLE_METRICS=true -WHITEBOXAI_METRICS_PORT=9090 -WHITEBOXAI_LOG_LEVEL=info +WHITEBOXXAI_ENABLE_METRICS=true +WHITEBOXXAI_METRICS_PORT=9090 +WHITEBOXXAI_LOG_LEVEL=info ``` ### Configuration File (config.yaml) ```yaml -whiteboxai: +whiteboxxai: # API Settings api: - base_url: ${WHITEBOXAI_BASE_URL} - api_key: ${WHITEBOXAI_API_KEY} + base_url: ${WHITEBOXXAI_BASE_URL} + api_key: ${WHITEBOXXAI_API_KEY} timeout: 30 max_retries: 3 @@ -155,15 +155,15 @@ whiteboxai: ```python import os -from whiteboxai import WhiteBoxAI -from whiteboxai.config import Config +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.config import Config # Load config based on environment env = os.getenv('ENVIRONMENT', 'development') config = Config.from_file(f'config.{env}.yaml') # Initialize client with config -client = WhiteBoxAI(config=config) +client = WhiteBoxXAI(config=config) ``` ### Secrets Management @@ -183,8 +183,8 @@ def get_secret(secret_name): raise e # Get API key from secrets manager -api_key = get_secret('whiteboxai/api-key') -client = WhiteBoxAI(api_key=api_key) +api_key = get_secret('whiteboxxai/api-key') +client = WhiteBoxXAI(api_key=api_key) ``` #### HashiCorp Vault @@ -199,8 +199,8 @@ def get_vault_secret(path): return secret['data']['data'] # Get API key from Vault -secrets = get_vault_secret('whiteboxai/production') -client = WhiteBoxAI(api_key=secrets['api_key']) +secrets = get_vault_secret('whiteboxxai/production') +client = WhiteBoxXAI(api_key=secrets['api_key']) ``` --- @@ -212,7 +212,7 @@ client = WhiteBoxAI(api_key=secrets['api_key']) Always use batch logging for production: ```python -from whiteboxai.monitor import ModelMonitor +from whiteboxxai.monitor import ModelMonitor from collections import deque import threading import time @@ -267,7 +267,7 @@ Use async logging to prevent blocking: ```python import asyncio -from whiteboxai import AsyncWhiteBoxAI +from whiteboxxai import AsyncWhiteBoxXAI async def predict_and_log(model, X): # Make prediction (sync) @@ -283,7 +283,7 @@ async def predict_and_log(model, X): return prediction # Initialize async client -client = AsyncWhiteBoxAI(api_key='your-api-key') +client = AsyncWhiteBoxXAI(api_key='your-api-key') # Run asyncio.run(predict_and_log(model, X)) @@ -334,7 +334,7 @@ class CachedMonitor: Use connection pooling for database connections: ```python -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI from urllib3.util.retry import Retry from requests.adapters import HTTPAdapter @@ -354,7 +354,7 @@ session.mount('http://', adapter) session.mount('https://', adapter) # Use custom session -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key='your-api-key', session=session ) @@ -368,20 +368,20 @@ client = WhiteBoxAI( ```python # DON'T: Hardcode API keys -client = WhiteBoxAI(api_key='sk-abc123...') # BAD! +client = WhiteBoxXAI(api_key='sk-abc123...') # BAD! # DO: Use environment variables -client = WhiteBoxAI() # Reads from WHITEBOXAI_API_KEY +client = WhiteBoxXAI() # Reads from WHITEBOXXAI_API_KEY # DO: Use secrets management from your_secrets import get_secret -client = WhiteBoxAI(api_key=get_secret('whiteboxai-api-key')) +client = WhiteBoxXAI(api_key=get_secret('whiteboxxai-api-key')) ``` ### 2. PII Detection and Masking ```python -from whiteboxai.security import PIIDetector, DataMasker +from whiteboxxai.security import PIIDetector, DataMasker # Enable PII detection pii_detector = PIIDetector() @@ -410,7 +410,7 @@ def safe_log_prediction(inputs, prediction): ```python # Verify SSL certificates in production -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key='your-api-key', verify_ssl=True, cert='/path/to/cert.pem' # Optional custom cert @@ -420,7 +420,7 @@ client = WhiteBoxAI( ### 4. Rate Limiting ```python -from whiteboxai.utils import RateLimiter +from whiteboxxai.utils import RateLimiter # Implement client-side rate limiting limiter = RateLimiter( @@ -444,7 +444,7 @@ def log_prediction(inputs, prediction): ### 1. Circuit Breaker Pattern ```python -from whiteboxai.resilience import CircuitBreaker +from whiteboxxai.resilience import CircuitBreaker circuit_breaker = CircuitBreaker( failure_threshold=5, @@ -468,7 +468,7 @@ try: except CircuitBreaker.CircuitBreakerOpen: # Circuit is open, skip logging prediction = model.predict(inputs) - logger.warning("WhiteBoxAI circuit breaker open, prediction not logged") + logger.warning("WhiteBoxXAI circuit breaker open, prediction not logged") ``` ### 2. Fallback Strategies @@ -497,13 +497,13 @@ class LocalLogger: ### 3. Health Checks ```python -from whiteboxai.health import HealthCheck +from whiteboxxai.health import HealthCheck health_check = HealthCheck(client) # Check API health if not health_check.is_healthy(): - logger.error("WhiteBoxAI API is unhealthy") + logger.error("WhiteBoxXAI API is unhealthy") # Switch to degraded mode # Expose health endpoint for orchestrators @@ -511,7 +511,7 @@ if not health_check.is_healthy(): def health(): return { "status": "healthy" if health_check.is_healthy() else "degraded", - "whiteboxai": health_check.check_api(), + "whiteboxxai": health_check.check_api(), "model": health_check.check_model(model_id) } ``` @@ -524,21 +524,21 @@ def health(): ```python from prometheus_client import Counter, Histogram, Gauge -from whiteboxai.metrics import PrometheusExporter +from whiteboxxai.metrics import PrometheusExporter # Define metrics predictions_logged = Counter( - 'whiteboxai_predictions_logged_total', - 'Total predictions logged to WhiteBoxAI' + 'whiteboxxai_predictions_logged_total', + 'Total predictions logged to WhiteBoxXAI' ) logging_latency = Histogram( - 'whiteboxai_logging_latency_seconds', + 'whiteboxxai_logging_latency_seconds', 'Latency of logging operations' ) api_errors = Counter( - 'whiteboxai_api_errors_total', + 'whiteboxxai_api_errors_total', 'Total API errors', ['error_type'] ) @@ -599,26 +599,26 @@ def log_prediction_with_context(inputs, prediction): ```bash #!/bin/bash -# backup-whiteboxai.sh +# backup-whiteboxxai.sh # Backup configuration DATE=$(date +%Y%m%d_%H%M%S) -BACKUP_DIR="/backups/whiteboxai/${DATE}" +BACKUP_DIR="/backups/whiteboxxai/${DATE}" # Backup database -pg_dump whiteboxai_production > "${BACKUP_DIR}/database.sql" +pg_dump whiteboxxai_production > "${BACKUP_DIR}/database.sql" # Backup model metadata python -c " -from whiteboxai import WhiteBoxAI -client = WhiteBoxAI() +from whiteboxxai import WhiteBoxXAI +client = WhiteBoxXAI() models = client.list_models() with open('${BACKUP_DIR}/models.json', 'w') as f: json.dump(models, f) " # Upload to S3 -aws s3 sync ${BACKUP_DIR} s3://whiteboxai-backups/${DATE} +aws s3 sync ${BACKUP_DIR} s3://whiteboxxai-backups/${DATE} ``` ### Recovery Procedures @@ -687,7 +687,7 @@ class ScalableMonitor: ### Kubernetes Deployment ```yaml -# whiteboxai-monitor.yaml +# whiteboxxai-monitor.yaml apiVersion: apps/v1 kind: Deployment metadata: @@ -706,13 +706,13 @@ spec: - name: ml-service image: your-ml-service:latest env: - - name: WHITEBOXAI_API_KEY + - name: WHITEBOXXAI_API_KEY valueFrom: secretKeyRef: - name: whiteboxai-secrets + name: whiteboxxai-secrets key: api-key - - name: WHITEBOXAI_BASE_URL - value: "https://api.whiteboxai.internal" + - name: WHITEBOXXAI_BASE_URL + value: "https://api.whiteboxxai.internal" resources: requests: memory: "512Mi" @@ -726,10 +726,10 @@ spec: ## Resources -- [WhiteBoxAI API Documentation](https://docs.whiteboxai.com/api) -- [SDK Reference](https://docs.whiteboxai.com/sdk) -- [Best Practices](https://docs.whiteboxai.com/best-practices) -- [Support](mailto:whiteboxai-support@kpmg.com) +- [WhiteBoxXAI API Documentation](https://docs.whiteboxxai.com/api) +- [SDK Reference](https://docs.whiteboxxai.com/sdk) +- [Best Practices](https://docs.whiteboxxai.com/best-practices) +- [Support](mailto:whiteboxxai-support@whiteboxxai.com) --- diff --git a/docs/PYTORCH_INTEGRATION.md b/docs/PYTORCH_INTEGRATION.md index b69798f..c91c61c 100644 --- a/docs/PYTORCH_INTEGRATION.md +++ b/docs/PYTORCH_INTEGRATION.md @@ -1,6 +1,6 @@ # PyTorch Integration Guide -Complete guide for integrating WhiteBoxAI monitoring with PyTorch models. +Complete guide for integrating WhiteBoxXAI monitoring with PyTorch models. ## Table of Contents @@ -21,11 +21,11 @@ Complete guide for integrating WhiteBoxAI monitoring with PyTorch models. ```python import torch import torch.nn as nn -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.pytorch import TorchMonitor -# Initialize WhiteBoxAI -client = WhiteBoxAI(api_key='your-api-key') +# Initialize WhiteBoxXAI +client = WhiteBoxXAI(api_key='your-api-key') # Define your PyTorch model class SimpleNN(nn.Module): @@ -59,7 +59,7 @@ prediction = monitor.predict(X) ## Supported Models -WhiteBoxAI supports all PyTorch model types: +WhiteBoxXAI supports all PyTorch model types: ### Sequential Models @@ -131,7 +131,7 @@ monitor = TorchMonitor( ### Basic Training with Monitoring ```python -from whiteboxai.integrations.pytorch import TrainingMonitor +from whiteboxxai.integrations.pytorch import TrainingMonitor # Create training monitor training_monitor = TrainingMonitor( @@ -208,7 +208,7 @@ training_monitor.log_validation( ### Early Stopping with Drift Detection ```python -from whiteboxai.monitoring import EarlyStopping +from whiteboxxai.monitoring import EarlyStopping early_stopping = EarlyStopping( monitor=training_monitor, @@ -252,7 +252,7 @@ traced_model = torch.jit.trace(model, example_input) # Save traced model torch.jit.save(traced_model, "model_traced.pt") -# Register with WhiteBoxAI +# Register with WhiteBoxXAI monitor = TorchMonitor( model=traced_model, client=client, @@ -270,7 +270,7 @@ scripted_model = torch.jit.script(model) # Save scripted model torch.jit.save(scripted_model, "model_scripted.pt") -# Register with WhiteBoxAI +# Register with WhiteBoxXAI monitor = TorchMonitor( model=scripted_model, client=client, @@ -372,7 +372,7 @@ for data, target in monitored_loader: ### GPU Metrics Tracking ```python -from whiteboxai.monitoring import GPUMonitor +from whiteboxxai.monitoring import GPUMonitor # Enable GPU monitoring gpu_monitor = GPUMonitor() @@ -473,7 +473,7 @@ def train_ddp(rank, world_size): model = SimpleNN().to(rank) ddp_model = DDP(model, device_ids=[rank]) - # Only rank 0 logs to WhiteBoxAI + # Only rank 0 logs to WhiteBoxXAI if rank == 0: monitor = TorchMonitor( model=ddp_model.module, @@ -520,11 +520,11 @@ def save_checkpoint(model, optimizer, epoch, filepath): 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), - 'whiteboxai_model_id': monitor.model_id + 'whiteboxxai_model_id': monitor.model_id } torch.save(checkpoint, filepath) - # Log checkpoint to WhiteBoxAI + # Log checkpoint to WhiteBoxXAI monitor.log_checkpoint( epoch=epoch, checkpoint_path=filepath @@ -534,7 +534,7 @@ def load_checkpoint(filepath, model, optimizer): checkpoint = torch.load(filepath) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) - return checkpoint['epoch'], checkpoint['whiteboxai_model_id'] + return checkpoint['epoch'], checkpoint['whiteboxxai_model_id'] ``` ### 2. Gradient Monitoring @@ -638,8 +638,8 @@ import torch.optim as optim from sklearn.datasets import make_classification from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.pytorch import TorchMonitor # Generate synthetic data X, y = make_classification( @@ -688,8 +688,8 @@ model = BinaryClassifier(input_dim=20) optimizer = optim.Adam(model.parameters(), lr=0.001) criterion = nn.BCELoss() -# Initialize WhiteBoxAI monitor -client = WhiteBoxAI(api_key='your-api-key') +# Initialize WhiteBoxXAI monitor +client = WhiteBoxXAI(api_key='your-api-key') monitor = TorchMonitor( model=model, client=client, @@ -1014,7 +1014,7 @@ monitor = TorchMonitor( # Save only state dict torch.save(model.state_dict(), 'model_weights.pth') -# Register with WhiteBoxAI +# Register with WhiteBoxXAI monitor.register_model( model_path='model_weights.pth', save_weights_only=True @@ -1026,8 +1026,8 @@ monitor.register_model( ## Resources - [PyTorch Documentation](https://pytorch.org/docs/) -- [WhiteBoxAI API Reference](https://docs.whiteboxai.com/api) -- [Best Practices Guide](https://docs.whiteboxai.com/best-practices) +- [WhiteBoxXAI API Reference](https://docs.whiteboxxai.com/api) +- [Best Practices Guide](https://docs.whiteboxxai.com/best-practices) --- diff --git a/docs/SDK_MIGRATION_GUIDE.md b/docs/SDK_MIGRATION_GUIDE.md new file mode 100644 index 0000000..c7da9e0 --- /dev/null +++ b/docs/SDK_MIGRATION_GUIDE.md @@ -0,0 +1,1214 @@ +# WhiteBoxXAI SDK - Separate Repository Migration Guide + +This document provides comprehensive instructions for migrating the WhiteBoxXAI Python SDK to its own repository for easier distribution and better developer experience. + +--- + +## Why Separate the SDK? + +### Benefits of Separate SDK Repository + +#### 1. **Easier Distribution & Installation** +```bash +# Users can install directly from PyPI +pip install whiteboxai + +# Or from GitHub +pip install git+https://github.com/AgentaFlow/whiteboxai-python-sdk.git +``` + +#### 2. **Independent Versioning** +- SDK can have its own semantic versioning (v1.2.3) +- Backend API changes don't force SDK updates +- Users can pin to stable SDK versions while you iterate on backend +- Clear separation of concerns between platform and client library + +#### 3. **Cleaner Developer Experience** +- SDK contributors don't need to clone entire platform +- Smaller repository = faster clones +- Separate issue tracking for SDK vs platform +- Independent CI/CD pipelines +- Focused documentation and examples + +#### 4. **Better Package Management** +- Cleaner package structure for PyPI +- Simplified dependency management +- Better testing isolation +- Industry-standard approach + +#### 5. **PyPI Publishing** +Separate repo makes PyPI publishing cleaner: +```bash +# Build and publish +python -m build +twine upload dist/* +``` + +--- + +## Recommended Repository Structure + +### Main Platform Repository (Current) +``` +whitebox-local-demo/ +├── backend/ # FastAPI application +├── frontend/ # React/Node.js UI +├── docker-compose.yml # Platform deployment +├── README.md # Platform overview +├── docs/ # Platform documentation +├── tests/ # Platform tests +└── infrastructure/ # Infrastructure as Code (Bicep, Azure) +``` + +### SDK Repository (NEW) +``` +whiteboxai-python-sdk/ +├── src/ +│ └── whiteboxai/ +│ ├── __init__.py +│ ├── __version__.py +│ ├── client.py # WhiteBoxAI client +│ ├── monitor.py # ModelMonitor +│ ├── config.py # Configuration +│ ├── decorators.py # @monitor_model, @monitor_prediction +│ ├── privacy.py # PIIDetector, DataMasker +│ ├── offline.py # OfflineQueue, OfflineManager +│ ├── integrations/ +│ │ ├── __init__.py +│ │ ├── sklearn.py # Scikit-learn integration +│ │ ├── pytorch.py # PyTorch integration +│ │ ├── tensorflow.py # TensorFlow/Keras integration +│ │ ├── transformers.py # Hugging Face integration +│ │ ├── langchain.py # LangChain integration +│ │ └── boosting.py # XGBoost/LightGBM integration +│ ├── models/ # Pydantic models +│ │ ├── __init__.py +│ │ ├── prediction.py +│ │ ├── model.py +│ │ └── response.py +│ └── utils/ +│ ├── __init__.py +│ ├── logging.py +│ └── validation.py +├── tests/ +│ ├── __init__.py +│ ├── unit/ +│ │ ├── test_client.py +│ │ ├── test_monitor.py +│ │ ├── test_privacy.py +│ │ └── test_offline.py +│ ├── integration/ +│ │ ├── test_sklearn.py +│ │ ├── test_pytorch.py +│ │ ├── test_tensorflow.py +│ │ └── test_api.py +│ └── e2e/ +│ └── test_workflow.py +├── examples/ +│ ├── README.md +│ ├── basic_monitoring.py +│ ├── sklearn_integration.py +│ ├── pytorch_integration.py +│ ├── tensorflow_integration.py +│ ├── transformers_integration.py +│ ├── langchain_integration.py +│ ├── offline_mode_example.py +│ ├── privacy_example.py +│ └── notebooks/ +│ ├── getting_started.ipynb +│ ├── advanced_monitoring.ipynb +│ └── llm_monitoring.ipynb +├── docs/ +│ ├── index.md +│ ├── getting-started.md +│ ├── installation.md +│ ├── api-reference.md +│ ├── integrations/ +│ │ ├── sklearn.md +│ │ ├── pytorch.md +│ │ ├── tensorflow.md +│ │ ├── transformers.md +│ │ ├── langchain.md +│ │ └── boosting.md +│ ├── features/ +│ │ ├── offline-mode.md +│ │ ├── privacy-filters.md +│ │ ├── decorators.md +│ │ └── async-support.md +│ ├── guides/ +│ │ ├── production-deployment.md +│ │ ├── best-practices.md +│ │ └── troubleshooting.md +│ └── changelog.md +├── .github/ +│ └── workflows/ +│ ├── test.yml # Run tests on PR +│ ├── publish.yml # Publish to PyPI on release +│ ├── docs.yml # Build and deploy docs +│ └── lint.yml # Code quality checks +├── pyproject.toml # Modern Python packaging +├── setup.py # Legacy support +├── README.md # SDK-focused README +├── CHANGELOG.md # Version history +├── LICENSE # MIT License +├── MANIFEST.in # Package includes +├── .gitignore +├── .pre-commit-config.yaml # Pre-commit hooks +└── mkdocs.yml # Documentation config +``` + +--- + +## Migration Strategy + +### Phase 1: Create New Repository + +#### Step 1.1: Create Repository on GitHub +```bash +# On GitHub, create new repository: +# Name: whiteboxxai-python-sdk +# Description: Official Python SDK for WhiteBoxXAI - AI Observability & Explainability Platform +# Visibility: Public +# Initialize with: README, .gitignore (Python), License (MIT) +``` + +#### Step 1.2: Clone and Setup Local Repository +```bash +# Clone the new repository +git clone https://github.com/AgentaFlow/whiteboxxai-python-sdk.git +cd whiteboxxai-python-sdk + +# Create branch for migration +git checkout -b feature/initial-migration +``` + +#### Step 1.3: Copy SDK Files +```bash +# Copy SDK source code +cp -r ../whitebox-local-demo/sdk/whiteboxxai ./src/ + +# Copy examples +cp -r ../whitebox-local-demo/examples/sdk_examples ./examples/ + +# Copy tests +cp -r ../whitebox-local-demo/tests/sdk ./tests/ + +# Copy documentation +cp ../whitebox-local-demo/sdk/README.md ./README.md +cp -r ../whitebox-local-demo/docs/SDK_*.md ./docs/ +``` + +#### Step 1.4: Restructure to Modern Python Layout +```bash +# Create src layout (recommended by PyPA) +mkdir -p src/whiteboxxai +mv whiteboxxai/* src/whiteboxxai/ + +# Create version file +cat > src/whiteboxxai/__version__.py << 'EOF' +"""WhiteBoxXAI SDK version information.""" + +__version__ = "0.1.0" +__author__ = "AgentaFlow" +__email__ = "support@whiteboxxai.com" +__license__ = "MIT" +EOF + +# Update __init__.py to include version +cat >> src/whiteboxxai/__init__.py << 'EOF' + +from .__version__ import __version__, __author__, __email__ + +__all__ = [ + "WhiteBoxXAI", + "ModelMonitor", + "__version__", +] +EOF +``` + +#### Step 1.5: Create Modern pyproject.toml +```bash +cat > pyproject.toml << 'EOF' +[build-system] +requires = ["setuptools>=68.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "whiteboxxai" +version = "0.1.0" +description = "Official Python SDK for WhiteBoxXAI - AI Observability & Explainability Platform" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "MIT"} +authors = [ + {name = "AgentaFlow", email = "support@whiteboxxai.com"} +] +keywords = [ + "ai", + "ml", + "machine-learning", + "observability", + "explainability", + "xai", + "monitoring", + "mlops", + "model-monitoring", + "drift-detection", + "bias-detection" +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", +] + +dependencies = [ + "httpx>=0.24.0", + "pydantic>=2.0.0", + "pydantic-settings>=2.0.0", + "numpy>=1.20.0", + "tenacity>=8.0.0", +] + +[project.optional-dependencies] +sklearn = [ + "scikit-learn>=1.0.0", + "shap>=0.41.0", +] +pytorch = [ + "torch>=1.9.0", +] +tensorflow = [ + "tensorflow>=2.6.0", +] +transformers = [ + "transformers>=4.20.0", + "torch>=1.9.0", +] +langchain = [ + "langchain>=0.1.0", +] +boosting = [ + "xgboost>=1.5.0", + "lightgbm>=3.3.0", +] +all = [ + "whiteboxxai[sklearn,pytorch,tensorflow,transformers,langchain,boosting]", +] +dev = [ + "pytest>=7.0", + "pytest-asyncio>=0.21.0", + "pytest-cov>=4.0.0", + "black>=23.0.0", + "ruff>=0.1.0", + "mypy>=1.0.0", + "pre-commit>=3.0.0", + "build>=0.10.0", + "twine>=4.0.0", +] +docs = [ + "mkdocs>=1.5.0", + "mkdocs-material>=9.0.0", + "mkdocstrings[python]>=0.24.0", +] + +[project.urls] +Homepage = "https://whiteboxxai.com" +Documentation = "https://docs.whiteboxxai.com/sdk" +Repository = "https://github.com/AgentaFlow/whiteboxxai-python-sdk" +Issues = "https://github.com/AgentaFlow/whiteboxxai-python-sdk/issues" +Changelog = "https://github.com/AgentaFlow/whiteboxxai-python-sdk/blob/main/CHANGELOG.md" + +[tool.setuptools] +packages = ["whiteboxxai"] +package-dir = {"" = "src"} + +[tool.setuptools.package-data] +whiteboxxai = ["py.typed"] + +[tool.black] +line-length = 100 +target-version = ['py39', 'py310', 'py311'] +include = '\.pyi?$' + +[tool.ruff] +line-length = 100 +target-version = "py39" +select = ["E", "F", "I", "N", "W", "UP"] +ignore = ["E501"] + +[tool.mypy] +python_version = "3.9" +warn_return_any = true +warn_unused_configs = true +disallow_untyped_defs = true +disallow_incomplete_defs = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = "-v --cov=whiteboxxai --cov-report=term-missing --cov-report=html" + +[tool.coverage.run] +source = ["src/whiteboxxai"] +omit = ["*/tests/*", "*/test_*.py"] + +[tool.coverage.report] +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "raise AssertionError", + "raise NotImplementedError", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", +] +EOF +``` + +#### Step 1.6: Create setup.py (for backward compatibility) +```bash +cat > setup.py << 'EOF' +""" +WhiteBoxXAI Python SDK Setup +For modern builds, use: python -m build +""" +from setuptools import setup + +# All configuration is in pyproject.toml +# This file exists for backward compatibility +setup() +EOF +``` + +#### Step 1.7: Create MANIFEST.in +```bash +cat > MANIFEST.in << 'EOF' +include README.md +include LICENSE +include CHANGELOG.md +include pyproject.toml +include setup.py +recursive-include src/whiteboxxai *.py +recursive-include src/whiteboxxai py.typed +recursive-include tests *.py +recursive-include examples *.py *.ipynb +recursive-include docs *.md +prune tests/__pycache__ +prune src/whiteboxxai/__pycache__ +EOF +``` + +#### Step 1.8: Create CHANGELOG.md +```bash +cat > CHANGELOG.md << 'EOF' +# Changelog + +All notable changes to the WhiteBoxXAI Python SDK will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [0.1.0] - 2026-01-05 + +### Added +- Initial release of WhiteBoxXAI Python SDK +- Core client and model monitoring functionality +- Framework integrations: + - Scikit-learn + - PyTorch + - TensorFlow/Keras + - Hugging Face Transformers + - LangChain + - XGBoost/LightGBM +- Privacy features (PII detection and masking) +- Offline mode with SQLite queue +- Decorator-based monitoring +- Async/sync interfaces +- Comprehensive examples and documentation + +[0.1.0]: https://github.com/AgentaFlow/whiteboxxai-python-sdk/releases/tag/v0.1.0 +EOF +``` + +--- + +### Phase 2: Setup CI/CD + +#### Step 2.1: Create GitHub Actions for Testing +```bash +mkdir -p .github/workflows + +cat > .github/workflows/test.yml << 'EOF' +name: Tests + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + python-version: ['3.9', '3.10', '3.11', '3.12'] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run tests + run: | + pytest --cov=whiteboxxai --cov-report=xml + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v3 + with: + file: ./coverage.xml + flags: unittests + name: codecov-umbrella +EOF +``` + +#### Step 2.2: Create GitHub Actions for Publishing +```bash +cat > .github/workflows/publish.yml << 'EOF' +name: Publish to PyPI + +on: + release: + types: [published] + +jobs: + build-and-publish: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: python -m build + + - name: Check package + run: twine check dist/* + + - name: Publish to Test PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_API_TOKEN }} + run: twine upload --repository testpypi dist/* + + - name: Publish to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: twine upload dist/* +EOF +``` + +#### Step 2.3: Create GitHub Actions for Linting +```bash +cat > .github/workflows/lint.yml << 'EOF' +name: Lint + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + lint: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install black ruff mypy + + - name: Run Black + run: black --check src/ tests/ + + - name: Run Ruff + run: ruff check src/ tests/ + + - name: Run MyPy + run: mypy src/whiteboxxai +EOF +``` + +#### Step 2.4: Create GitHub Actions for Documentation +```bash +cat > .github/workflows/docs.yml << 'EOF' +name: Build Docs + +on: + push: + branches: [ main ] + +jobs: + build-docs: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[docs]" + + - name: Build documentation + run: mkdocs build + + - name: Deploy to GitHub Pages + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./site +EOF +``` + +--- + +### Phase 3: Update Main Platform Repository + +#### Step 3.1: Update Main README +```markdown +# WhiteBoxXAI Platform + +AI Observability & Explainability Platform for production ML models. + +## Components + +- **Platform**: FastAPI backend + React frontend (this repository) +- **Python SDK**: [whiteboxxai-python-sdk](https://github.com/AgentaFlow/whiteboxxai-python-sdk) +- **Documentation**: [docs.whiteboxxai.com](https://docs.whiteboxxai.com) + +## Quick Start + +### 1. Install SDK +```bash +pip install whiteboxxai +``` + +### 2. Monitor Your Models +```python +from whiteboxxai import WhiteBoxXAI, ModelMonitor + +client = WhiteBoxXAI(api_key="your-api-key") +monitor = ModelMonitor(client) + +# Register model +model_id = monitor.register_model( + name="fraud_detection", + model_type="classification" +) + +# Log predictions +monitor.log_prediction( + inputs={"amount": 100.0}, + output={"fraud": False} +) +``` + +### 3. View in Dashboard +Visit http://localhost:3000 to see real-time monitoring, drift detection, and explainability. + +## For Developers + +See [DEVELOPMENT.md](DEVELOPMENT.md) for platform development setup. + +## Documentation + +- **SDK Documentation**: https://github.com/AgentaFlow/whiteboxxai-python-sdk +- **Platform Documentation**: [docs/](docs/) +- **API Reference**: [API_REFERENCE.md](docs/API_REFERENCE.md) + +## License + +MIT License - see [LICENSE](LICENSE) +``` + +#### Step 3.2: Remove SDK Directory from Main Repo +```bash +cd whitebox-local-demo + +# Create branch +git checkout -b feature/remove-sdk-directory + +# Remove SDK directory +git rm -r sdk/ + +# Update .gitignore if needed +echo "" >> .gitignore +echo "# SDK is now in separate repository" >> .gitignore +echo "sdk/" >> .gitignore + +# Commit changes +git add . +git commit -m "Move SDK to separate repository + +SDK has been moved to https://github.com/AgentaFlow/whiteboxxai-python-sdk +for better distribution and independent versioning. + +Users should now install via: pip install whiteboxxai" + +# Push and create PR +git push origin feature/remove-sdk-directory +``` + +#### Step 3.3: Update docker-compose.yml (if SDK is used in tests) +```yaml +# docker-compose.yml +services: + backend: + # ... existing config ... + volumes: + - ./backend:/app/backend + # Remove SDK volume mount if it exists + environment: + # Add if testing against local SDK + - INSTALL_LOCAL_SDK=false +``` + +#### Step 3.4: Update Documentation Links +Update all references to SDK in the main repository: +- `README.md` → Point to SDK repository +- `docs/GETTING_STARTED.md` → Update SDK installation +- `docs/SDK_DOCUMENTATION.md` → Add deprecation notice pointing to SDK repo +- `examples/` → Update import statements and installation instructions + +--- + +### Phase 4: Publishing to PyPI + +#### Step 4.1: Register on PyPI +```bash +# Register accounts (if not already done) +# 1. Go to https://pypi.org/account/register/ +# 2. Go to https://test.pypi.org/account/register/ +# 3. Set up 2FA for security + +# Generate API tokens +# PyPI → Account Settings → API tokens → Add API token +# Scope: Entire account or specific project +# Copy token and save securely +``` + +#### Step 4.2: Test Build Locally +```bash +cd whiteboxxai-python-sdk + +# Install build tools +pip install build twine + +# Build package +python -m build + +# Check build +twine check dist/* + +# Should output: +# Checking dist/whiteboxxai-0.1.0-py3-none-any.whl: PASSED +# Checking dist/whiteboxxai-0.1.0.tar.gz: PASSED +``` + +#### Step 4.3: Publish to Test PyPI +```bash +# Upload to Test PyPI first +twine upload --repository testpypi dist/* + +# Enter username: __token__ +# Enter password: + +# Test installation from Test PyPI +pip install --index-url https://test.pypi.org/simple/ whiteboxxai + +# Test that it works +python -c "from whiteboxxai import WhiteBoxXAI; print('Success!')" +``` + +#### Step 4.4: Publish to Production PyPI +```bash +# Upload to production PyPI +twine upload dist/* + +# Enter username: __token__ +# Enter password: + +# Test installation +pip install whiteboxxai + +# Verify +python -c "from whiteboxxai import WhiteBoxXAI, __version__; print(__version__)" +``` + +#### Step 4.5: Setup GitHub Secrets for Automated Publishing +```bash +# On GitHub repository settings: +# Settings → Secrets and variables → Actions → New repository secret + +# Add secrets: +# - TEST_PYPI_API_TOKEN: +# - PYPI_API_TOKEN: + +# Now releases will automatically publish to PyPI +``` + +--- + +### Phase 5: Documentation Setup + +#### Step 5.1: Create MkDocs Configuration +```bash +cat > mkdocs.yml << 'EOF' +site_name: WhiteBoxXAI Python SDK +site_description: Official Python SDK for WhiteBoxXAI - AI Observability & Explainability +site_author: AgentaFlow +site_url: https://docs.whiteboxxai.com/sdk + +repo_name: AgentaFlow/whiteboxxai-python-sdk +repo_url: https://github.com/AgentaFlow/whiteboxxai-python-sdk + +theme: + name: material + palette: + - scheme: default + primary: indigo + accent: indigo + features: + - navigation.tabs + - navigation.sections + - navigation.expand + - search.suggest + - search.highlight + - content.code.copy + +plugins: + - search + - mkdocstrings: + handlers: + python: + options: + show_source: true + show_root_heading: true + +nav: + - Home: index.md + - Getting Started: + - Installation: getting-started/installation.md + - Quick Start: getting-started/quickstart.md + - Configuration: getting-started/configuration.md + - Integrations: + - Scikit-learn: integrations/sklearn.md + - PyTorch: integrations/pytorch.md + - TensorFlow: integrations/tensorflow.md + - Hugging Face: integrations/transformers.md + - LangChain: integrations/langchain.md + - XGBoost/LightGBM: integrations/boosting.md + - Features: + - Offline Mode: features/offline-mode.md + - Privacy Filters: features/privacy.md + - Decorators: features/decorators.md + - Async Support: features/async.md + - Guides: + - Production Deployment: guides/production.md + - Best Practices: guides/best-practices.md + - Troubleshooting: guides/troubleshooting.md + - API Reference: + - Client: api/client.md + - Monitor: api/monitor.md + - Integrations: api/integrations.md + - Changelog: changelog.md + +markdown_extensions: + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.superfences + - pymdownx.inlinehilite + - pymdownx.snippets + - admonition + - pymdownx.details + - pymdownx.tabbed: + alternate_style: true + - attr_list + - md_in_html +EOF +``` + +#### Step 5.2: Create Documentation Index +```bash +mkdir -p docs/getting-started docs/integrations docs/features docs/guides docs/api + +cat > docs/index.md << 'EOF' +# WhiteBoxXAI Python SDK + +Official Python SDK for integrating WhiteBoxXAI monitoring into your ML applications. + +## Features + +- 🚀 **Easy Integration** - Monitor models with just a few lines of code +- 📊 **Framework Support** - Native integrations for Scikit-learn, PyTorch, TensorFlow, XGBoost, and more +- 🎯 **Decorator-based Monitoring** - Zero-code-change monitoring with decorators +- ⚡ **Async/Sync Interfaces** - Support for both synchronous and asynchronous workflows +- 🔒 **Privacy-First** - Built-in PII detection and data masking +- 💾 **Offline Mode** - Queue predictions when API is unavailable +- 📈 **Drift Detection** - Automatic model and data drift monitoring + +## Quick Start + +```python +from whiteboxxai import WhiteBoxXAI, ModelMonitor + +# Initialize client +client = WhiteBoxXAI(api_key="your-api-key") + +# Create monitor +monitor = ModelMonitor(client) + +# Register model +model_id = monitor.register_model( + name="fraud_detection", + model_type="classification", + framework="sklearn" +) + +# Log predictions +monitor.log_prediction( + inputs={"amount": 100.0, "merchant": "store_123"}, + output={"fraud_probability": 0.15, "prediction": "legitimate"} +) +``` + +## Installation + +```bash +pip install whiteboxxai + +# With specific framework support +pip install whiteboxxai[sklearn] +pip install whiteboxxai[pytorch] +pip install whiteboxxai[all] # All integrations +``` + +## Next Steps + +- [Installation Guide](getting-started/installation.md) +- [Quick Start Tutorial](getting-started/quickstart.md) +- [Framework Integrations](integrations/sklearn.md) +- [API Reference](api/client.md) +EOF +``` + +--- + +### Phase 6: Pre-commit Hooks and Code Quality + +#### Step 6.1: Create .pre-commit-config.yaml +```bash +cat > .pre-commit-config.yaml << 'EOF' +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-added-large-files + - id: check-json + - id: check-toml + - id: check-merge-conflict + - id: debug-statements + + - repo: https://github.com/psf/black + rev: 23.12.1 + hooks: + - id: black + language_version: python3.11 + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.9 + hooks: + - id: ruff + args: [--fix, --exit-non-zero-on-fix] + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.8.0 + hooks: + - id: mypy + additional_dependencies: [pydantic>=2.0.0] +EOF + +# Install pre-commit +pip install pre-commit +pre-commit install +``` + +--- + +## Version Release Process + +### Creating a New Release + +#### 1. Update Version Number +```bash +# Update version in src/whiteboxxai/__version__.py +# Update version in pyproject.toml +# Update CHANGELOG.md with new version +``` + +#### 2. Commit Changes +```bash +git add . +git commit -m "chore: bump version to 0.2.0" +git push origin main +``` + +#### 3. Create Git Tag +```bash +git tag -a v0.2.0 -m "Release v0.2.0" +git push origin v0.2.0 +``` + +#### 4. Create GitHub Release +```bash +# On GitHub: +# Releases → Draft a new release +# Tag: v0.2.0 +# Title: WhiteBoxXAI Python SDK v0.2.0 +# Description: Copy from CHANGELOG.md +# Publish release + +# This will automatically trigger the publish workflow +# and upload to PyPI +``` + +--- + +## Maintenance and Development + +### Local Development Setup + +```bash +# Clone repository +git clone https://github.com/AgentaFlow/whiteboxxai-python-sdk.git +cd whiteboxxai-python-sdk + +# Create virtual environment +python -m venv venv +source venv/bin/activate # On Windows: venv\Scripts\activate + +# Install in editable mode with dev dependencies +pip install -e ".[dev,all]" + +# Install pre-commit hooks +pre-commit install + +# Run tests +pytest + +# Run linting +black src/ tests/ +ruff check src/ tests/ +mypy src/whiteboxxai + +# Build documentation locally +mkdocs serve +# Visit http://localhost:8000 +``` + +### Testing Against Local Platform + +```bash +# Terminal 1: Run platform +cd whitebox-local-demo +docker-compose up + +# Terminal 2: Test SDK +cd whiteboxxai-python-sdk +export WHITEBOXXAI_BASE_URL=http://localhost:8000 +pytest tests/integration/ +``` + +--- + +## Best Practices + +### 1. Semantic Versioning +- **MAJOR** (1.0.0): Breaking API changes +- **MINOR** (0.1.0): New features, backward compatible +- **PATCH** (0.0.1): Bug fixes, backward compatible + +### 2. Changelog Management +- Update CHANGELOG.md for every release +- Use categories: Added, Changed, Deprecated, Removed, Fixed, Security +- Include migration guides for breaking changes + +### 3. Backward Compatibility +- Deprecate before removing features +- Use warnings for deprecated functionality +- Maintain compatibility for at least 2 minor versions + +### 4. Documentation +- Update docs with every feature addition +- Include code examples in docstrings +- Keep README.md concise, detailed docs in MkDocs + +### 5. Testing +- Maintain >80% code coverage +- Test all integrations separately +- Include integration tests with mock API +- Add end-to-end tests for critical workflows + +--- + +## Support and Community + +### Documentation +- **SDK Docs**: https://docs.whiteboxxai.com/sdk +- **Platform Docs**: https://docs.whiteboxxai.com +- **API Reference**: https://api.whiteboxxai.com/docs + +### Community +- **GitHub Issues**: https://github.com/AgentaFlow/whiteboxxai-python-sdk/issues +- **Discussions**: https://github.com/AgentaFlow/whiteboxxai-python-sdk/discussions +- **Discord**: https://discord.gg/whiteboxxai (if available) + +### Contributing +See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines. + +--- + +## Success Metrics + +Track these metrics after migration: + +### Distribution Metrics +- PyPI downloads per month +- GitHub stars and forks +- Issue resolution time +- PR merge time + +### Quality Metrics +- Test coverage percentage +- Number of open issues +- Number of open PRs +- Code quality scores + +### Community Metrics +- Number of contributors +- Community engagement (issues, PRs, discussions) +- Documentation page views +- SDK adoption rate + +--- + +## Comparison: Before vs After + +| Aspect | Before (Monorepo) | After (Separate Repo) | +|--------|-------------------|----------------------| +| Installation | Clone entire platform | `pip install whiteboxxai` | +| Updates | Coupled with platform | Independent versioning | +| Contributing | Need platform knowledge | SDK-focused contributions | +| CI/CD | Shared with platform | Dedicated pipelines | +| Documentation | Mixed with platform | SDK-specific docs | +| Testing | Platform tests | Isolated SDK tests | +| Release Cycle | Tied to platform | Independent releases | +| Repository Size | Large (~500MB+) | Small (~5MB) | +| Clone Time | Minutes | Seconds | +| Focus | Divided | SDK-specific | + +--- + +## Industry Examples + +### Similar Approaches by Leading Companies + +**Stripe**: +- Platform: Multiple microservices +- SDKs: `stripe-python`, `stripe-go`, `stripe-ruby`, `stripe-node` +- Each SDK in separate repository + +**Sentry**: +- Platform: Sentry error tracking service +- SDKs: `sentry-python`, `sentry-javascript`, `sentry-java` +- Independent versioning and release cycles + +**OpenAI**: +- Platform: OpenAI API service +- SDKs: `openai-python`, `openai-node` +- Clean separation of concerns + +**AWS**: +- Platform: AWS Cloud Services +- SDKs: Hundreds of separate SDK repositories +- Language-specific repositories + +--- + +## Conclusion + +Migrating the WhiteBoxXAI Python SDK to its own repository is a strategic decision that will: + +1. ✅ Improve user experience with simpler installation +2. ✅ Enable independent versioning and release cycles +3. ✅ Reduce repository complexity and size +4. ✅ Attract more SDK-focused contributors +5. ✅ Follow industry best practices +6. ✅ Prepare for future multi-language SDK support +7. ✅ Enable professional PyPI distribution +8. ✅ Improve documentation and discoverability + +The migration is straightforward and can be completed in a few days with minimal disruption to users. + +--- + +**Last Updated**: January 5, 2026 +**Version**: 1.0 +**Status**: Ready for Implementation diff --git a/docs/SKLEARN_INTEGRATION.md b/docs/SKLEARN_INTEGRATION.md index 3c74a93..e7ab317 100644 --- a/docs/SKLEARN_INTEGRATION.md +++ b/docs/SKLEARN_INTEGRATION.md @@ -1,6 +1,6 @@ # Scikit-learn Integration Guide -This guide shows you how to integrate WhiteBoxAI monitoring with scikit-learn models. +This guide shows you how to integrate WhiteBoxXAI monitoring with scikit-learn models. ## Table of Contents @@ -18,13 +18,13 @@ This guide shows you how to integrate WhiteBoxAI monitoring with scikit-learn mo ## Quick Start ```python -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.sklearn import SklearnMonitor from sklearn.ensemble import RandomForestClassifier from sklearn.datasets import load_iris # Initialize client -client = WhiteBoxAI(api_key='your-api-key') +client = WhiteBoxXAI(api_key='your-api-key') # Train your model X, y = load_iris(return_X_y=True) @@ -47,7 +47,7 @@ predictions = monitored_model.predict(X) ## Supported Model Types -WhiteBoxAI supports all major scikit-learn model types: +WhiteBoxXAI supports all major scikit-learn model types: ### Classification Models - LogisticRegression @@ -81,10 +81,10 @@ WhiteBoxAI supports all major scikit-learn model types: ## Automatic Model Detection -WhiteBoxAI automatically detects model types and configuration: +WhiteBoxXAI automatically detects model types and configuration: ```python -from whiteboxai.integrations.sklearn import SklearnWrapper +from whiteboxxai.integrations.sklearn import SklearnWrapper # Automatic detection wrapper = SklearnWrapper(client=client) @@ -185,7 +185,7 @@ Monitor models during cross-validation: ```python from sklearn.model_selection import cross_val_score, cross_validate -from whiteboxai.integrations.sklearn import monitor_cv +from whiteboxxai.integrations.sklearn import monitor_cv # Simple cross-validation with monitoring model = RandomForestClassifier() @@ -245,7 +245,7 @@ outer_scores = monitor_cv( Track hyperparameter search experiments: ```python -from whiteboxai.integrations.sklearn import MonitoredGridSearchCV +from whiteboxxai.integrations.sklearn import MonitoredGridSearchCV # Define parameter grid param_grid = { @@ -389,8 +389,8 @@ monitored_v2 = SklearnMonitor( from sklearn.datasets import make_classification from sklearn.model_selection import train_test_split from sklearn.ensemble import RandomForestClassifier -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.sklearn import SklearnMonitor import pandas as pd # Generate data @@ -400,7 +400,7 @@ X_df = pd.DataFrame(X, columns=feature_names) X_train, X_test, y_train, y_test = train_test_split(X_df, y, test_size=0.2) # Initialize client -client = WhiteBoxAI(api_key='your-api-key') +client = WhiteBoxXAI(api_key='your-api-key') # Train model model = RandomForestClassifier(n_estimators=100, random_state=42) @@ -564,10 +564,10 @@ monitored_model = SklearnMonitor( ## Resources -- [WhiteBoxAI SDK Documentation](https://docs.whiteboxai.com/sdk) +- [WhiteBoxXAI SDK Documentation](https://docs.whiteboxxai.com/sdk) - [Scikit-learn Documentation](https://scikit-learn.org) -- [API Reference](https://docs.whiteboxai.com/api) -- [Support](mailto:whiteboxai-support@kpmg.com) +- [API Reference](https://docs.whiteboxxai.com/api) +- [Support](mailto:whiteboxxai-support@whiteboxxai.com) --- diff --git a/docs/TENSORFLOW_INTEGRATION.md b/docs/TENSORFLOW_INTEGRATION.md index 7974e09..328b75a 100644 --- a/docs/TENSORFLOW_INTEGRATION.md +++ b/docs/TENSORFLOW_INTEGRATION.md @@ -1,6 +1,6 @@ # TensorFlow/Keras Integration Guide -Complete guide for integrating WhiteBoxAI monitoring with TensorFlow and Keras models. +Complete guide for integrating WhiteBoxXAI monitoring with TensorFlow and Keras models. ## Table of Contents @@ -22,11 +22,11 @@ Complete guide for integrating WhiteBoxAI monitoring with TensorFlow and Keras m ```python import tensorflow as tf from tensorflow import keras -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.tensorflow import KerasMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.tensorflow import KerasMonitor -# Initialize WhiteBoxAI -client = WhiteBoxAI(api_key='your-api-key') +# Initialize WhiteBoxXAI +client = WhiteBoxXAI(api_key='your-api-key') # Define Keras model model = keras.Sequential([ @@ -111,13 +111,13 @@ monitor = KerasMonitor(model, client, "custom_model") ## Keras Callbacks -### WhiteBoxAI Training Callback +### WhiteBoxXAI Training Callback ```python -from whiteboxai.integrations.tensorflow import WhiteBoxAICallback +from whiteboxxai.integrations.tensorflow import WhiteBoxXAICallback # Create callback -callback = WhiteBoxAICallback( +callback = WhiteBoxXAICallback( monitor=monitor, log_frequency=1, # Log every epoch log_weights=False, @@ -183,7 +183,7 @@ class PredictionCallback(keras.callbacks.Callback): # Make predictions predictions = self.model.predict(self.X_val, verbose=0) - # Log to WhiteBoxAI + # Log to WhiteBoxXAI self.monitor.log_batch( inputs=self.X_val, predictions=predictions, @@ -254,7 +254,7 @@ monitor.set_baseline(X_train, y_train) # Callbacks callbacks = [ - WhiteBoxAICallback(monitor, log_frequency=1), + WhiteBoxXAICallback(monitor, log_frequency=1), keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True), keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5) ] @@ -307,7 +307,7 @@ model.fit( X_train, y_train, epochs=50, callbacks=[ - WhiteBoxAICallback(monitor), + WhiteBoxXAICallback(monitor), lr_schedule, LRLogger(monitor) ] @@ -324,7 +324,7 @@ model.fit( # Save model in SavedModel format model.save('saved_model/my_model') -# Register saved model with WhiteBoxAI +# Register saved model with WhiteBoxXAI monitor.register_saved_model( model_path='saved_model/my_model', metadata={ @@ -415,7 +415,7 @@ model.compile( # Train with monitoring monitor = KerasMonitor(model, client, "mnist_classifier") -callback = WhiteBoxAICallback(monitor) +callback = WhiteBoxXAICallback(monitor) model.fit( ds_train, @@ -496,7 +496,7 @@ model.fit( X_train, y_train, epochs=50, batch_size=32 * strategy.num_replicas_in_sync, - callbacks=[WhiteBoxAICallback(monitor)] + callbacks=[WhiteBoxXAICallback(monitor)] ) ``` @@ -521,7 +521,7 @@ with strategy.scope(): monitor = KerasMonitor(model, client, "tpu_model") # Train -model.fit(ds_train, epochs=10, callbacks=[WhiteBoxAICallback(monitor)]) +model.fit(ds_train, epochs=10, callbacks=[WhiteBoxXAICallback(monitor)]) ``` --- @@ -536,7 +536,7 @@ version = 1 export_path = f'serving/my_model/{version}' model.save(export_path, save_format='tf') -# Register with WhiteBoxAI +# Register with WhiteBoxXAI monitor.register_serving_model( model_path=export_path, version=version, @@ -571,7 +571,7 @@ def predict_with_serving(inputs): predictions = response.json()['predictions'] - # Log to WhiteBoxAI + # Log to WhiteBoxXAI monitor.log_batch( inputs=inputs, predictions=predictions, @@ -619,7 +619,7 @@ model.fit( callbacks=[ checkpoint_cb, CheckpointLogger(monitor), - WhiteBoxAICallback(monitor) + WhiteBoxXAICallback(monitor) ] ) ``` @@ -634,12 +634,12 @@ tensorboard_cb = keras.callbacks.TensorBoard( write_graph=True ) -# Combine with WhiteBoxAI +# Combine with WhiteBoxXAI model.fit( X_train, y_train, epochs=50, callbacks=[ - WhiteBoxAICallback(monitor), + WhiteBoxXAICallback(monitor), tensorboard_cb ] ) @@ -675,7 +675,7 @@ model.compile( # Train with monitoring monitor = KerasMonitor(model, client, "mixed_precision_model") -model.fit(X_train, y_train, callbacks=[WhiteBoxAICallback(monitor)]) +model.fit(X_train, y_train, callbacks=[WhiteBoxXAICallback(monitor)]) ``` ### 4. Data Augmentation @@ -696,7 +696,7 @@ model.fit( datagen.flow(X_train, y_train, batch_size=32), epochs=50, validation_data=(X_val, y_val), - callbacks=[WhiteBoxAICallback(monitor)] + callbacks=[WhiteBoxXAICallback(monitor)] ) ``` @@ -709,8 +709,8 @@ model.fit( ```python import tensorflow as tf from tensorflow import keras -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.tensorflow import KerasMonitor, WhiteBoxAICallback +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.tensorflow import KerasMonitor, WhiteBoxXAICallback # Load CIFAR-10 (X_train, y_train), (X_test, y_test) = keras.datasets.cifar10.load_data() @@ -750,7 +750,7 @@ model.compile( ) # Initialize monitoring -client = WhiteBoxAI(api_key='your-api-key') +client = WhiteBoxXAI(api_key='your-api-key') monitor = KerasMonitor( model=model, client=client, @@ -765,7 +765,7 @@ monitor.set_baseline(X_train[:1000], y_train[:1000]) # Callbacks callbacks = [ - WhiteBoxAICallback(monitor, log_frequency=1), + WhiteBoxXAICallback(monitor, log_frequency=1), keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True), keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5, min_lr=1e-6) ] @@ -863,7 +863,7 @@ history = model.fit( epochs=50, batch_size=32, validation_split=0.2, - callbacks=[WhiteBoxAICallback(monitor)], + callbacks=[WhiteBoxXAICallback(monitor)], verbose=1 ) @@ -888,7 +888,7 @@ print(f"LSTM model ID: {monitor.model_id}") ```python # Correct -model.fit(X, y, callbacks=[WhiteBoxAICallback(monitor)]) +model.fit(X, y, callbacks=[WhiteBoxXAICallback(monitor)]) # Incorrect (callback not passed) model.fit(X, y) @@ -924,8 +924,8 @@ model.save('model_dir') - [TensorFlow Documentation](https://www.tensorflow.org/api_docs) - [Keras Guide](https://keras.io/guides/) -- [WhiteBoxAI API Reference](https://docs.whiteboxai.com/api) -- [Best Practices](https://docs.whiteboxai.com/best-practices) +- [WhiteBoxXAI API Reference](https://docs.whiteboxxai.com/api) +- [Best Practices](https://docs.whiteboxxai.com/best-practices) --- diff --git a/docs/api-reference.md b/docs/api-reference.md index 2f53317..7284b6c 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -1,21 +1,21 @@ # API Reference -Complete API documentation for WhiteBoxAI SDK. +Complete API documentation for WhiteBoxXAI SDK. -## WhiteBoxAI Client +## WhiteBoxXAI Client -Main client for interacting with the WhiteBoxAI API. +Main client for interacting with the WhiteBoxXAI API. ### Constructor ```python -WhiteBoxAI( +WhiteBoxXAI( api_key: str = None, - base_url: str = "https://api.whiteboxai.io", + base_url: str = "https://api.whiteboxxai.com", timeout: int = 30, max_retries: int = 3, enable_offline: bool = False, - offline_dir: str = "./whiteboxai_offline", + offline_dir: str = "./whiteboxxai_offline", offline_max_queue_size: int = 10000, offline_auto_sync: bool = False, offline_sync_interval: int = 60, @@ -29,7 +29,7 @@ WhiteBoxAI( ``` **Parameters:** -- `api_key` (str): API key for authentication. Can be set via EXPLAINAI_API_KEY env var +- `api_key` (str): API key for authentication. Can be set via WHITEBOXXAI_API_KEY env var - `base_url` (str): Base URL for API endpoint - `timeout` (int): Request timeout in seconds - `max_retries` (int): Maximum number of retry attempts @@ -86,7 +86,7 @@ Simplified monitoring interface. ```python ModelMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, model_id: int = None, sampling_rate: float = 1.0, enable_explanations: bool = False @@ -188,10 +188,10 @@ Monitor predictions with custom extractors. ### SklearnMonitor ```python -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai.integrations.sklearn import SklearnMonitor monitor = SklearnMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, model=None, model_name: str = None ) @@ -204,10 +204,10 @@ monitor.predict(model, X, y=None, log=True) ### TorchMonitor ```python -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai.integrations.pytorch import TorchMonitor monitor = TorchMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, model=None, model_name: str = None ) @@ -219,26 +219,26 @@ monitor.wrap_model(model) ### KerasMonitor ```python -from whiteboxai.integrations.tensorflow import KerasMonitor, WhiteBoxAICallback +from whiteboxxai.integrations.tensorflow import KerasMonitor, WhiteBoxXAICallback monitor = KerasMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, model=None, model_name: str = None ) monitor.register_from_model(model_type: str, **kwargs) -> int -callback = WhiteBoxAICallback(monitor, log_frequency=1) +callback = WhiteBoxXAICallback(monitor, log_frequency=1) monitor.predict(X, log=True) ``` ### TransformersMonitor ```python -from whiteboxai.integrations.transformers import TransformersMonitor +from whiteboxxai.integrations.transformers import TransformersMonitor monitor = TransformersMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, pipeline=None, model_name: str = None ) @@ -250,10 +250,10 @@ monitor.predict(text, log=True) ### LangChainMonitor ```python -from whiteboxai.integrations.langchain import LangChainMonitor +from whiteboxxai.integrations.langchain import LangChainMonitor monitor = LangChainMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, application_name: str, track_tokens: bool = True, track_cost: bool = True @@ -266,10 +266,10 @@ callback = monitor.create_callback_handler() ### XGBoostMonitor / LightGBMMonitor ```python -from whiteboxai.integrations.boosting import XGBoostMonitor, LightGBMMonitor +from whiteboxxai.integrations.boosting import XGBoostMonitor, LightGBMMonitor monitor = XGBoostMonitor( - client: WhiteBoxAI, + client: WhiteBoxXAI, model_name: str, track_feature_importance: bool = True, importance_type: str = "gain" @@ -282,8 +282,8 @@ monitor.predict(model, X, y=None, log=True) ## Exceptions ```python -from whiteboxai.exceptions import ( - WhiteBoxAIError, # Base exception +from whiteboxxai.exceptions import ( + WhiteBoxXAIError, # Base exception AuthenticationError, # Authentication failed APIError, # API request error ValidationError, # Validation error @@ -294,7 +294,7 @@ from whiteboxai.exceptions import ( ## Privacy ```python -from whiteboxai.privacy import mask_data +from whiteboxxai.privacy import mask_data masked = mask_data( data: dict, diff --git a/docs/getting-started.md b/docs/getting-started.md index 11b8e31..c9cebad 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -1,11 +1,11 @@ -# Getting Started with WhiteBoxAI SDK +# Getting Started with WhiteBoxXAI SDK ## Installation -Install the WhiteBoxAI SDK using pip: +Install the WhiteBoxXAI SDK using pip: ```bash -pip install whiteboxai-sdk +pip install whitebox-xai-sdk ``` ### Optional Dependencies @@ -14,16 +14,16 @@ Install with specific framework support: ```bash # Scikit-learn support -pip install whiteboxai-sdk[sklearn] +pip install whitebox-xai-sdk[sklearn] # PyTorch support -pip install whiteboxai-sdk[pytorch] +pip install whitebox-xai-sdk[pytorch] # TensorFlow support -pip install whiteboxai-sdk[tensorflow] +pip install whitebox-xai-sdk[tensorflow] # All integrations -pip install whiteboxai-sdk[all] +pip install whitebox-xai-sdk[all] ``` ## Quick Start @@ -31,21 +31,21 @@ pip install whiteboxai-sdk[all] ### 1. Initialize the Client ```python -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") ``` You can also set the API key via environment variable: ```bash -export EXPLAINAI_API_KEY=your-api-key +export WHITEBOXXAI_API_KEY=your-api-key ``` ### 2. Create a Monitor ```python -from whiteboxai import ModelMonitor +from whiteboxxai import ModelMonitor monitor = ModelMonitor(client) ``` @@ -81,9 +81,9 @@ monitor.log_prediction( Configure the client with various options: ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", - base_url="https://api.whiteboxai.io", + base_url="https://api.whiteboxxai.com", timeout=30, max_retries=3, enable_caching=True, @@ -95,5 +95,5 @@ client = WhiteBoxAI( ## Support For help and support: -- Documentation: https://whitebox.agentaflow.com +- Documentation: https://docs.whiteboxxai.com - Issues: https://github.com/AgentaFlow/whitebox-python-sdk/issues diff --git a/docs/index.md b/docs/index.md index 1f43b70..e9b7bbc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,6 @@ -# WhiteBoxAI Python SDK +# WhiteBoxXAI Python SDK -Official Python SDK for integrating WhiteBoxAI monitoring into your ML applications. +Official Python SDK for integrating WhiteBoxXAI monitoring into your ML applications. ## Features @@ -18,14 +18,14 @@ Official Python SDK for integrating WhiteBoxAI monitoring into your ML applicati ## Installation ```bash -pip install whiteboxai-sdk +pip install whitebox-xai-sdk # With specific framework support -pip install whiteboxai-sdk[sklearn] -pip install whiteboxai-sdk[pytorch] -pip install whiteboxai-sdk[langchain] -pip install whiteboxai-sdk[crewai] -pip install whiteboxai-sdk[all] # All integrations +pip install whitebox-xai-sdk[sklearn] +pip install whitebox-xai-sdk[pytorch] +pip install whitebox-xai-sdk[langchain] +pip install whitebox-xai-sdk[crewai] +pip install whitebox-xai-sdk[all] # All integrations ``` ## Quick Start @@ -33,10 +33,10 @@ pip install whiteboxai-sdk[all] # All integrations ### Basic Usage ```python -from whiteboxai import WhiteBoxAI, ModelMonitor +from whiteboxxai import WhiteBoxXAI, ModelMonitor # Initialize client -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") # Create monitor monitor = ModelMonitor(client) @@ -58,13 +58,13 @@ monitor.log_prediction( ### Git Integration ```python -from whiteboxai import WhiteBoxAI, detect_git_context +from whiteboxxai import WhiteBoxXAI, detect_git_context # Auto-detect Git context git_context = detect_git_context() # Initialize with Git context -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") model_id = client.models.register( name="my_model", **git_context.to_dict() # Include Git metadata @@ -74,7 +74,7 @@ model_id = client.models.register( ### CrewAI Multi-Agent Monitoring ```python -from whiteboxai.integrations import CrewAIMonitor +from whiteboxxai.integrations import CrewAIMonitor from crewai import Agent, Task, Crew # Initialize monitor @@ -99,7 +99,7 @@ summary = monitor.complete_monitoring(outputs={"result": result}) ### LangChain Multi-Agent Monitoring ```python -from whiteboxai.integrations import LangGraphMultiAgentMonitor +from whiteboxxai.integrations import LangGraphMultiAgentMonitor # Create monitor monitor = LangGraphMultiAgentMonitor( @@ -128,7 +128,7 @@ summary = monitor.complete_monitoring(outputs={"result": result}) ### Scikit-learn ```python -from whiteboxai.integrations import SklearnMonitor +from whiteboxxai.integrations import SklearnMonitor from sklearn.ensemble import RandomForestClassifier # Wrap your model @@ -144,7 +144,7 @@ predictions = wrapped_model.predict(X_test) ### PyTorch ```python -from whiteboxai.integrations import TorchMonitor +from whiteboxxai.integrations import TorchMonitor import torch.nn as nn # Monitor your model @@ -160,7 +160,7 @@ for epoch in range(num_epochs): ### TensorFlow/Keras ```python -from whiteboxai.integrations import KerasMonitor +from whiteboxxai.integrations import KerasMonitor # Add callback monitor = KerasMonitor(client=client, model_id=model_id) @@ -174,7 +174,7 @@ model.fit( ### LangChain ```python -from whiteboxai.integrations import LangChainMonitor +from whiteboxxai.integrations import LangChainMonitor # Monitor chain execution monitor = LangChainMonitor(client=client) diff --git a/docs/integrations.md b/docs/integrations.md index 23b2621..534131a 100644 --- a/docs/integrations.md +++ b/docs/integrations.md @@ -1,6 +1,6 @@ # Framework Integrations -WhiteBoxAI SDK provides native integrations for popular ML frameworks. +WhiteBoxXAI SDK provides native integrations for popular ML frameworks. ## Scikit-learn @@ -8,15 +8,15 @@ Monitor scikit-learn models with automatic feature extraction: ```python from sklearn.ensemble import RandomForestClassifier -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.sklearn import SklearnMonitor # Train your model model = RandomForestClassifier() model.fit(X_train, y_train) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = SklearnMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -33,14 +33,14 @@ Monitor PyTorch models with tensor tracking: ```python import torch.nn as nn -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.pytorch import TorchMonitor # Define your model model = nn.Sequential(...) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = TorchMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -56,19 +56,19 @@ Monitor TensorFlow models with training callbacks: ```python from tensorflow import keras -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.tensorflow import KerasMonitor, WhiteBoxAICallback +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.tensorflow import KerasMonitor, WhiteBoxXAICallback # Build your model model = keras.Sequential([...]) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = KerasMonitor(client, model=model, model_name="keras_model") monitor.register_from_model(model_type="regression") # Train with monitoring -callback = WhiteBoxAICallback(monitor, log_frequency=1) +callback = WhiteBoxXAICallback(monitor, log_frequency=1) model.fit(X_train, y_train, callbacks=[callback], epochs=50) ``` @@ -80,14 +80,14 @@ Monitor transformers pipelines and models: ```python from transformers import pipeline -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.transformers import TransformersMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.transformers import TransformersMonitor # Load model classifier = pipeline("sentiment-analysis") # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = TransformersMonitor( client=client, pipeline=classifier, @@ -106,11 +106,11 @@ Monitor LangChain applications and chains: ```python from langchain.chains import LLMChain -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.langchain import LangChainMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.langchain import LangChainMonitor # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = LangChainMonitor( client=client, application_name="qa_bot", @@ -131,15 +131,15 @@ Monitor gradient boosting models: ```python import xgboost as xgb -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.boosting import XGBoostMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.boosting import XGBoostMonitor # Train model model = xgb.XGBClassifier() model.fit(X_train, y_train) # Setup monitoring -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = XGBoostMonitor( client=client, model_name="fraud_detector", @@ -155,7 +155,7 @@ predictions = monitor.predict(model, X_test, y_test) Create custom integrations by extending the base `ModelMonitor` class: ```python -from whiteboxai import ModelMonitor +from whiteboxxai import ModelMonitor class CustomMonitor(ModelMonitor): def __init__(self, client, model, **kwargs): diff --git a/docs/offline-mode.md b/docs/offline-mode.md index b30d4fb..89b69fc 100644 --- a/docs/offline-mode.md +++ b/docs/offline-mode.md @@ -1,6 +1,6 @@ # Offline Mode -WhiteBoxAI SDK supports offline mode for robust operation when network connectivity is unreliable. +WhiteBoxXAI SDK supports offline mode for robust operation when network connectivity is unreliable. ## Overview @@ -14,13 +14,13 @@ Offline mode provides: ## Quick Start ```python -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI # Enable offline mode with auto-sync -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_offline=True, - offline_dir="./whiteboxai_offline", + offline_dir="./whiteboxxai_offline", offline_auto_sync=True, offline_sync_interval=60 ) @@ -31,7 +31,7 @@ client = WhiteBoxAI( ### Basic Configuration ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_offline=True, # Enable offline mode offline_dir="./offline_queue", # Storage directory @@ -42,7 +42,7 @@ client = WhiteBoxAI( ### Auto-Sync Configuration ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key="your-api-key", enable_offline=True, offline_auto_sync=True, # Enable background sync @@ -123,10 +123,10 @@ except Exception as e: For production deployments: ```python -client = WhiteBoxAI( +client = WhiteBoxXAI( api_key=os.getenv("WHITEBOXAI_API_KEY"), enable_offline=True, - offline_dir="/var/lib/whiteboxai/queue", + offline_dir="/var/lib/whiteboxxai/queue", offline_max_queue_size=50000, offline_auto_sync=True, offline_sync_interval=30, # More frequent sync in production diff --git a/examples/async_monitoring.py b/examples/async_monitoring.py index 85c9c82..28346b2 100644 --- a/examples/async_monitoring.py +++ b/examples/async_monitoring.py @@ -7,14 +7,13 @@ import asyncio import numpy as np - -from whiteboxai import ModelMonitor, WhiteBoxAI +from whiteboxxai import ModelMonitor, WhiteBoxXAI async def register_and_log(): """Register model and log predictions asynchronously.""" # Create client with async context manager - async with WhiteBoxAI(api_key="your-api-key") as client: + async with WhiteBoxXAI(api_key="your-api-key") as client: monitor = ModelMonitor(client) # Register model asynchronously @@ -38,7 +37,8 @@ async def register_and_log(): # Log batch predictions print("\nLogging batch predictions...") predictions = [ - {"inputs": {"amount": 50.0}, "output": {"fraud_prob": 0.05}} for _ in range(100) + {"inputs": {"amount": 50.0}, "output": {"fraud_prob": 0.05}} + for _ in range(100) ] await monitor.alog_batch(predictions) @@ -52,7 +52,7 @@ async def register_and_log(): async def parallel_logging(): """Log predictions in parallel for better throughput.""" - async with WhiteBoxAI(api_key="your-api-key") as client: + async with WhiteBoxXAI(api_key="your-api-key") as client: monitor = ModelMonitor(client, model_id=123) # Create multiple prediction logging tasks @@ -71,7 +71,7 @@ async def parallel_logging(): async def drift_detection(): """Detect drift asynchronously.""" - async with WhiteBoxAI(api_key="your-api-key") as client: + async with WhiteBoxXAI(api_key="your-api-key") as client: monitor = ModelMonitor(client, model_id=123) # Set baseline diff --git a/examples/basic_monitoring.py b/examples/basic_monitoring.py index 1a107b6..13de9bf 100644 --- a/examples/basic_monitoring.py +++ b/examples/basic_monitoring.py @@ -4,12 +4,12 @@ This example demonstrates basic model registration and prediction logging. """ -from whiteboxai import ModelMonitor, WhiteBoxAI +from whiteboxxai import ModelMonitor, WhiteBoxXAI def main(): # Initialize client with API key - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") # Create monitor instance monitor = ModelMonitor(client) diff --git a/examples/boosting_example.py b/examples/boosting_example.py index 75d8674..b307e70 100644 --- a/examples/boosting_example.py +++ b/examples/boosting_example.py @@ -2,7 +2,7 @@ XGBoost and LightGBM Integration Examples This script demonstrates how to monitor gradient boosting models -(XGBoost and LightGBM) using WhiteBoxAI. +(XGBoost and LightGBM) using WhiteBoxXAI. Examples include: 1. XGBoost binary classification @@ -18,9 +18,9 @@ from sklearn.metrics import accuracy_score, mean_squared_error, r2_score from sklearn.model_selection import train_test_split -# WhiteBoxAI imports -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.boosting import ( +# WhiteBoxXAI imports +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.boosting import ( LightGBMMonitor, XGBoostMonitor, wrap_lightgbm_model, @@ -46,10 +46,12 @@ def example_xgboost_classification(): X, y = make_classification( n_samples=1000, n_features=20, n_informative=15, n_redundant=5, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Create monitor monitor = XGBoostMonitor( @@ -61,11 +63,13 @@ def example_xgboost_classification(): # Train XGBoost model print("Training XGBoost classifier...") - model = xgb.XGBClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) + model = xgb.XGBClassifier( + n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 + ) model.fit(X_train, y_train) - # Register model with WhiteBoxAI - print("Registering model with WhiteBoxAI...") + # Register model with WhiteBoxXAI + print("Registering model with WhiteBoxXAI...") model_id = monitor.register_from_model( model=model, X_train=X_train, @@ -85,13 +89,13 @@ def example_xgboost_classification(): # Calculate metrics accuracy = accuracy_score(y_test, predictions) print(f"Accuracy: {accuracy:.4f}") - print("Predictions logged to WhiteBoxAI") + print(f"Predictions logged to WhiteBoxXAI") # Get feature importance importance = model.feature_importances_ top_features = np.argsort(importance)[-5:][::-1] print(f"\nTop 5 features: {top_features.tolist()}") - print("Feature importance tracked in metadata") + print(f"Feature importance tracked in metadata") def example_xgboost_regression(): @@ -112,19 +116,25 @@ def example_xgboost_regression(): X, y = make_regression( n_samples=1000, n_features=10, n_informative=8, noise=10, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Create monitor monitor = XGBoostMonitor( - client=client, model_name="xgboost_price_predictor", track_feature_importance=True + client=client, + model_name="xgboost_price_predictor", + track_feature_importance=True, ) # Train XGBoost regressor print("Training XGBoost regressor...") - model = xgb.XGBRegressor(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) + model = xgb.XGBRegressor( + n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 + ) model.fit(X_train, y_train) # Wrap model for automatic monitoring @@ -140,7 +150,7 @@ def example_xgboost_regression(): r2 = r2_score(y_test, predictions) print(f"MSE: {mse:.4f}") print(f"R² Score: {r2:.4f}") - print("Predictions automatically logged via wrapper") + print(f"Predictions automatically logged via wrapper") def example_lightgbm_classification(): @@ -161,10 +171,12 @@ def example_lightgbm_classification(): X, y = make_classification( n_samples=1000, n_features=20, n_informative=15, n_redundant=5, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Create monitor monitor = LightGBMMonitor( @@ -176,11 +188,13 @@ def example_lightgbm_classification(): # Train LightGBM model print("Training LightGBM classifier...") - model = lgb.LGBMClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) + model = lgb.LGBMClassifier( + n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 + ) model.fit(X_train, y_train) - # Register model with WhiteBoxAI - print("Registering model with WhiteBoxAI...") + # Register model with WhiteBoxXAI + print("Registering model with WhiteBoxXAI...") model_id = monitor.register_from_model( model=model, X_train=X_train, @@ -197,11 +211,13 @@ def example_lightgbm_classification(): print("\nMaking predictions...") predictions = monitor.predict(model, X_test, y_test) - # Get probabilities and calculate metrics + # Get probabilities probabilities = model.predict_proba(X_test) + + # Calculate metrics accuracy = accuracy_score(y_test, predictions) print(f"Accuracy: {accuracy:.4f}") - print(f"Predictions logged with {len(probabilities)} probability estimates") + print(f"Predictions logged with probabilities") # Feature importance importance = model.feature_importances_ @@ -227,24 +243,32 @@ def example_lightgbm_regression(): X, y = make_regression( n_samples=1000, n_features=10, n_informative=8, noise=10, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Create monitor monitor = LightGBMMonitor( - client=client, model_name="lightgbm_sales_predictor", track_feature_importance=True + client=client, + model_name="lightgbm_sales_predictor", + track_feature_importance=True, ) # Train LightGBM regressor print("Training LightGBM regressor...") - model = lgb.LGBMRegressor(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) + model = lgb.LGBMRegressor( + n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 + ) model.fit(X_train, y_train) # Wrap model for automatic monitoring print("Wrapping model for automatic monitoring...") - wrapped_model = wrap_lightgbm_model(model=model, monitor=monitor, auto_register=True) + wrapped_model = wrap_lightgbm_model( + model=model, monitor=monitor, auto_register=True + ) # Predictions automatically logged print("\nMaking predictions (auto-logged)...") @@ -255,7 +279,7 @@ def example_lightgbm_regression(): r2 = r2_score(y_test, predictions) print(f"MSE: {mse:.4f}") print(f"R² Score: {r2:.4f}") - print("Predictions automatically logged via wrapper") + print(f"Predictions automatically logged via wrapper") def example_feature_importance_tracking(): @@ -276,16 +300,20 @@ def example_feature_importance_tracking(): # Generate synthetic data with feature names import pandas as pd - X, y = make_classification(n_samples=1000, n_features=10, n_informative=7, random_state=42) + X, y = make_classification( + n_samples=1000, n_features=10, n_informative=7, random_state=42 + ) # Create DataFrame with feature names feature_names = [f"feature_{i}" for i in range(10)] X_df = pd.DataFrame(X, columns=feature_names) - X_train, X_test, y_train, y_test = train_test_split(X_df, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X_df, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Train XGBoost model print("Training XGBoost model...") @@ -305,7 +333,9 @@ def example_feature_importance_tracking(): # Get importance importance_dict = monitor._get_feature_importance(xgb_model) if importance_dict: - sorted_features = sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:5] + sorted_features = sorted( + importance_dict.items(), key=lambda x: x[1], reverse=True + )[:5] for feat, score in sorted_features: print(f" {feat}: {score:.4f}") @@ -327,7 +357,9 @@ def example_feature_importance_tracking(): # Get importance importance_dict = monitor._get_feature_importance(lgb_model) if importance_dict: - sorted_features = sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:5] + sorted_features = sorted( + importance_dict.items(), key=lambda x: x[1], reverse=True + )[:5] for feat, score in sorted_features: print(f" {feat}: {score:.4f}") @@ -348,15 +380,21 @@ def example_model_comparison(): return # Generate synthetic data - X, y = make_classification(n_samples=1000, n_features=20, n_informative=15, random_state=42) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X, y = make_classification( + n_samples=1000, n_features=20, n_informative=15, random_state=42 + ) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="demo-api-key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="demo-api-key") # Train and monitor XGBoost print("Training XGBoost model...") - xgb_model = xgb.XGBClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) + xgb_model = xgb.XGBClassifier( + n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 + ) xgb_model.fit(X_train, y_train) xgb_monitor = XGBoostMonitor(client=client, model_name="xgb_comparison") @@ -383,12 +421,12 @@ def example_model_comparison(): print(f"XGBoost Accuracy: {xgb_accuracy:.4f}") print(f"LightGBM Accuracy: {lgb_accuracy:.4f}") print(f"\nBetter model: {'XGBoost' if xgb_accuracy > lgb_accuracy else 'LightGBM'}") - print("Both models logged to WhiteBoxAI for detailed analysis") + print("Both models logged to WhiteBoxXAI for detailed analysis") if __name__ == "__main__": print("\n" + "=" * 60) - print("WhiteBoxAI - Gradient Boosting Integration Examples") + print("WhiteBoxXAI - Gradient Boosting Integration Examples") print("=" * 60) # Run examples diff --git a/examples/decorator_monitoring.py b/examples/decorator_monitoring.py index 27a19de..22765d2 100644 --- a/examples/decorator_monitoring.py +++ b/examples/decorator_monitoring.py @@ -6,11 +6,10 @@ import numpy as np from sklearn.ensemble import RandomForestClassifier - -from whiteboxai import ModelMonitor, WhiteBoxAI, monitor_model, monitor_prediction +from whiteboxxai import ModelMonitor, WhiteBoxXAI, monitor_model, monitor_prediction # Global monitor instance -client = WhiteBoxAI(api_key="your-api-key") +client = WhiteBoxXAI(api_key="your-api-key") monitor = ModelMonitor(client, model_id=123) @@ -18,6 +17,8 @@ def predict_fraud(features): """Predict fraud probability.""" # Simulate model prediction + model = RandomForestClassifier() + # ... (assume model is trained) prediction = np.random.choice([0, 1]) probability = np.random.random() @@ -68,7 +69,9 @@ def main(): print("\n=== Custom Extractors ===") # Custom input/output extraction - result = score_transaction(data={"amount": 100.0, "velocity": 5.0, "location_risk": 0.3}) + result = score_transaction( + data={"amount": 100.0, "velocity": 5.0, "location_risk": 0.3} + ) print(f"Transaction score: {result}") print("\n=== Class Method Decorator ===") diff --git a/examples/langchain_example.py b/examples/langchain_example.py index 208c38b..a85162d 100644 --- a/examples/langchain_example.py +++ b/examples/langchain_example.py @@ -1,17 +1,45 @@ """ LangChain Integration Example -This example demonstrates how to use WhiteBoxAI to monitor LangChain applications. +This example demonstrates how to use WhiteBoxXAI to monitor LangChain applications. """ +import ast +import operator import os import time -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.langchain import LangChainMonitor, wrap_langchain_chain +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.langchain import LangChainMonitor, wrap_langchain_chain # Optional: Set API key -os.environ["WHITEBOXAI_API_KEY"] = "your-api-key-here" +os.environ["WHITEBOXXAI_API_KEY"] = "your-api-key-here" + +_SAFE_OPERATORS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.Pow: operator.pow, + ast.USub: operator.neg, +} + + +def safe_eval_arithmetic(expression: str) -> float: + """Evaluate a simple arithmetic expression without eval()'s ability to run + arbitrary code - important here since a tool like this may be called with + an LLM-influenced (and therefore untrusted) argument.""" + + def _eval(node): + if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)): + return node.value + if isinstance(node, ast.BinOp) and type(node.op) in _SAFE_OPERATORS: + return _SAFE_OPERATORS[type(node.op)](_eval(node.left), _eval(node.right)) + if isinstance(node, ast.UnaryOp) and type(node.op) in _SAFE_OPERATORS: + return _SAFE_OPERATORS[type(node.op)](_eval(node.operand)) + raise ValueError(f"Unsupported expression: {expression}") + + return _eval(ast.parse(expression, mode="eval").body) def example_simple_chain(): @@ -24,24 +52,30 @@ def example_simple_chain(): print("Simple LLM Chain Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Create monitor monitor = LangChainMonitor( - client=client, application_name="simple_qa_chain", track_tokens=True, track_cost=True + client=client, + application_name="simple_qa_chain", + track_tokens=True, + track_cost=True, ) # Register application app_id = monitor.register_application( - name="Simple Q&A Chain", version="1.0.0", description="Basic question-answering chain" + name="Simple Q&A Chain", + version="1.0.0", + description="Basic question-answering chain", ) print(f"✓ Application registered with ID: {app_id}") # Create chain llm = OpenAI(temperature=0.7) prompt = PromptTemplate( - input_variables=["question"], template="Answer the following question: {question}" + input_variables=["question"], + template="Answer the following question: {question}", ) chain = LLMChain(llm=llm, prompt=prompt) @@ -74,8 +108,8 @@ def example_sequential_chain(): print("Sequential Chain Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Create monitor monitor = LangChainMonitor( @@ -92,7 +126,8 @@ def example_sequential_chain(): # Chain 1: Generate topic prompt1 = PromptTemplate( - input_variables=["subject"], template="Generate a creative topic about {subject}" + input_variables=["subject"], + template="Generate a creative topic about {subject}", ) chain1 = LLMChain(llm=llm, prompt=prompt1, output_key="topic") @@ -116,7 +151,7 @@ def example_sequential_chain(): print("\nRunning sequential chain...") result = wrapped_chain({"subject": "artificial intelligence"}) - print("\n Subject: artificial intelligence") + print(f"\n Subject: artificial intelligence") print(f" Topic: {result['topic'].strip()}") print(f" Paragraph: {result['paragraph'].strip()[:100]}...") @@ -127,13 +162,14 @@ def example_agent(): """Example using an agent with tools.""" from langchain.agents import AgentType, Tool, initialize_agent from langchain.llms import OpenAI + from langchain.utilities import SerpAPIWrapper print("\n" + "=" * 60) print("Agent Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Create monitor monitor = LangChainMonitor( @@ -143,7 +179,9 @@ def example_agent(): # Register application monitor.register_application( - name="Search Agent", version="1.0.0", description="Agent with search capabilities" + name="Search Agent", + version="1.0.0", + description="Agent with search capabilities", ) print("✓ Application registered") @@ -155,7 +193,7 @@ def search_tool(query: str) -> str: def calculator_tool(expression: str) -> str: """Mock calculator tool.""" try: - return str(eval(expression)) + return str(safe_eval_arithmetic(expression)) except Exception: return "Invalid expression" @@ -190,13 +228,18 @@ def calculator_tool(expression: str) -> str: def example_rag_chain(): """Example using a RAG (Retrieval-Augmented Generation) chain.""" + from langchain.chains import RetrievalQA + from langchain.embeddings import OpenAIEmbeddings + from langchain.llms import OpenAI + from langchain.text_splitter import CharacterTextSplitter + from langchain.vectorstores import FAISS print("\n" + "=" * 60) print("RAG Chain Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Create monitor monitor = LangChainMonitor( @@ -206,7 +249,9 @@ def example_rag_chain(): # Register application monitor.register_application( - name="RAG Q&A System", version="1.0.0", description="Question answering with retrieval" + name="RAG Q&A System", + version="1.0.0", + description="Question answering with retrieval", ) print("✓ Application registered") @@ -270,8 +315,8 @@ def example_manual_logging(): print("Manual Logging Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Create monitor monitor = LangChainMonitor( @@ -322,7 +367,7 @@ def example_manual_logging(): def main(): """Run all examples.""" print("\n" + "=" * 60) - print("WhiteBoxAI - LangChain Integration Examples") + print("WhiteBoxXAI - LangChain Integration Examples") print("=" * 60) try: diff --git a/examples/offline_mode_example.py b/examples/offline_mode_example.py index e6a7297..41a0278 100644 --- a/examples/offline_mode_example.py +++ b/examples/offline_mode_example.py @@ -1,18 +1,20 @@ """ -Offline Mode Examples for WhiteBoxAI SDK +Offline Mode Examples for WhiteBoxXAI SDK This example demonstrates how to use offline mode for queueing operations when the API is unavailable and syncing when connection is restored. """ import os +import time +from typing import Dict, List +import numpy as np from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split - -from whiteboxai import WhiteBoxAI -from whiteboxai.offline import OperationPriority, OperationType +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.offline import OperationPriority, OperationType def example_1_basic_offline_mode(): @@ -22,8 +24,8 @@ def example_1_basic_offline_mode(): print("=" * 80) # Initialize client with offline mode enabled - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=True, # Auto-sync every 60 seconds @@ -35,7 +37,7 @@ def example_1_basic_offline_mode(): # Check offline status status = client.get_offline_status() - print("\nOffline Status:") + print(f"\nOffline Status:") print(f" Queue size: {status['queue_size']}") print(f" Statistics: {status['statistics']}") @@ -53,8 +55,8 @@ def example_2_manual_sync(): print("=" * 80) # Initialize with auto-sync disabled for manual control - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=False, # Disable auto-sync @@ -92,7 +94,7 @@ def example_2_manual_sync(): # Manually trigger sync when connection is available print("\nTriggering manual sync...") result = client.sync_offline_queue(batch_size=50) - print("Sync result:") + print(f"Sync result:") print(f" Synced: {result['synced']}") print(f" Failed: {result['failed']}") print(f" Pending: {result['pending']}") @@ -107,8 +109,8 @@ def example_3_priority_based_syncing(): print("Example 3: Priority-Based Syncing") print("=" * 80) - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=False, @@ -119,7 +121,9 @@ def example_3_priority_based_syncing(): # Low priority - batch logging client._offline_manager._queue.enqueue( - OperationType.LOG_BATCH, {"model_id": "model_123", "batch": []}, OperationPriority.LOW + OperationType.LOG_BATCH, + {"model_id": "model_123", "batch": []}, + OperationPriority.LOW, ) print(" ✓ Queued LOW priority: batch logging") @@ -163,8 +167,8 @@ def example_4_queue_management(): print("Example 4: Queue Management") print("=" * 80) - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=False, @@ -173,7 +177,7 @@ def example_4_queue_management(): # Get queue statistics status = client.get_offline_status() - print("\nInitial Queue Status:") + print(f"\nInitial Queue Status:") print(f" Total: {status['statistics']['total']}") print(f" Pending: {status['statistics']['pending']}") print(f" Completed: {status['statistics']['completed']}") @@ -202,7 +206,9 @@ def example_5_ml_model_with_offline(): X, y = make_classification( n_samples=1000, n_features=10, n_informative=8, n_redundant=2, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) # Train model print("\nTraining Random Forest model...") @@ -211,8 +217,8 @@ def example_5_ml_model_with_offline(): print(" ✓ Model trained") # Initialize client with offline mode - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./ml_offline_queue", offline_auto_sync=True, @@ -236,7 +242,7 @@ def example_5_ml_model_with_offline(): # Check queue status status = client.get_offline_status() - print("\nQueue Status:") + print(f"\nQueue Status:") print(f" Pending operations: {status['statistics']['pending']}") print(f" Completed: {status['statistics']['completed']}") @@ -250,8 +256,8 @@ def example_6_error_handling(): print("Example 6: Error Handling and Retry") print("=" * 80) - client = WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + client = WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=False, @@ -300,19 +306,19 @@ def example_7_context_manager(): print("Example 7: Context Manager Usage") print("=" * 80) - print("\nUsing WhiteBoxAI client as context manager:") + print("\nUsing WhiteBoxXAI client as context manager:") print(" - Automatically starts auto-sync") print(" - Properly stops sync on exit") print(" - Ensures resource cleanup") - with WhiteBoxAI( - api_key=os.getenv("WHITEBOXAI_API_KEY", "test_key"), + with WhiteBoxXAI( + api_key=os.getenv("WHITEBOXXAI_API_KEY", "test_key"), enable_offline=True, offline_dir="./offline_queue", offline_auto_sync=True, offline_sync_interval=60, ) as client: - print("\n✓ Client initialized with offline mode") + print(f"\n✓ Client initialized with offline mode") print(f" Auto-sync running: {client._offline_manager._sync_running}") status = client.get_offline_status() @@ -328,7 +334,7 @@ def example_7_context_manager(): def run_all_examples(): """Run all offline mode examples.""" print("\n" + "=" * 80) - print("WhiteBoxAI SDK - Offline Mode Examples") + print("WhiteBoxXAI SDK - Offline Mode Examples") print("=" * 80 + "\n") examples = [ @@ -357,8 +363,8 @@ def run_all_examples(): if __name__ == "__main__": # Set environment variable if not set - if "WHITEBOXAI_API_KEY" not in os.environ: - print("Note: WHITEBOXAI_API_KEY not set, using 'test_key' for examples") - print("Set your API key: export WHITEBOXAI_API_KEY=your_key_here\n") + if "WHITEBOXXAI_API_KEY" not in os.environ: + print("Note: WHITEBOXXAI_API_KEY not set, using 'test_key' for examples") + print("Set your API key: export WHITEBOXXAI_API_KEY=your_key_here\n") run_all_examples() diff --git a/examples/pytorch_integration.py b/examples/pytorch_integration.py index 7ec426e..cf2d5f7 100644 --- a/examples/pytorch_integration.py +++ b/examples/pytorch_integration.py @@ -8,9 +8,8 @@ import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset - -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.pytorch import TorchMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.pytorch import TorchMonitor class SimpleClassifier(nn.Module): @@ -76,8 +75,8 @@ def main(): train_model(model, train_loader, epochs=5) # Setup monitoring - print("\nSetting up WhiteBoxAI monitoring...") - client = WhiteBoxAI(api_key="your-api-key") + print("\nSetting up WhiteBoxXAI monitoring...") + client = WhiteBoxXAI(api_key="your-api-key") monitor = TorchMonitor(client, model=model) # Register model diff --git a/examples/sklearn_integration.py b/examples/sklearn_integration.py index 72c8b63..812e4fd 100644 --- a/examples/sklearn_integration.py +++ b/examples/sklearn_integration.py @@ -4,12 +4,12 @@ This example demonstrates monitoring scikit-learn models. """ +import numpy as np from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split - -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.sklearn import SklearnMonitor +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.sklearn import SklearnMonitor def main(): @@ -24,7 +24,9 @@ def main(): ) # Split data - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) # Train model print("Training Random Forest model...") @@ -35,8 +37,8 @@ def main(): print(f"Model accuracy: {accuracy:.3f}") # Setup monitoring - print("\nSetting up WhiteBoxAI monitoring...") - client = WhiteBoxAI(api_key="your-api-key") + print("\nSetting up WhiteBoxXAI monitoring...") + client = WhiteBoxXAI(api_key="your-api-key") monitor = SklearnMonitor(client, model=model) # Register model with automatic metadata extraction diff --git a/examples/tensorflow_example.py b/examples/tensorflow_example.py index 728f173..6a8776c 100644 --- a/examples/tensorflow_example.py +++ b/examples/tensorflow_example.py @@ -1,22 +1,24 @@ """ TensorFlow/Keras Integration Example -This example demonstrates how to use WhiteBoxAI with TensorFlow/Keras models. +This example demonstrates how to use WhiteBoxXAI with TensorFlow/Keras models. """ +import numpy as np from sklearn.datasets import make_classification from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler # TensorFlow imports try: + import tensorflow as tf from tensorflow import keras except ImportError: print("TensorFlow not installed. Install with: pip install tensorflow") exit(1) -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.tensorflow import KerasMonitor, WhiteBoxAICallback +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.tensorflow import KerasMonitor, WhiteBoxXAICallback def main(): @@ -27,11 +29,18 @@ def main(): # Generate synthetic classification data print("\n1. Generating synthetic data...") X, y = make_classification( - n_samples=1000, n_features=20, n_informative=15, n_redundant=5, n_classes=2, random_state=42 + n_samples=1000, + n_features=20, + n_informative=15, + n_redundant=5, + n_classes=2, + random_state=42, ) # Split data - X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) + X_train, X_test, y_train, y_test = train_test_split( + X, y, test_size=0.2, random_state=42 + ) # Standardize features scaler = StandardScaler() @@ -63,9 +72,9 @@ def main(): print(f" Model built with {model.count_params():,} parameters") - # Initialize WhiteBoxAI - print("\n3. Initializing WhiteBoxAI monitoring...") - client = WhiteBoxAI(api_key="demo-api-key", base_url="http://localhost:8000") + # Initialize WhiteBoxXAI + print("\n3. Initializing WhiteBoxXAI monitoring...") + client = WhiteBoxXAI(api_key="demo-api-key", base_url="http://localhost:8000") # Create Keras monitor monitor = KerasMonitor( @@ -87,9 +96,9 @@ def main(): monitor.set_baseline(X_train, y_train) print(" ✓ Baseline set") - # Create WhiteBoxAI callback - print("\n5. Training model with WhiteBoxAI monitoring...") - callback = WhiteBoxAICallback( + # Create WhiteBoxXAI callback + print("\n5. Training model with WhiteBoxXAI monitoring...") + callback = WhiteBoxXAICallback( monitor=monitor, log_frequency=5, log_validation=True # Log every 5 epochs ) @@ -140,7 +149,9 @@ def main(): prob = predictions[i][0] pred_class = 1 if prob > 0.5 else 0 actual = y_test[i] - print(f" Sample {i+1}: Predicted={pred_class} (prob={prob:.3f}), Actual={actual}") + print( + f" Sample {i+1}: Predicted={pred_class} (prob={prob:.3f}), Actual={actual}" + ) # Save model print("\n9. Saving model...") @@ -151,7 +162,7 @@ def main(): monitor.register_saved_model( model_path="models/keras_binary_classifier", metadata={"format": "SavedModel"} ) - print(" ✓ SavedModel registered with WhiteBoxAI") + print(" ✓ SavedModel registered with WhiteBoxXAI") # Check drift print("\n10. Checking for data drift...") diff --git a/examples/transformers_example.py b/examples/transformers_example.py index a5e23ff..9cc9c7a 100644 --- a/examples/transformers_example.py +++ b/examples/transformers_example.py @@ -1,16 +1,19 @@ """ Hugging Face Transformers Integration Example -This example demonstrates how to use WhiteBoxAI to monitor Hugging Face Transformers models. +This example demonstrates how to use WhiteBoxXAI to monitor Hugging Face Transformers models. """ import os -from whiteboxai import WhiteBoxAI -from whiteboxai.integrations.transformers import TransformersMonitor, wrap_transformers_pipeline +from whiteboxxai import WhiteBoxXAI +from whiteboxxai.integrations.transformers import ( + TransformersMonitor, + wrap_transformers_pipeline, +) # Optional: Set API key -os.environ["WHITEBOXAI_API_KEY"] = "your-api-key-here" +os.environ["WHITEBOXXAI_API_KEY"] = "your-api-key-here" def example_sentiment_analysis(): @@ -21,8 +24,8 @@ def example_sentiment_analysis(): print("Sentiment Analysis Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Load sentiment analysis pipeline classifier = pipeline("sentiment-analysis") @@ -78,17 +81,21 @@ def example_ner(): print("Named Entity Recognition Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Load NER pipeline ner_pipeline = pipeline("ner", aggregation_strategy="simple") # Create monitor - monitor = TransformersMonitor(client=client, pipeline=ner_pipeline, model_name="ner_model_v1") + monitor = TransformersMonitor( + client=client, pipeline=ner_pipeline, model_name="ner_model_v1" + ) # Register model - model_id = monitor.register_from_model(name="BERT NER Model", version="1.0.0", task="ner") + model_id = monitor.register_from_model( + name="BERT NER Model", version="1.0.0", task="ner" + ) print(f"✓ Model registered with ID: {model_id}") # Test text @@ -100,7 +107,9 @@ def example_ner(): print("\nDetected entities:") for entity in result: - print(f" - {entity['word']}: {entity['entity_group']} (score: {entity['score']:.3f})") + print( + f" - {entity['word']}: {entity['entity_group']} (score: {entity['score']:.3f})" + ) print("\n✓ NER monitoring complete!") @@ -115,14 +124,16 @@ def example_text_generation(): print("Text Generation Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Load text generation pipeline (using smaller GPT-2 for demo) generator = pipeline("text-generation", model="gpt2") # Create monitor - monitor = TransformersMonitor(client=client, pipeline=generator, model_name="gpt2_generator") + monitor = TransformersMonitor( + client=client, pipeline=generator, model_name="gpt2_generator" + ) # Register model model_id = monitor.register_from_model( @@ -160,8 +171,8 @@ def example_wrapped_pipeline(): print("Wrapped Pipeline Example (Auto-logging)") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Load pipeline classifier = pipeline("sentiment-analysis") @@ -202,14 +213,16 @@ def example_batch_prediction(): print("Batch Prediction Example") print("=" * 60) - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key=os.getenv("WHITEBOXAI_API_KEY")) + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key=os.getenv("WHITEBOXXAI_API_KEY")) # Load pipeline classifier = pipeline("sentiment-analysis") # Create monitor - monitor = TransformersMonitor(client=client, pipeline=classifier, model_name="batch_classifier") + monitor = TransformersMonitor( + client=client, pipeline=classifier, model_name="batch_classifier" + ) # Register model monitor.register_from_model(name="Batch Sentiment Classifier") @@ -239,7 +252,7 @@ def example_batch_prediction(): def main(): """Run all examples.""" print("\n" + "=" * 60) - print("WhiteBoxAI - Hugging Face Transformers Integration Examples") + print("WhiteBoxXAI - Hugging Face Transformers Integration Examples") print("=" * 60) try: diff --git a/mkdocs.yml b/mkdocs.yml index 016139d..e4c0d91 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ -site_name: WhiteBoxAI Python SDK -site_description: Official Python SDK for WhiteBoxAI - AI Observability & Explainability Platform +site_name: WhiteBoxXAI Python SDK +site_description: Official Python SDK for WhiteBoxXAI - AI Observability & Explainability Platform site_author: AgentaFlow site_url: https://github.com/AgentaFlow/whitebox-python-sdk @@ -41,7 +41,7 @@ plugins: - mkdocstrings: handlers: python: - paths: [src] + paths: [.] options: docstring_style: google show_source: true @@ -87,6 +87,6 @@ extra: - icon: fontawesome/brands/github link: https://github.com/AgentaFlow/whitebox-python-sdk - icon: fontawesome/brands/python - link: https://pypi.org/project/whiteboxai-sdk/ + link: https://pypi.org/project/whitebox-xai-sdk/ copyright: Copyright © 2026 AgentaFlow diff --git a/pyproject.toml b/pyproject.toml index 9822b92..d475d85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,16 +3,16 @@ requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "whiteboxai-sdk" -version = "0.2.0" -description = "Official Python SDK for WhiteBoxAI - ML Monitoring and Observability" +name = "whitebox-xai-sdk" +version = "1.0.0" +description = "Python SDK for WhiteBox XAI - AI Observability & Explainability Platform" readme = "README.md" authors = [ - {name = "AgentaFlow", email = "support@agentaflow.com"} + {name = "WhiteBoxXAI Team", email = "support@whiteboxxai.com"} ] license = {text = "MIT"} classifiers = [ - "Development Status :: 3 - Alpha", + "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", "Intended Audience :: Science/Research", "License :: OSI Approved :: MIT License", @@ -20,119 +20,71 @@ classifiers = [ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", ] -keywords = ["machine-learning", "explainability", "monitoring", "observability", "xai", "mlops", "whiteboxai"] +keywords = ["machine-learning", "explainability", "monitoring", "observability", "xai", "mlops"] requires-python = ">=3.9" dependencies = [ "httpx>=0.25.0", "pydantic>=2.0.0", - "python-dotenv>=1.0.0", - "numpy>=1.24.0", - "pandas>=1.3.0", "tenacity>=8.0.0", + "numpy>=1.20.0", + "pandas>=1.3.0", ] [project.optional-dependencies] dev = [ - "pytest>=7.4.0", + "pytest>=7.0.0", "pytest-asyncio>=0.21.0", - "pytest-cov>=4.1.0", - "pytest-mock>=3.12.0", - "black>=23.11.0", + "pytest-cov>=4.0.0", + "black>=23.0.0", "isort>=5.12.0", - "flake8>=6.1.0", - "pylint>=3.0.0", - "mypy>=1.7.0", - "bandit>=1.7.5", + "flake8>=6.0.0", + "mypy>=1.0.0", ] -git = ["gitpython>=3.1.0"] -sklearn = ["scikit-learn>=1.3.0"] +sklearn = ["scikit-learn>=1.0.0"] pytorch = ["torch>=2.0.0"] -tensorflow = ["tensorflow>=2.13.0"] -transformers = ["transformers>=4.30.0"] -langchain = ["langchain>=0.0.200"] -crewai = ["crewai>=0.1.0"] -boosting = ["xgboost>=1.7.0", "lightgbm>=4.0.0"] +tensorflow = ["tensorflow>=2.10.0"] +xgboost = ["xgboost>=1.7.0"] +lightgbm = ["lightgbm>=4.0.0"] +huggingface = ["transformers>=4.30.0"] +langchain = ["langchain>=0.1.0"] all = [ - "gitpython>=3.1.0", - "scikit-learn>=1.3.0", + "scikit-learn>=1.0.0", "torch>=2.0.0", - "tensorflow>=2.13.0", - "transformers>=4.30.0", - "langchain>=0.0.200", - "crewai>=0.1.0", + "tensorflow>=2.10.0", "xgboost>=1.7.0", "lightgbm>=4.0.0", + "transformers>=4.30.0", + "langchain>=0.1.0", ] +[project.urls] +Homepage = "https://whiteboxxai.com" +Documentation = "https://docs.whiteboxxai.com/sdk" + +[tool.setuptools.packages.find] +where = ["."] +include = ["whiteboxxai*"] + [tool.black] line-length = 100 target-version = ['py39'] -include = '\.pyi?$' -extend-exclude = ''' -/( - # directories - \.eggs - | \.git - | \.hg - | \.mypy_cache - | \.tox - | \.venv - | build - | dist -)/ -''' [tool.isort] profile = "black" line_length = 100 -multi_line_output = 3 -include_trailing_comma = true -force_grid_wrap = 0 -use_parentheses = true -ensure_newline_before_comments = true [tool.mypy] python_version = "3.9" warn_return_any = true warn_unused_configs = true -disallow_untyped_defs = false -disallow_incomplete_defs = false -check_untyped_defs = true -no_implicit_optional = true -warn_redundant_casts = true -warn_unused_ignores = true -warn_no_return = true +disallow_untyped_defs = true -[tool.pylint.messages_control] -disable = "C0330, C0326" - -[tool.pylint.format] -max-line-length = "120" - -[tool.coverage.run] -source = ["src/whiteboxai"] -omit = ["*/tests/*", "*/test_*.py"] - -[tool.coverage.report] -exclude_lines = [ - "pragma: no cover", - "def __repr__", - "raise AssertionError", - "raise NotImplementedError", - "if __name__ == .__main__.:", - "if TYPE_CHECKING:", - "@abstractmethod", -] - -[project.urls] -Homepage = "https://whitebox.agentaflow.com" -Documentation = "https://whitebox.agentaflow.com/docs" -Repository = "https://github.com/AgentaFlow/whitebox-python-sdk" -"Bug Tracker" = "https://github.com/AgentaFlow/whitebox-python-sdk/issues" - -[tool.setuptools.packages.find] -where = ["src"] -include = ["whiteboxai*"] +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = "-v --cov=whiteboxxai --cov-report=html --cov-report=term" diff --git a/requirements-dev.txt b/requirements-dev.txt deleted file mode 100644 index a87bedf..0000000 --- a/requirements-dev.txt +++ /dev/null @@ -1,26 +0,0 @@ -# Development dependencies for WhiteBoxAI SDK - -# Testing -pytest>=7.4.0 -pytest-asyncio>=0.21.0 -pytest-cov>=4.1.0 -pytest-mock>=3.12.0 - -# Code quality -black>=23.11.0 -isort>=5.12.0 -flake8>=6.1.0 -pylint>=3.0.0 -mypy>=1.7.0 - -# Security -bandit>=1.7.5 - -# Documentation -mkdocs>=1.5.0 -mkdocs-material>=9.0.0 -mkdocstrings[python]>=0.24.0 - -# Build tools -build>=1.0.0 -twine>=4.0.0 diff --git a/setup.py b/setup.py deleted file mode 100644 index 23da3c8..0000000 --- a/setup.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -WhiteBoxAI Python SDK - ML Monitoring and Observability -""" - -from pathlib import Path - -from setuptools import find_packages, setup - -# Read the contents of README file -this_directory = Path(__file__).parent -long_description = (this_directory / "README.md").read_text(encoding="utf-8") - - -# Read requirements -def read_requirements(filename): - """Read requirements from file.""" - filepath = this_directory / filename - if not filepath.exists(): - return [] - with open(filepath, encoding="utf-8") as f: - return [ - line.strip() - for line in f - if line.strip() and not line.startswith("#") and not line.startswith("-r") - ] - - -# Core requirements -install_requires = [ - "httpx>=0.24.0", - "pydantic>=2.0.0", - "python-dotenv>=1.0.0", - "numpy>=1.24.0", -] - -# Optional dependencies for integrations -extras_require = { - "sklearn": ["scikit-learn>=1.3.0"], - "pytorch": ["torch>=2.0.0"], - "tensorflow": ["tensorflow>=2.13.0"], - "transformers": ["transformers>=4.30.0"], - "langchain": ["langchain>=0.0.200"], - "boosting": ["xgboost>=1.7.0", "lightgbm>=4.0.0"], - "all": [ - "scikit-learn>=1.3.0", - "torch>=2.0.0", - "tensorflow>=2.13.0", - "transformers>=4.30.0", - "langchain>=0.0.200", - "xgboost>=1.7.0", - "lightgbm>=4.0.0", - ], - "dev": [ - "pytest>=7.4.0", - "pytest-asyncio>=0.21.0", - "pytest-cov>=4.1.0", - "black>=23.11.0", - "isort>=5.12.0", - "flake8>=6.1.0", - "pylint>=3.0.0", - "mypy>=1.7.0", - "bandit>=1.7.5", - ], -} - -setup( - name="whiteboxai-sdk", - version="0.1.0", - author="AgentaFlow", - author_email="support@agentaflow.com", - description="Official Python SDK for WhiteBoxAI ML monitoring and observability", - long_description=long_description, - long_description_content_type="text/markdown", - url="https://github.com/AgentaFlow/whitebox-python-sdk", - project_urls={ - "Bug Reports": "https://github.com/AgentaFlow/whitebox-python-sdk/issues", - "Source": "https://github.com/AgentaFlow/whitebox-python-sdk", - "Documentation": "https://whitebox.agentaflow.com", - }, - packages=find_packages(where="src"), - package_dir={"": "src"}, - classifiers=[ - "Development Status :: 3 - Alpha", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", - "License :: OSI Approved :: MIT License", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Operating System :: OS Independent", - ], - python_requires=">=3.9", - install_requires=install_requires, - extras_require=extras_require, - include_package_data=True, - zip_safe=False, -) diff --git a/src/whiteboxai/__init__.py b/src/whiteboxai/__init__.py deleted file mode 100644 index 5fbf541..0000000 --- a/src/whiteboxai/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -""" -WhiteBoxAI Python SDK - -Official Python SDK for WhiteBoxAI - AI Observability & Explainability Platform. -""" - -__version__ = "0.2.0" -__author__ = "WhiteBoxAI Team" -__license__ = "MIT" - -from whiteboxai.client import WhiteBoxAI -from whiteboxai.decorators import monitor_model, monitor_prediction -from whiteboxai.git_utils import GitContext, detect_git_context, validate_git_context -from whiteboxai.monitor import ModelMonitor - -__all__ = [ - "WhiteBoxAI", - "ModelMonitor", - "monitor_model", - "monitor_prediction", - "GitContext", - "detect_git_context", - "validate_git_context", -] diff --git a/src/whiteboxai/__version__.py b/src/whiteboxai/__version__.py deleted file mode 100644 index afa047f..0000000 --- a/src/whiteboxai/__version__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Version information for WhiteBoxAI SDK.""" - -__version__ = "0.2.0" -__author__ = "AgentaFlow" -__email__ = "support@agentaflow.com" -__license__ = "MIT" diff --git a/src/whiteboxai/exceptions.py b/src/whiteboxai/exceptions.py deleted file mode 100644 index 7932c7f..0000000 --- a/src/whiteboxai/exceptions.py +++ /dev/null @@ -1,57 +0,0 @@ -""" -SDK Exceptions -""" - - -class WhiteBoxAIError(Exception): - """Base exception for WhiteBoxAI SDK.""" - - pass - - -class APIError(WhiteBoxAIError): - """API request error.""" - - pass - - -class AuthenticationError(WhiteBoxAIError): - """Authentication error.""" - - pass - - -class RateLimitError(WhiteBoxAIError): - """Rate limit exceeded error.""" - - pass - - -class ValidationError(WhiteBoxAIError): - """Validation error.""" - - pass - - -class NotFoundError(WhiteBoxAIError): - """Resource not found error.""" - - pass - - -class ConfigurationError(WhiteBoxAIError): - """Configuration error.""" - - pass - - -class IntegrationError(WhiteBoxAIError): - """Framework integration error.""" - - pass - - -class CacheError(WhiteBoxAIError): - """Cache operation error.""" - - pass diff --git a/src/whiteboxai/models/__init__.py b/src/whiteboxai/models/__init__.py deleted file mode 100644 index 5fbfb4d..0000000 --- a/src/whiteboxai/models/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Models package for WhiteBoxAI SDK.""" diff --git a/src/whiteboxai/resources.py b/src/whiteboxai/resources.py deleted file mode 100644 index d0908a4..0000000 --- a/src/whiteboxai/resources.py +++ /dev/null @@ -1,460 +0,0 @@ -""" -API Resources - -Resource classes for interacting with different WhiteBoxAI API endpoints. -""" - -from typing import TYPE_CHECKING, Any, Dict, List, Optional - -if TYPE_CHECKING: - from whiteboxai.client import WhiteBoxAI - - -class BaseResource: - """Base class for API resources.""" - - def __init__(self, client: "WhiteBoxAI"): - """Initialize resource with client.""" - self.client = client - - -class ModelsResource(BaseResource): - """Models API resource.""" - - def register( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Register a new model. - - Args: - name: Model name - model_type: Model type (classification, regression, etc.) - framework: ML framework (sklearn, pytorch, tensorflow, etc.) - version: Model version - **kwargs: Additional model metadata - - Returns: - Registered model data - """ - data = { - "name": name, - "model_type": model_type, - "framework": framework, - "version": version, - **kwargs, - } - return self.client.request("POST", "/api/v1/models", data=data) - - async def aregister( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of register().""" - data = { - "name": name, - "model_type": model_type, - "framework": framework, - "version": version, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/models", data=data) - - def get(self, model_id: int) -> Dict[str, Any]: - """Get model by ID.""" - return self.client.request("GET", f"/api/v1/models/{model_id}") - - async def aget(self, model_id: int) -> Dict[str, Any]: - """Async version of get().""" - return await self.client.arequest("GET", f"/api/v1/models/{model_id}") - - def list(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: - """List registered models.""" - params = {"limit": limit, "offset": offset} - return self.client.request("GET", "/api/v1/models", params=params) - - async def alist(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: - """Async version of list().""" - params = {"limit": limit, "offset": offset} - return await self.client.arequest("GET", "/api/v1/models", params=params) - - def update(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: - """Update model metadata.""" - return self.client.request("PUT", f"/api/v1/models/{model_id}", data=kwargs) - - async def aupdate(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: - """Async version of update().""" - return await self.client.arequest("PUT", f"/api/v1/models/{model_id}", data=kwargs) - - -class PredictionsResource(BaseResource): - """Predictions API resource.""" - - def log( - self, - model_id: int, - inputs: Any, - outputs: Any, - explain: bool = False, - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """ - Log a single prediction. - - Args: - model_id: Model ID - inputs: Input features/data - outputs: Model prediction/output - explain: Whether to generate explanation - metadata: Additional prediction metadata - - Returns: - Logged prediction data - """ - data = { - "model_id": model_id, - "inputs": inputs, - "outputs": outputs, - "explain": explain, - "metadata": metadata or {}, - } - return self.client.request("POST", "/api/v1/predictions/log", data=data) - - async def alog( - self, - model_id: int, - inputs: Any, - outputs: Any, - explain: bool = False, - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Async version of log().""" - data = { - "model_id": model_id, - "inputs": inputs, - "outputs": outputs, - "explain": explain, - "metadata": metadata or {}, - } - return await self.client.arequest("POST", "/api/v1/predictions/log", data=data) - - def log_batch( - self, - model_id: int, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """ - Log multiple predictions in batch. - - Args: - model_id: Model ID - predictions: List of prediction dictionaries - - Returns: - Batch logging result - """ - data = {"model_id": model_id, "predictions": predictions} - return self.client.request("POST", "/api/v1/predictions/log/batch", data=data) - - async def alog_batch( - self, - model_id: int, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Async version of log_batch().""" - data = {"model_id": model_id, "predictions": predictions} - return await self.client.arequest("POST", "/api/v1/predictions/log/batch", data=data) - - -class ExplanationsResource(BaseResource): - """Explanations API resource.""" - - def generate( - self, - prediction_id: int, - method: str = "shap", - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Generate explanation for a prediction. - - Args: - prediction_id: Prediction ID - method: Explanation method (shap, lime) - **kwargs: Method-specific parameters - - Returns: - Explanation data - """ - data = {"prediction_id": prediction_id, "method": method, **kwargs} - return self.client.request("POST", "/api/v1/explanations/generate", data=data) - - async def agenerate( - self, - prediction_id: int, - method: str = "shap", - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of generate().""" - data = {"prediction_id": prediction_id, "method": method, **kwargs} - return await self.client.arequest("POST", "/api/v1/explanations/generate", data=data) - - def get(self, explanation_id: int) -> Dict[str, Any]: - """Get explanation by ID.""" - return self.client.request("GET", f"/api/v1/explanations/{explanation_id}") - - async def aget(self, explanation_id: int) -> Dict[str, Any]: - """Async version of get().""" - return await self.client.arequest("GET", f"/api/v1/explanations/{explanation_id}") - - -class DriftResource(BaseResource): - """Drift detection API resource.""" - - def detect( - self, - model_id: int, - reference_data: Optional[Any] = None, - current_data: Optional[Any] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Detect drift for a model. - - Args: - model_id: Model ID - reference_data: Reference/baseline data - current_data: Current data to compare - **kwargs: Detection parameters - - Returns: - Drift detection results - """ - data = { - "model_id": model_id, - "reference_data": reference_data, - "current_data": current_data, - **kwargs, - } - return self.client.request("POST", "/api/v1/drift/detect", data=data) - - async def adetect( - self, - model_id: int, - reference_data: Optional[Any] = None, - current_data: Optional[Any] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of detect().""" - data = { - "model_id": model_id, - "reference_data": reference_data, - "current_data": current_data, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/drift/detect", data=data) - - def get_report(self, model_id: int, report_id: int) -> Dict[str, Any]: - """Get drift report.""" - return self.client.request("GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}") - - async def aget_report(self, model_id: int, report_id: int) -> Dict[str, Any]: - """Async version of get_report().""" - return await self.client.arequest( - "GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}" - ) - - -class AlertsResource(BaseResource): - """Alerts API resource.""" - - def create( - self, - name: str, - alert_type: str, - conditions: Dict[str, Any], - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Create an alert rule. - - Args: - name: Alert name - alert_type: Alert type - conditions: Alert conditions - **kwargs: Additional alert configuration - - Returns: - Created alert data - """ - data = { - "name": name, - "alert_type": alert_type, - "conditions": conditions, - **kwargs, - } - return self.client.request("POST", "/api/v1/alerts", data=data) - - async def acreate( - self, - name: str, - alert_type: str, - conditions: Dict[str, Any], - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of create().""" - data = { - "name": name, - "alert_type": alert_type, - "conditions": conditions, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/alerts", data=data) - - def list(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: - """List alert rules.""" - params = {"model_id": model_id} if model_id else {} - return self.client.request("GET", "/api/v1/alerts", params=params) - - async def alist(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: - """Async version of list().""" - params = {"model_id": model_id} if model_id else {} - return await self.client.arequest("GET", "/api/v1/alerts", params=params) - - -class AgentWorkflowsResource(BaseResource): - """Agent Workflows API resource for multi-agent monitoring.""" - - def create( - self, - name: str, - framework: str, - inputs: Optional[Dict[str, Any]] = None, - meta_data: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Create a new agent workflow.""" - data = { - "name": name, - "framework": framework, - "inputs": inputs or {}, - "meta_data": meta_data or {}, - } - return self.client.request("POST", "/api/v1/workflows/multi-agent", data=data) - - def start( - self, - workflow_id: str, - inputs: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Start an agent workflow.""" - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/start", - data={"inputs": inputs or {}}, - ) - - def complete( - self, - workflow_id: str, - outputs: Optional[Dict[str, Any]] = None, - status: str = "completed", - ) -> Dict[str, Any]: - """Complete an agent workflow.""" - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/complete", - data={"outputs": outputs or {}, "status": status}, - ) - - def register_agent( - self, - workflow_id: str, - name: str, - role: Optional[str] = None, - model_name: Optional[str] = None, - tools: Optional[List[str]] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Register an agent in a workflow.""" - data = { - "name": name, - "role": role or name, - "model_name": model_name, - "tools": tools or [], - **kwargs, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/agents", - data=data, - ) - - def create_execution( - self, - workflow_id: str, - agent_name: str, - status: str, - inputs: Optional[Dict[str, Any]] = None, - outputs: Optional[Dict[str, Any]] = None, - duration_ms: Optional[int] = None, - llm_call_count: Optional[int] = None, - tool_call_count: Optional[int] = None, - tokens_used: Optional[int] = None, - cost: Optional[float] = None, - ) -> Dict[str, Any]: - """Log an agent execution within a workflow.""" - data = { - "agent_name": agent_name, - "status": status, - "inputs": inputs, - "outputs": outputs, - "duration_ms": duration_ms, - "llm_call_count": llm_call_count, - "tool_call_count": tool_call_count, - "tokens_used": tokens_used, - "cost": cost, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/executions", - data=data, - ) - - def create_interaction( - self, - workflow_id: str, - from_agent: str, - to_agent: str, - interaction_type: str, - message: Optional[str] = None, - meta_data: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Log an agent-to-agent interaction.""" - data = { - "from_agent": from_agent, - "to_agent": to_agent, - "interaction_type": interaction_type, - "message": message, - "meta_data": meta_data or {}, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/interactions", - data=data, - ) - - def get_analytics(self, workflow_id: str) -> Dict[str, Any]: - """Get analytics for a workflow.""" - return self.client.request( - "GET", - f"/api/v1/workflows/multi-agent/{workflow_id}/analytics", - ) diff --git a/tests/conftest.py b/tests/conftest.py index 5cb5804..6c3f46f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,7 +2,7 @@ import pytest -from whiteboxai import WhiteBoxAI +from whiteboxxai import WhiteBoxXAI @pytest.fixture @@ -13,8 +13,8 @@ def mock_api_key(): @pytest.fixture def client(mock_api_key): - """Create a WhiteBoxAI client for testing.""" - return WhiteBoxAI(api_key=mock_api_key, base_url="http://localhost:8000", timeout=10) + """Create a WhiteBoxXAI client for testing.""" + return WhiteBoxXAI(api_key=mock_api_key, base_url="http://localhost:8000", timeout=10) @pytest.fixture diff --git a/tests/integration/test_crewai_integration.py b/tests/integration/test_crewai_integration.py new file mode 100644 index 0000000..e43e193 --- /dev/null +++ b/tests/integration/test_crewai_integration.py @@ -0,0 +1,455 @@ +""" +Unit tests for CrewAI integration. + +Tests CrewAI monitoring with mocked API calls. +""" + +from unittest.mock import MagicMock, Mock, patch +from uuid import uuid4 + +import pytest +from whiteboxxai.integrations.crewai_monitor import CrewAIMonitor, monitor_crew + + +@pytest.fixture +def mock_client(): + """Create mock WhiteBoxXAI client.""" + client = Mock() + client.request = Mock() + return client + + +@pytest.fixture +def mock_crew(): + """Create mock CrewAI Crew.""" + # Mock agents + agent1 = Mock() + agent1.role = "Research Analyst" + agent1.goal = "Find accurate information" + agent1.backstory = "Expert researcher" + agent1.tools = [Mock(__class__=Mock(__name__="SearchTool"))] + agent1.llm = Mock(model_name="gpt-4") + agent1.verbose = True + agent1.allow_delegation = False + agent1.max_iter = 15 + + agent2 = Mock() + agent2.role = "Content Writer" + agent2.goal = "Write engaging content" + agent2.backstory = "Professional writer" + agent2.tools = [] + agent2.llm = Mock(model_name="gpt-4") + agent2.verbose = False + agent2.allow_delegation = True + agent2.max_iter = 10 + + # Mock tasks + task1 = Mock() + task1.description = "Research AI safety regulations" + task1.expected_output = "List of regulations" + task1.agent = agent1 + task1.tools = [] + task1.async_execution = False + + task2 = Mock() + task2.description = "Write article about AI safety" + task2.expected_output = "1500 word article" + task2.agent = agent2 + task2.tools = [] + task2.async_execution = False + + # Mock crew + crew = Mock() + crew.agents = [agent1, agent2] + crew.tasks = [task1, task2] + crew.process = "sequential" + + return crew + + +class TestCrewAIMonitor: + """Test CrewAI monitor functionality.""" + + def test_init(self, mock_client): + """Test monitor initialization.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + monitor = CrewAIMonitor(api_key="test_api_key") + + assert monitor.client == mock_client + assert monitor.workflow_id is None + assert monitor.agent_map == {} + assert monitor.task_map == {} + assert monitor.execution_map == {} + + def test_start_monitoring(self, mock_client, mock_crew): + """Test starting workflow monitoring.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + agent1_id = str(uuid4()) + agent2_id = str(uuid4()) + task1_id = str(uuid4()) + task2_id = str(uuid4()) + + # Mock API responses + mock_client.request.side_effect = [ + {"id": workflow_id}, # Create workflow + {"id": agent1_id}, # Register agent 1 + {"id": agent2_id}, # Register agent 2 + {"id": task1_id}, # Register task 1 + {"id": task2_id}, # Register task 2 + {}, # Start workflow + ] + + monitor = CrewAIMonitor(api_key="test_api_key") + result_id = monitor.start_monitoring( + crew=mock_crew, + workflow_name="Test Workflow", + metadata={"project": "test"}, + ) + + assert result_id == workflow_id + assert monitor.workflow_id == workflow_id + assert len(monitor.agent_map) == 2 + assert len(monitor.task_map) == 2 + + # Verify API calls + calls = mock_client.request.call_args_list + + # Check create workflow call + assert calls[0][0][0] == "POST" + assert calls[0][0][1] == "/api/v1/workflows/multi-agent/start" + assert calls[0][1]["data"]["name"] == "Test Workflow" + assert calls[0][1]["data"]["framework"] == "crewai" + assert calls[0][1]["data"]["metadata"]["project"] == "test" + + # Check start workflow call + assert calls[5][0][0] == "POST" + assert ( + calls[5][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/start" + ) + + def test_register_agent(self, mock_client, mock_crew): + """Test agent registration.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + agent_id = str(uuid4()) + + mock_client.request.return_value = {"id": agent_id} + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + + result_id = monitor._register_agent(mock_crew.agents[0]) + + assert result_id == agent_id + assert monitor.agent_map[id(mock_crew.agents[0])] == agent_id + + # Verify API call + call_args = mock_client.request.call_args + assert call_args[0][0] == "POST" + assert ( + call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/agents" + ) + + agent_data = call_args[1]["data"] + assert agent_data["name"] == "Research Analyst" + assert agent_data["role"] == "Research Analyst" + assert agent_data["agent_type"] == "crewai_agent" + assert agent_data["goal"] == "Find accurate information" + assert agent_data["backstory"] == "Expert researcher" + assert agent_data["tools"] == ["SearchTool"] + assert agent_data["llm_provider"] == "gpt" + assert agent_data["model_name"] == "gpt-4" + + def test_register_task(self, mock_client, mock_crew): + """Test task registration.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + task_id = str(uuid4()) + agent_id = str(uuid4()) + + mock_client.request.return_value = {"id": task_id} + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + monitor.agent_map[id(mock_crew.tasks[0].agent)] = agent_id + + result_id = monitor._register_task(mock_crew.tasks[0]) + + assert result_id == task_id + assert monitor.task_map[id(mock_crew.tasks[0])] == task_id + + # Verify API call + call_args = mock_client.request.call_args + assert call_args[0][0] == "POST" + assert ( + call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/tasks" + ) + + task_data = call_args[1]["data"] + assert task_data["task_name"] == "Research AI safety regulations" + assert task_data["description"] == "Research AI safety regulations" + assert task_data["expected_output"] == "List of regulations" + assert task_data["agent_id"] == agent_id + + def test_log_agent_execution(self, mock_client, mock_crew): + """Test logging agent execution.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + agent_id = str(uuid4()) + + mock_client.request.return_value = {} + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + monitor.agent_map[id(mock_crew.agents[0])] = agent_id + + monitor.log_agent_execution( + agent=mock_crew.agents[0], + inputs={"query": "test"}, + outputs={"result": "data"}, + tokens_used=100, + cost=0.002, + ) + + # Verify API call + call_args = mock_client.request.call_args + assert call_args[0][0] == "POST" + assert ( + call_args[0][1] + == f"/api/v1/workflows/multi-agent/{workflow_id}/executions" + ) + + execution_data = call_args[1]["data"] + assert execution_data["agent_id"] == agent_id + assert execution_data["inputs"] == {"query": "test"} + assert execution_data["outputs"] == {"result": "data"} + assert execution_data["tokens_used"] == 100 + assert execution_data["cost"] == 0.002 + assert execution_data["status"] == "completed" + + def test_log_task_completion(self, mock_client, mock_crew): + """Test logging task completion.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + task_id = str(uuid4()) + + mock_client.request.return_value = {} + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.task_map[id(mock_crew.tasks[0])] = task_id + + monitor.log_task_completion( + task=mock_crew.tasks[0], + status="completed", + output_data={"result": "success"}, + ) + + # Verify API call + call_args = mock_client.request.call_args + assert call_args[0][0] == "PATCH" + assert call_args[0][1] == f"/api/v1/workflows/multi-agent/tasks/{task_id}" + + update_data = call_args[1]["data"] + assert update_data["status"] == "completed" + assert update_data["output_data"] == {"result": "success"} + + def test_log_interaction(self, mock_client, mock_crew): + """Test logging agent interaction.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + agent1_id = str(uuid4()) + agent2_id = str(uuid4()) + + mock_client.request.return_value = {} + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + monitor.agent_map[id(mock_crew.agents[0])] = agent1_id + monitor.agent_map[id(mock_crew.agents[1])] = agent2_id + + monitor.log_interaction( + from_agent=mock_crew.agents[0], + to_agent=mock_crew.agents[1], + interaction_type="delegation", + message="Please write article", + ) + + # Verify API call + call_args = mock_client.request.call_args + assert call_args[0][0] == "POST" + assert ( + call_args[0][1] + == f"/api/v1/workflows/multi-agent/{workflow_id}/interactions" + ) + + interaction_data = call_args[1]["data"] + assert interaction_data["interaction_type"] == "delegation" + assert interaction_data["from_agent_id"] == agent1_id + assert interaction_data["to_agent_id"] == agent2_id + assert interaction_data["message"] == "Please write article" + + def test_complete_monitoring(self, mock_client): + """Test completing workflow monitoring.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + + mock_client.request.side_effect = [ + {}, # Complete workflow + {"total_tokens": 5000, "total_cost": 0.05}, # Analytics + {"total_cost": 0.05, "agents": []}, # Cost breakdown + ] + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + + result = monitor.complete_monitoring( + status="completed", outputs={"article": "content"} + ) + + assert result["workflow_id"] == workflow_id + assert result["status"] == "completed" + assert "analytics" in result + + # Verify complete call + calls = mock_client.request.call_args_list + assert calls[0][0][0] == "POST" + assert ( + calls[0][0][1] + == f"/api/v1/workflows/multi-agent/{workflow_id}/complete" + ) + + complete_data = calls[0][1]["data"] + assert complete_data["status"] == "completed" + assert complete_data["outputs"] == {"article": "content"} + + def test_get_analytics(self, mock_client): + """Test getting workflow analytics.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + + mock_client.request.side_effect = [ + {"total_tokens": 5000, "total_cost": 0.05}, + {"total_cost": 0.05, "agents": []}, + ] + + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = workflow_id + + analytics = monitor.get_analytics() + + assert "metrics" in analytics + assert "cost_breakdown" in analytics + + # Verify calls + calls = mock_client.request.call_args_list + assert ( + calls[0][0][1] + == f"/api/v1/workflows/multi-agent/{workflow_id}/analytics" + ) + assert ( + calls[1][0][1] + == f"/api/v1/workflows/multi-agent/{workflow_id}/cost-breakdown" + ) + + +class TestMonitorCrewHelper: + """Test monitor_crew convenience function.""" + + def test_monitor_crew(self, mock_client, mock_crew): + """Test monitor_crew helper function.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + workflow_id = str(uuid4()) + + mock_client.request.side_effect = [ + {"id": workflow_id}, # Create workflow + {"id": str(uuid4())}, # Agent 1 + {"id": str(uuid4())}, # Agent 2 + {"id": str(uuid4())}, # Task 1 + {"id": str(uuid4())}, # Task 2 + {}, # Start workflow + ] + + monitor = monitor_crew( + crew=mock_crew, + workflow_name="Test Workflow", + api_key="test_api_key", + metadata={"test": "data"}, + ) + + assert isinstance(monitor, CrewAIMonitor) + assert monitor.workflow_id == workflow_id + + +class TestErrorHandling: + """Test error handling.""" + + def test_start_monitoring_error(self, mock_client, mock_crew): + """Test error handling in start_monitoring.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + mock_client.request.side_effect = Exception("API Error") + + monitor = CrewAIMonitor(api_key="test_api_key") + + with pytest.raises(Exception, match="API Error"): + monitor.start_monitoring(crew=mock_crew, workflow_name="Test Workflow") + + def test_log_execution_unregistered_agent(self, mock_client, mock_crew): + """Test logging execution for unregistered agent.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + monitor = CrewAIMonitor(api_key="test_api_key") + monitor.workflow_id = str(uuid4()) + + # Should not raise exception, just log warning + monitor.log_agent_execution( + agent=mock_crew.agents[0], inputs={"test": "data"} + ) + + # No API call should be made + mock_client.request.assert_not_called() + + def test_complete_monitoring_no_workflow(self, mock_client): + """Test completing monitoring without active workflow.""" + with patch( + "whiteboxxai.integrations.crewai_monitor.WhiteBoxXAI", + return_value=mock_client, + ): + monitor = CrewAIMonitor(api_key="test_api_key") + + with pytest.raises(Exception): + monitor.complete_monitoring(status="completed") diff --git a/tests/integration/test_langchain_integration.py b/tests/integration/test_langchain_integration.py new file mode 100644 index 0000000..0fb4e2d --- /dev/null +++ b/tests/integration/test_langchain_integration.py @@ -0,0 +1,452 @@ +"""Unit tests for LangChain multi-agent integration.""" + +from datetime import datetime +from unittest.mock import MagicMock, Mock, call, patch + +import pytest + +# NOTE: this requires the real `langchain` package (see the `langchain` extra +# in sdk/pyproject.toml) to be installed. langchain_agents.py subclasses +# langchain's real BaseCallbackHandler at class-definition time, so faking +# out sys.modules["langchain.callbacks.base"] etc. here (as this file +# previously did) makes that inheritance operate on a Mock instead of a +# real class, producing undefined/garbage behavior. +from whiteboxxai.integrations.langchain_agents import ( + LangGraphMultiAgentMonitor, + MultiAgentCallbackHandler, + monitor_langchain_agent, +) + + +@pytest.fixture +def mock_client(): + """Create a mock WhiteBoxXAI client.""" + client = Mock() + client.agent_workflows = Mock() + return client + + +@pytest.fixture +def mock_workflow_response(): + """Mock workflow creation response.""" + return {"id": "workflow_123", "status": "pending"} + + +@pytest.fixture +def callback_handler(mock_client, mock_workflow_response): + """Create a MultiAgentCallbackHandler instance.""" + return MultiAgentCallbackHandler( + client=mock_client, + workflow_id="workflow_123", + agent_name="test_agent", + agent_role="Test Agent", + ) + + +class TestMultiAgentCallbackHandler: + """Tests for MultiAgentCallbackHandler.""" + + def test_init(self, mock_client): + """Test callback handler initialization.""" + handler = MultiAgentCallbackHandler( + client=mock_client, + workflow_id="workflow_123", + agent_name="test_agent", + agent_role="Test Role", + ) + + assert handler.client == mock_client + assert handler.workflow_id == "workflow_123" + assert handler.agent_name == "test_agent" + assert handler.agent_role == "Test Role" + assert handler.track_tokens is True + assert handler.track_costs is True + assert handler.llm_call_count == 0 + assert handler.tool_call_count == 0 + + def test_on_chain_start(self, callback_handler): + """Test chain start callback.""" + inputs = {"input": "test query"} + + callback_handler.on_chain_start( + serialized={"name": "test_chain"}, inputs=inputs + ) + + assert callback_handler.execution_inputs == inputs + assert callback_handler.execution_start_time is not None + assert callback_handler.llm_call_count == 0 + assert callback_handler.tool_call_count == 0 + + def test_on_chain_end(self, callback_handler, mock_client): + """Test chain end callback.""" + callback_handler.execution_start_time = datetime.utcnow() + callback_handler.execution_inputs = {"input": "test"} + callback_handler.llm_call_count = 2 + callback_handler.tool_call_count = 1 + callback_handler.total_tokens = 150 + callback_handler.total_cost = 0.003 + + mock_client.agent_workflows.create_execution.return_value = {"id": "exec_123"} + + outputs = {"output": "test result"} + callback_handler.on_chain_end(outputs=outputs) + + # Verify execution was logged + mock_client.agent_workflows.create_execution.assert_called_once() + call_args = mock_client.agent_workflows.create_execution.call_args + + assert call_args.kwargs["workflow_id"] == "workflow_123" + assert call_args.kwargs["agent_name"] == "test_agent" + assert call_args.kwargs["status"] == "completed" + assert call_args.kwargs["inputs"] == {"input": "test"} + assert call_args.kwargs["outputs"] == outputs + assert call_args.kwargs["llm_call_count"] == 2 + assert call_args.kwargs["tool_call_count"] == 1 + assert call_args.kwargs["tokens_used"] == 150 + assert call_args.kwargs["cost"] == 0.003 + + # Verify state was reset + assert callback_handler.execution_start_time is None + + def test_on_chain_error(self, callback_handler, mock_client): + """Test chain error callback.""" + callback_handler.execution_start_time = datetime.utcnow() + callback_handler.execution_inputs = {"input": "test"} + callback_handler.llm_call_count = 1 + + error = ValueError("Test error") + callback_handler.on_chain_error(error=error) + + # Verify failed execution was logged + mock_client.agent_workflows.create_execution.assert_called_once() + call_args = mock_client.agent_workflows.create_execution.call_args + + assert call_args.kwargs["status"] == "failed" + assert "error" in call_args.kwargs["outputs"] + assert "Test error" in call_args.kwargs["outputs"]["error"] + + def test_on_llm_start(self, callback_handler): + """Test LLM start callback.""" + initial_count = callback_handler.llm_call_count + + callback_handler.on_llm_start( + serialized={"name": "test_llm"}, prompts=["test prompt"] + ) + + assert callback_handler.llm_call_count == initial_count + 1 + + def test_on_llm_end_with_tokens(self, callback_handler): + """Test LLM end callback with token tracking.""" + # Mock LLMResult with token usage + mock_result = Mock() + mock_result.llm_output = { + "token_usage": { + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + } + } + + callback_handler.on_llm_end(response=mock_result) + + assert callback_handler.total_tokens == 100 + assert callback_handler.total_cost > 0 # Should estimate cost + + def test_on_agent_action(self, callback_handler, mock_client): + """Test agent action (tool call) callback.""" + # Mock AgentAction + mock_action = Mock() + mock_action.tool = "search" + mock_action.tool_input = "AI safety" + mock_action.log = "Searching for AI safety" + + callback_handler.on_agent_action(action=mock_action) + + # Verify tool call was logged as interaction + mock_client.agent_workflows.create_interaction.assert_called_once() + call_args = mock_client.agent_workflows.create_interaction.call_args + + assert call_args.kwargs["workflow_id"] == "workflow_123" + assert call_args.kwargs["from_agent"] == "test_agent" + assert call_args.kwargs["to_agent"] == "tool" + assert call_args.kwargs["interaction_type"] == "tool_call" + assert "search" in call_args.kwargs["message"] + + # Verify tool count incremented + assert callback_handler.tool_call_count == 1 + + def test_on_tool_end(self, callback_handler, mock_client): + """Test tool end callback.""" + output = "Search results for AI safety" + + callback_handler.on_tool_end(output=output) + + # Verify tool result was logged as interaction + mock_client.agent_workflows.create_interaction.assert_called_once() + call_args = mock_client.agent_workflows.create_interaction.call_args + + assert call_args.kwargs["from_agent"] == "tool" + assert call_args.kwargs["to_agent"] == "test_agent" + assert call_args.kwargs["interaction_type"] == "response" + assert output in call_args.kwargs["message"] + + def test_on_tool_error(self, callback_handler, mock_client): + """Test tool error callback.""" + error = RuntimeError("Tool failed") + + callback_handler.on_tool_error(error=error) + + # Verify tool error was logged + mock_client.agent_workflows.create_interaction.assert_called_once() + call_args = mock_client.agent_workflows.create_interaction.call_args + + assert call_args.kwargs["interaction_type"] == "response" + assert "Tool error" in call_args.kwargs["message"] + assert "RuntimeError" in call_args.kwargs["meta_data"]["error_type"] + + +class TestLangGraphMultiAgentMonitor: + """Tests for LangGraphMultiAgentMonitor.""" + + def test_init(self, mock_client): + """Test monitor initialization.""" + monitor = LangGraphMultiAgentMonitor( + client=mock_client, + workflow_name="Test Workflow", + meta_data={"key": "value"}, + ) + + assert monitor.client == mock_client + assert monitor.workflow_name == "Test Workflow" + assert monitor.workflow_meta_data == {"key": "value"} + assert monitor.workflow_id is None + assert len(monitor.callbacks) == 0 + + def test_start_monitoring(self, mock_client, mock_workflow_response): + """Test starting workflow monitoring.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + monitor = LangGraphMultiAgentMonitor( + client=mock_client, workflow_name="Test Workflow" + ) + + inputs = {"query": "test"} + workflow_id = monitor.start_monitoring(inputs=inputs) + + # Verify workflow was created + mock_client.agent_workflows.create.assert_called_once_with( + name="Test Workflow", framework="langchain", inputs=inputs, meta_data={} + ) + + # Verify workflow was started + mock_client.agent_workflows.start.assert_called_once_with("workflow_123") + + assert workflow_id == "workflow_123" + assert monitor.workflow_id == "workflow_123" + assert monitor.start_time is not None + + def test_register_agent(self, mock_client, mock_workflow_response): + """Test agent registration.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + monitor = LangGraphMultiAgentMonitor( + client=mock_client, workflow_name="Test Workflow" + ) + monitor.start_monitoring() + + monitor.register_agent( + agent_name="researcher", + role="Research Agent", + model_name="gpt-4", + tools=["search", "scrape"], + ) + + # Verify agent was registered + mock_client.agent_workflows.register_agent.assert_called_once_with( + workflow_id="workflow_123", + name="researcher", + role="Research Agent", + model_name="gpt-4", + tools=["search", "scrape"], + ) + + def test_get_callbacks(self, mock_client, mock_workflow_response): + """Test getting callbacks for an agent.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + monitor = LangGraphMultiAgentMonitor( + client=mock_client, workflow_name="Test Workflow" + ) + monitor.start_monitoring() + + callbacks = monitor.get_callbacks("researcher", agent_role="Research Agent") + + assert len(callbacks) == 1 + assert isinstance(callbacks[0], MultiAgentCallbackHandler) + assert callbacks[0].agent_name == "researcher" + assert callbacks[0].agent_role == "Research Agent" + + # Getting callbacks again should return the same instance + callbacks2 = monitor.get_callbacks("researcher") + assert callbacks2[0] is callbacks[0] + + def test_log_handoff(self, mock_client, mock_workflow_response): + """Test logging agent-to-agent handoff.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + monitor = LangGraphMultiAgentMonitor( + client=mock_client, workflow_name="Test Workflow" + ) + monitor.start_monitoring() + + monitor.log_handoff( + from_agent="supervisor", + to_agent="researcher", + message="Please research this topic", + meta_data={"topic": "AI safety"}, + ) + + # Verify handoff was logged + mock_client.agent_workflows.create_interaction.assert_called_once_with( + workflow_id="workflow_123", + from_agent="supervisor", + to_agent="researcher", + interaction_type="handoff", + message="Please research this topic", + meta_data={"topic": "AI safety"}, + ) + + def test_complete_monitoring(self, mock_client, mock_workflow_response): + """Test completing workflow monitoring.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + mock_client.agent_workflows.get_analytics.return_value = { + "total_cost": 0.15, + "total_tokens": 750, + } + + monitor = LangGraphMultiAgentMonitor( + client=mock_client, workflow_name="Test Workflow" + ) + monitor.start_monitoring() + + outputs = {"result": "completed successfully"} + summary = monitor.complete_monitoring(outputs=outputs, status="completed") + + # Verify workflow was completed + mock_client.agent_workflows.complete.assert_called_once_with( + workflow_id="workflow_123", outputs=outputs, status="completed" + ) + + # Verify analytics were retrieved + mock_client.agent_workflows.get_analytics.assert_called_once_with( + "workflow_123" + ) + + assert summary["workflow_id"] == "workflow_123" + assert summary["status"] == "completed" + assert summary["outputs"] == outputs + assert "analytics" in summary + assert summary["analytics"]["total_cost"] == 0.15 + + +class TestMonitorLangChainAgent: + """Tests for monitor_langchain_agent helper function.""" + + def test_successful_execution(self, mock_client, mock_workflow_response): + """Test successful agent execution monitoring.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + # Mock agent executor + mock_executor = Mock() + mock_executor.run.return_value = "Agent result" + + result_dict = monitor_langchain_agent( + client=mock_client, + agent_executor=mock_executor, + workflow_name="Test Task", + agent_name="test_agent", + inputs={"input": "test query"}, + ) + + # Verify workflow was created + mock_client.agent_workflows.create.assert_called_once() + + # Verify workflow was started + mock_client.agent_workflows.start.assert_called_once_with("workflow_123") + + # Verify agent was run with callback + mock_executor.run.assert_called_once() + assert "callbacks" in mock_executor.run.call_args.kwargs + + # Verify workflow was completed + mock_client.agent_workflows.complete.assert_called_once() + complete_args = mock_client.agent_workflows.complete.call_args + assert complete_args.args[0] == "workflow_123" + assert complete_args.kwargs["outputs"] == {"result": "Agent result"} + + # Verify result + assert result_dict["result"] == "Agent result" + assert result_dict["workflow_id"] == "workflow_123" + assert result_dict["status"] == "completed" + + def test_failed_execution(self, mock_client, mock_workflow_response): + """Test failed agent execution monitoring.""" + mock_client.agent_workflows.create.return_value = mock_workflow_response + + # Mock agent executor that raises error + mock_executor = Mock() + mock_executor.run.side_effect = RuntimeError("Agent failed") + + result_dict = monitor_langchain_agent( + client=mock_client, + agent_executor=mock_executor, + workflow_name="Test Task", + agent_name="test_agent", + ) + + # Verify workflow was marked as failed + complete_args = mock_client.agent_workflows.complete.call_args + assert complete_args.kwargs["status"] == "failed" + assert "error" in complete_args.kwargs["outputs"] + + # Verify result + assert result_dict["result"] is None + assert result_dict["status"] == "failed" + assert "Agent failed" in result_dict["error"] + + +class TestErrorHandling: + """Tests for error handling.""" + + def test_callback_api_error_handling(self, mock_client): + """Test that API errors don't crash the callback.""" + mock_client.agent_workflows.create_execution.side_effect = Exception( + "API Error" + ) + + handler = MultiAgentCallbackHandler( + client=mock_client, workflow_id="workflow_123", agent_name="test" + ) + + handler.execution_start_time = datetime.utcnow() + handler.execution_inputs = {} + + # Should not raise exception + handler.on_chain_end(outputs={"result": "test"}) + + # State should still be reset + assert handler.execution_start_time is None + + def test_register_agent_before_start(self, mock_client): + """Test that registering agent before start raises error.""" + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test") + + with pytest.raises(ValueError, match="Must call start_monitoring"): + monitor.register_agent("test_agent") + + def test_get_callbacks_before_start(self, mock_client): + """Test that getting callbacks before start raises error.""" + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test") + + with pytest.raises(ValueError, match="Must call start_monitoring"): + monitor.get_callbacks("test_agent") diff --git a/tests/integration/test_sklearn.py b/tests/integration/test_sklearn.py index e895047..184e5ae 100644 --- a/tests/integration/test_sklearn.py +++ b/tests/integration/test_sklearn.py @@ -8,7 +8,7 @@ from sklearn.datasets import make_classification # noqa: E402 from sklearn.ensemble import RandomForestClassifier # noqa: E402 -from whiteboxai.integrations.sklearn import SklearnMonitor # noqa: E402 +from whiteboxxai.integrations.sklearn import SklearnMonitor # noqa: E402 class TestSklearnIntegration: diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 3486e4a..7d34e1b 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1,30 +1,229 @@ -"""Unit tests for WhiteBoxAI client.""" +""" +Tests for WhiteBoxXAI SDK Client +Tests for the main SDK client class. +""" + +from unittest.mock import Mock, patch + +import httpx import pytest -from whiteboxai import WhiteBoxAI -from whiteboxai.exceptions import AuthenticationError +from whiteboxxai.client import WhiteBoxXAI +from whiteboxxai.config import Config +from whiteboxxai.exceptions import ( + APIError, + AuthenticationError, + RateLimitError, + ValidationError, +) + + +class TestWhiteBoxXAIClient: + """Tests for WhiteBoxXAI client initialization and configuration.""" + + def test_client_initialization(self): + """Test basic client initialization.""" + client = WhiteBoxXAI(api_key="test_key") + assert client.config.api_key == "test_key" + assert client.config.base_url == "https://api.whiteboxxai.com" + assert client.config.timeout == 30 + assert client.config.max_retries == 3 + + def test_client_custom_config(self): + """Test client with custom configuration.""" + client = WhiteBoxXAI( + api_key="test_key", + base_url="https://custom.api.com", + timeout=60, + max_retries=5, + ) + assert client.config.base_url == "https://custom.api.com" + assert client.config.timeout == 60 + assert client.config.max_retries == 5 + + def test_client_resources_initialized(self): + """Test that resource managers are initialized.""" + client = WhiteBoxXAI(api_key="test_key") + assert client.models is not None + assert client.predictions is not None + assert client.explanations is not None + assert client.drift is not None + assert client.alerts is not None + + def test_sync_client_property(self): + """Test synchronous client property.""" + client = WhiteBoxXAI(api_key="test_key") + sync_client = client.sync_client + assert isinstance(sync_client, httpx.Client) + assert sync_client.base_url == "https://api.whiteboxxai.com" + + def test_async_client_property(self): + """Test asynchronous client property.""" + client = WhiteBoxXAI(api_key="test_key") + async_client = client.async_client + assert isinstance(async_client, httpx.AsyncClient) + assert async_client.base_url == "https://api.whiteboxxai.com" + + def test_headers_include_api_key(self): + """Test that headers include API key.""" + client = WhiteBoxXAI(api_key="test_secret_key") + headers = client._get_headers() + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test_secret_key" + + def test_client_close(self): + """Test client cleanup.""" + client = WhiteBoxXAI(api_key="test_key") + # Access clients to create them + _ = client.sync_client + _ = client.async_client + + # Close should not raise + client.close() + + def test_context_manager(self): + """Test client as context manager.""" + with WhiteBoxXAI(api_key="test_key") as client: + assert client.config.api_key == "test_key" + + +class TestClientHTTPMethods: + """Tests for client HTTP request methods.""" + + @patch("httpx.Client.request") + def test_request_success(self, mock_request): + """Test successful HTTP request.""" + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"id": "123", "name": "test"} + mock_request.return_value = mock_response + + client = WhiteBoxXAI(api_key="test_key") + response = client.request("GET", "/models/123") + + assert response == {"id": "123", "name": "test"} + mock_request.assert_called_once() + + @patch("httpx.Client.request") + def test_request_authentication_error(self, mock_request): + """Test authentication error handling.""" + mock_response = Mock() + mock_response.status_code = 401 + mock_response.json.return_value = {"error": "Unauthorized"} + mock_request.return_value = mock_response + + client = WhiteBoxXAI(api_key="invalid_key") + with pytest.raises(AuthenticationError): + client.request("GET", "/models") + + @patch("httpx.Client.request") + def test_request_validation_error(self, mock_request): + """Test validation error handling.""" + mock_response = Mock() + mock_response.status_code = 422 + mock_response.json.return_value = {"detail": "Validation failed"} + mock_request.return_value = mock_response + + client = WhiteBoxXAI(api_key="test_key") + with pytest.raises(ValidationError): + client.request("POST", "/models", data={"invalid": "data"}) + + @patch("httpx.Client.request") + def test_request_rate_limit_error(self, mock_request): + """Test rate limit error handling.""" + mock_response = Mock() + mock_response.status_code = 429 + mock_response.json.return_value = {"error": "Rate limit exceeded"} + mock_request.return_value = mock_response + + client = WhiteBoxXAI(api_key="test_key") + with pytest.raises(RateLimitError): + client.request("GET", "/models") + + @patch("httpx.Client.request") + def test_request_server_error(self, mock_request): + """Test server error handling.""" + mock_response = Mock() + mock_response.status_code = 500 + mock_response.json.return_value = {"error": "Internal server error"} + mock_request.return_value = mock_response + + client = WhiteBoxXAI(api_key="test_key") + with pytest.raises(APIError): + client.request("GET", "/models") + + @patch("httpx.Client.request") + def test_request_with_retry(self, mock_request): + """Test request retry on transient failures.""" + # First call fails, second succeeds + mock_response_error = Mock() + mock_response_error.status_code = 503 + mock_response_error.json.return_value = {"error": "Service unavailable"} + + mock_response_success = Mock() + mock_response_success.status_code = 200 + mock_response_success.json.return_value = {"id": "123"} + + mock_request.side_effect = [mock_response_error, mock_response_success] + + client = WhiteBoxXAI(api_key="test_key", max_retries=2) + # Note: Retry behavior depends on implementation + # This test verifies the interface exists + + +class TestClientConfiguration: + """Tests for client configuration.""" + + def test_config_validation(self): + """Test configuration validation.""" + config = Config(api_key="test_key", base_url="https://api.test.com") + assert config.api_key == "test_key" + assert config.base_url == "https://api.test.com" + + def test_config_defaults(self): + """Test configuration defaults.""" + config = Config(api_key="test_key") + assert config.timeout == 30 + assert config.max_retries == 3 + assert config.base_url == "https://api.whiteboxxai.com" + + def test_invalid_api_key(self): + """Test that empty API key raises error.""" + with pytest.raises((ValueError, ValidationError)): + Config(api_key="") + +class TestClientResourceIntegration: + """Tests for client resource integration.""" -class TestWhiteBoxAIClient: - """Test cases for WhiteBoxAI client.""" + def test_models_resource_access(self): + """Test accessing models resource.""" + client = WhiteBoxXAI(api_key="test_key") + assert hasattr(client.models, "register") + assert hasattr(client.models, "list") + assert hasattr(client.models, "get") - def test_client_initialization(self, mock_api_key): - """Test client can be initialized with API key.""" - client = WhiteBoxAI(api_key=mock_api_key) - assert client is not None + def test_predictions_resource_access(self): + """Test accessing predictions resource.""" + client = WhiteBoxXAI(api_key="test_key") + assert hasattr(client.predictions, "log") + assert hasattr(client.predictions, "query") - def test_client_requires_api_key(self): - """Test client raises error without API key.""" - with pytest.raises((ValueError, AuthenticationError)): - WhiteBoxAI(api_key=None) + def test_explanations_resource_access(self): + """Test accessing explanations resource.""" + client = WhiteBoxXAI(api_key="test_key") + assert hasattr(client.explanations, "generate") + assert hasattr(client.explanations, "get") - def test_client_with_custom_base_url(self, mock_api_key): - """Test client accepts custom base URL.""" - client = WhiteBoxAI(api_key=mock_api_key, base_url="https://custom.api.example.com") - assert client is not None + def test_drift_resource_access(self): + """Test accessing drift resource.""" + client = WhiteBoxXAI(api_key="test_key") + assert hasattr(client.drift, "detect") + assert hasattr(client.drift, "get_reports") - def test_offline_mode_configuration(self, mock_api_key): - """Test offline mode can be enabled.""" - client = WhiteBoxAI(api_key=mock_api_key, enable_offline=True, offline_dir="./test_offline") - assert client is not None + def test_alerts_resource_access(self): + """Test accessing alerts resource.""" + client = WhiteBoxXAI(api_key="test_key") + assert hasattr(client.alerts, "list") + assert hasattr(client.alerts, "create") diff --git a/tests/unit/test_config_exceptions.py b/tests/unit/test_config_exceptions.py new file mode 100644 index 0000000..485acd7 --- /dev/null +++ b/tests/unit/test_config_exceptions.py @@ -0,0 +1,301 @@ +""" +Tests for SDK Config and Exceptions + +Tests for SDK configuration and exception handling. +""" + +import pytest + +from whiteboxxai.config import Config +from whiteboxxai.exceptions import ( + APIError, + AuthenticationError, + NotFoundError, + RateLimitError, + ValidationError, + WhiteBoxXAIError, +) + + +class TestConfig: + """Tests for Config class.""" + + def test_config_basic(self): + """Test basic configuration.""" + config = Config(api_key="test-key-123") + assert config.api_key == "test-key-123" + assert config.base_url == "https://api.whiteboxxai.com" + assert config.timeout == 30 + assert config.max_retries == 3 + + def test_config_custom_base_url(self): + """Test configuration with custom base URL.""" + config = Config(api_key="test-key", base_url="https://custom.api.com") + assert config.base_url == "https://custom.api.com" + + def test_config_custom_timeout(self): + """Test configuration with custom timeout.""" + config = Config(api_key="test-key", timeout=60) + assert config.timeout == 60 + + def test_config_custom_max_retries(self): + """Test configuration with custom max retries.""" + config = Config(api_key="test-key", max_retries=5) + assert config.max_retries == 5 + + def test_config_additional_params(self): + """Test configuration with additional parameters.""" + config = Config(api_key="test-key", custom_param="value", another_param=123) + assert config.api_key == "test-key" + # Additional params should be accessible + + def test_config_empty_api_key(self): + """Test that empty API key raises error.""" + with pytest.raises(ValueError): + Config(api_key="") + + def test_config_none_api_key(self): + """Test that None API key raises error.""" + with pytest.raises((ValueError, TypeError)): + Config(api_key=None) + + def test_config_invalid_timeout(self): + """Test that invalid timeout raises error.""" + with pytest.raises((ValueError, TypeError)): + Config(api_key="test-key", timeout=-1) + + def test_config_invalid_max_retries(self): + """Test that invalid max_retries raises error.""" + with pytest.raises((ValueError, TypeError)): + Config(api_key="test-key", max_retries=-1) + + def test_config_base_url_validation(self): + """Test base URL validation.""" + # Valid URLs + config1 = Config(api_key="test-key", base_url="https://api.example.com") + assert config1.base_url.startswith("https://") + + # Invalid URL (no protocol) - should handle gracefully or raise + try: + config2 = Config(api_key="test-key", base_url="invalid-url") + # If it accepts, it should add protocol or validate differently + except ValueError: + # Acceptable to reject invalid URLs + pass + + +class TestWhiteBoxXAIError: + """Tests for base WhiteBoxXAIError exception.""" + + def test_base_error_creation(self): + """Test creating base error.""" + error = WhiteBoxXAIError("Something went wrong") + assert str(error) == "Something went wrong" + + def test_base_error_with_details(self): + """Test base error with additional details.""" + error = WhiteBoxXAIError("Error occurred", details={"code": "ERR001"}) + assert "Error occurred" in str(error) + + def test_base_error_inheritance(self): + """Test that base error inherits from Exception.""" + error = WhiteBoxXAIError("Test") + assert isinstance(error, Exception) + + +class TestAuthenticationError: + """Tests for AuthenticationError exception.""" + + def test_authentication_error_creation(self): + """Test creating authentication error.""" + error = AuthenticationError("Invalid API key") + assert "Invalid API key" in str(error) + + def test_authentication_error_inheritance(self): + """Test that authentication error inherits from WhiteBoxXAIError.""" + error = AuthenticationError("Test") + assert isinstance(error, WhiteBoxXAIError) + assert isinstance(error, Exception) + + def test_authentication_error_with_response(self): + """Test authentication error with response details.""" + error = AuthenticationError( + "Unauthorized", status_code=401, response={"error": "invalid_token"} + ) + assert "Unauthorized" in str(error) + + +class TestValidationError: + """Tests for ValidationError exception.""" + + def test_validation_error_creation(self): + """Test creating validation error.""" + error = ValidationError("Invalid input data") + assert "Invalid input data" in str(error) + + def test_validation_error_inheritance(self): + """Test that validation error inherits from WhiteBoxXAIError.""" + error = ValidationError("Test") + assert isinstance(error, WhiteBoxXAIError) + + def test_validation_error_with_fields(self): + """Test validation error with field details.""" + error = ValidationError( + "Validation failed", fields={"email": ["Invalid email format"]} + ) + assert "Validation failed" in str(error) + + +class TestNotFoundError: + """Tests for NotFoundError exception.""" + + def test_not_found_error_creation(self): + """Test creating not found error.""" + error = NotFoundError("Model not found") + assert "Model not found" in str(error) + + def test_not_found_error_inheritance(self): + """Test that not found error inherits from WhiteBoxXAIError.""" + error = NotFoundError("Test") + assert isinstance(error, WhiteBoxXAIError) + + def test_not_found_error_with_resource(self): + """Test not found error with resource details.""" + error = NotFoundError( + "Resource not found", resource_type="model", resource_id="123" + ) + assert "Resource not found" in str(error) + + +class TestRateLimitError: + """Tests for RateLimitError exception.""" + + def test_rate_limit_error_creation(self): + """Test creating rate limit error.""" + error = RateLimitError("Rate limit exceeded") + assert "Rate limit exceeded" in str(error) + + def test_rate_limit_error_inheritance(self): + """Test that rate limit error inherits from WhiteBoxXAIError.""" + error = RateLimitError("Test") + assert isinstance(error, WhiteBoxXAIError) + + def test_rate_limit_error_with_retry(self): + """Test rate limit error with retry information.""" + error = RateLimitError( + "Too many requests", retry_after=60, limit=100, remaining=0 + ) + assert "Too many requests" in str(error) + + +class TestAPIError: + """Tests for APIError exception.""" + + def test_api_error_creation(self): + """Test creating API error.""" + error = APIError("API request failed") + assert "API request failed" in str(error) + + def test_api_error_inheritance(self): + """Test that API error inherits from WhiteBoxXAIError.""" + error = APIError("Test") + assert isinstance(error, WhiteBoxXAIError) + + def test_api_error_with_status(self): + """Test API error with status code.""" + error = APIError( + "Server error", status_code=500, response={"error": "Internal error"} + ) + assert "Server error" in str(error) + + def test_api_error_with_request_id(self): + """Test API error with request ID.""" + error = APIError("Request failed", request_id="req-abc-123") + assert "Request failed" in str(error) + + +class TestExceptionRaising: + """Tests for raising exceptions in different scenarios.""" + + def test_raise_authentication_error(self): + """Test raising authentication error.""" + with pytest.raises(AuthenticationError) as exc_info: + raise AuthenticationError("Invalid credentials") + assert "Invalid credentials" in str(exc_info.value) + + def test_raise_validation_error(self): + """Test raising validation error.""" + with pytest.raises(ValidationError) as exc_info: + raise ValidationError("Invalid data format") + assert "Invalid data format" in str(exc_info.value) + + def test_raise_not_found_error(self): + """Test raising not found error.""" + with pytest.raises(NotFoundError) as exc_info: + raise NotFoundError("Resource not found") + assert "Resource not found" in str(exc_info.value) + + def test_raise_rate_limit_error(self): + """Test raising rate limit error.""" + with pytest.raises(RateLimitError) as exc_info: + raise RateLimitError("Rate limit exceeded") + assert "Rate limit exceeded" in str(exc_info.value) + + def test_raise_api_error(self): + """Test raising API error.""" + with pytest.raises(APIError) as exc_info: + raise APIError("API error occurred") + assert "API error occurred" in str(exc_info.value) + + +class TestExceptionCatching: + """Tests for catching exception hierarchy.""" + + def test_catch_specific_exception(self): + """Test catching specific exception type.""" + try: + raise ValidationError("Validation failed") + except ValidationError as e: + assert isinstance(e, ValidationError) + + def test_catch_base_exception(self): + """Test catching via base WhiteBoxXAIError.""" + try: + raise AuthenticationError("Auth failed") + except WhiteBoxXAIError as e: + assert isinstance(e, WhiteBoxXAIError) + assert isinstance(e, AuthenticationError) + + def test_catch_multiple_exceptions(self): + """Test catching multiple exception types.""" + for error_class in [ + AuthenticationError, + ValidationError, + NotFoundError, + RateLimitError, + APIError, + ]: + try: + raise error_class("Test error") + except WhiteBoxXAIError as e: + assert isinstance(e, WhiteBoxXAIError) + + +class TestExceptionAttributes: + """Tests for exception attributes and methods.""" + + def test_error_message_attribute(self): + """Test that error message is accessible.""" + error = ValidationError("Invalid input") + assert error.args[0] == "Invalid input" + + def test_error_string_representation(self): + """Test error string representation.""" + error = APIError("Request failed") + assert str(error) == "Request failed" + + def test_error_repr(self): + """Test error repr.""" + error = NotFoundError("Not found") + repr_str = repr(error) + assert "NotFoundError" in repr_str or "Not found" in repr_str diff --git a/tests/unit/test_decorators.py b/tests/unit/test_decorators.py new file mode 100644 index 0000000..c2bef32 --- /dev/null +++ b/tests/unit/test_decorators.py @@ -0,0 +1,297 @@ +""" +Tests for SDK Decorators + +Tests for monitoring decorators. +""" + +import time +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from whiteboxxai.decorators import ( + _extract_inputs, + _extract_output, + monitor_model, + monitor_performance, +) +from whiteboxxai.monitor import ModelMonitor + + +class TestMonitorModelDecorator: + """Tests for monitor_model decorator.""" + + def test_decorator_basic_usage(self): + """Test basic decorator usage.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, input_keys=["x"], explain=False) + def predict(x): + return x * 2 + + result = predict(x=5) + assert result == 10 + # Verify monitor was called + mock_monitor.log_prediction.assert_called_once() + + def test_decorator_extracts_inputs(self): + """Test that decorator extracts inputs correctly.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, input_keys=["features"], explain=False) + def predict(features): + return {"prediction": sum(features)} + + result = predict(features=[1, 2, 3]) + + # Verify inputs were extracted + call_args = mock_monitor.log_prediction.call_args + assert call_args[1]["inputs"] == {"features": [1, 2, 3]} + + def test_decorator_extracts_output(self): + """Test that decorator extracts output correctly.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model( + mock_monitor, input_keys=["x"], output_key="prediction", explain=False + ) + def predict(x): + return {"prediction": x * 2, "confidence": 0.95} + + result = predict(x=5) + + # Verify output was extracted + call_args = mock_monitor.log_prediction.call_args + assert call_args[1]["output"] == 10 # x * 2 + + def test_decorator_with_explain(self): + """Test decorator with explain=True.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, input_keys=["x"], explain=True) + def predict(x): + return x * 2 + + result = predict(x=5) + + # Verify explain was passed + call_args = mock_monitor.log_prediction.call_args + assert call_args[1]["explain"] is True + + def test_decorator_tracks_inference_time(self): + """Test that decorator tracks inference time.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, input_keys=["x"], explain=False) + def predict(x): + time.sleep(0.1) # Simulate processing + return x * 2 + + result = predict(x=5) + + # Verify metadata includes inference time + call_args = mock_monitor.log_prediction.call_args + metadata = call_args[1]["metadata"] + assert "inference_time_ms" in metadata + assert metadata["inference_time_ms"] >= 100 # At least 100ms + + def test_decorator_preserves_function_metadata(self): + """Test that decorator preserves function metadata.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, input_keys=["x"], explain=False) + def my_predict_function(x): + """My prediction function.""" + return x * 2 + + assert my_predict_function.__name__ == "my_predict_function" + assert "My prediction function" in my_predict_function.__doc__ + + @pytest.mark.asyncio + async def test_decorator_async_function(self): + """Test decorator with async function.""" + mock_monitor = Mock(spec=ModelMonitor) + mock_monitor.alog_prediction = AsyncMock() + + @monitor_model(mock_monitor, input_keys=["x"], explain=False) + async def predict_async(x): + return x * 2 + + result = await predict_async(x=5) + assert result == 10 + # Verify async log was called + mock_monitor.alog_prediction.assert_called_once() + + +class TestExtractInputs: + """Tests for _extract_inputs helper function.""" + + def test_extract_from_kwargs(self): + """Test extracting inputs from kwargs.""" + + def dummy_func(x, y): + pass + + inputs = _extract_inputs( + dummy_func, + args=(), + kwargs={"x": 1, "y": 2}, + input_keys=["x", "y"], + ) + assert inputs == {"x": 1, "y": 2} + + def test_extract_from_args(self): + """Test extracting inputs from positional args.""" + + def dummy_func(x, y): + pass + + inputs = _extract_inputs( + dummy_func, + args=(1, 2), + kwargs={}, + input_keys=["x", "y"], + ) + assert inputs == {"x": 1, "y": 2} + + def test_extract_mixed_args_kwargs(self): + """Test extracting from both args and kwargs.""" + + def dummy_func(x, y, z): + pass + + inputs = _extract_inputs( + dummy_func, + args=(1,), + kwargs={"y": 2, "z": 3}, + input_keys=["x", "y", "z"], + ) + assert inputs == {"x": 1, "y": 2, "z": 3} + + def test_extract_without_input_keys(self): + """Test extraction without input_keys (use all args).""" + + def dummy_func(x, y): + pass + + inputs = _extract_inputs( + dummy_func, + args=(1, 2), + kwargs={}, + input_keys=None, + ) + # Should extract all arguments + assert "x" in inputs or len(inputs) > 0 + + +class TestExtractOutput: + """Tests for _extract_output helper function.""" + + def test_extract_with_key(self): + """Test extracting output with specified key.""" + result = {"prediction": 42, "confidence": 0.95} + output = _extract_output(result, output_key="prediction") + assert output == 42 + + def test_extract_without_key(self): + """Test extracting output without key (use whole result).""" + result = {"prediction": 42} + output = _extract_output(result, output_key=None) + assert output == {"prediction": 42} + + def test_extract_from_simple_value(self): + """Test extracting from simple value.""" + result = 42 + output = _extract_output(result, output_key=None) + assert output == 42 + + def test_extract_missing_key(self): + """Test extracting with missing key.""" + result = {"prediction": 42} + output = _extract_output(result, output_key="nonexistent") + # Should handle gracefully + assert output is None or output == result + + +class TestMonitorPerformanceDecorator: + """Tests for monitor_performance decorator.""" + + def test_performance_tracking(self): + """Test that performance is tracked.""" + + @monitor_performance(threshold_ms=10) + def compute_intensive_task(): + time.sleep(0.05) + return "done" + + result = compute_intensive_task() + assert result == "done" + + def test_performance_with_errors(self): + """Test performance tracking with errors.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_performance(mock_monitor) + def failing_task(): + raise ValueError("Test error") + + with pytest.raises(ValueError): + failing_task() + # Verify error was logged + # (Implementation-specific assertion) + + +class TestDecoratorEdgeCases: + """Tests for decorator edge cases.""" + + def test_decorator_with_no_args_function(self): + """Test decorator on function with no arguments.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, explain=False) + def predict_constant(): + return 42 + + result = predict_constant() + assert result == 42 + mock_monitor.log_prediction.assert_called_once() + + def test_decorator_with_varargs(self): + """Test decorator on function with *args.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, explain=False) + def predict_sum(*numbers): + return sum(numbers) + + result = predict_sum(1, 2, 3, 4) + assert result == 10 + mock_monitor.log_prediction.assert_called_once() + + def test_decorator_with_kwargs_only(self): + """Test decorator on function with **kwargs.""" + mock_monitor = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor, explain=False) + def predict_with_options(**options): + return options.get("value", 0) * 2 + + result = predict_with_options(value=5) + assert result == 10 + mock_monitor.log_prediction.assert_called_once() + + def test_decorator_multiple_decorators(self): + """Test stacking multiple decorators.""" + mock_monitor1 = Mock(spec=ModelMonitor) + mock_monitor2 = Mock(spec=ModelMonitor) + + @monitor_model(mock_monitor1, input_keys=["x"], explain=False) + @monitor_model(mock_monitor2, input_keys=["x"], explain=False) + def predict(x): + return x * 2 + + result = predict(x=5) + assert result == 10 + # Both monitors should be called + mock_monitor1.log_prediction.assert_called_once() + mock_monitor2.log_prediction.assert_called_once() diff --git a/tests/unit/test_monitor.py b/tests/unit/test_monitor.py index 3e05595..fa8bbce 100644 --- a/tests/unit/test_monitor.py +++ b/tests/unit/test_monitor.py @@ -1,26 +1,327 @@ -"""Unit tests for ModelMonitor.""" +""" +Tests for SDK Monitor Module + +Tests for model monitoring functionality. +""" + +from unittest.mock import AsyncMock, Mock, patch import pytest -from whiteboxai import ModelMonitor +from whiteboxxai.client import WhiteBoxXAI +from whiteboxxai.monitor import ModelMonitor + + +class TestModelMonitorInitialization: + """Tests for ModelMonitor initialization.""" + + def test_monitor_initialization(self): + """Test basic monitor initialization.""" + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor(client=mock_client, model_id="model-123") + assert monitor.model_id == "model-123" + assert monitor.client == mock_client + + def test_monitor_with_config(self): + """Test monitor with custom configuration.""" + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor( + client=mock_client, + model_id="model-123", + auto_explain=True, + buffer_size=100, + ) + assert monitor.auto_explain is True + assert monitor.buffer_size == 100 + + +class TestPredictionLogging: + """Tests for prediction logging.""" + + @patch("whiteboxxai.monitor.ModelMonitor.log_prediction") + def test_log_prediction_basic(self, mock_log): + """Test basic prediction logging.""" + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + inputs = {"feature1": 1.0, "feature2": 2.0} + output = 0.85 + + monitor.log_prediction(inputs=inputs, output=output) + # Verify log was called + mock_log.assert_called_once_with(inputs=inputs, output=output) -class TestModelMonitor: - """Test cases for ModelMonitor.""" + def test_log_prediction_with_metadata(self): + """Test prediction logging with metadata.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log = Mock() - def test_monitor_initialization(self, client): - """Test monitor can be initialized with client.""" - monitor = ModelMonitor(client) - assert monitor is not None + monitor = ModelMonitor(client=mock_client, model_id="model-123") - def test_monitor_with_sampling(self, client): - """Test monitor with sampling rate.""" - monitor = ModelMonitor(client, sampling_rate=0.5) - assert monitor is not None + inputs = [1.0, 2.0, 3.0] + output = 0.75 + metadata = {"version": "1.0", "region": "us-east"} + + monitor.log_prediction(inputs=inputs, output=output, metadata=metadata) + + # Verify metadata was included + call_args = mock_client.predictions.log.call_args + assert call_args is not None + + def test_log_prediction_with_explain(self): + """Test prediction logging with explanation.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.explanations = Mock() + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + inputs = [1.0, 2.0] + output = 0.9 + + monitor.log_prediction(inputs=inputs, output=output, explain=True) + + # Verify explanation was requested + # (Implementation-specific verification) @pytest.mark.asyncio - async def test_async_operations(self, client): - """Test async monitor operations.""" - monitor = ModelMonitor(client) - # Add async tests when implementing - assert monitor is not None + async def test_async_log_prediction(self): + """Test async prediction logging.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.alog = AsyncMock() + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + inputs = [1.0, 2.0] + output = 0.9 + + await monitor.alog_prediction(inputs=inputs, output=output) + + # Verify async log was called + mock_client.predictions.alog.assert_called_once() + + +class TestBatchPredictionLogging: + """Tests for batch prediction logging.""" + + def test_log_batch_predictions(self): + """Test logging batch of predictions.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log_batch = Mock() + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + batch = [ + {"inputs": [1.0, 2.0], "output": 0.8}, + {"inputs": [2.0, 3.0], "output": 0.9}, + ] + + monitor.log_batch(batch) + + # Verify batch log was called + mock_client.predictions.log_batch.assert_called_once() + + def test_log_batch_with_buffer(self): + """Test batch logging with buffering.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log_batch = Mock() + monitor = ModelMonitor(client=mock_client, model_id="model-123", buffer_size=3) + + # Log predictions below buffer size + monitor.log_prediction(inputs=[1.0], output=0.5) + monitor.log_prediction(inputs=[2.0], output=0.6) + + # Buffer should not flush yet + # (Implementation-specific verification) + + # Log one more to exceed buffer + monitor.log_prediction(inputs=[3.0], output=0.7) + + # Should have flushed + # (Implementation-specific verification) + + +class TestDriftMonitoring: + """Tests for drift monitoring functionality.""" + + def test_detect_drift(self): + """Test drift detection.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.drift = Mock() + mock_client.drift.detect = Mock(return_value={"drift_detected": True}) + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + result = monitor.detect_drift(window_size=100) + + # Verify drift detection was called + mock_client.drift.detect.assert_called_once() + assert result["drift_detected"] is True + + def test_get_drift_reports(self): + """Test getting drift reports.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.drift = Mock() + mock_client.drift.get_reports = Mock( + return_value=[{"id": "report-1", "drift_score": 0.8}] + ) + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + reports = monitor.get_drift_reports(limit=10) + + # Verify reports were retrieved + assert len(reports) == 1 + assert reports[0]["drift_score"] == 0.8 + + +class TestAlertManagement: + """Tests for alert management.""" + + def test_create_alert_rule(self): + """Test creating alert rule.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.alerts = Mock() + mock_client.alerts.create = Mock(return_value={"id": "rule-1"}) + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + rule = monitor.create_alert_rule( + metric="accuracy", threshold=0.8, condition="below" + ) + + # Verify alert rule was created + assert rule["id"] == "rule-1" + + def test_list_active_alerts(self): + """Test listing active alerts.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.alerts = Mock() + mock_client.alerts.list = Mock( + return_value=[ + {"id": "alert-1", "severity": "high"}, + {"id": "alert-2", "severity": "medium"}, + ] + ) + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + alerts = monitor.get_active_alerts() + + # Verify alerts were retrieved + assert len(alerts) == 2 + assert alerts[0]["severity"] == "high" + + +class TestMonitorContextManager: + """Tests for monitor as context manager.""" + + def test_context_manager(self): + """Test monitor as context manager.""" + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + with monitor: + # Use monitor + pass + + # Verify cleanup was called + # (Implementation-specific verification) + + def test_context_manager_flushes_buffer(self): + """Test that context manager flushes buffer on exit.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log_batch = Mock() + monitor = ModelMonitor(client=mock_client, model_id="model-123", buffer_size=10) + + with monitor: + # Log some predictions + monitor.log_prediction(inputs=[1.0], output=0.5) + monitor.log_prediction(inputs=[2.0], output=0.6) + + # Buffer should be flushed on exit + # (Implementation-specific verification) + + +class TestMonitorStatistics: + """Tests for monitor statistics.""" + + def test_get_prediction_count(self): + """Test getting prediction count.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log = Mock() + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + # Log some predictions + monitor.log_prediction(inputs=[1.0], output=0.5) + monitor.log_prediction(inputs=[2.0], output=0.6) + + count = monitor.get_prediction_count() + assert count >= 2 + + @pytest.mark.skip( + reason="No MetricsResource/client.metrics exists in whiteboxxai.client yet; " + "this needs a real backend-routes-driven implementation, not a guess." + ) + def test_get_performance_metrics(self): + """Test getting performance metrics.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.metrics = Mock() + mock_client.metrics.get = Mock( + return_value={"accuracy": 0.95, "precision": 0.93} + ) + + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + metrics = monitor.get_performance_metrics() + + # Verify metrics were retrieved + assert "accuracy" in metrics + assert metrics["accuracy"] == 0.95 + + +class TestMonitorEdgeCases: + """Tests for monitor edge cases.""" + + def test_monitor_with_invalid_model_id(self): + """Test that logging with an unregistered (model_id=None) monitor raises. + + model_id=None is a legitimate construction state (matches the + default and the class's own "register later" workflow via + register_model()) so construction itself must not raise; only + operations that require a registered model should. + """ + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor(client=mock_client, model_id=None) + + with pytest.raises(ValueError): + monitor.log_prediction(inputs=[1.0], output=0.5) + + def test_log_prediction_without_output(self): + """Test logging prediction without output.""" + mock_client = Mock(spec=WhiteBoxXAI) + monitor = ModelMonitor(client=mock_client, model_id="model-123") + + with pytest.raises((ValueError, TypeError)): + monitor.log_prediction(inputs=[1.0, 2.0], output=None) + + def test_buffer_overflow_handling(self): + """Test handling of buffer overflow.""" + mock_client = Mock(spec=WhiteBoxXAI) + mock_client.predictions = Mock() + mock_client.predictions.log_batch = Mock() + monitor = ModelMonitor(client=mock_client, model_id="model-123", buffer_size=2) + + # Log more than buffer size + for i in range(10): + monitor.log_prediction(inputs=[float(i)], output=0.5) + + # Should handle overflow gracefully + # (Implementation-specific verification) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py new file mode 100644 index 0000000..a43459a --- /dev/null +++ b/tests/unit/test_utils.py @@ -0,0 +1,277 @@ +""" +Tests for SDK Utils Module + +Tests for SDK utility functions. +""" + +import numpy as np +import pytest + +from whiteboxxai.utils import ( + batch_iterator, + compute_metrics, + deserialize_numpy, + format_features, + serialize_numpy, + validate_features, + validate_model_type, + validate_prediction, +) + + +class TestFeatureValidation: + """Tests for feature validation functions.""" + + def test_validate_features_list(self): + """Test validating list of features.""" + features = [1.0, 2.0, 3.0] + result = validate_features(features) + assert result is True + + def test_validate_features_numpy(self): + """Test validating numpy array features.""" + features = np.array([1.0, 2.0, 3.0]) + result = validate_features(features) + assert result is True + + def test_validate_features_dict(self): + """Test validating dictionary features.""" + features = {"feature1": 1.0, "feature2": 2.0} + result = validate_features(features) + assert result is True + + def test_validate_features_invalid(self): + """Test validating invalid features.""" + with pytest.raises(ValueError): + validate_features("invalid") + + def test_validate_features_empty(self): + """Test validating empty features.""" + with pytest.raises(ValueError): + validate_features([]) + + +class TestFeatureFormatting: + """Tests for feature formatting functions.""" + + def test_format_features_list(self): + """Test formatting list features.""" + features = [1, 2, 3] + formatted = format_features(features) + assert isinstance(formatted, list) + assert formatted == [1, 2, 3] + + def test_format_features_numpy(self): + """Test formatting numpy array features.""" + features = np.array([1.0, 2.0, 3.0]) + formatted = format_features(features) + assert isinstance(formatted, list) + assert formatted == [1.0, 2.0, 3.0] + + def test_format_features_dict(self): + """Test formatting dictionary features.""" + features = {"a": 1, "b": 2} + formatted = format_features(features) + assert isinstance(formatted, dict) + assert formatted == {"a": 1, "b": 2} + + def test_format_features_2d_array(self): + """Test formatting 2D numpy array.""" + features = np.array([[1, 2], [3, 4]]) + formatted = format_features(features) + assert isinstance(formatted, list) + assert len(formatted) == 2 + + +class TestModelTypeValidation: + """Tests for model type validation.""" + + def test_validate_model_type_classification(self): + """Test validating classification model type.""" + result = validate_model_type("classification") + assert result is True + + def test_validate_model_type_regression(self): + """Test validating regression model type.""" + result = validate_model_type("regression") + assert result is True + + def test_validate_model_type_invalid(self): + """Test validating invalid model type.""" + with pytest.raises(ValueError): + validate_model_type("invalid_type") + + def test_validate_model_type_case_insensitive(self): + """Test that validation is case-insensitive.""" + result = validate_model_type("CLASSIFICATION") + assert result is True + + +class TestPredictionValidation: + """Tests for prediction validation.""" + + def test_validate_prediction_number(self): + """Test validating numeric prediction.""" + result = validate_prediction(42.5) + assert result is True + + def test_validate_prediction_string(self): + """Test validating string prediction (class label).""" + result = validate_prediction("fraud") + assert result is True + + def test_validate_prediction_list(self): + """Test validating list prediction (probabilities).""" + result = validate_prediction([0.2, 0.8]) + assert result is True + + def test_validate_prediction_numpy(self): + """Test validating numpy array prediction.""" + result = validate_prediction(np.array([0.2, 0.8])) + assert result is True + + def test_validate_prediction_none(self): + """Test validating None prediction.""" + with pytest.raises(ValueError): + validate_prediction(None) + + +class TestNumpySerialization: + """Tests for numpy serialization functions.""" + + def test_serialize_numpy_array(self): + """Test serializing numpy array.""" + arr = np.array([1, 2, 3]) + serialized = serialize_numpy(arr) + assert isinstance(serialized, list) + assert serialized == [1, 2, 3] + + def test_serialize_2d_array(self): + """Test serializing 2D numpy array.""" + arr = np.array([[1, 2], [3, 4]]) + serialized = serialize_numpy(arr) + assert isinstance(serialized, list) + assert len(serialized) == 2 + assert serialized[0] == [1, 2] + + def test_serialize_non_numpy(self): + """Test serializing non-numpy object.""" + data = [1, 2, 3] + serialized = serialize_numpy(data) + assert serialized == [1, 2, 3] + + def test_deserialize_to_numpy(self): + """Test deserializing to numpy array.""" + data = [1, 2, 3] + arr = deserialize_numpy(data) + assert isinstance(arr, np.ndarray) + assert np.array_equal(arr, np.array([1, 2, 3])) + + def test_roundtrip_serialization(self): + """Test roundtrip serialization/deserialization.""" + original = np.array([1.0, 2.0, 3.0]) + serialized = serialize_numpy(original) + deserialized = deserialize_numpy(serialized) + assert np.array_equal(original, deserialized) + + +class TestMetricsComputation: + """Tests for metrics computation.""" + + def test_compute_metrics_classification(self): + """Test computing classification metrics.""" + y_true = [0, 1, 1, 0, 1] + y_pred = [0, 1, 0, 0, 1] + metrics = compute_metrics(y_true, y_pred, task="classification") + assert "accuracy" in metrics + assert "precision" in metrics + assert "recall" in metrics + assert 0 <= metrics["accuracy"] <= 1 + + def test_compute_metrics_regression(self): + """Test computing regression metrics.""" + y_true = [1.0, 2.0, 3.0, 4.0] + y_pred = [1.1, 2.1, 2.9, 4.2] + metrics = compute_metrics(y_true, y_pred, task="regression") + assert "mae" in metrics + assert "mse" in metrics + assert "rmse" in metrics + assert metrics["mae"] >= 0 + + def test_compute_metrics_invalid_task(self): + """Test computing metrics with invalid task.""" + y_true = [1, 2, 3] + y_pred = [1, 2, 3] + with pytest.raises(ValueError): + compute_metrics(y_true, y_pred, task="invalid") + + def test_compute_metrics_mismatched_lengths(self): + """Test computing metrics with mismatched array lengths.""" + y_true = [1, 2, 3] + y_pred = [1, 2] + with pytest.raises(ValueError): + compute_metrics(y_true, y_pred, task="classification") + + +class TestBatchIterator: + """Tests for batch iterator utility.""" + + def test_batch_iterator_basic(self): + """Test basic batch iteration.""" + data = list(range(10)) + batches = list(batch_iterator(data, batch_size=3)) + assert len(batches) == 4 # 3, 3, 3, 1 + assert batches[0] == [0, 1, 2] + assert batches[1] == [3, 4, 5] + assert batches[2] == [6, 7, 8] + assert batches[3] == [9] + + def test_batch_iterator_exact_size(self): + """Test batch iteration with exact batch size.""" + data = list(range(9)) + batches = list(batch_iterator(data, batch_size=3)) + assert len(batches) == 3 + assert all(len(batch) == 3 for batch in batches) + + def test_batch_iterator_single_batch(self): + """Test batch iteration with size larger than data.""" + data = [1, 2, 3] + batches = list(batch_iterator(data, batch_size=10)) + assert len(batches) == 1 + assert batches[0] == [1, 2, 3] + + def test_batch_iterator_numpy(self): + """Test batch iteration with numpy array.""" + data = np.array(range(10)) + batches = list(batch_iterator(data, batch_size=4)) + assert len(batches) == 3 + assert len(batches[0]) == 4 + + +class TestUtilsEdgeCases: + """Tests for utils edge cases.""" + + def test_format_features_with_nans(self): + """Test formatting features with NaN values.""" + features = [1.0, np.nan, 3.0] + formatted = format_features(features) + assert len(formatted) == 3 + assert np.isnan(formatted[1]) + + def test_serialize_numpy_with_inf(self): + """Test serializing numpy with infinity.""" + arr = np.array([1.0, np.inf, -np.inf]) + serialized = serialize_numpy(arr) + assert len(serialized) == 3 + + def test_validate_features_with_none_values(self): + """Test validating features with None values.""" + features = [1.0, None, 3.0] + # Should handle gracefully or raise ValueError + try: + result = validate_features(features) + # If it passes, it should still be True + assert result is True + except ValueError: + # Acceptable to reject None values + pass diff --git a/whiteboxxai/__init__.py b/whiteboxxai/__init__.py new file mode 100644 index 0000000..937c39c --- /dev/null +++ b/whiteboxxai/__init__.py @@ -0,0 +1,31 @@ +""" +WhiteBoxXAI Python SDK + +Official Python SDK for WhiteBoxXAI - AI Observability & Explainability Platform. +""" + +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("whitebox-xai-sdk") +except PackageNotFoundError: + # Package isn't installed (e.g. running directly from a source checkout). + __version__ = "0.0.0.dev0" + +__author__ = "WhiteBoxXAI Team" +__license__ = "MIT" + +from whiteboxxai.client import WhiteBoxXAI +from whiteboxxai.decorators import monitor_model, monitor_prediction +from whiteboxxai.git_utils import GitContext, detect_git_context, validate_git_context +from whiteboxxai.monitor import ModelMonitor + +__all__ = [ + "WhiteBoxXAI", + "ModelMonitor", + "monitor_model", + "monitor_prediction", + "GitContext", + "detect_git_context", + "validate_git_context", +] diff --git a/src/whiteboxai/cache.py b/whiteboxxai/cache.py similarity index 92% rename from src/whiteboxai/cache.py rename to whiteboxxai/cache.py index 8930c10..d5edc56 100644 --- a/src/whiteboxai/cache.py +++ b/whiteboxxai/cache.py @@ -114,4 +114,6 @@ def cache_key(*args, **kwargs) -> str: "kwargs": sorted(kwargs.items()), } key_str = json.dumps(key_data, sort_keys=True) - return hashlib.md5(key_str.encode()).hexdigest() + # Cache key, not a security control - usedforsecurity=False keeps this fast + # and avoids flagging MD5 as if it protected something sensitive. + return hashlib.md5(key_str.encode(), usedforsecurity=False).hexdigest() diff --git a/src/whiteboxai/client.py b/whiteboxxai/client.py similarity index 83% rename from src/whiteboxai/client.py rename to whiteboxxai/client.py index 705c027..9170316 100644 --- a/src/whiteboxai/client.py +++ b/whiteboxxai/client.py @@ -1,24 +1,28 @@ """ -WhiteBoxAI Client +WhiteBoxXAI Client -Main client class for interacting with the WhiteBoxAI API. +Main client class for interacting with the WhiteBoxXAI API. """ import asyncio import logging -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional, Union from urllib.parse import urljoin import httpx from tenacity import retry, stop_after_attempt, wait_exponential - -from whiteboxai.config import Config -from whiteboxai.exceptions import APIError, AuthenticationError, RateLimitError, ValidationError -from whiteboxai.resources import ( - AgentWorkflowsResource, +from whiteboxxai.config import Config +from whiteboxxai.exceptions import ( + APIError, + AuthenticationError, + RateLimitError, + ValidationError, +) +from whiteboxxai.resources import ( AlertsResource, DriftResource, ExplanationsResource, + FairnessResource, ModelsResource, PredictionsResource, ) @@ -26,32 +30,32 @@ logger = logging.getLogger(__name__) -class WhiteBoxAI: +class WhiteBoxXAI: """ - Main client for WhiteBoxAI SDK. + Main client for WhiteBoxXAI SDK. - Provides access to all WhiteBoxAI API resources including models, predictions, - explanations, drift detection, and alerting. + Provides access to all WhiteBoxXAI API resources including models, predictions, + explanations, drift detection, fairness/bias audits, and alerting. Args: - api_key: WhiteBoxAI API key - base_url: Base URL for WhiteBoxAI API (default: https://api.whiteboxai.io) + api_key: WhiteBoxXAI API key + base_url: Base URL for WhiteBoxXAI API (default: https://api.whiteboxxai.com) timeout: Request timeout in seconds (default: 30) max_retries: Maximum number of retry attempts (default: 3) enable_offline: Enable offline mode (queues operations when API unavailable) - offline_dir: Directory for offline queue storage (default: ./whiteboxai_offline) + offline_dir: Directory for offline queue storage (default: ./whiteboxxai_offline) offline_max_queue_size: Maximum operations to queue (default: 10000, 0 = unlimited) offline_auto_sync: Enable automatic syncing (default: True) offline_sync_interval: Seconds between sync attempts (default: 60) **kwargs: Additional configuration options Example: - >>> client = WhiteBoxAI(api_key="your_api_key") + >>> client = WhiteBoxXAI(api_key="your_api_key") >>> model = client.models.register(name="fraud_detection", model_type="classification") >>> client.predictions.log(model_id=model.id, inputs=features, outputs=prediction) With offline mode: - >>> client = WhiteBoxAI( + >>> client = WhiteBoxXAI( ... api_key="your_api_key", ... enable_offline=True, ... offline_dir="./offline_queue" @@ -63,17 +67,17 @@ class WhiteBoxAI: def __init__( self, api_key: str, - base_url: str = "https://api.whiteboxai.io", + base_url: str = "https://api.whiteboxxai.com", timeout: int = 30, max_retries: int = 3, enable_offline: bool = False, - offline_dir: str = "./whiteboxai_offline", + offline_dir: str = "./whiteboxxai_offline", offline_max_queue_size: int = 10000, offline_auto_sync: bool = True, offline_sync_interval: int = 60, **kwargs: Any, ): - """Initialize WhiteBoxAI client.""" + """Initialize WhiteBoxXAI client.""" self.config = Config( api_key=api_key, base_url=base_url, @@ -89,7 +93,7 @@ def __init__( # Initialize offline mode self._offline_manager = None if enable_offline: - from whiteboxai.offline import OfflineManager + from whiteboxxai.offline import OfflineManager self._offline_manager = OfflineManager( offline_dir=offline_dir, @@ -105,8 +109,8 @@ def __init__( self.predictions = PredictionsResource(self) self.explanations = ExplanationsResource(self) self.drift = DriftResource(self) + self.fairness = FairnessResource(self) self.alerts = AlertsResource(self) - self.agent_workflows = AgentWorkflowsResource(self) @property def sync_client(self) -> httpx.Client: @@ -135,7 +139,7 @@ def _get_headers(self) -> Dict[str, str]: return { "Authorization": f"Bearer {self.config.api_key}", "Content-Type": "application/json", - "User-Agent": f"whiteboxai-python-sdk/{self.config.sdk_version}", + "User-Agent": f"whiteboxxai-python-sdk/{self.config.sdk_version}", } @retry( @@ -239,13 +243,24 @@ def _handle_response(self, response: httpx.Response) -> Dict[str, Any]: APIError: For other error status codes """ if response.status_code == 401: - raise AuthenticationError("Invalid API key or authentication failed") + raise AuthenticationError( + "Invalid API key or authentication failed", + status_code=response.status_code, + ) elif response.status_code == 429: - raise RateLimitError("Rate limit exceeded. Please retry later.") + retry_after = response.headers.get("Retry-After") + raise RateLimitError( + "Rate limit exceeded. Please retry later.", + retry_after=int(retry_after) + if isinstance(retry_after, str) and retry_after.isdigit() + else None, + ) elif response.status_code == 422: try: error_data = response.json() - raise ValidationError(f"Validation error: {error_data}") + raise ValidationError(f"Validation error: {error_data}", fields=error_data) + except ValidationError: + raise except Exception: raise ValidationError("Validation error occurred") elif response.status_code >= 400: @@ -253,8 +268,14 @@ def _handle_response(self, response: httpx.Response) -> Dict[str, Any]: error_data = response.json() message = error_data.get("detail", response.text) except Exception: + error_data = None message = response.text - raise APIError(f"API error ({response.status_code}): {message}") + raise APIError( + f"API error ({response.status_code}): {message}", + status_code=response.status_code, + response=error_data, + request_id=response.headers.get("X-Request-ID"), + ) try: return response.json() @@ -337,7 +358,7 @@ async def aclose(self) -> None: if self._sync_client: self._sync_client.close() - def __enter__(self) -> "WhiteBoxAI": + def __enter__(self) -> "WhiteBoxXAI": """Context manager entry.""" return self @@ -345,7 +366,7 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: """Context manager exit.""" self.close() - async def __aenter__(self) -> "WhiteBoxAI": + async def __aenter__(self) -> "WhiteBoxXAI": """Async context manager entry.""" return self diff --git a/src/whiteboxai/config.py b/whiteboxxai/config.py similarity index 77% rename from src/whiteboxai/config.py rename to whiteboxxai/config.py index b6c7b1d..a8ef590 100644 --- a/src/whiteboxai/config.py +++ b/whiteboxxai/config.py @@ -5,17 +5,19 @@ import os from typing import Any, Optional +from whiteboxxai import __version__ as _sdk_version + class Config: """ SDK configuration class. - Manages configuration settings for the WhiteBoxAI SDK including API keys, + Manages configuration settings for the WhiteBoxXAI SDK including API keys, URLs, timeouts, and feature flags. Args: - api_key: WhiteBoxAI API key - base_url: Base URL for WhiteBoxAI API + api_key: WhiteBoxXAI API key + base_url: Base URL for WhiteBoxXAI API timeout: Request timeout in seconds max_retries: Maximum number of retry attempts **kwargs: Additional configuration options @@ -31,17 +33,25 @@ def __init__( ): """Initialize configuration.""" # API key (from parameter or environment) - self.api_key = api_key or os.getenv("EXPLAINAI_API_KEY") + self.api_key = api_key or os.getenv("WHITEBOXXAI_API_KEY") if not self.api_key: raise ValueError( "API key is required. Provide via 'api_key' parameter or " - "EXPLAINAI_API_KEY environment variable." + "WHITEBOXXAI_API_KEY environment variable." ) # Base URL - self.base_url = base_url or os.getenv("EXPLAINAI_BASE_URL") or "https://api.whiteboxai.io" + self.base_url = ( + base_url + or os.getenv("WHITEBOXXAI_BASE_URL") + or "https://api.whiteboxxai.com" + ) # Request settings + if timeout <= 0: + raise ValueError("timeout must be a positive number of seconds.") + if max_retries < 0: + raise ValueError("max_retries must be zero or a positive integer.") self.timeout = timeout self.max_retries = max_retries @@ -67,7 +77,7 @@ def __init__( self.async_enabled = kwargs.get("async_enabled", True) # SDK metadata - self.sdk_version = "0.2.0" + self.sdk_version = _sdk_version def to_dict(self) -> dict: """Convert configuration to dictionary.""" diff --git a/src/whiteboxai/decorators.py b/whiteboxxai/decorators.py similarity index 95% rename from src/whiteboxai/decorators.py rename to whiteboxxai/decorators.py index 21f2a13..92a9033 100644 --- a/src/whiteboxai/decorators.py +++ b/whiteboxxai/decorators.py @@ -9,7 +9,7 @@ import time from typing import Any, Callable, Dict, Optional -from whiteboxai.monitor import ModelMonitor +from whiteboxxai.monitor import ModelMonitor def monitor_model( @@ -29,9 +29,9 @@ def monitor_model( Example: ```python - from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_model + from whiteboxxai import WhiteBoxXAI, ModelMonitor, monitor_model - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = ModelMonitor(client, model_id=123) @monitor_model(monitor, input_keys=["features"], explain=True) @@ -126,9 +126,9 @@ def monitor_prediction( Example: ```python - from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_prediction + from whiteboxxai import WhiteBoxXAI, ModelMonitor, monitor_prediction - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = ModelMonitor(client, model_id=123) def extract_inputs(args, kwargs): @@ -254,7 +254,9 @@ def _extract_inputs( return _default_input_extractor(func, args, kwargs) -def _default_input_extractor(func: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]: +def _default_input_extractor( + func: Callable, args: tuple, kwargs: dict +) -> Dict[str, Any]: """Default input extractor.""" sig = inspect.signature(func) bound_args = sig.bind(*args, **kwargs) @@ -278,7 +280,7 @@ def monitor_performance(threshold_ms: Optional[float] = None): Example: ```python - from whiteboxai.decorators import monitor_performance + from whiteboxxai.decorators import monitor_performance @monitor_performance(threshold_ms=1000) def slow_function(): diff --git a/whiteboxxai/exceptions.py b/whiteboxxai/exceptions.py new file mode 100644 index 0000000..f4d011f --- /dev/null +++ b/whiteboxxai/exceptions.py @@ -0,0 +1,109 @@ +""" +SDK Exceptions +""" + +from typing import Any, Dict, Optional + + +class WhiteBoxXAIError(Exception): + """Base exception for WhiteBoxXAI SDK.""" + + def __init__(self, message: str, details: Optional[Dict[str, Any]] = None): + super().__init__(message) + self.message = message + self.details = details or {} + + +class APIError(WhiteBoxXAIError): + """API request error.""" + + def __init__( + self, + message: str, + status_code: Optional[int] = None, + response: Optional[Dict[str, Any]] = None, + request_id: Optional[str] = None, + **kwargs: Any, + ): + super().__init__(message, **kwargs) + self.status_code = status_code + self.response = response + self.request_id = request_id + + +class AuthenticationError(WhiteBoxXAIError): + """Authentication error.""" + + def __init__( + self, + message: str, + status_code: Optional[int] = None, + response: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + super().__init__(message, **kwargs) + self.status_code = status_code + self.response = response + + +class RateLimitError(WhiteBoxXAIError): + """Rate limit exceeded error.""" + + def __init__( + self, + message: str, + retry_after: Optional[int] = None, + limit: Optional[int] = None, + remaining: Optional[int] = None, + **kwargs: Any, + ): + super().__init__(message, **kwargs) + self.retry_after = retry_after + self.limit = limit + self.remaining = remaining + + +class ValidationError(WhiteBoxXAIError): + """Validation error.""" + + def __init__( + self, + message: str, + fields: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ): + super().__init__(message, **kwargs) + self.fields = fields or {} + + +class NotFoundError(WhiteBoxXAIError): + """Resource not found error.""" + + def __init__( + self, + message: str, + resource_type: Optional[str] = None, + resource_id: Optional[str] = None, + **kwargs: Any, + ): + super().__init__(message, **kwargs) + self.resource_type = resource_type + self.resource_id = resource_id + + +class ConfigurationError(WhiteBoxXAIError): + """Configuration error.""" + + pass + + +class IntegrationError(WhiteBoxXAIError): + """Framework integration error.""" + + pass + + +class CacheError(WhiteBoxXAIError): + """Cache operation error.""" + + pass diff --git a/src/whiteboxai/git_utils.py b/whiteboxxai/git_utils.py similarity index 85% rename from src/whiteboxai/git_utils.py rename to whiteboxxai/git_utils.py index 925817a..22b79ab 100644 --- a/src/whiteboxai/git_utils.py +++ b/whiteboxxai/git_utils.py @@ -7,11 +7,17 @@ import logging import os +import shutil import subprocess +from pathlib import Path from typing import Dict, Optional logger = logging.getLogger(__name__) +# Resolved once to an absolute path so a malicious "git" earlier on PATH +# can't get picked up instead of the real executable. +_GIT_EXECUTABLE = shutil.which("git") or "git" + class GitContext: """Git repository context information.""" @@ -95,7 +101,9 @@ def detect_git_context(path: Optional[str] = None) -> Optional[GitContext]: repository_url = remote.url # Convert SSH URLs to HTTPS if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace("git@github.com:", "https://github.com/") + repository_url = repository_url.replace( + "git@github.com:", "https://github.com/" + ) if repository_url.endswith(".git"): repository_url = repository_url[:-4] except Exception as e: @@ -171,8 +179,10 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo cwd = path or os.getcwd() # Check if Git is available - subprocess.run( - ["git", "--version"], + # Every subprocess.run call below uses a fully static argv list (no + # shell=True, no interpolated/untrusted strings) - nosec B603. + subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "--version"], capture_output=True, check=True, cwd=cwd, @@ -181,8 +191,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get repository URL repository_url = None try: - result = subprocess.run( - ["git", "remote", "get-url", "origin"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "remote", "get-url", "origin"], capture_output=True, text=True, check=True, @@ -192,7 +202,9 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Convert SSH to HTTPS if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace("git@github.com:", "https://github.com/") + repository_url = repository_url.replace( + "git@github.com:", "https://github.com/" + ) if repository_url.endswith(".git"): repository_url = repository_url[:-4] except subprocess.CalledProcessError: @@ -201,8 +213,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get commit SHA commit_sha = None try: - result = subprocess.run( - ["git", "rev-parse", "HEAD"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "rev-parse", "HEAD"], capture_output=True, text=True, check=True, @@ -215,8 +227,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get commit message commit_message = None try: - result = subprocess.run( - ["git", "log", "-1", "--pretty=%B"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "log", "-1", "--pretty=%B"], capture_output=True, text=True, check=True, @@ -229,8 +241,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get commit author commit_author = None try: - result = subprocess.run( - ["git", "log", "-1", "--pretty=%an <%ae>"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "log", "-1", "--pretty=%an <%ae>"], capture_output=True, text=True, check=True, @@ -243,8 +255,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get branch branch = None try: - result = subprocess.run( - ["git", "rev-parse", "--abbrev-ref", "HEAD"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "rev-parse", "--abbrev-ref", "HEAD"], capture_output=True, text=True, check=True, @@ -259,8 +271,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Get tag tag = None try: - result = subprocess.run( - ["git", "describe", "--exact-match", "--tags", "HEAD"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "describe", "--exact-match", "--tags", "HEAD"], capture_output=True, text=True, check=True, @@ -273,8 +285,8 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Check if dirty is_dirty = False try: - result = subprocess.run( - ["git", "status", "--porcelain"], + result = subprocess.run( # nosec B603 + [_GIT_EXECUTABLE, "status", "--porcelain"], capture_output=True, text=True, check=True, diff --git a/src/whiteboxai/integrations/__init__.py b/whiteboxxai/integrations/__init__.py similarity index 56% rename from src/whiteboxai/integrations/__init__.py rename to whiteboxxai/integrations/__init__.py index c69a821..806953d 100644 --- a/src/whiteboxai/integrations/__init__.py +++ b/whiteboxxai/integrations/__init__.py @@ -1,4 +1,4 @@ -"""Framework integrations for WhiteBoxAI SDK.""" +"""Framework integrations for WhiteBoxXAI SDK.""" # Scikit-learn integration try: @@ -21,12 +21,12 @@ # TensorFlow/Keras integration try: - from .tensorflow import KerasMonitor, WhiteBoxAICallback, wrap_keras_model + from .tensorflow import KerasMonitor, WhiteBoxXAICallback, wrap_keras_model if "__all__" in dir(): - __all__.extend(["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"]) + __all__.extend(["KerasMonitor", "WhiteBoxXAICallback", "wrap_keras_model"]) else: - __all__ = ["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"] + __all__ = ["KerasMonitor", "WhiteBoxXAICallback", "wrap_keras_model"] except ImportError: pass @@ -40,7 +40,11 @@ if "__all__" in dir(): __all__.extend( - ["TransformersMonitor", "TransformersPipelineWrapper", "wrap_transformers_pipeline"] + [ + "TransformersMonitor", + "TransformersPipelineWrapper", + "wrap_transformers_pipeline", + ] ) else: __all__ = [ @@ -53,25 +57,50 @@ # LangChain integration try: - from .langchain import LangChainMonitor, WhiteBoxAICallbackHandler, wrap_langchain_chain + from .langchain import ( + LangChainMonitor, + WhiteBoxXAICallbackHandler, + wrap_langchain_chain, + ) if "__all__" in dir(): - __all__.extend(["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"]) + __all__.extend( + ["LangChainMonitor", "WhiteBoxXAICallbackHandler", "wrap_langchain_chain"] + ) else: - __all__ = ["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"] + __all__ = [ + "LangChainMonitor", + "WhiteBoxXAICallbackHandler", + "wrap_langchain_chain", + ] except ImportError: pass # XGBoost/LightGBM integration try: - from .boosting import LightGBMMonitor, XGBoostMonitor, wrap_lightgbm_model, wrap_xgboost_model + from .boosting import ( + LightGBMMonitor, + XGBoostMonitor, + wrap_lightgbm_model, + wrap_xgboost_model, + ) if "__all__" in dir(): __all__.extend( - ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] + [ + "XGBoostMonitor", + "LightGBMMonitor", + "wrap_xgboost_model", + "wrap_lightgbm_model", + ] ) else: - __all__ = ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] + __all__ = [ + "XGBoostMonitor", + "LightGBMMonitor", + "wrap_xgboost_model", + "wrap_lightgbm_model", + ] except ImportError: pass @@ -86,7 +115,7 @@ except ImportError: pass -# LangChain Multi-Agent integration +# LangChain integration try: from .langchain_agents import ( LangGraphMultiAgentMonitor, @@ -96,7 +125,11 @@ if "__all__" in dir(): __all__.extend( - ["MultiAgentCallbackHandler", "LangGraphMultiAgentMonitor", "monitor_langchain_agent"] + [ + "MultiAgentCallbackHandler", + "LangGraphMultiAgentMonitor", + "monitor_langchain_agent", + ] ) else: __all__ = [ diff --git a/src/whiteboxai/integrations/boosting.py b/whiteboxxai/integrations/boosting.py similarity index 91% rename from src/whiteboxai/integrations/boosting.py rename to whiteboxxai/integrations/boosting.py index 210155b..d29a6b8 100644 --- a/src/whiteboxai/integrations/boosting.py +++ b/whiteboxxai/integrations/boosting.py @@ -1,5 +1,5 @@ """ -XGBoost and LightGBM integration for WhiteBoxAI. +XGBoost and LightGBM integration for WhiteBoxXAI. This module provides monitoring capabilities for gradient boosting models from XGBoost and LightGBM frameworks. It supports both frameworks through @@ -10,10 +10,10 @@ Basic XGBoost monitoring: >>> import xgboost as xgb - >>> from whiteboxai import WhiteBoxAI - >>> from whiteboxai.integrations.boosting import XGBoostMonitor + >>> from whiteboxxai import WhiteBoxXAI + >>> from whiteboxxai.integrations.boosting import XGBoostMonitor >>> - >>> client = WhiteBoxAI(api_key="your-api-key") + >>> client = WhiteBoxXAI(api_key="your-api-key") >>> monitor = XGBoostMonitor(client=client, model_name="fraud_detector") >>> >>> model = xgb.XGBClassifier() @@ -25,7 +25,7 @@ LightGBM monitoring: >>> import lightgbm as lgb - >>> from whiteboxai.integrations.boosting import LightGBMMonitor + >>> from whiteboxxai.integrations.boosting import LightGBMMonitor >>> >>> monitor = LightGBMMonitor(client=client, model_name="churn_predictor") >>> @@ -36,12 +36,14 @@ >>> predictions = monitor.predict(model, X_test, y_test) """ +import logging import warnings -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union import numpy as np +from whiteboxxai.monitor import ModelMonitor -from ..core import ModelMonitor +logger = logging.getLogger(__name__) # Optional imports - graceful degradation try: @@ -69,7 +71,7 @@ class XGBoostMonitor(ModelMonitor): logging, feature importance tracking, and automatic model registration. Attributes: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance model_name: Name of the model being monitored track_feature_importance: Whether to track feature importance with predictions importance_type: Type of importance to track ('weight', 'gain', 'cover', 'total_gain', 'total_cover') @@ -97,14 +99,16 @@ def __init__( Initialize XGBoost monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance model_name: Name of the model track_feature_importance: Whether to track feature importance importance_type: Type of importance ('weight', 'gain', 'cover', 'total_gain', 'total_cover') **kwargs: Additional arguments passed to ModelMonitor """ if not XGBOOST_AVAILABLE: - raise ImportError("XGBoost is not installed. Install with: pip install xgboost") + raise ImportError( + "XGBoost is not installed. Install with: pip install xgboost" + ) super().__init__(client=client, model_name=model_name, **kwargs) self.track_feature_importance = track_feature_importance @@ -120,7 +124,7 @@ def register_from_model( metadata: Optional[Dict[str, Any]] = None, ) -> str: """ - Register XGBoost model with WhiteBoxAI. + Register XGBoost model with WhiteBoxXAI. Automatically extracts model metadata including feature names, number of features, number of trees, and feature importance. @@ -213,7 +217,7 @@ def predict( metadata: Optional[Dict[str, Any]] = None, ) -> np.ndarray: """ - Make predictions and log to WhiteBoxAI. + Make predictions and log to WhiteBoxXAI. Args: model: XGBoost model @@ -233,8 +237,8 @@ def predict( if hasattr(model, "predict_proba"): try: probabilities = model.predict_proba(X) - except Exception: - pass + except Exception as e: + logger.debug(f"predict_proba failed, continuing without probabilities: {e}") # Log predictions if log_predictions: @@ -286,7 +290,9 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try get_booster for sklearn API if hasattr(model, "get_booster"): booster = model.get_booster() - importance_dict = booster.get_score(importance_type=self.importance_type) + importance_dict = booster.get_score( + importance_type=self.importance_type + ) return {k: float(v) for k, v in importance_dict.items()} return None @@ -302,7 +308,7 @@ class LightGBMMonitor(ModelMonitor): logging, feature importance tracking, and automatic model registration. Attributes: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance model_name: Name of the model being monitored track_feature_importance: Whether to track feature importance with predictions importance_type: Type of importance to track ('split' or 'gain') @@ -330,14 +336,16 @@ def __init__( Initialize LightGBM monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance model_name: Name of the model track_feature_importance: Whether to track feature importance importance_type: Type of importance ('split' or 'gain') **kwargs: Additional arguments passed to ModelMonitor """ if not LIGHTGBM_AVAILABLE: - raise ImportError("LightGBM is not installed. Install with: pip install lightgbm") + raise ImportError( + "LightGBM is not installed. Install with: pip install lightgbm" + ) super().__init__(client=client, model_name=model_name, **kwargs) self.track_feature_importance = track_feature_importance @@ -353,7 +361,7 @@ def register_from_model( metadata: Optional[Dict[str, Any]] = None, ) -> str: """ - Register LightGBM model with WhiteBoxAI. + Register LightGBM model with WhiteBoxXAI. Automatically extracts model metadata including feature names, number of features, number of trees, and feature importance. @@ -448,7 +456,7 @@ def predict( metadata: Optional[Dict[str, Any]] = None, ) -> np.ndarray: """ - Make predictions and log to WhiteBoxAI. + Make predictions and log to WhiteBoxXAI. Args: model: LightGBM model @@ -468,8 +476,8 @@ def predict( if hasattr(model, "predict_proba"): try: probabilities = model.predict_proba(X) - except Exception: - pass + except Exception as e: + logger.debug(f"predict_proba failed, continuing without probabilities: {e}") # Log predictions if log_predictions: @@ -515,7 +523,9 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try feature_importance method (native LightGBM) if hasattr(model, "feature_importance"): - importances = model.feature_importance(importance_type=self.importance_type) + importances = model.feature_importance( + importance_type=self.importance_type + ) if self._feature_names and len(importances) == len(self._feature_names): return dict(zip(self._feature_names, importances.tolist())) else: @@ -524,7 +534,9 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try booster if hasattr(model, "booster_"): booster = model.booster_ - importances = booster.feature_importance(importance_type=self.importance_type) + importances = booster.feature_importance( + importance_type=self.importance_type + ) feature_names = booster.feature_name() if len(importances) == len(feature_names): return dict(zip(feature_names, importances.tolist())) @@ -535,7 +547,9 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Unified wrapper functions -def wrap_xgboost_model(model: Any, monitor: XGBoostMonitor, auto_register: bool = True) -> Any: +def wrap_xgboost_model( + model: Any, monitor: XGBoostMonitor, auto_register: bool = True +) -> Any: """ Wrap an XGBoost model for automatic monitoring. @@ -601,7 +615,9 @@ def wrapped_predict_proba(X, *args, **kwargs): return model -def wrap_lightgbm_model(model: Any, monitor: LightGBMMonitor, auto_register: bool = True) -> Any: +def wrap_lightgbm_model( + model: Any, monitor: LightGBMMonitor, auto_register: bool = True +) -> Any: """ Wrap a LightGBM model for automatic monitoring. diff --git a/src/whiteboxai/integrations/crewai_monitor.py b/whiteboxxai/integrations/crewai_monitor.py similarity index 79% rename from src/whiteboxai/integrations/crewai_monitor.py rename to whiteboxxai/integrations/crewai_monitor.py index 9ab35cd..fb8970a 100644 --- a/src/whiteboxai/integrations/crewai_monitor.py +++ b/whiteboxxai/integrations/crewai_monitor.py @@ -1,12 +1,17 @@ """ -CrewAI Integration for WhiteBoxAI +CrewAI Integration for WhiteBoxXAI Monitor CrewAI multi-agent workflows with automatic tracking of agents, tasks, interactions, and costs. """ import logging -from typing import Any, Dict, Optional +import time +import uuid +from datetime import datetime +from typing import Any, Dict, List, Optional + +from whiteboxxai import WhiteBoxXAI logger = logging.getLogger(__name__) @@ -23,7 +28,7 @@ class CrewAIMonitor: - Token usage and costs Example: - >>> from whiteboxai.integrations import CrewAIMonitor + >>> from whiteboxxai.integrations import CrewAIMonitor >>> from crewai import Agent, Task, Crew >>> >>> monitor = CrewAIMonitor(api_key="your_api_key") @@ -75,27 +80,28 @@ class CrewAIMonitor: """ def __init__( - self, api_key: str, api_url: Optional[str] = None, organization_id: Optional[str] = None + self, + api_key: str, + api_url: Optional[str] = None, + organization_id: Optional[str] = None, ): """ Initialize CrewAI monitor. Args: - api_key: WhiteBoxAI API key - api_url: WhiteBoxAI API URL (optional, defaults to production) + api_key: WhiteBoxXAI API key + api_url: WhiteBoxXAI API URL (optional, defaults to production) organization_id: Organization ID (optional, extracted from token if not provided) """ - from whiteboxai import WhiteBoxAI - self.client = ( - WhiteBoxAI(api_key=api_key, base_url=api_url) + WhiteBoxXAI(api_key=api_key, base_url=api_url) if api_url - else WhiteBoxAI(api_key=api_key) + else WhiteBoxXAI(api_key=api_key) ) self.organization_id = organization_id self.workflow_id = None - self.agent_map = {} # Maps CrewAI agent to WhiteBoxAI agent_id - self.task_map = {} # Maps CrewAI task to WhiteBoxAI task_id + self.agent_map = {} # Maps CrewAI agent to WhiteBoxXAI agent_id + self.task_map = {} # Maps CrewAI task to WhiteBoxXAI task_id self.execution_map = {} # Maps agent to execution_id logger.info("CrewAI Monitor initialized") @@ -145,12 +151,16 @@ def start_monitoring( "inputs": { "agent_count": len(crew.agents), "task_count": len(crew.tasks), - "process": str(crew.process) if hasattr(crew, "process") else "sequential", + "process": str(crew.process) + if hasattr(crew, "process") + else "sequential", } } self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/start", data=start_data + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/start", + data=start_data, ) return self.workflow_id @@ -161,7 +171,7 @@ def start_monitoring( def _register_agent(self, crew_agent: Any) -> str: """ - Register a CrewAI agent with WhiteBoxAI. + Register a CrewAI agent with WhiteBoxXAI. Args: crew_agent: CrewAI Agent instance @@ -176,17 +186,19 @@ def _register_agent(self, crew_agent: Any) -> str: "agent_type": "crewai_agent", "goal": getattr(crew_agent, "goal", None), "backstory": getattr(crew_agent, "backstory", None), - "tools": [tool.__class__.__name__ for tool in getattr(crew_agent, "tools", [])], - "llm_provider": ( - getattr(getattr(crew_agent, "llm", None), "model_name", "unknown").split("/")[0] - if hasattr(crew_agent, "llm") - else None - ), - "model_name": ( - getattr(getattr(crew_agent, "llm", None), "model_name", None) - if hasattr(crew_agent, "llm") - else None - ), + "tools": [ + tool.__class__.__name__ for tool in getattr(crew_agent, "tools", []) + ], + "llm_provider": getattr( + getattr(crew_agent, "llm", None), "model_name", "unknown" + ).split("-")[0] + if hasattr(crew_agent, "llm") + else None, + "model_name": getattr( + getattr(crew_agent, "llm", None), "model_name", None + ) + if hasattr(crew_agent, "llm") + else None, "metadata": { "verbose": getattr(crew_agent, "verbose", False), "allow_delegation": getattr(crew_agent, "allow_delegation", False), @@ -195,7 +207,9 @@ def _register_agent(self, crew_agent: Any) -> str: } response = self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/agents", data=agent_data + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/agents", + data=agent_data, ) agent_id = response.get("id") @@ -211,7 +225,7 @@ def _register_agent(self, crew_agent: Any) -> str: def _register_task(self, crew_task: Any) -> str: """ - Register a CrewAI task with WhiteBoxAI. + Register a CrewAI task with WhiteBoxXAI. Args: crew_task: CrewAI Task instance @@ -231,19 +245,26 @@ def _register_task(self, crew_task: Any) -> str: "expected_output": getattr(crew_task, "expected_output", None), "agent_id": agent_id, "context": { - "tools": [tool.__class__.__name__ for tool in getattr(crew_task, "tools", [])], + "tools": [ + tool.__class__.__name__ + for tool in getattr(crew_task, "tools", []) + ], "async_execution": getattr(crew_task, "async_execution", False), }, } response = self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/tasks", data=task_data + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/tasks", + data=task_data, ) task_id = response.get("id") self.task_map[id(crew_task)] = task_id - logger.debug(f"Registered task: {task_data['task_name'][:50]}... ({task_id})") + logger.debug( + f"Registered task: {task_data['task_name'][:50]}... ({task_id})" + ) return task_id @@ -272,7 +293,7 @@ def log_agent_execution( try: agent_id = self.agent_map.get(id(agent)) if not agent_id: - logger.warning("Agent not registered, skipping execution log") + logger.warning(f"Agent not registered, skipping execution log") return execution_data = { @@ -314,7 +335,7 @@ def log_task_completion( try: task_id = self.task_map.get(id(task)) if not task_id: - logger.warning("Task not registered, skipping completion log") + logger.warning(f"Task not registered, skipping completion log") return update_data = { @@ -324,7 +345,9 @@ def log_task_completion( } self.client.request( - "PATCH", f"/api/v1/workflows/multi-agent/tasks/{task_id}", data=update_data + "PATCH", + f"/api/v1/workflows/multi-agent/tasks/{task_id}", + data=update_data, ) logger.debug(f"Logged task completion: {task_id} ({status})") @@ -393,10 +416,17 @@ def complete_monitoring( Returns: Workflow summary with analytics """ + if not self.workflow_id: + raise ValueError("No active workflow to complete. Call start_monitoring() first.") + try: - complete_data = {"status": status, "outputs": outputs, "error_message": error_message} + complete_data = { + "status": status, + "outputs": outputs, + "error_message": error_message, + } - self.client.request( + response = self.client.request( "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/complete", data=complete_data, @@ -407,7 +437,11 @@ def complete_monitoring( # Get analytics analytics = self.get_analytics() - return {"workflow_id": self.workflow_id, "status": status, "analytics": analytics} + return { + "workflow_id": self.workflow_id, + "status": status, + "analytics": analytics, + } except Exception as e: logger.error(f"Error completing monitoring: {str(e)}") @@ -429,7 +463,8 @@ def get_analytics(self) -> Dict[str, Any]: ) cost_breakdown = self.client.request( - "GET", f"/api/v1/workflows/multi-agent/{self.workflow_id}/cost-breakdown" + "GET", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/cost-breakdown", ) return {"metrics": analytics, "cost_breakdown": cost_breakdown} @@ -453,8 +488,8 @@ def monitor_crew( Args: crew: CrewAI Crew instance workflow_name: Workflow name - api_key: WhiteBoxAI API key - api_url: WhiteBoxAI API URL (optional) + api_key: WhiteBoxXAI API key + api_url: WhiteBoxXAI API URL (optional) metadata: Workflow metadata (optional) Returns: diff --git a/src/whiteboxai/integrations/langchain.py b/whiteboxxai/integrations/langchain.py similarity index 94% rename from src/whiteboxai/integrations/langchain.py rename to whiteboxxai/integrations/langchain.py index 4112d89..7aa005e 100644 --- a/src/whiteboxai/integrations/langchain.py +++ b/whiteboxxai/integrations/langchain.py @@ -4,9 +4,12 @@ Integration for monitoring LangChain applications including chains, agents, and RAG pipelines. """ +from __future__ import annotations + +import json import time import warnings -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Union try: from langchain.agents import AgentExecutor @@ -20,8 +23,11 @@ BaseCallbackHandler = object Chain = object AgentExecutor = object + AgentAction = object + AgentFinish = object + LLMResult = object -from whiteboxai.monitor import ModelMonitor +from whiteboxxai.monitor import ModelMonitor class LangChainMonitor(ModelMonitor): @@ -40,11 +46,11 @@ class LangChainMonitor(ModelMonitor): ```python from langchain.chains import LLMChain from langchain.llms import OpenAI - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.langchain import LangChainMonitor + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.langchain import LangChainMonitor # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = LangChainMonitor( client=client, application_name="my_langchain_app" @@ -78,14 +84,16 @@ def __init__( Initialize LangChain monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance application_name: Name of the LangChain application track_tokens: Track token usage track_cost: Track API costs **kwargs: Additional arguments for ModelMonitor """ if not LANGCHAIN_AVAILABLE: - raise ImportError("langchain is not installed. Install with: pip install langchain") + raise ImportError( + "langchain is not installed. Install with: pip install langchain" + ) super().__init__(client, **kwargs) self._application_name = application_name @@ -101,7 +109,7 @@ def register_application( **kwargs: Any, ) -> int: """ - Register LangChain application with WhiteBoxAI. + Register LangChain application with WhiteBoxXAI. Args: name: Application name @@ -137,12 +145,12 @@ def register_application( self._registered = True return self.model_id - def create_callback_handler(self) -> "WhiteBoxAICallbackHandler": + def create_callback_handler(self) -> "WhiteBoxXAICallbackHandler": """ Create a LangChain callback handler for automatic logging. Returns: - WhiteBoxAICallbackHandler instance + WhiteBoxXAICallbackHandler instance Example: ```python @@ -153,7 +161,7 @@ def create_callback_handler(self) -> "WhiteBoxAICallbackHandler": if not self._registered: self.register_application() - return WhiteBoxAICallbackHandler(monitor=self) + return WhiteBoxXAICallbackHandler(monitor=self) def log_chain_execution( self, @@ -332,9 +340,9 @@ def log_rag_retrieval( ) -class WhiteBoxAICallbackHandler(BaseCallbackHandler): +class WhiteBoxXAICallbackHandler(BaseCallbackHandler): """ - LangChain callback handler for automatic logging to WhiteBoxAI. + LangChain callback handler for automatic logging to WhiteBoxXAI. Example: ```python @@ -506,7 +514,9 @@ def on_llm_end( self.monitor.log_llm_call( prompt="", # Prompt tracked separately response=response_text, - model=kwargs.get("invocation_params", {}).get("model_name", "unknown"), + model=kwargs.get("invocation_params", {}).get( + "model_name", "unknown" + ), tokens_used=tokens_used, latency=latency, ) @@ -631,6 +641,6 @@ def logged_call(*args, **kwargs): __all__ = [ "LangChainMonitor", - "WhiteBoxAICallbackHandler", + "WhiteBoxXAICallbackHandler", "wrap_langchain_chain", ] diff --git a/src/whiteboxai/integrations/langchain_agents.py b/whiteboxxai/integrations/langchain_agents.py similarity index 88% rename from src/whiteboxai/integrations/langchain_agents.py rename to whiteboxxai/integrations/langchain_agents.py index ec782da..894ba0f 100644 --- a/src/whiteboxai/integrations/langchain_agents.py +++ b/whiteboxxai/integrations/langchain_agents.py @@ -1,5 +1,5 @@ """ -LangChain Multi-Agent Integration for WhiteBoxAI +LangChain Multi-Agent Integration for WhiteBoxXAI Enhanced callback handler for monitoring multi-agent LangChain workflows including: - LangGraph multi-agent patterns @@ -13,11 +13,13 @@ from langchain.callbacks.base import BaseCallbackHandler from langchain.schema import AgentAction, AgentFinish, LLMResult +from langchain.schema.document import Document +from langchain.schema.output import ChatGeneration, Generation try: - from whiteboxai import WhiteBoxAI + from whiteboxxai import WhiteBoxXAI except ImportError: - WhiteBoxAI = None + WhiteBoxXAI = None class MultiAgentCallbackHandler(BaseCallbackHandler): @@ -33,16 +35,16 @@ class MultiAgentCallbackHandler(BaseCallbackHandler): Example: ```python from langchain.agents import AgentExecutor, create_react_agent - from whiteboxai.integrations import MultiAgentCallbackHandler + from whiteboxxai.integrations import MultiAgentCallbackHandler - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="your_key") + # Initialize WhiteBoxXAI client + client = WhiteBoxXAI(api_key="your_key") # Create workflow workflow_id = client.agent_workflows.create( name="Research Workflow", framework="langchain" - ).get("id") + ).id # Start workflow client.agent_workflows.start(workflow_id) @@ -72,7 +74,7 @@ class MultiAgentCallbackHandler(BaseCallbackHandler): def __init__( self, - client: "WhiteBoxAI", + client: "WhiteBoxXAI", workflow_id: str, agent_name: str = "main", agent_role: Optional[str] = None, @@ -82,16 +84,17 @@ def __init__( """Initialize the callback handler. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance workflow_id: ID of the workflow to track agent_name: Name of the current agent agent_role: Role/description of the agent track_tokens: Whether to track token usage track_costs: Whether to estimate costs """ - if WhiteBoxAI is None: + if WhiteBoxXAI is None: raise ImportError( - "whiteboxai package not installed. " "Install with: pip install whiteboxai" + "whiteboxxai package not installed. " + "Install with: pip install whiteboxxai" ) self.client = client @@ -150,7 +153,9 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None: # Reset state self.execution_start_time = None - def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: + def on_chain_error( + self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any + ) -> None: """Run when chain errors.""" if self.execution_start_time: duration_ms = int( @@ -174,7 +179,9 @@ def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: A self.execution_start_time = None - def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None: + def on_llm_start( + self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any + ) -> None: """Run when LLM starts.""" self.llm_call_count += 1 @@ -219,7 +226,9 @@ def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None: # This is called when the agent completes its reasoning pass - def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None: + def on_tool_start( + self, serialized: Dict[str, Any], input_str: str, **kwargs: Any + ) -> None: """Run when tool starts.""" pass @@ -238,7 +247,9 @@ def on_tool_end(self, output: str, **kwargs: Any) -> None: except Exception as e: print(f"Warning: Failed to log tool result: {e}") - def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: + def on_tool_error( + self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any + ) -> None: """Run when tool errors.""" try: self.client.agent_workflows.create_interaction( @@ -268,7 +279,7 @@ class LangGraphMultiAgentMonitor: Example: ```python from langgraph.graph import StateGraph - from whiteboxai.integrations import LangGraphMultiAgentMonitor + from whiteboxxai.integrations import LangGraphMultiAgentMonitor # Create monitor monitor = LangGraphMultiAgentMonitor( @@ -297,18 +308,22 @@ class LangGraphMultiAgentMonitor: """ def __init__( - self, client: "WhiteBoxAI", workflow_name: str, meta_data: Optional[Dict[str, Any]] = None + self, + client: "WhiteBoxXAI", + workflow_name: str, + meta_data: Optional[Dict[str, Any]] = None, ): """Initialize the LangGraph monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance workflow_name: Name for the workflow meta_data: Additional meta_data to attach """ - if WhiteBoxAI is None: + if WhiteBoxXAI is None: raise ImportError( - "whiteboxai package not installed. " "Install with: pip install whiteboxai" + "whiteboxxai package not installed. " + "Install with: pip install whiteboxxai" ) self.client = client @@ -455,11 +470,15 @@ def complete_monitoring( } except Exception as e: print(f"Warning: Failed to retrieve analytics: {e}") - return {"workflow_id": self.workflow_id, "status": status, "outputs": outputs} + return { + "workflow_id": self.workflow_id, + "status": status, + "outputs": outputs, + } def monitor_langchain_agent( - client: "WhiteBoxAI", + client: "WhiteBoxXAI", agent_executor: Any, workflow_name: str, agent_name: str = "main", @@ -469,7 +488,7 @@ def monitor_langchain_agent( """Helper function to monitor a single LangChain agent execution. Args: - client: WhiteBoxAI client + client: WhiteBoxXAI client agent_executor: LangChain AgentExecutor instance workflow_name: Name for the workflow agent_name: Name of the agent @@ -482,7 +501,7 @@ def monitor_langchain_agent( Example: ```python from langchain.agents import AgentExecutor, create_react_agent - from whiteboxai.integrations import monitor_langchain_agent + from whiteboxxai.integrations import monitor_langchain_agent result_dict = monitor_langchain_agent( client=client, @@ -520,6 +539,13 @@ def monitor_langchain_agent( return {"result": result, "workflow_id": workflow_id, "status": "completed"} except Exception as e: # Log failure - client.agent_workflows.complete(workflow_id, outputs={"error": str(e)}, status="failed") + client.agent_workflows.complete( + workflow_id, outputs={"error": str(e)}, status="failed" + ) - return {"result": None, "workflow_id": workflow_id, "status": "failed", "error": str(e)} + return { + "result": None, + "workflow_id": workflow_id, + "status": "failed", + "error": str(e), + } diff --git a/src/whiteboxai/integrations/pytorch.py b/whiteboxxai/integrations/pytorch.py similarity index 85% rename from src/whiteboxai/integrations/pytorch.py rename to whiteboxxai/integrations/pytorch.py index 16b7ddf..98c9993 100644 --- a/src/whiteboxai/integrations/pytorch.py +++ b/whiteboxxai/integrations/pytorch.py @@ -4,7 +4,9 @@ Integration for monitoring PyTorch models. """ -from typing import Any, Callable, Dict, Optional +from __future__ import annotations + +from typing import Any, Callable, Dict, Optional, Union try: import torch @@ -13,10 +15,18 @@ TORCH_AVAILABLE = True except ImportError: TORCH_AVAILABLE = False - nn = object torch = None -from whiteboxai.monitor import ModelMonitor + class _NNStub: + """Stand-in for `torch.nn` so `nn.Module` stays resolvable (as a type + hint and as a base class) when PyTorch isn't installed.""" + + Module = object + + nn = _NNStub() + +import numpy as np +from whiteboxxai.monitor import ModelMonitor class TorchMonitor(ModelMonitor): @@ -27,8 +37,8 @@ class TorchMonitor(ModelMonitor): ```python import torch import torch.nn as nn - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.pytorch import TorchMonitor + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.pytorch import TorchMonitor # Define model model = nn.Sequential( @@ -38,7 +48,7 @@ class TorchMonitor(ModelMonitor): ) # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = TorchMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -53,7 +63,9 @@ class TorchMonitor(ModelMonitor): def __init__(self, client, model: Optional[nn.Module] = None, **kwargs): """Initialize PyTorch monitor.""" if not TORCH_AVAILABLE: - raise ImportError("PyTorch is not installed. Install with: pip install torch") + raise ImportError( + "PyTorch is not installed. Install with: pip install torch" + ) super().__init__(client, **kwargs) self.model = model @@ -142,7 +154,9 @@ def _extract_model_metadata(self) -> Dict[str, Any]: # Count parameters total_params = sum(p.numel() for p in self.model.parameters()) - trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + trainable_params = sum( + p.numel() for p in self.model.parameters() if p.requires_grad + ) metadata["total_parameters"] = total_params metadata["trainable_parameters"] = trainable_params @@ -164,7 +178,7 @@ class TorchWrapper(nn.Module): """ Wrapper for PyTorch models with automatic monitoring. - This wrapper intercepts forward calls and logs predictions to WhiteBoxAI. + This wrapper intercepts forward calls and logs predictions to WhiteBoxXAI. """ def __init__(self, model: nn.Module, monitor: TorchMonitor): @@ -184,7 +198,9 @@ def forward(self, x: torch.Tensor, log: bool = True) -> torch.Tensor: return output - def _log_batch_predictions(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None: + def _log_batch_predictions( + self, inputs: torch.Tensor, outputs: torch.Tensor + ) -> None: """Log batch of predictions.""" # Convert to numpy/lists inputs_np = inputs.detach().cpu().numpy() @@ -224,10 +240,10 @@ def monitor_forward(monitor: TorchMonitor, input_extractor: Optional[Callable] = Example: ```python - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.pytorch import TorchMonitor, monitor_forward + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.pytorch import TorchMonitor, monitor_forward - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = TorchMonitor(client, model_id=123) class MyModel(nn.Module): @@ -249,7 +265,7 @@ def wrapper(self, *args, **kwargs): inputs = args[0] if args else None # Log prediction - if inputs is not None and isinstance(inputs, torch.Tensor): + if TORCH_AVAILABLE and inputs is not None and isinstance(inputs, torch.Tensor): inputs_np = inputs.detach().cpu().numpy() outputs_np = output.detach().cpu().numpy() diff --git a/src/whiteboxai/integrations/sklearn.py b/whiteboxxai/integrations/sklearn.py similarity index 92% rename from src/whiteboxai/integrations/sklearn.py rename to whiteboxxai/integrations/sklearn.py index a1e9b4f..ccf89c1 100644 --- a/src/whiteboxai/integrations/sklearn.py +++ b/whiteboxxai/integrations/sklearn.py @@ -16,8 +16,7 @@ BaseEstimator = object import numpy as np - -from whiteboxai.monitor import ModelMonitor +from whiteboxxai.monitor import ModelMonitor class SklearnMonitor(ModelMonitor): @@ -27,15 +26,15 @@ class SklearnMonitor(ModelMonitor): Example: ```python from sklearn.ensemble import RandomForestClassifier - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.sklearn import SklearnMonitor + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.sklearn import SklearnMonitor # Train model model = RandomForestClassifier() model.fit(X_train, y_train) # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = SklearnMonitor(client, model=model) monitor.register_from_model(model_type="classification") @@ -170,7 +169,7 @@ class SklearnWrapper: Wrapper for sklearn models with automatic monitoring. This wrapper intercepts predict/predict_proba calls and logs predictions - to WhiteBoxAI automatically. + to WhiteBoxXAI automatically. """ def __init__(self, model: BaseEstimator, monitor: SklearnMonitor): @@ -202,8 +201,12 @@ def _log_batch_predictions(self, inputs: np.ndarray, outputs: np.ndarray) -> Non predictions = [] for i in range(len(inputs)): pred = { - "inputs": inputs[i].tolist() if isinstance(inputs[i], np.ndarray) else inputs[i], - "output": outputs[i].tolist() if isinstance(outputs[i], np.ndarray) else outputs[i], + "inputs": inputs[i].tolist() + if isinstance(inputs[i], np.ndarray) + else inputs[i], + "output": outputs[i].tolist() + if isinstance(outputs[i], np.ndarray) + else outputs[i], } predictions.append(pred) diff --git a/src/whiteboxai/integrations/tensorflow.py b/whiteboxxai/integrations/tensorflow.py similarity index 87% rename from src/whiteboxai/integrations/tensorflow.py rename to whiteboxxai/integrations/tensorflow.py index b617673..b3dad14 100644 --- a/src/whiteboxai/integrations/tensorflow.py +++ b/whiteboxxai/integrations/tensorflow.py @@ -4,8 +4,13 @@ Integration for monitoring TensorFlow and Keras models. """ +from __future__ import annotations + +import logging import warnings -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, List, Optional, Union + +logger = logging.getLogger(__name__) try: import tensorflow as tf @@ -14,11 +19,19 @@ TENSORFLOW_AVAILABLE = True except ImportError: TENSORFLOW_AVAILABLE = False - keras = None + tf = None -import numpy as np + class _KerasStub: + """Stand-in for `keras` so `keras.callbacks.Callback` stays + resolvable as a base class when TensorFlow isn't installed.""" -from whiteboxai.monitor import ModelMonitor + class callbacks: + Callback = object + + keras = _KerasStub() + +import numpy as np +from whiteboxxai.monitor import ModelMonitor class KerasMonitor(ModelMonitor): @@ -28,8 +41,8 @@ class KerasMonitor(ModelMonitor): Example: ```python from tensorflow import keras - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.tensorflow import KerasMonitor + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.tensorflow import KerasMonitor # Build model model = keras.Sequential([ @@ -39,14 +52,14 @@ class KerasMonitor(ModelMonitor): model.compile(optimizer='adam', loss='mse') # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = KerasMonitor(client, model=model, model_name="keras_model") # Train with monitoring callback - from whiteboxai.integrations.tensorflow import WhiteBoxAICallback + from whiteboxxai.integrations.tensorflow import WhiteBoxXAICallback model.fit(X_train, y_train, - callbacks=[WhiteBoxAICallback(monitor)]) + callbacks=[WhiteBoxXAICallback(monitor)]) # Make predictions with automatic logging predictions = monitor.predict(X_test) @@ -65,14 +78,16 @@ def __init__( Initialize Keras monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance model: Keras model to monitor model_name: Name for the model model_type: Type of model (classification, regression) **kwargs: Additional arguments for ModelMonitor """ if not TENSORFLOW_AVAILABLE: - raise ImportError("TensorFlow is not installed. Install with: pip install tensorflow") + raise ImportError( + "TensorFlow is not installed. Install with: pip install tensorflow" + ) super().__init__(client, **kwargs) self.model = model @@ -89,7 +104,7 @@ def register_from_model( **kwargs: Any, ) -> int: """ - Register Keras model with WhiteBoxAI. + Register Keras model with WhiteBoxXAI. Args: name: Model name (default: model class name) @@ -127,8 +142,8 @@ def register_from_model( try: config = self.model.get_config() metadata["model_config"] = str(config) - except Exception: - pass + except Exception as e: + logger.debug(f"Could not read model config: {e}") # Get input/output shapes try: @@ -136,8 +151,8 @@ def register_from_model( metadata["input_shape"] = str(self.model.input_shape) if hasattr(self.model, "output_shape"): metadata["output_shape"] = str(self.model.output_shape) - except Exception: - pass + except Exception as e: + logger.debug(f"Could not read model input/output shape: {e}") # Register model self.model_id = self.client.register_model( @@ -199,9 +214,9 @@ def predict( # Single prediction self.log_prediction( inputs=inputs_np[0] if len(inputs_np.shape) > 1 else inputs_np, - prediction=( - predictions_np[0] if len(predictions_np.shape) > 1 else predictions_np - ), + prediction=predictions_np[0] + if len(predictions_np.shape) > 1 + else predictions_np, actual=actuals[0] if actuals is not None else None, **kwargs, ) @@ -298,7 +313,7 @@ def register_saved_model( metadata: Optional[Dict[str, Any]] = None, ) -> None: """ - Register a SavedModel with WhiteBoxAI. + Register a SavedModel with WhiteBoxXAI. Args: model_path: Path to SavedModel directory @@ -314,14 +329,14 @@ def register_saved_model( self.register_from_model(**model_metadata) -class WhiteBoxAICallback(keras.callbacks.Callback): +class WhiteBoxXAICallback(keras.callbacks.Callback): """ Keras callback for automatic monitoring during training. Example: ```python monitor = KerasMonitor(client, model=model, model_name="my_model") - callback = WhiteBoxAICallback(monitor, log_frequency=1) + callback = WhiteBoxXAICallback(monitor, log_frequency=1) model.fit(X_train, y_train, validation_data=(X_val, y_val), @@ -393,7 +408,9 @@ def on_train_end(self, logs: Optional[Dict[str, Any]] = None): self.monitor.log_custom_metric( "training_complete", { - "final_metrics": {k: float(v) if v is not None else None for k, v in logs.items()}, + "final_metrics": { + k: float(v) if v is not None else None for k, v in logs.items() + }, }, ) @@ -408,7 +425,7 @@ class TorchMonitor(ModelMonitor): def __init__(self, *args, **kwargs): warnings.warn( "TorchMonitor in tensorflow module is deprecated. " - "Use whiteboxai.integrations.pytorch.TorchMonitor instead.", + "Use whiteboxxai.integrations.pytorch.TorchMonitor instead.", DeprecationWarning, stacklevel=2, ) @@ -441,7 +458,7 @@ def logged_predict(x, *args, **kwargs): # Make prediction predictions = original_predict(x, *args, **kwargs) - # Log to WhiteBoxAI + # Log to WhiteBoxXAI try: if isinstance(x, tf.Tensor): x_np = x.numpy() @@ -469,7 +486,7 @@ def logged_predict(x, *args, **kwargs): __all__ = [ "KerasMonitor", - "WhiteBoxAICallback", + "WhiteBoxXAICallback", "TorchMonitor", # Deprecated "wrap_keras_model", ] diff --git a/src/whiteboxai/integrations/transformers.py b/whiteboxxai/integrations/transformers.py similarity index 94% rename from src/whiteboxai/integrations/transformers.py rename to whiteboxxai/integrations/transformers.py index 550a22d..1fa9b35 100644 --- a/src/whiteboxai/integrations/transformers.py +++ b/whiteboxxai/integrations/transformers.py @@ -5,11 +5,11 @@ """ import warnings -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union try: import transformers - from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer + from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer, pipeline TRANSFORMERS_AVAILABLE = True except ImportError: @@ -18,7 +18,8 @@ PreTrainedTokenizer = object Pipeline = object -from whiteboxai.monitor import ModelMonitor +import numpy as np +from whiteboxxai.monitor import ModelMonitor class TransformersMonitor(ModelMonitor): @@ -37,14 +38,14 @@ class TransformersMonitor(ModelMonitor): Example: ```python from transformers import pipeline - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.transformers import TransformersMonitor + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.transformers import TransformersMonitor # Load model classifier = pipeline("sentiment-analysis") # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = TransformersMonitor( client=client, pipeline=classifier, @@ -70,7 +71,7 @@ def __init__( Initialize Transformers monitor. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance pipeline: Hugging Face pipeline (recommended) model: PreTrainedModel (if not using pipeline) tokenizer: PreTrainedTokenizer (if not using pipeline) @@ -104,7 +105,7 @@ def register_from_model( **kwargs: Any, ) -> int: """ - Register Transformers model with WhiteBoxAI. + Register Transformers model with WhiteBoxXAI. Args: name: Model name (default: model class name or pipeline task) @@ -244,7 +245,9 @@ def predict( # Single prediction self.log_prediction_transformers( input_text=inputs, - prediction=predictions if isinstance(predictions, dict) else predictions[0], + prediction=predictions + if isinstance(predictions, dict) + else predictions[0], actual=actuals, ) @@ -272,7 +275,7 @@ def log_prediction_transformers( else: pred_value = prediction - # Log to WhiteBoxAI + # Log to WhiteBoxXAI self.log_prediction( inputs={"text": input_text}, prediction=pred_value, @@ -401,11 +404,11 @@ class TransformersPipelineWrapper: Example: ```python from transformers import pipeline - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.transformers import TransformersPipelineWrapper + from whiteboxxai import WhiteBoxXAI + from whiteboxxai.integrations.transformers import TransformersPipelineWrapper classifier = pipeline("sentiment-analysis") - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") # Wrap pipeline wrapped = TransformersPipelineWrapper(classifier, client) @@ -427,7 +430,7 @@ def __init__( Args: pipeline: Hugging Face pipeline - client: WhiteBoxAI client + client: WhiteBoxXAI client model_name: Model name auto_register: Auto-register model on first prediction """ @@ -500,7 +503,7 @@ def logged_call(*args, **kwargs): # Make prediction result = original_call(*args, **kwargs) - # Log to WhiteBoxAI + # Log to WhiteBoxXAI try: inputs = args[0] if args else kwargs.get("inputs") diff --git a/src/whiteboxai/monitor.py b/whiteboxxai/monitor.py similarity index 54% rename from src/whiteboxai/monitor.py rename to whiteboxxai/monitor.py index 0fb1460..805ecff 100644 --- a/src/whiteboxai/monitor.py +++ b/whiteboxxai/monitor.py @@ -4,12 +4,13 @@ Simplified monitoring interface for ML models. """ +import time from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np if TYPE_CHECKING: - from whiteboxai.client import WhiteBoxAI + from whiteboxxai.client import WhiteBoxXAI class ModelMonitor: @@ -18,9 +19,9 @@ class ModelMonitor: Example: ```python - from whiteboxai import WhiteBoxAI, ModelMonitor + from whiteboxxai import WhiteBoxXAI, ModelMonitor - client = WhiteBoxAI(api_key="your-api-key") + client = WhiteBoxXAI(api_key="your-api-key") monitor = ModelMonitor(client, model_id=123) # Log prediction @@ -39,29 +40,44 @@ class ModelMonitor: def __init__( self, - client: "WhiteBoxAI", + client: "WhiteBoxXAI", model_id: Optional[int] = None, model_name: Optional[str] = None, auto_explain: bool = False, sampling_rate: float = 1.0, + buffer_size: Optional[int] = None, ): """ Initialize model monitor. Args: - client: WhiteBoxAI client instance - model_id: Model ID (if already registered) + client: WhiteBoxXAI client instance + model_id: Model ID (if already registered; can be set later via + register_model()) model_name: Model name (for registration) auto_explain: Automatically generate explanations sampling_rate: Prediction sampling rate (0.0-1.0) + buffer_size: If set, log_prediction() buffers locally and sends + predictions as a batch once the buffer reaches this size (or + on flush()/context-manager exit), instead of sending + immediately on every call. """ self.client = client self.model_id = model_id self.model_name = model_name self.auto_explain = auto_explain self.sampling_rate = sampling_rate + self.buffer_size = buffer_size + self._buffer: List[Dict[str, Any]] = [] + self._prediction_count = 0 self._baseline_data: Optional[np.ndarray] = None + def __enter__(self) -> "ModelMonitor": + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + self.flush() + def register_model( self, name: str, @@ -131,24 +147,43 @@ def log_prediction( metadata: Additional metadata Returns: - Prediction data if sampled, None if skipped + Prediction data if sent immediately, None if buffered or skipped """ + if output is None: + raise ValueError("output cannot be None.") + if not self._should_sample(): return None if self.model_id is None: raise ValueError("Model not registered. Call register_model() first.") + self._prediction_count += 1 + + if self.buffer_size: + self._buffer.append( + {"input_data": inputs, "output_data": output, "metadata": metadata} + ) + if len(self._buffer) >= self.buffer_size: + self.flush() + return None + explain = explain if explain is not None else self.auto_explain - return self.client.predictions.log( + result = self.client.predictions.log( model_id=self.model_id, - inputs=inputs, - outputs=output, - explain=explain, + input_data=inputs, + output_data=output, metadata=metadata, ) + if explain and isinstance(result, dict) and result.get("id"): + result["explanation"] = self.client.explanations.generate( + prediction_id=result["id"] + ) + + return result + async def alog_prediction( self, inputs: Any, @@ -157,22 +192,54 @@ async def alog_prediction( metadata: Optional[Dict[str, Any]] = None, ) -> Optional[Dict[str, Any]]: """Async version of log_prediction().""" + if output is None: + raise ValueError("output cannot be None.") + if not self._should_sample(): return None if self.model_id is None: raise ValueError("Model not registered. Call register_model() first.") + self._prediction_count += 1 + + if self.buffer_size: + self._buffer.append( + {"input_data": inputs, "output_data": output, "metadata": metadata} + ) + if len(self._buffer) >= self.buffer_size: + await self.aflush() + return None + explain = explain if explain is not None else self.auto_explain - return await self.client.predictions.alog( + result = await self.client.predictions.alog( model_id=self.model_id, - inputs=inputs, - outputs=output, - explain=explain, + input_data=inputs, + output_data=output, metadata=metadata, ) + if explain and isinstance(result, dict) and result.get("id"): + result["explanation"] = await self.client.explanations.agenerate( + prediction_id=result["id"] + ) + + return result + + @staticmethod + def _normalize_batch_item(prediction: Dict[str, Any]) -> Dict[str, Any]: + """Map the SDK's `inputs`/`output` shorthand keys to the + `input_data`/`output_data` keys the predictions API expects.""" + if "inputs" in prediction or "output" in prediction: + normalized = dict(prediction) + if "inputs" in normalized: + normalized["input_data"] = normalized.pop("inputs") + if "output" in normalized: + normalized["output_data"] = normalized.pop("output") + return normalized + return prediction + def log_batch( self, predictions: List[Dict[str, Any]], @@ -181,7 +248,8 @@ def log_batch( Log multiple predictions in batch. Args: - predictions: List of prediction dictionaries + predictions: List of prediction dictionaries (each with + `inputs`/`output` or `input_data`/`output_data` keys) Returns: Batch logging result @@ -193,6 +261,8 @@ def log_batch( if self.sampling_rate < 1.0: predictions = self._sample_predictions(predictions) + predictions = [self._normalize_batch_item(p) for p in predictions] + return self.client.predictions.log_batch( model_id=self.model_id, predictions=predictions, @@ -210,11 +280,81 @@ async def alog_batch( if self.sampling_rate < 1.0: predictions = self._sample_predictions(predictions) + predictions = [self._normalize_batch_item(p) for p in predictions] + + return await self.client.predictions.alog_batch( + model_id=self.model_id, + predictions=predictions, + ) + + def flush(self) -> Optional[Dict[str, Any]]: + """Send any buffered predictions as a batch and clear the buffer.""" + if not self._buffer: + return None + predictions, self._buffer = self._buffer, [] + return self.client.predictions.log_batch( + model_id=self.model_id, + predictions=predictions, + ) + + async def aflush(self) -> Optional[Dict[str, Any]]: + """Async version of flush().""" + if not self._buffer: + return None + predictions, self._buffer = self._buffer, [] return await self.client.predictions.alog_batch( model_id=self.model_id, predictions=predictions, ) + def get_prediction_count(self) -> int: + """Return the number of predictions logged (sent or buffered) by + this monitor instance so far.""" + return self._prediction_count + + def get_drift_reports( + self, limit: int = 10, skip: int = 0 + ) -> List[Dict[str, Any]]: + """Get persisted drift reports for this model, most recent first.""" + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + return self.client.drift.get_reports( + model_id=self.model_id, limit=limit, skip=skip + ) + + def create_alert_rule( + self, + metric: str, + threshold: float, + condition: str, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Create an alert rule for this model. + + Args: + metric: Metric to monitor (e.g. "accuracy") + threshold: Threshold value that triggers the alert + condition: Comparison condition (e.g. "below", "above") + **kwargs: Additional alert configuration (name, etc.) + + Returns: + Created alert rule data + """ + name = kwargs.pop("name", f"{self.model_id}_{metric}_{condition}_{threshold}") + return self.client.alerts.create( + name=name, + alert_type="threshold", + conditions={"metric": metric, "threshold": threshold, "condition": condition}, + model_id=self.model_id, + **kwargs, + ) + + def get_active_alerts(self) -> List[Dict[str, Any]]: + """Get alert rules for this model.""" + return self.client.alerts.list(model_id=self.model_id) + def set_baseline(self, data: np.ndarray) -> None: """ Set baseline data for drift detection. @@ -244,9 +384,9 @@ def detect_drift( return self.client.drift.detect( model_id=self.model_id, - reference_data=( - self._baseline_data.tolist() if self._baseline_data is not None else None - ), + reference_data=self._baseline_data.tolist() + if self._baseline_data is not None + else None, current_data=current_data.tolist() if current_data is not None else None, **kwargs, ) @@ -262,9 +402,9 @@ async def adetect_drift( return await self.client.drift.adetect( model_id=self.model_id, - reference_data=( - self._baseline_data.tolist() if self._baseline_data is not None else None - ), + reference_data=self._baseline_data.tolist() + if self._baseline_data is not None + else None, current_data=current_data.tolist() if current_data is not None else None, **kwargs, ) @@ -275,7 +415,9 @@ def _should_sample(self) -> bool: return True return np.random.random() < self.sampling_rate - def _sample_predictions(self, predictions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + def _sample_predictions( + self, predictions: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: """Sample predictions based on sampling rate.""" n_samples = int(len(predictions) * self.sampling_rate) if n_samples == 0: diff --git a/src/whiteboxai/offline.py b/whiteboxxai/offline.py similarity index 91% rename from src/whiteboxai/offline.py rename to whiteboxxai/offline.py index 903a77b..28d1d5e 100644 --- a/src/whiteboxai/offline.py +++ b/whiteboxxai/offline.py @@ -1,5 +1,5 @@ """ -Offline mode support for WhiteBoxAI SDK. +Offline mode support for WhiteBoxXAI SDK. This module provides offline queueing capabilities for the SDK, allowing operations to be queued locally when the API is unavailable and synced @@ -15,11 +15,11 @@ Example: Enable offline mode: - >>> from whiteboxai import WhiteBoxAI - >>> client = WhiteBoxAI( + >>> from whiteboxxai import WhiteBoxXAI + >>> client = WhiteBoxXAI( ... api_key="your-api-key", ... enable_offline=True, - ... offline_dir="./whiteboxai_offline" + ... offline_dir="./whiteboxxai_offline" ... ) >>> >>> # Operations are queued when offline @@ -34,9 +34,11 @@ import os import sqlite3 import threading +import time +from datetime import datetime from enum import Enum from pathlib import Path -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple logger = logging.getLogger(__name__) @@ -72,7 +74,9 @@ class OfflineQueue: auto_sync: Whether to automatically sync when connection available """ - def __init__(self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True): + def __init__( + self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True + ): """ Initialize offline queue. @@ -143,7 +147,9 @@ def enqueue( if self.max_queue_size > 0: current_size = self.get_queue_size() if current_size >= self.max_queue_size: - raise ValueError(f"Queue is full ({current_size}/{self.max_queue_size})") + raise ValueError( + f"Queue is full ({current_size}/{self.max_queue_size})" + ) # Insert operation with sqlite3.connect(self.db_path) as conn: @@ -160,7 +166,9 @@ def enqueue( logger.info(f"Queued operation {op_id}: {operation_type.value}") return op_id - def dequeue(self, limit: int = 100) -> List[Tuple[int, OperationType, Dict[str, Any]]]: + def dequeue( + self, limit: int = 100 + ) -> List[Tuple[int, OperationType, Dict[str, Any]]]: """ Get pending operations from queue. @@ -189,7 +197,9 @@ def dequeue(self, limit: int = 100) -> List[Tuple[int, OperationType, Dict[str, operations = [] for row in cursor.fetchall(): op_id, op_type, data_json = row - operations.append((op_id, OperationType(op_type), json.loads(data_json))) + operations.append( + (op_id, OperationType(op_type), json.loads(data_json)) + ) return operations @@ -278,7 +288,9 @@ def get_queue_size(self, status: str = "pending") -> int: """ with sqlite3.connect(self.db_path) as conn: if status: - cursor = conn.execute("SELECT COUNT(*) FROM queue WHERE status = ?", (status,)) + cursor = conn.execute( + "SELECT COUNT(*) FROM queue WHERE status = ?", (status,) + ) else: cursor = conn.execute("SELECT COUNT(*) FROM queue") @@ -330,7 +342,9 @@ def clear_completed(self, older_than_days: int = 7): conn.commit() if deleted > 0: - logger.info(f"Cleared {deleted} completed operations older than {older_than_days} days") + logger.info( + f"Cleared {deleted} completed operations older than {older_than_days} days" + ) def clear_all(self): """Clear all operations from queue (use with caution).""" @@ -377,7 +391,7 @@ def get_failed_operations(self) -> List[Dict[str, Any]]: class OfflineManager: """ - Manages offline mode for WhiteBoxAI SDK. + Manages offline mode for WhiteBoxXAI SDK. Handles automatic queueing when offline and syncing when online. @@ -389,7 +403,7 @@ class OfflineManager: def __init__( self, - offline_dir: str = "./whiteboxai_offline", + offline_dir: str = "./whiteboxxai_offline", max_queue_size: int = 10000, auto_sync: bool = True, sync_interval: int = 60, @@ -417,17 +431,17 @@ def __init__( self.max_retries = max_retries self._sync_thread = None self._stop_sync = threading.Event() - self._client = None # Will be set by WhiteBoxAI client + self._client = None # Will be set by WhiteBoxXAI client if auto_sync: self.start_auto_sync() def set_client(self, client): """ - Set the WhiteBoxAI client for syncing. + Set the WhiteBoxXAI client for syncing. Args: - client: WhiteBoxAI client instance + client: WhiteBoxXAI client instance """ self._client = client @@ -435,7 +449,9 @@ def start_auto_sync(self): """Start automatic sync thread.""" if self._sync_thread is None or not self._sync_thread.is_alive(): self._stop_sync.clear() - self._sync_thread = threading.Thread(target=self._auto_sync_loop, daemon=True) + self._sync_thread = threading.Thread( + target=self._auto_sync_loop, daemon=True + ) self._sync_thread.start() logger.info("Started automatic sync thread") diff --git a/src/whiteboxai/privacy.py b/whiteboxxai/privacy.py similarity index 95% rename from src/whiteboxai/privacy.py rename to whiteboxxai/privacy.py index 86c97b8..745b308 100644 --- a/src/whiteboxai/privacy.py +++ b/whiteboxxai/privacy.py @@ -5,7 +5,7 @@ """ import re -from typing import Any, Dict, List, Optional, Pattern +from typing import Any, Dict, List, Optional, Pattern, Union class PIIDetector: @@ -24,7 +24,9 @@ def __init__(self): """Initialize PII detector with regex patterns.""" self.patterns: Dict[str, Pattern] = { "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"), - "phone": re.compile(r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), + "phone": re.compile( + r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b" + ), "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "credit_card": re.compile(r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b"), "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), @@ -204,7 +206,9 @@ def mask_pii(text: str, mask_char: str = "*") -> str: return _pii_detector.mask(text, mask_char=mask_char) -def mask_data(data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True) -> Any: +def mask_data( + data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True +) -> Any: """ Mask sensitive data using global masker. diff --git a/whiteboxxai/resources.py b/whiteboxxai/resources.py new file mode 100644 index 0000000..3d0cc42 --- /dev/null +++ b/whiteboxxai/resources.py @@ -0,0 +1,872 @@ +""" +API Resources + +Resource classes for interacting with different WhiteBoxXAI API endpoints. +""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +if TYPE_CHECKING: + from whiteboxxai.client import WhiteBoxXAI + +logger = logging.getLogger(__name__) + + +class BaseResource: + """Base class for API resources.""" + + def __init__(self, client: "WhiteBoxXAI"): + """Initialize resource with client.""" + self.client = client + + +class ModelsResource(BaseResource): + """Models API resource (backed by backend/api/v1/models.py).""" + + def _build_register_data( + self, + name: str, + model_type: str, + framework: Optional[str], + version: Optional[str], + auto_detect_git: bool, + require_clean_git: bool, + kwargs: Dict[str, Any], + ) -> Dict[str, Any]: + data = { + "name": name, + "model_type": model_type, + "framework": framework, + "version": version, + **kwargs, + } + + if auto_detect_git: + from whiteboxxai.git_utils import detect_git_context, validate_git_context + + git_context = detect_git_context() + if git_context: + if validate_git_context(git_context, require_clean=require_clean_git): + git_data = git_context.to_dict() + for key, value in git_data.items(): + if key not in data and value is not None: + data[key] = value + logger.info(f"Auto-detected Git context: {git_context}") + else: + logger.warning("Git context validation failed") + else: + logger.warning( + "auto_detect_git=True but no Git repository found. " + "Proceeding without Git metadata." + ) + + return data + + def register( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + auto_detect_git: bool = False, + require_clean_git: bool = False, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Register a new model. + + Args: + name: Model name + model_type: Model type (classification, regression, etc.) + framework: ML framework (sklearn, pytorch, tensorflow, etc.) + version: Model version + auto_detect_git: Automatically detect Git repository context (commit, branch, etc.) + require_clean_git: Require clean Git working directory (no uncommitted changes) + **kwargs: Additional model metadata (features, baseline_metrics, tags, etc.) + + Returns: + Registered model data + + Example: + >>> client.models.register( + ... name="fraud_detector", + ... model_type="classification", + ... framework="sklearn", + ... features=["amount", "merchant_category"], + ... baseline_metrics={"accuracy": 0.94}, + ... ) + """ + data = self._build_register_data( + name, + model_type, + framework, + version, + auto_detect_git, + require_clean_git, + kwargs, + ) + return self.client.request("POST", "/api/v1/models/", data=data) + + async def aregister( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + auto_detect_git: bool = False, + require_clean_git: bool = False, + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of register().""" + data = self._build_register_data( + name, + model_type, + framework, + version, + auto_detect_git, + require_clean_git, + kwargs, + ) + return await self.client.arequest("POST", "/api/v1/models/", data=data) + + def get(self, model_id: str) -> Dict[str, Any]: + """Get model by ID.""" + return self.client.request("GET", f"/api/v1/models/{model_id}") + + async def aget(self, model_id: str) -> Dict[str, Any]: + """Async version of get().""" + return await self.client.arequest("GET", f"/api/v1/models/{model_id}") + + def list( + self, + skip: int = 0, + limit: int = 100, + status_filter: Optional[str] = None, + model_type: Optional[str] = None, + owner_id: Optional[str] = None, + tags: Optional[str] = None, + search: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """List registered models with optional filters.""" + params: Dict[str, Any] = {"skip": skip, "limit": limit} + if status_filter is not None: + params["status_filter"] = status_filter + if model_type is not None: + params["model_type"] = model_type + if owner_id is not None: + params["owner_id"] = owner_id + if tags is not None: + params["tags"] = tags + if search is not None: + params["search"] = search + return self.client.request("GET", "/api/v1/models/", params=params) + + async def alist( + self, + skip: int = 0, + limit: int = 100, + status_filter: Optional[str] = None, + model_type: Optional[str] = None, + owner_id: Optional[str] = None, + tags: Optional[str] = None, + search: Optional[str] = None, + ) -> List[Dict[str, Any]]: + """Async version of list().""" + params: Dict[str, Any] = {"skip": skip, "limit": limit} + if status_filter is not None: + params["status_filter"] = status_filter + if model_type is not None: + params["model_type"] = model_type + if owner_id is not None: + params["owner_id"] = owner_id + if tags is not None: + params["tags"] = tags + if search is not None: + params["search"] = search + return await self.client.arequest("GET", "/api/v1/models/", params=params) + + def update(self, model_id: str, **kwargs: Any) -> Dict[str, Any]: + """Update model metadata (partial update).""" + return self.client.request("PATCH", f"/api/v1/models/{model_id}", data=kwargs) + + async def aupdate(self, model_id: str, **kwargs: Any) -> Dict[str, Any]: + """Async version of update().""" + return await self.client.arequest( + "PATCH", f"/api/v1/models/{model_id}", data=kwargs + ) + + def update_status(self, model_id: str, new_status: str) -> Dict[str, Any]: + """Update model status (ACTIVE, INACTIVE, DEPRECATED, ARCHIVED).""" + return self.client.request( + "PATCH", + f"/api/v1/models/{model_id}/status", + params={"new_status": new_status}, + ) + + async def aupdate_status(self, model_id: str, new_status: str) -> Dict[str, Any]: + """Async version of update_status().""" + return await self.client.arequest( + "PATCH", + f"/api/v1/models/{model_id}/status", + params={"new_status": new_status}, + ) + + def get_versions( + self, model_name: str, skip: int = 0, limit: int = 100 + ) -> List[Dict[str, Any]]: + """Get all versions of a model by name.""" + return self.client.request( + "GET", + f"/api/v1/models/{model_name}/versions", + params={"skip": skip, "limit": limit}, + ) + + async def aget_versions( + self, model_name: str, skip: int = 0, limit: int = 100 + ) -> List[Dict[str, Any]]: + """Async version of get_versions().""" + return await self.client.arequest( + "GET", + f"/api/v1/models/{model_name}/versions", + params={"skip": skip, "limit": limit}, + ) + + def get_latest(self, model_name: str) -> Dict[str, Any]: + """Get the latest version of a model by name.""" + return self.client.request("GET", f"/api/v1/models/{model_name}/latest") + + async def aget_latest(self, model_name: str) -> Dict[str, Any]: + """Async version of get_latest().""" + return await self.client.arequest("GET", f"/api/v1/models/{model_name}/latest") + + def update_baseline( + self, + model_id: str, + baseline_metrics: Dict[str, Any], + baseline_data_hash: Optional[str] = None, + baseline_data_count: Optional[int] = None, + ) -> Dict[str, Any]: + """Update model baseline metrics/profile (e.g. after retraining).""" + params: Dict[str, Any] = {} + if baseline_data_hash is not None: + params["baseline_data_hash"] = baseline_data_hash + if baseline_data_count is not None: + params["baseline_data_count"] = baseline_data_count + return self.client.request( + "PATCH", + f"/api/v1/models/{model_id}/baseline", + data=baseline_metrics, + params=params, + ) + + async def aupdate_baseline( + self, + model_id: str, + baseline_metrics: Dict[str, Any], + baseline_data_hash: Optional[str] = None, + baseline_data_count: Optional[int] = None, + ) -> Dict[str, Any]: + """Async version of update_baseline().""" + params: Dict[str, Any] = {} + if baseline_data_hash is not None: + params["baseline_data_hash"] = baseline_data_hash + if baseline_data_count is not None: + params["baseline_data_count"] = baseline_data_count + return await self.client.arequest( + "PATCH", + f"/api/v1/models/{model_id}/baseline", + data=baseline_metrics, + params=params, + ) + + def archive(self, model_id: str) -> Dict[str, Any]: + """Archive a model (soft delete).""" + return self.client.request("POST", f"/api/v1/models/{model_id}/archive") + + async def aarchive(self, model_id: str) -> Dict[str, Any]: + """Async version of archive().""" + return await self.client.arequest("POST", f"/api/v1/models/{model_id}/archive") + + def restore(self, model_id: str) -> Dict[str, Any]: + """Restore an archived model to active status.""" + return self.client.request("POST", f"/api/v1/models/{model_id}/restore") + + async def arestore(self, model_id: str) -> Dict[str, Any]: + """Async version of restore().""" + return await self.client.arequest("POST", f"/api/v1/models/{model_id}/restore") + + def delete(self, model_id: str) -> Dict[str, Any]: + """Permanently delete a model (hard delete, cannot be undone).""" + return self.client.request("DELETE", f"/api/v1/models/{model_id}") + + async def adelete(self, model_id: str) -> Dict[str, Any]: + """Async version of delete().""" + return await self.client.arequest("DELETE", f"/api/v1/models/{model_id}") + + +class PredictionsResource(BaseResource): + """Predictions API resource (backed by backend/api/v1/predictions.py).""" + + def log( + self, + model_id: str, + input_data: Dict[str, Any], + output_data: Dict[str, Any], + prediction_id: Optional[str] = None, + latency_ms: Optional[float] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """ + Log a single prediction. + + Args: + model_id: Model ID + input_data: Input features/data + output_data: Model prediction/output + prediction_id: Optional external prediction ID + latency_ms: Prediction latency in milliseconds + metadata: Additional prediction metadata + + Returns: + Logged prediction data + """ + data: Dict[str, Any] = { + "model_id": model_id, + "input_data": input_data, + "output_data": output_data, + } + if prediction_id is not None: + data["prediction_id"] = prediction_id + if latency_ms is not None: + data["latency_ms"] = latency_ms + if metadata is not None: + data["metadata"] = metadata + return self.client.request("POST", "/api/v1/predictions/log", data=data) + + async def alog( + self, + model_id: str, + input_data: Dict[str, Any], + output_data: Dict[str, Any], + prediction_id: Optional[str] = None, + latency_ms: Optional[float] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Async version of log().""" + data: Dict[str, Any] = { + "model_id": model_id, + "input_data": input_data, + "output_data": output_data, + } + if prediction_id is not None: + data["prediction_id"] = prediction_id + if latency_ms is not None: + data["latency_ms"] = latency_ms + if metadata is not None: + data["metadata"] = metadata + return await self.client.arequest("POST", "/api/v1/predictions/log", data=data) + + def log_batch( + self, + model_id: str, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """ + Log multiple predictions in batch (max 1000 per call). + + Args: + model_id: Model ID + predictions: List of dicts, each with input_data/output_data + (and optionally prediction_id/latency_ms/metadata) + + Returns: + Batch logging summary: {total, logged, failed, errors} + """ + data = {"model_id": model_id, "predictions": predictions} + return self.client.request("POST", "/api/v1/predictions/log/batch", data=data) + + async def alog_batch( + self, + model_id: str, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """Async version of log_batch().""" + data = {"model_id": model_id, "predictions": predictions} + return await self.client.arequest( + "POST", "/api/v1/predictions/log/batch", data=data + ) + + def get(self, prediction_id: str) -> Dict[str, Any]: + """Get a prediction by its ID.""" + return self.client.request("GET", f"/api/v1/predictions/{prediction_id}") + + async def aget(self, prediction_id: str) -> Dict[str, Any]: + """Async version of get().""" + return await self.client.arequest("GET", f"/api/v1/predictions/{prediction_id}") + + def query( + self, + model_id: str, + start_time: Optional[str] = None, + end_time: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> List[Dict[str, Any]]: + """ + Query predictions with filters. + + Args: + model_id: Model ID (required by the backend) + start_time: ISO-8601 timestamp, filter predictions after this time + end_time: ISO-8601 timestamp, filter predictions before this time + limit: Maximum results to return + offset: Number of results to skip + """ + data: Dict[str, Any] = {"model_id": model_id, "limit": limit, "offset": offset} + if start_time is not None: + data["start_time"] = start_time + if end_time is not None: + data["end_time"] = end_time + return self.client.request("POST", "/api/v1/predictions/query", data=data) + + async def aquery( + self, + model_id: str, + start_time: Optional[str] = None, + end_time: Optional[str] = None, + limit: int = 100, + offset: int = 0, + ) -> List[Dict[str, Any]]: + """Async version of query().""" + data: Dict[str, Any] = {"model_id": model_id, "limit": limit, "offset": offset} + if start_time is not None: + data["start_time"] = start_time + if end_time is not None: + data["end_time"] = end_time + return await self.client.arequest( + "POST", "/api/v1/predictions/query", data=data + ) + + def get_stats( + self, + model_id: str, + start_time: Optional[str] = None, + end_time: Optional[str] = None, + ) -> Dict[str, Any]: + """Get prediction statistics for a model.""" + params: Dict[str, Any] = {} + if start_time is not None: + params["start_time"] = start_time + if end_time is not None: + params["end_time"] = end_time + return self.client.request( + "GET", f"/api/v1/predictions/models/{model_id}/stats", params=params + ) + + async def aget_stats( + self, + model_id: str, + start_time: Optional[str] = None, + end_time: Optional[str] = None, + ) -> Dict[str, Any]: + """Async version of get_stats().""" + params: Dict[str, Any] = {} + if start_time is not None: + params["start_time"] = start_time + if end_time is not None: + params["end_time"] = end_time + return await self.client.arequest( + "GET", f"/api/v1/predictions/models/{model_id}/stats", params=params + ) + + def get_recent(self, model_id: str, limit: int = 10) -> List[Dict[str, Any]]: + """Get the most recent predictions for a model (max 100).""" + return self.client.request( + "GET", + f"/api/v1/predictions/models/{model_id}/recent", + params={"limit": limit}, + ) + + async def aget_recent(self, model_id: str, limit: int = 10) -> List[Dict[str, Any]]: + """Async version of get_recent().""" + return await self.client.arequest( + "GET", + f"/api/v1/predictions/models/{model_id}/recent", + params={"limit": limit}, + ) + + +class ExplanationsResource(BaseResource): + """Explanations API resource.""" + + def generate( + self, + prediction_id: int, + method: str = "shap", + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Generate explanation for a prediction. + + Args: + prediction_id: Prediction ID + method: Explanation method (shap, lime) + **kwargs: Method-specific parameters + + Returns: + Explanation data + """ + data = {"prediction_id": prediction_id, "method": method, **kwargs} + return self.client.request("POST", "/api/v1/explanations/generate", data=data) + + async def agenerate( + self, + prediction_id: int, + method: str = "shap", + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of generate().""" + data = {"prediction_id": prediction_id, "method": method, **kwargs} + return await self.client.arequest( + "POST", "/api/v1/explanations/generate", data=data + ) + + def get(self, explanation_id: int) -> Dict[str, Any]: + """Get explanation by ID.""" + return self.client.request("GET", f"/api/v1/explanations/{explanation_id}") + + async def aget(self, explanation_id: int) -> Dict[str, Any]: + """Async version of get().""" + return await self.client.arequest( + "GET", f"/api/v1/explanations/{explanation_id}" + ) + + +class DriftResource(BaseResource): + """Drift detection API resource (backed by backend/api/v1/drift.py).""" + + def detect( + self, + model_id: str, + window_size: int = 1000, + feature_names: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """ + Detect data drift for a model (ephemeral, not persisted). + + Args: + model_id: Model ID + window_size: Number of recent predictions to analyze (100-10000) + feature_names: Optional subset of features to analyze + + Returns: + Drift detection results + """ + params: Dict[str, Any] = {"window_size": window_size} + if feature_names is not None: + params["feature_names"] = feature_names + return self.client.request( + "POST", f"/api/v1/drift/detect/{model_id}", params=params + ) + + async def adetect( + self, + model_id: str, + window_size: int = 1000, + feature_names: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """Async version of detect().""" + params: Dict[str, Any] = {"window_size": window_size} + if feature_names is not None: + params["feature_names"] = feature_names + return await self.client.arequest( + "POST", f"/api/v1/drift/detect/{model_id}", params=params + ) + + def create_report(self, model_id: str, window_size: int = 1000) -> Dict[str, Any]: + """Run drift analysis and persist the result as a drift report.""" + return self.client.request( + "POST", + "/api/v1/drift/reports", + params={"model_id": model_id, "window_size": window_size}, + ) + + async def acreate_report( + self, model_id: str, window_size: int = 1000 + ) -> Dict[str, Any]: + """Async version of create_report().""" + return await self.client.arequest( + "POST", + "/api/v1/drift/reports", + params={"model_id": model_id, "window_size": window_size}, + ) + + def get_reports( + self, model_id: str, limit: int = 10, skip: int = 0 + ) -> List[Dict[str, Any]]: + """List drift reports for a model, most recent first.""" + return self.client.request( + "GET", + f"/api/v1/drift/reports/{model_id}", + params={"limit": limit, "skip": skip}, + ) + + async def aget_reports( + self, model_id: str, limit: int = 10, skip: int = 0 + ) -> List[Dict[str, Any]]: + """Async version of get_reports().""" + return await self.client.arequest( + "GET", + f"/api/v1/drift/reports/{model_id}", + params={"limit": limit, "skip": skip}, + ) + + def get_report(self, model_id: str, report_id: str) -> Dict[str, Any]: + """Get a specific drift report with per-feature statistics.""" + return self.client.request( + "GET", f"/api/v1/drift/reports/{model_id}/{report_id}" + ) + + async def aget_report(self, model_id: str, report_id: str) -> Dict[str, Any]: + """Async version of get_report().""" + return await self.client.arequest( + "GET", f"/api/v1/drift/reports/{model_id}/{report_id}" + ) + + def get_trend(self, model_id: str, days: int = 7) -> Dict[str, Any]: + """Get drift trend over time for a model.""" + return self.client.request( + "GET", f"/api/v1/drift/trend/{model_id}", params={"days": days} + ) + + async def aget_trend(self, model_id: str, days: int = 7) -> Dict[str, Any]: + """Async version of get_trend().""" + return await self.client.arequest( + "GET", f"/api/v1/drift/trend/{model_id}", params={"days": days} + ) + + +class FairnessResource(BaseResource): + """Bias/fairness auditing API resource (backed by backend/api/v1/fairness.py).""" + + def audit( + self, + model_id: str, + sensitive_attributes: List[str], + y_true: List[int], + y_pred: List[int], + group_data: Dict[str, List[Any]], + y_prob: Optional[List[float]] = None, + fairness_thresholds: Optional[Dict[str, float]] = None, + metrics_to_compute: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """ + Run a bias audit on model predictions. + + Note: unlike other endpoints, this is not a data-source query — the + caller must already have the true labels, predicted labels, and + per-attribute group assignments in hand (e.g. pulled via + PredictionsResource.query() and joined with demographic data). + + Args: + model_id: Model ID (UUID string) + sensitive_attributes: Protected attributes to analyze (e.g. ["gender", "race"]) + y_true: Ground-truth labels + y_pred: Predicted labels + group_data: Per-attribute group assignments, e.g. {"gender": ["M", "F", ...]} + y_prob: Optional predicted probabilities + fairness_thresholds: Optional custom thresholds per fairness metric type + metrics_to_compute: Optional subset of fairness metrics to compute (all if None) + + Returns: + Complete bias audit results with metrics, group comparisons, and recommendations + """ + data: Dict[str, Any] = { + "model_id": model_id, + "sensitive_attributes": sensitive_attributes, + "y_true": y_true, + "y_pred": y_pred, + "group_data": group_data, + } + if y_prob is not None: + data["y_prob"] = y_prob + if fairness_thresholds is not None: + data["fairness_thresholds"] = fairness_thresholds + if metrics_to_compute is not None: + data["metrics_to_compute"] = metrics_to_compute + return self.client.request("POST", "/api/v1/fairness/audit", data=data) + + async def aaudit( + self, + model_id: str, + sensitive_attributes: List[str], + y_true: List[int], + y_pred: List[int], + group_data: Dict[str, List[Any]], + y_prob: Optional[List[float]] = None, + fairness_thresholds: Optional[Dict[str, float]] = None, + metrics_to_compute: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """Async version of audit().""" + data: Dict[str, Any] = { + "model_id": model_id, + "sensitive_attributes": sensitive_attributes, + "y_true": y_true, + "y_pred": y_pred, + "group_data": group_data, + } + if y_prob is not None: + data["y_prob"] = y_prob + if fairness_thresholds is not None: + data["fairness_thresholds"] = fairness_thresholds + if metrics_to_compute is not None: + data["metrics_to_compute"] = metrics_to_compute + return await self.client.arequest("POST", "/api/v1/fairness/audit", data=data) + + def get_audit(self, audit_id: str) -> Dict[str, Any]: + """Get bias audit results by ID.""" + return self.client.request("GET", f"/api/v1/fairness/audits/{audit_id}") + + async def aget_audit(self, audit_id: str) -> Dict[str, Any]: + """Async version of get_audit().""" + return await self.client.arequest("GET", f"/api/v1/fairness/audits/{audit_id}") + + def list_audits( + self, + model_id: Optional[str] = None, + limit: int = 50, + offset: int = 0, + ) -> List[Dict[str, Any]]: + """List bias audits, optionally filtered by model.""" + params: Dict[str, Any] = {"limit": limit, "offset": offset} + if model_id is not None: + params["model_id"] = model_id + return self.client.request("GET", "/api/v1/fairness/audits", params=params) + + async def alist_audits( + self, + model_id: Optional[str] = None, + limit: int = 50, + offset: int = 0, + ) -> List[Dict[str, Any]]: + """Async version of list_audits().""" + params: Dict[str, Any] = {"limit": limit, "offset": offset} + if model_id is not None: + params["model_id"] = model_id + return await self.client.arequest( + "GET", "/api/v1/fairness/audits", params=params + ) + + def get_bias_history(self, model_id: str, days: int = 30) -> Dict[str, Any]: + """Get historical bias metrics/trend for a model.""" + return self.client.request( + "GET", + f"/api/v1/fairness/models/{model_id}/bias-history", + params={"days": days}, + ) + + async def aget_bias_history(self, model_id: str, days: int = 30) -> Dict[str, Any]: + """Async version of get_bias_history().""" + return await self.client.arequest( + "GET", + f"/api/v1/fairness/models/{model_id}/bias-history", + params={"days": days}, + ) + + def get_metric_history( + self, model_id: str, metric_type: str, days: int = 30 + ) -> Dict[str, Any]: + """Get history for a specific fairness metric type.""" + return self.client.request( + "GET", + f"/api/v1/fairness/models/{model_id}/metrics/{metric_type}/history", + params={"days": days}, + ) + + async def aget_metric_history( + self, model_id: str, metric_type: str, days: int = 30 + ) -> Dict[str, Any]: + """Async version of get_metric_history().""" + return await self.client.arequest( + "GET", + f"/api/v1/fairness/models/{model_id}/metrics/{metric_type}/history", + params={"days": days}, + ) + + def get_latest_audit(self, model_id: str) -> Dict[str, Any]: + """Get the most recent bias audit for a model.""" + return self.client.request( + "GET", f"/api/v1/fairness/models/{model_id}/latest-audit" + ) + + async def aget_latest_audit(self, model_id: str) -> Dict[str, Any]: + """Async version of get_latest_audit().""" + return await self.client.arequest( + "GET", f"/api/v1/fairness/models/{model_id}/latest-audit" + ) + + +class AlertsResource(BaseResource): + """Alerts API resource. + + NOTE: as of this writing, backend/api/v1/alerts.py is not registered in + backend/api/v1/__init__.py (SQLAlchemy model collisions with the + canonical alert tables -- see the NOTE comment there), so every method + on this resource currently calls a route that does not exist on a live + WhiteBoxXAI backend. Left in place for interface stability; do not build + new functionality on top of it until the backend router is wired in. + """ + + def create( + self, + name: str, + alert_type: str, + conditions: Dict[str, Any], + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Create an alert rule. + + Args: + name: Alert name + alert_type: Alert type + conditions: Alert conditions + **kwargs: Additional alert configuration + + Returns: + Created alert data + """ + data = { + "name": name, + "alert_type": alert_type, + "conditions": conditions, + **kwargs, + } + return self.client.request("POST", "/api/v1/alerts", data=data) + + async def acreate( + self, + name: str, + alert_type: str, + conditions: Dict[str, Any], + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of create().""" + data = { + "name": name, + "alert_type": alert_type, + "conditions": conditions, + **kwargs, + } + return await self.client.arequest("POST", "/api/v1/alerts", data=data) + + def list(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: + """List alert rules.""" + params = {"model_id": model_id} if model_id else {} + return self.client.request("GET", "/api/v1/alerts", params=params) + + async def alist(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: + """Async version of list().""" + params = {"model_id": model_id} if model_id else {} + return await self.client.arequest("GET", "/api/v1/alerts", params=params) diff --git a/src/whiteboxai/utils.py b/whiteboxxai/utils.py similarity index 74% rename from src/whiteboxai/utils.py rename to whiteboxxai/utils.py index 5c41860..764c2db 100644 --- a/src/whiteboxai/utils.py +++ b/whiteboxxai/utils.py @@ -1,7 +1,7 @@ """ SDK Utilities -General utility functions for the WhiteBoxAI SDK. +General utility functions for the WhiteBoxXAI SDK. """ import json @@ -47,7 +47,10 @@ def validate_model_type(model_type: str) -> bool: model_type: Model type string Returns: - True if valid, False otherwise + True if valid + + Raises: + ValueError: If model_type is not a recognized type """ valid_types = { "classification", @@ -61,7 +64,11 @@ def validate_model_type(model_type: str) -> bool: "computer_vision", "other", } - return model_type.lower() in valid_types + if model_type.lower() not in valid_types: + raise ValueError( + f"Invalid model type: {model_type!r}. Must be one of {sorted(valid_types)}." + ) + return True def extract_feature_names(data: Union[np.ndarray, pd.DataFrame, Dict]) -> List[str]: @@ -227,7 +234,7 @@ def validate_features(features: Any) -> bool: raise ValueError(f"Unsupported features type: {type(features)}") -def format_features(features: Any) -> Dict[str, Any]: +def format_features(features: Any) -> Union[List[Any], Dict[str, Any]]: """ Format features for API transmission. @@ -235,14 +242,15 @@ def format_features(features: Any) -> Dict[str, Any]: features: Features to format Returns: - Formatted features dictionary + Formatted features as a list (list/tuple/numpy array input) or a + dictionary (dict input, returned as-is). """ if isinstance(features, dict): return features elif isinstance(features, (list, tuple, np.ndarray)): - return {"features": serialize_data(features)} + return serialize_data(features) else: - return {"features": features} + return serialize_data(features) def validate_prediction(prediction: Any) -> bool: @@ -263,17 +271,21 @@ def validate_prediction(prediction: Any) -> bool: return True -def serialize_numpy(arr: np.ndarray) -> List: +def serialize_numpy(arr: Any) -> List: """ - Serialize numpy array to list. + Serialize numpy array (or array-like) to list. Args: - arr: Numpy array + arr: Numpy array, or any array-like object (e.g. a plain list) Returns: List representation """ - return arr.tolist() + if isinstance(arr, np.ndarray): + return arr.tolist() + if hasattr(arr, "tolist"): + return arr.tolist() + return list(arr) def deserialize_numpy(data: List) -> np.ndarray: @@ -289,29 +301,63 @@ def deserialize_numpy(data: List) -> np.ndarray: return np.array(data) -def compute_metrics(y_true: Any, y_pred: Any) -> Dict[str, float]: +def compute_metrics( + y_true: Any, y_pred: Any, task: str = "classification" +) -> Dict[str, float]: """ Compute basic metrics for predictions. Args: y_true: True labels y_pred: Predicted labels + task: Type of task - "classification" or "regression" Returns: Dictionary of metrics + + Raises: + ValueError: If task is not recognized, or y_true/y_pred lengths differ """ + if task not in ("classification", "regression"): + raise ValueError( + f"Unsupported task: {task!r}. Must be 'classification' or 'regression'." + ) + y_true = np.array(y_true) y_pred = np.array(y_pred) - metrics = {} + if y_true.shape[0] != y_pred.shape[0]: + raise ValueError( + "y_true and y_pred must have the same length " + f"(got {y_true.shape[0]} and {y_pred.shape[0]})." + ) - # Basic accuracy for classification - if y_true.dtype == y_pred.dtype and np.issubdtype(y_true.dtype, np.integer): + metrics: Dict[str, float] = {} + + if task == "classification": metrics["accuracy"] = float(np.mean(y_true == y_pred)) + + precisions = [] + recalls = [] + for cls in np.unique(y_true): + predicted_positive = np.sum(y_pred == cls) + actual_positive = np.sum(y_true == cls) + true_positive = np.sum((y_pred == cls) & (y_true == cls)) + + precisions.append( + true_positive / predicted_positive if predicted_positive > 0 else 0.0 + ) + recalls.append( + true_positive / actual_positive if actual_positive > 0 else 0.0 + ) + + metrics["precision"] = float(np.mean(precisions)) + metrics["recall"] = float(np.mean(recalls)) else: - # MSE for regression - metrics["mse"] = float(np.mean((y_true - y_pred) ** 2)) + mse = float(np.mean((y_true - y_pred) ** 2)) + metrics["mse"] = mse metrics["mae"] = float(np.mean(np.abs(y_true - y_pred))) + metrics["rmse"] = float(np.sqrt(mse)) return metrics From 3f20fcd1404cccfdf86733f363005732393a4898 Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 16:05:38 -0700 Subject: [PATCH 2/6] feature: update sdk --- __init__.py | 1 + requirements-dev.txt | 26 + setup.py | 100 +++ src/whiteboxai/__init__.py | 24 + src/whiteboxai/__version__.py | 6 + src/whiteboxai/cache.py | 117 +++ src/whiteboxai/client.py | 354 ++++++++++ src/whiteboxai/config.py | 90 +++ src/whiteboxai/decorators.py | 323 +++++++++ src/whiteboxai/exceptions.py | 57 ++ src/whiteboxai/git_utils.py | 331 +++++++++ src/whiteboxai/integrations/__init__.py | 112 +++ src/whiteboxai/integrations/boosting.py | 667 ++++++++++++++++++ src/whiteboxai/integrations/crewai_monitor.py | 474 +++++++++++++ src/whiteboxai/integrations/langchain.py | 636 +++++++++++++++++ .../integrations/langchain_agents.py | 525 ++++++++++++++ src/whiteboxai/integrations/pytorch.py | 268 +++++++ src/whiteboxai/integrations/sklearn.py | 219 ++++++ src/whiteboxai/integrations/tensorflow.py | 475 +++++++++++++ src/whiteboxai/integrations/transformers.py | 531 ++++++++++++++ src/whiteboxai/models/__init__.py | 1 + src/whiteboxai/monitor.py | 284 ++++++++ src/whiteboxai/offline.py | 541 ++++++++++++++ src/whiteboxai/privacy.py | 219 ++++++ src/whiteboxai/resources.py | 460 ++++++++++++ src/whiteboxai/utils.py | 337 +++++++++ 26 files changed, 7178 insertions(+) create mode 100644 __init__.py create mode 100644 requirements-dev.txt create mode 100644 setup.py create mode 100644 src/whiteboxai/__init__.py create mode 100644 src/whiteboxai/__version__.py create mode 100644 src/whiteboxai/cache.py create mode 100644 src/whiteboxai/client.py create mode 100644 src/whiteboxai/config.py create mode 100644 src/whiteboxai/decorators.py create mode 100644 src/whiteboxai/exceptions.py create mode 100644 src/whiteboxai/git_utils.py create mode 100644 src/whiteboxai/integrations/__init__.py create mode 100644 src/whiteboxai/integrations/boosting.py create mode 100644 src/whiteboxai/integrations/crewai_monitor.py create mode 100644 src/whiteboxai/integrations/langchain.py create mode 100644 src/whiteboxai/integrations/langchain_agents.py create mode 100644 src/whiteboxai/integrations/pytorch.py create mode 100644 src/whiteboxai/integrations/sklearn.py create mode 100644 src/whiteboxai/integrations/tensorflow.py create mode 100644 src/whiteboxai/integrations/transformers.py create mode 100644 src/whiteboxai/models/__init__.py create mode 100644 src/whiteboxai/monitor.py create mode 100644 src/whiteboxai/offline.py create mode 100644 src/whiteboxai/privacy.py create mode 100644 src/whiteboxai/resources.py create mode 100644 src/whiteboxai/utils.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b2b2619 --- /dev/null +++ b/__init__.py @@ -0,0 +1 @@ +# WhiteBoxAI SDK diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..a87bedf --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,26 @@ +# Development dependencies for WhiteBoxAI SDK + +# Testing +pytest>=7.4.0 +pytest-asyncio>=0.21.0 +pytest-cov>=4.1.0 +pytest-mock>=3.12.0 + +# Code quality +black>=23.11.0 +isort>=5.12.0 +flake8>=6.1.0 +pylint>=3.0.0 +mypy>=1.7.0 + +# Security +bandit>=1.7.5 + +# Documentation +mkdocs>=1.5.0 +mkdocs-material>=9.0.0 +mkdocstrings[python]>=0.24.0 + +# Build tools +build>=1.0.0 +twine>=4.0.0 diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..23da3c8 --- /dev/null +++ b/setup.py @@ -0,0 +1,100 @@ +""" +WhiteBoxAI Python SDK - ML Monitoring and Observability +""" + +from pathlib import Path + +from setuptools import find_packages, setup + +# Read the contents of README file +this_directory = Path(__file__).parent +long_description = (this_directory / "README.md").read_text(encoding="utf-8") + + +# Read requirements +def read_requirements(filename): + """Read requirements from file.""" + filepath = this_directory / filename + if not filepath.exists(): + return [] + with open(filepath, encoding="utf-8") as f: + return [ + line.strip() + for line in f + if line.strip() and not line.startswith("#") and not line.startswith("-r") + ] + + +# Core requirements +install_requires = [ + "httpx>=0.24.0", + "pydantic>=2.0.0", + "python-dotenv>=1.0.0", + "numpy>=1.24.0", +] + +# Optional dependencies for integrations +extras_require = { + "sklearn": ["scikit-learn>=1.3.0"], + "pytorch": ["torch>=2.0.0"], + "tensorflow": ["tensorflow>=2.13.0"], + "transformers": ["transformers>=4.30.0"], + "langchain": ["langchain>=0.0.200"], + "boosting": ["xgboost>=1.7.0", "lightgbm>=4.0.0"], + "all": [ + "scikit-learn>=1.3.0", + "torch>=2.0.0", + "tensorflow>=2.13.0", + "transformers>=4.30.0", + "langchain>=0.0.200", + "xgboost>=1.7.0", + "lightgbm>=4.0.0", + ], + "dev": [ + "pytest>=7.4.0", + "pytest-asyncio>=0.21.0", + "pytest-cov>=4.1.0", + "black>=23.11.0", + "isort>=5.12.0", + "flake8>=6.1.0", + "pylint>=3.0.0", + "mypy>=1.7.0", + "bandit>=1.7.5", + ], +} + +setup( + name="whiteboxai-sdk", + version="0.1.0", + author="AgentaFlow", + author_email="support@agentaflow.com", + description="Official Python SDK for WhiteBoxAI ML monitoring and observability", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/AgentaFlow/whitebox-python-sdk", + project_urls={ + "Bug Reports": "https://github.com/AgentaFlow/whitebox-python-sdk/issues", + "Source": "https://github.com/AgentaFlow/whitebox-python-sdk", + "Documentation": "https://whitebox.agentaflow.com", + }, + packages=find_packages(where="src"), + package_dir={"": "src"}, + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Operating System :: OS Independent", + ], + python_requires=">=3.9", + install_requires=install_requires, + extras_require=extras_require, + include_package_data=True, + zip_safe=False, +) diff --git a/src/whiteboxai/__init__.py b/src/whiteboxai/__init__.py new file mode 100644 index 0000000..5fbf541 --- /dev/null +++ b/src/whiteboxai/__init__.py @@ -0,0 +1,24 @@ +""" +WhiteBoxAI Python SDK + +Official Python SDK for WhiteBoxAI - AI Observability & Explainability Platform. +""" + +__version__ = "0.2.0" +__author__ = "WhiteBoxAI Team" +__license__ = "MIT" + +from whiteboxai.client import WhiteBoxAI +from whiteboxai.decorators import monitor_model, monitor_prediction +from whiteboxai.git_utils import GitContext, detect_git_context, validate_git_context +from whiteboxai.monitor import ModelMonitor + +__all__ = [ + "WhiteBoxAI", + "ModelMonitor", + "monitor_model", + "monitor_prediction", + "GitContext", + "detect_git_context", + "validate_git_context", +] diff --git a/src/whiteboxai/__version__.py b/src/whiteboxai/__version__.py new file mode 100644 index 0000000..afa047f --- /dev/null +++ b/src/whiteboxai/__version__.py @@ -0,0 +1,6 @@ +"""Version information for WhiteBoxAI SDK.""" + +__version__ = "0.2.0" +__author__ = "AgentaFlow" +__email__ = "support@agentaflow.com" +__license__ = "MIT" diff --git a/src/whiteboxai/cache.py b/src/whiteboxai/cache.py new file mode 100644 index 0000000..8930c10 --- /dev/null +++ b/src/whiteboxai/cache.py @@ -0,0 +1,117 @@ +""" +Local Caching + +Simple TTL-based cache for API responses. +""" + +import hashlib +import json +import time +from collections import OrderedDict +from typing import Any, Optional + + +class TTLCache: + """ + Time-to-live cache with LRU eviction. + + Features: + - TTL-based expiration + - LRU eviction when max size reached + - Thread-safe operations + """ + + def __init__(self, max_size: int = 1000, ttl: int = 3600): + """ + Initialize cache. + + Args: + max_size: Maximum number of items to cache + ttl: Time-to-live in seconds + """ + self.max_size = max_size + self.ttl = ttl + self._cache: OrderedDict = OrderedDict() + self._timestamps: dict = {} + + def get(self, key: str) -> Optional[Any]: + """ + Get value from cache. + + Args: + key: Cache key + + Returns: + Cached value or None if not found/expired + """ + if key not in self._cache: + return None + + # Check expiration + timestamp = self._timestamps.get(key, 0) + if time.time() - timestamp > self.ttl: + # Expired, remove from cache + del self._cache[key] + del self._timestamps[key] + return None + + # Move to end (most recently used) + self._cache.move_to_end(key) + return self._cache[key] + + def set(self, key: str, value: Any) -> None: + """ + Set value in cache. + + Args: + key: Cache key + value: Value to cache + """ + # Remove if already exists + if key in self._cache: + del self._cache[key] + del self._timestamps[key] + + # Check size limit + if len(self._cache) >= self.max_size: + # Remove least recently used + oldest_key = next(iter(self._cache)) + del self._cache[oldest_key] + del self._timestamps[oldest_key] + + # Add to cache + self._cache[key] = value + self._timestamps[key] = time.time() + + def clear(self) -> None: + """Clear all cached items.""" + self._cache.clear() + self._timestamps.clear() + + def __contains__(self, key: str) -> bool: + """Check if key exists and is not expired.""" + return self.get(key) is not None + + def __len__(self) -> int: + """Get number of cached items.""" + return len(self._cache) + + +def cache_key(*args, **kwargs) -> str: + """ + Generate cache key from arguments. + + Args: + *args: Positional arguments + **kwargs: Keyword arguments + + Returns: + Cache key string + """ + # Create deterministic key from arguments + key_data = { + "args": args, + "kwargs": sorted(kwargs.items()), + } + key_str = json.dumps(key_data, sort_keys=True) + return hashlib.md5(key_str.encode()).hexdigest() diff --git a/src/whiteboxai/client.py b/src/whiteboxai/client.py new file mode 100644 index 0000000..705c027 --- /dev/null +++ b/src/whiteboxai/client.py @@ -0,0 +1,354 @@ +""" +WhiteBoxAI Client + +Main client class for interacting with the WhiteBoxAI API. +""" + +import asyncio +import logging +from typing import Any, Dict, Optional +from urllib.parse import urljoin + +import httpx +from tenacity import retry, stop_after_attempt, wait_exponential + +from whiteboxai.config import Config +from whiteboxai.exceptions import APIError, AuthenticationError, RateLimitError, ValidationError +from whiteboxai.resources import ( + AgentWorkflowsResource, + AlertsResource, + DriftResource, + ExplanationsResource, + ModelsResource, + PredictionsResource, +) + +logger = logging.getLogger(__name__) + + +class WhiteBoxAI: + """ + Main client for WhiteBoxAI SDK. + + Provides access to all WhiteBoxAI API resources including models, predictions, + explanations, drift detection, and alerting. + + Args: + api_key: WhiteBoxAI API key + base_url: Base URL for WhiteBoxAI API (default: https://api.whiteboxai.io) + timeout: Request timeout in seconds (default: 30) + max_retries: Maximum number of retry attempts (default: 3) + enable_offline: Enable offline mode (queues operations when API unavailable) + offline_dir: Directory for offline queue storage (default: ./whiteboxai_offline) + offline_max_queue_size: Maximum operations to queue (default: 10000, 0 = unlimited) + offline_auto_sync: Enable automatic syncing (default: True) + offline_sync_interval: Seconds between sync attempts (default: 60) + **kwargs: Additional configuration options + + Example: + >>> client = WhiteBoxAI(api_key="your_api_key") + >>> model = client.models.register(name="fraud_detection", model_type="classification") + >>> client.predictions.log(model_id=model.id, inputs=features, outputs=prediction) + + With offline mode: + >>> client = WhiteBoxAI( + ... api_key="your_api_key", + ... enable_offline=True, + ... offline_dir="./offline_queue" + ... ) + >>> # Operations queued when offline + >>> client.predictions.log(...) # Queued if API unavailable + """ + + def __init__( + self, + api_key: str, + base_url: str = "https://api.whiteboxai.io", + timeout: int = 30, + max_retries: int = 3, + enable_offline: bool = False, + offline_dir: str = "./whiteboxai_offline", + offline_max_queue_size: int = 10000, + offline_auto_sync: bool = True, + offline_sync_interval: int = 60, + **kwargs: Any, + ): + """Initialize WhiteBoxAI client.""" + self.config = Config( + api_key=api_key, + base_url=base_url, + timeout=timeout, + max_retries=max_retries, + **kwargs, + ) + + # Initialize HTTP clients + self._sync_client: Optional[httpx.Client] = None + self._async_client: Optional[httpx.AsyncClient] = None + + # Initialize offline mode + self._offline_manager = None + if enable_offline: + from whiteboxai.offline import OfflineManager + + self._offline_manager = OfflineManager( + offline_dir=offline_dir, + max_queue_size=offline_max_queue_size, + auto_sync=offline_auto_sync, + sync_interval=offline_sync_interval, + ) + self._offline_manager.set_client(self) + logger.info("Offline mode enabled") + + # Initialize resource managers + self.models = ModelsResource(self) + self.predictions = PredictionsResource(self) + self.explanations = ExplanationsResource(self) + self.drift = DriftResource(self) + self.alerts = AlertsResource(self) + self.agent_workflows = AgentWorkflowsResource(self) + + @property + def sync_client(self) -> httpx.Client: + """Get or create synchronous HTTP client.""" + if self._sync_client is None: + self._sync_client = httpx.Client( + base_url=self.config.base_url, + headers=self._get_headers(), + timeout=self.config.timeout, + ) + return self._sync_client + + @property + def async_client(self) -> httpx.AsyncClient: + """Get or create asynchronous HTTP client.""" + if self._async_client is None: + self._async_client = httpx.AsyncClient( + base_url=self.config.base_url, + headers=self._get_headers(), + timeout=self.config.timeout, + ) + return self._async_client + + def _get_headers(self) -> Dict[str, str]: + """Get request headers with authentication.""" + return { + "Authorization": f"Bearer {self.config.api_key}", + "Content-Type": "application/json", + "User-Agent": f"whiteboxai-python-sdk/{self.config.sdk_version}", + } + + @retry( + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=10), + reraise=True, + ) + def request( + self, + method: str, + endpoint: str, + data: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Make synchronous HTTP request with retry logic. + + Args: + method: HTTP method (GET, POST, PUT, DELETE, etc.) + endpoint: API endpoint path + data: Request body data + params: Query parameters + **kwargs: Additional request options + + Returns: + Response data as dictionary + + Raises: + APIError: For general API errors + AuthenticationError: For authentication failures + RateLimitError: For rate limit errors + ValidationError: For validation errors + """ + url = urljoin(self.config.base_url, endpoint) + + try: + response = self.sync_client.request( + method=method, url=url, json=data, params=params, **kwargs + ) + return self._handle_response(response) + except httpx.HTTPError as e: + raise APIError(f"Request failed: {str(e)}") + + @retry( + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=2, max=10), + reraise=True, + ) + async def arequest( + self, + method: str, + endpoint: str, + data: Optional[Dict[str, Any]] = None, + params: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Make asynchronous HTTP request with retry logic. + + Args: + method: HTTP method (GET, POST, PUT, DELETE, etc.) + endpoint: API endpoint path + data: Request body data + params: Query parameters + **kwargs: Additional request options + + Returns: + Response data as dictionary + + Raises: + APIError: For general API errors + AuthenticationError: For authentication failures + RateLimitError: For rate limit errors + ValidationError: For validation errors + """ + url = urljoin(self.config.base_url, endpoint) + + try: + response = await self.async_client.request( + method=method, url=url, json=data, params=params, **kwargs + ) + return self._handle_response(response) + except httpx.HTTPError as e: + raise APIError(f"Request failed: {str(e)}") + + def _handle_response(self, response: httpx.Response) -> Dict[str, Any]: + """ + Handle HTTP response and raise appropriate exceptions. + + Args: + response: HTTP response + + Returns: + Response data as dictionary + + Raises: + AuthenticationError: For 401 status codes + RateLimitError: For 429 status codes + ValidationError: For 422 status codes + APIError: For other error status codes + """ + if response.status_code == 401: + raise AuthenticationError("Invalid API key or authentication failed") + elif response.status_code == 429: + raise RateLimitError("Rate limit exceeded. Please retry later.") + elif response.status_code == 422: + try: + error_data = response.json() + raise ValidationError(f"Validation error: {error_data}") + except Exception: + raise ValidationError("Validation error occurred") + elif response.status_code >= 400: + try: + error_data = response.json() + message = error_data.get("detail", response.text) + except Exception: + message = response.text + raise APIError(f"API error ({response.status_code}): {message}") + + try: + return response.json() + except Exception: + return {"status": "success"} + + def is_offline_enabled(self) -> bool: + """Check if offline mode is enabled.""" + return self._offline_manager is not None + + def sync_offline_queue(self, batch_size: int = 100) -> Dict[str, int]: + """ + Manually sync offline queue with API. + + Args: + batch_size: Number of operations to sync per batch + + Returns: + Dictionary with sync statistics + + Raises: + ValueError: If offline mode is not enabled + """ + if not self._offline_manager: + raise ValueError("Offline mode is not enabled") + + return self._offline_manager.sync(batch_size=batch_size) + + def get_offline_status(self) -> Dict[str, Any]: + """ + Get offline mode status. + + Returns: + Status dictionary with queue stats + + Raises: + ValueError: If offline mode is not enabled + """ + if not self._offline_manager: + raise ValueError("Offline mode is not enabled") + + return self._offline_manager.get_status() + + def cleanup_offline_queue(self, older_than_days: int = 7): + """ + Clean up old completed operations from offline queue. + + Args: + older_than_days: Remove operations older than this many days + + Raises: + ValueError: If offline mode is not enabled + """ + if not self._offline_manager: + raise ValueError("Offline mode is not enabled") + + self._offline_manager.cleanup(older_than_days=older_than_days) + + def close(self) -> None: + """Close HTTP clients and cleanup resources.""" + # Stop offline sync if enabled + if self._offline_manager: + self._offline_manager.stop_auto_sync() + logger.info("Offline mode stopped") + + if self._sync_client: + self._sync_client.close() + if self._async_client: + asyncio.run(self._async_client.aclose()) + + async def aclose(self) -> None: + """Asynchronously close HTTP clients and cleanup resources.""" + # Stop offline sync if enabled + if self._offline_manager: + self._offline_manager.stop_auto_sync() + logger.info("Offline mode stopped") + + if self._async_client: + await self._async_client.aclose() + if self._sync_client: + self._sync_client.close() + + def __enter__(self) -> "WhiteBoxAI": + """Context manager entry.""" + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Context manager exit.""" + self.close() + + async def __aenter__(self) -> "WhiteBoxAI": + """Async context manager entry.""" + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Async context manager exit.""" + await self.aclose() diff --git a/src/whiteboxai/config.py b/src/whiteboxai/config.py new file mode 100644 index 0000000..b6c7b1d --- /dev/null +++ b/src/whiteboxai/config.py @@ -0,0 +1,90 @@ +""" +SDK Configuration +""" + +import os +from typing import Any, Optional + + +class Config: + """ + SDK configuration class. + + Manages configuration settings for the WhiteBoxAI SDK including API keys, + URLs, timeouts, and feature flags. + + Args: + api_key: WhiteBoxAI API key + base_url: Base URL for WhiteBoxAI API + timeout: Request timeout in seconds + max_retries: Maximum number of retry attempts + **kwargs: Additional configuration options + """ + + def __init__( + self, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + timeout: int = 30, + max_retries: int = 3, + **kwargs: Any, + ): + """Initialize configuration.""" + # API key (from parameter or environment) + self.api_key = api_key or os.getenv("EXPLAINAI_API_KEY") + if not self.api_key: + raise ValueError( + "API key is required. Provide via 'api_key' parameter or " + "EXPLAINAI_API_KEY environment variable." + ) + + # Base URL + self.base_url = base_url or os.getenv("EXPLAINAI_BASE_URL") or "https://api.whiteboxai.io" + + # Request settings + self.timeout = timeout + self.max_retries = max_retries + + # Feature flags + self.enable_caching = kwargs.get("enable_caching", True) + self.enable_offline_mode = kwargs.get("enable_offline_mode", False) + self.enable_privacy_filters = kwargs.get("enable_privacy_filters", True) + self.enable_sampling = kwargs.get("enable_sampling", True) + + # Sampling settings + self.sampling_rate = kwargs.get("sampling_rate", 1.0) # 100% by default + + # Cache settings + self.cache_ttl = kwargs.get("cache_ttl", 3600) # 1 hour default + self.cache_max_size = kwargs.get("cache_max_size", 1000) + + # Privacy settings + self.detect_pii = kwargs.get("detect_pii", True) + self.mask_pii = kwargs.get("mask_pii", True) + + # Performance settings + self.batch_size = kwargs.get("batch_size", 100) + self.async_enabled = kwargs.get("async_enabled", True) + + # SDK metadata + self.sdk_version = "0.2.0" + + def to_dict(self) -> dict: + """Convert configuration to dictionary.""" + return { + "base_url": self.base_url, + "timeout": self.timeout, + "max_retries": self.max_retries, + "enable_caching": self.enable_caching, + "enable_offline_mode": self.enable_offline_mode, + "enable_privacy_filters": self.enable_privacy_filters, + "enable_sampling": self.enable_sampling, + "sampling_rate": self.sampling_rate, + "cache_ttl": self.cache_ttl, + "cache_max_size": self.cache_max_size, + "detect_pii": self.detect_pii, + "mask_pii": self.mask_pii, + "batch_size": self.batch_size, + "async_enabled": self.async_enabled, + "sdk_version": self.sdk_version, + } diff --git a/src/whiteboxai/decorators.py b/src/whiteboxai/decorators.py new file mode 100644 index 0000000..21f2a13 --- /dev/null +++ b/src/whiteboxai/decorators.py @@ -0,0 +1,323 @@ +""" +Monitoring Decorators + +Decorators for automatic model and prediction monitoring. +""" + +import functools +import inspect +import time +from typing import Any, Callable, Dict, Optional + +from whiteboxai.monitor import ModelMonitor + + +def monitor_model( + monitor: ModelMonitor, + input_keys: Optional[list] = None, + output_key: Optional[str] = None, + explain: bool = False, +): + """ + Decorator to automatically monitor model predictions. + + Args: + monitor: ModelMonitor instance + input_keys: Keys to extract inputs from function args/kwargs + output_key: Key to extract output from function return value + explain: Generate explanations for predictions + + Example: + ```python + from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_model + + client = WhiteBoxAI(api_key="your-api-key") + monitor = ModelMonitor(client, model_id=123) + + @monitor_model(monitor, input_keys=["features"], explain=True) + def predict(features): + # Your prediction logic + return model.predict(features) + + # Predictions are automatically logged + result = predict(features=[1.0, 2.0, 3.0]) + ``` + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs): + start_time = time.time() + + # Call original function + result = func(*args, **kwargs) + + # Extract inputs + inputs = _extract_inputs(func, args, kwargs, input_keys) + + # Extract output + output = _extract_output(result, output_key) + + # Calculate inference time + inference_time = time.time() - start_time + + # Log prediction + metadata = {"inference_time_ms": inference_time * 1000} + monitor.log_prediction( + inputs=inputs, + output=output, + explain=explain, + metadata=metadata, + ) + + return result + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + start_time = time.time() + + # Call original function + result = await func(*args, **kwargs) + + # Extract inputs + inputs = _extract_inputs(func, args, kwargs, input_keys) + + # Extract output + output = _extract_output(result, output_key) + + # Calculate inference time + inference_time = time.time() - start_time + + # Log prediction + metadata = {"inference_time_ms": inference_time * 1000} + await monitor.alog_prediction( + inputs=inputs, + output=output, + explain=explain, + metadata=metadata, + ) + + return result + + # Return appropriate wrapper based on function type + if inspect.iscoroutinefunction(func): + return async_wrapper + return wrapper + + return decorator + + +def monitor_prediction( + monitor: ModelMonitor, + input_extractor: Optional[Callable] = None, + output_extractor: Optional[Callable] = None, + explain: bool = False, + metadata_extractor: Optional[Callable] = None, +): + """ + Decorator to monitor individual predictions with custom extractors. + + Args: + monitor: ModelMonitor instance + input_extractor: Function to extract inputs from args/kwargs + output_extractor: Function to extract output from result + explain: Generate explanations for predictions + metadata_extractor: Function to extract additional metadata + + Example: + ```python + from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_prediction + + client = WhiteBoxAI(api_key="your-api-key") + monitor = ModelMonitor(client, model_id=123) + + def extract_inputs(args, kwargs): + return {"features": kwargs.get("data")} + + def extract_output(result): + return result["prediction"] + + @monitor_prediction( + monitor, + input_extractor=extract_inputs, + output_extractor=extract_output, + explain=True + ) + def predict(data): + # Your prediction logic + return {"prediction": 0.85, "confidence": 0.92} + + result = predict(data=[1.0, 2.0, 3.0]) + ``` + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs): + start_time = time.time() + + # Call original function + result = func(*args, **kwargs) + + # Extract inputs + if input_extractor: + inputs = input_extractor(args, kwargs) + else: + inputs = _default_input_extractor(func, args, kwargs) + + # Extract output + if output_extractor: + output = output_extractor(result) + else: + output = result + + # Extract metadata + metadata = {"inference_time_ms": (time.time() - start_time) * 1000} + if metadata_extractor: + custom_metadata = metadata_extractor(args, kwargs, result) + metadata.update(custom_metadata) + + # Log prediction + monitor.log_prediction( + inputs=inputs, + output=output, + explain=explain, + metadata=metadata, + ) + + return result + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + start_time = time.time() + + # Call original function + result = await func(*args, **kwargs) + + # Extract inputs + if input_extractor: + inputs = input_extractor(args, kwargs) + else: + inputs = _default_input_extractor(func, args, kwargs) + + # Extract output + if output_extractor: + output = output_extractor(result) + else: + output = result + + # Extract metadata + metadata = {"inference_time_ms": (time.time() - start_time) * 1000} + if metadata_extractor: + custom_metadata = metadata_extractor(args, kwargs, result) + metadata.update(custom_metadata) + + # Log prediction + await monitor.alog_prediction( + inputs=inputs, + output=output, + explain=explain, + metadata=metadata, + ) + + return result + + # Return appropriate wrapper based on function type + if inspect.iscoroutinefunction(func): + return async_wrapper + return wrapper + + return decorator + + +def _extract_inputs( + func: Callable, + args: tuple, + kwargs: dict, + input_keys: Optional[list] = None, +) -> Any: + """Extract inputs from function arguments.""" + if input_keys: + # Extract specified keys + sig = inspect.signature(func) + bound_args = sig.bind(*args, **kwargs) + bound_args.apply_defaults() + + inputs = {} + for key in input_keys: + if key in bound_args.arguments: + inputs[key] = bound_args.arguments[key] + + return inputs + else: + # Use all arguments + return _default_input_extractor(func, args, kwargs) + + +def _default_input_extractor(func: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]: + """Default input extractor.""" + sig = inspect.signature(func) + bound_args = sig.bind(*args, **kwargs) + bound_args.apply_defaults() + return dict(bound_args.arguments) + + +def _extract_output(result: Any, output_key: Optional[str] = None) -> Any: + """Extract output from function result.""" + if output_key and isinstance(result, dict): + return result.get(output_key) + return result + + +def monitor_performance(threshold_ms: Optional[float] = None): + """ + Decorator to monitor function performance. + + Args: + threshold_ms: Optional performance threshold in milliseconds + + Example: + ```python + from whiteboxai.decorators import monitor_performance + + @monitor_performance(threshold_ms=1000) + def slow_function(): + # Your code here + pass + ``` + """ + + def decorator(func: Callable) -> Callable: + @functools.wraps(func) + def wrapper(*args, **kwargs): + start_time = time.time() + result = func(*args, **kwargs) + elapsed_ms = (time.time() - start_time) * 1000 + + if threshold_ms and elapsed_ms > threshold_ms: + print( + f"Warning: {func.__name__} took {elapsed_ms:.2f}ms " + f"(threshold: {threshold_ms}ms)" + ) + + return result + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs): + start_time = time.time() + result = await func(*args, **kwargs) + elapsed_ms = (time.time() - start_time) * 1000 + + if threshold_ms and elapsed_ms > threshold_ms: + print( + f"Warning: {func.__name__} took {elapsed_ms:.2f}ms " + f"(threshold: {threshold_ms}ms)" + ) + + return result + + if inspect.iscoroutinefunction(func): + return async_wrapper + return wrapper + + return decorator diff --git a/src/whiteboxai/exceptions.py b/src/whiteboxai/exceptions.py new file mode 100644 index 0000000..7932c7f --- /dev/null +++ b/src/whiteboxai/exceptions.py @@ -0,0 +1,57 @@ +""" +SDK Exceptions +""" + + +class WhiteBoxAIError(Exception): + """Base exception for WhiteBoxAI SDK.""" + + pass + + +class APIError(WhiteBoxAIError): + """API request error.""" + + pass + + +class AuthenticationError(WhiteBoxAIError): + """Authentication error.""" + + pass + + +class RateLimitError(WhiteBoxAIError): + """Rate limit exceeded error.""" + + pass + + +class ValidationError(WhiteBoxAIError): + """Validation error.""" + + pass + + +class NotFoundError(WhiteBoxAIError): + """Resource not found error.""" + + pass + + +class ConfigurationError(WhiteBoxAIError): + """Configuration error.""" + + pass + + +class IntegrationError(WhiteBoxAIError): + """Framework integration error.""" + + pass + + +class CacheError(WhiteBoxAIError): + """Cache operation error.""" + + pass diff --git a/src/whiteboxai/git_utils.py b/src/whiteboxai/git_utils.py new file mode 100644 index 0000000..925817a --- /dev/null +++ b/src/whiteboxai/git_utils.py @@ -0,0 +1,331 @@ +""" +Git Auto-Detection Utilities + +Utilities for automatically detecting Git context (repository, commit, branch) +for model registration. +""" + +import logging +import os +import subprocess +from typing import Dict, Optional + +logger = logging.getLogger(__name__) + + +class GitContext: + """Git repository context information.""" + + def __init__( + self, + repository_url: Optional[str] = None, + commit_sha: Optional[str] = None, + commit_message: Optional[str] = None, + commit_author: Optional[str] = None, + branch: Optional[str] = None, + tag: Optional[str] = None, + is_dirty: bool = False, + ): + """Initialize Git context.""" + self.repository_url = repository_url + self.commit_sha = commit_sha + self.commit_message = commit_message + self.commit_author = commit_author + self.branch = branch + self.tag = tag + self.is_dirty = is_dirty + + def to_dict(self) -> Dict[str, Optional[str]]: + """Convert to dictionary for API submission.""" + return { + "github_repository_url": self.repository_url, + "github_commit_hash": self.commit_sha, + "github_commit_message": self.commit_message, + "github_commit_author": self.commit_author, + "github_branch": self.branch, + "github_tag": self.tag, + } + + def __repr__(self) -> str: + """String representation.""" + parts = [] + if self.repository_url: + parts.append(f"repo={self.repository_url}") + if self.commit_sha: + parts.append(f"commit={self.commit_sha[:7]}") + if self.branch: + parts.append(f"branch={self.branch}") + if self.tag: + parts.append(f"tag={self.tag}") + return f"GitContext({', '.join(parts)})" + + +def detect_git_context(path: Optional[str] = None) -> Optional[GitContext]: + """ + Auto-detect Git context from the current directory or specified path. + + Args: + path: Path to check for Git repository (defaults to current directory) + + Returns: + GitContext object if Git repository found, None otherwise + + Example: + >>> context = detect_git_context() + >>> if context: + ... print(f"Repository: {context.repository_url}") + ... print(f"Commit: {context.commit_sha}") + ... print(f"Branch: {context.branch}") + """ + try: + import git + except ImportError: + logger.warning("GitPython not installed. Install with: pip install gitpython") + return _detect_git_context_subprocess(path) + + try: + # Find repository + search_path = path or os.getcwd() + repo = git.Repo(search_path, search_parent_directories=True) + + # Get repository URL + repository_url = None + try: + remote = repo.remote("origin") + repository_url = remote.url + # Convert SSH URLs to HTTPS + if repository_url.startswith("git@github.com:"): + repository_url = repository_url.replace("git@github.com:", "https://github.com/") + if repository_url.endswith(".git"): + repository_url = repository_url[:-4] + except Exception as e: + logger.debug(f"Could not get remote URL: {e}") + + # Get current commit + commit_sha = None + commit_message = None + commit_author = None + + try: + head_commit = repo.head.commit + commit_sha = head_commit.hexsha + commit_message = head_commit.message.strip() + commit_author = str(head_commit.author) + except Exception as e: + logger.debug(f"Could not get commit info: {e}") + + # Get current branch + branch = None + try: + if not repo.head.is_detached: + branch = repo.active_branch.name + except Exception as e: + logger.debug(f"Could not get branch: {e}") + + # Get current tag (if on a tag) + tag = None + try: + tags = [tag for tag in repo.tags if tag.commit == repo.head.commit] + if tags: + tag = tags[0].name + except Exception as e: + logger.debug(f"Could not get tag: {e}") + + # Check if working directory is dirty + is_dirty = repo.is_dirty(untracked_files=True) + if is_dirty: + logger.warning("Working directory has uncommitted changes") + + context = GitContext( + repository_url=repository_url, + commit_sha=commit_sha, + commit_message=commit_message, + commit_author=commit_author, + branch=branch, + tag=tag, + is_dirty=is_dirty, + ) + + logger.info(f"Detected Git context: {context}") + return context + + except git.InvalidGitRepositoryError: + logger.debug(f"No Git repository found at {search_path}") + return None + except Exception as e: + logger.error(f"Error detecting Git context: {e}") + return None + + +def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitContext]: + """ + Fallback Git detection using subprocess (when GitPython not available). + + Args: + path: Path to check for Git repository + + Returns: + GitContext object or None + """ + try: + cwd = path or os.getcwd() + + # Check if Git is available + subprocess.run( + ["git", "--version"], + capture_output=True, + check=True, + cwd=cwd, + ) + + # Get repository URL + repository_url = None + try: + result = subprocess.run( + ["git", "remote", "get-url", "origin"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + repository_url = result.stdout.strip() + + # Convert SSH to HTTPS + if repository_url.startswith("git@github.com:"): + repository_url = repository_url.replace("git@github.com:", "https://github.com/") + if repository_url.endswith(".git"): + repository_url = repository_url[:-4] + except subprocess.CalledProcessError: + pass + + # Get commit SHA + commit_sha = None + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + commit_sha = result.stdout.strip() + except subprocess.CalledProcessError: + pass + + # Get commit message + commit_message = None + try: + result = subprocess.run( + ["git", "log", "-1", "--pretty=%B"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + commit_message = result.stdout.strip() + except subprocess.CalledProcessError: + pass + + # Get commit author + commit_author = None + try: + result = subprocess.run( + ["git", "log", "-1", "--pretty=%an <%ae>"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + commit_author = result.stdout.strip() + except subprocess.CalledProcessError: + pass + + # Get branch + branch = None + try: + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + branch_output = result.stdout.strip() + if branch_output != "HEAD": # Not in detached HEAD + branch = branch_output + except subprocess.CalledProcessError: + pass + + # Get tag + tag = None + try: + result = subprocess.run( + ["git", "describe", "--exact-match", "--tags", "HEAD"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + tag = result.stdout.strip() + except subprocess.CalledProcessError: + pass + + # Check if dirty + is_dirty = False + try: + result = subprocess.run( + ["git", "status", "--porcelain"], + capture_output=True, + text=True, + check=True, + cwd=cwd, + ) + is_dirty = bool(result.stdout.strip()) + except subprocess.CalledProcessError: + pass + + context = GitContext( + repository_url=repository_url, + commit_sha=commit_sha, + commit_message=commit_message, + commit_author=commit_author, + branch=branch, + tag=tag, + is_dirty=is_dirty, + ) + + logger.info(f"Detected Git context (subprocess): {context}") + return context + + except (subprocess.CalledProcessError, FileNotFoundError): + logger.debug("Git not available or not a Git repository") + return None + except Exception as e: + logger.error(f"Error detecting Git context with subprocess: {e}") + return None + + +def validate_git_context(context: GitContext, require_clean: bool = False) -> bool: + """ + Validate Git context before model registration. + + Args: + context: GitContext object + require_clean: Whether to require a clean working directory + + Returns: + True if valid, False otherwise + """ + if not context: + logger.warning("No Git context provided") + return False + + if not context.commit_sha: + logger.error("Git context missing commit SHA") + return False + + if require_clean and context.is_dirty: + logger.error("Working directory has uncommitted changes (require_clean=True)") + return False + + return True diff --git a/src/whiteboxai/integrations/__init__.py b/src/whiteboxai/integrations/__init__.py new file mode 100644 index 0000000..c69a821 --- /dev/null +++ b/src/whiteboxai/integrations/__init__.py @@ -0,0 +1,112 @@ +"""Framework integrations for WhiteBoxAI SDK.""" + +# Scikit-learn integration +try: + from .sklearn import SklearnMonitor, SklearnWrapper + + __all__ = ["SklearnMonitor", "SklearnWrapper"] +except ImportError: + pass + +# PyTorch integration +try: + from .pytorch import TorchMonitor, TorchWrapper + + if "__all__" in dir(): + __all__.extend(["TorchMonitor", "TorchWrapper"]) + else: + __all__ = ["TorchMonitor", "TorchWrapper"] +except ImportError: + pass + +# TensorFlow/Keras integration +try: + from .tensorflow import KerasMonitor, WhiteBoxAICallback, wrap_keras_model + + if "__all__" in dir(): + __all__.extend(["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"]) + else: + __all__ = ["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"] +except ImportError: + pass + +# Hugging Face Transformers integration +try: + from .transformers import ( + TransformersMonitor, + TransformersPipelineWrapper, + wrap_transformers_pipeline, + ) + + if "__all__" in dir(): + __all__.extend( + ["TransformersMonitor", "TransformersPipelineWrapper", "wrap_transformers_pipeline"] + ) + else: + __all__ = [ + "TransformersMonitor", + "TransformersPipelineWrapper", + "wrap_transformers_pipeline", + ] +except ImportError: + pass + +# LangChain integration +try: + from .langchain import LangChainMonitor, WhiteBoxAICallbackHandler, wrap_langchain_chain + + if "__all__" in dir(): + __all__.extend(["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"]) + else: + __all__ = ["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"] +except ImportError: + pass + +# XGBoost/LightGBM integration +try: + from .boosting import LightGBMMonitor, XGBoostMonitor, wrap_lightgbm_model, wrap_xgboost_model + + if "__all__" in dir(): + __all__.extend( + ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] + ) + else: + __all__ = ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] +except ImportError: + pass + +# CrewAI integration +try: + from .crewai_monitor import CrewAIMonitor, monitor_crew + + if "__all__" in dir(): + __all__.extend(["CrewAIMonitor", "monitor_crew"]) + else: + __all__ = ["CrewAIMonitor", "monitor_crew"] +except ImportError: + pass + +# LangChain Multi-Agent integration +try: + from .langchain_agents import ( + LangGraphMultiAgentMonitor, + MultiAgentCallbackHandler, + monitor_langchain_agent, + ) + + if "__all__" in dir(): + __all__.extend( + ["MultiAgentCallbackHandler", "LangGraphMultiAgentMonitor", "monitor_langchain_agent"] + ) + else: + __all__ = [ + "MultiAgentCallbackHandler", + "LangGraphMultiAgentMonitor", + "monitor_langchain_agent", + ] +except ImportError: + pass + +# Ensure __all__ exists even if all imports fail +if "__all__" not in dir(): + __all__ = [] diff --git a/src/whiteboxai/integrations/boosting.py b/src/whiteboxai/integrations/boosting.py new file mode 100644 index 0000000..210155b --- /dev/null +++ b/src/whiteboxai/integrations/boosting.py @@ -0,0 +1,667 @@ +""" +XGBoost and LightGBM integration for WhiteBoxAI. + +This module provides monitoring capabilities for gradient boosting models +from XGBoost and LightGBM frameworks. It supports both frameworks through +a unified interface with automatic feature importance tracking, prediction +logging, and performance monitoring. + +Example: + Basic XGBoost monitoring: + + >>> import xgboost as xgb + >>> from whiteboxai import WhiteBoxAI + >>> from whiteboxai.integrations.boosting import XGBoostMonitor + >>> + >>> client = WhiteBoxAI(api_key="your-api-key") + >>> monitor = XGBoostMonitor(client=client, model_name="fraud_detector") + >>> + >>> model = xgb.XGBClassifier() + >>> model.fit(X_train, y_train) + >>> + >>> monitor.register_from_model(model, X_train, y_train) + >>> predictions = monitor.predict(model, X_test, y_test) + + LightGBM monitoring: + + >>> import lightgbm as lgb + >>> from whiteboxai.integrations.boosting import LightGBMMonitor + >>> + >>> monitor = LightGBMMonitor(client=client, model_name="churn_predictor") + >>> + >>> model = lgb.LGBMClassifier() + >>> model.fit(X_train, y_train) + >>> + >>> monitor.register_from_model(model, X_train, y_train) + >>> predictions = monitor.predict(model, X_test, y_test) +""" + +import warnings +from typing import Any, Dict, List, Optional + +import numpy as np + +from ..core import ModelMonitor + +# Optional imports - graceful degradation +try: + import xgboost as xgb + + XGBOOST_AVAILABLE = True +except ImportError: + XGBOOST_AVAILABLE = False + xgb = None + +try: + import lightgbm as lgb + + LIGHTGBM_AVAILABLE = True +except ImportError: + LIGHTGBM_AVAILABLE = False + lgb = None + + +class XGBoostMonitor(ModelMonitor): + """ + Monitor for XGBoost models. + + Provides comprehensive monitoring for XGBoost models including prediction + logging, feature importance tracking, and automatic model registration. + + Attributes: + client: WhiteBoxAI client instance + model_name: Name of the model being monitored + track_feature_importance: Whether to track feature importance with predictions + importance_type: Type of importance to track ('weight', 'gain', 'cover', 'total_gain', 'total_cover') + + Example: + >>> monitor = XGBoostMonitor( + ... client=client, + ... model_name="xgb_classifier", + ... track_feature_importance=True, + ... importance_type="gain" + ... ) + >>> monitor.register_from_model(model, X_train, y_train) + >>> predictions = monitor.predict(model, X_test, y_test) + """ + + def __init__( + self, + client: Any, + model_name: str = "xgboost_model", + track_feature_importance: bool = True, + importance_type: str = "gain", + **kwargs, + ): + """ + Initialize XGBoost monitor. + + Args: + client: WhiteBoxAI client instance + model_name: Name of the model + track_feature_importance: Whether to track feature importance + importance_type: Type of importance ('weight', 'gain', 'cover', 'total_gain', 'total_cover') + **kwargs: Additional arguments passed to ModelMonitor + """ + if not XGBOOST_AVAILABLE: + raise ImportError("XGBoost is not installed. Install with: pip install xgboost") + + super().__init__(client=client, model_name=model_name, **kwargs) + self.track_feature_importance = track_feature_importance + self.importance_type = importance_type + self._feature_names: Optional[List[str]] = None + + def register_from_model( + self, + model: Any, + X_train: Optional[np.ndarray] = None, + y_train: Optional[np.ndarray] = None, + model_type: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Register XGBoost model with WhiteBoxAI. + + Automatically extracts model metadata including feature names, + number of features, number of trees, and feature importance. + + Args: + model: XGBoost model (Booster, XGBClassifier, XGBRegressor, etc.) + X_train: Training features for baseline + y_train: Training labels for baseline + model_type: Model type override (auto-detected if None) + metadata: Additional metadata + + Returns: + Model ID from registration + """ + # Extract model information + model_metadata = metadata or {} + + # Get feature names + if hasattr(model, "feature_names_in_"): + self._feature_names = list(model.feature_names_in_) + model_metadata["feature_names"] = self._feature_names + elif hasattr(model, "feature_names"): + self._feature_names = model.feature_names + model_metadata["feature_names"] = self._feature_names + elif X_train is not None and hasattr(X_train, "columns"): + self._feature_names = list(X_train.columns) + model_metadata["feature_names"] = self._feature_names + + # Get number of features + if hasattr(model, "n_features_in_"): + model_metadata["num_features"] = model.n_features_in_ + elif self._feature_names: + model_metadata["num_features"] = len(self._feature_names) + + # Get number of trees/boosting rounds + if hasattr(model, "n_estimators"): + model_metadata["num_trees"] = model.n_estimators + elif hasattr(model, "best_iteration"): + model_metadata["num_trees"] = model.best_iteration + + # Get XGBoost parameters + if hasattr(model, "get_params"): + params = model.get_params() + model_metadata["xgboost_params"] = { + k: str(v) for k, v in params.items() if v is not None + } + + # Get feature importance + if self.track_feature_importance: + try: + importance = self._get_feature_importance(model) + if importance: + model_metadata["feature_importance"] = importance + except Exception as e: + warnings.warn(f"Failed to extract feature importance: {e}") + + # Detect model type + if model_type is None: + if hasattr(model, "_estimator_type"): + model_type = model._estimator_type + elif isinstance(model, xgb.XGBClassifier): + model_type = "classification" + elif isinstance(model, xgb.XGBRegressor): + model_type = "regression" + elif isinstance(model, xgb.XGBRanker): + model_type = "ranking" + else: + model_type = "classification" # Default + + model_metadata["framework"] = "xgboost" + model_metadata["xgboost_version"] = xgb.__version__ + + # Register model + model_id = self.register_model( + name=self.model_name, model_type=model_type, metadata=model_metadata + ) + + # Set baseline if training data provided + if X_train is not None and y_train is not None: + self.set_baseline(X_train, y_train) + + return model_id + + def predict( + self, + model: Any, + X: np.ndarray, + y_true: Optional[np.ndarray] = None, + log_predictions: bool = True, + metadata: Optional[Dict[str, Any]] = None, + ) -> np.ndarray: + """ + Make predictions and log to WhiteBoxAI. + + Args: + model: XGBoost model + X: Input features + y_true: True labels (optional) + log_predictions: Whether to log predictions + metadata: Additional metadata + + Returns: + Model predictions + """ + # Make predictions + predictions = model.predict(X) + + # Get prediction probabilities if classification + probabilities = None + if hasattr(model, "predict_proba"): + try: + probabilities = model.predict_proba(X) + except Exception: + pass + + # Log predictions + if log_predictions: + pred_metadata = metadata or {} + + # Add feature importance for this batch if enabled + if self.track_feature_importance: + try: + importance = self._get_feature_importance(model) + if importance: + pred_metadata["feature_importance"] = importance + except Exception as e: + warnings.warn(f"Failed to extract feature importance: {e}") + + self.log_predictions( + inputs=X, + predictions=predictions, + actuals=y_true, + probabilities=probabilities, + metadata=pred_metadata, + ) + + return predictions + + def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: + """ + Extract feature importance from XGBoost model. + + Args: + model: XGBoost model + + Returns: + Dictionary mapping feature names to importance scores + """ + try: + # Try sklearn-style feature_importances_ + if hasattr(model, "feature_importances_"): + importances = model.feature_importances_ + if self._feature_names and len(importances) == len(self._feature_names): + return dict(zip(self._feature_names, importances.tolist())) + else: + return {f"f{i}": float(v) for i, v in enumerate(importances)} + + # Try get_score method (native XGBoost) + if hasattr(model, "get_score"): + importance_dict = model.get_score(importance_type=self.importance_type) + return {k: float(v) for k, v in importance_dict.items()} + + # Try get_booster for sklearn API + if hasattr(model, "get_booster"): + booster = model.get_booster() + importance_dict = booster.get_score(importance_type=self.importance_type) + return {k: float(v) for k, v in importance_dict.items()} + + return None + except Exception: + return None + + +class LightGBMMonitor(ModelMonitor): + """ + Monitor for LightGBM models. + + Provides comprehensive monitoring for LightGBM models including prediction + logging, feature importance tracking, and automatic model registration. + + Attributes: + client: WhiteBoxAI client instance + model_name: Name of the model being monitored + track_feature_importance: Whether to track feature importance with predictions + importance_type: Type of importance to track ('split' or 'gain') + + Example: + >>> monitor = LightGBMMonitor( + ... client=client, + ... model_name="lgbm_classifier", + ... track_feature_importance=True, + ... importance_type="gain" + ... ) + >>> monitor.register_from_model(model, X_train, y_train) + >>> predictions = monitor.predict(model, X_test, y_test) + """ + + def __init__( + self, + client: Any, + model_name: str = "lightgbm_model", + track_feature_importance: bool = True, + importance_type: str = "gain", + **kwargs, + ): + """ + Initialize LightGBM monitor. + + Args: + client: WhiteBoxAI client instance + model_name: Name of the model + track_feature_importance: Whether to track feature importance + importance_type: Type of importance ('split' or 'gain') + **kwargs: Additional arguments passed to ModelMonitor + """ + if not LIGHTGBM_AVAILABLE: + raise ImportError("LightGBM is not installed. Install with: pip install lightgbm") + + super().__init__(client=client, model_name=model_name, **kwargs) + self.track_feature_importance = track_feature_importance + self.importance_type = importance_type + self._feature_names: Optional[List[str]] = None + + def register_from_model( + self, + model: Any, + X_train: Optional[np.ndarray] = None, + y_train: Optional[np.ndarray] = None, + model_type: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Register LightGBM model with WhiteBoxAI. + + Automatically extracts model metadata including feature names, + number of features, number of trees, and feature importance. + + Args: + model: LightGBM model (Booster, LGBMClassifier, LGBMRegressor, etc.) + X_train: Training features for baseline + y_train: Training labels for baseline + model_type: Model type override (auto-detected if None) + metadata: Additional metadata + + Returns: + Model ID from registration + """ + # Extract model information + model_metadata = metadata or {} + + # Get feature names + if hasattr(model, "feature_name_"): + self._feature_names = model.feature_name_ + model_metadata["feature_names"] = self._feature_names + elif hasattr(model, "feature_names_in_"): + self._feature_names = list(model.feature_names_in_) + model_metadata["feature_names"] = self._feature_names + elif X_train is not None and hasattr(X_train, "columns"): + self._feature_names = list(X_train.columns) + model_metadata["feature_names"] = self._feature_names + + # Get number of features + if hasattr(model, "n_features_in_"): + model_metadata["num_features"] = model.n_features_in_ + elif self._feature_names: + model_metadata["num_features"] = len(self._feature_names) + + # Get number of trees + if hasattr(model, "n_estimators"): + model_metadata["num_trees"] = model.n_estimators + elif hasattr(model, "best_iteration_"): + model_metadata["num_trees"] = model.best_iteration_ + elif hasattr(model, "num_trees"): + model_metadata["num_trees"] = model.num_trees() + + # Get LightGBM parameters + if hasattr(model, "get_params"): + params = model.get_params() + model_metadata["lightgbm_params"] = { + k: str(v) for k, v in params.items() if v is not None + } + + # Get feature importance + if self.track_feature_importance: + try: + importance = self._get_feature_importance(model) + if importance: + model_metadata["feature_importance"] = importance + except Exception as e: + warnings.warn(f"Failed to extract feature importance: {e}") + + # Detect model type + if model_type is None: + if hasattr(model, "_estimator_type"): + model_type = model._estimator_type + elif isinstance(model, lgb.LGBMClassifier): + model_type = "classification" + elif isinstance(model, lgb.LGBMRegressor): + model_type = "regression" + elif isinstance(model, lgb.LGBMRanker): + model_type = "ranking" + else: + model_type = "classification" # Default + + model_metadata["framework"] = "lightgbm" + model_metadata["lightgbm_version"] = lgb.__version__ + + # Register model + model_id = self.register_model( + name=self.model_name, model_type=model_type, metadata=model_metadata + ) + + # Set baseline if training data provided + if X_train is not None and y_train is not None: + self.set_baseline(X_train, y_train) + + return model_id + + def predict( + self, + model: Any, + X: np.ndarray, + y_true: Optional[np.ndarray] = None, + log_predictions: bool = True, + metadata: Optional[Dict[str, Any]] = None, + ) -> np.ndarray: + """ + Make predictions and log to WhiteBoxAI. + + Args: + model: LightGBM model + X: Input features + y_true: True labels (optional) + log_predictions: Whether to log predictions + metadata: Additional metadata + + Returns: + Model predictions + """ + # Make predictions + predictions = model.predict(X) + + # Get prediction probabilities if classification + probabilities = None + if hasattr(model, "predict_proba"): + try: + probabilities = model.predict_proba(X) + except Exception: + pass + + # Log predictions + if log_predictions: + pred_metadata = metadata or {} + + # Add feature importance for this batch if enabled + if self.track_feature_importance: + try: + importance = self._get_feature_importance(model) + if importance: + pred_metadata["feature_importance"] = importance + except Exception as e: + warnings.warn(f"Failed to extract feature importance: {e}") + + self.log_predictions( + inputs=X, + predictions=predictions, + actuals=y_true, + probabilities=probabilities, + metadata=pred_metadata, + ) + + return predictions + + def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: + """ + Extract feature importance from LightGBM model. + + Args: + model: LightGBM model + + Returns: + Dictionary mapping feature names to importance scores + """ + try: + # Try sklearn-style feature_importances_ + if hasattr(model, "feature_importances_"): + importances = model.feature_importances_ + if self._feature_names and len(importances) == len(self._feature_names): + return dict(zip(self._feature_names, importances.tolist())) + else: + return {f"f{i}": float(v) for i, v in enumerate(importances)} + + # Try feature_importance method (native LightGBM) + if hasattr(model, "feature_importance"): + importances = model.feature_importance(importance_type=self.importance_type) + if self._feature_names and len(importances) == len(self._feature_names): + return dict(zip(self._feature_names, importances.tolist())) + else: + return {f"f{i}": float(v) for i, v in enumerate(importances)} + + # Try booster + if hasattr(model, "booster_"): + booster = model.booster_ + importances = booster.feature_importance(importance_type=self.importance_type) + feature_names = booster.feature_name() + if len(importances) == len(feature_names): + return dict(zip(feature_names, importances.tolist())) + + return None + except Exception: + return None + + +# Unified wrapper functions +def wrap_xgboost_model(model: Any, monitor: XGBoostMonitor, auto_register: bool = True) -> Any: + """ + Wrap an XGBoost model for automatic monitoring. + + This creates a wrapper that automatically logs predictions when + the model's predict method is called. + + Args: + model: XGBoost model to wrap + monitor: XGBoostMonitor instance + auto_register: Whether to auto-register the model + + Returns: + Wrapped model with monitoring + + Example: + >>> monitor = XGBoostMonitor(client=client, model_name="my_model") + >>> wrapped_model = wrap_xgboost_model(model, monitor) + >>> predictions = wrapped_model.predict(X_test) # Auto-logged + """ + if auto_register and not monitor._model_id: + monitor.register_from_model(model) + + # Store original methods + original_predict = model.predict + if hasattr(model, "predict_proba"): + original_predict_proba = model.predict_proba + else: + original_predict_proba = None + + # Wrap predict method + def wrapped_predict(X, *args, **kwargs): + predictions = original_predict(X, *args, **kwargs) + + # Log predictions + try: + monitor.log_predictions(inputs=X, predictions=predictions) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return predictions + + # Wrap predict_proba if available + if original_predict_proba: + + def wrapped_predict_proba(X, *args, **kwargs): + probabilities = original_predict_proba(X, *args, **kwargs) + + # Log with probabilities + try: + predictions = np.argmax(probabilities, axis=1) + monitor.log_predictions( + inputs=X, predictions=predictions, probabilities=probabilities + ) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return probabilities + + model.predict_proba = wrapped_predict_proba + + model.predict = wrapped_predict + + return model + + +def wrap_lightgbm_model(model: Any, monitor: LightGBMMonitor, auto_register: bool = True) -> Any: + """ + Wrap a LightGBM model for automatic monitoring. + + This creates a wrapper that automatically logs predictions when + the model's predict method is called. + + Args: + model: LightGBM model to wrap + monitor: LightGBMMonitor instance + auto_register: Whether to auto-register the model + + Returns: + Wrapped model with monitoring + + Example: + >>> monitor = LightGBMMonitor(client=client, model_name="my_model") + >>> wrapped_model = wrap_lightgbm_model(model, monitor) + >>> predictions = wrapped_model.predict(X_test) # Auto-logged + """ + if auto_register and not monitor._model_id: + monitor.register_from_model(model) + + # Store original methods + original_predict = model.predict + if hasattr(model, "predict_proba"): + original_predict_proba = model.predict_proba + else: + original_predict_proba = None + + # Wrap predict method + def wrapped_predict(X, *args, **kwargs): + predictions = original_predict(X, *args, **kwargs) + + # Log predictions + try: + monitor.log_predictions(inputs=X, predictions=predictions) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return predictions + + # Wrap predict_proba if available + if original_predict_proba: + + def wrapped_predict_proba(X, *args, **kwargs): + probabilities = original_predict_proba(X, *args, **kwargs) + + # Log with probabilities + try: + predictions = np.argmax(probabilities, axis=1) + monitor.log_predictions( + inputs=X, predictions=predictions, probabilities=probabilities + ) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return probabilities + + model.predict_proba = wrapped_predict_proba + + model.predict = wrapped_predict + + return model diff --git a/src/whiteboxai/integrations/crewai_monitor.py b/src/whiteboxai/integrations/crewai_monitor.py new file mode 100644 index 0000000..9ab35cd --- /dev/null +++ b/src/whiteboxai/integrations/crewai_monitor.py @@ -0,0 +1,474 @@ +""" +CrewAI Integration for WhiteBoxAI + +Monitor CrewAI multi-agent workflows with automatic tracking of agents, +tasks, interactions, and costs. +""" + +import logging +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + + +class CrewAIMonitor: + """ + Monitor CrewAI multi-agent workflows. + + Automatically tracks: + - Workflow lifecycle (start, completion) + - Agent definitions and executions + - Task assignments and completions + - Agent-to-agent interactions + - Token usage and costs + + Example: + >>> from whiteboxai.integrations import CrewAIMonitor + >>> from crewai import Agent, Task, Crew + >>> + >>> monitor = CrewAIMonitor(api_key="your_api_key") + >>> + >>> # Define agents + >>> researcher = Agent( + ... role="Research Analyst", + ... goal="Find accurate information", + ... backstory="Expert researcher", + ... tools=[search_tool] + ... ) + >>> + >>> writer = Agent( + ... role="Content Writer", + ... goal="Write engaging content", + ... backstory="Professional writer", + ... tools=[writing_tool] + ... ) + >>> + >>> # Create tasks + >>> research_task = Task( + ... description="Research topic X", + ... agent=researcher + ... ) + >>> + >>> writing_task = Task( + ... description="Write article based on research", + ... agent=writer + ... ) + >>> + >>> # Create and monitor crew + >>> crew = Crew( + ... agents=[researcher, writer], + ... tasks=[research_task, writing_task], + ... process=Process.sequential + ... ) + >>> + >>> workflow_id = monitor.start_monitoring( + ... crew=crew, + ... workflow_name="Article Generation", + ... metadata={"topic": "AI Safety"} + ... ) + >>> + >>> # Execute crew (automatically monitored) + >>> result = crew.kickoff() + >>> + >>> # Complete monitoring + >>> monitor.complete_monitoring(outputs={"article": result}) + """ + + def __init__( + self, api_key: str, api_url: Optional[str] = None, organization_id: Optional[str] = None + ): + """ + Initialize CrewAI monitor. + + Args: + api_key: WhiteBoxAI API key + api_url: WhiteBoxAI API URL (optional, defaults to production) + organization_id: Organization ID (optional, extracted from token if not provided) + """ + from whiteboxai import WhiteBoxAI + + self.client = ( + WhiteBoxAI(api_key=api_key, base_url=api_url) + if api_url + else WhiteBoxAI(api_key=api_key) + ) + self.organization_id = organization_id + self.workflow_id = None + self.agent_map = {} # Maps CrewAI agent to WhiteBoxAI agent_id + self.task_map = {} # Maps CrewAI task to WhiteBoxAI task_id + self.execution_map = {} # Maps agent to execution_id + + logger.info("CrewAI Monitor initialized") + + def start_monitoring( + self, + crew: Any, # crewai.Crew + workflow_name: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Start monitoring a CrewAI crew workflow. + + Args: + crew: CrewAI Crew instance + workflow_name: Name for the workflow + metadata: Optional workflow metadata + + Returns: + Workflow ID (UUID string) + """ + try: + # Create workflow + workflow_data = { + "name": workflow_name, + "framework": "crewai", + "metadata": metadata or {}, + } + + response = self.client.request( + "POST", "/api/v1/workflows/multi-agent/start", data=workflow_data + ) + + self.workflow_id = response.get("id") + logger.info(f"Started monitoring CrewAI workflow: {self.workflow_id}") + + # Register agents + for agent in crew.agents: + self._register_agent(agent) + + # Register tasks + for task in crew.tasks: + self._register_task(task) + + # Start workflow + start_data = { + "inputs": { + "agent_count": len(crew.agents), + "task_count": len(crew.tasks), + "process": str(crew.process) if hasattr(crew, "process") else "sequential", + } + } + + self.client.request( + "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/start", data=start_data + ) + + return self.workflow_id + + except Exception as e: + logger.error(f"Error starting CrewAI monitoring: {str(e)}") + raise + + def _register_agent(self, crew_agent: Any) -> str: + """ + Register a CrewAI agent with WhiteBoxAI. + + Args: + crew_agent: CrewAI Agent instance + + Returns: + Agent ID (UUID string) + """ + try: + agent_data = { + "name": getattr(crew_agent, "role", "Unknown Agent"), + "role": getattr(crew_agent, "role", None), + "agent_type": "crewai_agent", + "goal": getattr(crew_agent, "goal", None), + "backstory": getattr(crew_agent, "backstory", None), + "tools": [tool.__class__.__name__ for tool in getattr(crew_agent, "tools", [])], + "llm_provider": ( + getattr(getattr(crew_agent, "llm", None), "model_name", "unknown").split("/")[0] + if hasattr(crew_agent, "llm") + else None + ), + "model_name": ( + getattr(getattr(crew_agent, "llm", None), "model_name", None) + if hasattr(crew_agent, "llm") + else None + ), + "metadata": { + "verbose": getattr(crew_agent, "verbose", False), + "allow_delegation": getattr(crew_agent, "allow_delegation", False), + "max_iter": getattr(crew_agent, "max_iter", None), + }, + } + + response = self.client.request( + "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/agents", data=agent_data + ) + + agent_id = response.get("id") + self.agent_map[id(crew_agent)] = agent_id + + logger.debug(f"Registered agent: {agent_data['name']} ({agent_id})") + + return agent_id + + except Exception as e: + logger.error(f"Error registering agent: {str(e)}") + raise + + def _register_task(self, crew_task: Any) -> str: + """ + Register a CrewAI task with WhiteBoxAI. + + Args: + crew_task: CrewAI Task instance + + Returns: + Task ID (UUID string) + """ + try: + # Get agent ID for task's assigned agent + agent_id = None + if hasattr(crew_task, "agent") and crew_task.agent: + agent_id = self.agent_map.get(id(crew_task.agent)) + + task_data = { + "task_name": getattr(crew_task, "description", "Unknown Task")[:255], + "description": getattr(crew_task, "description", None), + "expected_output": getattr(crew_task, "expected_output", None), + "agent_id": agent_id, + "context": { + "tools": [tool.__class__.__name__ for tool in getattr(crew_task, "tools", [])], + "async_execution": getattr(crew_task, "async_execution", False), + }, + } + + response = self.client.request( + "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/tasks", data=task_data + ) + + task_id = response.get("id") + self.task_map[id(crew_task)] = task_id + + logger.debug(f"Registered task: {task_data['task_name'][:50]}... ({task_id})") + + return task_id + + except Exception as e: + logger.error(f"Error registering task: {str(e)}") + raise + + def log_agent_execution( + self, + agent: Any, + inputs: Optional[Dict] = None, + outputs: Optional[Dict] = None, + tokens_used: int = 0, + cost: float = 0.0, + ) -> None: + """ + Log an agent execution. + + Args: + agent: CrewAI Agent instance + inputs: Execution inputs + outputs: Execution outputs + tokens_used: Tokens consumed + cost: Execution cost + """ + try: + agent_id = self.agent_map.get(id(agent)) + if not agent_id: + logger.warning("Agent not registered, skipping execution log") + return + + execution_data = { + "agent_id": agent_id, + "inputs": inputs, + "outputs": outputs, + "tokens_used": tokens_used, + "cost": cost, + "status": "completed", + } + + self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/executions", + data=execution_data, + ) + + logger.debug(f"Logged execution for agent: {agent_id}") + + except Exception as e: + logger.error(f"Error logging agent execution: {str(e)}") + + def log_task_completion( + self, + task: Any, + status: str = "completed", + output_data: Optional[Dict] = None, + error_message: Optional[str] = None, + ) -> None: + """ + Log task completion. + + Args: + task: CrewAI Task instance + status: Task status (completed, failed) + output_data: Task output + error_message: Error message if failed + """ + try: + task_id = self.task_map.get(id(task)) + if not task_id: + logger.warning("Task not registered, skipping completion log") + return + + update_data = { + "status": status, + "output_data": output_data, + "error_message": error_message, + } + + self.client.request( + "PATCH", f"/api/v1/workflows/multi-agent/tasks/{task_id}", data=update_data + ) + + logger.debug(f"Logged task completion: {task_id} ({status})") + + except Exception as e: + logger.error(f"Error logging task completion: {str(e)}") + + def log_interaction( + self, + from_agent: Any, + to_agent: Any, + interaction_type: str = "delegation", + message: Optional[str] = None, + ) -> None: + """ + Log agent-to-agent interaction. + + Args: + from_agent: Source agent + to_agent: Target agent + interaction_type: Type of interaction (delegation, handoff, query, feedback) + message: Interaction message + """ + try: + from_agent_id = self.agent_map.get(id(from_agent)) + to_agent_id = self.agent_map.get(id(to_agent)) + + if not from_agent_id or not to_agent_id: + logger.warning("Agents not registered, skipping interaction log") + return + + interaction_data = { + "interaction_type": interaction_type, + "from_agent_id": from_agent_id, + "to_agent_id": to_agent_id, + "message": message, + } + + self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/interactions", + data=interaction_data, + ) + + logger.debug( + f"Logged interaction: {from_agent_id} -> {to_agent_id} ({interaction_type})" + ) + + except Exception as e: + logger.error(f"Error logging interaction: {str(e)}") + + def complete_monitoring( + self, + status: str = "completed", + outputs: Optional[Dict] = None, + error_message: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Complete workflow monitoring. + + Args: + status: Final workflow status (completed, failed, cancelled) + outputs: Workflow outputs + error_message: Error message if failed + + Returns: + Workflow summary with analytics + """ + try: + complete_data = {"status": status, "outputs": outputs, "error_message": error_message} + + self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{self.workflow_id}/complete", + data=complete_data, + ) + + logger.info(f"Completed monitoring workflow: {self.workflow_id} ({status})") + + # Get analytics + analytics = self.get_analytics() + + return {"workflow_id": self.workflow_id, "status": status, "analytics": analytics} + + except Exception as e: + logger.error(f"Error completing monitoring: {str(e)}") + raise + + def get_analytics(self) -> Dict[str, Any]: + """ + Get workflow analytics. + + Returns: + Dict with metrics, cost breakdown, and bottlenecks + """ + try: + if not self.workflow_id: + raise ValueError("No active workflow") + + analytics = self.client.request( + "GET", f"/api/v1/workflows/multi-agent/{self.workflow_id}/analytics" + ) + + cost_breakdown = self.client.request( + "GET", f"/api/v1/workflows/multi-agent/{self.workflow_id}/cost-breakdown" + ) + + return {"metrics": analytics, "cost_breakdown": cost_breakdown} + + except Exception as e: + logger.error(f"Error getting analytics: {str(e)}") + return {} + + +# Helper function for easy usage +def monitor_crew( + crew: Any, + workflow_name: str, + api_key: str, + api_url: Optional[str] = None, + metadata: Optional[Dict] = None, +) -> CrewAIMonitor: + """ + Convenience function to start monitoring a CrewAI crew. + + Args: + crew: CrewAI Crew instance + workflow_name: Workflow name + api_key: WhiteBoxAI API key + api_url: WhiteBoxAI API URL (optional) + metadata: Workflow metadata (optional) + + Returns: + CrewAIMonitor instance with workflow started + + Example: + >>> monitor = monitor_crew( + ... crew=my_crew, + ... workflow_name="Research & Writing", + ... api_key="your_api_key" + ... ) + >>> result = my_crew.kickoff() + >>> monitor.complete_monitoring(outputs={"result": result}) + """ + monitor = CrewAIMonitor(api_key=api_key, api_url=api_url) + monitor.start_monitoring(crew=crew, workflow_name=workflow_name, metadata=metadata) + return monitor diff --git a/src/whiteboxai/integrations/langchain.py b/src/whiteboxai/integrations/langchain.py new file mode 100644 index 0000000..4112d89 --- /dev/null +++ b/src/whiteboxai/integrations/langchain.py @@ -0,0 +1,636 @@ +""" +LangChain Integration + +Integration for monitoring LangChain applications including chains, agents, and RAG pipelines. +""" + +import time +import warnings +from typing import Any, Dict, List, Optional + +try: + from langchain.agents import AgentExecutor + from langchain.callbacks.base import BaseCallbackHandler + from langchain.chains.base import Chain + from langchain.schema import AgentAction, AgentFinish, LLMResult + + LANGCHAIN_AVAILABLE = True +except ImportError: + LANGCHAIN_AVAILABLE = False + BaseCallbackHandler = object + Chain = object + AgentExecutor = object + +from whiteboxai.monitor import ModelMonitor + + +class LangChainMonitor(ModelMonitor): + """ + Monitor for LangChain applications. + + Supports: + - Chain executions + - Agent runs + - LLM calls + - Tool usage + - RAG pipelines + - Memory operations + + Example: + ```python + from langchain.chains import LLMChain + from langchain.llms import OpenAI + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.langchain import LangChainMonitor + + # Setup monitoring + client = WhiteBoxAI(api_key="your-api-key") + monitor = LangChainMonitor( + client=client, + application_name="my_langchain_app" + ) + + # Register application + monitor.register_application( + name="LangChain Q&A", + version="1.0.0" + ) + + # Create callback handler + callback = monitor.create_callback_handler() + + # Use with chain + llm = OpenAI() + chain = LLMChain(llm=llm, prompt=prompt) + result = chain.run(input="Question", callbacks=[callback]) + ``` + """ + + def __init__( + self, + client, + application_name: Optional[str] = None, + track_tokens: bool = True, + track_cost: bool = True, + **kwargs, + ): + """ + Initialize LangChain monitor. + + Args: + client: WhiteBoxAI client instance + application_name: Name of the LangChain application + track_tokens: Track token usage + track_cost: Track API costs + **kwargs: Additional arguments for ModelMonitor + """ + if not LANGCHAIN_AVAILABLE: + raise ImportError("langchain is not installed. Install with: pip install langchain") + + super().__init__(client, **kwargs) + self._application_name = application_name + self._track_tokens = track_tokens + self._track_cost = track_cost + self._registered = False + + def register_application( + self, + name: Optional[str] = None, + version: Optional[str] = None, + description: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register LangChain application with WhiteBoxAI. + + Args: + name: Application name + version: Application version + description: Application description + **kwargs: Additional metadata + + Returns: + Model ID (application ID) + """ + if name is None: + name = self._application_name or "langchain_app" + + # Prepare metadata + metadata = { + "framework": "langchain", + "application_type": "langchain", + "track_tokens": self._track_tokens, + "track_cost": self._track_cost, + **kwargs, + } + + # Register as a model (application) + self.model_id = self.client.register_model( + name=name, + model_type="generation", # LangChain apps are typically generative + framework="langchain", + version=version, + description=description, + metadata=metadata, + ) + + self._registered = True + return self.model_id + + def create_callback_handler(self) -> "WhiteBoxAICallbackHandler": + """ + Create a LangChain callback handler for automatic logging. + + Returns: + WhiteBoxAICallbackHandler instance + + Example: + ```python + callback = monitor.create_callback_handler() + chain.run(input="Question", callbacks=[callback]) + ``` + """ + if not self._registered: + self.register_application() + + return WhiteBoxAICallbackHandler(monitor=self) + + def log_chain_execution( + self, + chain_name: str, + inputs: Dict[str, Any], + outputs: Dict[str, Any], + execution_time: float, + llm_calls: Optional[List[Dict]] = None, + tool_calls: Optional[List[Dict]] = None, + **kwargs: Any, + ) -> None: + """ + Log a chain execution. + + Args: + chain_name: Name of the chain + inputs: Chain inputs + outputs: Chain outputs + execution_time: Time taken to execute (seconds) + llm_calls: List of LLM calls made + tool_calls: List of tool calls made + **kwargs: Additional metadata + """ + metadata = { + "chain_name": chain_name, + "execution_time": execution_time, + "num_llm_calls": len(llm_calls) if llm_calls else 0, + "num_tool_calls": len(tool_calls) if tool_calls else 0, + **kwargs, + } + + if llm_calls: + metadata["llm_calls"] = llm_calls + if tool_calls: + metadata["tool_calls"] = tool_calls + + # Log as prediction + self.log_prediction( + inputs=inputs, + prediction=outputs, + metadata=metadata, + ) + + def log_agent_execution( + self, + agent_name: str, + inputs: Dict[str, Any], + outputs: Dict[str, Any], + execution_time: float, + steps: List[Dict], + **kwargs: Any, + ) -> None: + """ + Log an agent execution. + + Args: + agent_name: Name of the agent + inputs: Agent inputs + outputs: Agent outputs + execution_time: Time taken to execute (seconds) + steps: List of agent steps/actions + **kwargs: Additional metadata + """ + metadata = { + "agent_name": agent_name, + "execution_time": execution_time, + "num_steps": len(steps), + "steps": steps, + **kwargs, + } + + # Log as prediction + self.log_prediction( + inputs=inputs, + prediction=outputs, + metadata=metadata, + ) + + def log_llm_call( + self, + prompt: str, + response: str, + model: str, + tokens_used: Optional[int] = None, + cost: Optional[float] = None, + latency: Optional[float] = None, + **kwargs: Any, + ) -> None: + """ + Log an LLM call. + + Args: + prompt: Input prompt + response: LLM response + model: Model name + tokens_used: Number of tokens used + cost: API cost + latency: Response latency (seconds) + **kwargs: Additional metadata + """ + metadata = { + "model": model, + "tokens_used": tokens_used, + "cost": cost, + "latency": latency, + **kwargs, + } + + self.log_prediction( + inputs={"prompt": prompt}, + prediction={"response": response}, + metadata=metadata, + ) + + def log_tool_call( + self, + tool_name: str, + tool_input: Any, + tool_output: Any, + execution_time: Optional[float] = None, + **kwargs: Any, + ) -> None: + """ + Log a tool call. + + Args: + tool_name: Name of the tool + tool_input: Tool input + tool_output: Tool output + execution_time: Time taken to execute (seconds) + **kwargs: Additional metadata + """ + metadata = { + "tool_name": tool_name, + "execution_time": execution_time, + **kwargs, + } + + self.log_prediction( + inputs={"tool_input": tool_input}, + prediction={"tool_output": tool_output}, + metadata=metadata, + ) + + def log_rag_retrieval( + self, + query: str, + documents: List[Dict], + num_retrieved: int, + retrieval_time: Optional[float] = None, + relevance_scores: Optional[List[float]] = None, + **kwargs: Any, + ) -> None: + """ + Log a RAG retrieval operation. + + Args: + query: Query text + documents: Retrieved documents + num_retrieved: Number of documents retrieved + retrieval_time: Time taken to retrieve (seconds) + relevance_scores: Relevance scores for documents + **kwargs: Additional metadata + """ + metadata = { + "num_retrieved": num_retrieved, + "retrieval_time": retrieval_time, + "relevance_scores": relevance_scores, + **kwargs, + } + + self.log_prediction( + inputs={"query": query}, + prediction={"documents": documents}, + metadata=metadata, + ) + + +class WhiteBoxAICallbackHandler(BaseCallbackHandler): + """ + LangChain callback handler for automatic logging to WhiteBoxAI. + + Example: + ```python + monitor = LangChainMonitor(client, application_name="my_app") + monitor.register_application() + callback = monitor.create_callback_handler() + + # Use with chain + chain.run(input="Question", callbacks=[callback]) + + # Use with agent + agent.run(input="Task", callbacks=[callback]) + ``` + """ + + def __init__(self, monitor: LangChainMonitor): + """ + Initialize callback handler. + + Args: + monitor: LangChainMonitor instance + """ + super().__init__() + self.monitor = monitor + self._chain_start_times = {} + self._agent_start_times = {} + self._llm_start_times = {} + self._tool_start_times = {} + self._chain_inputs = {} + self._agent_inputs = {} + self._agent_steps = {} + + def on_chain_start( + self, + serialized: Dict[str, Any], + inputs: Dict[str, Any], + **kwargs: Any, + ) -> None: + """Called when a chain starts running.""" + run_id = kwargs.get("run_id", id(serialized)) + self._chain_start_times[run_id] = time.time() + self._chain_inputs[run_id] = inputs + + def on_chain_end( + self, + outputs: Dict[str, Any], + **kwargs: Any, + ) -> None: + """Called when a chain ends running.""" + run_id = kwargs.get("run_id") + if run_id in self._chain_start_times: + execution_time = time.time() - self._chain_start_times[run_id] + inputs = self._chain_inputs.get(run_id, {}) + + # Get chain name + chain_name = kwargs.get("name", "unknown_chain") + + # Log chain execution + try: + self.monitor.log_chain_execution( + chain_name=chain_name, + inputs=inputs, + outputs=outputs, + execution_time=execution_time, + ) + except Exception as e: + warnings.warn(f"Failed to log chain execution: {e}") + + # Cleanup + del self._chain_start_times[run_id] + if run_id in self._chain_inputs: + del self._chain_inputs[run_id] + + def on_chain_error( + self, + error: Exception, + **kwargs: Any, + ) -> None: + """Called when a chain errors.""" + run_id = kwargs.get("run_id") + if run_id in self._chain_start_times: + # Cleanup on error + del self._chain_start_times[run_id] + if run_id in self._chain_inputs: + del self._chain_inputs[run_id] + + def on_agent_action( + self, + action: AgentAction, + **kwargs: Any, + ) -> None: + """Called when an agent takes an action.""" + run_id = kwargs.get("run_id") + if run_id not in self._agent_steps: + self._agent_steps[run_id] = [] + + self._agent_steps[run_id].append( + { + "tool": action.tool, + "tool_input": action.tool_input, + "log": action.log, + } + ) + + def on_agent_finish( + self, + finish: AgentFinish, + **kwargs: Any, + ) -> None: + """Called when an agent finishes running.""" + run_id = kwargs.get("run_id") + if run_id in self._agent_start_times: + execution_time = time.time() - self._agent_start_times[run_id] + inputs = self._agent_inputs.get(run_id, {}) + steps = self._agent_steps.get(run_id, []) + + # Log agent execution + try: + self.monitor.log_agent_execution( + agent_name=kwargs.get("name", "unknown_agent"), + inputs=inputs, + outputs=finish.return_values, + execution_time=execution_time, + steps=steps, + ) + except Exception as e: + warnings.warn(f"Failed to log agent execution: {e}") + + # Cleanup + del self._agent_start_times[run_id] + if run_id in self._agent_inputs: + del self._agent_inputs[run_id] + if run_id in self._agent_steps: + del self._agent_steps[run_id] + + def on_llm_start( + self, + serialized: Dict[str, Any], + prompts: List[str], + **kwargs: Any, + ) -> None: + """Called when an LLM starts generating.""" + run_id = kwargs.get("run_id", id(serialized)) + self._llm_start_times[run_id] = time.time() + + def on_llm_end( + self, + response: LLMResult, + **kwargs: Any, + ) -> None: + """Called when an LLM ends generating.""" + run_id = kwargs.get("run_id") + if run_id in self._llm_start_times: + latency = time.time() - self._llm_start_times[run_id] + + # Extract response details + if response.generations and len(response.generations) > 0: + generation = response.generations[0][0] + response_text = generation.text + + # Extract token usage if available + tokens_used = None + if response.llm_output and "token_usage" in response.llm_output: + tokens_used = response.llm_output["token_usage"].get("total_tokens") + + # Log LLM call (if configured) + if self.monitor._track_tokens or self.monitor._track_cost: + try: + self.monitor.log_llm_call( + prompt="", # Prompt tracked separately + response=response_text, + model=kwargs.get("invocation_params", {}).get("model_name", "unknown"), + tokens_used=tokens_used, + latency=latency, + ) + except Exception as e: + warnings.warn(f"Failed to log LLM call: {e}") + + # Cleanup + del self._llm_start_times[run_id] + + def on_llm_error( + self, + error: Exception, + **kwargs: Any, + ) -> None: + """Called when an LLM errors.""" + run_id = kwargs.get("run_id") + if run_id in self._llm_start_times: + del self._llm_start_times[run_id] + + def on_tool_start( + self, + serialized: Dict[str, Any], + input_str: str, + **kwargs: Any, + ) -> None: + """Called when a tool starts running.""" + run_id = kwargs.get("run_id", id(serialized)) + self._tool_start_times[run_id] = time.time() + + def on_tool_end( + self, + output: str, + **kwargs: Any, + ) -> None: + """Called when a tool ends running.""" + run_id = kwargs.get("run_id") + if run_id in self._tool_start_times: + execution_time = time.time() - self._tool_start_times[run_id] + + # Log tool call + try: + self.monitor.log_tool_call( + tool_name=kwargs.get("name", "unknown_tool"), + tool_input=kwargs.get("input_str", ""), + tool_output=output, + execution_time=execution_time, + ) + except Exception as e: + warnings.warn(f"Failed to log tool call: {e}") + + # Cleanup + del self._tool_start_times[run_id] + + def on_tool_error( + self, + error: Exception, + **kwargs: Any, + ) -> None: + """Called when a tool errors.""" + run_id = kwargs.get("run_id") + if run_id in self._tool_start_times: + del self._tool_start_times[run_id] + + +def wrap_langchain_chain( + chain: Chain, + monitor: LangChainMonitor, +) -> Chain: + """ + Wrap a LangChain chain to automatically log executions. + + Args: + chain: LangChain chain to wrap + monitor: LangChainMonitor instance + + Returns: + Wrapped chain that logs executions + + Example: + ```python + chain = LLMChain(llm=llm, prompt=prompt) + wrapped_chain = wrap_langchain_chain(chain, monitor) + + # Executions automatically logged + result = wrapped_chain.run(input="Question") + ``` + """ + if not monitor._registered: + monitor.register_application() + + # Create callback handler + callback = monitor.create_callback_handler() + + # Store original run method + original_run = chain.run + original_call = chain.__call__ + + def logged_run(*args, **kwargs): + # Add callback to kwargs + if "callbacks" not in kwargs: + kwargs["callbacks"] = [] + kwargs["callbacks"].append(callback) + + # Run chain + return original_run(*args, **kwargs) + + def logged_call(*args, **kwargs): + # Add callback to kwargs + if "callbacks" not in kwargs: + kwargs["callbacks"] = [] + kwargs["callbacks"].append(callback) + + # Call chain + return original_call(*args, **kwargs) + + # Replace methods + chain.run = logged_run + chain.__call__ = logged_call + + return chain + + +__all__ = [ + "LangChainMonitor", + "WhiteBoxAICallbackHandler", + "wrap_langchain_chain", +] diff --git a/src/whiteboxai/integrations/langchain_agents.py b/src/whiteboxai/integrations/langchain_agents.py new file mode 100644 index 0000000..ec782da --- /dev/null +++ b/src/whiteboxai/integrations/langchain_agents.py @@ -0,0 +1,525 @@ +""" +LangChain Multi-Agent Integration for WhiteBoxAI + +Enhanced callback handler for monitoring multi-agent LangChain workflows including: +- LangGraph multi-agent patterns +- Agent supervisors and coordinators +- Tool usage and agent handoffs +- Agent-to-agent communication +""" + +from datetime import datetime +from typing import Any, Dict, List, Optional, Union + +from langchain.callbacks.base import BaseCallbackHandler +from langchain.schema import AgentAction, AgentFinish, LLMResult + +try: + from whiteboxai import WhiteBoxAI +except ImportError: + WhiteBoxAI = None + + +class MultiAgentCallbackHandler(BaseCallbackHandler): + """Enhanced callback handler for multi-agent LangChain workflows. + + This handler tracks: + - Agent executions and decisions + - Tool calls and results + - Agent-to-agent handoffs + - LLM calls per agent + - Workflow-level metrics + + Example: + ```python + from langchain.agents import AgentExecutor, create_react_agent + from whiteboxai.integrations import MultiAgentCallbackHandler + + # Initialize WhiteBoxAI client + client = WhiteBoxAI(api_key="your_key") + + # Create workflow + workflow_id = client.agent_workflows.create( + name="Research Workflow", + framework="langchain" + ).get("id") + + # Start workflow + client.agent_workflows.start(workflow_id) + + # Create callback + callback = MultiAgentCallbackHandler( + client=client, + workflow_id=workflow_id, + agent_name="researcher" + ) + + # Use with agent + agent_executor = AgentExecutor( + agent=agent, + tools=tools, + callbacks=[callback] + ) + result = agent_executor.run("Research AI safety") + + # Complete workflow + client.agent_workflows.complete( + workflow_id, + outputs={"result": result} + ) + ``` + """ + + def __init__( + self, + client: "WhiteBoxAI", + workflow_id: str, + agent_name: str = "main", + agent_role: Optional[str] = None, + track_tokens: bool = True, + track_costs: bool = True, + ): + """Initialize the callback handler. + + Args: + client: WhiteBoxAI client instance + workflow_id: ID of the workflow to track + agent_name: Name of the current agent + agent_role: Role/description of the agent + track_tokens: Whether to track token usage + track_costs: Whether to estimate costs + """ + if WhiteBoxAI is None: + raise ImportError( + "whiteboxai package not installed. " "Install with: pip install whiteboxai" + ) + + self.client = client + self.workflow_id = workflow_id + self.agent_name = agent_name + self.agent_role = agent_role or agent_name + self.track_tokens = track_tokens + self.track_costs = track_costs + + # Tracking state + self.current_execution_id: Optional[str] = None + self.execution_start_time: Optional[datetime] = None + self.llm_call_count = 0 + self.tool_call_count = 0 + self.total_tokens = 0 + self.total_cost = 0.0 + self.execution_inputs: Optional[Dict[str, Any]] = None + + def on_chain_start( + self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any + ) -> None: + """Run when chain starts.""" + # Start agent execution + self.execution_start_time = datetime.utcnow() + self.execution_inputs = inputs + self.llm_call_count = 0 + self.tool_call_count = 0 + self.total_tokens = 0 + self.total_cost = 0.0 + + def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None: + """Run when chain ends successfully.""" + if self.execution_start_time: + duration_ms = int( + (datetime.utcnow() - self.execution_start_time).total_seconds() * 1000 + ) + + # Log agent execution + try: + response = self.client.agent_workflows.create_execution( + workflow_id=self.workflow_id, + agent_name=self.agent_name, + status="completed", + inputs=self.execution_inputs, + outputs=outputs, + duration_ms=duration_ms, + llm_call_count=self.llm_call_count, + tool_call_count=self.tool_call_count, + tokens_used=self.total_tokens if self.track_tokens else None, + cost=self.total_cost if self.track_costs else None, + ) + self.current_execution_id = response.get("id") + except Exception as e: + print(f"Warning: Failed to log execution: {e}") + + # Reset state + self.execution_start_time = None + + def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: + """Run when chain errors.""" + if self.execution_start_time: + duration_ms = int( + (datetime.utcnow() - self.execution_start_time).total_seconds() * 1000 + ) + + # Log failed execution + try: + self.client.agent_workflows.create_execution( + workflow_id=self.workflow_id, + agent_name=self.agent_name, + status="failed", + inputs=self.execution_inputs, + outputs={"error": str(error)}, + duration_ms=duration_ms, + llm_call_count=self.llm_call_count, + tool_call_count=self.tool_call_count, + ) + except Exception as e: + print(f"Warning: Failed to log error: {e}") + + self.execution_start_time = None + + def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None: + """Run when LLM starts.""" + self.llm_call_count += 1 + + def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + """Run when LLM ends.""" + # Track tokens if available + if self.track_tokens and hasattr(response, "llm_output"): + llm_output = response.llm_output or {} + token_usage = llm_output.get("token_usage", {}) + + total = token_usage.get("total_tokens", 0) + self.total_tokens += total + + # Estimate cost if tracking + if self.track_costs and total > 0: + # Rough estimate: $0.002 per 1K tokens (GPT-3.5 pricing) + self.total_cost += (total / 1000) * 0.002 + + def on_agent_action(self, action: AgentAction, **kwargs: Any) -> None: + """Run when agent takes an action (tool call).""" + self.tool_call_count += 1 + + # Log tool call as interaction + try: + self.client.agent_workflows.create_interaction( + workflow_id=self.workflow_id, + from_agent=self.agent_name, + to_agent="tool", + interaction_type="tool_call", + message=f"Tool: {action.tool}, Input: {action.tool_input}", + meta_data={ + "tool": action.tool, + "tool_input": action.tool_input, + "log": action.log, + }, + ) + except Exception as e: + print(f"Warning: Failed to log tool call: {e}") + + def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None: + """Run when agent finishes execution.""" + # This is called when the agent completes its reasoning + pass + + def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None: + """Run when tool starts.""" + pass + + def on_tool_end(self, output: str, **kwargs: Any) -> None: + """Run when tool ends.""" + # Log tool result as interaction + try: + self.client.agent_workflows.create_interaction( + workflow_id=self.workflow_id, + from_agent="tool", + to_agent=self.agent_name, + interaction_type="response", + message=f"Tool result: {output[:500]}", # Truncate long outputs + meta_data={"output": output}, + ) + except Exception as e: + print(f"Warning: Failed to log tool result: {e}") + + def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: + """Run when tool errors.""" + try: + self.client.agent_workflows.create_interaction( + workflow_id=self.workflow_id, + from_agent="tool", + to_agent=self.agent_name, + interaction_type="response", + message=f"Tool error: {str(error)}", + meta_data={"error": str(error), "error_type": type(error).__name__}, + ) + except Exception as e: + print(f"Warning: Failed to log tool error: {e}") + + def on_text(self, text: str, **kwargs: Any) -> None: + """Run on arbitrary text.""" + pass + + +class LangGraphMultiAgentMonitor: + """Monitor for LangGraph multi-agent workflows. + + Provides higher-level monitoring for LangGraph patterns like: + - Agent supervisors + - Agent networks + - Sequential/parallel agent execution + + Example: + ```python + from langgraph.graph import StateGraph + from whiteboxai.integrations import LangGraphMultiAgentMonitor + + # Create monitor + monitor = LangGraphMultiAgentMonitor( + client=client, + workflow_name="Multi-Agent Research" + ) + + # Start monitoring + workflow_id = monitor.start_monitoring() + + # Register agents + monitor.register_agent("supervisor", role="Coordinates other agents") + monitor.register_agent("researcher", role="Gathers information") + monitor.register_agent("writer", role="Writes content") + + # Execute graph with callbacks + graph = StateGraph(...) + result = graph.invoke( + inputs, + config={"callbacks": [monitor.get_callbacks("supervisor")]} + ) + + # Complete monitoring + monitor.complete_monitoring(outputs={"result": result}) + ``` + """ + + def __init__( + self, client: "WhiteBoxAI", workflow_name: str, meta_data: Optional[Dict[str, Any]] = None + ): + """Initialize the LangGraph monitor. + + Args: + client: WhiteBoxAI client instance + workflow_name: Name for the workflow + meta_data: Additional meta_data to attach + """ + if WhiteBoxAI is None: + raise ImportError( + "whiteboxai package not installed. " "Install with: pip install whiteboxai" + ) + + self.client = client + self.workflow_name = workflow_name + self.workflow_meta_data = meta_data or {} + self.workflow_id: Optional[str] = None + self.callbacks: Dict[str, MultiAgentCallbackHandler] = {} + self.start_time: Optional[datetime] = None + + def start_monitoring(self, inputs: Optional[Dict[str, Any]] = None) -> str: + """Start workflow monitoring. + + Args: + inputs: Initial workflow inputs + + Returns: + workflow_id: ID of the created workflow + """ + self.start_time = datetime.utcnow() + + # Create workflow + response = self.client.agent_workflows.create( + name=self.workflow_name, + framework="langchain", + inputs=inputs, + meta_data=self.workflow_meta_data, + ) + self.workflow_id = response.get("id") + + # Start workflow + self.client.agent_workflows.start(self.workflow_id) + + return self.workflow_id + + def register_agent( + self, + agent_name: str, + role: Optional[str] = None, + model_name: Optional[str] = None, + tools: Optional[List[str]] = None, + **kwargs, + ) -> None: + """Register an agent in the workflow. + + Args: + agent_name: Name of the agent + role: Agent's role/goal + model_name: LLM model used + tools: List of tool names + **kwargs: Additional agent configuration + """ + if not self.workflow_id: + raise ValueError("Must call start_monitoring() first") + + self.client.agent_workflows.register_agent( + workflow_id=self.workflow_id, + name=agent_name, + role=role or agent_name, + model_name=model_name, + tools=tools, + **kwargs, + ) + + def get_callbacks( + self, agent_name: str, agent_role: Optional[str] = None + ) -> List[BaseCallbackHandler]: + """Get callbacks for a specific agent. + + Args: + agent_name: Name of the agent + agent_role: Optional role description + + Returns: + List of callback handlers + """ + if not self.workflow_id: + raise ValueError("Must call start_monitoring() first") + + if agent_name not in self.callbacks: + self.callbacks[agent_name] = MultiAgentCallbackHandler( + client=self.client, + workflow_id=self.workflow_id, + agent_name=agent_name, + agent_role=agent_role, + ) + + return [self.callbacks[agent_name]] + + def log_handoff( + self, + from_agent: str, + to_agent: str, + message: str, + meta_data: Optional[Dict[str, Any]] = None, + ) -> None: + """Log an agent-to-agent handoff. + + Args: + from_agent: Agent passing control + to_agent: Agent receiving control + message: Handoff message/context + meta_data: Additional meta_data + """ + if not self.workflow_id: + raise ValueError("Must call start_monitoring() first") + + self.client.agent_workflows.create_interaction( + workflow_id=self.workflow_id, + from_agent=from_agent, + to_agent=to_agent, + interaction_type="handoff", + message=message, + meta_data=meta_data, + ) + + def complete_monitoring( + self, outputs: Optional[Dict[str, Any]] = None, status: str = "completed" + ) -> Dict[str, Any]: + """Complete workflow monitoring. + + Args: + outputs: Final workflow outputs + status: Workflow status (completed/failed) + + Returns: + Summary with analytics + """ + if not self.workflow_id: + raise ValueError("Must call start_monitoring() first") + + # Complete workflow + self.client.agent_workflows.complete( + workflow_id=self.workflow_id, outputs=outputs, status=status + ) + + # Get analytics + try: + analytics = self.client.agent_workflows.get_analytics(self.workflow_id) + return { + "workflow_id": self.workflow_id, + "status": status, + "outputs": outputs, + "analytics": analytics, + } + except Exception as e: + print(f"Warning: Failed to retrieve analytics: {e}") + return {"workflow_id": self.workflow_id, "status": status, "outputs": outputs} + + +def monitor_langchain_agent( + client: "WhiteBoxAI", + agent_executor: Any, + workflow_name: str, + agent_name: str = "main", + inputs: Optional[Dict[str, Any]] = None, + **run_kwargs, +) -> Dict[str, Any]: + """Helper function to monitor a single LangChain agent execution. + + Args: + client: WhiteBoxAI client + agent_executor: LangChain AgentExecutor instance + workflow_name: Name for the workflow + agent_name: Name of the agent + inputs: Inputs to the agent + **run_kwargs: Additional arguments to pass to agent.run() + + Returns: + Dict with result and workflow_id + + Example: + ```python + from langchain.agents import AgentExecutor, create_react_agent + from whiteboxai.integrations import monitor_langchain_agent + + result_dict = monitor_langchain_agent( + client=client, + agent_executor=agent_executor, + workflow_name="Research Task", + agent_name="researcher", + inputs={"input": "Research AI safety"} + ) + + print(f"Result: {result_dict['result']}") + print(f"Workflow ID: {result_dict['workflow_id']}") + ``` + """ + # Create workflow + response = client.agent_workflows.create( + name=workflow_name, framework="langchain", inputs=inputs + ) + workflow_id = response.get("id") + + # Start workflow + client.agent_workflows.start(workflow_id) + + # Create callback + callback = MultiAgentCallbackHandler( + client=client, workflow_id=workflow_id, agent_name=agent_name + ) + + try: + # Run agent with callback + result = agent_executor.run(callbacks=[callback], **run_kwargs) + + # Complete workflow + client.agent_workflows.complete(workflow_id, outputs={"result": result}) + + return {"result": result, "workflow_id": workflow_id, "status": "completed"} + except Exception as e: + # Log failure + client.agent_workflows.complete(workflow_id, outputs={"error": str(e)}, status="failed") + + return {"result": None, "workflow_id": workflow_id, "status": "failed", "error": str(e)} diff --git a/src/whiteboxai/integrations/pytorch.py b/src/whiteboxai/integrations/pytorch.py new file mode 100644 index 0000000..16b7ddf --- /dev/null +++ b/src/whiteboxai/integrations/pytorch.py @@ -0,0 +1,268 @@ +""" +PyTorch Integration + +Integration for monitoring PyTorch models. +""" + +from typing import Any, Callable, Dict, Optional + +try: + import torch + import torch.nn as nn + + TORCH_AVAILABLE = True +except ImportError: + TORCH_AVAILABLE = False + nn = object + torch = None + +from whiteboxai.monitor import ModelMonitor + + +class TorchMonitor(ModelMonitor): + """ + Monitor for PyTorch models. + + Example: + ```python + import torch + import torch.nn as nn + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.pytorch import TorchMonitor + + # Define model + model = nn.Sequential( + nn.Linear(10, 64), + nn.ReLU(), + nn.Linear(64, 2) + ) + + # Setup monitoring + client = WhiteBoxAI(api_key="your-api-key") + monitor = TorchMonitor(client, model=model) + monitor.register_from_model(model_type="classification") + + # Wrap model for automatic monitoring + monitored_model = monitor.wrap_model(model) + + # Predictions are automatically logged + outputs = monitored_model(inputs) + ``` + """ + + def __init__(self, client, model: Optional[nn.Module] = None, **kwargs): + """Initialize PyTorch monitor.""" + if not TORCH_AVAILABLE: + raise ImportError("PyTorch is not installed. Install with: pip install torch") + + super().__init__(client, **kwargs) + self.model = model + + def register_from_model( + self, + name: Optional[str] = None, + model_type: str = "classification", + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register model from PyTorch module. + + Args: + name: Model name (default: class name) + model_type: Model type (classification, regression, etc.) + version: Model version + **kwargs: Additional metadata + + Returns: + Model ID + """ + if self.model is None: + raise ValueError("No model provided") + + # Generate name from model class + if name is None: + name = self.model.__class__.__name__ + + # Extract model metadata + metadata = self._extract_model_metadata() + metadata.update(kwargs) + + return self.register_model( + name=name, + model_type=model_type, + framework="pytorch", + version=version, + **metadata, + ) + + async def aregister_from_model( + self, + name: Optional[str] = None, + model_type: str = "classification", + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """Async version of register_from_model().""" + if self.model is None: + raise ValueError("No model provided") + + if name is None: + name = self.model.__class__.__name__ + + metadata = self._extract_model_metadata() + metadata.update(kwargs) + + return await self.aregister_model( + name=name, + model_type=model_type, + framework="pytorch", + version=version, + **metadata, + ) + + def wrap_model(self, model: nn.Module) -> "TorchWrapper": + """ + Wrap PyTorch model for automatic monitoring. + + Args: + model: PyTorch module + + Returns: + Wrapped model + """ + return TorchWrapper(model, self) + + def _extract_model_metadata(self) -> Dict[str, Any]: + """Extract metadata from PyTorch model.""" + metadata = {} + + if self.model is None: + return metadata + + # Count parameters + total_params = sum(p.numel() for p in self.model.parameters()) + trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + + metadata["total_parameters"] = total_params + metadata["trainable_parameters"] = trainable_params + + # Get model architecture + metadata["architecture"] = str(self.model) + + # Get PyTorch version + metadata["pytorch_version"] = torch.__version__ + + # Get device + device = next(self.model.parameters()).device + metadata["device"] = str(device) + + return metadata + + +class TorchWrapper(nn.Module): + """ + Wrapper for PyTorch models with automatic monitoring. + + This wrapper intercepts forward calls and logs predictions to WhiteBoxAI. + """ + + def __init__(self, model: nn.Module, monitor: TorchMonitor): + """Initialize wrapper.""" + super().__init__() + self.model = model + self.monitor = monitor + + def forward(self, x: torch.Tensor, log: bool = True) -> torch.Tensor: + """Forward pass with automatic logging.""" + # Call original forward + output = self.model(x) + + # Log predictions if enabled + if log: + self._log_batch_predictions(x, output) + + return output + + def _log_batch_predictions(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None: + """Log batch of predictions.""" + # Convert to numpy/lists + inputs_np = inputs.detach().cpu().numpy() + outputs_np = outputs.detach().cpu().numpy() + + # Create predictions list + predictions = [] + for i in range(len(inputs_np)): + pred = { + "inputs": inputs_np[i].tolist(), + "output": outputs_np[i].tolist(), + } + predictions.append(pred) + + # Log batch + try: + self.monitor.log_batch(predictions) + except Exception as e: + # Don't fail prediction if logging fails + print(f"Warning: Failed to log predictions: {e}") + + def __getattr__(self, name: str): + """Delegate attributes to wrapped model.""" + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + +def monitor_forward(monitor: TorchMonitor, input_extractor: Optional[Callable] = None): + """ + Decorator to monitor PyTorch forward passes. + + Args: + monitor: TorchMonitor instance + input_extractor: Function to extract inputs from args + + Example: + ```python + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.pytorch import TorchMonitor, monitor_forward + + client = WhiteBoxAI(api_key="your-api-key") + monitor = TorchMonitor(client, model_id=123) + + class MyModel(nn.Module): + @monitor_forward(monitor) + def forward(self, x): + return self.layers(x) + ``` + """ + + def decorator(forward_fn: Callable) -> Callable: + def wrapper(self, *args, **kwargs): + # Call original forward + output = forward_fn(self, *args, **kwargs) + + # Extract inputs + if input_extractor: + inputs = input_extractor(args, kwargs) + else: + inputs = args[0] if args else None + + # Log prediction + if inputs is not None and isinstance(inputs, torch.Tensor): + inputs_np = inputs.detach().cpu().numpy() + outputs_np = output.detach().cpu().numpy() + + try: + monitor.log_prediction( + inputs=inputs_np.tolist(), + output=outputs_np.tolist(), + ) + except Exception as e: + print(f"Warning: Failed to log prediction: {e}") + + return output + + return wrapper + + return decorator diff --git a/src/whiteboxai/integrations/sklearn.py b/src/whiteboxai/integrations/sklearn.py new file mode 100644 index 0000000..a1e9b4f --- /dev/null +++ b/src/whiteboxai/integrations/sklearn.py @@ -0,0 +1,219 @@ +""" +Scikit-learn Integration + +Integration for monitoring scikit-learn models. +""" + +from typing import Any, Dict, Optional + +try: + import sklearn + from sklearn.base import BaseEstimator + + SKLEARN_AVAILABLE = True +except ImportError: + SKLEARN_AVAILABLE = False + BaseEstimator = object + +import numpy as np + +from whiteboxai.monitor import ModelMonitor + + +class SklearnMonitor(ModelMonitor): + """ + Monitor for scikit-learn models. + + Example: + ```python + from sklearn.ensemble import RandomForestClassifier + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.sklearn import SklearnMonitor + + # Train model + model = RandomForestClassifier() + model.fit(X_train, y_train) + + # Setup monitoring + client = WhiteBoxAI(api_key="your-api-key") + monitor = SklearnMonitor(client, model=model) + monitor.register_from_model(model_type="classification") + + # Wrap model for automatic monitoring + monitored_model = monitor.wrap_model(model) + + # Predictions are automatically logged + predictions = monitored_model.predict(X_test) + ``` + """ + + def __init__(self, client, model: Optional[BaseEstimator] = None, **kwargs): + """Initialize sklearn monitor.""" + if not SKLEARN_AVAILABLE: + raise ImportError( + "scikit-learn is not installed. Install with: pip install scikit-learn" + ) + + super().__init__(client, **kwargs) + self.model = model + + def register_from_model( + self, + name: Optional[str] = None, + model_type: str = "classification", + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register model from sklearn estimator. + + Args: + name: Model name (default: class name) + model_type: Model type (classification, regression, clustering) + version: Model version + **kwargs: Additional metadata + + Returns: + Model ID + """ + if self.model is None: + raise ValueError("No model provided") + + # Generate name from model class + if name is None: + name = self.model.__class__.__name__ + + # Extract model metadata + metadata = self._extract_model_metadata() + metadata.update(kwargs) + + return self.register_model( + name=name, + model_type=model_type, + framework="sklearn", + version=version, + **metadata, + ) + + async def aregister_from_model( + self, + name: Optional[str] = None, + model_type: str = "classification", + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """Async version of register_from_model().""" + if self.model is None: + raise ValueError("No model provided") + + if name is None: + name = self.model.__class__.__name__ + + metadata = self._extract_model_metadata() + metadata.update(kwargs) + + return await self.aregister_model( + name=name, + model_type=model_type, + framework="sklearn", + version=version, + **metadata, + ) + + def wrap_model(self, model: BaseEstimator) -> "SklearnWrapper": + """ + Wrap sklearn model for automatic monitoring. + + Args: + model: Sklearn estimator + + Returns: + Wrapped model + """ + return SklearnWrapper(model, self) + + def _extract_model_metadata(self) -> Dict[str, Any]: + """Extract metadata from sklearn model.""" + metadata = {} + + if self.model is None: + return metadata + + # Get model parameters + if hasattr(self.model, "get_params"): + metadata["parameters"] = self.model.get_params() + + # Get feature names + if hasattr(self.model, "feature_names_in_"): + metadata["feature_names"] = self.model.feature_names_in_.tolist() + + # Get feature importances (for tree-based models) + if hasattr(self.model, "feature_importances_"): + metadata["feature_importances"] = self.model.feature_importances_.tolist() + + # Get classes (for classifiers) + if hasattr(self.model, "classes_"): + metadata["classes"] = self.model.classes_.tolist() + + # Get number of features + if hasattr(self.model, "n_features_in_"): + metadata["n_features"] = self.model.n_features_in_ + + # Get sklearn version + metadata["sklearn_version"] = sklearn.__version__ + + return metadata + + +class SklearnWrapper: + """ + Wrapper for sklearn models with automatic monitoring. + + This wrapper intercepts predict/predict_proba calls and logs predictions + to WhiteBoxAI automatically. + """ + + def __init__(self, model: BaseEstimator, monitor: SklearnMonitor): + """Initialize wrapper.""" + self.model = model + self.monitor = monitor + + def predict(self, X: np.ndarray, **kwargs) -> np.ndarray: + """Predict with automatic logging.""" + predictions = self.model.predict(X, **kwargs) + + # Log predictions + self._log_batch_predictions(X, predictions) + + return predictions + + def predict_proba(self, X: np.ndarray, **kwargs) -> np.ndarray: + """Predict probabilities with automatic logging.""" + probas = self.model.predict_proba(X, **kwargs) + + # Log predictions + self._log_batch_predictions(X, probas) + + return probas + + def _log_batch_predictions(self, inputs: np.ndarray, outputs: np.ndarray) -> None: + """Log batch of predictions.""" + # Convert to list format + predictions = [] + for i in range(len(inputs)): + pred = { + "inputs": inputs[i].tolist() if isinstance(inputs[i], np.ndarray) else inputs[i], + "output": outputs[i].tolist() if isinstance(outputs[i], np.ndarray) else outputs[i], + } + predictions.append(pred) + + # Log batch + try: + self.monitor.log_batch(predictions) + except Exception as e: + # Don't fail prediction if logging fails + print(f"Warning: Failed to log predictions: {e}") + + def __getattr__(self, name: str): + """Delegate other methods to wrapped model.""" + return getattr(self.model, name) diff --git a/src/whiteboxai/integrations/tensorflow.py b/src/whiteboxai/integrations/tensorflow.py new file mode 100644 index 0000000..b617673 --- /dev/null +++ b/src/whiteboxai/integrations/tensorflow.py @@ -0,0 +1,475 @@ +""" +TensorFlow/Keras Integration + +Integration for monitoring TensorFlow and Keras models. +""" + +import warnings +from typing import Any, Dict, Optional, Union + +try: + import tensorflow as tf + from tensorflow import keras + + TENSORFLOW_AVAILABLE = True +except ImportError: + TENSORFLOW_AVAILABLE = False + keras = None + +import numpy as np + +from whiteboxai.monitor import ModelMonitor + + +class KerasMonitor(ModelMonitor): + """ + Monitor for TensorFlow/Keras models. + + Example: + ```python + from tensorflow import keras + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.tensorflow import KerasMonitor + + # Build model + model = keras.Sequential([ + keras.layers.Dense(64, activation='relu'), + keras.layers.Dense(1) + ]) + model.compile(optimizer='adam', loss='mse') + + # Setup monitoring + client = WhiteBoxAI(api_key="your-api-key") + monitor = KerasMonitor(client, model=model, model_name="keras_model") + + # Train with monitoring callback + from whiteboxai.integrations.tensorflow import WhiteBoxAICallback + + model.fit(X_train, y_train, + callbacks=[WhiteBoxAICallback(monitor)]) + + # Make predictions with automatic logging + predictions = monitor.predict(X_test) + ``` + """ + + def __init__( + self, + client, + model: Optional["keras.Model"] = None, + model_name: Optional[str] = None, + model_type: str = "regression", + **kwargs, + ): + """ + Initialize Keras monitor. + + Args: + client: WhiteBoxAI client instance + model: Keras model to monitor + model_name: Name for the model + model_type: Type of model (classification, regression) + **kwargs: Additional arguments for ModelMonitor + """ + if not TENSORFLOW_AVAILABLE: + raise ImportError("TensorFlow is not installed. Install with: pip install tensorflow") + + super().__init__(client, **kwargs) + self.model = model + self._model_name = model_name + self._model_type = model_type + self._registered = False + + def register_from_model( + self, + name: Optional[str] = None, + model_type: Optional[str] = None, + version: Optional[str] = None, + framework_version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register Keras model with WhiteBoxAI. + + Args: + name: Model name (default: model class name) + model_type: Model type (classification, regression) + version: Model version + framework_version: TensorFlow version + **kwargs: Additional metadata + + Returns: + Model ID + """ + if self.model is None: + raise ValueError("No model provided") + + # Auto-detect model name + if name is None: + name = self._model_name or self.model.__class__.__name__ + + # Use provided or instance model type + model_type = model_type or self._model_type + + # Get TensorFlow version + if framework_version is None: + framework_version = tf.__version__ + + # Extract model metadata + metadata = { + "framework": "tensorflow", + "framework_version": framework_version, + "model_class": self.model.__class__.__name__, + **kwargs, + } + + # Get model configuration + try: + config = self.model.get_config() + metadata["model_config"] = str(config) + except Exception: + pass + + # Get input/output shapes + try: + if hasattr(self.model, "input_shape"): + metadata["input_shape"] = str(self.model.input_shape) + if hasattr(self.model, "output_shape"): + metadata["output_shape"] = str(self.model.output_shape) + except Exception: + pass + + # Register model + self.model_id = self.client.register_model( + name=name, + model_type=model_type, + framework="tensorflow", + version=version, + metadata=metadata, + ) + + self._registered = True + return self.model_id + + def predict( + self, + inputs: Union[np.ndarray, tf.Tensor], + log: bool = True, + actuals: Optional[np.ndarray] = None, + **kwargs: Any, + ) -> np.ndarray: + """ + Make predictions and optionally log them. + + Args: + inputs: Input data (numpy array or TensorFlow tensor) + log: Whether to log the prediction + actuals: Actual values (if available) + **kwargs: Additional metadata to log + + Returns: + Predictions as numpy array + """ + if self.model is None: + raise ValueError("No model provided") + + # Ensure model is registered + if not self._registered: + self.register_from_model() + + # Convert to numpy if needed + if isinstance(inputs, tf.Tensor): + inputs_np = inputs.numpy() + else: + inputs_np = inputs + + # Make prediction + predictions = self.model.predict(inputs, verbose=0) + + # Convert predictions to numpy + if isinstance(predictions, tf.Tensor): + predictions_np = predictions.numpy() + else: + predictions_np = predictions + + # Log if requested + if log: + # Determine if batch or single + if len(inputs_np.shape) == 1 or inputs_np.shape[0] == 1: + # Single prediction + self.log_prediction( + inputs=inputs_np[0] if len(inputs_np.shape) > 1 else inputs_np, + prediction=( + predictions_np[0] if len(predictions_np.shape) > 1 else predictions_np + ), + actual=actuals[0] if actuals is not None else None, + **kwargs, + ) + else: + # Batch prediction + self.log_batch( + inputs=inputs_np, + predictions=predictions_np, + actuals=actuals, + **kwargs, + ) + + return predictions_np + + def set_baseline( + self, + baseline_data: Union[np.ndarray, tf.Tensor], + baseline_labels: Optional[Union[np.ndarray, tf.Tensor]] = None, + ) -> None: + """ + Set baseline data for drift detection. + + Args: + baseline_data: Baseline input data + baseline_labels: Baseline labels (optional) + """ + # Convert to numpy + if isinstance(baseline_data, tf.Tensor): + baseline_data = baseline_data.numpy() + if baseline_labels is not None and isinstance(baseline_labels, tf.Tensor): + baseline_labels = baseline_labels.numpy() + + # Make baseline predictions if labels not provided + if baseline_labels is None and self.model is not None: + baseline_predictions = self.model.predict(baseline_data, verbose=0) + if isinstance(baseline_predictions, tf.Tensor): + baseline_predictions = baseline_predictions.numpy() + else: + baseline_predictions = baseline_labels + + # Set baseline through parent class + super().set_baseline(baseline_data, baseline_predictions) + + def log_epoch( + self, + epoch: int, + train_loss: Optional[float] = None, + val_loss: Optional[float] = None, + **metrics: Any, + ) -> None: + """ + Log metrics from a training epoch. + + Args: + epoch: Epoch number + train_loss: Training loss + val_loss: Validation loss + **metrics: Additional metrics (accuracy, etc.) + """ + metric_data = { + "epoch": epoch, + "train_loss": train_loss, + "val_loss": val_loss, + **metrics, + } + + self.log_custom_metric("training_epoch", metric_data) + + def log_checkpoint( + self, + epoch: int, + checkpoint_path: str, + metrics: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Log a model checkpoint. + + Args: + epoch: Epoch number + checkpoint_path: Path to saved checkpoint + metrics: Metrics at checkpoint + """ + checkpoint_data = { + "epoch": epoch, + "checkpoint_path": checkpoint_path, + "metrics": metrics or {}, + } + + self.log_custom_metric("checkpoint", checkpoint_data) + + def register_saved_model( + self, + model_path: str, + metadata: Optional[Dict[str, Any]] = None, + ) -> None: + """ + Register a SavedModel with WhiteBoxAI. + + Args: + model_path: Path to SavedModel directory + metadata: Additional metadata + """ + model_metadata = { + "model_path": model_path, + "format": "SavedModel", + **(metadata or {}), + } + + if not self._registered: + self.register_from_model(**model_metadata) + + +class WhiteBoxAICallback(keras.callbacks.Callback): + """ + Keras callback for automatic monitoring during training. + + Example: + ```python + monitor = KerasMonitor(client, model=model, model_name="my_model") + callback = WhiteBoxAICallback(monitor, log_frequency=1) + + model.fit(X_train, y_train, + validation_data=(X_val, y_val), + callbacks=[callback], + epochs=50) + ``` + """ + + def __init__( + self, + monitor: KerasMonitor, + log_frequency: int = 1, + log_weights: bool = False, + log_gradients: bool = False, + log_validation: bool = True, + ): + """ + Initialize callback. + + Args: + monitor: KerasMonitor instance + log_frequency: Log metrics every N epochs + log_weights: Whether to log model weights + log_gradients: Whether to log gradients + log_validation: Whether to log validation predictions + """ + super().__init__() + self.monitor = monitor + self.log_frequency = log_frequency + self.log_weights = log_weights + self.log_gradients = log_gradients + self.log_validation = log_validation + + # Ensure model is registered + if not self.monitor._registered: + self.monitor.register_from_model() + + def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None): + """Called at the end of each epoch.""" + if logs is None: + logs = {} + + # Log metrics at specified frequency + if (epoch + 1) % self.log_frequency == 0: + # Extract train and val metrics + train_loss = logs.get("loss") + val_loss = logs.get("val_loss") + + # Extract additional metrics + metrics = {} + for key, value in logs.items(): + if key not in ["loss", "val_loss"]: + metrics[key] = float(value) if value is not None else None + + # Log epoch metrics + self.monitor.log_epoch( + epoch=epoch, + train_loss=float(train_loss) if train_loss is not None else None, + val_loss=float(val_loss) if val_loss is not None else None, + **metrics, + ) + + def on_train_end(self, logs: Optional[Dict[str, Any]] = None): + """Called at the end of training.""" + if logs is None: + logs = {} + + # Log final metrics + self.monitor.log_custom_metric( + "training_complete", + { + "final_metrics": {k: float(v) if v is not None else None for k, v in logs.items()}, + }, + ) + + +class TorchMonitor(ModelMonitor): + """ + Monitor for PyTorch models (kept for backwards compatibility). + + Note: This is a placeholder. Use the PyTorch integration module instead. + """ + + def __init__(self, *args, **kwargs): + warnings.warn( + "TorchMonitor in tensorflow module is deprecated. " + "Use whiteboxai.integrations.pytorch.TorchMonitor instead.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) + + +def wrap_keras_model( + model: "keras.Model", + monitor: KerasMonitor, +) -> "keras.Model": + """ + Wrap a Keras model to automatically log predictions. + + Args: + model: Keras model to wrap + monitor: KerasMonitor instance + + Returns: + Wrapped model that logs predictions + + Example: + ```python + wrapped_model = wrap_keras_model(model, monitor) + predictions = wrapped_model.predict(X_test) # Automatically logged + ``` + """ + original_predict = model.predict + + def logged_predict(x, *args, **kwargs): + # Make prediction + predictions = original_predict(x, *args, **kwargs) + + # Log to WhiteBoxAI + try: + if isinstance(x, tf.Tensor): + x_np = x.numpy() + else: + x_np = x + + if isinstance(predictions, tf.Tensor): + pred_np = predictions.numpy() + else: + pred_np = predictions + + monitor.log_batch( + inputs=x_np, + predictions=pred_np, + ) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return predictions + + # Replace predict method + model.predict = logged_predict + return model + + +__all__ = [ + "KerasMonitor", + "WhiteBoxAICallback", + "TorchMonitor", # Deprecated + "wrap_keras_model", +] diff --git a/src/whiteboxai/integrations/transformers.py b/src/whiteboxai/integrations/transformers.py new file mode 100644 index 0000000..550a22d --- /dev/null +++ b/src/whiteboxai/integrations/transformers.py @@ -0,0 +1,531 @@ +""" +Hugging Face Transformers Integration + +Integration for monitoring Hugging Face Transformers models. +""" + +import warnings +from typing import Any, Dict, List, Optional, Union + +try: + import transformers + from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer + + TRANSFORMERS_AVAILABLE = True +except ImportError: + TRANSFORMERS_AVAILABLE = False + PreTrainedModel = object + PreTrainedTokenizer = object + Pipeline = object + +from whiteboxai.monitor import ModelMonitor + + +class TransformersMonitor(ModelMonitor): + """ + Monitor for Hugging Face Transformers models. + + Supports: + - Text classification + - Named entity recognition (NER) + - Question answering + - Text generation + - Translation + - Summarization + - Sentiment analysis + + Example: + ```python + from transformers import pipeline + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.transformers import TransformersMonitor + + # Load model + classifier = pipeline("sentiment-analysis") + + # Setup monitoring + client = WhiteBoxAI(api_key="your-api-key") + monitor = TransformersMonitor( + client=client, + pipeline=classifier, + model_name="sentiment_classifier" + ) + + # Make predictions with automatic logging + result = monitor.predict("I love this product!", log=True) + ``` + """ + + def __init__( + self, + client, + pipeline: Optional[Pipeline] = None, + model: Optional[PreTrainedModel] = None, + tokenizer: Optional[PreTrainedTokenizer] = None, + model_name: Optional[str] = None, + task: Optional[str] = None, + **kwargs, + ): + """ + Initialize Transformers monitor. + + Args: + client: WhiteBoxAI client instance + pipeline: Hugging Face pipeline (recommended) + model: PreTrainedModel (if not using pipeline) + tokenizer: PreTrainedTokenizer (if not using pipeline) + model_name: Name for the model + task: Task type (text-classification, ner, qa, etc.) + **kwargs: Additional arguments for ModelMonitor + """ + if not TRANSFORMERS_AVAILABLE: + raise ImportError( + "transformers is not installed. Install with: pip install transformers" + ) + + super().__init__(client, **kwargs) + self.pipeline = pipeline + self.model = model + self.tokenizer = tokenizer + self._model_name = model_name + self._task = task + self._registered = False + + # Auto-detect task from pipeline + if pipeline is not None and task is None: + self._task = getattr(pipeline, "task", None) + + def register_from_model( + self, + name: Optional[str] = None, + task: Optional[str] = None, + version: Optional[str] = None, + framework_version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register Transformers model with WhiteBoxAI. + + Args: + name: Model name (default: model class name or pipeline task) + task: Task type (text-classification, ner, qa, etc.) + version: Model version + framework_version: Transformers version + **kwargs: Additional metadata + + Returns: + Model ID + """ + if self.pipeline is None and self.model is None: + raise ValueError("No model or pipeline provided") + + # Auto-detect model name + if name is None: + if self._model_name: + name = self._model_name + elif self.pipeline is not None: + name = f"{self.pipeline.task}_pipeline" + else: + name = self.model.__class__.__name__ + + # Use provided or detected task + task = task or self._task or "text-classification" + + # Get Transformers version + if framework_version is None: + framework_version = transformers.__version__ + + # Extract model metadata + metadata = { + "framework": "transformers", + "framework_version": framework_version, + "task": task, + **kwargs, + } + + # Get model information from pipeline or model + if self.pipeline is not None: + metadata["pipeline_type"] = self.pipeline.task + if hasattr(self.pipeline, "model"): + metadata["model_class"] = self.pipeline.model.__class__.__name__ + if hasattr(self.pipeline.model, "config"): + config = self.pipeline.model.config + metadata["model_type"] = getattr(config, "model_type", None) + metadata["num_parameters"] = getattr(config, "num_parameters", None) + metadata["vocab_size"] = getattr(config, "vocab_size", None) + elif self.model is not None: + metadata["model_class"] = self.model.__class__.__name__ + if hasattr(self.model, "config"): + config = self.model.config + metadata["model_type"] = getattr(config, "model_type", None) + metadata["num_parameters"] = getattr(config, "num_parameters", None) + + # Map task to model_type + model_type_mapping = { + "text-classification": "classification", + "sentiment-analysis": "classification", + "ner": "classification", + "token-classification": "classification", + "question-answering": "qa", + "text-generation": "generation", + "translation": "generation", + "summarization": "generation", + } + model_type = model_type_mapping.get(task, "classification") + + # Register model + self.model_id = self.client.register_model( + name=name, + model_type=model_type, + framework="transformers", + version=version, + metadata=metadata, + ) + + self._registered = True + return self.model_id + + def predict( + self, + inputs: Union[str, List[str]], + log: bool = True, + actuals: Optional[Union[str, List[str]]] = None, + **kwargs: Any, + ) -> Union[Dict, List[Dict]]: + """ + Make predictions and optionally log them. + + Args: + inputs: Input text or list of texts + log: Whether to log the prediction + actuals: Actual values (if available) + **kwargs: Additional arguments for the pipeline/model + + Returns: + Predictions + """ + if self.pipeline is None and self.model is None: + raise ValueError("No model or pipeline provided") + + # Ensure model is registered + if not self._registered: + self.register_from_model() + + # Make prediction using pipeline + if self.pipeline is not None: + predictions = self.pipeline(inputs, **kwargs) + else: + # Use model + tokenizer + if self.tokenizer is None: + raise ValueError("Tokenizer required when using model directly") + + # Tokenize inputs + encoded = self.tokenizer( + inputs, return_tensors="pt", padding=True, truncation=True, **kwargs + ) + + # Get predictions + outputs = self.model(**encoded) + predictions = outputs.logits.detach().cpu().numpy() + + # Log if requested + if log: + # Determine if batch or single + is_batch = isinstance(inputs, list) + + if is_batch: + # Batch prediction + self.log_batch_transformers( + inputs=inputs, + predictions=predictions, + actuals=actuals, + ) + else: + # Single prediction + self.log_prediction_transformers( + input_text=inputs, + prediction=predictions if isinstance(predictions, dict) else predictions[0], + actual=actuals, + ) + + return predictions + + def log_prediction_transformers( + self, + input_text: str, + prediction: Union[Dict, Any], + actual: Optional[Any] = None, + **kwargs: Any, + ) -> None: + """ + Log a single Transformers prediction. + + Args: + input_text: Input text + prediction: Prediction result (dict or value) + actual: Actual value (if available) + **kwargs: Additional metadata + """ + # Extract prediction value based on task + if isinstance(prediction, dict): + pred_value = self._extract_prediction_value(prediction) + else: + pred_value = prediction + + # Log to WhiteBoxAI + self.log_prediction( + inputs={"text": input_text}, + prediction=pred_value, + actual=actual, + metadata={ + "task": self._task, + "raw_prediction": str(prediction), + **kwargs, + }, + ) + + def log_batch_transformers( + self, + inputs: List[str], + predictions: List[Union[Dict, Any]], + actuals: Optional[List[Any]] = None, + **kwargs: Any, + ) -> None: + """ + Log batch of Transformers predictions. + + Args: + inputs: List of input texts + predictions: List of predictions + actuals: List of actual values (if available) + **kwargs: Additional metadata + """ + # Prepare batch data + batch_inputs = [{"text": text} for text in inputs] + batch_predictions = [ + self._extract_prediction_value(pred) if isinstance(pred, dict) else pred + for pred in predictions + ] + + # Log batch + self.log_batch( + inputs=batch_inputs, + predictions=batch_predictions, + actuals=actuals, + metadata={ + "task": self._task, + **kwargs, + }, + ) + + def _extract_prediction_value(self, prediction: Dict) -> Any: + """Extract the main prediction value from a pipeline output.""" + if isinstance(prediction, list) and len(prediction) > 0: + prediction = prediction[0] + + if isinstance(prediction, dict): + # Common keys in transformers outputs + if "label" in prediction: + return prediction["label"] + elif "answer" in prediction: + return prediction["answer"] + elif "generated_text" in prediction: + return prediction["generated_text"] + elif "translation_text" in prediction: + return prediction["translation_text"] + elif "summary_text" in prediction: + return prediction["summary_text"] + elif "score" in prediction: + return prediction["score"] + + return prediction + + def set_baseline( + self, + baseline_texts: List[str], + baseline_labels: Optional[List[Any]] = None, + ) -> None: + """ + Set baseline data for drift detection. + + Args: + baseline_texts: Baseline input texts + baseline_labels: Baseline labels (optional) + """ + # Make baseline predictions if labels not provided + if baseline_labels is None and self.pipeline is not None: + baseline_predictions = self.pipeline(baseline_texts) + baseline_labels = [ + self._extract_prediction_value(pred) for pred in baseline_predictions + ] + + # Convert to format expected by parent class + baseline_inputs = [{"text": text} for text in baseline_texts] + + # Set baseline through parent class + super().set_baseline(baseline_inputs, baseline_labels) + + def log_generation_metrics( + self, + prompt: str, + generated_text: str, + num_tokens: Optional[int] = None, + generation_time: Optional[float] = None, + **kwargs: Any, + ) -> None: + """ + Log metrics for text generation tasks. + + Args: + prompt: Input prompt + generated_text: Generated text + num_tokens: Number of tokens generated + generation_time: Time taken to generate (seconds) + **kwargs: Additional metrics + """ + metrics = { + "prompt": prompt, + "generated_text": generated_text, + "num_tokens": num_tokens, + "generation_time": generation_time, + **kwargs, + } + + self.log_custom_metric("text_generation", metrics) + + +class TransformersPipelineWrapper: + """ + Wrapper for Hugging Face pipelines with automatic logging. + + Example: + ```python + from transformers import pipeline + from whiteboxai import WhiteBoxAI + from whiteboxai.integrations.transformers import TransformersPipelineWrapper + + classifier = pipeline("sentiment-analysis") + client = WhiteBoxAI(api_key="your-api-key") + + # Wrap pipeline + wrapped = TransformersPipelineWrapper(classifier, client) + + # Predictions are automatically logged + result = wrapped("I love this!") + ``` + """ + + def __init__( + self, + pipeline: Pipeline, + client, + model_name: Optional[str] = None, + auto_register: bool = True, + ): + """ + Initialize pipeline wrapper. + + Args: + pipeline: Hugging Face pipeline + client: WhiteBoxAI client + model_name: Model name + auto_register: Auto-register model on first prediction + """ + self.pipeline = pipeline + self.monitor = TransformersMonitor( + client=client, + pipeline=pipeline, + model_name=model_name, + ) + self._auto_register = auto_register + + def __call__(self, *args, **kwargs): + """Make prediction with automatic logging.""" + # Auto-register if needed + if self._auto_register and not self.monitor._registered: + self.monitor.register_from_model() + + # Get inputs + inputs = args[0] if args else kwargs.get("inputs") + + # Make prediction + result = self.pipeline(*args, **kwargs) + + # Log prediction + try: + if isinstance(inputs, list): + self.monitor.log_batch_transformers( + inputs=inputs, + predictions=result, + ) + else: + self.monitor.log_prediction_transformers( + input_text=inputs, + prediction=result, + ) + except Exception as e: + warnings.warn(f"Failed to log prediction: {e}") + + return result + + +def wrap_transformers_pipeline( + pipeline: Pipeline, + monitor: TransformersMonitor, +) -> Pipeline: + """ + Wrap a Hugging Face pipeline to automatically log predictions. + + Args: + pipeline: Hugging Face pipeline to wrap + monitor: TransformersMonitor instance + + Returns: + Wrapped pipeline that logs predictions + + Example: + ```python + from transformers import pipeline + + classifier = pipeline("sentiment-analysis") + wrapped = wrap_transformers_pipeline(classifier, monitor) + + # Predictions automatically logged + result = wrapped("Great product!") + ``` + """ + original_call = pipeline.__call__ + + def logged_call(*args, **kwargs): + # Make prediction + result = original_call(*args, **kwargs) + + # Log to WhiteBoxAI + try: + inputs = args[0] if args else kwargs.get("inputs") + + if isinstance(inputs, list): + monitor.log_batch_transformers( + inputs=inputs, + predictions=result, + ) + else: + monitor.log_prediction_transformers( + input_text=inputs, + prediction=result, + ) + except Exception as e: + warnings.warn(f"Failed to log predictions: {e}") + + return result + + # Replace __call__ method + pipeline.__call__ = logged_call + return pipeline + + +__all__ = [ + "TransformersMonitor", + "TransformersPipelineWrapper", + "wrap_transformers_pipeline", +] diff --git a/src/whiteboxai/models/__init__.py b/src/whiteboxai/models/__init__.py new file mode 100644 index 0000000..5fbfb4d --- /dev/null +++ b/src/whiteboxai/models/__init__.py @@ -0,0 +1 @@ +"""Models package for WhiteBoxAI SDK.""" diff --git a/src/whiteboxai/monitor.py b/src/whiteboxai/monitor.py new file mode 100644 index 0000000..0fb1460 --- /dev/null +++ b/src/whiteboxai/monitor.py @@ -0,0 +1,284 @@ +""" +Model Monitoring + +Simplified monitoring interface for ML models. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import numpy as np + +if TYPE_CHECKING: + from whiteboxai.client import WhiteBoxAI + + +class ModelMonitor: + """ + Simplified model monitoring interface. + + Example: + ```python + from whiteboxai import WhiteBoxAI, ModelMonitor + + client = WhiteBoxAI(api_key="your-api-key") + monitor = ModelMonitor(client, model_id=123) + + # Log prediction + monitor.log_prediction( + inputs={"feature1": 1.0, "feature2": 2.0}, + output=0.85 + ) + + # Log batch + monitor.log_batch([ + {"inputs": {...}, "output": 0.85}, + {"inputs": {...}, "output": 0.92}, + ]) + ``` + """ + + def __init__( + self, + client: "WhiteBoxAI", + model_id: Optional[int] = None, + model_name: Optional[str] = None, + auto_explain: bool = False, + sampling_rate: float = 1.0, + ): + """ + Initialize model monitor. + + Args: + client: WhiteBoxAI client instance + model_id: Model ID (if already registered) + model_name: Model name (for registration) + auto_explain: Automatically generate explanations + sampling_rate: Prediction sampling rate (0.0-1.0) + """ + self.client = client + self.model_id = model_id + self.model_name = model_name + self.auto_explain = auto_explain + self.sampling_rate = sampling_rate + self._baseline_data: Optional[np.ndarray] = None + + def register_model( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """ + Register a new model. + + Args: + name: Model name + model_type: Model type (classification, regression, etc.) + framework: ML framework + version: Model version + **kwargs: Additional metadata + + Returns: + Model ID + """ + result = self.client.models.register( + name=name, + model_type=model_type, + framework=framework, + version=version, + **kwargs, + ) + self.model_id = result["id"] + self.model_name = name + return self.model_id + + async def aregister_model( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + **kwargs: Any, + ) -> int: + """Async version of register_model().""" + result = await self.client.models.aregister( + name=name, + model_type=model_type, + framework=framework, + version=version, + **kwargs, + ) + self.model_id = result["id"] + self.model_name = name + return self.model_id + + def log_prediction( + self, + inputs: Any, + output: Any, + explain: Optional[bool] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Optional[Dict[str, Any]]: + """ + Log a single prediction. + + Args: + inputs: Input features/data + output: Model prediction/output + explain: Generate explanation (default: use auto_explain) + metadata: Additional metadata + + Returns: + Prediction data if sampled, None if skipped + """ + if not self._should_sample(): + return None + + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + explain = explain if explain is not None else self.auto_explain + + return self.client.predictions.log( + model_id=self.model_id, + inputs=inputs, + outputs=output, + explain=explain, + metadata=metadata, + ) + + async def alog_prediction( + self, + inputs: Any, + output: Any, + explain: Optional[bool] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> Optional[Dict[str, Any]]: + """Async version of log_prediction().""" + if not self._should_sample(): + return None + + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + explain = explain if explain is not None else self.auto_explain + + return await self.client.predictions.alog( + model_id=self.model_id, + inputs=inputs, + outputs=output, + explain=explain, + metadata=metadata, + ) + + def log_batch( + self, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """ + Log multiple predictions in batch. + + Args: + predictions: List of prediction dictionaries + + Returns: + Batch logging result + """ + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + # Apply sampling + if self.sampling_rate < 1.0: + predictions = self._sample_predictions(predictions) + + return self.client.predictions.log_batch( + model_id=self.model_id, + predictions=predictions, + ) + + async def alog_batch( + self, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """Async version of log_batch().""" + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + # Apply sampling + if self.sampling_rate < 1.0: + predictions = self._sample_predictions(predictions) + + return await self.client.predictions.alog_batch( + model_id=self.model_id, + predictions=predictions, + ) + + def set_baseline(self, data: np.ndarray) -> None: + """ + Set baseline data for drift detection. + + Args: + data: Baseline data array + """ + self._baseline_data = data + + def detect_drift( + self, + current_data: Optional[np.ndarray] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Detect drift for the model. + + Args: + current_data: Current data to compare against baseline + **kwargs: Additional detection parameters + + Returns: + Drift detection results + """ + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + return self.client.drift.detect( + model_id=self.model_id, + reference_data=( + self._baseline_data.tolist() if self._baseline_data is not None else None + ), + current_data=current_data.tolist() if current_data is not None else None, + **kwargs, + ) + + async def adetect_drift( + self, + current_data: Optional[np.ndarray] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of detect_drift().""" + if self.model_id is None: + raise ValueError("Model not registered. Call register_model() first.") + + return await self.client.drift.adetect( + model_id=self.model_id, + reference_data=( + self._baseline_data.tolist() if self._baseline_data is not None else None + ), + current_data=current_data.tolist() if current_data is not None else None, + **kwargs, + ) + + def _should_sample(self) -> bool: + """Check if prediction should be sampled.""" + if self.sampling_rate >= 1.0: + return True + return np.random.random() < self.sampling_rate + + def _sample_predictions(self, predictions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Sample predictions based on sampling rate.""" + n_samples = int(len(predictions) * self.sampling_rate) + if n_samples == 0: + n_samples = 1 # Always log at least one + indices = np.random.choice(len(predictions), n_samples, replace=False) + return [predictions[i] for i in indices] diff --git a/src/whiteboxai/offline.py b/src/whiteboxai/offline.py new file mode 100644 index 0000000..903a77b --- /dev/null +++ b/src/whiteboxai/offline.py @@ -0,0 +1,541 @@ +""" +Offline mode support for WhiteBoxAI SDK. + +This module provides offline queueing capabilities for the SDK, allowing +operations to be queued locally when the API is unavailable and synced +when connectivity is restored. + +Features: +- Persistent queue storage (SQLite-based) +- Automatic retry on connection restoration +- Configurable queue limits +- Operation prioritization +- Failed operation tracking + +Example: + Enable offline mode: + + >>> from whiteboxai import WhiteBoxAI + >>> client = WhiteBoxAI( + ... api_key="your-api-key", + ... enable_offline=True, + ... offline_dir="./whiteboxai_offline" + ... ) + >>> + >>> # Operations are queued when offline + >>> client.predict(...) # Queued if API unavailable + >>> + >>> # Manually sync + >>> client.sync_offline_queue() +""" + +import json +import logging +import os +import sqlite3 +import threading +from enum import Enum +from pathlib import Path +from typing import Any, Dict, List, Tuple + +logger = logging.getLogger(__name__) + + +class OperationType(Enum): + """Types of operations that can be queued.""" + + PREDICT = "predict" + REGISTER_MODEL = "register_model" + UPDATE_BASELINE = "update_baseline" + LOG_BATCH = "log_batch" + + +class OperationPriority(Enum): + """Priority levels for queued operations.""" + + LOW = 1 + NORMAL = 2 + HIGH = 3 + CRITICAL = 4 + + +class OfflineQueue: + """ + Persistent queue for offline operations. + + Uses SQLite for storage to ensure operations are not lost even if + the application crashes or restarts. + + Attributes: + db_path: Path to the SQLite database file + max_queue_size: Maximum number of operations to queue + auto_sync: Whether to automatically sync when connection available + """ + + def __init__(self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True): + """ + Initialize offline queue. + + Args: + db_path: Path to SQLite database file + max_queue_size: Maximum operations to queue (0 = unlimited) + auto_sync: Enable automatic syncing + """ + self.db_path = db_path + self.max_queue_size = max_queue_size + self.auto_sync = auto_sync + self._lock = threading.Lock() + self._ensure_db() + + def _ensure_db(self): + """Create database and tables if they don't exist.""" + # Ensure directory exists + os.makedirs(os.path.dirname(self.db_path), exist_ok=True) + + with sqlite3.connect(self.db_path) as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS queue ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + operation_type TEXT NOT NULL, + priority INTEGER NOT NULL, + data TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + retry_count INTEGER DEFAULT 0, + last_error TEXT, + status TEXT DEFAULT 'pending' + ) + """ + ) + + # Create index for efficient querying + conn.execute( + """ + CREATE INDEX IF NOT EXISTS idx_status_priority + ON queue(status, priority DESC, created_at ASC) + """ + ) + + conn.commit() + + def enqueue( + self, + operation_type: OperationType, + data: Dict[str, Any], + priority: OperationPriority = OperationPriority.NORMAL, + ) -> int: + """ + Add an operation to the queue. + + Args: + operation_type: Type of operation + data: Operation data (will be JSON serialized) + priority: Operation priority + + Returns: + Operation ID + + Raises: + ValueError: If queue is full + """ + with self._lock: + # Check queue size + if self.max_queue_size > 0: + current_size = self.get_queue_size() + if current_size >= self.max_queue_size: + raise ValueError(f"Queue is full ({current_size}/{self.max_queue_size})") + + # Insert operation + with sqlite3.connect(self.db_path) as conn: + cursor = conn.execute( + """ + INSERT INTO queue (operation_type, priority, data) + VALUES (?, ?, ?) + """, + (operation_type.value, priority.value, json.dumps(data)), + ) + conn.commit() + op_id = cursor.lastrowid + + logger.info(f"Queued operation {op_id}: {operation_type.value}") + return op_id + + def dequeue(self, limit: int = 100) -> List[Tuple[int, OperationType, Dict[str, Any]]]: + """ + Get pending operations from queue. + + Returns operations ordered by priority (highest first) and + creation time (oldest first). + + Args: + limit: Maximum number of operations to return + + Returns: + List of (id, operation_type, data) tuples + """ + with self._lock: + with sqlite3.connect(self.db_path) as conn: + cursor = conn.execute( + """ + SELECT id, operation_type, data + FROM queue + WHERE status = 'pending' + ORDER BY priority DESC, created_at ASC + LIMIT ? + """, + (limit,), + ) + + operations = [] + for row in cursor.fetchall(): + op_id, op_type, data_json = row + operations.append((op_id, OperationType(op_type), json.loads(data_json))) + + return operations + + def mark_success(self, operation_id: int): + """ + Mark operation as successfully completed. + + Args: + operation_id: Operation ID + """ + with self._lock: + with sqlite3.connect(self.db_path) as conn: + conn.execute( + """ + UPDATE queue + SET status = 'completed' + WHERE id = ? + """, + (operation_id,), + ) + conn.commit() + + logger.debug(f"Marked operation {operation_id} as completed") + + def mark_failure(self, operation_id: int, error: str, max_retries: int = 3): + """ + Mark operation as failed and increment retry count. + + Args: + operation_id: Operation ID + error: Error message + max_retries: Maximum retry attempts + """ + with self._lock: + with sqlite3.connect(self.db_path) as conn: + # Increment retry count + cursor = conn.execute( + """ + UPDATE queue + SET retry_count = retry_count + 1, + last_error = ? + WHERE id = ? + """, + (error, operation_id), + ) + + # Check if max retries exceeded + cursor = conn.execute( + """ + SELECT retry_count FROM queue WHERE id = ? + """, + (operation_id,), + ) + retry_count = cursor.fetchone()[0] + + if retry_count >= max_retries: + conn.execute( + """ + UPDATE queue + SET status = 'failed' + WHERE id = ? + """, + (operation_id,), + ) + logger.error( + f"Operation {operation_id} permanently failed after " + f"{retry_count} retries: {error}" + ) + else: + logger.warning( + f"Operation {operation_id} failed (retry {retry_count}/" + f"{max_retries}): {error}" + ) + + conn.commit() + + def get_queue_size(self, status: str = "pending") -> int: + """ + Get number of operations in queue. + + Args: + status: Filter by status ('pending', 'completed', 'failed', or None for all) + + Returns: + Number of operations + """ + with sqlite3.connect(self.db_path) as conn: + if status: + cursor = conn.execute("SELECT COUNT(*) FROM queue WHERE status = ?", (status,)) + else: + cursor = conn.execute("SELECT COUNT(*) FROM queue") + + return cursor.fetchone()[0] + + def get_statistics(self) -> Dict[str, int]: + """ + Get queue statistics. + + Returns: + Dictionary with counts by status + """ + with sqlite3.connect(self.db_path) as conn: + cursor = conn.execute( + """ + SELECT status, COUNT(*) as count + FROM queue + GROUP BY status + """ + ) + + stats = {"total": 0, "pending": 0, "completed": 0, "failed": 0} + + for row in cursor.fetchall(): + status, count = row + stats[status] = count + stats["total"] += count + + return stats + + def clear_completed(self, older_than_days: int = 7): + """ + Remove completed operations older than specified days. + + Args: + older_than_days: Remove operations older than this many days + """ + with self._lock: + with sqlite3.connect(self.db_path) as conn: + cursor = conn.execute( + """ + DELETE FROM queue + WHERE status = 'completed' + AND created_at < datetime('now', '-' || ? || ' days') + """, + (older_than_days,), + ) + deleted = cursor.rowcount + conn.commit() + + if deleted > 0: + logger.info(f"Cleared {deleted} completed operations older than {older_than_days} days") + + def clear_all(self): + """Clear all operations from queue (use with caution).""" + with self._lock: + with sqlite3.connect(self.db_path) as conn: + conn.execute("DELETE FROM queue") + conn.commit() + + logger.warning("Cleared all operations from queue") + + def get_failed_operations(self) -> List[Dict[str, Any]]: + """ + Get all permanently failed operations. + + Returns: + List of failed operation details + """ + with sqlite3.connect(self.db_path) as conn: + cursor = conn.execute( + """ + SELECT id, operation_type, data, created_at, retry_count, last_error + FROM queue + WHERE status = 'failed' + ORDER BY created_at DESC + """ + ) + + failed = [] + for row in cursor.fetchall(): + op_id, op_type, data_json, created_at, retry_count, last_error = row + failed.append( + { + "id": op_id, + "operation_type": op_type, + "data": json.loads(data_json), + "created_at": created_at, + "retry_count": retry_count, + "last_error": last_error, + } + ) + + return failed + + +class OfflineManager: + """ + Manages offline mode for WhiteBoxAI SDK. + + Handles automatic queueing when offline and syncing when online. + + Attributes: + queue: OfflineQueue instance + sync_interval: Seconds between automatic sync attempts + max_retries: Maximum retry attempts for failed operations + """ + + def __init__( + self, + offline_dir: str = "./whiteboxai_offline", + max_queue_size: int = 10000, + auto_sync: bool = True, + sync_interval: int = 60, + max_retries: int = 3, + ): + """ + Initialize offline manager. + + Args: + offline_dir: Directory for offline storage + max_queue_size: Maximum operations to queue + auto_sync: Enable automatic syncing + sync_interval: Seconds between sync attempts + max_retries: Maximum retry attempts + """ + self.offline_dir = Path(offline_dir) + self.offline_dir.mkdir(parents=True, exist_ok=True) + + db_path = str(self.offline_dir / "queue.db") + self.queue = OfflineQueue( + db_path=db_path, max_queue_size=max_queue_size, auto_sync=auto_sync + ) + + self.sync_interval = sync_interval + self.max_retries = max_retries + self._sync_thread = None + self._stop_sync = threading.Event() + self._client = None # Will be set by WhiteBoxAI client + + if auto_sync: + self.start_auto_sync() + + def set_client(self, client): + """ + Set the WhiteBoxAI client for syncing. + + Args: + client: WhiteBoxAI client instance + """ + self._client = client + + def start_auto_sync(self): + """Start automatic sync thread.""" + if self._sync_thread is None or not self._sync_thread.is_alive(): + self._stop_sync.clear() + self._sync_thread = threading.Thread(target=self._auto_sync_loop, daemon=True) + self._sync_thread.start() + logger.info("Started automatic sync thread") + + def stop_auto_sync(self): + """Stop automatic sync thread.""" + if self._sync_thread and self._sync_thread.is_alive(): + self._stop_sync.set() + self._sync_thread.join(timeout=5) + logger.info("Stopped automatic sync thread") + + def _auto_sync_loop(self): + """Background thread for automatic syncing.""" + while not self._stop_sync.is_set(): + try: + if self._client: + self.sync(batch_size=100) + except Exception as e: + logger.error(f"Error in auto-sync: {e}") + + # Wait for next sync interval + self._stop_sync.wait(self.sync_interval) + + def sync(self, batch_size: int = 100) -> Dict[str, int]: + """ + Sync queued operations with API. + + Args: + batch_size: Number of operations to sync per batch + + Returns: + Dictionary with sync statistics + """ + if not self._client: + logger.warning("No client set, cannot sync") + return {"synced": 0, "failed": 0, "pending": self.queue.get_queue_size()} + + stats = {"synced": 0, "failed": 0} + + # Get pending operations + operations = self.queue.dequeue(limit=batch_size) + + if not operations: + return {**stats, "pending": 0} + + logger.info(f"Syncing {len(operations)} operations...") + + for op_id, op_type, data in operations: + try: + # Execute operation + if op_type == OperationType.PREDICT: + self._client._api_predict(**data) + elif op_type == OperationType.REGISTER_MODEL: + self._client._api_register_model(**data) + elif op_type == OperationType.UPDATE_BASELINE: + self._client._api_update_baseline(**data) + elif op_type == OperationType.LOG_BATCH: + self._client._api_log_batch(**data) + + # Mark success + self.queue.mark_success(op_id) + stats["synced"] += 1 + + except Exception as e: + # Mark failure + error_msg = str(e) + self.queue.mark_failure(op_id, error_msg, self.max_retries) + stats["failed"] += 1 + + stats["pending"] = self.queue.get_queue_size() + + if stats["synced"] > 0: + logger.info( + f"Synced {stats['synced']} operations, {stats['failed']} failed, {stats['pending']} pending" + ) + + return stats + + def get_status(self) -> Dict[str, Any]: + """ + Get offline mode status. + + Returns: + Status dictionary with queue stats and sync info + """ + stats = self.queue.get_statistics() + + return { + "enabled": True, + "auto_sync": self.queue.auto_sync, + "queue_stats": stats, + "sync_interval": self.sync_interval, + "offline_dir": str(self.offline_dir), + "auto_sync_running": self._sync_thread and self._sync_thread.is_alive(), + } + + def cleanup(self, older_than_days: int = 7): + """ + Clean up old completed operations. + + Args: + older_than_days: Remove operations older than this many days + """ + self.queue.clear_completed(older_than_days) diff --git a/src/whiteboxai/privacy.py b/src/whiteboxai/privacy.py new file mode 100644 index 0000000..86c97b8 --- /dev/null +++ b/src/whiteboxai/privacy.py @@ -0,0 +1,219 @@ +""" +Privacy Utilities + +PII detection and data masking utilities. +""" + +import re +from typing import Any, Dict, List, Optional, Pattern + + +class PIIDetector: + """ + Detect personally identifiable information (PII) in data. + + Supports: + - Email addresses + - Phone numbers + - Credit card numbers + - Social security numbers + - IP addresses + """ + + def __init__(self): + """Initialize PII detector with regex patterns.""" + self.patterns: Dict[str, Pattern] = { + "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"), + "phone": re.compile(r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), + "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), + "credit_card": re.compile(r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b"), + "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), + } + + def detect(self, text: str) -> List[Dict[str, Any]]: + """ + Detect PII in text. + + Args: + text: Text to scan for PII + + Returns: + List of detected PII items with type and location + """ + detections = [] + + for pii_type, pattern in self.patterns.items(): + for match in pattern.finditer(text): + detections.append( + { + "type": pii_type, + "value": match.group(), + "start": match.start(), + "end": match.end(), + } + ) + + return detections + + def mask( + self, + text: str, + mask_char: str = "*", + preserve_length: bool = True, + ) -> str: + """ + Mask PII in text. + + Args: + text: Text to mask + mask_char: Character to use for masking + preserve_length: Whether to preserve original length + + Returns: + Masked text + """ + detections = self.detect(text) + + # Sort by position (reverse order to avoid index shifting) + detections.sort(key=lambda x: x["start"], reverse=True) + + # Mask each detection + for detection in detections: + start = detection["start"] + end = detection["end"] + original = detection["value"] + + if preserve_length: + masked = mask_char * len(original) + else: + masked = f"[{detection['type'].upper()}]" + + text = text[:start] + masked + text[end:] + + return text + + +class DataMasker: + """ + Mask sensitive data in structured formats (dict, list). + """ + + def __init__(self, pii_detector: Optional[PIIDetector] = None): + """ + Initialize data masker. + + Args: + pii_detector: PII detector instance (creates new if None) + """ + self.pii_detector = pii_detector or PIIDetector() + self.sensitive_keys = { + "password", + "passwd", + "pwd", + "secret", + "token", + "api_key", + "apikey", + "access_token", + "refresh_token", + } + + def mask( + self, + data: Any, + mask_pii: bool = True, + mask_sensitive_keys: bool = True, + ) -> Any: + """ + Mask sensitive data. + + Args: + data: Data to mask (dict, list, or primitive) + mask_pii: Whether to mask PII in strings + mask_sensitive_keys: Whether to mask values of sensitive keys + + Returns: + Masked data + """ + if isinstance(data, dict): + return self._mask_dict(data, mask_pii, mask_sensitive_keys) + elif isinstance(data, list): + return self._mask_list(data, mask_pii, mask_sensitive_keys) + elif isinstance(data, str): + if mask_pii: + return self.pii_detector.mask(data) + return data + else: + return data + + def _mask_dict( + self, + data: dict, + mask_pii: bool, + mask_sensitive_keys: bool, + ) -> dict: + """Mask sensitive data in dictionary.""" + masked = {} + for key, value in data.items(): + # Check if key is sensitive + if mask_sensitive_keys and key.lower() in self.sensitive_keys: + masked[key] = "***MASKED***" + else: + masked[key] = self.mask(value, mask_pii, mask_sensitive_keys) + return masked + + def _mask_list( + self, + data: list, + mask_pii: bool, + mask_sensitive_keys: bool, + ) -> list: + """Mask sensitive data in list.""" + return [self.mask(item, mask_pii, mask_sensitive_keys) for item in data] + + +# Global instances +_pii_detector = PIIDetector() +_data_masker = DataMasker(_pii_detector) + + +def detect_pii(text: str) -> List[Dict[str, Any]]: + """ + Detect PII in text using global detector. + + Args: + text: Text to scan + + Returns: + List of detected PII items + """ + return _pii_detector.detect(text) + + +def mask_pii(text: str, mask_char: str = "*") -> str: + """ + Mask PII in text using global detector. + + Args: + text: Text to mask + mask_char: Character to use for masking + + Returns: + Masked text + """ + return _pii_detector.mask(text, mask_char=mask_char) + + +def mask_data(data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True) -> Any: + """ + Mask sensitive data using global masker. + + Args: + data: Data to mask + mask_pii: Whether to mask PII + mask_sensitive_keys: Whether to mask sensitive keys + + Returns: + Masked data + """ + return _data_masker.mask(data, mask_pii, mask_sensitive_keys) diff --git a/src/whiteboxai/resources.py b/src/whiteboxai/resources.py new file mode 100644 index 0000000..d0908a4 --- /dev/null +++ b/src/whiteboxai/resources.py @@ -0,0 +1,460 @@ +""" +API Resources + +Resource classes for interacting with different WhiteBoxAI API endpoints. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +if TYPE_CHECKING: + from whiteboxai.client import WhiteBoxAI + + +class BaseResource: + """Base class for API resources.""" + + def __init__(self, client: "WhiteBoxAI"): + """Initialize resource with client.""" + self.client = client + + +class ModelsResource(BaseResource): + """Models API resource.""" + + def register( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Register a new model. + + Args: + name: Model name + model_type: Model type (classification, regression, etc.) + framework: ML framework (sklearn, pytorch, tensorflow, etc.) + version: Model version + **kwargs: Additional model metadata + + Returns: + Registered model data + """ + data = { + "name": name, + "model_type": model_type, + "framework": framework, + "version": version, + **kwargs, + } + return self.client.request("POST", "/api/v1/models", data=data) + + async def aregister( + self, + name: str, + model_type: str, + framework: Optional[str] = None, + version: Optional[str] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of register().""" + data = { + "name": name, + "model_type": model_type, + "framework": framework, + "version": version, + **kwargs, + } + return await self.client.arequest("POST", "/api/v1/models", data=data) + + def get(self, model_id: int) -> Dict[str, Any]: + """Get model by ID.""" + return self.client.request("GET", f"/api/v1/models/{model_id}") + + async def aget(self, model_id: int) -> Dict[str, Any]: + """Async version of get().""" + return await self.client.arequest("GET", f"/api/v1/models/{model_id}") + + def list(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: + """List registered models.""" + params = {"limit": limit, "offset": offset} + return self.client.request("GET", "/api/v1/models", params=params) + + async def alist(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: + """Async version of list().""" + params = {"limit": limit, "offset": offset} + return await self.client.arequest("GET", "/api/v1/models", params=params) + + def update(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: + """Update model metadata.""" + return self.client.request("PUT", f"/api/v1/models/{model_id}", data=kwargs) + + async def aupdate(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: + """Async version of update().""" + return await self.client.arequest("PUT", f"/api/v1/models/{model_id}", data=kwargs) + + +class PredictionsResource(BaseResource): + """Predictions API resource.""" + + def log( + self, + model_id: int, + inputs: Any, + outputs: Any, + explain: bool = False, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """ + Log a single prediction. + + Args: + model_id: Model ID + inputs: Input features/data + outputs: Model prediction/output + explain: Whether to generate explanation + metadata: Additional prediction metadata + + Returns: + Logged prediction data + """ + data = { + "model_id": model_id, + "inputs": inputs, + "outputs": outputs, + "explain": explain, + "metadata": metadata or {}, + } + return self.client.request("POST", "/api/v1/predictions/log", data=data) + + async def alog( + self, + model_id: int, + inputs: Any, + outputs: Any, + explain: bool = False, + metadata: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Async version of log().""" + data = { + "model_id": model_id, + "inputs": inputs, + "outputs": outputs, + "explain": explain, + "metadata": metadata or {}, + } + return await self.client.arequest("POST", "/api/v1/predictions/log", data=data) + + def log_batch( + self, + model_id: int, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """ + Log multiple predictions in batch. + + Args: + model_id: Model ID + predictions: List of prediction dictionaries + + Returns: + Batch logging result + """ + data = {"model_id": model_id, "predictions": predictions} + return self.client.request("POST", "/api/v1/predictions/log/batch", data=data) + + async def alog_batch( + self, + model_id: int, + predictions: List[Dict[str, Any]], + ) -> Dict[str, Any]: + """Async version of log_batch().""" + data = {"model_id": model_id, "predictions": predictions} + return await self.client.arequest("POST", "/api/v1/predictions/log/batch", data=data) + + +class ExplanationsResource(BaseResource): + """Explanations API resource.""" + + def generate( + self, + prediction_id: int, + method: str = "shap", + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Generate explanation for a prediction. + + Args: + prediction_id: Prediction ID + method: Explanation method (shap, lime) + **kwargs: Method-specific parameters + + Returns: + Explanation data + """ + data = {"prediction_id": prediction_id, "method": method, **kwargs} + return self.client.request("POST", "/api/v1/explanations/generate", data=data) + + async def agenerate( + self, + prediction_id: int, + method: str = "shap", + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of generate().""" + data = {"prediction_id": prediction_id, "method": method, **kwargs} + return await self.client.arequest("POST", "/api/v1/explanations/generate", data=data) + + def get(self, explanation_id: int) -> Dict[str, Any]: + """Get explanation by ID.""" + return self.client.request("GET", f"/api/v1/explanations/{explanation_id}") + + async def aget(self, explanation_id: int) -> Dict[str, Any]: + """Async version of get().""" + return await self.client.arequest("GET", f"/api/v1/explanations/{explanation_id}") + + +class DriftResource(BaseResource): + """Drift detection API resource.""" + + def detect( + self, + model_id: int, + reference_data: Optional[Any] = None, + current_data: Optional[Any] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Detect drift for a model. + + Args: + model_id: Model ID + reference_data: Reference/baseline data + current_data: Current data to compare + **kwargs: Detection parameters + + Returns: + Drift detection results + """ + data = { + "model_id": model_id, + "reference_data": reference_data, + "current_data": current_data, + **kwargs, + } + return self.client.request("POST", "/api/v1/drift/detect", data=data) + + async def adetect( + self, + model_id: int, + reference_data: Optional[Any] = None, + current_data: Optional[Any] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of detect().""" + data = { + "model_id": model_id, + "reference_data": reference_data, + "current_data": current_data, + **kwargs, + } + return await self.client.arequest("POST", "/api/v1/drift/detect", data=data) + + def get_report(self, model_id: int, report_id: int) -> Dict[str, Any]: + """Get drift report.""" + return self.client.request("GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}") + + async def aget_report(self, model_id: int, report_id: int) -> Dict[str, Any]: + """Async version of get_report().""" + return await self.client.arequest( + "GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}" + ) + + +class AlertsResource(BaseResource): + """Alerts API resource.""" + + def create( + self, + name: str, + alert_type: str, + conditions: Dict[str, Any], + **kwargs: Any, + ) -> Dict[str, Any]: + """ + Create an alert rule. + + Args: + name: Alert name + alert_type: Alert type + conditions: Alert conditions + **kwargs: Additional alert configuration + + Returns: + Created alert data + """ + data = { + "name": name, + "alert_type": alert_type, + "conditions": conditions, + **kwargs, + } + return self.client.request("POST", "/api/v1/alerts", data=data) + + async def acreate( + self, + name: str, + alert_type: str, + conditions: Dict[str, Any], + **kwargs: Any, + ) -> Dict[str, Any]: + """Async version of create().""" + data = { + "name": name, + "alert_type": alert_type, + "conditions": conditions, + **kwargs, + } + return await self.client.arequest("POST", "/api/v1/alerts", data=data) + + def list(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: + """List alert rules.""" + params = {"model_id": model_id} if model_id else {} + return self.client.request("GET", "/api/v1/alerts", params=params) + + async def alist(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: + """Async version of list().""" + params = {"model_id": model_id} if model_id else {} + return await self.client.arequest("GET", "/api/v1/alerts", params=params) + + +class AgentWorkflowsResource(BaseResource): + """Agent Workflows API resource for multi-agent monitoring.""" + + def create( + self, + name: str, + framework: str, + inputs: Optional[Dict[str, Any]] = None, + meta_data: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Create a new agent workflow.""" + data = { + "name": name, + "framework": framework, + "inputs": inputs or {}, + "meta_data": meta_data or {}, + } + return self.client.request("POST", "/api/v1/workflows/multi-agent", data=data) + + def start( + self, + workflow_id: str, + inputs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Start an agent workflow.""" + return self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{workflow_id}/start", + data={"inputs": inputs or {}}, + ) + + def complete( + self, + workflow_id: str, + outputs: Optional[Dict[str, Any]] = None, + status: str = "completed", + ) -> Dict[str, Any]: + """Complete an agent workflow.""" + return self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{workflow_id}/complete", + data={"outputs": outputs or {}, "status": status}, + ) + + def register_agent( + self, + workflow_id: str, + name: str, + role: Optional[str] = None, + model_name: Optional[str] = None, + tools: Optional[List[str]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Register an agent in a workflow.""" + data = { + "name": name, + "role": role or name, + "model_name": model_name, + "tools": tools or [], + **kwargs, + } + return self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{workflow_id}/agents", + data=data, + ) + + def create_execution( + self, + workflow_id: str, + agent_name: str, + status: str, + inputs: Optional[Dict[str, Any]] = None, + outputs: Optional[Dict[str, Any]] = None, + duration_ms: Optional[int] = None, + llm_call_count: Optional[int] = None, + tool_call_count: Optional[int] = None, + tokens_used: Optional[int] = None, + cost: Optional[float] = None, + ) -> Dict[str, Any]: + """Log an agent execution within a workflow.""" + data = { + "agent_name": agent_name, + "status": status, + "inputs": inputs, + "outputs": outputs, + "duration_ms": duration_ms, + "llm_call_count": llm_call_count, + "tool_call_count": tool_call_count, + "tokens_used": tokens_used, + "cost": cost, + } + return self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{workflow_id}/executions", + data=data, + ) + + def create_interaction( + self, + workflow_id: str, + from_agent: str, + to_agent: str, + interaction_type: str, + message: Optional[str] = None, + meta_data: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Log an agent-to-agent interaction.""" + data = { + "from_agent": from_agent, + "to_agent": to_agent, + "interaction_type": interaction_type, + "message": message, + "meta_data": meta_data or {}, + } + return self.client.request( + "POST", + f"/api/v1/workflows/multi-agent/{workflow_id}/interactions", + data=data, + ) + + def get_analytics(self, workflow_id: str) -> Dict[str, Any]: + """Get analytics for a workflow.""" + return self.client.request( + "GET", + f"/api/v1/workflows/multi-agent/{workflow_id}/analytics", + ) diff --git a/src/whiteboxai/utils.py b/src/whiteboxai/utils.py new file mode 100644 index 0000000..5c41860 --- /dev/null +++ b/src/whiteboxai/utils.py @@ -0,0 +1,337 @@ +""" +SDK Utilities + +General utility functions for the WhiteBoxAI SDK. +""" + +import json +from typing import Any, Dict, List, Union + +import numpy as np +import pandas as pd + + +def serialize_data(data: Any) -> Any: + """ + Serialize data for API transmission. + + Handles numpy arrays, pandas DataFrames, and other common data types. + + Args: + data: Data to serialize + + Returns: + JSON-serializable data + """ + if isinstance(data, np.ndarray): + return data.tolist() + elif isinstance(data, pd.DataFrame): + return data.to_dict(orient="records") + elif isinstance(data, pd.Series): + return data.to_dict() + elif isinstance(data, (np.integer, np.floating)): + return data.item() + elif isinstance(data, dict): + return {k: serialize_data(v) for k, v in data.items()} + elif isinstance(data, (list, tuple)): + return [serialize_data(item) for item in data] + else: + return data + + +def validate_model_type(model_type: str) -> bool: + """ + Validate model type. + + Args: + model_type: Model type string + + Returns: + True if valid, False otherwise + """ + valid_types = { + "classification", + "regression", + "clustering", + "ranking", + "recommendation", + "anomaly_detection", + "time_series", + "nlp", + "computer_vision", + "other", + } + return model_type.lower() in valid_types + + +def extract_feature_names(data: Union[np.ndarray, pd.DataFrame, Dict]) -> List[str]: + """ + Extract feature names from data. + + Args: + data: Input data + + Returns: + List of feature names + """ + if isinstance(data, pd.DataFrame): + return data.columns.tolist() + elif isinstance(data, dict): + return list(data.keys()) + elif isinstance(data, np.ndarray): + # Generate default feature names + n_features = data.shape[1] if len(data.shape) > 1 else 1 + return [f"feature_{i}" for i in range(n_features)] + else: + return [] + + +def format_bytes(size: int) -> str: + """ + Format byte size as human-readable string. + + Args: + size: Size in bytes + + Returns: + Formatted string (e.g., "1.5 MB") + """ + for unit in ["B", "KB", "MB", "GB", "TB"]: + if size < 1024.0: + return f"{size:.1f} {unit}" + size /= 1024.0 + return f"{size:.1f} PB" + + +def truncate_string(text: str, max_length: int = 100, suffix: str = "...") -> str: + """ + Truncate string to maximum length. + + Args: + text: String to truncate + max_length: Maximum length + suffix: Suffix to append if truncated + + Returns: + Truncated string + """ + if len(text) <= max_length: + return text + return text[: max_length - len(suffix)] + suffix + + +def dict_to_query_params(params: Dict[str, Any]) -> str: + """ + Convert dictionary to URL query parameters. + + Args: + params: Parameter dictionary + + Returns: + Query string + """ + parts = [] + for key, value in params.items(): + if value is not None: + if isinstance(value, (list, tuple)): + for item in value: + parts.append(f"{key}={item}") + else: + parts.append(f"{key}={value}") + return "&".join(parts) + + +def deep_merge(dict1: Dict, dict2: Dict) -> Dict: + """ + Deep merge two dictionaries. + + Args: + dict1: First dictionary + dict2: Second dictionary (takes precedence) + + Returns: + Merged dictionary + """ + result = dict1.copy() + for key, value in dict2.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = deep_merge(result[key], value) + else: + result[key] = value + return result + + +def safe_json_loads(text: str, default: Any = None) -> Any: + """ + Safely load JSON with fallback. + + Args: + text: JSON string + default: Default value if parsing fails + + Returns: + Parsed JSON or default value + """ + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError): + return default + + +def is_numeric(value: Any) -> bool: + """ + Check if value is numeric. + + Args: + value: Value to check + + Returns: + True if numeric, False otherwise + """ + try: + float(value) + return True + except (ValueError, TypeError): + return False + + +def validate_features(features: Any) -> bool: + """ + Validate features for model input. + + Args: + features: Features to validate + + Returns: + True if valid + + Raises: + ValueError: If features are invalid + """ + if features is None: + raise ValueError("Features cannot be None") + + if isinstance(features, (list, tuple)): + if len(features) == 0: + raise ValueError("Features cannot be empty") + return True + elif isinstance(features, np.ndarray): + if features.size == 0: + raise ValueError("Features cannot be empty") + return True + elif isinstance(features, dict): + if len(features) == 0: + raise ValueError("Features cannot be empty") + return True + else: + raise ValueError(f"Unsupported features type: {type(features)}") + + +def format_features(features: Any) -> Dict[str, Any]: + """ + Format features for API transmission. + + Args: + features: Features to format + + Returns: + Formatted features dictionary + """ + if isinstance(features, dict): + return features + elif isinstance(features, (list, tuple, np.ndarray)): + return {"features": serialize_data(features)} + else: + return {"features": features} + + +def validate_prediction(prediction: Any) -> bool: + """ + Validate model prediction. + + Args: + prediction: Prediction to validate + + Returns: + True if valid + + Raises: + ValueError: If prediction is invalid + """ + if prediction is None: + raise ValueError("Prediction cannot be None") + return True + + +def serialize_numpy(arr: np.ndarray) -> List: + """ + Serialize numpy array to list. + + Args: + arr: Numpy array + + Returns: + List representation + """ + return arr.tolist() + + +def deserialize_numpy(data: List) -> np.ndarray: + """ + Deserialize list to numpy array. + + Args: + data: List data + + Returns: + Numpy array + """ + return np.array(data) + + +def compute_metrics(y_true: Any, y_pred: Any) -> Dict[str, float]: + """ + Compute basic metrics for predictions. + + Args: + y_true: True labels + y_pred: Predicted labels + + Returns: + Dictionary of metrics + """ + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + metrics = {} + + # Basic accuracy for classification + if y_true.dtype == y_pred.dtype and np.issubdtype(y_true.dtype, np.integer): + metrics["accuracy"] = float(np.mean(y_true == y_pred)) + else: + # MSE for regression + metrics["mse"] = float(np.mean((y_true - y_pred) ** 2)) + metrics["mae"] = float(np.mean(np.abs(y_true - y_pred))) + + return metrics + + +def batch_iterator(iterable: Any, batch_size: int): + """ + Iterate over data in batches. + + Args: + iterable: Data to iterate over + batch_size: Size of each batch + + Yields: + Batches of data + """ + batch = [] + for item in iterable: + batch.append(item) + if len(batch) == batch_size: + yield batch + batch = [] + if batch: + yield batch From 31c4fea927a848133b9950614ad39358ab6548b0 Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 16:20:23 -0700 Subject: [PATCH 3/6] Apply black/isort formatting to the ported SDK code Fixes Code Quality Checks (black --check, isort --check-only), which was never run against this content before the previous commit. Co-Authored-By: Claude Sonnet 5 --- examples/async_monitoring.py | 4 +- examples/boosting_example.py | 64 +++++-------------- examples/decorator_monitoring.py | 5 +- examples/offline_mode_example.py | 5 +- examples/pytorch_integration.py | 1 + examples/sklearn_integration.py | 5 +- examples/tensorflow_example.py | 8 +-- examples/transformers_example.py | 25 ++------ tests/integration/test_crewai_integration.py | 46 ++++--------- .../integration/test_langchain_integration.py | 36 +++-------- tests/unit/test_client.py | 7 +- tests/unit/test_config_exceptions.py | 16 ++--- tests/unit/test_decorators.py | 4 +- tests/unit/test_monitor.py | 12 +--- whiteboxxai/client.py | 8 +-- whiteboxxai/config.py | 4 +- whiteboxxai/decorators.py | 4 +- whiteboxxai/git_utils.py | 8 +-- whiteboxxai/integrations/__init__.py | 17 +---- whiteboxxai/integrations/boosting.py | 29 +++------ whiteboxxai/integrations/crewai_monitor.py | 21 ++---- whiteboxxai/integrations/langchain.py | 8 +-- whiteboxxai/integrations/langchain_agents.py | 26 ++------ whiteboxxai/integrations/pytorch.py | 13 ++-- whiteboxxai/integrations/sklearn.py | 9 +-- whiteboxxai/integrations/tensorflow.py | 9 +-- whiteboxxai/integrations/transformers.py | 5 +- whiteboxxai/monitor.py | 24 ++----- whiteboxxai/offline.py | 28 ++------ whiteboxxai/privacy.py | 8 +-- whiteboxxai/resources.py | 64 +++++-------------- whiteboxxai/utils.py | 16 ++--- 32 files changed, 140 insertions(+), 399 deletions(-) diff --git a/examples/async_monitoring.py b/examples/async_monitoring.py index 28346b2..b138a6f 100644 --- a/examples/async_monitoring.py +++ b/examples/async_monitoring.py @@ -7,6 +7,7 @@ import asyncio import numpy as np + from whiteboxxai import ModelMonitor, WhiteBoxXAI @@ -37,8 +38,7 @@ async def register_and_log(): # Log batch predictions print("\nLogging batch predictions...") predictions = [ - {"inputs": {"amount": 50.0}, "output": {"fraud_prob": 0.05}} - for _ in range(100) + {"inputs": {"amount": 50.0}, "output": {"fraud_prob": 0.05}} for _ in range(100) ] await monitor.alog_batch(predictions) diff --git a/examples/boosting_example.py b/examples/boosting_example.py index b307e70..27d4f7c 100644 --- a/examples/boosting_example.py +++ b/examples/boosting_example.py @@ -46,9 +46,7 @@ def example_xgboost_classification(): X, y = make_classification( n_samples=1000, n_features=20, n_informative=15, n_redundant=5, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") @@ -63,9 +61,7 @@ def example_xgboost_classification(): # Train XGBoost model print("Training XGBoost classifier...") - model = xgb.XGBClassifier( - n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 - ) + model = xgb.XGBClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) model.fit(X_train, y_train) # Register model with WhiteBoxXAI @@ -116,9 +112,7 @@ def example_xgboost_regression(): X, y = make_regression( n_samples=1000, n_features=10, n_informative=8, noise=10, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") @@ -132,9 +126,7 @@ def example_xgboost_regression(): # Train XGBoost regressor print("Training XGBoost regressor...") - model = xgb.XGBRegressor( - n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 - ) + model = xgb.XGBRegressor(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) model.fit(X_train, y_train) # Wrap model for automatic monitoring @@ -171,9 +163,7 @@ def example_lightgbm_classification(): X, y = make_classification( n_samples=1000, n_features=20, n_informative=15, n_redundant=5, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") @@ -188,9 +178,7 @@ def example_lightgbm_classification(): # Train LightGBM model print("Training LightGBM classifier...") - model = lgb.LGBMClassifier( - n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 - ) + model = lgb.LGBMClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) model.fit(X_train, y_train) # Register model with WhiteBoxXAI @@ -243,9 +231,7 @@ def example_lightgbm_regression(): X, y = make_regression( n_samples=1000, n_features=10, n_informative=8, noise=10, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") @@ -259,16 +245,12 @@ def example_lightgbm_regression(): # Train LightGBM regressor print("Training LightGBM regressor...") - model = lgb.LGBMRegressor( - n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 - ) + model = lgb.LGBMRegressor(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) model.fit(X_train, y_train) # Wrap model for automatic monitoring print("Wrapping model for automatic monitoring...") - wrapped_model = wrap_lightgbm_model( - model=model, monitor=monitor, auto_register=True - ) + wrapped_model = wrap_lightgbm_model(model=model, monitor=monitor, auto_register=True) # Predictions automatically logged print("\nMaking predictions (auto-logged)...") @@ -300,17 +282,13 @@ def example_feature_importance_tracking(): # Generate synthetic data with feature names import pandas as pd - X, y = make_classification( - n_samples=1000, n_features=10, n_informative=7, random_state=42 - ) + X, y = make_classification(n_samples=1000, n_features=10, n_informative=7, random_state=42) # Create DataFrame with feature names feature_names = [f"feature_{i}" for i in range(10)] X_df = pd.DataFrame(X, columns=feature_names) - X_train, X_test, y_train, y_test = train_test_split( - X_df, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X_df, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") @@ -333,9 +311,7 @@ def example_feature_importance_tracking(): # Get importance importance_dict = monitor._get_feature_importance(xgb_model) if importance_dict: - sorted_features = sorted( - importance_dict.items(), key=lambda x: x[1], reverse=True - )[:5] + sorted_features = sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:5] for feat, score in sorted_features: print(f" {feat}: {score:.4f}") @@ -357,9 +333,7 @@ def example_feature_importance_tracking(): # Get importance importance_dict = monitor._get_feature_importance(lgb_model) if importance_dict: - sorted_features = sorted( - importance_dict.items(), key=lambda x: x[1], reverse=True - )[:5] + sorted_features = sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)[:5] for feat, score in sorted_features: print(f" {feat}: {score:.4f}") @@ -380,21 +354,15 @@ def example_model_comparison(): return # Generate synthetic data - X, y = make_classification( - n_samples=1000, n_features=20, n_informative=15, random_state=42 - ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X, y = make_classification(n_samples=1000, n_features=20, n_informative=15, random_state=42) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Initialize WhiteBoxXAI client client = WhiteBoxXAI(api_key="demo-api-key") # Train and monitor XGBoost print("Training XGBoost model...") - xgb_model = xgb.XGBClassifier( - n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42 - ) + xgb_model = xgb.XGBClassifier(n_estimators=100, max_depth=5, learning_rate=0.1, random_state=42) xgb_model.fit(X_train, y_train) xgb_monitor = XGBoostMonitor(client=client, model_name="xgb_comparison") diff --git a/examples/decorator_monitoring.py b/examples/decorator_monitoring.py index 22765d2..9ffa04d 100644 --- a/examples/decorator_monitoring.py +++ b/examples/decorator_monitoring.py @@ -6,6 +6,7 @@ import numpy as np from sklearn.ensemble import RandomForestClassifier + from whiteboxxai import ModelMonitor, WhiteBoxXAI, monitor_model, monitor_prediction # Global monitor instance @@ -69,9 +70,7 @@ def main(): print("\n=== Custom Extractors ===") # Custom input/output extraction - result = score_transaction( - data={"amount": 100.0, "velocity": 5.0, "location_risk": 0.3} - ) + result = score_transaction(data={"amount": 100.0, "velocity": 5.0, "location_risk": 0.3}) print(f"Transaction score: {result}") print("\n=== Class Method Decorator ===") diff --git a/examples/offline_mode_example.py b/examples/offline_mode_example.py index 41a0278..fdc1e8d 100644 --- a/examples/offline_mode_example.py +++ b/examples/offline_mode_example.py @@ -13,6 +13,7 @@ from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split + from whiteboxxai import WhiteBoxXAI from whiteboxxai.offline import OperationPriority, OperationType @@ -206,9 +207,7 @@ def example_5_ml_model_with_offline(): X, y = make_classification( n_samples=1000, n_features=10, n_informative=8, n_redundant=2, random_state=42 ) - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Train model print("\nTraining Random Forest model...") diff --git a/examples/pytorch_integration.py b/examples/pytorch_integration.py index cf2d5f7..ccf3e5b 100644 --- a/examples/pytorch_integration.py +++ b/examples/pytorch_integration.py @@ -8,6 +8,7 @@ import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset + from whiteboxxai import WhiteBoxXAI from whiteboxxai.integrations.pytorch import TorchMonitor diff --git a/examples/sklearn_integration.py b/examples/sklearn_integration.py index 812e4fd..e61b138 100644 --- a/examples/sklearn_integration.py +++ b/examples/sklearn_integration.py @@ -8,6 +8,7 @@ from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split + from whiteboxxai import WhiteBoxXAI from whiteboxxai.integrations.sklearn import SklearnMonitor @@ -24,9 +25,7 @@ def main(): ) # Split data - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Train model print("Training Random Forest model...") diff --git a/examples/tensorflow_example.py b/examples/tensorflow_example.py index 6a8776c..11b08ee 100644 --- a/examples/tensorflow_example.py +++ b/examples/tensorflow_example.py @@ -38,9 +38,7 @@ def main(): ) # Split data - X_train, X_test, y_train, y_test = train_test_split( - X, y, test_size=0.2, random_state=42 - ) + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42) # Standardize features scaler = StandardScaler() @@ -149,9 +147,7 @@ def main(): prob = predictions[i][0] pred_class = 1 if prob > 0.5 else 0 actual = y_test[i] - print( - f" Sample {i+1}: Predicted={pred_class} (prob={prob:.3f}), Actual={actual}" - ) + print(f" Sample {i+1}: Predicted={pred_class} (prob={prob:.3f}), Actual={actual}") # Save model print("\n9. Saving model...") diff --git a/examples/transformers_example.py b/examples/transformers_example.py index 9cc9c7a..9cba12b 100644 --- a/examples/transformers_example.py +++ b/examples/transformers_example.py @@ -7,10 +7,7 @@ import os from whiteboxxai import WhiteBoxXAI -from whiteboxxai.integrations.transformers import ( - TransformersMonitor, - wrap_transformers_pipeline, -) +from whiteboxxai.integrations.transformers import TransformersMonitor, wrap_transformers_pipeline # Optional: Set API key os.environ["WHITEBOXXAI_API_KEY"] = "your-api-key-here" @@ -88,14 +85,10 @@ def example_ner(): ner_pipeline = pipeline("ner", aggregation_strategy="simple") # Create monitor - monitor = TransformersMonitor( - client=client, pipeline=ner_pipeline, model_name="ner_model_v1" - ) + monitor = TransformersMonitor(client=client, pipeline=ner_pipeline, model_name="ner_model_v1") # Register model - model_id = monitor.register_from_model( - name="BERT NER Model", version="1.0.0", task="ner" - ) + model_id = monitor.register_from_model(name="BERT NER Model", version="1.0.0", task="ner") print(f"✓ Model registered with ID: {model_id}") # Test text @@ -107,9 +100,7 @@ def example_ner(): print("\nDetected entities:") for entity in result: - print( - f" - {entity['word']}: {entity['entity_group']} (score: {entity['score']:.3f})" - ) + print(f" - {entity['word']}: {entity['entity_group']} (score: {entity['score']:.3f})") print("\n✓ NER monitoring complete!") @@ -131,9 +122,7 @@ def example_text_generation(): generator = pipeline("text-generation", model="gpt2") # Create monitor - monitor = TransformersMonitor( - client=client, pipeline=generator, model_name="gpt2_generator" - ) + monitor = TransformersMonitor(client=client, pipeline=generator, model_name="gpt2_generator") # Register model model_id = monitor.register_from_model( @@ -220,9 +209,7 @@ def example_batch_prediction(): classifier = pipeline("sentiment-analysis") # Create monitor - monitor = TransformersMonitor( - client=client, pipeline=classifier, model_name="batch_classifier" - ) + monitor = TransformersMonitor(client=client, pipeline=classifier, model_name="batch_classifier") # Register model monitor.register_from_model(name="Batch Sentiment Classifier") diff --git a/tests/integration/test_crewai_integration.py b/tests/integration/test_crewai_integration.py index e43e193..675c4ee 100644 --- a/tests/integration/test_crewai_integration.py +++ b/tests/integration/test_crewai_integration.py @@ -8,6 +8,7 @@ from uuid import uuid4 import pytest + from whiteboxxai.integrations.crewai_monitor import CrewAIMonitor, monitor_crew @@ -130,9 +131,7 @@ def test_start_monitoring(self, mock_client, mock_crew): # Check start workflow call assert calls[5][0][0] == "POST" - assert ( - calls[5][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/start" - ) + assert calls[5][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/start" def test_register_agent(self, mock_client, mock_crew): """Test agent registration.""" @@ -156,9 +155,7 @@ def test_register_agent(self, mock_client, mock_crew): # Verify API call call_args = mock_client.request.call_args assert call_args[0][0] == "POST" - assert ( - call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/agents" - ) + assert call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/agents" agent_data = call_args[1]["data"] assert agent_data["name"] == "Research Analyst" @@ -194,9 +191,7 @@ def test_register_task(self, mock_client, mock_crew): # Verify API call call_args = mock_client.request.call_args assert call_args[0][0] == "POST" - assert ( - call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/tasks" - ) + assert call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/tasks" task_data = call_args[1]["data"] assert task_data["task_name"] == "Research AI safety regulations" @@ -230,10 +225,7 @@ def test_log_agent_execution(self, mock_client, mock_crew): # Verify API call call_args = mock_client.request.call_args assert call_args[0][0] == "POST" - assert ( - call_args[0][1] - == f"/api/v1/workflows/multi-agent/{workflow_id}/executions" - ) + assert call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/executions" execution_data = call_args[1]["data"] assert execution_data["agent_id"] == agent_id @@ -298,10 +290,7 @@ def test_log_interaction(self, mock_client, mock_crew): # Verify API call call_args = mock_client.request.call_args assert call_args[0][0] == "POST" - assert ( - call_args[0][1] - == f"/api/v1/workflows/multi-agent/{workflow_id}/interactions" - ) + assert call_args[0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/interactions" interaction_data = call_args[1]["data"] assert interaction_data["interaction_type"] == "delegation" @@ -326,9 +315,7 @@ def test_complete_monitoring(self, mock_client): monitor = CrewAIMonitor(api_key="test_api_key") monitor.workflow_id = workflow_id - result = monitor.complete_monitoring( - status="completed", outputs={"article": "content"} - ) + result = monitor.complete_monitoring(status="completed", outputs={"article": "content"}) assert result["workflow_id"] == workflow_id assert result["status"] == "completed" @@ -337,10 +324,7 @@ def test_complete_monitoring(self, mock_client): # Verify complete call calls = mock_client.request.call_args_list assert calls[0][0][0] == "POST" - assert ( - calls[0][0][1] - == f"/api/v1/workflows/multi-agent/{workflow_id}/complete" - ) + assert calls[0][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/complete" complete_data = calls[0][1]["data"] assert complete_data["status"] == "completed" @@ -369,14 +353,8 @@ def test_get_analytics(self, mock_client): # Verify calls calls = mock_client.request.call_args_list - assert ( - calls[0][0][1] - == f"/api/v1/workflows/multi-agent/{workflow_id}/analytics" - ) - assert ( - calls[1][0][1] - == f"/api/v1/workflows/multi-agent/{workflow_id}/cost-breakdown" - ) + assert calls[0][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/analytics" + assert calls[1][0][1] == f"/api/v1/workflows/multi-agent/{workflow_id}/cost-breakdown" class TestMonitorCrewHelper: @@ -436,9 +414,7 @@ def test_log_execution_unregistered_agent(self, mock_client, mock_crew): monitor.workflow_id = str(uuid4()) # Should not raise exception, just log warning - monitor.log_agent_execution( - agent=mock_crew.agents[0], inputs={"test": "data"} - ) + monitor.log_agent_execution(agent=mock_crew.agents[0], inputs={"test": "data"}) # No API call should be made mock_client.request.assert_not_called() diff --git a/tests/integration/test_langchain_integration.py b/tests/integration/test_langchain_integration.py index 0fb4e2d..2282fdd 100644 --- a/tests/integration/test_langchain_integration.py +++ b/tests/integration/test_langchain_integration.py @@ -68,9 +68,7 @@ def test_on_chain_start(self, callback_handler): """Test chain start callback.""" inputs = {"input": "test query"} - callback_handler.on_chain_start( - serialized={"name": "test_chain"}, inputs=inputs - ) + callback_handler.on_chain_start(serialized={"name": "test_chain"}, inputs=inputs) assert callback_handler.execution_inputs == inputs assert callback_handler.execution_start_time is not None @@ -129,9 +127,7 @@ def test_on_llm_start(self, callback_handler): """Test LLM start callback.""" initial_count = callback_handler.llm_call_count - callback_handler.on_llm_start( - serialized={"name": "test_llm"}, prompts=["test prompt"] - ) + callback_handler.on_llm_start(serialized={"name": "test_llm"}, prompts=["test prompt"]) assert callback_handler.llm_call_count == initial_count + 1 @@ -226,9 +222,7 @@ def test_start_monitoring(self, mock_client, mock_workflow_response): """Test starting workflow monitoring.""" mock_client.agent_workflows.create.return_value = mock_workflow_response - monitor = LangGraphMultiAgentMonitor( - client=mock_client, workflow_name="Test Workflow" - ) + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test Workflow") inputs = {"query": "test"} workflow_id = monitor.start_monitoring(inputs=inputs) @@ -249,9 +243,7 @@ def test_register_agent(self, mock_client, mock_workflow_response): """Test agent registration.""" mock_client.agent_workflows.create.return_value = mock_workflow_response - monitor = LangGraphMultiAgentMonitor( - client=mock_client, workflow_name="Test Workflow" - ) + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test Workflow") monitor.start_monitoring() monitor.register_agent( @@ -274,9 +266,7 @@ def test_get_callbacks(self, mock_client, mock_workflow_response): """Test getting callbacks for an agent.""" mock_client.agent_workflows.create.return_value = mock_workflow_response - monitor = LangGraphMultiAgentMonitor( - client=mock_client, workflow_name="Test Workflow" - ) + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test Workflow") monitor.start_monitoring() callbacks = monitor.get_callbacks("researcher", agent_role="Research Agent") @@ -294,9 +284,7 @@ def test_log_handoff(self, mock_client, mock_workflow_response): """Test logging agent-to-agent handoff.""" mock_client.agent_workflows.create.return_value = mock_workflow_response - monitor = LangGraphMultiAgentMonitor( - client=mock_client, workflow_name="Test Workflow" - ) + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test Workflow") monitor.start_monitoring() monitor.log_handoff( @@ -324,9 +312,7 @@ def test_complete_monitoring(self, mock_client, mock_workflow_response): "total_tokens": 750, } - monitor = LangGraphMultiAgentMonitor( - client=mock_client, workflow_name="Test Workflow" - ) + monitor = LangGraphMultiAgentMonitor(client=mock_client, workflow_name="Test Workflow") monitor.start_monitoring() outputs = {"result": "completed successfully"} @@ -338,9 +324,7 @@ def test_complete_monitoring(self, mock_client, mock_workflow_response): ) # Verify analytics were retrieved - mock_client.agent_workflows.get_analytics.assert_called_once_with( - "workflow_123" - ) + mock_client.agent_workflows.get_analytics.assert_called_once_with("workflow_123") assert summary["workflow_id"] == "workflow_123" assert summary["status"] == "completed" @@ -420,9 +404,7 @@ class TestErrorHandling: def test_callback_api_error_handling(self, mock_client): """Test that API errors don't crash the callback.""" - mock_client.agent_workflows.create_execution.side_effect = Exception( - "API Error" - ) + mock_client.agent_workflows.create_execution.side_effect = Exception("API Error") handler = MultiAgentCallbackHandler( client=mock_client, workflow_id="workflow_123", agent_name="test" diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 7d34e1b..e374ebe 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -11,12 +11,7 @@ from whiteboxxai.client import WhiteBoxXAI from whiteboxxai.config import Config -from whiteboxxai.exceptions import ( - APIError, - AuthenticationError, - RateLimitError, - ValidationError, -) +from whiteboxxai.exceptions import APIError, AuthenticationError, RateLimitError, ValidationError class TestWhiteBoxXAIClient: diff --git a/tests/unit/test_config_exceptions.py b/tests/unit/test_config_exceptions.py index 485acd7..e7bbd10 100644 --- a/tests/unit/test_config_exceptions.py +++ b/tests/unit/test_config_exceptions.py @@ -140,9 +140,7 @@ def test_validation_error_inheritance(self): def test_validation_error_with_fields(self): """Test validation error with field details.""" - error = ValidationError( - "Validation failed", fields={"email": ["Invalid email format"]} - ) + error = ValidationError("Validation failed", fields={"email": ["Invalid email format"]}) assert "Validation failed" in str(error) @@ -161,9 +159,7 @@ def test_not_found_error_inheritance(self): def test_not_found_error_with_resource(self): """Test not found error with resource details.""" - error = NotFoundError( - "Resource not found", resource_type="model", resource_id="123" - ) + error = NotFoundError("Resource not found", resource_type="model", resource_id="123") assert "Resource not found" in str(error) @@ -182,9 +178,7 @@ def test_rate_limit_error_inheritance(self): def test_rate_limit_error_with_retry(self): """Test rate limit error with retry information.""" - error = RateLimitError( - "Too many requests", retry_after=60, limit=100, remaining=0 - ) + error = RateLimitError("Too many requests", retry_after=60, limit=100, remaining=0) assert "Too many requests" in str(error) @@ -203,9 +197,7 @@ def test_api_error_inheritance(self): def test_api_error_with_status(self): """Test API error with status code.""" - error = APIError( - "Server error", status_code=500, response={"error": "Internal error"} - ) + error = APIError("Server error", status_code=500, response={"error": "Internal error"}) assert "Server error" in str(error) def test_api_error_with_request_id(self): diff --git a/tests/unit/test_decorators.py b/tests/unit/test_decorators.py index c2bef32..9611047 100644 --- a/tests/unit/test_decorators.py +++ b/tests/unit/test_decorators.py @@ -52,9 +52,7 @@ def test_decorator_extracts_output(self): """Test that decorator extracts output correctly.""" mock_monitor = Mock(spec=ModelMonitor) - @monitor_model( - mock_monitor, input_keys=["x"], output_key="prediction", explain=False - ) + @monitor_model(mock_monitor, input_keys=["x"], output_key="prediction", explain=False) def predict(x): return {"prediction": x * 2, "confidence": 0.95} diff --git a/tests/unit/test_monitor.py b/tests/unit/test_monitor.py index fa8bbce..ed77de1 100644 --- a/tests/unit/test_monitor.py +++ b/tests/unit/test_monitor.py @@ -167,9 +167,7 @@ def test_get_drift_reports(self): """Test getting drift reports.""" mock_client = Mock(spec=WhiteBoxXAI) mock_client.drift = Mock() - mock_client.drift.get_reports = Mock( - return_value=[{"id": "report-1", "drift_score": 0.8}] - ) + mock_client.drift.get_reports = Mock(return_value=[{"id": "report-1", "drift_score": 0.8}]) monitor = ModelMonitor(client=mock_client, model_id="model-123") @@ -191,9 +189,7 @@ def test_create_alert_rule(self): monitor = ModelMonitor(client=mock_client, model_id="model-123") - rule = monitor.create_alert_rule( - metric="accuracy", threshold=0.8, condition="below" - ) + rule = monitor.create_alert_rule(metric="accuracy", threshold=0.8, condition="below") # Verify alert rule was created assert rule["id"] == "rule-1" @@ -274,9 +270,7 @@ def test_get_performance_metrics(self): """Test getting performance metrics.""" mock_client = Mock(spec=WhiteBoxXAI) mock_client.metrics = Mock() - mock_client.metrics.get = Mock( - return_value={"accuracy": 0.95, "precision": 0.93} - ) + mock_client.metrics.get = Mock(return_value={"accuracy": 0.95, "precision": 0.93}) monitor = ModelMonitor(client=mock_client, model_id="model-123") diff --git a/whiteboxxai/client.py b/whiteboxxai/client.py index 9170316..3af65a9 100644 --- a/whiteboxxai/client.py +++ b/whiteboxxai/client.py @@ -11,13 +11,9 @@ import httpx from tenacity import retry, stop_after_attempt, wait_exponential + from whiteboxxai.config import Config -from whiteboxxai.exceptions import ( - APIError, - AuthenticationError, - RateLimitError, - ValidationError, -) +from whiteboxxai.exceptions import APIError, AuthenticationError, RateLimitError, ValidationError from whiteboxxai.resources import ( AlertsResource, DriftResource, diff --git a/whiteboxxai/config.py b/whiteboxxai/config.py index a8ef590..7ecd22e 100644 --- a/whiteboxxai/config.py +++ b/whiteboxxai/config.py @@ -42,9 +42,7 @@ def __init__( # Base URL self.base_url = ( - base_url - or os.getenv("WHITEBOXXAI_BASE_URL") - or "https://api.whiteboxxai.com" + base_url or os.getenv("WHITEBOXXAI_BASE_URL") or "https://api.whiteboxxai.com" ) # Request settings diff --git a/whiteboxxai/decorators.py b/whiteboxxai/decorators.py index 92a9033..a2120c1 100644 --- a/whiteboxxai/decorators.py +++ b/whiteboxxai/decorators.py @@ -254,9 +254,7 @@ def _extract_inputs( return _default_input_extractor(func, args, kwargs) -def _default_input_extractor( - func: Callable, args: tuple, kwargs: dict -) -> Dict[str, Any]: +def _default_input_extractor(func: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]: """Default input extractor.""" sig = inspect.signature(func) bound_args = sig.bind(*args, **kwargs) diff --git a/whiteboxxai/git_utils.py b/whiteboxxai/git_utils.py index 22b79ab..789e5b9 100644 --- a/whiteboxxai/git_utils.py +++ b/whiteboxxai/git_utils.py @@ -101,9 +101,7 @@ def detect_git_context(path: Optional[str] = None) -> Optional[GitContext]: repository_url = remote.url # Convert SSH URLs to HTTPS if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace( - "git@github.com:", "https://github.com/" - ) + repository_url = repository_url.replace("git@github.com:", "https://github.com/") if repository_url.endswith(".git"): repository_url = repository_url[:-4] except Exception as e: @@ -202,9 +200,7 @@ def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitCo # Convert SSH to HTTPS if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace( - "git@github.com:", "https://github.com/" - ) + repository_url = repository_url.replace("git@github.com:", "https://github.com/") if repository_url.endswith(".git"): repository_url = repository_url[:-4] except subprocess.CalledProcessError: diff --git a/whiteboxxai/integrations/__init__.py b/whiteboxxai/integrations/__init__.py index 806953d..f4da8df 100644 --- a/whiteboxxai/integrations/__init__.py +++ b/whiteboxxai/integrations/__init__.py @@ -57,16 +57,10 @@ # LangChain integration try: - from .langchain import ( - LangChainMonitor, - WhiteBoxXAICallbackHandler, - wrap_langchain_chain, - ) + from .langchain import LangChainMonitor, WhiteBoxXAICallbackHandler, wrap_langchain_chain if "__all__" in dir(): - __all__.extend( - ["LangChainMonitor", "WhiteBoxXAICallbackHandler", "wrap_langchain_chain"] - ) + __all__.extend(["LangChainMonitor", "WhiteBoxXAICallbackHandler", "wrap_langchain_chain"]) else: __all__ = [ "LangChainMonitor", @@ -78,12 +72,7 @@ # XGBoost/LightGBM integration try: - from .boosting import ( - LightGBMMonitor, - XGBoostMonitor, - wrap_lightgbm_model, - wrap_xgboost_model, - ) + from .boosting import LightGBMMonitor, XGBoostMonitor, wrap_lightgbm_model, wrap_xgboost_model if "__all__" in dir(): __all__.extend( diff --git a/whiteboxxai/integrations/boosting.py b/whiteboxxai/integrations/boosting.py index d29a6b8..bfc5ae2 100644 --- a/whiteboxxai/integrations/boosting.py +++ b/whiteboxxai/integrations/boosting.py @@ -41,6 +41,7 @@ from typing import Any, Dict, List, Optional, Union import numpy as np + from whiteboxxai.monitor import ModelMonitor logger = logging.getLogger(__name__) @@ -106,9 +107,7 @@ def __init__( **kwargs: Additional arguments passed to ModelMonitor """ if not XGBOOST_AVAILABLE: - raise ImportError( - "XGBoost is not installed. Install with: pip install xgboost" - ) + raise ImportError("XGBoost is not installed. Install with: pip install xgboost") super().__init__(client=client, model_name=model_name, **kwargs) self.track_feature_importance = track_feature_importance @@ -290,9 +289,7 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try get_booster for sklearn API if hasattr(model, "get_booster"): booster = model.get_booster() - importance_dict = booster.get_score( - importance_type=self.importance_type - ) + importance_dict = booster.get_score(importance_type=self.importance_type) return {k: float(v) for k, v in importance_dict.items()} return None @@ -343,9 +340,7 @@ def __init__( **kwargs: Additional arguments passed to ModelMonitor """ if not LIGHTGBM_AVAILABLE: - raise ImportError( - "LightGBM is not installed. Install with: pip install lightgbm" - ) + raise ImportError("LightGBM is not installed. Install with: pip install lightgbm") super().__init__(client=client, model_name=model_name, **kwargs) self.track_feature_importance = track_feature_importance @@ -523,9 +518,7 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try feature_importance method (native LightGBM) if hasattr(model, "feature_importance"): - importances = model.feature_importance( - importance_type=self.importance_type - ) + importances = model.feature_importance(importance_type=self.importance_type) if self._feature_names and len(importances) == len(self._feature_names): return dict(zip(self._feature_names, importances.tolist())) else: @@ -534,9 +527,7 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Try booster if hasattr(model, "booster_"): booster = model.booster_ - importances = booster.feature_importance( - importance_type=self.importance_type - ) + importances = booster.feature_importance(importance_type=self.importance_type) feature_names = booster.feature_name() if len(importances) == len(feature_names): return dict(zip(feature_names, importances.tolist())) @@ -547,9 +538,7 @@ def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: # Unified wrapper functions -def wrap_xgboost_model( - model: Any, monitor: XGBoostMonitor, auto_register: bool = True -) -> Any: +def wrap_xgboost_model(model: Any, monitor: XGBoostMonitor, auto_register: bool = True) -> Any: """ Wrap an XGBoost model for automatic monitoring. @@ -615,9 +604,7 @@ def wrapped_predict_proba(X, *args, **kwargs): return model -def wrap_lightgbm_model( - model: Any, monitor: LightGBMMonitor, auto_register: bool = True -) -> Any: +def wrap_lightgbm_model(model: Any, monitor: LightGBMMonitor, auto_register: bool = True) -> Any: """ Wrap a LightGBM model for automatic monitoring. diff --git a/whiteboxxai/integrations/crewai_monitor.py b/whiteboxxai/integrations/crewai_monitor.py index fb8970a..2872641 100644 --- a/whiteboxxai/integrations/crewai_monitor.py +++ b/whiteboxxai/integrations/crewai_monitor.py @@ -151,9 +151,7 @@ def start_monitoring( "inputs": { "agent_count": len(crew.agents), "task_count": len(crew.tasks), - "process": str(crew.process) - if hasattr(crew, "process") - else "sequential", + "process": str(crew.process) if hasattr(crew, "process") else "sequential", } } @@ -186,17 +184,13 @@ def _register_agent(self, crew_agent: Any) -> str: "agent_type": "crewai_agent", "goal": getattr(crew_agent, "goal", None), "backstory": getattr(crew_agent, "backstory", None), - "tools": [ - tool.__class__.__name__ for tool in getattr(crew_agent, "tools", []) - ], + "tools": [tool.__class__.__name__ for tool in getattr(crew_agent, "tools", [])], "llm_provider": getattr( getattr(crew_agent, "llm", None), "model_name", "unknown" ).split("-")[0] if hasattr(crew_agent, "llm") else None, - "model_name": getattr( - getattr(crew_agent, "llm", None), "model_name", None - ) + "model_name": getattr(getattr(crew_agent, "llm", None), "model_name", None) if hasattr(crew_agent, "llm") else None, "metadata": { @@ -245,10 +239,7 @@ def _register_task(self, crew_task: Any) -> str: "expected_output": getattr(crew_task, "expected_output", None), "agent_id": agent_id, "context": { - "tools": [ - tool.__class__.__name__ - for tool in getattr(crew_task, "tools", []) - ], + "tools": [tool.__class__.__name__ for tool in getattr(crew_task, "tools", [])], "async_execution": getattr(crew_task, "async_execution", False), }, } @@ -262,9 +253,7 @@ def _register_task(self, crew_task: Any) -> str: task_id = response.get("id") self.task_map[id(crew_task)] = task_id - logger.debug( - f"Registered task: {task_data['task_name'][:50]}... ({task_id})" - ) + logger.debug(f"Registered task: {task_data['task_name'][:50]}... ({task_id})") return task_id diff --git a/whiteboxxai/integrations/langchain.py b/whiteboxxai/integrations/langchain.py index 7aa005e..9267a39 100644 --- a/whiteboxxai/integrations/langchain.py +++ b/whiteboxxai/integrations/langchain.py @@ -91,9 +91,7 @@ def __init__( **kwargs: Additional arguments for ModelMonitor """ if not LANGCHAIN_AVAILABLE: - raise ImportError( - "langchain is not installed. Install with: pip install langchain" - ) + raise ImportError("langchain is not installed. Install with: pip install langchain") super().__init__(client, **kwargs) self._application_name = application_name @@ -514,9 +512,7 @@ def on_llm_end( self.monitor.log_llm_call( prompt="", # Prompt tracked separately response=response_text, - model=kwargs.get("invocation_params", {}).get( - "model_name", "unknown" - ), + model=kwargs.get("invocation_params", {}).get("model_name", "unknown"), tokens_used=tokens_used, latency=latency, ) diff --git a/whiteboxxai/integrations/langchain_agents.py b/whiteboxxai/integrations/langchain_agents.py index 894ba0f..6972c0c 100644 --- a/whiteboxxai/integrations/langchain_agents.py +++ b/whiteboxxai/integrations/langchain_agents.py @@ -93,8 +93,7 @@ def __init__( """ if WhiteBoxXAI is None: raise ImportError( - "whiteboxxai package not installed. " - "Install with: pip install whiteboxxai" + "whiteboxxai package not installed. " "Install with: pip install whiteboxxai" ) self.client = client @@ -153,9 +152,7 @@ def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None: # Reset state self.execution_start_time = None - def on_chain_error( - self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any - ) -> None: + def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: """Run when chain errors.""" if self.execution_start_time: duration_ms = int( @@ -179,9 +176,7 @@ def on_chain_error( self.execution_start_time = None - def on_llm_start( - self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any - ) -> None: + def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None: """Run when LLM starts.""" self.llm_call_count += 1 @@ -226,9 +221,7 @@ def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None: # This is called when the agent completes its reasoning pass - def on_tool_start( - self, serialized: Dict[str, Any], input_str: str, **kwargs: Any - ) -> None: + def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None: """Run when tool starts.""" pass @@ -247,9 +240,7 @@ def on_tool_end(self, output: str, **kwargs: Any) -> None: except Exception as e: print(f"Warning: Failed to log tool result: {e}") - def on_tool_error( - self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any - ) -> None: + def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: """Run when tool errors.""" try: self.client.agent_workflows.create_interaction( @@ -322,8 +313,7 @@ def __init__( """ if WhiteBoxXAI is None: raise ImportError( - "whiteboxxai package not installed. " - "Install with: pip install whiteboxxai" + "whiteboxxai package not installed. " "Install with: pip install whiteboxxai" ) self.client = client @@ -539,9 +529,7 @@ def monitor_langchain_agent( return {"result": result, "workflow_id": workflow_id, "status": "completed"} except Exception as e: # Log failure - client.agent_workflows.complete( - workflow_id, outputs={"error": str(e)}, status="failed" - ) + client.agent_workflows.complete(workflow_id, outputs={"error": str(e)}, status="failed") return { "result": None, diff --git a/whiteboxxai/integrations/pytorch.py b/whiteboxxai/integrations/pytorch.py index 98c9993..83a53d8 100644 --- a/whiteboxxai/integrations/pytorch.py +++ b/whiteboxxai/integrations/pytorch.py @@ -26,6 +26,7 @@ class _NNStub: nn = _NNStub() import numpy as np + from whiteboxxai.monitor import ModelMonitor @@ -63,9 +64,7 @@ class TorchMonitor(ModelMonitor): def __init__(self, client, model: Optional[nn.Module] = None, **kwargs): """Initialize PyTorch monitor.""" if not TORCH_AVAILABLE: - raise ImportError( - "PyTorch is not installed. Install with: pip install torch" - ) + raise ImportError("PyTorch is not installed. Install with: pip install torch") super().__init__(client, **kwargs) self.model = model @@ -154,9 +153,7 @@ def _extract_model_metadata(self) -> Dict[str, Any]: # Count parameters total_params = sum(p.numel() for p in self.model.parameters()) - trainable_params = sum( - p.numel() for p in self.model.parameters() if p.requires_grad - ) + trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) metadata["total_parameters"] = total_params metadata["trainable_parameters"] = trainable_params @@ -198,9 +195,7 @@ def forward(self, x: torch.Tensor, log: bool = True) -> torch.Tensor: return output - def _log_batch_predictions( - self, inputs: torch.Tensor, outputs: torch.Tensor - ) -> None: + def _log_batch_predictions(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None: """Log batch of predictions.""" # Convert to numpy/lists inputs_np = inputs.detach().cpu().numpy() diff --git a/whiteboxxai/integrations/sklearn.py b/whiteboxxai/integrations/sklearn.py index ccf89c1..f63fce0 100644 --- a/whiteboxxai/integrations/sklearn.py +++ b/whiteboxxai/integrations/sklearn.py @@ -16,6 +16,7 @@ BaseEstimator = object import numpy as np + from whiteboxxai.monitor import ModelMonitor @@ -201,12 +202,8 @@ def _log_batch_predictions(self, inputs: np.ndarray, outputs: np.ndarray) -> Non predictions = [] for i in range(len(inputs)): pred = { - "inputs": inputs[i].tolist() - if isinstance(inputs[i], np.ndarray) - else inputs[i], - "output": outputs[i].tolist() - if isinstance(outputs[i], np.ndarray) - else outputs[i], + "inputs": inputs[i].tolist() if isinstance(inputs[i], np.ndarray) else inputs[i], + "output": outputs[i].tolist() if isinstance(outputs[i], np.ndarray) else outputs[i], } predictions.append(pred) diff --git a/whiteboxxai/integrations/tensorflow.py b/whiteboxxai/integrations/tensorflow.py index b3dad14..5b159d8 100644 --- a/whiteboxxai/integrations/tensorflow.py +++ b/whiteboxxai/integrations/tensorflow.py @@ -31,6 +31,7 @@ class callbacks: keras = _KerasStub() import numpy as np + from whiteboxxai.monitor import ModelMonitor @@ -85,9 +86,7 @@ def __init__( **kwargs: Additional arguments for ModelMonitor """ if not TENSORFLOW_AVAILABLE: - raise ImportError( - "TensorFlow is not installed. Install with: pip install tensorflow" - ) + raise ImportError("TensorFlow is not installed. Install with: pip install tensorflow") super().__init__(client, **kwargs) self.model = model @@ -408,9 +407,7 @@ def on_train_end(self, logs: Optional[Dict[str, Any]] = None): self.monitor.log_custom_metric( "training_complete", { - "final_metrics": { - k: float(v) if v is not None else None for k, v in logs.items() - }, + "final_metrics": {k: float(v) if v is not None else None for k, v in logs.items()}, }, ) diff --git a/whiteboxxai/integrations/transformers.py b/whiteboxxai/integrations/transformers.py index 1fa9b35..87fae0a 100644 --- a/whiteboxxai/integrations/transformers.py +++ b/whiteboxxai/integrations/transformers.py @@ -19,6 +19,7 @@ Pipeline = object import numpy as np + from whiteboxxai.monitor import ModelMonitor @@ -245,9 +246,7 @@ def predict( # Single prediction self.log_prediction_transformers( input_text=inputs, - prediction=predictions - if isinstance(predictions, dict) - else predictions[0], + prediction=predictions if isinstance(predictions, dict) else predictions[0], actual=actuals, ) diff --git a/whiteboxxai/monitor.py b/whiteboxxai/monitor.py index 805ecff..fa0c54d 100644 --- a/whiteboxxai/monitor.py +++ b/whiteboxxai/monitor.py @@ -161,9 +161,7 @@ def log_prediction( self._prediction_count += 1 if self.buffer_size: - self._buffer.append( - {"input_data": inputs, "output_data": output, "metadata": metadata} - ) + self._buffer.append({"input_data": inputs, "output_data": output, "metadata": metadata}) if len(self._buffer) >= self.buffer_size: self.flush() return None @@ -178,9 +176,7 @@ def log_prediction( ) if explain and isinstance(result, dict) and result.get("id"): - result["explanation"] = self.client.explanations.generate( - prediction_id=result["id"] - ) + result["explanation"] = self.client.explanations.generate(prediction_id=result["id"]) return result @@ -204,9 +200,7 @@ async def alog_prediction( self._prediction_count += 1 if self.buffer_size: - self._buffer.append( - {"input_data": inputs, "output_data": output, "metadata": metadata} - ) + self._buffer.append({"input_data": inputs, "output_data": output, "metadata": metadata}) if len(self._buffer) >= self.buffer_size: await self.aflush() return None @@ -312,16 +306,12 @@ def get_prediction_count(self) -> int: this monitor instance so far.""" return self._prediction_count - def get_drift_reports( - self, limit: int = 10, skip: int = 0 - ) -> List[Dict[str, Any]]: + def get_drift_reports(self, limit: int = 10, skip: int = 0) -> List[Dict[str, Any]]: """Get persisted drift reports for this model, most recent first.""" if self.model_id is None: raise ValueError("Model not registered. Call register_model() first.") - return self.client.drift.get_reports( - model_id=self.model_id, limit=limit, skip=skip - ) + return self.client.drift.get_reports(model_id=self.model_id, limit=limit, skip=skip) def create_alert_rule( self, @@ -415,9 +405,7 @@ def _should_sample(self) -> bool: return True return np.random.random() < self.sampling_rate - def _sample_predictions( - self, predictions: List[Dict[str, Any]] - ) -> List[Dict[str, Any]]: + def _sample_predictions(self, predictions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """Sample predictions based on sampling rate.""" n_samples = int(len(predictions) * self.sampling_rate) if n_samples == 0: diff --git a/whiteboxxai/offline.py b/whiteboxxai/offline.py index 28d1d5e..b7a30fc 100644 --- a/whiteboxxai/offline.py +++ b/whiteboxxai/offline.py @@ -74,9 +74,7 @@ class OfflineQueue: auto_sync: Whether to automatically sync when connection available """ - def __init__( - self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True - ): + def __init__(self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True): """ Initialize offline queue. @@ -147,9 +145,7 @@ def enqueue( if self.max_queue_size > 0: current_size = self.get_queue_size() if current_size >= self.max_queue_size: - raise ValueError( - f"Queue is full ({current_size}/{self.max_queue_size})" - ) + raise ValueError(f"Queue is full ({current_size}/{self.max_queue_size})") # Insert operation with sqlite3.connect(self.db_path) as conn: @@ -166,9 +162,7 @@ def enqueue( logger.info(f"Queued operation {op_id}: {operation_type.value}") return op_id - def dequeue( - self, limit: int = 100 - ) -> List[Tuple[int, OperationType, Dict[str, Any]]]: + def dequeue(self, limit: int = 100) -> List[Tuple[int, OperationType, Dict[str, Any]]]: """ Get pending operations from queue. @@ -197,9 +191,7 @@ def dequeue( operations = [] for row in cursor.fetchall(): op_id, op_type, data_json = row - operations.append( - (op_id, OperationType(op_type), json.loads(data_json)) - ) + operations.append((op_id, OperationType(op_type), json.loads(data_json))) return operations @@ -288,9 +280,7 @@ def get_queue_size(self, status: str = "pending") -> int: """ with sqlite3.connect(self.db_path) as conn: if status: - cursor = conn.execute( - "SELECT COUNT(*) FROM queue WHERE status = ?", (status,) - ) + cursor = conn.execute("SELECT COUNT(*) FROM queue WHERE status = ?", (status,)) else: cursor = conn.execute("SELECT COUNT(*) FROM queue") @@ -342,9 +332,7 @@ def clear_completed(self, older_than_days: int = 7): conn.commit() if deleted > 0: - logger.info( - f"Cleared {deleted} completed operations older than {older_than_days} days" - ) + logger.info(f"Cleared {deleted} completed operations older than {older_than_days} days") def clear_all(self): """Clear all operations from queue (use with caution).""" @@ -449,9 +437,7 @@ def start_auto_sync(self): """Start automatic sync thread.""" if self._sync_thread is None or not self._sync_thread.is_alive(): self._stop_sync.clear() - self._sync_thread = threading.Thread( - target=self._auto_sync_loop, daemon=True - ) + self._sync_thread = threading.Thread(target=self._auto_sync_loop, daemon=True) self._sync_thread.start() logger.info("Started automatic sync thread") diff --git a/whiteboxxai/privacy.py b/whiteboxxai/privacy.py index 745b308..4cb7990 100644 --- a/whiteboxxai/privacy.py +++ b/whiteboxxai/privacy.py @@ -24,9 +24,7 @@ def __init__(self): """Initialize PII detector with regex patterns.""" self.patterns: Dict[str, Pattern] = { "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"), - "phone": re.compile( - r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b" - ), + "phone": re.compile(r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), "credit_card": re.compile(r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b"), "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), @@ -206,9 +204,7 @@ def mask_pii(text: str, mask_char: str = "*") -> str: return _pii_detector.mask(text, mask_char=mask_char) -def mask_data( - data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True -) -> Any: +def mask_data(data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True) -> Any: """ Mask sensitive data using global masker. diff --git a/whiteboxxai/resources.py b/whiteboxxai/resources.py index 3d0cc42..af1d9f0 100644 --- a/whiteboxxai/resources.py +++ b/whiteboxxai/resources.py @@ -192,9 +192,7 @@ def update(self, model_id: str, **kwargs: Any) -> Dict[str, Any]: async def aupdate(self, model_id: str, **kwargs: Any) -> Dict[str, Any]: """Async version of update().""" - return await self.client.arequest( - "PATCH", f"/api/v1/models/{model_id}", data=kwargs - ) + return await self.client.arequest("PATCH", f"/api/v1/models/{model_id}", data=kwargs) def update_status(self, model_id: str, new_status: str) -> Dict[str, Any]: """Update model status (ACTIVE, INACTIVE, DEPRECATED, ARCHIVED).""" @@ -393,9 +391,7 @@ async def alog_batch( ) -> Dict[str, Any]: """Async version of log_batch().""" data = {"model_id": model_id, "predictions": predictions} - return await self.client.arequest( - "POST", "/api/v1/predictions/log/batch", data=data - ) + return await self.client.arequest("POST", "/api/v1/predictions/log/batch", data=data) def get(self, prediction_id: str) -> Dict[str, Any]: """Get a prediction by its ID.""" @@ -444,9 +440,7 @@ async def aquery( data["start_time"] = start_time if end_time is not None: data["end_time"] = end_time - return await self.client.arequest( - "POST", "/api/v1/predictions/query", data=data - ) + return await self.client.arequest("POST", "/api/v1/predictions/query", data=data) def get_stats( self, @@ -528,9 +522,7 @@ async def agenerate( ) -> Dict[str, Any]: """Async version of generate().""" data = {"prediction_id": prediction_id, "method": method, **kwargs} - return await self.client.arequest( - "POST", "/api/v1/explanations/generate", data=data - ) + return await self.client.arequest("POST", "/api/v1/explanations/generate", data=data) def get(self, explanation_id: int) -> Dict[str, Any]: """Get explanation by ID.""" @@ -538,9 +530,7 @@ def get(self, explanation_id: int) -> Dict[str, Any]: async def aget(self, explanation_id: int) -> Dict[str, Any]: """Async version of get().""" - return await self.client.arequest( - "GET", f"/api/v1/explanations/{explanation_id}" - ) + return await self.client.arequest("GET", f"/api/v1/explanations/{explanation_id}") class DriftResource(BaseResource): @@ -566,9 +556,7 @@ def detect( params: Dict[str, Any] = {"window_size": window_size} if feature_names is not None: params["feature_names"] = feature_names - return self.client.request( - "POST", f"/api/v1/drift/detect/{model_id}", params=params - ) + return self.client.request("POST", f"/api/v1/drift/detect/{model_id}", params=params) async def adetect( self, @@ -580,9 +568,7 @@ async def adetect( params: Dict[str, Any] = {"window_size": window_size} if feature_names is not None: params["feature_names"] = feature_names - return await self.client.arequest( - "POST", f"/api/v1/drift/detect/{model_id}", params=params - ) + return await self.client.arequest("POST", f"/api/v1/drift/detect/{model_id}", params=params) def create_report(self, model_id: str, window_size: int = 1000) -> Dict[str, Any]: """Run drift analysis and persist the result as a drift report.""" @@ -592,9 +578,7 @@ def create_report(self, model_id: str, window_size: int = 1000) -> Dict[str, Any params={"model_id": model_id, "window_size": window_size}, ) - async def acreate_report( - self, model_id: str, window_size: int = 1000 - ) -> Dict[str, Any]: + async def acreate_report(self, model_id: str, window_size: int = 1000) -> Dict[str, Any]: """Async version of create_report().""" return await self.client.arequest( "POST", @@ -602,9 +586,7 @@ async def acreate_report( params={"model_id": model_id, "window_size": window_size}, ) - def get_reports( - self, model_id: str, limit: int = 10, skip: int = 0 - ) -> List[Dict[str, Any]]: + def get_reports(self, model_id: str, limit: int = 10, skip: int = 0) -> List[Dict[str, Any]]: """List drift reports for a model, most recent first.""" return self.client.request( "GET", @@ -624,21 +606,15 @@ async def aget_reports( def get_report(self, model_id: str, report_id: str) -> Dict[str, Any]: """Get a specific drift report with per-feature statistics.""" - return self.client.request( - "GET", f"/api/v1/drift/reports/{model_id}/{report_id}" - ) + return self.client.request("GET", f"/api/v1/drift/reports/{model_id}/{report_id}") async def aget_report(self, model_id: str, report_id: str) -> Dict[str, Any]: """Async version of get_report().""" - return await self.client.arequest( - "GET", f"/api/v1/drift/reports/{model_id}/{report_id}" - ) + return await self.client.arequest("GET", f"/api/v1/drift/reports/{model_id}/{report_id}") def get_trend(self, model_id: str, days: int = 7) -> Dict[str, Any]: """Get drift trend over time for a model.""" - return self.client.request( - "GET", f"/api/v1/drift/trend/{model_id}", params={"days": days} - ) + return self.client.request("GET", f"/api/v1/drift/trend/{model_id}", params={"days": days}) async def aget_trend(self, model_id: str, days: int = 7) -> Dict[str, Any]: """Async version of get_trend().""" @@ -754,9 +730,7 @@ async def alist_audits( params: Dict[str, Any] = {"limit": limit, "offset": offset} if model_id is not None: params["model_id"] = model_id - return await self.client.arequest( - "GET", "/api/v1/fairness/audits", params=params - ) + return await self.client.arequest("GET", "/api/v1/fairness/audits", params=params) def get_bias_history(self, model_id: str, days: int = 30) -> Dict[str, Any]: """Get historical bias metrics/trend for a model.""" @@ -774,9 +748,7 @@ async def aget_bias_history(self, model_id: str, days: int = 30) -> Dict[str, An params={"days": days}, ) - def get_metric_history( - self, model_id: str, metric_type: str, days: int = 30 - ) -> Dict[str, Any]: + def get_metric_history(self, model_id: str, metric_type: str, days: int = 30) -> Dict[str, Any]: """Get history for a specific fairness metric type.""" return self.client.request( "GET", @@ -796,15 +768,11 @@ async def aget_metric_history( def get_latest_audit(self, model_id: str) -> Dict[str, Any]: """Get the most recent bias audit for a model.""" - return self.client.request( - "GET", f"/api/v1/fairness/models/{model_id}/latest-audit" - ) + return self.client.request("GET", f"/api/v1/fairness/models/{model_id}/latest-audit") async def aget_latest_audit(self, model_id: str) -> Dict[str, Any]: """Async version of get_latest_audit().""" - return await self.client.arequest( - "GET", f"/api/v1/fairness/models/{model_id}/latest-audit" - ) + return await self.client.arequest("GET", f"/api/v1/fairness/models/{model_id}/latest-audit") class AlertsResource(BaseResource): diff --git a/whiteboxxai/utils.py b/whiteboxxai/utils.py index 764c2db..372b72c 100644 --- a/whiteboxxai/utils.py +++ b/whiteboxxai/utils.py @@ -301,9 +301,7 @@ def deserialize_numpy(data: List) -> np.ndarray: return np.array(data) -def compute_metrics( - y_true: Any, y_pred: Any, task: str = "classification" -) -> Dict[str, float]: +def compute_metrics(y_true: Any, y_pred: Any, task: str = "classification") -> Dict[str, float]: """ Compute basic metrics for predictions. @@ -319,9 +317,7 @@ def compute_metrics( ValueError: If task is not recognized, or y_true/y_pred lengths differ """ if task not in ("classification", "regression"): - raise ValueError( - f"Unsupported task: {task!r}. Must be 'classification' or 'regression'." - ) + raise ValueError(f"Unsupported task: {task!r}. Must be 'classification' or 'regression'.") y_true = np.array(y_true) y_pred = np.array(y_pred) @@ -344,12 +340,8 @@ def compute_metrics( actual_positive = np.sum(y_true == cls) true_positive = np.sum((y_pred == cls) & (y_true == cls)) - precisions.append( - true_positive / predicted_positive if predicted_positive > 0 else 0.0 - ) - recalls.append( - true_positive / actual_positive if actual_positive > 0 else 0.0 - ) + precisions.append(true_positive / predicted_positive if predicted_positive > 0 else 0.0) + recalls.append(true_positive / actual_positive if actual_positive > 0 else 0.0) metrics["precision"] = float(np.mean(precisions)) metrics["recall"] = float(np.mean(recalls)) From 7b80133f82b50b7d48c69775c0deba7cdfaecd10 Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 16:21:49 -0700 Subject: [PATCH 4/6] Revert "feature: update sdk" This reverts commit 3f20fcd1404cccfdf86733f363005732393a4898. --- __init__.py | 1 - requirements-dev.txt | 26 - setup.py | 100 --- src/whiteboxai/__init__.py | 24 - src/whiteboxai/__version__.py | 6 - src/whiteboxai/cache.py | 117 --- src/whiteboxai/client.py | 354 ---------- src/whiteboxai/config.py | 90 --- src/whiteboxai/decorators.py | 323 --------- src/whiteboxai/exceptions.py | 57 -- src/whiteboxai/git_utils.py | 331 --------- src/whiteboxai/integrations/__init__.py | 112 --- src/whiteboxai/integrations/boosting.py | 667 ------------------ src/whiteboxai/integrations/crewai_monitor.py | 474 ------------- src/whiteboxai/integrations/langchain.py | 636 ----------------- .../integrations/langchain_agents.py | 525 -------------- src/whiteboxai/integrations/pytorch.py | 268 ------- src/whiteboxai/integrations/sklearn.py | 219 ------ src/whiteboxai/integrations/tensorflow.py | 475 ------------- src/whiteboxai/integrations/transformers.py | 531 -------------- src/whiteboxai/models/__init__.py | 1 - src/whiteboxai/monitor.py | 284 -------- src/whiteboxai/offline.py | 541 -------------- src/whiteboxai/privacy.py | 219 ------ src/whiteboxai/resources.py | 460 ------------ src/whiteboxai/utils.py | 337 --------- 26 files changed, 7178 deletions(-) delete mode 100644 __init__.py delete mode 100644 requirements-dev.txt delete mode 100644 setup.py delete mode 100644 src/whiteboxai/__init__.py delete mode 100644 src/whiteboxai/__version__.py delete mode 100644 src/whiteboxai/cache.py delete mode 100644 src/whiteboxai/client.py delete mode 100644 src/whiteboxai/config.py delete mode 100644 src/whiteboxai/decorators.py delete mode 100644 src/whiteboxai/exceptions.py delete mode 100644 src/whiteboxai/git_utils.py delete mode 100644 src/whiteboxai/integrations/__init__.py delete mode 100644 src/whiteboxai/integrations/boosting.py delete mode 100644 src/whiteboxai/integrations/crewai_monitor.py delete mode 100644 src/whiteboxai/integrations/langchain.py delete mode 100644 src/whiteboxai/integrations/langchain_agents.py delete mode 100644 src/whiteboxai/integrations/pytorch.py delete mode 100644 src/whiteboxai/integrations/sklearn.py delete mode 100644 src/whiteboxai/integrations/tensorflow.py delete mode 100644 src/whiteboxai/integrations/transformers.py delete mode 100644 src/whiteboxai/models/__init__.py delete mode 100644 src/whiteboxai/monitor.py delete mode 100644 src/whiteboxai/offline.py delete mode 100644 src/whiteboxai/privacy.py delete mode 100644 src/whiteboxai/resources.py delete mode 100644 src/whiteboxai/utils.py diff --git a/__init__.py b/__init__.py deleted file mode 100644 index b2b2619..0000000 --- a/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# WhiteBoxAI SDK diff --git a/requirements-dev.txt b/requirements-dev.txt deleted file mode 100644 index a87bedf..0000000 --- a/requirements-dev.txt +++ /dev/null @@ -1,26 +0,0 @@ -# Development dependencies for WhiteBoxAI SDK - -# Testing -pytest>=7.4.0 -pytest-asyncio>=0.21.0 -pytest-cov>=4.1.0 -pytest-mock>=3.12.0 - -# Code quality -black>=23.11.0 -isort>=5.12.0 -flake8>=6.1.0 -pylint>=3.0.0 -mypy>=1.7.0 - -# Security -bandit>=1.7.5 - -# Documentation -mkdocs>=1.5.0 -mkdocs-material>=9.0.0 -mkdocstrings[python]>=0.24.0 - -# Build tools -build>=1.0.0 -twine>=4.0.0 diff --git a/setup.py b/setup.py deleted file mode 100644 index 23da3c8..0000000 --- a/setup.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -WhiteBoxAI Python SDK - ML Monitoring and Observability -""" - -from pathlib import Path - -from setuptools import find_packages, setup - -# Read the contents of README file -this_directory = Path(__file__).parent -long_description = (this_directory / "README.md").read_text(encoding="utf-8") - - -# Read requirements -def read_requirements(filename): - """Read requirements from file.""" - filepath = this_directory / filename - if not filepath.exists(): - return [] - with open(filepath, encoding="utf-8") as f: - return [ - line.strip() - for line in f - if line.strip() and not line.startswith("#") and not line.startswith("-r") - ] - - -# Core requirements -install_requires = [ - "httpx>=0.24.0", - "pydantic>=2.0.0", - "python-dotenv>=1.0.0", - "numpy>=1.24.0", -] - -# Optional dependencies for integrations -extras_require = { - "sklearn": ["scikit-learn>=1.3.0"], - "pytorch": ["torch>=2.0.0"], - "tensorflow": ["tensorflow>=2.13.0"], - "transformers": ["transformers>=4.30.0"], - "langchain": ["langchain>=0.0.200"], - "boosting": ["xgboost>=1.7.0", "lightgbm>=4.0.0"], - "all": [ - "scikit-learn>=1.3.0", - "torch>=2.0.0", - "tensorflow>=2.13.0", - "transformers>=4.30.0", - "langchain>=0.0.200", - "xgboost>=1.7.0", - "lightgbm>=4.0.0", - ], - "dev": [ - "pytest>=7.4.0", - "pytest-asyncio>=0.21.0", - "pytest-cov>=4.1.0", - "black>=23.11.0", - "isort>=5.12.0", - "flake8>=6.1.0", - "pylint>=3.0.0", - "mypy>=1.7.0", - "bandit>=1.7.5", - ], -} - -setup( - name="whiteboxai-sdk", - version="0.1.0", - author="AgentaFlow", - author_email="support@agentaflow.com", - description="Official Python SDK for WhiteBoxAI ML monitoring and observability", - long_description=long_description, - long_description_content_type="text/markdown", - url="https://github.com/AgentaFlow/whitebox-python-sdk", - project_urls={ - "Bug Reports": "https://github.com/AgentaFlow/whitebox-python-sdk/issues", - "Source": "https://github.com/AgentaFlow/whitebox-python-sdk", - "Documentation": "https://whitebox.agentaflow.com", - }, - packages=find_packages(where="src"), - package_dir={"": "src"}, - classifiers=[ - "Development Status :: 3 - Alpha", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", - "License :: OSI Approved :: MIT License", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Operating System :: OS Independent", - ], - python_requires=">=3.9", - install_requires=install_requires, - extras_require=extras_require, - include_package_data=True, - zip_safe=False, -) diff --git a/src/whiteboxai/__init__.py b/src/whiteboxai/__init__.py deleted file mode 100644 index 5fbf541..0000000 --- a/src/whiteboxai/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -""" -WhiteBoxAI Python SDK - -Official Python SDK for WhiteBoxAI - AI Observability & Explainability Platform. -""" - -__version__ = "0.2.0" -__author__ = "WhiteBoxAI Team" -__license__ = "MIT" - -from whiteboxai.client import WhiteBoxAI -from whiteboxai.decorators import monitor_model, monitor_prediction -from whiteboxai.git_utils import GitContext, detect_git_context, validate_git_context -from whiteboxai.monitor import ModelMonitor - -__all__ = [ - "WhiteBoxAI", - "ModelMonitor", - "monitor_model", - "monitor_prediction", - "GitContext", - "detect_git_context", - "validate_git_context", -] diff --git a/src/whiteboxai/__version__.py b/src/whiteboxai/__version__.py deleted file mode 100644 index afa047f..0000000 --- a/src/whiteboxai/__version__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Version information for WhiteBoxAI SDK.""" - -__version__ = "0.2.0" -__author__ = "AgentaFlow" -__email__ = "support@agentaflow.com" -__license__ = "MIT" diff --git a/src/whiteboxai/cache.py b/src/whiteboxai/cache.py deleted file mode 100644 index 8930c10..0000000 --- a/src/whiteboxai/cache.py +++ /dev/null @@ -1,117 +0,0 @@ -""" -Local Caching - -Simple TTL-based cache for API responses. -""" - -import hashlib -import json -import time -from collections import OrderedDict -from typing import Any, Optional - - -class TTLCache: - """ - Time-to-live cache with LRU eviction. - - Features: - - TTL-based expiration - - LRU eviction when max size reached - - Thread-safe operations - """ - - def __init__(self, max_size: int = 1000, ttl: int = 3600): - """ - Initialize cache. - - Args: - max_size: Maximum number of items to cache - ttl: Time-to-live in seconds - """ - self.max_size = max_size - self.ttl = ttl - self._cache: OrderedDict = OrderedDict() - self._timestamps: dict = {} - - def get(self, key: str) -> Optional[Any]: - """ - Get value from cache. - - Args: - key: Cache key - - Returns: - Cached value or None if not found/expired - """ - if key not in self._cache: - return None - - # Check expiration - timestamp = self._timestamps.get(key, 0) - if time.time() - timestamp > self.ttl: - # Expired, remove from cache - del self._cache[key] - del self._timestamps[key] - return None - - # Move to end (most recently used) - self._cache.move_to_end(key) - return self._cache[key] - - def set(self, key: str, value: Any) -> None: - """ - Set value in cache. - - Args: - key: Cache key - value: Value to cache - """ - # Remove if already exists - if key in self._cache: - del self._cache[key] - del self._timestamps[key] - - # Check size limit - if len(self._cache) >= self.max_size: - # Remove least recently used - oldest_key = next(iter(self._cache)) - del self._cache[oldest_key] - del self._timestamps[oldest_key] - - # Add to cache - self._cache[key] = value - self._timestamps[key] = time.time() - - def clear(self) -> None: - """Clear all cached items.""" - self._cache.clear() - self._timestamps.clear() - - def __contains__(self, key: str) -> bool: - """Check if key exists and is not expired.""" - return self.get(key) is not None - - def __len__(self) -> int: - """Get number of cached items.""" - return len(self._cache) - - -def cache_key(*args, **kwargs) -> str: - """ - Generate cache key from arguments. - - Args: - *args: Positional arguments - **kwargs: Keyword arguments - - Returns: - Cache key string - """ - # Create deterministic key from arguments - key_data = { - "args": args, - "kwargs": sorted(kwargs.items()), - } - key_str = json.dumps(key_data, sort_keys=True) - return hashlib.md5(key_str.encode()).hexdigest() diff --git a/src/whiteboxai/client.py b/src/whiteboxai/client.py deleted file mode 100644 index 705c027..0000000 --- a/src/whiteboxai/client.py +++ /dev/null @@ -1,354 +0,0 @@ -""" -WhiteBoxAI Client - -Main client class for interacting with the WhiteBoxAI API. -""" - -import asyncio -import logging -from typing import Any, Dict, Optional -from urllib.parse import urljoin - -import httpx -from tenacity import retry, stop_after_attempt, wait_exponential - -from whiteboxai.config import Config -from whiteboxai.exceptions import APIError, AuthenticationError, RateLimitError, ValidationError -from whiteboxai.resources import ( - AgentWorkflowsResource, - AlertsResource, - DriftResource, - ExplanationsResource, - ModelsResource, - PredictionsResource, -) - -logger = logging.getLogger(__name__) - - -class WhiteBoxAI: - """ - Main client for WhiteBoxAI SDK. - - Provides access to all WhiteBoxAI API resources including models, predictions, - explanations, drift detection, and alerting. - - Args: - api_key: WhiteBoxAI API key - base_url: Base URL for WhiteBoxAI API (default: https://api.whiteboxai.io) - timeout: Request timeout in seconds (default: 30) - max_retries: Maximum number of retry attempts (default: 3) - enable_offline: Enable offline mode (queues operations when API unavailable) - offline_dir: Directory for offline queue storage (default: ./whiteboxai_offline) - offline_max_queue_size: Maximum operations to queue (default: 10000, 0 = unlimited) - offline_auto_sync: Enable automatic syncing (default: True) - offline_sync_interval: Seconds between sync attempts (default: 60) - **kwargs: Additional configuration options - - Example: - >>> client = WhiteBoxAI(api_key="your_api_key") - >>> model = client.models.register(name="fraud_detection", model_type="classification") - >>> client.predictions.log(model_id=model.id, inputs=features, outputs=prediction) - - With offline mode: - >>> client = WhiteBoxAI( - ... api_key="your_api_key", - ... enable_offline=True, - ... offline_dir="./offline_queue" - ... ) - >>> # Operations queued when offline - >>> client.predictions.log(...) # Queued if API unavailable - """ - - def __init__( - self, - api_key: str, - base_url: str = "https://api.whiteboxai.io", - timeout: int = 30, - max_retries: int = 3, - enable_offline: bool = False, - offline_dir: str = "./whiteboxai_offline", - offline_max_queue_size: int = 10000, - offline_auto_sync: bool = True, - offline_sync_interval: int = 60, - **kwargs: Any, - ): - """Initialize WhiteBoxAI client.""" - self.config = Config( - api_key=api_key, - base_url=base_url, - timeout=timeout, - max_retries=max_retries, - **kwargs, - ) - - # Initialize HTTP clients - self._sync_client: Optional[httpx.Client] = None - self._async_client: Optional[httpx.AsyncClient] = None - - # Initialize offline mode - self._offline_manager = None - if enable_offline: - from whiteboxai.offline import OfflineManager - - self._offline_manager = OfflineManager( - offline_dir=offline_dir, - max_queue_size=offline_max_queue_size, - auto_sync=offline_auto_sync, - sync_interval=offline_sync_interval, - ) - self._offline_manager.set_client(self) - logger.info("Offline mode enabled") - - # Initialize resource managers - self.models = ModelsResource(self) - self.predictions = PredictionsResource(self) - self.explanations = ExplanationsResource(self) - self.drift = DriftResource(self) - self.alerts = AlertsResource(self) - self.agent_workflows = AgentWorkflowsResource(self) - - @property - def sync_client(self) -> httpx.Client: - """Get or create synchronous HTTP client.""" - if self._sync_client is None: - self._sync_client = httpx.Client( - base_url=self.config.base_url, - headers=self._get_headers(), - timeout=self.config.timeout, - ) - return self._sync_client - - @property - def async_client(self) -> httpx.AsyncClient: - """Get or create asynchronous HTTP client.""" - if self._async_client is None: - self._async_client = httpx.AsyncClient( - base_url=self.config.base_url, - headers=self._get_headers(), - timeout=self.config.timeout, - ) - return self._async_client - - def _get_headers(self) -> Dict[str, str]: - """Get request headers with authentication.""" - return { - "Authorization": f"Bearer {self.config.api_key}", - "Content-Type": "application/json", - "User-Agent": f"whiteboxai-python-sdk/{self.config.sdk_version}", - } - - @retry( - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=1, min=2, max=10), - reraise=True, - ) - def request( - self, - method: str, - endpoint: str, - data: Optional[Dict[str, Any]] = None, - params: Optional[Dict[str, Any]] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Make synchronous HTTP request with retry logic. - - Args: - method: HTTP method (GET, POST, PUT, DELETE, etc.) - endpoint: API endpoint path - data: Request body data - params: Query parameters - **kwargs: Additional request options - - Returns: - Response data as dictionary - - Raises: - APIError: For general API errors - AuthenticationError: For authentication failures - RateLimitError: For rate limit errors - ValidationError: For validation errors - """ - url = urljoin(self.config.base_url, endpoint) - - try: - response = self.sync_client.request( - method=method, url=url, json=data, params=params, **kwargs - ) - return self._handle_response(response) - except httpx.HTTPError as e: - raise APIError(f"Request failed: {str(e)}") - - @retry( - stop=stop_after_attempt(3), - wait=wait_exponential(multiplier=1, min=2, max=10), - reraise=True, - ) - async def arequest( - self, - method: str, - endpoint: str, - data: Optional[Dict[str, Any]] = None, - params: Optional[Dict[str, Any]] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Make asynchronous HTTP request with retry logic. - - Args: - method: HTTP method (GET, POST, PUT, DELETE, etc.) - endpoint: API endpoint path - data: Request body data - params: Query parameters - **kwargs: Additional request options - - Returns: - Response data as dictionary - - Raises: - APIError: For general API errors - AuthenticationError: For authentication failures - RateLimitError: For rate limit errors - ValidationError: For validation errors - """ - url = urljoin(self.config.base_url, endpoint) - - try: - response = await self.async_client.request( - method=method, url=url, json=data, params=params, **kwargs - ) - return self._handle_response(response) - except httpx.HTTPError as e: - raise APIError(f"Request failed: {str(e)}") - - def _handle_response(self, response: httpx.Response) -> Dict[str, Any]: - """ - Handle HTTP response and raise appropriate exceptions. - - Args: - response: HTTP response - - Returns: - Response data as dictionary - - Raises: - AuthenticationError: For 401 status codes - RateLimitError: For 429 status codes - ValidationError: For 422 status codes - APIError: For other error status codes - """ - if response.status_code == 401: - raise AuthenticationError("Invalid API key or authentication failed") - elif response.status_code == 429: - raise RateLimitError("Rate limit exceeded. Please retry later.") - elif response.status_code == 422: - try: - error_data = response.json() - raise ValidationError(f"Validation error: {error_data}") - except Exception: - raise ValidationError("Validation error occurred") - elif response.status_code >= 400: - try: - error_data = response.json() - message = error_data.get("detail", response.text) - except Exception: - message = response.text - raise APIError(f"API error ({response.status_code}): {message}") - - try: - return response.json() - except Exception: - return {"status": "success"} - - def is_offline_enabled(self) -> bool: - """Check if offline mode is enabled.""" - return self._offline_manager is not None - - def sync_offline_queue(self, batch_size: int = 100) -> Dict[str, int]: - """ - Manually sync offline queue with API. - - Args: - batch_size: Number of operations to sync per batch - - Returns: - Dictionary with sync statistics - - Raises: - ValueError: If offline mode is not enabled - """ - if not self._offline_manager: - raise ValueError("Offline mode is not enabled") - - return self._offline_manager.sync(batch_size=batch_size) - - def get_offline_status(self) -> Dict[str, Any]: - """ - Get offline mode status. - - Returns: - Status dictionary with queue stats - - Raises: - ValueError: If offline mode is not enabled - """ - if not self._offline_manager: - raise ValueError("Offline mode is not enabled") - - return self._offline_manager.get_status() - - def cleanup_offline_queue(self, older_than_days: int = 7): - """ - Clean up old completed operations from offline queue. - - Args: - older_than_days: Remove operations older than this many days - - Raises: - ValueError: If offline mode is not enabled - """ - if not self._offline_manager: - raise ValueError("Offline mode is not enabled") - - self._offline_manager.cleanup(older_than_days=older_than_days) - - def close(self) -> None: - """Close HTTP clients and cleanup resources.""" - # Stop offline sync if enabled - if self._offline_manager: - self._offline_manager.stop_auto_sync() - logger.info("Offline mode stopped") - - if self._sync_client: - self._sync_client.close() - if self._async_client: - asyncio.run(self._async_client.aclose()) - - async def aclose(self) -> None: - """Asynchronously close HTTP clients and cleanup resources.""" - # Stop offline sync if enabled - if self._offline_manager: - self._offline_manager.stop_auto_sync() - logger.info("Offline mode stopped") - - if self._async_client: - await self._async_client.aclose() - if self._sync_client: - self._sync_client.close() - - def __enter__(self) -> "WhiteBoxAI": - """Context manager entry.""" - return self - - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - """Context manager exit.""" - self.close() - - async def __aenter__(self) -> "WhiteBoxAI": - """Async context manager entry.""" - return self - - async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: - """Async context manager exit.""" - await self.aclose() diff --git a/src/whiteboxai/config.py b/src/whiteboxai/config.py deleted file mode 100644 index b6c7b1d..0000000 --- a/src/whiteboxai/config.py +++ /dev/null @@ -1,90 +0,0 @@ -""" -SDK Configuration -""" - -import os -from typing import Any, Optional - - -class Config: - """ - SDK configuration class. - - Manages configuration settings for the WhiteBoxAI SDK including API keys, - URLs, timeouts, and feature flags. - - Args: - api_key: WhiteBoxAI API key - base_url: Base URL for WhiteBoxAI API - timeout: Request timeout in seconds - max_retries: Maximum number of retry attempts - **kwargs: Additional configuration options - """ - - def __init__( - self, - api_key: Optional[str] = None, - base_url: Optional[str] = None, - timeout: int = 30, - max_retries: int = 3, - **kwargs: Any, - ): - """Initialize configuration.""" - # API key (from parameter or environment) - self.api_key = api_key or os.getenv("EXPLAINAI_API_KEY") - if not self.api_key: - raise ValueError( - "API key is required. Provide via 'api_key' parameter or " - "EXPLAINAI_API_KEY environment variable." - ) - - # Base URL - self.base_url = base_url or os.getenv("EXPLAINAI_BASE_URL") or "https://api.whiteboxai.io" - - # Request settings - self.timeout = timeout - self.max_retries = max_retries - - # Feature flags - self.enable_caching = kwargs.get("enable_caching", True) - self.enable_offline_mode = kwargs.get("enable_offline_mode", False) - self.enable_privacy_filters = kwargs.get("enable_privacy_filters", True) - self.enable_sampling = kwargs.get("enable_sampling", True) - - # Sampling settings - self.sampling_rate = kwargs.get("sampling_rate", 1.0) # 100% by default - - # Cache settings - self.cache_ttl = kwargs.get("cache_ttl", 3600) # 1 hour default - self.cache_max_size = kwargs.get("cache_max_size", 1000) - - # Privacy settings - self.detect_pii = kwargs.get("detect_pii", True) - self.mask_pii = kwargs.get("mask_pii", True) - - # Performance settings - self.batch_size = kwargs.get("batch_size", 100) - self.async_enabled = kwargs.get("async_enabled", True) - - # SDK metadata - self.sdk_version = "0.2.0" - - def to_dict(self) -> dict: - """Convert configuration to dictionary.""" - return { - "base_url": self.base_url, - "timeout": self.timeout, - "max_retries": self.max_retries, - "enable_caching": self.enable_caching, - "enable_offline_mode": self.enable_offline_mode, - "enable_privacy_filters": self.enable_privacy_filters, - "enable_sampling": self.enable_sampling, - "sampling_rate": self.sampling_rate, - "cache_ttl": self.cache_ttl, - "cache_max_size": self.cache_max_size, - "detect_pii": self.detect_pii, - "mask_pii": self.mask_pii, - "batch_size": self.batch_size, - "async_enabled": self.async_enabled, - "sdk_version": self.sdk_version, - } diff --git a/src/whiteboxai/decorators.py b/src/whiteboxai/decorators.py deleted file mode 100644 index 21f2a13..0000000 --- a/src/whiteboxai/decorators.py +++ /dev/null @@ -1,323 +0,0 @@ -""" -Monitoring Decorators - -Decorators for automatic model and prediction monitoring. -""" - -import functools -import inspect -import time -from typing import Any, Callable, Dict, Optional - -from whiteboxai.monitor import ModelMonitor - - -def monitor_model( - monitor: ModelMonitor, - input_keys: Optional[list] = None, - output_key: Optional[str] = None, - explain: bool = False, -): - """ - Decorator to automatically monitor model predictions. - - Args: - monitor: ModelMonitor instance - input_keys: Keys to extract inputs from function args/kwargs - output_key: Key to extract output from function return value - explain: Generate explanations for predictions - - Example: - ```python - from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_model - - client = WhiteBoxAI(api_key="your-api-key") - monitor = ModelMonitor(client, model_id=123) - - @monitor_model(monitor, input_keys=["features"], explain=True) - def predict(features): - # Your prediction logic - return model.predict(features) - - # Predictions are automatically logged - result = predict(features=[1.0, 2.0, 3.0]) - ``` - """ - - def decorator(func: Callable) -> Callable: - @functools.wraps(func) - def wrapper(*args, **kwargs): - start_time = time.time() - - # Call original function - result = func(*args, **kwargs) - - # Extract inputs - inputs = _extract_inputs(func, args, kwargs, input_keys) - - # Extract output - output = _extract_output(result, output_key) - - # Calculate inference time - inference_time = time.time() - start_time - - # Log prediction - metadata = {"inference_time_ms": inference_time * 1000} - monitor.log_prediction( - inputs=inputs, - output=output, - explain=explain, - metadata=metadata, - ) - - return result - - @functools.wraps(func) - async def async_wrapper(*args, **kwargs): - start_time = time.time() - - # Call original function - result = await func(*args, **kwargs) - - # Extract inputs - inputs = _extract_inputs(func, args, kwargs, input_keys) - - # Extract output - output = _extract_output(result, output_key) - - # Calculate inference time - inference_time = time.time() - start_time - - # Log prediction - metadata = {"inference_time_ms": inference_time * 1000} - await monitor.alog_prediction( - inputs=inputs, - output=output, - explain=explain, - metadata=metadata, - ) - - return result - - # Return appropriate wrapper based on function type - if inspect.iscoroutinefunction(func): - return async_wrapper - return wrapper - - return decorator - - -def monitor_prediction( - monitor: ModelMonitor, - input_extractor: Optional[Callable] = None, - output_extractor: Optional[Callable] = None, - explain: bool = False, - metadata_extractor: Optional[Callable] = None, -): - """ - Decorator to monitor individual predictions with custom extractors. - - Args: - monitor: ModelMonitor instance - input_extractor: Function to extract inputs from args/kwargs - output_extractor: Function to extract output from result - explain: Generate explanations for predictions - metadata_extractor: Function to extract additional metadata - - Example: - ```python - from whiteboxai import WhiteBoxAI, ModelMonitor, monitor_prediction - - client = WhiteBoxAI(api_key="your-api-key") - monitor = ModelMonitor(client, model_id=123) - - def extract_inputs(args, kwargs): - return {"features": kwargs.get("data")} - - def extract_output(result): - return result["prediction"] - - @monitor_prediction( - monitor, - input_extractor=extract_inputs, - output_extractor=extract_output, - explain=True - ) - def predict(data): - # Your prediction logic - return {"prediction": 0.85, "confidence": 0.92} - - result = predict(data=[1.0, 2.0, 3.0]) - ``` - """ - - def decorator(func: Callable) -> Callable: - @functools.wraps(func) - def wrapper(*args, **kwargs): - start_time = time.time() - - # Call original function - result = func(*args, **kwargs) - - # Extract inputs - if input_extractor: - inputs = input_extractor(args, kwargs) - else: - inputs = _default_input_extractor(func, args, kwargs) - - # Extract output - if output_extractor: - output = output_extractor(result) - else: - output = result - - # Extract metadata - metadata = {"inference_time_ms": (time.time() - start_time) * 1000} - if metadata_extractor: - custom_metadata = metadata_extractor(args, kwargs, result) - metadata.update(custom_metadata) - - # Log prediction - monitor.log_prediction( - inputs=inputs, - output=output, - explain=explain, - metadata=metadata, - ) - - return result - - @functools.wraps(func) - async def async_wrapper(*args, **kwargs): - start_time = time.time() - - # Call original function - result = await func(*args, **kwargs) - - # Extract inputs - if input_extractor: - inputs = input_extractor(args, kwargs) - else: - inputs = _default_input_extractor(func, args, kwargs) - - # Extract output - if output_extractor: - output = output_extractor(result) - else: - output = result - - # Extract metadata - metadata = {"inference_time_ms": (time.time() - start_time) * 1000} - if metadata_extractor: - custom_metadata = metadata_extractor(args, kwargs, result) - metadata.update(custom_metadata) - - # Log prediction - await monitor.alog_prediction( - inputs=inputs, - output=output, - explain=explain, - metadata=metadata, - ) - - return result - - # Return appropriate wrapper based on function type - if inspect.iscoroutinefunction(func): - return async_wrapper - return wrapper - - return decorator - - -def _extract_inputs( - func: Callable, - args: tuple, - kwargs: dict, - input_keys: Optional[list] = None, -) -> Any: - """Extract inputs from function arguments.""" - if input_keys: - # Extract specified keys - sig = inspect.signature(func) - bound_args = sig.bind(*args, **kwargs) - bound_args.apply_defaults() - - inputs = {} - for key in input_keys: - if key in bound_args.arguments: - inputs[key] = bound_args.arguments[key] - - return inputs - else: - # Use all arguments - return _default_input_extractor(func, args, kwargs) - - -def _default_input_extractor(func: Callable, args: tuple, kwargs: dict) -> Dict[str, Any]: - """Default input extractor.""" - sig = inspect.signature(func) - bound_args = sig.bind(*args, **kwargs) - bound_args.apply_defaults() - return dict(bound_args.arguments) - - -def _extract_output(result: Any, output_key: Optional[str] = None) -> Any: - """Extract output from function result.""" - if output_key and isinstance(result, dict): - return result.get(output_key) - return result - - -def monitor_performance(threshold_ms: Optional[float] = None): - """ - Decorator to monitor function performance. - - Args: - threshold_ms: Optional performance threshold in milliseconds - - Example: - ```python - from whiteboxai.decorators import monitor_performance - - @monitor_performance(threshold_ms=1000) - def slow_function(): - # Your code here - pass - ``` - """ - - def decorator(func: Callable) -> Callable: - @functools.wraps(func) - def wrapper(*args, **kwargs): - start_time = time.time() - result = func(*args, **kwargs) - elapsed_ms = (time.time() - start_time) * 1000 - - if threshold_ms and elapsed_ms > threshold_ms: - print( - f"Warning: {func.__name__} took {elapsed_ms:.2f}ms " - f"(threshold: {threshold_ms}ms)" - ) - - return result - - @functools.wraps(func) - async def async_wrapper(*args, **kwargs): - start_time = time.time() - result = await func(*args, **kwargs) - elapsed_ms = (time.time() - start_time) * 1000 - - if threshold_ms and elapsed_ms > threshold_ms: - print( - f"Warning: {func.__name__} took {elapsed_ms:.2f}ms " - f"(threshold: {threshold_ms}ms)" - ) - - return result - - if inspect.iscoroutinefunction(func): - return async_wrapper - return wrapper - - return decorator diff --git a/src/whiteboxai/exceptions.py b/src/whiteboxai/exceptions.py deleted file mode 100644 index 7932c7f..0000000 --- a/src/whiteboxai/exceptions.py +++ /dev/null @@ -1,57 +0,0 @@ -""" -SDK Exceptions -""" - - -class WhiteBoxAIError(Exception): - """Base exception for WhiteBoxAI SDK.""" - - pass - - -class APIError(WhiteBoxAIError): - """API request error.""" - - pass - - -class AuthenticationError(WhiteBoxAIError): - """Authentication error.""" - - pass - - -class RateLimitError(WhiteBoxAIError): - """Rate limit exceeded error.""" - - pass - - -class ValidationError(WhiteBoxAIError): - """Validation error.""" - - pass - - -class NotFoundError(WhiteBoxAIError): - """Resource not found error.""" - - pass - - -class ConfigurationError(WhiteBoxAIError): - """Configuration error.""" - - pass - - -class IntegrationError(WhiteBoxAIError): - """Framework integration error.""" - - pass - - -class CacheError(WhiteBoxAIError): - """Cache operation error.""" - - pass diff --git a/src/whiteboxai/git_utils.py b/src/whiteboxai/git_utils.py deleted file mode 100644 index 925817a..0000000 --- a/src/whiteboxai/git_utils.py +++ /dev/null @@ -1,331 +0,0 @@ -""" -Git Auto-Detection Utilities - -Utilities for automatically detecting Git context (repository, commit, branch) -for model registration. -""" - -import logging -import os -import subprocess -from typing import Dict, Optional - -logger = logging.getLogger(__name__) - - -class GitContext: - """Git repository context information.""" - - def __init__( - self, - repository_url: Optional[str] = None, - commit_sha: Optional[str] = None, - commit_message: Optional[str] = None, - commit_author: Optional[str] = None, - branch: Optional[str] = None, - tag: Optional[str] = None, - is_dirty: bool = False, - ): - """Initialize Git context.""" - self.repository_url = repository_url - self.commit_sha = commit_sha - self.commit_message = commit_message - self.commit_author = commit_author - self.branch = branch - self.tag = tag - self.is_dirty = is_dirty - - def to_dict(self) -> Dict[str, Optional[str]]: - """Convert to dictionary for API submission.""" - return { - "github_repository_url": self.repository_url, - "github_commit_hash": self.commit_sha, - "github_commit_message": self.commit_message, - "github_commit_author": self.commit_author, - "github_branch": self.branch, - "github_tag": self.tag, - } - - def __repr__(self) -> str: - """String representation.""" - parts = [] - if self.repository_url: - parts.append(f"repo={self.repository_url}") - if self.commit_sha: - parts.append(f"commit={self.commit_sha[:7]}") - if self.branch: - parts.append(f"branch={self.branch}") - if self.tag: - parts.append(f"tag={self.tag}") - return f"GitContext({', '.join(parts)})" - - -def detect_git_context(path: Optional[str] = None) -> Optional[GitContext]: - """ - Auto-detect Git context from the current directory or specified path. - - Args: - path: Path to check for Git repository (defaults to current directory) - - Returns: - GitContext object if Git repository found, None otherwise - - Example: - >>> context = detect_git_context() - >>> if context: - ... print(f"Repository: {context.repository_url}") - ... print(f"Commit: {context.commit_sha}") - ... print(f"Branch: {context.branch}") - """ - try: - import git - except ImportError: - logger.warning("GitPython not installed. Install with: pip install gitpython") - return _detect_git_context_subprocess(path) - - try: - # Find repository - search_path = path or os.getcwd() - repo = git.Repo(search_path, search_parent_directories=True) - - # Get repository URL - repository_url = None - try: - remote = repo.remote("origin") - repository_url = remote.url - # Convert SSH URLs to HTTPS - if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace("git@github.com:", "https://github.com/") - if repository_url.endswith(".git"): - repository_url = repository_url[:-4] - except Exception as e: - logger.debug(f"Could not get remote URL: {e}") - - # Get current commit - commit_sha = None - commit_message = None - commit_author = None - - try: - head_commit = repo.head.commit - commit_sha = head_commit.hexsha - commit_message = head_commit.message.strip() - commit_author = str(head_commit.author) - except Exception as e: - logger.debug(f"Could not get commit info: {e}") - - # Get current branch - branch = None - try: - if not repo.head.is_detached: - branch = repo.active_branch.name - except Exception as e: - logger.debug(f"Could not get branch: {e}") - - # Get current tag (if on a tag) - tag = None - try: - tags = [tag for tag in repo.tags if tag.commit == repo.head.commit] - if tags: - tag = tags[0].name - except Exception as e: - logger.debug(f"Could not get tag: {e}") - - # Check if working directory is dirty - is_dirty = repo.is_dirty(untracked_files=True) - if is_dirty: - logger.warning("Working directory has uncommitted changes") - - context = GitContext( - repository_url=repository_url, - commit_sha=commit_sha, - commit_message=commit_message, - commit_author=commit_author, - branch=branch, - tag=tag, - is_dirty=is_dirty, - ) - - logger.info(f"Detected Git context: {context}") - return context - - except git.InvalidGitRepositoryError: - logger.debug(f"No Git repository found at {search_path}") - return None - except Exception as e: - logger.error(f"Error detecting Git context: {e}") - return None - - -def _detect_git_context_subprocess(path: Optional[str] = None) -> Optional[GitContext]: - """ - Fallback Git detection using subprocess (when GitPython not available). - - Args: - path: Path to check for Git repository - - Returns: - GitContext object or None - """ - try: - cwd = path or os.getcwd() - - # Check if Git is available - subprocess.run( - ["git", "--version"], - capture_output=True, - check=True, - cwd=cwd, - ) - - # Get repository URL - repository_url = None - try: - result = subprocess.run( - ["git", "remote", "get-url", "origin"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - repository_url = result.stdout.strip() - - # Convert SSH to HTTPS - if repository_url.startswith("git@github.com:"): - repository_url = repository_url.replace("git@github.com:", "https://github.com/") - if repository_url.endswith(".git"): - repository_url = repository_url[:-4] - except subprocess.CalledProcessError: - pass - - # Get commit SHA - commit_sha = None - try: - result = subprocess.run( - ["git", "rev-parse", "HEAD"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - commit_sha = result.stdout.strip() - except subprocess.CalledProcessError: - pass - - # Get commit message - commit_message = None - try: - result = subprocess.run( - ["git", "log", "-1", "--pretty=%B"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - commit_message = result.stdout.strip() - except subprocess.CalledProcessError: - pass - - # Get commit author - commit_author = None - try: - result = subprocess.run( - ["git", "log", "-1", "--pretty=%an <%ae>"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - commit_author = result.stdout.strip() - except subprocess.CalledProcessError: - pass - - # Get branch - branch = None - try: - result = subprocess.run( - ["git", "rev-parse", "--abbrev-ref", "HEAD"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - branch_output = result.stdout.strip() - if branch_output != "HEAD": # Not in detached HEAD - branch = branch_output - except subprocess.CalledProcessError: - pass - - # Get tag - tag = None - try: - result = subprocess.run( - ["git", "describe", "--exact-match", "--tags", "HEAD"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - tag = result.stdout.strip() - except subprocess.CalledProcessError: - pass - - # Check if dirty - is_dirty = False - try: - result = subprocess.run( - ["git", "status", "--porcelain"], - capture_output=True, - text=True, - check=True, - cwd=cwd, - ) - is_dirty = bool(result.stdout.strip()) - except subprocess.CalledProcessError: - pass - - context = GitContext( - repository_url=repository_url, - commit_sha=commit_sha, - commit_message=commit_message, - commit_author=commit_author, - branch=branch, - tag=tag, - is_dirty=is_dirty, - ) - - logger.info(f"Detected Git context (subprocess): {context}") - return context - - except (subprocess.CalledProcessError, FileNotFoundError): - logger.debug("Git not available or not a Git repository") - return None - except Exception as e: - logger.error(f"Error detecting Git context with subprocess: {e}") - return None - - -def validate_git_context(context: GitContext, require_clean: bool = False) -> bool: - """ - Validate Git context before model registration. - - Args: - context: GitContext object - require_clean: Whether to require a clean working directory - - Returns: - True if valid, False otherwise - """ - if not context: - logger.warning("No Git context provided") - return False - - if not context.commit_sha: - logger.error("Git context missing commit SHA") - return False - - if require_clean and context.is_dirty: - logger.error("Working directory has uncommitted changes (require_clean=True)") - return False - - return True diff --git a/src/whiteboxai/integrations/__init__.py b/src/whiteboxai/integrations/__init__.py deleted file mode 100644 index c69a821..0000000 --- a/src/whiteboxai/integrations/__init__.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Framework integrations for WhiteBoxAI SDK.""" - -# Scikit-learn integration -try: - from .sklearn import SklearnMonitor, SklearnWrapper - - __all__ = ["SklearnMonitor", "SklearnWrapper"] -except ImportError: - pass - -# PyTorch integration -try: - from .pytorch import TorchMonitor, TorchWrapper - - if "__all__" in dir(): - __all__.extend(["TorchMonitor", "TorchWrapper"]) - else: - __all__ = ["TorchMonitor", "TorchWrapper"] -except ImportError: - pass - -# TensorFlow/Keras integration -try: - from .tensorflow import KerasMonitor, WhiteBoxAICallback, wrap_keras_model - - if "__all__" in dir(): - __all__.extend(["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"]) - else: - __all__ = ["KerasMonitor", "WhiteBoxAICallback", "wrap_keras_model"] -except ImportError: - pass - -# Hugging Face Transformers integration -try: - from .transformers import ( - TransformersMonitor, - TransformersPipelineWrapper, - wrap_transformers_pipeline, - ) - - if "__all__" in dir(): - __all__.extend( - ["TransformersMonitor", "TransformersPipelineWrapper", "wrap_transformers_pipeline"] - ) - else: - __all__ = [ - "TransformersMonitor", - "TransformersPipelineWrapper", - "wrap_transformers_pipeline", - ] -except ImportError: - pass - -# LangChain integration -try: - from .langchain import LangChainMonitor, WhiteBoxAICallbackHandler, wrap_langchain_chain - - if "__all__" in dir(): - __all__.extend(["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"]) - else: - __all__ = ["LangChainMonitor", "WhiteBoxAICallbackHandler", "wrap_langchain_chain"] -except ImportError: - pass - -# XGBoost/LightGBM integration -try: - from .boosting import LightGBMMonitor, XGBoostMonitor, wrap_lightgbm_model, wrap_xgboost_model - - if "__all__" in dir(): - __all__.extend( - ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] - ) - else: - __all__ = ["XGBoostMonitor", "LightGBMMonitor", "wrap_xgboost_model", "wrap_lightgbm_model"] -except ImportError: - pass - -# CrewAI integration -try: - from .crewai_monitor import CrewAIMonitor, monitor_crew - - if "__all__" in dir(): - __all__.extend(["CrewAIMonitor", "monitor_crew"]) - else: - __all__ = ["CrewAIMonitor", "monitor_crew"] -except ImportError: - pass - -# LangChain Multi-Agent integration -try: - from .langchain_agents import ( - LangGraphMultiAgentMonitor, - MultiAgentCallbackHandler, - monitor_langchain_agent, - ) - - if "__all__" in dir(): - __all__.extend( - ["MultiAgentCallbackHandler", "LangGraphMultiAgentMonitor", "monitor_langchain_agent"] - ) - else: - __all__ = [ - "MultiAgentCallbackHandler", - "LangGraphMultiAgentMonitor", - "monitor_langchain_agent", - ] -except ImportError: - pass - -# Ensure __all__ exists even if all imports fail -if "__all__" not in dir(): - __all__ = [] diff --git a/src/whiteboxai/integrations/boosting.py b/src/whiteboxai/integrations/boosting.py deleted file mode 100644 index 210155b..0000000 --- a/src/whiteboxai/integrations/boosting.py +++ /dev/null @@ -1,667 +0,0 @@ -""" -XGBoost and LightGBM integration for WhiteBoxAI. - -This module provides monitoring capabilities for gradient boosting models -from XGBoost and LightGBM frameworks. It supports both frameworks through -a unified interface with automatic feature importance tracking, prediction -logging, and performance monitoring. - -Example: - Basic XGBoost monitoring: - - >>> import xgboost as xgb - >>> from whiteboxai import WhiteBoxAI - >>> from whiteboxai.integrations.boosting import XGBoostMonitor - >>> - >>> client = WhiteBoxAI(api_key="your-api-key") - >>> monitor = XGBoostMonitor(client=client, model_name="fraud_detector") - >>> - >>> model = xgb.XGBClassifier() - >>> model.fit(X_train, y_train) - >>> - >>> monitor.register_from_model(model, X_train, y_train) - >>> predictions = monitor.predict(model, X_test, y_test) - - LightGBM monitoring: - - >>> import lightgbm as lgb - >>> from whiteboxai.integrations.boosting import LightGBMMonitor - >>> - >>> monitor = LightGBMMonitor(client=client, model_name="churn_predictor") - >>> - >>> model = lgb.LGBMClassifier() - >>> model.fit(X_train, y_train) - >>> - >>> monitor.register_from_model(model, X_train, y_train) - >>> predictions = monitor.predict(model, X_test, y_test) -""" - -import warnings -from typing import Any, Dict, List, Optional - -import numpy as np - -from ..core import ModelMonitor - -# Optional imports - graceful degradation -try: - import xgboost as xgb - - XGBOOST_AVAILABLE = True -except ImportError: - XGBOOST_AVAILABLE = False - xgb = None - -try: - import lightgbm as lgb - - LIGHTGBM_AVAILABLE = True -except ImportError: - LIGHTGBM_AVAILABLE = False - lgb = None - - -class XGBoostMonitor(ModelMonitor): - """ - Monitor for XGBoost models. - - Provides comprehensive monitoring for XGBoost models including prediction - logging, feature importance tracking, and automatic model registration. - - Attributes: - client: WhiteBoxAI client instance - model_name: Name of the model being monitored - track_feature_importance: Whether to track feature importance with predictions - importance_type: Type of importance to track ('weight', 'gain', 'cover', 'total_gain', 'total_cover') - - Example: - >>> monitor = XGBoostMonitor( - ... client=client, - ... model_name="xgb_classifier", - ... track_feature_importance=True, - ... importance_type="gain" - ... ) - >>> monitor.register_from_model(model, X_train, y_train) - >>> predictions = monitor.predict(model, X_test, y_test) - """ - - def __init__( - self, - client: Any, - model_name: str = "xgboost_model", - track_feature_importance: bool = True, - importance_type: str = "gain", - **kwargs, - ): - """ - Initialize XGBoost monitor. - - Args: - client: WhiteBoxAI client instance - model_name: Name of the model - track_feature_importance: Whether to track feature importance - importance_type: Type of importance ('weight', 'gain', 'cover', 'total_gain', 'total_cover') - **kwargs: Additional arguments passed to ModelMonitor - """ - if not XGBOOST_AVAILABLE: - raise ImportError("XGBoost is not installed. Install with: pip install xgboost") - - super().__init__(client=client, model_name=model_name, **kwargs) - self.track_feature_importance = track_feature_importance - self.importance_type = importance_type - self._feature_names: Optional[List[str]] = None - - def register_from_model( - self, - model: Any, - X_train: Optional[np.ndarray] = None, - y_train: Optional[np.ndarray] = None, - model_type: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """ - Register XGBoost model with WhiteBoxAI. - - Automatically extracts model metadata including feature names, - number of features, number of trees, and feature importance. - - Args: - model: XGBoost model (Booster, XGBClassifier, XGBRegressor, etc.) - X_train: Training features for baseline - y_train: Training labels for baseline - model_type: Model type override (auto-detected if None) - metadata: Additional metadata - - Returns: - Model ID from registration - """ - # Extract model information - model_metadata = metadata or {} - - # Get feature names - if hasattr(model, "feature_names_in_"): - self._feature_names = list(model.feature_names_in_) - model_metadata["feature_names"] = self._feature_names - elif hasattr(model, "feature_names"): - self._feature_names = model.feature_names - model_metadata["feature_names"] = self._feature_names - elif X_train is not None and hasattr(X_train, "columns"): - self._feature_names = list(X_train.columns) - model_metadata["feature_names"] = self._feature_names - - # Get number of features - if hasattr(model, "n_features_in_"): - model_metadata["num_features"] = model.n_features_in_ - elif self._feature_names: - model_metadata["num_features"] = len(self._feature_names) - - # Get number of trees/boosting rounds - if hasattr(model, "n_estimators"): - model_metadata["num_trees"] = model.n_estimators - elif hasattr(model, "best_iteration"): - model_metadata["num_trees"] = model.best_iteration - - # Get XGBoost parameters - if hasattr(model, "get_params"): - params = model.get_params() - model_metadata["xgboost_params"] = { - k: str(v) for k, v in params.items() if v is not None - } - - # Get feature importance - if self.track_feature_importance: - try: - importance = self._get_feature_importance(model) - if importance: - model_metadata["feature_importance"] = importance - except Exception as e: - warnings.warn(f"Failed to extract feature importance: {e}") - - # Detect model type - if model_type is None: - if hasattr(model, "_estimator_type"): - model_type = model._estimator_type - elif isinstance(model, xgb.XGBClassifier): - model_type = "classification" - elif isinstance(model, xgb.XGBRegressor): - model_type = "regression" - elif isinstance(model, xgb.XGBRanker): - model_type = "ranking" - else: - model_type = "classification" # Default - - model_metadata["framework"] = "xgboost" - model_metadata["xgboost_version"] = xgb.__version__ - - # Register model - model_id = self.register_model( - name=self.model_name, model_type=model_type, metadata=model_metadata - ) - - # Set baseline if training data provided - if X_train is not None and y_train is not None: - self.set_baseline(X_train, y_train) - - return model_id - - def predict( - self, - model: Any, - X: np.ndarray, - y_true: Optional[np.ndarray] = None, - log_predictions: bool = True, - metadata: Optional[Dict[str, Any]] = None, - ) -> np.ndarray: - """ - Make predictions and log to WhiteBoxAI. - - Args: - model: XGBoost model - X: Input features - y_true: True labels (optional) - log_predictions: Whether to log predictions - metadata: Additional metadata - - Returns: - Model predictions - """ - # Make predictions - predictions = model.predict(X) - - # Get prediction probabilities if classification - probabilities = None - if hasattr(model, "predict_proba"): - try: - probabilities = model.predict_proba(X) - except Exception: - pass - - # Log predictions - if log_predictions: - pred_metadata = metadata or {} - - # Add feature importance for this batch if enabled - if self.track_feature_importance: - try: - importance = self._get_feature_importance(model) - if importance: - pred_metadata["feature_importance"] = importance - except Exception as e: - warnings.warn(f"Failed to extract feature importance: {e}") - - self.log_predictions( - inputs=X, - predictions=predictions, - actuals=y_true, - probabilities=probabilities, - metadata=pred_metadata, - ) - - return predictions - - def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: - """ - Extract feature importance from XGBoost model. - - Args: - model: XGBoost model - - Returns: - Dictionary mapping feature names to importance scores - """ - try: - # Try sklearn-style feature_importances_ - if hasattr(model, "feature_importances_"): - importances = model.feature_importances_ - if self._feature_names and len(importances) == len(self._feature_names): - return dict(zip(self._feature_names, importances.tolist())) - else: - return {f"f{i}": float(v) for i, v in enumerate(importances)} - - # Try get_score method (native XGBoost) - if hasattr(model, "get_score"): - importance_dict = model.get_score(importance_type=self.importance_type) - return {k: float(v) for k, v in importance_dict.items()} - - # Try get_booster for sklearn API - if hasattr(model, "get_booster"): - booster = model.get_booster() - importance_dict = booster.get_score(importance_type=self.importance_type) - return {k: float(v) for k, v in importance_dict.items()} - - return None - except Exception: - return None - - -class LightGBMMonitor(ModelMonitor): - """ - Monitor for LightGBM models. - - Provides comprehensive monitoring for LightGBM models including prediction - logging, feature importance tracking, and automatic model registration. - - Attributes: - client: WhiteBoxAI client instance - model_name: Name of the model being monitored - track_feature_importance: Whether to track feature importance with predictions - importance_type: Type of importance to track ('split' or 'gain') - - Example: - >>> monitor = LightGBMMonitor( - ... client=client, - ... model_name="lgbm_classifier", - ... track_feature_importance=True, - ... importance_type="gain" - ... ) - >>> monitor.register_from_model(model, X_train, y_train) - >>> predictions = monitor.predict(model, X_test, y_test) - """ - - def __init__( - self, - client: Any, - model_name: str = "lightgbm_model", - track_feature_importance: bool = True, - importance_type: str = "gain", - **kwargs, - ): - """ - Initialize LightGBM monitor. - - Args: - client: WhiteBoxAI client instance - model_name: Name of the model - track_feature_importance: Whether to track feature importance - importance_type: Type of importance ('split' or 'gain') - **kwargs: Additional arguments passed to ModelMonitor - """ - if not LIGHTGBM_AVAILABLE: - raise ImportError("LightGBM is not installed. Install with: pip install lightgbm") - - super().__init__(client=client, model_name=model_name, **kwargs) - self.track_feature_importance = track_feature_importance - self.importance_type = importance_type - self._feature_names: Optional[List[str]] = None - - def register_from_model( - self, - model: Any, - X_train: Optional[np.ndarray] = None, - y_train: Optional[np.ndarray] = None, - model_type: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """ - Register LightGBM model with WhiteBoxAI. - - Automatically extracts model metadata including feature names, - number of features, number of trees, and feature importance. - - Args: - model: LightGBM model (Booster, LGBMClassifier, LGBMRegressor, etc.) - X_train: Training features for baseline - y_train: Training labels for baseline - model_type: Model type override (auto-detected if None) - metadata: Additional metadata - - Returns: - Model ID from registration - """ - # Extract model information - model_metadata = metadata or {} - - # Get feature names - if hasattr(model, "feature_name_"): - self._feature_names = model.feature_name_ - model_metadata["feature_names"] = self._feature_names - elif hasattr(model, "feature_names_in_"): - self._feature_names = list(model.feature_names_in_) - model_metadata["feature_names"] = self._feature_names - elif X_train is not None and hasattr(X_train, "columns"): - self._feature_names = list(X_train.columns) - model_metadata["feature_names"] = self._feature_names - - # Get number of features - if hasattr(model, "n_features_in_"): - model_metadata["num_features"] = model.n_features_in_ - elif self._feature_names: - model_metadata["num_features"] = len(self._feature_names) - - # Get number of trees - if hasattr(model, "n_estimators"): - model_metadata["num_trees"] = model.n_estimators - elif hasattr(model, "best_iteration_"): - model_metadata["num_trees"] = model.best_iteration_ - elif hasattr(model, "num_trees"): - model_metadata["num_trees"] = model.num_trees() - - # Get LightGBM parameters - if hasattr(model, "get_params"): - params = model.get_params() - model_metadata["lightgbm_params"] = { - k: str(v) for k, v in params.items() if v is not None - } - - # Get feature importance - if self.track_feature_importance: - try: - importance = self._get_feature_importance(model) - if importance: - model_metadata["feature_importance"] = importance - except Exception as e: - warnings.warn(f"Failed to extract feature importance: {e}") - - # Detect model type - if model_type is None: - if hasattr(model, "_estimator_type"): - model_type = model._estimator_type - elif isinstance(model, lgb.LGBMClassifier): - model_type = "classification" - elif isinstance(model, lgb.LGBMRegressor): - model_type = "regression" - elif isinstance(model, lgb.LGBMRanker): - model_type = "ranking" - else: - model_type = "classification" # Default - - model_metadata["framework"] = "lightgbm" - model_metadata["lightgbm_version"] = lgb.__version__ - - # Register model - model_id = self.register_model( - name=self.model_name, model_type=model_type, metadata=model_metadata - ) - - # Set baseline if training data provided - if X_train is not None and y_train is not None: - self.set_baseline(X_train, y_train) - - return model_id - - def predict( - self, - model: Any, - X: np.ndarray, - y_true: Optional[np.ndarray] = None, - log_predictions: bool = True, - metadata: Optional[Dict[str, Any]] = None, - ) -> np.ndarray: - """ - Make predictions and log to WhiteBoxAI. - - Args: - model: LightGBM model - X: Input features - y_true: True labels (optional) - log_predictions: Whether to log predictions - metadata: Additional metadata - - Returns: - Model predictions - """ - # Make predictions - predictions = model.predict(X) - - # Get prediction probabilities if classification - probabilities = None - if hasattr(model, "predict_proba"): - try: - probabilities = model.predict_proba(X) - except Exception: - pass - - # Log predictions - if log_predictions: - pred_metadata = metadata or {} - - # Add feature importance for this batch if enabled - if self.track_feature_importance: - try: - importance = self._get_feature_importance(model) - if importance: - pred_metadata["feature_importance"] = importance - except Exception as e: - warnings.warn(f"Failed to extract feature importance: {e}") - - self.log_predictions( - inputs=X, - predictions=predictions, - actuals=y_true, - probabilities=probabilities, - metadata=pred_metadata, - ) - - return predictions - - def _get_feature_importance(self, model: Any) -> Optional[Dict[str, float]]: - """ - Extract feature importance from LightGBM model. - - Args: - model: LightGBM model - - Returns: - Dictionary mapping feature names to importance scores - """ - try: - # Try sklearn-style feature_importances_ - if hasattr(model, "feature_importances_"): - importances = model.feature_importances_ - if self._feature_names and len(importances) == len(self._feature_names): - return dict(zip(self._feature_names, importances.tolist())) - else: - return {f"f{i}": float(v) for i, v in enumerate(importances)} - - # Try feature_importance method (native LightGBM) - if hasattr(model, "feature_importance"): - importances = model.feature_importance(importance_type=self.importance_type) - if self._feature_names and len(importances) == len(self._feature_names): - return dict(zip(self._feature_names, importances.tolist())) - else: - return {f"f{i}": float(v) for i, v in enumerate(importances)} - - # Try booster - if hasattr(model, "booster_"): - booster = model.booster_ - importances = booster.feature_importance(importance_type=self.importance_type) - feature_names = booster.feature_name() - if len(importances) == len(feature_names): - return dict(zip(feature_names, importances.tolist())) - - return None - except Exception: - return None - - -# Unified wrapper functions -def wrap_xgboost_model(model: Any, monitor: XGBoostMonitor, auto_register: bool = True) -> Any: - """ - Wrap an XGBoost model for automatic monitoring. - - This creates a wrapper that automatically logs predictions when - the model's predict method is called. - - Args: - model: XGBoost model to wrap - monitor: XGBoostMonitor instance - auto_register: Whether to auto-register the model - - Returns: - Wrapped model with monitoring - - Example: - >>> monitor = XGBoostMonitor(client=client, model_name="my_model") - >>> wrapped_model = wrap_xgboost_model(model, monitor) - >>> predictions = wrapped_model.predict(X_test) # Auto-logged - """ - if auto_register and not monitor._model_id: - monitor.register_from_model(model) - - # Store original methods - original_predict = model.predict - if hasattr(model, "predict_proba"): - original_predict_proba = model.predict_proba - else: - original_predict_proba = None - - # Wrap predict method - def wrapped_predict(X, *args, **kwargs): - predictions = original_predict(X, *args, **kwargs) - - # Log predictions - try: - monitor.log_predictions(inputs=X, predictions=predictions) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return predictions - - # Wrap predict_proba if available - if original_predict_proba: - - def wrapped_predict_proba(X, *args, **kwargs): - probabilities = original_predict_proba(X, *args, **kwargs) - - # Log with probabilities - try: - predictions = np.argmax(probabilities, axis=1) - monitor.log_predictions( - inputs=X, predictions=predictions, probabilities=probabilities - ) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return probabilities - - model.predict_proba = wrapped_predict_proba - - model.predict = wrapped_predict - - return model - - -def wrap_lightgbm_model(model: Any, monitor: LightGBMMonitor, auto_register: bool = True) -> Any: - """ - Wrap a LightGBM model for automatic monitoring. - - This creates a wrapper that automatically logs predictions when - the model's predict method is called. - - Args: - model: LightGBM model to wrap - monitor: LightGBMMonitor instance - auto_register: Whether to auto-register the model - - Returns: - Wrapped model with monitoring - - Example: - >>> monitor = LightGBMMonitor(client=client, model_name="my_model") - >>> wrapped_model = wrap_lightgbm_model(model, monitor) - >>> predictions = wrapped_model.predict(X_test) # Auto-logged - """ - if auto_register and not monitor._model_id: - monitor.register_from_model(model) - - # Store original methods - original_predict = model.predict - if hasattr(model, "predict_proba"): - original_predict_proba = model.predict_proba - else: - original_predict_proba = None - - # Wrap predict method - def wrapped_predict(X, *args, **kwargs): - predictions = original_predict(X, *args, **kwargs) - - # Log predictions - try: - monitor.log_predictions(inputs=X, predictions=predictions) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return predictions - - # Wrap predict_proba if available - if original_predict_proba: - - def wrapped_predict_proba(X, *args, **kwargs): - probabilities = original_predict_proba(X, *args, **kwargs) - - # Log with probabilities - try: - predictions = np.argmax(probabilities, axis=1) - monitor.log_predictions( - inputs=X, predictions=predictions, probabilities=probabilities - ) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return probabilities - - model.predict_proba = wrapped_predict_proba - - model.predict = wrapped_predict - - return model diff --git a/src/whiteboxai/integrations/crewai_monitor.py b/src/whiteboxai/integrations/crewai_monitor.py deleted file mode 100644 index 9ab35cd..0000000 --- a/src/whiteboxai/integrations/crewai_monitor.py +++ /dev/null @@ -1,474 +0,0 @@ -""" -CrewAI Integration for WhiteBoxAI - -Monitor CrewAI multi-agent workflows with automatic tracking of agents, -tasks, interactions, and costs. -""" - -import logging -from typing import Any, Dict, Optional - -logger = logging.getLogger(__name__) - - -class CrewAIMonitor: - """ - Monitor CrewAI multi-agent workflows. - - Automatically tracks: - - Workflow lifecycle (start, completion) - - Agent definitions and executions - - Task assignments and completions - - Agent-to-agent interactions - - Token usage and costs - - Example: - >>> from whiteboxai.integrations import CrewAIMonitor - >>> from crewai import Agent, Task, Crew - >>> - >>> monitor = CrewAIMonitor(api_key="your_api_key") - >>> - >>> # Define agents - >>> researcher = Agent( - ... role="Research Analyst", - ... goal="Find accurate information", - ... backstory="Expert researcher", - ... tools=[search_tool] - ... ) - >>> - >>> writer = Agent( - ... role="Content Writer", - ... goal="Write engaging content", - ... backstory="Professional writer", - ... tools=[writing_tool] - ... ) - >>> - >>> # Create tasks - >>> research_task = Task( - ... description="Research topic X", - ... agent=researcher - ... ) - >>> - >>> writing_task = Task( - ... description="Write article based on research", - ... agent=writer - ... ) - >>> - >>> # Create and monitor crew - >>> crew = Crew( - ... agents=[researcher, writer], - ... tasks=[research_task, writing_task], - ... process=Process.sequential - ... ) - >>> - >>> workflow_id = monitor.start_monitoring( - ... crew=crew, - ... workflow_name="Article Generation", - ... metadata={"topic": "AI Safety"} - ... ) - >>> - >>> # Execute crew (automatically monitored) - >>> result = crew.kickoff() - >>> - >>> # Complete monitoring - >>> monitor.complete_monitoring(outputs={"article": result}) - """ - - def __init__( - self, api_key: str, api_url: Optional[str] = None, organization_id: Optional[str] = None - ): - """ - Initialize CrewAI monitor. - - Args: - api_key: WhiteBoxAI API key - api_url: WhiteBoxAI API URL (optional, defaults to production) - organization_id: Organization ID (optional, extracted from token if not provided) - """ - from whiteboxai import WhiteBoxAI - - self.client = ( - WhiteBoxAI(api_key=api_key, base_url=api_url) - if api_url - else WhiteBoxAI(api_key=api_key) - ) - self.organization_id = organization_id - self.workflow_id = None - self.agent_map = {} # Maps CrewAI agent to WhiteBoxAI agent_id - self.task_map = {} # Maps CrewAI task to WhiteBoxAI task_id - self.execution_map = {} # Maps agent to execution_id - - logger.info("CrewAI Monitor initialized") - - def start_monitoring( - self, - crew: Any, # crewai.Crew - workflow_name: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> str: - """ - Start monitoring a CrewAI crew workflow. - - Args: - crew: CrewAI Crew instance - workflow_name: Name for the workflow - metadata: Optional workflow metadata - - Returns: - Workflow ID (UUID string) - """ - try: - # Create workflow - workflow_data = { - "name": workflow_name, - "framework": "crewai", - "metadata": metadata or {}, - } - - response = self.client.request( - "POST", "/api/v1/workflows/multi-agent/start", data=workflow_data - ) - - self.workflow_id = response.get("id") - logger.info(f"Started monitoring CrewAI workflow: {self.workflow_id}") - - # Register agents - for agent in crew.agents: - self._register_agent(agent) - - # Register tasks - for task in crew.tasks: - self._register_task(task) - - # Start workflow - start_data = { - "inputs": { - "agent_count": len(crew.agents), - "task_count": len(crew.tasks), - "process": str(crew.process) if hasattr(crew, "process") else "sequential", - } - } - - self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/start", data=start_data - ) - - return self.workflow_id - - except Exception as e: - logger.error(f"Error starting CrewAI monitoring: {str(e)}") - raise - - def _register_agent(self, crew_agent: Any) -> str: - """ - Register a CrewAI agent with WhiteBoxAI. - - Args: - crew_agent: CrewAI Agent instance - - Returns: - Agent ID (UUID string) - """ - try: - agent_data = { - "name": getattr(crew_agent, "role", "Unknown Agent"), - "role": getattr(crew_agent, "role", None), - "agent_type": "crewai_agent", - "goal": getattr(crew_agent, "goal", None), - "backstory": getattr(crew_agent, "backstory", None), - "tools": [tool.__class__.__name__ for tool in getattr(crew_agent, "tools", [])], - "llm_provider": ( - getattr(getattr(crew_agent, "llm", None), "model_name", "unknown").split("/")[0] - if hasattr(crew_agent, "llm") - else None - ), - "model_name": ( - getattr(getattr(crew_agent, "llm", None), "model_name", None) - if hasattr(crew_agent, "llm") - else None - ), - "metadata": { - "verbose": getattr(crew_agent, "verbose", False), - "allow_delegation": getattr(crew_agent, "allow_delegation", False), - "max_iter": getattr(crew_agent, "max_iter", None), - }, - } - - response = self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/agents", data=agent_data - ) - - agent_id = response.get("id") - self.agent_map[id(crew_agent)] = agent_id - - logger.debug(f"Registered agent: {agent_data['name']} ({agent_id})") - - return agent_id - - except Exception as e: - logger.error(f"Error registering agent: {str(e)}") - raise - - def _register_task(self, crew_task: Any) -> str: - """ - Register a CrewAI task with WhiteBoxAI. - - Args: - crew_task: CrewAI Task instance - - Returns: - Task ID (UUID string) - """ - try: - # Get agent ID for task's assigned agent - agent_id = None - if hasattr(crew_task, "agent") and crew_task.agent: - agent_id = self.agent_map.get(id(crew_task.agent)) - - task_data = { - "task_name": getattr(crew_task, "description", "Unknown Task")[:255], - "description": getattr(crew_task, "description", None), - "expected_output": getattr(crew_task, "expected_output", None), - "agent_id": agent_id, - "context": { - "tools": [tool.__class__.__name__ for tool in getattr(crew_task, "tools", [])], - "async_execution": getattr(crew_task, "async_execution", False), - }, - } - - response = self.client.request( - "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/tasks", data=task_data - ) - - task_id = response.get("id") - self.task_map[id(crew_task)] = task_id - - logger.debug(f"Registered task: {task_data['task_name'][:50]}... ({task_id})") - - return task_id - - except Exception as e: - logger.error(f"Error registering task: {str(e)}") - raise - - def log_agent_execution( - self, - agent: Any, - inputs: Optional[Dict] = None, - outputs: Optional[Dict] = None, - tokens_used: int = 0, - cost: float = 0.0, - ) -> None: - """ - Log an agent execution. - - Args: - agent: CrewAI Agent instance - inputs: Execution inputs - outputs: Execution outputs - tokens_used: Tokens consumed - cost: Execution cost - """ - try: - agent_id = self.agent_map.get(id(agent)) - if not agent_id: - logger.warning("Agent not registered, skipping execution log") - return - - execution_data = { - "agent_id": agent_id, - "inputs": inputs, - "outputs": outputs, - "tokens_used": tokens_used, - "cost": cost, - "status": "completed", - } - - self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{self.workflow_id}/executions", - data=execution_data, - ) - - logger.debug(f"Logged execution for agent: {agent_id}") - - except Exception as e: - logger.error(f"Error logging agent execution: {str(e)}") - - def log_task_completion( - self, - task: Any, - status: str = "completed", - output_data: Optional[Dict] = None, - error_message: Optional[str] = None, - ) -> None: - """ - Log task completion. - - Args: - task: CrewAI Task instance - status: Task status (completed, failed) - output_data: Task output - error_message: Error message if failed - """ - try: - task_id = self.task_map.get(id(task)) - if not task_id: - logger.warning("Task not registered, skipping completion log") - return - - update_data = { - "status": status, - "output_data": output_data, - "error_message": error_message, - } - - self.client.request( - "PATCH", f"/api/v1/workflows/multi-agent/tasks/{task_id}", data=update_data - ) - - logger.debug(f"Logged task completion: {task_id} ({status})") - - except Exception as e: - logger.error(f"Error logging task completion: {str(e)}") - - def log_interaction( - self, - from_agent: Any, - to_agent: Any, - interaction_type: str = "delegation", - message: Optional[str] = None, - ) -> None: - """ - Log agent-to-agent interaction. - - Args: - from_agent: Source agent - to_agent: Target agent - interaction_type: Type of interaction (delegation, handoff, query, feedback) - message: Interaction message - """ - try: - from_agent_id = self.agent_map.get(id(from_agent)) - to_agent_id = self.agent_map.get(id(to_agent)) - - if not from_agent_id or not to_agent_id: - logger.warning("Agents not registered, skipping interaction log") - return - - interaction_data = { - "interaction_type": interaction_type, - "from_agent_id": from_agent_id, - "to_agent_id": to_agent_id, - "message": message, - } - - self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{self.workflow_id}/interactions", - data=interaction_data, - ) - - logger.debug( - f"Logged interaction: {from_agent_id} -> {to_agent_id} ({interaction_type})" - ) - - except Exception as e: - logger.error(f"Error logging interaction: {str(e)}") - - def complete_monitoring( - self, - status: str = "completed", - outputs: Optional[Dict] = None, - error_message: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Complete workflow monitoring. - - Args: - status: Final workflow status (completed, failed, cancelled) - outputs: Workflow outputs - error_message: Error message if failed - - Returns: - Workflow summary with analytics - """ - try: - complete_data = {"status": status, "outputs": outputs, "error_message": error_message} - - self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{self.workflow_id}/complete", - data=complete_data, - ) - - logger.info(f"Completed monitoring workflow: {self.workflow_id} ({status})") - - # Get analytics - analytics = self.get_analytics() - - return {"workflow_id": self.workflow_id, "status": status, "analytics": analytics} - - except Exception as e: - logger.error(f"Error completing monitoring: {str(e)}") - raise - - def get_analytics(self) -> Dict[str, Any]: - """ - Get workflow analytics. - - Returns: - Dict with metrics, cost breakdown, and bottlenecks - """ - try: - if not self.workflow_id: - raise ValueError("No active workflow") - - analytics = self.client.request( - "GET", f"/api/v1/workflows/multi-agent/{self.workflow_id}/analytics" - ) - - cost_breakdown = self.client.request( - "GET", f"/api/v1/workflows/multi-agent/{self.workflow_id}/cost-breakdown" - ) - - return {"metrics": analytics, "cost_breakdown": cost_breakdown} - - except Exception as e: - logger.error(f"Error getting analytics: {str(e)}") - return {} - - -# Helper function for easy usage -def monitor_crew( - crew: Any, - workflow_name: str, - api_key: str, - api_url: Optional[str] = None, - metadata: Optional[Dict] = None, -) -> CrewAIMonitor: - """ - Convenience function to start monitoring a CrewAI crew. - - Args: - crew: CrewAI Crew instance - workflow_name: Workflow name - api_key: WhiteBoxAI API key - api_url: WhiteBoxAI API URL (optional) - metadata: Workflow metadata (optional) - - Returns: - CrewAIMonitor instance with workflow started - - Example: - >>> monitor = monitor_crew( - ... crew=my_crew, - ... workflow_name="Research & Writing", - ... api_key="your_api_key" - ... ) - >>> result = my_crew.kickoff() - >>> monitor.complete_monitoring(outputs={"result": result}) - """ - monitor = CrewAIMonitor(api_key=api_key, api_url=api_url) - monitor.start_monitoring(crew=crew, workflow_name=workflow_name, metadata=metadata) - return monitor diff --git a/src/whiteboxai/integrations/langchain.py b/src/whiteboxai/integrations/langchain.py deleted file mode 100644 index 4112d89..0000000 --- a/src/whiteboxai/integrations/langchain.py +++ /dev/null @@ -1,636 +0,0 @@ -""" -LangChain Integration - -Integration for monitoring LangChain applications including chains, agents, and RAG pipelines. -""" - -import time -import warnings -from typing import Any, Dict, List, Optional - -try: - from langchain.agents import AgentExecutor - from langchain.callbacks.base import BaseCallbackHandler - from langchain.chains.base import Chain - from langchain.schema import AgentAction, AgentFinish, LLMResult - - LANGCHAIN_AVAILABLE = True -except ImportError: - LANGCHAIN_AVAILABLE = False - BaseCallbackHandler = object - Chain = object - AgentExecutor = object - -from whiteboxai.monitor import ModelMonitor - - -class LangChainMonitor(ModelMonitor): - """ - Monitor for LangChain applications. - - Supports: - - Chain executions - - Agent runs - - LLM calls - - Tool usage - - RAG pipelines - - Memory operations - - Example: - ```python - from langchain.chains import LLMChain - from langchain.llms import OpenAI - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.langchain import LangChainMonitor - - # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") - monitor = LangChainMonitor( - client=client, - application_name="my_langchain_app" - ) - - # Register application - monitor.register_application( - name="LangChain Q&A", - version="1.0.0" - ) - - # Create callback handler - callback = monitor.create_callback_handler() - - # Use with chain - llm = OpenAI() - chain = LLMChain(llm=llm, prompt=prompt) - result = chain.run(input="Question", callbacks=[callback]) - ``` - """ - - def __init__( - self, - client, - application_name: Optional[str] = None, - track_tokens: bool = True, - track_cost: bool = True, - **kwargs, - ): - """ - Initialize LangChain monitor. - - Args: - client: WhiteBoxAI client instance - application_name: Name of the LangChain application - track_tokens: Track token usage - track_cost: Track API costs - **kwargs: Additional arguments for ModelMonitor - """ - if not LANGCHAIN_AVAILABLE: - raise ImportError("langchain is not installed. Install with: pip install langchain") - - super().__init__(client, **kwargs) - self._application_name = application_name - self._track_tokens = track_tokens - self._track_cost = track_cost - self._registered = False - - def register_application( - self, - name: Optional[str] = None, - version: Optional[str] = None, - description: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register LangChain application with WhiteBoxAI. - - Args: - name: Application name - version: Application version - description: Application description - **kwargs: Additional metadata - - Returns: - Model ID (application ID) - """ - if name is None: - name = self._application_name or "langchain_app" - - # Prepare metadata - metadata = { - "framework": "langchain", - "application_type": "langchain", - "track_tokens": self._track_tokens, - "track_cost": self._track_cost, - **kwargs, - } - - # Register as a model (application) - self.model_id = self.client.register_model( - name=name, - model_type="generation", # LangChain apps are typically generative - framework="langchain", - version=version, - description=description, - metadata=metadata, - ) - - self._registered = True - return self.model_id - - def create_callback_handler(self) -> "WhiteBoxAICallbackHandler": - """ - Create a LangChain callback handler for automatic logging. - - Returns: - WhiteBoxAICallbackHandler instance - - Example: - ```python - callback = monitor.create_callback_handler() - chain.run(input="Question", callbacks=[callback]) - ``` - """ - if not self._registered: - self.register_application() - - return WhiteBoxAICallbackHandler(monitor=self) - - def log_chain_execution( - self, - chain_name: str, - inputs: Dict[str, Any], - outputs: Dict[str, Any], - execution_time: float, - llm_calls: Optional[List[Dict]] = None, - tool_calls: Optional[List[Dict]] = None, - **kwargs: Any, - ) -> None: - """ - Log a chain execution. - - Args: - chain_name: Name of the chain - inputs: Chain inputs - outputs: Chain outputs - execution_time: Time taken to execute (seconds) - llm_calls: List of LLM calls made - tool_calls: List of tool calls made - **kwargs: Additional metadata - """ - metadata = { - "chain_name": chain_name, - "execution_time": execution_time, - "num_llm_calls": len(llm_calls) if llm_calls else 0, - "num_tool_calls": len(tool_calls) if tool_calls else 0, - **kwargs, - } - - if llm_calls: - metadata["llm_calls"] = llm_calls - if tool_calls: - metadata["tool_calls"] = tool_calls - - # Log as prediction - self.log_prediction( - inputs=inputs, - prediction=outputs, - metadata=metadata, - ) - - def log_agent_execution( - self, - agent_name: str, - inputs: Dict[str, Any], - outputs: Dict[str, Any], - execution_time: float, - steps: List[Dict], - **kwargs: Any, - ) -> None: - """ - Log an agent execution. - - Args: - agent_name: Name of the agent - inputs: Agent inputs - outputs: Agent outputs - execution_time: Time taken to execute (seconds) - steps: List of agent steps/actions - **kwargs: Additional metadata - """ - metadata = { - "agent_name": agent_name, - "execution_time": execution_time, - "num_steps": len(steps), - "steps": steps, - **kwargs, - } - - # Log as prediction - self.log_prediction( - inputs=inputs, - prediction=outputs, - metadata=metadata, - ) - - def log_llm_call( - self, - prompt: str, - response: str, - model: str, - tokens_used: Optional[int] = None, - cost: Optional[float] = None, - latency: Optional[float] = None, - **kwargs: Any, - ) -> None: - """ - Log an LLM call. - - Args: - prompt: Input prompt - response: LLM response - model: Model name - tokens_used: Number of tokens used - cost: API cost - latency: Response latency (seconds) - **kwargs: Additional metadata - """ - metadata = { - "model": model, - "tokens_used": tokens_used, - "cost": cost, - "latency": latency, - **kwargs, - } - - self.log_prediction( - inputs={"prompt": prompt}, - prediction={"response": response}, - metadata=metadata, - ) - - def log_tool_call( - self, - tool_name: str, - tool_input: Any, - tool_output: Any, - execution_time: Optional[float] = None, - **kwargs: Any, - ) -> None: - """ - Log a tool call. - - Args: - tool_name: Name of the tool - tool_input: Tool input - tool_output: Tool output - execution_time: Time taken to execute (seconds) - **kwargs: Additional metadata - """ - metadata = { - "tool_name": tool_name, - "execution_time": execution_time, - **kwargs, - } - - self.log_prediction( - inputs={"tool_input": tool_input}, - prediction={"tool_output": tool_output}, - metadata=metadata, - ) - - def log_rag_retrieval( - self, - query: str, - documents: List[Dict], - num_retrieved: int, - retrieval_time: Optional[float] = None, - relevance_scores: Optional[List[float]] = None, - **kwargs: Any, - ) -> None: - """ - Log a RAG retrieval operation. - - Args: - query: Query text - documents: Retrieved documents - num_retrieved: Number of documents retrieved - retrieval_time: Time taken to retrieve (seconds) - relevance_scores: Relevance scores for documents - **kwargs: Additional metadata - """ - metadata = { - "num_retrieved": num_retrieved, - "retrieval_time": retrieval_time, - "relevance_scores": relevance_scores, - **kwargs, - } - - self.log_prediction( - inputs={"query": query}, - prediction={"documents": documents}, - metadata=metadata, - ) - - -class WhiteBoxAICallbackHandler(BaseCallbackHandler): - """ - LangChain callback handler for automatic logging to WhiteBoxAI. - - Example: - ```python - monitor = LangChainMonitor(client, application_name="my_app") - monitor.register_application() - callback = monitor.create_callback_handler() - - # Use with chain - chain.run(input="Question", callbacks=[callback]) - - # Use with agent - agent.run(input="Task", callbacks=[callback]) - ``` - """ - - def __init__(self, monitor: LangChainMonitor): - """ - Initialize callback handler. - - Args: - monitor: LangChainMonitor instance - """ - super().__init__() - self.monitor = monitor - self._chain_start_times = {} - self._agent_start_times = {} - self._llm_start_times = {} - self._tool_start_times = {} - self._chain_inputs = {} - self._agent_inputs = {} - self._agent_steps = {} - - def on_chain_start( - self, - serialized: Dict[str, Any], - inputs: Dict[str, Any], - **kwargs: Any, - ) -> None: - """Called when a chain starts running.""" - run_id = kwargs.get("run_id", id(serialized)) - self._chain_start_times[run_id] = time.time() - self._chain_inputs[run_id] = inputs - - def on_chain_end( - self, - outputs: Dict[str, Any], - **kwargs: Any, - ) -> None: - """Called when a chain ends running.""" - run_id = kwargs.get("run_id") - if run_id in self._chain_start_times: - execution_time = time.time() - self._chain_start_times[run_id] - inputs = self._chain_inputs.get(run_id, {}) - - # Get chain name - chain_name = kwargs.get("name", "unknown_chain") - - # Log chain execution - try: - self.monitor.log_chain_execution( - chain_name=chain_name, - inputs=inputs, - outputs=outputs, - execution_time=execution_time, - ) - except Exception as e: - warnings.warn(f"Failed to log chain execution: {e}") - - # Cleanup - del self._chain_start_times[run_id] - if run_id in self._chain_inputs: - del self._chain_inputs[run_id] - - def on_chain_error( - self, - error: Exception, - **kwargs: Any, - ) -> None: - """Called when a chain errors.""" - run_id = kwargs.get("run_id") - if run_id in self._chain_start_times: - # Cleanup on error - del self._chain_start_times[run_id] - if run_id in self._chain_inputs: - del self._chain_inputs[run_id] - - def on_agent_action( - self, - action: AgentAction, - **kwargs: Any, - ) -> None: - """Called when an agent takes an action.""" - run_id = kwargs.get("run_id") - if run_id not in self._agent_steps: - self._agent_steps[run_id] = [] - - self._agent_steps[run_id].append( - { - "tool": action.tool, - "tool_input": action.tool_input, - "log": action.log, - } - ) - - def on_agent_finish( - self, - finish: AgentFinish, - **kwargs: Any, - ) -> None: - """Called when an agent finishes running.""" - run_id = kwargs.get("run_id") - if run_id in self._agent_start_times: - execution_time = time.time() - self._agent_start_times[run_id] - inputs = self._agent_inputs.get(run_id, {}) - steps = self._agent_steps.get(run_id, []) - - # Log agent execution - try: - self.monitor.log_agent_execution( - agent_name=kwargs.get("name", "unknown_agent"), - inputs=inputs, - outputs=finish.return_values, - execution_time=execution_time, - steps=steps, - ) - except Exception as e: - warnings.warn(f"Failed to log agent execution: {e}") - - # Cleanup - del self._agent_start_times[run_id] - if run_id in self._agent_inputs: - del self._agent_inputs[run_id] - if run_id in self._agent_steps: - del self._agent_steps[run_id] - - def on_llm_start( - self, - serialized: Dict[str, Any], - prompts: List[str], - **kwargs: Any, - ) -> None: - """Called when an LLM starts generating.""" - run_id = kwargs.get("run_id", id(serialized)) - self._llm_start_times[run_id] = time.time() - - def on_llm_end( - self, - response: LLMResult, - **kwargs: Any, - ) -> None: - """Called when an LLM ends generating.""" - run_id = kwargs.get("run_id") - if run_id in self._llm_start_times: - latency = time.time() - self._llm_start_times[run_id] - - # Extract response details - if response.generations and len(response.generations) > 0: - generation = response.generations[0][0] - response_text = generation.text - - # Extract token usage if available - tokens_used = None - if response.llm_output and "token_usage" in response.llm_output: - tokens_used = response.llm_output["token_usage"].get("total_tokens") - - # Log LLM call (if configured) - if self.monitor._track_tokens or self.monitor._track_cost: - try: - self.monitor.log_llm_call( - prompt="", # Prompt tracked separately - response=response_text, - model=kwargs.get("invocation_params", {}).get("model_name", "unknown"), - tokens_used=tokens_used, - latency=latency, - ) - except Exception as e: - warnings.warn(f"Failed to log LLM call: {e}") - - # Cleanup - del self._llm_start_times[run_id] - - def on_llm_error( - self, - error: Exception, - **kwargs: Any, - ) -> None: - """Called when an LLM errors.""" - run_id = kwargs.get("run_id") - if run_id in self._llm_start_times: - del self._llm_start_times[run_id] - - def on_tool_start( - self, - serialized: Dict[str, Any], - input_str: str, - **kwargs: Any, - ) -> None: - """Called when a tool starts running.""" - run_id = kwargs.get("run_id", id(serialized)) - self._tool_start_times[run_id] = time.time() - - def on_tool_end( - self, - output: str, - **kwargs: Any, - ) -> None: - """Called when a tool ends running.""" - run_id = kwargs.get("run_id") - if run_id in self._tool_start_times: - execution_time = time.time() - self._tool_start_times[run_id] - - # Log tool call - try: - self.monitor.log_tool_call( - tool_name=kwargs.get("name", "unknown_tool"), - tool_input=kwargs.get("input_str", ""), - tool_output=output, - execution_time=execution_time, - ) - except Exception as e: - warnings.warn(f"Failed to log tool call: {e}") - - # Cleanup - del self._tool_start_times[run_id] - - def on_tool_error( - self, - error: Exception, - **kwargs: Any, - ) -> None: - """Called when a tool errors.""" - run_id = kwargs.get("run_id") - if run_id in self._tool_start_times: - del self._tool_start_times[run_id] - - -def wrap_langchain_chain( - chain: Chain, - monitor: LangChainMonitor, -) -> Chain: - """ - Wrap a LangChain chain to automatically log executions. - - Args: - chain: LangChain chain to wrap - monitor: LangChainMonitor instance - - Returns: - Wrapped chain that logs executions - - Example: - ```python - chain = LLMChain(llm=llm, prompt=prompt) - wrapped_chain = wrap_langchain_chain(chain, monitor) - - # Executions automatically logged - result = wrapped_chain.run(input="Question") - ``` - """ - if not monitor._registered: - monitor.register_application() - - # Create callback handler - callback = monitor.create_callback_handler() - - # Store original run method - original_run = chain.run - original_call = chain.__call__ - - def logged_run(*args, **kwargs): - # Add callback to kwargs - if "callbacks" not in kwargs: - kwargs["callbacks"] = [] - kwargs["callbacks"].append(callback) - - # Run chain - return original_run(*args, **kwargs) - - def logged_call(*args, **kwargs): - # Add callback to kwargs - if "callbacks" not in kwargs: - kwargs["callbacks"] = [] - kwargs["callbacks"].append(callback) - - # Call chain - return original_call(*args, **kwargs) - - # Replace methods - chain.run = logged_run - chain.__call__ = logged_call - - return chain - - -__all__ = [ - "LangChainMonitor", - "WhiteBoxAICallbackHandler", - "wrap_langchain_chain", -] diff --git a/src/whiteboxai/integrations/langchain_agents.py b/src/whiteboxai/integrations/langchain_agents.py deleted file mode 100644 index ec782da..0000000 --- a/src/whiteboxai/integrations/langchain_agents.py +++ /dev/null @@ -1,525 +0,0 @@ -""" -LangChain Multi-Agent Integration for WhiteBoxAI - -Enhanced callback handler for monitoring multi-agent LangChain workflows including: -- LangGraph multi-agent patterns -- Agent supervisors and coordinators -- Tool usage and agent handoffs -- Agent-to-agent communication -""" - -from datetime import datetime -from typing import Any, Dict, List, Optional, Union - -from langchain.callbacks.base import BaseCallbackHandler -from langchain.schema import AgentAction, AgentFinish, LLMResult - -try: - from whiteboxai import WhiteBoxAI -except ImportError: - WhiteBoxAI = None - - -class MultiAgentCallbackHandler(BaseCallbackHandler): - """Enhanced callback handler for multi-agent LangChain workflows. - - This handler tracks: - - Agent executions and decisions - - Tool calls and results - - Agent-to-agent handoffs - - LLM calls per agent - - Workflow-level metrics - - Example: - ```python - from langchain.agents import AgentExecutor, create_react_agent - from whiteboxai.integrations import MultiAgentCallbackHandler - - # Initialize WhiteBoxAI client - client = WhiteBoxAI(api_key="your_key") - - # Create workflow - workflow_id = client.agent_workflows.create( - name="Research Workflow", - framework="langchain" - ).get("id") - - # Start workflow - client.agent_workflows.start(workflow_id) - - # Create callback - callback = MultiAgentCallbackHandler( - client=client, - workflow_id=workflow_id, - agent_name="researcher" - ) - - # Use with agent - agent_executor = AgentExecutor( - agent=agent, - tools=tools, - callbacks=[callback] - ) - result = agent_executor.run("Research AI safety") - - # Complete workflow - client.agent_workflows.complete( - workflow_id, - outputs={"result": result} - ) - ``` - """ - - def __init__( - self, - client: "WhiteBoxAI", - workflow_id: str, - agent_name: str = "main", - agent_role: Optional[str] = None, - track_tokens: bool = True, - track_costs: bool = True, - ): - """Initialize the callback handler. - - Args: - client: WhiteBoxAI client instance - workflow_id: ID of the workflow to track - agent_name: Name of the current agent - agent_role: Role/description of the agent - track_tokens: Whether to track token usage - track_costs: Whether to estimate costs - """ - if WhiteBoxAI is None: - raise ImportError( - "whiteboxai package not installed. " "Install with: pip install whiteboxai" - ) - - self.client = client - self.workflow_id = workflow_id - self.agent_name = agent_name - self.agent_role = agent_role or agent_name - self.track_tokens = track_tokens - self.track_costs = track_costs - - # Tracking state - self.current_execution_id: Optional[str] = None - self.execution_start_time: Optional[datetime] = None - self.llm_call_count = 0 - self.tool_call_count = 0 - self.total_tokens = 0 - self.total_cost = 0.0 - self.execution_inputs: Optional[Dict[str, Any]] = None - - def on_chain_start( - self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any - ) -> None: - """Run when chain starts.""" - # Start agent execution - self.execution_start_time = datetime.utcnow() - self.execution_inputs = inputs - self.llm_call_count = 0 - self.tool_call_count = 0 - self.total_tokens = 0 - self.total_cost = 0.0 - - def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None: - """Run when chain ends successfully.""" - if self.execution_start_time: - duration_ms = int( - (datetime.utcnow() - self.execution_start_time).total_seconds() * 1000 - ) - - # Log agent execution - try: - response = self.client.agent_workflows.create_execution( - workflow_id=self.workflow_id, - agent_name=self.agent_name, - status="completed", - inputs=self.execution_inputs, - outputs=outputs, - duration_ms=duration_ms, - llm_call_count=self.llm_call_count, - tool_call_count=self.tool_call_count, - tokens_used=self.total_tokens if self.track_tokens else None, - cost=self.total_cost if self.track_costs else None, - ) - self.current_execution_id = response.get("id") - except Exception as e: - print(f"Warning: Failed to log execution: {e}") - - # Reset state - self.execution_start_time = None - - def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: - """Run when chain errors.""" - if self.execution_start_time: - duration_ms = int( - (datetime.utcnow() - self.execution_start_time).total_seconds() * 1000 - ) - - # Log failed execution - try: - self.client.agent_workflows.create_execution( - workflow_id=self.workflow_id, - agent_name=self.agent_name, - status="failed", - inputs=self.execution_inputs, - outputs={"error": str(error)}, - duration_ms=duration_ms, - llm_call_count=self.llm_call_count, - tool_call_count=self.tool_call_count, - ) - except Exception as e: - print(f"Warning: Failed to log error: {e}") - - self.execution_start_time = None - - def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None: - """Run when LLM starts.""" - self.llm_call_count += 1 - - def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: - """Run when LLM ends.""" - # Track tokens if available - if self.track_tokens and hasattr(response, "llm_output"): - llm_output = response.llm_output or {} - token_usage = llm_output.get("token_usage", {}) - - total = token_usage.get("total_tokens", 0) - self.total_tokens += total - - # Estimate cost if tracking - if self.track_costs and total > 0: - # Rough estimate: $0.002 per 1K tokens (GPT-3.5 pricing) - self.total_cost += (total / 1000) * 0.002 - - def on_agent_action(self, action: AgentAction, **kwargs: Any) -> None: - """Run when agent takes an action (tool call).""" - self.tool_call_count += 1 - - # Log tool call as interaction - try: - self.client.agent_workflows.create_interaction( - workflow_id=self.workflow_id, - from_agent=self.agent_name, - to_agent="tool", - interaction_type="tool_call", - message=f"Tool: {action.tool}, Input: {action.tool_input}", - meta_data={ - "tool": action.tool, - "tool_input": action.tool_input, - "log": action.log, - }, - ) - except Exception as e: - print(f"Warning: Failed to log tool call: {e}") - - def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None: - """Run when agent finishes execution.""" - # This is called when the agent completes its reasoning - pass - - def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None: - """Run when tool starts.""" - pass - - def on_tool_end(self, output: str, **kwargs: Any) -> None: - """Run when tool ends.""" - # Log tool result as interaction - try: - self.client.agent_workflows.create_interaction( - workflow_id=self.workflow_id, - from_agent="tool", - to_agent=self.agent_name, - interaction_type="response", - message=f"Tool result: {output[:500]}", # Truncate long outputs - meta_data={"output": output}, - ) - except Exception as e: - print(f"Warning: Failed to log tool result: {e}") - - def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None: - """Run when tool errors.""" - try: - self.client.agent_workflows.create_interaction( - workflow_id=self.workflow_id, - from_agent="tool", - to_agent=self.agent_name, - interaction_type="response", - message=f"Tool error: {str(error)}", - meta_data={"error": str(error), "error_type": type(error).__name__}, - ) - except Exception as e: - print(f"Warning: Failed to log tool error: {e}") - - def on_text(self, text: str, **kwargs: Any) -> None: - """Run on arbitrary text.""" - pass - - -class LangGraphMultiAgentMonitor: - """Monitor for LangGraph multi-agent workflows. - - Provides higher-level monitoring for LangGraph patterns like: - - Agent supervisors - - Agent networks - - Sequential/parallel agent execution - - Example: - ```python - from langgraph.graph import StateGraph - from whiteboxai.integrations import LangGraphMultiAgentMonitor - - # Create monitor - monitor = LangGraphMultiAgentMonitor( - client=client, - workflow_name="Multi-Agent Research" - ) - - # Start monitoring - workflow_id = monitor.start_monitoring() - - # Register agents - monitor.register_agent("supervisor", role="Coordinates other agents") - monitor.register_agent("researcher", role="Gathers information") - monitor.register_agent("writer", role="Writes content") - - # Execute graph with callbacks - graph = StateGraph(...) - result = graph.invoke( - inputs, - config={"callbacks": [monitor.get_callbacks("supervisor")]} - ) - - # Complete monitoring - monitor.complete_monitoring(outputs={"result": result}) - ``` - """ - - def __init__( - self, client: "WhiteBoxAI", workflow_name: str, meta_data: Optional[Dict[str, Any]] = None - ): - """Initialize the LangGraph monitor. - - Args: - client: WhiteBoxAI client instance - workflow_name: Name for the workflow - meta_data: Additional meta_data to attach - """ - if WhiteBoxAI is None: - raise ImportError( - "whiteboxai package not installed. " "Install with: pip install whiteboxai" - ) - - self.client = client - self.workflow_name = workflow_name - self.workflow_meta_data = meta_data or {} - self.workflow_id: Optional[str] = None - self.callbacks: Dict[str, MultiAgentCallbackHandler] = {} - self.start_time: Optional[datetime] = None - - def start_monitoring(self, inputs: Optional[Dict[str, Any]] = None) -> str: - """Start workflow monitoring. - - Args: - inputs: Initial workflow inputs - - Returns: - workflow_id: ID of the created workflow - """ - self.start_time = datetime.utcnow() - - # Create workflow - response = self.client.agent_workflows.create( - name=self.workflow_name, - framework="langchain", - inputs=inputs, - meta_data=self.workflow_meta_data, - ) - self.workflow_id = response.get("id") - - # Start workflow - self.client.agent_workflows.start(self.workflow_id) - - return self.workflow_id - - def register_agent( - self, - agent_name: str, - role: Optional[str] = None, - model_name: Optional[str] = None, - tools: Optional[List[str]] = None, - **kwargs, - ) -> None: - """Register an agent in the workflow. - - Args: - agent_name: Name of the agent - role: Agent's role/goal - model_name: LLM model used - tools: List of tool names - **kwargs: Additional agent configuration - """ - if not self.workflow_id: - raise ValueError("Must call start_monitoring() first") - - self.client.agent_workflows.register_agent( - workflow_id=self.workflow_id, - name=agent_name, - role=role or agent_name, - model_name=model_name, - tools=tools, - **kwargs, - ) - - def get_callbacks( - self, agent_name: str, agent_role: Optional[str] = None - ) -> List[BaseCallbackHandler]: - """Get callbacks for a specific agent. - - Args: - agent_name: Name of the agent - agent_role: Optional role description - - Returns: - List of callback handlers - """ - if not self.workflow_id: - raise ValueError("Must call start_monitoring() first") - - if agent_name not in self.callbacks: - self.callbacks[agent_name] = MultiAgentCallbackHandler( - client=self.client, - workflow_id=self.workflow_id, - agent_name=agent_name, - agent_role=agent_role, - ) - - return [self.callbacks[agent_name]] - - def log_handoff( - self, - from_agent: str, - to_agent: str, - message: str, - meta_data: Optional[Dict[str, Any]] = None, - ) -> None: - """Log an agent-to-agent handoff. - - Args: - from_agent: Agent passing control - to_agent: Agent receiving control - message: Handoff message/context - meta_data: Additional meta_data - """ - if not self.workflow_id: - raise ValueError("Must call start_monitoring() first") - - self.client.agent_workflows.create_interaction( - workflow_id=self.workflow_id, - from_agent=from_agent, - to_agent=to_agent, - interaction_type="handoff", - message=message, - meta_data=meta_data, - ) - - def complete_monitoring( - self, outputs: Optional[Dict[str, Any]] = None, status: str = "completed" - ) -> Dict[str, Any]: - """Complete workflow monitoring. - - Args: - outputs: Final workflow outputs - status: Workflow status (completed/failed) - - Returns: - Summary with analytics - """ - if not self.workflow_id: - raise ValueError("Must call start_monitoring() first") - - # Complete workflow - self.client.agent_workflows.complete( - workflow_id=self.workflow_id, outputs=outputs, status=status - ) - - # Get analytics - try: - analytics = self.client.agent_workflows.get_analytics(self.workflow_id) - return { - "workflow_id": self.workflow_id, - "status": status, - "outputs": outputs, - "analytics": analytics, - } - except Exception as e: - print(f"Warning: Failed to retrieve analytics: {e}") - return {"workflow_id": self.workflow_id, "status": status, "outputs": outputs} - - -def monitor_langchain_agent( - client: "WhiteBoxAI", - agent_executor: Any, - workflow_name: str, - agent_name: str = "main", - inputs: Optional[Dict[str, Any]] = None, - **run_kwargs, -) -> Dict[str, Any]: - """Helper function to monitor a single LangChain agent execution. - - Args: - client: WhiteBoxAI client - agent_executor: LangChain AgentExecutor instance - workflow_name: Name for the workflow - agent_name: Name of the agent - inputs: Inputs to the agent - **run_kwargs: Additional arguments to pass to agent.run() - - Returns: - Dict with result and workflow_id - - Example: - ```python - from langchain.agents import AgentExecutor, create_react_agent - from whiteboxai.integrations import monitor_langchain_agent - - result_dict = monitor_langchain_agent( - client=client, - agent_executor=agent_executor, - workflow_name="Research Task", - agent_name="researcher", - inputs={"input": "Research AI safety"} - ) - - print(f"Result: {result_dict['result']}") - print(f"Workflow ID: {result_dict['workflow_id']}") - ``` - """ - # Create workflow - response = client.agent_workflows.create( - name=workflow_name, framework="langchain", inputs=inputs - ) - workflow_id = response.get("id") - - # Start workflow - client.agent_workflows.start(workflow_id) - - # Create callback - callback = MultiAgentCallbackHandler( - client=client, workflow_id=workflow_id, agent_name=agent_name - ) - - try: - # Run agent with callback - result = agent_executor.run(callbacks=[callback], **run_kwargs) - - # Complete workflow - client.agent_workflows.complete(workflow_id, outputs={"result": result}) - - return {"result": result, "workflow_id": workflow_id, "status": "completed"} - except Exception as e: - # Log failure - client.agent_workflows.complete(workflow_id, outputs={"error": str(e)}, status="failed") - - return {"result": None, "workflow_id": workflow_id, "status": "failed", "error": str(e)} diff --git a/src/whiteboxai/integrations/pytorch.py b/src/whiteboxai/integrations/pytorch.py deleted file mode 100644 index 16b7ddf..0000000 --- a/src/whiteboxai/integrations/pytorch.py +++ /dev/null @@ -1,268 +0,0 @@ -""" -PyTorch Integration - -Integration for monitoring PyTorch models. -""" - -from typing import Any, Callable, Dict, Optional - -try: - import torch - import torch.nn as nn - - TORCH_AVAILABLE = True -except ImportError: - TORCH_AVAILABLE = False - nn = object - torch = None - -from whiteboxai.monitor import ModelMonitor - - -class TorchMonitor(ModelMonitor): - """ - Monitor for PyTorch models. - - Example: - ```python - import torch - import torch.nn as nn - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.pytorch import TorchMonitor - - # Define model - model = nn.Sequential( - nn.Linear(10, 64), - nn.ReLU(), - nn.Linear(64, 2) - ) - - # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") - monitor = TorchMonitor(client, model=model) - monitor.register_from_model(model_type="classification") - - # Wrap model for automatic monitoring - monitored_model = monitor.wrap_model(model) - - # Predictions are automatically logged - outputs = monitored_model(inputs) - ``` - """ - - def __init__(self, client, model: Optional[nn.Module] = None, **kwargs): - """Initialize PyTorch monitor.""" - if not TORCH_AVAILABLE: - raise ImportError("PyTorch is not installed. Install with: pip install torch") - - super().__init__(client, **kwargs) - self.model = model - - def register_from_model( - self, - name: Optional[str] = None, - model_type: str = "classification", - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register model from PyTorch module. - - Args: - name: Model name (default: class name) - model_type: Model type (classification, regression, etc.) - version: Model version - **kwargs: Additional metadata - - Returns: - Model ID - """ - if self.model is None: - raise ValueError("No model provided") - - # Generate name from model class - if name is None: - name = self.model.__class__.__name__ - - # Extract model metadata - metadata = self._extract_model_metadata() - metadata.update(kwargs) - - return self.register_model( - name=name, - model_type=model_type, - framework="pytorch", - version=version, - **metadata, - ) - - async def aregister_from_model( - self, - name: Optional[str] = None, - model_type: str = "classification", - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """Async version of register_from_model().""" - if self.model is None: - raise ValueError("No model provided") - - if name is None: - name = self.model.__class__.__name__ - - metadata = self._extract_model_metadata() - metadata.update(kwargs) - - return await self.aregister_model( - name=name, - model_type=model_type, - framework="pytorch", - version=version, - **metadata, - ) - - def wrap_model(self, model: nn.Module) -> "TorchWrapper": - """ - Wrap PyTorch model for automatic monitoring. - - Args: - model: PyTorch module - - Returns: - Wrapped model - """ - return TorchWrapper(model, self) - - def _extract_model_metadata(self) -> Dict[str, Any]: - """Extract metadata from PyTorch model.""" - metadata = {} - - if self.model is None: - return metadata - - # Count parameters - total_params = sum(p.numel() for p in self.model.parameters()) - trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) - - metadata["total_parameters"] = total_params - metadata["trainable_parameters"] = trainable_params - - # Get model architecture - metadata["architecture"] = str(self.model) - - # Get PyTorch version - metadata["pytorch_version"] = torch.__version__ - - # Get device - device = next(self.model.parameters()).device - metadata["device"] = str(device) - - return metadata - - -class TorchWrapper(nn.Module): - """ - Wrapper for PyTorch models with automatic monitoring. - - This wrapper intercepts forward calls and logs predictions to WhiteBoxAI. - """ - - def __init__(self, model: nn.Module, monitor: TorchMonitor): - """Initialize wrapper.""" - super().__init__() - self.model = model - self.monitor = monitor - - def forward(self, x: torch.Tensor, log: bool = True) -> torch.Tensor: - """Forward pass with automatic logging.""" - # Call original forward - output = self.model(x) - - # Log predictions if enabled - if log: - self._log_batch_predictions(x, output) - - return output - - def _log_batch_predictions(self, inputs: torch.Tensor, outputs: torch.Tensor) -> None: - """Log batch of predictions.""" - # Convert to numpy/lists - inputs_np = inputs.detach().cpu().numpy() - outputs_np = outputs.detach().cpu().numpy() - - # Create predictions list - predictions = [] - for i in range(len(inputs_np)): - pred = { - "inputs": inputs_np[i].tolist(), - "output": outputs_np[i].tolist(), - } - predictions.append(pred) - - # Log batch - try: - self.monitor.log_batch(predictions) - except Exception as e: - # Don't fail prediction if logging fails - print(f"Warning: Failed to log predictions: {e}") - - def __getattr__(self, name: str): - """Delegate attributes to wrapped model.""" - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.model, name) - - -def monitor_forward(monitor: TorchMonitor, input_extractor: Optional[Callable] = None): - """ - Decorator to monitor PyTorch forward passes. - - Args: - monitor: TorchMonitor instance - input_extractor: Function to extract inputs from args - - Example: - ```python - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.pytorch import TorchMonitor, monitor_forward - - client = WhiteBoxAI(api_key="your-api-key") - monitor = TorchMonitor(client, model_id=123) - - class MyModel(nn.Module): - @monitor_forward(monitor) - def forward(self, x): - return self.layers(x) - ``` - """ - - def decorator(forward_fn: Callable) -> Callable: - def wrapper(self, *args, **kwargs): - # Call original forward - output = forward_fn(self, *args, **kwargs) - - # Extract inputs - if input_extractor: - inputs = input_extractor(args, kwargs) - else: - inputs = args[0] if args else None - - # Log prediction - if inputs is not None and isinstance(inputs, torch.Tensor): - inputs_np = inputs.detach().cpu().numpy() - outputs_np = output.detach().cpu().numpy() - - try: - monitor.log_prediction( - inputs=inputs_np.tolist(), - output=outputs_np.tolist(), - ) - except Exception as e: - print(f"Warning: Failed to log prediction: {e}") - - return output - - return wrapper - - return decorator diff --git a/src/whiteboxai/integrations/sklearn.py b/src/whiteboxai/integrations/sklearn.py deleted file mode 100644 index a1e9b4f..0000000 --- a/src/whiteboxai/integrations/sklearn.py +++ /dev/null @@ -1,219 +0,0 @@ -""" -Scikit-learn Integration - -Integration for monitoring scikit-learn models. -""" - -from typing import Any, Dict, Optional - -try: - import sklearn - from sklearn.base import BaseEstimator - - SKLEARN_AVAILABLE = True -except ImportError: - SKLEARN_AVAILABLE = False - BaseEstimator = object - -import numpy as np - -from whiteboxai.monitor import ModelMonitor - - -class SklearnMonitor(ModelMonitor): - """ - Monitor for scikit-learn models. - - Example: - ```python - from sklearn.ensemble import RandomForestClassifier - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.sklearn import SklearnMonitor - - # Train model - model = RandomForestClassifier() - model.fit(X_train, y_train) - - # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") - monitor = SklearnMonitor(client, model=model) - monitor.register_from_model(model_type="classification") - - # Wrap model for automatic monitoring - monitored_model = monitor.wrap_model(model) - - # Predictions are automatically logged - predictions = monitored_model.predict(X_test) - ``` - """ - - def __init__(self, client, model: Optional[BaseEstimator] = None, **kwargs): - """Initialize sklearn monitor.""" - if not SKLEARN_AVAILABLE: - raise ImportError( - "scikit-learn is not installed. Install with: pip install scikit-learn" - ) - - super().__init__(client, **kwargs) - self.model = model - - def register_from_model( - self, - name: Optional[str] = None, - model_type: str = "classification", - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register model from sklearn estimator. - - Args: - name: Model name (default: class name) - model_type: Model type (classification, regression, clustering) - version: Model version - **kwargs: Additional metadata - - Returns: - Model ID - """ - if self.model is None: - raise ValueError("No model provided") - - # Generate name from model class - if name is None: - name = self.model.__class__.__name__ - - # Extract model metadata - metadata = self._extract_model_metadata() - metadata.update(kwargs) - - return self.register_model( - name=name, - model_type=model_type, - framework="sklearn", - version=version, - **metadata, - ) - - async def aregister_from_model( - self, - name: Optional[str] = None, - model_type: str = "classification", - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """Async version of register_from_model().""" - if self.model is None: - raise ValueError("No model provided") - - if name is None: - name = self.model.__class__.__name__ - - metadata = self._extract_model_metadata() - metadata.update(kwargs) - - return await self.aregister_model( - name=name, - model_type=model_type, - framework="sklearn", - version=version, - **metadata, - ) - - def wrap_model(self, model: BaseEstimator) -> "SklearnWrapper": - """ - Wrap sklearn model for automatic monitoring. - - Args: - model: Sklearn estimator - - Returns: - Wrapped model - """ - return SklearnWrapper(model, self) - - def _extract_model_metadata(self) -> Dict[str, Any]: - """Extract metadata from sklearn model.""" - metadata = {} - - if self.model is None: - return metadata - - # Get model parameters - if hasattr(self.model, "get_params"): - metadata["parameters"] = self.model.get_params() - - # Get feature names - if hasattr(self.model, "feature_names_in_"): - metadata["feature_names"] = self.model.feature_names_in_.tolist() - - # Get feature importances (for tree-based models) - if hasattr(self.model, "feature_importances_"): - metadata["feature_importances"] = self.model.feature_importances_.tolist() - - # Get classes (for classifiers) - if hasattr(self.model, "classes_"): - metadata["classes"] = self.model.classes_.tolist() - - # Get number of features - if hasattr(self.model, "n_features_in_"): - metadata["n_features"] = self.model.n_features_in_ - - # Get sklearn version - metadata["sklearn_version"] = sklearn.__version__ - - return metadata - - -class SklearnWrapper: - """ - Wrapper for sklearn models with automatic monitoring. - - This wrapper intercepts predict/predict_proba calls and logs predictions - to WhiteBoxAI automatically. - """ - - def __init__(self, model: BaseEstimator, monitor: SklearnMonitor): - """Initialize wrapper.""" - self.model = model - self.monitor = monitor - - def predict(self, X: np.ndarray, **kwargs) -> np.ndarray: - """Predict with automatic logging.""" - predictions = self.model.predict(X, **kwargs) - - # Log predictions - self._log_batch_predictions(X, predictions) - - return predictions - - def predict_proba(self, X: np.ndarray, **kwargs) -> np.ndarray: - """Predict probabilities with automatic logging.""" - probas = self.model.predict_proba(X, **kwargs) - - # Log predictions - self._log_batch_predictions(X, probas) - - return probas - - def _log_batch_predictions(self, inputs: np.ndarray, outputs: np.ndarray) -> None: - """Log batch of predictions.""" - # Convert to list format - predictions = [] - for i in range(len(inputs)): - pred = { - "inputs": inputs[i].tolist() if isinstance(inputs[i], np.ndarray) else inputs[i], - "output": outputs[i].tolist() if isinstance(outputs[i], np.ndarray) else outputs[i], - } - predictions.append(pred) - - # Log batch - try: - self.monitor.log_batch(predictions) - except Exception as e: - # Don't fail prediction if logging fails - print(f"Warning: Failed to log predictions: {e}") - - def __getattr__(self, name: str): - """Delegate other methods to wrapped model.""" - return getattr(self.model, name) diff --git a/src/whiteboxai/integrations/tensorflow.py b/src/whiteboxai/integrations/tensorflow.py deleted file mode 100644 index b617673..0000000 --- a/src/whiteboxai/integrations/tensorflow.py +++ /dev/null @@ -1,475 +0,0 @@ -""" -TensorFlow/Keras Integration - -Integration for monitoring TensorFlow and Keras models. -""" - -import warnings -from typing import Any, Dict, Optional, Union - -try: - import tensorflow as tf - from tensorflow import keras - - TENSORFLOW_AVAILABLE = True -except ImportError: - TENSORFLOW_AVAILABLE = False - keras = None - -import numpy as np - -from whiteboxai.monitor import ModelMonitor - - -class KerasMonitor(ModelMonitor): - """ - Monitor for TensorFlow/Keras models. - - Example: - ```python - from tensorflow import keras - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.tensorflow import KerasMonitor - - # Build model - model = keras.Sequential([ - keras.layers.Dense(64, activation='relu'), - keras.layers.Dense(1) - ]) - model.compile(optimizer='adam', loss='mse') - - # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") - monitor = KerasMonitor(client, model=model, model_name="keras_model") - - # Train with monitoring callback - from whiteboxai.integrations.tensorflow import WhiteBoxAICallback - - model.fit(X_train, y_train, - callbacks=[WhiteBoxAICallback(monitor)]) - - # Make predictions with automatic logging - predictions = monitor.predict(X_test) - ``` - """ - - def __init__( - self, - client, - model: Optional["keras.Model"] = None, - model_name: Optional[str] = None, - model_type: str = "regression", - **kwargs, - ): - """ - Initialize Keras monitor. - - Args: - client: WhiteBoxAI client instance - model: Keras model to monitor - model_name: Name for the model - model_type: Type of model (classification, regression) - **kwargs: Additional arguments for ModelMonitor - """ - if not TENSORFLOW_AVAILABLE: - raise ImportError("TensorFlow is not installed. Install with: pip install tensorflow") - - super().__init__(client, **kwargs) - self.model = model - self._model_name = model_name - self._model_type = model_type - self._registered = False - - def register_from_model( - self, - name: Optional[str] = None, - model_type: Optional[str] = None, - version: Optional[str] = None, - framework_version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register Keras model with WhiteBoxAI. - - Args: - name: Model name (default: model class name) - model_type: Model type (classification, regression) - version: Model version - framework_version: TensorFlow version - **kwargs: Additional metadata - - Returns: - Model ID - """ - if self.model is None: - raise ValueError("No model provided") - - # Auto-detect model name - if name is None: - name = self._model_name or self.model.__class__.__name__ - - # Use provided or instance model type - model_type = model_type or self._model_type - - # Get TensorFlow version - if framework_version is None: - framework_version = tf.__version__ - - # Extract model metadata - metadata = { - "framework": "tensorflow", - "framework_version": framework_version, - "model_class": self.model.__class__.__name__, - **kwargs, - } - - # Get model configuration - try: - config = self.model.get_config() - metadata["model_config"] = str(config) - except Exception: - pass - - # Get input/output shapes - try: - if hasattr(self.model, "input_shape"): - metadata["input_shape"] = str(self.model.input_shape) - if hasattr(self.model, "output_shape"): - metadata["output_shape"] = str(self.model.output_shape) - except Exception: - pass - - # Register model - self.model_id = self.client.register_model( - name=name, - model_type=model_type, - framework="tensorflow", - version=version, - metadata=metadata, - ) - - self._registered = True - return self.model_id - - def predict( - self, - inputs: Union[np.ndarray, tf.Tensor], - log: bool = True, - actuals: Optional[np.ndarray] = None, - **kwargs: Any, - ) -> np.ndarray: - """ - Make predictions and optionally log them. - - Args: - inputs: Input data (numpy array or TensorFlow tensor) - log: Whether to log the prediction - actuals: Actual values (if available) - **kwargs: Additional metadata to log - - Returns: - Predictions as numpy array - """ - if self.model is None: - raise ValueError("No model provided") - - # Ensure model is registered - if not self._registered: - self.register_from_model() - - # Convert to numpy if needed - if isinstance(inputs, tf.Tensor): - inputs_np = inputs.numpy() - else: - inputs_np = inputs - - # Make prediction - predictions = self.model.predict(inputs, verbose=0) - - # Convert predictions to numpy - if isinstance(predictions, tf.Tensor): - predictions_np = predictions.numpy() - else: - predictions_np = predictions - - # Log if requested - if log: - # Determine if batch or single - if len(inputs_np.shape) == 1 or inputs_np.shape[0] == 1: - # Single prediction - self.log_prediction( - inputs=inputs_np[0] if len(inputs_np.shape) > 1 else inputs_np, - prediction=( - predictions_np[0] if len(predictions_np.shape) > 1 else predictions_np - ), - actual=actuals[0] if actuals is not None else None, - **kwargs, - ) - else: - # Batch prediction - self.log_batch( - inputs=inputs_np, - predictions=predictions_np, - actuals=actuals, - **kwargs, - ) - - return predictions_np - - def set_baseline( - self, - baseline_data: Union[np.ndarray, tf.Tensor], - baseline_labels: Optional[Union[np.ndarray, tf.Tensor]] = None, - ) -> None: - """ - Set baseline data for drift detection. - - Args: - baseline_data: Baseline input data - baseline_labels: Baseline labels (optional) - """ - # Convert to numpy - if isinstance(baseline_data, tf.Tensor): - baseline_data = baseline_data.numpy() - if baseline_labels is not None and isinstance(baseline_labels, tf.Tensor): - baseline_labels = baseline_labels.numpy() - - # Make baseline predictions if labels not provided - if baseline_labels is None and self.model is not None: - baseline_predictions = self.model.predict(baseline_data, verbose=0) - if isinstance(baseline_predictions, tf.Tensor): - baseline_predictions = baseline_predictions.numpy() - else: - baseline_predictions = baseline_labels - - # Set baseline through parent class - super().set_baseline(baseline_data, baseline_predictions) - - def log_epoch( - self, - epoch: int, - train_loss: Optional[float] = None, - val_loss: Optional[float] = None, - **metrics: Any, - ) -> None: - """ - Log metrics from a training epoch. - - Args: - epoch: Epoch number - train_loss: Training loss - val_loss: Validation loss - **metrics: Additional metrics (accuracy, etc.) - """ - metric_data = { - "epoch": epoch, - "train_loss": train_loss, - "val_loss": val_loss, - **metrics, - } - - self.log_custom_metric("training_epoch", metric_data) - - def log_checkpoint( - self, - epoch: int, - checkpoint_path: str, - metrics: Optional[Dict[str, Any]] = None, - ) -> None: - """ - Log a model checkpoint. - - Args: - epoch: Epoch number - checkpoint_path: Path to saved checkpoint - metrics: Metrics at checkpoint - """ - checkpoint_data = { - "epoch": epoch, - "checkpoint_path": checkpoint_path, - "metrics": metrics or {}, - } - - self.log_custom_metric("checkpoint", checkpoint_data) - - def register_saved_model( - self, - model_path: str, - metadata: Optional[Dict[str, Any]] = None, - ) -> None: - """ - Register a SavedModel with WhiteBoxAI. - - Args: - model_path: Path to SavedModel directory - metadata: Additional metadata - """ - model_metadata = { - "model_path": model_path, - "format": "SavedModel", - **(metadata or {}), - } - - if not self._registered: - self.register_from_model(**model_metadata) - - -class WhiteBoxAICallback(keras.callbacks.Callback): - """ - Keras callback for automatic monitoring during training. - - Example: - ```python - monitor = KerasMonitor(client, model=model, model_name="my_model") - callback = WhiteBoxAICallback(monitor, log_frequency=1) - - model.fit(X_train, y_train, - validation_data=(X_val, y_val), - callbacks=[callback], - epochs=50) - ``` - """ - - def __init__( - self, - monitor: KerasMonitor, - log_frequency: int = 1, - log_weights: bool = False, - log_gradients: bool = False, - log_validation: bool = True, - ): - """ - Initialize callback. - - Args: - monitor: KerasMonitor instance - log_frequency: Log metrics every N epochs - log_weights: Whether to log model weights - log_gradients: Whether to log gradients - log_validation: Whether to log validation predictions - """ - super().__init__() - self.monitor = monitor - self.log_frequency = log_frequency - self.log_weights = log_weights - self.log_gradients = log_gradients - self.log_validation = log_validation - - # Ensure model is registered - if not self.monitor._registered: - self.monitor.register_from_model() - - def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None): - """Called at the end of each epoch.""" - if logs is None: - logs = {} - - # Log metrics at specified frequency - if (epoch + 1) % self.log_frequency == 0: - # Extract train and val metrics - train_loss = logs.get("loss") - val_loss = logs.get("val_loss") - - # Extract additional metrics - metrics = {} - for key, value in logs.items(): - if key not in ["loss", "val_loss"]: - metrics[key] = float(value) if value is not None else None - - # Log epoch metrics - self.monitor.log_epoch( - epoch=epoch, - train_loss=float(train_loss) if train_loss is not None else None, - val_loss=float(val_loss) if val_loss is not None else None, - **metrics, - ) - - def on_train_end(self, logs: Optional[Dict[str, Any]] = None): - """Called at the end of training.""" - if logs is None: - logs = {} - - # Log final metrics - self.monitor.log_custom_metric( - "training_complete", - { - "final_metrics": {k: float(v) if v is not None else None for k, v in logs.items()}, - }, - ) - - -class TorchMonitor(ModelMonitor): - """ - Monitor for PyTorch models (kept for backwards compatibility). - - Note: This is a placeholder. Use the PyTorch integration module instead. - """ - - def __init__(self, *args, **kwargs): - warnings.warn( - "TorchMonitor in tensorflow module is deprecated. " - "Use whiteboxai.integrations.pytorch.TorchMonitor instead.", - DeprecationWarning, - stacklevel=2, - ) - super().__init__(*args, **kwargs) - - -def wrap_keras_model( - model: "keras.Model", - monitor: KerasMonitor, -) -> "keras.Model": - """ - Wrap a Keras model to automatically log predictions. - - Args: - model: Keras model to wrap - monitor: KerasMonitor instance - - Returns: - Wrapped model that logs predictions - - Example: - ```python - wrapped_model = wrap_keras_model(model, monitor) - predictions = wrapped_model.predict(X_test) # Automatically logged - ``` - """ - original_predict = model.predict - - def logged_predict(x, *args, **kwargs): - # Make prediction - predictions = original_predict(x, *args, **kwargs) - - # Log to WhiteBoxAI - try: - if isinstance(x, tf.Tensor): - x_np = x.numpy() - else: - x_np = x - - if isinstance(predictions, tf.Tensor): - pred_np = predictions.numpy() - else: - pred_np = predictions - - monitor.log_batch( - inputs=x_np, - predictions=pred_np, - ) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return predictions - - # Replace predict method - model.predict = logged_predict - return model - - -__all__ = [ - "KerasMonitor", - "WhiteBoxAICallback", - "TorchMonitor", # Deprecated - "wrap_keras_model", -] diff --git a/src/whiteboxai/integrations/transformers.py b/src/whiteboxai/integrations/transformers.py deleted file mode 100644 index 550a22d..0000000 --- a/src/whiteboxai/integrations/transformers.py +++ /dev/null @@ -1,531 +0,0 @@ -""" -Hugging Face Transformers Integration - -Integration for monitoring Hugging Face Transformers models. -""" - -import warnings -from typing import Any, Dict, List, Optional, Union - -try: - import transformers - from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer - - TRANSFORMERS_AVAILABLE = True -except ImportError: - TRANSFORMERS_AVAILABLE = False - PreTrainedModel = object - PreTrainedTokenizer = object - Pipeline = object - -from whiteboxai.monitor import ModelMonitor - - -class TransformersMonitor(ModelMonitor): - """ - Monitor for Hugging Face Transformers models. - - Supports: - - Text classification - - Named entity recognition (NER) - - Question answering - - Text generation - - Translation - - Summarization - - Sentiment analysis - - Example: - ```python - from transformers import pipeline - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.transformers import TransformersMonitor - - # Load model - classifier = pipeline("sentiment-analysis") - - # Setup monitoring - client = WhiteBoxAI(api_key="your-api-key") - monitor = TransformersMonitor( - client=client, - pipeline=classifier, - model_name="sentiment_classifier" - ) - - # Make predictions with automatic logging - result = monitor.predict("I love this product!", log=True) - ``` - """ - - def __init__( - self, - client, - pipeline: Optional[Pipeline] = None, - model: Optional[PreTrainedModel] = None, - tokenizer: Optional[PreTrainedTokenizer] = None, - model_name: Optional[str] = None, - task: Optional[str] = None, - **kwargs, - ): - """ - Initialize Transformers monitor. - - Args: - client: WhiteBoxAI client instance - pipeline: Hugging Face pipeline (recommended) - model: PreTrainedModel (if not using pipeline) - tokenizer: PreTrainedTokenizer (if not using pipeline) - model_name: Name for the model - task: Task type (text-classification, ner, qa, etc.) - **kwargs: Additional arguments for ModelMonitor - """ - if not TRANSFORMERS_AVAILABLE: - raise ImportError( - "transformers is not installed. Install with: pip install transformers" - ) - - super().__init__(client, **kwargs) - self.pipeline = pipeline - self.model = model - self.tokenizer = tokenizer - self._model_name = model_name - self._task = task - self._registered = False - - # Auto-detect task from pipeline - if pipeline is not None and task is None: - self._task = getattr(pipeline, "task", None) - - def register_from_model( - self, - name: Optional[str] = None, - task: Optional[str] = None, - version: Optional[str] = None, - framework_version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register Transformers model with WhiteBoxAI. - - Args: - name: Model name (default: model class name or pipeline task) - task: Task type (text-classification, ner, qa, etc.) - version: Model version - framework_version: Transformers version - **kwargs: Additional metadata - - Returns: - Model ID - """ - if self.pipeline is None and self.model is None: - raise ValueError("No model or pipeline provided") - - # Auto-detect model name - if name is None: - if self._model_name: - name = self._model_name - elif self.pipeline is not None: - name = f"{self.pipeline.task}_pipeline" - else: - name = self.model.__class__.__name__ - - # Use provided or detected task - task = task or self._task or "text-classification" - - # Get Transformers version - if framework_version is None: - framework_version = transformers.__version__ - - # Extract model metadata - metadata = { - "framework": "transformers", - "framework_version": framework_version, - "task": task, - **kwargs, - } - - # Get model information from pipeline or model - if self.pipeline is not None: - metadata["pipeline_type"] = self.pipeline.task - if hasattr(self.pipeline, "model"): - metadata["model_class"] = self.pipeline.model.__class__.__name__ - if hasattr(self.pipeline.model, "config"): - config = self.pipeline.model.config - metadata["model_type"] = getattr(config, "model_type", None) - metadata["num_parameters"] = getattr(config, "num_parameters", None) - metadata["vocab_size"] = getattr(config, "vocab_size", None) - elif self.model is not None: - metadata["model_class"] = self.model.__class__.__name__ - if hasattr(self.model, "config"): - config = self.model.config - metadata["model_type"] = getattr(config, "model_type", None) - metadata["num_parameters"] = getattr(config, "num_parameters", None) - - # Map task to model_type - model_type_mapping = { - "text-classification": "classification", - "sentiment-analysis": "classification", - "ner": "classification", - "token-classification": "classification", - "question-answering": "qa", - "text-generation": "generation", - "translation": "generation", - "summarization": "generation", - } - model_type = model_type_mapping.get(task, "classification") - - # Register model - self.model_id = self.client.register_model( - name=name, - model_type=model_type, - framework="transformers", - version=version, - metadata=metadata, - ) - - self._registered = True - return self.model_id - - def predict( - self, - inputs: Union[str, List[str]], - log: bool = True, - actuals: Optional[Union[str, List[str]]] = None, - **kwargs: Any, - ) -> Union[Dict, List[Dict]]: - """ - Make predictions and optionally log them. - - Args: - inputs: Input text or list of texts - log: Whether to log the prediction - actuals: Actual values (if available) - **kwargs: Additional arguments for the pipeline/model - - Returns: - Predictions - """ - if self.pipeline is None and self.model is None: - raise ValueError("No model or pipeline provided") - - # Ensure model is registered - if not self._registered: - self.register_from_model() - - # Make prediction using pipeline - if self.pipeline is not None: - predictions = self.pipeline(inputs, **kwargs) - else: - # Use model + tokenizer - if self.tokenizer is None: - raise ValueError("Tokenizer required when using model directly") - - # Tokenize inputs - encoded = self.tokenizer( - inputs, return_tensors="pt", padding=True, truncation=True, **kwargs - ) - - # Get predictions - outputs = self.model(**encoded) - predictions = outputs.logits.detach().cpu().numpy() - - # Log if requested - if log: - # Determine if batch or single - is_batch = isinstance(inputs, list) - - if is_batch: - # Batch prediction - self.log_batch_transformers( - inputs=inputs, - predictions=predictions, - actuals=actuals, - ) - else: - # Single prediction - self.log_prediction_transformers( - input_text=inputs, - prediction=predictions if isinstance(predictions, dict) else predictions[0], - actual=actuals, - ) - - return predictions - - def log_prediction_transformers( - self, - input_text: str, - prediction: Union[Dict, Any], - actual: Optional[Any] = None, - **kwargs: Any, - ) -> None: - """ - Log a single Transformers prediction. - - Args: - input_text: Input text - prediction: Prediction result (dict or value) - actual: Actual value (if available) - **kwargs: Additional metadata - """ - # Extract prediction value based on task - if isinstance(prediction, dict): - pred_value = self._extract_prediction_value(prediction) - else: - pred_value = prediction - - # Log to WhiteBoxAI - self.log_prediction( - inputs={"text": input_text}, - prediction=pred_value, - actual=actual, - metadata={ - "task": self._task, - "raw_prediction": str(prediction), - **kwargs, - }, - ) - - def log_batch_transformers( - self, - inputs: List[str], - predictions: List[Union[Dict, Any]], - actuals: Optional[List[Any]] = None, - **kwargs: Any, - ) -> None: - """ - Log batch of Transformers predictions. - - Args: - inputs: List of input texts - predictions: List of predictions - actuals: List of actual values (if available) - **kwargs: Additional metadata - """ - # Prepare batch data - batch_inputs = [{"text": text} for text in inputs] - batch_predictions = [ - self._extract_prediction_value(pred) if isinstance(pred, dict) else pred - for pred in predictions - ] - - # Log batch - self.log_batch( - inputs=batch_inputs, - predictions=batch_predictions, - actuals=actuals, - metadata={ - "task": self._task, - **kwargs, - }, - ) - - def _extract_prediction_value(self, prediction: Dict) -> Any: - """Extract the main prediction value from a pipeline output.""" - if isinstance(prediction, list) and len(prediction) > 0: - prediction = prediction[0] - - if isinstance(prediction, dict): - # Common keys in transformers outputs - if "label" in prediction: - return prediction["label"] - elif "answer" in prediction: - return prediction["answer"] - elif "generated_text" in prediction: - return prediction["generated_text"] - elif "translation_text" in prediction: - return prediction["translation_text"] - elif "summary_text" in prediction: - return prediction["summary_text"] - elif "score" in prediction: - return prediction["score"] - - return prediction - - def set_baseline( - self, - baseline_texts: List[str], - baseline_labels: Optional[List[Any]] = None, - ) -> None: - """ - Set baseline data for drift detection. - - Args: - baseline_texts: Baseline input texts - baseline_labels: Baseline labels (optional) - """ - # Make baseline predictions if labels not provided - if baseline_labels is None and self.pipeline is not None: - baseline_predictions = self.pipeline(baseline_texts) - baseline_labels = [ - self._extract_prediction_value(pred) for pred in baseline_predictions - ] - - # Convert to format expected by parent class - baseline_inputs = [{"text": text} for text in baseline_texts] - - # Set baseline through parent class - super().set_baseline(baseline_inputs, baseline_labels) - - def log_generation_metrics( - self, - prompt: str, - generated_text: str, - num_tokens: Optional[int] = None, - generation_time: Optional[float] = None, - **kwargs: Any, - ) -> None: - """ - Log metrics for text generation tasks. - - Args: - prompt: Input prompt - generated_text: Generated text - num_tokens: Number of tokens generated - generation_time: Time taken to generate (seconds) - **kwargs: Additional metrics - """ - metrics = { - "prompt": prompt, - "generated_text": generated_text, - "num_tokens": num_tokens, - "generation_time": generation_time, - **kwargs, - } - - self.log_custom_metric("text_generation", metrics) - - -class TransformersPipelineWrapper: - """ - Wrapper for Hugging Face pipelines with automatic logging. - - Example: - ```python - from transformers import pipeline - from whiteboxai import WhiteBoxAI - from whiteboxai.integrations.transformers import TransformersPipelineWrapper - - classifier = pipeline("sentiment-analysis") - client = WhiteBoxAI(api_key="your-api-key") - - # Wrap pipeline - wrapped = TransformersPipelineWrapper(classifier, client) - - # Predictions are automatically logged - result = wrapped("I love this!") - ``` - """ - - def __init__( - self, - pipeline: Pipeline, - client, - model_name: Optional[str] = None, - auto_register: bool = True, - ): - """ - Initialize pipeline wrapper. - - Args: - pipeline: Hugging Face pipeline - client: WhiteBoxAI client - model_name: Model name - auto_register: Auto-register model on first prediction - """ - self.pipeline = pipeline - self.monitor = TransformersMonitor( - client=client, - pipeline=pipeline, - model_name=model_name, - ) - self._auto_register = auto_register - - def __call__(self, *args, **kwargs): - """Make prediction with automatic logging.""" - # Auto-register if needed - if self._auto_register and not self.monitor._registered: - self.monitor.register_from_model() - - # Get inputs - inputs = args[0] if args else kwargs.get("inputs") - - # Make prediction - result = self.pipeline(*args, **kwargs) - - # Log prediction - try: - if isinstance(inputs, list): - self.monitor.log_batch_transformers( - inputs=inputs, - predictions=result, - ) - else: - self.monitor.log_prediction_transformers( - input_text=inputs, - prediction=result, - ) - except Exception as e: - warnings.warn(f"Failed to log prediction: {e}") - - return result - - -def wrap_transformers_pipeline( - pipeline: Pipeline, - monitor: TransformersMonitor, -) -> Pipeline: - """ - Wrap a Hugging Face pipeline to automatically log predictions. - - Args: - pipeline: Hugging Face pipeline to wrap - monitor: TransformersMonitor instance - - Returns: - Wrapped pipeline that logs predictions - - Example: - ```python - from transformers import pipeline - - classifier = pipeline("sentiment-analysis") - wrapped = wrap_transformers_pipeline(classifier, monitor) - - # Predictions automatically logged - result = wrapped("Great product!") - ``` - """ - original_call = pipeline.__call__ - - def logged_call(*args, **kwargs): - # Make prediction - result = original_call(*args, **kwargs) - - # Log to WhiteBoxAI - try: - inputs = args[0] if args else kwargs.get("inputs") - - if isinstance(inputs, list): - monitor.log_batch_transformers( - inputs=inputs, - predictions=result, - ) - else: - monitor.log_prediction_transformers( - input_text=inputs, - prediction=result, - ) - except Exception as e: - warnings.warn(f"Failed to log predictions: {e}") - - return result - - # Replace __call__ method - pipeline.__call__ = logged_call - return pipeline - - -__all__ = [ - "TransformersMonitor", - "TransformersPipelineWrapper", - "wrap_transformers_pipeline", -] diff --git a/src/whiteboxai/models/__init__.py b/src/whiteboxai/models/__init__.py deleted file mode 100644 index 5fbfb4d..0000000 --- a/src/whiteboxai/models/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Models package for WhiteBoxAI SDK.""" diff --git a/src/whiteboxai/monitor.py b/src/whiteboxai/monitor.py deleted file mode 100644 index 0fb1460..0000000 --- a/src/whiteboxai/monitor.py +++ /dev/null @@ -1,284 +0,0 @@ -""" -Model Monitoring - -Simplified monitoring interface for ML models. -""" - -from typing import TYPE_CHECKING, Any, Dict, List, Optional - -import numpy as np - -if TYPE_CHECKING: - from whiteboxai.client import WhiteBoxAI - - -class ModelMonitor: - """ - Simplified model monitoring interface. - - Example: - ```python - from whiteboxai import WhiteBoxAI, ModelMonitor - - client = WhiteBoxAI(api_key="your-api-key") - monitor = ModelMonitor(client, model_id=123) - - # Log prediction - monitor.log_prediction( - inputs={"feature1": 1.0, "feature2": 2.0}, - output=0.85 - ) - - # Log batch - monitor.log_batch([ - {"inputs": {...}, "output": 0.85}, - {"inputs": {...}, "output": 0.92}, - ]) - ``` - """ - - def __init__( - self, - client: "WhiteBoxAI", - model_id: Optional[int] = None, - model_name: Optional[str] = None, - auto_explain: bool = False, - sampling_rate: float = 1.0, - ): - """ - Initialize model monitor. - - Args: - client: WhiteBoxAI client instance - model_id: Model ID (if already registered) - model_name: Model name (for registration) - auto_explain: Automatically generate explanations - sampling_rate: Prediction sampling rate (0.0-1.0) - """ - self.client = client - self.model_id = model_id - self.model_name = model_name - self.auto_explain = auto_explain - self.sampling_rate = sampling_rate - self._baseline_data: Optional[np.ndarray] = None - - def register_model( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """ - Register a new model. - - Args: - name: Model name - model_type: Model type (classification, regression, etc.) - framework: ML framework - version: Model version - **kwargs: Additional metadata - - Returns: - Model ID - """ - result = self.client.models.register( - name=name, - model_type=model_type, - framework=framework, - version=version, - **kwargs, - ) - self.model_id = result["id"] - self.model_name = name - return self.model_id - - async def aregister_model( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> int: - """Async version of register_model().""" - result = await self.client.models.aregister( - name=name, - model_type=model_type, - framework=framework, - version=version, - **kwargs, - ) - self.model_id = result["id"] - self.model_name = name - return self.model_id - - def log_prediction( - self, - inputs: Any, - output: Any, - explain: Optional[bool] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> Optional[Dict[str, Any]]: - """ - Log a single prediction. - - Args: - inputs: Input features/data - output: Model prediction/output - explain: Generate explanation (default: use auto_explain) - metadata: Additional metadata - - Returns: - Prediction data if sampled, None if skipped - """ - if not self._should_sample(): - return None - - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - explain = explain if explain is not None else self.auto_explain - - return self.client.predictions.log( - model_id=self.model_id, - inputs=inputs, - outputs=output, - explain=explain, - metadata=metadata, - ) - - async def alog_prediction( - self, - inputs: Any, - output: Any, - explain: Optional[bool] = None, - metadata: Optional[Dict[str, Any]] = None, - ) -> Optional[Dict[str, Any]]: - """Async version of log_prediction().""" - if not self._should_sample(): - return None - - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - explain = explain if explain is not None else self.auto_explain - - return await self.client.predictions.alog( - model_id=self.model_id, - inputs=inputs, - outputs=output, - explain=explain, - metadata=metadata, - ) - - def log_batch( - self, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """ - Log multiple predictions in batch. - - Args: - predictions: List of prediction dictionaries - - Returns: - Batch logging result - """ - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - # Apply sampling - if self.sampling_rate < 1.0: - predictions = self._sample_predictions(predictions) - - return self.client.predictions.log_batch( - model_id=self.model_id, - predictions=predictions, - ) - - async def alog_batch( - self, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Async version of log_batch().""" - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - # Apply sampling - if self.sampling_rate < 1.0: - predictions = self._sample_predictions(predictions) - - return await self.client.predictions.alog_batch( - model_id=self.model_id, - predictions=predictions, - ) - - def set_baseline(self, data: np.ndarray) -> None: - """ - Set baseline data for drift detection. - - Args: - data: Baseline data array - """ - self._baseline_data = data - - def detect_drift( - self, - current_data: Optional[np.ndarray] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Detect drift for the model. - - Args: - current_data: Current data to compare against baseline - **kwargs: Additional detection parameters - - Returns: - Drift detection results - """ - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - return self.client.drift.detect( - model_id=self.model_id, - reference_data=( - self._baseline_data.tolist() if self._baseline_data is not None else None - ), - current_data=current_data.tolist() if current_data is not None else None, - **kwargs, - ) - - async def adetect_drift( - self, - current_data: Optional[np.ndarray] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of detect_drift().""" - if self.model_id is None: - raise ValueError("Model not registered. Call register_model() first.") - - return await self.client.drift.adetect( - model_id=self.model_id, - reference_data=( - self._baseline_data.tolist() if self._baseline_data is not None else None - ), - current_data=current_data.tolist() if current_data is not None else None, - **kwargs, - ) - - def _should_sample(self) -> bool: - """Check if prediction should be sampled.""" - if self.sampling_rate >= 1.0: - return True - return np.random.random() < self.sampling_rate - - def _sample_predictions(self, predictions: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Sample predictions based on sampling rate.""" - n_samples = int(len(predictions) * self.sampling_rate) - if n_samples == 0: - n_samples = 1 # Always log at least one - indices = np.random.choice(len(predictions), n_samples, replace=False) - return [predictions[i] for i in indices] diff --git a/src/whiteboxai/offline.py b/src/whiteboxai/offline.py deleted file mode 100644 index 903a77b..0000000 --- a/src/whiteboxai/offline.py +++ /dev/null @@ -1,541 +0,0 @@ -""" -Offline mode support for WhiteBoxAI SDK. - -This module provides offline queueing capabilities for the SDK, allowing -operations to be queued locally when the API is unavailable and synced -when connectivity is restored. - -Features: -- Persistent queue storage (SQLite-based) -- Automatic retry on connection restoration -- Configurable queue limits -- Operation prioritization -- Failed operation tracking - -Example: - Enable offline mode: - - >>> from whiteboxai import WhiteBoxAI - >>> client = WhiteBoxAI( - ... api_key="your-api-key", - ... enable_offline=True, - ... offline_dir="./whiteboxai_offline" - ... ) - >>> - >>> # Operations are queued when offline - >>> client.predict(...) # Queued if API unavailable - >>> - >>> # Manually sync - >>> client.sync_offline_queue() -""" - -import json -import logging -import os -import sqlite3 -import threading -from enum import Enum -from pathlib import Path -from typing import Any, Dict, List, Tuple - -logger = logging.getLogger(__name__) - - -class OperationType(Enum): - """Types of operations that can be queued.""" - - PREDICT = "predict" - REGISTER_MODEL = "register_model" - UPDATE_BASELINE = "update_baseline" - LOG_BATCH = "log_batch" - - -class OperationPriority(Enum): - """Priority levels for queued operations.""" - - LOW = 1 - NORMAL = 2 - HIGH = 3 - CRITICAL = 4 - - -class OfflineQueue: - """ - Persistent queue for offline operations. - - Uses SQLite for storage to ensure operations are not lost even if - the application crashes or restarts. - - Attributes: - db_path: Path to the SQLite database file - max_queue_size: Maximum number of operations to queue - auto_sync: Whether to automatically sync when connection available - """ - - def __init__(self, db_path: str, max_queue_size: int = 10000, auto_sync: bool = True): - """ - Initialize offline queue. - - Args: - db_path: Path to SQLite database file - max_queue_size: Maximum operations to queue (0 = unlimited) - auto_sync: Enable automatic syncing - """ - self.db_path = db_path - self.max_queue_size = max_queue_size - self.auto_sync = auto_sync - self._lock = threading.Lock() - self._ensure_db() - - def _ensure_db(self): - """Create database and tables if they don't exist.""" - # Ensure directory exists - os.makedirs(os.path.dirname(self.db_path), exist_ok=True) - - with sqlite3.connect(self.db_path) as conn: - conn.execute( - """ - CREATE TABLE IF NOT EXISTS queue ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - operation_type TEXT NOT NULL, - priority INTEGER NOT NULL, - data TEXT NOT NULL, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - retry_count INTEGER DEFAULT 0, - last_error TEXT, - status TEXT DEFAULT 'pending' - ) - """ - ) - - # Create index for efficient querying - conn.execute( - """ - CREATE INDEX IF NOT EXISTS idx_status_priority - ON queue(status, priority DESC, created_at ASC) - """ - ) - - conn.commit() - - def enqueue( - self, - operation_type: OperationType, - data: Dict[str, Any], - priority: OperationPriority = OperationPriority.NORMAL, - ) -> int: - """ - Add an operation to the queue. - - Args: - operation_type: Type of operation - data: Operation data (will be JSON serialized) - priority: Operation priority - - Returns: - Operation ID - - Raises: - ValueError: If queue is full - """ - with self._lock: - # Check queue size - if self.max_queue_size > 0: - current_size = self.get_queue_size() - if current_size >= self.max_queue_size: - raise ValueError(f"Queue is full ({current_size}/{self.max_queue_size})") - - # Insert operation - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - INSERT INTO queue (operation_type, priority, data) - VALUES (?, ?, ?) - """, - (operation_type.value, priority.value, json.dumps(data)), - ) - conn.commit() - op_id = cursor.lastrowid - - logger.info(f"Queued operation {op_id}: {operation_type.value}") - return op_id - - def dequeue(self, limit: int = 100) -> List[Tuple[int, OperationType, Dict[str, Any]]]: - """ - Get pending operations from queue. - - Returns operations ordered by priority (highest first) and - creation time (oldest first). - - Args: - limit: Maximum number of operations to return - - Returns: - List of (id, operation_type, data) tuples - """ - with self._lock: - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - SELECT id, operation_type, data - FROM queue - WHERE status = 'pending' - ORDER BY priority DESC, created_at ASC - LIMIT ? - """, - (limit,), - ) - - operations = [] - for row in cursor.fetchall(): - op_id, op_type, data_json = row - operations.append((op_id, OperationType(op_type), json.loads(data_json))) - - return operations - - def mark_success(self, operation_id: int): - """ - Mark operation as successfully completed. - - Args: - operation_id: Operation ID - """ - with self._lock: - with sqlite3.connect(self.db_path) as conn: - conn.execute( - """ - UPDATE queue - SET status = 'completed' - WHERE id = ? - """, - (operation_id,), - ) - conn.commit() - - logger.debug(f"Marked operation {operation_id} as completed") - - def mark_failure(self, operation_id: int, error: str, max_retries: int = 3): - """ - Mark operation as failed and increment retry count. - - Args: - operation_id: Operation ID - error: Error message - max_retries: Maximum retry attempts - """ - with self._lock: - with sqlite3.connect(self.db_path) as conn: - # Increment retry count - cursor = conn.execute( - """ - UPDATE queue - SET retry_count = retry_count + 1, - last_error = ? - WHERE id = ? - """, - (error, operation_id), - ) - - # Check if max retries exceeded - cursor = conn.execute( - """ - SELECT retry_count FROM queue WHERE id = ? - """, - (operation_id,), - ) - retry_count = cursor.fetchone()[0] - - if retry_count >= max_retries: - conn.execute( - """ - UPDATE queue - SET status = 'failed' - WHERE id = ? - """, - (operation_id,), - ) - logger.error( - f"Operation {operation_id} permanently failed after " - f"{retry_count} retries: {error}" - ) - else: - logger.warning( - f"Operation {operation_id} failed (retry {retry_count}/" - f"{max_retries}): {error}" - ) - - conn.commit() - - def get_queue_size(self, status: str = "pending") -> int: - """ - Get number of operations in queue. - - Args: - status: Filter by status ('pending', 'completed', 'failed', or None for all) - - Returns: - Number of operations - """ - with sqlite3.connect(self.db_path) as conn: - if status: - cursor = conn.execute("SELECT COUNT(*) FROM queue WHERE status = ?", (status,)) - else: - cursor = conn.execute("SELECT COUNT(*) FROM queue") - - return cursor.fetchone()[0] - - def get_statistics(self) -> Dict[str, int]: - """ - Get queue statistics. - - Returns: - Dictionary with counts by status - """ - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - SELECT status, COUNT(*) as count - FROM queue - GROUP BY status - """ - ) - - stats = {"total": 0, "pending": 0, "completed": 0, "failed": 0} - - for row in cursor.fetchall(): - status, count = row - stats[status] = count - stats["total"] += count - - return stats - - def clear_completed(self, older_than_days: int = 7): - """ - Remove completed operations older than specified days. - - Args: - older_than_days: Remove operations older than this many days - """ - with self._lock: - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - DELETE FROM queue - WHERE status = 'completed' - AND created_at < datetime('now', '-' || ? || ' days') - """, - (older_than_days,), - ) - deleted = cursor.rowcount - conn.commit() - - if deleted > 0: - logger.info(f"Cleared {deleted} completed operations older than {older_than_days} days") - - def clear_all(self): - """Clear all operations from queue (use with caution).""" - with self._lock: - with sqlite3.connect(self.db_path) as conn: - conn.execute("DELETE FROM queue") - conn.commit() - - logger.warning("Cleared all operations from queue") - - def get_failed_operations(self) -> List[Dict[str, Any]]: - """ - Get all permanently failed operations. - - Returns: - List of failed operation details - """ - with sqlite3.connect(self.db_path) as conn: - cursor = conn.execute( - """ - SELECT id, operation_type, data, created_at, retry_count, last_error - FROM queue - WHERE status = 'failed' - ORDER BY created_at DESC - """ - ) - - failed = [] - for row in cursor.fetchall(): - op_id, op_type, data_json, created_at, retry_count, last_error = row - failed.append( - { - "id": op_id, - "operation_type": op_type, - "data": json.loads(data_json), - "created_at": created_at, - "retry_count": retry_count, - "last_error": last_error, - } - ) - - return failed - - -class OfflineManager: - """ - Manages offline mode for WhiteBoxAI SDK. - - Handles automatic queueing when offline and syncing when online. - - Attributes: - queue: OfflineQueue instance - sync_interval: Seconds between automatic sync attempts - max_retries: Maximum retry attempts for failed operations - """ - - def __init__( - self, - offline_dir: str = "./whiteboxai_offline", - max_queue_size: int = 10000, - auto_sync: bool = True, - sync_interval: int = 60, - max_retries: int = 3, - ): - """ - Initialize offline manager. - - Args: - offline_dir: Directory for offline storage - max_queue_size: Maximum operations to queue - auto_sync: Enable automatic syncing - sync_interval: Seconds between sync attempts - max_retries: Maximum retry attempts - """ - self.offline_dir = Path(offline_dir) - self.offline_dir.mkdir(parents=True, exist_ok=True) - - db_path = str(self.offline_dir / "queue.db") - self.queue = OfflineQueue( - db_path=db_path, max_queue_size=max_queue_size, auto_sync=auto_sync - ) - - self.sync_interval = sync_interval - self.max_retries = max_retries - self._sync_thread = None - self._stop_sync = threading.Event() - self._client = None # Will be set by WhiteBoxAI client - - if auto_sync: - self.start_auto_sync() - - def set_client(self, client): - """ - Set the WhiteBoxAI client for syncing. - - Args: - client: WhiteBoxAI client instance - """ - self._client = client - - def start_auto_sync(self): - """Start automatic sync thread.""" - if self._sync_thread is None or not self._sync_thread.is_alive(): - self._stop_sync.clear() - self._sync_thread = threading.Thread(target=self._auto_sync_loop, daemon=True) - self._sync_thread.start() - logger.info("Started automatic sync thread") - - def stop_auto_sync(self): - """Stop automatic sync thread.""" - if self._sync_thread and self._sync_thread.is_alive(): - self._stop_sync.set() - self._sync_thread.join(timeout=5) - logger.info("Stopped automatic sync thread") - - def _auto_sync_loop(self): - """Background thread for automatic syncing.""" - while not self._stop_sync.is_set(): - try: - if self._client: - self.sync(batch_size=100) - except Exception as e: - logger.error(f"Error in auto-sync: {e}") - - # Wait for next sync interval - self._stop_sync.wait(self.sync_interval) - - def sync(self, batch_size: int = 100) -> Dict[str, int]: - """ - Sync queued operations with API. - - Args: - batch_size: Number of operations to sync per batch - - Returns: - Dictionary with sync statistics - """ - if not self._client: - logger.warning("No client set, cannot sync") - return {"synced": 0, "failed": 0, "pending": self.queue.get_queue_size()} - - stats = {"synced": 0, "failed": 0} - - # Get pending operations - operations = self.queue.dequeue(limit=batch_size) - - if not operations: - return {**stats, "pending": 0} - - logger.info(f"Syncing {len(operations)} operations...") - - for op_id, op_type, data in operations: - try: - # Execute operation - if op_type == OperationType.PREDICT: - self._client._api_predict(**data) - elif op_type == OperationType.REGISTER_MODEL: - self._client._api_register_model(**data) - elif op_type == OperationType.UPDATE_BASELINE: - self._client._api_update_baseline(**data) - elif op_type == OperationType.LOG_BATCH: - self._client._api_log_batch(**data) - - # Mark success - self.queue.mark_success(op_id) - stats["synced"] += 1 - - except Exception as e: - # Mark failure - error_msg = str(e) - self.queue.mark_failure(op_id, error_msg, self.max_retries) - stats["failed"] += 1 - - stats["pending"] = self.queue.get_queue_size() - - if stats["synced"] > 0: - logger.info( - f"Synced {stats['synced']} operations, {stats['failed']} failed, {stats['pending']} pending" - ) - - return stats - - def get_status(self) -> Dict[str, Any]: - """ - Get offline mode status. - - Returns: - Status dictionary with queue stats and sync info - """ - stats = self.queue.get_statistics() - - return { - "enabled": True, - "auto_sync": self.queue.auto_sync, - "queue_stats": stats, - "sync_interval": self.sync_interval, - "offline_dir": str(self.offline_dir), - "auto_sync_running": self._sync_thread and self._sync_thread.is_alive(), - } - - def cleanup(self, older_than_days: int = 7): - """ - Clean up old completed operations. - - Args: - older_than_days: Remove operations older than this many days - """ - self.queue.clear_completed(older_than_days) diff --git a/src/whiteboxai/privacy.py b/src/whiteboxai/privacy.py deleted file mode 100644 index 86c97b8..0000000 --- a/src/whiteboxai/privacy.py +++ /dev/null @@ -1,219 +0,0 @@ -""" -Privacy Utilities - -PII detection and data masking utilities. -""" - -import re -from typing import Any, Dict, List, Optional, Pattern - - -class PIIDetector: - """ - Detect personally identifiable information (PII) in data. - - Supports: - - Email addresses - - Phone numbers - - Credit card numbers - - Social security numbers - - IP addresses - """ - - def __init__(self): - """Initialize PII detector with regex patterns.""" - self.patterns: Dict[str, Pattern] = { - "email": re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"), - "phone": re.compile(r"\b(?:\+?1[-.]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"), - "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"), - "credit_card": re.compile(r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b"), - "ip_address": re.compile(r"\b(?:\d{1,3}\.){3}\d{1,3}\b"), - } - - def detect(self, text: str) -> List[Dict[str, Any]]: - """ - Detect PII in text. - - Args: - text: Text to scan for PII - - Returns: - List of detected PII items with type and location - """ - detections = [] - - for pii_type, pattern in self.patterns.items(): - for match in pattern.finditer(text): - detections.append( - { - "type": pii_type, - "value": match.group(), - "start": match.start(), - "end": match.end(), - } - ) - - return detections - - def mask( - self, - text: str, - mask_char: str = "*", - preserve_length: bool = True, - ) -> str: - """ - Mask PII in text. - - Args: - text: Text to mask - mask_char: Character to use for masking - preserve_length: Whether to preserve original length - - Returns: - Masked text - """ - detections = self.detect(text) - - # Sort by position (reverse order to avoid index shifting) - detections.sort(key=lambda x: x["start"], reverse=True) - - # Mask each detection - for detection in detections: - start = detection["start"] - end = detection["end"] - original = detection["value"] - - if preserve_length: - masked = mask_char * len(original) - else: - masked = f"[{detection['type'].upper()}]" - - text = text[:start] + masked + text[end:] - - return text - - -class DataMasker: - """ - Mask sensitive data in structured formats (dict, list). - """ - - def __init__(self, pii_detector: Optional[PIIDetector] = None): - """ - Initialize data masker. - - Args: - pii_detector: PII detector instance (creates new if None) - """ - self.pii_detector = pii_detector or PIIDetector() - self.sensitive_keys = { - "password", - "passwd", - "pwd", - "secret", - "token", - "api_key", - "apikey", - "access_token", - "refresh_token", - } - - def mask( - self, - data: Any, - mask_pii: bool = True, - mask_sensitive_keys: bool = True, - ) -> Any: - """ - Mask sensitive data. - - Args: - data: Data to mask (dict, list, or primitive) - mask_pii: Whether to mask PII in strings - mask_sensitive_keys: Whether to mask values of sensitive keys - - Returns: - Masked data - """ - if isinstance(data, dict): - return self._mask_dict(data, mask_pii, mask_sensitive_keys) - elif isinstance(data, list): - return self._mask_list(data, mask_pii, mask_sensitive_keys) - elif isinstance(data, str): - if mask_pii: - return self.pii_detector.mask(data) - return data - else: - return data - - def _mask_dict( - self, - data: dict, - mask_pii: bool, - mask_sensitive_keys: bool, - ) -> dict: - """Mask sensitive data in dictionary.""" - masked = {} - for key, value in data.items(): - # Check if key is sensitive - if mask_sensitive_keys and key.lower() in self.sensitive_keys: - masked[key] = "***MASKED***" - else: - masked[key] = self.mask(value, mask_pii, mask_sensitive_keys) - return masked - - def _mask_list( - self, - data: list, - mask_pii: bool, - mask_sensitive_keys: bool, - ) -> list: - """Mask sensitive data in list.""" - return [self.mask(item, mask_pii, mask_sensitive_keys) for item in data] - - -# Global instances -_pii_detector = PIIDetector() -_data_masker = DataMasker(_pii_detector) - - -def detect_pii(text: str) -> List[Dict[str, Any]]: - """ - Detect PII in text using global detector. - - Args: - text: Text to scan - - Returns: - List of detected PII items - """ - return _pii_detector.detect(text) - - -def mask_pii(text: str, mask_char: str = "*") -> str: - """ - Mask PII in text using global detector. - - Args: - text: Text to mask - mask_char: Character to use for masking - - Returns: - Masked text - """ - return _pii_detector.mask(text, mask_char=mask_char) - - -def mask_data(data: Any, mask_pii: bool = True, mask_sensitive_keys: bool = True) -> Any: - """ - Mask sensitive data using global masker. - - Args: - data: Data to mask - mask_pii: Whether to mask PII - mask_sensitive_keys: Whether to mask sensitive keys - - Returns: - Masked data - """ - return _data_masker.mask(data, mask_pii, mask_sensitive_keys) diff --git a/src/whiteboxai/resources.py b/src/whiteboxai/resources.py deleted file mode 100644 index d0908a4..0000000 --- a/src/whiteboxai/resources.py +++ /dev/null @@ -1,460 +0,0 @@ -""" -API Resources - -Resource classes for interacting with different WhiteBoxAI API endpoints. -""" - -from typing import TYPE_CHECKING, Any, Dict, List, Optional - -if TYPE_CHECKING: - from whiteboxai.client import WhiteBoxAI - - -class BaseResource: - """Base class for API resources.""" - - def __init__(self, client: "WhiteBoxAI"): - """Initialize resource with client.""" - self.client = client - - -class ModelsResource(BaseResource): - """Models API resource.""" - - def register( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Register a new model. - - Args: - name: Model name - model_type: Model type (classification, regression, etc.) - framework: ML framework (sklearn, pytorch, tensorflow, etc.) - version: Model version - **kwargs: Additional model metadata - - Returns: - Registered model data - """ - data = { - "name": name, - "model_type": model_type, - "framework": framework, - "version": version, - **kwargs, - } - return self.client.request("POST", "/api/v1/models", data=data) - - async def aregister( - self, - name: str, - model_type: str, - framework: Optional[str] = None, - version: Optional[str] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of register().""" - data = { - "name": name, - "model_type": model_type, - "framework": framework, - "version": version, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/models", data=data) - - def get(self, model_id: int) -> Dict[str, Any]: - """Get model by ID.""" - return self.client.request("GET", f"/api/v1/models/{model_id}") - - async def aget(self, model_id: int) -> Dict[str, Any]: - """Async version of get().""" - return await self.client.arequest("GET", f"/api/v1/models/{model_id}") - - def list(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: - """List registered models.""" - params = {"limit": limit, "offset": offset} - return self.client.request("GET", "/api/v1/models", params=params) - - async def alist(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]: - """Async version of list().""" - params = {"limit": limit, "offset": offset} - return await self.client.arequest("GET", "/api/v1/models", params=params) - - def update(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: - """Update model metadata.""" - return self.client.request("PUT", f"/api/v1/models/{model_id}", data=kwargs) - - async def aupdate(self, model_id: int, **kwargs: Any) -> Dict[str, Any]: - """Async version of update().""" - return await self.client.arequest("PUT", f"/api/v1/models/{model_id}", data=kwargs) - - -class PredictionsResource(BaseResource): - """Predictions API resource.""" - - def log( - self, - model_id: int, - inputs: Any, - outputs: Any, - explain: bool = False, - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """ - Log a single prediction. - - Args: - model_id: Model ID - inputs: Input features/data - outputs: Model prediction/output - explain: Whether to generate explanation - metadata: Additional prediction metadata - - Returns: - Logged prediction data - """ - data = { - "model_id": model_id, - "inputs": inputs, - "outputs": outputs, - "explain": explain, - "metadata": metadata or {}, - } - return self.client.request("POST", "/api/v1/predictions/log", data=data) - - async def alog( - self, - model_id: int, - inputs: Any, - outputs: Any, - explain: bool = False, - metadata: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Async version of log().""" - data = { - "model_id": model_id, - "inputs": inputs, - "outputs": outputs, - "explain": explain, - "metadata": metadata or {}, - } - return await self.client.arequest("POST", "/api/v1/predictions/log", data=data) - - def log_batch( - self, - model_id: int, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """ - Log multiple predictions in batch. - - Args: - model_id: Model ID - predictions: List of prediction dictionaries - - Returns: - Batch logging result - """ - data = {"model_id": model_id, "predictions": predictions} - return self.client.request("POST", "/api/v1/predictions/log/batch", data=data) - - async def alog_batch( - self, - model_id: int, - predictions: List[Dict[str, Any]], - ) -> Dict[str, Any]: - """Async version of log_batch().""" - data = {"model_id": model_id, "predictions": predictions} - return await self.client.arequest("POST", "/api/v1/predictions/log/batch", data=data) - - -class ExplanationsResource(BaseResource): - """Explanations API resource.""" - - def generate( - self, - prediction_id: int, - method: str = "shap", - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Generate explanation for a prediction. - - Args: - prediction_id: Prediction ID - method: Explanation method (shap, lime) - **kwargs: Method-specific parameters - - Returns: - Explanation data - """ - data = {"prediction_id": prediction_id, "method": method, **kwargs} - return self.client.request("POST", "/api/v1/explanations/generate", data=data) - - async def agenerate( - self, - prediction_id: int, - method: str = "shap", - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of generate().""" - data = {"prediction_id": prediction_id, "method": method, **kwargs} - return await self.client.arequest("POST", "/api/v1/explanations/generate", data=data) - - def get(self, explanation_id: int) -> Dict[str, Any]: - """Get explanation by ID.""" - return self.client.request("GET", f"/api/v1/explanations/{explanation_id}") - - async def aget(self, explanation_id: int) -> Dict[str, Any]: - """Async version of get().""" - return await self.client.arequest("GET", f"/api/v1/explanations/{explanation_id}") - - -class DriftResource(BaseResource): - """Drift detection API resource.""" - - def detect( - self, - model_id: int, - reference_data: Optional[Any] = None, - current_data: Optional[Any] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Detect drift for a model. - - Args: - model_id: Model ID - reference_data: Reference/baseline data - current_data: Current data to compare - **kwargs: Detection parameters - - Returns: - Drift detection results - """ - data = { - "model_id": model_id, - "reference_data": reference_data, - "current_data": current_data, - **kwargs, - } - return self.client.request("POST", "/api/v1/drift/detect", data=data) - - async def adetect( - self, - model_id: int, - reference_data: Optional[Any] = None, - current_data: Optional[Any] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of detect().""" - data = { - "model_id": model_id, - "reference_data": reference_data, - "current_data": current_data, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/drift/detect", data=data) - - def get_report(self, model_id: int, report_id: int) -> Dict[str, Any]: - """Get drift report.""" - return self.client.request("GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}") - - async def aget_report(self, model_id: int, report_id: int) -> Dict[str, Any]: - """Async version of get_report().""" - return await self.client.arequest( - "GET", f"/api/v1/drift/models/{model_id}/reports/{report_id}" - ) - - -class AlertsResource(BaseResource): - """Alerts API resource.""" - - def create( - self, - name: str, - alert_type: str, - conditions: Dict[str, Any], - **kwargs: Any, - ) -> Dict[str, Any]: - """ - Create an alert rule. - - Args: - name: Alert name - alert_type: Alert type - conditions: Alert conditions - **kwargs: Additional alert configuration - - Returns: - Created alert data - """ - data = { - "name": name, - "alert_type": alert_type, - "conditions": conditions, - **kwargs, - } - return self.client.request("POST", "/api/v1/alerts", data=data) - - async def acreate( - self, - name: str, - alert_type: str, - conditions: Dict[str, Any], - **kwargs: Any, - ) -> Dict[str, Any]: - """Async version of create().""" - data = { - "name": name, - "alert_type": alert_type, - "conditions": conditions, - **kwargs, - } - return await self.client.arequest("POST", "/api/v1/alerts", data=data) - - def list(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: - """List alert rules.""" - params = {"model_id": model_id} if model_id else {} - return self.client.request("GET", "/api/v1/alerts", params=params) - - async def alist(self, model_id: Optional[int] = None) -> List[Dict[str, Any]]: - """Async version of list().""" - params = {"model_id": model_id} if model_id else {} - return await self.client.arequest("GET", "/api/v1/alerts", params=params) - - -class AgentWorkflowsResource(BaseResource): - """Agent Workflows API resource for multi-agent monitoring.""" - - def create( - self, - name: str, - framework: str, - inputs: Optional[Dict[str, Any]] = None, - meta_data: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Create a new agent workflow.""" - data = { - "name": name, - "framework": framework, - "inputs": inputs or {}, - "meta_data": meta_data or {}, - } - return self.client.request("POST", "/api/v1/workflows/multi-agent", data=data) - - def start( - self, - workflow_id: str, - inputs: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Start an agent workflow.""" - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/start", - data={"inputs": inputs or {}}, - ) - - def complete( - self, - workflow_id: str, - outputs: Optional[Dict[str, Any]] = None, - status: str = "completed", - ) -> Dict[str, Any]: - """Complete an agent workflow.""" - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/complete", - data={"outputs": outputs or {}, "status": status}, - ) - - def register_agent( - self, - workflow_id: str, - name: str, - role: Optional[str] = None, - model_name: Optional[str] = None, - tools: Optional[List[str]] = None, - **kwargs: Any, - ) -> Dict[str, Any]: - """Register an agent in a workflow.""" - data = { - "name": name, - "role": role or name, - "model_name": model_name, - "tools": tools or [], - **kwargs, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/agents", - data=data, - ) - - def create_execution( - self, - workflow_id: str, - agent_name: str, - status: str, - inputs: Optional[Dict[str, Any]] = None, - outputs: Optional[Dict[str, Any]] = None, - duration_ms: Optional[int] = None, - llm_call_count: Optional[int] = None, - tool_call_count: Optional[int] = None, - tokens_used: Optional[int] = None, - cost: Optional[float] = None, - ) -> Dict[str, Any]: - """Log an agent execution within a workflow.""" - data = { - "agent_name": agent_name, - "status": status, - "inputs": inputs, - "outputs": outputs, - "duration_ms": duration_ms, - "llm_call_count": llm_call_count, - "tool_call_count": tool_call_count, - "tokens_used": tokens_used, - "cost": cost, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/executions", - data=data, - ) - - def create_interaction( - self, - workflow_id: str, - from_agent: str, - to_agent: str, - interaction_type: str, - message: Optional[str] = None, - meta_data: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: - """Log an agent-to-agent interaction.""" - data = { - "from_agent": from_agent, - "to_agent": to_agent, - "interaction_type": interaction_type, - "message": message, - "meta_data": meta_data or {}, - } - return self.client.request( - "POST", - f"/api/v1/workflows/multi-agent/{workflow_id}/interactions", - data=data, - ) - - def get_analytics(self, workflow_id: str) -> Dict[str, Any]: - """Get analytics for a workflow.""" - return self.client.request( - "GET", - f"/api/v1/workflows/multi-agent/{workflow_id}/analytics", - ) diff --git a/src/whiteboxai/utils.py b/src/whiteboxai/utils.py deleted file mode 100644 index 5c41860..0000000 --- a/src/whiteboxai/utils.py +++ /dev/null @@ -1,337 +0,0 @@ -""" -SDK Utilities - -General utility functions for the WhiteBoxAI SDK. -""" - -import json -from typing import Any, Dict, List, Union - -import numpy as np -import pandas as pd - - -def serialize_data(data: Any) -> Any: - """ - Serialize data for API transmission. - - Handles numpy arrays, pandas DataFrames, and other common data types. - - Args: - data: Data to serialize - - Returns: - JSON-serializable data - """ - if isinstance(data, np.ndarray): - return data.tolist() - elif isinstance(data, pd.DataFrame): - return data.to_dict(orient="records") - elif isinstance(data, pd.Series): - return data.to_dict() - elif isinstance(data, (np.integer, np.floating)): - return data.item() - elif isinstance(data, dict): - return {k: serialize_data(v) for k, v in data.items()} - elif isinstance(data, (list, tuple)): - return [serialize_data(item) for item in data] - else: - return data - - -def validate_model_type(model_type: str) -> bool: - """ - Validate model type. - - Args: - model_type: Model type string - - Returns: - True if valid, False otherwise - """ - valid_types = { - "classification", - "regression", - "clustering", - "ranking", - "recommendation", - "anomaly_detection", - "time_series", - "nlp", - "computer_vision", - "other", - } - return model_type.lower() in valid_types - - -def extract_feature_names(data: Union[np.ndarray, pd.DataFrame, Dict]) -> List[str]: - """ - Extract feature names from data. - - Args: - data: Input data - - Returns: - List of feature names - """ - if isinstance(data, pd.DataFrame): - return data.columns.tolist() - elif isinstance(data, dict): - return list(data.keys()) - elif isinstance(data, np.ndarray): - # Generate default feature names - n_features = data.shape[1] if len(data.shape) > 1 else 1 - return [f"feature_{i}" for i in range(n_features)] - else: - return [] - - -def format_bytes(size: int) -> str: - """ - Format byte size as human-readable string. - - Args: - size: Size in bytes - - Returns: - Formatted string (e.g., "1.5 MB") - """ - for unit in ["B", "KB", "MB", "GB", "TB"]: - if size < 1024.0: - return f"{size:.1f} {unit}" - size /= 1024.0 - return f"{size:.1f} PB" - - -def truncate_string(text: str, max_length: int = 100, suffix: str = "...") -> str: - """ - Truncate string to maximum length. - - Args: - text: String to truncate - max_length: Maximum length - suffix: Suffix to append if truncated - - Returns: - Truncated string - """ - if len(text) <= max_length: - return text - return text[: max_length - len(suffix)] + suffix - - -def dict_to_query_params(params: Dict[str, Any]) -> str: - """ - Convert dictionary to URL query parameters. - - Args: - params: Parameter dictionary - - Returns: - Query string - """ - parts = [] - for key, value in params.items(): - if value is not None: - if isinstance(value, (list, tuple)): - for item in value: - parts.append(f"{key}={item}") - else: - parts.append(f"{key}={value}") - return "&".join(parts) - - -def deep_merge(dict1: Dict, dict2: Dict) -> Dict: - """ - Deep merge two dictionaries. - - Args: - dict1: First dictionary - dict2: Second dictionary (takes precedence) - - Returns: - Merged dictionary - """ - result = dict1.copy() - for key, value in dict2.items(): - if key in result and isinstance(result[key], dict) and isinstance(value, dict): - result[key] = deep_merge(result[key], value) - else: - result[key] = value - return result - - -def safe_json_loads(text: str, default: Any = None) -> Any: - """ - Safely load JSON with fallback. - - Args: - text: JSON string - default: Default value if parsing fails - - Returns: - Parsed JSON or default value - """ - try: - return json.loads(text) - except (json.JSONDecodeError, TypeError): - return default - - -def is_numeric(value: Any) -> bool: - """ - Check if value is numeric. - - Args: - value: Value to check - - Returns: - True if numeric, False otherwise - """ - try: - float(value) - return True - except (ValueError, TypeError): - return False - - -def validate_features(features: Any) -> bool: - """ - Validate features for model input. - - Args: - features: Features to validate - - Returns: - True if valid - - Raises: - ValueError: If features are invalid - """ - if features is None: - raise ValueError("Features cannot be None") - - if isinstance(features, (list, tuple)): - if len(features) == 0: - raise ValueError("Features cannot be empty") - return True - elif isinstance(features, np.ndarray): - if features.size == 0: - raise ValueError("Features cannot be empty") - return True - elif isinstance(features, dict): - if len(features) == 0: - raise ValueError("Features cannot be empty") - return True - else: - raise ValueError(f"Unsupported features type: {type(features)}") - - -def format_features(features: Any) -> Dict[str, Any]: - """ - Format features for API transmission. - - Args: - features: Features to format - - Returns: - Formatted features dictionary - """ - if isinstance(features, dict): - return features - elif isinstance(features, (list, tuple, np.ndarray)): - return {"features": serialize_data(features)} - else: - return {"features": features} - - -def validate_prediction(prediction: Any) -> bool: - """ - Validate model prediction. - - Args: - prediction: Prediction to validate - - Returns: - True if valid - - Raises: - ValueError: If prediction is invalid - """ - if prediction is None: - raise ValueError("Prediction cannot be None") - return True - - -def serialize_numpy(arr: np.ndarray) -> List: - """ - Serialize numpy array to list. - - Args: - arr: Numpy array - - Returns: - List representation - """ - return arr.tolist() - - -def deserialize_numpy(data: List) -> np.ndarray: - """ - Deserialize list to numpy array. - - Args: - data: List data - - Returns: - Numpy array - """ - return np.array(data) - - -def compute_metrics(y_true: Any, y_pred: Any) -> Dict[str, float]: - """ - Compute basic metrics for predictions. - - Args: - y_true: True labels - y_pred: Predicted labels - - Returns: - Dictionary of metrics - """ - y_true = np.array(y_true) - y_pred = np.array(y_pred) - - metrics = {} - - # Basic accuracy for classification - if y_true.dtype == y_pred.dtype and np.issubdtype(y_true.dtype, np.integer): - metrics["accuracy"] = float(np.mean(y_true == y_pred)) - else: - # MSE for regression - metrics["mse"] = float(np.mean((y_true - y_pred) ** 2)) - metrics["mae"] = float(np.mean(np.abs(y_true - y_pred))) - - return metrics - - -def batch_iterator(iterable: Any, batch_size: int): - """ - Iterate over data in batches. - - Args: - iterable: Data to iterate over - batch_size: Size of each batch - - Yields: - Batches of data - """ - batch = [] - for item in iterable: - batch.append(item) - if len(batch) == batch_size: - yield batch - batch = [] - if batch: - yield batch From dfd7308567a416998929b465fa0d2677642b0892 Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 16:35:00 -0700 Subject: [PATCH 5/6] Fix flake8 findings (unused imports/variables, f-strings, import order) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes ~70 flake8 findings so Code Quality Checks passes fully: unused imports (F401), unused local variables (F841), f-strings without placeholders (F541), a redefined-but-unused pipeline import (F811) in the transformers integration, and two E402 import-order findings in tensorflow.py (noqa'd — the imports genuinely must follow the TORCH_AVAILABLE-style try/except guard). Co-Authored-By: Claude Sonnet 5 --- examples/boosting_example.py | 13 +++++-------- examples/decorator_monitoring.py | 4 +--- examples/langchain_example.py | 9 +-------- examples/offline_mode_example.py | 13 +++++-------- examples/sklearn_integration.py | 1 - examples/tensorflow_example.py | 2 -- tests/integration/test_crewai_integration.py | 2 +- tests/integration/test_langchain_integration.py | 2 +- tests/unit/test_client.py | 2 +- tests/unit/test_config_exceptions.py | 2 +- tests/unit/test_decorators.py | 10 +++++----- whiteboxxai/client.py | 2 +- whiteboxxai/git_utils.py | 1 - whiteboxxai/integrations/boosting.py | 2 +- whiteboxxai/integrations/crewai_monitor.py | 11 ++++------- whiteboxxai/integrations/langchain.py | 3 +-- whiteboxxai/integrations/langchain_agents.py | 2 -- whiteboxxai/integrations/pytorch.py | 4 +--- whiteboxxai/integrations/tensorflow.py | 6 +++--- whiteboxxai/integrations/transformers.py | 6 ++---- whiteboxxai/monitor.py | 1 - whiteboxxai/offline.py | 4 +--- whiteboxxai/privacy.py | 2 +- 23 files changed, 36 insertions(+), 68 deletions(-) diff --git a/examples/boosting_example.py b/examples/boosting_example.py index 27d4f7c..e8e499e 100644 --- a/examples/boosting_example.py +++ b/examples/boosting_example.py @@ -85,13 +85,13 @@ def example_xgboost_classification(): # Calculate metrics accuracy = accuracy_score(y_test, predictions) print(f"Accuracy: {accuracy:.4f}") - print(f"Predictions logged to WhiteBoxXAI") + print("Predictions logged to WhiteBoxXAI") # Get feature importance importance = model.feature_importances_ top_features = np.argsort(importance)[-5:][::-1] print(f"\nTop 5 features: {top_features.tolist()}") - print(f"Feature importance tracked in metadata") + print("Feature importance tracked in metadata") def example_xgboost_regression(): @@ -142,7 +142,7 @@ def example_xgboost_regression(): r2 = r2_score(y_test, predictions) print(f"MSE: {mse:.4f}") print(f"R² Score: {r2:.4f}") - print(f"Predictions automatically logged via wrapper") + print("Predictions automatically logged via wrapper") def example_lightgbm_classification(): @@ -199,13 +199,10 @@ def example_lightgbm_classification(): print("\nMaking predictions...") predictions = monitor.predict(model, X_test, y_test) - # Get probabilities - probabilities = model.predict_proba(X_test) - # Calculate metrics accuracy = accuracy_score(y_test, predictions) print(f"Accuracy: {accuracy:.4f}") - print(f"Predictions logged with probabilities") + print("Predictions logged with probabilities") # Feature importance importance = model.feature_importances_ @@ -261,7 +258,7 @@ def example_lightgbm_regression(): r2 = r2_score(y_test, predictions) print(f"MSE: {mse:.4f}") print(f"R² Score: {r2:.4f}") - print(f"Predictions automatically logged via wrapper") + print("Predictions automatically logged via wrapper") def example_feature_importance_tracking(): diff --git a/examples/decorator_monitoring.py b/examples/decorator_monitoring.py index 9ffa04d..80b87cc 100644 --- a/examples/decorator_monitoring.py +++ b/examples/decorator_monitoring.py @@ -17,9 +17,7 @@ @monitor_model(monitor, input_keys=["features"], explain=True) def predict_fraud(features): """Predict fraud probability.""" - # Simulate model prediction - model = RandomForestClassifier() - # ... (assume model is trained) + # Simulate model prediction (assume a RandomForestClassifier is trained) prediction = np.random.choice([0, 1]) probability = np.random.random() diff --git a/examples/langchain_example.py b/examples/langchain_example.py index a85162d..99070f6 100644 --- a/examples/langchain_example.py +++ b/examples/langchain_example.py @@ -151,7 +151,7 @@ def example_sequential_chain(): print("\nRunning sequential chain...") result = wrapped_chain({"subject": "artificial intelligence"}) - print(f"\n Subject: artificial intelligence") + print("\n Subject: artificial intelligence") print(f" Topic: {result['topic'].strip()}") print(f" Paragraph: {result['paragraph'].strip()[:100]}...") @@ -162,7 +162,6 @@ def example_agent(): """Example using an agent with tools.""" from langchain.agents import AgentType, Tool, initialize_agent from langchain.llms import OpenAI - from langchain.utilities import SerpAPIWrapper print("\n" + "=" * 60) print("Agent Example") @@ -228,12 +227,6 @@ def calculator_tool(expression: str) -> str: def example_rag_chain(): """Example using a RAG (Retrieval-Augmented Generation) chain.""" - from langchain.chains import RetrievalQA - from langchain.embeddings import OpenAIEmbeddings - from langchain.llms import OpenAI - from langchain.text_splitter import CharacterTextSplitter - from langchain.vectorstores import FAISS - print("\n" + "=" * 60) print("RAG Chain Example") print("=" * 60) diff --git a/examples/offline_mode_example.py b/examples/offline_mode_example.py index fdc1e8d..5262f07 100644 --- a/examples/offline_mode_example.py +++ b/examples/offline_mode_example.py @@ -6,10 +6,7 @@ """ import os -import time -from typing import Dict, List -import numpy as np from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split @@ -38,7 +35,7 @@ def example_1_basic_offline_mode(): # Check offline status status = client.get_offline_status() - print(f"\nOffline Status:") + print("\nOffline Status:") print(f" Queue size: {status['queue_size']}") print(f" Statistics: {status['statistics']}") @@ -95,7 +92,7 @@ def example_2_manual_sync(): # Manually trigger sync when connection is available print("\nTriggering manual sync...") result = client.sync_offline_queue(batch_size=50) - print(f"Sync result:") + print("Sync result:") print(f" Synced: {result['synced']}") print(f" Failed: {result['failed']}") print(f" Pending: {result['pending']}") @@ -178,7 +175,7 @@ def example_4_queue_management(): # Get queue statistics status = client.get_offline_status() - print(f"\nInitial Queue Status:") + print("\nInitial Queue Status:") print(f" Total: {status['statistics']['total']}") print(f" Pending: {status['statistics']['pending']}") print(f" Completed: {status['statistics']['completed']}") @@ -241,7 +238,7 @@ def example_5_ml_model_with_offline(): # Check queue status status = client.get_offline_status() - print(f"\nQueue Status:") + print("\nQueue Status:") print(f" Pending operations: {status['statistics']['pending']}") print(f" Completed: {status['statistics']['completed']}") @@ -317,7 +314,7 @@ def example_7_context_manager(): offline_auto_sync=True, offline_sync_interval=60, ) as client: - print(f"\n✓ Client initialized with offline mode") + print("\n✓ Client initialized with offline mode") print(f" Auto-sync running: {client._offline_manager._sync_running}") status = client.get_offline_status() diff --git a/examples/sklearn_integration.py b/examples/sklearn_integration.py index e61b138..d4b5773 100644 --- a/examples/sklearn_integration.py +++ b/examples/sklearn_integration.py @@ -4,7 +4,6 @@ This example demonstrates monitoring scikit-learn models. """ -import numpy as np from sklearn.datasets import make_classification from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import train_test_split diff --git a/examples/tensorflow_example.py b/examples/tensorflow_example.py index 11b08ee..ddda463 100644 --- a/examples/tensorflow_example.py +++ b/examples/tensorflow_example.py @@ -4,14 +4,12 @@ This example demonstrates how to use WhiteBoxXAI with TensorFlow/Keras models. """ -import numpy as np from sklearn.datasets import make_classification from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler # TensorFlow imports try: - import tensorflow as tf from tensorflow import keras except ImportError: print("TensorFlow not installed. Install with: pip install tensorflow") diff --git a/tests/integration/test_crewai_integration.py b/tests/integration/test_crewai_integration.py index 675c4ee..bd58d43 100644 --- a/tests/integration/test_crewai_integration.py +++ b/tests/integration/test_crewai_integration.py @@ -4,7 +4,7 @@ Tests CrewAI monitoring with mocked API calls. """ -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import Mock, patch from uuid import uuid4 import pytest diff --git a/tests/integration/test_langchain_integration.py b/tests/integration/test_langchain_integration.py index 2282fdd..6b4a942 100644 --- a/tests/integration/test_langchain_integration.py +++ b/tests/integration/test_langchain_integration.py @@ -1,7 +1,7 @@ """Unit tests for LangChain multi-agent integration.""" from datetime import datetime -from unittest.mock import MagicMock, Mock, call, patch +from unittest.mock import Mock import pytest diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index e374ebe..4bed52f 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -162,7 +162,7 @@ def test_request_with_retry(self, mock_request): mock_request.side_effect = [mock_response_error, mock_response_success] - client = WhiteBoxXAI(api_key="test_key", max_retries=2) + WhiteBoxXAI(api_key="test_key", max_retries=2) # Note: Retry behavior depends on implementation # This test verifies the interface exists diff --git a/tests/unit/test_config_exceptions.py b/tests/unit/test_config_exceptions.py index e7bbd10..3f6d87f 100644 --- a/tests/unit/test_config_exceptions.py +++ b/tests/unit/test_config_exceptions.py @@ -77,7 +77,7 @@ def test_config_base_url_validation(self): # Invalid URL (no protocol) - should handle gracefully or raise try: - config2 = Config(api_key="test-key", base_url="invalid-url") + Config(api_key="test-key", base_url="invalid-url") # If it accepts, it should add protocol or validate differently except ValueError: # Acceptable to reject invalid URLs diff --git a/tests/unit/test_decorators.py b/tests/unit/test_decorators.py index 9611047..31b3d85 100644 --- a/tests/unit/test_decorators.py +++ b/tests/unit/test_decorators.py @@ -5,7 +5,7 @@ """ import time -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, Mock import pytest @@ -42,7 +42,7 @@ def test_decorator_extracts_inputs(self): def predict(features): return {"prediction": sum(features)} - result = predict(features=[1, 2, 3]) + predict(features=[1, 2, 3]) # Verify inputs were extracted call_args = mock_monitor.log_prediction.call_args @@ -56,7 +56,7 @@ def test_decorator_extracts_output(self): def predict(x): return {"prediction": x * 2, "confidence": 0.95} - result = predict(x=5) + predict(x=5) # Verify output was extracted call_args = mock_monitor.log_prediction.call_args @@ -70,7 +70,7 @@ def test_decorator_with_explain(self): def predict(x): return x * 2 - result = predict(x=5) + predict(x=5) # Verify explain was passed call_args = mock_monitor.log_prediction.call_args @@ -85,7 +85,7 @@ def predict(x): time.sleep(0.1) # Simulate processing return x * 2 - result = predict(x=5) + predict(x=5) # Verify metadata includes inference time call_args = mock_monitor.log_prediction.call_args diff --git a/whiteboxxai/client.py b/whiteboxxai/client.py index 3af65a9..8fa2d99 100644 --- a/whiteboxxai/client.py +++ b/whiteboxxai/client.py @@ -6,7 +6,7 @@ import asyncio import logging -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Optional from urllib.parse import urljoin import httpx diff --git a/whiteboxxai/git_utils.py b/whiteboxxai/git_utils.py index 789e5b9..f996dce 100644 --- a/whiteboxxai/git_utils.py +++ b/whiteboxxai/git_utils.py @@ -9,7 +9,6 @@ import os import shutil import subprocess -from pathlib import Path from typing import Dict, Optional logger = logging.getLogger(__name__) diff --git a/whiteboxxai/integrations/boosting.py b/whiteboxxai/integrations/boosting.py index bfc5ae2..1362160 100644 --- a/whiteboxxai/integrations/boosting.py +++ b/whiteboxxai/integrations/boosting.py @@ -38,7 +38,7 @@ import logging import warnings -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional import numpy as np diff --git a/whiteboxxai/integrations/crewai_monitor.py b/whiteboxxai/integrations/crewai_monitor.py index 2872641..81bc78b 100644 --- a/whiteboxxai/integrations/crewai_monitor.py +++ b/whiteboxxai/integrations/crewai_monitor.py @@ -6,10 +6,7 @@ """ import logging -import time -import uuid -from datetime import datetime -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Optional from whiteboxxai import WhiteBoxXAI @@ -282,7 +279,7 @@ def log_agent_execution( try: agent_id = self.agent_map.get(id(agent)) if not agent_id: - logger.warning(f"Agent not registered, skipping execution log") + logger.warning("Agent not registered, skipping execution log") return execution_data = { @@ -324,7 +321,7 @@ def log_task_completion( try: task_id = self.task_map.get(id(task)) if not task_id: - logger.warning(f"Task not registered, skipping completion log") + logger.warning("Task not registered, skipping completion log") return update_data = { @@ -415,7 +412,7 @@ def complete_monitoring( "error_message": error_message, } - response = self.client.request( + self.client.request( "POST", f"/api/v1/workflows/multi-agent/{self.workflow_id}/complete", data=complete_data, diff --git a/whiteboxxai/integrations/langchain.py b/whiteboxxai/integrations/langchain.py index 9267a39..ae00904 100644 --- a/whiteboxxai/integrations/langchain.py +++ b/whiteboxxai/integrations/langchain.py @@ -6,10 +6,9 @@ from __future__ import annotations -import json import time import warnings -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional try: from langchain.agents import AgentExecutor diff --git a/whiteboxxai/integrations/langchain_agents.py b/whiteboxxai/integrations/langchain_agents.py index 6972c0c..a757701 100644 --- a/whiteboxxai/integrations/langchain_agents.py +++ b/whiteboxxai/integrations/langchain_agents.py @@ -13,8 +13,6 @@ from langchain.callbacks.base import BaseCallbackHandler from langchain.schema import AgentAction, AgentFinish, LLMResult -from langchain.schema.document import Document -from langchain.schema.output import ChatGeneration, Generation try: from whiteboxxai import WhiteBoxXAI diff --git a/whiteboxxai/integrations/pytorch.py b/whiteboxxai/integrations/pytorch.py index 83a53d8..a8b1600 100644 --- a/whiteboxxai/integrations/pytorch.py +++ b/whiteboxxai/integrations/pytorch.py @@ -6,7 +6,7 @@ from __future__ import annotations -from typing import Any, Callable, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional try: import torch @@ -25,8 +25,6 @@ class _NNStub: nn = _NNStub() -import numpy as np - from whiteboxxai.monitor import ModelMonitor diff --git a/whiteboxxai/integrations/tensorflow.py b/whiteboxxai/integrations/tensorflow.py index 5b159d8..127d40a 100644 --- a/whiteboxxai/integrations/tensorflow.py +++ b/whiteboxxai/integrations/tensorflow.py @@ -8,7 +8,7 @@ import logging import warnings -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, Optional, Union logger = logging.getLogger(__name__) @@ -30,9 +30,9 @@ class callbacks: keras = _KerasStub() -import numpy as np +import numpy as np # noqa: E402 -from whiteboxxai.monitor import ModelMonitor +from whiteboxxai.monitor import ModelMonitor # noqa: E402 class KerasMonitor(ModelMonitor): diff --git a/whiteboxxai/integrations/transformers.py b/whiteboxxai/integrations/transformers.py index 87fae0a..4e4db44 100644 --- a/whiteboxxai/integrations/transformers.py +++ b/whiteboxxai/integrations/transformers.py @@ -5,11 +5,11 @@ """ import warnings -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union try: import transformers - from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer, pipeline + from transformers import Pipeline, PreTrainedModel, PreTrainedTokenizer TRANSFORMERS_AVAILABLE = True except ImportError: @@ -18,8 +18,6 @@ PreTrainedTokenizer = object Pipeline = object -import numpy as np - from whiteboxxai.monitor import ModelMonitor diff --git a/whiteboxxai/monitor.py b/whiteboxxai/monitor.py index fa0c54d..ad0b2b5 100644 --- a/whiteboxxai/monitor.py +++ b/whiteboxxai/monitor.py @@ -4,7 +4,6 @@ Simplified monitoring interface for ML models. """ -import time from typing import TYPE_CHECKING, Any, Dict, List, Optional import numpy as np diff --git a/whiteboxxai/offline.py b/whiteboxxai/offline.py index b7a30fc..e2b11f6 100644 --- a/whiteboxxai/offline.py +++ b/whiteboxxai/offline.py @@ -34,11 +34,9 @@ import os import sqlite3 import threading -import time -from datetime import datetime from enum import Enum from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Tuple logger = logging.getLogger(__name__) diff --git a/whiteboxxai/privacy.py b/whiteboxxai/privacy.py index 4cb7990..86c97b8 100644 --- a/whiteboxxai/privacy.py +++ b/whiteboxxai/privacy.py @@ -5,7 +5,7 @@ """ import re -from typing import Any, Dict, List, Optional, Pattern, Union +from typing import Any, Dict, List, Optional, Pattern class PIIDetector: From 75e7431b20bfaa84e5e40b3f84c86ae59cf62b6e Mon Sep 17 00:00:00 2001 From: DeWitt Gibson Date: Tue, 14 Jul 2026 16:44:28 -0700 Subject: [PATCH 6/6] Fix broken relative links in SDK_MIGRATION_GUIDE.md for strict mkdocs build The guide's embedded example content linked to files from the main platform repo's structure (DEVELOPMENT.md, CONTRIBUTING.md, a proposed docs/ layout) that don't exist in this repo, failing `mkdocs build --strict`. De-linked them to plain text since they're illustrative references, not real navigable docs here. Co-Authored-By: Claude Sonnet 5 --- docs/SDK_MIGRATION_GUIDE.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/SDK_MIGRATION_GUIDE.md b/docs/SDK_MIGRATION_GUIDE.md index c7da9e0..e948f3e 100644 --- a/docs/SDK_MIGRATION_GUIDE.md +++ b/docs/SDK_MIGRATION_GUIDE.md @@ -662,13 +662,13 @@ Visit http://localhost:3000 to see real-time monitoring, drift detection, and ex ## For Developers -See [DEVELOPMENT.md](DEVELOPMENT.md) for platform development setup. +See `DEVELOPMENT.md` in the main platform repository for platform development setup. ## Documentation - **SDK Documentation**: https://github.com/AgentaFlow/whiteboxxai-python-sdk -- **Platform Documentation**: [docs/](docs/) -- **API Reference**: [API_REFERENCE.md](docs/API_REFERENCE.md) +- **Platform Documentation**: `docs/` in the main platform repository +- **API Reference**: `API_REFERENCE.md` in the main platform repository ## License @@ -939,10 +939,10 @@ pip install whiteboxxai[all] # All integrations ## Next Steps -- [Installation Guide](getting-started/installation.md) -- [Quick Start Tutorial](getting-started/quickstart.md) -- [Framework Integrations](integrations/sklearn.md) -- [API Reference](api/client.md) +- Installation Guide: `getting-started/installation.md` +- Quick Start Tutorial: `getting-started/quickstart.md` +- Framework Integrations: `integrations/sklearn.md` +- API Reference: `api/client.md` EOF ``` @@ -1121,7 +1121,7 @@ pytest tests/integration/ - **Discord**: https://discord.gg/whiteboxxai (if available) ### Contributing -See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution guidelines. +See `CONTRIBUTING.md` for contribution guidelines. ---