bridgekit 0.3.7__tar.gz → 0.3.8__tar.gz

{bridgekit-0.3.7 → bridgekit-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bridgekit
-Version: 0.3.7
+Version: 0.3.8
 Summary: AI tools that make you a better data scientist, not a redundant one.
 License: MIT
 Project-URL: Homepage, https://usebridgekit.com
@@ -136,6 +136,9 @@ onboarding users to reporting as a growth lever.
 """
 
 print(evaluate(text))
+
+# Override for longer analyses
+print(evaluate(text, max_tokens=2048))
 ```
 
 **Output:**
@@ -190,6 +193,9 @@ Supports `.txt`, `.md`, `.pdf`, `.docx`, `.pptx`, and `.ipynb` files.
 from bridgekit import ask
 
 print(ask("what drove churn in Q3?", source="reports/"))
+
+# Override for longer responses
+print(ask("what drove churn in Q3?", source="reports/", max_tokens=2048))
 ```
 
 **From raw text:**
@@ -234,7 +240,7 @@ print(plan(
 ))
 ```
 
-`data_description` and `goal` are optional — the more context you provide, the more tailored the recommendation.
+`data_description`, `goal`, and `max_tokens` are optional — the more context you provide, the more tailored the recommendation.
 
 **`goal` examples:** `"causal inference"`, `"prediction"`, `"segmentation"`, `"hypothesis testing"`, `"exploration"`
 
@@ -297,6 +303,9 @@ print(redteam(text))
 # Or specify a stakeholder
 print(redteam(text, stakeholder="VP of Engineering"))
 print(redteam(text, stakeholder="VP of Marketing"))
+
+# Override for longer responses
+print(redteam(text, max_tokens=2048))
 ```
 
 Same writeup, different attack angles:
@@ -414,7 +423,7 @@ Bridgekit automatically detects the provider from model names:
 - Models starting with "gemini" → Google Gemini
 
 **Default models by provider:**
-- Anthropic: `claude-3-5-sonnet-20241022`
+- Anthropic: `claude-opus-4-8`
 - OpenAI: `gpt-4o`
 - Gemini: `gemini-1.5-pro`
 

{bridgekit-0.3.7 → bridgekit-0.3.8}/README.md

@@ -104,6 +104,9 @@ onboarding users to reporting as a growth lever.
 """
 
 print(evaluate(text))
+
+# Override for longer analyses
+print(evaluate(text, max_tokens=2048))
 ```
 
 **Output:**
@@ -158,6 +161,9 @@ Supports `.txt`, `.md`, `.pdf`, `.docx`, `.pptx`, and `.ipynb` files.
 from bridgekit import ask
 
 print(ask("what drove churn in Q3?", source="reports/"))
+
+# Override for longer responses
+print(ask("what drove churn in Q3?", source="reports/", max_tokens=2048))
 ```
 
 **From raw text:**
@@ -202,7 +208,7 @@ print(plan(
 ))
 ```
 
-`data_description` and `goal` are optional — the more context you provide, the more tailored the recommendation.
+`data_description`, `goal`, and `max_tokens` are optional — the more context you provide, the more tailored the recommendation.
 
 **`goal` examples:** `"causal inference"`, `"prediction"`, `"segmentation"`, `"hypothesis testing"`, `"exploration"`
 
@@ -265,6 +271,9 @@ print(redteam(text))
 # Or specify a stakeholder
 print(redteam(text, stakeholder="VP of Engineering"))
 print(redteam(text, stakeholder="VP of Marketing"))
+
+# Override for longer responses
+print(redteam(text, max_tokens=2048))
 ```
 
 Same writeup, different attack angles:
@@ -382,7 +391,7 @@ Bridgekit automatically detects the provider from model names:
 - Models starting with "gemini" → Google Gemini
 
 **Default models by provider:**
-- Anthropic: `claude-3-5-sonnet-20241022`
+- Anthropic: `claude-opus-4-8`
 - OpenAI: `gpt-4o`
 - Gemini: `gemini-1.5-pro`
 

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/__init__.py

@@ -3,5 +3,5 @@ from .search import ask
 from .planner import plan
 from .redteam import redteam
 
-__version__ = "0.3.7"
+__version__ = "0.3.8"
 __all__ = ["evaluate", "ask", "plan", "redteam"]

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/cli.py

@@ -9,7 +9,7 @@ from .search import ask
 
 def _add_provider_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--provider", help='AI provider: "anthropic", "openai", or "gemini"')
-    parser.add_argument("--model", help="Specific model to use (e.g. claude-opus-4-6, gpt-4o)")
+    parser.add_argument("--model", help="Specific model to use (e.g. claude-opus-4-8, gpt-4o)")
 
 
 def _cmd_plan(args: argparse.Namespace) -> None:

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/config.py

@@ -11,7 +11,7 @@ class Provider(Enum):
 
 # Default models for each provider
 DEFAULT_MODELS = {
-    Provider.ANTHROPIC: "claude-opus-4-6",
+    Provider.ANTHROPIC: "claude-opus-4-8",
     Provider.OPENAI: "gpt-4o",
     Provider.GEMINI: "gemini-1.5-pro"
 }

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/planner.py

@@ -29,7 +29,7 @@ ALTERNATIVES
 """
 
 
-def plan(question: str, data_description: str = None, goal: str = None, provider: str = None, model: str = None, system_prompt: str = None) -> str:
+def plan(question: str, data_description: str = None, goal: str = None, provider: str = None, model: str = None, system_prompt: str = None, max_tokens: int = 1024) -> str:
     """
     Recommend the right analytical approach for your problem.
 
@@ -42,6 +42,7 @@ def plan(question: str, data_description: str = None, goal: str = None, provider
                           If not specified, defaults to "anthropic" or infers from model.
         model:            Optional. The specific model to use. If not specified, uses the provider's default.
         system_prompt:    Optional. A custom system prompt to override the default planner persona.
+        max_tokens:       Optional. Maximum tokens in the response. Defaults to 1024.
 
     Returns:
         A structured analytical plan covering the recommended approach, assumptions,
@@ -66,5 +67,5 @@ def plan(question: str, data_description: str = None, goal: str = None, provider
         system_prompt=system_prompt or SYSTEM_PROMPT,
         user_message=user_message,
         model=model,
-        max_tokens=1024
+        max_tokens=max_tokens
     )

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/redteam.py

@@ -39,7 +39,7 @@ HARDEST QUESTION TO ANSWER
 """
 
 
-def redteam(text: str, stakeholder: str = None, provider: str = None, model: str = None, system_prompt: str = None) -> str:
+def redteam(text: str, stakeholder: str = None, provider: str = None, model: str = None, system_prompt: str = None, max_tokens: int = 1024) -> str:
     """
     Red-team a data science analysis writeup from the perspective of a skeptical stakeholder.
 
@@ -53,6 +53,7 @@ def redteam(text: str, stakeholder: str = None, provider: str = None, model: str
         model:         Optional. The specific model to use. If not specified, uses the provider's default.
         system_prompt: Optional. A custom system prompt to fully override the default red team persona.
                        When provided, the stakeholder parameter is ignored.
+        max_tokens:    Optional. Maximum tokens in the response. Defaults to 1024.
 
     Returns:
         The 3-5 hardest critiques the stakeholder would make, plus the single
@@ -81,5 +82,5 @@ def redteam(text: str, stakeholder: str = None, provider: str = None, model: str
         system_prompt=system_prompt,
         user_message=user_message,
         model=model,
-        max_tokens=1024
+        max_tokens=max_tokens
     )

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/reviewer.py

@@ -42,7 +42,7 @@ BOTTOM LINE
 [one sentence]
 """
 
-def evaluate(text: str, provider: str = None, model: str = None, system_prompt: str = None) -> str:
+def evaluate(text: str, provider: str = None, model: str = None, system_prompt: str = None, max_tokens: int = 1024) -> str:
     """
     Evaluate a data science analysis writeup and return structured feedback.
 
@@ -52,6 +52,7 @@ def evaluate(text: str, provider: str = None, model: str = None, system_prompt:
                  If not specified, defaults to "anthropic" or infers from model.
         model: Optional. The specific model to use. If not specified, uses the provider's default.
         system_prompt: Optional. A custom system prompt to override the default reviewer persona.
+        max_tokens: Optional. Maximum tokens in the response. Defaults to 1024.
 
     Returns:
         Structured feedback across four dimensions.
@@ -71,5 +72,5 @@ def evaluate(text: str, provider: str = None, model: str = None, system_prompt:
         system_prompt=system_prompt or SYSTEM_PROMPT,
         user_message=user_message,
         model=model,
-        max_tokens=1024
+        max_tokens=max_tokens
     )

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit/search.py

@@ -56,7 +56,7 @@ DEFAULT_SYSTEM_PROMPT = (
 )
 
 
-def ask(question: str, source: str = None, text: str = None, provider: str = None, model: str = None, system_prompt: str = None) -> str:
+def ask(question: str, source: str = None, text: str = None, provider: str = None, model: str = None, system_prompt: str = None, max_tokens: int = 1024) -> str:
     """
     Ask a question across a collection of analysis documents or raw text.
 
@@ -68,6 +68,7 @@ def ask(question: str, source: str = None, text: str = None, provider: str = Non
                        If not specified, defaults to "anthropic" or infers from model.
         model:         Optional. The specific model to use. If not specified, uses the provider's default.
         system_prompt: Optional. A custom system prompt to override the default answering persona.
+        max_tokens:    Optional. Maximum tokens in the response. Defaults to 1024.
 
     Returns:
         An answer grounded in the provided documents.
@@ -121,5 +122,5 @@ def ask(question: str, source: str = None, text: str = None, provider: str = Non
         system_prompt=system_prompt or DEFAULT_SYSTEM_PROMPT,
         user_message=user_message,
         model=model,
-        max_tokens=1024
+        max_tokens=max_tokens
     )

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bridgekit
-Version: 0.3.7
+Version: 0.3.8
 Summary: AI tools that make you a better data scientist, not a redundant one.
 License: MIT
 Project-URL: Homepage, https://usebridgekit.com
@@ -136,6 +136,9 @@ onboarding users to reporting as a growth lever.
 """
 
 print(evaluate(text))
+
+# Override for longer analyses
+print(evaluate(text, max_tokens=2048))
 ```
 
 **Output:**
@@ -190,6 +193,9 @@ Supports `.txt`, `.md`, `.pdf`, `.docx`, `.pptx`, and `.ipynb` files.
 from bridgekit import ask
 
 print(ask("what drove churn in Q3?", source="reports/"))
+
+# Override for longer responses
+print(ask("what drove churn in Q3?", source="reports/", max_tokens=2048))
 ```
 
 **From raw text:**
@@ -234,7 +240,7 @@ print(plan(
 ))
 ```
 
-`data_description` and `goal` are optional — the more context you provide, the more tailored the recommendation.
+`data_description`, `goal`, and `max_tokens` are optional — the more context you provide, the more tailored the recommendation.
 
 **`goal` examples:** `"causal inference"`, `"prediction"`, `"segmentation"`, `"hypothesis testing"`, `"exploration"`
 
@@ -297,6 +303,9 @@ print(redteam(text))
 # Or specify a stakeholder
 print(redteam(text, stakeholder="VP of Engineering"))
 print(redteam(text, stakeholder="VP of Marketing"))
+
+# Override for longer responses
+print(redteam(text, max_tokens=2048))
 ```
 
 Same writeup, different attack angles:
@@ -414,7 +423,7 @@ Bridgekit automatically detects the provider from model names:
 - Models starting with "gemini" → Google Gemini
 
 **Default models by provider:**
-- Anthropic: `claude-3-5-sonnet-20241022`
+- Anthropic: `claude-opus-4-8`
 - OpenAI: `gpt-4o`
 - Gemini: `gemini-1.5-pro`
 

{bridgekit-0.3.7 → bridgekit-0.3.8}/bridgekit.egg-info/SOURCES.txt

@@ -15,6 +15,7 @@ bridgekit.egg-info/dependency_links.txt
 bridgekit.egg-info/entry_points.txt
 bridgekit.egg-info/requires.txt
 bridgekit.egg-info/top_level.txt
+tests/test_cli.py
 tests/test_config.py
 tests/test_planner.py
 tests/test_providers.py

{bridgekit-0.3.7 → bridgekit-0.3.8}/pyproject.toml

@@ -7,7 +7,7 @@ include = ["bridgekit*"]
 
 [project]
 name = "bridgekit"
-version = "0.3.7"
+version = "0.3.8"
 description = "AI tools that make you a better data scientist, not a redundant one."
 readme = "README.md"
 requires-python = ">=3.9"

bridgekit-0.3.8/tests/test_cli.py

@@ -0,0 +1,175 @@
+import sys
+import pytest
+from unittest.mock import patch
+
+
+FAKE_PLAN = "BRIDGEKIT ANALYSIS PLAN\n─────\nRECOMMENDED APPROACH\nUse a t-test."
+FAKE_REVIEW = "BRIDGEKIT ANALYSIS REVIEW\n─────\n1. CLARITY\n✅ STRONG Clear writing."
+FAKE_REDTEAM = "BRIDGEKIT RED TEAM\n─────\nCRITIQUE 1: Sample Size\nHARDEST QUESTION TO ANSWER\nWhat is n?"
+FAKE_SEARCH = "Based on the documents, the answer is 42."
+
+
+class TestPlanCommand:
+    def test_basic_question(self, capsys):
+        with patch("bridgekit.cli.plan", return_value=FAKE_PLAN) as mock_plan:
+            with patch("sys.argv", ["bridgekit", "plan", "should I use a t-test?"]):
+                from bridgekit.cli import main
+                main()
+        mock_plan.assert_called_once_with(
+            question="should I use a t-test?",
+            data_description=None,
+            goal=None,
+            provider=None,
+            model=None,
+        )
+        assert FAKE_PLAN in capsys.readouterr().out
+
+    def test_with_data_and_goal(self, capsys):
+        with patch("bridgekit.cli.plan", return_value=FAKE_PLAN) as mock_plan:
+            with patch("sys.argv", ["bridgekit", "plan", "my question",
+                                    "--data", "50 rows", "--goal", "compare means"]):
+                from bridgekit.cli import main
+                main()
+        mock_plan.assert_called_once_with(
+            question="my question",
+            data_description="50 rows",
+            goal="compare means",
+            provider=None,
+            model=None,
+        )
+
+    def test_with_provider_and_model(self):
+        with patch("bridgekit.cli.plan", return_value=FAKE_PLAN) as mock_plan:
+            with patch("sys.argv", ["bridgekit", "plan", "my question",
+                                    "--provider", "openai", "--model", "gpt-4o"]):
+                from bridgekit.cli import main
+                main()
+        mock_plan.assert_called_once_with(
+            question="my question",
+            data_description=None,
+            goal=None,
+            provider="openai",
+            model="gpt-4o",
+        )
+
+    def test_missing_question_exits(self):
+        with patch("sys.argv", ["bridgekit", "plan"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit):
+                main()
+
+    def test_environment_error_exits(self, capsys):
+        with patch("bridgekit.cli.plan", side_effect=EnvironmentError("ANTHROPIC_API_KEY not found")):
+            with patch("sys.argv", ["bridgekit", "plan", "my question"]):
+                from bridgekit.cli import main
+                with pytest.raises(SystemExit) as exc:
+                    main()
+        assert exc.value.code == 1
+        assert "ANTHROPIC_API_KEY" in capsys.readouterr().err
+
+
+class TestReviewCommand:
+    def test_basic_text(self, capsys):
+        with patch("bridgekit.cli.evaluate", return_value=FAKE_REVIEW) as mock_evaluate:
+            with patch("sys.argv", ["bridgekit", "review", "my analysis text"]):
+                from bridgekit.cli import main
+                main()
+        mock_evaluate.assert_called_once_with(
+            text="my analysis text",
+            provider=None,
+            model=None,
+        )
+        assert FAKE_REVIEW in capsys.readouterr().out
+
+    def test_missing_text_exits(self):
+        with patch("sys.argv", ["bridgekit", "review"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit):
+                main()
+
+
+class TestRedteamCommand:
+    def test_basic_text(self, capsys):
+        with patch("bridgekit.cli.redteam", return_value=FAKE_REDTEAM) as mock_redteam:
+            with patch("sys.argv", ["bridgekit", "redteam", "my analysis text"]):
+                from bridgekit.cli import main
+                main()
+        mock_redteam.assert_called_once_with(
+            text="my analysis text",
+            stakeholder=None,
+            provider=None,
+            model=None,
+        )
+        assert FAKE_REDTEAM in capsys.readouterr().out
+
+    def test_with_stakeholder(self):
+        with patch("bridgekit.cli.redteam", return_value=FAKE_REDTEAM) as mock_redteam:
+            with patch("sys.argv", ["bridgekit", "redteam", "my analysis text",
+                                    "--stakeholder", "VP of Finance"]):
+                from bridgekit.cli import main
+                main()
+        mock_redteam.assert_called_once_with(
+            text="my analysis text",
+            stakeholder="VP of Finance",
+            provider=None,
+            model=None,
+        )
+
+    def test_missing_text_exits(self):
+        with patch("sys.argv", ["bridgekit", "redteam"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit):
+                main()
+
+
+class TestSearchCommand:
+    def test_with_source(self, capsys):
+        with patch("bridgekit.cli.ask", return_value=FAKE_SEARCH) as mock_ask:
+            with patch("sys.argv", ["bridgekit", "search", "my question",
+                                    "--source", "./my_docs"]):
+                from bridgekit.cli import main
+                main()
+        mock_ask.assert_called_once_with(
+            question="my question",
+            source="./my_docs",
+            text=None,
+            provider=None,
+            model=None,
+        )
+        assert FAKE_SEARCH in capsys.readouterr().out
+
+    def test_with_text(self):
+        with patch("bridgekit.cli.ask", return_value=FAKE_SEARCH) as mock_ask:
+            with patch("sys.argv", ["bridgekit", "search", "my question",
+                                    "--text", "some raw text"]):
+                from bridgekit.cli import main
+                main()
+        mock_ask.assert_called_once_with(
+            question="my question",
+            source=None,
+            text="some raw text",
+            provider=None,
+            model=None,
+        )
+
+    def test_missing_source_and_text_exits(self, capsys):
+        with patch("sys.argv", ["bridgekit", "search", "my question"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit) as exc:
+                main()
+        assert exc.value.code == 1
+        assert "error" in capsys.readouterr().err
+
+    def test_missing_question_exits(self):
+        with patch("sys.argv", ["bridgekit", "search"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit):
+                main()
+
+
+class TestNoCommand:
+    def test_no_subcommand_exits(self):
+        with patch("sys.argv", ["bridgekit"]):
+            from bridgekit.cli import main
+            with pytest.raises(SystemExit):
+                main()

{bridgekit-0.3.7 → bridgekit-0.3.8}/tests/test_planner.py

@@ -192,3 +192,33 @@ class TestPlanOptionalParameters:
                 content = str(messages_arg)
                 assert "5,000 users split 50/50." in content
                 assert "causal inference" in content
+
+
+class TestPlanMaxTokens:
+    """plan() should pass max_tokens through to the API."""
+
+    def test_default_max_tokens_is_1024(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.planner import plan
+                plan("Does our new onboarding flow increase upgrade rates?")
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 1024
+
+    def test_custom_max_tokens_reaches_api(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.planner import plan
+                plan("Does our new onboarding flow increase upgrade rates?", max_tokens=2048)
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 2048

{bridgekit-0.3.7 → bridgekit-0.3.8}/tests/test_redteam.py

@@ -153,3 +153,33 @@ class TestRedteamCustomSystemPrompt:
 
                 call_kwargs = mock_client.messages.create.call_args
                 assert call_kwargs.kwargs.get("system") == custom_prompt
+
+
+class TestRedteamMaxTokens:
+    """redteam() should pass max_tokens through to the API."""
+
+    def test_default_max_tokens_is_1024(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.redteam import redteam
+                redteam("Some analysis text.")
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 1024
+
+    def test_custom_max_tokens_reaches_api(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.redteam import redteam
+                redteam("Some analysis text.", max_tokens=2048)
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 2048

{bridgekit-0.3.7 → bridgekit-0.3.8}/tests/test_reviewer.py

@@ -176,3 +176,33 @@ class TestEvaluateCustomSystemPrompt:
 
                 call_kwargs = mock_client.messages.create.call_args
                 assert call_kwargs.kwargs.get("system") == custom_prompt
+
+
+class TestEvaluateMaxTokens:
+    """evaluate() should pass max_tokens through to the API."""
+
+    def test_default_max_tokens_is_1024(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.reviewer import evaluate
+                evaluate("Some analysis text.")
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 1024
+
+    def test_custom_max_tokens_reaches_api(self):
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic:
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_RESPONSE)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.reviewer import evaluate
+                evaluate("Some analysis text.", max_tokens=2048)
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 2048

{bridgekit-0.3.7 → bridgekit-0.3.8}/tests/test_search.py

@@ -257,3 +257,45 @@ class TestAskWithSourceFolder:
                     from bridgekit.search import ask
                     with pytest.raises(ValueError, match="No content found"):
                         ask("What happened?", source=tmpdir)
+
+
+class TestAskMaxTokens:
+    """ask() should pass max_tokens through to the API."""
+
+    def test_default_max_tokens_is_1024(self):
+        mock_chromadb, mock_ef = _make_mock_chromadb()
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic, \
+                 patch("chromadb.Client", mock_chromadb.Client), \
+                 patch(
+                     "chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction",
+                     mock_ef,
+                 ):
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_ANSWER)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.search import ask
+                ask("What was the conversion rate?", text="The conversion rate increased by 12%.")
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 1024
+
+    def test_custom_max_tokens_reaches_api(self):
+        mock_chromadb, mock_ef = _make_mock_chromadb()
+        with patch.dict(os.environ, {"ANTHROPIC_API_KEY": "test-key"}):
+            with patch("anthropic.Anthropic") as MockAnthropic, \
+                 patch("chromadb.Client", mock_chromadb.Client), \
+                 patch(
+                     "chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction",
+                     mock_ef,
+                 ):
+                mock_client = MagicMock()
+                mock_client.messages.create.return_value = _make_mock_message(FAKE_ANSWER)
+                MockAnthropic.return_value = mock_client
+
+                from bridgekit.search import ask
+                ask("What was the conversion rate?", text="The conversion rate increased by 12%.", max_tokens=2048)
+
+                call_kwargs = mock_client.messages.create.call_args
+                assert call_kwargs.kwargs.get("max_tokens") == 2048