bridgekit 0.3.5__tar.gz → 0.3.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {bridgekit-0.3.5 → bridgekit-0.3.7}/PKG-INFO +84 -3
  2. {bridgekit-0.3.5 → bridgekit-0.3.7}/README.md +79 -2
  3. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit/__init__.py +1 -1
  4. bridgekit-0.3.7/bridgekit/cli.py +105 -0
  5. bridgekit-0.3.7/bridgekit/config.py +83 -0
  6. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit/planner.py +17 -17
  7. bridgekit-0.3.7/bridgekit/providers.py +112 -0
  8. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit/redteam.py +33 -30
  9. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit/reviewer.py +18 -17
  10. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit/search.py +30 -23
  11. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit.egg-info/PKG-INFO +84 -3
  12. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit.egg-info/SOURCES.txt +5 -0
  13. bridgekit-0.3.7/bridgekit.egg-info/entry_points.txt +2 -0
  14. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit.egg-info/requires.txt +6 -0
  15. {bridgekit-0.3.5 → bridgekit-0.3.7}/pyproject.toml +6 -1
  16. {bridgekit-0.3.5 → bridgekit-0.3.7}/tests/test_planner.py +14 -0
  17. bridgekit-0.3.7/tests/test_providers.py +71 -0
  18. bridgekit-0.3.7/tests/test_redteam.py +155 -0
  19. {bridgekit-0.3.5 → bridgekit-0.3.7}/tests/test_reviewer.py +18 -0
  20. {bridgekit-0.3.5 → bridgekit-0.3.7}/tests/test_search.py +21 -0
  21. bridgekit-0.3.5/bridgekit/config.py +0 -18
  22. {bridgekit-0.3.5 → bridgekit-0.3.7}/LICENSE +0 -0
  23. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit.egg-info/dependency_links.txt +0 -0
  24. {bridgekit-0.3.5 → bridgekit-0.3.7}/bridgekit.egg-info/top_level.txt +0 -0
  25. {bridgekit-0.3.5 → bridgekit-0.3.7}/setup.cfg +0 -0
  26. {bridgekit-0.3.5 → bridgekit-0.3.7}/tests/test_config.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bridgekit
3
- Version: 0.3.5
3
+ Version: 0.3.7
4
4
  Summary: AI tools that make you a better data scientist, not a redundant one.
5
5
  License: MIT
6
6
  Project-URL: Homepage, https://usebridgekit.com
@@ -21,6 +21,10 @@ Requires-Dist: pypdf>=3.0.0
21
21
  Requires-Dist: python-docx>=1.0.0
22
22
  Requires-Dist: python-pptx>=0.6.0
23
23
  Requires-Dist: nbformat>=5.0.0
24
+ Provides-Extra: openai
25
+ Requires-Dist: openai>=1.0.0; extra == "openai"
26
+ Provides-Extra: gemini
27
+ Requires-Dist: google-generativeai>=0.3.0; extra == "gemini"
24
28
  Provides-Extra: dev
25
29
  Requires-Dist: pytest>=7.0.0; extra == "dev"
26
30
  Requires-Dist: pytest-mock>=3.0.0; extra == "dev"
@@ -63,12 +67,26 @@ pip install bridgekit
63
67
  !pip install bridgekit
64
68
  ```
65
69
 
66
- Requires an Anthropic API key:
70
+ Requires an API key for your chosen provider:
67
71
 
72
+ **Anthropic (default):**
68
73
  ```bash
74
+ pip install bridgekit
69
75
  export ANTHROPIC_API_KEY=your_key_here
