diff --git a/README.md b/README.md index 1024e10..9560275 100644 --- a/README.md +++ b/README.md @@ -357,3 +357,44 @@ python -m build Tests use deterministic service/model fixtures. They establish integration behavior, not detection effectiveness, customer savings, or production capacity. + + +### Runnable Ollama provider example + +[`examples/ollama_budget.py`](examples/ollama_budget.py) wraps a real +[`ollama.AsyncClient.generate`](https://github.com/ollama/ollama-python) call +with the existing published WebDecoy SDK. Install the example dependencies: + +```sh +python -m pip install -r examples/requirements-ollama.txt +ollama pull qwen2.5:0.5b +python examples/ollama_budget.py +``` + +Run Ollama locally first. Set `WEBDECOY_KEY`, `WEBDECOY_PROPERTY_ID` and a stable +server-only `WEBDECOY_SUBJECT_SECRET` (at least 32 characters). The demo uses +fixed local identities and observe mode. In an application, call `generate()` +after authentication, tenant authorization and request admission; supply only +server-derived account/organization context and pass the admission decision ID +for correlation. Handle budget denial before starting a response. Do not expose +the demo identities as public authentication. + +The example sends prompts only to loopback Ollama and accounting metadata to +WebDecoy. It uses one fixed local model with zero monetary rates, no fallback, +raw generation, a 1,024-byte input bound, a 2,048-token context setting and +128-token output setting. It reserves the full input context plus output allowance. +These depend on the provider honoring its settings; byte length is not an exact +token count. Overruns remain visible and are not a guaranteed billing ceiling. + +This example is non-streaming. It validates final provider/model/count fields; +missing usage retains the reservation, while explicit zero can settle it. Errors +and cancellation never retry inference or refund unknown work; settlement failure +preserves the returned text. Cancellation does not prove the server stopped work. +For streaming ownership, see the budget lifecycle guidance above. LangChain, +LangGraph and automatic HTTP streaming integration are follow-up work. + +CI uses the real pinned Ollama 0.6.3 Python SDK with an in-memory HTTP transport +and synthetic final responses. It verifies serialization, usage normalization, +denial before the provider, privacy, cancellation and accounting failures. It +does not execute a model or validate live provider billing or model quality. +Ollama is an example dependency, not a dependency of the WebDecoy package. diff --git a/examples/ollama_budget.py b/examples/ollama_budget.py new file mode 100644 index 0000000..bf0e7ed --- /dev/null +++ b/examples/ollama_budget.py @@ -0,0 +1,85 @@ +"""One non-streaming local Ollama attempt, after application auth/admission.""" + +import asyncio +import os + +from ollama import AsyncClient + +from webdecoy_ai_protection import ( + Budget, + BudgetCall, + BudgetCompletion, + BudgetLimits, + BudgetPrice, + BudgetSubject, + Client, + ollama_budget_usage, +) + +MODEL = "qwen2.5:0.5b" +CONTEXT_TOKENS = 2048 +OUTPUT_TOKENS = 128 + + +async def generate(protection, provider, trusted_context, prompt, request_id=None): + """Preserve output and accounting status; do not retry on accounting failure.""" + if not isinstance(prompt, str) or not 1 <= len(prompt.encode("utf-8")) <= 1024: + raise ValueError("prompt must contain 1–1024 UTF-8 bytes") + + async def attempt(runtime): + if runtime.provider != "ollama" or runtime.model != MODEL: + raise ValueError("example requires the configured local model") + # Raw generation excludes model templates/history/tools. Reserve the + # entire configured input context, not a character-based token estimate. + response = await provider.generate( + model=runtime.model, + prompt=prompt, + raw=True, + stream=False, + options={"num_ctx": runtime.max_input_tokens, "num_predict": runtime.max_output_tokens}, + ) + usage = ollama_budget_usage( + response.model, response.done, response.prompt_eval_count, response.eval_count + ) + return BudgetCompletion(response.response, usage) + + return await protection.run_budget( + trusted_context, + BudgetCall("local", CONTEXT_TOKENS, OUTPUT_TOKENS, request_id), + attempt, + ) + + +async def main(): + budget = Budget( + rule_id="ollama_example_v1", + subject_secret=os.environ["WEBDECOY_SUBJECT_SECRET"], + subject=lambda ctx: BudgetSubject(ctx["account_id"], ctx["organization_id"]), + window_seconds=3600, + limits=BudgetLimits(account_tokens=100_000), + prices={"local": BudgetPrice("ollama", MODEL, 0, 0)}, + mode="observe", + failure_mode="open", + ) + # Explicit loopback host and local model: no paid-provider fallback. + provider = AsyncClient(host="http://127.0.0.1:11434", timeout=60, trust_env=False) + async with Client( + api_key=os.environ["WEBDECOY_KEY"], + property_id=os.environ["WEBDECOY_PROPERTY_ID"], + budget=budget, + ) as protection: + result = await generate( + protection, + provider, + {"account_id": "local-demo", "organization_id": "local-demo"}, + "Say hello in one short sentence.", + ) + print(result.reason) + if result.allowed: + print(result.value) + else: + print(f"Denied: HTTP {result.status}; retry after {result.retry_after_seconds}s") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/requirements-ollama.txt b/examples/requirements-ollama.txt new file mode 100644 index 0000000..6a1da65 --- /dev/null +++ b/examples/requirements-ollama.txt @@ -0,0 +1,2 @@ +webdecoy-ai-protection==0.1.0a1 +ollama==0.6.3 diff --git a/pyproject.toml b/pyproject.toml index 01dc777..5e145bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ Issues = "https://github.com/WebDecoy/ai-protection-python/issues" [project.optional-dependencies] fastapi = ["fastapi>=0.115,<1"] -dev = ["fastapi>=0.115,<1", "uvicorn>=0.34,<1", "build>=1.2,<2", "ruff>=0.11,<1"] +dev = ["fastapi>=0.115,<1", "uvicorn>=0.34,<1", "build>=1.2,<2", "ruff>=0.11,<1", "ollama==0.6.3"] [tool.hatch.build.targets.wheel] packages = ["src/webdecoy_ai_protection"] diff --git a/tests/test_ollama_example.py b/tests/test_ollama_example.py new file mode 100644 index 0000000..e117ee9 --- /dev/null +++ b/tests/test_ollama_example.py @@ -0,0 +1,139 @@ +"""Exercise the real Ollama SDK HTTP serialization and response parsing offline.""" + +import asyncio +import json +import unittest + +import httpx +from ollama import AsyncClient, ResponseError +from test_budget import BudgetFixture, policy + +from examples.ollama_budget import MODEL, generate +from webdecoy_ai_protection import BudgetPrice + + +class OllamaExampleTests(unittest.IsolatedAsyncioTestCase): + async def run_example(self, payload=None, *, override=None, error=False, cancel=False): + fixture = BudgetFixture() + fixture.override.update(override or {}) + calls = [] + + async def handle(request): + calls.append(json.loads(request.content)) + self.assertEqual(request.url.path, "/api/generate") + self.assertEqual(len(fixture.budget_calls), 1) # reserve before inference + if cancel: + raise asyncio.CancelledError() + if error: + return httpx.Response(500, json={"error": "provider unavailable"}) + return httpx.Response( + 200, + json=payload + or { + "model": MODEL, + "done": True, + "response": "private answer", + "prompt_eval_count": 12, + "eval_count": 4, + }, + ) + + provider = AsyncClient(host="http://localhost:11434", transport=httpx.MockTransport(handle)) + config = policy(prices={"local": BudgetPrice("ollama", MODEL, 0, 0)}, mode="enforce") + async with fixture.client(budget=config) as client: + try: + result = await generate(client, provider, None, "private prompt") + except (ResponseError, asyncio.CancelledError) as exc: + result = exc + await client.flush() + return result, fixture, calls + + async def test_real_sdk_request_and_final_usage_privacy(self): + result, fixture, calls = await self.run_example() + self.assertEqual(result.value, "private answer") + self.assertEqual(result.reason, "budget_settled") + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["model"], MODEL) + self.assertTrue(calls[0]["raw"]) + self.assertFalse(calls[0]["stream"]) + self.assertEqual(calls[0]["options"], {"num_ctx": 2048, "num_predict": 128}) + self.assertEqual(fixture.budget_calls[0]["tokens"], 2176) + self.assertEqual(fixture.budget_calls[1]["tokens"], 16) + wire = json.dumps([fixture.budget_calls, fixture.usage]) + self.assertNotIn("private prompt", wire) + self.assertNotIn("private answer", wire) + + async def test_denial_never_calls_provider(self): + result, fixture, calls = await self.run_example( + override={ + "reserve": { + "schema": 1, + "allowed": False, + "granted": False, + "overrun": False, + "reason": "budget_exceeded", + "retry_after_seconds": 30, + } + } + ) + self.assertEqual(result.status, 429) + self.assertEqual(calls, []) + self.assertEqual(len(fixture.budget_calls), 1) + + async def test_unknown_or_wrong_model_retains_reservation(self): + for fields in ( + {}, + {"prompt_eval_count": 12}, + {"model": "other", "prompt_eval_count": 12, "eval_count": 4}, + {"done": False, "prompt_eval_count": 12, "eval_count": 4}, + ): + with self.subTest(fields=fields): + result, fixture, calls = await self.run_example( + { + "model": MODEL, + "done": True, + "response": "answer", + **fields, + } + ) + self.assertEqual(result.value, "answer") + self.assertEqual(result.reason, "budget_usage_unknown") + self.assertEqual(len(fixture.budget_calls), 1) + self.assertEqual(len(calls), 1) + + async def test_explicit_zero_settles(self): + result, fixture, _ = await self.run_example( + { + "model": MODEL, + "done": True, + "response": "", + "prompt_eval_count": 0, + "eval_count": 0, + } + ) + self.assertEqual(result.reason, "budget_settled") + self.assertEqual(fixture.budget_calls[-1]["tokens"], 0) + + async def test_accounting_outage_keeps_output_without_retry(self): + result, fixture, calls = await self.run_example(override={"settle": 503}) + self.assertEqual(result.value, "private answer") + self.assertEqual(result.reason, "budget_settlement_unavailable") + self.assertEqual(len(calls), 1) + self.assertEqual(len(fixture.budget_calls), 2) + + async def test_error_and_cancellation_never_settle_or_retry(self): + for options, exception in ( + ({"error": True}, ResponseError), + ({"cancel": True}, asyncio.CancelledError), + ): + result, fixture, calls = await self.run_example(**options) + self.assertIsInstance(result, exception) + self.assertEqual(len(calls), 1) + self.assertEqual(len(fixture.budget_calls), 1) + + async def test_input_bound_before_reservation(self): + fixture = BudgetFixture() + async with fixture.client(budget=policy()) as client: + with self.assertRaises(ValueError): + await generate(client, None, None, "a" * 1025) + self.assertEqual(fixture.budget_calls, [])