70
76
  ```
71
77
 
78
+ **OpenAI:**
79
+ ```bash
80
+ pip install bridgekit[openai]
81
+ export OPENAI_API_KEY=your_key_here
82
+ ```
83
+
84
+ **Google Gemini:**
85
+ ```bash
86
+ pip install bridgekit[gemini]
87
+ export GOOGLE_API_KEY=your_key_here
88
+ ```
89
+
72
90
  ---
73
91
 
74
92
  ## Getting Started
@@ -369,6 +387,69 @@ willing to commit to — and what's your confidence interval on that estimate?"
369
387
 
370
388
  ---
371
389
 
390
+ ## Multi-Provider Support
391
+
392
+ Bridgekit now supports multiple AI providers so you're not locked into one API. You can use Anthropic, OpenAI, or Google Gemini models with any tool.
393
+
394
+ **Using different providers:**
395
+
396
+ ```python
397
+ from bridgekit import evaluate, plan, ask, redteam
398
+
399
+ # Use OpenAI (default model: gpt-4o)
400
+ print(evaluate("Your analysis here", provider="openai"))
401
+
402
+ # Use Google Gemini (default model: gemini-1.5-pro)
403
+ print(plan("Your question here", provider="gemini"))
404
+
405
+ # Use specific model
406
+ print(redteam("Your analysis here", model="gpt-4-turbo"))
407
+ print(ask("Your question here", source="reports/", model="claude-3-opus-20240229"))
408
+ ```
409
+
410
+ **Provider auto-detection:**
411
+ Bridgekit automatically detects the provider from model names:
412
+ - Models starting with "claude" → Anthropic
413
+ - Models starting with "gpt" → OpenAI
414
+ - Models starting with "gemini" → Google Gemini
415
+
416
+ **Default models by provider:**
417
+ - Anthropic: `claude-3-5-sonnet-20241022`
418
+ - OpenAI: `gpt-4o`
419
+ - Gemini: `gemini-1.5-pro`
420
+
421
+ All tools support the same `provider` and `model` parameters:
422
+ - `evaluate(text, provider=None, model=None, system_prompt=None)`
423
+ - `plan(question, provider=None, model=None, ..., system_prompt=None)`
424
+ - `ask(question, provider=None, model=None, ..., system_prompt=None)`
425
+ - `redteam(text, provider=None, model=None, ..., system_prompt=None)`
426
+
427
+ ---
428
+
429
+ ## Custom System Prompts
430
+
431
+ Every tool accepts an optional `system_prompt` parameter to override the default persona. Use this to adapt the tone or focus to a specific domain without changing anything else.
432
+
433
+ ```python
434
+ from bridgekit import evaluate, plan, ask, redteam
435
+
436
+ # Narrow the reviewer to a specific domain
437
+ print(evaluate("my analysis", system_prompt="You are a skeptical PhD statistician focused only on methodology"))
438
+
439
+ # Tailor the planner to a specific industry
440
+ print(plan("my question", system_prompt="You are a data scientist specializing in healthcare analytics"))
441
+
442
+ # Replace the red team persona entirely
443
+ print(redteam("my analysis", system_prompt="You are a hostile regulator looking for compliance violations"))
444
+
445
+ # Change the answering style for ask
446
+ print(ask("my question", text="...", system_prompt="You are a financial analyst. Answer only in terms of revenue impact."))
447
+ ```
448
+
449
+ When `system_prompt` is not provided, each tool uses its built-in default — existing behavior is unchanged.
450
+
451
+ ---
452
+
372
453
  ## Why not just use Claude?
373
454
 
374
455
  You could. But you'd need to know what to ask, how to frame it, and what a good answer looks like. Bridgekit has that baked in — it knows you're a data scientist presenting findings, so it asks the right questions automatically. No prompt engineering required. Just paste your work and run it.
@@ -389,7 +470,7 @@ Bridgekit is a suite, not a one-off. Four tools are live — more are coming:
389
470
 
390
471
  - **Stakeholder translator** — turn your technical findings into a narrative a non-technical audience will actually follow
391
472
  - **Assumption checker** — state your analytical assumptions, get the ones you missed
392
- - **Multi-model support** — use any LLM provider (OpenAI, Gemini, open source models via OpenRouter) instead of being tied to Anthropic
473
+ - **More specialized tools** — focused on specific data science workflows and challenges
393
474
 
394
475
  Each tool is small, focused, and built for the way data scientists actually work.
395
476
 
@@ -35,12 +35,26 @@ pip install bridgekit
35
35
  !pip install bridgekit
36
36
  ```
37
37
 
38
- Requires an Anthropic API key:
38
+ Requires an API key for your chosen provider:
39
39
 
40
+ **Anthropic (default):**
40
41
  ```bash
42
+ pip install bridgekit
41
43
  export ANTHROPIC_API_KEY=your_key_here
42
44
  ```
43
45
 
46
+ **OpenAI:**
47
+ ```bash
48
+ pip install bridgekit[openai]
49
+ export OPENAI_API_KEY=your_key_here
50
+ ```
51
+
52
+ **Google Gemini:**
53
+ ```bash
54
+ pip install bridgekit[gemini]
55
+ export GOOGLE_API_KEY=your_key_here
56
+ ```
57
+
44
58
  ---
45
59
 
46
60
  ## Getting Started
@@ -341,6 +355,69 @@ willing to commit to — and what's your confidence interval on that estimate?"
341
355
 
342
356
  ---
343
357
 
358
+ ## Multi-Provider Support
359
+
360
+ Bridgekit now supports multiple AI providers so you're not locked into one API. You can use Anthropic, OpenAI, or Google Gemini models with any tool.
361
+
362
+ **Using different providers:**
363
+
364
+ ```python
365
+ from bridgekit import evaluate, plan, ask, redteam
366
+
367
+ # Use OpenAI (default model: gpt-4o)
368
+ print(evaluate("Your analysis here", provider="openai"))
369
+
370
+ # Use Google Gemini (default model: gemini-1.5-pro)
371
+ print(plan("Your question here", provider="gemini"))
372
+
373
+ # Use specific model
374
+ print(redteam("Your analysis here", model="gpt-4-turbo"))
375
+ print(ask("Your question here", source="reports/", model="claude-3-opus-20240229"))
376
+ ```
377
+
378
+ **Provider auto-detection:**
379
+ Bridgekit automatically detects the provider from model names:
380
+ - Models starting with "claude" → Anthropic
381
+ - Models starting with "gpt" → OpenAI
382
+ - Models starting with "gemini" → Google Gemini
383
+
384
+ **Default models by provider:**
385
+ - Anthropic: `claude-3-5-sonnet-20241022`
386
+ - OpenAI: `gpt-4o`
387
+ - Gemini: `gemini-1.5-pro`
388
+
389
+ All tools support the same `provider` and `model` parameters:
390
+ - `evaluate(text, provider=None, model=None, system_prompt=None)`
391
+ - `plan(question, provider=None, model=None, ..., system_prompt=None)`
392
+ - `ask(question, provider=None, model=None, ..., system_prompt=None)`
393
+ - `redteam(text, provider=None, model=None, ..., system_prompt=None)`
394
+
395
+ ---
396
+
397
+ ## Custom System Prompts
398
+
399
+ Every tool accepts an optional `system_prompt` parameter to override the default persona. Use this to adapt the tone or focus to a specific domain without changing anything else.
400
+
401
+ ```python
402
+ from bridgekit import evaluate, plan, ask, redteam
403
+
404
+ # Narrow the reviewer to a specific domain
405
+ print(evaluate("my analysis", system_prompt="You are a skeptical PhD statistician focused only on methodology"))
406
+
407
+ # Tailor the planner to a specific industry
408
+ print(plan("my question", system_prompt="You are a data scientist specializing in healthcare analytics"))
409
+
410
+ # Replace the red team persona entirely
411
+ print(redteam("my analysis", system_prompt="You are a hostile regulator looking for compliance violations"))
412
+
413
+ # Change the answering style for ask
414
+ print(ask("my question", text="...", system_prompt="You are a financial analyst. Answer only in terms of revenue impact."))
415
+ ```
416
+
417
+ When `system_prompt` is not provided, each tool uses its built-in default — existing behavior is unchanged.
418
+
419
+ ---
420
+
344
421
  ## Why not just use Claude?
345
422
 
346
423
  You could. But you'd need to know what to ask, how to frame it, and what a good answer looks like. Bridgekit has that baked in — it knows you're a data scientist presenting findings, so it asks the right questions automatically. No prompt engineering required. Just paste your work and run it.
@@ -361,7 +438,7 @@ Bridgekit is a suite, not a one-off. Four tools are live — more are coming:
361
438
 
362
439
  - **Stakeholder translator** — turn your technical findings into a narrative a non-technical audience will actually follow
363
440
  - **Assumption checker** — state your analytical assumptions, get the ones you missed
364
- - **Multi-model support** — use any LLM provider (OpenAI, Gemini, open source models via OpenRouter) instead of being tied to Anthropic
441
+ - **More specialized tools** — focused on specific data science workflows and challenges
365
442
 
366
443
  Each tool is small, focused, and built for the way data scientists actually work.
367
444
 
@@ -3,5 +3,5 @@ from .search import ask
3
3
  from .planner import plan
4
4
  from .redteam import redteam
5
5
 
6
- __version__ = "0.3.4"
6
+ __version__ = "0.3.7"
7
7
  __all__ = ["evaluate", "ask", "plan", "redteam"]
@@ -0,0 +1,105 @@
1
+ import argparse
2
+ import sys
3
+
4
+ from .planner import plan
5
+ from .reviewer import evaluate
6
+ from .redteam import redteam
7
+ from .search import ask
8
+
9
+
10
+ def _add_provider_args(parser: argparse.ArgumentParser) -> None:
11
+ parser.add_argument("--provider", help='AI provider: "anthropic", "openai", or "gemini"')
12
+ parser.add_argument("--model", help="Specific model to use (e.g. claude-opus-4-6, gpt-4o)")
13
+
14
+
15
+ def _cmd_plan(args: argparse.Namespace) -> None:
16
+ result = plan(
17
+ question=args.question,
18
+ data_description=args.data,
19
+ goal=args.goal,
20
+ provider=args.provider,
21
+ model=args.model,
22
+ )
23
+ print(result)
24
+
25
+
26
+ def _cmd_review(args: argparse.Namespace) -> None:
27
+ result = evaluate(
28
+ text=args.text,
29
+ provider=args.provider,
30
+ model=args.model,
31
+ )
32
+ print(result)
33
+
34
+
35
+ def _cmd_redteam(args: argparse.Namespace) -> None:
36
+ result = redteam(
37
+ text=args.text,
38
+ stakeholder=args.stakeholder,
39
+ provider=args.provider,
40
+ model=args.model,
41
+ )
42
+ print(result)
43
+
44
+
45
+ def _cmd_search(args: argparse.Namespace) -> None:
46
+ if not args.source and not args.text:
47
+ print("error: provide --source or --text", file=sys.stderr)
48
+ sys.exit(1)
49
+ result = ask(
50
+ question=args.question,
51
+ source=args.source,
52
+ text=args.text,
53
+ provider=args.provider,
54
+ model=args.model,
55
+ )
56
+ print(result)
57
+
58
+
59
+ def main() -> None:
60
+ parser = argparse.ArgumentParser(
61
+ prog="bridgekit",
62
+ description="AI tools for data scientists",
63
+ )
64
+ sub = parser.add_subparsers(dest="command", metavar="COMMAND")
65
+ sub.required = True
66
+
67
+ # plan
68
+ p_plan = sub.add_parser("plan", help="Recommend the right analytical approach")
69
+ p_plan.add_argument("question", help="The analytical question you want to answer")
70
+ p_plan.add_argument("--data", metavar="DESCRIPTION", help="Description of your available data")
71
+ p_plan.add_argument("--goal", help='Goal of the analysis (e.g. "prediction", "hypothesis testing")')
72
+ _add_provider_args(p_plan)
73
+ p_plan.set_defaults(func=_cmd_plan)
74
+
75
+ # review
76
+ p_review = sub.add_parser("review", help="Evaluate a data science analysis writeup")
77
+ p_review.add_argument("text", help="The analysis text to review")
78
+ _add_provider_args(p_review)
79
+ p_review.set_defaults(func=_cmd_review)
80
+
81
+ # redteam
82
+ p_redteam = sub.add_parser("redteam", help="Red-team an analysis from a skeptical stakeholder")
83
+ p_redteam.add_argument("text", help="The analysis text to red-team")
84
+ p_redteam.add_argument("--stakeholder", help='Stakeholder role (e.g. "VP of Finance")')
85
+ _add_provider_args(p_redteam)
86
+ p_redteam.set_defaults(func=_cmd_redteam)
87
+
88
+ # search
89
+ p_search = sub.add_parser("search", help="Ask a question across documents or text")
90
+ p_search.add_argument("question", help="The question to answer")
91
+ p_search.add_argument("--source", metavar="PATH", help="Folder of documents to search")
92
+ p_search.add_argument("--text", help="Raw text to search instead of a folder")
93
+ _add_provider_args(p_search)
94
+ p_search.set_defaults(func=_cmd_search)
95
+
96
+ args = parser.parse_args()
97
+ try:
98
+ args.func(args)
99
+ except (ValueError, EnvironmentError) as e:
100
+ print(f"error: {e}", file=sys.stderr)
101
+ sys.exit(1)
102
+
103
+
104
+ if __name__ == "__main__":
105
+ main()
@@ -0,0 +1,83 @@
1
+ import os
2
+ from enum import Enum
3
+ from typing import Optional
4
+
5
+
6
+ class Provider(Enum):
7
+ ANTHROPIC = "anthropic"
8
+ OPENAI = "openai"
9
+ GEMINI = "gemini"
10
+
11
+
12
+ # Default models for each provider
13
+ DEFAULT_MODELS = {
14
+ Provider.ANTHROPIC: "claude-opus-4-6",
15
+ Provider.OPENAI: "gpt-4o",
16
+ Provider.GEMINI: "gemini-1.5-pro"
17
+ }
18
+
19
+ # Legacy support
20
+ DEFAULT_MODEL = DEFAULT_MODELS[Provider.ANTHROPIC]
21
+
22
+
23
+ def require_api_key(provider: Provider = Provider.ANTHROPIC) -> str:
24
+ """Return the API key for the specified provider from the environment, or raise a clear error.
25
+
26
+ Each Bridgekit tool calls this before constructing a client so
27
+ users get the same friendly message instead of whatever the SDK surfaces
28
+ when the key is missing.
29
+ """
30
+ if provider == Provider.ANTHROPIC:
31
+ api_key = os.environ.get("ANTHROPIC_API_KEY")
32
+ if not api_key:
33
+ raise EnvironmentError(
34
+ "ANTHROPIC_API_KEY not found. Set it with: export ANTHROPIC_API_KEY=your_key_here"
35
+ )
36
+ return api_key
37
+ elif provider == Provider.OPENAI:
38
+ api_key = os.environ.get("OPENAI_API_KEY")
39
+ if not api_key:
40
+ raise EnvironmentError(
41
+ "OPENAI_API_KEY not found. Set it with: export OPENAI_API_KEY=your_key_here"
42
+ )
43
+ return api_key
44
+ elif provider == Provider.GEMINI:
45
+ api_key = os.environ.get("GOOGLE_API_KEY")
46
+ if not api_key:
47
+ raise EnvironmentError(
48
+ "GOOGLE_API_KEY not found. Set it with: export GOOGLE_API_KEY=your_key_here"
49
+ )
50
+ return api_key
51
+ else:
52
+ raise ValueError(f"Unsupported provider: {provider}")
53
+
54
+
55
+ def require_anthropic_api_key() -> str:
56
+ """Legacy function for backward compatibility."""
57
+ return require_api_key(Provider.ANTHROPIC)
58
+
59
+
60
+ def parse_provider(provider: Optional[str] = None, model: Optional[str] = None) -> Provider:
61
+ """Parse provider from provider string or model name."""
62
+ if provider:
63
+ try:
64
+ return Provider(provider.lower())
65
+ except ValueError:
66
+ raise ValueError(f"Unsupported provider: {provider}. Supported providers: {[p.value for p in Provider]}")
67
+
68
+ if model:
69
+ # Infer provider from model name
70
+ if model.startswith("claude"):
71
+ return Provider.ANTHROPIC
72
+ elif model.startswith("gpt"):
73
+ return Provider.OPENAI
74
+ elif model.startswith("gemini"):
75
+ return Provider.GEMINI
76
+
77
+ # Default to Anthropic for backward compatibility
78
+ return Provider.ANTHROPIC
79
+
80
+
81
+ def get_default_model(provider: Provider) -> str:
82
+ """Get the default model for a provider."""
83
+ return DEFAULT_MODELS.get(provider, DEFAULT_MODELS[Provider.ANTHROPIC])
@@ -1,5 +1,5 @@
1
- import anthropic
2
- from .config import DEFAULT_MODEL, require_anthropic_api_key
1
+ from .config import DEFAULT_MODEL, parse_provider, get_default_model
2
+ from .providers import create_message
3
3
 
4
4
  SYSTEM_PROMPT = """You are a senior statistician and data scientist advising a colleague on the right analytical approach for their problem.
5
5
 
@@ -29,7 +29,7 @@ ALTERNATIVES
29
29
  """
30
30
 
31
31
 
32
- def plan(question: str, data_description: str = None, goal: str = None) -> str:
32
+ def plan(question: str, data_description: str = None, goal: str = None, provider: str = None, model: str = None, system_prompt: str = None) -> str:
33
33
  """
34
34
  Recommend the right analytical approach for your problem.
35
35
 
@@ -38,6 +38,10 @@ def plan(question: str, data_description: str = None, goal: str = None) -> str:
38
38
  data_description: Optional. A plain text description of your available data.
39
39
  goal: Optional. The goal of your analysis (e.g. "causal inference",
40
40
  "prediction", "segmentation", "hypothesis testing", "exploration").
41
+ provider: Optional. The AI provider to use ("anthropic", "openai", "gemini").
42
+ If not specified, defaults to "anthropic" or infers from model.
43
+ model: Optional. The specific model to use. If not specified, uses the provider's default.
44
+ system_prompt: Optional. A custom system prompt to override the default planner persona.
41
45
 
42
46
  Returns:
43
47
  A structured analytical plan covering the recommended approach, assumptions,
@@ -46,7 +50,10 @@ def plan(question: str, data_description: str = None, goal: str = None) -> str:
46
50
  if not question or not question.strip():
47
51
  raise ValueError("Question cannot be empty.")
48
52
 
49
- api_key = require_anthropic_api_key()
53
+ # Parse provider and determine model
54
+ provider_enum = parse_provider(provider, model)
55
+ if model is None:
56
+ model = get_default_model(provider_enum)
50
57
 
51
58
  user_message = f"Question: {question}"
52
59
  if data_description:
@@ -54,17 +61,10 @@ def plan(question: str, data_description: str = None, goal: str = None) -> str:
54
61
  if goal:
55
62
  user_message += f"\n\nGoal: {goal}"
56
63
 
57
- client = anthropic.Anthropic(api_key=api_key)
58
- message = client.messages.create(
59
- model=DEFAULT_MODEL,
60
- max_tokens=1024,
61
- system=SYSTEM_PROMPT,
62
- messages=[
63
- {
64
- "role": "user",
65
- "content": user_message
66
- }
67
- ]
64
+ return create_message(
65
+ provider=provider_enum,
66
+ system_prompt=system_prompt or SYSTEM_PROMPT,
67
+ user_message=user_message,
68
+ model=model,
69
+ max_tokens=1024
68
70
  )
69
-
70
- return message.content[0].text
@@ -0,0 +1,112 @@
1
+ """Provider client factory for Bridgekit multi-provider support."""
2
+
3
+ from typing import Union, Dict, Any, List
4
+ from .config import Provider, require_api_key, get_default_model
5
+
6
+
7
+ class BaseProviderClient:
8
+ """Base class for provider clients."""
9
+
10
+ def __init__(self, api_key: str):
11
+ self.api_key = api_key
12
+
13
+ def create_message(self, system_prompt: str, user_message: str, model: str, max_tokens: int = 1024) -> str:
14
+ """Create a message with the given parameters."""
15
+ raise NotImplementedError("Subclasses must implement create_message")
16
+
17
+
18
+ class AnthropicClient(BaseProviderClient):
19
+ """Anthropic provider client."""
20
+
21
+ def __init__(self, api_key: str):
22
+ super().__init__(api_key)
23
+ import anthropic
24
+ self.client = anthropic.Anthropic(api_key=api_key)
25
+
26
+ def create_message(self, system_prompt: str, user_message: str, model: str, max_tokens: int = 1024) -> str:
27
+ """Create a message using Anthropic's API."""
28
+ message = self.client.messages.create(
29
+ model=model,
30
+ max_tokens=max_tokens,
31
+ system=system_prompt,
32
+ messages=[{
33
+ "role": "user",
34
+ "content": user_message
35
+ }]
36
+ )
37
+ return message.content[0].text
38
+
39
+
40
+ class OpenAIClient(BaseProviderClient):
41
+ """OpenAI provider client."""
42
+
43
+ def __init__(self, api_key: str):
44
+ super().__init__(api_key)
45
+ import openai
46
+ self.client = openai.OpenAI(api_key=api_key)
47
+
48
+ def create_message(self, system_prompt: str, user_message: str, model: str, max_tokens: int = 1024) -> str:
49
+ """Create a message using OpenAI's API."""
50
+ messages = [
51
+ {"role": "system", "content": system_prompt},
52
+ {"role": "user", "content": user_message}
53
+ ]
54
+
55
+ response = self.client.chat.completions.create(
56
+ model=model,
57
+ messages=messages,
58
+ max_tokens=max_tokens
59
+ )
60
+
61
+ return response.choices[0].message.content
62
+
63
+
64
+ class GeminiClient(BaseProviderClient):
65
+ """Google Gemini provider client."""
66
+
67
+ def __init__(self, api_key: str):
68
+ super().__init__(api_key)
69
+ import google.generativeai as genai
70
+ genai.configure(api_key=api_key)
71
+
72
+ def create_message(self, system_prompt: str, user_message: str, model: str, max_tokens: int = 1024) -> str:
73
+ """Create a message using Google's Gemini API."""
74
+ import google.generativeai as genai
75
+
76
+ # Configure the model with system instruction
77
+ model_instance = genai.GenerativeModel(
78
+ model_name=model,
79
+ system_instruction=system_prompt,
80
+ generation_config=genai.types.GenerationConfig(
81
+ max_output_tokens=max_tokens,
82
+ )
83
+ )
84
+
85
+ response = model_instance.generate_content(user_message)
86
+ return response.text
87
+
88
+
89
+ def create_client(provider: Provider, model: str = None) -> BaseProviderClient:
90
+ """Create a client for the specified provider."""
91
+ if model is None:
92
+ model = get_default_model(provider)
93
+
94
+ api_key = require_api_key(provider)
95
+
96
+ if provider == Provider.ANTHROPIC:
97
+ return AnthropicClient(api_key)
98
+ elif provider == Provider.OPENAI:
99
+ return OpenAIClient(api_key)
100
+ elif provider == Provider.GEMINI:
101
+ return GeminiClient(api_key)
102
+ else:
103
+ raise ValueError(f"Unsupported provider: {provider}")
104
+
105
+
106
+ def create_message(provider: Provider, system_prompt: str, user_message: str, model: str = None, max_tokens: int = 1024) -> str:
107
+ """Create a message using the specified provider."""
108
+ if model is None:
109
+ model = get_default_model(provider)
110
+
111
+ client = create_client(provider, model)
112
+ return client.create_message(system_prompt, user_message, model, max_tokens)
@@ -1,5 +1,5 @@
1
- import anthropic
2
- from .config import DEFAULT_MODEL, require_anthropic_api_key
1
+ from .config import DEFAULT_MODEL, parse_provider, get_default_model
2
+ from .providers import create_message
3
3
 
4
4
  DEFAULT_STAKEHOLDER = "a skeptical senior executive with no tolerance for weak methodology, unsupported claims, or vague business impact"
5
5
 
@@ -39,15 +39,20 @@ HARDEST QUESTION TO ANSWER
39
39
  """
40
40
 
41
41
 
42
- def redteam(text: str, stakeholder: str = None) -> str:
42
+ def redteam(text: str, stakeholder: str = None, provider: str = None, model: str = None, system_prompt: str = None) -> str:
43
43
  """
44
44
  Red-team a data science analysis writeup from the perspective of a skeptical stakeholder.
45
45
 
46
46
  Args:
47
- text: Your analysis writeup as a plain string.
48
- stakeholder: Optional. The skeptical stakeholder role (e.g. "VP of Finance",
49
- "skeptical board member", "Chief Revenue Officer").
50
- Defaults to a generic skeptical senior executive.
47
+ text: Your analysis writeup as a plain string.
48
+ stakeholder: Optional. The skeptical stakeholder role (e.g. "VP of Finance",
49
+ "skeptical board member", "Chief Revenue Officer").
50
+ Defaults to a generic skeptical senior executive.
51
+ provider: Optional. The AI provider to use ("anthropic", "openai", "gemini").
52
+ If not specified, defaults to "anthropic" or infers from model.
53
+ model: Optional. The specific model to use. If not specified, uses the provider's default.
54
+ system_prompt: Optional. A custom system prompt to fully override the default red team persona.
55
+ When provided, the stakeholder parameter is ignored.
51
56
 
52
57
  Returns:
53
58
  The 3-5 hardest critiques the stakeholder would make, plus the single
@@ -56,27 +61,25 @@ def redteam(text: str, stakeholder: str = None) -> str:
56
61
  if not text or not text.strip():
57
62
  raise ValueError("Text cannot be empty.")
58
63
 
59
- api_key = require_anthropic_api_key()
60
-
61
- stakeholder_label = stakeholder if stakeholder else "Skeptical Senior Executive"
62
- stakeholder_desc = stakeholder if stakeholder else DEFAULT_STAKEHOLDER
63
-
64
- system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
65
- stakeholder=stakeholder_desc,
66
- stakeholder_label=stakeholder_label
64
+ # Parse provider and determine model
65
+ provider_enum = parse_provider(provider, model)
66
+ if model is None:
67
+ model = get_default_model(provider_enum)
68
+
69
+ if system_prompt is None:
70
+ stakeholder_label = stakeholder if stakeholder else "Skeptical Senior Executive"
71
+ stakeholder_desc = stakeholder if stakeholder else DEFAULT_STAKEHOLDER
72
+ system_prompt = SYSTEM_PROMPT_TEMPLATE.format(
73
+ stakeholder=stakeholder_desc,
74
+ stakeholder_label=stakeholder_label
75
+ )
76
+
77
+ user_message = f"Red-team this analysis writeup:\n\n{text}"
78
+
79
+ return create_message(
80
+ provider=provider_enum,
81
+ system_prompt=system_prompt,
82
+ user_message=user_message,
83
+ model=model,
84
+ max_tokens=1024
67
85
  )
68
-
69
- client = anthropic.Anthropic(api_key=api_key)
70
- message = client.messages.create(
71
- model=DEFAULT_MODEL,
72
- max_tokens=1024,
73
- system=system_prompt,
74
- messages=[
75
- {
76
- "role": "user",
77
- "content": f"Red-team this analysis writeup:\n\n{text}"
78
- }
79
- ]
80
- )
81
-
82
- return message.content[0].text