From b05b41a5374a0f6a2acdf08dbc0b96a75fb365c6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Sat, 29 Aug 2026 23:35:34 +0400 Subject: [PATCH 1/4] Implement capability-oriented chat runtime --- .env.example | 8 +- .github/scripts/run-modal-migration.sh | 12 + .github/scripts/sync-modal-secret.sh | 4 +- .github/workflows/deploy.yml | 8 +- .github/workflows/pr-beta-deploy.yml | 10 +- AGENTS.md | 3 + Makefile | 29 +- backend/alembic.ini | 146 ++ backend/api/main.py | 9 +- backend/capabilities/__init__.py | 27 + backend/capabilities/application.py | 183 ++ backend/capabilities/artifacts.py | 163 ++ backend/capabilities/chart.py | 236 ++ backend/capabilities/compatibility.py | 50 + backend/capabilities/composition.py | 66 + backend/capabilities/context.py | 322 +++ backend/capabilities/contracts.py | 147 ++ backend/capabilities/executor.py | 391 +++ backend/capabilities/follow_up.py | 206 ++ backend/capabilities/household.py | 2113 +++++++++++++++++ backend/capabilities/household_input.py | 532 +++++ backend/capabilities/input_resolution.py | 41 + backend/capabilities/policy_information.py | 298 +++ backend/capabilities/policy_reform.py | 584 +++++ backend/capabilities/registry.py | 135 ++ backend/capabilities/relevance.py | 154 ++ backend/capabilities/repository.py | 104 + backend/capabilities/society.py | 441 ++++ backend/capabilities/tracing.py | 294 +++ backend/chat/__init__.py | 21 +- backend/chat/activity.py | 97 + backend/chat/artifact_context.py | 93 + backend/chat/capability_runtime.py | 21 + backend/chat/capability_service.py | 1328 +++++++++++ backend/chat/events.py | 50 +- backend/chat/model_port.py | 163 ++ backend/chat/model_selection.py | 199 -- backend/chat/narration.py | 120 + backend/chat/orchestrator.py | 731 ------ backend/chat/public_service.py | 60 +- backend/chat/routes.py | 2 +- backend/chat/schemas.py | 2 + backend/chat/system_blocks.py | 76 - backend/chat/turn_input.py | 4 + backend/config/__init__.py | 2 +- backend/config/models.py | 2 +- backend/config/sampling.py | 2 +- backend/conversation_context/__init__.py | 97 + .../conversation_context/change_pipeline.py | 1190 ++++++++++ .../conversation_context/engine_projection.py | 105 + .../conversation_context/household_view.py | 290 +++ backend/conversation_context/models.py | 600 +++++ backend/conversation_context/projection.py | 55 + backend/conversation_context/quantities.py | 177 ++ backend/conversation_context/reducer.py | 911 +++++++ backend/conversation_context/registry.py | 362 +++ backend/conversation_context/repository.py | 24 + backend/conversation_context/tools.py | 592 +++++ .../variable_resolution.py | 988 ++++++++ backend/conversations/__init__.py | 10 +- backend/conversations/models.py | 44 +- backend/conversations/store.py | 3 + backend/engine/discovery.py | 38 +- backend/eval/deployed_runner.py | 74 +- backend/eval/loaders.py | 2 - backend/eval/runner.py | 173 +- backend/eval/schemas.py | 100 +- backend/eval/service.py | 64 +- backend/gateway/__init__.py | 20 - backend/gateway/assessment.py | 453 ---- backend/gateway/catalogue.py | 234 -- backend/gateway/clarifications.py | 138 -- backend/gateway/execution.py | 168 -- backend/gateway/intent.py | 321 --- backend/gateway/policy.py | 381 --- backend/gateway/proposals.py | 431 ---- backend/gateway/runtime.py | 850 ------- backend/gateway/trace.py | 156 -- backend/migrations/README | 3 + backend/migrations/env.py | 92 + backend/migrations/script.py.mako | 28 + ...pre_branch_conversation_schema_baseline.py | 47 + .../0002_add_capability_persistence_tables.py | 120 + ...14_add_conversation_context_persistence.py | 42 + ...592837_add_invocation_debug_projections.py | 34 + backend/mypy.ini | 44 + backend/observability/fastapi.py | 3 - backend/observability/segments.py | 6 - backend/persistence/__init__.py | 4 + backend/persistence/capability_repository.py | 408 ++++ backend/persistence/context_repository.py | 101 + backend/persistence/deletion.py | 29 + backend/persistence/idempotency.py | 222 ++ backend/persistence/rows.py | 132 + backend/persistence/schema.py | 68 + backend/persistence/trace_repository.py | 114 + backend/prompts/__init__.py | 32 +- backend/prompts/gateway.py | 182 -- backend/prompts/system.py | 181 -- backend/requirements-test.txt | 1 + backend/requirements.txt | 5 +- backend/tests/test_agent_tools.py | 51 - backend/tests/test_anthropic_sdk_contract.py | 712 ++++++ backend/tests/test_api.py | 411 +--- backend/tests/test_capability_artifacts.py | 184 ++ backend/tests/test_capability_chat_service.py | 865 +++++++ backend/tests/test_capability_composition.py | 402 ++++ backend/tests/test_capability_persistence.py | 485 ++++ backend/tests/test_capability_public_api.py | 275 +++ backend/tests/test_chat_events.py | 33 +- backend/tests/test_chat_orchestrator.py | 462 ---- backend/tests/test_chat_public_service.py | 86 +- backend/tests/test_conversation_context.py | 1687 +++++++++++++ backend/tests/test_conversation_models.py | 89 +- backend/tests/test_database_migrations.py | 155 ++ backend/tests/test_deployed_evaluation.py | 145 +- backend/tests/test_discovery.py | 52 + backend/tests/test_eval_service.py | 148 +- backend/tests/test_evaluation.py | 122 +- backend/tests/test_fact_resolution.py | 1222 ++++++++++ backend/tests/test_gateway.py | 945 -------- backend/tests/test_gateway_assessment.py | 353 --- backend/tests/test_gateway_catalogue.py | 848 ------- backend/tests/test_gateway_clarifications.py | 153 -- backend/tests/test_gateway_execution_plan.py | 167 -- backend/tests/test_gateway_intent.py | 154 -- backend/tests/test_gateway_proposals.py | 238 -- backend/tests/test_gateway_trace.py | 70 - backend/tests/test_household_capability.py | 1445 +++++++++++ .../tests/test_household_chat_integration.py | 1746 ++++++++++++++ .../tests/test_invocation_observability.py | 310 +++ backend/tests/test_modal_app.py | 43 +- backend/tests/test_model_selection.py | 114 - backend/tests/test_policy_capabilities.py | 475 ++++ backend/tests/test_prompts.py | 176 +- backend/tests/test_society_capabilities.py | 509 ++++ backend/tests/test_typed_tools.py | 337 +++ backend/tools/analysis_support.py | 329 +++ backend/tools/context.py | 2 - backend/tools/contracts.py | 96 + backend/tools/definitions.py | 9 +- backend/tools/dispatch.py | 19 +- backend/tools/registry.py | 72 + backend/tools/typed_dispatch.py | 120 + backend/tools/typed_models.py | 199 ++ docker-compose.yml | 24 +- docs-site/_toc.yml | 4 - docs-site/architecture.md | 64 +- docs-site/backend/api.md | 13 +- docs-site/backend/billing.md | 8 +- docs-site/backend/chat.md | 91 +- docs-site/backend/gateway.md | 82 - docs-site/backend/overview.md | 30 +- docs-site/design/conversation-memory-layer.md | 298 --- docs-site/getting-started.md | 12 +- docs-site/intro.md | 21 +- docs/design/conversation-memory-layer.md | 298 --- .../engineering/capability-runtime-rollout.md | 103 + docs/engineering/skills/ai-evals.md | 21 +- .../engineering/skills/database-migrations.md | 80 + docs/engineering/skills/testing.md | 139 ++ docs/engineering/skills/uk-chat-runtime.md | 895 ++++--- evals/README.md | 21 +- evals/cases/answer/capability_runtime.yaml | 83 + evals/cases/answer/core.yaml | 751 ------ evals/cases/answer/live.yaml | 51 - evals/cases/gateway/plan.yaml | 166 -- .../policyengine_uk.generated.yaml | 52 +- evals/cases/tool_loop/capability_runtime.yaml | 299 +++ evals/cases/tool_loop/core.yaml | 365 --- evals/cases/tool_loop/uk_population_live.yaml | 7 - .../cases/trajectory/capability_runtime.yaml | 160 ++ evals/cases/trajectory/core.yaml | 971 -------- evals/cases/trajectory/live.yaml | 20 - .../capability_outputs/chart_completed.json | 53 + .../capability_outputs/follow_up_rerun.json | 37 + .../capability_outputs/follow_up_reused.json | 37 + .../household_completed.json | 107 + .../capability_outputs/society_completed.json | 91 + frontend/src/app/ChatPage.test.tsx | 115 +- frontend/src/app/ChatPage.tsx | 433 +--- .../src/app/api/proxy/backend-url.test.ts | 4 +- frontend/src/app/api/proxy/backend-url.ts | 2 +- frontend/src/components/DebugSetting.test.tsx | 41 + frontend/src/components/DebugSetting.tsx | 77 + .../components/InvocationActivity.test.tsx | 180 ++ .../src/components/InvocationActivity.tsx | 399 ++++ frontend/src/utils/useLocalStorage.test.ts | 73 + frontend/src/utils/useLocalStorage.ts | 78 + modal_app.py | 22 +- supabase/migrations/README.md | 15 + 191 files changed, 30740 insertions(+), 13306 deletions(-) create mode 100755 .github/scripts/run-modal-migration.sh create mode 100644 backend/alembic.ini create mode 100644 backend/capabilities/__init__.py create mode 100644 backend/capabilities/application.py create mode 100644 backend/capabilities/artifacts.py create mode 100644 backend/capabilities/chart.py create mode 100644 backend/capabilities/compatibility.py create mode 100644 backend/capabilities/composition.py create mode 100644 backend/capabilities/context.py create mode 100644 backend/capabilities/contracts.py create mode 100644 backend/capabilities/executor.py create mode 100644 backend/capabilities/follow_up.py create mode 100644 backend/capabilities/household.py create mode 100644 backend/capabilities/household_input.py create mode 100644 backend/capabilities/input_resolution.py create mode 100644 backend/capabilities/policy_information.py create mode 100644 backend/capabilities/policy_reform.py create mode 100644 backend/capabilities/registry.py create mode 100644 backend/capabilities/relevance.py create mode 100644 backend/capabilities/repository.py create mode 100644 backend/capabilities/society.py create mode 100644 backend/capabilities/tracing.py create mode 100644 backend/chat/activity.py create mode 100644 backend/chat/artifact_context.py create mode 100644 backend/chat/capability_runtime.py create mode 100644 backend/chat/capability_service.py create mode 100644 backend/chat/model_port.py delete mode 100644 backend/chat/model_selection.py create mode 100644 backend/chat/narration.py delete mode 100644 backend/chat/orchestrator.py delete mode 100644 backend/chat/system_blocks.py create mode 100644 backend/conversation_context/__init__.py create mode 100644 backend/conversation_context/change_pipeline.py create mode 100644 backend/conversation_context/engine_projection.py create mode 100644 backend/conversation_context/household_view.py create mode 100644 backend/conversation_context/models.py create mode 100644 backend/conversation_context/projection.py create mode 100644 backend/conversation_context/quantities.py create mode 100644 backend/conversation_context/reducer.py create mode 100644 backend/conversation_context/registry.py create mode 100644 backend/conversation_context/repository.py create mode 100644 backend/conversation_context/tools.py create mode 100644 backend/conversation_context/variable_resolution.py delete mode 100644 backend/gateway/__init__.py delete mode 100644 backend/gateway/assessment.py delete mode 100644 backend/gateway/catalogue.py delete mode 100644 backend/gateway/clarifications.py delete mode 100644 backend/gateway/execution.py delete mode 100644 backend/gateway/intent.py delete mode 100644 backend/gateway/policy.py delete mode 100644 backend/gateway/proposals.py delete mode 100644 backend/gateway/runtime.py delete mode 100644 backend/gateway/trace.py create mode 100644 backend/migrations/README create mode 100644 backend/migrations/env.py create mode 100644 backend/migrations/script.py.mako create mode 100644 backend/migrations/versions/0001_pre_branch_conversation_schema_baseline.py create mode 100644 backend/migrations/versions/0002_add_capability_persistence_tables.py create mode 100644 backend/migrations/versions/9526d8c80914_add_conversation_context_persistence.py create mode 100644 backend/migrations/versions/d97a20592837_add_invocation_debug_projections.py create mode 100644 backend/mypy.ini create mode 100644 backend/persistence/__init__.py create mode 100644 backend/persistence/capability_repository.py create mode 100644 backend/persistence/context_repository.py create mode 100644 backend/persistence/deletion.py create mode 100644 backend/persistence/idempotency.py create mode 100644 backend/persistence/rows.py create mode 100644 backend/persistence/schema.py create mode 100644 backend/persistence/trace_repository.py delete mode 100644 backend/prompts/gateway.py delete mode 100644 backend/prompts/system.py create mode 100644 backend/tests/test_anthropic_sdk_contract.py create mode 100644 backend/tests/test_capability_artifacts.py create mode 100644 backend/tests/test_capability_chat_service.py create mode 100644 backend/tests/test_capability_composition.py create mode 100644 backend/tests/test_capability_persistence.py create mode 100644 backend/tests/test_capability_public_api.py delete mode 100644 backend/tests/test_chat_orchestrator.py create mode 100644 backend/tests/test_conversation_context.py create mode 100644 backend/tests/test_database_migrations.py create mode 100644 backend/tests/test_fact_resolution.py delete mode 100644 backend/tests/test_gateway.py delete mode 100644 backend/tests/test_gateway_assessment.py delete mode 100644 backend/tests/test_gateway_catalogue.py delete mode 100644 backend/tests/test_gateway_clarifications.py delete mode 100644 backend/tests/test_gateway_execution_plan.py delete mode 100644 backend/tests/test_gateway_intent.py delete mode 100644 backend/tests/test_gateway_proposals.py delete mode 100644 backend/tests/test_gateway_trace.py create mode 100644 backend/tests/test_household_capability.py create mode 100644 backend/tests/test_household_chat_integration.py create mode 100644 backend/tests/test_invocation_observability.py delete mode 100644 backend/tests/test_model_selection.py create mode 100644 backend/tests/test_policy_capabilities.py create mode 100644 backend/tests/test_society_capabilities.py create mode 100644 backend/tests/test_typed_tools.py create mode 100644 backend/tools/analysis_support.py create mode 100644 backend/tools/contracts.py create mode 100644 backend/tools/typed_dispatch.py create mode 100644 backend/tools/typed_models.py delete mode 100644 docs-site/backend/gateway.md delete mode 100644 docs-site/design/conversation-memory-layer.md delete mode 100644 docs/design/conversation-memory-layer.md create mode 100644 docs/engineering/capability-runtime-rollout.md create mode 100644 docs/engineering/skills/database-migrations.md create mode 100644 evals/cases/answer/capability_runtime.yaml delete mode 100644 evals/cases/answer/core.yaml delete mode 100644 evals/cases/answer/live.yaml delete mode 100644 evals/cases/gateway/plan.yaml create mode 100644 evals/cases/tool_loop/capability_runtime.yaml delete mode 100644 evals/cases/tool_loop/core.yaml create mode 100644 evals/cases/trajectory/capability_runtime.yaml delete mode 100644 evals/cases/trajectory/core.yaml delete mode 100644 evals/cases/trajectory/live.yaml create mode 100644 evals/fixtures/tool_outputs/capability_outputs/chart_completed.json create mode 100644 evals/fixtures/tool_outputs/capability_outputs/follow_up_rerun.json create mode 100644 evals/fixtures/tool_outputs/capability_outputs/follow_up_reused.json create mode 100644 evals/fixtures/tool_outputs/capability_outputs/household_completed.json create mode 100644 evals/fixtures/tool_outputs/capability_outputs/society_completed.json create mode 100644 frontend/src/components/DebugSetting.test.tsx create mode 100644 frontend/src/components/DebugSetting.tsx create mode 100644 frontend/src/components/InvocationActivity.test.tsx create mode 100644 frontend/src/components/InvocationActivity.tsx create mode 100644 frontend/src/utils/useLocalStorage.test.ts create mode 100644 frontend/src/utils/useLocalStorage.ts create mode 100644 supabase/migrations/README.md diff --git a/.env.example b/.env.example index 233875bb..c8b2d891 100644 --- a/.env.example +++ b/.env.example @@ -2,11 +2,6 @@ ANTHROPIC_API_KEY=your_api_key_here ANTHROPIC_MODEL=claude-sonnet-4-6 -# Signs resumable gateway proposals carried in assistant history. Generate a -# stable random value of at least 32 bytes; this is integrity protection, not -# encryption. -GATEWAY_PROPOSAL_SIGNING_KEY=replace_with_a_stable_random_signing_key - # Token protecting the server's deployed evaluation route. Use a long random # value and configure the same value as EVAL_RUN_TOKEN when invoking evals. UK_CHAT_EVAL_TOKEN=replace_with_a_long_random_token @@ -39,6 +34,9 @@ BILLING_ENABLED=false DB_NAME=microsim DB_USERNAME=postgres DB_PASSWORD=postgres +# Host migration commands require an explicit PostgreSQL URL with schema-change +# permission. Docker Compose derives its migration URL from the values above. +# ALEMBIC_DATABASE_URL=postgresql://postgres:postgres@localhost:5433/microsim # Supabase (frontend auth) NEXT_PUBLIC_SUPABASE_URL=https://your-project.supabase.co diff --git a/.github/scripts/run-modal-migration.sh b/.github/scripts/run-modal-migration.sh new file mode 100755 index 00000000..5f79ca1b --- /dev/null +++ b/.github/scripts/run-modal-migration.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash + +set -euo pipefail + +if [[ -n "${MODAL_APP_NAME:-}" ]]; then + : "${MODAL_SECRET_NAME:?MODAL_SECRET_NAME is required with MODAL_APP_NAME}" + export POLICYENGINE_UK_CHAT_MODAL_APP_NAME="$MODAL_APP_NAME" + export POLICYENGINE_UK_CHAT_MODAL_SECRET_NAME="$MODAL_SECRET_NAME" +fi + +migration_run_name="${MODAL_APP_NAME:-policyengine-uk-chat}-migration" +modal run --name "$migration_run_name" modal_app.py::migrate diff --git a/.github/scripts/sync-modal-secret.sh b/.github/scripts/sync-modal-secret.sh index 30df5d64..ec64fe33 100755 --- a/.github/scripts/sync-modal-secret.sh +++ b/.github/scripts/sync-modal-secret.sh @@ -7,7 +7,7 @@ set -euo pipefail : "${HOSTNAMES:?HOSTNAMES is required}" : "${PUBLIC_BASE_URL:?PUBLIC_BASE_URL is required}" : "${UK_CHAT_EVAL_TOKEN:?UK_CHAT_EVAL_TOKEN is required}" -: "${GATEWAY_PROPOSAL_SIGNING_KEY:?GATEWAY_PROPOSAL_SIGNING_KEY is required}" +: "${ALEMBIC_DATABASE_URL:?ALEMBIC_DATABASE_URL is required}" billing_enabled="${BILLING_ENABLED:-false}" @@ -17,11 +17,11 @@ secret_args=( "ANTHROPIC_COMPLEX_MODEL=claude-sonnet-4-6" "ANTHROPIC_TITLE_MODEL=claude-haiku-4-5" "ANTHROPIC_DEFAULT_MODEL=claude-haiku-4-5" - "GATEWAY_PROPOSAL_SIGNING_KEY=$GATEWAY_PROPOSAL_SIGNING_KEY" "UK_CHAT_EVAL_TOKEN=$UK_CHAT_EVAL_TOKEN" "POLICYENGINE_UK_DATA_TOKEN=$POLICYENGINE_UK_DATA_TOKEN" "HUGGING_FACE_TOKEN=$HUGGING_FACE_TOKEN" "DATABASE_URL=$DATABASE_URL" + "ALEMBIC_DATABASE_URL=$ALEMBIC_DATABASE_URL" "BILLING_ENABLED=$billing_enabled" "OBSERVABILITY_ENVIRONMENT=$OBSERVABILITY_ENVIRONMENT" "OBSERVABILITY_GOOGLE_CLOUD_PROJECT=$OBSERVABILITY_GOOGLE_CLOUD_PROJECT" diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index cba639e7..3f1e6494 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -108,12 +108,12 @@ jobs: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - GATEWAY_PROPOSAL_SIGNING_KEY: ${{ secrets.GATEWAY_PROPOSAL_SIGNING_KEY }} UK_CHAT_EVAL_TOKEN: ${{ secrets.UK_CHAT_EVAL_TOKEN }} POLICYENGINE_UK_DATA_TOKEN: ${{ secrets.POLICYENGINE_UK_DATA_TOKEN }} # TEMPORARY: data access for the policyengine.py engine override. HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }} DATABASE_URL: ${{ secrets.POLICYENGINE_UK_CHAT_DATABASE_URL }} + ALEMBIC_DATABASE_URL: ${{ secrets.POLICYENGINE_UK_CHAT_DATABASE_URL }} BILLING_ENABLED: "false" SUPABASE_URL: ${{ secrets.SUPABASE_URL }} SUPABASE_SERVICE_ROLE_KEY: ${{ secrets.SUPABASE_SERVICE_ROLE_KEY }} @@ -130,6 +130,12 @@ jobs: PUBLIC_BASE_URL: https://policyengine-uk-chat.vercel.app/uk/chat run: .github/scripts/sync-modal-secret.sh + - name: Upgrade database schema + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + run: .github/scripts/run-modal-migration.sh + - name: Deploy to Modal env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} diff --git a/.github/workflows/pr-beta-deploy.yml b/.github/workflows/pr-beta-deploy.yml index 5b1d9eed..552899ad 100644 --- a/.github/workflows/pr-beta-deploy.yml +++ b/.github/workflows/pr-beta-deploy.yml @@ -64,11 +64,11 @@ jobs: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} - GATEWAY_PROPOSAL_SIGNING_KEY: ${{ secrets.GATEWAY_PROPOSAL_SIGNING_KEY }} UK_CHAT_EVAL_TOKEN: ${{ secrets.UK_CHAT_EVAL_TOKEN }} POLICYENGINE_UK_DATA_TOKEN: ${{ secrets.POLICYENGINE_UK_DATA_TOKEN }} HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }} DATABASE_URL: ${{ secrets.POLICYENGINE_UK_CHAT_DATABASE_URL }} + ALEMBIC_DATABASE_URL: ${{ secrets.POLICYENGINE_UK_CHAT_DATABASE_URL }} BILLING_ENABLED: "false" SUPABASE_URL: ${{ secrets.SUPABASE_URL }} SUPABASE_SERVICE_ROLE_KEY: ${{ secrets.SUPABASE_SERVICE_ROLE_KEY }} @@ -84,6 +84,14 @@ jobs: PUBLIC_BASE_URL: ${{ steps.names.outputs.frontend_url }}/uk/chat run: .github/scripts/sync-modal-secret.sh + - name: Upgrade database schema + env: + MODAL_APP_NAME: ${{ env.MODAL_PREVIEW_APP_NAME }} + MODAL_SECRET_NAME: ${{ env.MODAL_PREVIEW_SECRET_NAME }} + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + run: .github/scripts/run-modal-migration.sh + - name: Deploy backend preview to Modal id: modal_deploy env: diff --git a/AGENTS.md b/AGENTS.md index 304b1d54..034e7640 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,6 +15,9 @@ When adding, moving, or reviewing manual AI eval cases or harness code, read When changing the chat model pathway, system prompts, tool definitions, or calculation boundaries, read `docs/engineering/skills/uk-chat-runtime.md`. +When changing SQLModel schema, Alembic revisions, database adoption, or +migration behavior, read `docs/engineering/skills/database-migrations.md`. + Keep this file thin. Do not duplicate durable engineering guidance here; update the canonical docs first, then adjust this adapter only when an entry point needs to point at new guidance. diff --git a/Makefile b/Makefile index f3b5b3b3..b7d34ef9 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: up down build logs restart shell-backend shell-frontend test test-backend test-frontend sync-policyengine-uk-evals check-policyengine-uk-evals eval-ai-offline eval-ai-live eval-ai-live-uk-population eval-ai-deployed-uk-population +.PHONY: up down build logs restart shell-backend shell-frontend migrate migration-check migration-current test test-backend test-frontend typecheck-backend sync-policyengine-uk-evals check-policyengine-uk-evals eval-ai-offline eval-ai-live eval-ai-deployed-uk-population # Start all services in dev mode (live reload) up: @@ -39,6 +39,15 @@ shell-backend: shell-frontend: docker compose exec frontend sh +migrate: + python -m alembic -c backend/alembic.ini upgrade head + +migration-check: + python -m alembic -c backend/alembic.ini check + +migration-current: + python -m alembic -c backend/alembic.ini current + # One-time setup: copy .env.example to .env init: @if [ ! -f .env ]; then cp .env.example .env && echo "Created .env — fill in your ANTHROPIC_API_KEY"; else echo ".env already exists"; fi @@ -47,26 +56,26 @@ init: test: test-backend test-frontend test-backend: - PYTHONPATH=backend python -m pytest backend/tests --cov --cov-config=.coveragerc --cov-report=term-missing --cov-report=xml --cov-fail-under=80 + PYTHONPATH=backend:. python -m pytest backend/tests --cov --cov-config=.coveragerc --cov-report=term-missing --cov-report=xml --cov-fail-under=80 test-frontend: cd frontend && npm run test:coverage cd frontend && npm run build +typecheck-backend: + PYTHONPATH=backend:. python -m mypy --config-file backend/mypy.ini + sync-policyengine-uk-evals: - PYTHONPATH=backend python -m eval.sync_policyengine_uk --sync + PYTHONPATH=backend:. python -m eval.sync_policyengine_uk --sync check-policyengine-uk-evals: - PYTHONPATH=backend python -m eval.sync_policyengine_uk --check + PYTHONPATH=backend:. python -m eval.sync_policyengine_uk --check eval-ai-offline: check-policyengine-uk-evals - PYTHONPATH=backend python -m eval.run --mode offline + PYTHONPATH=backend:. python -m eval.run --mode offline eval-ai-live: check-policyengine-uk-evals - PYTHONPATH=backend python -m eval.run --mode live --provider anthropic - -eval-ai-live-uk-population: check-policyengine-uk-evals - RUN_DATA_EVALS=1 PYTHONPATH=backend python -m eval.run --suite tool_loop --mode live --provider anthropic + PYTHONPATH=backend:. python -m eval.run --mode live --provider anthropic eval-ai-deployed-uk-population: - PYTHONPATH=backend python -m eval.run_deployed --case-file evals/cases/tool_loop/uk_population_live.yaml + PYTHONPATH=backend:. python -m eval.run_deployed --case-file evals/cases/tool_loop/uk_population_live.yaml diff --git a/backend/alembic.ini b/backend/alembic.ini new file mode 100644 index 00000000..66e6e6cc --- /dev/null +++ b/backend/alembic.ini @@ -0,0 +1,146 @@ +# Alembic configuration for SQLModel-owned UK Chat tables. + +[alembic] +# path to migration scripts. +# this is typically a path given in POSIX (e.g. forward slashes) +# format, relative to the token %(here)s which refers to the location of this +# ini file +script_location = %(here)s/migrations + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s +# Or organize into date-based subdirectories (requires recursive_version_locations = true) +# file_template = %%(year)d/%%(month).2d/%%(day).2d_%%(hour).2d%%(minute).2d_%%(second).2d_%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. for multiple paths, the path separator +# is defined by "path_separator" below. +prepend_sys_path = %(here)s + + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the tzdata library which can be installed by adding +# `alembic[tz]` to the pip requirements. +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to /versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "path_separator" +# below. +# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions + +# path_separator; This indicates what character is used to split lists of file +# paths, including version_locations and prepend_sys_path within configparser +# files such as alembic.ini. +# The default rendered in new alembic.ini files is "os", which uses os.pathsep +# to provide os-dependent path splitting. +# +# Note that in order to support legacy alembic.ini files, this default does NOT +# take place if path_separator is not present in alembic.ini. If this +# option is omitted entirely, fallback logic is as follows: +# +# 1. Parsing of the version_locations option falls back to using the legacy +# "version_path_separator" key, which if absent then falls back to the legacy +# behavior of splitting on spaces and/or commas. +# 2. Parsing of the prepend_sys_path option falls back to the legacy +# behavior of splitting on spaces, commas, or colons. +# +# Valid values for path_separator are: +# +# path_separator = : +# path_separator = ; +# path_separator = space +# path_separator = newline +# +# Use os.pathsep. Default configuration used for new projects. +path_separator = os + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +# env.py requires ALEMBIC_DATABASE_URL and never reads a committed URL. + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module +# hooks = ruff +# ruff.type = module +# ruff.module = ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Alternatively, use the exec runner to execute a binary found on your PATH +# hooks = ruff +# ruff.type = exec +# ruff.executable = ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Logging configuration. This is also consumed by the user-maintained +# env.py script only. +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/backend/api/main.py b/backend/api/main.py index 610def4e..e05c37a9 100644 --- a/backend/api/main.py +++ b/backend/api/main.py @@ -13,7 +13,9 @@ import billing import chat -import conversations +from chat.activity import router as chat_activity_router +from conversations.routes import router as conversations_router +from persistence.schema import verify_database_schema from eval.routes import router as eval_router from api.errors import NaNSafeJSONResponse, rate_limit_handler from api.rate_limit import limiter @@ -38,7 +40,7 @@ @asynccontextmanager async def lifespan(_app: FastAPI): - conversations.ensure_table() + verify_database_schema() yield shutdown_observability() @@ -65,7 +67,8 @@ async def lifespan(_app: FastAPI): app.include_router(billing.router) app.include_router(chat.router) -app.include_router(conversations.router) +app.include_router(chat_activity_router) +app.include_router(conversations_router) app.include_router(eval_router) init_observability(app, service_role="api") diff --git a/backend/capabilities/__init__.py b/backend/capabilities/__init__.py new file mode 100644 index 00000000..881a9c8a --- /dev/null +++ b/backend/capabilities/__init__.py @@ -0,0 +1,27 @@ +"""Capability-oriented composition for the UK Chat runtime.""" + +from capabilities.contracts import ( + Accepted, + ArtifactContract, + Capability, + CapabilityDependency, + CapabilityOutcome, + CapabilitySpec, + Completed, + Failed, + NeedsInput, + Unsupported, +) + +__all__ = [ + "Accepted", + "ArtifactContract", + "Capability", + "CapabilityDependency", + "CapabilityOutcome", + "CapabilitySpec", + "Completed", + "Failed", + "NeedsInput", + "Unsupported", +] diff --git a/backend/capabilities/application.py b/backend/capabilities/application.py new file mode 100644 index 00000000..7dc0331d --- /dev/null +++ b/backend/capabilities/application.py @@ -0,0 +1,183 @@ +"""Concrete application composition for the capability-oriented chat runtime.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass, replace +from functools import lru_cache +from typing import Any +from uuid import uuid4 + +from capabilities.chart import SocietyChartCapability +from capabilities.composition import RuntimeComposition, compose_runtime +from capabilities.follow_up import AnalysisFollowUpCapability +from capabilities.household import ( + AssembleHouseholdCandidateTool, + HouseholdAnalysisCapability, + HouseholdAnalysisDraft, +) +from capabilities.policy_information import PolicyInformationCapability +from capabilities.policy_reform import ( + AnthropicReformCandidateResolver, + PolicyReformCapability, + PolicyReformInput, + ResolveReformTool, +) +from capabilities.relevance import ( + AnthropicRelevanceAssessor, + AssessRelevanceTool, + ConversationRelevanceCapability, +) +from capabilities.society import SocietyAnalysisCapability, SocietyAnalysisInput +from capabilities.tracing import InvocationTracer +from chat.artifact_context import RepositoryArtifactSummarySource +from chat.capability_service import ChatTurnService +from chat.events import CancellationProbe, ChatEvent +from chat.model_port import AnthropicConversationModel +from chat.turn_input import ChatTurnInput +from persistence.capability_repository import ( + PartialInputRegistry, + RepositoryArtifactAccess, + SQLConversationCapabilityRepository, +) +from persistence.context_repository import SQLConversationContextRepository +from persistence.idempotency import SQLIdempotencyRepository +from persistence.trace_repository import SQLInvocationTraceRepository +from tools.analysis_support import build_analysis_support_tools +from tools.context import TurnResultStore +from tools.typed_dispatch import build_dispatch_tools +from conversation_context.reducer import ContextReducer +from conversation_context.change_pipeline import ContextChangeApplier, ContextChangeValidator +from conversation_context.engine_projection import HouseholdEngineFactProjector +from conversation_context.registry import build_default_fact_registry +from conversation_context.tools import ( + AnthropicContextProposalReviewer, + AnthropicContextInterpreter, + ApplyContextChangeTool, + ProposeContextChangeTool, + ReduceContextPatchTool, + ValidateContextChangeTool, +) +from conversation_context.variable_resolution import ( + AnthropicVariableMapper, + ContextChangeResolver, + ResolveContextChangeTool, +) + + +@dataclass(frozen=True, slots=True) +class CapabilityChatApplication: + """Long-lived composition plus request-scoped execution construction.""" + + composition: RuntimeComposition + service: ChatTurnService + artifacts: RepositoryArtifactAccess + + async def run( + self, + turn: ChatTurnInput, + *, + is_cancelled: CancellationProbe, + ) -> AsyncIterator[ChatEvent]: + effective_turn = ( + turn if turn.turn_id else replace(turn, turn_id=uuid4().hex) + ) + context = self.composition.executor.context( + request_id=uuid4().hex, + conversation_id=effective_turn.session_id, + turn_id=effective_turn.turn_id, + is_cancelled=is_cancelled, + artifacts=self.artifacts, + result_store=TurnResultStore(), + ) + async for event in self.service.run( + effective_turn, + is_cancelled=is_cancelled, + context=context, + ): + yield event + + +def build_capability_chat_application(*, engine: Any | None = None) -> CapabilityChatApplication: + """Assemble every concrete object and validate the dependency graph once.""" + + partial_inputs = PartialInputRegistry() + partial_inputs.register( + "policy_reform", + schema_version="1", + model=PolicyReformInput, + ) + partial_inputs.register( + "household_analysis", + schema_version="1", + model=HouseholdAnalysisDraft, + ) + partial_inputs.register( + "society_analysis", + schema_version="1", + model=SocietyAnalysisInput, + ) + artifact_repository = SQLConversationCapabilityRepository( + engine=engine, + partial_inputs=partial_inputs, + ) + trace_repository = SQLInvocationTraceRepository(engine=engine) + tracer = InvocationTracer(sink=trace_repository) + fact_registry = build_default_fact_registry() + context_repository = SQLConversationContextRepository(engine=engine) + context_reducer = ContextReducer(fact_registry) + + tools = ( + *build_dispatch_tools(), + *build_analysis_support_tools(), + AssessRelevanceTool(AnthropicRelevanceAssessor()), + ProposeContextChangeTool(AnthropicContextInterpreter()), + ValidateContextChangeTool( + ContextChangeValidator(context_reducer, fact_registry), + AnthropicContextProposalReviewer(), + ), + ApplyContextChangeTool(ContextChangeApplier(context_repository)), + ReduceContextPatchTool(context_reducer), + ResolveContextChangeTool( + ContextChangeResolver( + fact_registry, + AnthropicVariableMapper(), + ) + ), + ResolveReformTool(AnthropicReformCandidateResolver()), + AssembleHouseholdCandidateTool(), + ) + capabilities = ( + ConversationRelevanceCapability(), + PolicyInformationCapability(), + PolicyReformCapability(), + HouseholdAnalysisCapability(HouseholdEngineFactProjector(fact_registry)), + SocietyAnalysisCapability(), + AnalysisFollowUpCapability(), + SocietyChartCapability(), + ) + composition = compose_runtime( + tools=tools, + capabilities=capabilities, + tracer=tracer, + ) + idempotency = SQLIdempotencyRepository(engine=engine) + service = ChatTurnService( + executor=composition.executor, + capabilities=composition.capabilities, + model=AnthropicConversationModel(), + idempotency=idempotency, + artifact_summaries=RepositoryArtifactSummarySource(artifact_repository), + context_repository=context_repository, + fact_registry=fact_registry, + ) + return CapabilityChatApplication( + composition=composition, + service=service, + artifacts=RepositoryArtifactAccess(artifact_repository), + ) + + +@lru_cache(maxsize=1) +def get_capability_chat_application() -> CapabilityChatApplication: + return build_capability_chat_application() diff --git a/backend/capabilities/artifacts.py b/backend/capabilities/artifacts.py new file mode 100644 index 00000000..8ef4e224 --- /dev/null +++ b/backend/capabilities/artifacts.py @@ -0,0 +1,163 @@ +"""Immutable typed values that capabilities may transfer across chat turns.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Literal, TypeAlias +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, Field + + +Scalar: TypeAlias = str | int | float | bool | None + + +def _artifact_id() -> str: + return uuid4().hex + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +class ImmutableModel(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class ArtifactProvenance(ImmutableModel): + conversation_id: str + turn_id: str + capability_id: str + capability_version: str + invocation_id: str + sources: tuple[str, ...] = () + + +class PolicyChange(ImmutableModel): + parameter_path: str + value: Scalar + effective_date: str | None = None + + +class HouseholdValue(ImmutableModel): + field_id: str + label: str + value: Scalar + source: Literal["user", "artifact", "default"] + subject_entity_id: str | None = None + engine_variable: str | None = None + period: Literal["annual", "monthly", "weekly", "four_weekly"] | None = None + + +class HouseholdEntityPosition(ImmutableModel): + entity_id: str + engine_position: str + + +class AggregateDimension(ImmutableModel): + name: str + value: str + + +class AggregateValue(ImmutableModel): + output_id: str + metric_id: str + label: str + value: float | int | None + unit: str + dimensions: tuple[AggregateDimension, ...] = () + + +class ArtifactBase(ImmutableModel): + artifact_id: str = Field(default_factory=_artifact_id) + schema_version: str = "1" + created_at: datetime = Field(default_factory=_now) + provenance: ArtifactProvenance + + +class PolicyScenarioRef(ArtifactBase): + artifact_type: Literal["policy_scenario"] = "policy_scenario" + year: int + scenario_revision: str + catalogue_version: str + calculation_engine_version: str + baseline: bool + verified_changes: tuple[PolicyChange, ...] = () + + +class HouseholdRef(ArtifactBase): + artifact_type: Literal["household"] = "household" + year: int + household_revision: str + catalogue_version: str + calculation_engine_version: str + values: tuple[HouseholdValue, ...] + context_scope_id: str | None = None + context_revision: int | None = None + entity_positions: tuple[HouseholdEntityPosition, ...] = () + + +class HouseholdResultRef(ArtifactBase): + artifact_type: Literal["household_result"] = "household_result" + year: int + household_artifact_id: str + policy_scenario_artifact_id: str + scenario_revision: str + calculation_engine_version: str + outputs: tuple[AggregateValue, ...] + context_scope_id: str | None = None + context_revision: int | None = None + + +class RequestedOutputIssue(ImmutableModel): + request: str + kind: Literal["ambiguous", "unsupported"] + guidance: str + + +class SocietyAnalysisResultRef(ArtifactBase): + artifact_type: Literal["society_analysis_result"] = "society_analysis_result" + year: int + policy_scenario_artifact_id: str + scenario_revision: str + catalogue_version: str + dataset_version: str + calculation_engine_version: str + default_profile_version: str + calculated_output_ids: tuple[str, ...] + outputs: tuple[AggregateValue, ...] + requested_output_issues: tuple[RequestedOutputIssue, ...] = () + + +class ChartPresentation(ImmutableModel): + chart_type: str + title: str + serialized_spec: str + + +class ChartArtifactRef(ArtifactBase): + artifact_type: Literal["chart"] = "chart" + source_result_artifact_id: str + source_result_schema_version: str + year: int + scenario_revision: str + calculation_engine_version: str + presentation: ChartPresentation + + +TransferableArtifact: TypeAlias = ( + PolicyScenarioRef + | HouseholdRef + | HouseholdResultRef + | SocietyAnalysisResultRef + | ChartArtifactRef +) + + +ARTIFACT_MODELS: dict[str, type[ArtifactBase]] = { + "policy_scenario": PolicyScenarioRef, + "household": HouseholdRef, + "household_result": HouseholdResultRef, + "society_analysis_result": SocietyAnalysisResultRef, + "chart": ChartArtifactRef, +} diff --git a/backend/capabilities/chart.py b/backend/capabilities/chart.py new file mode 100644 index 00000000..5cb01f1f --- /dev/null +++ b/backend/capabilities/chart.py @@ -0,0 +1,236 @@ +"""Deterministic chart creation from typed population result artifacts.""" + +from __future__ import annotations + +import json + +from pydantic import BaseModel, ConfigDict + +from capabilities.artifacts import ( + ArtifactProvenance, + ChartArtifactRef, + ChartPresentation, + PolicyScenarioRef, + SocietyAnalysisResultRef, +) +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilityDependency, + CapabilitySpec, + Completed, + NeedsInput, + Unsupported, +) +from capabilities.society import SocietyAnalysisInput, SocietyAnalysisOutput +from tools.contracts import CallerType, Visibility +from tools.typed_models import SafeToolOutput + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class SocietyChartInput(StrictModel): + referenced_result_id: str | None = None + requested_output: str | None = None + title: str | None = None + analysis: SocietyAnalysisInput | None = None + + +class SocietyChartOutput(StrictModel): + chart: ChartArtifactRef + source_result: SocietyAnalysisResultRef + + +class SocietyChartCapability(Capability[SocietyChartInput, SocietyChartOutput]): + spec = CapabilitySpec( + identifier="society_chart", + version="1", + description=( + "Create a deterministic chart from a compatible population result, " + "running population analysis first only when required inputs are complete." + ), + required_use="Use when the user asks to chart a population-analysis result.", + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=SocietyChartInput, + output_model=SocietyChartOutput, + accepted_artifacts=( + ArtifactContract( + artifact_type="society_analysis_result", + schema_version="1", + ), + ), + produced_artifacts=( + ArtifactContract(artifact_type="chart", schema_version="1"), + ), + dependencies=( + CapabilityDependency( + capability_id="society_analysis", + artifact=ArtifactContract( + artifact_type="society_analysis_result", + schema_version="1", + ), + ), + ), + tool_dependencies=("generate_chart",), + ) + + async def run(self, capability_input: SocietyChartInput, context): + results = list(await context.find_artifacts(SocietyAnalysisResultRef)) + if capability_input.referenced_result_id is not None: + results = [ + result + for result in results + if result.artifact_id == capability_input.referenced_result_id + ] + result = results[0] if len(results) == 1 else None + if result is None and capability_input.analysis is not None: + analysis_input = capability_input.analysis.model_copy( + update={ + "requested_outputs": tuple( + dict.fromkeys( + [ + *capability_input.analysis.requested_outputs, + *( + [capability_input.requested_output] + if capability_input.requested_output + else [] + ), + ] + ) + ) + } + ) + analysis = await context.invoke_capability( + "society_analysis", + analysis_input, + ) + if not isinstance(analysis, Completed) or not isinstance( + analysis.value, + SocietyAnalysisOutput, + ): + return analysis + result = analysis.value.result + if result is None: + return NeedsInput( + prompt=( + "Which population result should I chart, or what policy scenario " + "should I calculate first?" + ), + missing_fields=("referenced_result_id", "analysis"), + partial_input=capability_input.model_dump(mode="json", exclude_none=True), + ) + + output_id = self._output_id(capability_input, result) + if output_id not in result.calculated_output_ids: + scenarios = await context.find_artifacts(PolicyScenarioRef) + scenario = next( + ( + item + for item in scenarios + if item.artifact_id == result.policy_scenario_artifact_id + and item.scenario_revision == result.scenario_revision + ), + None, + ) + if scenario is None or capability_input.requested_output is None: + return Unsupported( + reason=( + "The retained result does not contain the requested chart metric " + "and its verified scenario is unavailable for rerun." + ) + ) + rerun = await context.invoke_capability( + "society_analysis", + { + "referenced_policy_scenario_id": scenario.artifact_id, + "year": result.year, + "requested_outputs": [capability_input.requested_output], + }, + ) + if not isinstance(rerun, Completed) or not isinstance( + rerun.value, + SocietyAnalysisOutput, + ): + return rerun + result = rerun.value.result + output_id = self._output_id(capability_input, result) + if output_id not in result.calculated_output_ids: + return Unsupported(reason="The requested chart metric is unsupported.") + selected = [output for output in result.outputs if output.output_id == output_id] + if not selected: + return Unsupported(reason="The requested chart metric has no retained values.") + rows = [ + { + "label": self._row_label(output), + "value": output.value, + } + for output in selected + if output.value is not None + ] + if not rows: + return Unsupported(reason="The requested chart metric has no numeric values.") + title = capability_input.title or output_id.replace("_", " ").title() + chart_result = await context.invoke_tool( + "generate_chart", + { + "chart_kind": "generic_bar", + "data": rows, + "title": title, + "x_field": "label", + "y_fields": ["value"], + "source": "PolicyEngine UK", + }, + ) + if not isinstance(chart_result, SafeToolOutput): + raise TypeError("Chart generation returned an incompatible output.") + spec = chart_result.root.get("spec") + markdown = chart_result.root.get("chart_markdown") + if not isinstance(spec, dict) or not isinstance(markdown, str): + return Unsupported(reason="Chart generation did not return a supported artifact.") + chart = ChartArtifactRef( + provenance=ArtifactProvenance( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + capability_id=self.spec.identifier, + capability_version=self.spec.version, + invocation_id=context.capability_invocation_id, + sources=(result.artifact_id, output_id), + ), + source_result_artifact_id=result.artifact_id, + source_result_schema_version=result.schema_version, + year=result.year, + scenario_revision=result.scenario_revision, + calculation_engine_version=result.calculation_engine_version, + presentation=ChartPresentation( + chart_type="generic_bar", + title=title, + serialized_spec=json.dumps(spec, sort_keys=True), + ), + ) + chart = await context.save_artifact(chart) + return Completed(value=SocietyChartOutput(chart=chart, source_result=result)) + + @staticmethod + def _output_id(capability_input, result): + if capability_input.requested_output: + normalized = capability_input.requested_output.casefold().replace(" ", "_") + aliases = { + "budget": "budgetary_impact", + "winners_and_losers": "winners_losers", + "deciles": "decile_impacts", + "poverty_rate": "poverty", + } + return aliases.get(normalized, normalized) + if "decile_impacts" in result.calculated_output_ids: + return "decile_impacts" + return result.calculated_output_ids[0] + + @staticmethod + def _row_label(output): + dimensions = ", ".join( + f"{dimension.name} {dimension.value}" for dimension in output.dimensions + ) + return dimensions or output.label diff --git a/backend/capabilities/compatibility.py b/backend/capabilities/compatibility.py new file mode 100644 index 00000000..821205d8 --- /dev/null +++ b/backend/capabilities/compatibility.py @@ -0,0 +1,50 @@ +"""Central compatibility checks for transferable capability artifacts.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict + +from capabilities.artifacts import ArtifactBase + + +class ArtifactRequirements(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + artifact_type: str + schema_version: str + year: int | None = None + scenario_revision: str | None = None + catalogue_version: str | None = None + dataset_version: str | None = None + calculation_engine_version: str | None = None + + +class ArtifactCompatibility(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + compatible: bool + issues: tuple[str, ...] + + +def check_artifact_compatibility( + artifact: ArtifactBase, + requirements: ArtifactRequirements, +) -> ArtifactCompatibility: + issues: list[str] = [] + fields = ( + "artifact_type", + "schema_version", + "year", + "scenario_revision", + "catalogue_version", + "dataset_version", + "calculation_engine_version", + ) + for field in fields: + required = getattr(requirements, field) + if required is None: + continue + actual = getattr(artifact, field, None) + if actual != required: + issues.append(f"{field}: expected {required!r}, got {actual!r}") + return ArtifactCompatibility(compatible=not issues, issues=tuple(issues)) diff --git a/backend/capabilities/composition.py b/backend/capabilities/composition.py new file mode 100644 index 00000000..ada55b8a --- /dev/null +++ b/backend/capabilities/composition.py @@ -0,0 +1,66 @@ +"""Explicit startup composition for typed tools and capabilities.""" + +from __future__ import annotations + +from dataclasses import dataclass +from collections.abc import Iterable +from typing import Any + +from capabilities.contracts import Capability +from capabilities.executor import InvocationExecutor +from capabilities.registry import CapabilityRegistry +from capabilities.tracing import InvocationTracer +from tools.contracts import Tool +from tools.registry import ToolRegistry + + +@dataclass(frozen=True, slots=True) +class RuntimeComposition: + tools: ToolRegistry + capabilities: CapabilityRegistry + tracer: InvocationTracer + executor: InvocationExecutor + + +def compose_runtime( + *, + tools: Iterable[Tool[Any, Any]], + capabilities: Iterable[Capability[Any, Any]], + tracer: InvocationTracer | None = None, +) -> RuntimeComposition: + tool_registry = ToolRegistry() + for tool in tools: + tool_registry.register(tool) + tool_registry.validate() + + capability_registry = CapabilityRegistry() + for capability in capabilities: + capability_registry.register(capability) + capability_registry.validate() + + registered_tool_ids = {spec.identifier for spec in tool_registry.specs()} + for tool_spec in tool_registry.specs(): + missing = set(tool_spec.tool_dependencies) - registered_tool_ids + if missing: + raise ValueError( + f"Tool {tool_spec.identifier} requires unknown tools: {sorted(missing)}." + ) + for capability_spec in capability_registry.specs(): + missing = set(capability_spec.tool_dependencies) - registered_tool_ids + if missing: + raise ValueError( + f"Capability {capability_spec.identifier} requires unknown tools: {sorted(missing)}." + ) + + invocation_tracer = tracer or InvocationTracer() + executor = InvocationExecutor( + tools=tool_registry, + capabilities=capability_registry, + tracer=invocation_tracer, + ) + return RuntimeComposition( + tools=tool_registry, + capabilities=capability_registry, + tracer=invocation_tracer, + executor=executor, + ) diff --git a/backend/capabilities/context.py b/backend/capabilities/context.py new file mode 100644 index 00000000..86ebf4ea --- /dev/null +++ b/backend/capabilities/context.py @@ -0,0 +1,322 @@ +"""Request-scoped context supplied to capability implementations.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field, replace +from threading import Lock +from typing import TYPE_CHECKING, Protocol, TypeVar + +from pydantic import BaseModel + +from capabilities.artifacts import ArtifactBase +from tools.contracts import CallerType +from tools.context import TurnResultStore + + +if TYPE_CHECKING: + from capabilities.contracts import CapabilityOutcome + from capabilities.executor import InvocationExecutor + from capabilities.repository import WaitingCapabilityInvocation + from conversation_context.models import ConversationContext + + +ArtifactT = TypeVar("ArtifactT", bound=ArtifactBase) +CancellationProbe = Callable[[], Awaitable[bool]] + + +@dataclass(frozen=True, slots=True) +class ModelUsageSnapshot: + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + def since(self, earlier: "ModelUsageSnapshot") -> "ModelUsageSnapshot": + return ModelUsageSnapshot( + input_tokens=self.input_tokens - earlier.input_tokens, + output_tokens=self.output_tokens - earlier.output_tokens, + cache_creation_input_tokens=( + self.cache_creation_input_tokens + - earlier.cache_creation_input_tokens + ), + cache_read_input_tokens=( + self.cache_read_input_tokens - earlier.cache_read_input_tokens + ), + ) + + +@dataclass(slots=True) +class ModelUsageLedger: + _input_tokens: int = 0 + _output_tokens: int = 0 + _cache_creation_input_tokens: int = 0 + _cache_read_input_tokens: int = 0 + _lock: Lock = field(default_factory=Lock) + + def record( + self, + *, + input_tokens: int = 0, + output_tokens: int = 0, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ) -> None: + with self._lock: + self._input_tokens += input_tokens + self._output_tokens += output_tokens + self._cache_creation_input_tokens += cache_creation_input_tokens + self._cache_read_input_tokens += cache_read_input_tokens + + def snapshot(self) -> ModelUsageSnapshot: + with self._lock: + return ModelUsageSnapshot( + input_tokens=self._input_tokens, + output_tokens=self._output_tokens, + cache_creation_input_tokens=self._cache_creation_input_tokens, + cache_read_input_tokens=self._cache_read_input_tokens, + ) + + +class ArtifactAccess(Protocol): + async def find_artifacts( + self, + *, + conversation_id: str, + artifact_model: type[ArtifactT], + ) -> tuple[ArtifactT, ...]: ... + + async def save_artifact( + self, + *, + conversation_id: str, + artifact: ArtifactT, + ) -> ArtifactT: ... + + async def save_waiting(self, invocation: object) -> object: ... + + async def list_waiting( + self, + *, + conversation_id: str, + capability_id: str, + ) -> tuple["WaitingCapabilityInvocation", ...]: ... + + async def update_waiting( + self, + *, + invocation_id: str, + partial_input: BaseModel, + ) -> "WaitingCapabilityInvocation": ... + + async def remove_waiting(self, *, invocation_id: str) -> None: ... + + +class EmptyArtifactAccess: + async def find_artifacts( + self, + *, + conversation_id: str, + artifact_model: type[ArtifactT], + ) -> tuple[ArtifactT, ...]: + del conversation_id, artifact_model + return () + + async def save_artifact( + self, + *, + conversation_id: str, + artifact: ArtifactT, + ) -> ArtifactT: + del conversation_id + return artifact + + async def save_waiting(self, invocation: object) -> object: + return invocation + + async def list_waiting( + self, + *, + conversation_id: str, + capability_id: str, + ) -> tuple["WaitingCapabilityInvocation", ...]: + del conversation_id, capability_id + return () + + async def update_waiting( + self, + *, + invocation_id: str, + partial_input: BaseModel, + ) -> "WaitingCapabilityInvocation": + del invocation_id, partial_input + raise KeyError("No waiting capability invocation is available.") + + async def remove_waiting(self, *, invocation_id: str) -> None: + del invocation_id + + +@dataclass(frozen=True, slots=True) +class CapabilityContext: + request_id: str + conversation_id: str + turn_id: str + is_cancelled: CancellationProbe + artifacts: ArtifactAccess + result_store: TurnResultStore + model_usage: ModelUsageLedger + _executor: "InvocationExecutor" + current_user_message: str = "" + conversation_context: "ConversationContext | None" = None + _capability_id: str | None = None + _capability_invocation_id: str | None = None + _capability_version: str | None = None + _tool_id: str | None = None + + def for_capability( + self, + capability_id: str, + invocation_id: str | None = None, + capability_version: str | None = None, + ) -> "CapabilityContext": + return replace( + self, + _capability_id=capability_id, + _capability_invocation_id=invocation_id, + _capability_version=capability_version, + _tool_id=None, + ) + + def for_tool(self, tool_id: str) -> "CapabilityContext": + return replace(self, _tool_id=tool_id) + + def with_current_user_message(self, message: str) -> "CapabilityContext": + """Bind the exact current user text as request-scoped evidence.""" + + return replace(self, current_user_message=message) + + def with_conversation_context( + self, + conversation_context: "ConversationContext", + ) -> "CapabilityContext": + """Bind the validated typed context revision for this turn.""" + + return replace(self, conversation_context=conversation_context) + + async def cancelled(self) -> bool: + return await self.is_cancelled() + + async def find_artifacts( + self, + artifact_model: type[ArtifactT], + ) -> tuple[ArtifactT, ...]: + return await self.artifacts.find_artifacts( + conversation_id=self.conversation_id, + artifact_model=artifact_model, + ) + + async def save_artifact(self, artifact: ArtifactT) -> ArtifactT: + saved = await self.artifacts.save_artifact( + conversation_id=self.conversation_id, + artifact=artifact, + ) + if not isinstance(saved, type(artifact)): + raise TypeError("Artifact repository returned an incompatible model.") + return saved + + @property + def capability_invocation_id(self) -> str: + if self._capability_invocation_id is None: + raise RuntimeError("Capability invocation identity is unavailable.") + return self._capability_invocation_id + + async def persist_waiting( + self, + partial_input: BaseModel, + *, + input_schema_version: str = "1", + ) -> object: + if self._capability_id is None or self._capability_version is None: + raise RuntimeError("Capability identity is unavailable for waiting input.") + from capabilities.repository import WaitingCapabilityInvocation + + invocation = WaitingCapabilityInvocation( + invocation_id=self.capability_invocation_id, + conversation_id=self.conversation_id, + capability_id=self._capability_id, + capability_version=self._capability_version, + input_schema_version=input_schema_version, + partial_input=partial_input, + source_turn_id=self.turn_id, + context_scope_id=getattr(partial_input, "context_scope_id", None), + context_revision=getattr(partial_input, "context_revision", None), + requirements=tuple( + getattr(partial_input, "fact_requirements", ()) + ), + ) + return await self.artifacts.save_waiting(invocation) + + async def waiting_invocations( + self, + capability_id: str, + ) -> tuple["WaitingCapabilityInvocation", ...]: + return await self.artifacts.list_waiting( + conversation_id=self.conversation_id, + capability_id=capability_id, + ) + + async def update_waiting( + self, + invocation_id: str, + partial_input: BaseModel, + ) -> "WaitingCapabilityInvocation": + return await self.artifacts.update_waiting( + invocation_id=invocation_id, + partial_input=partial_input, + ) + + async def remove_waiting(self, invocation_id: str) -> None: + await self.artifacts.remove_waiting(invocation_id=invocation_id) + + async def invoke_tool(self, identifier: str, tool_input: object) -> BaseModel: + if self._tool_id is not None: + caller = CallerType.TOOL + elif self._capability_id is not None: + caller = CallerType.CAPABILITY + else: + raise RuntimeError("Nested tool invocation requires a caller identity.") + return await self._executor.invoke_tool( + identifier, + tool_input, + caller=caller, + context=self, + ) + + async def invoke_capability( + self, + identifier: str, + capability_input: object, + ) -> "CapabilityOutcome[BaseModel]": + if self._capability_id is None: + raise RuntimeError("Nested capability invocation requires a capability identity.") + return await self._executor.invoke_capability( + identifier, + capability_input, + caller=CallerType.CAPABILITY, + context=self, + ) + + def record_model_usage( + self, + *, + input_tokens: int = 0, + output_tokens: int = 0, + cache_creation_input_tokens: int = 0, + cache_read_input_tokens: int = 0, + ) -> None: + self.model_usage.record( + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_creation_input_tokens=cache_creation_input_tokens, + cache_read_input_tokens=cache_read_input_tokens, + ) diff --git a/backend/capabilities/contracts.py b/backend/capabilities/contracts.py new file mode 100644 index 00000000..ccbcdd98 --- /dev/null +++ b/backend/capabilities/contracts.py @@ -0,0 +1,147 @@ +"""Typed contracts and outcomes for conversational capabilities.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Generic, Literal, TypeAlias, TypeVar + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from tools.contracts import CallerType, Visibility +from conversation_context.models import ( + CapabilityInvocationReference, + FactRequirement, +) + + +InputT = TypeVar("InputT", bound=BaseModel) +OutputT = TypeVar("OutputT", bound=BaseModel) + + +class ArtifactContract(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + artifact_type: str + schema_version: str + + @field_validator("artifact_type", "schema_version") + @classmethod + def require_non_empty(cls, value: str) -> str: + if not value.strip(): + raise ValueError("must not be empty") + return value + + def is_compatible_with(self, produced: "ArtifactContract") -> bool: + return ( + self.artifact_type == produced.artifact_type + and self.schema_version == produced.schema_version + ) + + +class CapabilityDependency(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + capability_id: str + artifact: ArtifactContract | None = None + + +class CapabilitySpec(BaseModel): + """Immutable registration and composition metadata for one capability.""" + + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True, extra="forbid") + + identifier: str + version: str + description: str + required_use: str + visibility: Visibility + allowed_callers: frozenset[CallerType] + input_model: type[BaseModel] + output_model: type[BaseModel] + accepted_artifacts: tuple[ArtifactContract, ...] = () + produced_artifacts: tuple[ArtifactContract, ...] = () + tool_dependencies: tuple[str, ...] = () + dependencies: tuple[CapabilityDependency, ...] = () + + @field_validator("identifier", "version", "description", "required_use") + @classmethod + def require_non_empty(cls, value: str) -> str: + if not value.strip(): + raise ValueError("must not be empty") + return value + + @field_validator("allowed_callers") + @classmethod + def require_allowed_caller( + cls, value: frozenset[CallerType] + ) -> frozenset[CallerType]: + if not value: + raise ValueError("at least one allowed caller is required") + return value + + +class Completed(BaseModel, Generic[OutputT]): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: Literal["completed"] = "completed" + value: OutputT + + +class NeedsInput(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: Literal["needs_input"] = "needs_input" + prompt: str + missing_fields: tuple[str, ...] = () + partial_input: dict[str, object] = Field(default_factory=dict) + fact_requirements: tuple[FactRequirement, ...] = () + capability_invocation: CapabilityInvocationReference | None = None + + +class Unsupported(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: Literal["unsupported"] = "unsupported" + reason: str + + +class Failed(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + status: Literal["failed"] = "failed" + safe_message: str + error_code: str + + +class Accepted(BaseModel): + """Reserved for a future explicitly asynchronous capability.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + status: Literal["accepted"] = "accepted" + operation_id: str + + +CapabilityOutcome: TypeAlias = ( + Completed[OutputT] | NeedsInput | Unsupported | Failed | Accepted +) + + +class Capability(ABC, Generic[InputT, OutputT]): + """A cohesive component that coordinates typed tools and prerequisites.""" + + spec: CapabilitySpec + + @abstractmethod + async def run( + self, + capability_input: InputT, + context: "CapabilityContext", + ) -> CapabilityOutcome[OutputT]: + """Execute the capability-owned sequence.""" + + def trace_summary(self, status: str) -> str: + return f"{self.spec.identifier} {status}" + + +from capabilities.context import CapabilityContext # noqa: E402 diff --git a/backend/capabilities/executor.py b/backend/capabilities/executor.py new file mode 100644 index 00000000..2bae5b77 --- /dev/null +++ b/backend/capabilities/executor.py @@ -0,0 +1,391 @@ +"""Authorized validation, dispatch, cancellation, and trace recording.""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import TYPE_CHECKING, Any, Protocol + +from pydantic import BaseModel, ValidationError + +from capabilities.context import ( + ArtifactAccess, + CancellationProbe, + CapabilityContext, + EmptyArtifactAccess, + ModelUsageLedger, +) +from capabilities.contracts import ( + Accepted, + Capability, + CapabilityOutcome, + Completed, + Failed, + NeedsInput, + Unsupported, +) +from capabilities.registry import CapabilityRegistry +from capabilities.tracing import ( + debug_projection, + InvocationKind, + InvocationStatus, + InvocationTracer, +) +from tools.contracts import CallerType +from tools.context import TurnResultStore +from tools.registry import ToolRegistry + +if TYPE_CHECKING: + from conversation_context.models import ConversationContext + + +_PARENT_INVOCATION: ContextVar[str | None] = ContextVar( + "capability_parent_invocation", + default=None, +) + + +class InvocationCancelled(Exception): + pass + + +class InvocationTraceValues(Protocol): + """Select the validated values represented by one invocation trace.""" + + def input_value( + self, + *, + raw_input: object, + validated_input: BaseModel, + ) -> object: ... + + def output_value(self, validated_output: object) -> object: ... + + def failed_output(self) -> object: ... + + def cancelled_output(self) -> object: ... + + +class RegisteredBoundaryTraceValues: + """Represent the typed input and output seen by a registered operation.""" + + def input_value( + self, + *, + raw_input: object, + validated_input: BaseModel, + ) -> object: + del raw_input + return validated_input + + def output_value(self, validated_output: object) -> object: + return validated_output + + def failed_output(self) -> object: + return {"status": "failed"} + + def cancelled_output(self) -> object: + return {"status": "cancelled"} + + +_REGISTERED_BOUNDARY_TRACE_VALUES = RegisteredBoundaryTraceValues() + + +class InvocationExecutor: + """Execute only registered, authorized operations with typed boundaries.""" + + def __init__( + self, + *, + tools: ToolRegistry, + capabilities: CapabilityRegistry, + tracer: InvocationTracer, + ) -> None: + self._tools = tools + self._capabilities = capabilities + self._tracer = tracer + + @property + def tracer(self) -> InvocationTracer: + """Expose only sanitized trace projections to the chat adapter.""" + + return self._tracer + + def context( + self, + *, + request_id: str, + conversation_id: str, + turn_id: str, + is_cancelled: CancellationProbe, + artifacts: ArtifactAccess | None = None, + result_store: TurnResultStore | None = None, + model_usage: ModelUsageLedger | None = None, + conversation_context: "ConversationContext | None" = None, + ) -> CapabilityContext: + return CapabilityContext( + request_id=request_id, + conversation_id=conversation_id, + turn_id=turn_id, + is_cancelled=is_cancelled, + artifacts=artifacts or EmptyArtifactAccess(), + result_store=result_store or TurnResultStore(), + model_usage=model_usage or ModelUsageLedger(), + conversation_context=conversation_context, + _executor=self, + ) + + async def invoke_tool( + self, + identifier: str, + raw_input: object, + *, + caller: CallerType, + context: CapabilityContext, + ) -> BaseModel: + tool = self._tools.get(identifier, caller=caller) + self._validate_declared_tool_call(identifier, caller, context) + tool_input = self._validate(tool.spec.input_model, raw_input, identifier, "input") + record = self._tracer.start( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + parent_invocation_id=_PARENT_INVOCATION.get(), + kind=InvocationKind.TOOL, + identifier=identifier, + version=tool.spec.version, + visibility=tool.spec.visibility, + summary=self._trace_summary("tool", identifier, "started"), + debug_input=debug_projection(tool_input), + ) + token = _PARENT_INVOCATION.set(record.invocation_id) + try: + await self._check_cancelled(context) + raw_output = await tool.run(tool_input, context.for_tool(identifier)) + output = self._validate( + tool.spec.output_model, + raw_output, + identifier, + "output", + ) + await self._check_cancelled(context) + except InvocationCancelled: + self._tracer.finish( + record.invocation_id, + status=InvocationStatus.CANCELLED, + summary=self._trace_summary("tool", identifier, "cancelled"), + debug_output={"status": "cancelled"}, + ) + raise + except Exception: + self._tracer.finish( + record.invocation_id, + status=InvocationStatus.FAILED, + summary=self._trace_summary("tool", identifier, "failed"), + debug_output={"status": "failed"}, + ) + raise + else: + self._tracer.finish( + record.invocation_id, + status=InvocationStatus.COMPLETED, + summary=self._trace_summary("tool", identifier, "completed"), + debug_output=debug_projection(output), + ) + return output + finally: + _PARENT_INVOCATION.reset(token) + + async def invoke_capability( + self, + identifier: str, + raw_input: object, + *, + caller: CallerType, + context: CapabilityContext, + trace_values: InvocationTraceValues = _REGISTERED_BOUNDARY_TRACE_VALUES, + ) -> CapabilityOutcome[BaseModel]: + capability = self._capabilities.get(identifier, caller=caller) + self._validate_declared_capability_call(identifier, caller, context) + capability_input = self._validate( + capability.spec.input_model, + raw_input, + identifier, + "input", + ) + record = self._tracer.start( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + parent_invocation_id=_PARENT_INVOCATION.get(), + kind=InvocationKind.CAPABILITY, + identifier=identifier, + version=capability.spec.version, + visibility=capability.spec.visibility, + summary=self._trace_summary("capability", identifier, "started"), + debug_input=debug_projection( + trace_values.input_value( + raw_input=raw_input, + validated_input=capability_input, + ) + ), + ) + token = _PARENT_INVOCATION.set(record.invocation_id) + try: + await self._check_cancelled(context) + outcome = await capability.run( + capability_input, + context.for_capability( + identifier, + record.invocation_id, + capability.spec.version, + ), + ) + validated = self._validate_outcome(capability, outcome) + await self._check_cancelled(context) + trace_status = self._outcome_status(validated) + except InvocationCancelled: + self._tracer.finish( + record.invocation_id, + status=InvocationStatus.CANCELLED, + summary=self._trace_summary("capability", identifier, "cancelled"), + debug_output=debug_projection(trace_values.cancelled_output()), + ) + raise + except Exception: + self._tracer.finish( + record.invocation_id, + status=InvocationStatus.FAILED, + summary=self._trace_summary("capability", identifier, "failed"), + debug_output=debug_projection(trace_values.failed_output()), + ) + raise + else: + self._tracer.finish( + record.invocation_id, + status=trace_status, + summary=self._trace_summary( + "capability", + identifier, + validated.status, + ), + debug_output=debug_projection( + trace_values.output_value(validated) + ), + ) + return validated + finally: + _PARENT_INVOCATION.reset(token) + + @staticmethod + def _validate( + model: type[BaseModel], + value: object, + identifier: str, + boundary: str, + ) -> BaseModel: + try: + if isinstance(value, model): + return value + return model.model_validate(value) + except ValidationError as exc: + raise TypeError( + f"Invalid {boundary} for registered operation {identifier}." + ) from exc + + @staticmethod + def _validate_outcome( + capability: Capability[Any, Any], + outcome: object, + ) -> CapabilityOutcome[BaseModel]: + if isinstance(outcome, Completed): + value = InvocationExecutor._validate( + capability.spec.output_model, + outcome.value, + capability.spec.identifier, + "completed output", + ) + return Completed(value=value) + if isinstance(outcome, NeedsInput): + unknown = set(outcome.partial_input) - set( + capability.spec.input_model.model_fields + ) + if unknown: + raise TypeError( + f"Capability {capability.spec.identifier} returned unknown partial " + f"input fields: {sorted(unknown)}." + ) + return outcome + if isinstance(outcome, (Unsupported, Failed)): + return outcome + if isinstance(outcome, Accepted): + raise TypeError( + "Accepted is reserved for a future explicitly asynchronous capability." + ) + raise TypeError( + f"Capability {capability.spec.identifier} returned an invalid outcome." + ) + + @staticmethod + def _outcome_status(outcome: CapabilityOutcome[BaseModel]) -> InvocationStatus: + if isinstance(outcome, Completed): + return InvocationStatus.COMPLETED + if isinstance(outcome, NeedsInput): + return InvocationStatus.NEEDS_INPUT + if isinstance(outcome, Unsupported): + return InvocationStatus.UNSUPPORTED + return InvocationStatus.FAILED + + async def _check_cancelled(self, context: CapabilityContext) -> None: + if await context.cancelled(): + raise InvocationCancelled + + @staticmethod + def _trace_summary(kind: str, identifier: str, status: str) -> str: + """Build trace text exclusively from registered metadata and status.""" + + return f"{kind} {identifier} {status}" + + def _validate_declared_tool_call( + self, + identifier: str, + caller: CallerType, + context: CapabilityContext, + ) -> None: + if caller is CallerType.TOOL: + source_id = context._tool_id + if source_id is None: + raise PermissionError("Tool caller identity is required.") + tool_source = self._tools.registered(source_id) + if identifier not in tool_source.spec.tool_dependencies: + raise PermissionError( + f"Tool {source_id} did not declare tool dependency {identifier}." + ) + return + if caller is not CallerType.CAPABILITY: + return + source_id = context._capability_id + if source_id is None: + raise PermissionError("Capability caller identity is required.") + capability_source = self._capabilities.registered(source_id) + if identifier not in capability_source.spec.tool_dependencies: + raise PermissionError( + f"Capability {source_id} did not declare tool dependency {identifier}." + ) + + def _validate_declared_capability_call( + self, + identifier: str, + caller: CallerType, + context: CapabilityContext, + ) -> None: + if caller is not CallerType.CAPABILITY: + return + source_id = context._capability_id + if source_id is None: + raise PermissionError("Capability caller identity is required.") + source = self._capabilities.registered(source_id) + declared = { + dependency.capability_id for dependency in source.spec.dependencies + } + if identifier not in declared: + raise PermissionError( + f"Capability {source_id} did not declare capability dependency {identifier}." + ) diff --git a/backend/capabilities/follow_up.py b/backend/capabilities/follow_up.py new file mode 100644 index 00000000..1bf2e352 --- /dev/null +++ b/backend/capabilities/follow_up.py @@ -0,0 +1,206 @@ +"""Typed household and population result follow-up capability.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict + +from capabilities.artifacts import ( + HouseholdResultRef, + PolicyScenarioRef, + SocietyAnalysisResultRef, +) +from capabilities.compatibility import ArtifactRequirements, check_artifact_compatibility +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilityDependency, + CapabilitySpec, + Completed, + NeedsInput, + Unsupported, +) +from capabilities.society import ( + SOCIETY_DEFAULT_PROFILE_VERSION, + SocietyAnalysisOutput, + current_dataset_version, +) +from tools.analysis_support import ( + ExtractResultFindingsOutput, + NumericalFact, + SelectSupportedOutputsOutput, +) +from tools.contracts import CallerType, Visibility + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class AnalysisFollowUpInput(StrictModel): + question: str + referenced_result_id: str | None = None + requested_outputs: tuple[str, ...] = () + + +class AnalysisFollowUpOutput(StrictModel): + source_artifact_id: str + result: HouseholdResultRef | SocietyAnalysisResultRef + reran_provider: bool = False + narration_facts: tuple[NumericalFact, ...] + + +class AnalysisFollowUpCapability( + Capability[AnalysisFollowUpInput, AnalysisFollowUpOutput] +): + spec = CapabilitySpec( + identifier="analysis_follow_up", + version="1", + description=( + "Explain or extend a compatible typed household or population result " + "without parsing earlier assistant prose." + ), + required_use=( + "Use for later questions whose authoritative basis is an existing " + "household or population result." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=AnalysisFollowUpInput, + output_model=AnalysisFollowUpOutput, + accepted_artifacts=( + ArtifactContract(artifact_type="household_result", schema_version="1"), + ArtifactContract( + artifact_type="society_analysis_result", + schema_version="1", + ), + ), + dependencies=( + CapabilityDependency( + capability_id="society_analysis", + artifact=ArtifactContract( + artifact_type="society_analysis_result", + schema_version="1", + ), + ), + ), + tool_dependencies=("select_supported_outputs", "extract_result_findings"), + ) + + async def run(self, capability_input: AnalysisFollowUpInput, context): + candidates = [ + *await context.find_artifacts(HouseholdResultRef), + *await context.find_artifacts(SocietyAnalysisResultRef), + ] + if capability_input.referenced_result_id is not None: + candidates = [ + item + for item in candidates + if item.artifact_id == capability_input.referenced_result_id + ] + if not candidates: + return NeedsInput( + prompt="Which completed household or population result should I use?", + missing_fields=("referenced_result_id",), + partial_input=capability_input.model_dump(mode="json", exclude_none=True), + ) + if len(candidates) > 1: + distinctions = ", ".join( + f"{item.artifact_type} {item.artifact_id} ({item.year})" + for item in candidates + ) + return NeedsInput( + prompt=f"Please choose one prior result: {distinctions}.", + missing_fields=("referenced_result_id",), + partial_input=capability_input.model_dump(mode="json", exclude_none=True), + ) + source = candidates[0] + result = source + reran = False + if isinstance(source, SocietyAnalysisResultRef): + compatibility = check_artifact_compatibility( + source, + ArtifactRequirements( + artifact_type="society_analysis_result", + schema_version="1", + year=source.year, + scenario_revision=source.scenario_revision, + dataset_version=current_dataset_version(), + calculation_engine_version=source.calculation_engine_version, + ), + ) + if not compatibility.compatible: + return Unsupported( + reason="The retained population result is not compatible with this operation." + ) + if source.default_profile_version != SOCIETY_DEFAULT_PROFILE_VERSION: + return Unsupported( + reason="The retained population result uses an unsupported default profile." + ) + selected = await context.invoke_tool( + "select_supported_outputs", + {"requested_outputs": list(capability_input.requested_outputs)}, + ) + if not isinstance(selected, SelectSupportedOutputsOutput): + raise TypeError("Follow-up output selection returned an incompatible result.") + missing = set(selected.output_ids) - set(source.calculated_output_ids) + if missing: + scenario = await self._scenario(source, context) + if scenario is None: + return Unsupported( + reason=( + "The requested aggregate was not retained and its verified " + "policy scenario is unavailable for rerun." + ) + ) + rerun = await context.invoke_capability( + "society_analysis", + { + "referenced_policy_scenario_id": scenario.artifact_id, + "year": source.year, + "requested_outputs": list(capability_input.requested_outputs), + }, + ) + if not isinstance(rerun, Completed) or not isinstance( + rerun.value, + SocietyAnalysisOutput, + ): + return rerun + result = rerun.value.result + reran = True + + extracted = await context.invoke_tool( + "extract_result_findings", + { + "outputs": [ + output.model_dump(mode="json") for output in result.outputs + ] + }, + ) + if not isinstance(extracted, ExtractResultFindingsOutput): + raise TypeError("Follow-up finding extraction returned an incompatible result.") + facts = tuple( + NumericalFact(label=finding.label, value=finding.value, unit=finding.unit) + for finding in extracted.findings + if finding.value is not None + ) + return Completed( + value=AnalysisFollowUpOutput( + source_artifact_id=source.artifact_id, + result=result, + reran_provider=reran, + narration_facts=facts, + ) + ) + + @staticmethod + async def _scenario(result, context): + scenarios = await context.find_artifacts(PolicyScenarioRef) + return next( + ( + scenario + for scenario in scenarios + if scenario.artifact_id == result.policy_scenario_artifact_id + and scenario.scenario_revision == result.scenario_revision + ), + None, + ) diff --git a/backend/capabilities/household.py b/backend/capabilities/household.py new file mode 100644 index 00000000..bc65b2bb --- /dev/null +++ b/backend/capabilities/household.py @@ -0,0 +1,2113 @@ +"""Household input assembly, explicit assumptions, validation, and calculation.""" + +from __future__ import annotations + +import hashlib +import json +import re +from enum import Enum +from typing import Literal + +from pydantic import AliasChoices, Field, JsonValue + +from capabilities.artifacts import ( + AggregateDimension, + AggregateValue, + ArtifactProvenance, + HouseholdRef, + HouseholdEntityPosition, + HouseholdResultRef, + HouseholdValue, + PolicyScenarioRef, +) +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilityDependency, + CapabilitySpec, + Completed, + Failed, + NeedsInput, + Unsupported, +) +from capabilities.input_resolution import InputSource, resolve_policy_year +from capabilities.household_input import ( + AmountFrequency, + HouseholdEvidence, + HouseholdEvidenceAmbiguity, + HouseholdEvidenceAmbiguityKind, + HouseholdEvidenceAssembler, + HouseholdEvidenceResult, + HouseholdInvocationDefaults, + HouseholdCalculationRequirements, + HouseholdInputResolver, + HouseholdInputCompleteness, + PeriodicAmount, + PersonEvidence, + StrictModel, +) +from capabilities.policy_reform import PolicyReformOutput +from conversation_context.household_view import HouseholdContextView +from conversation_context.engine_projection import ( + HouseholdEngineFactProjector, + HouseholdEngineInputs, +) +from conversation_context.models import ( + CapabilityInvocationReference, + ExplicitAbsenceAssertion, + PendingQuestionStatus, + FactRequirement, +) +from tools.analysis_support import ( + ExtractResultFindingsOutput, + NumericalFact, +) +from tools.contracts import CallerType, Tool, ToolCallContext, ToolSpec, Visibility +from tools.typed_models import SafeToolOutput + + +class HouseholdAssemblyStatus(str, Enum): + READY = "ready" + NEEDS_INPUT = "needs_input" + FAILED = "failed" + + +class HouseholdAssumption(StrictModel): + field_id: str + label: str + assumed_value: JsonValue + plain_statement: str + label_source: Literal["catalogue", "household_contract", "system_default"] + material: bool = True + + +class HouseholdCandidate(StrictModel): + people: tuple[dict[str, JsonValue], ...] + benunit: dict[str, JsonValue] + household: dict[str, JsonValue] + field_values: tuple[HouseholdValue, ...] + entity_positions: tuple[HouseholdEntityPosition, ...] = () + input_narration_facts: tuple[NumericalFact, ...] = () + + +class AssembleHouseholdInput(StrictModel): + description: str + requirements: HouseholdCalculationRequirements = Field( + default_factory=HouseholdCalculationRequirements + ) + existing_values: tuple[HouseholdValue, ...] = () + retained_evidence: HouseholdEvidence = Field(default_factory=HouseholdEvidence) + invocation_defaults: HouseholdInvocationDefaults = Field( + default_factory=HouseholdInvocationDefaults + ) + retained_ambiguities: tuple[HouseholdEvidenceAmbiguity, ...] = () + pending_fields: tuple[str, ...] = () + + +class AssembleHouseholdOutput(StrictModel): + status: HouseholdAssemblyStatus + evidence: HouseholdEvidence = Field(default_factory=HouseholdEvidence) + invocation_defaults: HouseholdInvocationDefaults = Field( + default_factory=HouseholdInvocationDefaults + ) + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...] = () + candidate: HouseholdCandidate | None = None + assumptions: tuple[HouseholdAssumption, ...] = () + questions: tuple[str, ...] = () + missing_fields: tuple[str, ...] = () + fact_requirements: tuple[FactRequirement, ...] = () + error: str | None = None + + +_CATALOGUE_FIELDS = { + "age": "Age", + "employment_income": "Employment income", + "self_employment_income": "Self-employment income", + "pension_income": "Pension income", + "is_married": "Married or in a civil partnership", + "childcare_expenses": "Childcare expenses", + "rent": "Rent", + "council_tax": "Council Tax", +} + +_HOUSEHOLD_OUTPUT_ALIASES = { + "benefits": "household_benefits", + "benefit entitlement": "household_benefits", + "benefit entitlements": "household_benefits", + "household benefits": "household_benefits", + "housing benefit": "housing_benefit", + "income tax": "income_tax", + "national insurance": "national_insurance", + "national insurance contribution": "national_insurance", + "national insurance contributions": "national_insurance", + "total tax": "household_tax", + "total taxes": "household_tax", + "household tax": "household_tax", + "net income": "household_net_income", + "household net income": "household_net_income", + "universal credit": "universal_credit", +} + +_TAX_ONLY_HOUSEHOLD_OUTPUTS = frozenset( + {"income_tax", "national_insurance", "household_tax"} +) + +_HOUSEHOLD_OUTPUT_GROUPS = { + "tax": ( + "income_tax", + "national_insurance", + "household_tax", + ), + "taxes": ( + "income_tax", + "national_insurance", + "household_tax", + ), +} + + +class AssembleHouseholdCandidateTool( + Tool[AssembleHouseholdInput, AssembleHouseholdOutput] +): + spec = ToolSpec( + identifier="assemble_household_candidate", + version="1", + description=( + "Extract household evidence, apply documented safe defaults, and return " + "every material assumption with a plain-language label." + ), + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.CAPABILITY}), + input_model=AssembleHouseholdInput, + output_model=AssembleHouseholdOutput, + tool_dependencies=("get_variable",), + ) + + def __init__( + self, + assembler: HouseholdEvidenceAssembler | None = None, + resolver: HouseholdInputResolver | None = None, + ) -> None: + self._assembler = assembler + self._resolver = resolver or HouseholdInputResolver() + + async def run(self, tool_input: AssembleHouseholdInput, context: ToolCallContext): + assembled = ( + await self._assembler.assemble( + description=tool_input.description, + retained_evidence=tool_input.retained_evidence, + retained_ambiguities=tool_input.retained_ambiguities, + ) + if self._assembler is not None + else HouseholdEvidenceResult(evidence=HouseholdEvidence()) + ) + context.record_model_usage(**assembled.usage.model_dump()) + artifact_evidence = self._evidence_from_values(tool_input.existing_values) + defaulted_evidence = self._merge_evidence( + tool_input.invocation_defaults.evidence, + artifact_evidence, + ) + retained_evidence = self._merge_evidence( + defaulted_evidence, + tool_input.retained_evidence, + ) + current_evidence = self._as_current_user_evidence( + assembled.evidence, + assembled.ambiguities, + ) + evidence = self._merge_evidence(retained_evidence, current_evidence) + ambiguities = self._merge_ambiguities( + tool_input.retained_ambiguities, + assembled.ambiguities, + current_evidence, + ) + evidence, ambiguities = self._resolver.apply_documented_defaults( + evidence, + ambiguities, + current_user_message=tool_input.description, + requirements=tool_input.requirements, + ) + invocation_defaults = HouseholdInvocationDefaults( + evidence=self._merge_evidence( + tool_input.invocation_defaults.evidence, + self._default_only_evidence(evidence), + ) + ) + resolution = self._resolver.resolve( + evidence, + ambiguities, + tool_input.requirements, + ) + if resolution.questions: + return AssembleHouseholdOutput( + status=HouseholdAssemblyStatus.NEEDS_INPUT, + evidence=evidence, + invocation_defaults=invocation_defaults, + ambiguities=ambiguities, + questions=resolution.questions, + missing_fields=resolution.missing_fields, + fact_requirements=resolution.fact_requirements, + ) + + labels = await self._catalogue_labels(context) + if labels is None: + return AssembleHouseholdOutput( + status=HouseholdAssemblyStatus.FAILED, + evidence=evidence, + invocation_defaults=invocation_defaults, + ambiguities=ambiguities, + error="A supported household input lacks an authoritative presentation label.", + ) + assumptions: list[HouseholdAssumption] = [] + values: list[HouseholdValue] = [] + adult_count = sum( + 1 + for person in evidence.people + if person.age is not None and person.age >= 16 + ) + child_count = sum( + 1 + for person in evidence.people + if person.age is not None and person.age < 16 + ) + input_facts: list[NumericalFact] = [ + NumericalFact( + label="Number of people in the household", + value=len(evidence.people), + unit="people", + ), + NumericalFact( + label="Number of adults in the household", + value=adult_count, + unit="adults", + ), + NumericalFact( + label="Number of children in the household", + value=child_count, + unit="children", + ), + ] + people: list[dict[str, JsonValue]] = [] + for index, person in enumerate(evidence.people): + person_data: dict[str, JsonValue] = {"age": person.age} + if person.age is not None: + input_facts.append( + NumericalFact( + label=( + "Your age" + if person.relationship_to_user == "self" + else ( + f"Age of {person.display_label}" + if person.display_label + else "Age" + ) + ), + value=person.age, + unit="years", + ) + ) + values.append( + HouseholdValue( + field_id=f"people[{index}].age", + label=labels["age"], + value=person.age, + source=person.sources.get("age", "user"), + subject_entity_id=person.entity_id, + engine_variable="age", + ) + ) + for field in ( + "employment_income", + "self_employment_income", + "pension_income", + ): + periodic_value = getattr(person, field) + if periodic_value is None: + annual_value = 0.0 + assumptions.append( + HouseholdAssumption( + field_id=f"people[{index}].{field}", + label=labels[field], + assumed_value=annual_value, + plain_statement=( + f"No {labels[field].casefold()}." + if len(evidence.people) == 1 + else ( + f"{self._person_subject(person, index)} has no " + f"{labels[field].casefold()}." + ) + ), + label_source="catalogue", + ) + ) + source = "default" + else: + annual_value = periodic_value.annual_value() + input_facts.extend( + self._periodic_amount_facts( + labels[field], + periodic_value, + ) + ) + source = person.sources.get(field, "user") + if source == "default": + assumptions.append( + HouseholdAssumption( + field_id=f"people[{index}].{field}", + label=labels[field], + assumed_value=annual_value, + plain_statement=( + f"The £{annual_value:,.2f} income amount is " + f"treated as annual {labels[field].casefold()}." + ), + label_source="catalogue", + ) + ) + person_data[field] = annual_value + values.append( + HouseholdValue( + field_id=f"people[{index}].{field}", + label=labels[field], + value=annual_value, + source=source, + subject_entity_id=person.entity_id, + engine_variable=field, + period="annual", + ) + ) + people.append(person_data) + + if child_count == 0 and evidence.has_children is None: + assumptions.append( + HouseholdAssumption( + field_id="household.children", + label="Children", + assumed_value=0, + plain_statement="The household has no children.", + label_source="household_contract", + ) + ) + is_married = evidence.is_married + if is_married is None: + is_married = False + assumptions.append( + HouseholdAssumption( + field_id="benunit.is_married", + label=labels["is_married"], + assumed_value=False, + plain_statement="The household has no partner or spouse.", + label_source="catalogue", + ) + ) + values.append( + HouseholdValue( + field_id="benunit.is_married", + label=labels["is_married"], + value=is_married, + source=evidence.sources.get("is_married", "default"), + engine_variable="is_married", + ) + ) + + benunit: dict[str, JsonValue] = {"is_married": is_married} + household: dict[str, JsonValue] = {} + primary_adult_index = next( + index + for index, person in enumerate(evidence.people) + if person.age is not None and person.age >= 16 + ) + childcare_expenses = evidence.childcare_expenses + if childcare_expenses is None: + annual_childcare_expenses = 0.0 + assumptions.append( + HouseholdAssumption( + field_id=( + f"people[{primary_adult_index}].childcare_expenses" + ), + label=labels["childcare_expenses"], + assumed_value=annual_childcare_expenses, + plain_statement="The household has no childcare expenses.", + label_source="catalogue", + ) + ) + childcare_source = "default" + else: + annual_childcare_expenses = childcare_expenses.annual_value() + input_facts.extend( + self._periodic_amount_facts( + labels["childcare_expenses"], + childcare_expenses, + ) + ) + childcare_source = evidence.sources.get("childcare_expenses", "user") + people[primary_adult_index]["childcare_expenses"] = annual_childcare_expenses + values.append( + HouseholdValue( + field_id=f"people[{primary_adult_index}].childcare_expenses", + label=labels["childcare_expenses"], + value=annual_childcare_expenses, + source=childcare_source, + subject_entity_id=evidence.people[primary_adult_index].entity_id, + engine_variable="childcare_expenses", + period="annual", + ) + ) + + for field, target in ( + ("rent", household), + ("council_tax", household), + ): + periodic_value = getattr(evidence, field) + if periodic_value is None: + if tool_input.requirements.require_housing_costs: + return AssembleHouseholdOutput( + status=HouseholdAssemblyStatus.FAILED, + evidence=evidence, + ambiguities=ambiguities, + error=( + f"Household resolution left consequential {labels[field]} " + "information unresolved." + ), + ) + periodic_value = PeriodicAmount( + amount=0, + frequency=AmountFrequency.ANNUAL, + ) + annual_value = periodic_value.annual_value() + source = evidence.sources.get(field, "default") + if source != "default": + input_facts.extend( + self._periodic_amount_facts(labels[field], periodic_value) + ) + target[field] = annual_value + values.append( + HouseholdValue( + field_id=( + f"benunit.{field}" + if target is benunit + else f"household.{field}" + ), + label=labels[field], + value=annual_value, + source=source, + engine_variable=field, + period="annual", + ) + ) + + country = evidence.country or "ENGLAND" + if evidence.country is None: + assumptions.append( + HouseholdAssumption( + field_id="household.country", + label="Country", + assumed_value=country, + plain_statement="The household lives in England.", + label_source="household_contract", + ) + ) + household["country"] = country + values.append( + HouseholdValue( + field_id="household.country", + label="Country", + value=country, + source=evidence.sources.get("country", "default"), + engine_variable="country", + ) + ) + return AssembleHouseholdOutput( + status=HouseholdAssemblyStatus.READY, + evidence=evidence, + invocation_defaults=invocation_defaults, + ambiguities=ambiguities, + candidate=HouseholdCandidate( + people=tuple(people), + benunit=benunit, + household=household, + field_values=tuple(values), + entity_positions=tuple( + HouseholdEntityPosition( + entity_id=person.entity_id, + engine_position=f"people[{index}]", + ) + for index, person in enumerate(evidence.people) + if person.entity_id is not None + ), + input_narration_facts=tuple(input_facts), + ), + assumptions=tuple(assumptions), + ) + + @classmethod + def _evidence_from_values( + cls, + values: tuple[HouseholdValue, ...], + ) -> HouseholdEvidence: + people: dict[int, dict[str, object]] = {} + household_values: dict[str, object] = {} + household_sources: dict[str, str] = {} + person_pattern = re.compile(r"people\[(\d+)]\.(.+)") + for value in values: + if value.source == "default": + continue + person_match = person_pattern.fullmatch(value.field_id) + if person_match is not None: + person_index = int(person_match.group(1)) + field = person_match.group(2) + if field == "childcare_expenses": + periodic = cls._annual_amount(value.value) + if periodic is not None: + household_values["childcare_expenses"] = periodic + household_sources["childcare_expenses"] = "artifact" + continue + if field not in { + "age", + "employment_income", + "self_employment_income", + "pension_income", + }: + continue + person = people.setdefault(person_index, {"sources": {}}) + if field == "age" and isinstance(value.value, (int, float)): + person[field] = int(value.value) + elif field != "age": + periodic = cls._annual_amount(value.value) + if periodic is None: + continue + person[field] = periodic + person["sources"][field] = "artifact" + continue + field = value.field_id.removeprefix("household.").removeprefix( + "benunit." + ) + if field in {"rent", "council_tax"}: + periodic = cls._annual_amount(value.value) + if periodic is not None: + household_values[field] = periodic + household_sources[field] = "artifact" + elif field == "is_married" and isinstance(value.value, bool): + household_values[field] = value.value + household_sources[field] = "artifact" + elif field == "country" and value.value in { + "ENGLAND", + "SCOTLAND", + "WALES", + "NORTHERN_IRELAND", + }: + household_values[field] = value.value + household_sources[field] = "artifact" + + people_tuple = tuple( + PersonEvidence.model_validate(people[index]) + for index in sorted(people) + ) + if people_tuple: + has_children = any( + person.age is not None and person.age < 16 for person in people_tuple + ) + household_values["has_children"] = has_children + household_sources["has_children"] = "artifact" + return HouseholdEvidence.model_validate( + { + "people": people_tuple, + **household_values, + "sources": household_sources, + } + ) + + @staticmethod + def _annual_amount(value: JsonValue) -> PeriodicAmount | None: + if isinstance(value, (int, float)) and not isinstance(value, bool): + return PeriodicAmount( + amount=float(value), + frequency=AmountFrequency.ANNUAL, + ) + return None + + @staticmethod + def _default_only_evidence(evidence: HouseholdEvidence) -> HouseholdEvidence: + people: list[PersonEvidence] = [] + last_default_index = -1 + for index, person in enumerate(evidence.people): + values: dict[str, object] = { + "entity_id": person.entity_id, + "display_label": person.display_label, + "relationship_to_user": person.relationship_to_user, + } + sources: dict[str, str] = {} + for field in ( + "age", + "employment_income", + "self_employment_income", + "pension_income", + ): + if person.sources.get(field) != "default": + continue + values[field] = getattr(person, field) + sources[field] = "default" + values["sources"] = sources + people.append(PersonEvidence.model_validate(values)) + if sources: + last_default_index = index + + household_values: dict[str, object] = {} + household_sources: dict[str, str] = {} + for field in ( + "has_children", + "is_married", + "childcare_expenses", + "rent", + "council_tax", + "country", + ): + if evidence.sources.get(field) != "default": + continue + household_values[field] = getattr(evidence, field) + household_sources[field] = "default" + return HouseholdEvidence.model_validate( + { + "people": tuple(people[: last_default_index + 1]), + **household_values, + "sources": household_sources, + } + ) + + @classmethod + def _merge_evidence( + cls, + retained: HouseholdEvidence, + current: HouseholdEvidence, + ) -> HouseholdEvidence: + people_count = max(len(retained.people), len(current.people)) + people: list[PersonEvidence] = [] + for index in range(people_count): + retained_person = ( + retained.people[index] + if index < len(retained.people) + else PersonEvidence() + ) + current_person = ( + current.people[index] + if index < len(current.people) + else PersonEvidence() + ) + updates: dict[str, object] = {} + sources = dict(retained_person.sources) + for identity_field in ( + "entity_id", + "display_label", + "relationship_to_user", + ): + current_identity = getattr(current_person, identity_field) + if current_identity is not None: + updates[identity_field] = current_identity + for field in ( + "age", + "employment_income", + "self_employment_income", + "pension_income", + ): + current_value = getattr(current_person, field) + if current_value is not None: + updates[field] = current_value + sources[field] = current_person.sources.get(field, "user") + updates["sources"] = sources + people.append(retained_person.model_copy(update=updates)) + + household_updates: dict[str, object] = {"people": tuple(people)} + household_sources = dict(retained.sources) + for field in ( + "has_children", + "is_married", + "childcare_expenses", + "rent", + "council_tax", + "country", + ): + current_value = getattr(current, field) + if current_value is not None: + household_updates[field] = current_value + household_sources[field] = current.sources.get(field, "user") + household_updates["sources"] = household_sources + return retained.model_copy(update=household_updates) + + @staticmethod + def _as_current_user_evidence( + evidence: HouseholdEvidence, + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...], + ) -> HouseholdEvidence: + people = [] + for index, person in enumerate(evidence.people): + person_updates: dict[str, object] = {} + for ambiguity in ambiguities: + if ambiguity.kind not in { + HouseholdEvidenceAmbiguityKind.INCOME_OWNER, + HouseholdEvidenceAmbiguityKind.INCOME_FREQUENCY, + }: + continue + if ambiguity.field not in { + "employment_income", + "self_employment_income", + "pension_income", + }: + continue + if ambiguity.person_indices and index not in ambiguity.person_indices: + continue + person_updates[ambiguity.field] = None + normalized = person.model_copy(update=person_updates) + people.append( + normalized.model_copy( + update={ + "sources": { + field: "user" + for field in ( + "age", + "employment_income", + "self_employment_income", + "pension_income", + ) + if getattr(normalized, field) is not None + } + } + ) + ) + sources = { + field: "user" + for field in ( + "has_children", + "is_married", + "childcare_expenses", + "rent", + "council_tax", + "country", + ) + if getattr(evidence, field) is not None + } + return evidence.model_copy( + update={ + "people": tuple(people), + "sources": sources, + } + ) + + @classmethod + def _merge_ambiguities( + cls, + retained: tuple[HouseholdEvidenceAmbiguity, ...], + current: tuple[HouseholdEvidenceAmbiguity, ...], + current_evidence: HouseholdEvidence, + ) -> tuple[HouseholdEvidenceAmbiguity, ...]: + current_keys = {cls._ambiguity_key(item) for item in current} + merged = [ + item + for item in retained + if cls._ambiguity_key(item) not in current_keys + and not cls._ambiguity_resolved(item, current_evidence) + ] + merged.extend(current) + return tuple( + { + cls._ambiguity_key(item): item + for item in merged + }.values() + ) + + @staticmethod + def _ambiguity_key(ambiguity: HouseholdEvidenceAmbiguity): + return ambiguity.kind, ambiguity.field, ambiguity.person_indices + + @staticmethod + def _ambiguity_resolved( + ambiguity: HouseholdEvidenceAmbiguity, + current_evidence: HouseholdEvidence, + ) -> bool: + if ambiguity.kind is HouseholdEvidenceAmbiguityKind.ADULT_RELATIONSHIP: + return current_evidence.is_married is not None + if ambiguity.kind not in { + HouseholdEvidenceAmbiguityKind.INCOME_OWNER, + HouseholdEvidenceAmbiguityKind.INCOME_FREQUENCY, + }: + return False + if ambiguity.field not in { + "employment_income", + "self_employment_income", + "pension_income", + }: + return False + candidate_indices = ( + ambiguity.person_indices + if ambiguity.person_indices + else tuple(range(len(current_evidence.people))) + ) + return any( + index < len(current_evidence.people) + and getattr(current_evidence.people[index], ambiguity.field) is not None + for index in candidate_indices + ) + + @staticmethod + def _periodic_amount_facts( + label: str, + value: PeriodicAmount, + ) -> tuple[NumericalFact, ...]: + unit = { + AmountFrequency.WEEKLY: "GBP/week", + AmountFrequency.MONTHLY: "GBP/month", + AmountFrequency.ANNUAL: "GBP/year", + }[value.frequency] + facts = [NumericalFact(label=label, value=value.amount, unit=unit)] + annual_value = value.annual_value() + if value.frequency is not AmountFrequency.ANNUAL: + facts.append( + NumericalFact( + label=f"Annual {label.casefold()}", + value=annual_value, + unit="GBP/year", + ) + ) + return tuple(facts) + + async def _catalogue_labels(self, context): + labels = {} + for field, fallback in _CATALOGUE_FIELDS.items(): + result = await context.invoke_tool("get_variable", {"name": field}) + if not isinstance(result, SafeToolOutput): + return None + variable = result.root.get("variable") + if not isinstance(variable, dict): + return None + label = variable.get("label") + if not isinstance(label, str) or not label.strip(): + return None + labels[field] = label + return labels + + @staticmethod + def _person_subject(person: PersonEvidence, index: int) -> str: + if person.display_label: + return person.display_label[:1].upper() + person.display_label[1:] + if person.relationship_to_user == "self": + return "You" + del index + return "Another household member" + + +class HouseholdAnalysisInput(StrictModel): + description: str + year: int | None = None + referenced_household_id: str | None = Field( + default=None, + description=( + "Compatible retained HouseholdRef artifact identifier. Supply it when " + "the user clearly corrects, extends, or recalculates the same household; " + "omit it for a clearly separate household." + ), + ) + referenced_policy_scenario_id: str | None = None + reform_instruction: str | None = None + requested_outputs: tuple[str, ...] = Field( + default=(), + description=( + "Every household metric the user explicitly asks to calculate, using " + "their ordinary wording; for example, income tax, National Insurance, " + "Universal Credit, or net income." + ), + ) + start_new_invocation: bool = Field( + default=False, + description=( + "True only when the user clearly starts a separate household calculation " + "instead of answering the single pending household clarification." + ), + ) + + +class HouseholdAnalysisDraft(HouseholdAnalysisInput): + invocation_id: str | None = Field( + default=None, + validation_alias=AliasChoices("invocation_id", "resuming_invocation_id"), + ) + context_scope_id: str | None = None + context_revision: int | None = None + fact_requirements: tuple[FactRequirement, ...] = () + evidence: HouseholdEvidence = Field(default_factory=HouseholdEvidence) + invocation_defaults: HouseholdInvocationDefaults = Field( + default_factory=HouseholdInvocationDefaults + ) + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...] = () + pending_fields: tuple[str, ...] = () + authoritative_messages: tuple[str, ...] = () + unresolved_sterling_mentions: tuple[str, ...] = () + + +class HouseholdOutputIssue(StrictModel): + request: str + guidance: str + + +class HouseholdAnalysisOutput(StrictModel): + result: HouseholdResultRef + assumptions: tuple[HouseholdAssumption, ...] + year_source: InputSource + output_issues: tuple[HouseholdOutputIssue, ...] = () + narration_facts: tuple[NumericalFact, ...] = () + assumption_statements: tuple[str, ...] = () + narration_requirement: str = ( + "Report every calculated value in result.outputs and explain every item in " + "output_issues before offering optional follow-up analysis. Do not claim an " + "amount, entitlement, or ineligibility for a benefit absent from " + "result.outputs." + ) + narration_fallback: str + + +class HouseholdAnalysisCapability( + Capability[HouseholdAnalysisInput, HouseholdAnalysisOutput] +): + spec = CapabilitySpec( + identifier="household_analysis", + version="1", + description=( + "Assemble a described household, state every material default in plain " + "language, validate it, and calculate deterministic household impacts." + ), + required_use=( + "Must be invoked immediately for tax amounts, benefit amounts, benefit " + "entitlements, or policy impacts for a described or retained household, " + "even when household details are incomplete. The conversational model " + "must not ask its own household-input questions before this capability runs." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=HouseholdAnalysisInput, + output_model=HouseholdAnalysisOutput, + accepted_artifacts=( + ArtifactContract(artifact_type="household", schema_version="1"), + ArtifactContract(artifact_type="policy_scenario", schema_version="1"), + ), + produced_artifacts=( + ArtifactContract(artifact_type="household", schema_version="1"), + ArtifactContract(artifact_type="household_result", schema_version="1"), + ), + dependencies=( + CapabilityDependency( + capability_id="policy_reform", + artifact=ArtifactContract( + artifact_type="policy_scenario", + schema_version="1", + ), + ), + ), + tool_dependencies=( + "assemble_household_candidate", + "validate_household", + "run_household_simulation", + "search_variables", + "get_variable", + "extract_result_findings", + ), + ) + + def __init__( + self, + engine_fact_projector: HouseholdEngineFactProjector | None = None, + ) -> None: + self._engine_fact_projector = engine_fact_projector + self._input_completeness = HouseholdInputCompleteness() + + async def run(self, capability_input: HouseholdAnalysisInput, context): + waiting, selection_outcome = await self._select_waiting( + capability_input, + context, + ) + if selection_outcome is not None: + return selection_outcome + retained_draft = ( + HouseholdAnalysisDraft.model_validate(waiting.partial_input.model_dump()) + if waiting is not None + else None + ) + effective_input = self._merge_capability_input( + retained_draft, + capability_input, + ) + context_view = ( + HouseholdContextView(context.conversation_context) + if context.conversation_context is not None + else None + ) + pending_resolutions = ( + context_view.pending_fact_resolutions() + if context_view is not None + else () + ) + if pending_resolutions: + return NeedsInput( + prompt=self._clarification_prompt( + tuple(item.prompt for item in pending_resolutions) + ), + missing_fields=tuple( + f"fact_resolution:{item.proposal_id}" + for item in pending_resolutions + ), + partial_input=( + self._public_partial_input(retained_draft) + if retained_draft is not None + else effective_input.model_dump(mode="json", exclude_none=True) + ), + ) + existing_household = await self._find_household(effective_input, context) + scenario = await self._find_scenario(effective_input, context) + defaulted_to_current_policy = ( + scenario is None and effective_input.reform_instruction is None + ) + resolved_year = resolve_policy_year( + explicit_year=( + effective_input.year + if effective_input.year is not None + else context_view.policy_year() if context_view is not None else None + ), + referenced_year=( + scenario.year + if scenario is not None + else existing_household.year if existing_household is not None else None + ), + ) + requested_requests = self._effective_requested_output_requests( + effective_input.requested_outputs, + context_view=context_view, + current_user_message=context.current_user_message, + ) + requested, issues = await self._requested_outputs( + requested_requests, + context, + ) + requirements = self._calculation_requirements(requested) + context_evidence = context_view.evidence() if context_view is not None else None + if context_view is not None: + requirements = requirements.model_copy( + update={ + "context_scope_id": context_view.scope_id, + "household_entity_id": context_view.household_entity_id, + } + ) + assembly = await context.invoke_tool( + "assemble_household_candidate", + { + "description": ( + context.current_user_message or capability_input.description + ), + "requirements": requirements.model_dump(mode="json"), + "existing_values": ( + [value.model_dump(mode="json") for value in existing_household.values] + if existing_household is not None + else [] + ), + "retained_evidence": ( + context_evidence.model_dump(mode="json") + if context_evidence is not None + else retained_draft.evidence.model_dump(mode="json") + if retained_draft is not None + else {} + ), + "invocation_defaults": ( + retained_draft.invocation_defaults.model_dump(mode="json") + if retained_draft is not None + else {} + ), + "retained_ambiguities": ( + [ + item.model_dump(mode="json") + for item in retained_draft.ambiguities + ] + if retained_draft is not None + else [] + ), + "pending_fields": ( + list(retained_draft.pending_fields) + if retained_draft is not None + else [] + ), + }, + ) + if not isinstance(assembly, AssembleHouseholdOutput): + raise TypeError("Household assembly returned an incompatible output.") + if assembly.status is HouseholdAssemblyStatus.NEEDS_INPUT: + draft = HouseholdAnalysisDraft.model_validate( + { + **effective_input.model_dump(mode="json"), + "requested_outputs": requested, + "evidence": assembly.evidence.model_dump(mode="json"), + "invocation_defaults": assembly.invocation_defaults.model_dump( + mode="json" + ), + "authoritative_messages": self._authoritative_messages( + retained_draft, + context.current_user_message, + ), + "ambiguities": [ + item.model_dump(mode="json") + for item in assembly.ambiguities + ], + "pending_fields": list(assembly.missing_fields), + "fact_requirements": [ + item.model_dump(mode="json") + for item in assembly.fact_requirements + ], + } + ) + draft = await self._persist_draft(draft, waiting, context) + return NeedsInput( + prompt=self._clarification_prompt(assembly.questions), + missing_fields=assembly.missing_fields, + partial_input=self._public_partial_input(draft), + fact_requirements=assembly.fact_requirements, + capability_invocation=self._capability_invocation(draft), + ) + if assembly.status is HouseholdAssemblyStatus.FAILED or assembly.candidate is None: + return Failed( + safe_message=assembly.error or "Household assembly failed.", + error_code="household_assembly_contract", + ) + + draft = HouseholdAnalysisDraft.model_validate( + { + **effective_input.model_dump(mode="json"), + "requested_outputs": requested, + "evidence": assembly.evidence.model_dump(mode="json"), + "invocation_defaults": assembly.invocation_defaults.model_dump( + mode="json" + ), + "authoritative_messages": self._authoritative_messages( + retained_draft, + context.current_user_message, + ), + "ambiguities": [ + item.model_dump(mode="json") for item in assembly.ambiguities + ], + "pending_fields": [], + "fact_requirements": [], + } + ) + + unresolved_mentions = self._input_completeness.unresolved_mentions( + authoritative_messages=draft.authoritative_messages, + verified_amounts=tuple( + float(fact.value) + for fact in assembly.candidate.input_narration_facts + if fact.unit.startswith("GBP/") + ) + + assembly.invocation_defaults.sterling_amounts(), + excluded_texts=( + (effective_input.reform_instruction or ""), + ), + ) + if unresolved_mentions: + mention_text = self._natural_list( + tuple(item.text for item in unresolved_mentions) + ) + draft = draft.model_copy( + update={ + "pending_fields": tuple( + f"unresolved_sterling:{item.amount}" + for item in unresolved_mentions + ), + "unresolved_sterling_mentions": tuple( + item.text for item in unresolved_mentions + ), + } + ) + draft = await self._persist_draft(draft, waiting, context) + return NeedsInput( + prompt=( + f"I could not connect {mention_text} to a validated household " + "input. Please restate what each amount represents, who receives " + "or pays it, and whether it is weekly, monthly, or annual." + ), + missing_fields=tuple(draft.pending_fields), + partial_input=self._public_partial_input(draft), + ) + + if effective_input.reform_instruction: + scenario_outcome = await context.invoke_capability( + "policy_reform", + { + "instruction": effective_input.reform_instruction, + "year": resolved_year.year, + "referenced_policy_scenario_id": ( + scenario.artifact_id if scenario is not None else None + ), + }, + ) + if not isinstance(scenario_outcome, Completed) or not isinstance( + scenario_outcome.value, + PolicyReformOutput, + ): + return await self._forward_reform_outcome( + scenario_outcome, + draft, + waiting, + context, + ) + scenario = scenario_outcome.value.scenario + elif scenario is None: + scenario_outcome = await context.invoke_capability( + "policy_reform", + {"instruction": "current law", "year": resolved_year.year}, + ) + if not isinstance(scenario_outcome, Completed) or not isinstance( + scenario_outcome.value, + PolicyReformOutput, + ): + return await self._forward_reform_outcome( + scenario_outcome, + draft, + waiting, + context, + ) + scenario = scenario_outcome.value.scenario + + candidate = assembly.candidate + if ( + self._engine_fact_projector is not None + and context_view is not None + and context.conversation_context is not None + ): + candidate = self._merge_engine_facts( + candidate, + self._engine_fact_projector.project( + context.conversation_context, + scope_id=context_view.scope_id, + person_entity_ids=context_view.person_entity_ids, + household_entity_id=context_view.household_entity_id, + ), + ) + reform = { + change.parameter_path: change.value for change in scenario.verified_changes + } + simulation_input = { + "people": list(candidate.people), + "benunit": candidate.benunit, + "household": candidate.household, + "year": resolved_year.year, + "reform": reform or None, + "extra_variables": list(requested), + } + validation = await context.invoke_tool("validate_household", simulation_input) + if not isinstance(validation, SafeToolOutput): + raise TypeError("Household validation returned an incompatible output.") + if validation.root.get("valid") is not True: + missing_fields = self._validation_fields(validation.root) + draft = draft.model_copy(update={"pending_fields": missing_fields}) + draft = await self._persist_draft(draft, waiting, context) + return NeedsInput( + prompt=self._clarification_prompt( + (self._validation_prompt(validation.root),), + ), + missing_fields=missing_fields, + partial_input=self._public_partial_input(draft), + ) + if waiting is not None: + await context.remove_waiting(waiting.invocation_id) + household_ref = HouseholdRef( + provenance=self._provenance(context, "validated household"), + year=resolved_year.year, + household_revision=self._household_revision(candidate), + catalogue_version=scenario.catalogue_version, + calculation_engine_version=scenario.calculation_engine_version, + values=candidate.field_values, + context_scope_id=( + context.conversation_context.focus.scope_id + if context.conversation_context is not None + else None + ), + context_revision=( + context.conversation_context.revision + if context.conversation_context is not None + else None + ), + entity_positions=candidate.entity_positions, + ) + household_ref = await context.save_artifact(household_ref) + simulation = await context.invoke_tool("run_household_simulation", simulation_input) + if not isinstance(simulation, SafeToolOutput): + raise TypeError("Household simulation returned an incompatible output.") + if "error" in simulation.root: + return Failed( + safe_message="The deterministic household calculation failed.", + error_code="household_simulation_failed", + ) + outputs = self._extract_outputs(simulation.root, requested) + result = HouseholdResultRef( + provenance=self._provenance(context, "household simulation"), + year=resolved_year.year, + household_artifact_id=household_ref.artifact_id, + policy_scenario_artifact_id=scenario.artifact_id, + scenario_revision=scenario.scenario_revision, + calculation_engine_version=scenario.calculation_engine_version, + outputs=outputs, + context_scope_id=household_ref.context_scope_id, + context_revision=household_ref.context_revision, + ) + result = await context.save_artifact(result) + extracted = await context.invoke_tool( + "extract_result_findings", + {"outputs": [output.model_dump(mode="json") for output in outputs]}, + ) + if not isinstance(extracted, ExtractResultFindingsOutput): + raise TypeError("Household finding extraction returned an incompatible output.") + output_facts = tuple( + fact + for finding in extracted.findings + if finding.value is not None + for fact in self._output_narration_facts(finding) + ) + facts = ( + *candidate.input_narration_facts, + NumericalFact( + label="Policy year", + value=resolved_year.year, + unit="year", + ), + *output_facts, + ) + completed_assumptions = self._completed_assumptions( + assembly.assumptions, + year=resolved_year.year, + year_source=resolved_year.source, + defaulted_to_current_policy=defaulted_to_current_policy, + ) + return Completed( + value=HouseholdAnalysisOutput( + result=result, + assumptions=completed_assumptions, + year_source=resolved_year.source, + output_issues=issues, + narration_facts=facts, + assumption_statements=tuple( + assumption.plain_statement + for assumption in completed_assumptions + ), + narration_fallback=self._narration_fallback( + outputs, + completed_assumptions, + issues, + ), + ) + ) + + @staticmethod + def _merge_engine_facts( + candidate: HouseholdCandidate, + projection: HouseholdEngineInputs, + ) -> HouseholdCandidate: + """Add catalogue-backed facts while retaining assembled explicit values.""" + + people = [dict(item) for item in candidate.people] + positions = { + item.entity_id: int( + item.engine_position.removeprefix("people[").removesuffix("]") + ) + for item in candidate.entity_positions + if item.engine_position.startswith("people[") + and item.engine_position.endswith("]") + and item.engine_position.removeprefix("people[") + .removesuffix("]") + .isdigit() + } + for person in projection.people: + index = positions.get(person.entity_id) + if index is None or index >= len(people): + continue + for variable_name, value in person.values.items(): + people[index].setdefault(variable_name, value) + return candidate.model_copy( + update={ + "people": tuple(people), + "benunit": {**projection.benunit, **candidate.benunit}, + "household": {**projection.household, **candidate.household}, + } + ) + + @staticmethod + def _output_narration_facts(finding) -> tuple[NumericalFact, ...]: + facts = [ + NumericalFact( + label=finding.label, + value=finding.value, + unit=finding.unit, + ) + ] + if finding.unit == "GBP/year": + facts.extend( + ( + NumericalFact( + label=f"Monthly {finding.label.casefold()}", + value=finding.value / 12, + unit="GBP/month", + ), + NumericalFact( + label=f"Weekly {finding.label.casefold()}", + value=finding.value / 52, + unit="GBP/week", + ), + ) + ) + return tuple(facts) + + @staticmethod + def _narration_fallback(outputs, assumptions, issues) -> str: + result_lines = [] + for output in outputs: + if output.value is None: + result_lines.append(f"- {output.label}: unavailable") + elif output.unit == "GBP/year": + result_lines.append( + f"- {output.label}: £{output.value:,.2f} per year" + ) + else: + result_lines.append( + f"- {output.label}: {output.value:g} {output.unit}" + ) + if result_lines: + paragraphs = ["### Results\n\n" + "\n".join(result_lines)] + else: + paragraphs = ["The calculation did not return a supported household output."] + if assumptions: + statements = "\n".join( + f"- {item.plain_statement}" for item in assumptions + ) + paragraphs.append(f"### Assumptions used\n\n{statements}") + if issues: + issue_text = " ".join( + f"I could not calculate {item.request}: {item.guidance}" + for item in issues + ) + paragraphs.append(issue_text) + return "\n\n".join(paragraphs) + + @staticmethod + def _completed_assumptions( + household_assumptions, + *, + year, + year_source, + defaulted_to_current_policy, + ): + assumptions = list(household_assumptions) + if year_source is InputSource.SERVER_DEFAULT: + assumptions.insert( + 0, + HouseholdAssumption( + field_id="policy.year", + label="Policy year", + assumed_value=year, + plain_statement=f"Policy year: {year}.", + label_source="system_default", + ), + ) + if defaulted_to_current_policy: + assumptions.insert( + 1 if year_source is InputSource.SERVER_DEFAULT else 0, + HouseholdAssumption( + field_id="policy.scenario", + label="Policy scenario", + assumed_value="current_policy", + plain_statement="Policy scenario: current policy.", + label_source="system_default", + ), + ) + return tuple(assumptions) + + async def _forward_reform_outcome( + self, + outcome, + draft, + waiting, + context, + ): + if not isinstance(outcome, NeedsInput): + return outcome + draft = draft.model_copy(update={"pending_fields": ("reform_instruction",)}) + draft = await self._persist_draft(draft, waiting, context) + return NeedsInput( + prompt=outcome.prompt, + missing_fields=("reform_instruction",), + partial_input=self._public_partial_input( + draft, + exclude_defaults=True, + ), + ) + + @classmethod + async def _select_waiting(cls, capability_input, context): + waiting = await context.waiting_invocations(cls.spec.identifier) + if capability_input.start_new_invocation: + return None, None + active_scope_id = ( + context.conversation_context.focus.scope_id + if context.conversation_context is not None + else None + ) + compatible = tuple( + item + for item in waiting + if cls._waiting_scope_id(item) in {None, active_scope_id} + ) + pending_questions = ( + tuple( + question + for question in context.conversation_context.pending_questions + if question.capability_id == cls.spec.identifier + and question.capability_invocation is not None + and question.capability_invocation.context_scope_id + in {active_scope_id, None} + ) + if context.conversation_context is not None + else () + ) + compatible_by_id = {item.invocation_id: item for item in compatible} + linked = tuple( + compatible_by_id[question.capability_invocation.invocation_id] + for question in pending_questions + if question.capability_invocation is not None + and question.capability_invocation.invocation_id in compatible_by_id + ) + if len(linked) == 1: + return linked[0], None + if len(linked) > 1 and context.conversation_context is not None: + answered = tuple( + compatible_by_id[question.capability_invocation.invocation_id] + for question in pending_questions + if question.capability_invocation is not None + and question.capability_invocation.invocation_id in compatible_by_id + and ( + question.status is PendingQuestionStatus.ANSWER_RECEIVED + or cls._requirements_satisfied( + context.conversation_context, + question.requirements, + ) + ) + ) + if len(answered) == 1: + return answered[0], None + if not linked and len(compatible) == 1: + # Repair compatibility for version-one contexts that lost or never + # stored the pending-question link. + return compatible[0], None + candidates = linked or compatible + if len(candidates) > 1: + choices = "; ".join( + getattr(item.partial_input, "description", "household calculation") + for item in candidates + ) + return None, NeedsInput( + prompt=( + "More than one household calculation is waiting for information. " + f"Which one do you want to continue? {choices}" + ), + missing_fields=("pending_household_selection",), + partial_input={}, + ) + return None, None + + @staticmethod + def _waiting_scope_id(waiting) -> str | None: + return getattr(waiting.partial_input, "context_scope_id", None) + + @staticmethod + def _requirements_satisfied(context, requirements) -> bool: + if not requirements: + return False + for requirement in requirements: + subject_ids: tuple[str, ...] + if requirement.subject_entity_id is not None: + subject_ids = (requirement.subject_entity_id,) + elif requirement.subject_kind is not None: + subject_ids = tuple( + entity.entity_id + for entity in context.entities + if entity.kind is requirement.subject_kind + ) + else: + return False + matching = tuple( + fact + for subject_id in subject_ids + if ( + fact := context.active_fact( + requirement.fact_key, + subject_id, + requirement.scope_id, + ) + ) + is not None + ) + if not matching: + return False + if ( + not requirement.allow_explicit_absence + and all( + isinstance(fact.assertion, ExplicitAbsenceAssertion) + for fact in matching + ) + ): + return False + return True + + @staticmethod + def _merge_capability_input(retained_draft, current): + if retained_draft is None: + return current + return HouseholdAnalysisInput( + description=current.description, + year=current.year if current.year is not None else retained_draft.year, + referenced_household_id=( + current.referenced_household_id + if current.referenced_household_id is not None + else retained_draft.referenced_household_id + ), + referenced_policy_scenario_id=( + current.referenced_policy_scenario_id + if current.referenced_policy_scenario_id is not None + else retained_draft.referenced_policy_scenario_id + ), + reform_instruction=( + current.reform_instruction + if current.reform_instruction is not None + else retained_draft.reform_instruction + ), + requested_outputs=( + retained_draft.requested_outputs + ), + start_new_invocation=False, + ) + + @staticmethod + async def _persist_draft(draft, waiting, context): + invocation_id = ( + waiting.invocation_id + if waiting is not None + else context.capability_invocation_id + ) + stored_draft = draft.model_copy( + update={ + "invocation_id": invocation_id, + "context_scope_id": ( + context.conversation_context.focus.scope_id + if context.conversation_context is not None + else draft.context_scope_id + ), + "context_revision": ( + context.conversation_context.revision + if context.conversation_context is not None + else draft.context_revision + ), + "start_new_invocation": False, + } + ) + if context.conversation_context is not None: + stored_draft = stored_draft.model_copy( + update={ + "evidence": HouseholdEvidence(), + "ambiguities": (), + } + ) + if waiting is None: + await context.persist_waiting(stored_draft) + else: + await context.update_waiting(invocation_id, stored_draft) + return stored_draft + + @staticmethod + def _capability_invocation( + draft: HouseholdAnalysisDraft, + ) -> CapabilityInvocationReference | None: + if ( + draft.invocation_id is None + or draft.context_scope_id is None + or draft.context_revision is None + ): + return None + return CapabilityInvocationReference( + invocation_id=draft.invocation_id, + capability_id=HouseholdAnalysisCapability.spec.identifier, + capability_version=HouseholdAnalysisCapability.spec.version, + context_scope_id=draft.context_scope_id, + context_revision=draft.context_revision, + ) + + @staticmethod + def _public_partial_input(draft, *, exclude_defaults=False): + return draft.model_dump( + mode="json", + exclude={ + "invocation_id", + "context_scope_id", + "context_revision", + "fact_requirements", + "evidence", + "invocation_defaults", + "ambiguities", + "pending_fields", + "authoritative_messages", + "unresolved_sterling_mentions", + }, + exclude_none=True, + exclude_defaults=exclude_defaults, + ) + + @staticmethod + def _authoritative_messages(retained_draft, current_message): + retained = ( + retained_draft.authoritative_messages + if retained_draft is not None + else () + ) + if not current_message or current_message in retained: + return retained + return (*retained, current_message) + + @staticmethod + def _natural_list(items): + if len(items) == 1: + return items[0] + if len(items) == 2: + return f"{items[0]} and {items[1]}" + return f"{', '.join(items[:-1])}, and {items[-1]}" + + @staticmethod + def _clarification_prompt(questions): + return " ".join(dict.fromkeys(questions)) + + @staticmethod + async def _find_household(capability_input, context): + artifacts = await context.find_artifacts(HouseholdRef) + if capability_input.referenced_household_id is not None: + return next( + ( + item + for item in artifacts + if item.artifact_id == capability_input.referenced_household_id + ), + None, + ) + if context.conversation_context is None: + return None + scope_id = context.conversation_context.focus.scope_id + compatible = tuple( + item for item in artifacts if item.context_scope_id == scope_id + ) + if not compatible: + return None + return max( + compatible, + key=lambda item: (item.context_revision or -1, item.created_at), + ) + + @staticmethod + async def _find_scenario(capability_input, context): + if capability_input.referenced_policy_scenario_id is None: + return None + artifacts = await context.find_artifacts(PolicyScenarioRef) + return next( + ( + item + for item in artifacts + if item.artifact_id == capability_input.referenced_policy_scenario_id + ), + None, + ) + + @staticmethod + async def _requested_outputs(requests, context): + selected = [] + issues = [] + for request in requests: + normalized = " ".join(request.casefold().replace("_", " ").split()) + grouped_identifiers = _HOUSEHOLD_OUTPUT_GROUPS.get(normalized) + if grouped_identifiers is not None: + missing_group_member = False + for identifier in grouped_identifiers: + result = await context.invoke_tool( + "get_variable", + {"name": identifier}, + ) + variable = ( + result.root.get("variable") + if isinstance(result, SafeToolOutput) + else None + ) + if not ( + isinstance(variable, dict) + and variable.get("name") == identifier + ): + missing_group_member = True + issues.append( + HouseholdOutputIssue( + request=request, + guidance=( + "A required tax output is unavailable in the " + "authoritative variable catalogue." + ), + ) + ) + break + if missing_group_member: + continue + selected.extend( + identifier + for identifier in grouped_identifiers + if identifier not in selected + ) + continue + aliased_identifier = _HOUSEHOLD_OUTPUT_ALIASES.get(normalized) + if aliased_identifier is not None: + result = await context.invoke_tool( + "get_variable", + {"name": aliased_identifier}, + ) + variable = ( + result.root.get("variable") + if isinstance(result, SafeToolOutput) + else None + ) + if ( + isinstance(variable, dict) + and variable.get("name") == aliased_identifier + ): + if aliased_identifier not in selected: + selected.append(aliased_identifier) + continue + issues.append( + HouseholdOutputIssue( + request=request, + guidance=( + "The mapped household output is unavailable in the " + "authoritative variable catalogue." + ), + ) + ) + continue + result = await context.invoke_tool( + "search_variables", + {"query": request, "limit": 10}, + ) + rows = ( + result.root.get("variables") + if isinstance(result, SafeToolOutput) + else None + ) + matches = [ + row + for row in rows or [] + if isinstance(row, dict) and isinstance(row.get("name"), str) + ] + if not matches: + issues.append( + HouseholdOutputIssue( + request=request, + guidance="No authoritative household output matched this request.", + ) + ) + continue + exact = [ + row + for row in matches + if normalized + in { + str(row["name"]).casefold().replace("_", " "), + str(row.get("label", "")).casefold(), + } + ] + chosen = exact[0] if len(exact) == 1 else None + if chosen is None: + issues.append( + HouseholdOutputIssue( + request=request, + guidance=( + "No unambiguous authoritative household output exactly " + "matched this request; please clarify." + ), + ) + ) + continue + if chosen["name"] not in selected: + selected.append(chosen["name"]) + if not selected: + selected.append("household_net_income") + return tuple(selected), tuple(issues) + + @staticmethod + def _effective_requested_output_requests( + explicit_requests: tuple[str, ...], + *, + context_view: HouseholdContextView | None, + current_user_message: str | None, + ) -> tuple[str, ...]: + combined = list(dict.fromkeys(explicit_requests)) + detected = HouseholdAnalysisCapability._requested_outputs_from_message( + current_user_message + ) + combined.extend(item for item in detected if item not in combined) + if combined: + return tuple(combined) + if context_view is not None: + retained = context_view.requested_outputs() + if retained: + return retained + return () + + @staticmethod + def _requested_outputs_from_message( + current_user_message: str | None, + ) -> tuple[str, ...]: + if not current_user_message: + return () + normalized = " ".join( + current_user_message.casefold().replace("_", " ").split() + ) + matched: list[str] = [] + specific_tax_requested = False + for label, output_id in sorted( + _HOUSEHOLD_OUTPUT_ALIASES.items(), + key=lambda item: len(item[0]), + reverse=True, + ): + if re.search(rf"\b{re.escape(label)}\b", normalized) is None: + continue + if output_id not in matched: + matched.append(output_id) + if output_id in _TAX_ONLY_HOUSEHOLD_OUTPUTS: + specific_tax_requested = True + + if ( + not specific_tax_requested + and re.search(r"\btaxes?\b", normalized) is not None + ): + matched.extend( + output_id + for output_id in ( + "income_tax", + "national_insurance", + "household_tax", + ) + if output_id not in matched + ) + return tuple(matched) + + @staticmethod + def _calculation_requirements( + requested_output_ids: tuple[str, ...], + ) -> HouseholdCalculationRequirements: + tax_only = bool(requested_output_ids) and set(requested_output_ids).issubset( + _TAX_ONLY_HOUSEHOLD_OUTPUTS + ) + return HouseholdCalculationRequirements( + requested_output_ids=requested_output_ids, + require_housing_costs=not tax_only, + ) + + @classmethod + def _extract_outputs(cls, payload, requested): + reform_applied = payload.get("reform_applied") is True + outputs = [] + for variable in requested: + if reform_applied: + baseline = cls._find_number(payload.get("baseline"), variable) + reform = cls._find_number(payload.get("reform"), variable) + for metric, value in ( + ("baseline", baseline), + ("reform", reform), + ( + "change", + reform - baseline + if reform is not None and baseline is not None + else None, + ), + ): + outputs.append( + AggregateValue( + output_id=variable, + metric_id=metric, + label=variable.replace("_", " ").title(), + value=value, + unit="GBP/year", + dimensions=( + AggregateDimension(name="scenario", value=metric), + ), + ) + ) + else: + outputs.append( + AggregateValue( + output_id=variable, + metric_id="current_law", + label=variable.replace("_", " ").title(), + value=cls._find_number(payload, variable), + unit="GBP/year", + ) + ) + return tuple(outputs) + + @classmethod + def _find_number(cls, value, key): + if isinstance(value, dict): + if key in value and isinstance(value[key], (int, float)): + return value[key] + for nested in value.values(): + found = cls._find_number(nested, key) + if found is not None: + return found + if isinstance(value, list): + for nested in value: + found = cls._find_number(nested, key) + if found is not None: + return found + return None + + @staticmethod + def _validation_prompt(payload): + errors = payload.get("errors") + if isinstance(errors, list) and errors: + messages = tuple( + error["message"] + for error in errors + if isinstance(error, dict) + and isinstance(error.get("message"), str) + and error["message"].strip() + ) + if messages: + return " ".join(dict.fromkeys(messages)) + return "Please correct the household details before calculation." + + @staticmethod + def _validation_fields(payload): + errors = payload.get("errors") + if not isinstance(errors, list): + return ("household",) + fields = tuple( + error["path"] + for error in errors + if isinstance(error, dict) + and isinstance(error.get("path"), str) + and error["path"].strip() + ) + return tuple(dict.fromkeys(fields)) or ("household",) + + @staticmethod + def _household_revision(candidate): + payload = candidate.model_dump_json() + return hashlib.sha256(payload.encode()).hexdigest()[:16] + + @classmethod + def _provenance(cls, context, source): + return ArtifactProvenance( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + capability_id=cls.spec.identifier, + capability_version=cls.spec.version, + invocation_id=context.capability_invocation_id, + sources=(source,), + ) diff --git a/backend/capabilities/household_input.py b/backend/capabilities/household_input.py new file mode 100644 index 00000000..5f074267 --- /dev/null +++ b/backend/capabilities/household_input.py @@ -0,0 +1,532 @@ +"""Typed household evidence and deterministic clarification resolution.""" + +from __future__ import annotations + +import re +from decimal import Decimal +from enum import Enum +from typing import Literal, Protocol + +from pydantic import BaseModel, ConfigDict, Field, model_validator +from conversation_context.models import EntityKind, FactRequirement + + +PersonEvidenceField = Literal[ + "age", + "employment_income", + "self_employment_income", + "pension_income", +] +HouseholdEvidenceField = Literal[ + "is_married", + "has_children", + "childcare_expenses", + "rent", + "council_tax", + "country", +] + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class AmountFrequency(str, Enum): + WEEKLY = "weekly" + MONTHLY = "monthly" + ANNUAL = "annual" + + +class PeriodicAmount(StrictModel): + amount: float + frequency: AmountFrequency + + def annual_value(self) -> float: + multiplier = { + AmountFrequency.WEEKLY: 52, + AmountFrequency.MONTHLY: 12, + AmountFrequency.ANNUAL: 1, + }[self.frequency] + return self.amount * multiplier + + +class PersonEvidence(StrictModel): + entity_id: str | None = None + display_label: str | None = None + relationship_to_user: str | None = None + age: int | None = None + employment_income: PeriodicAmount | None = None + self_employment_income: PeriodicAmount | None = None + pension_income: PeriodicAmount | None = None + sources: dict[PersonEvidenceField, Literal["user", "artifact", "default"]] = ( + Field(default_factory=dict) + ) + + +class HouseholdEvidence(StrictModel): + people: tuple[PersonEvidence, ...] = Field( + default=(), + description="Every described adult and child, kept in conversational order.", + ) + has_children: bool | None = Field( + default=None, + description=( + "True when the current wording says the household has children; false " + "when it explicitly says there are no children or describes all members " + "as adults." + ), + ) + is_married: bool | None = Field( + default=None, + description=( + "True when multiple adults are described as a couple, married, or in a " + "civil partnership; false when they are explicitly unrelated adults or " + "separate benefit units. Null when their relationship is not supplied." + ), + ) + childcare_expenses: PeriodicAmount | None = None + rent: PeriodicAmount | None = None + council_tax: PeriodicAmount | None = None + country: Literal["ENGLAND", "SCOTLAND", "WALES", "NORTHERN_IRELAND"] | None = None + sources: dict[ + HouseholdEvidenceField, + Literal["user", "artifact", "default"], + ] = Field( + default_factory=dict + ) + + +class HouseholdInvocationDefaults(StrictModel): + """Invocation-owned values that are reported as defaults, never user facts.""" + + evidence: HouseholdEvidence = Field(default_factory=HouseholdEvidence) + + @model_validator(mode="after") + def only_contains_defaults(self) -> "HouseholdInvocationDefaults": + for person in self.evidence.people: + if any(source != "default" for source in person.sources.values()): + raise ValueError("invocation defaults cannot contain accepted facts") + if any(source != "default" for source in self.evidence.sources.values()): + raise ValueError("invocation defaults cannot contain accepted facts") + return self + + def sterling_amounts(self) -> tuple[float, ...]: + """Return monetary values introduced by documented invocation defaults.""" + + amounts: list[float] = [] + for person in self.evidence.people: + for field in ( + "employment_income", + "self_employment_income", + "pension_income", + ): + value = getattr(person, field) + if value is not None: + amounts.extend((value.amount, value.annual_value())) + for field in ("childcare_expenses", "rent", "council_tax"): + value = getattr(self.evidence, field) + if value is not None: + amounts.extend((value.amount, value.annual_value())) + return tuple(dict.fromkeys(amounts)) + + +class SterlingMention(StrictModel): + """One exact monetary mention retained from an authoritative user message.""" + + text: str + amount: Decimal + message: str + + +class HouseholdInputCompleteness: + """Reject household calculations that silently lose monetary input mentions.""" + + _sterling_pattern = re.compile( + r"£\s*(\d{1,3}(?:,\d{3})*(?:\.\d+)?|\d+(?:\.\d+)?)" + ) + + def unresolved_mentions( + self, + *, + authoritative_messages: tuple[str, ...], + verified_amounts: tuple[float, ...], + excluded_texts: tuple[str, ...] = (), + ) -> tuple[SterlingMention, ...]: + verified = {Decimal(str(value)) for value in verified_amounts} + excluded = { + Decimal(raw.replace(",", "")) + for text in excluded_texts + for raw in self._sterling_pattern.findall(text) + } + unresolved: list[SterlingMention] = [] + for message in authoritative_messages: + for match in self._sterling_pattern.finditer(message): + amount = Decimal(match.group(1).replace(",", "")) + if amount in verified or amount in excluded: + continue + mention = SterlingMention( + text=match.group(0), + amount=amount, + message=message, + ) + if all( + item.amount != mention.amount or item.message != mention.message + for item in unresolved + ): + unresolved.append(mention) + return tuple(unresolved) + + +class HouseholdAssemblerUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class HouseholdEvidenceAmbiguityKind(str, Enum): + INCOME_OWNER = "income_owner" + INCOME_FREQUENCY = "income_frequency" + ADULT_RELATIONSHIP = "adult_relationship" + CONFLICTING_EVIDENCE = "conflicting_evidence" + + +class HouseholdEvidenceAmbiguity(StrictModel): + kind: HouseholdEvidenceAmbiguityKind + field: Literal[ + "employment_income", + "self_employment_income", + "pension_income", + "household_structure", + ] | None = None + person_indices: tuple[int, ...] = () + amount: float | None = None + frequency: AmountFrequency | None = None + + +class HouseholdEvidenceResult(StrictModel): + evidence: HouseholdEvidence + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...] = () + usage: HouseholdAssemblerUsage = Field(default_factory=HouseholdAssemblerUsage) + + +class HouseholdEvidenceAssembler(Protocol): + async def assemble( + self, + *, + description: str, + retained_evidence: HouseholdEvidence, + retained_ambiguities: tuple[HouseholdEvidenceAmbiguity, ...], + ) -> HouseholdEvidenceResult: ... + + +class HouseholdResolution(StrictModel): + missing_fields: tuple[str, ...] = () + questions: tuple[str, ...] = () + fact_requirements: tuple[FactRequirement, ...] = () + + +class HouseholdCalculationRequirements(StrictModel): + """Resolved input requirements for the outputs requested in one calculation.""" + + requested_output_ids: tuple[str, ...] = () + require_housing_costs: bool = True + context_scope_id: str = "scope:primary-household" + household_entity_id: str = "household:primary" + + +class HouseholdInputResolver: + """Own deterministic household requirements and comprehensive questions.""" + + _first_person_pattern = re.compile(r"\b(?:i|i'm|i am|my|me)\b", re.IGNORECASE) + _income_word_pattern = re.compile( + r"\b(?:income|earn|earns|earning|salary|wage|wages)\b", + re.IGNORECASE, + ) + _specified_income_kind_pattern = re.compile( + r"\b(?:self[- ]?employ|pension|dividend|rental|capital gain)", + re.IGNORECASE, + ) + _sterling_amount_pattern = re.compile( + r"£\s*(\d{1,3}(?:,\d{3})*(?:\.\d+)?|\d+(?:\.\d+)?)" + ) + _tax_only_outputs = frozenset( + {"income_tax", "national_insurance", "household_tax"} + ) + + def apply_documented_defaults( + self, + evidence: HouseholdEvidence, + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...], + *, + current_user_message: str, + requirements: HouseholdCalculationRequirements, + ) -> tuple[HouseholdEvidence, tuple[HouseholdEvidenceAmbiguity, ...]]: + """Apply narrow, reportable defaults before deciding what remains missing.""" + + output_ids = set(requirements.requested_output_ids) + if not output_ids or not output_ids.issubset(self._tax_only_outputs): + return evidence, ambiguities + if not self._first_person_pattern.search(current_user_message): + return evidence, ambiguities + if not self._income_word_pattern.search(current_user_message): + return evidence, ambiguities + if self._specified_income_kind_pattern.search(current_user_message): + return evidence, ambiguities + matches = self._sterling_amount_pattern.findall(current_user_message) + if len(matches) != 1 or len(evidence.people) > 1: + return evidence, ambiguities + + people = list(evidence.people or (PersonEvidence(),)) + person = people[0] + if any( + getattr(person, field) is not None + for field in ( + "employment_income", + "self_employment_income", + "pension_income", + ) + ): + return evidence, ambiguities + amount = float(matches[0].replace(",", "")) + sources = dict(person.sources) + sources["employment_income"] = "default" + people[0] = person.model_copy( + update={ + "employment_income": PeriodicAmount( + amount=amount, + frequency=AmountFrequency.ANNUAL, + ), + "sources": sources, + } + ) + remaining_ambiguities = tuple( + ambiguity + for ambiguity in ambiguities + if ambiguity.kind + not in { + HouseholdEvidenceAmbiguityKind.INCOME_OWNER, + HouseholdEvidenceAmbiguityKind.INCOME_FREQUENCY, + } + ) + return ( + evidence.model_copy(update={"people": tuple(people)}), + remaining_ambiguities, + ) + + def resolve( + self, + evidence: HouseholdEvidence, + ambiguities: tuple[HouseholdEvidenceAmbiguity, ...] = (), + requirements: HouseholdCalculationRequirements | None = None, + ) -> HouseholdResolution: + requirements = requirements or HouseholdCalculationRequirements() + missing_fields: list[str] = [] + questions: list[str] = [] + + if not evidence.people: + missing_fields.append("people") + questions.append( + "Who lives in the household? Please give each person's age, say how " + "multiple adults are related (including whether they form a couple or " + "civil partnership), and identify which person receives each income " + "amount and whether it is weekly, monthly, or annual." + ) + else: + missing_age_indices = tuple( + index + 1 + for index, person in enumerate(evidence.people) + if person.age is None + ) + if missing_age_indices: + missing_fields.extend( + f"people[{index - 1}].age" for index in missing_age_indices + ) + people_text = self._natural_list( + tuple( + self._person_label(evidence.people[index - 1], index) + for index in missing_age_indices + ) + ) + relationship_text = ( + " If more than one is an adult, also say whether the adults form " + "one couple or civil partnership." + if len(evidence.people) > 1 and evidence.is_married is None + else "" + ) + if relationship_text: + missing_fields.append("benunit.is_married_if_multiple_adults") + if len(evidence.people) == 1: + age_question = "What age should I use for this calculation?" + elif len(missing_age_indices) == 1 and people_text == "you": + age_question = "What is your age?" + else: + age_question = ( + f"What is the age of {people_text}?" + if len(missing_age_indices) == 1 + else f"What ages should I use for {people_text}?" + ) + questions.append(f"{age_question}{relationship_text}") + + known_adult_count = sum( + 1 + for person in evidence.people + if person.age is not None and person.age >= 16 + ) + all_ages_known = not missing_age_indices + if all_ages_known and known_adult_count == 0: + missing_fields.append("adult") + questions.append( + "The household needs at least one adult. Which person is an adult, " + "and what is their age?" + ) + if ( + known_adult_count > 1 + and evidence.is_married is None + and not missing_age_indices + ): + missing_fields.append("benunit.is_married") + questions.append("Do the adults form one couple or civil partnership?") + if evidence.has_children is True and not any( + person.age is not None and person.age < 16 + for person in evidence.people + ): + missing_fields.append("children.ages") + questions.append("What are the ages of the household's children?") + + for ambiguity in ambiguities: + if ambiguity.kind is HouseholdEvidenceAmbiguityKind.INCOME_OWNER: + missing_fields.append("income.owner") + questions.append("Which person receives each income amount you gave?") + elif ambiguity.kind is HouseholdEvidenceAmbiguityKind.INCOME_FREQUENCY: + missing_fields.append("income.frequency") + questions.append( + "For each income amount, is it weekly, monthly, or annual?" + ) + elif ambiguity.kind is HouseholdEvidenceAmbiguityKind.ADULT_RELATIONSHIP: + missing_fields.append("benunit.is_married") + questions.append("Do the adults form one couple or civil partnership?") + elif ambiguity.kind is HouseholdEvidenceAmbiguityKind.CONFLICTING_EVIDENCE: + missing_fields.append("conflicting_evidence") + questions.append( + "Which of the conflicting household details should I use?" + ) + + if requirements.require_housing_costs: + missing_housing_fields: list[str] = [] + if evidence.rent is None: + missing_fields.append("household.rent") + missing_housing_fields.append("rent") + if evidence.council_tax is None: + missing_fields.append("household.council_tax") + missing_housing_fields.append("Council Tax") + if len(missing_housing_fields) == 2: + questions.append( + "Does the household pay rent or Council Tax? For each that applies, " + "please give the amount and say whether it is weekly, monthly, or " + "annual; otherwise say that it does not apply." + ) + elif missing_housing_fields: + cost = missing_housing_fields[0] + questions.append( + f"Does the household pay {cost}? If so, how much, and is that weekly, " + "monthly, or annual? Otherwise say that it does not apply." + ) + + return HouseholdResolution( + missing_fields=tuple(dict.fromkeys(missing_fields)), + questions=tuple(dict.fromkeys(questions)), + fact_requirements=self._fact_requirements( + tuple(dict.fromkeys(missing_fields)), + evidence, + requirements, + ), + ) + + @staticmethod + def _fact_requirements( + missing_fields: tuple[str, ...], + evidence: HouseholdEvidence, + requirements: HouseholdCalculationRequirements, + ) -> tuple[FactRequirement, ...]: + result: list[FactRequirement] = [] + for field in missing_fields: + fact_key: str + subject_id: str | None = requirements.household_entity_id + subject_kind: EntityKind | None = None + expected = "boolean" + allow_absence = False + if field == "people": + fact_key = "household.members" + expected = "entity_references" + elif field.startswith("people[") and field.endswith("].age"): + try: + index = int(field.removeprefix("people[").split("]", 1)[0]) + except ValueError: + continue + if index >= len(evidence.people): + continue + fact_key = "person.age" + subject_id = evidence.people[index].entity_id + subject_kind = EntityKind.PERSON if subject_id is None else None + expected = "integer" + elif field in {"adult", "children.ages"}: + fact_key = "person.age" + subject_id = None + subject_kind = EntityKind.PERSON + expected = "integer" + elif field.startswith("benunit.is_married"): + fact_key = "household.is_married" + elif field == "household.rent": + fact_key = "household.rent" + expected = "money" + allow_absence = True + elif field == "household.council_tax": + fact_key = "household.council_tax" + expected = "money" + allow_absence = True + elif field in {"income.owner", "income.frequency"}: + fact_key = "person.employment_income" + subject_id = None + subject_kind = EntityKind.PERSON + expected = "money" + elif field == "conflicting_evidence": + continue + else: + continue + result.append( + FactRequirement( + requirement_id=f"household:{field}", + fact_key=fact_key, + subject_entity_id=subject_id, + subject_kind=subject_kind, + scope_id=requirements.context_scope_id, + expected_value_kind=expected, + allow_explicit_absence=allow_absence, + reason=f"Required to resolve household input {field}.", + ) + ) + return tuple( + { + (item.fact_key, item.subject_entity_id, item.requirement_id): item + for item in result + }.values() + ) + + @staticmethod + def _natural_list(values: tuple[str, ...]) -> str: + if len(values) == 1: + return values[0] + if len(values) == 2: + return f"{values[0]} and {values[1]}" + return f"{', '.join(values[:-1])}, and {values[-1]}" + + @staticmethod + def _person_label(person: PersonEvidence, position: int) -> str: + if person.display_label: + return person.display_label + if person.relationship_to_user == "self": + return "you" + if person.relationship_to_user: + return f"your {person.relationship_to_user.replace('_', ' ')}" + return "the other person" if position == 2 else f"person {position}" diff --git a/backend/capabilities/input_resolution.py b/backend/capabilities/input_resolution.py new file mode 100644 index 00000000..f4554a92 --- /dev/null +++ b/backend/capabilities/input_resolution.py @@ -0,0 +1,41 @@ +"""Small deterministic input-precedence helpers shared by capabilities.""" + +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, ConfigDict + + +DEFAULT_POLICY_YEAR = 2026 + + +class InputSource(str, Enum): + CURRENT_REQUEST = "current_request" + REFERENCED_ARTIFACT = "referenced_artifact" + SERVER_DEFAULT = "server_default" + + +class ResolvedYear(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + year: int + source: InputSource + + +def resolve_policy_year( + *, + explicit_year: int | None, + referenced_year: int | None, +) -> ResolvedYear: + if explicit_year is not None: + return ResolvedYear(year=explicit_year, source=InputSource.CURRENT_REQUEST) + if referenced_year is not None: + return ResolvedYear( + year=referenced_year, + source=InputSource.REFERENCED_ARTIFACT, + ) + return ResolvedYear( + year=DEFAULT_POLICY_YEAR, + source=InputSource.SERVER_DEFAULT, + ) diff --git a/backend/capabilities/policy_information.py b/backend/capabilities/policy_information.py new file mode 100644 index 00000000..d729ba8d --- /dev/null +++ b/backend/capabilities/policy_information.py @@ -0,0 +1,298 @@ +"""Authoritative policy catalogue and calculation-method capability.""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, ConfigDict, JsonValue + +from capabilities.artifacts import PolicyScenarioRef +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilitySpec, + Completed, + Unsupported, +) +from capabilities.input_resolution import InputSource, resolve_policy_year +from tools.analysis_support import NumericalFact +from tools.contracts import CallerType, Visibility +from tools.typed_models import SafeToolOutput + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class PolicyInformationInput(StrictModel): + question: str + search_terms: tuple[str, ...] = () + year: int | None = None + referenced_policy_scenario_id: str | None = None + parameter_path: str | None = None + variable_name: str | None = None + + +class PolicyCatalogueFact(StrictModel): + kind: Literal["parameter", "variable"] + identifier: str + label: str | None = None + description: str | None = None + unit: str | None = None + entity: str | None = None + definition_period: str | None = None + quantity_type: str | None = None + reference: JsonValue | None = None + defined_for: str | None = None + min_value: int | float | None = None + max_value: int | float | None = None + is_period_size_independent: bool | None = None + metadata: JsonValue | None = None + value: JsonValue | None = None + source_tool: str + + +class PolicyInformationOutput(StrictModel): + year: int + year_source: InputSource + facts: tuple[PolicyCatalogueFact, ...] + source_context: tuple[str, ...] + limitations: tuple[str, ...] = () + narration_facts: tuple[NumericalFact, ...] = () + + +class PolicyInformationCapability( + Capability[PolicyInformationInput, PolicyInformationOutput] +): + spec = CapabilitySpec( + identifier="policy_information", + version="1", + description=( + "Retrieve authoritative UK policy parameters, variables, scope, and " + "calculation-method metadata from ordinary language or exact identifiers." + ), + required_use=( + "Required for questions about how a government policy value is formulated, " + "scoped, or calculated." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=PolicyInformationInput, + output_model=PolicyInformationOutput, + accepted_artifacts=( + ArtifactContract(artifact_type="policy_scenario", schema_version="1"), + ), + tool_dependencies=( + "search_parameters", + "get_parameter", + "search_variables", + "get_variable", + ), + ) + + async def run(self, capability_input: PolicyInformationInput, context): + referenced_year = await self._referenced_year(capability_input, context) + resolved_year = resolve_policy_year( + explicit_year=capability_input.year, + referenced_year=referenced_year, + ) + facts: list[PolicyCatalogueFact] = [] + sources: list[str] = [] + + if capability_input.parameter_path: + result = await context.invoke_tool( + "get_parameter", + { + "path": capability_input.parameter_path, + "year": resolved_year.year, + }, + ) + facts.extend(self._parameter_facts(result, "get_parameter")) + sources.append("get_parameter") + if capability_input.variable_name: + result = await context.invoke_tool( + "get_variable", + {"name": capability_input.variable_name}, + ) + facts.extend(self._variable_facts(result, "get_variable")) + sources.append("get_variable") + + for query in self._queries(capability_input): + if not capability_input.parameter_path: + result = await context.invoke_tool( + "search_parameters", + {"query": query, "limit": 10}, + ) + facts.extend(self._parameter_facts(result, "search_parameters")) + sources.append("search_parameters") + if not capability_input.variable_name: + result = await context.invoke_tool( + "search_variables", + {"query": query, "limit": 10}, + ) + facts.extend(self._variable_facts(result, "search_variables")) + sources.append("search_variables") + if facts: + break + + facts = list( + { + (fact.kind, fact.identifier): fact + for fact in facts + }.values() + ) + if not facts: + return Unsupported( + reason=( + "The authoritative policy catalogue did not return enough " + "information to support this explanation." + ) + ) + numeric = tuple( + NumericalFact( + label=fact.label or fact.identifier, + value=fact.value, + unit=fact.unit or "value", + ) + for fact in facts + if fact.kind == "parameter" + and fact.source_tool == "get_parameter" + and isinstance(fact.value, (int, float)) + and not isinstance(fact.value, bool) + ) + return Completed( + value=PolicyInformationOutput( + year=resolved_year.year, + year_source=resolved_year.source, + facts=tuple(facts), + source_context=tuple(sources), + narration_facts=numeric, + ) + ) + + @staticmethod + async def _referenced_year(capability_input, context) -> int | None: + reference = capability_input.referenced_policy_scenario_id + if reference is None: + return None + scenarios = await context.find_artifacts(PolicyScenarioRef) + scenario = next( + (item for item in scenarios if item.artifact_id == reference), + None, + ) + return scenario.year if scenario is not None else None + + @staticmethod + def _payload(result, expected: str) -> dict[str, JsonValue]: + if not isinstance(result, SafeToolOutput): + raise TypeError(f"{expected} returned an incompatible result.") + return result.root + + @staticmethod + def _queries(capability_input: PolicyInformationInput) -> tuple[str, ...]: + if capability_input.search_terms: + return tuple(dict.fromkeys(capability_input.search_terms))[:4] + tokens = [ + token.strip(".,?!:;()[]{}\"'") + for token in capability_input.question.casefold().split() + ] + stop = { + "how", + "is", + "are", + "a", + "an", + "the", + "for", + "of", + "to", + "value", + "amount", + "calculated", + "determined", + } + significant = [token for token in tokens if token and token not in stop] + phrases = [capability_input.question] + phrases.extend( + " ".join(significant[index : index + size]) + for size in (3, 2) + for index in range(max(0, len(significant) - size + 1)) + ) + return tuple(dict.fromkeys(phrases))[:4] + + @classmethod + def _parameter_facts(cls, result, source) -> list[PolicyCatalogueFact]: + payload = cls._payload(result, source) + rows = payload.get("parameters") + if rows is None and isinstance(payload.get("parameter"), dict): + rows = [payload["parameter"]] + if not isinstance(rows, list): + return [] + return [ + PolicyCatalogueFact( + kind="parameter", + identifier=row["path"], + label=row.get("label"), + description=row.get("description"), + unit=row.get("unit"), + value=row.get("value"), + source_tool=source, + ) + for row in rows + if isinstance(row, dict) and isinstance(row.get("path"), str) + ] + + @classmethod + def _variable_facts(cls, result, source) -> list[PolicyCatalogueFact]: + payload = cls._payload(result, source) + rows = payload.get("variables") + if rows is None and isinstance(payload.get("variable"), dict): + rows = [payload["variable"]] + if not isinstance(rows, list): + return [] + facts: list[PolicyCatalogueFact] = [] + for row in rows: + if not isinstance(row, dict): + continue + identifier = row.get("name") + if not isinstance(identifier, str): + continue + facts.append( + PolicyCatalogueFact( + kind="variable", + identifier=identifier, + label=cls._optional_string(row.get("label")), + description=cls._optional_string(row.get("description")), + unit=cls._optional_string(row.get("unit")), + entity=cls._optional_string(row.get("entity")), + definition_period=cls._optional_string( + row.get("definition_period") + ), + quantity_type=cls._optional_string(row.get("quantity_type")), + reference=row.get("reference"), + defined_for=cls._optional_string(row.get("defined_for")), + min_value=cls._optional_number(row.get("min_value")), + max_value=cls._optional_number(row.get("max_value")), + is_period_size_independent=cls._optional_bool( + row.get("is_period_size_independent") + ), + metadata=row.get("metadata"), + value=row.get("default_value"), + source_tool=source, + ) + ) + return facts + + @staticmethod + def _optional_string(value: JsonValue | None) -> str | None: + return value if isinstance(value, str) else None + + @staticmethod + def _optional_number(value: JsonValue | None) -> int | float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return value + + @staticmethod + def _optional_bool(value: JsonValue | None) -> bool | None: + return value if isinstance(value, bool) else None diff --git a/backend/capabilities/policy_reform.py b/backend/capabilities/policy_reform.py new file mode 100644 index 00000000..72cb194a --- /dev/null +++ b/backend/capabilities/policy_reform.py @@ -0,0 +1,584 @@ +"""Catalogue-constrained reform resolution and verified scenario creation.""" + +from __future__ import annotations + +import hashlib +import json +from enum import Enum +from importlib.metadata import PackageNotFoundError, version +from typing import Literal, Protocol + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from capabilities.artifacts import ( + ArtifactProvenance, + PolicyChange, + PolicyScenarioRef, + Scalar, +) +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilitySpec, + Completed, + Failed, + NeedsInput, + Unsupported, +) +from capabilities.input_resolution import InputSource, resolve_policy_year +from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_async_client +from tools.contracts import CallerType, Tool, ToolCallContext, ToolSpec, Visibility +from tools.typed_models import SafeToolOutput + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class ReformResolutionKind(str, Enum): + RESOLVED = "resolved" + NEEDS_CLARIFICATION = "needs_clarification" + NO_REFORM = "no_reform" + UNSUPPORTED = "unsupported" + FAILED = "failed" + + +class ReformMeaning(StrictModel): + target: str + operation: Literal["set", "increase", "decrease", "abolish"] + value: Scalar + unit: str | None + effective_date: str | None + population: str + jurisdiction: str + + +class ResolverUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class ReformResolutionDecision(StrictModel): + outcome: ReformResolutionKind + summary: str + clarification: str | None = None + reform: dict[str, Scalar] = Field(default_factory=dict) + meaning: ReformMeaning | None = None + usage: ResolverUsage = Field(default_factory=ResolverUsage) + + +class ResolveReformInput(StrictModel): + instruction: str + year: int + existing_scenario_id: str | None = None + + +class ReformCatalogueBinding(StrictModel): + parameter_path: str + label: str + + +class ResolveReformOutput(StrictModel): + outcome: ReformResolutionKind + year: int + summary: str + clarification: str | None = None + changes: tuple[PolicyChange, ...] = () + bindings: tuple[ReformCatalogueBinding, ...] = () + catalogue_version: str + calculation_engine_version: str + + +class ReformCandidateResolver(Protocol): + async def resolve( + self, + *, + instruction: str, + year: int, + candidates: tuple[dict[str, JsonValue], ...], + ) -> ReformResolutionDecision: ... + + async def correct_representation( + self, + *, + instruction: str, + year: int, + candidates: tuple[dict[str, JsonValue], ...], + previous: ReformResolutionDecision, + validation_errors: tuple[str, ...], + ) -> ReformResolutionDecision: ... + + +class AnthropicReformCandidateResolver: + async def resolve(self, *, instruction, year, candidates): + return await self._call( + instruction=instruction, + year=year, + candidates=candidates, + correction=None, + ) + + async def correct_representation( + self, + *, + instruction, + year, + candidates, + previous, + validation_errors, + ): + return await self._call( + instruction=instruction, + year=year, + candidates=candidates, + correction={ + "previous": previous.model_dump(mode="json", exclude={"usage"}), + "validation_errors": validation_errors, + }, + ) + + async def _call(self, *, instruction, year, candidates, correction): + client = get_async_client() + tool = { + "name": "submit_reform_resolution", + "description": "Return one bounded reform-resolution outcome.", + "input_schema": ReformResolutionDecision.model_json_schema(), + } + response = await client.messages.create( + model=DEFAULT_FAST_MODEL, + max_tokens=2_000, + temperature=DEFAULT_TEMPERATURE, + system=( + "Resolve ordinary UK policy wording only against the supplied catalogue " + "entries. A resolved outcome must use only supplied parameter paths and " + "must state target, operation, value, unit, effective date, population, " + "and jurisdiction without inventing consequential meaning. The reform " + "field must be a flat JSON object whose key is the exact supplied " + "parameter path and whose value is the final scalar parameter value; " + "never put path, operation, from, or to fields inside reform. Return " + "needs_clarification for semantic ambiguity and unsupported when the " + "engine cannot represent the instruction. On a correction request, fix " + "only serialization or type shape and preserve all ReformMeaning fields." + ), + messages=[ + { + "role": "user", + "content": json.dumps( + { + "instruction": instruction, + "year": year, + "catalogue_candidates": candidates, + "representation_correction": correction, + }, + default=str, + ), + } + ], + tools=[tool], + tool_choice={"type": "tool", "name": "submit_reform_resolution"}, + ) + block = next( + ( + item + for item in response.content + if getattr(item, "type", None) == "tool_use" + and getattr(item, "name", None) == "submit_reform_resolution" + ), + None, + ) + if block is None: + raise RuntimeError("Reform resolver did not return structured output.") + usage = getattr(response, "usage", None) + payload = dict(block.input) + payload["usage"] = { + "input_tokens": getattr(usage, "input_tokens", 0), + "output_tokens": getattr(usage, "output_tokens", 0), + "cache_creation_input_tokens": getattr( + usage, "cache_creation_input_tokens", 0 + ), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", 0), + } + return ReformResolutionDecision.model_validate(payload) + + +def _package_version(package: str) -> str: + try: + return version(package) + except PackageNotFoundError: + return "unavailable" + + +class ResolveReformTool(Tool[ResolveReformInput, ResolveReformOutput]): + spec = ToolSpec( + identifier="resolve_reform", + version="1", + description=( + "Search the authoritative catalogue, construct one reform candidate, " + "validate it, and return a bounded resolution or clarification." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.CAPABILITY}), + input_model=ResolveReformInput, + output_model=ResolveReformOutput, + tool_dependencies=( + "list_reform_targets", + "get_parameter", + "validate_reform", + ), + ) + + def __init__(self, resolver: ReformCandidateResolver) -> None: + self._resolver = resolver + + async def run(self, tool_input: ResolveReformInput, context: ToolCallContext): + if self._requests_baseline(tool_input.instruction): + return self._output( + tool_input, + outcome=ReformResolutionKind.NO_REFORM, + summary="Use the current-law baseline.", + ) + search = await context.invoke_tool( + "list_reform_targets", + {"query": tool_input.instruction, "limit": 20}, + ) + candidates = await self._catalogue_candidates(search, tool_input.year, context) + if not candidates: + return self._output( + tool_input, + outcome=ReformResolutionKind.UNSUPPORTED, + summary="No authoritative reform target matched the instruction.", + ) + decision = await self._resolver.resolve( + instruction=tool_input.instruction, + year=tool_input.year, + candidates=candidates, + ) + context.record_model_usage(**decision.usage.model_dump()) + if decision.outcome is not ReformResolutionKind.RESOLVED: + return self._from_decision(tool_input, decision) + + decision = self._normalize_candidate_mapping(decision, candidates) + checked = self._check_meaning(decision, candidates) + if checked is not None: + return self._output( + tool_input, + outcome=ReformResolutionKind.NEEDS_CLARIFICATION, + summary=checked, + clarification=checked, + ) + validation = await self._validate(decision.reform, tool_input.year, context) + if not validation.get("valid"): + errors = self._validation_errors(validation) + if self._representation_only(errors): + corrected = await self._resolver.correct_representation( + instruction=tool_input.instruction, + year=tool_input.year, + candidates=candidates, + previous=decision, + validation_errors=errors, + ) + context.record_model_usage(**corrected.usage.model_dump()) + if corrected.meaning != decision.meaning: + return self._output( + tool_input, + outcome=ReformResolutionKind.NEEDS_CLARIFICATION, + summary="The proposed correction changed the policy meaning.", + clarification=( + "Please restate the intended target, operation, value, and unit." + ), + ) + decision = corrected + validation = await self._validate( + decision.reform, + tool_input.year, + context, + ) + if not validation.get("valid"): + return self._output( + tool_input, + outcome=ReformResolutionKind.FAILED, + summary="The reform candidate failed deterministic validation.", + ) + + normalized = validation.get("normalized_reform") + if not isinstance(normalized, dict) or set(normalized) != set(decision.reform): + return self._output( + tool_input, + outcome=ReformResolutionKind.FAILED, + summary="Deterministic validation changed the reform target set.", + ) + by_path = {candidate["path"]: candidate for candidate in candidates} + changes = tuple( + PolicyChange(parameter_path=path, value=value) + for path, value in normalized.items() + if isinstance(value, (str, int, float, bool)) or value is None + ) + if len(changes) != len(normalized): + return self._output( + tool_input, + outcome=ReformResolutionKind.UNSUPPORTED, + summary="This reform encoding is not yet transferable as a scalar change.", + ) + return self._output( + tool_input, + outcome=ReformResolutionKind.RESOLVED, + summary=decision.summary, + changes=changes, + bindings=tuple( + ReformCatalogueBinding( + parameter_path=path, + label=str(by_path[path].get("label") or path), + ) + for path in normalized + ), + ) + + async def _catalogue_candidates(self, result, year, context): + if not isinstance(result, SafeToolOutput): + raise TypeError("Reform catalogue returned an incompatible output.") + rows = result.root.get("targets") + if not isinstance(rows, list): + return () + enriched = [] + for row in rows[:20]: + if not isinstance(row, dict) or not isinstance(row.get("path"), str): + continue + item = dict(row) + detail = await context.invoke_tool( + "get_parameter", + {"path": row["path"], "year": year}, + ) + if isinstance(detail, SafeToolOutput) and isinstance( + detail.root.get("parameter"), dict + ): + item.update(detail.root["parameter"]) + enriched.append(item) + return tuple(enriched) + + @staticmethod + async def _validate(reform, year, context): + result = await context.invoke_tool( + "validate_reform", + {"reform": reform, "year": year}, + ) + if not isinstance(result, SafeToolOutput): + raise TypeError("Reform validation returned an incompatible output.") + return result.root + + @staticmethod + def _normalize_candidate_mapping(decision, candidates): + meaning = decision.meaning + if meaning is None: + return decision + paths = {candidate["path"] for candidate in candidates} + if meaning.target not in paths or set(decision.reform) == {meaning.target}: + return decision + final_value = None + if meaning.operation == "set": + final_value = meaning.value + elif decision.reform.get("to") == meaning.value: + final_value = meaning.value + if final_value is None: + return decision + return decision.model_copy( + update={"reform": {meaning.target: final_value}} + ) + + @staticmethod + def _check_meaning(decision, candidates) -> str | None: + if decision.meaning is None: + return "The reform target, operation, value, or scope is incomplete." + if not decision.reform: + return "The resolved reform contains no policy change." + paths = {candidate["path"] for candidate in candidates} + unknown = set(decision.reform) - paths + if unknown: + return "The proposed reform used a parameter absent from catalogue results." + if decision.meaning.target not in decision.reform: + return "The stated reform target does not match the constructed reform." + return None + + @staticmethod + def _validation_errors(validation) -> tuple[str, ...]: + errors = validation.get("errors") + if not isinstance(errors, list): + return ("unknown validation failure",) + return tuple( + str(item.get("message", item)) if isinstance(item, dict) else str(item) + for item in errors + ) + + @staticmethod + def _representation_only(errors: tuple[str, ...]) -> bool: + markers = ("object", "mapping", "type", "format", "serialization", "shape") + return bool(errors) and all( + any(marker in error.casefold() for marker in markers) for error in errors + ) + + @staticmethod + def _requests_baseline(instruction: str) -> bool: + normalized = " ".join(instruction.casefold().split()) + return normalized in { + "current law", + "current policy", + "no reform", + "baseline", + "use the baseline", + } + + @staticmethod + def _output( + tool_input, + *, + outcome, + summary, + clarification=None, + changes=(), + bindings=(), + ): + return ResolveReformOutput( + outcome=outcome, + year=tool_input.year, + summary=summary, + clarification=clarification, + changes=changes, + bindings=bindings, + catalogue_version=_package_version("policyengine-uk"), + calculation_engine_version=_package_version("policyengine"), + ) + + @classmethod + def _from_decision(cls, tool_input, decision): + return cls._output( + tool_input, + outcome=decision.outcome, + summary=decision.summary, + clarification=decision.clarification, + ) + + +class PolicyReformInput(StrictModel): + instruction: str + year: int | None = None + referenced_policy_scenario_id: str | None = None + + +class PolicyReformOutput(StrictModel): + scenario: PolicyScenarioRef + year_source: InputSource + resolution_summary: str + + +class PolicyReformCapability(Capability[PolicyReformInput, PolicyReformOutput]): + spec = CapabilitySpec( + identifier="policy_reform", + version="1", + description=( + "Resolve an ordinary-language UK policy change into a deterministically " + "validated scenario reference." + ), + required_use=( + "Use when another capability needs a verified reform or the user asks to " + "construct or validate a policy change." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=PolicyReformInput, + output_model=PolicyReformOutput, + accepted_artifacts=( + ArtifactContract(artifact_type="policy_scenario", schema_version="1"), + ), + produced_artifacts=( + ArtifactContract(artifact_type="policy_scenario", schema_version="1"), + ), + tool_dependencies=("resolve_reform",), + ) + + async def run(self, capability_input: PolicyReformInput, context): + referenced = await self._referenced_scenario(capability_input, context) + resolved_year = resolve_policy_year( + explicit_year=capability_input.year, + referenced_year=referenced.year if referenced is not None else None, + ) + result = await context.invoke_tool( + "resolve_reform", + { + "instruction": capability_input.instruction, + "year": resolved_year.year, + "existing_scenario_id": ( + referenced.artifact_id if referenced is not None else None + ), + }, + ) + if not isinstance(result, ResolveReformOutput): + raise TypeError("Reform resolution returned an incompatible output.") + if result.outcome is ReformResolutionKind.NEEDS_CLARIFICATION: + partial = PolicyReformInput( + instruction=capability_input.instruction, + year=resolved_year.year, + referenced_policy_scenario_id=( + referenced.artifact_id if referenced is not None else None + ), + ) + await context.persist_waiting(partial) + return NeedsInput( + prompt=result.clarification or result.summary, + missing_fields=("instruction",), + partial_input=partial.model_dump(mode="json", exclude_none=True), + ) + if result.outcome is ReformResolutionKind.UNSUPPORTED: + return Unsupported(reason=result.summary) + if result.outcome is ReformResolutionKind.FAILED: + return Failed( + safe_message=result.summary, + error_code="reform_resolution_failed", + ) + baseline = result.outcome is ReformResolutionKind.NO_REFORM + scenario = PolicyScenarioRef( + provenance=ArtifactProvenance( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + capability_id=self.spec.identifier, + capability_version=self.spec.version, + invocation_id=context.capability_invocation_id, + sources=(resolved_year.source.value, "resolve_reform"), + ), + year=resolved_year.year, + scenario_revision=self._revision(resolved_year.year, result.changes), + catalogue_version=result.catalogue_version, + calculation_engine_version=result.calculation_engine_version, + baseline=baseline, + verified_changes=result.changes, + ) + saved = await context.save_artifact(scenario) + return Completed( + value=PolicyReformOutput( + scenario=saved, + year_source=resolved_year.source, + resolution_summary=result.summary, + ) + ) + + @staticmethod + async def _referenced_scenario(capability_input, context): + artifact_id = capability_input.referenced_policy_scenario_id + if artifact_id is None: + return None + scenarios = await context.find_artifacts(PolicyScenarioRef) + return next( + (scenario for scenario in scenarios if scenario.artifact_id == artifact_id), + None, + ) + + @staticmethod + def _revision(year: int, changes: tuple[PolicyChange, ...]) -> str: + payload = json.dumps( + { + "year": year, + "changes": [change.model_dump(mode="json") for change in changes], + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest()[:16] diff --git a/backend/capabilities/registry.py b/backend/capabilities/registry.py new file mode 100644 index 00000000..848cb3ba --- /dev/null +++ b/backend/capabilities/registry.py @@ -0,0 +1,135 @@ +"""Registry and startup validation for typed capabilities.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel + +from capabilities.contracts import Capability, CapabilitySpec +from tools.contracts import CallerType, Visibility + + +class CapabilityRegistry: + def __init__(self) -> None: + self._capabilities: dict[str, Capability[Any, Any]] = {} + + def register(self, capability: Capability[Any, Any]) -> None: + spec = getattr(capability, "spec", None) + if not isinstance(spec, CapabilitySpec): + raise TypeError("Capabilities must declare a valid CapabilitySpec.") + if spec.identifier in self._capabilities: + raise ValueError(f"Duplicate capability registration: {spec.identifier}") + if not issubclass(spec.input_model, BaseModel) or not issubclass( + spec.output_model, BaseModel + ): + raise TypeError( + f"Capability {spec.identifier} input and output declarations must be Pydantic models." + ) + self._capabilities[spec.identifier] = capability + + def get( + self, + identifier: str, + *, + caller: CallerType, + ) -> Capability[Any, Any]: + try: + capability = self._capabilities[identifier] + except KeyError as exc: + raise KeyError(f"Unknown capability: {identifier}") from exc + if caller not in capability.spec.allowed_callers: + raise PermissionError( + f"Caller {caller.value} may not invoke capability {identifier}." + ) + return capability + + def descriptions_for( + self, + caller: CallerType, + *, + include_private: bool = False, + ) -> tuple[dict[str, object], ...]: + descriptions: list[dict[str, object]] = [] + for capability in self._capabilities.values(): + spec = capability.spec + if caller not in spec.allowed_callers: + continue + if spec.visibility is Visibility.PRIVATE and not include_private: + continue + descriptions.append( + { + "identifier": spec.identifier, + "version": spec.version, + "description": spec.description, + "required_use": spec.required_use, + "input_schema": spec.input_model.model_json_schema(), + } + ) + return tuple(descriptions) + + def specs(self) -> tuple[CapabilitySpec, ...]: + return tuple(capability.spec for capability in self._capabilities.values()) + + def registered(self, identifier: str) -> Capability[Any, Any]: + try: + return self._capabilities[identifier] + except KeyError as exc: + raise KeyError(f"Unknown capability: {identifier}") from exc + + def validate(self) -> None: + self._validate_dependencies() + self._validate_cycles() + + def _validate_dependencies(self) -> None: + for consumer in self._capabilities.values(): + for dependency in consumer.spec.dependencies: + provider = self._capabilities.get(dependency.capability_id) + if provider is None: + raise ValueError( + f"Capability {consumer.spec.identifier} requires unknown capability " + f"{dependency.capability_id}." + ) + if dependency.artifact is None: + continue + if not any( + dependency.artifact.is_compatible_with(produced) + for produced in provider.spec.produced_artifacts + ): + raise ValueError( + f"Capability {consumer.spec.identifier} requires incompatible " + f"artifact {dependency.artifact.artifact_type} " + f"from {dependency.capability_id}." + ) + if not any( + accepted.is_compatible_with(dependency.artifact) + for accepted in consumer.spec.accepted_artifacts + ): + raise ValueError( + f"Capability {consumer.spec.identifier} does not declare accepted " + f"artifact {dependency.artifact.artifact_type}." + ) + + def _validate_cycles(self) -> None: + visiting: list[str] = [] + visited: set[str] = set() + + def visit(identifier: str) -> None: + if identifier in visited: + return + if identifier in visiting: + start = visiting.index(identifier) + cycle = visiting[start:] + [identifier] + raise ValueError( + "Capability dependency cycle: " + " -> ".join(cycle) + ) + visiting.append(identifier) + capability = self._capabilities[identifier] + for dependency in capability.spec.dependencies: + if dependency.capability_id in self._capabilities: + visit(dependency.capability_id) + visiting.pop() + visited.add(identifier) + + for identifier in self._capabilities: + visit(identifier) diff --git a/backend/capabilities/relevance.py b/backend/capabilities/relevance.py new file mode 100644 index 00000000..5c3d365d --- /dev/null +++ b/backend/capabilities/relevance.py @@ -0,0 +1,154 @@ +"""Every-turn UK policy relevance assessment with bounded outcomes.""" + +from __future__ import annotations + +from enum import Enum +from typing import Protocol + +from pydantic import BaseModel, ConfigDict + +from capabilities.contracts import Capability, CapabilitySpec, Completed +from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_async_client +from tools.contracts import CallerType, Tool, ToolCallContext, ToolSpec, Visibility +from conversation_context.projection import ContextProjection + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class ConversationExcerpt(StrictModel): + role: str + content: str + + +class RelevanceResult(str, Enum): + RELEVANT = "relevant" + UNCERTAIN = "uncertain" + CLEARLY_OUT_OF_SCOPE = "clearly_out_of_scope" + + +class RelevanceUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class AssessRelevanceInput(StrictModel): + current_message: str + conversation: tuple[ConversationExcerpt, ...] + context: ContextProjection | None = None + + +class RelevanceAssessment(StrictModel): + result: RelevanceResult + explanation: str + usage: RelevanceUsage = RelevanceUsage() + + +class RelevanceAssessor(Protocol): + async def assess(self, request: AssessRelevanceInput) -> RelevanceAssessment: ... + + +class AnthropicRelevanceAssessor: + """One forced structured model call whose output cannot select capabilities.""" + + async def assess(self, request: AssessRelevanceInput) -> RelevanceAssessment: + client = get_async_client() + tool = { + "name": "submit_relevance_assessment", + "description": "Return only the bounded UK Chat relevance result.", + "input_schema": RelevanceAssessment.model_json_schema(), + } + response = await client.messages.create( + model=DEFAULT_FAST_MODEL, + max_tokens=500, + temperature=DEFAULT_TEMPERATURE, + system=( + "Assess whether the latest turn concerns supported UK tax, benefit, " + "government-policy, household-impact, or population-impact discussion. " + "Return relevant for supported UK content, uncertain when context could " + "make it relevant, and clearly_out_of_scope only for an explicitly " + "unsupported jurisdiction or unrelated request. Do not select another " + "capability, infer calculation inputs, or propose policy values." + ), + messages=[ + { + "role": "user", + "content": request.model_dump_json(), + } + ], + tools=[tool], + tool_choice={"type": "tool", "name": "submit_relevance_assessment"}, + ) + block = next( + ( + item + for item in response.content + if getattr(item, "type", None) == "tool_use" + and getattr(item, "name", None) == "submit_relevance_assessment" + ), + None, + ) + if block is None: + raise RuntimeError("Relevance assessor did not return structured output.") + usage = getattr(response, "usage", None) + payload = dict(block.input) + payload["usage"] = { + "input_tokens": getattr(usage, "input_tokens", 0), + "output_tokens": getattr(usage, "output_tokens", 0), + "cache_creation_input_tokens": getattr( + usage, + "cache_creation_input_tokens", + 0, + ), + "cache_read_input_tokens": getattr(usage, "cache_read_input_tokens", 0), + } + return RelevanceAssessment.model_validate(payload) + + +class AssessRelevanceTool(Tool[AssessRelevanceInput, RelevanceAssessment]): + spec = ToolSpec( + identifier="assess_relevance", + version="1", + description="Assess only whether the current turn is within UK Chat scope.", + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.CAPABILITY}), + input_model=AssessRelevanceInput, + output_model=RelevanceAssessment, + ) + + def __init__(self, assessor: RelevanceAssessor) -> None: + self._assessor = assessor + + async def run( + self, + tool_input: AssessRelevanceInput, + context: ToolCallContext, + ) -> RelevanceAssessment: + result = await self._assessor.assess(tool_input) + context.record_model_usage(**result.usage.model_dump()) + return result + + +class ConversationRelevanceCapability( + Capability[AssessRelevanceInput, RelevanceAssessment] +): + spec = CapabilitySpec( + identifier="conversation_relevance", + version="1", + description="Assess the scope of each user turn before normal conversation.", + required_use="Run once for every user turn.", + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=AssessRelevanceInput, + output_model=RelevanceAssessment, + tool_dependencies=("assess_relevance",), + ) + + async def run(self, capability_input: AssessRelevanceInput, context): + result = await context.invoke_tool("assess_relevance", capability_input) + if not isinstance(result, RelevanceAssessment): + raise TypeError("Relevance tool returned an incompatible output.") + return Completed(value=result) diff --git a/backend/capabilities/repository.py b/backend/capabilities/repository.py new file mode 100644 index 00000000..a668a715 --- /dev/null +++ b/backend/capabilities/repository.py @@ -0,0 +1,104 @@ +"""Repository contracts for cross-turn capability data.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Protocol + +from pydantic import BaseModel, ConfigDict, Field + +from capabilities.artifacts import ArtifactBase, TransferableArtifact +from capabilities.tracing import InvocationRecord +from conversation_context.models import ( + CapabilityInvocationReference, + FactRequirement, +) + + +class WaitingCapabilityInvocation(BaseModel): + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True, extra="forbid") + + invocation_id: str + conversation_id: str + capability_id: str + capability_version: str + input_schema_version: str + partial_input: BaseModel + source_turn_id: str + context_scope_id: str | None = None + context_revision: int | None = None + requirements: tuple[FactRequirement, ...] = () + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + def reference(self) -> CapabilityInvocationReference | None: + if self.context_scope_id is None or self.context_revision is None: + return None + return CapabilityInvocationReference( + invocation_id=self.invocation_id, + capability_id=self.capability_id, + capability_version=self.capability_version, + context_scope_id=self.context_scope_id, + context_revision=self.context_revision, + ) + + +class ConversationCapabilityRepository(Protocol): + def save_artifact( + self, + conversation_id: str, + artifact: TransferableArtifact, + ) -> ArtifactBase: ... + + def get_artifact(self, conversation_id: str, artifact_id: str) -> ArtifactBase: ... + + def find_artifacts( + self, + conversation_id: str, + artifact_model: type[ArtifactBase], + ) -> tuple[ArtifactBase, ...]: ... + + def create_waiting( + self, invocation: WaitingCapabilityInvocation + ) -> WaitingCapabilityInvocation: ... + + def get_waiting(self, invocation_id: str) -> WaitingCapabilityInvocation: ... + + def list_waiting( + self, + conversation_id: str, + *, + capability_id: str | None = None, + ) -> tuple[WaitingCapabilityInvocation, ...]: ... + + def update_waiting( + self, + invocation_id: str, + partial_input: BaseModel, + ) -> WaitingCapabilityInvocation: ... + + def resume_waiting( + self, + invocation_id: str, + updates: dict[str, object], + ) -> WaitingCapabilityInvocation: ... + + def branch_waiting( + self, + invocation_id: str, + new_invocation_id: str, + source_turn_id: str, + ) -> WaitingCapabilityInvocation: ... + + def remove_waiting(self, invocation_id: str) -> None: ... + + +class InvocationTraceRepository(Protocol): + def save(self, record: InvocationRecord) -> None: ... + + def list_for_conversation( + self, + conversation_id: str, + *, + include_private: bool, + ) -> tuple[InvocationRecord, ...]: ... diff --git a/backend/capabilities/society.py b/backend/capabilities/society.py new file mode 100644 index 00000000..bcdf8e4a --- /dev/null +++ b/backend/capabilities/society.py @@ -0,0 +1,441 @@ +"""Deterministic population analysis with mandatory default aggregates.""" + +from __future__ import annotations + +from importlib.metadata import PackageNotFoundError, version + +from pydantic import BaseModel, ConfigDict + +from capabilities.artifacts import ( + AggregateDimension, + AggregateValue, + ArtifactProvenance, + PolicyScenarioRef, + RequestedOutputIssue, + SocietyAnalysisResultRef, +) +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilityDependency, + CapabilitySpec, + Completed, + Failed, + NeedsInput, +) +from capabilities.input_resolution import InputSource, resolve_policy_year +from capabilities.policy_reform import PolicyReformOutput +from engine.constants import UK_CHAT_DATASET +from tools.analysis_support import ( + ExtractResultFindingsOutput, + NumericalFact, + SelectSupportedOutputsOutput, +) +from tools.contracts import CallerType, Visibility +from tools.typed_models import SafeToolOutput + + +SOCIETY_DEFAULT_PROFILE_VERSION = "1" +SOCIETY_DEFAULT_OUTPUTS = ( + "budgetary_impact", + "winners_losers", + "decile_impacts", +) +SOCIETY_DEFAULT_DECILE_CONCEPT = "household_net_income" +SOCIETY_EXTRA_VARIABLES_BY_OUTPUT: dict[str, dict[str, tuple[str, ...]]] = { + "budgetary_impact": {}, + "program_statistics": {}, + "decile_impacts": {}, + "winners_losers": {}, + "poverty": {}, + "inequality": {}, +} + + +def _package_version(package: str) -> str: + try: + return version(package) + except PackageNotFoundError: + return "unavailable" + + +def current_dataset_version() -> str: + return UK_CHAT_DATASET.uri.rsplit("@", 1)[-1] + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class SocietyAnalysisInput(StrictModel): + reform_instruction: str | None = None + referenced_policy_scenario_id: str | None = None + year: int | None = None + requested_outputs: tuple[str, ...] = () + + +class SocietyAnalysisOutput(StrictModel): + result: SocietyAnalysisResultRef + year_source: InputSource + narration_facts: tuple[NumericalFact, ...] + required_output_ids: tuple[str, ...] = SOCIETY_DEFAULT_OUTPUTS + narration_requirement: str = ( + "Present budgetary impact, winners/losers/unchanged, and income-decile " + "impacts, plus every successfully calculated requested output and any issue." + ) + + +_DERIVATIVE_TOOL = { + "budgetary_impact": "compute_budgetary_impact", + "program_statistics": "compute_program_breakdown", + "decile_impacts": "compute_decile_impacts", + "winners_losers": "compute_winners_losers", + "poverty": "compute_poverty_metrics", + "inequality": "compute_inequality_metrics", +} + +_DERIVATIVE_EXECUTION_PRIORITY = ( + "decile_impacts", + "winners_losers", + "budgetary_impact", +) + + +class SocietyAnalysisCapability(Capability[SocietyAnalysisInput, SocietyAnalysisOutput]): + spec = CapabilitySpec( + identifier="society_analysis", + version="1", + description=( + "Run deterministic UK population analysis with mandatory budget, " + "winners/losers, and household-income-decile aggregates plus supported " + "requested outputs." + ), + required_use=( + "Required for population-wide reform costs, distributional effects, " + "poverty, inequality, winners/losers, or aggregate benefit impacts." + ), + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL, CallerType.CAPABILITY}), + input_model=SocietyAnalysisInput, + output_model=SocietyAnalysisOutput, + accepted_artifacts=( + ArtifactContract(artifact_type="policy_scenario", schema_version="1"), + ), + produced_artifacts=( + ArtifactContract( + artifact_type="society_analysis_result", + schema_version="1", + ), + ), + dependencies=( + CapabilityDependency( + capability_id="policy_reform", + artifact=ArtifactContract( + artifact_type="policy_scenario", + schema_version="1", + ), + ), + ), + tool_dependencies=( + "select_supported_outputs", + "run_society_simulation", + "compute_budgetary_impact", + "compute_program_breakdown", + "compute_decile_impacts", + "compute_winners_losers", + "compute_poverty_metrics", + "compute_inequality_metrics", + "extract_result_findings", + ), + ) + + async def run(self, capability_input: SocietyAnalysisInput, context): + scenario = await self._scenario(capability_input, context) + resolved_year = resolve_policy_year( + explicit_year=capability_input.year, + referenced_year=scenario.year if scenario is not None else None, + ) + if capability_input.reform_instruction: + scenario_outcome = await context.invoke_capability( + "policy_reform", + { + "instruction": capability_input.reform_instruction, + "year": resolved_year.year, + "referenced_policy_scenario_id": ( + scenario.artifact_id if scenario is not None else None + ), + }, + ) + if not isinstance(scenario_outcome, Completed) or not isinstance( + scenario_outcome.value, + PolicyReformOutput, + ): + return await self._forward_reform_outcome( + scenario_outcome, + capability_input, + context, + ) + scenario = scenario_outcome.value.scenario + elif scenario is None: + scenario_outcome = await context.invoke_capability( + "policy_reform", + {"instruction": "current law", "year": resolved_year.year}, + ) + if not isinstance(scenario_outcome, Completed) or not isinstance( + scenario_outcome.value, + PolicyReformOutput, + ): + return await self._forward_reform_outcome( + scenario_outcome, + capability_input, + context, + ) + scenario = scenario_outcome.value.scenario + + selected = await context.invoke_tool( + "select_supported_outputs", + {"requested_outputs": list(capability_input.requested_outputs)}, + ) + if not isinstance(selected, SelectSupportedOutputsOutput): + raise TypeError("Society output selection returned an incompatible output.") + reform = { + change.parameter_path: change.value for change in scenario.verified_changes + } + simulation_input = { + "year": resolved_year.year, + "reform": reform or None, + } + extra_variables = self._extra_variables(selected.output_ids) + if extra_variables: + simulation_input["extra_variables"] = extra_variables + simulation = await context.invoke_tool( + "run_society_simulation", + simulation_input, + ) + if not isinstance(simulation, SafeToolOutput): + raise TypeError("Society simulation returned an incompatible output.") + simulation_id = simulation.root.get("result_id") + if not isinstance(simulation_id, str) or "error" in simulation.root: + return Failed( + safe_message="The deterministic population calculation failed.", + error_code="society_simulation_failed", + ) + + aggregate_values: list[AggregateValue] = [] + for output_id in self._execution_order(selected.output_ids): + tool_id = _DERIVATIVE_TOOL.get(output_id) + if tool_id is None: + continue + tool_input = {"simulation_id": simulation_id} + if output_id == "decile_impacts": + tool_input["decile_concept"] = SOCIETY_DEFAULT_DECILE_CONCEPT + derivative = await context.invoke_tool(tool_id, tool_input) + if not isinstance(derivative, SafeToolOutput): + raise TypeError(f"{tool_id} returned an incompatible output.") + if "error" in derivative.root: + if output_id in SOCIETY_DEFAULT_OUTPUTS: + return Failed( + safe_message=f"The required {output_id} calculation failed.", + error_code="society_default_output_failed", + ) + continue + aggregate_values.extend(self._flatten(output_id, derivative.root)) + + calculated_ids = tuple( + output_id + for output_id in selected.output_ids + if any(value.output_id == output_id for value in aggregate_values) + ) + missing_defaults = set(SOCIETY_DEFAULT_OUTPUTS) - set(calculated_ids) + if missing_defaults: + return Failed( + safe_message=( + "The population calculation did not produce the complete default " + "aggregate profile." + ), + error_code="society_default_profile_incomplete", + ) + issues = tuple( + RequestedOutputIssue( + request=issue.request, + kind=issue.kind, + guidance=issue.guidance, + ) + for issue in selected.issues + ) + result = SocietyAnalysisResultRef( + provenance=ArtifactProvenance( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + capability_id=self.spec.identifier, + capability_version=self.spec.version, + invocation_id=context.capability_invocation_id, + sources=("fixed dataset", "verified scenario", "aggregate derivatives"), + ), + year=resolved_year.year, + policy_scenario_artifact_id=scenario.artifact_id, + scenario_revision=scenario.scenario_revision, + catalogue_version=scenario.catalogue_version, + dataset_version=current_dataset_version(), + calculation_engine_version=scenario.calculation_engine_version, + default_profile_version=SOCIETY_DEFAULT_PROFILE_VERSION, + calculated_output_ids=calculated_ids, + outputs=tuple(aggregate_values), + requested_output_issues=issues, + ) + result = await context.save_artifact(result) + extracted = await context.invoke_tool( + "extract_result_findings", + { + "outputs": [ + output.model_dump(mode="json") for output in result.outputs + ] + }, + ) + if not isinstance(extracted, ExtractResultFindingsOutput): + raise TypeError("Society finding extraction returned an incompatible output.") + narration_facts = tuple( + NumericalFact(label=finding.label, value=finding.value, unit=finding.unit) + for finding in extracted.findings + if finding.value is not None + ) + return Completed( + value=SocietyAnalysisOutput( + result=result, + year_source=resolved_year.source, + narration_facts=narration_facts, + ) + ) + + @staticmethod + async def _forward_reform_outcome(outcome, capability_input, context): + if not isinstance(outcome, NeedsInput): + return outcome + await context.persist_waiting(capability_input) + return NeedsInput( + prompt=outcome.prompt, + missing_fields=("reform_instruction",), + partial_input=capability_input.model_dump( + mode="json", + exclude_none=True, + exclude_defaults=True, + ), + ) + + @staticmethod + async def _scenario(capability_input, context): + if capability_input.referenced_policy_scenario_id is None: + return None + scenarios = await context.find_artifacts(PolicyScenarioRef) + return next( + ( + scenario + for scenario in scenarios + if scenario.artifact_id + == capability_input.referenced_policy_scenario_id + ), + None, + ) + + @classmethod + def _flatten(cls, output_id: str, payload: dict) -> tuple[AggregateValue, ...]: + values: list[AggregateValue] = [] + ignored = { + "status", + "simulation_id", + "result_id", + "year", + "quantiles", + "decile_concept", + } + + def visit(value, path: tuple[str, ...], dimensions=()): + if isinstance(value, bool) or value is None: + return + if isinstance(value, (int, float)): + metric = ".".join(path) or "value" + values.append( + AggregateValue( + output_id=output_id, + metric_id=metric, + label=cls._label(output_id, path), + value=value, + unit=cls._unit(output_id, path), + dimensions=dimensions, + ) + ) + return + if isinstance(value, dict): + local_dimensions = list(dimensions) + for key in ("decile", "group", "age_group", "poverty_type"): + if key in value and isinstance(value[key], (str, int)): + local_dimensions.append( + AggregateDimension(name=key, value=str(value[key])) + ) + for key, nested in value.items(): + if key in ignored or key in { + "decile", + "group", + "age_group", + "poverty_type", + "label", + "description", + }: + continue + visit(nested, (*path, str(key)), tuple(local_dimensions)) + return + if isinstance(value, list): + for item in value: + visit(item, path, dimensions) + + visit(payload, ()) + return tuple(values) + + @staticmethod + def _label(output_id, path): + suffix = " ".join(path).replace("_", " ").title() + prefix = output_id.replace("_", " ").title() + return f"{prefix}: {suffix}" if suffix else prefix + + @staticmethod + def _unit(output_id, path): + metric = " ".join(path).casefold() + if any(word in metric for word in ("rate", "share", "relative", "percent")): + return "ratio" + if output_id == "winners_losers" and any( + word in metric for word in ("winner", "loser", "unchanged") + ): + return "people" + if output_id in {"budgetary_impact", "program_statistics", "decile_impacts"}: + return "GBP/year" + return "number" + + @staticmethod + def _execution_order(output_ids): + priority = { + output_id: index + for index, output_id in enumerate(_DERIVATIVE_EXECUTION_PRIORITY) + } + return tuple( + sorted( + output_ids, + key=lambda output_id: priority.get( + output_id, + len(priority), + ), + ) + ) + + @staticmethod + def _extra_variables(output_ids): + by_entity: dict[str, list[str]] = {} + for output_id in output_ids: + for entity, variables in SOCIETY_EXTRA_VARIABLES_BY_OUTPUT.get( + output_id, + {}, + ).items(): + selected = by_entity.setdefault(entity, []) + for variable in variables: + if variable not in selected: + selected.append(variable) + return by_entity diff --git a/backend/capabilities/tracing.py b/backend/capabilities/tracing.py new file mode 100644 index 00000000..c3199604 --- /dev/null +++ b/backend/capabilities/tracing.py @@ -0,0 +1,294 @@ +"""Sanitized invocation records shared by execution and observability.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from enum import Enum +import re +from threading import Lock +from time import monotonic +from typing import Literal, Protocol +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from tools.contracts import Visibility + + +class InvocationKind(str, Enum): + TOOL = "tool" + CAPABILITY = "capability" + + +class InvocationStatus(str, Enum): + RUNNING = "running" + COMPLETED = "completed" + NEEDS_INPUT = "needs_input" + UNSUPPORTED = "unsupported" + FAILED = "failed" + CANCELLED = "cancelled" + + +class InvocationRecord(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + conversation_id: str + turn_id: str + invocation_id: str + parent_invocation_id: str | None = None + sequence: int + kind: InvocationKind + identifier: str + version: str + visibility: Visibility + started_at: datetime + completed_at: datetime | None = None + duration_ms: int | None = Field(default=None, ge=0) + status: InvocationStatus + summary: str + debug_input: JsonValue | None = None + debug_output: JsonValue | None = None + + +class InvocationTraceSink(Protocol): + """Persistence operations needed by the request-independent tracer.""" + + def save(self, record: InvocationRecord) -> None: ... + + def last_sequence(self, conversation_id: str) -> int: ... + + +class InvocationTraceEvent(BaseModel): + """One allowlisted start or finish update suitable for public projection.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + event_index: int + phase: Literal["started", "finished"] + record: InvocationRecord + + +_CONTROL_CHARACTERS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]") +_SECRET_VALUE = re.compile( + r"(?i)(?:bearer\s+|sk-(?:ant-)?|hf_)[a-z0-9_\-]{8,}" +) +_MAX_SUMMARY_LENGTH = 240 +_SECRET_KEYS = frozenset( + { + "api_key", + "authorization", + "credential", + "password", + "secret", + "token", + } +) +_SECRET_KEY_SUFFIXES = ( + "_access_token", + "_api_key", + "_auth_token", + "_credential", + "_password", + "_secret", + "_service_role_key", +) +_REQUEST_LOCAL_KEYS = frozenset( + { + "request_id", + "result_handle", + "result_id", + "simulation_id", + } +) +_ROW_LEVEL_KEYS = frozenset( + { + "microdata", + "records", + "row_data", + "row_level_data", + "survey_records", + } +) +def sanitize_trace_summary(summary: str) -> str: + """Restrict retained summaries to short, single-line metadata text.""" + + normalized = " ".join(_CONTROL_CHARACTERS.sub(" ", summary).split()) + if not normalized: + return "Invocation status updated." + return normalized[:_MAX_SUMMARY_LENGTH] + + +def _safe_string(value: str) -> str: + return _SECRET_VALUE.sub("[redacted secret]", value) + + +def _is_secret_key(key: str) -> bool: + return key in _SECRET_KEYS or key.endswith(_SECRET_KEY_SUFFIXES) + + +def _project_value(value: object) -> JsonValue: + if value is None or isinstance(value, (bool, int, float)): + return value + if isinstance(value, str): + return _safe_string(value) + if isinstance(value, dict): + projected: dict[str, JsonValue] = {} + for raw_key, item in value.items(): + key = str(raw_key) + normalized = key.casefold() + if _is_secret_key(normalized): + projected[key] = "[redacted secret]" + elif normalized in _REQUEST_LOCAL_KEYS: + projected[key] = "[request-local identifier omitted]" + elif normalized in _ROW_LEVEL_KEYS: + projected[key] = "[record-level data omitted]" + else: + projected[key] = _project_value(item) + return projected + if isinstance(value, (list, tuple)): + return [_project_value(item) for item in value] + return f"[{type(value).__name__} omitted]" + + +def debug_projection(value: BaseModel | object) -> JsonValue: + """Preserve validated structured values while removing prohibited trace data.""" + + serialized = ( + value.model_dump(mode="json") if isinstance(value, BaseModel) else value + ) + return _project_value(serialized) + + +class InvocationTracer: + """Record every invocation using a fixed metadata-only schema.""" + + def __init__(self, *, sink: InvocationTraceSink | None = None) -> None: + self._sink = sink + self._records: dict[str, InvocationRecord] = {} + self._started: dict[str, float] = {} + self._sequence: dict[str, int] = {} + self._events: list[InvocationTraceEvent] = [] + self._lock = Lock() + + def start( + self, + *, + conversation_id: str, + turn_id: str, + parent_invocation_id: str | None, + kind: InvocationKind, + identifier: str, + version: str, + visibility: Visibility, + summary: str, + debug_input: JsonValue | None = None, + ) -> InvocationRecord: + with self._lock: + if conversation_id not in self._sequence: + self._sequence[conversation_id] = ( + self._sink.last_sequence(conversation_id) + if self._sink is not None + else 0 + ) + sequence = self._sequence[conversation_id] + 1 + self._sequence[conversation_id] = sequence + invocation_id = uuid4().hex + record = InvocationRecord( + conversation_id=conversation_id, + turn_id=turn_id, + invocation_id=invocation_id, + parent_invocation_id=parent_invocation_id, + sequence=sequence, + kind=kind, + identifier=identifier, + version=version, + visibility=visibility, + started_at=datetime.now(timezone.utc), + status=InvocationStatus.RUNNING, + summary=sanitize_trace_summary(summary), + debug_input=debug_input, + ) + self._records[invocation_id] = record + self._started[invocation_id] = monotonic() + self._append_event("started", record) + if self._sink is not None: + self._sink.save(record) + return record + + def finish( + self, + invocation_id: str, + *, + status: InvocationStatus, + summary: str, + debug_output: JsonValue | None = None, + ) -> InvocationRecord: + completed_at = datetime.now(timezone.utc) + with self._lock: + current = self._records[invocation_id] + started = self._started.pop(invocation_id) + updated = current.model_copy( + update={ + "completed_at": completed_at, + "duration_ms": max(0, round((monotonic() - started) * 1000)), + "status": status, + "summary": sanitize_trace_summary(summary), + "debug_output": debug_output, + } + ) + self._records[invocation_id] = updated + self._append_event("finished", updated) + if self._sink is not None: + self._sink.save(updated) + return updated + + def records( + self, + conversation_id: str, + *, + include_private: bool, + ) -> tuple[InvocationRecord, ...]: + with self._lock: + records = tuple( + record + for record in self._records.values() + if record.conversation_id == conversation_id + and (include_private or record.visibility is Visibility.PUBLIC) + ) + return tuple(sorted(records, key=lambda record: record.sequence)) + + def events_for_turn( + self, + *, + conversation_id: str, + turn_id: str, + after_event_index: int = 0, + include_private: bool, + ) -> tuple[InvocationTraceEvent, ...]: + """Return ordered trace updates created after a caller-owned cursor.""" + + with self._lock: + return tuple( + event + for event in self._events + if event.event_index > after_event_index + and event.record.conversation_id == conversation_id + and event.record.turn_id == turn_id + and ( + include_private + or event.record.visibility is Visibility.PUBLIC + ) + ) + + def _append_event( + self, + phase: Literal["started", "finished"], + record: InvocationRecord, + ) -> None: + self._events.append( + InvocationTraceEvent( + event_index=len(self._events) + 1, + phase=phase, + record=record, + ) + ) diff --git a/backend/chat/__init__.py b/backend/chat/__init__.py index 8e8bf349..c6f890b4 100644 --- a/backend/chat/__init__.py +++ b/backend/chat/__init__.py @@ -1,7 +1,20 @@ -"""The chat turn: orchestration, system-block assembly, model selection, -titles, and follow-up suggestions. Re-exports the router for app wiring. -""" +"""Chat package with a lazily imported FastAPI adapter.""" + +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from fastapi import APIRouter + + router: APIRouter + + +def __getattr__(name: str): + if name != "router": + raise AttributeError(name) + from chat.routes import router + + return router -from chat.routes import router __all__ = ["router"] diff --git a/backend/chat/activity.py b/backend/chat/activity.py new file mode 100644 index 00000000..693d0d8a --- /dev/null +++ b/backend/chat/activity.py @@ -0,0 +1,97 @@ +"""Conversation-authorized projections of sanitized invocation activity.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel, ConfigDict, JsonValue +from sqlmodel import Session, select + +from capabilities.tracing import InvocationKind, InvocationStatus +from conversations.models import ChatConversation, get_engine +from persistence.trace_repository import SQLInvocationTraceRepository +from tools.contracts import Visibility + + +router = APIRouter(prefix="/chat", tags=["chatbot"]) + + +class InvocationActivityItem(BaseModel): + """The complete and exclusive field allowlist exposed by the activity API.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + turn_id: str + invocation_id: str + parent_invocation_id: str | None = None + sequence: int + kind: InvocationKind + identifier: str + version: str + visibility: Visibility + started_at: datetime + completed_at: datetime | None = None + duration_ms: int | None = None + status: InvocationStatus + summary: str + debug_input: JsonValue | None = None + debug_output: JsonValue | None = None + + +class ConversationActivityResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + projection: Literal["normal", "debug"] + invocations: tuple[InvocationActivityItem, ...] + + +def _authorize_conversation_read(session_id: str, user_id: str | None) -> None: + engine = get_engine() + with Session(engine) as session: + row = session.exec( + select(ChatConversation).where(ChatConversation.session_id == session_id) + ).one_or_none() + if row is None: + raise HTTPException(status_code=404, detail="Conversation not found") + if row.user_id != user_id: + raise HTTPException(status_code=403, detail="Not your conversation") + + +@router.get( + "/{session_id}/activity", + response_model=ConversationActivityResponse, + response_model_exclude_none=True, +) +def get_conversation_activity( + session_id: str, + *, + user_id: str | None = None, + debug: bool = False, +) -> ConversationActivityResponse: + """Return public activity normally and all sanitized activity in debug mode.""" + + _authorize_conversation_read(session_id, user_id) + records = SQLInvocationTraceRepository().list_for_conversation( + session_id, + include_private=debug, + ) + return ConversationActivityResponse( + projection="debug" if debug else "normal", + invocations=tuple( + InvocationActivityItem.model_validate( + record.model_dump( + exclude={ + "conversation_id", + *( + () + if debug + else ("debug_input", "debug_output") + ), + } + ) + ) + for record in records + ), + ) diff --git a/backend/chat/artifact_context.py b/backend/chat/artifact_context.py new file mode 100644 index 00000000..a9d0b8b4 --- /dev/null +++ b/backend/chat/artifact_context.py @@ -0,0 +1,93 @@ +"""Sanitized model-context summaries of compatible typed artifacts.""" + +from __future__ import annotations + +import asyncio + +from capabilities.artifacts import ( + ArtifactBase, + ChartArtifactRef, + HouseholdRef, + HouseholdResultRef, + PolicyScenarioRef, + SocietyAnalysisResultRef, +) +from persistence.capability_repository import SQLConversationCapabilityRepository + + +def sanitized_artifact_summary(artifact: ArtifactBase) -> dict[str, object]: + base: dict[str, object] = { + "artifact_id": artifact.artifact_id, + "artifact_type": getattr(artifact, "artifact_type"), + "schema_version": artifact.schema_version, + } + if isinstance(artifact, PolicyScenarioRef): + return { + **base, + "year": artifact.year, + "scenario_revision": artifact.scenario_revision, + "baseline": artifact.baseline, + "change_count": len(artifact.verified_changes), + } + if isinstance(artifact, HouseholdRef): + return { + **base, + "year": artifact.year, + "household_revision": artifact.household_revision, + "context_scope_id": artifact.context_scope_id, + "context_revision": artifact.context_revision, + "validated_inputs": tuple( + {"label": value.label, "value": value.value, "source": value.source} + for value in artifact.values + ), + } + if isinstance(artifact, HouseholdResultRef): + return { + **base, + "year": artifact.year, + "scenario_revision": artifact.scenario_revision, + "context_scope_id": artifact.context_scope_id, + "context_revision": artifact.context_revision, + "outputs": tuple(output.model_dump(mode="json") for output in artifact.outputs), + } + if isinstance(artifact, SocietyAnalysisResultRef): + return { + **base, + "year": artifact.year, + "scenario_revision": artifact.scenario_revision, + "default_profile_version": artifact.default_profile_version, + "outputs": tuple(output.model_dump(mode="json") for output in artifact.outputs), + "requested_output_issues": tuple( + issue.model_dump(mode="json") + for issue in artifact.requested_output_issues + ), + } + if isinstance(artifact, ChartArtifactRef): + return { + **base, + "year": artifact.year, + "scenario_revision": artifact.scenario_revision, + "chart_type": artifact.presentation.chart_type, + "title": artifact.presentation.title, + "source_result_artifact_id": artifact.source_result_artifact_id, + } + raise TypeError(f"Unsupported artifact model: {type(artifact).__name__}") + + +class RepositoryArtifactSummarySource: + def __init__(self, repository: SQLConversationCapabilityRepository) -> None: + self._repository = repository + + async def __call__( + self, + conversation_id: str, + ) -> tuple[dict[str, object], ...]: + artifacts = await asyncio.to_thread( + self._repository.list_artifacts, + conversation_id, + ) + return tuple( + sanitized_artifact_summary(artifact) + for artifact in artifacts + if artifact.schema_version == "1" + ) diff --git a/backend/chat/capability_runtime.py b/backend/chat/capability_runtime.py new file mode 100644 index 00000000..3d9c824d --- /dev/null +++ b/backend/chat/capability_runtime.py @@ -0,0 +1,21 @@ +"""Capability-oriented chat entry point behind the existing public adapter.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator + +from capabilities.application import get_capability_chat_application +from chat.events import CancellationProbe, ChatEvent +from chat.turn_input import ChatTurnInput + + +async def run_capability_chat_turn( + turn: ChatTurnInput, + *, + is_cancelled: CancellationProbe, +) -> AsyncIterator[ChatEvent]: + """Run one turn using the validated long-lived application composition.""" + + application = get_capability_chat_application() + async for event in application.run(turn, is_cancelled=is_cancelled): + yield event diff --git a/backend/chat/capability_service.py b/backend/chat/capability_service.py new file mode 100644 index 00000000..001346a9 --- /dev/null +++ b/backend/chat/capability_service.py @@ -0,0 +1,1328 @@ +"""Model-led chat flow with optional and required typed capabilities.""" + +from __future__ import annotations + +import asyncio +import json +import logging +from collections.abc import AsyncIterator, Awaitable, Callable +from typing import TypeVar +from uuid import uuid4 + +from pydantic import BaseModel + +from capabilities.context import ( + CapabilityContext, + ModelUsageSnapshot, +) +from capabilities.contracts import Completed +from capabilities.executor import ( + InvocationCancelled, + InvocationExecutor, + InvocationTraceValues, +) +from capabilities.registry import CapabilityRegistry +from capabilities.relevance import ( + AssessRelevanceInput, + ConversationExcerpt, + RelevanceAssessment, + RelevanceResult, +) +from conversation_context.models import ( + CapabilityInvocationReference, + ContextPatch, + ContextReduction, + ConversationContext, + FactRequirement, + PendingQuestion, + PendingQuestionStatus, + ReplacePendingQuestionsOperation, +) +from conversation_context.change_pipeline import ( + ContextChangeProposal, + ContextValidationOutcome, + ContextValidationStatus, + ContextValidationIssue, + ValidateContextChangeInput, +) +from conversation_context.projection import ContextProjection, project_context +from conversation_context.registry import FactDefinitionRegistry +from conversation_context.repository import ConversationContextRepository +from conversation_context.tools import ( + ApplyContextChangeInput, + ContextProposalStatus, + ProposeContextChangeInput, + ProposeContextChangeOutput, + ContextConversationExcerpt, + ReduceContextPatchInput, +) +from conversation_context.variable_resolution import ( + ResolveContextChangeInput, + ResolveContextChangeOutput, +) +from chat.events import ( + CancellationProbe, + ChatEvent, + ChatUsage, + InvocationActivity, + TextChunk, + TurnCancelled, + TurnCompleted, + TurnFailed, +) +from chat.model_port import ( + ConversationModel, + ConversationModelRequest, + ConversationModelResponse, + ModelCapabilityCall, + ModelUsage, +) +from chat.narration import ClarificationNarrationGuard, NumericalNarrationVerifier +from chat.turn_input import ChatTurnInput +from persistence.idempotency import ( + IdempotencyDecision, + SQLIdempotencyRepository, + request_fingerprint, +) +from tools.analysis_support import NumericalFact, VerifyNumericalResponseOutput + + +logger = logging.getLogger(__name__) +from tools.contracts import CallerType, Visibility + + +MANDATORY_CAPABILITY_CONTRACT = """ +The conversation is primary. Invoke zero or more public capabilities only when useful, +except for these required authoritative mappings: +- Government policy formulation, scope, formula, or calculation-method questions must + invoke policy_information before answering. +- A tax amount, benefit amount, benefit entitlement, or policy impact for a described + or retained household must invoke household_analysis, unless a compatible household result is + processed by analysis_follow_up. Invoke it even when required household details are + missing: deciding what is missing is part of the capability's job. +- Population-wide reform or benefit impacts must invoke society_analysis, unless a + compatible society result is processed by analysis_follow_up. +- A request spanning classes must invoke every applicable capability. +Never answer one of these required classes only from model memory. Outside these +classes, answer directly when no capability is needed. Capability results supply +authoritative facts but do not impose a fixed narrative layout. +For a household tax, benefit-impact, or entitlement request, your first action must be to +invoke household_analysis with the household evidence already available in the +conversation. Include every explicitly requested calculation metric in the +requested_outputs field; do not omit a named metric merely because it is common or +appears in the description. A user-facing response that asks for household details before this call +is invalid. Do not conduct a separate informal household intake first, even when you +already know more facts would be useful. The capability is the sole authority for +which household details require clarification and which defaults are safe. If a later +turn answers its clarification, invoke household_analysis with that new answer. The +server selects compatible waiting input from typed conversation context; never copy +or invent an internal invocation identifier. +Set start_new_invocation only when the user clearly starts a separate household +calculation while another household clarification remains pending. +The runtime associates retained household artifacts with stable context scopes and +selects the latest compatible artifact for the same household when it is active. Pass an explicit +referenced_household_id only when the user clearly selects a non-active historical +branch; never choose an unrelated retained household. +When household_analysis completes, follow its narration_requirement. Include every +string in its assumption_statements exactly once as a Markdown bullet under one +`Assumptions used` heading. Do not weave these statements into result prose, repeat +them in another section, or prefix every item with "I assumed". Report the completed +values in result.outputs in the same turn. Do not ask the user to choose outputs +after household_analysis has returned a completed result; any further output choice +is an optional follow-up after the completed values have been explained. +When society_analysis completes, cover every category in required_output_ids and +every successfully calculated requested output, and explain each requested-output +issue. Follow its narration_requirement, but choose a natural layout for the answer. +Do not infer policy mechanisms, eligibility rules, or causal explanations from a +calculation result alone. Use facts returned by policy_information for those claims; +otherwise describe only the calculated values, supplied inputs, explicit assumptions, +and clearly identified interpretation. + +When a capability returns needs_input, ask its supplied prompt as a conversational +clarification. Do not call that capability again in the same turn without new user +evidence, do not describe needs_input as a backend failure, and do not estimate the +requested result. When a capability returns failed or unsupported, explain only its +safe message or reason and do not substitute model-estimated facts or numbers. +""".strip() + + +ArtifactSummarySource = Callable[[str], Awaitable[tuple[dict[str, object], ...]]] +ActivityTaskResult = TypeVar("ActivityTaskResult") + + +async def _no_artifact_summaries( + conversation_id: str, +) -> tuple[dict[str, object], ...]: + del conversation_id + return () + + +def _with_model_response_guidance( + result: dict[str, object], +) -> dict[str, object]: + guided = dict(result) + status = guided.get("status") + if status == "needs_input": + guided["response_guidance"] = ( + "Ask the supplied prompt as a concise conversational clarification. " + "Do not retry this capability until a later user turn supplies new " + "evidence, and do not estimate the requested result." + ) + elif status == "unsupported": + guided["response_guidance"] = ( + "Explain the supplied reason without substituting an estimated result." + ) + elif status == "failed": + guided.setdefault( + "response_guidance", + "Explain only the safe message and do not estimate the requested result.", + ) + elif status == "completed": + guided["response_guidance"] = ( + "Answer the current request now using the completed value and its " + "verified outputs. Do not ask for input or an output choice after a " + "completed result; offer further analysis only after reporting what " + "was calculated." + ) + return guided + + +def capability_result_for_model( + result: dict[str, object], +) -> dict[str, object]: + """Project a capability outcome to the JSON the conversational model sees.""" + + projected = dict(result) + projected.pop("capability_invocation", None) + value = projected.get("value") + if isinstance(value, dict): + visible_value = dict(value) + # These fields support server-side numerical verification and its final + # fail-safe. Showing them to the model turns the fail-safe into a prose + # template and duplicates values already present in the typed result. + visible_value.pop("narration_facts", None) + visible_value.pop("narration_fallback", None) + projected["value"] = visible_value + return _with_model_response_guidance(projected) + + +def _model_capability_failure() -> dict[str, object]: + return { + "status": "failed", + "safe_message": ( + "This operation could not be completed, but other parts of " + "the request can continue." + ), + "error_code": "capability_invocation_failed", + } + + +class ModelCapabilityTraceValues(InvocationTraceValues): + """Represent exactly the capability JSON exchanged with the chat model.""" + + def input_value( + self, + *, + raw_input: object, + validated_input: BaseModel, + ) -> object: + del validated_input + return raw_input + + def output_value(self, validated_output: object) -> object: + if not isinstance(validated_output, BaseModel): + raise TypeError("Model capability trace output must be validated.") + return capability_result_for_model( + validated_output.model_dump(mode="json") + ) + + def failed_output(self) -> object: + return capability_result_for_model(_model_capability_failure()) + + def cancelled_output(self) -> object: + return {"status": "cancelled"} + + +_MODEL_CAPABILITY_TRACE_VALUES = ModelCapabilityTraceValues() + + +class ChatTurnService: + def __init__( + self, + *, + executor: InvocationExecutor, + capabilities: CapabilityRegistry, + model: ConversationModel, + idempotency: SQLIdempotencyRepository | None = None, + artifact_summaries: ArtifactSummarySource = _no_artifact_summaries, + context_repository: ConversationContextRepository | None = None, + fact_registry: FactDefinitionRegistry | None = None, + max_iterations: int = 20, + ) -> None: + self._executor = executor + self._capabilities = capabilities + self._model = model + self._idempotency = idempotency + self._artifact_summaries = artifact_summaries + self._context_repository = context_repository + self._fact_registry = fact_registry + self._max_iterations = max_iterations + self._narration = NumericalNarrationVerifier(executor) + self._clarification_narration = ClarificationNarrationGuard() + + async def run( + self, + turn: ChatTurnInput, + *, + is_cancelled: CancellationProbe, + context: CapabilityContext, + ) -> AsyncIterator[ChatEvent]: + context = context.with_current_user_message( + self._last_user_text(turn.messages) + ) + turn_id = turn.turn_id or uuid4().hex + fingerprint = request_fingerprint( + { + "messages": turn.messages, + "charts_mode": turn.charts_mode, + } + ) + if self._idempotency is not None: + receipt = self._idempotency.begin_turn( + conversation_id=turn.session_id, + turn_id=turn_id, + fingerprint=fingerprint, + ) + if receipt.decision is IdempotencyDecision.CONFLICT: + yield TurnFailed( + content="This turn identifier was already used for different input.", + session_id=turn.session_id, + stop_reason="idempotency_conflict", + usage=ChatUsage(), + turn_id=turn_id, + ) + return + if receipt.decision is IdempotencyDecision.IN_PROGRESS: + yield TurnCompleted( + content="This turn is already being processed.", + session_id=turn.session_id, + model=None, + route="capability", + outcome="in_progress", + stop_reason="idempotent_in_progress", + usage=ChatUsage(), + turn_id=turn_id, + ) + return + if receipt.decision is IdempotencyDecision.REPLAY: + replay_outcome = receipt.outcome or {} + content = str(replay_outcome.get("content", "")) + replay_model_value = replay_outcome.get("model") + replay_model = ( + replay_model_value + if isinstance(replay_model_value, str) + else None + ) + if content: + yield TextChunk(content) + yield TurnCompleted( + content=content, + session_id=turn.session_id, + model=replay_model, + route="capability", + outcome="replay", + stop_reason="idempotent_replay", + usage=ChatUsage(), + turn_id=turn_id, + ) + return + + usage = ChatUsage() + typed_context = ( + self._context_repository.load(turn.session_id) + if self._context_repository is not None + else ConversationContext.initial(turn.session_id) + ) + if self._fact_registry is not None: + self._fact_registry.restore_engine_definitions(typed_context) + context = context.with_conversation_context(typed_context) + usage_baseline = context.model_usage.snapshot() + model_name: str | None = None + trace_cursor = 0 + try: + if await is_cancelled(): + raise InvocationCancelled + relevance_task = asyncio.create_task( + self._assess_relevance(turn, context) + ) + async for event, trace_cursor in self._activity_while_running( + relevance_task, + context=context, + cursor=trace_cursor, + include_private=turn.debug, + is_cancelled=is_cancelled, + ): + if event is not None: + yield event + relevance = await relevance_task + if relevance.result is RelevanceResult.CLEARLY_OUT_OF_SCOPE: + usage = self._add_context_usage(usage, context, usage_baseline) + content = ( + "I can help with UK tax, benefit, and government-policy questions, " + "but this request is outside that supported scope." + ) + yield TextChunk(content) + completed = TurnCompleted( + content=content, + session_id=turn.session_id, + model=None, + route="capability", + outcome="out_of_scope", + stop_reason="out_of_scope", + usage=usage, + turn_id=turn_id, + ) + self._complete_turn(turn_id, fingerprint, completed) + yield completed + return + + if self._context_repository is not None and self._fact_registry is not None: + context_change_task = asyncio.create_task( + self._process_context_change(turn, context) + ) + async for event, trace_cursor in self._activity_while_running( + context_change_task, + context=context, + cursor=trace_cursor, + include_private=turn.debug, + is_cancelled=is_cancelled, + ): + if event is not None: + yield event + context_validation = await context_change_task + typed_context = context_validation.context + context = context.with_conversation_context(typed_context) + else: + context_validation = ContextValidationOutcome( + status=ContextValidationStatus.NO_CHANGE, + previous_revision=typed_context.revision, + context=typed_context, + ) + + conversation = [dict(message) for message in turn.messages] + capability_definitions = self._capabilities.descriptions_for( + CallerType.MODEL, + include_private=False, + ) + if context_validation.status is ContextValidationStatus.NEEDS_CLARIFICATION: + capability_definitions = () + artifact_summaries = await self._artifact_summaries(turn.session_id) + system = self._system_prompt( + artifact_summaries, + project_context(typed_context), + context_validation.issues, + ) + narration_facts: list[NumericalFact] = [] + assumption_statements: list[str] = [] + unresolved_fallbacks: list[str] = [] + completed_fallbacks: list[str] = [] + blocked_capabilities: dict[str, dict[str, object]] = {} + + for _iteration in range(self._max_iterations): + if await is_cancelled(): + raise InvocationCancelled + available_capability_definitions = tuple( + definition + for definition in capability_definitions + if definition["identifier"] not in blocked_capabilities + ) + response = await self._model.respond( + ConversationModelRequest( + messages=tuple(conversation), + system=system, + capabilities=available_capability_definitions, + ) + ) + offered_capability_ids = { + str(definition["identifier"]) + for definition in available_capability_definitions + } + permitted_calls = tuple( + call + for call in response.capability_calls + if call.capability_id in offered_capability_ids + ) + if permitted_calls != response.capability_calls: + fallback_text = response.text + if not fallback_text and not unresolved_fallbacks: + fallback_text = ( + "I couldn't safely apply all of that yet. Could you clarify " + "the values and which people or scenario they apply to?" + ) + response = response.model_copy( + update={ + "text": fallback_text, + "capability_calls": permitted_calls, + } + ) + model_name = response.model or model_name + usage = self._add_usage(usage, response.usage) + if not response.capability_calls: + deterministic_fallback = self._join_fallbacks( + [*completed_fallbacks, *unresolved_fallbacks] + ) + if unresolved_fallbacks and not narration_facts: + final_text = self._clarification_narration.finalize( + draft=response.text, + deterministic_fallback=deterministic_fallback, + ) + correction_usage = None + else: + final_text, correction_usage = await self._finalize_narration( + response, + tuple(narration_facts), + context, + deterministic_fallback=deterministic_fallback, + allow_redraft=not unresolved_fallbacks, + ) + if correction_usage is not None: + usage = self._add_usage(usage, correction_usage) + final_text = self._ensure_assumption_list( + final_text, + assumption_statements, + ) + activity, trace_cursor = self._activity_since( + context, + trace_cursor, + include_private=turn.debug, + ) + for event in activity: + yield event + usage = self._add_context_usage(usage, context, usage_baseline) + if final_text: + yield TextChunk(final_text) + completed = TurnCompleted( + content=final_text, + session_id=turn.session_id, + model=model_name, + route="capability", + outcome="completed", + stop_reason=response.stop_reason, + usage=usage, + turn_id=turn_id, + ) + self._complete_turn(turn_id, fingerprint, completed) + yield completed + return + + assistant_blocks: list[dict[str, object]] = [] + if response.text: + assistant_blocks.append({"type": "text", "text": response.text}) + for call in response.capability_calls: + assistant_blocks.append( + { + "type": "tool_use", + "id": call.call_id, + "name": call.capability_id, + "input": call.input, + } + ) + conversation.append({"role": "assistant", "content": assistant_blocks}) + result_blocks: list[dict[str, object]] = [] + for call in response.capability_calls: + invocation_task = asyncio.create_task( + self._invoke_model_capability( + call, + turn=turn, + turn_id=turn_id, + context=context, + blocked_result=blocked_capabilities.get( + call.capability_id + ), + ) + ) + async for event, trace_cursor in self._activity_while_running( + invocation_task, + context=context, + cursor=trace_cursor, + include_private=turn.debug, + is_cancelled=is_cancelled, + ): + if event is not None: + yield event + result, _status = await invocation_task + updated_context = await self._sync_pending_context( + capability_id=call.capability_id, + result=result, + context=context, + ) + if updated_context is not None: + context = context.with_conversation_context(updated_context) + system = self._system_prompt( + artifact_summaries, + project_context(updated_context), + context_validation.issues, + ) + narration_facts.extend(self._narration_facts(result)) + assumption_statements.extend( + self._assumption_statements(result) + ) + completed_fallback = self._completed_fallback(result) + if completed_fallback is not None: + completed_fallbacks.append(completed_fallback) + fallback = self._outcome_fallback(result) + if fallback is not None: + unresolved_fallbacks.append(fallback) + if _status in {"needs_input", "unsupported"}: + blocked_capabilities[call.capability_id] = result + result_blocks.append( + { + "type": "tool_result", + "tool_use_id": call.call_id, + "content": json.dumps( + capability_result_for_model(result), + default=str, + ), + } + ) + conversation.append({"role": "user", "content": result_blocks}) + + raise RuntimeError("Capability model loop exceeded its iteration limit.") + except InvocationCancelled: + activity, trace_cursor = self._activity_since( + context, + trace_cursor, + include_private=turn.debug, + ) + for event in activity: + yield event + usage = self._add_context_usage(usage, context, usage_baseline) + if self._idempotency is not None: + self._idempotency.fail_turn(turn_id=turn_id, fingerprint=fingerprint) + yield TurnCancelled( + session_id=turn.session_id, + model=model_name, + route="capability", + usage=usage, + turn_id=turn_id, + ) + except Exception: + logger.exception("Chat turn processing failed") + activity, trace_cursor = self._activity_since( + context, + trace_cursor, + include_private=turn.debug, + ) + for event in activity: + yield event + usage = self._add_context_usage(usage, context, usage_baseline) + if self._idempotency is not None: + self._idempotency.fail_turn(turn_id=turn_id, fingerprint=fingerprint) + yield TurnFailed( + content="Something went wrong while generating this response.", + session_id=turn.session_id, + stop_reason="error", + usage=usage, + billable=bool(usage.input_tokens or usage.output_tokens), + turn_id=turn_id, + ) + + async def _activity_while_running( + self, + task: asyncio.Task[ActivityTaskResult], + *, + context: CapabilityContext, + cursor: int, + include_private: bool, + is_cancelled: CancellationProbe, + ) -> AsyncIterator[tuple[InvocationActivity | None, int]]: + """Emit invocation updates while an operation is still executing.""" + + while True: + await asyncio.wait((task,), timeout=0.05) + activity, next_cursor = self._activity_since( + context, + cursor, + include_private=include_private, + ) + if activity: + for event in activity: + yield event, next_cursor + elif next_cursor != cursor: + # Advance across private records omitted from the projection. + yield None, next_cursor + cursor = next_cursor + if task.done(): + return + if await is_cancelled(): + task.cancel() + raise InvocationCancelled + + def _activity_since( + self, + context: CapabilityContext, + cursor: int, + *, + include_private: bool, + ) -> tuple[tuple[InvocationActivity, ...], int]: + all_events = self._executor.tracer.events_for_turn( + conversation_id=context.conversation_id, + turn_id=context.turn_id, + after_event_index=cursor, + include_private=True, + ) + if not all_events: + return (), cursor + projected = tuple( + InvocationActivity( + phase=event.phase, + record=( + event.record + if include_private + else event.record.model_copy( + update={"debug_input": None, "debug_output": None} + ) + ), + ) + for event in all_events + if include_private or event.record.visibility is Visibility.PUBLIC + ) + return projected, all_events[-1].event_index + + async def _assess_relevance( + self, + turn: ChatTurnInput, + context: CapabilityContext, + ) -> RelevanceAssessment: + current = self._last_user_text(turn.messages) + excerpts = tuple( + ConversationExcerpt( + role=str(message.get("role", "")), + content=self._content_text(message.get("content")), + ) + for message in turn.messages + ) + outcome = await self._executor.invoke_capability( + "conversation_relevance", + AssessRelevanceInput( + current_message=current, + conversation=excerpts, + context=( + project_context(context.conversation_context) + if context.conversation_context is not None + else None + ), + ), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(outcome, Completed) or not isinstance( + outcome.value, + RelevanceAssessment, + ): + raise TypeError("Conversation relevance returned an incompatible outcome.") + return outcome.value + + async def _process_context_change( + self, + turn: ChatTurnInput, + context: CapabilityContext, + ) -> ContextValidationOutcome: + if ( + self._context_repository is None + or self._fact_registry is None + or context.conversation_context is None + ): + raise RuntimeError("Conversation context processing is not configured.") + current = self._last_user_text(turn.messages) + excerpts = tuple( + ContextConversationExcerpt( + role=str(message.get("role", "")), + content=self._content_text(message.get("content")), + ) + for message in turn.messages + ) + prior = context.conversation_context + repair_issues: tuple[ContextValidationIssue, ...] = () + previous_proposal: ContextChangeProposal | None = None + for proposal_attempt in range(2): + raw_proposal = await self._executor.invoke_tool( + "propose_context_change", + ProposeContextChangeInput( + current_message=current, + conversation=excerpts, + context=project_context(prior), + fact_definitions=self._fact_registry.definitions(), + previous_proposal=previous_proposal, + repair_issues=repair_issues, + ), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(raw_proposal, ProposeContextChangeOutput): + raise TypeError("Context interpretation returned an incompatible output.") + if raw_proposal.status is ContextProposalStatus.NEEDS_CLARIFICATION: + return self._context_issue_outcome(prior, raw_proposal.issues) + + proposal = ContextChangeProposal( + expected_revision=raw_proposal.expected_revision, + candidate_entities=raw_proposal.candidate_entities, + changes=raw_proposal.changes, + focus=raw_proposal.focus, + ) + validation = await self._executor.invoke_tool( + "validate_context_change", + ValidateContextChangeInput( + context=prior, + proposal=proposal, + claims_resolved=False, + turn_id=context.turn_id, + evidence=current, + ), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(validation, ContextValidationOutcome): + raise TypeError("Context validation returned an incompatible output.") + if validation.status is ContextValidationStatus.NEEDS_CLARIFICATION: + if proposal_attempt == 0: + previous_proposal = proposal + repair_issues = validation.issues + continue + return validation + + validated = validation + if validation.claims_to_resolve: + try: + raw_resolution = await self._executor.invoke_tool( + "resolve_context_change", + ResolveContextChangeInput( + context=validation.context, + proposal=proposal, + validation_issues=validation.issues, + claims=validation.claims_to_resolve, + turn_id=context.turn_id, + evidence=current, + ), + caller=CallerType.RUNTIME, + context=context, + ) + except (TypeError, ValueError, RuntimeError): + issues = ( + ContextValidationIssue( + code="fact_resolution_failed", + path=("claims",), + message=( + "The current-message fact claim could not be " + "validated against an authoritative variable." + ), + evidence=current, + ), + ) + if proposal_attempt == 0: + previous_proposal = proposal + repair_issues = issues + continue + return self._context_issue_outcome(prior, issues) + if not isinstance(raw_resolution, ResolveContextChangeOutput): + raise TypeError("Fact-claim resolution returned an incompatible output.") + raw_validated = await self._executor.invoke_tool( + "validate_context_change", + ValidateContextChangeInput( + context=prior, + proposal=proposal, + resolution_patch=raw_resolution.patch, + claims_resolved=True, + turn_id=context.turn_id, + evidence=current, + ), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(raw_validated, ContextValidationOutcome): + raise TypeError("Context validation returned an incompatible output.") + validated = raw_validated + if validated.status is ContextValidationStatus.NEEDS_CLARIFICATION: + if proposal_attempt == 0: + previous_proposal = proposal + repair_issues = validated.issues + continue + return validated + + if validated.committable: + applied = await self._executor.invoke_tool( + "apply_context_change", + ApplyContextChangeInput(outcome=validated), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(applied, ConversationContext): + raise TypeError("Context application returned an incompatible output.") + validated = validated.model_copy(update={"context": applied}) + return validated + return self._context_issue_outcome(prior, repair_issues) + + @staticmethod + def _context_issue_outcome( + prior: ConversationContext, + issues: tuple[ContextValidationIssue, ...], + ) -> ContextValidationOutcome: + return ContextValidationOutcome( + status=ContextValidationStatus.NEEDS_CLARIFICATION, + previous_revision=prior.revision, + context=prior, + issues=issues, + ) + + async def _sync_pending_context( + self, + *, + capability_id: str, + result: dict[str, object], + context: CapabilityContext, + ) -> ConversationContext | None: + if ( + self._context_repository is None + or context.conversation_context is None + ): + return None + current = context.conversation_context + next_questions = current.pending_questions + if result.get("status") == "needs_input": + raw_reference = result.get("capability_invocation") + try: + reference = CapabilityInvocationReference.model_validate(raw_reference) + except Exception: + # Outcomes without a durable continuation must not clear an earlier + # valid pending question. + return None + if reference.capability_id != capability_id: + raise TypeError( + "Pending capability reference does not match the invoked capability." + ) + raw_requirements = result.get("fact_requirements") + requirements: list[FactRequirement] = [] + if isinstance(raw_requirements, list): + for item in raw_requirements: + try: + requirements.append(FactRequirement.model_validate(item)) + except Exception: + continue + prompt = result.get("prompt") + if requirements and isinstance(prompt, str) and prompt.strip(): + waiting = await context.waiting_invocations(capability_id) + matched = next( + ( + item + for item in waiting + if item.invocation_id == reference.invocation_id + ), + None, + ) + if matched is None: + raise RuntimeError( + "A pending question cannot reference a missing waiting invocation." + ) + if matched.reference() != reference: + raise RuntimeError( + "Pending context and waiting invocation metadata do not match." + ) + if matched.requirements != tuple(requirements): + raise RuntimeError( + "Pending context and waiting invocation requirements do not match." + ) + existing = next( + ( + question + for question in current.pending_questions + if question.capability_invocation is not None + and question.capability_invocation.invocation_id + == reference.invocation_id + ), + None, + ) + retained = tuple( + question + for question in current.pending_questions + if question is not existing + and not ( + question.capability_id == capability_id + and question.capability_invocation is None + and question.requirements == tuple(requirements) + ) + ) + next_questions = ( + *retained, + PendingQuestion( + question_id=( + existing.question_id if existing is not None else uuid4().hex + ), + capability_id=capability_id, + capability_invocation=reference, + prompt=prompt, + requirements=tuple(requirements), + created_turn_id=context.turn_id, + status=PendingQuestionStatus.AWAITING_ANSWER, + ), + ) + elif result.get("status") == "completed": + waiting_ids = { + item.invocation_id + for item in await context.waiting_invocations(capability_id) + } + next_questions = tuple( + question + for question in current.pending_questions + if not ( + question.capability_id == capability_id + and ( + ( + question.capability_invocation is not None + and question.capability_invocation.invocation_id + not in waiting_ids + ) + or ( + question.capability_invocation is None + and not waiting_ids + ) + ) + ) + ) + if next_questions == current.pending_questions: + return None + raw_reduction = await self._executor.invoke_tool( + "reduce_context_patch", + ReduceContextPatchInput( + context=current, + patch=ContextPatch( + expected_revision=current.revision, + operations=( + ReplacePendingQuestionsOperation( + questions=next_questions, + ), + ), + ), + turn_id=context.turn_id, + evidence="Capability requirement update.", + ), + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(raw_reduction, ContextReduction): + raise TypeError("Pending context reduction returned an incompatible output.") + self._context_repository.save( + raw_reduction.context, + expected_revision=raw_reduction.previous_revision, + ) + return raw_reduction.context + + async def _invoke_model_capability( + self, + call: ModelCapabilityCall, + *, + turn: ChatTurnInput, + turn_id: str, + context: CapabilityContext, + blocked_result: dict[str, object] | None = None, + ) -> tuple[dict[str, object], str]: + if blocked_result is not None: + replayed_result = dict(blocked_result) + status_value = replayed_result.get("status") + status = status_value if isinstance(status_value, str) else "failed" + return replayed_result, status + fingerprint = request_fingerprint( + {"capability": call.capability_id, "input": call.input} + ) + if self._idempotency is not None: + receipt = self._idempotency.begin_call( + conversation_id=turn.session_id, + turn_id=turn_id, + call_id=call.call_id, + operation_id=call.capability_id, + fingerprint=fingerprint, + ) + if receipt.decision is IdempotencyDecision.CONFLICT: + return {"status": "conflict"}, "conflict" + if receipt.decision is IdempotencyDecision.IN_PROGRESS: + return {"status": "in_progress"}, "failed" + if receipt.decision is IdempotencyDecision.REPLAY: + replay = receipt.outcome or {"status": "replay"} + return self._with_response_guidance(replay), "replay" + result: dict[str, object] + try: + outcome = await self._executor.invoke_capability( + call.capability_id, + call.input, + caller=CallerType.MODEL, + context=context, + trace_values=_MODEL_CAPABILITY_TRACE_VALUES, + ) + except InvocationCancelled: + raise + except Exception: + result = _model_capability_failure() + status = "failed" + else: + result = outcome.model_dump(mode="json") + status = outcome.status + result = self._with_response_guidance(result) + if self._idempotency is not None: + self._idempotency.complete_call( + call_id=call.call_id, + fingerprint=fingerprint, + outcome=capability_result_for_model(result), + ) + return result, status + + async def _finalize_narration( + self, + response: ConversationModelResponse, + facts: tuple[NumericalFact, ...], + context: CapabilityContext, + deterministic_fallback: str | None = None, + allow_redraft: bool = True, + ) -> tuple[str, ModelUsage | None]: + if not facts and deterministic_fallback is None: + return response.text, None + + correction_usage: ModelUsage | None = None + + async def redraft( + draft: str, + result: VerifyNumericalResponseOutput, + ) -> str: + nonlocal correction_usage + correction = await self._model.redraft_numerical( + draft=draft, + unsupported_claims=tuple( + claim.text for claim in result.unsupported_claims + ), + fact_summary=result.deterministic_fact_summary, + ) + correction_usage = correction.usage + return correction.text + + text = await self._narration.finalize( + draft=response.text, + facts=facts, + context=context, + redraft=redraft, + deterministic_fallback=deterministic_fallback, + allow_redraft=allow_redraft, + ) + return text, correction_usage + + @staticmethod + def _with_response_guidance( + result: dict[str, object], + ) -> dict[str, object]: + return _with_model_response_guidance(result) + + @staticmethod + def _outcome_fallback(result: dict[str, object]) -> str | None: + status = result.get("status") + if status == "needs_input": + prompt = result.get("prompt") + if isinstance(prompt, str) and prompt.strip(): + return prompt.strip() + if status == "unsupported": + reason = result.get("reason") + if isinstance(reason, str) and reason.strip(): + return reason.strip() + if status == "failed": + safe_message = result.get("safe_message") + if isinstance(safe_message, str) and safe_message.strip(): + return safe_message.strip() + return None + + @staticmethod + def _completed_fallback(result: dict[str, object]) -> str | None: + value = result.get("value") + if not isinstance(value, dict): + return None + fallback = value.get("narration_fallback") + if isinstance(fallback, str) and fallback.strip(): + return fallback.strip() + return None + + @staticmethod + def _join_fallbacks(fallbacks: list[str]) -> str | None: + unique = tuple(dict.fromkeys(fallbacks)) + return "\n\n".join(unique) if unique else None + + @staticmethod + def _assumption_statements(result: dict[str, object]) -> tuple[str, ...]: + value = result.get("value") + if not isinstance(value, dict): + return () + statements = value.get("assumption_statements") + if not isinstance(statements, list): + return () + return tuple( + item.strip() + for item in statements + if isinstance(item, str) and item.strip() + ) + + @staticmethod + def _ensure_assumption_list( + response: str, + statements: list[str], + ) -> str: + unique = tuple(dict.fromkeys(statements)) + if not unique: + return response + lines = response.splitlines() + heading_index = next( + ( + index + for index, line in enumerate(lines) + if line.lstrip("#").strip().strip("*_`").strip().casefold() + == "assumptions used" + ), + None, + ) + if heading_index is None: + section = [ + "### Assumptions used", + "", + *(f"- {item}" for item in unique), + ] + section_text = "\n".join(section) + prefix = response.rstrip() + return f"{prefix}\n\n{section_text}" if prefix else section_text + + section_end = next( + ( + index + for index in range(heading_index + 1, len(lines)) + if lines[index].lstrip().startswith("#") + ), + len(lines), + ) + existing_bullets = { + line.strip()[2:].strip().casefold() + for line in lines[heading_index + 1 : section_end] + if line.strip().startswith("- ") + } + missing = [item for item in unique if item.casefold() not in existing_bullets] + if not missing: + return response + lines[section_end:section_end] = [f"- {item}" for item in missing] + return "\n".join(lines) + + def _complete_turn( + self, + turn_id: str, + fingerprint: str, + completed: TurnCompleted, + ) -> None: + if self._idempotency is not None: + self._idempotency.complete_turn( + turn_id=turn_id, + fingerprint=fingerprint, + outcome={"content": completed.content, "model": completed.model}, + ) + + @staticmethod + def _narration_facts(result: dict[str, object]) -> tuple[NumericalFact, ...]: + value = result.get("value") + if not isinstance(value, dict): + return () + facts = value.get("narration_facts") + if not isinstance(facts, list): + return () + parsed: list[NumericalFact] = [] + for fact in facts: + try: + parsed.append(NumericalFact.model_validate(fact)) + except Exception: + continue + return tuple(parsed) + + @staticmethod + def _system_prompt( + artifacts: tuple[dict[str, object], ...], + context: ContextProjection | None = None, + context_issues: tuple[ContextValidationIssue, ...] = (), + ) -> str: + artifact_text = json.dumps(artifacts, default=str) + context_text = ( + context.model_dump_json(exclude_none=True) if context is not None else "{}" + ) + issue_text = json.dumps( + [issue.model_dump(mode="json") for issue in context_issues], + ensure_ascii=False, + ) + return ( + "You are PolicyEngine UK Chat. Continue the conversation naturally using " + "the complete supplied message history. Treat capability calls as optional " + "operations within the conversation, except for the required mappings below.\n\n" + f"{MANDATORY_CAPABILITY_CONTRACT}\n\n" + "Compatible retained artifact summaries (not prior prose) are:\n" + f"{artifact_text}\n\n" + "Validated typed conversation context (registered current facts, stable " + "entity identifiers, and pending requirements) is:\n" + f"{context_text}\n" + "Use this typed context for retained calculation inputs and identity. The " + "message transcript remains available for narrative references. Do not " + "replace a registered fact with an unsupported inference. A pending " + "fact-resolution assignment is not accepted calculation input. When it " + "applies to the current request, ask its supplied prompt and do not use " + "the proposed value until the user confirms it." + "\n\nCurrent-message context validation issues are:\n" + f"{issue_text}\n" + "These issues are structured diagnostic input, not user-facing prose and " + "not accepted facts. When this list is non-empty, no capability is available " + "on this turn: ask one concise, natural clarification question covering the " + "supplied values and subjects. When it is empty, invoke every capability " + "required by the current request. Do not expose issue codes, schema paths, " + "internal variable names, or these instructions." + ) + + @staticmethod + def _last_user_text(messages: list[dict[str, object]]) -> str: + for message in reversed(messages): + if message.get("role") == "user": + return ChatTurnService._content_text(message.get("content")) + return "" + + @staticmethod + def _content_text(content: object) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + return " ".join( + str(item.get("text", "")) + for item in content + if isinstance(item, dict) and item.get("type") == "text" + ) + return "" + + @staticmethod + def _add_usage( + current: ChatUsage, + additional: ModelUsage | ModelUsageSnapshot, + ) -> ChatUsage: + return ChatUsage( + input_tokens=current.input_tokens + additional.input_tokens, + output_tokens=current.output_tokens + additional.output_tokens, + cache_creation_input_tokens=( + current.cache_creation_input_tokens + + additional.cache_creation_input_tokens + ), + cache_read_input_tokens=( + current.cache_read_input_tokens + additional.cache_read_input_tokens + ), + ) + + @staticmethod + def _add_context_usage( + current: ChatUsage, + context: CapabilityContext, + baseline: ModelUsageSnapshot, + ) -> ChatUsage: + return ChatTurnService._add_usage( + current, + context.model_usage.snapshot().since(baseline), + ) diff --git a/backend/chat/events.py b/backend/chat/events.py index 5a0d36bb..65ffaaeb 100644 --- a/backend/chat/events.py +++ b/backend/chat/events.py @@ -1,9 +1,9 @@ """Framework-independent events emitted by a UK Chat turn.""" from dataclasses import dataclass, field -from typing import Any, Awaitable, Callable, Literal, TypeAlias +from typing import Awaitable, Callable, Literal, TypeAlias -from gateway.trace import GatewayTrace +from capabilities.tracing import InvocationRecord CancellationProbe: TypeAlias = Callable[[], Awaitable[bool]] @@ -25,13 +25,6 @@ def as_dict(self) -> dict[str, int]: } -@dataclass(frozen=True, slots=True) -class ToolStarted: - tool_name: str - tool_id: str - type: Literal["tool_start"] = field(default="tool_start", init=False) - - @dataclass(frozen=True, slots=True) class TextChunk: content: str @@ -39,25 +32,15 @@ class TextChunk: @dataclass(frozen=True, slots=True) -class ToolUsed: - tool_name: str - tool_id: str - tool_input: dict[str, Any] - type: Literal["tool_use"] = field(default="tool_use", init=False) - - -@dataclass(frozen=True, slots=True) -class ThinkingCompleted: - type: Literal["thinking_done"] = field(default="thinking_done", init=False) +class InvocationActivity: + """One sanitized invocation state change with optional structured debug values.""" - -@dataclass(frozen=True, slots=True) -class ToolCompleted: - tool_name: str - tool_id: str - status: Literal["success", "error"] - output: Any - type: Literal["tool_result"] = field(default="tool_result", init=False) + phase: Literal["started", "finished"] + record: InvocationRecord + type: Literal["invocation_activity"] = field( + default="invocation_activity", + init=False, + ) @dataclass(frozen=True, slots=True) @@ -69,7 +52,7 @@ class TurnCompleted: outcome: str | None stop_reason: str | None usage: ChatUsage - gateway_trace: GatewayTrace | None = None + turn_id: str | None = None type: Literal["done"] = field(default="done", init=False) @@ -86,7 +69,7 @@ class TurnFailed: stop_reason: str usage: ChatUsage billable: bool = False - gateway_trace: GatewayTrace | None = None + turn_id: str | None = None type: Literal["error"] = field(default="error", init=False) @@ -96,16 +79,13 @@ class TurnCancelled: model: str | None route: str usage: ChatUsage - gateway_trace: GatewayTrace | None = None + turn_id: str | None = None type: Literal["cancelled"] = field(default="cancelled", init=False) ChatEvent: TypeAlias = ( - ToolStarted - | TextChunk - | ToolUsed - | ThinkingCompleted - | ToolCompleted + TextChunk + | InvocationActivity | TurnCompleted | SuggestionsGenerated | TurnFailed diff --git a/backend/chat/model_port.py b/backend/chat/model_port.py new file mode 100644 index 00000000..c7bcd2ab --- /dev/null +++ b/backend/chat/model_port.py @@ -0,0 +1,163 @@ +"""Provider boundary for capability-oriented conversational model calls.""" + +from __future__ import annotations + +from typing import Any, Protocol + +from pydantic import BaseModel, ConfigDict, Field + +from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_async_client + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class ModelUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + +class ModelCapabilityCall(StrictModel): + call_id: str + capability_id: str + input: dict[str, object] = Field(default_factory=dict) + + +class ConversationModelRequest(StrictModel): + messages: tuple[dict[str, object], ...] + system: str + capabilities: tuple[dict[str, object], ...] + + +class ConversationModelResponse(StrictModel): + text: str = "" + capability_calls: tuple[ModelCapabilityCall, ...] = () + stop_reason: str | None = None + model: str | None = None + usage: ModelUsage = Field(default_factory=ModelUsage) + + +class ConversationModel(Protocol): + async def respond( + self, + request: ConversationModelRequest, + ) -> ConversationModelResponse: ... + + async def redraft_numerical( + self, + *, + draft: str, + unsupported_claims: tuple[str, ...], + fact_summary: str, + ) -> ConversationModelResponse: ... + + +class AnthropicConversationModel: + def __init__(self, model: str = DEFAULT_FAST_MODEL) -> None: + self._model = model + + async def respond( + self, + request: ConversationModelRequest, + ) -> ConversationModelResponse: + client = get_async_client() # type: ignore[no-untyped-call] + kwargs: dict[str, Any] = { + "model": self._model, + "max_tokens": 16_000, + "temperature": DEFAULT_TEMPERATURE, + "system": request.system, + "messages": list(request.messages), + } + if request.capabilities: + kwargs["tools"] = [ + { + "name": item["identifier"], + "description": ( + f"{item['description']} Required-use rule: {item['required_use']}" + ), + "input_schema": item["input_schema"], + } + for item in request.capabilities + ] + response = await client.messages.create(**kwargs) + text = "".join( + getattr(block, "text", "") + for block in response.content + if getattr(block, "type", None) == "text" + ) + calls = tuple( + ModelCapabilityCall( + call_id=block.id, + capability_id=block.name, + input=dict(block.input) if isinstance(block.input, dict) else {}, + ) + for block in response.content + if getattr(block, "type", None) == "tool_use" + ) + return ConversationModelResponse( + text=text, + capability_calls=calls, + stop_reason=getattr(response, "stop_reason", None), + model=self._model, + usage=self._usage(response), + ) + + async def redraft_numerical( + self, + *, + draft: str, + unsupported_claims: tuple[str, ...], + fact_summary: str, + ) -> ConversationModelResponse: + client = get_async_client() # type: ignore[no-untyped-call] + response = await client.messages.create( + model=self._model, + max_tokens=4_000, + temperature=DEFAULT_TEMPERATURE, + system=( + "Correct only the unsupported numerical claims in the draft while " + "preserving its natural wording and Markdown structure. Every number " + "in the corrected answer must repeat a verified fact, allowing only " + "equivalent units, scale, sign, or rounding. Do not calculate a new " + "total, difference, rate, or other derived value. Remove an unsupported " + "claim when no verified fact can replace it, and do not repeat any " + "expression listed as unsupported. Return only the corrected answer." + ), + messages=[ + { + "role": "user", + "content": ( + f"Draft:\n{draft}\n\nUnsupported expressions: " + f"{', '.join(unsupported_claims)}\n\nVerified facts:\n{fact_summary}" + ), + } + ], + ) + text = "".join( + getattr(block, "text", "") + for block in response.content + if getattr(block, "type", None) == "text" + ) + return ConversationModelResponse( + text=text, + stop_reason=getattr(response, "stop_reason", None), + model=self._model, + usage=self._usage(response), + ) + + @staticmethod + def _usage(response: Any) -> ModelUsage: + usage = getattr(response, "usage", None) + return ModelUsage( + input_tokens=getattr(usage, "input_tokens", 0), + output_tokens=getattr(usage, "output_tokens", 0), + cache_creation_input_tokens=getattr( + usage, + "cache_creation_input_tokens", + 0, + ), + cache_read_input_tokens=getattr(usage, "cache_read_input_tokens", 0), + ) diff --git a/backend/chat/model_selection.py b/backend/chat/model_selection.py deleted file mode 100644 index 0aa74283..00000000 --- a/backend/chat/model_selection.py +++ /dev/null @@ -1,199 +0,0 @@ -"""Chat model selection and small turn-classification helpers.""" - -import logging -import re -from typing import Any, List - -from config import ( - DEFAULT_COMPLEX_MODEL, - DEFAULT_FAST_MODEL, - DEFAULT_REASONING_MODEL, - FAST_MODEL_MAX_INPUT_TOKENS, -) -from prompts import SYSTEM_PROMPT - -logger = logging.getLogger(__name__) - - -_REFORM_CAPABLE_TOOLS = { - "validate_reform", - "run_household_simulation", - "run_society_simulation", - "compute_budgetary_impact", - "compute_decile_impacts", - "compute_inequality_metrics", - "compute_poverty_metrics", - "compute_winners_losers", -} - -_DISTRIBUTIONAL_OUTPUTS = { - "decile_impact", - "inequality_impact", - "marginal_rate", - "poverty_impact", - "winners_losers", -} - -_DISTRIBUTIONAL_TEXT_RE = re.compile( - r"\b(?:deciles?|quintiles?|distributional|winners?|losers?|poverty|" - r"inequality|gini|marginal tax|effective tax|marginal rate|effective rate)\b", - re.IGNORECASE, -) - -_POLICY_CONTEXT_RE = re.compile( - r"\b(?:income tax|tax(?:es)?|taxable|national insurance|ni|universal credit|" - r"uc|child benefit|pension credit|housing benefit|council tax|vat|" - r"capital gains tax|inheritance tax|benefit(?:s)?|allowance(?:s)?|" - r"personal allowance|threshold(?:s)?|basic rate|higher rate|" - r"additional rate|tax rate|marginal rate|effective rate|taper rate|" - r"band(?:s)?|taper(?:s)?|lha)\b", - re.IGNORECASE, -) - -_REFORM_INTENT_RE = re.compile( - r"\b(?:reform|raise|increase|decrease|cut|reduce|lower|change|replace|" - r"freeze|uprate|abolish|scrap|introduce|set)\b", - re.IGNORECASE, -) - -_POLICY_NUMERIC_CHANGE_RE = re.compile( - r"(?:\bby\s+\d+(?:\.\d+)?\s*%)" - r"|(?:\bfrom\s+\d+(?:\.\d+)?\s*%\s*to\s+\d+(?:\.\d+)?\s*%)" - r"|(?:\b\d+\s*pp\b)", - re.IGNORECASE, -) - - -# Upper-bound token cost of one attached image. Anthropic bills an image at -# roughly (width * height) / 750 tokens, capped near this for a full-size image. -# We don't have dimensions at selection time, so use the cap as a safe estimate -# rather than counting the base64 payload as if it were text (which inflated a -# ~500 KB photo to ~170k "tokens" and misrouted trivial questions to the -# expensive model). -_IMAGE_TOKEN_ESTIMATE = 1600 - - -def _estimate_message_tokens(messages: List[dict]) -> int: - char_count = 0 - image_count = 0 - for block in messages: - content = block.get("content", "") - if isinstance(content, str): - char_count += len(content) - elif isinstance(content, list): - for part in content: - if isinstance(part, dict) and part.get("type") == "image": - image_count += 1 - elif isinstance(part, dict) and part.get("type") == "text": - char_count += len(str(part.get("text", ""))) - else: - # tool_use / tool_result / other blocks — count serialised size. - char_count += len(str(part)) - else: - char_count += len(str(content)) - return char_count // 4 + image_count * _IMAGE_TOKEN_ESTIMATE - - -def _slot_value(slot: Any, attr: str) -> Any: - if isinstance(slot, dict): - return slot.get(attr) - return getattr(slot, attr, None) - - -def _detect_gateway_reasoning_signal(verdict: Any) -> str | None: - """Return a structured reform/distributional signal from the gateway.""" - if verdict is None: - return None - tool = getattr(verdict, "tool", None) - slots = getattr(verdict, "slots", []) or [] - for slot in slots: - name = str(_slot_value(slot, "name") or "") - kind = str(_slot_value(slot, "kind") or "tool_input") - source = str(_slot_value(slot, "source") or "") - if ( - tool in _REFORM_CAPABLE_TOOLS - and name == "reform" - and source in {"prompt", "default"} - ): - return f"gateway:{tool}:reform" - if kind == "output" and name in _DISTRIBUTIONAL_OUTPUTS: - return f"gateway:{tool}:output:{name}" - return None - - -def _detect_reform_signal(text: str) -> str | None: - """Return a narrow text reform/distributional signal, if any. - - Opening turns should prefer the gateway's structured plan. This text check - is a fallback for follow-ups and gateway fail-safe turns, so it requires - policy context before generic change language can trigger the reasoning - model. - """ - if not text: - return None - distributional = _DISTRIBUTIONAL_TEXT_RE.search(text) - if distributional: - return distributional.group(0) - - if not _POLICY_CONTEXT_RE.search(text): - return None - - reform_intent = _REFORM_INTENT_RE.search(text) - if reform_intent: - return reform_intent.group(0) - - numeric_change = _POLICY_NUMERIC_CHANGE_RE.search(text) - if numeric_change: - return numeric_change.group(0) - return None - - -def select_chat_model( - messages: List[dict], - *, - charts_mode: bool = False, - gateway_verdict: Any = None, -) -> str: - signal = _detect_gateway_reasoning_signal(gateway_verdict) - if signal is None: - signal = _detect_reform_signal(last_user_text(messages)) - if signal: - logger.info("[MODEL] Routed to reasoning model (signal=%r)", signal) - return DEFAULT_REASONING_MODEL - - if charts_mode: - logger.info("[MODEL] Routed to reasoning model (charts_mode=True)") - return DEFAULT_REASONING_MODEL - - estimated_input_tokens = ( - _estimate_message_tokens(messages) - + len(SYSTEM_PROMPT) // 4 - ) - if estimated_input_tokens > FAST_MODEL_MAX_INPUT_TOKENS: - return DEFAULT_COMPLEX_MODEL - return DEFAULT_FAST_MODEL - - -def last_user_text(conversation: List[dict]) -> str: - """Latest user message as plain text (flattening any image+text content).""" - for msg in reversed(conversation): - if msg.get("role") != "user": - continue - content = msg.get("content", "") - if isinstance(content, str): - return content - if isinstance(content, list): # [image block, text block, ...] - return " ".join( - str(b.get("text", "")) for b in content if isinstance(b, dict) - ).strip() - return "" - - -def is_followup(conversation: List[dict]) -> bool: - """True once the conversation contains an assistant turn — i.e. this is not - the opening user message. The gateway runs only on the opening turn; a - single-message classifier can't see the context a follow-up depends on, and - a user's reply to a partial/needs_plan prompt should flow straight to - compute (which it does, because that turn now has a prior assistant reply). - """ - return any(msg.get("role") == "assistant" for msg in conversation) diff --git a/backend/chat/narration.py b/backend/chat/narration.py new file mode 100644 index 00000000..1b97f92d --- /dev/null +++ b/backend/chat/narration.py @@ -0,0 +1,120 @@ +"""Free-form narration safeguards for verified quantitative facts.""" + +from __future__ import annotations + +import re +from collections.abc import Awaitable, Callable + +from capabilities.context import CapabilityContext +from capabilities.executor import InvocationExecutor +from tools.analysis_support import ( + NumericalFact, + VerifyNumericalResponseOutput, +) +from tools.contracts import CallerType + + +Redraft = Callable[[str, VerifyNumericalResponseOutput], Awaitable[str]] + + +class ClarificationNarrationGuard: + """Permit natural questions while rejecting unsupported quantitative answers.""" + + _ordered_list_prefix = re.compile(r"(?m)^\s*\d+[.)]\s+") + _substantive_number = re.compile( + r"(?:[£$€]\s*\d|\d[\d,]*(?:\.\d+)?\s*(?:%|percent\b)|" + r"(? str: + without_list_ordinals = self._ordered_list_prefix.sub("", draft) + if self._substantive_number.search(without_list_ordinals): + return deterministic_fallback or "" + return draft or deterministic_fallback or "" + + +class NumericalNarrationVerifier: + """Allow free-form prose, one correction, then a verified fact summary.""" + + def __init__(self, executor: InvocationExecutor) -> None: + self._executor = executor + + async def finalize( + self, + *, + draft: str, + facts: tuple[NumericalFact, ...], + context: CapabilityContext, + redraft: Redraft, + deterministic_fallback: str | None = None, + allow_redraft: bool = True, + ) -> str: + first = await self._verify(draft, facts, context) + if first.supported: + return draft + if not allow_redraft and deterministic_fallback is not None: + return deterministic_fallback + + corrected = await redraft(draft, first) + second = await self._verify(corrected, facts, context) + if second.supported: + return corrected + sanitized = self._without_unsupported_sentences( + corrected, + tuple(claim.text for claim in second.unsupported_claims), + ) + if sanitized and sanitized != corrected: + sanitized_result = await self._verify(sanitized, facts, context) + if sanitized_result.supported: + return sanitized + return deterministic_fallback or second.deterministic_fact_summary + + @staticmethod + def _without_unsupported_sentences( + draft: str, + unsupported_expressions: tuple[str, ...], + ) -> str: + """Remove only prose units that still contain unsupported numbers.""" + + expressions = tuple( + expression.strip() + for expression in unsupported_expressions + if expression.strip() + ) + if not expressions: + return draft + kept_lines: list[str] = [] + for line in draft.splitlines(): + if not any(expression in line for expression in expressions): + kept_lines.append(line) + continue + indentation = line[: len(line) - len(line.lstrip())] + sentences = re.split(r"(?<=[.!?])\s+", line.strip()) + retained = [ + sentence + for sentence in sentences + if not any(expression in sentence for expression in expressions) + ] + if retained: + kept_lines.append(indentation + " ".join(retained)) + return re.sub(r"\n{3,}", "\n\n", "\n".join(kept_lines)).strip() + + async def _verify( + self, + draft: str, + facts: tuple[NumericalFact, ...], + context: CapabilityContext, + ) -> VerifyNumericalResponseOutput: + output = await self._executor.invoke_tool( + "verify_numerical_response", + { + "draft": draft, + "facts": [fact.model_dump() for fact in facts], + }, + caller=CallerType.RUNTIME, + context=context, + ) + if not isinstance(output, VerifyNumericalResponseOutput): + raise TypeError("Numerical verifier returned an incompatible output.") + return output diff --git a/backend/chat/orchestrator.py b/backend/chat/orchestrator.py deleted file mode 100644 index e68bc269..00000000 --- a/backend/chat/orchestrator.py +++ /dev/null @@ -1,731 +0,0 @@ -"""Framework-independent model/tool orchestration for one UK Chat turn. - -Builds the per-turn plan (gateway), selects the model, streams from Anthropic, -runs tools in parallel, records usage, and emits typed internal events. HTTP, -SSE projection, and billing are handled by the public and eval adapters. -""" - -import asyncio -import json -import logging -from contextlib import asynccontextmanager -from dataclasses import replace -from functools import partial -from typing import Any, Dict, List - -import httpx -from policyengine_observability import annotate -from policyengine_observability import asegment -from policyengine_observability import mark_ttft_attribute -from policyengine_observability import operation -from policyengine_observability import record_error -from policyengine_observability import record_event -from policyengine_observability import segment - -from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_async_client -from gateway import run_gateway, serialise_plan_for_system -from gateway.clarifications import render_clarification -from gateway.proposals import ( - ProposalSigningError, - append_proposal_marker, - proposal_payload_from_verdict, - resume_gateway_proposal, - strip_proposal_markers_from_conversation, -) -from gateway.trace import gateway_trace_from_verdict -from observability.segments import SegmentName -from tools.context import new_tool_context -from tools.dispatch import execute_tool - -from chat.model_selection import is_followup, last_user_text, select_chat_model -from chat.events import ( - CancellationProbe, - ChatUsage, - SuggestionsGenerated, - TextChunk, - ThinkingCompleted, - ToolCompleted, - ToolStarted, - ToolUsed, - TurnCancelled, - TurnCompleted, - TurnFailed, -) -from chat.suggestions import generate_followup_suggestions -from chat.system_blocks import ( - build_lightweight_system_blocks, - build_system_blocks, - serialise_tool_result, - tool_defs_for_anthropic, -) -from chat.turn_input import ChatTurnInput - -logger = logging.getLogger(__name__) - -# Soft cap on tool-use iterations within a single /chat/message stream. -# An "iteration" is one round-trip to Anthropic that may include tool calls. -# We lowered this from 60 to keep runaway agents from hanging the Vercel -# proxy (which times out the SSE connection and surfaces as "Failed to fetch" -# in the browser). When hit, we emit a user-facing fallback message and a -# `done` event with stop_reason="iteration_cap" instead of cutting off mid-stream. -# NOTE: this is per-request. The /chat/message "continue" flow re-enters this -# loop with a fresh budget, but the prior tool transcript is already in the -# conversation so the model resumes mid-thought rather than restarting. -MAX_ITERATIONS = 30 -MAX_TOOL_RESULT_CHARS = 15000 - - -def _serialise_tool_result_for_model(tool_result: Any) -> str: - """Return bounded, valid JSON for a tool result sent back to the model.""" - - result_json = serialise_tool_result(tool_result) - if len(result_json) <= MAX_TOOL_RESULT_CHARS: - return result_json - - if isinstance(tool_result, dict): - data_key = next( - ( - key - for key, value in tool_result.items() - if isinstance(value, list) and len(value) > 5 - ), - None, - ) - if data_key: - from engine.serialization import explore_tabular_data - - data_array = tool_result[data_key] - processed = { - **{key: value for key, value in tool_result.items() if key != data_key}, - "note": ( - f"Large '{data_key}' array ({len(data_array)} rows) - " - "showing first 20 with column metadata" - ), - "exploration": explore_tabular_data(data_array), - data_key: data_array[:20], - } - result_json = serialise_tool_result(processed) - if len(result_json) <= MAX_TOOL_RESULT_CHARS: - return result_json - - fallback = { - "status": ( - "error" - if tool_result.get("error") - else tool_result.get("status", "success") - ), - "result_id": tool_result.get("result_id"), - "note": ( - "Tool result exceeded the model context limit. Use a narrower " - "discovery query or request a more specific derivative output." - ), - } - return serialise_tool_result( - {key: value for key, value in fallback.items() if value is not None} - ) - - return serialise_tool_result( - { - "status": "truncated", - "note": ( - "Tool result exceeded the model context limit. Use a narrower " - "query or request a more specific output." - ), - } - ) - - -def _user_facing_error_message(session_id: str) -> str: - """Generic text for terminal SSE `error` events. - - Raw exception strings (Anthropic SDK, httpx, Supabase) can embed internal - URLs, file paths, and provider payloads, so they must never be streamed to - the client. The full exception and traceback stay in the server logs; the - session id gives users a correlation reference to quote when reporting the - problem. - """ - return ( - "Something went wrong while generating this response. Please try " - f"again — if the problem persists, quote session {session_id} when " - "reporting it." - ) - - -async def run_chat_turn( - turn: ChatTurnInput, - *, - is_cancelled: CancellationProbe, -): - """Run one model/tool turn and emit framework-independent events.""" - - ttft_recorded = False - terminal_stop_reason: str | None = None - captured_exception: Exception | None = None - captured_traceback = "" - conversation = turn.messages.copy() - iteration = 0 - total_input_tokens = 0 - total_output_tokens = 0 - total_cache_read_input_tokens = 0 - total_cache_creation_input_tokens = 0 - recent_tool_calls: List[str] = [] - tool_call_counts: Dict[str, int] = {} - last_tool_error: str | None = None - verdict = None - route = "compute" - model: str | None = None - tool_context = new_tool_context(turn_id=turn.session_id) - - def usage() -> ChatUsage: - return ChatUsage( - input_tokens=total_input_tokens, - output_tokens=total_output_tokens, - cache_creation_input_tokens=total_cache_creation_input_tokens, - cache_read_input_tokens=total_cache_read_input_tokens, - ) - - def annotate_turn(stop_reason: str | None) -> None: - annotate( - model=model, - stop_reason=stop_reason, - iterations=iteration, - tool_calls=sum(tool_call_counts.values()), - input_tokens=total_input_tokens, - output_tokens=total_output_tokens, - cache_read_input_tokens=total_cache_read_input_tokens, - cache_creation_input_tokens=total_cache_creation_input_tokens, - ) - - def cancelled_event() -> TurnCancelled: - nonlocal terminal_stop_reason - terminal_stop_reason = "client_disconnected" - annotate_turn(terminal_stop_reason) - record_event( - "chat.client_disconnected", - session_id=turn.session_id, - route=route, - model=model, - iterations=iteration, - tool_calls=sum(tool_call_counts.values()), - ) - return TurnCancelled( - session_id=turn.session_id, - model=model, - route=route, - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) - - @asynccontextmanager - async def capture_turn_errors(): - nonlocal captured_exception - nonlocal captured_traceback - nonlocal terminal_stop_reason - - try: - yield - except Exception as exc: - import traceback - - captured_exception = exc - captured_traceback = traceback.format_exc() - if terminal_stop_reason is None: - terminal_stop_reason = "error" - try: - annotate_turn(terminal_stop_reason) - record_error(exc, handled=True, status_code=500) - except Exception: - logger.exception("[CHAT] Failed to record observability for chat turn error") - - try: - async with operation( - "chat.turn", - flavor="chat", - session_id=turn.session_id, - ), capture_turn_errors(): - annotate(session_id=turn.session_id, charts_mode=turn.charts_mode) - if await is_cancelled(): - yield cancelled_event() - return - - loop = asyncio.get_running_loop() - proposal_error: str | None = None - if is_followup(conversation): - try: - verdict = await loop.run_in_executor( - None, - partial( - resume_gateway_proposal, - conversation, - session_id=turn.session_id, - ), - ) - except ProposalSigningError: - logger.warning( - "[GATEWAY] Session %s supplied an invalid proposal marker", - turn.session_id, - ) - proposal_error = ( - "I couldn’t verify the earlier reform proposal. Please " - "restate the policy change you want me to model." - ) - except Exception as exc: - verdict = getattr(exc, "gateway_verdict", None) - raise - else: - try: - async with asegment(SegmentName.GATEWAY_CLASSIFY): - verdict = await loop.run_in_executor( - None, run_gateway, last_user_text(conversation) - ) - except Exception as exc: - verdict = getattr(exc, "gateway_verdict", None) - raise - - if verdict is not None: - route = verdict.route - - clarification: str | None = proposal_error - if verdict is not None and verdict.outcome == "needs_plan": - clarification = render_clarification(verdict) - if clarification is None: - logger.warning( - "[GATEWAY] Unrenderable reasons; failing open to compute: %s", - verdict.gating_reasons, - ) - unresolved_reform = ( - verdict.reform_intent is not None - and verdict.reform_assessment is None - ) - verdict = replace( - verdict, - outcome="ready", - route="compute", - gating_reasons=[], - ) - route = "compute" - if unresolved_reform: - tool_context.require_approved_reform = True - tool_context.approved_reform = None - elif any( - reason.code == "confirm_reform" - for reason in verdict.gating_reasons - ): - clarification = append_proposal_marker( - clarification, - proposal_payload_from_verdict(verdict), - session_id=turn.session_id, - source_prompt=last_user_text(conversation), - ) - - if clarification is not None: - route = "lightweight" - if await is_cancelled(): - yield cancelled_event() - return - mark_ttft_attribute() - ttft_recorded = True - terminal_stop_reason = "gateway_clarification" - annotate_turn(terminal_stop_reason) - yield TextChunk(clarification) - yield TurnCompleted( - content=clarification, - session_id=turn.session_id, - model=None, - route=route, - outcome="needs_plan", - stop_reason=terminal_stop_reason, - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) - return - - if verdict is not None: - if ( - verdict.execution_plan is not None - and verdict.execution_plan.approved_reform is not None - ): - tool_context.approved_reform = dict( - verdict.execution_plan.approved_reform - ) - tool_context.require_approved_reform = True - - conversation = strip_proposal_markers_from_conversation(conversation) - client = get_async_client() - - annotate(gateway_route=route) - if verdict is not None: - annotate(gateway_outcome=verdict.outcome, gateway_tool=verdict.tool) - - with segment( - SegmentName.MODEL_SELECT, - route=route, - charts_mode=turn.charts_mode, - ): - model = ( - DEFAULT_FAST_MODEL - if route == "lightweight" - else select_chat_model( - conversation, - charts_mode=turn.charts_mode, - gateway_verdict=verdict, - ) - ) - - if route == "lightweight": - tools = [] - with segment(SegmentName.SYSTEM_BUILD, route=route): - system_blocks = build_lightweight_system_blocks(verdict) - else: - with segment(SegmentName.TOOL_SCHEMA_BUILD): - tools = tool_defs_for_anthropic() - gateway_plan = None - if verdict is not None: - with segment(SegmentName.GATEWAY_PLAN_SERIALIZE): - gateway_plan = serialise_plan_for_system(verdict) - with segment( - SegmentName.SYSTEM_BUILD, - route=route, - charts_mode=turn.charts_mode, - ): - system_blocks = build_system_blocks( - charts_mode=turn.charts_mode, - gateway_plan=gateway_plan, - ) - annotate(model=model) - - logger.info( - f"[CHAT] Session {turn.session_id}: {len(conversation)} messages" - f"{' [CHARTS MODE]' if turn.charts_mode else ''}" - f"{f' [GATEWAY {verdict.outcome}]' if verdict is not None else ''}" - ) - - while iteration < MAX_ITERATIONS: - if await is_cancelled(): - yield cancelled_event() - return - - iteration += 1 - async with asegment( - SegmentName.MODEL_ITERATION, - iteration=iteration, - model=model, - ): - tool_uses = [] - assistant_content = "" - last_stop_reason: str | None = None - max_retries = 2 - - for attempt in range(max_retries + 1): - try: - stream_kwargs: Dict[str, Any] = { - "model": model, - "max_tokens": 16000, - "temperature": DEFAULT_TEMPERATURE, - "system": system_blocks, - "messages": conversation, - } - if tools: - stream_kwargs["tools"] = tools - - async with ( - asegment( - SegmentName.MODEL_STREAM, - iteration=iteration, - model=model, - ), - client.messages.stream(**stream_kwargs) as stream, - ): - announced_tools: set[str] = set() - async for event in stream: - event_type = type(event).__name__ - if event_type == "RawContentBlockStartEvent": - block = event.content_block - if ( - block.type == "tool_use" - and block.id not in announced_tools - ): - announced_tools.add(block.id) - yield ToolStarted(block.name, block.id) - elif event_type == "RawContentBlockDeltaEvent": - delta = event.delta - if delta.type == "text_delta" and delta.text: - if not ttft_recorded: - mark_ttft_attribute() - ttft_recorded = True - assistant_content += delta.text - yield TextChunk(delta.text) - elif event_type == "RawMessageStartEvent": - event_usage = getattr(event.message, "usage", None) - if event_usage: - total_input_tokens += getattr( - event_usage, "input_tokens", 0 - ) - total_cache_read_input_tokens += getattr( - event_usage, - "cache_read_input_tokens", - 0, - ) - total_cache_creation_input_tokens += getattr( - event_usage, - "cache_creation_input_tokens", - 0, - ) - elif event_type == "RawMessageDeltaEvent": - event_usage = getattr(event, "usage", None) - if event_usage: - total_output_tokens += getattr( - event_usage, "output_tokens", 0 - ) - - final = await stream.get_final_message() - last_stop_reason = getattr(final, "stop_reason", None) - for block in final.content: - if block.type != "tool_use": - continue - if not tools: - logger.warning( - "[CHAT] Dropping unexpected tool_use with no tools " - f"sent: {block.name}" - ) - continue - tool_input = ( - block.input if isinstance(block.input, dict) else {} - ) - tool_use = { - "id": block.id, - "name": block.name, - "input": tool_input, - } - tool_uses.append(tool_use) - yield ToolUsed( - tool_name=block.name, - tool_id=block.id, - tool_input=tool_input, - ) - break - except ( - httpx.ReadError, - httpx.RemoteProtocolError, - httpx.ConnectError, - ) as exc: - logger.warning( - "[CHAT] Anthropic stream error " - f"(attempt {attempt + 1}/{max_retries + 1}): {exc}" - ) - if attempt == max_retries: - raise - tool_uses = [] - assistant_content = "" - await asyncio.sleep(1) - - if tool_uses and assistant_content.strip(): - yield ThinkingCompleted() - - if not tool_uses: - terminal_stop_reason = last_stop_reason - annotate_turn(terminal_stop_reason) - yield TurnCompleted( - content=assistant_content, - session_id=turn.session_id, - model=model, - route=route, - outcome=verdict.outcome if verdict else None, - stop_reason=terminal_stop_reason, - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) - if route == "compute" and last_stop_reason in ( - "end_turn", - "stop_sequence", - None, - ): - last_user_message = next( - ( - message["content"] - for message in reversed(turn.messages) - if message.get("role") == "user" - ), - "", - ) - if isinstance(last_user_message, list): - last_user_message = " ".join( - str(block.get("text", "")) - for block in last_user_message - if isinstance(block, dict) - ) - async with asegment(SegmentName.SUGGESTIONS): - suggestions = await generate_followup_suggestions( - last_user_message=str(last_user_message), - assistant_answer=assistant_content, - ) - if suggestions: - yield SuggestionsGenerated(suggestions) - return - - signature = ",".join( - sorted( - f"{item['name']}:{json.dumps(item['input'], sort_keys=True)}" - for item in tool_uses - ) - ) - recent_tool_calls.append(signature) - if len(recent_tool_calls) > 3: - recent_tool_calls.pop(0) - if len(set(recent_tool_calls)) == 1: - terminal_stop_reason = "loop_detected" - annotate_turn(terminal_stop_reason) - yield TurnFailed( - content=( - "Agent appears to be stuck in a loop. " - "Please try rephrasing your question." - ), - session_id=turn.session_id, - stop_reason=terminal_stop_reason, - usage=usage(), - billable=True, - gateway_trace=gateway_trace_from_verdict(verdict), - ) - return - - assistant_message: Dict[str, Any] = { - "role": "assistant", - "content": [], - } - if assistant_content: - assistant_message["content"].append( - {"type": "text", "text": assistant_content} - ) - for tool_use in tool_uses: - assistant_message["content"].append( - { - "type": "tool_use", - "id": tool_use["id"], - "name": tool_use["name"], - "input": tool_use["input"], - } - ) - conversation.append(assistant_message) - - async def execute_tool_async(tool_use): - loop = asyncio.get_event_loop() - async with asegment( - SegmentName.TOOL_EXECUTE, - iteration=iteration, - tool=tool_use["name"], - ): - result = await loop.run_in_executor( - None, - partial( - execute_tool, - tool_use["name"], - tool_use["input"], - context=tool_context, - ), - ) - return tool_use, result - - tasks = [ - asyncio.ensure_future(execute_tool_async(tool_use)) - for tool_use in tool_uses - ] - completed_tools = {} - for future in asyncio.as_completed(tasks): - tool_use, result = await future - if await is_cancelled(): - yield cancelled_event() - return - completed_tools[tool_use["id"]] = result - tool_call_counts[tool_use["name"]] = ( - tool_call_counts.get(tool_use["name"], 0) + 1 - ) - if isinstance(result, dict): - error = result.get("error") or result.get("stderr") - if error: - error_text = str(error).strip() - if error_text: - last_tool_error = error_text.splitlines()[-1][:120] - status = ( - "error" - if isinstance(result, dict) and result.get("error") - else "success" - ) - yield ToolCompleted( - tool_name=tool_use["name"], - tool_id=tool_use["id"], - status=status, - output=result, - ) - - tool_results = [] - for tool_use in tool_uses: - tool_results.append( - { - "type": "tool_result", - "tool_use_id": tool_use["id"], - "content": _serialise_tool_result_for_model( - completed_tools[tool_use["id"]] - ), - } - ) - conversation.append({"role": "user", "content": tool_results}) - else: - if tool_call_counts: - tried = "ran " + ", ".join( - f"`{name}` {count}×" - for name, count in tool_call_counts.items() - ) - else: - tried = "didn't complete any tool calls" - error_clause = ( - f', last attempt errored with "{last_tool_error}"' - if last_tool_error - else "" - ) - fallback = ( - "\n\nI'm spending more iterations than expected on this without " - f"converging. Here's what I tried: {tried}{error_clause}. " - "Could you (a) rephrase the question or (b) try a more specific " - "scenario?" - ) - if len(fallback) > 600: - fallback = fallback[:597] + "..." - terminal_stop_reason = "iteration_cap" - annotate_turn(terminal_stop_reason) - yield TextChunk(fallback) - yield TurnCompleted( - content=assistant_content + fallback, - session_id=turn.session_id, - model=model, - route=route, - outcome=verdict.outcome if verdict else None, - stop_reason=terminal_stop_reason, - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) - - if captured_exception is not None: - logger.error( - f"[CHAT] Session {turn.session_id} exception: " - f"{captured_exception}\n{captured_traceback}" - ) - yield TurnFailed( - content=_user_facing_error_message(turn.session_id), - session_id=turn.session_id, - stop_reason="error", - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) - except Exception as exc: - import traceback - - logger.error( - f"[CHAT] Session {turn.session_id} exception: {exc}\n{traceback.format_exc()}" - ) - yield TurnFailed( - content=_user_facing_error_message(turn.session_id), - session_id=turn.session_id, - stop_reason="error", - usage=usage(), - gateway_trace=gateway_trace_from_verdict(verdict), - ) diff --git a/backend/chat/public_service.py b/backend/chat/public_service.py index 2d85fe97..7a8c0642 100644 --- a/backend/chat/public_service.py +++ b/backend/chat/public_service.py @@ -15,22 +15,17 @@ CancellationProbe, SuggestionsGenerated, TextChunk, - ThinkingCompleted, - ToolCompleted, - ToolStarted, - ToolUsed, + InvocationActivity, TurnCancelled, TurnCompleted, TurnFailed, ) -from chat.orchestrator import run_chat_turn +from chat.capability_runtime import run_capability_chat_turn from chat.schemas import ChatRequest -from chat.system_blocks import serialise_tool_result from chat.turn_input import ChatTurnInput, prepare_turn_input logger = logging.getLogger(__name__) -MAX_PUBLIC_TOOL_RESULT_CHARS = 5000 class InsufficientCredit(Exception): @@ -47,10 +42,21 @@ def _record_turn_usage( session_id: str, model: str | None, usage: ChatUsage, + turn_id: str | None = None, ) -> dict | None: if not billing_enabled(): return None + if turn_id is not None: + try: + from persistence.idempotency import SQLIdempotencyRepository + + if not SQLIdempotencyRepository().claim_billing(turn_id): + return None + except Exception as exc: + logger.warning("[CHAT] Failed to claim idempotent billing: %s", exc) + return None + try: with segment(SegmentName.BILLING_RECORD_USAGE): return billing.record_usage( @@ -65,37 +71,18 @@ def _record_turn_usage( def _public_payload(event: ChatEvent, billing: dict | None = None) -> dict | None: - if isinstance(event, ToolStarted): - return { - "type": event.type, - "tool_name": event.tool_name, - "tool_id": event.tool_id, - } if isinstance(event, TextChunk): return {"type": event.type, "content": event.content} - if isinstance(event, ToolUsed): - return { - "type": event.type, - "tool_name": event.tool_name, - "tool_id": event.tool_id, - "tool_input": event.tool_input, - "status": "pending", - } - if isinstance(event, ThinkingCompleted): - return {"type": event.type} - if isinstance(event, ToolCompleted): - result = serialise_tool_result(event.output) - summary = ( - result[:MAX_PUBLIC_TOOL_RESULT_CHARS] + "..." - if len(result) > MAX_PUBLIC_TOOL_RESULT_CHARS - else result - ) + if isinstance(event, InvocationActivity): + record = event.record return { "type": event.type, - "tool_name": event.tool_name, - "tool_id": event.tool_id, - "status": event.status, - "result_summary": summary, + "phase": event.phase, + "invocation": record.model_dump( + mode="json", + exclude={"conversation_id"}, + exclude_none=True, + ), } if isinstance(event, TurnCompleted): return { @@ -125,7 +112,7 @@ async def _stream_public_chat( *, is_cancelled: CancellationProbe, ) -> AsyncIterator[str]: - async for event in run_chat_turn(turn, is_cancelled=is_cancelled): + async for event in run_capability_chat_turn(turn, is_cancelled=is_cancelled): billing = None if isinstance(event, TurnCompleted): billing = _record_turn_usage( @@ -133,6 +120,7 @@ async def _stream_public_chat( session_id=event.session_id, model=event.model, usage=event.usage, + turn_id=event.turn_id, ) elif isinstance(event, TurnCancelled): _record_turn_usage( @@ -140,6 +128,7 @@ async def _stream_public_chat( session_id=event.session_id, model=event.model, usage=event.usage, + turn_id=event.turn_id, ) elif isinstance(event, TurnFailed) and event.billable: _record_turn_usage( @@ -147,6 +136,7 @@ async def _stream_public_chat( session_id=event.session_id, model=None, usage=event.usage, + turn_id=event.turn_id, ) payload = _public_payload(event, billing) diff --git a/backend/chat/routes.py b/backend/chat/routes.py index 29ca7c4d..9be3b991 100644 --- a/backend/chat/routes.py +++ b/backend/chat/routes.py @@ -23,7 +23,7 @@ def generate_title(request: TitleRequest): @limiter.limit(CHAT_IP_LIMIT) async def chat_message(request: Request, chat_request: ChatRequest): # `request` is the Starlette Request that slowapi's @limiter.limit decorators - # require; the parsed body is `chat_request`. Delegate to the orchestrator. + # require; the parsed body is `chat_request`. Delegate to the chat service. try: stream = await start_public_chat( chat_request, diff --git a/backend/chat/schemas.py b/backend/chat/schemas.py index b16dcc94..8baf8f38 100644 --- a/backend/chat/schemas.py +++ b/backend/chat/schemas.py @@ -13,8 +13,10 @@ class ChatMessage(BaseModel): class ChatRequest(BaseModel): messages: List[ChatMessage] session_id: str | None = None + turn_id: str | None = None user_id: str | None = None charts_mode: bool = False + debug: bool = False # Optional image attached to the latest user message. Sent as raw base64 # (no `data:image/...;base64,` prefix) plus a media type like `image/png`. # When present, the backend converts the latest user message into a diff --git a/backend/chat/system_blocks.py b/backend/chat/system_blocks.py deleted file mode 100644 index 58966c91..00000000 --- a/backend/chat/system_blocks.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Assembly of model system blocks and the lightweight gateway system.""" - -import json -from typing import Any, List - -from gateway import gateway_writer_directive -from prompts import ( - CHARTS_MODE_DIRECTIVE, - DEFAULT_SCOPE_DESCRIPTOR, - SYSTEM_PROMPT, - lightweight_system, -) -from tools.definitions import TOOL_DEFINITIONS - -SCOPE_DESCRIPTOR = DEFAULT_SCOPE_DESCRIPTOR -LIGHTWEIGHT_SYSTEM = lightweight_system(SCOPE_DESCRIPTOR) - - -def tool_defs_for_anthropic(): - """Convert our TOOL_DEFINITIONS to Anthropic SDK format. - Mark the last tool with cache_control so the system prompt + all tools - are cached across requests (prompt caching).""" - defs = [] - for i, t in enumerate(TOOL_DEFINITIONS): - d = { - "name": t["name"], - "description": t["description"], - "input_schema": t["input_schema"], - } - if i == len(TOOL_DEFINITIONS) - 1: - d["cache_control"] = {"type": "ephemeral"} - defs.append(d) - return defs - - -def serialise_tool_result(result: Any) -> str: - return json.dumps(result, ensure_ascii=False, default=str) - - -def build_system_blocks( - charts_mode: bool = False, - gateway_plan: str | None = None, -) -> List[dict]: - """System prompt + optional per-turn directives. - - The system prompt is marked with cache_control so it persists across - requests. Per-turn directives are appended after the cache breakpoint so - toggling them does not invalidate the cached block. - """ - blocks: List[dict] = [{ - "type": "text", - "text": SYSTEM_PROMPT, - "cache_control": {"type": "ephemeral"}, - }] - if charts_mode: - blocks.append({"type": "text", "text": CHARTS_MODE_DIRECTIVE}) - if gateway_plan: - blocks.append({"type": "text", "text": gateway_plan}) - return blocks - - -def build_lightweight_system_blocks(verdict) -> List[dict]: - """Lean model payload for irrelevant, out-of-scope, or partial turns. - - `needs_plan` turns terminate through the deterministic clarification - renderer before this function can be called. - """ - blocks: List[dict] = [{ - "type": "text", - "text": LIGHTWEIGHT_SYSTEM, - "cache_control": {"type": "ephemeral"}, - }] - directive = gateway_writer_directive(verdict) - if directive: - blocks.append({"type": "text", "text": directive}) - return blocks diff --git a/backend/chat/turn_input.py b/backend/chat/turn_input.py index d0f0baca..f0315871 100644 --- a/backend/chat/turn_input.py +++ b/backend/chat/turn_input.py @@ -18,7 +18,9 @@ class InvalidChatRequest(ValueError): class ChatTurnInput: messages: list[dict[str, Any]] session_id: str + turn_id: str = "" charts_mode: bool = False + debug: bool = False def prepare_turn_input(chat_request: ChatRequest) -> ChatTurnInput: @@ -66,5 +68,7 @@ def prepare_turn_input(chat_request: ChatRequest) -> ChatTurnInput: return ChatTurnInput( messages=deduplicated, session_id=chat_request.session_id or str(uuid.uuid4()), + turn_id=chat_request.turn_id or str(uuid.uuid4()), charts_mode=chat_request.charts_mode, + debug=chat_request.debug, ) diff --git a/backend/config/__init__.py b/backend/config/__init__.py index eb3ed9b3..773257fc 100644 --- a/backend/config/__init__.py +++ b/backend/config/__init__.py @@ -3,7 +3,7 @@ Split by sub-responsibility (models / sampling / clients) so each file owns one concern, but this package re-exports the public surface so callers do `from config import X` regardless of which sub-module owns X. Kept import-light -(no heavy deps) so the gateway and the eval harness can import it without +(no heavy deps) so the capability runtime and evaluation harness can import it without dragging in the chat route. """ diff --git a/backend/config/models.py b/backend/config/models.py index a8bf9f99..d8b86e6d 100644 --- a/backend/config/models.py +++ b/backend/config/models.py @@ -2,7 +2,7 @@ import os -# Fast model: the default chat model under the token threshold, plus the gateway +# Fast model: the default chat model under the token threshold, plus capability # classifier, titling, and follow-up suggestions. DEFAULT_FAST_MODEL = os.environ.get("ANTHROPIC_FAST_MODEL", "claude-haiku-4-5") diff --git a/backend/config/sampling.py b/backend/config/sampling.py index 6a406041..1af5d45c 100644 --- a/backend/config/sampling.py +++ b/backend/config/sampling.py @@ -2,7 +2,7 @@ import os -# 0 = deterministic, which is what the compute loop, titling, the gateway +# 0 = deterministic, which is what the compute loop, titling, capability # classifier, and the evals all want. DEFAULT_TEMPERATURE = float(os.environ.get("ANTHROPIC_TEMPERATURE", "0")) diff --git a/backend/conversation_context/__init__.py b/backend/conversation_context/__init__.py new file mode 100644 index 00000000..ebba0801 --- /dev/null +++ b/backend/conversation_context/__init__.py @@ -0,0 +1,97 @@ +"""Typed conversational facts shared by the chat runtime and capabilities.""" + +from conversation_context.models import ( + BooleanFactValue, + ClaimedMoneyValue, + ContextEntity, + ContextEntityCandidate, + ContextFact, + ContextFocusCandidate, + ContextPatch, + ContextScope, + ConversationContext, + EntityKind, + EntityReferenceFactValue, + EntityReferencesFactValue, + ExplicitAbsenceAssertion, + FactClaim, + FactDecision, + FactDecisionStatus, + FactClaimRelationship, + FactResolutionAssignment, + FactResolutionStatus, + FactProvenance, + FactRequirement, + IntegerFactValue, + MoneyFactValue, + MoneyPeriod, + PendingQuestion, + PendingQuestionStatus, + PendingFactResolution, + PendingFactResolutionResponse, + PresentAssertion, + SetFactOperation, + TextFactValue, + TextSetFactValue, +) +from conversation_context.registry import ( + FactDefinition, + FactDefinitionRegistry, + FactUpdatePolicy, + build_default_fact_registry, +) +from conversation_context.reducer import ContextReducer +from conversation_context.change_pipeline import ( + ContextChangeApplier, + ContextChangeProposal, + ContextChangeValidator, + ContextValidationOutcome, + ContextValidationStatus, + ContextValidationIssue, +) + +__all__ = [ + "BooleanFactValue", + "ClaimedMoneyValue", + "ContextEntity", + "ContextEntityCandidate", + "ContextFact", + "ContextFocusCandidate", + "ContextPatch", + "ContextChangeApplier", + "ContextChangeProposal", + "ContextChangeValidator", + "ContextReducer", + "ContextScope", + "ConversationContext", + "ContextValidationOutcome", + "ContextValidationStatus", + "ContextValidationIssue", + "EntityKind", + "EntityReferenceFactValue", + "EntityReferencesFactValue", + "ExplicitAbsenceAssertion", + "FactClaim", + "FactDecision", + "FactDecisionStatus", + "FactClaimRelationship", + "FactResolutionAssignment", + "FactResolutionStatus", + "FactDefinition", + "FactDefinitionRegistry", + "FactUpdatePolicy", + "FactProvenance", + "FactRequirement", + "IntegerFactValue", + "MoneyFactValue", + "MoneyPeriod", + "PendingQuestion", + "PendingQuestionStatus", + "PendingFactResolution", + "PendingFactResolutionResponse", + "PresentAssertion", + "SetFactOperation", + "TextFactValue", + "TextSetFactValue", + "build_default_fact_registry", +] diff --git a/backend/conversation_context/change_pipeline.py b/backend/conversation_context/change_pipeline.py new file mode 100644 index 00000000..3832e29f --- /dev/null +++ b/backend/conversation_context/change_pipeline.py @@ -0,0 +1,1190 @@ +"""Validation and atomic application of model-proposed context changes.""" + +from __future__ import annotations + +from decimal import Decimal +from enum import Enum +import re + +from pydantic import BaseModel, ConfigDict, ValidationError + +from conversation_context.models import ( + ClaimedMoneyValue, + ConfirmPendingFactResolutionOperation, + ContextEntityCandidate, + ContextFocusCandidate, + ContextOperation, + ContextPatch, + ContextReduction, + ConversationContext, + EnsureEntityOperation, + ExplicitAbsenceAssertion, + FactClaim, + FactClaimFieldUpdate, + FactClaimRelationship, + FactAssertion, + FactDecision, + FactDecisionStatus, + FactResolutionStatus, + IntegerFactValue, + MoneyFactValue, + MoneyPeriod, + PendingFactResolution, + PendingFactResolutionResponse, + PendingResolutionAction, + PresentAssertion, + ProposedContextChange, + SetFactOperation, + SetFocusOperation, + TextFactValue, + TextSetFactValue, +) +from conversation_context.reducer import ContextReducer +from conversation_context.registry import FactDefinitionRegistry +from conversation_context.repository import ConversationContextRepository +from conversation_context.quantities import MonetaryExpressionParser + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class ContextChangeProposal(StrictModel): + """One model-authored interpretation with no persistence operations.""" + + expected_revision: int + candidate_entities: tuple[ContextEntityCandidate, ...] = () + changes: tuple[ProposedContextChange, ...] = () + focus: ContextFocusCandidate | None = None + + @property + def claims(self) -> tuple[FactClaim, ...]: + return tuple(change for change in self.changes if isinstance(change, FactClaim)) + + @property + def proposal_responses(self) -> tuple[PendingFactResolutionResponse, ...]: + return tuple( + change + for change in self.changes + if isinstance(change, PendingFactResolutionResponse) + ) + + def change_index(self, change_id: str) -> int | None: + return next( + ( + index + for index, change in enumerate(self.changes) + if ( + isinstance(change, FactClaim) + and change.claim_id == change_id + ) + or ( + isinstance(change, PendingFactResolutionResponse) + and change.response_id == change_id + ) + ), + None, + ) + + +class ContextValidationIssue(StrictModel): + """Machine-readable feedback for interpretation repair and debug output.""" + + code: str + message: str + path: tuple[str, ...] = () + claim_index: int | None = None + operation_index: int | None = None + evidence: str | None = None + + +class SemanticClaimReview(StrictModel): + """One model-authored semantic verdict retained in validation debug output.""" + + claim_id: str + supported: bool + reason: str + evidence: str + + +class ContextValidationStatus(str, Enum): + RESOLUTION_REQUIRED = "resolution_required" + READY_TO_COMMIT = "ready_to_commit" + NEEDS_CLARIFICATION = "needs_clarification" + NO_CHANGE = "no_change" + + +class ContextValidationOutcome(StrictModel): + """The complete deterministic result for one proposed context change.""" + + status: ContextValidationStatus + previous_revision: int + context: ConversationContext + generated_operations: tuple[ContextOperation, ...] = () + claims_to_resolve: tuple[FactClaim, ...] = () + decisions: tuple[FactDecision, ...] = () + issues: tuple[ContextValidationIssue, ...] = () + semantic_reviews: tuple[SemanticClaimReview, ...] = () + + @property + def committable(self) -> bool: + return self.status is ContextValidationStatus.READY_TO_COMMIT + + +class ValidateContextChangeInput(StrictModel): + context: ConversationContext + proposal: ContextChangeProposal + resolution_patch: ContextPatch | None = None + claims_resolved: bool = False + turn_id: str + evidence: str + + +class ContextChangeValidator: + """Validate one complete model proposal without choosing semantic meaning.""" + + _accepted_statuses = { + FactDecisionStatus.ACCEPTED, + FactDecisionStatus.IGNORED, + FactDecisionStatus.SUPERSEDED, + } + + def __init__( + self, + reducer: ContextReducer, + registry: FactDefinitionRegistry, + monetary_parser: MonetaryExpressionParser | None = None, + ) -> None: + self._reducer = reducer + self._registry = registry + self._monetary_parser = monetary_parser or MonetaryExpressionParser() + + def validate( + self, + validation_input: ValidateContextChangeInput, + *, + semantic_issues: tuple[ContextValidationIssue, ...] = (), + ) -> ContextValidationOutcome: + prior = validation_input.context + proposal = validation_input.proposal + if proposal.expected_revision != prior.revision: + return self._invalid( + prior, + issues=( + ContextValidationIssue( + code="stale_proposal_revision", + path=("proposal", "expected_revision"), + message=( + "The interpreted context revision does not match the " + "currently loaded conversation context." + ), + ), + ), + ) + + grounding_issues = self._grounding_issues( + validation_input.evidence, + proposal, + ) + + generated, claims_to_resolve, planning_issues = self._plan( + prior, + proposal, + current_message=validation_input.evidence, + ) + + generated_patch = ContextPatch( + expected_revision=prior.revision, + operations=generated, + ) + provisional = self._reduce( + prior, + generated_patch, + turn_id=validation_input.turn_id, + evidence=validation_input.evidence, + ) + if isinstance(provisional, tuple): + return self._invalid( + prior, + generated_operations=generated, + claims_to_resolve=claims_to_resolve, + issues=provisional, + ) + direct_issues = self._decision_issues(provisional.decisions) + aggregate_issues = self._aggregate_issues(provisional.context) + validation_issues = ( + *semantic_issues, + *grounding_issues, + *planning_issues, + *direct_issues, + *aggregate_issues, + ) + if validation_issues: + return self._invalid( + prior, + generated_operations=generated, + claims_to_resolve=claims_to_resolve, + decisions=provisional.decisions, + issues=validation_issues, + ) + + if claims_to_resolve and not validation_input.claims_resolved: + resolution_issues = tuple( + self._claim_issue( + self._claim_change_index(prior, proposal, claim), + claim, + "authoritative_resolution_required", + ( + "The model-proposed claim requires an authoritative semantic " + "mapping before it can be validated as a context change." + ), + ) + for claim in claims_to_resolve + ) + return ContextValidationOutcome( + status=ContextValidationStatus.RESOLUTION_REQUIRED, + previous_revision=prior.revision, + context=provisional.context, + generated_operations=generated, + claims_to_resolve=claims_to_resolve, + decisions=provisional.decisions, + issues=resolution_issues, + ) + + if claims_to_resolve and validation_input.resolution_patch is None: + return self._invalid( + prior, + generated_operations=generated, + claims_to_resolve=claims_to_resolve, + decisions=provisional.decisions, + issues=( + ContextValidationIssue( + code="fact_claims_not_resolved", + path=("proposal", "changes"), + message=( + "The proposal contains claims requiring authoritative " + "resolution but no validated resolution operations." + ), + ), + ), + ) + + resolution_patch = validation_input.resolution_patch + if resolution_patch is not None: + if resolution_patch.expected_revision != provisional.context.revision: + return self._invalid( + prior, + generated_operations=generated, + claims_to_resolve=claims_to_resolve, + decisions=provisional.decisions, + issues=( + ContextValidationIssue( + code="stale_resolution_revision", + path=("resolution_patch", "expected_revision"), + message=( + "Fact-resolution operations were not derived from the " + "validated provisional context." + ), + ), + ), + ) + combined_operations = (*generated, *resolution_patch.operations) + combined_patch = ContextPatch( + expected_revision=prior.revision, + operations=combined_operations, + ) + combined = self._reduce( + prior, + combined_patch, + turn_id=validation_input.turn_id, + evidence=validation_input.evidence, + ) + if isinstance(combined, tuple): + return self._invalid( + prior, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + issues=combined, + ) + else: + combined_operations = generated + combined = provisional + + decision_issues = self._decision_issues(combined.decisions) + if decision_issues: + return self._invalid( + prior, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + decisions=combined.decisions, + issues=decision_issues, + ) + final_issues = self._aggregate_issues(combined.context) + if final_issues: + return self._invalid( + prior, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + decisions=combined.decisions, + issues=final_issues, + ) + if combined.context.revision == prior.revision: + return ContextValidationOutcome( + status=ContextValidationStatus.NO_CHANGE, + previous_revision=prior.revision, + context=prior, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + decisions=combined.decisions, + ) + if combined.context.revision != prior.revision + 1: + return self._invalid( + prior, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + decisions=combined.decisions, + issues=( + ContextValidationIssue( + code="invalid_validated_revision", + path=("context", "revision"), + message="A current-message update must create exactly one revision.", + ), + ), + ) + return ContextValidationOutcome( + status=ContextValidationStatus.READY_TO_COMMIT, + previous_revision=prior.revision, + context=combined.context, + generated_operations=combined_operations, + claims_to_resolve=claims_to_resolve, + decisions=combined.decisions, + ) + + def _plan( + self, + context: ConversationContext, + proposal: ContextChangeProposal, + *, + current_message: str, + ) -> tuple[ + tuple[ContextOperation, ...], + tuple[FactClaim, ...], + tuple[ContextValidationIssue, ...], + ]: + operations: list[ContextOperation] = [ + EnsureEntityOperation( + reference=candidate.reference, + kind=candidate.kind, + aliases=candidate.aliases, + relationship_to_user=candidate.relationship_to_user, + ) + for candidate in proposal.candidate_entities + ] + claims_to_resolve: list[FactClaim] = [] + issues: list[ContextValidationIssue] = [] + accepted_claim_ids: set[str] = set() + planned_facts: dict[tuple[str, str, str, str], FactAssertion] = {} + active_scope = context.focus.scope_id + + pending_resolutions = { + pending.proposal_id: pending for pending in context.pending_fact_resolutions + } + seen_response_ids: set[str] = set() + valid_responses: list[PendingFactResolutionResponse] = [] + for response in proposal.proposal_responses: + index = proposal.change_index(response.response_id) + assert index is not None + valid = True + if response.proposal_id in seen_response_ids: + valid = False + issues.append( + ContextValidationIssue( + code="duplicate_pending_resolution_response", + message=( + "The proposal responds to the same pending fact resolution " + "more than once." + ), + path=("proposal", "changes", str(index)), + ) + ) + elif response.proposal_id not in pending_resolutions: + valid = False + issues.append( + ContextValidationIssue( + code="unknown_pending_resolution_response", + message=( + "The proposal response does not identify an existing " + "pending fact-resolution proposal." + ), + path=("proposal", "changes", str(index)), + ) + ) + else: + pending = pending_resolutions[response.proposal_id] + if ( + response.action is not PendingResolutionAction.SUPPLY + and not self._text_is_grounded(response.evidence, current_message) + ): + valid = False + issues.append( + ContextValidationIssue( + code="uncited_pending_resolution_response", + message=( + "The pending-resolution response evidence is absent " + "from the exact current message." + ), + path=("proposal", "changes", str(index), "evidence"), + evidence=response.evidence, + ) + ) + if response.action is PendingResolutionAction.SUPPLY: + merged, response_issues = self._merge_pending_claim( + pending, + response, + current_message=current_message, + change_index=index, + ) + issues.extend(response_issues) + valid = valid and merged is not None and not response_issues + if valid and merged is not None: + claims_to_resolve.append(merged) + elif pending.status is not FactResolutionStatus.AWAITING_CONFIRMATION: + valid = False + issues.append( + ContextValidationIssue( + code="pending_resolution_not_confirmable", + message=( + "This pending record requests clarification. Accept or " + "reject applies only to a calculated assignment awaiting " + "confirmation; supply the missing source-claim fields." + ), + path=("proposal", "changes", str(index)), + ) + ) + if valid: + valid_responses.append(response) + seen_response_ids.add(response.proposal_id) + + for claim in proposal.claims: + index = proposal.change_index(claim.claim_id) + assert index is not None + scope_id = claim.scope_id or active_scope + if scope_id is None: + issues.append( + self._claim_issue( + index, + claim, + "missing_claim_scope", + "The fact claim has no active or explicit context scope.", + ) + ) + continue + if claim.relationship is not FactClaimRelationship.DIRECT: + claims_to_resolve.append(claim) + continue + definition_key = claim.definition_key + definition_version = claim.definition_version + inherited_period = None + if isinstance(claim.value, ClaimedMoneyValue) and claim.value.period is None: + inherited = self._periodless_direct_binding(context, claim, scope_id) + if inherited is None: + claims_to_resolve.append(claim) + continue + definition_key, definition_version, inherited_period = inherited + if definition_key is None: + if isinstance(claim.value, ClaimedMoneyValue): + claims_to_resolve.append(claim) + else: + issues.append( + self._claim_issue( + index, + claim, + "unmapped_fact_claim", + ( + "A non-monetary direct claim must select one supplied " + "registered fact definition." + ), + ) + ) + continue + try: + self._registry.get(definition_key, definition_version) + except KeyError: + issues.append( + self._claim_issue( + index, + claim, + "unregistered_fact_claim", + "The fact claim selects an unregistered fact definition.", + ) + ) + continue + assertion: FactAssertion + if claim.explicit_absence: + assertion = ExplicitAbsenceAssertion() + elif isinstance(claim.value, ClaimedMoneyValue): + period = claim.value.period or inherited_period + if period is None: + issues.append( + self._claim_issue( + index, + claim, + "unresolved_fact_period", + ( + "The monetary fact claim has no validated period and " + "must remain available for authoritative resolution." + ), + ) + ) + continue + assertion = PresentAssertion( + value=MoneyFactValue( + amount=claim.value.amount, + period=period, + currency=claim.value.currency, + ) + ) + elif claim.value is not None: + assertion = PresentAssertion(value=claim.value) + else: + issues.append( + self._claim_issue( + index, + claim, + "missing_fact_claim_value", + "The fact claim contains neither a value nor explicit absence.", + ) + ) + continue + fact_target = ( + definition_key, + definition_version, + self._resolve_subject_reference( + context, + claim.subject_references[0], + ) + or claim.subject_references[0].strip().casefold(), + scope_id, + ) + previous_assertion = planned_facts.get(fact_target) + if previous_assertion is not None: + if previous_assertion != assertion: + issues.append( + self._claim_issue( + index, + claim, + "conflicting_fact_claims", + ( + "Two claims in the same proposal assign different values " + "to the same registered fact and subject." + ), + ) + ) + else: + accepted_claim_ids.add(claim.claim_id) + continue + planned_facts[fact_target] = assertion + accepted_claim_ids.add(claim.claim_id) + operations.append( + SetFactOperation( + definition_key=definition_key, + definition_version=definition_version, + subject_reference=claim.subject_references[0], + scope_id=scope_id, + assertion=assertion, + correction=claim.correction, + ) + ) + + resolved_claim_ids = {claim.claim_id for claim in claims_to_resolve} + issue_indexes = { + int(issue.path[2]) + for issue in issues + if len(issue.path) >= 3 + and issue.path[:2] == ("proposal", "changes") + and issue.path[2].isdigit() + } + for claim in proposal.claims: + index = proposal.change_index(claim.claim_id) + assert index is not None + if ( + claim.claim_id not in accepted_claim_ids + and claim.claim_id not in resolved_claim_ids + and index not in issue_indexes + ): + issues.append( + self._claim_issue( + index, + claim, + "unaccounted_fact_claim", + ( + "Validation did not produce an operation, authoritative " + "resolution request, or claim-specific issue." + ), + ) + ) + + if proposal.focus is not None: + operations.append( + SetFocusOperation( + scope_id=proposal.focus.scope_id, + entity_references=proposal.focus.entity_references, + ) + ) + operations.extend( + ConfirmPendingFactResolutionOperation( + proposal_id=response.proposal_id, + accepted=response.action is PendingResolutionAction.ACCEPT, + ) + for response in valid_responses + if response.action + in {PendingResolutionAction.ACCEPT, PendingResolutionAction.REJECT} + ) + return tuple(operations), tuple(claims_to_resolve), tuple(issues) + + def _merge_pending_claim( + self, + pending: PendingFactResolution, + response: PendingFactResolutionResponse, + *, + current_message: str, + change_index: int, + ) -> tuple[FactClaim | None, tuple[ContextValidationIssue, ...]]: + base_path = ("proposal", "changes", str(change_index)) + issues: list[ContextValidationIssue] = [] + if pending.status is not FactResolutionStatus.NEEDS_CLARIFICATION: + issues.append( + ContextValidationIssue( + code="pending_resolution_not_supplementable", + message=( + "Only a pending resolution that requests clarification can " + "receive source-claim field updates." + ), + path=base_path, + evidence=response.evidence, + ) + ) + if pending.source_claim is None: + issues.append( + ContextValidationIssue( + code="pending_resolution_missing_source_claim", + message=( + "The referenced pending resolution predates retained source " + "claims and cannot accept a partial field update." + ), + path=base_path, + evidence=response.evidence, + ) + ) + if not self._text_is_grounded(response.evidence, current_message): + issues.append( + ContextValidationIssue( + code="uncited_pending_resolution_response", + message=( + "The pending-resolution response evidence is absent from the " + "exact current message." + ), + path=(*base_path, "evidence"), + evidence=response.evidence, + ) + ) + if issues or pending.source_claim is None: + return None, tuple(issues) + + claim_data = pending.source_claim.model_dump(mode="python") + seen_paths: set[tuple[str, ...]] = set() + immutable_roots = {"kind", "claim_id", "evidence"} + for update_index, update in enumerate(response.updates): + update_path = (*base_path, "updates", str(update_index)) + if update.path in seen_paths: + issues.append( + ContextValidationIssue( + code="duplicate_fact_claim_field_update", + message="A retained fact-claim field is updated more than once.", + path=(*update_path, "path"), + evidence=update.evidence, + ) + ) + continue + seen_paths.add(update.path) + if update.path[0] in immutable_roots: + issues.append( + ContextValidationIssue( + code="immutable_fact_claim_field_update", + message=( + "A pending-resolution response cannot replace claim " + "identity or original source evidence." + ), + path=(*update_path, "path"), + evidence=update.evidence, + ) + ) + continue + if not self._text_is_grounded(update.evidence, current_message): + issues.append( + ContextValidationIssue( + code="uncited_fact_claim_field_update", + message=( + "The field update's cited evidence is absent from the " + "exact current message." + ), + path=(*update_path, "evidence"), + evidence=update.evidence, + ) + ) + continue + target: object = claim_data + for field in update.path[:-1]: + if not isinstance(target, dict) or field not in target: + target = None + break + target = target[field] + final_field = update.path[-1] + if not isinstance(target, dict) or final_field not in target: + issues.append( + ContextValidationIssue( + code="unknown_fact_claim_field_update", + message=( + "The supplied path does not identify a field in the " + "retained fact-claim schema." + ), + path=(*update_path, "path"), + evidence=update.evidence, + ) + ) + continue + if not self._supplement_value_is_grounded( + update, + current_message=current_message, + ): + issues.append( + ContextValidationIssue( + code="ungrounded_fact_claim_field_update", + message=( + "The supplied field value is not supported by the exact " + "current message." + ), + path=(*update_path, "value"), + evidence=update.evidence, + ) + ) + continue + target[final_field] = update.value + + if issues: + return None, tuple(issues) + try: + return FactClaim.model_validate(claim_data), () + except ValidationError as exc: + return None, tuple( + ContextValidationIssue( + code="invalid_merged_fact_claim", + message=error["msg"], + path=(*base_path, "updates", *(str(item) for item in error["loc"])), + evidence=response.evidence, + ) + for error in exc.errors(include_url=False, include_input=False) + ) + + def _supplement_value_is_grounded( + self, + update: FactClaimFieldUpdate, + *, + current_message: str, + ) -> bool: + value = update.value + if update.path[-1] == "amount": + try: + amount = Decimal(str(value)) + except Exception: + return False + return amount in { + expression.amount + for expression in self._monetary_parser.extract(current_message) + } + if update.path[-1] == "period" and isinstance(value, str): + period_tokens = { + "annual": ("annual", "annually", "year", "yearly", "per year"), + "monthly": ("month", "monthly", "per month"), + "four_weekly": ( + "four weekly", + "four-weekly", + "every four weeks", + "per four weeks", + ), + "weekly": ("week", "weekly", "per week"), + } + normalized_message = current_message.casefold() + return any( + token in normalized_message + for token in period_tokens.get(value.casefold(), ()) + ) + if isinstance(value, bool): + accepted = ("yes", "true", "does", "has", "is") if value else ( + "no", + "false", + "doesn't", + "does not", + "hasn't", + "has not", + "isn't", + "is not", + ) + return any(token in current_message.casefold() for token in accepted) + if isinstance(value, (int, float)): + normalized = " ".join(re.findall(r"[a-z0-9]+", current_message.casefold())) + return re.search(rf"(? bool: + normalized_evidence = " ".join(re.findall(r"[a-z0-9]+", evidence.casefold())) + normalized_message = " ".join( + re.findall(r"[a-z0-9]+", current_message.casefold()) + ) + return bool(normalized_evidence and normalized_evidence in normalized_message) + + @staticmethod + def _claim_change_index( + context: ConversationContext, + proposal: ContextChangeProposal, + claim: FactClaim, + ) -> int: + direct_index = proposal.change_index(claim.claim_id) + if direct_index is not None: + return direct_index + response_ids = { + response.proposal_id: response.response_id + for response in proposal.proposal_responses + if response.action is PendingResolutionAction.SUPPLY + } + response_id = next( + ( + response_ids[pending.proposal_id] + for pending in context.pending_fact_resolutions + if pending.claim_id == claim.claim_id + and pending.proposal_id in response_ids + ), + None, + ) + response_index = proposal.change_index(response_id) if response_id else None + return response_index if response_index is not None else 0 + + def _periodless_direct_binding( + self, + context: ConversationContext, + claim: FactClaim, + scope_id: str, + ) -> tuple[str, str, MoneyPeriod] | None: + subject_id = self._resolve_subject_reference( + context, + claim.subject_references[0], + ) + if subject_id is None: + return None + bindings: set[tuple[str, str, MoneyPeriod]] = set() + if claim.definition_key is not None: + active = context.active_fact( + claim.definition_key, + subject_id, + scope_id, + ) + if active is not None and isinstance(active.assertion, PresentAssertion): + value = active.assertion.value + if isinstance(value, MoneyFactValue): + bindings.add( + (claim.definition_key, claim.definition_version, value.period) + ) + for proposal in context.pending_fact_resolutions: + if ( + proposal.scope_id != scope_id + or proposal.period is None + or proposal.variable_name is None + or subject_id not in proposal.referenced_entity_ids + ): + continue + definition = self._registry.find_by_engine_binding( + proposal.variable_name, + entity=proposal.variable_entity, + ) + if definition is not None and ( + claim.definition_key is None + or definition.key == claim.definition_key + ): + bindings.add( + ( + definition.key, + definition.version, + proposal.period, + ) + ) + if len(bindings) != 1: + return None + return next(iter(bindings)) + + @staticmethod + def _resolve_subject_reference( + context: ConversationContext, + reference: str, + ) -> str | None: + normalized = reference.strip().casefold() + matches = { + entity.entity_id + for entity in context.entities + if normalized + in { + entity.entity_id.casefold(), + (entity.relationship_to_user or "").casefold(), + *(alias.casefold() for alias in entity.aliases), + } + } + if len(matches) != 1: + return None + return next(iter(matches)) + + def _grounding_issues( + self, + current_message: str, + proposal: ContextChangeProposal, + ) -> tuple[ContextValidationIssue, ...]: + """Check that a model proposal is grounded in the exact current message.""" + + claims = proposal.claims + expressions = self._monetary_parser.extract(current_message) + message_amounts = {expression.amount for expression in expressions} + claim_amounts: set[Decimal] = set() + for claim in claims: + if isinstance(claim.value, ClaimedMoneyValue): + claim_amounts.add(claim.value.amount) + elif isinstance(claim.value, TextFactValue): + claim_amounts.update( + expression.amount + for expression in self._monetary_parser.extract(claim.value.value) + ) + elif isinstance(claim.value, TextSetFactValue): + claim_amounts.update( + expression.amount + for value in claim.value.values + for expression in self._monetary_parser.extract(value) + ) + for response in proposal.proposal_responses: + for update in response.updates: + if update.path[-1] != "amount": + continue + try: + claim_amounts.add(Decimal(str(update.value))) + except Exception: + continue + current_evidence = " ".join( + re.findall(r"[a-z0-9]+", current_message.casefold()) + ) + issues = [ + ContextValidationIssue( + code="uncited_fact_claim", + path=( + "proposal", + "changes", + str(proposal.change_index(claim.claim_id) or 0), + "evidence", + ), + claim_index=index, + message=( + "The claim's cited evidence is absent from the exact current " + "message. Remove the copied claim or cite current-message text." + ), + evidence=claim.evidence, + ) + for index, claim in enumerate(claims) + if not self._claim_is_grounded( + claim, + current_evidence=current_evidence, + message_amounts=message_amounts, + ) + ] + issues.extend( + ContextValidationIssue( + code="missing_monetary_fact_claim", + path=("proposal", "changes"), + message=( + "The proposal does not preserve this normalized monetary value " + "from the current message. Classify it in the same claim list " + "without inventing its meaning." + ), + evidence=expression.text, + ) + for expression in expressions + if expression.amount not in claim_amounts + ) + issues.extend( + ContextValidationIssue( + code="uncited_monetary_fact_claim", + path=( + "proposal", + "changes", + str(proposal.change_index(claim.claim_id) or 0), + "value", + "amount", + ), + claim_index=index, + message=( + "This monetary claim value is absent from the exact current " + "message. Remove the copied or calculated claim." + ), + evidence=claim.evidence, + ) + for index, claim in enumerate(claims) + if isinstance(claim.value, ClaimedMoneyValue) + and claim.value.amount not in message_amounts + ) + return tuple(issues) + + @staticmethod + def _claim_is_grounded( + claim: FactClaim, + *, + current_evidence: str, + message_amounts: set[Decimal], + ) -> bool: + claim_evidence = " ".join( + re.findall(r"[a-z0-9]+", claim.evidence.casefold()) + ) + if claim_evidence and claim_evidence in current_evidence: + return True + current_tokens = current_evidence.split() + if ( + current_evidence + and len(current_tokens) <= 4 + and current_evidence in claim_evidence + ): + return True + value = claim.value + if isinstance(value, ClaimedMoneyValue): + return value.amount in message_amounts + if isinstance(value, IntegerFactValue): + return re.search(rf"(? ContextValidationIssue: + return ContextValidationIssue( + code=code, + message=message, + path=("proposal", "changes", str(index)), + claim_index=index, + evidence=claim.evidence, + ) + + def _reduce( + self, + context: ConversationContext, + patch: ContextPatch, + *, + turn_id: str, + evidence: str, + ) -> ContextReduction | tuple[ContextValidationIssue, ...]: + try: + return self._reducer.reduce( + context, + patch, + turn_id=turn_id, + evidence=evidence, + ) + except (TypeError, ValueError) as exc: + return ( + ContextValidationIssue( + code="context_reduction_failed", + message=str(exc), + path=("generated_operations",), + ), + ) + + @classmethod + def _decision_issues( + cls, + decisions: tuple[FactDecision, ...], + ) -> tuple[ContextValidationIssue, ...]: + return tuple( + ContextValidationIssue( + code=f"context_operation_{decision.status.value}", + message=decision.reason, + path=("generated_operations", str(decision.operation_index)), + operation_index=decision.operation_index, + ) + for decision in decisions + if decision.status not in cls._accepted_statuses + ) + + @staticmethod + def _aggregate_issues( + context: ConversationContext, + ) -> tuple[ContextValidationIssue, ...]: + try: + ConversationContext.model_validate(context.model_dump(mode="python")) + except ValidationError as exc: + return tuple( + ContextValidationIssue( + code="invalid_conversation_context", + message=error["msg"], + path=tuple(str(item) for item in error["loc"]), + ) + for error in exc.errors(include_url=False, include_input=False) + ) + return () + + @staticmethod + def _invalid( + prior: ConversationContext, + *, + generated_operations: tuple[ContextOperation, ...] = (), + claims_to_resolve: tuple[FactClaim, ...] = (), + decisions: tuple[FactDecision, ...] = (), + issues: tuple[ContextValidationIssue, ...], + ) -> ContextValidationOutcome: + return ContextValidationOutcome( + status=ContextValidationStatus.NEEDS_CLARIFICATION, + previous_revision=prior.revision, + context=prior, + generated_operations=generated_operations, + claims_to_resolve=claims_to_resolve, + decisions=decisions, + issues=issues, + ) + + +class ContextChangeApplier: + """Persist only a fully validated context change with optimistic revision control.""" + + def __init__(self, repository: ConversationContextRepository) -> None: + self._repository = repository + + def apply(self, outcome: ContextValidationOutcome) -> ConversationContext: + if not outcome.committable: + raise ValueError("Only a fully validated context change can be applied.") + return self._repository.save( + outcome.context, + expected_revision=outcome.previous_revision, + ) diff --git a/backend/conversation_context/engine_projection.py b/backend/conversation_context/engine_projection.py new file mode 100644 index 00000000..9841230a --- /dev/null +++ b/backend/conversation_context/engine_projection.py @@ -0,0 +1,105 @@ +"""Project accepted engine-backed facts into deterministic calculation inputs.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from conversation_context.models import ( + BooleanFactValue, + ConversationContext, + FactValue, + IntegerFactValue, + MoneyFactValue, + PresentAssertion, + TextFactValue, +) +from conversation_context.registry import FactDefinitionRegistry + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class PersonEngineInputs(StrictModel): + entity_id: str + values: dict[str, JsonValue] + + +class HouseholdEngineInputs(StrictModel): + people: tuple[PersonEngineInputs, ...] = () + benunit: dict[str, JsonValue] = Field(default_factory=dict) + household: dict[str, JsonValue] = Field(default_factory=dict) + + +class HouseholdEngineFactProjector: + """Map registered PolicyEngine bindings without enumerating variable names.""" + + def __init__(self, registry: FactDefinitionRegistry) -> None: + self._registry = registry + + def project( + self, + context: ConversationContext, + *, + scope_id: str, + person_entity_ids: tuple[str, ...], + household_entity_id: str, + ) -> HouseholdEngineInputs: + self._registry.restore_engine_definitions(context) + people: dict[str, dict[str, JsonValue]] = { + entity_id: {} for entity_id in person_entity_ids + } + benunit: dict[str, JsonValue] = {} + household: dict[str, JsonValue] = {} + for fact in context.active_facts(): + if fact.scope_id != scope_id or not isinstance( + fact.assertion, + PresentAssertion, + ): + continue + try: + definition = self._registry.get( + fact.definition_key, + fact.definition_version, + ) + except KeyError: + continue + binding = definition.engine_binding + if binding is None or "." not in binding: + continue + engine_entity, variable_name = binding.split(".", 1) + value = self._value(fact.assertion.value) + if value is None: + continue + if engine_entity == "person" and fact.subject_entity_id in people: + people[fact.subject_entity_id][variable_name] = value + elif ( + engine_entity == "benunit" + and fact.subject_entity_id == household_entity_id + ): + benunit[variable_name] = value + elif ( + engine_entity == "household" + and fact.subject_entity_id == household_entity_id + ): + household[variable_name] = value + return HouseholdEngineInputs( + people=tuple( + PersonEngineInputs(entity_id=entity_id, values=people[entity_id]) + for entity_id in person_entity_ids + ), + benunit=benunit, + household=household, + ) + + @staticmethod + def _value(value: FactValue) -> JsonValue | None: + if isinstance(value, BooleanFactValue): + return value.value + if isinstance(value, IntegerFactValue): + return value.value + if isinstance(value, MoneyFactValue): + return float(value.amount) + if isinstance(value, TextFactValue): + return value.value + return None diff --git a/backend/conversation_context/household_view.py b/backend/conversation_context/household_view.py new file mode 100644 index 00000000..ff177d7d --- /dev/null +++ b/backend/conversation_context/household_view.py @@ -0,0 +1,290 @@ +"""Projection from registered conversation facts to household input evidence.""" + +from __future__ import annotations + +from decimal import Decimal + +from capabilities.household_input import ( + AmountFrequency, + HouseholdEvidence, + PeriodicAmount, + PersonEvidence, +) +from conversation_context.models import ( + BooleanFactValue, + ConversationContext, + EntityKind, + EntityReferencesFactValue, + ExplicitAbsenceAssertion, + IntegerFactValue, + MoneyFactValue, + MoneyPeriod, + PendingFactResolution, + PresentAssertion, + TextFactValue, + TextSetFactValue, + FactAssertion, +) + + +class HouseholdContextView: + """Read only active facts from one stable household scope.""" + + def __init__(self, context: ConversationContext) -> None: + self._context = context + + @property + def scope_id(self) -> str: + if self._context.focus.scope_id is not None: + return self._context.focus.scope_id + return next(scope.scope_id for scope in self._context.scopes if scope.active) + + def evidence(self) -> HouseholdEvidence: + scope = next( + scope for scope in self._context.scopes if scope.scope_id == self.scope_id + ) + household = next( + entity + for entity in self._context.entities + if entity.kind is EntityKind.HOUSEHOLD + and entity.entity_id in scope.subject_entity_ids + ) + member_ids = self._member_ids(household.entity_id, scope.subject_entity_ids) + values: dict[str, object] = { + "people": tuple(self._person_evidence(item) for item in member_ids) + } + sources: dict[str, str] = {} + for fact_key, field in ( + ("household.has_children", "has_children"), + ("household.is_married", "is_married"), + ("household.rent", "rent"), + ("household.council_tax", "council_tax"), + ("household.country", "country"), + ): + fact = self._context.active_fact(fact_key, household.entity_id, self.scope_id) + if fact is None: + continue + converted = self._household_value(fact.assertion) + if converted is None and not isinstance( + fact.assertion, + ExplicitAbsenceAssertion, + ): + continue + values[field] = converted + sources[field] = "user" + + childcare = next( + ( + fact + for entity_id in member_ids + if ( + fact := self._context.active_fact( + "person.childcare_expenses", + entity_id, + self.scope_id, + ) + ) + is not None + ), + None, + ) + if childcare is not None: + values["childcare_expenses"] = self._household_value( + childcare.assertion + ) + sources["childcare_expenses"] = "user" + values["sources"] = sources + return HouseholdEvidence.model_validate(values) + + def policy_year(self) -> int | None: + fact = self._context.active_fact( + "analysis.policy_year", + self.household_entity_id, + self.scope_id, + ) + if ( + fact is not None + and isinstance(fact.assertion, PresentAssertion) + and isinstance(fact.assertion.value, IntegerFactValue) + ): + return fact.assertion.value.value + return None + + def requested_outputs(self) -> tuple[str, ...]: + fact = self._context.active_fact( + "analysis.requested_outputs", + self.household_entity_id, + self.scope_id, + ) + if ( + fact is not None + and isinstance(fact.assertion, PresentAssertion) + and isinstance(fact.assertion.value, TextSetFactValue) + ): + return fact.assertion.value.values + return () + + def pending_fact_resolutions(self) -> tuple[PendingFactResolution, ...]: + """Return unresolved variable proposals that affect this household scope.""" + + scope = next( + item for item in self._context.scopes if item.scope_id == self.scope_id + ) + subjects = set(scope.subject_entity_ids) + return tuple( + proposal + for proposal in self._context.pending_fact_resolutions + if proposal.scope_id == self.scope_id + and bool(set(proposal.referenced_entity_ids) & subjects) + ) + + @property + def person_entity_ids(self) -> tuple[str, ...]: + """Return stable person identifiers in household calculation order.""" + + scope = next( + scope for scope in self._context.scopes if scope.scope_id == self.scope_id + ) + return self._member_ids(self.household_entity_id, scope.subject_entity_ids) + + @property + def household_entity_id(self) -> str: + scope = next( + scope for scope in self._context.scopes if scope.scope_id == self.scope_id + ) + return next( + entity.entity_id + for entity in self._context.entities + if entity.kind is EntityKind.HOUSEHOLD + and entity.entity_id in scope.subject_entity_ids + ) + + def _member_ids( + self, + household_id: str, + scope_subjects: tuple[str, ...], + ) -> tuple[str, ...]: + membership = self._context.active_fact( + "household.members", + household_id, + self.scope_id, + ) + if ( + membership is not None + and isinstance(membership.assertion, PresentAssertion) + and isinstance( + membership.assertion.value, + EntityReferencesFactValue, + ) + ): + candidates = membership.assertion.value.entity_ids + else: + candidates = scope_subjects + people = [ + entity + for entity in self._context.entities + if entity.kind is EntityKind.PERSON and entity.entity_id in candidates + ] + return tuple( + entity.entity_id + for entity in sorted( + people, + key=lambda item: ( + item.relationship_to_user != "self", + item.created_turn_id or "", + item.entity_id, + ), + ) + ) + + def _person_evidence(self, entity_id: str) -> PersonEvidence: + entity = next( + item for item in self._context.entities if item.entity_id == entity_id + ) + display_label = self._display_label(entity_id, entity.relationship_to_user) + values: dict[str, object] = { + "entity_id": entity_id, + "display_label": display_label, + "relationship_to_user": entity.relationship_to_user, + } + sources: dict[str, str] = {} + for fact_key, field in ( + ("person.age", "age"), + ("person.employment_income", "employment_income"), + ("person.self_employment_income", "self_employment_income"), + ("person.pension_income", "pension_income"), + ): + fact = self._context.active_fact(fact_key, entity_id, self.scope_id) + if fact is None or not isinstance(fact.assertion, PresentAssertion): + continue + value = fact.assertion.value + if isinstance(value, IntegerFactValue): + values[field] = value.value + elif isinstance(value, MoneyFactValue): + values[field] = self._periodic(value) + else: + continue + sources[field] = "user" + values["sources"] = sources + return PersonEvidence.model_validate(values) + + def _display_label( + self, + entity_id: str, + relationship: str | None, + ) -> str: + name = self._context.active_fact( + "person.name", + entity_id, + self.scope_id, + ) + if ( + name is not None + and isinstance(name.assertion, PresentAssertion) + and isinstance(name.assertion.value, TextFactValue) + ): + return name.assertion.value.value + if relationship == "self": + return "you" + if relationship: + return f"your {relationship.replace('_', ' ')}" + return "the other person" + + @staticmethod + def _household_value( + assertion: FactAssertion, + ) -> bool | str | PeriodicAmount | None: + if isinstance(assertion, ExplicitAbsenceAssertion): + return PeriodicAmount(amount=0, frequency=AmountFrequency.ANNUAL) + if not isinstance(assertion, PresentAssertion): + return None + value = assertion.value + if isinstance(value, BooleanFactValue): + return value.value + if isinstance(value, MoneyFactValue): + return HouseholdContextView._periodic(value) + if isinstance(value, TextFactValue): + return value.value + return None + + @staticmethod + def _periodic(value: MoneyFactValue) -> PeriodicAmount: + if value.period is MoneyPeriod.WEEKLY: + return PeriodicAmount( + amount=float(value.amount), + frequency=AmountFrequency.WEEKLY, + ) + if value.period is MoneyPeriod.MONTHLY: + return PeriodicAmount( + amount=float(value.amount), + frequency=AmountFrequency.MONTHLY, + ) + annual = ( + value.amount * Decimal(13) + if value.period is MoneyPeriod.FOUR_WEEKLY + else value.amount + ) + return PeriodicAmount( + amount=float(annual), + frequency=AmountFrequency.ANNUAL, + ) diff --git a/backend/conversation_context/models.py b/backend/conversation_context/models.py new file mode 100644 index 00000000..d408aa71 --- /dev/null +++ b/backend/conversation_context/models.py @@ -0,0 +1,600 @@ +"""Immutable models for stable conversational entities and registered facts.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from decimal import Decimal +from enum import Enum +from typing import Annotated, Literal, TypeAlias +from uuid import uuid4 + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class EntityKind(str, Enum): + PERSON = "person" + HOUSEHOLD = "household" + POLICY_SCENARIO = "policy_scenario" + + +class ScopeKind(str, Enum): + HOUSEHOLD = "household" + POLICY = "policy" + CONVERSATION = "conversation" + + +class MoneyPeriod(str, Enum): + ANNUAL = "annual" + MONTHLY = "monthly" + FOUR_WEEKLY = "four_weekly" + WEEKLY = "weekly" + + +class FactClaimRelationship(str, Enum): + DIRECT = "direct" + SUM = "sum" + + +class FactResolutionStatus(str, Enum): + AWAITING_CONFIRMATION = "awaiting_confirmation" + NEEDS_CLARIFICATION = "needs_clarification" + + +class PendingResolutionAction(str, Enum): + ACCEPT = "accept" + REJECT = "reject" + SUPPLY = "supply" + + +class PendingQuestionStatus(str, Enum): + AWAITING_ANSWER = "awaiting_answer" + ANSWER_RECEIVED = "answer_received" + + +class ContextEntity(StrictModel): + entity_id: str + kind: EntityKind + aliases: tuple[str, ...] = () + relationship_to_user: str | None = None + created_turn_id: str | None = None + + +class ContextScope(StrictModel): + scope_id: str + kind: ScopeKind + subject_entity_ids: tuple[str, ...] = () + active: bool = True + + +class BooleanFactValue(StrictModel): + kind: Literal["boolean"] = "boolean" + value: bool + + +class IntegerFactValue(StrictModel): + kind: Literal["integer"] = "integer" + value: int + + +class MoneyFactValue(StrictModel): + kind: Literal["money"] = "money" + amount: Decimal + period: MoneyPeriod + currency: Literal["GBP"] = "GBP" + + +class ClaimedMoneyValue(StrictModel): + """A cited monetary value that is not yet an accepted engine-backed fact.""" + + kind: Literal["money"] = "money" + amount: Decimal + period: MoneyPeriod | None = None + currency: Literal["GBP"] = "GBP" + + +class TextFactValue(StrictModel): + kind: Literal["text"] = "text" + value: str + + +class TextSetFactValue(StrictModel): + kind: Literal["text_set"] = "text_set" + values: tuple[str, ...] + + +class EntityReferenceFactValue(StrictModel): + kind: Literal["entity_reference"] = "entity_reference" + entity_id: str + + +class EntityReferencesFactValue(StrictModel): + kind: Literal["entity_references"] = "entity_references" + entity_ids: tuple[str, ...] + + +FactValue: TypeAlias = Annotated[ + BooleanFactValue + | IntegerFactValue + | MoneyFactValue + | TextFactValue + | TextSetFactValue + | EntityReferenceFactValue + | EntityReferencesFactValue, + Field(discriminator="kind"), +] + + +ClaimValue: TypeAlias = Annotated[ + BooleanFactValue + | IntegerFactValue + | ClaimedMoneyValue + | TextFactValue + | TextSetFactValue + | EntityReferenceFactValue + | EntityReferencesFactValue, + Field(discriminator="kind"), +] + + +class PresentAssertion(StrictModel): + kind: Literal["present"] = "present" + value: FactValue + + +class ExplicitAbsenceAssertion(StrictModel): + kind: Literal["explicit_absence"] = "explicit_absence" + + +FactAssertion: TypeAlias = Annotated[ + PresentAssertion | ExplicitAbsenceAssertion, + Field(discriminator="kind"), +] + + +class FactProvenance(StrictModel): + turn_id: str + source: Literal["user", "artifact", "system"] = "user" + evidence: str | None = None + + +class ContextFact(StrictModel): + fact_id: str = Field(default_factory=lambda: uuid4().hex) + definition_key: str + definition_version: str = "1" + subject_entity_id: str + scope_id: str + assertion: FactAssertion + provenance: FactProvenance + introduced_revision: int + supersedes_fact_id: str | None = None + created_at: datetime = Field(default_factory=_now) + + +class FactRequirement(StrictModel): + requirement_id: str + fact_key: str + subject_entity_id: str | None = None + subject_kind: EntityKind | None = None + scope_id: str + expected_value_kind: str + allow_explicit_absence: bool = False + reason: str + + +class CapabilityInvocationReference(StrictModel): + """Typed link from conversational state to one waiting capability call.""" + + invocation_id: str + capability_id: str + capability_version: str + context_scope_id: str + context_revision: int + + +class FactClaim(StrictModel): + """One declarative interpretation of a user-stated contextual fact.""" + + kind: Literal["fact_claim"] = "fact_claim" + claim_id: str = Field(default_factory=lambda: uuid4().hex) + concept: str + subject_references: tuple[str, ...] + relationship: FactClaimRelationship = FactClaimRelationship.DIRECT + value: ClaimValue | None = None + explicit_absence: bool = False + definition_key: str | None = None + definition_version: str = "1" + scope_id: str | None = None + correction: bool = False + evidence: str + + @model_validator(mode="after") + def validate_claim(self) -> "FactClaim": + if not self.concept.strip(): + raise ValueError("fact claim concept must not be empty") + if not self.evidence.strip(): + raise ValueError("fact claim evidence must not be empty") + if not self.subject_references: + raise ValueError("fact claim requires at least one subject reference") + normalized_subjects = { + reference.strip().casefold() + for reference in self.subject_references + if reference.strip() + } + if len(normalized_subjects) != len(self.subject_references): + raise ValueError("fact claim subject references must be non-empty and unique") + if self.explicit_absence == (self.value is not None): + raise ValueError( + "fact claim requires exactly one of value or explicit_absence" + ) + if self.relationship is FactClaimRelationship.DIRECT: + if len(self.subject_references) != 1: + raise ValueError("a direct fact claim requires exactly one subject") + return self + if len(self.subject_references) < 2: + raise ValueError("an additive fact claim requires at least two subjects") + if self.explicit_absence or not isinstance(self.value, ClaimedMoneyValue): + raise ValueError("an additive fact claim requires one monetary value") + return self + + +class ContextEntityCandidate(StrictModel): + """A model-interpreted entity declaration without persistence authority.""" + + reference: str + kind: EntityKind + aliases: tuple[str, ...] = () + relationship_to_user: str | None = None + + +class ContextFocusCandidate(StrictModel): + """A model-interpreted focus change without persistence authority.""" + + scope_id: str | None = None + entity_references: tuple[str, ...] = () + + +class FactClaimFieldUpdate(StrictModel): + """One current-message value supplied for a field on a retained fact claim.""" + + path: tuple[str, ...] + value: JsonValue + evidence: str + + @model_validator(mode="after") + def validate_update(self) -> "FactClaimFieldUpdate": + if not self.path or any(not part.strip() for part in self.path): + raise ValueError("fact-claim update path must contain non-empty fields") + if not self.evidence.strip(): + raise ValueError("fact-claim update evidence must not be empty") + return self + + +class PendingFactResolutionResponse(StrictModel): + """A current-message action on one retained server-authored resolution.""" + + kind: Literal["pending_resolution_response"] = "pending_resolution_response" + response_id: str = Field(default_factory=lambda: uuid4().hex) + proposal_id: str + action: PendingResolutionAction + updates: tuple[FactClaimFieldUpdate, ...] = () + evidence: str + + @model_validator(mode="after") + def validate_response(self) -> "PendingFactResolutionResponse": + if not self.evidence.strip(): + raise ValueError("pending-resolution response evidence must not be empty") + if self.action is PendingResolutionAction.SUPPLY: + if not self.updates: + raise ValueError("a supply response requires at least one field update") + elif self.updates: + raise ValueError("accept and reject responses cannot contain field updates") + return self + + +ProposedContextChange: TypeAlias = Annotated[ + FactClaim | PendingFactResolutionResponse, + Field(discriminator="kind"), +] + + +class FactResolutionSupplement(StrictModel): + """Validated current-message fields added to one retained source claim.""" + + turn_id: str + evidence: str + updates: tuple[FactClaimFieldUpdate, ...] + + +class FactResolutionTerm(StrictModel): + variable_name: str + subject_entity_id: str + coefficient: Decimal = Decimal("1") + known_value: Decimal | None = None + + +class FactResolutionAssignment(StrictModel): + definition_key: str + definition_version: str = "1" + subject_entity_id: str + scope_id: str + assertion: PresentAssertion + correction: bool = False + + +class PendingFactResolution(StrictModel): + """A validated variable mapping or calculation awaiting user input.""" + + proposal_id: str = Field(default_factory=lambda: uuid4().hex) + claim_id: str + source_turn_id: str + source_claim: FactClaim | None = None + supplements: tuple[FactResolutionSupplement, ...] = () + scope_id: str + referenced_entity_ids: tuple[str, ...] + evidence: str + status: FactResolutionStatus + prompt: str + variable_name: str | None = None + variable_entity: str | None = None + variable_label: str | None = None + definition_period: str | None = None + mapping_confidence: str | None = None + relationship: FactClaimRelationship + expected_total: Decimal | None = None + period: MoneyPeriod | None = None + terms: tuple[FactResolutionTerm, ...] = () + assignments: tuple[FactResolutionAssignment, ...] = () + created_revision: int + + +class PendingQuestion(StrictModel): + question_id: str = Field(default_factory=lambda: uuid4().hex) + capability_id: str + capability_invocation: CapabilityInvocationReference | None = None + prompt: str + requirements: tuple[FactRequirement, ...] + created_turn_id: str + status: PendingQuestionStatus = PendingQuestionStatus.AWAITING_ANSWER + + @model_validator(mode="after") + def validate_capability_invocation(self) -> "PendingQuestion": + reference = self.capability_invocation + if reference is not None and reference.capability_id != self.capability_id: + raise ValueError( + "pending question capability does not match its invocation reference" + ) + return self + + +class ConversationFocus(StrictModel): + scope_id: str | None = None + entity_ids: tuple[str, ...] = () + artifact_ids: tuple[str, ...] = () + + +class EnsureEntityOperation(StrictModel): + operation: Literal["ensure_entity"] = "ensure_entity" + reference: str + kind: EntityKind + aliases: tuple[str, ...] = () + relationship_to_user: str | None = None + + +class SetFactOperation(StrictModel): + operation: Literal["set_fact"] = "set_fact" + definition_key: str + definition_version: str = "1" + subject_reference: str + scope_id: str + assertion: FactAssertion + correction: bool = False + + +class SetFocusOperation(StrictModel): + operation: Literal["set_focus"] = "set_focus" + scope_id: str | None = None + entity_references: tuple[str, ...] = () + + +class ReplacePendingQuestionsOperation(StrictModel): + operation: Literal["replace_pending_questions"] = "replace_pending_questions" + questions: tuple[PendingQuestion, ...] = () + + +class AddPendingFactResolutionOperation(StrictModel): + operation: Literal["add_pending_fact_resolution"] = "add_pending_fact_resolution" + proposal: PendingFactResolution + + +class ConfirmPendingFactResolutionOperation(StrictModel): + operation: Literal["confirm_pending_fact_resolution"] = ( + "confirm_pending_fact_resolution" + ) + proposal_id: str + accepted: bool + + +ContextOperation: TypeAlias = Annotated[ + EnsureEntityOperation + | SetFactOperation + | SetFocusOperation + | ReplacePendingQuestionsOperation + | AddPendingFactResolutionOperation + | ConfirmPendingFactResolutionOperation, + Field(discriminator="operation"), +] + + +class ContextPatch(StrictModel): + expected_revision: int + operations: tuple[ContextOperation, ...] = () + + +class FactDecisionStatus(str, Enum): + ACCEPTED = "accepted" + REJECTED = "rejected" + CONFLICTED = "conflicted" + IGNORED = "ignored" + SUPERSEDED = "superseded" + + +class FactDecision(StrictModel): + operation_index: int + status: FactDecisionStatus + operation: str + definition_key: str | None = None + subject_entity_id: str | None = None + fact_id: str | None = None + superseded_fact_id: str | None = None + reason: str + + +class ConversationContext(StrictModel): + schema_version: Literal["1"] = "1" + conversation_id: str + revision: int = 0 + entities: tuple[ContextEntity, ...] + scopes: tuple[ContextScope, ...] + facts: tuple[ContextFact, ...] = () + pending_questions: tuple[PendingQuestion, ...] = () + pending_fact_resolutions: tuple[PendingFactResolution, ...] = () + focus: ConversationFocus = Field(default_factory=ConversationFocus) + updated_at: datetime = Field(default_factory=_now) + + @classmethod + def initial(cls, conversation_id: str) -> "ConversationContext": + person_id = "person:self" + household_id = "household:primary" + scope_id = "scope:primary-household" + return cls( + conversation_id=conversation_id, + entities=( + ContextEntity( + entity_id=person_id, + kind=EntityKind.PERSON, + aliases=("I", "me", "myself", "the user"), + relationship_to_user="self", + ), + ContextEntity( + entity_id=household_id, + kind=EntityKind.HOUSEHOLD, + aliases=("my household", "the household"), + relationship_to_user="primary_household", + ), + ), + scopes=( + ContextScope( + scope_id=scope_id, + kind=ScopeKind.HOUSEHOLD, + subject_entity_ids=(household_id, person_id), + ), + ), + focus=ConversationFocus( + scope_id=scope_id, + entity_ids=(household_id, person_id), + ), + ) + + @model_validator(mode="after") + def validate_references(self) -> "ConversationContext": + entity_ids = {entity.entity_id for entity in self.entities} + if len(entity_ids) != len(self.entities): + raise ValueError("context entity identifiers must be unique") + scope_ids = {scope.scope_id for scope in self.scopes} + if len(scope_ids) != len(self.scopes): + raise ValueError("context scope identifiers must be unique") + for scope in self.scopes: + if not set(scope.subject_entity_ids) <= entity_ids: + raise ValueError("context scope references an unknown entity") + for fact in self.facts: + if fact.subject_entity_id not in entity_ids: + raise ValueError("context fact references an unknown entity") + if fact.scope_id not in scope_ids: + raise ValueError("context fact references an unknown scope") + proposal_ids = { + proposal.proposal_id for proposal in self.pending_fact_resolutions + } + if len(proposal_ids) != len(self.pending_fact_resolutions): + raise ValueError("pending fact-resolution identifiers must be unique") + for proposal in self.pending_fact_resolutions: + if proposal.scope_id not in scope_ids: + raise ValueError("pending fact resolution references an unknown scope") + if not set(proposal.referenced_entity_ids) <= entity_ids: + raise ValueError("pending fact resolution references an unknown entity") + for assignment in proposal.assignments: + if assignment.subject_entity_id not in entity_ids: + raise ValueError( + "fact-resolution assignment references an unknown entity" + ) + if assignment.scope_id not in scope_ids: + raise ValueError( + "fact-resolution assignment references an unknown scope" + ) + question_ids = {question.question_id for question in self.pending_questions} + if len(question_ids) != len(self.pending_questions): + raise ValueError("pending question identifiers must be unique") + invocation_ids: set[str] = set() + for question in self.pending_questions: + for requirement in question.requirements: + if requirement.scope_id not in scope_ids: + raise ValueError("pending question requirement references an unknown scope") + if ( + requirement.subject_entity_id is not None + and requirement.subject_entity_id not in entity_ids + ): + raise ValueError( + "pending question requirement references an unknown entity" + ) + reference = question.capability_invocation + if reference is None: + # Version-one contexts written by the initial implementation did + # not persist this link. ChatTurnService repairs an unambiguous + # legacy record before supplying context to a new turn. + continue + if reference.context_scope_id not in scope_ids: + raise ValueError("pending capability invocation references an unknown scope") + if reference.invocation_id in invocation_ids: + raise ValueError( + "a capability invocation cannot own more than one pending question" + ) + invocation_ids.add(reference.invocation_id) + return self + + def active_facts(self) -> tuple[ContextFact, ...]: + superseded = { + fact.supersedes_fact_id + for fact in self.facts + if fact.supersedes_fact_id is not None + } + return tuple(fact for fact in self.facts if fact.fact_id not in superseded) + + def active_fact( + self, + definition_key: str, + subject_entity_id: str, + scope_id: str, + ) -> ContextFact | None: + return next( + ( + fact + for fact in reversed(self.active_facts()) + if fact.definition_key == definition_key + and fact.subject_entity_id == subject_entity_id + and fact.scope_id == scope_id + ), + None, + ) + + +class ContextReduction(StrictModel): + previous_revision: int + context: ConversationContext + decisions: tuple[FactDecision, ...] diff --git a/backend/conversation_context/projection.py b/backend/conversation_context/projection.py new file mode 100644 index 00000000..154040b2 --- /dev/null +++ b/backend/conversation_context/projection.py @@ -0,0 +1,55 @@ +"""Small typed projections for model prompts and capability input views.""" + +from __future__ import annotations + +from pydantic import BaseModel, ConfigDict + +from conversation_context.models import ( + ContextEntity, + ContextFact, + ConversationContext, + FactRequirement, + PendingFactResolution, + PendingQuestionStatus, +) + + +class PendingQuestionProjection(BaseModel): + """Model-safe pending question without internal persistence identifiers.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + capability_id: str + prompt: str + requirements: tuple[FactRequirement, ...] + status: PendingQuestionStatus + + +class ContextProjection(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + revision: int + entities: tuple[ContextEntity, ...] + active_facts: tuple[ContextFact, ...] + pending_questions: tuple[PendingQuestionProjection, ...] + pending_fact_resolutions: tuple[PendingFactResolution, ...] + active_scope_id: str | None + + +def project_context(context: ConversationContext) -> ContextProjection: + return ContextProjection( + revision=context.revision, + entities=context.entities, + active_facts=context.active_facts(), + pending_questions=tuple( + PendingQuestionProjection( + capability_id=question.capability_id, + prompt=question.prompt, + requirements=question.requirements, + status=question.status, + ) + for question in context.pending_questions + ), + pending_fact_resolutions=context.pending_fact_resolutions, + active_scope_id=context.focus.scope_id, + ) diff --git a/backend/conversation_context/quantities.py b/backend/conversation_context/quantities.py new file mode 100644 index 00000000..ed99fc7c --- /dev/null +++ b/backend/conversation_context/quantities.py @@ -0,0 +1,177 @@ +"""Deterministic normalization for monetary expressions in user messages.""" + +from __future__ import annotations + +from dataclasses import dataclass +from decimal import Decimal +import re + + +@dataclass(frozen=True) +class MonetaryExpression: + """One normalized monetary-looking expression and its source span.""" + + text: str + amount: Decimal + start: int + end: int + + +class MonetaryExpressionParser: + """Normalize common written forms without assigning a policy concept.""" + + _numeric_pattern = re.compile( + r"(?£\s*|GBP\s+)?" + r"(?P" + r"\d{1,3}(?:[,.\u00a0 ]\d{3})+(?:[,.]\d{1,2})?" + r"|\d+(?:[,.]\d+)?" + r")" + r"(?:\s*(?Pk|thousand|grand|m|million))?" + r"(?:\s*(?PGBP|pounds?))?" + r"(?!\w)", + re.IGNORECASE, + ) + _number_word_pattern = re.compile( + r"\b(?P" + r"(?:(?:zero|one|two|three|four|five|six|seven|eight|nine|ten|" + r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|" + r"eighteen|nineteen|twenty|thirty|forty|fifty|sixty|seventy|" + r"eighty|ninety|hundred|and)[\s-]+)*" + r"(?:thousand|grand|million)" + r")\b", + re.IGNORECASE, + ) + _small_numbers = { + "zero": 0, + "one": 1, + "two": 2, + "three": 3, + "four": 4, + "five": 5, + "six": 6, + "seven": 7, + "eight": 8, + "nine": 9, + "ten": 10, + "eleven": 11, + "twelve": 12, + "thirteen": 13, + "fourteen": 14, + "fifteen": 15, + "sixteen": 16, + "seventeen": 17, + "eighteen": 18, + "nineteen": 19, + "twenty": 20, + "thirty": 30, + "forty": 40, + "fifty": 50, + "sixty": 60, + "seventy": 70, + "eighty": 80, + "ninety": 90, + } + _scale_multipliers = { + "k": Decimal("1000"), + "thousand": Decimal("1000"), + "grand": Decimal("1000"), + "m": Decimal("1000000"), + "million": Decimal("1000000"), + } + + def extract(self, text: str) -> tuple[MonetaryExpression, ...]: + expressions: list[MonetaryExpression] = [] + occupied: list[tuple[int, int]] = [] + for match in self._numeric_pattern.finditer(text): + number = match.group("number") + scale = (match.group("scale") or "").casefold() + has_currency = bool(match.group("prefix") or match.group("suffix")) + if not self._is_monetary_form( + number=number, + scale=scale, + has_currency=has_currency, + ): + continue + amount = self._parse_numeric(number) + if scale: + amount *= self._scale_multipliers[scale] + expressions.append( + MonetaryExpression( + text=match.group(0), + amount=amount, + start=match.start(), + end=match.end(), + ) + ) + occupied.append((match.start(), match.end())) + + for match in self._number_word_pattern.finditer(text): + if any(start < match.end() and match.start() < end for start, end in occupied): + continue + expressions.append( + MonetaryExpression( + text=match.group(0), + amount=self._parse_number_words(match.group("words")), + start=match.start(), + end=match.end(), + ) + ) + + return tuple(sorted(expressions, key=lambda item: item.start)) + + @staticmethod + def _is_monetary_form( + *, + number: str, + scale: str, + has_currency: bool, + ) -> bool: + compact = number.replace(" ", "").replace("\u00a0", "") + digits = compact.replace(",", "").replace(".", "") + return ( + has_currency + or bool(scale) + or "," in compact + or "." in compact + or len(digits) >= 5 + ) + + @staticmethod + def _parse_numeric(value: str) -> Decimal: + compact = value.replace(" ", "").replace("\u00a0", "") + if "," in compact and "." in compact: + decimal_separator = "," if compact.rfind(",") > compact.rfind(".") else "." + grouping_separator = "." if decimal_separator == "," else "," + return Decimal( + compact.replace(grouping_separator, "").replace( + decimal_separator, + ".", + ) + ) + separator = "," if "," in compact else "." if "." in compact else None + if separator is None: + return Decimal(compact) + groups = compact.split(separator) + if len(groups) > 1 and all(len(group) == 3 for group in groups[1:]): + return Decimal("".join(groups)) + return Decimal(compact.replace(separator, ".")) + + def _parse_number_words(self, value: str) -> Decimal: + tokens = value.casefold().replace("-", " ").split() + total = 0 + current = 0 + for token in tokens: + if token == "and": + continue + if token in self._small_numbers: + current += self._small_numbers[token] + elif token == "hundred": + current = max(current, 1) * 100 + elif token in {"thousand", "grand"}: + total += max(current, 1) * 1_000 + current = 0 + elif token == "million": + total += max(current, 1) * 1_000_000 + current = 0 + return Decimal(total + current) diff --git a/backend/conversation_context/reducer.py b/backend/conversation_context/reducer.py new file mode 100644 index 00000000..772e2b00 --- /dev/null +++ b/backend/conversation_context/reducer.py @@ -0,0 +1,911 @@ +"""Deterministic validation and application of proposed context patches.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from decimal import Decimal +from uuid import NAMESPACE_URL, uuid5 + +from conversation_context.models import ( + AddPendingFactResolutionOperation, + ConfirmPendingFactResolutionOperation, + ContextEntity, + ContextFact, + ContextPatch, + ContextReduction, + ConversationContext, + EntityReferenceFactValue, + EntityReferencesFactValue, + EnsureEntityOperation, + ExplicitAbsenceAssertion, + FactDecision, + FactDecisionStatus, + FactProvenance, + FactAssertion, + FactValue, + FactResolutionStatus, + MoneyFactValue, + MoneyPeriod, + PendingFactResolution, + PendingQuestion, + PendingQuestionStatus, + PresentAssertion, + ReplacePendingQuestionsOperation, + SetFactOperation, + SetFocusOperation, +) +from conversation_context.quantities import MonetaryExpressionParser +from conversation_context.registry import ( + FactDefinitionRegistry, + FactUpdatePolicy, + FactValueKind, +) + + +class ContextReducer: + def __init__(self, registry: FactDefinitionRegistry) -> None: + self._registry = registry + self._money_parser = MonetaryExpressionParser() + + def reduce( + self, + context: ConversationContext, + patch: ContextPatch, + *, + turn_id: str, + evidence: str, + ) -> ContextReduction: + if patch.expected_revision != context.revision: + raise ValueError( + "Context patch revision does not match the current context revision." + ) + + next_revision = context.revision + 1 + entities = list(context.entities) + scopes = list(context.scopes) + facts = list(context.facts) + pending = context.pending_questions + original_pending = context.pending_questions + pending_resolutions = list(context.pending_fact_resolutions) + original_pending_resolutions = context.pending_fact_resolutions + focus = context.focus + references = {entity.entity_id: entity.entity_id for entity in entities} + for entity in entities: + for alias in entity.aliases: + references.setdefault(alias.casefold(), entity.entity_id) + decisions: list[FactDecision] = [] + + for index, operation in enumerate(patch.operations): + if isinstance(operation, EnsureEntityOperation): + existing_id = self._resolve_entity(operation.reference, references) + if existing_id is not None: + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.IGNORED, + operation=operation.operation, + subject_entity_id=existing_id, + reason="The entity reference already resolves to a stable entity.", + ) + ) + continue + entity_id = "entity:" + uuid5( + NAMESPACE_URL, + f"{context.conversation_id}:{operation.reference.casefold()}", + ).hex + entity = ContextEntity( + entity_id=entity_id, + kind=operation.kind, + aliases=tuple(dict.fromkeys((operation.reference, *operation.aliases))), + relationship_to_user=operation.relationship_to_user, + created_turn_id=turn_id, + ) + entities.append(entity) + if context.focus.scope_id is not None: + scopes = [ + scope.model_copy( + update={ + "subject_entity_ids": tuple( + dict.fromkeys( + (*scope.subject_entity_ids, entity_id) + ) + ) + } + ) + if scope.scope_id == context.focus.scope_id + else scope + for scope in scopes + ] + references[operation.reference.casefold()] = entity_id + references[entity_id] = entity_id + for alias in operation.aliases: + references.setdefault(alias.casefold(), entity_id) + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.ACCEPTED, + operation=operation.operation, + subject_entity_id=entity_id, + reason="Created a stable context entity.", + ) + ) + continue + + if isinstance(operation, SetFocusOperation): + resolved = tuple( + resolved_id + for reference in operation.entity_references + if ( + resolved_id := self._resolve_entity( + reference, + references, + ) + ) + is not None + ) + if operation.scope_id is not None and not any( + scope.scope_id == operation.scope_id for scope in scopes + ): + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.REJECTED, + operation=operation.operation, + reason="Focus references an unknown context scope.", + ) + ) + continue + focus = focus.model_copy( + update={ + "scope_id": operation.scope_id, + "entity_ids": resolved, + } + ) + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.ACCEPTED, + operation=operation.operation, + reason="Updated conversational focus.", + ) + ) + continue + + if isinstance(operation, ReplacePendingQuestionsOperation): + pending = operation.questions + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.ACCEPTED, + operation=operation.operation, + reason="Replaced the typed pending-question set.", + ) + ) + continue + + if isinstance(operation, AddPendingFactResolutionOperation): + proposal_error = self._validate_resolution_proposal( + operation.proposal, + context=context.model_copy( + update={ + "entities": tuple(entities), + "scopes": tuple(scopes), + "facts": tuple(facts), + } + ), + ) + if proposal_error is not None: + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.REJECTED, + operation=operation.operation, + reason=proposal_error, + ) + ) + continue + if any( + item.proposal_id == operation.proposal.proposal_id + for item in pending_resolutions + ): + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.IGNORED, + operation=operation.operation, + reason="The same fact-resolution proposal is already pending.", + ) + ) + continue + pending_resolutions = [ + item + for item in pending_resolutions + if item.claim_id != operation.proposal.claim_id + ] + pending_resolutions.append(operation.proposal) + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.ACCEPTED, + operation=operation.operation, + reason="Stored a validated fact-resolution proposal.", + ) + ) + continue + + if isinstance(operation, ConfirmPendingFactResolutionOperation): + proposal = next( + ( + item + for item in pending_resolutions + if item.proposal_id == operation.proposal_id + ), + None, + ) + if proposal is None: + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.REJECTED, + operation=operation.operation, + reason="The fact-resolution proposal is not pending.", + ) + ) + continue + if not operation.accepted: + pending_resolutions = [ + item.model_copy( + update={ + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + "What exact amounts, periods, and household-member " + "assignments should I use instead?" + ), + } + ) + if item.proposal_id == operation.proposal_id + else item + for item in pending_resolutions + ] + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.ACCEPTED, + operation=operation.operation, + reason=( + "Rejected the calculated assignment and retained a " + "request for an explicit breakdown." + ), + ) + ) + continue + if ( + proposal.status is not FactResolutionStatus.AWAITING_CONFIRMATION + or len(proposal.assignments) != 1 + ): + decisions.append( + FactDecision( + operation_index=index, + status=FactDecisionStatus.REJECTED, + operation=operation.operation, + reason=( + "Only one validated calculated assignment awaiting " + "confirmation can be applied." + ), + ) + ) + continue + assignment = proposal.assignments[0] + self._ensure_resolution_definition(proposal) + set_operation = SetFactOperation( + definition_key=assignment.definition_key, + definition_version=assignment.definition_version, + subject_reference=assignment.subject_entity_id, + scope_id=assignment.scope_id, + assertion=assignment.assertion, + correction=assignment.correction, + ) + decision, fact = self._reduce_fact( + set_operation, + index=index, + context=context.model_copy( + update={ + "entities": tuple(entities), + "scopes": tuple(scopes), + "facts": tuple(facts), + } + ), + references=references, + next_revision=next_revision, + turn_id=turn_id, + evidence=evidence, + ) + decisions.append( + decision.model_copy(update={"operation": operation.operation}) + ) + if fact is not None: + facts.append(fact) + if decision.status in { + FactDecisionStatus.ACCEPTED, + FactDecisionStatus.IGNORED, + FactDecisionStatus.SUPERSEDED, + }: + pending_resolutions = [ + item + for item in pending_resolutions + if item.proposal_id != operation.proposal_id + ] + continue + + if isinstance(operation, SetFactOperation): + decision, fact = self._reduce_fact( + operation, + index=index, + context=context.model_copy( + update={ + "entities": tuple(entities), + "scopes": tuple(scopes), + "facts": tuple(facts), + } + ), + references=references, + next_revision=next_revision, + turn_id=turn_id, + evidence=evidence, + ) + decisions.append(decision) + if fact is not None: + facts.append(fact) + + resolution_context = context.model_copy( + update={ + "entities": tuple(entities), + "scopes": tuple(scopes), + "facts": tuple(facts), + } + ) + pending_resolutions = [ + proposal + for proposal in pending_resolutions + if not self._is_replaced_by_explicit_facts( + proposal, + resolution_context, + ) + ] + + pending = tuple( + question.model_copy( + update={ + "status": ( + PendingQuestionStatus.ANSWER_RECEIVED + if self._requirements_satisfied(resolution_context, question) + else PendingQuestionStatus.AWAITING_ANSWER + ) + } + ) + for question in pending + ) + + changed = any( + decision.status + in {FactDecisionStatus.ACCEPTED, FactDecisionStatus.SUPERSEDED} + for decision in decisions + ) or tuple(pending_resolutions) != original_pending_resolutions or pending != original_pending + updated = context.model_copy( + update={ + "revision": next_revision if changed else context.revision, + "entities": tuple(entities), + "scopes": tuple(scopes), + "facts": tuple(facts), + "pending_questions": pending, + "pending_fact_resolutions": tuple(pending_resolutions), + "focus": focus, + "updated_at": datetime.now(timezone.utc), + } + ) + return ContextReduction( + previous_revision=context.revision, + context=updated, + decisions=tuple(decisions), + ) + + @staticmethod + def _requirements_satisfied( + context: ConversationContext, + question: PendingQuestion, + ) -> bool: + if not question.requirements: + return False + for requirement in question.requirements: + subject_ids: tuple[str, ...] + if requirement.subject_entity_id is not None: + subject_ids = (requirement.subject_entity_id,) + elif requirement.subject_kind is not None: + subject_ids = tuple( + entity.entity_id + for entity in context.entities + if entity.kind is requirement.subject_kind + ) + else: + return False + matching = tuple( + fact + for subject_id in subject_ids + if ( + fact := context.active_fact( + requirement.fact_key, + subject_id, + requirement.scope_id, + ) + ) + is not None + ) + if not matching: + return False + if ( + not requirement.allow_explicit_absence + and all( + isinstance(fact.assertion, ExplicitAbsenceAssertion) + for fact in matching + ) + ): + return False + return True + + def _is_replaced_by_explicit_facts( + self, + proposal: PendingFactResolution, + context: ConversationContext, + ) -> bool: + """Drop a disputed proposal only when accepted facts satisfy its constraint.""" + + expected_total = proposal.expected_total + period = proposal.period + definition_key: str | None = None + if proposal.assignments: + definition_keys = { + assignment.definition_key for assignment in proposal.assignments + } + if len(definition_keys) == 1: + definition_key = next(iter(definition_keys)) + if definition_key is None and proposal.variable_name is not None: + definition = self._registry.find_by_engine_binding( + proposal.variable_name, + entity=proposal.variable_entity, + ) + if definition is not None: + definition_key = definition.key + if definition_key is not None and proposal.terms: + if period is None or expected_total is None: + return False + total = Decimal("0") + for term in proposal.terms: + fact = context.active_fact( + definition_key, + term.subject_entity_id, + proposal.scope_id, + ) + value = self._money_fact_value(fact, period) + if value is None: + value = term.known_value + if value is None: + return False + total += term.coefficient * value + return total == expected_total + + # Version-one reconciliation could write an incomplete proposal directly. + # Repair only when later explicit facts give every referenced entity the same + # registered monetary fact and exactly satisfy the stated total. The original + # implementation omitted the typed total and period on this failure path, so + # recover those only when the evidence and accepted facts are unambiguous. + if proposal.variable_name is not None or proposal.terms: + return False + candidate_keys: set[str] | None = None + active_facts = context.active_facts() + for entity_id in proposal.referenced_entity_ids: + keys = { + fact.definition_key + for fact in active_facts + if fact.subject_entity_id == entity_id + and fact.scope_id == proposal.scope_id + and fact.introduced_revision > proposal.created_revision + and isinstance(fact.assertion, PresentAssertion) + and isinstance(fact.assertion.value, MoneyFactValue) + } + candidate_keys = keys if candidate_keys is None else candidate_keys & keys + if candidate_keys is None or len(candidate_keys) != 1: + return False + common_key = next(iter(candidate_keys)) + facts: list[ContextFact] = [] + for entity_id in proposal.referenced_entity_ids: + fact = context.active_fact(common_key, entity_id, proposal.scope_id) + if fact is None: + return False + facts.append(fact) + + if period is None: + fact_periods = { + fact.assertion.value.period + for fact in facts + if isinstance(fact.assertion, PresentAssertion) + and isinstance(fact.assertion.value, MoneyFactValue) + } + if len(fact_periods) != 1: + return False + period = next(iter(fact_periods)) + if expected_total is None: + expressions = self._money_parser.extract(proposal.evidence) + if len(expressions) != 1: + return False + expected_total = expressions[0].amount + + total = Decimal("0") + for fact in facts: + value = self._money_fact_value(fact, period) + if value is None: + return False + total += value + return total == expected_total + + @staticmethod + def _money_fact_value( + fact: ContextFact | None, + period: MoneyPeriod, + ) -> Decimal | None: + if fact is None or not isinstance(fact.assertion, PresentAssertion): + return None + value = fact.assertion.value + if not isinstance(value, MoneyFactValue): + return None + periods_per_year = { + MoneyPeriod.ANNUAL: Decimal("1"), + MoneyPeriod.MONTHLY: Decimal("12"), + MoneyPeriod.FOUR_WEEKLY: Decimal("13"), + MoneyPeriod.WEEKLY: Decimal("52"), + } + annual = value.amount * periods_per_year[value.period] + return annual / periods_per_year[period] + + def _validate_resolution_proposal( + self, + proposal: PendingFactResolution, + *, + context: ConversationContext, + ) -> str | None: + entity_ids = {entity.entity_id for entity in context.entities} + scope_ids = {scope.scope_id for scope in context.scopes} + if proposal.scope_id not in scope_ids: + return "The fact-resolution proposal references an unknown scope." + if not set(proposal.referenced_entity_ids) <= entity_ids: + return "The fact-resolution proposal references an unknown entity." + if proposal.status is FactResolutionStatus.AWAITING_CONFIRMATION: + if len(proposal.assignments) != 1: + return "A confirmable resolution must contain exactly one assignment." + if ( + proposal.variable_name is None + or proposal.variable_entity is None + or proposal.variable_label is None + or proposal.expected_total is None + or proposal.period is None + or not proposal.terms + ): + return ( + "A confirmable resolution must contain its validated variable, " + "period, equation terms, and expected total." + ) + if any( + term.variable_name != proposal.variable_name + or term.subject_entity_id not in entity_ids + for term in proposal.terms + ): + return "A confirmable resolution contains an invalid equation term." + assignment = proposal.assignments[0] + unresolved_subjects = { + term.subject_entity_id + for term in proposal.terms + if term.known_value is None + } + if ( + proposal.relationship.value == "sum" + and unresolved_subjects != {assignment.subject_entity_id} + ) or ( + proposal.relationship.value == "direct" + and {term.subject_entity_id for term in proposal.terms} + != {assignment.subject_entity_id} + ): + return "The calculated assignment does not target the unresolved term." + assignment_value = assignment.assertion.value + if ( + not isinstance(assignment_value, MoneyFactValue) + or assignment_value.period is not proposal.period + ): + return "The calculated assignment does not use the validated period." + equation_total = sum( + ( + term.coefficient + * ( + assignment_value.amount + if term.subject_entity_id == assignment.subject_entity_id + else term.known_value or Decimal("0") + ) + for term in proposal.terms + ), + start=Decimal("0"), + ) + if equation_total != proposal.expected_total: + return "The calculated assignment does not satisfy the validated total." + self._ensure_resolution_definition(proposal) + for assignment in proposal.assignments: + if assignment.subject_entity_id not in entity_ids: + return "A calculated assignment references an unknown entity." + if assignment.scope_id not in scope_ids: + return "A calculated assignment references an unknown scope." + try: + definition = self._registry.get( + assignment.definition_key, + assignment.definition_version, + ) + except KeyError: + return "A calculated assignment references an unknown fact definition." + entity = next( + item + for item in context.entities + if item.entity_id == assignment.subject_entity_id + ) + if entity.kind not in definition.subject_kinds: + return "A calculated assignment targets an incompatible entity type." + error = definition.validate_value(assignment.assertion.value) + if error is not None: + return error + return None + + def _ensure_resolution_definition( + self, + proposal: PendingFactResolution, + ) -> None: + if not proposal.assignments: + return + assignment = proposal.assignments[0] + try: + self._registry.get( + assignment.definition_key, + assignment.definition_version, + ) + return + except KeyError: + pass + if ( + proposal.variable_name is None + or proposal.variable_entity is None + or proposal.variable_label is None + ): + return + self._registry.ensure_engine_definition( + variable_name=proposal.variable_name, + entity=proposal.variable_entity, + label=proposal.variable_label, + value_kind=FactValueKind.MONEY, + ) + + def _reduce_fact( + self, + operation: SetFactOperation, + *, + index: int, + context: ConversationContext, + references: dict[str, str], + next_revision: int, + turn_id: str, + evidence: str, + ) -> tuple[FactDecision, ContextFact | None]: + subject_id = self._resolve_entity(operation.subject_reference, references) + if subject_id is None: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + "The subject reference does not resolve to a stable entity.", + ), None + try: + definition = self._registry.get( + operation.definition_key, + operation.definition_version, + ) + except KeyError: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + "The fact definition is not registered.", + subject_id, + ), None + entity = next(item for item in context.entities if item.entity_id == subject_id) + if entity.kind not in definition.subject_kinds: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + "The fact definition does not permit this subject entity type.", + subject_id, + ), None + if not any(scope.scope_id == operation.scope_id for scope in context.scopes): + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + "The fact references an unknown context scope.", + subject_id, + ), None + if isinstance(operation.assertion, ExplicitAbsenceAssertion): + if not definition.allow_explicit_absence: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + "The fact definition does not permit explicit absence.", + subject_id, + ), None + elif isinstance(operation.assertion, PresentAssertion): + value_error = definition.validate_value(operation.assertion.value) + if value_error is not None: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + value_error, + subject_id, + ), None + reference_error = self._validate_value_references( + operation.assertion.value, + references, + ) + if reference_error is not None: + return self._decision( + index, + operation, + FactDecisionStatus.REJECTED, + reference_error, + subject_id, + ), None + + normalized_assertion = self._normalize_value_references( + operation.assertion, + references, + ) + + current = context.active_fact( + operation.definition_key, + subject_id, + operation.scope_id, + ) + if current is not None and current.assertion == normalized_assertion: + return self._decision( + index, + operation, + FactDecisionStatus.IGNORED, + "The same fact assertion is already active.", + subject_id, + fact_id=current.fact_id, + ), None + if ( + current is not None + and not operation.correction + and definition.update_policy + is FactUpdatePolicy.REQUIRE_EXPLICIT_CORRECTION + ): + return self._decision( + index, + operation, + FactDecisionStatus.CONFLICTED, + "A different fact is active and the proposal is not an explicit correction.", + subject_id, + fact_id=current.fact_id, + ), None + + fact = ContextFact( + definition_key=operation.definition_key, + definition_version=operation.definition_version, + subject_entity_id=subject_id, + scope_id=operation.scope_id, + assertion=normalized_assertion, + provenance=FactProvenance( + turn_id=turn_id, + evidence=evidence, + ), + introduced_revision=next_revision, + supersedes_fact_id=current.fact_id if current is not None else None, + ) + status = ( + FactDecisionStatus.SUPERSEDED + if current is not None + else FactDecisionStatus.ACCEPTED + ) + return self._decision( + index, + operation, + status, + "Accepted a validated fact assertion." + if current is None + else ( + "Accepted an explicit correction and superseded the active fact." + if operation.correction + else "Accepted a new explicit assertion and superseded the active fact." + ), + subject_id, + fact_id=fact.fact_id, + superseded_fact_id=current.fact_id if current is not None else None, + ), fact + + @staticmethod + def _resolve_entity(reference: str, references: dict[str, str]) -> str | None: + return references.get(reference) or references.get(reference.casefold()) + + @staticmethod + def _validate_value_references( + value: FactValue, + references: dict[str, str], + ) -> str | None: + if isinstance(value, EntityReferenceFactValue): + if ContextReducer._resolve_entity(value.entity_id, references) is None: + return "The fact value references an unknown entity." + if isinstance(value, EntityReferencesFactValue): + if any( + ContextReducer._resolve_entity(item, references) is None + for item in value.entity_ids + ): + return "The fact value references an unknown entity." + return None + + @staticmethod + def _normalize_value_references( + assertion: FactAssertion, + references: dict[str, str], + ) -> FactAssertion: + if not isinstance(assertion, PresentAssertion): + return assertion + value = assertion.value + if isinstance(value, EntityReferenceFactValue): + resolved_id = ContextReducer._resolve_entity(value.entity_id, references) + if resolved_id is not None: + return assertion.model_copy( + update={ + "value": value.model_copy( + update={"entity_id": resolved_id} + ) + } + ) + if isinstance(value, EntityReferencesFactValue): + resolved_ids = tuple( + ContextReducer._resolve_entity(item, references) or item + for item in value.entity_ids + ) + return assertion.model_copy( + update={ + "value": value.model_copy( + update={"entity_ids": resolved_ids} + ) + } + ) + return assertion + + @staticmethod + def _decision( + index: int, + operation: SetFactOperation, + status: FactDecisionStatus, + reason: str, + subject_id: str | None = None, + *, + fact_id: str | None = None, + superseded_fact_id: str | None = None, + ) -> FactDecision: + return FactDecision( + operation_index=index, + status=status, + operation=operation.operation, + definition_key=operation.definition_key, + subject_entity_id=subject_id, + fact_id=fact_id, + superseded_fact_id=superseded_fact_id, + reason=reason, + ) diff --git a/backend/conversation_context/registry.py b/backend/conversation_context/registry.py new file mode 100644 index 00000000..00aeb68f --- /dev/null +++ b/backend/conversation_context/registry.py @@ -0,0 +1,362 @@ +"""Registered definitions for facts the current runtime understands.""" + +from __future__ import annotations + +from decimal import Decimal +from enum import Enum +from typing import Any + +from pydantic import BaseModel, ConfigDict, field_validator, model_validator + +from conversation_context.models import ( + ConversationContext, + EntityKind, + FactValue, + MoneyFactValue, + PresentAssertion, +) + + +class FactValueKind(str, Enum): + BOOLEAN = "boolean" + INTEGER = "integer" + MONEY = "money" + TEXT = "text" + TEXT_SET = "text_set" + ENTITY_REFERENCE = "entity_reference" + ENTITY_REFERENCES = "entity_references" + + +class TemporalSemantics(str, Enum): + CURRENT = "current" + AS_OF_YEAR = "as_of_year" + SCENARIO = "scenario" + + +class FactUpdatePolicy(str, Enum): + REQUIRE_EXPLICIT_CORRECTION = "require_explicit_correction" + REPLACE_ON_EXPLICIT_ASSERTION = "replace_on_explicit_assertion" + + +class FactDefinition(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + key: str + version: str = "1" + value_kind: FactValueKind + subject_kinds: frozenset[EntityKind] + cardinality: str = "one_per_subject_scope" + temporal_semantics: TemporalSemantics = TemporalSemantics.CURRENT + update_policy: FactUpdatePolicy = FactUpdatePolicy.REQUIRE_EXPLICIT_CORRECTION + label: str + sensitivity: str = "personal" + allow_explicit_absence: bool = False + allowed_text_values: tuple[str, ...] = () + minimum: Decimal | None = None + maximum: Decimal | None = None + engine_binding: str | None = None + + @field_validator("key", "version", "label", "cardinality", "sensitivity") + @classmethod + def non_empty(cls, value: str) -> str: + if not value.strip(): + raise ValueError("must not be empty") + return value + + @model_validator(mode="after") + def validate_definition(self) -> "FactDefinition": + if not self.subject_kinds: + raise ValueError("at least one subject kind is required") + if self.allowed_text_values and self.value_kind is not FactValueKind.TEXT: + raise ValueError("allowed_text_values require a text fact") + return self + + def validate_value(self, value: FactValue) -> str | None: + if value.kind != self.value_kind.value: + return f"expected {self.value_kind.value} value, received {value.kind}" + scalar: Decimal | None = None + if isinstance(value, MoneyFactValue): + scalar = value.amount + elif value.kind == "integer": + scalar = Decimal(value.value) + if scalar is not None and self.minimum is not None and scalar < self.minimum: + return f"value must be at least {self.minimum}" + if scalar is not None and self.maximum is not None and scalar > self.maximum: + return f"value must be at most {self.maximum}" + if value.kind == "text" and self.allowed_text_values: + if value.value not in self.allowed_text_values: + return "value is not in the registered allowed set" + return None + + +class FactDefinitionRegistry: + def __init__(self, definitions: tuple[FactDefinition, ...] = ()) -> None: + self._definitions: dict[tuple[str, str], FactDefinition] = {} + for definition in definitions: + self.register(definition) + + def register(self, definition: FactDefinition) -> None: + identity = (definition.key, definition.version) + if identity in self._definitions: + raise ValueError( + f"Duplicate fact definition: {definition.key}@{definition.version}" + ) + self._definitions[identity] = definition + + def get(self, key: str, version: str = "1") -> FactDefinition: + try: + return self._definitions[(key, version)] + except KeyError as exc: + raise KeyError(f"Unknown fact definition: {key}@{version}") from exc + + def find_by_engine_binding( + self, + variable_name: str, + *, + entity: str | None = None, + ) -> FactDefinition | None: + """Return a definition backed by one exact PolicyEngine variable.""" + + matches: list[FactDefinition] = [] + for definition in self._definitions.values(): + binding = definition.engine_binding + if binding is None: + continue + binding_parts = binding.split(".", 1) + binding_entity = binding_parts[0] if len(binding_parts) == 2 else None + binding_name = binding_parts[-1] + if binding_name != variable_name: + continue + if entity is not None and binding_entity not in {None, entity}: + continue + matches.append(definition) + if len(matches) == 1: + return matches[0] + return None + + def ensure_engine_definition( + self, + *, + variable_name: str, + entity: str, + label: str, + value_kind: FactValueKind, + ) -> FactDefinition: + """Materialize registry metadata from one verified catalogue record.""" + + existing = self.find_by_engine_binding(variable_name, entity=entity) + if existing is not None: + return existing + subject_kind = { + "person": EntityKind.PERSON, + "household": EntityKind.HOUSEHOLD, + "benunit": EntityKind.HOUSEHOLD, + }.get(entity) + if subject_kind is None: + raise ValueError( + f"PolicyEngine entity {entity!r} has no conversation-entity mapping." + ) + definition = FactDefinition( + key=f"pe.{entity}.{variable_name}", + value_kind=value_kind, + subject_kinds=frozenset({subject_kind}), + label=label, + engine_binding=f"{entity}.{variable_name}", + ) + self.register(definition) + return definition + + def restore_engine_definitions(self, context: ConversationContext) -> None: + """Restore verified generated definitions referenced by persisted facts.""" + + for fact in context.facts: + try: + self.get(fact.definition_key, fact.definition_version) + continue + except KeyError: + pass + parts = fact.definition_key.split(".", 2) + if ( + len(parts) != 3 + or parts[0] != "pe" + or fact.definition_version != "1" + or not isinstance(fact.assertion, PresentAssertion) + ): + continue + entity, variable_name = parts[1:] + try: + value_kind = FactValueKind(fact.assertion.value.kind) + except ValueError: + continue + self.ensure_engine_definition( + variable_name=variable_name, + entity=entity, + label=variable_name.replace("_", " ").title(), + value_kind=value_kind, + ) + + def definitions(self) -> tuple[FactDefinition, ...]: + return tuple(self._definitions.values()) + + def model_projection(self) -> tuple[dict[str, object], ...]: + return tuple( + definition.model_dump(mode="json", exclude_none=True) + for definition in self.definitions() + ) + + +def _fact( + key: str, + value_kind: FactValueKind, + subjects: frozenset[EntityKind], + label: str, + **kwargs: Any, +) -> FactDefinition: + return FactDefinition( + key=key, + value_kind=value_kind, + subject_kinds=subjects, + label=label, + **kwargs, + ) + + +def build_default_fact_registry() -> FactDefinitionRegistry: + person = frozenset({EntityKind.PERSON}) + household = frozenset({EntityKind.HOUSEHOLD}) + scenario = frozenset({EntityKind.POLICY_SCENARIO}) + non_negative = {"minimum": Decimal("0")} + return FactDefinitionRegistry( + ( + _fact( + "person.name", + FactValueKind.TEXT, + person, + "Name", + engine_binding=None, + ), + _fact( + "person.age", + FactValueKind.INTEGER, + person, + "Age", + minimum=Decimal("0"), + maximum=Decimal("120"), + engine_binding="person.age", + ), + _fact( + "person.employment_income", + FactValueKind.MONEY, + person, + "Employment income", + **non_negative, + engine_binding="person.employment_income", + ), + _fact( + "person.self_employment_income", + FactValueKind.MONEY, + person, + "Self-employment income", + **non_negative, + engine_binding="person.self_employment_income", + ), + _fact( + "person.pension_income", + FactValueKind.MONEY, + person, + "Pension income", + **non_negative, + engine_binding="person.pension_income", + ), + _fact( + "person.childcare_expenses", + FactValueKind.MONEY, + person, + "Childcare expenses", + allow_explicit_absence=True, + **non_negative, + engine_binding="person.childcare_expenses", + ), + _fact( + "person.medical_expenses", + FactValueKind.MONEY, + person, + "Medical expenses", + allow_explicit_absence=True, + **non_negative, + ), + _fact( + "household.members", + FactValueKind.ENTITY_REFERENCES, + household, + "Household members", + ), + _fact( + "household.is_married", + FactValueKind.BOOLEAN, + household, + "Married or in a civil partnership", + engine_binding="benunit.is_married", + ), + _fact( + "household.has_children", + FactValueKind.BOOLEAN, + household, + "Has children", + ), + _fact( + "household.rent", + FactValueKind.MONEY, + household, + "Rent", + allow_explicit_absence=True, + **non_negative, + engine_binding="household.rent", + ), + _fact( + "household.council_tax", + FactValueKind.MONEY, + household, + "Council Tax", + allow_explicit_absence=True, + **non_negative, + engine_binding="household.council_tax", + ), + _fact( + "household.country", + FactValueKind.TEXT, + household, + "Country within the UK", + allowed_text_values=( + "ENGLAND", + "SCOTLAND", + "WALES", + "NORTHERN_IRELAND", + ), + engine_binding="household.country", + ), + _fact( + "analysis.policy_year", + FactValueKind.INTEGER, + frozenset({EntityKind.HOUSEHOLD, EntityKind.POLICY_SCENARIO}), + "Policy year", + minimum=Decimal("2000"), + maximum=Decimal("2100"), + temporal_semantics=TemporalSemantics.AS_OF_YEAR, + ), + _fact( + "analysis.requested_outputs", + FactValueKind.TEXT_SET, + frozenset({EntityKind.HOUSEHOLD, EntityKind.POLICY_SCENARIO}), + "Requested calculation outputs", + update_policy=FactUpdatePolicy.REPLACE_ON_EXPLICIT_ASSERTION, + ), + _fact( + "policy.reform_instruction", + FactValueKind.TEXT, + scenario, + "Policy reform instruction", + temporal_semantics=TemporalSemantics.SCENARIO, + ), + ) + ) diff --git a/backend/conversation_context/repository.py b/backend/conversation_context/repository.py new file mode 100644 index 00000000..c0c00e2e --- /dev/null +++ b/backend/conversation_context/repository.py @@ -0,0 +1,24 @@ +"""Persistence contract for the latest typed conversation context.""" + +from __future__ import annotations + +from typing import Protocol + +from conversation_context.models import ConversationContext + + +class ConversationContextConflict(ValueError): + """The stored revision changed after the caller loaded it.""" + + +class ConversationContextRepository(Protocol): + def load(self, conversation_id: str) -> ConversationContext: ... + + def save( + self, + context: ConversationContext, + *, + expected_revision: int, + ) -> ConversationContext: ... + + def delete(self, conversation_id: str) -> None: ... diff --git a/backend/conversation_context/tools.py b/backend/conversation_context/tools.py new file mode 100644 index 00000000..d6f2ce8f --- /dev/null +++ b/backend/conversation_context/tools.py @@ -0,0 +1,592 @@ +"""Private typed operations for model-assisted fact extraction and reduction.""" + +from __future__ import annotations + +from enum import Enum +import json +from typing import Protocol + +from pydantic import BaseModel, ConfigDict, Field, ValidationError + +from config import ( + DEFAULT_COMPLEX_MODEL, + DEFAULT_FAST_MODEL, + DEFAULT_TEMPERATURE, + get_async_client, +) +from conversation_context.change_pipeline import ( + ContextChangeApplier, + ContextChangeProposal, + ContextChangeValidator, + ContextValidationOutcome, + ContextValidationIssue, + SemanticClaimReview, + ValidateContextChangeInput, +) +from conversation_context.models import ( + ContextPatch, + ContextReduction, + ConversationContext, + FactClaim, +) +from conversation_context.projection import ContextProjection +from conversation_context.reducer import ContextReducer +from conversation_context.registry import FactDefinition +from tools.contracts import CallerType, Tool, ToolCallContext, ToolSpec, Visibility + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class ContextConversationExcerpt(StrictModel): + role: str + content: str + + +class ContextModelUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + def plus(self, other: "ContextModelUsage") -> "ContextModelUsage": + return ContextModelUsage( + input_tokens=self.input_tokens + other.input_tokens, + output_tokens=self.output_tokens + other.output_tokens, + cache_creation_input_tokens=( + self.cache_creation_input_tokens + + other.cache_creation_input_tokens + ), + cache_read_input_tokens=( + self.cache_read_input_tokens + other.cache_read_input_tokens + ), + ) + + +class ProposeContextChangeInput(StrictModel): + current_message: str + conversation: tuple[ContextConversationExcerpt, ...] + context: ContextProjection + fact_definitions: tuple[FactDefinition, ...] + previous_proposal: ContextChangeProposal | None = None + repair_issues: tuple[ContextValidationIssue, ...] = () + + +class ContextProposalStatus(str, Enum): + READY = "ready" + NEEDS_CLARIFICATION = "needs_clarification" + + +class ProposeContextChangeOutput(ContextChangeProposal): + status: ContextProposalStatus = ContextProposalStatus.READY + issues: tuple[ContextValidationIssue, ...] = () + provider_attempts: int = Field(default=1, ge=1, le=2) + usage: ContextModelUsage = Field(default_factory=ContextModelUsage) + + +class ContextInterpreter(Protocol): + async def propose( + self, + request: "ProposeContextChangeInput", + ) -> ProposeContextChangeOutput: ... + + +class ContextSemanticReview(StrictModel): + reviews: tuple[SemanticClaimReview, ...] = () + + +class ContextSemanticReviewOutput(ContextSemanticReview): + usage: ContextModelUsage = Field(default_factory=ContextModelUsage) + + +class ContextProposalReviewer(Protocol): + async def review( + self, + request: ValidateContextChangeInput, + ) -> ContextSemanticReviewOutput: ... + + +class AnthropicContextProposalReviewer: + """Check whether proposed claim semantics match the current user message.""" + + async def review( + self, + request: ValidateContextChangeInput, + ) -> ContextSemanticReviewOutput: + client = get_async_client() # type: ignore[no-untyped-call] + tool = { + "name": "submit_context_semantic_review", + "description": ( + "Return one semantic-fidelity verdict for every proposed claim." + ), + "input_schema": ContextSemanticReview.model_json_schema(), + } + system = ( + "Independently review each fact_claim in proposal.changes. Ignore " + "pending_resolution_response items; deterministic validation checks those " + "against the referenced retained record. " + "Use evidence as the exact current user message and the supplied known " + "entity identities and active scenario scope only to resolve subjects and " + "conversational references. You do not receive retained fact values. Evaluate " + "whether each claim's definition concept, grammatical subject, " + "subject_references, relationship, and value mean what the current message " + "says. Typed " + "conversation context represents the active scenario being discussed, not " + "only the user's real-world biography. Treat facts inside a hypothetical or " + "counterfactual request as proposed facts for that calculation scenario; do " + "not reject them merely because the user said 'what if' or used another " + "conditional construction. A direct claim is valid only when the message " + "assigns that value to exactly that one " + "subject. A total, comparison, or other relationship over several known " + "entities must remain one relational claim over every denoted entity; do not " + "accept a model-calculated allocation or a direct assignment to one member. " + "A relational claim is intentionally unresolved: accept it without a " + "per-entity allocation, and accept a registered per-entity definition_key as " + "the concept to resolve over all referenced subjects. Never reject a sum " + "claim merely because its distribution is not supplied. Explicit per-entity " + "amounts remain separate direct claims. Accept entity and membership facts " + "that are a faithful typed representation of the scenario stated in the " + "message. Report only a clear semantic contradiction, not a possible " + "alternative interpretation. Do not assess " + "registration, data types, periods, correction flags, arithmetic, capability " + "requirements, or persistence; deterministic validation handles those. Return " + "exactly one review for every proposal claim and copy that claim's exact " + "opaque claim_id into the review. Set " + "supported=true when the claim is semantically faithful and supported=false " + "only for a clear mismatch. The reason must agree with the boolean: if the " + "reason says a claim is valid, faithful, or accurate, supported must be true. " + "Include the exact supporting text as evidence." + ) + response = await client.messages.create( + model=DEFAULT_FAST_MODEL, + max_tokens=900, + temperature=DEFAULT_TEMPERATURE, + system=system, + messages=[ + { + "role": "user", + "content": json.dumps( + { + "current_message": request.evidence, + "known_entities": [ + entity.model_dump(mode="json") + for entity in request.context.entities + ], + "active_scope_id": request.context.focus.scope_id, + "proposal": request.proposal.model_dump(mode="json"), + }, + ensure_ascii=False, + ), + } + ], + tools=[tool], + tool_choice={ + "type": "tool", + "name": "submit_context_semantic_review", + }, + ) + block = next( + ( + item + for item in response.content + if getattr(item, "type", None) == "tool_use" + and getattr(item, "name", None) + == "submit_context_semantic_review" + ), + None, + ) + usage = AnthropicContextInterpreter._usage(response) + if block is None: + return ContextSemanticReviewOutput( + usage=usage, + ) + try: + review = ContextSemanticReview.model_validate(block.input) + except ValidationError: + return ContextSemanticReviewOutput( + usage=usage, + ) + return ContextSemanticReviewOutput(reviews=review.reviews, usage=usage) + + +class AnthropicContextInterpreter: + """Propose candidate facts without deciding capability requirements.""" + + async def propose( + self, + request: "ProposeContextChangeInput", + ) -> ProposeContextChangeOutput: + client = get_async_client() # type: ignore[no-untyped-call] + tool = { + "name": "submit_context_change", + "description": ( + "Return candidate entities and one ordered list of typed context changes." + ), + "input_schema": ContextChangeProposal.model_json_schema(), + } + system = ( + "Interpret only facts explicitly stated, corrected, or denied in the exact " + "current user message into one declarative context proposal. Return every " + "new assertion and pending-resolution response in the single ordered changes " + "list. Use kind=fact_claim for a new assertion and " + "kind=pending_resolution_response for an action on one exact retained " + "pending resolution. Use " + "conversation excerpts only to resolve pronouns, aliases, and which pending " + "requirements a short answer addresses; do not copy old transcript claims " + "that are absent from typed context. Return every supported assertion exactly " + "once as a fact_claim. Claims are declarative; never return a context patch, set-fact " + "operation, or a separate unresolved-claim collection. Use only supplied fact " + "definitions, entity identifiers, and scope identifiers. Declare a newly " + "mentioned person or policy scenario in candidate_entities and reference that " + "same candidate from claims. A direct claim has exactly one grammatical " + "subject. A sum claim has every person denoted by its grammatical subject and " + "at least two distinct subject references. A statement about several people " + "is a relationship over those people, not a direct fact about one person and " + "not an invented household-level version of a person fact. Mark correction " + "only when the current message clearly corrects an active value. Use " + "definition_key for every non-monetary claim; it must equal one exact key " + "from fact_definitions, because concept alone does not select a definition. " + "When a short answer satisfies a pending requirement, copy that requirement's " + "fact_key into definition_key, its subject_entity_id into subject_references, " + "and its scope_id into scope_id. " + "Set correction true when the current message explicitly tells you to use, " + "set, change, correct, or keep a value and typed context already contains a " + "different value for that same fact, subject, and scope. The flag describes " + "fact replacement; it does not require the user to say that an earlier value " + "was wrong. " + "explicit_absence only when the user says a registered optional value does not " + "apply; zero is a present numeric value. Leave an unstated monetary period null. " + "Validation may inherit it only from one compatible active fact " + "or pending fact-resolution constraint; otherwise the applicable capability " + "decides from the exact current message whether to apply a documented " + "non-persistent default or ask the user. " + "Do not infer an income source or frequency, apply a default, decide what " + "a capability requires, ask a question, select a capability, allocate a total, " + "or calculate anything. If the current message clearly accepts or rejects one " + "pending fact-resolution proposal whose status is awaiting_confirmation, put " + "one pending_resolution_response in changes with its exact proposal_id, " + "action=accept or action=reject, the shortest exact current-message evidence, " + "and no updates. Do not reconstruct its assignments. When a short answer " + "supplies one or more fields requested by a needs_clarification resolution, " + "return action=supply and only schema-addressed updates supported by the exact " + "current message. For example, an annual-period answer supplies path " + "[\"value\",\"period\"] and value \"annual\". Do not copy its retained " + "amount, subjects, relationship, mapping, or source evidence into the current " + "proposal. An explicit user-supplied per-person breakdown remains ordinary " + "direct fact_claim items; do not also emit a supply response for that " + "breakdown. Supply updates may address only fields of the retained source " + "fact claim, such as [\"value\",\"period\"] or " + "[\"value\",\"amount\"]. Never update resolver-owned terms, assignments, " + "known values, prompts, or mapping metadata. You cannot create a pending " + "fact-resolution proposal or replace " + "pending capability questions; those are server-owned. Treat calculation " + "outputs explicitly named in the current request as an " + "analysis.requested_outputs fact on the relevant household or policy " + "scenario; do not invent outputs that were not requested. A short answer such " + "as 'neither' may satisfy every compatible pending requirement. Set each " + "claim's evidence to the shortest exact contiguous quote from current_message " + "that supports it; for a one-token answer such as '27', use exactly that " + "token. A pending_resolution_response may identify only an exact proposal_id listed in " + "pending_fact_resolutions, never a pending capability identifier. Set " + "expected_revision to the supplied context revision. Return no operations " + "and no claims when the message establishes no contextual fact." + ) + issues = request.repair_issues + correction = self._repair_instruction(issues) if issues else "" + max_attempts = 1 if issues else 2 + usage_total = ContextModelUsage() + for attempt in range(max_attempts): + response = await client.messages.create( + model=(DEFAULT_COMPLEX_MODEL if issues else DEFAULT_FAST_MODEL), + max_tokens=1800, + temperature=DEFAULT_TEMPERATURE, + system=system, + messages=[ + { + "role": "user", + "content": request.model_dump_json() + correction, + } + ], + tools=[tool], + tool_choice={"type": "tool", "name": "submit_context_change"}, + ) + usage_total = usage_total.plus(self._usage(response)) + block = next( + ( + item + for item in response.content + if getattr(item, "type", None) == "tool_use" + and getattr(item, "name", None) == "submit_context_change" + ), + None, + ) + if block is None: + issues = ( + ContextValidationIssue( + code="missing_structured_context_output", + path=("submit_context_change",), + message=( + "The context interpreter did not return the required " + "structured tool output." + ), + evidence=request.current_message, + ), + ) + correction = self._repair_instruction(issues) + continue + try: + submission = ContextChangeProposal.model_validate(block.input) + except ValidationError as exc: + issues = tuple( + ContextValidationIssue( + code="invalid_context_submission", + path=tuple(str(item) for item in error["loc"]), + message=error["msg"], + evidence=request.current_message, + ) + for error in exc.errors(include_url=False, include_input=False) + ) + correction = self._repair_instruction(issues) + continue + return ProposeContextChangeOutput( + **submission.model_dump(mode="python"), + provider_attempts=attempt + 1, + usage=usage_total, + ) + return ProposeContextChangeOutput( + status=ContextProposalStatus.NEEDS_CLARIFICATION, + expected_revision=request.context.revision, + issues=issues, + provider_attempts=max_attempts, + usage=usage_total, + ) + + @staticmethod + def _repair_instruction(issues: tuple[ContextValidationIssue, ...]) -> str: + payload = [issue.model_dump(mode="json") for issue in issues] + return ( + "\n\nThe previous context-change proposal did not validate. Correct only " + "the previous_proposal using these machine-readable issues. Preserve its " + "valid changes and revise only invalid fields. Do not repeat a value rejected at the " + "cited path. In particular, an " + "unmapped_fact_claim requires definition_key to be set to one exact supplied " + "fact-definition key; concept does not satisfy it. An uncited evidence issue " + "requires a shortest exact quote from current_message. Preserve valid fields, " + "and correct a pending-resolution response only from its referenced record and " + "the exact current-message evidence. Never copy retained fields into a supply " + "response. If ordinary direct fact_claim items already express an explicit " + "per-subject breakdown, remove a redundant supply response entirely. A supply " + "response never updates resolver-owned terms, assignments, known values, " + "prompts, or mapping metadata. " + "For semantic_claim_mismatch, revise the cited claim to match the review. If a " + "direct claim was rejected because the message relates several entities, keep " + "the exact stated value as one relational claim over all denoted subjects; do " + "not calculate an allocation or drop a current-message monetary value. " + "A context_operation_conflicted issue means a generated fact would replace a " + "different active value without permission: set that claim's correction field " + "true only when the current message explicitly directs the new or retained " + "value; otherwise remove the unsupported claim. " + "use supplied stable entity references, and do not ask the user a question, " + "invent a default, or create a proposal_response unless its exact identifier " + "appears in pending_fact_resolutions:\n" + + json.dumps(payload, ensure_ascii=False) + ) + + @staticmethod + def _usage(response: object) -> ContextModelUsage: + usage = getattr(response, "usage", None) + return ContextModelUsage( + input_tokens=getattr(usage, "input_tokens", 0), + output_tokens=getattr(usage, "output_tokens", 0), + cache_creation_input_tokens=getattr( + usage, + "cache_creation_input_tokens", + 0, + ), + cache_read_input_tokens=getattr( + usage, + "cache_read_input_tokens", + 0, + ), + ) + + +class ProposeContextChangeTool( + Tool["ProposeContextChangeInput", ProposeContextChangeOutput] +): + spec = ToolSpec( + identifier="propose_context_change", + version="1", + description=( + "Interpret candidate entities and declarative fact claims from the exact " + "current message." + ), + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=ProposeContextChangeInput, + output_model=ProposeContextChangeOutput, + ) + + def __init__(self, interpreter: ContextInterpreter) -> None: + self._interpreter = interpreter + + async def run( + self, + tool_input: "ProposeContextChangeInput", + context: ToolCallContext, + ) -> ProposeContextChangeOutput: + result = await self._interpreter.propose(tool_input) + context.record_model_usage(**result.usage.model_dump()) + return result + + +class ReduceContextPatchInput(StrictModel): + context: ConversationContext + patch: ContextPatch + turn_id: str + evidence: str + + +class ReduceContextPatchTool(Tool[ReduceContextPatchInput, ContextReduction]): + spec = ToolSpec( + identifier="reduce_context_patch", + version="1", + description=( + "Validate candidate registered facts and produce the next context revision." + ), + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=ReduceContextPatchInput, + output_model=ContextReduction, + ) + + def __init__(self, reducer: ContextReducer) -> None: + self._reducer = reducer + + async def run( + self, + tool_input: ReduceContextPatchInput, + context: ToolCallContext, + ) -> ContextReduction: + del context + return self._reducer.reduce( + tool_input.context, + tool_input.patch, + turn_id=tool_input.turn_id, + evidence=tool_input.evidence, + ) + + +class ValidateContextChangeTool( + Tool[ValidateContextChangeInput, ContextValidationOutcome] +): + spec = ToolSpec( + identifier="validate_context_change", + version="1", + description=( + "Validate one complete current-message context update without partial " + "persistence." + ), + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=ValidateContextChangeInput, + output_model=ContextValidationOutcome, + ) + + def __init__( + self, + validator: ContextChangeValidator, + reviewer: ContextProposalReviewer | None = None, + ) -> None: + self._validator = validator + self._reviewer = reviewer + + async def run( + self, + tool_input: ValidateContextChangeInput, + context: ToolCallContext, + ) -> ContextValidationOutcome: + semantic_issues: tuple[ContextValidationIssue, ...] = () + semantic_reviews: tuple[SemanticClaimReview, ...] = () + if ( + self._reviewer is not None + and not tool_input.claims_resolved + and tool_input.proposal.claims + ): + review = await self._reviewer.review(tool_input) + context.record_model_usage(**review.usage.model_dump()) + semantic_reviews = review.reviews + claim_indexes = { + change.claim_id: index + for index, change in enumerate(tool_input.proposal.changes) + if isinstance(change, FactClaim) + } + expected_claim_ids = set(claim_indexes) + reviewed_claim_ids = {item.claim_id for item in semantic_reviews} + complete_review = ( + len(claim_indexes) == len(tool_input.proposal.claims) + and len(semantic_reviews) == len(reviewed_claim_ids) + and reviewed_claim_ids == expected_claim_ids + ) + semantic_issues = tuple( + ContextValidationIssue( + code="semantic_claim_mismatch", + message=item.reason, + path=( + "proposal", + "changes", + str(claim_indexes[item.claim_id]), + ), + claim_index=claim_indexes[item.claim_id], + evidence=item.evidence, + ) + for item in semantic_reviews + if not item.supported and item.claim_id in claim_indexes + ) + if not complete_review: + semantic_issues = ( + ContextValidationIssue( + code="semantic_review_failed", + message=( + "The semantic review did not return exactly one verdict for " + "every proposed claim." + ), + path=("proposal", "claims"), + evidence=tool_input.evidence, + ), + ) + outcome = self._validator.validate( + tool_input, + semantic_issues=semantic_issues, + ) + return outcome.model_copy( + update={"semantic_reviews": semantic_reviews} + ) + + +class ApplyContextChangeInput(StrictModel): + outcome: ContextValidationOutcome + + +class ApplyContextChangeTool(Tool[ApplyContextChangeInput, ConversationContext]): + spec = ToolSpec( + identifier="apply_context_change", + version="1", + description="Persist one fully validated context change atomically.", + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=ApplyContextChangeInput, + output_model=ConversationContext, + ) + + def __init__(self, applier: ContextChangeApplier) -> None: + self._applier = applier + + async def run( + self, + tool_input: ApplyContextChangeInput, + context: ToolCallContext, + ) -> ConversationContext: + del context + return self._applier.apply(tool_input.outcome) diff --git a/backend/conversation_context/variable_resolution.py b/backend/conversation_context/variable_resolution.py new file mode 100644 index 00000000..e5fb0640 --- /dev/null +++ b/backend/conversation_context/variable_resolution.py @@ -0,0 +1,988 @@ +"""Catalogue-backed resolution of declarative monetary fact claims.""" + +from __future__ import annotations + +from decimal import Decimal +from enum import Enum +import json +from typing import Literal, Protocol + +from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator + +from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_async_client +from conversation_context.change_pipeline import ContextChangeProposal, ContextValidationIssue +from conversation_context.models import ( + AddPendingFactResolutionOperation, + ClaimedMoneyValue, + ContextFact, + ContextOperation, + ContextPatch, + ConversationContext, + EntityKind, + FactClaim, + FactClaimRelationship, + FactResolutionAssignment, + FactResolutionSupplement, + FactResolutionStatus, + FactResolutionTerm, + MoneyFactValue, + MoneyPeriod, + PendingFactResolution, + PendingFactResolutionResponse, + PendingResolutionAction, + PresentAssertion, + SetFactOperation, +) +from conversation_context.registry import FactDefinitionRegistry, FactValueKind +from tools.contracts import CallerType, Tool, ToolCallContext, ToolSpec, Visibility +from tools.typed_models import SafeToolOutput + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class MappingConfidence(str, Enum): + HIGH = "high" + MEDIUM = "medium" + LOW = "low" + + +class MappingStatus(str, Enum): + MATCHED = "matched" + AMBIGUOUS = "ambiguous" + UNSUPPORTED = "unsupported" + + +class PolicyEngineVariableCandidate(StrictModel): + name: str + label: str | None = None + entity: str + description: str | None = None + definition_period: str | None = None + value_type: str | None = None + + +class VariableMappingSelection(StrictModel): + status: MappingStatus + variable_name: str | None = None + confidence: MappingConfidence = MappingConfidence.LOW + target_period: MoneyPeriod | None = None + + +class VariableMappingUsage(StrictModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + def plus(self, other: "VariableMappingUsage") -> "VariableMappingUsage": + return VariableMappingUsage( + input_tokens=self.input_tokens + other.input_tokens, + output_tokens=self.output_tokens + other.output_tokens, + cache_creation_input_tokens=( + self.cache_creation_input_tokens + + other.cache_creation_input_tokens + ), + cache_read_input_tokens=( + self.cache_read_input_tokens + other.cache_read_input_tokens + ), + ) + + +class VariableMappingResult(StrictModel): + selection: VariableMappingSelection + usage: VariableMappingUsage = Field(default_factory=VariableMappingUsage) + + +class VariableMapper(Protocol): + async def select( + self, + *, + claim: FactClaim, + candidates: tuple[PolicyEngineVariableCandidate, ...], + context: ConversationContext, + registry: FactDefinitionRegistry, + validation_issues: tuple[ContextValidationIssue, ...] = (), + ) -> VariableMappingResult: ... + + +class AnthropicVariableMapper: + """Select one exact returned variable without calculating assignments.""" + + async def select( + self, + *, + claim: FactClaim, + candidates: tuple[PolicyEngineVariableCandidate, ...], + context: ConversationContext, + registry: FactDefinitionRegistry, + validation_issues: tuple[ContextValidationIssue, ...] = (), + ) -> VariableMappingResult: + client = get_async_client() # type: ignore[no-untyped-call] + tool = { + "name": "submit_variable_mapping", + "description": "Select at most one exact returned PolicyEngine variable.", + "input_schema": VariableMappingSelection.model_json_schema(), + } + active_bindings = [] + for fact in context.active_facts(): + try: + definition = registry.get( + fact.definition_key, + fact.definition_version, + ) + except KeyError: + continue + if definition.engine_binding is None: + continue + active_bindings.append( + { + "definition_key": definition.key, + "engine_binding": definition.engine_binding, + "subject_entity_id": fact.subject_entity_id, + "assertion": fact.assertion.model_dump(mode="json"), + } + ) + system = ( + "Map the cited claim to at most one PolicyEngine variable from the exact " + "candidate list. Never invent or alter a variable name. Use the current " + "active engine bindings to interpret a concise continuation. Mark confidence " + "high only when the concept, entity, value type, and relationship align; " + "otherwise return ambiguous or unsupported. Select a target period only " + "when it follows from the returned variable definition period, an explicit " + "claim period, or a compatible active binding. Do not calculate, allocate, " + "apply defaults, or write a user-facing answer." + ) + payload = { + "claim": claim.model_dump(mode="json"), + "candidates": [item.model_dump(mode="json") for item in candidates], + "entities": [item.model_dump(mode="json") for item in context.entities], + "active_engine_bindings": active_bindings, + "validation_issues": [ + item.model_dump(mode="json") for item in validation_issues + ], + } + last_error: Exception | None = None + for _attempt in range(2): + response = await client.messages.create( + model=DEFAULT_FAST_MODEL, + max_tokens=500, + temperature=DEFAULT_TEMPERATURE, + system=system, + messages=[ + { + "role": "user", + "content": json.dumps(payload, separators=(",", ":")), + } + ], + tools=[tool], + tool_choice={"type": "tool", "name": "submit_variable_mapping"}, + ) + block = next( + ( + item + for item in response.content + if getattr(item, "type", None) == "tool_use" + and getattr(item, "name", None) == "submit_variable_mapping" + ), + None, + ) + if block is None: + last_error = RuntimeError( + "Variable mapper did not return structured output." + ) + continue + try: + selection = VariableMappingSelection.model_validate(block.input) + except ValidationError as exc: + last_error = exc + continue + usage = getattr(response, "usage", None) + return VariableMappingResult( + selection=selection, + usage=VariableMappingUsage( + input_tokens=getattr(usage, "input_tokens", 0), + output_tokens=getattr(usage, "output_tokens", 0), + cache_creation_input_tokens=getattr( + usage, + "cache_creation_input_tokens", + 0, + ), + cache_read_input_tokens=getattr( + usage, + "cache_read_input_tokens", + 0, + ), + ), + ) + raise RuntimeError( + "Variable mapper failed to return a valid selection after one retry." + ) from last_error + + +class FactResolutionDecision(StrictModel): + claim_id: str + status: FactResolutionStatus | Literal["resolved"] + candidates: tuple[PolicyEngineVariableCandidate, ...] = () + selection_source: Literal["proposal", "resolver_model"] + selection: VariableMappingSelection + proposal: PendingFactResolution | None = None + operation: SetFactOperation | None = None + + @model_validator(mode="after") + def validate_result(self) -> "FactResolutionDecision": + if self.status == "resolved": + if self.operation is None or self.proposal is not None: + raise ValueError( + "a resolved direct fact requires one operation and no pending proposal" + ) + elif self.proposal is None or self.operation is not None: + raise ValueError( + "an incomplete or confirmable resolution requires one pending proposal" + ) + return self + + +class ResolveContextChangeInput(StrictModel): + context: ConversationContext + proposal: ContextChangeProposal | None = None + validation_issues: tuple[ContextValidationIssue, ...] = () + claims: tuple[FactClaim, ...] + turn_id: str + evidence: str + + @model_validator(mode="after") + def validate_monetary_claims(self) -> "ResolveContextChangeInput": + if any(not isinstance(claim.value, ClaimedMoneyValue) for claim in self.claims): + raise ValueError( + "catalogue-backed fact resolution accepts only monetary fact claims" + ) + if self.proposal is not None: + proposal_claim_ids = {claim.claim_id for claim in self.proposal.claims} + supplied_proposal_ids = { + response.proposal_id + for response in self.proposal.proposal_responses + if response.action is PendingResolutionAction.SUPPLY + } + proposal_claim_ids.update( + pending.source_claim.claim_id + for pending in self.context.pending_fact_resolutions + if pending.proposal_id in supplied_proposal_ids + and pending.source_claim is not None + ) + if any(claim.claim_id not in proposal_claim_ids for claim in self.claims): + raise ValueError( + "every claim requiring resolution must come from the current " + "proposal or a referenced retained source claim" + ) + return self + + +class ResolveContextChangeOutput(StrictModel): + patch: ContextPatch + decisions: tuple[FactResolutionDecision, ...] + usage: VariableMappingUsage = Field(default_factory=VariableMappingUsage) + + +class FactConstraintIssue(str, Enum): + SUBJECT_REQUIRED = "subject_required" + ALLOCATION_REQUIRED = "allocation_required" + INCONSISTENT_TOTAL = "inconsistent_total" + + +class FactConstraintSolution(StrictModel): + terms: tuple[FactResolutionTerm, ...] + subject_entity_id: str | None = None + amount: Decimal | None = None + issue: FactConstraintIssue | None = None + + +def _solve_constraint( + *, + variable_name: str, + relationship: FactClaimRelationship, + entity_ids: tuple[str, ...], + known_values: dict[str, Decimal | None], + expected_total: Decimal, +) -> FactConstraintSolution: + terms = tuple( + FactResolutionTerm( + variable_name=variable_name, + subject_entity_id=entity_id, + known_value=known_values.get(entity_id), + ) + for entity_id in entity_ids + ) + if relationship is FactClaimRelationship.DIRECT: + if len(entity_ids) != 1: + return FactConstraintSolution( + terms=terms, + issue=FactConstraintIssue.SUBJECT_REQUIRED, + ) + return FactConstraintSolution( + terms=terms, + subject_entity_id=entity_ids[0], + amount=expected_total, + ) + + unknown = tuple( + entity_id + for entity_id in entity_ids + if known_values.get(entity_id) is None + ) + if len(unknown) != 1: + return FactConstraintSolution( + terms=terms, + issue=FactConstraintIssue.ALLOCATION_REQUIRED, + ) + known_total = sum( + ( + known + for known in known_values.values() + if known is not None + ), + start=Decimal("0"), + ) + amount = expected_total - known_total + if amount < 0: + return FactConstraintSolution( + terms=terms, + issue=FactConstraintIssue.INCONSISTENT_TOTAL, + ) + return FactConstraintSolution( + terms=terms, + subject_entity_id=unknown[0], + amount=amount, + ) + + +class ContextChangeResolver: + """Resolve model-proposed claims using model selection and checked calculations.""" + + _periods_per_year = { + MoneyPeriod.ANNUAL: Decimal("1"), + MoneyPeriod.MONTHLY: Decimal("12"), + MoneyPeriod.FOUR_WEEKLY: Decimal("13"), + MoneyPeriod.WEEKLY: Decimal("52"), + } + + def __init__( + self, + registry: FactDefinitionRegistry, + mapper: VariableMapper, + ) -> None: + self._registry = registry + self._mapper = mapper + + async def resolve( + self, + tool_input: "ResolveContextChangeInput", + candidate_sets: tuple[tuple[PolicyEngineVariableCandidate, ...], ...], + ) -> "ResolveContextChangeOutput": + if len(candidate_sets) != len(tool_input.claims): + raise ValueError("Each unresolved claim requires one candidate set.") + decisions: list[FactResolutionDecision] = [] + operations: list[ContextOperation] = [] + usage = VariableMappingUsage() + for claim, candidates in zip(tool_input.claims, candidate_sets, strict=True): + selection = self._proposal_selection(claim, candidates) + if selection is None: + mapped = await self._mapper.select( + claim=claim, + candidates=candidates, + context=tool_input.context, + registry=self._registry, + validation_issues=tool_input.validation_issues, + ) + selection = mapped.selection + selection_source: Literal["proposal", "resolver_model"] = ( + "resolver_model" + ) + usage = usage.plus(mapped.usage) + else: + selection_source = "proposal" + resolution = self._resolve_claim( + claim=claim, + selection=selection, + candidates=candidates, + context=tool_input.context, + context_proposal=tool_input.proposal, + turn_id=tool_input.turn_id, + evidence=tool_input.evidence, + ) + if isinstance(resolution, SetFactOperation): + decisions.append( + FactResolutionDecision( + claim_id=claim.claim_id, + status="resolved", + candidates=candidates, + selection_source=selection_source, + selection=selection, + operation=resolution, + ) + ) + operations.append(resolution) + else: + decisions.append( + FactResolutionDecision( + claim_id=claim.claim_id, + status=resolution.status, + candidates=candidates, + selection_source=selection_source, + selection=selection, + proposal=resolution, + ) + ) + operations.append( + AddPendingFactResolutionOperation(proposal=resolution) + ) + return ResolveContextChangeOutput( + patch=ContextPatch( + expected_revision=tool_input.context.revision, + operations=tuple(operations), + ), + decisions=tuple(decisions), + usage=usage, + ) + + def _proposal_selection( + self, + claim: FactClaim, + candidates: tuple[PolicyEngineVariableCandidate, ...], + ) -> VariableMappingSelection | None: + """Validate a semantic mapping already selected by the proposal model.""" + + if claim.definition_key is None: + return None + try: + definition = self._registry.get( + claim.definition_key, + claim.definition_version, + ) + except KeyError: + return None + if definition.engine_binding is None: + return None + binding_parts = definition.engine_binding.split(".", 1) + binding_entity = binding_parts[0] if len(binding_parts) == 2 else None + binding_name = binding_parts[-1] + selected = next( + ( + candidate + for candidate in candidates + if candidate.name == binding_name + and binding_entity in {None, candidate.entity} + ), + None, + ) + if selected is None: + return None + return VariableMappingSelection( + status=MappingStatus.MATCHED, + variable_name=selected.name, + confidence=MappingConfidence.HIGH, + target_period=claim.value.period + if isinstance(claim.value, ClaimedMoneyValue) + else None, + ) + + def _resolve_claim( + self, + *, + claim: FactClaim, + selection: VariableMappingSelection, + candidates: tuple[PolicyEngineVariableCandidate, ...], + context: ConversationContext, + context_proposal: ContextChangeProposal | None, + turn_id: str, + evidence: str, + ) -> PendingFactResolution | SetFactOperation: + claim_value = self._claim_value(claim) + scope_id = context.focus.scope_id or next( + scope.scope_id for scope in context.scopes if scope.active + ) + entity_ids = self._resolve_entities(claim.subject_references, context) + prior_resolution = self._supplemented_resolution( + claim, + context=context, + proposal=context_proposal, + ) + if prior_resolution is None: + source_turn_id = turn_id + source_claim = claim + supplements: tuple[FactResolutionSupplement, ...] = () + source_evidence = evidence + created_revision = context.revision + else: + prior, response = prior_resolution + source_turn_id = prior.source_turn_id + source_claim = prior.source_claim or claim + supplements = ( + *prior.supplements, + FactResolutionSupplement( + turn_id=turn_id, + evidence=response.evidence, + updates=response.updates, + ), + ) + source_evidence = prior.evidence + created_revision = prior.created_revision + base = { + "claim_id": claim.claim_id, + "source_turn_id": source_turn_id, + "source_claim": source_claim, + "supplements": supplements, + "scope_id": scope_id, + "referenced_entity_ids": entity_ids, + "evidence": source_evidence, + "relationship": claim.relationship, + "mapping_confidence": selection.confidence.value, + # Preserve the claim's native constraint even when catalogue mapping + # needs clarification. Later explicit facts can then satisfy it without + # depending on another model interpretation of the original message. + "expected_total": claim_value.amount, + "period": claim_value.period, + "created_revision": created_revision, + } + if len(entity_ids) != len(claim.subject_references): + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + "Which exact people or household should the stated value " + "cover?" + ), + } + ) + selected = next( + ( + item + for item in candidates + if item.name == selection.variable_name + ), + None, + ) + if ( + selection.status is not MappingStatus.MATCHED + or selection.confidence is not MappingConfidence.HIGH + or selected is None + ): + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + f"What does the {self._money(claim_value.amount)}" + f"{self._claim_period_suffix(claim_value.period)} amount " + "represent, and which household member or members does it " + "cover?" + ), + } + ) + + value_type = (selected.value_type or "").casefold() + if value_type and not any( + token in value_type for token in ("float", "int", "decimal", "number") + ): + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + f"{self._plain_label(selected)} does not accept a monetary " + "amount. What kind of income or expense did you mean?" + ), + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + } + ) + + variable_period = self._metadata_period(selected.definition_period) + target_period = variable_period or selection.target_period or claim_value.period + if target_period is None: + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + f"What period should I use for the {claim.concept} amount: " + "weekly, every four weeks, monthly, or annual?" + ), + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + } + ) + if not self._entity_matches(selected.entity, entity_ids, context): + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + f"Which person or household should the " + f"{self._plain_label(selected).casefold()} amount apply to?" + ), + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + "period": target_period, + } + ) + definition = self._registry.ensure_engine_definition( + variable_name=selected.name, + entity=selected.entity, + label=selected.label or selected.name.replace("_", " ").title(), + value_kind=FactValueKind.MONEY, + ) + if definition.value_kind is not FactValueKind.MONEY: + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": ( + f"{self._plain_label(selected)} does not accept a monetary " + "amount. What kind of income or expense did you mean?" + ), + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + "period": target_period, + } + ) + expected_total = self._convert( + claim_value.amount, + claim_value.period or target_period, + target_period, + ) + if ( + claim.relationship is FactClaimRelationship.DIRECT + and len(entity_ids) == 1 + and claim_value.period in {None, target_period} + ): + return SetFactOperation( + definition_key=definition.key, + definition_version=definition.version, + subject_reference=entity_ids[0], + scope_id=scope_id, + assertion=PresentAssertion( + value=MoneyFactValue( + amount=expected_total, + period=target_period, + currency=claim_value.currency, + ) + ), + correction=claim.correction, + ) + known_values: dict[str, Decimal | None] = {} + for entity_id in entity_ids: + fact = context.active_fact(definition.key, entity_id, scope_id) + known = self._money_value(fact, target_period) + known_values[entity_id] = known + solution = _solve_constraint( + variable_name=selected.name, + relationship=claim.relationship, + entity_ids=entity_ids, + known_values=known_values, + expected_total=expected_total, + ) + if solution.issue is not None: + prompt = { + FactConstraintIssue.SUBJECT_REQUIRED: ( + "Which one person or household should receive this value?" + ), + FactConstraintIssue.ALLOCATION_REQUIRED: ( + f"How should the {self._money(expected_total)} " + f"{self._period_label(target_period)} total be divided between " + "the household members?" + ), + FactConstraintIssue.INCONSISTENT_TOTAL: ( + f"The known {selected.label or selected.name} amounts exceed " + f"the stated total of {self._money(expected_total)}. Which " + "amounts should I use?" + ), + }[solution.issue] + return self._clarification( + base, + selected, + target_period, + expected_total, + solution.terms, + prompt, + ) + if solution.subject_entity_id is None or solution.amount is None: + raise RuntimeError("A solved fact constraint lacks its exact assignment.") + + subject_id = solution.subject_entity_id + assignment_amount = solution.amount + assignment = FactResolutionAssignment( + definition_key=definition.key, + subject_entity_id=subject_id, + scope_id=scope_id, + assertion=PresentAssertion( + value=MoneyFactValue( + amount=assignment_amount, + period=target_period, + currency=claim_value.currency, + ) + ), + correction=claim.correction, + ) + subject = self._entity_label(subject_id, context) + prompt = ( + f"Using the amounts already provided, the " + f"{self._money(expected_total)} {self._period_label(target_period)} " + f"{self._plain_label(selected).casefold()} total implies " + f"{self._money(assignment_amount)} {self._period_label(target_period)} " + f"for {subject}. " + "Is that the correct breakdown?" + ) + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.AWAITING_CONFIRMATION, + "prompt": prompt, + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + "expected_total": expected_total, + "period": target_period, + "terms": solution.terms, + "assignments": (assignment,), + } + ) + + @staticmethod + def _supplemented_resolution( + claim: FactClaim, + *, + context: ConversationContext, + proposal: ContextChangeProposal | None, + ) -> tuple[PendingFactResolution, PendingFactResolutionResponse] | None: + if proposal is None: + return None + responses = { + response.proposal_id: response + for response in proposal.proposal_responses + if response.action is PendingResolutionAction.SUPPLY + } + return next( + ( + (pending, responses[pending.proposal_id]) + for pending in context.pending_fact_resolutions + if pending.claim_id == claim.claim_id + and pending.proposal_id in responses + ), + None, + ) + + @staticmethod + def _clarification( + base: dict[str, object], + selected: PolicyEngineVariableCandidate, + period: MoneyPeriod, + expected_total: Decimal, + terms: tuple[FactResolutionTerm, ...], + prompt: str, + ) -> PendingFactResolution: + return PendingFactResolution.model_validate( + { + **base, + "status": FactResolutionStatus.NEEDS_CLARIFICATION, + "prompt": prompt, + "variable_name": selected.name, + "variable_entity": selected.entity, + "variable_label": selected.label, + "definition_period": selected.definition_period, + "expected_total": expected_total, + "period": period, + "terms": terms, + } + ) + + @staticmethod + def _resolve_entities( + references: tuple[str, ...], + context: ConversationContext, + ) -> tuple[str, ...]: + aliases: dict[str, str] = {} + for entity in context.entities: + aliases[entity.entity_id.casefold()] = entity.entity_id + if entity.relationship_to_user: + aliases[entity.relationship_to_user.casefold()] = entity.entity_id + for alias in entity.aliases: + aliases[alias.casefold()] = entity.entity_id + resolved = tuple( + dict.fromkeys( + aliases[reference.casefold()] + for reference in references + if reference.casefold() in aliases + ) + ) + return resolved + + @staticmethod + def _entity_matches( + variable_entity: str, + entity_ids: tuple[str, ...], + context: ConversationContext, + ) -> bool: + expected = { + "person": EntityKind.PERSON, + "household": EntityKind.HOUSEHOLD, + "benunit": EntityKind.HOUSEHOLD, + }.get(variable_entity) + if expected is None or not entity_ids: + return False + return all( + next(item for item in context.entities if item.entity_id == entity_id).kind + is expected + for entity_id in entity_ids + ) + + @classmethod + def _convert( + cls, + amount: Decimal, + source: MoneyPeriod, + target: MoneyPeriod, + ) -> Decimal: + annual = amount * cls._periods_per_year[source] + return annual / cls._periods_per_year[target] + + @classmethod + def _money_value( + cls, + fact: ContextFact | None, + period: MoneyPeriod, + ) -> Decimal | None: + if fact is None or not isinstance(fact.assertion, PresentAssertion): + return None + value = fact.assertion.value + if not isinstance(value, MoneyFactValue): + return None + return cls._convert(value.amount, value.period, period) + + @staticmethod + def _metadata_period(value: str | None) -> MoneyPeriod | None: + normalized = (value or "").casefold() + return { + "year": MoneyPeriod.ANNUAL, + "annual": MoneyPeriod.ANNUAL, + "month": MoneyPeriod.MONTHLY, + "monthly": MoneyPeriod.MONTHLY, + "week": MoneyPeriod.WEEKLY, + "weekly": MoneyPeriod.WEEKLY, + }.get(normalized) + + @staticmethod + def _claim_value(claim: FactClaim) -> ClaimedMoneyValue: + if not isinstance(claim.value, ClaimedMoneyValue): + raise ValueError("fact resolution requires a monetary fact claim") + return claim.value + + @staticmethod + def _plain_label(candidate: PolicyEngineVariableCandidate) -> str: + return ( + candidate.label + or candidate.name.replace("_", " ").replace("-", " ") + ).strip() + + @staticmethod + def _claim_period_suffix(period: MoneyPeriod | None) -> str: + if period is None: + return "" + return " " + { + MoneyPeriod.ANNUAL: "annual", + MoneyPeriod.MONTHLY: "monthly", + MoneyPeriod.FOUR_WEEKLY: "four-weekly", + MoneyPeriod.WEEKLY: "weekly", + }[period] + + @staticmethod + def _entity_label(entity_id: str, context: ConversationContext) -> str: + entity = next(item for item in context.entities if item.entity_id == entity_id) + if entity.relationship_to_user == "self": + return "you" + if entity.relationship_to_user: + return f"your {entity.relationship_to_user.replace('_', ' ')}" + return "the other household member" + + @staticmethod + def _money(value: Decimal) -> str: + return f"£{value:,.2f}".replace(".00", "") + + @staticmethod + def _period_label(period: MoneyPeriod) -> str: + return { + MoneyPeriod.ANNUAL: "per year", + MoneyPeriod.MONTHLY: "per month", + MoneyPeriod.FOUR_WEEKLY: "every four weeks", + MoneyPeriod.WEEKLY: "per week", + }[period] + + +class ResolveContextChangeTool( + Tool["ResolveContextChangeInput", "ResolveContextChangeOutput"] +): + spec = ToolSpec( + identifier="resolve_context_change", + version="1", + description=( + "Map monetary fact claims to exact PolicyEngine variables and propose " + "only validated, confirmable assignments." + ), + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=ResolveContextChangeInput, + output_model=ResolveContextChangeOutput, + tool_dependencies=("search_variables",), + ) + + def __init__(self, resolver: ContextChangeResolver) -> None: + self._resolver = resolver + + async def run( + self, + tool_input: "ResolveContextChangeInput", + context: ToolCallContext, + ) -> "ResolveContextChangeOutput": + candidate_sets: list[tuple[PolicyEngineVariableCandidate, ...]] = [] + for claim in tool_input.claims: + raw = await context.invoke_tool( + "search_variables", + {"query": claim.concept, "limit": 12}, + ) + rows = raw.root.get("variables") if isinstance(raw, SafeToolOutput) else None + if not isinstance(rows, list): + rows = [] + candidates: list[PolicyEngineVariableCandidate] = [] + for row in rows or []: + if not isinstance(row, dict): + continue + try: + candidates.append( + PolicyEngineVariableCandidate.model_validate( + { + "name": row.get("name"), + "label": row.get("label"), + "entity": row.get("entity"), + "description": row.get("description"), + "definition_period": row.get("definition_period"), + "value_type": row.get("value_type"), + } + ) + ) + except ValidationError: + continue + candidate_sets.append(tuple(candidates)) + result = await self._resolver.resolve(tool_input, tuple(candidate_sets)) + context.record_model_usage(**result.usage.model_dump()) + return result diff --git a/backend/conversations/__init__.py b/backend/conversations/__init__.py index b2c46947..b14488e4 100644 --- a/backend/conversations/__init__.py +++ b/backend/conversations/__init__.py @@ -1,9 +1 @@ -"""Stored chat history: persistence, sharing, and reporting. - -Re-exports the router (for app wiring) and ensure_table (called at startup). -""" - -from conversations.models import ensure_table -from conversations.routes import router - -__all__ = ["router", "ensure_table"] +"""Stored chat history package.""" diff --git a/backend/conversations/models.py b/backend/conversations/models.py index 87a76a6c..781421e0 100644 --- a/backend/conversations/models.py +++ b/backend/conversations/models.py @@ -1,24 +1,31 @@ -"""The conversations table model, engine, and schema bootstrap.""" +"""The SQLModel conversation table and runtime database engine.""" -import logging import os from datetime import datetime from typing import Optional +from sqlalchemy import Index from sqlmodel import Field, SQLModel, create_engine -logger = logging.getLogger(__name__) - class ChatConversation(SQLModel, table=True): __tablename__ = "chat_conversations" + __table_args__ = ( + Index( + "idx_chat_conversations_session_id_unique", + "session_id", + unique=True, + ), + Index("idx_chat_conversations_share_token", "share_token"), + ) + id: Optional[int] = Field(default=None, primary_key=True) - session_id: str = Field(index=True, unique=True) + session_id: str title: str messages: str # JSON string user_id: Optional[str] = None user_email: Optional[str] = None - share_token: Optional[str] = Field(default=None, index=True) + share_token: Optional[str] = None created_at: datetime updated_at: datetime @@ -34,28 +41,3 @@ def get_engine(): raise RuntimeError("DATABASE_URL not set") _engine = create_engine(url) return _engine - - -def ensure_table(): - try: - engine = get_engine() - SQLModel.metadata.create_all(engine) - # Add columns that may not exist yet on older databases - from sqlalchemy import text - with engine.connect() as conn: - for col, col_type in [("share_token", "TEXT"), ("user_email", "TEXT")]: - try: - conn.execute(text(f"ALTER TABLE chat_conversations ADD COLUMN {col} {col_type}")) - conn.commit() - logger.info(f"Added column {col} to chat_conversations") - except Exception: - conn.rollback() # Column already exists - try: - conn.execute(text("CREATE INDEX IF NOT EXISTS idx_chat_conversations_share_token ON chat_conversations (share_token)")) - conn.commit() - except Exception: - conn.rollback() - logger.info("Conversations table ensured successfully") - except Exception as e: - logger.error(f"Could not ensure conversations table: {e}") - import traceback; logger.error(traceback.format_exc()) diff --git a/backend/conversations/store.py b/backend/conversations/store.py index b43ba4ac..539585e6 100644 --- a/backend/conversations/store.py +++ b/backend/conversations/store.py @@ -175,5 +175,8 @@ def delete_conversation(conversation_id: int): row = session.get(ChatConversation, conversation_id) if not row: raise HTTPException(status_code=404, detail="Conversation not found") + from persistence.deletion import delete_capability_records + + delete_capability_records(session, row.session_id) session.delete(row) session.commit() diff --git a/backend/engine/discovery.py b/backend/engine/discovery.py index 0868dcee..54ab5fda 100644 --- a/backend/engine/discovery.py +++ b/backend/engine/discovery.py @@ -20,6 +20,28 @@ def _matches(query: str, *values: str | None) -> bool: return bool(get_close_matches(q, [v.lower() for v in values if v], n=1, cutoff=0.65)) +def _normalized_search_text(value: str | None) -> str: + return " ".join((value or "").casefold().replace("_", " ").split()) + + +def _variable_search_rank(query: str, item: dict[str, Any]) -> tuple[int, str]: + normalized_query = _normalized_search_text(query) + name = _normalized_search_text(item.get("name")) + label = _normalized_search_text(item.get("label")) + description = _normalized_search_text(item.get("description")) + if normalized_query in {name, label}: + rank = 0 + elif name.startswith(normalized_query) or label.startswith(normalized_query): + rank = 1 + elif normalized_query in name or normalized_query in label: + rank = 2 + elif normalized_query in description: + rank = 3 + else: + rank = 4 + return rank, name + + def _default_output_entities(model: Any, name: str) -> list[str]: return [ entity @@ -36,6 +58,16 @@ def _variable_item(name: str, variable: Any, model: Any) -> dict[str, Any]: "entity": getattr(variable, "entity", None), "description": getattr(variable, "description", None), "definition_period": getattr(variable, "definition_period", None), + "unit": getattr(variable, "unit", None), + "quantity_type": getattr(variable, "quantity_type", None), + "reference": json_safe(getattr(variable, "reference", None)), + "defined_for": getattr(variable, "defined_for", None), + "min_value": json_safe(getattr(variable, "min_value", None)), + "max_value": json_safe(getattr(variable, "max_value", None)), + "is_period_size_independent": getattr( + variable, "is_period_size_independent", None + ), + "metadata": json_safe(getattr(variable, "metadata", None)), "value_type": getattr(getattr(variable, "value_type", None), "__name__", None) or str(getattr(variable, "value_type", "")), "default_value": json_safe(getattr(variable, "default_value", None)), @@ -71,13 +103,13 @@ def search_variables( if not _matches(query, name, item.get("label"), item.get("description")): continue rows.append(item) - if len(rows) >= limit: - break + if query: + rows.sort(key=lambda item: _variable_search_rank(query, item)) return { "status": "success", "query": query, "entity": entity, - "variables": rows, + "variables": rows[:limit], } diff --git a/backend/eval/deployed_runner.py b/backend/eval/deployed_runner.py index 82cb63d0..a924d049 100644 --- a/backend/eval/deployed_runner.py +++ b/backend/eval/deployed_runner.py @@ -61,57 +61,7 @@ def _failed_trial( ) -def grade_gateway_expectation( - case: ToolLoopCase, - response: EvalChatResponse, -) -> list[str]: - """Grade routing and reform authorization before tool/answer quality.""" - - expectation = case.gateway_expect - if expectation is None: - return [] - errors: list[str] = [] - if response.route != expectation.route: - errors.append( - f"gateway route was {response.route!r}, expected {expectation.route!r}" - ) - if response.outcome != expectation.outcome: - errors.append( - f"gateway outcome was {response.outcome!r}, expected {expectation.outcome!r}" - ) - trace = response.gateway_trace - if trace is None: - errors.append("gateway trace was missing") - return errors - for name, expected in expectation.defaults_contains.items(): - if name not in trace.defaults_applied: - errors.append( - f"gateway default {name!r} was missing; expected {expected!r}" - ) - elif trace.defaults_applied[name] != expected: - errors.append( - f"gateway default {name!r} was {trace.defaults_applied[name]!r}, " - f"expected {expected!r}" - ) - minimum = expectation.min_reform_confidence - if minimum is not None: - if trace.reform_confidence is None: - errors.append( - "gateway reform confidence was missing; " - f"expected at least {minimum}" - ) - elif trace.reform_confidence < minimum: - errors.append( - f"gateway reform confidence was {trace.reform_confidence}, " - f"expected at least {minimum}" - ) - if expectation.require_parameter_binding and not trace.parameter_bindings: - errors.append("gateway produced no validated parameter binding") - return errors - - def _grade_response(case: ToolLoopCase, response: EvalChatResponse) -> CaseResult: - gateway_errors = grade_gateway_expectation(case, response) response_details = { "session_id": response.session_id, "model": response.model, @@ -119,11 +69,9 @@ def _grade_response(case: ToolLoopCase, response: EvalChatResponse) -> CaseResul "outcome": response.outcome, "stop_reason": response.stop_reason, "usage": response.usage.model_dump(), - "gateway_trace": ( - response.gateway_trace.model_dump() - if response.gateway_trace is not None - else None - ), + "invocation_trace": [ + trace.model_dump() for trace in response.invocation_trace + ], } if response.status != "completed": failed = _failed_trial( @@ -132,23 +80,23 @@ def _grade_response(case: ToolLoopCase, response: EvalChatResponse) -> CaseResul details={ **response_details, "text": response.content, - "tool_trace": [trace.model_dump() for trace in response.tool_trace], + "invocation_trace": [ + trace.model_dump() for trace in response.invocation_trace + ], }, ) - return failed.model_copy( - update={"errors": [*failed.errors, *gateway_errors]} - ) + return failed tool_calls = [ - ModelToolCall(id=trace.tool_id, name=trace.name, input=trace.input) - for trace in response.tool_trace + ModelToolCall(id=trace.invocation_id, name=trace.name, input=trace.input) + for trace in response.invocation_trace ] result = grade_tool_loop_case( case, text=response.content, tool_calls=tool_calls, - tool_outputs=[trace.output for trace in response.tool_trace], - errors=gateway_errors, + tool_outputs=[trace.output for trace in response.invocation_trace], + errors=[], ) return result.model_copy( update={"details": {**result.details, "deployed": response_details}} diff --git a/backend/eval/loaders.py b/backend/eval/loaders.py index 2816329a..032d8bfe 100644 --- a/backend/eval/loaders.py +++ b/backend/eval/loaders.py @@ -8,7 +8,6 @@ from eval.schemas import ( AnswerCase, EvalCase, - GatewayCase, ToolContractCase, ToolLoopCase, TrajectoryCase, @@ -20,7 +19,6 @@ "trajectory": TrajectoryCase, "answer": AnswerCase, "tool_loop": ToolLoopCase, - "gateway": GatewayCase, } diff --git a/backend/eval/runner.py b/backend/eval/runner.py index 5c394c65..4b19e28b 100644 --- a/backend/eval/runner.py +++ b/backend/eval/runner.py @@ -8,9 +8,18 @@ from pathlib import Path from typing import Any, Dict, Iterable, List -from prompts import CHARTS_MODE_DIRECTIVE, SYSTEM_PROMPT +from capabilities.chart import SocietyChartCapability +from capabilities.contracts import Completed, Failed, NeedsInput, Unsupported +from capabilities.follow_up import AnalysisFollowUpCapability +from capabilities.household import HouseholdAnalysisCapability +from capabilities.policy_information import PolicyInformationCapability +from capabilities.policy_reform import PolicyReformCapability +from capabilities.society import SocietyAnalysisCapability +from chat.capability_service import ( + MANDATORY_CAPABILITY_CONTRACT, + capability_result_for_model, +) from tools.context import new_tool_context -from tools.definitions import TOOL_DEFINITIONS from tools.dispatch import execute_tool from eval.graders import grade_output, grade_text, grade_tool_calls @@ -22,7 +31,6 @@ CaseResult, EvalCase, EvalReport, - GatewayCase, ModelTurn, ToolContractCase, ToolLoopCase, @@ -43,9 +51,13 @@ "trajectory": CASE_ROOT / "trajectory", "answer": CASE_ROOT / "answer", "tool_loop": CASE_ROOT / "tool_loop", - "gateway": CASE_ROOT / "gateway", } +CAPABILITY_CHARTS_MODE_DIRECTIVE = ( + "The user enabled chart presentation. Use society_chart when a chart is " + "requested and its typed prerequisites can be satisfied." +) + def _utc_now() -> str: return datetime.now(timezone.utc).isoformat(timespec="seconds") @@ -128,17 +140,36 @@ def _build_offline_client(cases: List[EvalCase]) -> FakeModelClient: return FakeModelClient(turns) -def _tool_specs_for_model() -> List[Dict[str, Any]]: +def _public_capabilities(): + return ( + PolicyInformationCapability(), + PolicyReformCapability(), + HouseholdAnalysisCapability(), + SocietyAnalysisCapability(), + AnalysisFollowUpCapability(), + SocietyChartCapability(), + ) + + +def _capability_specs_for_model() -> List[Dict[str, Any]]: return [ { - "name": tool["name"], - "description": tool["description"], - "input_schema": tool["input_schema"], + "name": capability.spec.identifier, + "description": ( + f"{capability.spec.description} Required-use rule: " + f"{capability.spec.required_use}" + ), + "input_schema": capability.spec.input_model.model_json_schema(), } - for tool in TOOL_DEFINITIONS + for capability in _public_capabilities() ] +def _operations_for_case(case: TrajectoryCase | ToolLoopCase) -> List[Dict[str, Any]]: + del case + return _capability_specs_for_model() + + def _messages_for_case(case: TrajectoryCase | ToolLoopCase) -> List[Dict[str, Any]]: if case.messages: return case.messages @@ -146,9 +177,12 @@ def _messages_for_case(case: TrajectoryCase | ToolLoopCase) -> List[Dict[str, An def _system_for_case(case: TrajectoryCase | ToolLoopCase) -> str: - sections = [SYSTEM_PROMPT] + sections = [ + "You are PolicyEngine UK Chat. Continue naturally using the supplied " + "conversation history.\n\n" + MANDATORY_CAPABILITY_CONTRACT + ] if case.charts_mode: - sections.append(CHARTS_MODE_DIRECTIVE) + sections.append(CAPABILITY_CHARTS_MODE_DIRECTIVE) return "\n\n".join(sections) @@ -177,7 +211,7 @@ def _run_trajectory(case: TrajectoryCase, client: ModelClient) -> CaseResult: case_id=case.id, messages=_messages_for_case(case), system=_system_for_case(case), - tools=_tool_specs_for_model(), + tools=_operations_for_case(case), ) except Exception as exc: return _result(case, "failed", 0.0, [f"{type(exc).__name__}: {exc}"]) @@ -197,12 +231,14 @@ def _tool_result_text(case: AnswerCase) -> str: output = call.output if call.output_fixture: output = _load_fixture(call.output_fixture) + _validate_capability_output(call.name, output or {}) + model_output = capability_result_for_model(output or {}) chunks.append( "\n".join( [ f"{index}. {call.name}", f"Input: {json.dumps(call.input, sort_keys=True)}", - f"Output: {json.dumps(output or {}, sort_keys=True)}", + f"Output: {json.dumps(model_output, sort_keys=True)}", ] ) ) @@ -218,7 +254,11 @@ def _run_answer(case: AnswerCase, client: ModelClient) -> CaseResult: {"role": "user", "content": case.prompt}, {"role": "user", "content": _tool_result_text(case)}, ], - system=SYSTEM_PROMPT, + system=( + "You are PolicyEngine UK Chat. Use validated capability output as " + "authoritative facts while writing natural prose.\n\n" + + MANDATORY_CAPABILITY_CONTRACT + ), tools=None, ) except Exception as exc: @@ -237,13 +277,43 @@ def _tool_use_id(case: ToolLoopCase, iteration: int, index: int, call_id: str) - return call_id or f"{case.id}-{iteration}-{index}" +def _load_frozen_output(call) -> Dict[str, Any]: + if call.output_fixture: + return _load_fixture(call.output_fixture) + return call.output or {} + + +def _validate_capability_output(name: str, output: Dict[str, Any]) -> None: + capability = next( + ( + item + for item in _public_capabilities() + if item.spec.identifier == name + ), + None, + ) + if capability is None: + raise ValueError(f"Unknown capability output fixture: {name}") + status = output.get("status") + if status == "completed": + Completed[capability.spec.output_model].model_validate(output) + elif status == "needs_input": + NeedsInput.model_validate(output) + elif status == "unsupported": + Unsupported.model_validate(output) + elif status == "failed": + Failed.model_validate(output) + else: + raise ValueError(f"Invalid capability outcome status for {name}: {status!r}") + + def _run_tool_loop(case: ToolLoopCase, client: ModelClient) -> CaseResult: messages: List[Dict[str, Any]] = _messages_for_case(case) - tool_context = new_tool_context(turn_id=case.id) tool_calls = [] tool_outputs: List[Dict[str, Any]] = [] final_text = "" errors: List[str] = [] + frozen_capability_outputs = list(case.capability_outputs) for iteration in range(1, case.max_iterations + 1): try: @@ -251,7 +321,7 @@ def _run_tool_loop(case: ToolLoopCase, client: ModelClient) -> CaseResult: case_id=case.id, messages=messages, system=_system_for_case(case), - tools=_tool_specs_for_model(), + tools=_operations_for_case(case), ) except Exception as exc: return _result(case, "failed", 0.0, [f"{type(exc).__name__}: {exc}"]) @@ -278,7 +348,17 @@ def _run_tool_loop(case: ToolLoopCase, client: ModelClient) -> CaseResult: } ) try: - output = execute_tool(call.name, call.input, context=tool_context) + if not frozen_capability_outputs: + raise ValueError( + f"No frozen capability output remains for {call.name}." + ) + frozen = frozen_capability_outputs.pop(0) + if frozen.name != call.name: + raise ValueError( + f"Expected frozen output for {frozen.name}, got {call.name}." + ) + output = _load_frozen_output(frozen) + _validate_capability_output(call.name, output) except Exception as exc: return _result( case, @@ -288,11 +368,16 @@ def _run_tool_loop(case: ToolLoopCase, client: ModelClient) -> CaseResult: {"text": final_text, "tool_calls": [call.model_dump() for call in tool_calls]}, ) tool_outputs.append(output) + model_output = capability_result_for_model(output) tool_results.append( { "type": "tool_result", "tool_use_id": tool_use_id, - "content": json.dumps(output, ensure_ascii=False, default=str), + "content": json.dumps( + model_output, + ensure_ascii=False, + default=str, + ), } ) @@ -315,55 +400,6 @@ def _run_tool_loop_trials(case: ToolLoopCase, client: ModelClient) -> CaseResult return aggregate_tool_loop_trials(case, trial_results) -def _run_gateway(case: GatewayCase) -> CaseResult: - """Live-only: run the gateway pre-pass and grade the verdict. Outcome is the - primary assertion; tool/forbidden_tool and per-slot expectations are - secondary (graded only when the case declares them). Binary 0/1 score to - match the other suites, with the full plan stashed in details for tuning.""" - from gateway import run_gateway - - try: - verdict = run_gateway(case.prompt) - except Exception as exc: - return _result(case, "failed", 0.0, [f"{type(exc).__name__}: {exc}"]) - - errors: List[str] = [] - if verdict.outcome != case.expected_outcome: - errors.append(f"expected outcome {case.expected_outcome!r}, got {verdict.outcome!r}") - if case.expected_tool and verdict.tool != case.expected_tool: - errors.append(f"expected tool {case.expected_tool!r}, got {verdict.tool!r}") - if case.forbidden_tool and verdict.tool == case.forbidden_tool: - errors.append(f"forbidden tool {case.forbidden_tool!r} was selected") - if case.expected_gating_slots: - got = set(verdict.gating_slots) - want = set(case.expected_gating_slots) - if got != want: - errors.append(f"gating slots: expected {sorted(want)}, got {sorted(got)}") - - by_name = {s.name: s for s in verdict.slots} - for exp in case.expected_slots: - got_slot = by_name.get(exp.slot) - if got_slot is None: - errors.append(f"slot {exp.slot!r} missing from plan") - continue - if exp.source is not None and got_slot.source != exp.source: - errors.append(f"slot {exp.slot!r} source: expected {exp.source!r}, got {got_slot.source!r}") - if exp.gates is not None and (exp.slot in verdict.gating_slots) != exp.gates: - errors.append(f"slot {exp.slot!r} gates: expected {exp.gates}, got {exp.slot in verdict.gating_slots}") - - details = { - "outcome": verdict.outcome, - "tool": verdict.tool, - "gating_slots": verdict.gating_slots, - "unmodellable_outputs": verdict.unmodellable_outputs, - "slots": [ - {"name": s.name, "kind": s.kind, "source": s.source, "value": s.value} - for s in verdict.slots - ], - } - return _result(case, "failed" if errors else "passed", 0.0 if errors else 1.0, errors, details) - - def run_eval( *, suites: List[str] | None = None, @@ -412,9 +448,6 @@ def run_eval( results.append(_run_tool_loop_trials(case, client)) else: results.append(_run_tool_loop(case, client)) - elif isinstance(case, GatewayCase): - results.append(_run_gateway(case)) - report = EvalReport( mode=mode, suites=selected_suites, diff --git a/backend/eval/schemas.py b/backend/eval/schemas.py index 90bc1152..e159909e 100644 --- a/backend/eval/schemas.py +++ b/backend/eval/schemas.py @@ -114,20 +114,6 @@ class AnswerCase(CaseBase): offline_response: Optional[ModelTurn] = None -class GatewayTraceExpectation(StrictModel): - route: Literal["compute", "lightweight"] - outcome: Literal[ - "irrelevant", - "out_of_scope", - "partial", - "needs_plan", - "ready", - ] - defaults_contains: Dict[str, Any] = Field(default_factory=dict) - min_reform_confidence: Optional[int] = Field(default=None, ge=0, le=100) - require_parameter_binding: bool = False - - class ToolLoopCase(CaseBase): suite: Literal["tool_loop"] = "tool_loop" prompt: str @@ -140,28 +126,8 @@ class ToolLoopCase(CaseBase): trials: int = Field(default=1, ge=1, le=10) pass_threshold: float = Field(default=1.0, ge=0.0, le=1.0) offline_responses: List[ModelTurn] = Field(default_factory=list) - gateway_expect: Optional[GatewayTraceExpectation] = None - - -class SlotExpectation(StrictModel): - slot: str - source: Optional[Literal["prompt", "default", "assumed", "runtime"]] = None - gates: Optional[bool] = None # whether this slot should trigger a question - - -class GatewayCase(CaseBase): - suite: Literal["gateway"] = "gateway" - prompt: str - expected_outcome: Literal[ - "irrelevant", "out_of_scope", "partial", "needs_plan", "ready" - ] - expected_tool: Optional[str] = None - forbidden_tool: Optional[str] = None - expected_gating_slots: List[str] = Field(default_factory=list) - expected_slots: List[SlotExpectation] = Field(default_factory=list) - - -EvalCase = ToolContractCase | TrajectoryCase | AnswerCase | ToolLoopCase | GatewayCase + capability_outputs: List[FrozenToolCall] = Field(default_factory=list) +EvalCase = ToolContractCase | TrajectoryCase | AnswerCase | ToolLoopCase class CaseResult(StrictModel): @@ -203,65 +169,29 @@ class EvalUsage(StrictModel): cache_read_input_tokens: int = 0 -class EvalToolTrace(StrictModel): - tool_id: str +class EvalInvocationTrace(StrictModel): + invocation_id: str + kind: Literal["capability", "tool"] name: str input: Dict[str, Any] = Field(default_factory=dict) - status: Literal["pending", "success", "error"] = "pending" + status: Literal[ + "running", + "completed", + "needs_input", + "unsupported", + "failed", + "cancelled", + ] = "running" output: Any = None -class EvalGatewaySlot(StrictModel): - name: str - kind: str - source: str - value: Optional[str] = None - - -class EvalGatewayReason(StrictModel): - code: str - slot: str - options: List[str] = Field(default_factory=list) - evidence: Optional[str] = None - - -class EvalGatewayBinding(StrictModel): - parameter_path: str - label: str - catalogue_evidence: str - - -class EvalGatewayAlternative(StrictModel): - summary: str - parameter_bindings: List[EvalGatewayBinding] = Field(default_factory=list) - reform: Dict[str, Any] = Field(default_factory=dict) - - -class EvalGatewayTrace(StrictModel): - selected_tool: Optional[str] = None - target_tool: Optional[str] = None - slots: List[EvalGatewaySlot] = Field(default_factory=list) - gating_reasons: List[EvalGatewayReason] = Field(default_factory=list) - defaults_applied: Dict[str, Any] = Field(default_factory=dict) - reform_confidence: Optional[int] = Field(default=None, ge=0, le=100) - reform_summary: Optional[str] = None - reform_search_queries: List[str] = Field(default_factory=list) - catalogue_version: Optional[str] = None - resolver_model: Optional[str] = None - parameter_bindings: List[EvalGatewayBinding] = Field(default_factory=list) - alternatives: List[EvalGatewayAlternative] = Field(default_factory=list) - catalogue_recovery_used: bool = False - proposal_resumed: bool = False - - class EvalChatResponse(StrictModel): status: Literal["completed", "failed"] content: str = "" session_id: str model: Optional[str] = None - route: str = "compute" + route: str = "capability" outcome: Optional[str] = None stop_reason: Optional[str] = None usage: EvalUsage = Field(default_factory=EvalUsage) - tool_trace: List[EvalToolTrace] = Field(default_factory=list) - gateway_trace: Optional[EvalGatewayTrace] = None + invocation_trace: List[EvalInvocationTrace] = Field(default_factory=list) diff --git a/backend/eval/service.py b/backend/eval/service.py index d8478909..85d8712c 100644 --- a/backend/eval/service.py +++ b/backend/eval/service.py @@ -1,22 +1,21 @@ -"""Structured evaluation adapter over the shared UK Chat turn engine.""" +"""Structured evaluation adapter over the capability chat runtime.""" -from dataclasses import asdict +from dataclasses import replace from typing import Any from pydantic import TypeAdapter from chat.events import ( CancellationProbe, - ToolCompleted, - ToolUsed, + InvocationActivity, TurnCancelled, TurnCompleted, TurnFailed, ) -from chat.orchestrator import run_chat_turn +from chat.capability_runtime import run_capability_chat_turn from chat.schemas import ChatRequest from chat.turn_input import prepare_turn_input -from eval.schemas import EvalChatResponse, EvalGatewayTrace, EvalToolTrace, EvalUsage +from eval.schemas import EvalChatResponse, EvalInvocationTrace, EvalUsage _ANY_ADAPTER = TypeAdapter(Any) @@ -30,12 +29,6 @@ def _json_safe(value: Any) -> Any: return _ANY_ADAPTER.dump_python(value, mode="json") -def _gateway_trace(value) -> EvalGatewayTrace | None: - if value is None: - return None - return EvalGatewayTrace.model_validate(_json_safe(asdict(value))) - - async def run_eval_chat( chat_request: ChatRequest, *, @@ -43,33 +36,27 @@ async def run_eval_chat( ) -> EvalChatResponse: """Run one deployed chat turn and retain its complete structured trace.""" - turn = prepare_turn_input(chat_request) - trace: list[EvalToolTrace] = [] + turn = replace(prepare_turn_input(chat_request), debug=True) + trace: list[EvalInvocationTrace] = [] trace_indexes: dict[str, int] = {} - stream = run_chat_turn(turn, is_cancelled=is_cancelled) + stream = run_capability_chat_turn(turn, is_cancelled=is_cancelled) try: async for event in stream: - if isinstance(event, ToolUsed): - trace_indexes[event.tool_id] = len(trace) - trace.append( - EvalToolTrace( - tool_id=event.tool_id, - name=event.tool_name, - input=_json_safe(event.tool_input), - ) - ) - elif isinstance(event, ToolCompleted): - index = trace_indexes.get(event.tool_id) - completed = EvalToolTrace( - tool_id=event.tool_id, - name=event.tool_name, - input=trace[index].input if index is not None else {}, - status=event.status, - output=_json_safe(event.output), + if isinstance(event, InvocationActivity): + record = event.record + index = trace_indexes.get(record.invocation_id) + input_value = _json_safe(record.debug_input) + completed = EvalInvocationTrace( + invocation_id=record.invocation_id, + kind=record.kind.value, + name=record.identifier, + input=input_value if isinstance(input_value, dict) else {}, + status=record.status.value, + output=_json_safe(record.debug_output), ) if index is None: - trace_indexes[event.tool_id] = len(trace) + trace_indexes[record.invocation_id] = len(trace) trace.append(completed) else: trace[index] = completed @@ -83,8 +70,7 @@ async def run_eval_chat( outcome=event.outcome, stop_reason=event.stop_reason, usage=_usage(event.usage), - tool_trace=trace, - gateway_trace=_gateway_trace(event.gateway_trace), + invocation_trace=trace, ) elif isinstance(event, TurnFailed): return EvalChatResponse( @@ -93,8 +79,7 @@ async def run_eval_chat( session_id=event.session_id, stop_reason=event.stop_reason, usage=_usage(event.usage), - tool_trace=trace, - gateway_trace=_gateway_trace(event.gateway_trace), + invocation_trace=trace, ) elif isinstance(event, TurnCancelled): return EvalChatResponse( @@ -104,8 +89,7 @@ async def run_eval_chat( route=event.route, stop_reason="client_disconnected", usage=_usage(event.usage), - tool_trace=trace, - gateway_trace=_gateway_trace(event.gateway_trace), + invocation_trace=trace, ) finally: await stream.aclose() @@ -115,5 +99,5 @@ async def run_eval_chat( content="Chat turn ended without a terminal event.", session_id=turn.session_id, stop_reason="missing_terminal_event", - tool_trace=trace, + invocation_trace=trace, ) diff --git a/backend/gateway/__init__.py b/backend/gateway/__init__.py deleted file mode 100644 index 8b1fa979..00000000 --- a/backend/gateway/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -"""The gateway: a cheap pre-pass that routes the opening turn. - -runtime.py builds the grounded plan and runs the forced-tool classifier; -policy.py is the pure, offline-testable gate (criticality + outcome). This -package re-exports the runtime surface so callers do `from gateway import X`. -""" - -from gateway.runtime import ( - GatewayVerdict, - gateway_writer_directive, - run_gateway, - serialise_plan_for_system, -) - -__all__ = [ - "GatewayVerdict", - "run_gateway", - "gateway_writer_directive", - "serialise_plan_for_system", -] diff --git a/backend/gateway/assessment.py b/backend/gateway/assessment.py deleted file mode 100644 index ab4b759e..00000000 --- a/backend/gateway/assessment.py +++ /dev/null @@ -1,453 +0,0 @@ -"""Bounded model-assisted construction of exact PolicyEngine reforms.""" - -from __future__ import annotations - -import json -import os -from dataclasses import dataclass -from importlib.metadata import PackageNotFoundError, version -from typing import Any, Callable - -from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_sync_client -from engine.discovery import get_parameter -from engine.reforms import search_reform_targets, validate_reform_dict -from gateway.intent import ReformIntent -from tools.definitions import DEFAULT_SIMULATION_YEAR, REFORM_SCHEMA - -AUTO_EXECUTE_REFORM_CONFIDENCE = 80 -MAX_REFORM_SEARCHES = 4 -REFORM_SEARCH_LIMIT = 20 -MAX_RESOLVER_ITERATIONS = 8 -MAX_ASSESSMENT_REPAIRS = 1 -REFORM_RESOLVER_MODEL = os.environ.get( - "POLICYENGINE_CHAT_REFORM_RESOLVER_MODEL", - DEFAULT_FAST_MODEL, -) -REFORM_RESOLVER_MAX_TOKENS = int( - os.environ.get("POLICYENGINE_CHAT_REFORM_RESOLVER_MAX_TOKENS", "2048") -) - - -class ReformAssessmentError(RuntimeError): - """The resolver did not produce a structurally safe assessment.""" - - -class GatewayCatalogueUnavailable(ReformAssessmentError): - """The current PolicyEngine parameter catalogue could not be queried.""" - - -@dataclass(frozen=True) -class ValidatedParameterBinding: - parameter_path: str - label: str - catalogue_evidence: str - - -@dataclass(frozen=True) -class ReformAlternative: - summary: str - parameter_bindings: tuple[ValidatedParameterBinding, ...] - reform: dict[str, Any] - - -@dataclass(frozen=True) -class ReformAssessment: - reform: dict[str, Any] | None - summary: str | None - confidence: int - parameter_bindings: tuple[ValidatedParameterBinding, ...] - alternatives: tuple[ReformAlternative, ...] - search_queries: tuple[str, ...] - catalogue_version: str - - -_SEARCH_TOOL = { - "name": "search_reform_targets", - "description": ( - "Search current PolicyEngine UK reformable parameters. Search before " - "constructing the reform and use only returned parameter paths." - ), - "input_schema": { - "type": "object", - "properties": {"query": {"type": "string", "minLength": 1}}, - "required": ["query"], - "additionalProperties": False, - }, -} - -_BINDING_SCHEMA = { - "type": "object", - "properties": { - "parameter_path": {"type": "string"}, - "label": {"type": "string"}, - }, - "required": ["parameter_path", "label"], - "additionalProperties": False, -} - -_ALTERNATIVE_SCHEMA = { - "type": "object", - "properties": { - "summary": {"type": "string"}, - "reform": REFORM_SCHEMA, - "bindings": {"type": "array", "items": _BINDING_SCHEMA}, - }, - "required": ["summary", "reform", "bindings"], - "additionalProperties": False, -} - -_ASSESSMENT_TOOL = { - "name": "emit_reform_assessment", - "description": ( - "Emit the best exact reform construction and calibrated confidence after " - "searching the current parameter catalogue." - ), - "input_schema": { - "type": "object", - "properties": { - "summary": {"type": "string"}, - "confidence": {"type": "integer", "minimum": 0, "maximum": 100}, - "reform": REFORM_SCHEMA, - "bindings": {"type": "array", "items": _BINDING_SCHEMA}, - "alternatives": { - "type": "array", - "maxItems": 3, - "items": _ALTERNATIVE_SCHEMA, - }, - }, - "required": ["summary", "confidence", "reform", "bindings", "alternatives"], - "additionalProperties": False, - }, -} - -_SYSTEM = """You resolve a grounded UK tax-benefit reform into exact PolicyEngine -parameter changes. Search before assessing. Use only parameter paths and labels -returned by search. Search results include the current-year value and unit so -you can turn relative wording into final values. Emit one best construction, -0-100 confidence, and up to three materially plausible alternatives. Confidence -means confidence that the construction exactly represents the user's wording, -not merely that the parameter exists. Never invent a path or label.""" - - -def current_catalogue_version() -> str: - try: - return version("policyengine-uk") - except PackageNotFoundError: - return "unknown" - - -def _search_with_values(query: str, limit: int) -> list[dict[str, Any]]: - rows = search_reform_targets(query=query, limit=limit) - enriched: list[dict[str, Any]] = [] - for row in rows: - item = dict(row) - detail = get_parameter(item["path"], DEFAULT_SIMULATION_YEAR) - parameter = detail.get("parameter") if isinstance(detail, dict) else None - if isinstance(parameter, dict): - item.update(parameter) - enriched.append(item) - return enriched - - -def _validate_reform(reform: dict[str, Any], year: int) -> dict[str, Any]: - return validate_reform_dict(reform, year=year) - - -def _tool_block(block: Any) -> dict[str, Any]: - return { - "type": "tool_use", - "id": str(getattr(block, "id", "resolver-tool")), - "name": str(getattr(block, "name", "")), - "input": getattr(block, "input", {}), - } - - -def _tool_result(tool_id: str, content: Any, *, is_error: bool = False) -> dict[str, Any]: - result = { - "type": "tool_result", - "tool_use_id": tool_id, - "content": json.dumps(content, default=str), - } - if is_error: - result["is_error"] = True - return result - - -def _binding_rows( - raw_bindings: Any, - reform: dict[str, Any], - candidates: dict[str, dict[str, Any]], -) -> tuple[ValidatedParameterBinding, ...]: - if not isinstance(raw_bindings, list): - raise ReformAssessmentError("assessment bindings must be a list") - bindings: list[ValidatedParameterBinding] = [] - seen: set[str] = set() - for item in raw_bindings: - if not isinstance(item, dict): - raise ReformAssessmentError("assessment binding must be an object") - path = item.get("parameter_path") - label = item.get("label") - if not isinstance(path, str) or path not in candidates: - raise ReformAssessmentError( - "assessment parameter path was not present in search results" - ) - expected_label = candidates[path].get("label") or path - if label != expected_label: - raise ReformAssessmentError("assessment binding label did not match catalogue label") - if path in seen: - continue - seen.add(path) - bindings.append( - ValidatedParameterBinding( - parameter_path=path, - label=label, - catalogue_evidence=str(candidates[path].get("query", "")), - ) - ) - if seen != set(reform): - raise ReformAssessmentError("assessment bindings must exactly cover reform paths") - return tuple(bindings) - - -def _direction_matches( - reform: dict[str, Any], - candidates: dict[str, dict[str, Any]], - intent: ReformIntent, -) -> bool: - comparable = [] - for path, proposed in reform.items(): - current = candidates[path].get("value") - if isinstance(current, (int, float)) and isinstance(proposed, (int, float)): - comparable.append((float(current), float(proposed))) - if not comparable: - return True - if intent.action in ("increase", "uprate"): - return all(proposed > current for current, proposed in comparable) - if intent.action == "decrease": - return all(proposed < current for current, proposed in comparable) - if intent.action == "multiply" and intent.amount == "2x": - return all(abs(proposed - current * 2) <= 1e-9 for current, proposed in comparable) - if intent.action == "abolish": - return all(proposed == 0 for _current, proposed in comparable) - return True - - -def _validated_construction( - raw: dict[str, Any], - *, - candidates: dict[str, dict[str, Any]], - intent: ReformIntent, - validate: Callable[[dict[str, Any], int], dict[str, Any]], -) -> tuple[dict[str, Any], tuple[ValidatedParameterBinding, ...], str]: - reform = raw.get("reform") - summary = raw.get("summary") - if not isinstance(reform, dict) or not reform: - raise ReformAssessmentError("assessment reform must be a non-empty object") - if not isinstance(summary, str) or not summary.strip(): - raise ReformAssessmentError("assessment summary must be non-empty") - unknown = set(reform).difference(candidates) - if unknown: - raise ReformAssessmentError( - "assessment parameter path was not present in search results" - ) - validation = validate(reform, DEFAULT_SIMULATION_YEAR) - if not validation.get("valid"): - raise ReformAssessmentError("assessment reform failed PolicyEngine validation") - normalized = validation.get("normalized_reform") - if not isinstance(normalized, dict) or set(normalized) != set(reform): - raise ReformAssessmentError("validated reform changed the proposed paths") - if not _direction_matches(normalized, candidates, intent): - raise ReformAssessmentError("assessment reform contradicts the requested direction") - bindings = _binding_rows(raw.get("bindings"), normalized, candidates) - return normalized, bindings, summary.strip() - - -def _parse_assessment( - raw: Any, - *, - candidates: dict[str, dict[str, Any]], - intent: ReformIntent, - validate: Callable[[dict[str, Any], int], dict[str, Any]], - searches: tuple[str, ...], - catalogue_version: str, -) -> ReformAssessment: - if not isinstance(raw, dict): - raise ReformAssessmentError("assessment output must be an object") - confidence = raw.get("confidence") - if isinstance(confidence, bool) or not isinstance(confidence, int) or not 0 <= confidence <= 100: - raise ReformAssessmentError("assessment confidence must be an integer from 0 to 100") - if raw.get("reform") == {}: - if raw.get("bindings") != [] or raw.get("alternatives") != []: - raise ReformAssessmentError("empty assessment cannot contain bindings or alternatives") - summary = raw.get("summary") - if not isinstance(summary, str) or not summary.strip(): - raise ReformAssessmentError("assessment summary must be non-empty") - return ReformAssessment( - reform=None, - summary=summary.strip(), - confidence=confidence, - parameter_bindings=(), - alternatives=(), - search_queries=searches, - catalogue_version=catalogue_version, - ) - reform, bindings, summary = _validated_construction( - raw, - candidates=candidates, - intent=intent, - validate=validate, - ) - alternatives: list[ReformAlternative] = [] - raw_alternatives = raw.get("alternatives") - if not isinstance(raw_alternatives, list): - raise ReformAssessmentError("assessment alternatives must be a list") - for alternative in raw_alternatives[:3]: - if not isinstance(alternative, dict): - continue - try: - alt_reform, alt_bindings, alt_summary = _validated_construction( - alternative, - candidates=candidates, - intent=intent, - validate=validate, - ) - except ReformAssessmentError: - # Alternatives are optional clarification aids, never executable - # authority. A malformed or directionally contradictory suggestion - # must not discard an otherwise valid best construction. - continue - alternatives.append(ReformAlternative(alt_summary, alt_bindings, alt_reform)) - return ReformAssessment( - reform=reform, - summary=summary, - confidence=confidence, - parameter_bindings=bindings, - alternatives=tuple(alternatives), - search_queries=searches, - catalogue_version=catalogue_version, - ) - - -def assess_reform_with_catalogue( - prompt: str, - reform_intent: ReformIntent, - *, - client: Any | None = None, - search: Callable[[str, int], list[dict[str, Any]]] = _search_with_values, - validate: Callable[[dict[str, Any], int], dict[str, Any]] = _validate_reform, - catalogue_version: str | None = None, -) -> ReformAssessment: - """Search, construct, validate, and score one exact reform proposal.""" - - client = client or get_sync_client() - resolved_version = catalogue_version or current_catalogue_version() - messages: list[dict[str, Any]] = [ - { - "role": "user", - "content": ( - f"USER REQUEST:\n{prompt[:4000]}\n\n" - f"GROUNDED REFORM INTENT:\n" - f"policy={reform_intent.policy_phrase}\n" - f"action={reform_intent.action}\n" - f"amount={reform_intent.amount}\n" - f"scope={reform_intent.scope}\n" - f"evidence={reform_intent.evidence}" - ), - } - ] - searches: list[str] = [] - candidates: dict[str, dict[str, Any]] = {} - repairs = 0 - last_error = "resolver did not emit an assessment" - - for _iteration in range(MAX_RESOLVER_ITERATIONS): - response = client.messages.create( - model=REFORM_RESOLVER_MODEL, - max_tokens=REFORM_RESOLVER_MAX_TOKENS, - temperature=DEFAULT_TEMPERATURE, - system=_SYSTEM, - tools=[_SEARCH_TOOL, _ASSESSMENT_TOOL], - tool_choice={"type": "any"}, - messages=messages, - ) - blocks = [block for block in response.content or [] if getattr(block, "type", None) == "tool_use"] - if not blocks: - last_error = "resolver returned no tool call" - continue - - assistant_blocks = [_tool_block(block) for block in blocks] - results: list[dict[str, Any]] = [] - emitted = None - emitted_id = "assessment" - for block in blocks: - name = getattr(block, "name", None) - tool_id = str(getattr(block, "id", "resolver-tool")) - tool_input = getattr(block, "input", {}) - if name == "search_reform_targets": - query = tool_input.get("query") if isinstance(tool_input, dict) else None - query = query.strip() if isinstance(query, str) else "" - key = query.casefold() - if not query: - results.append(_tool_result(tool_id, {"error": "query is required"}, is_error=True)) - continue - if key in {item.casefold() for item in searches}: - rows = [row for row in candidates.values() if row.get("query", "").casefold() == key] - results.append(_tool_result(tool_id, {"query": query, "targets": rows})) - continue - if len(searches) >= MAX_REFORM_SEARCHES: - results.append( - _tool_result( - tool_id, - {"error": f"search limit is {MAX_REFORM_SEARCHES}; emit the assessment now"}, - is_error=True, - ) - ) - continue - try: - rows = search(query, REFORM_SEARCH_LIMIT) - except Exception as exc: - raise GatewayCatalogueUnavailable(str(exc)) from exc - searches.append(query) - for row in rows: - path = row.get("path") if isinstance(row, dict) else None - if not isinstance(path, str): - continue - candidate = dict(row) - candidate["query"] = query - candidates[path] = candidate - results.append(_tool_result(tool_id, {"query": query, "targets": rows})) - elif name == "emit_reform_assessment": - emitted = tool_input - emitted_id = tool_id - else: - results.append(_tool_result(tool_id, {"error": "unknown resolver tool"}, is_error=True)) - - if emitted is not None: - try: - if not searches: - raise ReformAssessmentError("resolver must search before assessment") - return _parse_assessment( - emitted, - candidates=candidates, - intent=reform_intent, - validate=validate, - searches=tuple(searches), - catalogue_version=resolved_version, - ) - except ReformAssessmentError as exc: - last_error = str(exc) - if repairs >= MAX_ASSESSMENT_REPAIRS: - raise - repairs += 1 - results.append( - _tool_result( - emitted_id, - {"error": last_error, "instruction": "repair and emit a valid assessment"}, - is_error=True, - ) - ) - - messages.append({"role": "assistant", "content": assistant_blocks}) - messages.append({"role": "user", "content": results}) - - raise ReformAssessmentError(last_error) diff --git a/backend/gateway/catalogue.py b/backend/gateway/catalogue.py deleted file mode 100644 index 95a9b8ed..00000000 --- a/backend/gateway/catalogue.py +++ /dev/null @@ -1,234 +0,0 @@ -"""Deterministic policyengine.py catalogue evidence for opening-turn routing.""" - -from __future__ import annotations - -import re -from dataclasses import dataclass -from difflib import SequenceMatcher -from typing import Literal, Sequence - -from engine.discovery import search_variables -from engine.py_runtime import uk_model_version -from engine.reforms import search_reform_targets - -CatalogueKind = Literal["reform_target", "variable"] -CatalogueMatchType = Literal[ - "exact_identifier", - "exact_alias", - "exact_label", - "strong_phrase", - "fuzzy_suggestion", -] - -MAX_CATALOGUE_QUERIES = 4 -MATCH_LIMIT = 5 -CANDIDATE_LIMIT = 100 - - -@dataclass(frozen=True) -class CatalogueQuery: - """One concise model-concept lookup requested by the gateway classifier.""" - - kind: CatalogueKind - query: str - evidence: str | None = None - - -@dataclass(frozen=True) -class CatalogueMatch: - """One current policyengine.py parameter or variable match.""" - - kind: CatalogueKind - query: str - identifier: str - label: str - match_type: CatalogueMatchType = "strong_phrase" - score: float = 0.9 - - @property - def authoritative(self) -> bool: - return self.match_type != "fuzzy_suggestion" - - -@dataclass(frozen=True) -class CatalogueEvidence: - """Resolved opening-turn catalogue evidence. - - ``available`` is false only when the current model catalogue could not be - loaded. That is distinct from a successful lookup with no matches, which - is represented by ``unresolved_queries``. - """ - - available: bool - matches: tuple[CatalogueMatch, ...] = () - unresolved_queries: tuple[CatalogueQuery, ...] = () - - @property - def authoritative_matches(self) -> tuple[CatalogueMatch, ...]: - return tuple(match for match in self.matches if match.authoritative) - - @property - def suggestions(self) -> tuple[CatalogueMatch, ...]: - return tuple(match for match in self.matches if not match.authoritative) - - -def _normalise_match_text(value: str | None) -> str: - return " ".join(re.findall(r"[a-z0-9]+", (value or "").casefold())) - - -def _classify_match( - query: str, - *, - identifier: str, - label: str, - aliases: Sequence[str] = (), - description: str | None = None, -) -> tuple[CatalogueMatchType, float]: - """Classify deterministic lookup output by the evidence it actually gives. - - Fuzzy similarity remains useful for suggestions, but only exact and - sufficiently specific phrase matches can authorize catalogue recovery. - """ - - query_text = _normalise_match_text(query) - identifier_text = _normalise_match_text(identifier) - label_text = _normalise_match_text(label) - alias_texts = tuple(_normalise_match_text(alias) for alias in aliases) - - if query_text and query_text == identifier_text: - return "exact_identifier", 1.0 - if query_text and query_text in alias_texts: - return "exact_alias", 1.0 - if query_text and query_text == label_text: - return "exact_label", 1.0 - - query_tokens = set(query_text.split()) - strong_fields = (identifier_text, label_text, *alias_texts) - if len(query_tokens) >= 2 and any( - query_tokens.issubset(set(field.split())) for field in strong_fields if field - ): - return "strong_phrase", 0.9 - - comparison_fields = (*strong_fields, _normalise_match_text(description)) - score = max( - ( - SequenceMatcher(None, query_text, field).ratio() - for field in comparison_fields - if query_text and field - ), - default=0.0, - ) - return "fuzzy_suggestion", round(score, 4) - - -def _catalogue_match( - *, - kind: CatalogueKind, - query: str, - identifier: str, - label: str, - aliases: Sequence[str] = (), - description: str | None = None, -) -> CatalogueMatch: - match_type, score = _classify_match( - query, - identifier=identifier, - label=label, - aliases=aliases, - description=description, - ) - return CatalogueMatch( - kind=kind, - query=query, - identifier=identifier, - label=label, - match_type=match_type, - score=score, - ) - - -def _normalise_queries(queries: Sequence[CatalogueQuery]) -> tuple[CatalogueQuery, ...]: - """Bound and de-duplicate untrusted classifier output before model lookup.""" - - normalised: list[CatalogueQuery] = [] - seen: set[tuple[str, str]] = set() - for item in queries: - if item.kind not in ("reform_target", "variable"): - continue - query = item.query.strip() - if not query: - continue - key = (item.kind, query.casefold()) - if key in seen: - continue - seen.add(key) - normalised.append(CatalogueQuery(item.kind, query, item.evidence)) - if len(normalised) == MAX_CATALOGUE_QUERIES: - break - return tuple(normalised) - - -def resolve_catalogue_queries(queries: Sequence[CatalogueQuery]) -> CatalogueEvidence: - """Resolve gateway terms against the current policyengine.py model catalogue. - - This is intentionally an internal server lookup, not a model-facing tool. - It confirms that a named concept exists without deciding which candidate the - user intended or whether the request has enough information to execute. - """ - - queries = _normalise_queries(queries) - if not queries: - return CatalogueEvidence(available=True) - - try: - # Check availability once before calling helpers that search the same - # cached model. Availability remains explicit so routing evidence is - # never confused with a successful lookup that returned no matches. - uk_model_version() - matches: list[CatalogueMatch] = [] - unresolved: list[CatalogueQuery] = [] - for item in queries: - if item.kind == "reform_target": - rows = search_reform_targets(item.query, limit=CANDIDATE_LIMIT) - item_matches = [ - _catalogue_match( - kind=item.kind, - query=item.query, - identifier=row["path"], - label=row.get("label") or row["path"], - aliases=row.get("aliases") or (), - description=row.get("description"), - ) - for row in rows - ] - else: - response = search_variables(item.query, limit=CANDIDATE_LIMIT) - item_matches = [ - _catalogue_match( - kind=item.kind, - query=item.query, - identifier=row["name"], - label=row.get("label") or row["name"], - description=row.get("description"), - ) - for row in response["variables"] - ] - item_matches.sort( - key=lambda match: ( - not match.authoritative, - -match.score, - match.label.casefold(), - ) - ) - item_matches = item_matches[:MATCH_LIMIT] - matches.extend(item_matches) - if not any(match.authoritative for match in item_matches): - unresolved.append(item) - except Exception: # noqa: BLE001 - catalogue metadata must fail open - return CatalogueEvidence(available=False) - - return CatalogueEvidence( - available=True, - matches=tuple(matches), - unresolved_queries=tuple(unresolved), - ) diff --git a/backend/gateway/clarifications.py b/backend/gateway/clarifications.py deleted file mode 100644 index 3b105ab2..00000000 --- a/backend/gateway/clarifications.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Deterministic rendering for authorized gateway clarification reasons.""" - -from __future__ import annotations - -from typing import Any - -from gateway.policy import GatingReason - -MAX_CLARIFICATION_QUESTIONS = 3 - - -def _quoted_labels(labels: tuple[str, ...]) -> str: - quoted = [f"“{label}”" for label in labels] - if len(quoted) == 1: - return quoted[0] - if len(quoted) == 2: - return f"{quoted[0]} or {quoted[1]}" - return ", ".join(quoted[:-1]) + f", or {quoted[-1]}" - - -def render_gating_reason(reason: GatingReason) -> str | None: - """Render one reason without inspecting arbitrary slot names.""" - - if reason.code == "missing_reform": - return "What policy change would you like me to model?" - if reason.code == "missing_output": - return ( - "What result would you like to see—for example, budgetary impact, " - "poverty, decile impacts, or winners and losers?" - ) - if reason.code == "missing_household_composition": - return "What household composition should I model?" - if reason.code == "missing_tool": - return "What tax-benefit calculation would you like me to run?" - if reason.code == "catalogue_choice" and reason.options: - return ( - "Which supported parameter did you mean: " - + _quoted_labels(reason.options) - + "?" - ) - if reason.code == "catalogue_no_match": - return ( - "I couldn’t identify a supported PolicyEngine parameter for that reform. " - "Could you name the specific tax, benefit, rate, threshold, or allowance " - "you want to change?" - ) - return None - - -def _binding_labels(bindings: Any) -> tuple[str, ...]: - labels = [] - for binding in bindings or (): - label = getattr(binding, "label", None) - if isinstance(label, str) and label and label not in labels: - labels.append(label) - return tuple(labels) - - -def _proposal_description(intent: Any, bindings: Any) -> str | None: - labels = _binding_labels(bindings) - if not labels or intent is None: - return None - subject = _quoted_labels(labels).replace(" or ", " and ") - amount = getattr(intent, "amount", None) - action = getattr(intent, "action", None) - if action == "increase" and amount: - return f"increasing {subject} by {amount}" - if action == "decrease" and amount: - return f"decreasing {subject} by {amount}" - if action == "set" and amount: - return f"setting {subject} to {amount}" - if action == "abolish": - return f"abolishing {subject}" - if action == "freeze": - return f"freezing {subject}" - if action == "uprate" and amount: - return f"uprating {subject} by {amount}" - if action == "replace" and amount: - return f"replacing {subject} with {amount}" - if action == "multiply" and amount == "2x": - return f"doubling {subject}" - return None - - -def _render_confirmation(verdict: Any) -> str | None: - assessment = getattr(verdict, "reform_assessment", None) - intent = getattr(verdict, "reform_intent", None) - if assessment is None: - return None - proposal = _proposal_description(intent, assessment.parameter_bindings) - if proposal is None: - return None - rendered = f"I would model this as {proposal}. Is that what you intended?" - alternatives = [ - description - for alternative in assessment.alternatives - if ( - description := _proposal_description( - intent, - alternative.parameter_bindings, - ) - ) - ] - if alternatives: - label = ( - "Other plausible interpretation" - if len(alternatives) == 1 - else "Other plausible interpretations" - ) - rendered += "\n\n" + label + ": " + "; ".join(alternatives) + "." - return rendered - - -def render_clarification(verdict: Any) -> str | None: - """Render at most three stable questions, or fail closed with ``None``.""" - - rendered: list[str] = [] - seen: set[str] = set() - for reason in getattr(verdict, "gating_reasons", ()): - if reason.code == "confirm_reform": - question = _render_confirmation(verdict) - else: - question = render_gating_reason(reason) - if question is None: - return None - if question in seen: - continue - seen.add(question) - rendered.append(question) - if len(rendered) == MAX_CLARIFICATION_QUESTIONS: - break - if not rendered: - return None - if len(rendered) == 1: - return rendered[0] - return "\n".join( - f"{index}. {question}" for index, question in enumerate(rendered, start=1) - ) diff --git a/backend/gateway/execution.py b/backend/gateway/execution.py deleted file mode 100644 index 356534c2..00000000 --- a/backend/gateway/execution.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Dependency-aware execution plans produced from a ready gateway verdict.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Literal - -from engine.constants import UK_CHAT_DATASET -from engine.decile_concepts import DEFAULT_DECILE_CONCEPT -from gateway.assessment import ReformAssessment, ValidatedParameterBinding -from gateway.intent import ReformIntent -from gateway.policy import SlotFact -from tools.definitions import DEFAULT_SIMULATION_YEAR - -OUTPUT_TO_ANALYSIS_TOOL = { - "budgetary_impact": "compute_budgetary_impact", - "tax_revenue": "compute_budgetary_impact", - "benefit_spending": "compute_budgetary_impact", - "poverty_impact": "compute_poverty_metrics", - "inequality_impact": "compute_inequality_metrics", - "decile_impact": "compute_decile_impacts", - "winners_losers": "compute_winners_losers", -} - -TOOL_DEPENDENCIES = { - "compute_budgetary_impact": ("run_society_simulation",), - "compute_program_breakdown": ("run_society_simulation",), - "compute_decile_impacts": ("run_society_simulation",), - "compute_winners_losers": ("run_society_simulation",), - "compute_poverty_metrics": ("run_society_simulation",), - "compute_inequality_metrics": ("run_society_simulation",), - "aggregate_result": ("run_society_simulation",), -} - -_SOCIETY_TOOLS = {"run_society_simulation", *TOOL_DEPENDENCIES} - - -@dataclass(frozen=True) -class ExecutionInput: - name: str - value: str - source: Literal["prompt", "default"] - - -@dataclass(frozen=True) -class ExecutionConvention: - name: str - value: str - - -@dataclass(frozen=True) -class GatewayExecutionPlan: - target_tool: str | None - prerequisites: tuple[str, ...] - inputs: tuple[ExecutionInput, ...] - conventions: tuple[ExecutionConvention, ...] - parameter_bindings: tuple[ValidatedParameterBinding, ...] - approved_reform: dict[str, Any] | None = None - - -def analysis_tool_for_output( - selected_tool: str | None, - output: str | None, -) -> str | None: - """Upgrade society analysis to the derivative requested by the user.""" - - if selected_tool not in _SOCIETY_TOOLS: - return selected_tool - return OUTPUT_TO_ANALYSIS_TOOL.get(output or "", selected_tool) - - -def _slot_value(slots: list[SlotFact], name: str, *, kind: str | None = None) -> str | None: - return next( - ( - slot.value - for slot in slots - if slot.name == name - and (kind is None or slot.kind == kind) - and slot.value is not None - ), - None, - ) - - -def _output_value(slots: list[SlotFact]) -> str | None: - """Return the canonical output value regardless of classifier slot name.""" - - return next( - ( - slot.value - for slot in slots - if slot.kind == "output" and slot.value is not None - ), - None, - ) - - -def build_execution_plan( - selected_tool: str | None, - slots: list[SlotFact], - reform_intent: ReformIntent | None, - prompt: str, - reform_assessment: ReformAssessment | Any | None, -) -> GatewayExecutionPlan: - """Build an ordered, exact plan for the compute model to execute.""" - - del prompt # reserved for future deterministic convention selection - output = _output_value(slots) - target = analysis_tool_for_output(selected_tool, output) - is_society = target in _SOCIETY_TOOLS - if is_society and reform_intent is not None and ( - reform_assessment is None or reform_assessment.reform is None - ): - raise ValueError("society reform execution requires a validated assessment") - - inputs = [ - ExecutionInput(slot.name, slot.value, slot.source) - for slot in slots - if slot.kind == "tool_input" - and slot.source in ("prompt", "default") - and slot.value is not None - ] - if is_society and not any(item.name == "year" for item in inputs): - inputs.append( - ExecutionInput("year", str(DEFAULT_SIMULATION_YEAR), "default") - ) - - conventions: list[ExecutionConvention] = [] - if is_society: - conventions.extend( - [ - ExecutionConvention("comparator", "current law"), - ExecutionConvention("population", "full modelled population"), - ExecutionConvention( - "jurisdictions", - "applicable modelled UK jurisdictions", - ), - ExecutionConvention("method", "direct static microsimulation"), - ExecutionConvention("dataset", UK_CHAT_DATASET.label), - ] - ) - if target == "compute_decile_impacts": - conventions.append( - ExecutionConvention( - "decile_concept", - _slot_value(slots, "decile_concept") - or DEFAULT_DECILE_CONCEPT.value, - ) - ) - - bindings = ( - tuple(reform_assessment.parameter_bindings) - if reform_assessment is not None - else () - ) - approved_reform = ( - dict(reform_assessment.reform) - if reform_assessment is not None and reform_assessment.reform is not None - else None - ) - return GatewayExecutionPlan( - target_tool=target, - prerequisites=TOOL_DEPENDENCIES.get(target, ()), - inputs=tuple(inputs), - conventions=tuple(conventions), - parameter_bindings=bindings, - approved_reform=approved_reform, - ) diff --git a/backend/gateway/intent.py b/backend/gateway/intent.py deleted file mode 100644 index 35b338ae..00000000 --- a/backend/gateway/intent.py +++ /dev/null @@ -1,321 +0,0 @@ -"""Deterministic extraction of bounded output and reform intent.""" - -from __future__ import annotations - -import re -from dataclasses import dataclass, replace -from typing import Literal - -from gateway.policy import SlotFact, TOOL_SLOT_REQUIREMENT - -OutputKind = Literal[ - "budgetary_impact", - "tax_revenue", - "benefit_spending", - "poverty_impact", - "inequality_impact", - "decile_impact", - "winners_losers", -] - -ReformAction = Literal[ - "increase", - "decrease", - "set", - "abolish", - "freeze", - "uprate", - "replace", - "multiply", -] - -ReformScope = Literal["unspecified", "all", "every", "both"] - - -@dataclass(frozen=True) -class OutputIntent: - value: OutputKind - evidence: str - - -@dataclass(frozen=True) -class ReformIntent: - policy_phrase: str - action: ReformAction - amount: str | None - scope: ReformScope - evidence: str - - -_OUTPUT_PATTERNS: tuple[tuple[OutputKind, tuple[re.Pattern[str], ...]], ...] = ( - ( - "winners_losers", - ( - re.compile(r"\bgain(?:s|ed)?\s+or\s+los(?:e|es|t)\b", re.I), - re.compile(r"\bhow many\s+(?:people|households)\s+(?:would\s+)?(?:gain|lose|be affected)\b", re.I), - re.compile(r"\bhouseholds?\s+(?:would\s+)?gain\b", re.I), - re.compile(r"\baffected[- ]household counts?\b", re.I), - ), - ), - ( - "decile_impact", - ( - re.compile(r"\bby\s+(?:income\s+)?decile\b", re.I), - re.compile(r"\bdecile\s+impacts?\b", re.I), - re.compile(r"\bdistributional impact\s+by\s+decile\b", re.I), - ), - ), - ( - "poverty_impact", - (re.compile(r"\b(?:child\s+)?poverty\b", re.I),), - ), - ( - "inequality_impact", - ( - re.compile(r"\binequality\b", re.I), - re.compile(r"\bgini\b", re.I), - ), - ), - ( - "tax_revenue", - ( - re.compile(r"\bannual revenue\b", re.I), - re.compile(r"\btax revenue\b", re.I), - re.compile(r"\brevenue from\b", re.I), - ), - ), - ( - "benefit_spending", - ( - re.compile(r"\bbenefit spending\b", re.I), - re.compile(r"\bbenefit expenditure\b", re.I), - ), - ), - ( - "budgetary_impact", - ( - re.compile(r"\bbudgetary (?:cost|impact)\b", re.I), - re.compile(r"\bfiscal (?:cost|impact)\b", re.I), - re.compile(r"\bannual cost\b", re.I), - re.compile( - r"\bcost of\s+(?:increas\w*|rais\w*|reduc\w*|lower\w*|cut\w*|abolish\w*|freez\w*|uprat\w*|replac\w*|doubl\w*|sett?\w*)\b", - re.I, - ), - ), - ), -) - - -def output_from_prompt(prompt: str) -> OutputIntent | None: - """Return the highest-precedence directly modelled output in ``prompt``.""" - - for value, patterns in _OUTPUT_PATTERNS: - matches = [match for pattern in patterns if (match := pattern.search(prompt))] - if matches: - match = min(matches, key=lambda candidate: candidate.start()) - return OutputIntent(value=value, evidence=match.group(0)) - return None - - -_AMOUNT = ( - r"(?:" - r"£\s?\d[\d,]*(?:\.\d+)?(?:\s+per\s+(?:week|month|year))?" - r"|\d+(?:\.\d+)?\s*%" - r"|(?:one|two|three|four|five|six|seven|eight|nine|ten|\d+(?:\.\d+)?)" - r"\s*(?:percentage points?|pp)" - r")" -) -_POLICY = r"(?P[^?.!,;]+?)" -_SCOPE_TAIL = r"(?:\s+for\s+(?:all|every|both)\b[^?.!,;]*)?" - -_AMOUNT_ACTION = re.compile( - rf"\b(?Pincreas(?:e|es|ed|ing)|rais(?:e|es|ed|ing)|" - rf"reduc(?:e|es|ed|ing)|lower(?:s|ed|ing)?|cut(?:s|ting)?|" - rf"uprat(?:e|es|ed|ing))\s+{_POLICY}\s+(?:by|to)\s+" - rf"(?P{_AMOUNT}){_SCOPE_TAIL}", - re.I, -) -_FROM_TO_ACTION = re.compile( - rf"\b(?Pincreas(?:e|es|ed|ing)|rais(?:e|es|ed|ing)|" - rf"reduc(?:e|es|ed|ing)|lower(?:s|ed|ing)?|cut(?:s|ting)?)\s+" - rf"{_POLICY}\s+from\s+{_AMOUNT}\s+to\s+(?P{_AMOUNT})" - rf"{_SCOPE_TAIL}", - re.I, -) -_SET_ACTION = re.compile( - rf"\bset(?:s|ting)?\s+{_POLICY}\s+to\s+(?P{_AMOUNT}){_SCOPE_TAIL}", - re.I, -) -_REPLACE_ACTION = re.compile( - rf"\breplac(?:e|es|ed|ing)\s+{_POLICY}\s+with\s+(?P[^?.!,;]+)", - re.I, -) -_NO_AMOUNT_ACTION = re.compile( - r"\b(?Pabolish(?:es|ed|ing)?|scrap(?:s|ped|ping)?|" - r"freez(?:e|es|ing)|froze)\s+(?P[^?.!,;]+)", - re.I, -) -_MULTIPLY_ACTION = re.compile( - r"\b(?Pdouble|doubles|doubled|doubling)\s+(?P[^?.!,;]+)", - re.I, -) - -_DECREASE_VERBS = ("reduc", "lower", "cut") -_UPRATE_VERBS = ("uprat",) -_GENERIC_POLICIES = { - "it", - "this", - "that", - "them", - "the reform", - "a reform", - "reform", - "the policy", - "policy", - "the two reforms", - "two reforms", -} - - -def _scope(evidence: str) -> ReformScope: - lowered = evidence.casefold() - for scope in ("all", "every", "both"): - if re.search(rf"\b{scope}\b", lowered): - return scope # type: ignore[return-value] - return "unspecified" - - -def _clean_policy(value: str) -> str | None: - policy = value.strip(" \t\n\r-–—") - policy = re.sub(r"^(?:the|a|an)\s+", "", policy, flags=re.I) - policy = re.sub(r"^(?:all|every|both)\s+", "", policy, flags=re.I) - policy = policy.strip() - if not policy or policy.casefold() in _GENERIC_POLICIES: - return None - if len(re.findall(r"[A-Za-z]", policy)) < 3: - return None - return policy - - -def reform_intent_from_prompt(prompt: str) -> ReformIntent | None: - """Extract a complete, grounded natural-language reform operation.""" - - match = _MULTIPLY_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - if policy: - return ReformIntent(policy, "multiply", "2x", _scope(match.group(0)), match.group(0)) - - match = _REPLACE_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - amount = match.group("amount").strip() - if policy and amount: - return ReformIntent(policy, "replace", amount, _scope(match.group(0)), match.group(0)) - - match = _SET_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - if policy: - return ReformIntent( - policy, - "set", - match.group("amount").strip(), - _scope(match.group(0)), - match.group(0), - ) - - # A fully specified "from X to Y" proposal gives a final value. Treat it - # as a set operation so assessment does not compare the destination with a - # possibly different current catalogue value and reject the stated - # direction (for example a historical 18% -> 20% change). - match = _FROM_TO_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - if policy: - return ReformIntent( - policy, - "set", - match.group("amount").strip(), - _scope(match.group(0)), - match.group(0), - ) - - match = _AMOUNT_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - verb = match.group("verb").casefold() - if policy: - action: ReformAction - if verb.startswith(_DECREASE_VERBS): - action = "decrease" - elif verb.startswith(_UPRATE_VERBS): - action = "uprate" - else: - action = "increase" - return ReformIntent( - policy, - action, - match.group("amount").strip(), - _scope(match.group(0)), - match.group(0), - ) - - match = _NO_AMOUNT_ACTION.search(prompt) - if match: - policy = _clean_policy(match.group("policy")) - if policy: - verb = match.group("verb").casefold() - action = "freeze" if verb.startswith(("freez", "froze")) else "abolish" - return ReformIntent(policy, action, None, _scope(match.group(0)), match.group(0)) - return None - - -def upsert_output_slot(slots: list[SlotFact], intent: OutputIntent) -> list[SlotFact]: - """Ground a missing/assumed output without overriding explicit evidence.""" - - output_index = next( - (index for index, slot in enumerate(slots) if slot.kind == "output"), - None, - ) - grounded = SlotFact( - name="output", - source="prompt", - kind="output", - value=intent.value, - ) - if output_index is None: - return [*slots, grounded] - if slots[output_index].source == "prompt": - return list(slots) - updated = list(slots) - updated[output_index] = replace(grounded) - return updated - - -def upsert_prompt_year( - tool: str | None, - slots: list[SlotFact], - prompt: str, -) -> list[SlotFact]: - """Preserve one explicit prompt year for tools whose schema accepts it.""" - - if tool is None or (tool, "year") not in TOOL_SLOT_REQUIREMENT: - return list(slots) - years = list(dict.fromkeys(re.findall(r"\b(?:19|20)\d{2}\b", prompt))) - if len(years) != 1: - return list(slots) - grounded = SlotFact(name="year", source="prompt", value=years[0]) - year_index = next( - ( - index - for index, slot in enumerate(slots) - if slot.kind == "tool_input" and slot.name == "year" - ), - None, - ) - if year_index is None: - return [*slots, grounded] - updated = list(slots) - updated[year_index] = grounded - return updated diff --git a/backend/gateway/policy.py b/backend/gateway/policy.py deleted file mode 100644 index 58bb1da0..00000000 --- a/backend/gateway/policy.py +++ /dev/null @@ -1,381 +0,0 @@ -"""Gateway policy: per-slot criticality and the deterministic gate. - -The chat gateway splits responsibility deliberately: the *model grounds*, the -*server gates*. The gateway model call emits, per plan slot, only a value and a -``source`` (the grounding flag: ``prompt`` / ``default`` / ``assumed``). It does -NOT judge importance. This module owns the importance policy — a per-slot -``criticality`` — and the pure, deterministic ``gate()`` that turns a grounded -plan into one of five outcomes. Keeping this out of the model (and out of -``prompts/``/``chat/``) makes the gate auditable and unit-testable -offline, and lets the eval grader import it without dragging in the runtime. - -Before gating, the server completes the selected tool's schema slots that the -model omitted as ``assumed``. A slot *gates* (forces a clarifying question) iff -its ``source`` is ``assumed`` AND its criticality is high/medium AND it is not -model-inferable. ``needs_plan`` iff any slot gates; otherwise ``ready`` (once -admissibility — irrelevant / out_of_scope / partial — is decided). -""" - -from __future__ import annotations - -from dataclasses import dataclass, field, replace -from typing import List, Literal, Optional - -from tools.definitions import DEFAULT_SIMULATION_YEAR, TOOL_DEFINITIONS - -# The default simulation year as a string, for comparison against years parsed -# out of the prompt. Sourced from the single shared constant (which also feeds -# YEAR_SCHEMA) so it can't drift from the schema default. Only used to decide -# whether the prompt names a *non-default* year. -_DEFAULT_YEAR = str(DEFAULT_SIMULATION_YEAR) - -Criticality = Literal["high", "medium", "low"] -SlotSource = Literal["prompt", "default", "assumed", "runtime"] -GatingReasonCode = Literal[ - "missing_tool", - "missing_reform", - "missing_output", - "missing_household_composition", - "catalogue_choice", - "catalogue_no_match", - "confirm_reform", - "internal_slot", -] -DomainStatus = Literal["uk_or_unspecified", "explicit_non_uk", "unrelated"] -CapabilityStatus = Literal[ - "supported", - "catalogue_uncertain", - "explicitly_unmodellable", -] - - -@dataclass(frozen=True) -class DomainDecision: - """Validated domain classification from the user's original wording.""" - - status: DomainStatus = "uk_or_unspecified" - evidence: Optional[str] = None - - -@dataclass(frozen=True) -class CapabilityDecision: - """Validated statement about why the gateway could not select a tool.""" - - status: CapabilityStatus = "supported" - evidence: Optional[str] = None - - -@dataclass(frozen=True) -class SlotFact: - """One slot of the execution plan, as grounded by the gateway model.""" - - name: str - source: SlotSource - kind: str = "tool_input" # "tool_input" | "output" - value: Optional[str] = None - - -@dataclass(frozen=True) -class GatingReason: - code: GatingReasonCode - slot: str - options: tuple[str, ...] = () - evidence: str | None = None - - -@dataclass(frozen=True) -class GateResult: - outcome: str - gating_reasons: List[GatingReason] = field(default_factory=list) - - @property - def gating_slots(self) -> List[str]: - return [reason.slot for reason in self.gating_reasons] - - -# --------------------------------------------------------------------------- -# Slot inventory, derived from TOOL_DEFINITIONS so it can't drift from the -# schemas. Each (tool, slot) is classified required / defaulted / optional. -# --------------------------------------------------------------------------- - -def _build_slot_inventory() -> tuple[dict, dict]: - requirements: dict = {} - defaults: dict = {} - for tool in TOOL_DEFINITIONS: - name = tool["name"] - schema = tool.get("input_schema", {}) - required = set(schema.get("required", [])) - props = schema.get("properties", {}) or {} - for slot, spec in props.items(): - if slot in required: - requirements[(name, slot)] = "required" - elif isinstance(spec, dict) and "default" in spec: - requirements[(name, slot)] = "defaulted" - defaults[(name, slot)] = spec["default"] - else: - requirements[(name, slot)] = "optional" - return requirements, defaults - - -TOOL_SLOT_REQUIREMENT, TOOL_SLOT_DEFAULTS = _build_slot_inventory() - -_SOCIETY_DERIVATIVE_TOOLS = { - "compute_budgetary_impact", - "compute_program_breakdown", - "compute_decile_impacts", - "compute_winners_losers", - "compute_poverty_metrics", - "compute_inequality_metrics", - "aggregate_result", -} - -RUNTIME_PROVIDED_SLOTS = { - *((tool, "simulation_id") for tool in _SOCIETY_DERIVATIVE_TOOLS), - ("generate_chart", "result_id"), -} - -# Curated overrides where the schema's required/default flags don't match the -# real importance. Kept tiny and commented to limit drift. -_CRITICALITY_OVERRIDES: dict = { - # run_society_simulation has required=[] in the schema, but a society-wide - # simulation with no reform is just a baseline snapshot — almost never the - # intent. Treat the reform as load-bearing. - ("run_society_simulation", "reform"): "high", -} - -# Schema-required (or otherwise high) slots that the model can reliably INFER -# rather than ask the user about. Without this, every household calc would ask -# "which benefit unit?" — this is the primary bound on over-asking. -INFERABLE: set = { - ("run_household_simulation", "benunit"), - ("run_household_simulation", "household"), - ("run_household_simulation", "benefit_entitlement"), - ("get_parameter", "path"), - ("generate_chart", "chart_kind"), - ("generate_chart", "chart_output"), - ("generate_chart", "output"), - ("generate_chart", "title"), - ("generate_chart", "x_field"), - ("generate_chart", "y_fields"), - ("generate_chart", "data"), # comes from an upstream tool, not the user -} - - -def _missing_slot_fact(tool: str, name: str) -> SlotFact: - key = (tool, name) - if key in RUNTIME_PROVIDED_SLOTS: - return SlotFact(name=name, source="runtime") - if key in TOOL_SLOT_DEFAULTS: - return SlotFact(name=name, source="default", value=str(TOOL_SLOT_DEFAULTS[key])) - return SlotFact(name=name, source="assumed") - - -def complete_slots(tool: Optional[str], slots: List[SlotFact]) -> List[SlotFact]: - """Return a complete gateway slot state for a selected tool. - - The gateway model can ground any subset of the plan. Missing schema slots - must still reach ``gate()`` as ``assumed`` rather than being silently - treated as known. ``output`` is a synthetic, user-requested deliverable - slot, so it is added when the model did not name one. - """ - - completed = list(slots) - present_tool_inputs = { - slot.name for slot in slots if slot.kind == "tool_input" - } - if tool is not None: - for candidate_tool, name in TOOL_SLOT_REQUIREMENT: - if candidate_tool == tool and name not in present_tool_inputs: - completed.append(_missing_slot_fact(tool, name)) - - if tool is not None and not any(slot.kind == "output" for slot in slots): - completed.append(SlotFact(name="output", source="assumed", kind="output")) - return completed - -# Closed vocabulary for the synthetic "output" (deliverable) slot. The single -# source of truth for the output labels: the gateway runtime injects these into -# the classifier prompt (via gateway_system) so the model and this module can't -# drift apart on the label set. -OUTPUT_VOCAB = ( - "budgetary_impact", - "tax_revenue", - "benefit_spending", - "poverty_impact", - "inequality_impact", - "decile_impact", - "winners_losers", - "caseload", - "marginal_rate", - "net_income", - "benefit_entitlement", - "parameter_lookup", - "reform_validity", -) - -# Most safe defaults are expressed directly in a tool schema. Current law is -# also the documented baseline for a society simulation, even though a missing -# reform is represented by ``None`` rather than a JSON-schema ``default``. -_EXPLICIT_SAFE_DEFAULTS = { - ("run_society_simulation", "reform"), -} - - -def normalise_slot_grounding( - tool: Optional[str], - slots: List[SlotFact], -) -> List[SlotFact]: - """Reject empty or unsupported claims that a slot is already grounded.""" - - normalised: list[SlotFact] = [] - for slot in slots: - value = slot.value.strip() if isinstance(slot.value, str) else None - if slot.kind == "output": - output = value or slot.name - if slot.source == "prompt" and output in OUTPUT_VOCAB: - normalised.append(replace(slot, value=output)) - else: - normalised.append(replace(slot, source="assumed", value=None)) - continue - - key = (tool, slot.name) - if key in RUNTIME_PROVIDED_SLOTS: - normalised.append(replace(slot, source="runtime", value=None)) - continue - if key in TOOL_SLOT_DEFAULTS and slot.source != "prompt": - normalised.append( - replace( - slot, - source="default", - value=str(TOOL_SLOT_DEFAULTS[key]), - ) - ) - continue - has_schema_default = TOOL_SLOT_REQUIREMENT.get(key) == "defaulted" - can_default = has_schema_default or key in _EXPLICIT_SAFE_DEFAULTS - if slot.source == "prompt" and not value: - normalised.append(replace(slot, source="assumed", value=None)) - elif slot.source == "default" and not can_default: - normalised.append(replace(slot, source="assumed", value=None)) - else: - normalised.append(replace(slot, value=value)) - return normalised - - -# --------------------------------------------------------------------------- -# Context promotions: raise criticality when a default would be actively wrong, -# not merely absent. Deterministic keyword scans over the prompt. -# --------------------------------------------------------------------------- - -_REFORM_INTENT_KEYWORDS = ( - "reform", "raise", "cut", "increase", "decrease", "abolish", "scrap", - "introduce", "freeze", "uprate", "replace", "change the", "set the", -) - - -def _prompt_names_reform(prompt: str) -> bool: - p = prompt.lower() - return any(kw in p for kw in _REFORM_INTENT_KEYWORDS) - - -def _prompt_implies_nondefault_year(prompt: str) -> bool: - import re - - years = re.findall(r"\b(?:19|20)\d{2}\b", prompt) - return any(y != _DEFAULT_YEAR for y in years) - - -def _base_criticality(tool: Optional[str], slot: SlotFact) -> Criticality: - if slot.kind == "output": - return "high" - key = (tool, slot.name) - if key in _CRITICALITY_OVERRIDES: - return _CRITICALITY_OVERRIDES[key] - requirement = TOOL_SLOT_REQUIREMENT.get(key) - if requirement == "required": - return "high" - # defaulted, optional, or an unrecognised (possibly hallucinated) slot. - return "low" - - -def criticality(tool: Optional[str], slot: SlotFact, prompt: str = "") -> Criticality: - """Resolved criticality for a slot: static base + context promotions.""" - base = _base_criticality(tool, slot) - if slot.kind == "output": - return base - if slot.name == "reform" and _prompt_names_reform(prompt): - return "high" - if slot.name == "year" and base == "low" and _prompt_implies_nondefault_year(prompt): - return "medium" - return base - - -def is_inferable(tool: Optional[str], slot_name: str) -> bool: - return (tool, slot_name) in INFERABLE - - -def slot_gates(tool: Optional[str], slot: SlotFact, prompt: str = "") -> bool: - """True if this slot should force a clarifying question.""" - if slot.source != "assumed": - return False - if tool == "generate_chart" and slot.kind == "output": - return False - if is_inferable(tool, slot.name): - return False - return criticality(tool, slot, prompt) in ("high", "medium") - - -def _gating_reason(slot: SlotFact) -> GatingReason: - if slot.kind == "output": - return GatingReason(code="missing_output", slot=slot.name) - if slot.name == "reform": - return GatingReason(code="missing_reform", slot=slot.name) - if slot.name in {"people", "benunit", "household"}: - return GatingReason(code="missing_household_composition", slot=slot.name) - return GatingReason(code="internal_slot", slot=slot.name) - - -def gate( - in_domain: bool, - tool: Optional[str], - slots: List[SlotFact], - unmodellable_outputs: List[str], - prompt: str = "", - *, - explicitly_unmodellable: bool = False, - reform_intent: object | None = None, -) -> GateResult: - """Deterministically map a grounded plan to one of the five outcomes. - - Admissibility is decided first (irrelevant / out_of_scope / partial); only an - admissible, fully-grounded plan reaches ``ready``. Fail-safe directions are - the caller's job (run_gateway defaults to ``ready`` on any error); this - function is pure. - """ - if not in_domain: - return GateResult("irrelevant") - - if tool is None: - if unmodellable_outputs or explicitly_unmodellable: - # A refusal needs positive, prompt-grounded capability evidence. - return GateResult("out_of_scope") - # Failure to choose a tool is uncertainty, not proof of incapability. - return GateResult( - "needs_plan", - [GatingReason(code="missing_tool", slot="tool")], - ) - if unmodellable_outputs: - # Some of the ask is modellable, some isn't → confirm-first partial. - # Deliberately takes precedence over needs_plan: resolve scope first; - # under-specified inputs get clarified on the next turn if the user - # proceeds. Both are lightweight, so neither wrongly refuses. - return GateResult("partial") - - gating = [ - _gating_reason(slot) - for slot in slots - if slot_gates(tool, slot, prompt) - and not (slot.name == "reform" and reform_intent is not None) - ] - if gating: - return GateResult("needs_plan", gating) - return GateResult("ready") diff --git a/backend/gateway/proposals.py b/backend/gateway/proposals.py deleted file mode 100644 index 3597027f..00000000 --- a/backend/gateway/proposals.py +++ /dev/null @@ -1,431 +0,0 @@ -"""Stateless signed proposal markers carried in raw assistant chat history.""" - -from __future__ import annotations - -import base64 -import hashlib -import hmac -import json -import os -import re -import time -from dataclasses import asdict, is_dataclass -from typing import Any - -PROPOSAL_MARKER_VERSION = 1 -DEFAULT_PROPOSAL_TTL_SECONDS = 24 * 60 * 60 -_MARKER_RE = re.compile(r"") - - -class ProposalSigningError(ValueError): - """A proposal marker was absent, invalid, stale, or unverifiable.""" - - -class ProposalExpiredError(ProposalSigningError): - def __init__(self, envelope: dict[str, Any]): - self.envelope = envelope - super().__init__("proposal has expired") - - -def _key(value: str | None) -> bytes: - resolved = value or os.environ.get("GATEWAY_PROPOSAL_SIGNING_KEY", "") - encoded = resolved.encode("utf-8") - if len(encoded) < 32: - raise ProposalSigningError( - "GATEWAY_PROPOSAL_SIGNING_KEY must contain at least 32 bytes" - ) - return encoded - - -def _b64encode(value: bytes) -> str: - return base64.urlsafe_b64encode(value).decode("ascii").rstrip("=") - - -def _b64decode(value: str) -> bytes: - try: - return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4)) - except Exception as exc: - raise ProposalSigningError("proposal payload is not valid base64url") from exc - - -def _canonical(value: Any) -> bytes: - return json.dumps( - value, - sort_keys=True, - separators=(",", ":"), - ensure_ascii=False, - ).encode("utf-8") - - -def append_proposal_marker( - content: str, - proposal: dict[str, Any], - *, - session_id: str, - source_prompt: str, - signing_key: str | None = None, - now: int | None = None, - ttl_seconds: int = DEFAULT_PROPOSAL_TTL_SECONDS, -) -> str: - """Append a signed, non-rendered continuation payload to clarification text.""" - - issued_at = int(time.time()) if now is None else int(now) - envelope = { - "version": PROPOSAL_MARKER_VERSION, - "session_id": session_id, - "issued_at": issued_at, - "expires_at": issued_at + ttl_seconds, - "source_prompt_sha256": hashlib.sha256( - source_prompt.encode("utf-8") - ).hexdigest(), - "proposal": proposal, - } - payload = _b64encode(_canonical(envelope)) - signed = f"v1.{payload}".encode("ascii") - signature = _b64encode(hmac.new(_key(signing_key), signed, hashlib.sha256).digest()) - return content + f"\n\n" - - -def extract_proposal_marker(content: str) -> str: - matches = _MARKER_RE.findall(content or "") - if not matches: - raise ProposalSigningError("no signed proposal marker found") - return matches[-1] - - -def decode_proposal_marker( - marker: str, - *, - session_id: str, - signing_key: str | None = None, - now: int | None = None, -) -> dict[str, Any]: - """Verify and decode one marker without trusting client-provided history.""" - - parts = marker.split(":") - if len(parts) != 4 or parts[:2] != ["pe-proposal", "v1"]: - raise ProposalSigningError("unsupported proposal marker format") - payload, supplied_signature = parts[2:] - signed = f"v1.{payload}".encode("ascii") - expected = _b64encode(hmac.new(_key(signing_key), signed, hashlib.sha256).digest()) - if not hmac.compare_digest(supplied_signature, expected): - raise ProposalSigningError("proposal signature is invalid") - try: - envelope = json.loads(_b64decode(payload)) - except (json.JSONDecodeError, UnicodeDecodeError) as exc: - raise ProposalSigningError("proposal payload is not valid JSON") from exc - if not isinstance(envelope, dict) or envelope.get("version") != PROPOSAL_MARKER_VERSION: - raise ProposalSigningError("unsupported proposal payload version") - if envelope.get("session_id") != session_id: - raise ProposalSigningError("proposal belongs to another session") - current_time = int(time.time()) if now is None else int(now) - expires_at = envelope.get("expires_at") - if not isinstance(expires_at, int) or current_time > expires_at: - raise ProposalExpiredError(envelope) - if not isinstance(envelope.get("proposal"), dict): - raise ProposalSigningError("proposal content is invalid") - return envelope - - -def strip_proposal_markers(content: str) -> str: - return _MARKER_RE.sub("", content or "").rstrip() - - -def _plain(value: Any) -> Any: - if is_dataclass(value): - return {key: _plain(item) for key, item in asdict(value).items()} - if isinstance(value, dict): - return {str(key): _plain(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_plain(item) for item in value] - return value - - -def proposal_payload_from_verdict(verdict: Any) -> dict[str, Any]: - """Serialize the exact low-confidence construction needed for resumption.""" - - assessment = getattr(verdict, "reform_assessment", None) - intent = getattr(verdict, "reform_intent", None) - if assessment is None or assessment.reform is None or intent is None: - raise ProposalSigningError("verdict has no resumable reform proposal") - return { - "tool": verdict.tool, - "slots": [_plain(slot) for slot in verdict.slots], - "reform_intent": _plain(intent), - "assessment": _plain(assessment), - } - - -def _message_text(message: dict[str, Any]) -> str: - content = message.get("content", "") - if isinstance(content, str): - return content - if isinstance(content, list): - return "\n".join( - str(item.get("text", "")) - for item in content - if isinstance(item, dict) and item.get("type") == "text" - ) - return "" - - -def _active_proposal_message( - conversation: list[dict[str, Any]], -) -> tuple[int, str] | None: - """Return the proposal on the latest assistant turn, if there is one. - - Older markers remain in stateless chat history after they are consumed. - Looking only at the latest assistant message prevents a later ordinary - follow-up from reopening an already accepted proposal. - """ - - for index in range(len(conversation) - 1, -1, -1): - message = conversation[index] - if message.get("role") != "assistant": - continue - text = _message_text(message) - return (index, text) if _MARKER_RE.search(text) else None - return None - - -def _source_prompt(conversation: list[dict[str, Any]], before: int) -> str | None: - for index in range(before - 1, -1, -1): - if conversation[index].get("role") == "user": - prompt = _message_text(conversation[index]).strip() - if prompt: - return prompt - return None - - -def _latest_user(conversation: list[dict[str, Any]], after: int) -> str | None: - for message in reversed(conversation[after + 1 :]): - if message.get("role") == "user": - value = _message_text(message).strip() - if value: - return value - return None - - -_AFFIRMATIONS = { - "yes", - "yes please", - "correct", - "that's right", - "that is right", - "go ahead", - "run it", - "do it", - "proceed", -} -_ORDINALS = {"first": 0, "second": 1, "third": 2} - - -def _normalise_reply(value: str) -> str: - return " ".join(re.sub(r"[^a-z0-9' ]+", " ", value.casefold()).split()) - - -def _binding(value: dict[str, Any]): - from gateway.assessment import ValidatedParameterBinding - - return ValidatedParameterBinding( - parameter_path=value["parameter_path"], - label=value["label"], - catalogue_evidence=value.get("catalogue_evidence", ""), - ) - - -def _assessment(value: dict[str, Any], alternative_index: int | None = None): - from gateway.assessment import ReformAlternative, ReformAssessment - - alternatives = tuple( - ReformAlternative( - summary=item["summary"], - parameter_bindings=tuple( - _binding(binding) for binding in item["parameter_bindings"] - ), - reform=dict(item["reform"]), - ) - for item in value.get("alternatives", []) - ) - if alternative_index is not None: - if alternative_index >= len(alternatives): - raise ProposalSigningError("selected proposal alternative does not exist") - selected = alternatives[alternative_index] - return ReformAssessment( - reform=dict(selected.reform), - summary=selected.summary, - confidence=value["confidence"], - parameter_bindings=selected.parameter_bindings, - alternatives=alternatives, - search_queries=tuple(value.get("search_queries", [])), - catalogue_version=value["catalogue_version"], - ) - return ReformAssessment( - reform=dict(value["reform"]), - summary=value.get("summary"), - confidence=value["confidence"], - parameter_bindings=tuple( - _binding(binding) for binding in value["parameter_bindings"] - ), - alternatives=alternatives, - search_queries=tuple(value.get("search_queries", [])), - catalogue_version=value["catalogue_version"], - ) - - -def _alternative_from_reply(reply: str, assessment: dict[str, Any]) -> int | None: - normalized = _normalise_reply(reply) - for ordinal, index in _ORDINALS.items(): - if re.search(rf"\b{ordinal}(?: one| option)?\b", normalized): - return index - for index, alternative in enumerate(assessment.get("alternatives", [])): - labels = [ - binding.get("label", "") - for binding in alternative.get("parameter_bindings", []) - ] - if labels and all(label.casefold() in reply.casefold() for label in labels): - return index - return None - - -def _rebuild_ready_verdict( - proposal: dict[str, Any], - *, - source_prompt: str, - alternative_index: int | None, -): - from gateway.execution import build_execution_plan - from gateway.intent import ReformIntent - from gateway.policy import SlotFact - from gateway.runtime import GatewayVerdict - - intent = ReformIntent(**proposal["reform_intent"]) - slots = [SlotFact(**item) for item in proposal.get("slots", [])] - assessment = _assessment(proposal["assessment"], alternative_index) - execution = build_execution_plan( - proposal.get("tool"), - slots, - intent, - source_prompt, - assessment, - ) - return GatewayVerdict( - outcome="ready", - route="compute", - tool=proposal.get("tool"), - slots=slots, - reform_intent=intent, - reform_assessment=assessment, - execution_plan=execution, - proposal_resumed=True, - ) - - -def resume_gateway_proposal( - conversation: list[dict[str, Any]], - *, - session_id: str, - signing_key: str | None = None, -): - """Resume an exact proposal, reassess a correction, or return ``None``.""" - - found = _active_proposal_message(conversation) - if found is None: - return None - assistant_index, assistant_text = found - reply = _latest_user(conversation, assistant_index) - source_prompt = _source_prompt(conversation, assistant_index) - if reply is None or source_prompt is None: - raise ProposalSigningError("proposal history is incomplete") - marker = extract_proposal_marker(assistant_text) - try: - envelope = decode_proposal_marker( - marker, - session_id=session_id, - signing_key=signing_key, - ) - except ProposalExpiredError as exc: - envelope = exc.envelope - from gateway.runtime import run_gateway - - return run_gateway(source_prompt) - expected_hash = hashlib.sha256(source_prompt.encode("utf-8")).hexdigest() - if envelope.get("source_prompt_sha256") != expected_hash: - raise ProposalSigningError("proposal source prompt does not match history") - proposal = envelope["proposal"] - - from gateway.assessment import current_catalogue_version - - if proposal["assessment"].get("catalogue_version") != current_catalogue_version(): - from gateway.runtime import run_gateway - - return run_gateway(source_prompt) - - normalized = _normalise_reply(reply) - if normalized in _AFFIRMATIONS: - return _rebuild_ready_verdict( - proposal, - source_prompt=source_prompt, - alternative_index=None, - ) - alternative_index = _alternative_from_reply(reply, proposal["assessment"]) - if alternative_index is not None: - return _rebuild_ready_verdict( - proposal, - source_prompt=source_prompt, - alternative_index=alternative_index, - ) - - from gateway.intent import ReformIntent, reform_intent_from_prompt - from gateway.policy import GatingReason, SlotFact - from gateway.runtime import GatewayVerdict, run_gateway - - if reform_intent_from_prompt(reply) is not None: - output = next( - ( - slot.get("value") - for slot in proposal.get("slots", []) - if slot.get("kind") == "output" - ), - None, - ) - suffix = f" Requested output: {output}." if output else "" - return run_gateway(reply + suffix) - if normalized in {"no", "no thanks", "incorrect", "that's wrong", "that is wrong"}: - return GatewayVerdict( - outcome="needs_plan", - route="lightweight", - gating_reasons=[GatingReason("missing_reform", "reform")], - ) - return GatewayVerdict( - outcome="needs_plan", - route="lightweight", - tool=proposal.get("tool"), - slots=[SlotFact(**item) for item in proposal.get("slots", [])], - gating_reasons=[GatingReason("confirm_reform", "reform")], - reform_intent=ReformIntent(**proposal["reform_intent"]), - reform_assessment=_assessment(proposal["assessment"]), - ) - - -def strip_proposal_markers_from_conversation( - conversation: list[dict[str, Any]], -) -> list[dict[str, Any]]: - """Return a copy with hidden metadata removed before any model call.""" - - cleaned: list[dict[str, Any]] = [] - for message in conversation: - copy = dict(message) - content = copy.get("content") - if isinstance(content, str): - copy["content"] = strip_proposal_markers(content) - elif isinstance(content, list): - blocks = [] - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - block = dict(block) - block["text"] = strip_proposal_markers(str(block.get("text", ""))) - blocks.append(block) - copy["content"] = blocks - cleaned.append(copy) - return cleaned diff --git a/backend/gateway/runtime.py b/backend/gateway/runtime.py deleted file mode 100644 index 5336cd4a..00000000 --- a/backend/gateway/runtime.py +++ /dev/null @@ -1,850 +0,0 @@ -"""Gateway runtime: turn a user message into a grounded execution plan. - -One cheap forced-tool call asks a fast model to fill an execution plan (which -tool, which slots, each tagged with a grounding `source`). A second call is -allowed only when authoritative catalogue evidence can help recover a missing -tool. The deterministic `gate()` in `gateway.policy` maps each plan to one of -five outcomes. Only `ready` runs the full compute loop; `needs_plan` is rendered -deterministically and the other non-ready outcomes use the lean writer path. -Classifier transport/parse errors retain the existing `ready`/compute fallback, -but catalogue or reform-assessment failures are terminal so unverified reform -JSON can never reach simulation. - -Self-contained (builds its own sync client + system prompt) so the eval harness -can import and call `run_gateway` directly, mirroring the old `_route_scope`. -""" - -from __future__ import annotations - -import json -import logging -import os -import re -from dataclasses import dataclass, field, replace -from typing import List, Optional - -from config import DEFAULT_FAST_MODEL, DEFAULT_TEMPERATURE, get_sync_client -from gateway.assessment import ( - AUTO_EXECUTE_REFORM_CONFIDENCE, - GatewayCatalogueUnavailable, - ReformAssessment, - ReformAssessmentError, - assess_reform_with_catalogue, -) -from gateway.catalogue import ( - MAX_CATALOGUE_QUERIES, - CatalogueEvidence, - CatalogueQuery, - resolve_catalogue_queries, -) -from gateway.intent import ( - ReformIntent, - output_from_prompt, - reform_intent_from_prompt, - upsert_output_slot, - upsert_prompt_year, -) -from gateway.execution import GatewayExecutionPlan, build_execution_plan -from gateway.policy import ( - CapabilityDecision, - DomainDecision, - GatingReason, - OUTPUT_VOCAB, - SlotFact, - complete_slots, - gate, - normalise_slot_grounding, -) -from prompts import ( - DEFAULT_SCOPE_DESCRIPTOR, - GATEWAY_CATALOGUE_RECOVERY_DIRECTIVE, - GATEWAY_IRRELEVANT_DIRECTIVE, - GATEWAY_OUT_OF_SCOPE_DIRECTIVE, - GATEWAY_PARTIAL_CATALOGUE_DIRECTIVE, - GATEWAY_PARTIAL_DIRECTIVE, - gateway_system, -) -from tools.definitions import DEFAULT_SIMULATION_YEAR, TOOL_DEFINITIONS - -logger = logging.getLogger(__name__) - -GATEWAY_MODEL = os.environ.get("POLICYENGINE_CHAT_GATEWAY_MODEL", DEFAULT_FAST_MODEL) -GATEWAY_MAX_TOKENS = int(os.environ.get("POLICYENGINE_CHAT_GATEWAY_MAX_TOKENS", "1024")) -MAX_UNMODELLABLE_OUTPUTS = 4 - -# A refusal or partial result requires evidence of an effect the direct static -# microsimulation contract genuinely does not estimate. Exact quotation alone -# is insufficient: otherwise a classifier can turn ordinary uncertainty such -# as "compare these reforms" into a false capability refusal. -_UNMODELLABLE_EFFECT_RE = re.compile( - r"\b(?:inflation|gdp|gross domestic product|macroeconomic|economic growth|" - r"behaviou?ral(?: responses?| effects?)?|employment effects?|jobs? effects?|" - r"labou?r supply|stop working|work less|leave work|market reactions?|" - r"general equilibrium)\b", - re.I, -) -_BEHAVIOURAL_EFFECT_RE = re.compile( - r"\b(?:behaviou?ral|employment|jobs?|labou?r supply|stop working|work less|" - r"leave work)\b", - re.I, -) - -_TOOL_NAMES = [t["name"] for t in TOOL_DEFINITIONS] - - -def _build_tool_summary() -> str: - """One line per tool (name — purpose; required params), derived from the - tool schemas so it can't drift.""" - lines = [] - for t in TOOL_DEFINITIONS: - required = t.get("input_schema", {}).get("required", []) or [] - purpose = (t.get("description") or "").strip().split(". ")[0].rstrip(".") - req = ", ".join(required) if required else "none" - lines.append(f"- {t['name']} — {purpose}. Required: {req}.") - return "\n".join(lines) - - -TOOL_SUMMARY = _build_tool_summary() -SCOPE_DESCRIPTOR = DEFAULT_SCOPE_DESCRIPTOR -GATEWAY_SYSTEM = gateway_system( - SCOPE_DESCRIPTOR, TOOL_SUMMARY, ", ".join(OUTPUT_VOCAB), DEFAULT_SIMULATION_YEAR -) - - -@dataclass -class GatewayVerdict: - outcome: str - route: str # "compute" if outcome == "ready" else "lightweight" - tool: Optional[str] = None - slots: List[SlotFact] = field(default_factory=list) - gating_reasons: List[GatingReason] = field(default_factory=list) - unmodellable_outputs: List[str] = field(default_factory=list) - catalogue_evidence: CatalogueEvidence | None = None - domain: DomainDecision = field(default_factory=DomainDecision) - capability: CapabilityDecision = field(default_factory=CapabilityDecision) - reform_intent: ReformIntent | None = None - reform_assessment: ReformAssessment | None = None - catalogue_recovery_used: bool = False - execution_plan: GatewayExecutionPlan | None = None - proposal_resumed: bool = False - - @property - def gating_slots(self) -> List[str]: - return [reason.slot for reason in self.gating_reasons] - - -def _fail_safe() -> GatewayVerdict: - """The safe default: behave exactly like today (full compute background).""" - return GatewayVerdict(outcome="ready", route="compute") - - -# Forced-use tool that carries the structured plan. Local to the gateway — must -# NOT be added to TOOL_DEFINITIONS or it would leak into the compute loop. -_EMIT_PLAN_TOOL = { - "name": "emit_plan", - "description": "Emit the structured execution plan for the user's message.", - "input_schema": { - "type": "object", - "properties": { - "domain_status": { - "type": "string", - "enum": [ - "uk_or_unspecified", - "explicit_non_uk", - "unrelated", - ], - "description": ( - "Whether the request is UK tax-benefit work, explicitly " - "non-UK, or unrelated." - ), - }, - "domain_evidence": { - "type": "string", - "maxLength": 300, - "description": ( - "An exact quote supporting explicit_non_uk or unrelated; " - "omit for uk_or_unspecified." - ), - }, - "capability_status": { - "type": "string", - "enum": [ - "supported", - "catalogue_uncertain", - "explicitly_unmodellable", - ], - "description": ( - "Whether the requested work is supported, needs catalogue " - "confirmation, or explicitly asks only for an unmodellable " - "effect." - ), - }, - "capability_evidence": { - "type": "string", - "maxLength": 300, - "description": ( - "An exact quote supporting catalogue_uncertain or " - "explicitly_unmodellable; omit for supported." - ), - }, - "tool": { - "type": "string", - "enum": _TOOL_NAMES + ["none"], - "description": "Best-fitting tool for the modelled part, or 'none'.", - }, - "slots": { - "type": "array", - "items": { - "type": "object", - "properties": { - "name": {"type": "string"}, - "kind": {"type": "string", "enum": ["tool_input", "output"]}, - "value": {"type": "string"}, - "source": { - "type": "string", - "enum": ["prompt", "default", "assumed"], - }, - }, - "required": ["name", "kind", "source"], - }, - }, - "unmodellable_outputs": { - "type": "array", - "maxItems": MAX_UNMODELLABLE_OUTPUTS, - "description": ( - "Outputs explicitly requested by the user that the tool chain " - "cannot calculate. Every item must cite an exact quote from " - "the user's message that requests that output." - ), - "items": { - "type": "object", - "properties": { - "name": { - "type": "string", - "maxLength": 100, - "description": "Concise name of the unmodellable output.", - }, - "evidence": { - "type": "string", - "maxLength": 300, - "description": ( - "A short exact quote from the user's message that " - "explicitly requests this output." - ), - }, - }, - "required": ["name", "evidence"], - }, - }, - "catalogue_queries": { - "type": "array", - "maxItems": MAX_CATALOGUE_QUERIES, - "description": ( - "Short policyengine.py catalogue searches for named reform " - "measures or variable concepts. Every query must cite an " - "exact quote from the user's message containing the query. " - "Use an empty list when no catalogue concept is named." - ), - "items": { - "type": "object", - "properties": { - "kind": { - "type": "string", - "enum": ["reform_target", "variable"], - }, - "query": {"type": "string"}, - "evidence": { - "type": "string", - "maxLength": 300, - "description": ( - "An exact quote from the user's message that " - "contains this catalogue search term." - ), - }, - }, - "required": ["kind", "query", "evidence"], - }, - }, - "rationale": {"type": "string"}, - }, - "required": [ - "domain_status", - "capability_status", - "tool", - "slots", - "catalogue_queries", - ], - }, -} - - -def _catalogue_queries_from_plan( - plan: dict, - prompt: str, -) -> tuple[CatalogueQuery, ...]: - """Accept only catalogue searches grounded in an exact user quote.""" - - queries: list[CatalogueQuery] = [] - for item in plan.get("catalogue_queries") or []: - if not isinstance(item, dict): - continue - kind = item.get("kind") - selected_tool = plan.get("tool") - if selected_tool in {"get_parameter", "list_reform_targets"}: - kind = "reform_target" - elif selected_tool in {"get_variable", "search_variables"}: - kind = "variable" - query = item.get("query") - if kind not in ("reform_target", "variable") or not isinstance(query, str): - continue - query = query.strip() - evidence = _validated_quote(item.get("evidence"), prompt) - if ( - not query - or evidence is None - or _normalise_evidence_text(query) not in _normalise_evidence_text(evidence) - ): - continue - queries.append(CatalogueQuery(kind, query, evidence)) - return tuple(queries) - - -def _normalise_evidence_text(value: str) -> str: - """Normalise case and whitespace while preserving phrase boundaries.""" - - return " ".join(value.casefold().split()) - - -def _validated_quote(value: object, prompt: str) -> str | None: - """Return a prompt-grounded quote, or ``None`` for invented evidence.""" - - if not isinstance(value, str): - return None - quote = value.strip() - quote_text = _normalise_evidence_text(quote) - if not quote_text or quote_text not in _normalise_evidence_text(prompt): - return None - return quote - - -def _domain_from_plan(plan: dict, prompt: str) -> DomainDecision: - raw = plan.get("domain") - if isinstance(raw, dict): - status = raw.get("status") - raw_evidence = raw.get("evidence") - else: - status = plan.get("domain_status") - raw_evidence = plan.get("domain_evidence") - if status == "uk_or_unspecified": - return DomainDecision() - if status not in ("explicit_non_uk", "unrelated"): - return DomainDecision() - if status == "unrelated" and _UNMODELLABLE_EFFECT_RE.search(prompt): - return DomainDecision() - evidence = _validated_quote(raw_evidence, prompt) - if evidence is None: - return DomainDecision() - return DomainDecision(status=status, evidence=evidence) - - -def _capability_from_plan(plan: dict, prompt: str) -> CapabilityDecision: - raw = plan.get("capability") - if isinstance(raw, dict): - status = raw.get("status") - raw_evidence = raw.get("evidence") - else: - status = plan.get("capability_status") - raw_evidence = plan.get("capability_evidence") - if status == "supported": - return CapabilityDecision() - if status not in ("catalogue_uncertain", "explicitly_unmodellable"): - return CapabilityDecision() - evidence = _validated_quote(raw_evidence, prompt) - if evidence is None: - return CapabilityDecision() - if status == "explicitly_unmodellable" and not _UNMODELLABLE_EFFECT_RE.search( - evidence - ): - return CapabilityDecision() - return CapabilityDecision(status=status, evidence=evidence) - - -def _unmodellable_outputs_from_plan(plan: dict, prompt: str) -> list[str]: - """Accept only limitations backed by an exact phrase from the user. - - The classifier can explain capability boundaries, but it cannot promote a - merely possible behavioural or macroeconomic caveat into a requested output. - Requiring quoted prompt evidence keeps that distinction deterministic. - """ - - prompt_text = _normalise_evidence_text(prompt) - outputs: list[str] = [] - seen: set[str] = set() - raw_outputs = plan.get("unmodellable_outputs") - if not isinstance(raw_outputs, list): - return outputs - - for item in raw_outputs: - if not isinstance(item, dict): - continue - name = item.get("name") - evidence = item.get("evidence") - if not isinstance(name, str) or not isinstance(evidence, str): - continue - name = name.strip() - evidence_text = _normalise_evidence_text(evidence) - key = name.casefold() - if ( - not name - or not evidence_text - or evidence_text not in prompt_text - or not _UNMODELLABLE_EFFECT_RE.search(evidence) - or key in seen - ): - continue - seen.add(key) - outputs.append(name) - if len(outputs) == MAX_UNMODELLABLE_OUTPUTS: - break - return outputs - - -def apply_catalogue_evidence( - verdict: GatewayVerdict, - evidence: CatalogueEvidence, -) -> GatewayVerdict: - """Combine deterministic catalogue evidence with the model-grounded plan. - - Evidence confirms only modelability. It cannot turn an under-specified or - partial request into an executable one, so those outcomes are preserved. - """ - - verdict = replace(verdict, catalogue_evidence=evidence) - if verdict.outcome == "irrelevant": - return verdict - if not evidence.available: - return verdict - if evidence.unresolved_queries: - if verdict.reform_intent is not None: - return verdict - reasons = list(verdict.gating_reasons) - if "model_catalogue" not in verdict.gating_slots: - reasons.append( - GatingReason( - code="catalogue_no_match", - slot="model_catalogue", - evidence=", ".join(query.query for query in evidence.unresolved_queries), - ) - ) - if verdict.outcome in ("needs_plan", "partial"): - return replace(verdict, gating_reasons=reasons) - return replace( - verdict, - outcome="needs_plan", - route="lightweight", - gating_reasons=reasons, - ) - return verdict - - -def _verdict_from_plan( - plan: dict, - prompt: str, - catalogue_evidence: CatalogueEvidence, -) -> GatewayVerdict: - """Build a server-gated verdict from the model's grounded plan. The model's - own outcome is never trusted — the outcome is recomputed by gate().""" - domain = _domain_from_plan(plan, prompt) - capability = _capability_from_plan(plan, prompt) - effect_match = _UNMODELLABLE_EFFECT_RE.search(prompt) - if capability.status == "supported" and effect_match is not None: - capability = CapabilityDecision( - status="explicitly_unmodellable", - evidence=effect_match.group(0), - ) - in_domain = domain.status == "uk_or_unspecified" - output_intent = output_from_prompt(prompt) - reform_intent = reform_intent_from_prompt(prompt) - raw_tool = plan.get("tool") - tool = raw_tool if raw_tool in _TOOL_NAMES else None - policy_model_request = re.search( - r"\bmodel\b.*\b(?:tax|levy|benefit|allowance|credit|rate|threshold|reform|policy)\b", - prompt, - re.I, - ) - if ( - in_domain - and policy_model_request - and ( - tool in {"list_reform_targets", "search_parameters"} - or ( - tool is None - and capability.status != "catalogue_uncertain" - ) - ) - ): - tool = "run_society_simulation" - elif ( - in_domain - and tool is None - and reform_intent is not None - and capability.status != "catalogue_uncertain" - and not ( - output_intent is None and _BEHAVIOURAL_EFFECT_RE.search(prompt) - ) - ): - tool = "run_society_simulation" - - slots: List[SlotFact] = [] - for s in plan.get("slots") or []: - if not isinstance(s, dict) or "name" not in s: - continue - source = s.get("source", "assumed") - if source not in ("prompt", "default", "assumed"): - source = "assumed" - kind = s.get("kind", "tool_input") - if kind not in ("tool_input", "output"): - kind = "tool_input" - value = s.get("value") - slots.append( - SlotFact( - name=str(s["name"]), - source=source, - kind=kind, - value=value if isinstance(value, str) else None, - ) - ) - - if output_intent is not None: - slots = upsert_output_slot(slots, output_intent) - slots = upsert_prompt_year(tool, slots, prompt) - slots = normalise_slot_grounding(tool, slots) - if output_intent is not None: - # Classifier values such as "annual cost" can initially claim prompt - # grounding but fail the closed output vocabulary during normalization. - # Reapply the deterministic prompt intent to that now-assumed slot. - slots = upsert_output_slot(slots, output_intent) - slots = complete_slots(tool, slots) - unmodellable = _unmodellable_outputs_from_plan(plan, prompt) - if effect_match is not None and not unmodellable: - unmodellable.append(effect_match.group(0)) - result = gate( - in_domain, - tool, - slots, - unmodellable, - prompt, - explicitly_unmodellable=(capability.status == "explicitly_unmodellable"), - reform_intent=reform_intent, - ) - verdict = GatewayVerdict( - outcome=result.outcome, - route="compute" if result.outcome == "ready" else "lightweight", - tool=tool, - slots=slots, - gating_reasons=result.gating_reasons, - unmodellable_outputs=unmodellable, - domain=domain, - capability=capability, - reform_intent=reform_intent, - ) - return apply_catalogue_evidence(verdict, catalogue_evidence) - - -def _request_plan(client, last_user_message: str, system: str) -> dict | None: - """Request and extract one forced execution plan from the gateway model.""" - - response = client.messages.create( - model=GATEWAY_MODEL, - max_tokens=GATEWAY_MAX_TOKENS, - temperature=DEFAULT_TEMPERATURE, - system=system, - tools=[_EMIT_PLAN_TOOL], - tool_choice={"type": "tool", "name": "emit_plan"}, - messages=[{"role": "user", "content": last_user_message[:4000]}], - ) - for block in response.content or []: - if ( - getattr(block, "type", None) == "tool_use" - and getattr(block, "name", None) == "emit_plan" - ): - return block.input if isinstance(block.input, dict) else {} - return None - - -def _catalogue_recovery_system(evidence: CatalogueEvidence) -> str: - candidates = "\n".join( - f"- {match.kind}: {match.label} (`{match.identifier}`; {match.match_type})" - for match in evidence.authoritative_matches - ) - return ( - GATEWAY_SYSTEM - + "\n\n" - + GATEWAY_CATALOGUE_RECOVERY_DIRECTIVE - + "\n\nSERVER-VERIFIED CATALOGUE CANDIDATES:\n" - + candidates - ) - - -def _can_recover_with_catalogue(verdict: GatewayVerdict) -> bool: - """True only for a grounded UK capability uncertainty with strong evidence.""" - - evidence = verdict.catalogue_evidence - return bool( - verdict.domain.status == "uk_or_unspecified" - and verdict.capability.status == "catalogue_uncertain" - and verdict.tool is None - and not verdict.unmodellable_outputs - and evidence - and evidence.available - and evidence.authoritative_matches - and not evidence.unresolved_queries - ) - - -_SOCIETY_REFORM_TOOLS = { - "run_society_simulation", - "compute_budgetary_impact", - "compute_program_breakdown", - "compute_decile_impacts", - "compute_winners_losers", - "compute_poverty_metrics", - "compute_inequality_metrics", - "aggregate_result", -} - - -def _assess_ready_reform( - verdict: GatewayVerdict, - prompt: str, - client: object, -) -> GatewayVerdict: - if ( - verdict.outcome != "ready" - or verdict.tool not in _SOCIETY_REFORM_TOOLS - or verdict.reform_intent is None - ): - if verdict.outcome == "ready": - return replace( - verdict, - execution_plan=build_execution_plan( - verdict.tool, - verdict.slots, - verdict.reform_intent, - prompt, - verdict.reform_assessment, - ), - ) - return verdict - try: - assessment = assess_reform_with_catalogue( - prompt, - verdict.reform_intent, - client=client, - ) - except (GatewayCatalogueUnavailable, ReformAssessmentError) as exc: - exc.gateway_verdict = verdict - raise - verdict = replace(verdict, reform_assessment=assessment) - if assessment.reform is None: - return replace( - verdict, - outcome="needs_plan", - route="lightweight", - gating_reasons=[ - GatingReason( - code="catalogue_no_match", - slot="reform", - evidence=verdict.reform_intent.evidence, - ) - ], - ) - if assessment.confidence < AUTO_EXECUTE_REFORM_CONFIDENCE: - return replace( - verdict, - outcome="needs_plan", - route="lightweight", - gating_reasons=[ - GatingReason( - code="confirm_reform", - slot="reform", - options=tuple( - alternative.summary - for alternative in assessment.alternatives - ), - evidence=assessment.summary, - ) - ], - ) - return replace( - verdict, - execution_plan=build_execution_plan( - verdict.tool, - verdict.slots, - verdict.reform_intent, - prompt, - assessment, - ), - ) - - -def run_gateway(last_user_message: str) -> GatewayVerdict: - """Ground a plan, optionally recover once, then return the server gate. - - Fail-safe to ready/compute on empty input, classifier API errors, a missing - plan block, or an unparseable plan. Catalogue and exact-reform assessment - failures propagate as terminal errors. - """ - if not last_user_message or not last_user_message.strip(): - return _fail_safe() - try: - client = get_sync_client() - plan = _request_plan(client, last_user_message, GATEWAY_SYSTEM) - # No plan block, an empty plan, or one missing the routing decision is a - # parse failure — fall back to compute. A real refusal (`tool: "none"`, - # grounded negative domain/capability) is a well-formed plan and is NOT - # caught here, so a - # degenerate response can never masquerade as an out_of_scope refusal. - if not plan or "tool" not in plan: - return _fail_safe() - catalogue_evidence = resolve_catalogue_queries( - _catalogue_queries_from_plan(plan, last_user_message) - ) - verdict = _verdict_from_plan(plan, last_user_message, catalogue_evidence) - if not _can_recover_with_catalogue(verdict): - return _assess_ready_reform(verdict, last_user_message, client) - - verdict = replace(verdict, catalogue_recovery_used=True) - recovery_plan = _request_plan( - client, - last_user_message, - _catalogue_recovery_system(catalogue_evidence), - ) - if not recovery_plan or "tool" not in recovery_plan: - return verdict - recovered = _verdict_from_plan( - recovery_plan, - last_user_message, - catalogue_evidence, - ) - recovered = replace(recovered, catalogue_recovery_used=True) - return _assess_ready_reform(recovered, last_user_message, client) - except (GatewayCatalogueUnavailable, ReformAssessmentError): - raise - except Exception as e: # noqa: BLE001 — any failure falls back to full compute - logger.warning(f"[GATEWAY] failed; defaulting to ready/compute: {e}") - return _fail_safe() - - -_WRITER_DIRECTIVES = { - "irrelevant": GATEWAY_IRRELEVANT_DIRECTIVE, - "out_of_scope": GATEWAY_OUT_OF_SCOPE_DIRECTIVE, - "partial": GATEWAY_PARTIAL_DIRECTIVE, -} - - -def gateway_writer_directive(verdict: GatewayVerdict) -> str: - """Return model-writer instructions for non-clarification outcomes.""" - evidence = verdict.catalogue_evidence - if verdict.outcome == "partial" and evidence and evidence.unresolved_queries: - directive = GATEWAY_PARTIAL_CATALOGUE_DIRECTIVE - else: - directive = _WRITER_DIRECTIVES.get(verdict.outcome) - if directive is None: - return "" - parts = [directive] - if verdict.outcome == "partial" and verdict.unmodellable_outputs: - parts.append("Cannot model: " + ", ".join(verdict.unmodellable_outputs) + ".") - if evidence and evidence.unresolved_queries: - queries = ", ".join(query.query for query in evidence.unresolved_queries) - parts.append( - "The current model catalogue did not resolve: " - + queries - + ". Ask what supported policy measure or variable the user means; " - "do not say that it is unmodelled." - ) - return "\n\n".join(parts) - - -def serialise_plan_for_system(verdict: GatewayVerdict) -> str: - """Compact plan text injected into the compute system blocks for `ready`, so - the heavy model starts from the resolved tool + grounded args.""" - execution = verdict.execution_plan - if execution is not None: - lines = ["GATEWAY EXECUTION PLAN (validated opening-turn contract):"] - if execution.prerequisites: - ordered = [*execution.prerequisites, execution.target_tool] - lines.append("Required tool order: " + " -> ".join(ordered) + ".") - elif execution.target_tool: - lines.append(f"Requested target analysis: {execution.target_tool}.") - if execution.inputs: - lines.append( - "Resolved inputs: " - + "; ".join( - f"{item.name}={item.value} ({item.source})" - for item in execution.inputs - ) - + "." - ) - if verdict.reform_intent is not None: - lines.append( - "Exact user reform wording: " + verdict.reform_intent.evidence + "." - ) - if execution.approved_reform is not None: - lines.append( - "Use this exact validated reform JSON without modification: " - + json.dumps(execution.approved_reform, sort_keys=True) - + "." - ) - if execution.parameter_bindings: - lines.append("Validated parameter bindings:") - lines.extend( - f"- {binding.label}: `{binding.parameter_path}`" - for binding in execution.parameter_bindings - ) - if execution.conventions: - lines.append( - "Product conventions: " - + "; ".join( - f"{item.name}={item.value}" for item in execution.conventions - ) - + "." - ) - if "run_society_simulation" in execution.prerequisites: - lines.append( - "Runtime handoff: pass the result_id returned by " - "run_society_simulation as simulation_id to the target derivative." - ) - return "\n".join(lines) - - grounded = [ - f"{s.name}={s.value}" - for s in verdict.slots - if s.value and s.source in ("prompt", "default") - ] - lines = [] - if verdict.tool is not None: - lines.append( - f"GATEWAY PLAN (pre-resolved by a routing pass): tool={verdict.tool}." - ) - if grounded: - lines.append("Resolved inputs: " + "; ".join(grounded) + ".") - evidence = verdict.catalogue_evidence - if evidence and evidence.authoritative_matches: - lines.append( - "MODEL CATALOGUE EVIDENCE (verified current policyengine.py candidates):" - ) - lines.extend( - f"- {match.kind}: {match.label} (`{match.identifier}`)" - for match in evidence.authoritative_matches - ) - lines.append( - "Treat these as discovery candidates, not a resolution of user intent. " - "Use them internally and ask a concise clarification where needed." - ) - if not lines: - return "" - lines.append( - "Treat this as a starting point and verify it against the user's message." - ) - return "\n".join(lines) diff --git a/backend/gateway/trace.py b/backend/gateway/trace.py deleted file mode 100644 index 9ce31408..00000000 --- a/backend/gateway/trace.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Internal, structured observability for completed gateway decisions.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from typing import Any, TYPE_CHECKING - -from gateway.assessment import REFORM_RESOLVER_MODEL - -if TYPE_CHECKING: - from gateway.runtime import GatewayVerdict - - -@dataclass(frozen=True, slots=True) -class GatewayTraceSlot: - name: str - kind: str - source: str - value: str | None = None - - -@dataclass(frozen=True, slots=True) -class GatewayTraceReason: - code: str - slot: str - options: tuple[str, ...] = () - evidence: str | None = None - - -@dataclass(frozen=True, slots=True) -class GatewayTraceBinding: - parameter_path: str - label: str - catalogue_evidence: str - - -@dataclass(frozen=True, slots=True) -class GatewayTraceAlternative: - summary: str - parameter_bindings: tuple[GatewayTraceBinding, ...] - reform: dict[str, Any] - - -@dataclass(frozen=True, slots=True) -class GatewayTrace: - selected_tool: str | None = None - target_tool: str | None = None - slots: tuple[GatewayTraceSlot, ...] = () - gating_reasons: tuple[GatewayTraceReason, ...] = () - defaults_applied: dict[str, Any] = field(default_factory=dict) - reform_confidence: int | None = None - reform_summary: str | None = None - reform_search_queries: tuple[str, ...] = () - catalogue_version: str | None = None - resolver_model: str | None = None - parameter_bindings: tuple[GatewayTraceBinding, ...] = () - alternatives: tuple[GatewayTraceAlternative, ...] = () - catalogue_recovery_used: bool = False - proposal_resumed: bool = False - - -def _trace_value(value: str) -> Any: - """Retain text defaults while representing JSON scalar defaults naturally.""" - - try: - parsed = json.loads(value) - except (json.JSONDecodeError, TypeError): - return value - return parsed if isinstance(parsed, (int, float, bool)) or parsed is None else value - - -def _binding(value: Any) -> GatewayTraceBinding: - return GatewayTraceBinding( - parameter_path=value.parameter_path, - label=value.label, - catalogue_evidence=value.catalogue_evidence, - ) - - -def gateway_trace_from_verdict( - verdict: GatewayVerdict | None, -) -> GatewayTrace | None: - """Project a verdict into a stable internal/eval-safe trace.""" - - if verdict is None: - return None - execution = verdict.execution_plan - assessment = verdict.reform_assessment - defaults = { - slot.name: _trace_value(slot.value) - for slot in verdict.slots - if slot.source == "default" and slot.value is not None - } - if execution is not None: - defaults.update( - { - item.name: _trace_value(item.value) - for item in execution.inputs - if item.source == "default" - } - ) - return GatewayTrace( - selected_tool=verdict.tool, - target_tool=execution.target_tool if execution is not None else None, - slots=tuple( - GatewayTraceSlot( - name=slot.name, - kind=slot.kind, - source=slot.source, - value=slot.value, - ) - for slot in verdict.slots - ), - gating_reasons=tuple( - GatewayTraceReason( - code=reason.code, - slot=reason.slot, - options=tuple(reason.options), - evidence=reason.evidence, - ) - for reason in verdict.gating_reasons - ), - defaults_applied=defaults, - reform_confidence=(assessment.confidence if assessment is not None else None), - reform_summary=(assessment.summary if assessment is not None else None), - reform_search_queries=( - tuple(assessment.search_queries) if assessment is not None else () - ), - catalogue_version=( - assessment.catalogue_version if assessment is not None else None - ), - resolver_model=REFORM_RESOLVER_MODEL if assessment is not None else None, - parameter_bindings=( - tuple(_binding(value) for value in assessment.parameter_bindings) - if assessment is not None - else () - ), - alternatives=( - tuple( - GatewayTraceAlternative( - summary=alternative.summary, - parameter_bindings=tuple( - _binding(value) - for value in alternative.parameter_bindings - ), - reform=dict(alternative.reform), - ) - for alternative in assessment.alternatives - ) - if assessment is not None - else () - ), - catalogue_recovery_used=verdict.catalogue_recovery_used, - proposal_resumed=verdict.proposal_resumed, - ) diff --git a/backend/migrations/README b/backend/migrations/README new file mode 100644 index 00000000..5d2a0724 --- /dev/null +++ b/backend/migrations/README @@ -0,0 +1,3 @@ +Alembic owns only the SQLModel conversation, typed-context, and capability tables listed in +docs/engineering/skills/database-migrations.md. Revision files must be created +with the Alembic CLI's --autogenerate option. diff --git a/backend/migrations/env.py b/backend/migrations/env.py new file mode 100644 index 00000000..3412171e --- /dev/null +++ b/backend/migrations/env.py @@ -0,0 +1,92 @@ +"""Alembic environment for SQLModel-owned conversation and capability tables.""" + +import os +from logging.config import fileConfig + +from alembic import context +from sqlalchemy import create_engine, pool +from sqlalchemy.engine import make_url +from sqlmodel import SQLModel +from sqlmodel.sql.sqltypes import AutoString + +from conversations.models import ChatConversation # noqa: F401 +import persistence.rows # noqa: F401,E402 + +config = context.config + +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +target_metadata = SQLModel.metadata +managed_table_names = frozenset(target_metadata.tables) + + +def _database_url() -> str: + url = os.environ.get("ALEMBIC_DATABASE_URL") + if not url: + raise RuntimeError( + "ALEMBIC_DATABASE_URL is required for database migrations" + ) + if make_url(url).get_backend_name() != "postgresql": + raise RuntimeError( + "Alembic migrations require a disposable or deployed PostgreSQL database" + ) + return url + + +def _include_object( + object_, name: str | None, type_: str, reflected: bool, compare_to +) -> bool: + del compare_to + if type_ == "table": + return name in managed_table_names + table = getattr(object_, "table", None) + return table is None or table.name in managed_table_names + + +def _render_item(type_: str, object_, _autogen_context): + if type_ == "type" and isinstance(object_, AutoString): + return "sa.String()" + return False + + +def _configure(**kwargs) -> None: + context.configure( + target_metadata=target_metadata, + include_object=_include_object, + render_item=_render_item, + compare_type=True, + compare_server_default=True, + **kwargs, + ) + + +def run_migrations_offline() -> None: + """Run migrations without opening a database connection.""" + _configure( + url=_database_url(), + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations against the configured PostgreSQL database.""" + connectable = create_engine(_database_url(), poolclass=pool.NullPool) + + with connectable.connect() as connection: + _configure(connection=connection) + + with context.begin_transaction(): + context.run_migrations() + + connectable.dispose() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/backend/migrations/script.py.mako b/backend/migrations/script.py.mako new file mode 100644 index 00000000..11016301 --- /dev/null +++ b/backend/migrations/script.py.mako @@ -0,0 +1,28 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + """Upgrade schema.""" + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + """Downgrade schema.""" + ${downgrades if downgrades else "pass"} diff --git a/backend/migrations/versions/0001_pre_branch_conversation_schema_baseline.py b/backend/migrations/versions/0001_pre_branch_conversation_schema_baseline.py new file mode 100644 index 00000000..39ad9f75 --- /dev/null +++ b/backend/migrations/versions/0001_pre_branch_conversation_schema_baseline.py @@ -0,0 +1,47 @@ +"""pre-branch conversation schema baseline + +Revision ID: 0001 +Revises: +Create Date: 2026-08-26 14:08:18.772967 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '0001' +down_revision: Union[str, Sequence[str], None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('chat_conversations', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('session_id', sa.String(), nullable=False), + sa.Column('title', sa.String(), nullable=False), + sa.Column('messages', sa.String(), nullable=False), + sa.Column('user_id', sa.String(), nullable=True), + sa.Column('user_email', sa.String(), nullable=True), + sa.Column('share_token', sa.String(), nullable=True), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('id') + ) + op.create_index('idx_chat_conversations_session_id_unique', 'chat_conversations', ['session_id'], unique=True) + op.create_index('idx_chat_conversations_share_token', 'chat_conversations', ['share_token'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('idx_chat_conversations_share_token', table_name='chat_conversations') + op.drop_index('idx_chat_conversations_session_id_unique', table_name='chat_conversations') + op.drop_table('chat_conversations') + # ### end Alembic commands ### diff --git a/backend/migrations/versions/0002_add_capability_persistence_tables.py b/backend/migrations/versions/0002_add_capability_persistence_tables.py new file mode 100644 index 00000000..e18dbe8b --- /dev/null +++ b/backend/migrations/versions/0002_add_capability_persistence_tables.py @@ -0,0 +1,120 @@ +"""add capability persistence tables + +Revision ID: 0002 +Revises: 0001 +Create Date: 2026-08-26 14:08:46.133104 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '0002' +down_revision: Union[str, Sequence[str], None] = '0001' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('capability_artifacts', + sa.Column('artifact_id', sa.String(), nullable=False), + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('artifact_type', sa.String(), nullable=False), + sa.Column('schema_version', sa.String(), nullable=False), + sa.Column('payload_json', sa.String(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('artifact_id') + ) + op.create_index('idx_capability_artifacts_conversation', 'capability_artifacts', ['conversation_id'], unique=False) + op.create_index('idx_capability_artifacts_type', 'capability_artifacts', ['artifact_type'], unique=False) + op.create_table('capability_call_receipts', + sa.Column('call_id', sa.String(), nullable=False), + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('turn_id', sa.String(), nullable=False), + sa.Column('operation_id', sa.String(), nullable=False), + sa.Column('request_fingerprint', sa.String(), nullable=False), + sa.Column('status', sa.String(), nullable=False), + sa.Column('outcome_json', sa.String(), nullable=True), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('call_id') + ) + op.create_index('idx_capability_call_conversation', 'capability_call_receipts', ['conversation_id'], unique=False) + op.create_index('idx_capability_call_operation', 'capability_call_receipts', ['operation_id'], unique=False) + op.create_index('idx_capability_call_turn', 'capability_call_receipts', ['turn_id'], unique=False) + op.create_table('capability_invocation_traces', + sa.Column('invocation_id', sa.String(), nullable=False), + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('turn_id', sa.String(), nullable=False), + sa.Column('parent_invocation_id', sa.String(), nullable=True), + sa.Column('sequence', sa.Integer(), nullable=False), + sa.Column('kind', sa.String(), nullable=False), + sa.Column('identifier', sa.String(), nullable=False), + sa.Column('version', sa.String(), nullable=False), + sa.Column('visibility', sa.String(), nullable=False), + sa.Column('started_at', sa.DateTime(), nullable=False), + sa.Column('completed_at', sa.DateTime(), nullable=True), + sa.Column('duration_ms', sa.Integer(), nullable=True), + sa.Column('status', sa.String(), nullable=False), + sa.Column('summary', sa.String(), nullable=False), + sa.PrimaryKeyConstraint('invocation_id') + ) + op.create_index('idx_capability_trace_conversation_sequence', 'capability_invocation_traces', ['conversation_id', 'sequence'], unique=False) + op.create_index('idx_capability_trace_identifier', 'capability_invocation_traces', ['identifier'], unique=False) + op.create_index('idx_capability_trace_parent', 'capability_invocation_traces', ['parent_invocation_id'], unique=False) + op.create_index('idx_capability_trace_turn', 'capability_invocation_traces', ['turn_id'], unique=False) + op.create_table('capability_turn_receipts', + sa.Column('turn_id', sa.String(), nullable=False), + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('request_fingerprint', sa.String(), nullable=False), + sa.Column('status', sa.String(), nullable=False), + sa.Column('public_outcome_json', sa.String(), nullable=True), + sa.Column('billing_claimed', sa.Boolean(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('turn_id') + ) + op.create_index('idx_capability_turn_conversation', 'capability_turn_receipts', ['conversation_id'], unique=False) + op.create_table('waiting_capability_invocations', + sa.Column('invocation_id', sa.String(), nullable=False), + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('capability_id', sa.String(), nullable=False), + sa.Column('capability_version', sa.String(), nullable=False), + sa.Column('input_schema_version', sa.String(), nullable=False), + sa.Column('partial_input_json', sa.String(), nullable=False), + sa.Column('source_turn_id', sa.String(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('invocation_id') + ) + op.create_index('idx_waiting_capability_conversation', 'waiting_capability_invocations', ['conversation_id'], unique=False) + op.create_index('idx_waiting_capability_id', 'waiting_capability_invocations', ['capability_id'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('idx_waiting_capability_id', table_name='waiting_capability_invocations') + op.drop_index('idx_waiting_capability_conversation', table_name='waiting_capability_invocations') + op.drop_table('waiting_capability_invocations') + op.drop_index('idx_capability_turn_conversation', table_name='capability_turn_receipts') + op.drop_table('capability_turn_receipts') + op.drop_index('idx_capability_trace_turn', table_name='capability_invocation_traces') + op.drop_index('idx_capability_trace_parent', table_name='capability_invocation_traces') + op.drop_index('idx_capability_trace_identifier', table_name='capability_invocation_traces') + op.drop_index('idx_capability_trace_conversation_sequence', table_name='capability_invocation_traces') + op.drop_table('capability_invocation_traces') + op.drop_index('idx_capability_call_turn', table_name='capability_call_receipts') + op.drop_index('idx_capability_call_operation', table_name='capability_call_receipts') + op.drop_index('idx_capability_call_conversation', table_name='capability_call_receipts') + op.drop_table('capability_call_receipts') + op.drop_index('idx_capability_artifacts_type', table_name='capability_artifacts') + op.drop_index('idx_capability_artifacts_conversation', table_name='capability_artifacts') + op.drop_table('capability_artifacts') + # ### end Alembic commands ### diff --git a/backend/migrations/versions/9526d8c80914_add_conversation_context_persistence.py b/backend/migrations/versions/9526d8c80914_add_conversation_context_persistence.py new file mode 100644 index 00000000..813b7981 --- /dev/null +++ b/backend/migrations/versions/9526d8c80914_add_conversation_context_persistence.py @@ -0,0 +1,42 @@ +"""add conversation context persistence + +Revision ID: 9526d8c80914 +Revises: d97a20592837 +Create Date: 2026-08-27 21:45:35.140597 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = '9526d8c80914' +down_revision: Union[str, Sequence[str], None] = 'd97a20592837' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('conversation_contexts', + sa.Column('conversation_id', sa.String(), nullable=False), + sa.Column('schema_version', sa.String(), nullable=False), + sa.Column('revision', sa.Integer(), nullable=False), + sa.Column('payload_json', sa.String(), nullable=False), + sa.Column('created_at', sa.DateTime(), nullable=False), + sa.Column('updated_at', sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint('conversation_id') + ) + op.create_index('idx_conversation_context_updated', 'conversation_contexts', ['updated_at'], unique=False) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index('idx_conversation_context_updated', table_name='conversation_contexts') + op.drop_table('conversation_contexts') + # ### end Alembic commands ### diff --git a/backend/migrations/versions/d97a20592837_add_invocation_debug_projections.py b/backend/migrations/versions/d97a20592837_add_invocation_debug_projections.py new file mode 100644 index 00000000..26e8b106 --- /dev/null +++ b/backend/migrations/versions/d97a20592837_add_invocation_debug_projections.py @@ -0,0 +1,34 @@ +"""add invocation debug projections + +Revision ID: d97a20592837 +Revises: 0002 +Create Date: 2026-08-26 15:23:18.061675 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'd97a20592837' +down_revision: Union[str, Sequence[str], None] = '0002' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.add_column('capability_invocation_traces', sa.Column('debug_input_json', sa.String(), nullable=True)) + op.add_column('capability_invocation_traces', sa.Column('debug_output_json', sa.String(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column('capability_invocation_traces', 'debug_output_json') + op.drop_column('capability_invocation_traces', 'debug_input_json') + # ### end Alembic commands ### diff --git a/backend/mypy.ini b/backend/mypy.ini new file mode 100644 index 00000000..091574ba --- /dev/null +++ b/backend/mypy.ini @@ -0,0 +1,44 @@ +[mypy] +python_version = 3.13 +mypy_path = backend +strict = True +ignore_missing_imports = True +files = + backend/capabilities/artifacts.py, + backend/capabilities/application.py, + backend/capabilities/compatibility.py, + backend/capabilities/composition.py, + backend/capabilities/context.py, + backend/capabilities/contracts.py, + backend/capabilities/executor.py, + backend/capabilities/household_input.py, + backend/capabilities/input_resolution.py, + backend/capabilities/registry.py, + backend/capabilities/repository.py, + backend/capabilities/tracing.py, + backend/chat/artifact_context.py, + backend/chat/capability_runtime.py, + backend/chat/capability_service.py, + backend/chat/events.py, + backend/chat/model_port.py, + backend/chat/narration.py, + backend/chat/schemas.py, + backend/chat/turn_input.py, + backend/conversation_context/change_pipeline.py, + backend/conversation_context/household_view.py, + backend/conversation_context/engine_projection.py, + backend/conversation_context/models.py, + backend/conversation_context/projection.py, + backend/conversation_context/quantities.py, + backend/conversation_context/reducer.py, + backend/conversation_context/registry.py, + backend/conversation_context/repository.py, + backend/conversation_context/tools.py, + backend/conversation_context/variable_resolution.py, + backend/persistence/context_repository.py, + backend/tools/contracts.py, + backend/tools/registry.py, + backend/tools/typed_models.py + +[mypy-pydantic.*] +follow_untyped_imports = True diff --git a/backend/observability/fastapi.py b/backend/observability/fastapi.py index b8b05102..6145e639 100644 --- a/backend/observability/fastapi.py +++ b/backend/observability/fastapi.py @@ -24,9 +24,6 @@ "cloud_run_configuration", "cloud_run_revision", "cloud_run_service", - "gateway_outcome", - "gateway_route", - "gateway_tool", "google_cloud_project", "modal_app_name", "modal_environment", diff --git a/backend/observability/segments.py b/backend/observability/segments.py index ffc3733c..86307456 100644 --- a/backend/observability/segments.py +++ b/backend/observability/segments.py @@ -10,12 +10,6 @@ class SegmentName(StrEnum): UNKNOWN = UNKNOWN_SEGMENT - GATEWAY_CLASSIFY = "gateway.classify" - GATEWAY_PLAN_SERIALIZE = "gateway.plan_serialize" - MODEL_SELECT = "model.select" - SYSTEM_BUILD = "system.build" - TOOL_SCHEMA_BUILD = "tool_schema.build" - MODEL_ITERATION = "model.iteration" MODEL_STREAM = "model.stream" TOOL_EXECUTE = "tool.execute" BILLING_CHECK_BALANCE = "billing.check_balance" diff --git a/backend/persistence/__init__.py b/backend/persistence/__init__.py new file mode 100644 index 00000000..a0ede39f --- /dev/null +++ b/backend/persistence/__init__.py @@ -0,0 +1,4 @@ +"""SQL persistence adapters for capability-oriented chat state.""" +from persistence.context_repository import SQLConversationContextRepository + +__all__ = ["SQLConversationContextRepository"] diff --git a/backend/persistence/capability_repository.py b/backend/persistence/capability_repository.py new file mode 100644 index 00000000..95943961 --- /dev/null +++ b/backend/persistence/capability_repository.py @@ -0,0 +1,408 @@ +"""SQL repositories for typed artifacts and waiting capability input.""" + +from __future__ import annotations + +import asyncio +from datetime import datetime, timezone +from typing import Mapping, TypeVar, cast + +from pydantic import BaseModel, ValidationError +from sqlmodel import Session, select + +from capabilities.artifacts import ARTIFACT_MODELS, ArtifactBase, TransferableArtifact +from capabilities.repository import WaitingCapabilityInvocation +from conversations.models import get_engine +from persistence.rows import CapabilityArtifactRow, WaitingCapabilityInvocationRow + + +ArtifactT = TypeVar("ArtifactT", bound=ArtifactBase) + + +class InvalidPersistedRecord(ValueError): + pass + + +class PartialInputDefinition(BaseModel): + model_config = {"frozen": True, "arbitrary_types_allowed": True, "extra": "forbid"} + + schema_version: str + model: type[BaseModel] + + +class PartialInputRegistry: + def __init__( + self, + definitions: Mapping[str, PartialInputDefinition] | None = None, + ) -> None: + self._definitions = dict(definitions or {}) + + def register( + self, + capability_id: str, + *, + schema_version: str, + model: type[BaseModel], + ) -> None: + if capability_id in self._definitions: + raise ValueError(f"Duplicate partial-input registration: {capability_id}") + self._definitions[capability_id] = PartialInputDefinition( + schema_version=schema_version, + model=model, + ) + + def get(self, capability_id: str) -> PartialInputDefinition: + try: + return self._definitions[capability_id] + except KeyError as exc: + raise InvalidPersistedRecord( + f"No partial-input model is registered for {capability_id}." + ) from exc + + +class SQLConversationCapabilityRepository: + def __init__( + self, + *, + engine=None, + artifact_models: Mapping[str, type[ArtifactBase]] | None = None, + partial_inputs: PartialInputRegistry | None = None, + ) -> None: + self._engine = engine or get_engine() + self._artifact_models = dict(artifact_models or ARTIFACT_MODELS) + self._partial_inputs = partial_inputs or PartialInputRegistry() + + def save_artifact( + self, + conversation_id: str, + artifact: ArtifactT, + ) -> ArtifactT: + expected_model = self._artifact_models.get(artifact.artifact_type) + if expected_model is None or not isinstance(artifact, expected_model): + raise TypeError( + f"Artifact {artifact.artifact_type!r} is not registered with its declared type." + ) + row = CapabilityArtifactRow( + artifact_id=artifact.artifact_id, + conversation_id=conversation_id, + artifact_type=artifact.artifact_type, + schema_version=artifact.schema_version, + payload_json=artifact.model_dump_json(), + created_at=artifact.created_at, + ) + with Session(self._engine) as session: + if session.get(CapabilityArtifactRow, artifact.artifact_id) is not None: + raise ValueError(f"Artifact already exists: {artifact.artifact_id}") + session.add(row) + session.commit() + return artifact + + def get_artifact(self, conversation_id: str, artifact_id: str) -> ArtifactBase: + with Session(self._engine) as session: + row = session.get(CapabilityArtifactRow, artifact_id) + if row is None or row.conversation_id != conversation_id: + raise KeyError(f"Unknown artifact: {artifact_id}") + return self._decode_artifact(row) + + def find_artifacts( + self, + conversation_id: str, + artifact_model: type[ArtifactT], + ) -> tuple[ArtifactT, ...]: + artifact_type = artifact_model.model_fields["artifact_type"].default + with Session(self._engine) as session: + rows = session.exec( + select(CapabilityArtifactRow) + .where(CapabilityArtifactRow.conversation_id == conversation_id) + .where(CapabilityArtifactRow.artifact_type == artifact_type) + .order_by(CapabilityArtifactRow.created_at) + ).all() + artifacts = tuple(self._decode_artifact(row) for row in rows) + if not all(isinstance(artifact, artifact_model) for artifact in artifacts): + raise InvalidPersistedRecord( + f"Stored {artifact_type} artifact decoded to an incompatible model." + ) + return cast(tuple[ArtifactT, ...], artifacts) + + def list_artifacts(self, conversation_id: str) -> tuple[ArtifactBase, ...]: + with Session(self._engine) as session: + rows = session.exec( + select(CapabilityArtifactRow) + .where(CapabilityArtifactRow.conversation_id == conversation_id) + .order_by(CapabilityArtifactRow.created_at) + ).all() + return tuple(self._decode_artifact(row) for row in rows) + + def create_waiting( + self, + invocation: WaitingCapabilityInvocation, + ) -> WaitingCapabilityInvocation: + validated = self._validate_partial( + invocation.capability_id, + invocation.input_schema_version, + invocation.partial_input, + ) + stored = invocation.model_copy( + update={ + "partial_input": validated, + **self._waiting_metadata(validated), + } + ) + row = self._waiting_row(stored) + with Session(self._engine) as session: + if session.get(WaitingCapabilityInvocationRow, row.invocation_id) is not None: + raise ValueError(f"Waiting invocation already exists: {row.invocation_id}") + session.add(row) + session.commit() + return stored + + def get_waiting(self, invocation_id: str) -> WaitingCapabilityInvocation: + with Session(self._engine) as session: + row = session.get(WaitingCapabilityInvocationRow, invocation_id) + if row is None: + raise KeyError(f"Unknown waiting invocation: {invocation_id}") + return self._decode_waiting(row) + + def list_waiting( + self, + conversation_id: str, + *, + capability_id: str | None = None, + ) -> tuple[WaitingCapabilityInvocation, ...]: + statement = select(WaitingCapabilityInvocationRow).where( + WaitingCapabilityInvocationRow.conversation_id == conversation_id + ) + if capability_id is not None: + statement = statement.where( + WaitingCapabilityInvocationRow.capability_id == capability_id + ) + statement = statement.order_by(WaitingCapabilityInvocationRow.created_at) + with Session(self._engine) as session: + rows = session.exec(statement).all() + return tuple(self._decode_waiting(row) for row in rows) + + def update_waiting( + self, + invocation_id: str, + partial_input: BaseModel, + ) -> WaitingCapabilityInvocation: + current = self.get_waiting(invocation_id) + validated = self._validate_partial( + current.capability_id, + current.input_schema_version, + partial_input, + ) + now = datetime.now(timezone.utc) + with Session(self._engine) as session: + row = session.get(WaitingCapabilityInvocationRow, invocation_id) + if row is None: + raise KeyError(f"Unknown waiting invocation: {invocation_id}") + row.partial_input_json = validated.model_dump_json() + row.updated_at = now + session.add(row) + session.commit() + return current.model_copy( + update={ + "partial_input": validated, + "updated_at": now, + **self._waiting_metadata(validated), + } + ) + + def resume_waiting( + self, + invocation_id: str, + updates: dict[str, object], + ) -> WaitingCapabilityInvocation: + current = self.get_waiting(invocation_id) + definition = self._partial_inputs.get(current.capability_id) + merged = {**current.partial_input.model_dump(), **updates} + try: + resumed = definition.model.model_validate(merged) + except ValidationError as exc: + raise TypeError( + f"Invalid resumed input for capability {current.capability_id}." + ) from exc + return self.update_waiting(invocation_id, resumed) + + def branch_waiting( + self, + invocation_id: str, + new_invocation_id: str, + source_turn_id: str, + ) -> WaitingCapabilityInvocation: + current = self.get_waiting(invocation_id) + branched = current.model_copy( + update={ + "invocation_id": new_invocation_id, + "source_turn_id": source_turn_id, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } + ) + return self.create_waiting(branched) + + def remove_waiting(self, invocation_id: str) -> None: + with Session(self._engine) as session: + row = session.get(WaitingCapabilityInvocationRow, invocation_id) + if row is not None: + session.delete(row) + session.commit() + + def _decode_artifact(self, row: CapabilityArtifactRow) -> ArtifactBase: + model = self._artifact_models.get(row.artifact_type) + if model is None: + raise InvalidPersistedRecord( + f"Unknown persisted artifact type: {row.artifact_type}." + ) + try: + artifact = model.model_validate_json(row.payload_json) + except ValidationError as exc: + raise InvalidPersistedRecord( + f"Invalid persisted artifact {row.artifact_id}." + ) from exc + if artifact.schema_version != row.schema_version: + raise InvalidPersistedRecord( + f"Artifact envelope version mismatch for {row.artifact_id}." + ) + return artifact + + def _validate_partial( + self, + capability_id: str, + schema_version: str, + value: BaseModel, + ) -> BaseModel: + definition = self._partial_inputs.get(capability_id) + if definition.schema_version != schema_version: + raise InvalidPersistedRecord( + f"Partial-input version mismatch for {capability_id}." + ) + try: + return definition.model.model_validate(value.model_dump()) + except ValidationError as exc: + raise TypeError( + f"Invalid partial input for capability {capability_id}." + ) from exc + + def _decode_waiting( + self, + row: WaitingCapabilityInvocationRow, + ) -> WaitingCapabilityInvocation: + definition = self._partial_inputs.get(row.capability_id) + if definition.schema_version != row.input_schema_version: + raise InvalidPersistedRecord( + f"Partial-input envelope version mismatch for {row.invocation_id}." + ) + try: + partial_input = definition.model.model_validate_json(row.partial_input_json) + except ValidationError as exc: + raise InvalidPersistedRecord( + f"Invalid persisted waiting invocation {row.invocation_id}." + ) from exc + return WaitingCapabilityInvocation( + invocation_id=row.invocation_id, + conversation_id=row.conversation_id, + capability_id=row.capability_id, + capability_version=row.capability_version, + input_schema_version=row.input_schema_version, + partial_input=partial_input, + source_turn_id=row.source_turn_id, + **self._waiting_metadata(partial_input), + created_at=row.created_at, + updated_at=row.updated_at, + ) + + @staticmethod + def _waiting_metadata(partial_input: BaseModel) -> dict[str, object]: + return { + "context_scope_id": getattr(partial_input, "context_scope_id", None), + "context_revision": getattr(partial_input, "context_revision", None), + "requirements": tuple( + getattr(partial_input, "fact_requirements", ()) + ), + } + + @staticmethod + def _waiting_row( + invocation: WaitingCapabilityInvocation, + ) -> WaitingCapabilityInvocationRow: + return WaitingCapabilityInvocationRow( + invocation_id=invocation.invocation_id, + conversation_id=invocation.conversation_id, + capability_id=invocation.capability_id, + capability_version=invocation.capability_version, + input_schema_version=invocation.input_schema_version, + partial_input_json=invocation.partial_input.model_dump_json(), + source_turn_id=invocation.source_turn_id, + created_at=invocation.created_at, + updated_at=invocation.updated_at, + ) + + +class RepositoryArtifactAccess: + """Async request-context view over the synchronous SQL repository.""" + + def __init__(self, repository: SQLConversationCapabilityRepository) -> None: + self._repository = repository + + async def find_artifacts( + self, + *, + conversation_id: str, + artifact_model: type[ArtifactT], + ) -> tuple[ArtifactT, ...]: + return await asyncio.to_thread( + self._repository.find_artifacts, + conversation_id, + artifact_model, + ) + + async def save_artifact( + self, + *, + conversation_id: str, + artifact: ArtifactT, + ) -> ArtifactT: + return await asyncio.to_thread( + self._repository.save_artifact, + conversation_id, + artifact, + ) + + async def save_waiting(self, invocation: object) -> object: + if not isinstance(invocation, WaitingCapabilityInvocation): + raise TypeError("Waiting state must use WaitingCapabilityInvocation.") + return await asyncio.to_thread( + self._repository.create_waiting, + invocation, + ) + + async def list_waiting( + self, + *, + conversation_id: str, + capability_id: str, + ) -> tuple[WaitingCapabilityInvocation, ...]: + return await asyncio.to_thread( + self._repository.list_waiting, + conversation_id, + capability_id=capability_id, + ) + + async def update_waiting( + self, + *, + invocation_id: str, + partial_input: BaseModel, + ) -> WaitingCapabilityInvocation: + return await asyncio.to_thread( + self._repository.update_waiting, + invocation_id, + partial_input, + ) + + async def remove_waiting(self, *, invocation_id: str) -> None: + await asyncio.to_thread( + self._repository.remove_waiting, + invocation_id, + ) diff --git a/backend/persistence/context_repository.py b/backend/persistence/context_repository.py new file mode 100644 index 00000000..a6962ae0 --- /dev/null +++ b/backend/persistence/context_repository.py @@ -0,0 +1,101 @@ +"""SQL implementation of versioned typed conversation-context persistence.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, cast + +from pydantic import ValidationError +from sqlalchemy import Engine, update +from sqlmodel import Session + +from conversation_context.models import ConversationContext +from conversation_context.repository import ConversationContextConflict +from conversations.models import get_engine +from persistence.capability_repository import InvalidPersistedRecord +from persistence.rows import ConversationContextRow + + +class SQLConversationContextRepository: + def __init__(self, *, engine: Engine | None = None) -> None: + self._engine = engine or get_engine() # type: ignore[no-untyped-call] + + def load(self, conversation_id: str) -> ConversationContext: + with Session(self._engine) as session: + row = session.get(ConversationContextRow, conversation_id) + if row is None: + return ConversationContext.initial(conversation_id) + try: + context = ConversationContext.model_validate_json(row.payload_json) + except ValidationError as exc: + raise InvalidPersistedRecord( + f"Invalid persisted conversation context for {conversation_id}." + ) from exc + if ( + context.conversation_id != conversation_id + or context.schema_version != row.schema_version + or context.revision != row.revision + ): + raise InvalidPersistedRecord( + f"Conversation context envelope mismatch for {conversation_id}." + ) + return context + + def save( + self, + context: ConversationContext, + *, + expected_revision: int, + ) -> ConversationContext: + if context.revision < expected_revision: + raise ValueError("A context revision cannot move backwards.") + now = datetime.now(timezone.utc) + with Session(self._engine) as session: + row = session.get(ConversationContextRow, context.conversation_id) + if row is None: + if expected_revision != 0: + raise ConversationContextConflict( + "Conversation context was not present at the expected revision." + ) + session.add( + ConversationContextRow( + conversation_id=context.conversation_id, + schema_version=context.schema_version, + revision=context.revision, + payload_json=context.model_dump_json(), + created_at=now, + updated_at=now, + ) + ) + session.commit() + return context + + table = cast(Any, ConversationContextRow).__table__ + statement = ( + update(ConversationContextRow) + .where( + table.c.conversation_id == context.conversation_id + ) + .where(table.c.revision == expected_revision) + .values( + schema_version=context.schema_version, + revision=context.revision, + payload_json=context.model_dump_json(), + updated_at=now, + ) + ) + result = session.exec(statement) + if result.rowcount != 1: + session.rollback() + raise ConversationContextConflict( + "Conversation context changed after it was loaded." + ) + session.commit() + return context + + def delete(self, conversation_id: str) -> None: + with Session(self._engine) as session: + row = session.get(ConversationContextRow, conversation_id) + if row is not None: + session.delete(row) + session.commit() diff --git a/backend/persistence/deletion.py b/backend/persistence/deletion.py new file mode 100644 index 00000000..0b8adde9 --- /dev/null +++ b/backend/persistence/deletion.py @@ -0,0 +1,29 @@ +"""Conversation-owned capability record deletion.""" + +from sqlmodel import Session, delete + +from persistence.rows import ( + CapabilityArtifactRow, + CapabilityCallReceiptRow, + ConversationContextRow, + InvocationTraceRow, + TurnReceiptRow, + WaitingCapabilityInvocationRow, +) + + +CAPABILITY_CONVERSATION_ROWS = ( + CapabilityCallReceiptRow, + TurnReceiptRow, + InvocationTraceRow, + WaitingCapabilityInvocationRow, + CapabilityArtifactRow, + ConversationContextRow, +) + + +def delete_capability_records(session: Session, conversation_id: str) -> None: + for row_model in CAPABILITY_CONVERSATION_ROWS: + session.exec( + delete(row_model).where(row_model.conversation_id == conversation_id) + ) diff --git a/backend/persistence/idempotency.py b/backend/persistence/idempotency.py new file mode 100644 index 00000000..dc194e2c --- /dev/null +++ b/backend/persistence/idempotency.py @@ -0,0 +1,222 @@ +"""Idempotent turn and externally significant capability-call receipts.""" + +from __future__ import annotations + +import hashlib +import json +from datetime import datetime, timezone +from enum import Enum + +from pydantic import BaseModel, ConfigDict +from sqlalchemy import update +from sqlalchemy.exc import IntegrityError +from sqlmodel import Session + +from conversations.models import get_engine +from persistence.rows import CapabilityCallReceiptRow, TurnReceiptRow + + +class ReceiptStatus(str, Enum): + PROCESSING = "processing" + COMPLETED = "completed" + FAILED = "failed" + + +class IdempotencyDecision(str, Enum): + STARTED = "started" + IN_PROGRESS = "in_progress" + REPLAY = "replay" + CONFLICT = "conflict" + + +class IdempotencyResult(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + decision: IdempotencyDecision + status: ReceiptStatus + outcome: dict[str, object] | None = None + + +def request_fingerprint(payload: object) -> str: + canonical = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +class SQLIdempotencyRepository: + def __init__(self, *, engine=None) -> None: + self._engine = engine or get_engine() + + def begin_turn( + self, + *, + conversation_id: str, + turn_id: str, + fingerprint: str, + ) -> IdempotencyResult: + with Session(self._engine) as session: + row = session.get(TurnReceiptRow, turn_id) + if row is None: + row = TurnReceiptRow( + turn_id=turn_id, + conversation_id=conversation_id, + request_fingerprint=fingerprint, + status=ReceiptStatus.PROCESSING.value, + ) + session.add(row) + try: + session.commit() + except IntegrityError: + session.rollback() + row = session.get(TurnReceiptRow, turn_id) + if row is None: + raise + else: + return IdempotencyResult( + decision=IdempotencyDecision.STARTED, + status=ReceiptStatus.PROCESSING, + ) + return self._turn_result(row, conversation_id, fingerprint) + + def complete_turn( + self, + *, + turn_id: str, + fingerprint: str, + outcome: dict[str, object], + ) -> None: + with Session(self._engine) as session: + row = session.get(TurnReceiptRow, turn_id) + if row is None: + raise KeyError(f"Unknown turn receipt: {turn_id}") + if row.request_fingerprint != fingerprint: + raise ValueError("Turn fingerprint conflict.") + row.status = ReceiptStatus.COMPLETED.value + row.public_outcome_json = json.dumps(outcome, sort_keys=True) + row.updated_at = datetime.now(timezone.utc) + session.add(row) + session.commit() + + def fail_turn(self, *, turn_id: str, fingerprint: str) -> None: + with Session(self._engine) as session: + row = session.get(TurnReceiptRow, turn_id) + if row is None: + raise KeyError(f"Unknown turn receipt: {turn_id}") + if row.request_fingerprint != fingerprint: + raise ValueError("Turn fingerprint conflict.") + row.status = ReceiptStatus.FAILED.value + row.updated_at = datetime.now(timezone.utc) + session.add(row) + session.commit() + + def claim_billing(self, turn_id: str) -> bool: + with Session(self._engine) as session: + result = session.exec( + update(TurnReceiptRow) + .where(TurnReceiptRow.turn_id == turn_id) + .where(TurnReceiptRow.billing_claimed.is_(False)) + .values(billing_claimed=True, updated_at=datetime.now(timezone.utc)) + ) + session.commit() + return result.rowcount == 1 + + def begin_call( + self, + *, + conversation_id: str, + turn_id: str, + call_id: str, + operation_id: str, + fingerprint: str, + ) -> IdempotencyResult: + with Session(self._engine) as session: + row = session.get(CapabilityCallReceiptRow, call_id) + if row is None: + row = CapabilityCallReceiptRow( + call_id=call_id, + conversation_id=conversation_id, + turn_id=turn_id, + operation_id=operation_id, + request_fingerprint=fingerprint, + status=ReceiptStatus.PROCESSING.value, + ) + session.add(row) + try: + session.commit() + except IntegrityError: + session.rollback() + row = session.get(CapabilityCallReceiptRow, call_id) + if row is None: + raise + else: + return IdempotencyResult( + decision=IdempotencyDecision.STARTED, + status=ReceiptStatus.PROCESSING, + ) + if ( + row.conversation_id != conversation_id + or row.turn_id != turn_id + or row.operation_id != operation_id + or row.request_fingerprint != fingerprint + ): + return IdempotencyResult( + decision=IdempotencyDecision.CONFLICT, + status=ReceiptStatus(row.status), + ) + return self._result(row.status, row.outcome_json) + + def complete_call( + self, + *, + call_id: str, + fingerprint: str, + outcome: dict[str, object], + ) -> None: + with Session(self._engine) as session: + row = session.get(CapabilityCallReceiptRow, call_id) + if row is None: + raise KeyError(f"Unknown capability call receipt: {call_id}") + if row.request_fingerprint != fingerprint: + raise ValueError("Capability call fingerprint conflict.") + row.status = ReceiptStatus.COMPLETED.value + row.outcome_json = json.dumps(outcome, sort_keys=True) + row.updated_at = datetime.now(timezone.utc) + session.add(row) + session.commit() + + @staticmethod + def _turn_result( + row: TurnReceiptRow, + conversation_id: str, + fingerprint: str, + ) -> IdempotencyResult: + if ( + row.conversation_id != conversation_id + or row.request_fingerprint != fingerprint + ): + return IdempotencyResult( + decision=IdempotencyDecision.CONFLICT, + status=ReceiptStatus(row.status), + ) + return SQLIdempotencyRepository._result( + row.status, + row.public_outcome_json, + ) + + @staticmethod + def _result(status_value: str, outcome_json: str | None) -> IdempotencyResult: + status = ReceiptStatus(status_value) + if status is ReceiptStatus.PROCESSING: + decision = IdempotencyDecision.IN_PROGRESS + else: + decision = IdempotencyDecision.REPLAY + outcome = json.loads(outcome_json) if outcome_json is not None else None + return IdempotencyResult( + decision=decision, + status=status, + outcome=outcome, + ) diff --git a/backend/persistence/rows.py b/backend/persistence/rows.py new file mode 100644 index 00000000..979d2677 --- /dev/null +++ b/backend/persistence/rows.py @@ -0,0 +1,132 @@ +"""Additive SQLModel rows for capability state and invocation metadata.""" + +from __future__ import annotations + +from datetime import datetime, timezone + +from sqlalchemy import Column, Index, String +from sqlmodel import Field, SQLModel + + +def _now() -> datetime: + return datetime.now(timezone.utc) + + +def _conversation_column() -> Column: + return Column( + String, + nullable=False, + ) + + +class CapabilityArtifactRow(SQLModel, table=True): + __tablename__ = "capability_artifacts" + __table_args__ = ( + Index("idx_capability_artifacts_conversation", "conversation_id"), + Index("idx_capability_artifacts_type", "artifact_type"), + ) + + artifact_id: str = Field(primary_key=True) + conversation_id: str = Field(sa_column=_conversation_column()) + artifact_type: str + schema_version: str + payload_json: str + created_at: datetime = Field(default_factory=_now) + + +class ConversationContextRow(SQLModel, table=True): + __tablename__ = "conversation_contexts" + __table_args__ = ( + Index("idx_conversation_context_updated", "updated_at"), + ) + + conversation_id: str = Field(primary_key=True) + schema_version: str + revision: int + payload_json: str + created_at: datetime = Field(default_factory=_now) + updated_at: datetime = Field(default_factory=_now) + + +class WaitingCapabilityInvocationRow(SQLModel, table=True): + __tablename__ = "waiting_capability_invocations" + __table_args__ = ( + Index("idx_waiting_capability_conversation", "conversation_id"), + Index("idx_waiting_capability_id", "capability_id"), + ) + + invocation_id: str = Field(primary_key=True) + conversation_id: str = Field(sa_column=_conversation_column()) + capability_id: str + capability_version: str + input_schema_version: str + partial_input_json: str + source_turn_id: str + created_at: datetime = Field(default_factory=_now) + updated_at: datetime = Field(default_factory=_now) + + +class InvocationTraceRow(SQLModel, table=True): + __tablename__ = "capability_invocation_traces" + __table_args__ = ( + Index( + "idx_capability_trace_conversation_sequence", + "conversation_id", + "sequence", + ), + Index("idx_capability_trace_turn", "turn_id"), + Index("idx_capability_trace_parent", "parent_invocation_id"), + Index("idx_capability_trace_identifier", "identifier"), + ) + + invocation_id: str = Field(primary_key=True) + conversation_id: str = Field(sa_column=_conversation_column()) + turn_id: str + parent_invocation_id: str | None = None + sequence: int + kind: str + identifier: str + version: str + visibility: str + started_at: datetime + completed_at: datetime | None = None + duration_ms: int | None = None + status: str + summary: str + debug_input_json: str | None = None + debug_output_json: str | None = None + + +class TurnReceiptRow(SQLModel, table=True): + __tablename__ = "capability_turn_receipts" + __table_args__ = ( + Index("idx_capability_turn_conversation", "conversation_id"), + ) + + turn_id: str = Field(primary_key=True) + conversation_id: str = Field(sa_column=_conversation_column()) + request_fingerprint: str + status: str + public_outcome_json: str | None = None + billing_claimed: bool = False + created_at: datetime = Field(default_factory=_now) + updated_at: datetime = Field(default_factory=_now) + + +class CapabilityCallReceiptRow(SQLModel, table=True): + __tablename__ = "capability_call_receipts" + __table_args__ = ( + Index("idx_capability_call_conversation", "conversation_id"), + Index("idx_capability_call_turn", "turn_id"), + Index("idx_capability_call_operation", "operation_id"), + ) + + call_id: str = Field(primary_key=True) + conversation_id: str = Field(sa_column=_conversation_column()) + turn_id: str + operation_id: str + request_fingerprint: str + status: str + outcome_json: str | None = None + created_at: datetime = Field(default_factory=_now) + updated_at: datetime = Field(default_factory=_now) diff --git a/backend/persistence/schema.py b/backend/persistence/schema.py new file mode 100644 index 00000000..fd272e16 --- /dev/null +++ b/backend/persistence/schema.py @@ -0,0 +1,68 @@ +"""Read-only validation of the deployed Alembic schema revision.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +from alembic.config import Config +from alembic.runtime.migration import MigrationContext +from alembic.script import ScriptDirectory +from sqlalchemy.engine import Engine + +from conversations.models import get_engine + +logger = logging.getLogger(__name__) + +ALEMBIC_CONFIG_PATH = Path(__file__).resolve().parents[1] / "alembic.ini" + + +class DatabaseSchemaError(RuntimeError): + """Raised when a deployed database is not at the repository revision.""" + + +def expected_schema_revision( + config_path: Path = ALEMBIC_CONFIG_PATH, +) -> str: + """Return the repository's single Alembic head revision.""" + script = ScriptDirectory.from_config(Config(str(config_path))) + heads = script.get_heads() + if len(heads) != 1: + raise DatabaseSchemaError( + "Expected exactly one Alembic head revision; " + f"found {len(heads)}" + ) + return heads[0] + + +def current_schema_revision(engine: Engine) -> str | None: + """Read the database's current Alembic revision without changing schema.""" + with engine.connect() as connection: + return MigrationContext.configure(connection).get_current_revision() + + +def verify_database_schema(engine: Engine | None = None) -> str | None: + """Require deployed PostgreSQL to match the repository Alembic revision. + + SQLite is reserved for isolated repository tests, whose fixtures construct + temporary schemas directly from SQLModel metadata. + """ + database_engine = engine or get_engine() + if database_engine.dialect.name == "sqlite": + logger.info( + "Skipping Alembic revision validation for an isolated SQLite test database" + ) + return None + + expected = expected_schema_revision() + current = current_schema_revision(database_engine) + if current != expected: + found = current or "none" + raise DatabaseSchemaError( + "Database schema revision mismatch: " + f"expected {expected}, found {found}. " + "Run `alembic -c alembic.ini upgrade head` before starting the backend." + ) + + logger.info("Database schema revision %s verified", current) + return current diff --git a/backend/persistence/trace_repository.py b/backend/persistence/trace_repository.py new file mode 100644 index 00000000..6e826e03 --- /dev/null +++ b/backend/persistence/trace_repository.py @@ -0,0 +1,114 @@ +"""SQL storage for allowlisted invocation records.""" + +from __future__ import annotations + +from datetime import datetime, timezone +import json + +from sqlalchemy import func +from sqlmodel import Session, select + +from capabilities.tracing import InvocationRecord +from conversations.models import get_engine +from persistence.rows import InvocationTraceRow +from tools.contracts import Visibility + + +def _utc(value: datetime | None) -> datetime | None: + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=timezone.utc) + return value.astimezone(timezone.utc) + + +class SQLInvocationTraceRepository: + def __init__(self, *, engine=None) -> None: + self._engine = engine or get_engine() + + def save(self, record: InvocationRecord) -> None: + row = InvocationTraceRow( + invocation_id=record.invocation_id, + conversation_id=record.conversation_id, + turn_id=record.turn_id, + parent_invocation_id=record.parent_invocation_id, + sequence=record.sequence, + kind=record.kind.value, + identifier=record.identifier, + version=record.version, + visibility=record.visibility.value, + started_at=record.started_at, + completed_at=record.completed_at, + duration_ms=record.duration_ms, + status=record.status.value, + summary=record.summary, + debug_input_json=( + json.dumps(record.debug_input, separators=(",", ":")) + if record.debug_input is not None + else None + ), + debug_output_json=( + json.dumps(record.debug_output, separators=(",", ":")) + if record.debug_output is not None + else None + ), + ) + with Session(self._engine) as session: + session.merge(row) + session.commit() + + def last_sequence(self, conversation_id: str) -> int: + statement = select(func.max(InvocationTraceRow.sequence)).where( + InvocationTraceRow.conversation_id == conversation_id + ) + with Session(self._engine) as session: + value = session.exec(statement).one() + return int(value or 0) + + def list_for_conversation( + self, + conversation_id: str, + *, + include_private: bool, + ) -> tuple[InvocationRecord, ...]: + statement = select(InvocationTraceRow).where( + InvocationTraceRow.conversation_id == conversation_id + ) + if not include_private: + statement = statement.where( + InvocationTraceRow.visibility == Visibility.PUBLIC.value + ) + statement = statement.order_by(InvocationTraceRow.sequence) + with Session(self._engine) as session: + rows = session.exec(statement).all() + return tuple( + InvocationRecord.model_validate( + { + "conversation_id": row.conversation_id, + "turn_id": row.turn_id, + "invocation_id": row.invocation_id, + "parent_invocation_id": row.parent_invocation_id, + "sequence": row.sequence, + "kind": row.kind, + "identifier": row.identifier, + "version": row.version, + "visibility": row.visibility, + "started_at": _utc(row.started_at), + "completed_at": _utc(row.completed_at), + "duration_ms": row.duration_ms, + "status": row.status, + "summary": row.summary, + "debug_input": ( + json.loads(row.debug_input_json) + if row.debug_input_json is not None + else None + ), + "debug_output": ( + json.loads(row.debug_output_json) + if row.debug_output_json is not None + else None + ), + } + ) + for row in rows + ) diff --git a/backend/prompts/__init__.py b/backend/prompts/__init__.py index 3d504df3..3eeb7e6d 100644 --- a/backend/prompts/__init__.py +++ b/backend/prompts/__init__.py @@ -1,38 +1,8 @@ -"""Prompt text used by the UK chat backend. +"""Prompt text for chat titles and follow-up suggestions.""" -Split by audience (system / gateway / meta) so each file owns one set of -model-facing instructions, but this package re-exports the public surface so -callers keep doing `from prompts import X`. Keep the constants declarative: -routes assemble blocks and call models; this package owns the prompt text. - -For model-neutral engineering guidance around this runtime pathway, see -`docs/engineering/skills/uk-chat-runtime.md`. -""" - -from prompts.gateway import ( - DEFAULT_SCOPE_DESCRIPTOR, - GATEWAY_CATALOGUE_RECOVERY_DIRECTIVE, - GATEWAY_IRRELEVANT_DIRECTIVE, - GATEWAY_OUT_OF_SCOPE_DIRECTIVE, - GATEWAY_PARTIAL_CATALOGUE_DIRECTIVE, - GATEWAY_PARTIAL_DIRECTIVE, - gateway_system, - lightweight_system, -) from prompts.meta import SUGGESTION_SYSTEM, TITLE_SYSTEM -from prompts.system import CHARTS_MODE_DIRECTIVE, SYSTEM_PROMPT __all__ = [ - "SYSTEM_PROMPT", - "CHARTS_MODE_DIRECTIVE", - "DEFAULT_SCOPE_DESCRIPTOR", - "lightweight_system", - "gateway_system", - "GATEWAY_CATALOGUE_RECOVERY_DIRECTIVE", - "GATEWAY_IRRELEVANT_DIRECTIVE", - "GATEWAY_OUT_OF_SCOPE_DIRECTIVE", - "GATEWAY_PARTIAL_DIRECTIVE", - "GATEWAY_PARTIAL_CATALOGUE_DIRECTIVE", "SUGGESTION_SYSTEM", "TITLE_SYSTEM", ] diff --git a/backend/prompts/gateway.py b/backend/prompts/gateway.py deleted file mode 100644 index 1cdfb04d..00000000 --- a/backend/prompts/gateway.py +++ /dev/null @@ -1,182 +0,0 @@ -"""Gateway prompts: the scope descriptor, the lightweight no-computation system, -the gateway classifier instructions, and the per-outcome writer directives. -""" - -# This curated scope descriptor is the fallback used in local dev when an -# environment-specific descriptor is absent. Keep it compact — it is loaded -# into a cheap classifier prompt, not the full compute prompt. -DEFAULT_SCOPE_DESCRIPTOR = """ -This assistant models UK taxes and benefits with a microsimulation engine. -Modelled: income tax, National Insurance, Universal Credit, child benefit, -pension credit, tax credits, and related UK tax-and-benefit programmes, over the -pinned Enhanced FRS 2024-25 dataset for the supported tax years. -This is a summary, not an exhaustive list of every policyengine.py capability. -NOT modelled: macroeconomic / second-round effects (inflation, GDP, employment, -market reactions), behavioural response, non-UK policy, unannounced or future -Budgets, and legal or individual tax-filing advice. -""".strip() - - -_LIGHTWEIGHT_INSTRUCTIONS = """ -You are an expert assistant for a UK tax and benefit microsimulation platform. -This turn does not run the model, so you have no tools and no live parameter -data loaded. Respond briefly and directly to the user's message. - -Do NOT state specific quantitative figures, rates, or parameter values from -memory — you do not have the data loaded this turn. If a number is needed, say -you can compute it if the user asks. Use British English and stay factually -neutral: do not label policies good, bad, fair, regressive, progressive, -generous, or similar. -""".strip() - - -def lightweight_system(scope_descriptor: str) -> str: - """Lean no-computation system prompt, parameterised by the scope descriptor. - - Used as the base for model-written non-`ready` outcomes (irrelevant, - out_of_scope, partial); a per-outcome directive is appended at request - time. `needs_plan` is rendered deterministically without a model call. - """ - return _LIGHTWEIGHT_INSTRUCTIONS + "\n\n" + scope_descriptor.strip() - - -# The gateway is a cheap pre-pass that builds a structured execution plan before -# any expensive model call. It does NOT judge importance — it grounds each slot -# (prompt / default / assumed); the server applies criticality and gates. These -# instructions tell the model how to fill the plan via the forced `emit_plan` -# tool. The two fail-safe biases are stated explicitly. `{default_year}` is -# filled by gateway_system() from the caller-supplied engine default so the -# documented safe default can't drift from `DEFAULT_SIMULATION_YEAR`. -_GATEWAY_INSTRUCTIONS_TEMPLATE = """ -You are a routing pre-pass for a UK tax-and-benefit microsimulation assistant. -You do NOT answer the user. You build a short execution plan and emit it by -calling the `emit_plan` tool exactly once. Never write prose. - -Steps: -1. `domain_status`: classify the request as `uk_or_unspecified`, - `explicit_non_uk`, or `unrelated`. This is a UK product, so an unspecified - jurisdiction defaults to `uk_or_unspecified`. For either negative status, - include a short `domain_evidence` exact quote from the user's message that - proves the exclusion. General knowledge, chit-chat, coding, or explicitly - non-UK questions are excluded. -2. `capability_status`: use `supported` when the available tool chain fits. Use - `catalogue_uncertain` only when a named UK policy or variable may be supported - but needs current catalogue confirmation. Use `explicitly_unmodellable` only - when the user explicitly requests solely an unavailable effect. Both - non-supported statuses require a `capability_evidence` exact quote from the - message. -3. `tool`: pick the single best-fitting tool for the modelled part of the ask, - or "none" if nothing the engine computes applies (e.g. a pure macro/ - behavioural question). Use the tool list below. -4. `slots`: for the chosen tool, list its required and defaultable input slots, - plus one or more `output` slots naming what the user wants reported (use one - of the output labels listed below). For each slot set `value` and tag - `source`: - - "prompt": the user stated it or clearly implied it. - - "default": a documented safe default applies (year {default_year}; - baseline is current law). - - "assumed": you are guessing, or a documented default does not settle the - user's request. -5. `unmodellable_outputs`: include only outputs that the user explicitly asks for - and that the available tool chain cannot calculate. Each item must contain a - concise `name` and an `evidence` exact quote from the user's message that - explicitly requests it. Do not add behavioural, employment, take-up, - macroeconomic, market, or second-round effects merely because they could - affect the real-world result. Unless the user explicitly requests dynamic or - behavioural effects, interpret cost, revenue, spending, poverty, inequality, - decile, winners/losers, caseload, marginal-rate, and net-income questions as - requests for the direct static microsimulation result. Excluded secondary - effects are caveats for the final answer, not additional requested outputs, - and must not trigger `partial`. Leave the list empty when no explicitly - requested output is unmodellable. -6. `catalogue_queries`: for every named UK tax-benefit reform measure or model - variable concept, emit a short search term for the server to verify against - the current policyengine.py catalogue, plus an `evidence` exact quote from - the user's message containing that search term. Do not include rates, - amounts, or the whole user message. Never silently remove a foreign - jurisdiction from the quoted evidence. Use an empty list only when no such - concept is named. - The scope summary below is not an exhaustive list: never choose `tool="none"` - merely because a named policy is absent from it. - -Two fail-safe biases — apply them: -- Admissibility leans toward IN scope. When unsure whether a tool fits, pick a - tool and proceed rather than declaring "none"; a wrong refusal is worse than a - wrong compute. Use a negative domain or capability status only when the exact - quoted evidence clearly supports it. -- Grounding leans toward "assumed". When unsure whether the user actually - specified a slot, tag it "assumed" rather than "prompt" or "default", so the - server can ask instead of guessing on a load-bearing field. -""".strip() - - -def gateway_system( - scope_descriptor: str, tool_summary: str, output_labels: str, default_year: int -) -> str: - """Gateway classifier prompt, parameterised by the scope descriptor, a - compact tool summary, the output-slot labels, and the default simulation - year — all derived from the engine / config so they can't drift from a - hardcoded copy.""" - return ( - _GATEWAY_INSTRUCTIONS_TEMPLATE.format(default_year=default_year) - + "\n\nThe following directly modelled output labels are authoritative " - + "(use one per `output` slot): " - + output_labels - + "\n\nTools available (name — purpose; required params):\n" - + tool_summary.strip() - + "\n\nScope:\n" - + scope_descriptor.strip() - ) - - -# Appended only for the single recovery pass after the server has found -# authoritative catalogue candidates for a grounded user phrase. The runtime -# supplies the candidates below this directive; keeping the behavioural -# instructions here leaves model-facing policy in the prompts package. -GATEWAY_CATALOGUE_RECOVERY_DIRECTIVE = """ -The server found authoritative current-model catalogue candidates for the -user's grounded policy phrase. Rebuild the execution plan once using those -candidates and the original user message. - -- Catalogue candidates prove that related model capability exists; they do not - prove which candidate or reform the user intended. -- Select the best-fitting tool and ground its slots only from the original user - message, documented defaults, and the candidate metadata below. -- If a load-bearing choice remains ambiguous, mark that slot `assumed` so the - server asks a clarification instead of guessing. -- Preserve `explicit_non_uk`, `unrelated`, or `explicitly_unmodellable` only - when the original message contains exact quoted evidence for that decision. -- The catalogue lookup is already complete. Emit an empty `catalogue_queries` - list and do not request another lookup. -""".strip() - - -# Per-outcome writer directives. Appended to the lightweight system for the -# single no-tool turn that actually replies to the user on a non-`ready` -# outcome. The concrete slot names / unmodellable outputs are appended at -# request time by gateway.gateway_writer_directive(). -GATEWAY_IRRELEVANT_DIRECTIVE = """ -The user's message is outside UK tax and benefit policy. Decline in one or two -sentences and say what you can help with instead. Do not attempt to answer it. -""".strip() - -GATEWAY_OUT_OF_SCOPE_DIRECTIVE = """ -The user's question is about an effect this microsimulation does not model -(e.g. macroeconomic, inflation, behavioural, or non-UK). Say so clearly in one -or two sentences, and offer the closest modelled angle you could compute -instead (e.g. the direct fiscal or household-level effect). -""".strip() - -GATEWAY_PARTIAL_DIRECTIVE = """ -Part of the user's question is modellable and part is not. Briefly state the -part you CAN compute and the part you cannot (named below), then ask whether -they'd like you to run the modellable part. Do not run anything yet. -""".strip() - -GATEWAY_PARTIAL_CATALOGUE_DIRECTIVE = """ -Part of the user's question is modellable and part is not. Briefly state the -part you CAN compute and the part you cannot (named below). A named policy -measure or variable also needs clarification: ask what supported measure or -variable the user means before offering to run the modellable part. Do not run -anything yet. -""".strip() diff --git a/backend/prompts/system.py b/backend/prompts/system.py deleted file mode 100644 index 4a86f8c5..00000000 --- a/backend/prompts/system.py +++ /dev/null @@ -1,181 +0,0 @@ -"""The main compute system prompt and the chart-mode directive.""" - -from engine.constants import HOUSEHOLD_COUNTRY_IDS, UK_CHAT_DATASET -from tools.definitions import DEFAULT_SIMULATION_YEAR - -ROLE_AND_TASK = """ -You are an expert policy analysis assistant for a UK microsimulation platform. -You help users understand and analyse UK tax and benefit policy using the -policyengine.py UK model. -""" - -COMPUTATION_RULES = f""" -CRITICAL - ALWAYS COMPUTE WITH TOOLS: -- Never answer quantitative policy questions from memory. -- Every number in your answer must come directly from a tool result you just - computed in this turn. -- If the user does not provide a year, use the current calendar year - ({DEFAULT_SIMULATION_YEAR}). Preserve any year the user explicitly provides. -- Society-wide simulations always use UK Chat's pinned - `{UK_CHAT_DATASET.name}` dataset. The model cannot select another dataset. - Mention the dataset when it matters. -- If a question needs variables, parameters, model entities, reform - targets, household input variables, or supported outputs, use the discovery - tools first. Do not guess model names. -- Before a society simulation that needs variable-level outputs, call - `list_society_output_variables` unless its result is already available in the - conversation. For every required aggregate or filter variable not in that - default set, call `search_variables` or `get_variable` and wait for the - result before running the simulation. -- `extra_variables` only materializes existing policyengine-uk variables that - are absent from the default society outputs. It does not define new - variables, expressions, aliases, filters, or derived concepts. Omit default - variables from it, place each extra under the entity reported by variable - discovery, and omit the field entirely when no extra output is needed. -- Use `validate_reform` when drafting, debugging, or checking reform JSON. -- Use `validate_household` when checking whether a synthetic household is - shaped correctly. -- Use `run_household_simulation` for illustrative synthetic households. -- Use `run_society_simulation` for aggregate, society-wide reform analysis. -- After a society simulation, use derivative tools such as - `compute_budgetary_impact`, `compute_program_breakdown`, - `compute_decile_impacts`, `compute_winners_losers`, - `compute_poverty_metrics`, `compute_inequality_metrics`, or - `aggregate_result` for specific outputs. These tools use policyengine.py's - official weighted output classes; do not try to aggregate simulation rows. -- Do not run broad Python code for normal analysis. The model-facing tools are - the supported calculation interface. -""" - -DISCOVERY_RULES = """ -DISCOVERY AND VALIDATION: -- `list_entities` reports model entities. -- `search_variables` and `get_variable` verify exact model variables and report - whether they are default society outputs. -- `search_parameters` and `get_parameter` report parameters. -- `list_reform_targets` reports commonly supported reform paths. -- On an opening compute turn, `MODEL CATALOGUE EVIDENCE` is a server-generated - current policyengine.py discovery result. Treat matching paths and variables - as candidates, not a resolution of user intent; ask a concise clarification - before computing if the user's requested measure remains ambiguous. -- `list_household_input_variables` reports variables suitable for synthetic - household input. -- `list_society_output_variables` reports variables automatically materialized - by a policyengine.py society simulation, grouped by output entity. -- `list_supported_outputs` reports household, society, derivative, and chart - outputs available through this chat runtime. -- Validate before running when the user asks whether an input is valid, when - constructing a non-trivial reform, or when an earlier simulation fails. -""" - -REFORM_RULES = """ -REFORMS: -- Reforms are flat dictionaries keyed by policyengine.py parameter path, with - values applied from 1 January of the simulation year. -- Do not invent parameter paths. Search or inspect parameters first unless the - exact path is already present in the conversation or a tool result. -- For baseline/current-law questions, omit `reform`. -- If a reform is under-specified in a load-bearing way, ask a concise - clarifying question before computing. -""" - -_HOUSEHOLD_COUNTRY_IDS = ", ".join( - f"`{country_id}`" for country_id in HOUSEHOLD_COUNTRY_IDS -) - -MICRODATA_PRIVACY_RULES = f""" -MICRODATA PRIVACY AND ILLUSTRATIVE HOUSEHOLDS: -- Do not access, display, quote, or imply access to row-level survey microdata - or real households. -- The society simulation and derivative tools return aggregate outputs only. -- The household tool models exactly one household containing one benefit unit. - Do not combine unrelated adults or multiple benefit units in one call; use - separate illustrative calls or state the limitation. -- For the household `country` input, use one of - {_HOUSEHOLD_COUNTRY_IDS}; do not use ONS codes such as `E92000001`. -- If the user asks for examples of households from the dataset, explain that - this app cannot access or disclose real household records. -- For household examples, construct illustrative synthetic households and - label them synthetic, illustrative, or hypothetical. -""" - -ANALYTICAL_NOTES = """ -ANALYTICAL NOTES: -- Decile impacts are policyengine.py decile-level averages, not economy-wide means. -- For income-decile impacts, measure household net income and rank households - by that income by default, using `decile_concept="household_net_income"`. - Only when the user explicitly requests equivalised HBAI net income, use - `decile_concept="equivalised_hbai_net_income"`. policyengine.py forms - computed household income groups using person-weighted ranks. Households - with negative or non-finite values of that income concept are excluded from - the final reported deciles, consistently with country-package reporting. -- For wealth-decile impacts, use `decile_concept="wealth"` to group households - by wealth: group households by wealth and measure household net income. Do - not describe wealth deciles as income deciles. -- An empty decile has null income impacts, and a zero baseline mean has a null - relative change. These are missing results, not zero impacts: report them as - unavailable and do not describe them as no change. -- Poverty outputs report decimal rates and both absolute and relative changes. -- If a result is counterintuitive, explain the mechanism briefly. -- If something is not modelled well enough for a quantitative answer, say so - clearly and do not fabricate estimates. -- Use British English. -""" - -NEUTRALITY_RULES = """ -FACTUAL NEUTRALITY: -- Be factually neutral. -- Do not describe UK tax or benefit choices as good, bad, fair, unfair, - regressive, progressive, generous, punitive, or similar. -- Stick to mechanics and quantified effects. -- Describe who pays or receives more or less, by how much, over what period, - and under which dataset, year, and assumptions. -- If a distributional pattern matters, describe the measured direction - directly rather than applying value labels. -- Do not make policy recommendations unless the user explicitly asks for policy - design options. Even then, frame tradeoffs neutrally. -""" - -USER_FACING_STYLE = """ -USER-FACING STYLE: -- Prefer plain English in the prose answer. -- Avoid exposing internal parameter keys unless the user wants code-level - detail. -- Keep the answer grounded in tool outputs. -- Do not paste full raw tool JSON into the answer unless the user asks for it. -""" - -CHART_RULES = """ -CHARTS: -- When a visualisation would help, call `generate_chart` after the calculation - or derivative tool has produced the data. -- Prefer deterministic preset chart kinds for supported policy outputs: - `budget_waterfall`, `program_budget_waterfall`, `decile_absolute_bar`, - `decile_relative_bar`, `winners_losers_stacked_bar`, - `poverty_relative_bar`, `inequality_relative_bar`, or - `earnings_variation_line`. -- The tool returns a `chart_markdown` field containing a ```chart fenced JSON - block. Paste that block verbatim into your next text response so the - frontend can render it. -- Use factually neutral chart titles, labels, and captions. -""" - -SYSTEM_PROMPT_SECTIONS = ( - ROLE_AND_TASK, - COMPUTATION_RULES, - DISCOVERY_RULES, - REFORM_RULES, - MICRODATA_PRIVACY_RULES, - ANALYTICAL_NOTES, - NEUTRALITY_RULES, - USER_FACING_STYLE, - CHART_RULES, -) - -SYSTEM_PROMPT = "\n\n".join(section.strip() for section in SYSTEM_PROMPT_SECTIONS) - -CHARTS_MODE_DIRECTIVE = """ -The user has enabled chart mode. When the answer would benefit from a -visualisation, include a chart using the available chart tools alongside the -written explanation. Do not force charts on questions that are not chartable. -""".strip() diff --git a/backend/requirements-test.txt b/backend/requirements-test.txt index 9955decc..a1071106 100644 --- a/backend/requirements-test.txt +++ b/backend/requirements-test.txt @@ -1,2 +1,3 @@ pytest pytest-cov +mypy diff --git a/backend/requirements.txt b/backend/requirements.txt index 46c27a1f..fdd551cc 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -1,9 +1,10 @@ fastapi uvicorn[standard] +alembic sqlmodel psycopg2-binary -anthropic -policyengine[uk]==5.0.1 +anthropic==0.125.0 +policyengine[uk]==5.2.0 pandas pyyaml httpx diff --git a/backend/tests/test_agent_tools.py b/backend/tests/test_agent_tools.py index 4c077665..b08ad847 100644 --- a/backend/tests/test_agent_tools.py +++ b/backend/tests/test_agent_tools.py @@ -73,56 +73,6 @@ def test_simulation_schema_uses_current_year_and_fixed_dataset(): ).parameters -def test_society_simulation_rejects_reform_different_from_gateway_approval(monkeypatch): - approved = {"gov.hmrc.income_tax.rates.uk[0].rate": 0.21} - context = new_tool_context("guarded-simulation") - context.require_approved_reform = True - context.approved_reform = approved - called = False - - def build(**_kwargs): - nonlocal called - called = True - raise AssertionError("simulation must not run") - - monkeypatch.setattr(agent_tools, "build_society_simulation", build) - - result = agent_tools.run_society_simulation( - year=2026, - reform={"gov.hmrc.income_tax.rates.uk[0].rate": 0.22}, - _context=context, - ) - - assert result["error"] == "Gateway-approved reform mismatch" - assert called is False - - -def test_society_simulation_accepts_exact_gateway_approved_reform(monkeypatch): - approved = {"gov.hmrc.income_tax.rates.uk[0].rate": 0.21} - context = new_tool_context("guarded-simulation") - context.require_approved_reform = True - context.approved_reform = approved - - class Payload: - def metadata(self): - return {"status": "success"} - - monkeypatch.setattr( - agent_tools, - "build_society_simulation", - lambda **_kwargs: Payload(), - ) - - result = agent_tools.run_society_simulation( - year=2026, - reform=dict(approved), - _context=context, - ) - - assert result["status"] == "success" - assert result["result_id"].startswith("society_simulation_") - - def test_current_simulation_year_tracks_the_calendar(monkeypatch): class FutureDate(date): @classmethod @@ -680,7 +630,6 @@ def test_runtime_files_do_not_reference_compiled_package(): root / "api", root / "chat", root / "engine", - root / "gateway", root / "prompts", root / "tools", root / "Dockerfile", diff --git a/backend/tests/test_anthropic_sdk_contract.py b/backend/tests/test_anthropic_sdk_contract.py new file mode 100644 index 00000000..5ccd61f6 --- /dev/null +++ b/backend/tests/test_anthropic_sdk_contract.py @@ -0,0 +1,712 @@ +"""Compatibility checks for the Anthropic SDK calls used by UK Chat.""" + +import asyncio +import inspect +import json +from types import SimpleNamespace + +import pytest + +import chat.model_port as model_port +import conversation_context.tools as context_tools +import conversation_context.variable_resolution as variable_resolution +from chat.model_port import AnthropicConversationModel +from config.clients import get_async_client, get_sync_client +from conversation_context.models import ( + ClaimedMoneyValue, + ContextEntityCandidate, + ContextPatch, + ConversationContext, + EntityKind, + EnsureEntityOperation, + FactClaim, + FactClaimRelationship, + MoneyPeriod, + TextFactValue, +) +from conversation_context.projection import project_context +from conversation_context.reducer import ContextReducer +from conversation_context.registry import build_default_fact_registry +from conversation_context.change_pipeline import ( + ContextChangeProposal, + ContextChangeValidator, + ContextValidationIssue, + ContextValidationStatus, + ValidateContextChangeInput, +) +from conversation_context.tools import ( + AnthropicContextProposalReviewer, + AnthropicContextInterpreter, + ContextProposalStatus, + ProposeContextChangeInput, + ContextConversationExcerpt, +) +from conversation_context.variable_resolution import ( + AnthropicVariableMapper, + MappingConfidence, + MappingStatus, + PolicyEngineVariableCandidate, +) + + +def context_with_spouse(*, revision: int = 2) -> ConversationContext: + registry = build_default_fact_registry() + context = ContextReducer(registry).reduce( + ConversationContext.initial("conversation"), + ContextPatch( + expected_revision=0, + operations=( + EnsureEntityOperation( + reference="entity:spouse", + kind=EntityKind.PERSON, + aliases=("spouse",), + relationship_to_user="spouse", + ), + ), + ), + turn_id="turn-spouse", + evidence="I have a spouse.", + ).context + return context.model_copy(update={"revision": revision}) + + +def _usage(*, input_tokens: int = 3, output_tokens: int = 2): + return SimpleNamespace( + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_creation_input_tokens=0, + cache_read_input_tokens=0, + ) + + +def _request( + message: str, + *, + context: ConversationContext | None = None, +) -> ProposeContextChangeInput: + resolved_context = context or ConversationContext.initial("conversation") + return ProposeContextChangeInput( + current_message=message, + conversation=(ContextConversationExcerpt(role="user", content=message),), + context=project_context(resolved_context), + fact_definitions=build_default_fact_registry().definitions(), + ) + + +def test_anthropic_sdk_accepts_configured_sampling_parameter(monkeypatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key") + + sync_parameters = inspect.signature(get_sync_client().messages.create).parameters + async_parameters = inspect.signature(get_async_client().messages.stream).parameters + + assert "temperature" in sync_parameters + assert "temperature" in async_parameters + + +def test_numerical_redraft_removes_unverified_derived_values(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[SimpleNamespace(type="text", text="Verified values only.")], + stop_reason="end_turn", + usage=_usage(input_tokens=1, output_tokens=1), + ) + + monkeypatch.setattr( + model_port, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + + response = asyncio.run( + AnthropicConversationModel(model="test-model").redraft_numerical( + draft="Tax is £7,486 and inferred take-home pay is £42,514.", + unsupported_claims=("£42,514",), + fact_summary="Income Tax: 7486 GBP/year", + ) + ) + + assert response.text == "Verified values only." + system = calls[0]["system"] + assert "Do not calculate a new total, difference, rate" in system + assert "do not repeat any expression listed as unsupported" in system + + +def test_context_interpreter_uses_one_declarative_fact_claim_route(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input={ + "expected_revision": 0, + "changes": [ + { + "kind": "fact_claim", + "concept": "age", + "definition_key": "person.age", + "subject_references": ["person:self"], + "relationship": "direct", + "value": {"kind": "integer", "value": 42}, + "scope_id": "scope:primary-household", + "evidence": "I am 42", + } + ], + }, + ) + ], + usage=_usage(), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run(AnthropicContextInterpreter().propose(_request("I am 42"))) + + assert result.claims[0].definition_key == "person.age" + assert result.claims[0].subject_references == ("person:self",) + assert result.usage.input_tokens == 3 + assert calls[0]["tool_choice"] == { + "type": "tool", + "name": "submit_context_change", + } + schema = calls[0]["tools"][0]["input_schema"] + schema_text = json.dumps(schema) + assert "changes" in schema_text + assert "candidate_entities" in schema_text + assert "operations" not in schema["properties"] + assert "unresolved_claims" not in schema_text + assert "ContextPatch" not in schema_text + system = calls[0]["system"].casefold() + assert "every supported assertion exactly once" in system + assert "do not infer an income source" in system + assert "apply a default" in system + assert '"current_message":"I am 42"' in calls[0]["messages"][0]["content"] + + +def test_context_proposal_reviewer_uses_exact_claim_ids_without_retained_facts( + monkeypatch, +): + calls = [] + claim = FactClaim( + claim_id="opaque-claim-42", + concept="age", + definition_key="person.age", + subject_references=("person:self",), + value={"kind": "integer", "value": 42}, + evidence="I am 42", + ) + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_semantic_review", + input={ + "reviews": [ + { + "claim_id": "opaque-claim-42", + "supported": True, + "reason": "The message directly states the age.", + "evidence": "I am 42", + } + ] + }, + ) + ], + usage=_usage(), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + request = ValidateContextChangeInput( + context=ConversationContext.initial("conversation"), + proposal=ContextChangeProposal( + expected_revision=0, + changes=(claim,), + ), + turn_id="turn", + evidence="I am 42", + ) + + result = asyncio.run(AnthropicContextProposalReviewer().review(request)) + + assert result.reviews[0].claim_id == "opaque-claim-42" + assert result.reviews[0].supported is True + payload = json.loads(calls[0]["messages"][0]["content"]) + assert set(payload) == { + "current_message", + "known_entities", + "active_scope_id", + "proposal", + } + assert "facts" not in json.dumps(payload["known_entities"]) + assert calls[0]["tool_choice"] == { + "type": "tool", + "name": "submit_context_semantic_review", + } + + +def test_context_interpreter_returns_additive_and_direct_values_in_same_claim_list( + monkeypatch, +): + class FakeMessages: + async def create(self, **_kwargs): + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input={ + "expected_revision": 2, + "changes": [ + { + "kind": "fact_claim", + "concept": "age", + "definition_key": "person.age", + "subject_references": ["person:self"], + "relationship": "direct", + "value": {"kind": "integer", "value": 26}, + "scope_id": "scope:primary-household", + "evidence": "I am 26.", + }, + { + "kind": "fact_claim", + "concept": "employment income", + "subject_references": [ + "person:self", + "entity:spouse", + ], + "relationship": "sum", + "value": { + "kind": "money", + "amount": "70000", + "period": "annual", + "currency": "GBP", + }, + "evidence": "We earn £70,000 together.", + } + ], + }, + ) + ], + usage=_usage(), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run( + AnthropicContextInterpreter().propose( + _request( + "I am 26. We earn £70,000 together.", + context=context_with_spouse(), + ) + ) + ) + + assert result.expected_revision == 2 + assert len(result.claims) == 2 + assert result.claims[0].relationship is FactClaimRelationship.DIRECT + assert result.claims[0].definition_key == "person.age" + assert result.claims[1].relationship is FactClaimRelationship.SUM + assert isinstance(result.claims[1].value, ClaimedMoneyValue) + assert result.claims[1].value.amount == 70000 + + +def test_context_interpreter_repairs_invalid_relationship_cardinality(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + subjects = ["person:self"] + if len(calls) == 2: + subjects.append("entity:spouse") + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input={ + "expected_revision": 2, + "changes": [ + { + "kind": "fact_claim", + "concept": "employment income", + "subject_references": subjects, + "relationship": "sum", + "value": { + "kind": "money", + "amount": "70000", + "period": "annual", + }, + "evidence": "We earn £70,000 together.", + } + ], + }, + ) + ], + usage=_usage(), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run( + AnthropicContextInterpreter().propose( + _request("We earn £70,000 together.", context=context_with_spouse()) + ) + ) + + assert len(calls) == 2 + assert result.claims[0].subject_references == ("person:self", "entity:spouse") + repair_content = calls[1]["messages"][0]["content"] + assert '"code": "invalid_context_submission"' in repair_content + assert '"changes", "0"' in repair_content + assert "do not ask the user a question" in repair_content + + +def test_context_validator_rejects_a_missing_monetary_claim_without_phrase_rules(): + context = context_with_spouse() + validator = ContextChangeValidator( + ContextReducer(build_default_fact_registry()), + build_default_fact_registry(), + ) + + result = validator.validate( + ValidateContextChangeInput( + context=context, + proposal=ContextChangeProposal(expected_revision=context.revision), + turn_id="turn", + evidence="What if we made 70k?", + ) + ) + + assert [issue.code for issue in result.issues] == [ + "missing_monetary_fact_claim" + ] + assert result.issues[0].evidence == "70k" + + +@pytest.mark.parametrize( + "message", + ( + "70k", + "70,000", + "70.000", + "70 000", + "70\u00a0000", + "70 thousand", + "seventy thousand", + "GBP 70,000", + "£70k", + "0.07 million", + ), +) +def test_fact_claim_value_validation_normalizes_monetary_forms(message): + claim = FactClaim( + concept="income", + subject_references=("person:self", "entity:spouse"), + relationship=FactClaimRelationship.SUM, + value=ClaimedMoneyValue(amount=70000), + evidence=message, + ) + + context = context_with_spouse() + registry = build_default_fact_registry() + result = ContextChangeValidator(ContextReducer(registry), registry).validate( + ValidateContextChangeInput( + context=context, + proposal=ContextChangeProposal( + expected_revision=context.revision, + changes=(claim,), + ), + turn_id="turn", + evidence=message, + ) + ) + + assert result.status is ContextValidationStatus.RESOLUTION_REQUIRED + assert {issue.code for issue in result.issues} == { + "authoritative_resolution_required" + } + + +def test_fact_claim_value_validation_rejects_missing_and_copied_values(): + copied = FactClaim( + concept="income", + subject_references=("person:self",), + value=ClaimedMoneyValue(amount=50000), + evidence="copied from an earlier turn", + ) + + context = context_with_spouse() + registry = build_default_fact_registry() + result = ContextChangeValidator(ContextReducer(registry), registry).validate( + ValidateContextChangeInput( + context=context, + proposal=ContextChangeProposal( + expected_revision=context.revision, + changes=(copied,), + ), + turn_id="turn", + evidence="What if we made 70k?", + ) + ) + + assert [issue.code for issue in result.issues] == [ + "uncited_fact_claim", + "missing_monetary_fact_claim", + "uncited_monetary_fact_claim", + ] + + +def test_text_fact_can_preserve_an_embedded_monetary_reform_value(): + message = "Set the personal allowance to £15,000." + context = ConversationContext.initial("conversation") + registry = build_default_fact_registry() + + result = ContextChangeValidator(ContextReducer(registry), registry).validate( + ValidateContextChangeInput( + context=context, + proposal=ContextChangeProposal( + expected_revision=context.revision, + candidate_entities=( + ContextEntityCandidate( + reference="new:scenario", + kind=EntityKind.POLICY_SCENARIO, + aliases=("the reform",), + ), + ), + changes=( + FactClaim( + concept="policy reform instruction", + definition_key="policy.reform_instruction", + subject_references=("new:scenario",), + scope_id="scope:primary-household", + value=TextFactValue(value=message), + evidence=message, + ), + ), + ), + turn_id="turn", + evidence=message, + ) + ) + + assert result.status is ContextValidationStatus.READY_TO_COMMIT + assert result.issues == () + + +def test_context_interpreter_keeps_periodless_direct_money_declarative(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input={ + "expected_revision": 0, + "changes": [ + { + "kind": "fact_claim", + "concept": "income", + "subject_references": ["person:self"], + "relationship": "direct", + "value": { + "kind": "money", + "amount": "50000", + "period": None, + }, + "evidence": "£50,000 of income", + } + ], + }, + ) + ], + usage=_usage(), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run( + AnthropicContextInterpreter().propose( + _request("How much tax would I pay on £50,000 of income?") + ) + ) + + assert len(calls) == 1 + assert len(result.claims) == 1 + assert isinstance(result.claims[0].value, ClaimedMoneyValue) + assert result.claims[0].value.period is None + assert "leave an unstated monetary period null" in calls[0]["system"].casefold() + + +def test_context_interpreter_rejects_server_owned_operations_and_retries(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + value = {"expected_revision": 0} + if len(calls) == 1: + value["operations"] = [] + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input=value, + ) + ], + usage=_usage(input_tokens=1, output_tokens=1), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run(AnthropicContextInterpreter().propose(_request("Thanks."))) + + assert len(calls) == 2 + assert result.claims == () + assert result.candidate_entities == () + + +def test_context_interpreter_returns_structured_issues_after_invalid_retry(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_context_change", + input={"expected_revision": 0, "operations": []}, + ) + ], + usage=_usage(input_tokens=2, output_tokens=1), + ) + + monkeypatch.setattr( + context_tools, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run(AnthropicContextInterpreter().propose(_request("Thanks."))) + + assert len(calls) == 2 + assert result.status is ContextProposalStatus.NEEDS_CLARIFICATION + assert result.claims == () + assert result.provider_attempts == 2 + assert result.usage.input_tokens == 4 + assert result.issues[0].code == "invalid_context_submission" + assert result.issues[0].path == ("operations",) + + +def test_variable_mapper_can_select_only_an_exact_catalogue_candidate(monkeypatch): + calls = [] + + class FakeMessages: + async def create(self, **kwargs): + calls.append(kwargs) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + name="submit_variable_mapping", + input={ + "status": "matched", + "variable_name": "employment_income", + "confidence": "high", + "target_period": "annual", + }, + ) + ], + usage=_usage(input_tokens=4, output_tokens=2), + ) + + monkeypatch.setattr( + variable_resolution, + "get_async_client", + lambda: SimpleNamespace(messages=FakeMessages()), + ) + result = asyncio.run( + AnthropicVariableMapper().select( + claim=FactClaim( + claim_id="combined-income", + concept="employment income", + value=ClaimedMoneyValue( + amount=70000, + period=MoneyPeriod.ANNUAL, + ), + subject_references=("person:self", "entity:spouse"), + relationship=FactClaimRelationship.SUM, + evidence="We earn £70,000 together.", + ), + candidates=( + PolicyEngineVariableCandidate( + name="employment_income", + label="Employment income", + entity="person", + definition_period="year", + value_type="float", + ), + ), + context=context_with_spouse(revision=1), + registry=build_default_fact_registry(), + validation_issues=( + ContextValidationIssue( + code="authoritative_resolution_required", + message="Select an authoritative mapping.", + claim_index=0, + ), + ), + ) + ) + + assert result.selection.status is MappingStatus.MATCHED + assert result.selection.confidence is MappingConfidence.HIGH + assert result.selection.variable_name == "employment_income" + payload = json.loads(calls[0]["messages"][0]["content"]) + assert payload["validation_issues"] == [ + { + "code": "authoritative_resolution_required", + "message": "Select an authoritative mapping.", + "path": [], + "claim_index": 0, + "operation_index": None, + "evidence": None, + } + ] diff --git a/backend/tests/test_api.py b/backend/tests/test_api.py index 6c9a622e..7add9c2c 100644 --- a/backend/tests/test_api.py +++ b/backend/tests/test_api.py @@ -4,19 +4,14 @@ Run inside the backend container: pytest tests/ """ -import asyncio import json import os import threading from concurrent.futures import ThreadPoolExecutor -from types import SimpleNamespace import pytest from fastapi.testclient import TestClient -from policyengine_observability.runtime import EVENT_LOGGER -from policyengine_observability.runtime import OPERATION_LOGGER from api.main import app -from chat.orchestrator import MAX_TOOL_RESULT_CHARS, _serialise_tool_result_for_model client = TestClient(app) @@ -27,27 +22,6 @@ ) -def test_oversized_tool_results_remain_valid_json(): - result_json = _serialise_tool_result_for_model( - { - "status": "success", - "result_id": "result-1", - "payload": "x" * (MAX_TOOL_RESULT_CHARS + 1), - } - ) - - result = json.loads(result_json) - assert len(result_json) <= MAX_TOOL_RESULT_CHARS - assert result == { - "status": "success", - "result_id": "result-1", - "note": ( - "Tool result exceeded the model context limit. Use a narrower " - "discovery query or request a more specific derivative output." - ), - } - - # --------------------------------------------------------------------------- # Health # --------------------------------------------------------------------------- @@ -375,68 +349,6 @@ def parse_sse(response_text: str) -> list[dict]: return events -def _anthropic_event(name: str, **attrs): - event = type(name, (), {})() - for key, value in attrs.items(): - setattr(event, key, value) - return event - - -class _FakeAnthropicStream: - def __init__(self, *, chunks=None, final_content=None, stop_reason="end_turn"): - self._events = [ - _anthropic_event( - "RawContentBlockDeltaEvent", - delta=SimpleNamespace(type="text_delta", text=chunk), - ) - for chunk in (chunks or []) - ] - self._final = SimpleNamespace( - content=final_content or [], - stop_reason=stop_reason, - ) - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - def __aiter__(self): - return self._iter_events() - - async def _iter_events(self): - for event in self._events: - yield event - - async def get_final_message(self): - return self._final - - -class _FakeAnthropicMessages: - def __init__(self, streams): - self._streams = list(streams) - self.calls = [] - - def stream(self, **kwargs): - self.calls.append(kwargs) - return self._streams.pop(0) - - -class _FakeAnthropicClient: - def __init__(self, streams): - self.messages = _FakeAnthropicMessages(streams) - - -def _tool_use_block(name: str, tool_input: dict, tool_id: str = "tool-1"): - return SimpleNamespace( - type="tool_use", - id=tool_id, - name=name, - input=tool_input, - ) - - @requires_live_anthropic class TestChatMessage: def test_simple_chat_returns_sse(self): @@ -511,7 +423,7 @@ def test_chunk_events_contain_text(self): full_text = "".join(e["content"] for e in chunks) assert len(full_text) > 5 - def test_tool_use_for_simulation_query(self): + def test_invocation_activity_for_policy_question(self): with client.stream( "POST", "/chat/message", @@ -519,7 +431,7 @@ def test_tool_use_for_simulation_query(self): "messages": [ { "role": "user", - "content": "What is the current personal allowance? Use get_baseline_parameters.", + "content": "How is the current personal allowance determined?", } ], }, @@ -527,7 +439,7 @@ def test_tool_use_for_simulation_query(self): text = r.read().decode() events = parse_sse(text) types = [e["type"] for e in events] - assert "tool_start" in types or "tool_use" in types + assert "invocation_activity" in types def test_no_error_event_on_simple_query(self): with client.stream( @@ -557,313 +469,6 @@ def test_usage_in_done_event(self): assert done["usage"]["input_tokens"] > 0 -class TestChatRouteWithMockedAnthropic: - def test_disabled_billing_does_not_check_or_record_credit(self, monkeypatch): - import chat.orchestrator as chatbot - - async def no_suggestions(*_args, **_kwargs): - return [] - - def fail_if_called(*_args, **_kwargs): - raise AssertionError("disabled billing touched the credit store") - - fake_client = _FakeAnthropicClient( - [ - _FakeAnthropicStream( - chunks=["The application remains available without billing."], - final_content=[], - ) - ] - ) - monkeypatch.setenv("BILLING_ENABLED", "false") - monkeypatch.setattr(chatbot, "get_async_client", lambda: fake_client) - monkeypatch.setattr(chatbot, "generate_followup_suggestions", no_suggestions) - monkeypatch.setattr("billing.check_balance", fail_if_called) - monkeypatch.setattr("billing.record_usage", fail_if_called) - - with client.stream( - "POST", - "/chat/message", - json={ - "messages": [{"role": "user", "content": "Explain income tax."}], - "user_id": "user-1", - }, - ) as response: - assert response.status_code == 200 - text = response.read().decode() - - done = next(event for event in parse_sse(text) if event["type"] == "done") - assert done["cost_gbp"] is None - assert done["balance"] is None - - def test_chat_route_executes_tool_loop_and_returns_final_answer(self, monkeypatch): - import chat.orchestrator as chatbot - - operation_records = [] - - async def no_suggestions(*_args, **_kwargs): - return [] - - tool_input = { - "year": 2026, - "people": [ - { - "age": 35, - "employment_income": 30000, - } - ], - "benunit": {}, - "household": {}, - } - fake_client = _FakeAnthropicClient( - [ - _FakeAnthropicStream( - final_content=[ - _tool_use_block("run_household_simulation", tool_input), - ], - ), - _FakeAnthropicStream( - chunks=[ - "For this illustrative household, net income is £25119.60." - ], - final_content=[], - ), - ] - ) - executed = [] - - def fake_execute_tool(tool_name, received_input, context=None): - executed.append((tool_name, received_input)) - return { - "status": "success", - "household": [{"net_income": 25119.60}], - "result_id": "household_simulation_1", - } - - monkeypatch.setattr(chatbot, "get_async_client", lambda: fake_client) - monkeypatch.setattr(chatbot, "generate_followup_suggestions", no_suggestions) - monkeypatch.setattr(chatbot, "execute_tool", fake_execute_tool) - monkeypatch.setattr(OPERATION_LOGGER, "info", operation_records.append) - usage_calls = [] - - def fake_record_usage(**kwargs): - usage_calls.append(kwargs) - return {"cost_gbp": 0.0, "balance": 10.0} - - monkeypatch.setattr("billing.record_usage", fake_record_usage) - - with client.stream( - "POST", - "/chat/message", - json={ - "messages": [{"role": "user", "content": "Calculate this household."}] - }, - ) as response: - assert response.status_code == 200 - text = response.read().decode() - - events = parse_sse(text) - assert [ - event["type"] - for event in events - if event["type"] in {"tool_use", "tool_result", "done"} - ] == [ - "tool_use", - "tool_result", - "done", - ] - done = next(event for event in events if event["type"] == "done") - assert "£25119.60" in done["content"] - assert "timings" not in done - assert usage_calls - assert "timings" not in usage_calls[0] - assert "timings_ms" not in usage_calls[0] - assert "timing_counts" not in usage_calls[0] - assert executed == [("run_household_simulation", tool_input)] - assert "tools" in fake_client.messages.calls[0] - second_messages = fake_client.messages.calls[1]["messages"] - assert second_messages[-1]["content"][0]["type"] == "tool_result" - turn_log = next( - payload - for payload in map(json.loads, operation_records) - if payload.get("operation") == "chat.turn" - ) - assert turn_log["event"] == "operation_completed" - assert turn_log["gateway_route"] == "compute" - assert turn_log["gateway_outcome"] == "ready" - assert turn_log["model"] - assert turn_log["ttft_ms"] >= 0 - assert turn_log["timings_ms"]["gateway.classify"] >= 0 - assert turn_log["timings_ms"]["gateway.plan_serialize"] >= 0 - assert turn_log["timings_ms"]["model.select"] >= 0 - assert turn_log["timings_ms"]["system.build"] >= 0 - assert turn_log["timings_ms"]["tool_schema.build"] >= 0 - assert turn_log["timings_ms"]["model.iteration"] >= 0 - assert turn_log["timings_ms"]["model.stream"] >= 0 - assert turn_log["timings_ms"]["tool.execute"] >= 0 - assert turn_log["timings_ms"]["billing.record_usage"] >= 0 - assert turn_log["timing_counts"]["gateway.classify"] == 1 - assert turn_log["timing_counts"]["gateway.plan_serialize"] == 1 - assert turn_log["timing_counts"]["model.select"] == 1 - assert turn_log["timing_counts"]["system.build"] == 1 - assert turn_log["timing_counts"]["tool_schema.build"] == 1 - assert turn_log["timing_counts"]["model.iteration"] == 2 - assert turn_log["timing_counts"]["model.stream"] == 2 - assert turn_log["timing_counts"]["tool.execute"] == 1 - assert turn_log["timing_counts"]["billing.record_usage"] == 1 - - def test_chat_route_logs_client_disconnect(self, monkeypatch): - import chat.orchestrator as chatbot - from chat.public_service import start_public_chat - from chat.schemas import ChatRequest - - class DisconnectingRequest: - async def is_disconnected(self): - return True - - operation_records = [] - event_records = [] - monkeypatch.setattr(chatbot, "get_async_client", lambda: object()) - monkeypatch.setattr(chatbot, "is_followup", lambda _conversation: True) - monkeypatch.setattr(OPERATION_LOGGER, "info", operation_records.append) - monkeypatch.setattr(EVENT_LOGGER, "info", event_records.append) - - async def consume_stream(): - stream = await start_public_chat( - ChatRequest( - messages=[ - { - "role": "user", - "content": "Calculate this household.", - } - ], - session_id="disconnect-session", - ), - is_cancelled=DisconnectingRequest().is_disconnected, - ) - return [chunk async for chunk in stream] - - assert asyncio.run(consume_stream()) == [] - - turn_log = next( - payload - for payload in map(json.loads, operation_records) - if payload.get("operation") == "chat.turn" - ) - assert turn_log["stop_reason"] == "client_disconnected" - assert turn_log["session_id"] == "disconnect-session" - assert turn_log["iterations"] == 0 - assert "gateway.classify" not in turn_log["timing_counts"] - assert "model.select" not in turn_log["timing_counts"] - assert "system.build" not in turn_log["timing_counts"] - assert "tool_schema.build" not in turn_log["timing_counts"] - - disconnect_event = next( - payload - for payload in map(json.loads, event_records) - if payload.get("event") == "chat.client_disconnected" - ) - assert disconnect_event["session_id"] == "disconnect-session" - assert disconnect_event["iterations"] == 0 - assert disconnect_event["tool_calls"] == 0 - - def test_chat_route_logs_error_on_chat_turn(self, monkeypatch): - import chat.orchestrator as chatbot - from chat.public_service import start_public_chat - from chat.schemas import ChatRequest - - class ConnectedRequest: - async def is_disconnected(self): - return False - - operation_records = [] - - def raise_model_selection(*_args, **_kwargs): - raise RuntimeError("model selection failed") - - monkeypatch.setattr(chatbot, "get_async_client", lambda: object()) - monkeypatch.setattr(chatbot, "is_followup", lambda _conversation: True) - monkeypatch.setattr(chatbot, "select_chat_model", raise_model_selection) - # The installed observability version decides which level receives - # handled-error operation logs; capture all common levels. - monkeypatch.setattr(OPERATION_LOGGER, "info", operation_records.append) - monkeypatch.setattr(OPERATION_LOGGER, "warning", operation_records.append) - monkeypatch.setattr(OPERATION_LOGGER, "error", operation_records.append) - - async def consume_stream(): - stream = await start_public_chat( - ChatRequest( - messages=[ - { - "role": "user", - "content": "Calculate this household.", - } - ], - session_id="error-session", - ), - is_cancelled=ConnectedRequest().is_disconnected, - ) - return "".join([chunk async for chunk in stream]) - - events = parse_sse(asyncio.run(consume_stream())) - - # The SSE error event carries a generic message plus the session id as - # a correlation reference — never the raw exception text. - assert len(events) == 1 - error_event = events[0] - assert error_event["type"] == "error" - assert "model selection failed" not in error_event["content"] - assert "Something went wrong" in error_event["content"] - assert "error-session" in error_event["content"] - turn_log = next( - payload - for payload in map(json.loads, operation_records) - if payload.get("operation") == "chat.turn" - ) - assert turn_log["event"] == "operation_failed" - assert turn_log["stop_reason"] == "error" - assert turn_log["session_id"] == "error-session" - assert turn_log["iterations"] == 0 - assert turn_log["tool_calls"] == 0 - assert turn_log["timing_counts"]["model.select"] == 1 - - def test_chat_route_error_event_hides_exception_details(self, monkeypatch): - """Regression test for information disclosure: raw exception text - (Anthropic SDK / httpx / Supabase strings can embed internal URLs, - file paths, and provider payloads) must never reach the SSE `error` - event. Users get a generic message with the session id as a - correlation reference; the full detail stays in server logs.""" - import chat.orchestrator as chatbot - - class RaisingAnthropicMessages: - def stream(self, **_kwargs): - raise Exception("secret-internal-detail-xyz") - - fake_client = SimpleNamespace(messages=RaisingAnthropicMessages()) - - monkeypatch.setattr(chatbot, "get_async_client", lambda: fake_client) - monkeypatch.setattr(chatbot, "is_followup", lambda _conversation: True) - - with client.stream( - "POST", - "/chat/message", - json={ - "messages": [ - {"role": "user", "content": "Calculate this household."} - ], - "session_id": "leak-test-session", - }, - ) as response: - assert response.status_code == 200 - text = response.read().decode() - - assert "secret-internal-detail-xyz" not in text - events = parse_sse(text) - error_event = next(event for event in events if event["type"] == "error") - assert "Something went wrong" in error_event["content"] - assert "leak-test-session" in error_event["content"] - - # --------------------------------------------------------------------------- # Rate limiting # --------------------------------------------------------------------------- @@ -976,13 +581,21 @@ def test_chat_endpoint_exposes_starlette_request_to_slowapi(self): "the `request` parameter must be a starlette Request, not the body model" ) - def test_chat_message_does_not_500_from_rate_limit_decorator(self): + def test_chat_message_does_not_500_from_rate_limit_decorator(self, monkeypatch): """A single call must stream normally, not 500. conftest.py raises the limits far above test workload, so the only thing that can fail here is the rate-limit decorator itself — which is exactly the regression this guards. """ + async def fake_start(*_args, **_kwargs): + async def stream(): + yield 'data: {"type":"done","content":"hi"}\n\n' + + return stream() + + monkeypatch.setattr("chat.routes.start_public_chat", fake_start) + with client.stream( "POST", "/chat/message", diff --git a/backend/tests/test_capability_artifacts.py b/backend/tests/test_capability_artifacts.py new file mode 100644 index 00000000..053c9ddd --- /dev/null +++ b/backend/tests/test_capability_artifacts.py @@ -0,0 +1,184 @@ +from datetime import datetime, timezone + +import pytest +from pydantic import ValidationError + +from capabilities.artifacts import ( + AggregateValue, + ArtifactProvenance, + PolicyChange, + PolicyScenarioRef, + SocietyAnalysisResultRef, +) +from capabilities.compatibility import ( + ArtifactRequirements, + check_artifact_compatibility, +) +from chat.artifact_context import sanitized_artifact_summary + + +def provenance() -> ArtifactProvenance: + return ArtifactProvenance( + conversation_id="conversation-1", + turn_id="turn-1", + capability_id="policy_reform", + capability_version="1", + invocation_id="invocation-1", + sources=("current user instruction",), + ) + + +def test_policy_scenario_is_immutable_and_carries_compatibility_metadata(): + scenario = PolicyScenarioRef( + artifact_id="scenario-1", + created_at=datetime.now(timezone.utc), + provenance=provenance(), + year=2026, + scenario_revision="revision-1", + catalogue_version="catalogue-1", + calculation_engine_version="engine-1", + baseline=False, + verified_changes=( + PolicyChange( + parameter_path="gov.example.amount", + value=15_000, + effective_date="2026-01-01", + ), + ), + ) + + assert scenario.artifact_type == "policy_scenario" + assert scenario.verified_changes[0].value == 15_000 + with pytest.raises(ValidationError): + scenario.year = 2027 # type: ignore[misc] + with pytest.raises(TypeError): + scenario.verified_changes[0] = PolicyChange( # type: ignore[index] + parameter_path="changed", + value=1, + ) + + +def test_compatibility_checks_only_declared_consumer_requirements(): + result = SocietyAnalysisResultRef( + artifact_id="result-1", + provenance=provenance(), + year=2026, + policy_scenario_artifact_id="scenario-1", + scenario_revision="revision-1", + catalogue_version="catalogue-1", + dataset_version="dataset-1", + calculation_engine_version="engine-1", + default_profile_version="default-1", + calculated_output_ids=("budgetary_impact",), + outputs=( + AggregateValue( + output_id="budgetary_impact", + metric_id="net_cost", + label="Net budget cost", + value=1_000_000, + unit="GBP/year", + ), + ), + ) + + compatible = check_artifact_compatibility( + result, + ArtifactRequirements( + artifact_type="society_analysis_result", + schema_version="1", + year=2026, + scenario_revision="revision-1", + dataset_version="dataset-1", + calculation_engine_version="engine-1", + ), + ) + incompatible = check_artifact_compatibility( + result, + ArtifactRequirements( + artifact_type="society_analysis_result", + schema_version="1", + year=2025, + scenario_revision="revision-2", + calculation_engine_version="engine-2", + ), + ) + + assert compatible.compatible is True + assert compatible.issues == () + assert incompatible.compatible is False + assert {issue.split(":", 1)[0] for issue in incompatible.issues} == { + "year", + "scenario_revision", + "calculation_engine_version", + } + summary = sanitized_artifact_summary(result) + assert summary["artifact_type"] == "society_analysis_result" + assert summary["outputs"] == ( + { + "output_id": "budgetary_impact", + "metric_id": "net_cost", + "label": "Net budget cost", + "value": 1_000_000, + "unit": "GBP/year", + "dimensions": [], + }, + ) + assert "provenance" not in summary + + +@pytest.mark.parametrize( + ("field", "stale_value"), + [ + ("artifact_type", "household_result"), + ("schema_version", "2"), + ("year", 2025), + ("scenario_revision", "stale-revision"), + ("catalogue_version", "stale-catalogue"), + ("dataset_version", "stale-dataset"), + ("calculation_engine_version", "stale-engine"), + ], +) +def test_each_declared_compatibility_dimension_rejects_a_stale_reference( + field, + stale_value, +): + result = SocietyAnalysisResultRef( + artifact_id="result-property", + provenance=provenance(), + year=2026, + policy_scenario_artifact_id="scenario-property", + scenario_revision="revision-current", + catalogue_version="catalogue-current", + dataset_version="dataset-current", + calculation_engine_version="engine-current", + default_profile_version="1", + calculated_output_ids=("budgetary_impact",), + outputs=( + AggregateValue( + output_id="budgetary_impact", + metric_id="net_cost", + label="Net cost", + value=1, + unit="GBP/year", + ), + ), + ) + requirement_values = { + "artifact_type": "society_analysis_result", + "schema_version": "1", + "year": 2026, + "scenario_revision": "revision-current", + "catalogue_version": "catalogue-current", + "dataset_version": "dataset-current", + "calculation_engine_version": "engine-current", + } + requirement_values[field] = stale_value + + compatibility = check_artifact_compatibility( + result, + ArtifactRequirements.model_validate(requirement_values), + ) + + assert compatibility.compatible is False + assert len(compatibility.issues) == 1 + assert compatibility.issues[0].startswith(f"{field}:") diff --git a/backend/tests/test_capability_chat_service.py b/backend/tests/test_capability_chat_service.py new file mode 100644 index 00000000..165bc153 --- /dev/null +++ b/backend/tests/test_capability_chat_service.py @@ -0,0 +1,865 @@ +from __future__ import annotations + +import asyncio +import json + +from pydantic import BaseModel, ConfigDict +from sqlmodel import SQLModel, create_engine + +from capabilities.composition import compose_runtime +from capabilities.contracts import Capability, CapabilitySpec, Completed, NeedsInput +from capabilities.relevance import ( + AssessRelevanceTool, + ConversationRelevanceCapability, + RelevanceAssessment, + RelevanceResult, +) +from chat.capability_service import ChatTurnService +from chat.events import ( + InvocationActivity, + TextChunk, + TurnCancelled, + TurnCompleted, + TurnFailed, +) +from chat.model_port import ( + ConversationModelResponse, + ModelCapabilityCall, + ModelUsage, +) +from chat.turn_input import ChatTurnInput +from persistence.idempotency import SQLIdempotencyRepository +from tools.analysis_support import NumericalFact, VerifyNumericalResponseTool +from tools.contracts import CallerType, Visibility + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class EchoInput(StrictModel): + text: str + optional_note: str | None = None + + +class EchoOutput(StrictModel): + text: str + narration_facts: tuple[NumericalFact, ...] = () + assumption_statements: tuple[str, ...] = () + narration_fallback: str | None = None + + +class EchoCapability(Capability[EchoInput, EchoOutput]): + spec = CapabilitySpec( + identifier="echo_capability", + version="1", + description="Return a typed test fact.", + required_use="Use when the test asks for an echo.", + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.MODEL}), + input_model=EchoInput, + output_model=EchoOutput, + ) + + async def run(self, capability_input, context): + del context + facts = () + if capability_input.text == "numeric": + facts = ( + NumericalFact( + label="Net cost", + value=1_200_000_000, + unit="GBP/year", + ), + ) + assumptions = ( + ("The household has no childcare expenses.",) + if capability_input.text == "requirements" + else () + ) + return Completed( + value=EchoOutput( + text=capability_input.text, + narration_facts=facts, + assumption_statements=assumptions, + narration_fallback=( + "### Results\n\n- Net cost: £1.2 billion." + if capability_input.text == "numeric" + else None + ), + ) + ) + + +class ClarifyCapability(Capability[EchoInput, EchoOutput]): + spec = EchoCapability.spec.model_copy( + update={ + "identifier": "clarify_capability", + "description": "Request one missing test value.", + } + ) + + async def run(self, capability_input, context): + del capability_input, context + return NeedsInput( + prompt="Which amount?", + missing_fields=("text",), + ) + + +class FailingCapability(Capability[EchoInput, EchoOutput]): + spec = EchoCapability.spec.model_copy( + update={ + "identifier": "failing_capability", + "description": "Fail safely for a localized-failure test.", + } + ) + + async def run(self, capability_input, context): + del capability_input, context + raise RuntimeError("private provider failure") + + +class FakeRelevanceAssessor: + def __init__(self, result=RelevanceResult.RELEVANT): + self.result = result + self.requests = [] + + async def assess(self, request): + self.requests.append(request) + return RelevanceAssessment( + result=self.result, + explanation="bounded test result", + usage={"input_tokens": 2, "output_tokens": 1}, + ) + + +class FakeConversationModel: + def __init__(self, responses, redraft=""): + self.responses = list(responses) + self.requests = [] + self.redraft = redraft + self.redraft_calls = [] + + async def respond(self, request): + self.requests.append(request) + return self.responses.pop(0) + + async def redraft_numerical( + self, + *, + draft, + unsupported_claims, + fact_summary, + ): + self.redraft_calls.append((draft, unsupported_claims, fact_summary)) + return ConversationModelResponse( + text=self.redraft, + model="fake-model", + usage=ModelUsage(input_tokens=3, output_tokens=2), + ) + + +async def not_cancelled() -> bool: + return False + + +def _runtime(model, assessor, *, idempotency=None): + composition = compose_runtime( + tools=[AssessRelevanceTool(assessor), VerifyNumericalResponseTool()], + capabilities=[ + ConversationRelevanceCapability(), + EchoCapability(), + ClarifyCapability(), + FailingCapability(), + ], + ) + service = ChatTurnService( + executor=composition.executor, + capabilities=composition.capabilities, + model=model, + idempotency=idempotency, + ) + context = composition.executor.context( + request_id="request-1", + conversation_id="conversation-1", + turn_id="turn-1", + is_cancelled=not_cancelled, + ) + return composition, service, context + + +def _turn(content="Hello", *, turn_id="turn-1", debug=False): + return ChatTurnInput( + messages=[ + {"role": "user", "content": "Earlier question"}, + {"role": "assistant", "content": "Earlier answer"}, + {"role": "user", "content": content}, + ], + session_id="conversation-1", + turn_id=turn_id, + debug=debug, + ) + + +def _collect(service, turn, context, cancellation=not_cancelled): + async def collect(): + return [ + event + async for event in service.run( + turn, + is_cancelled=cancellation, + context=context, + ) + ] + + return asyncio.run(collect()) + + +def test_direct_answer_keeps_full_history_and_private_relevance_out_of_model_tools(): + assessor = FakeRelevanceAssessor() + model = FakeConversationModel( + [ + ConversationModelResponse( + text="Natural direct answer.", + model="fake-model", + stop_reason="end_turn", + usage=ModelUsage(input_tokens=5, output_tokens=4), + ) + ] + ) + _composition, service, context = _runtime(model, assessor) + + events = _collect(service, _turn(), context) + + assert [type(event) for event in events] == [TextChunk, TurnCompleted] + assert events[0].content == "Natural direct answer." + assert len(assessor.requests) == 1 + assert assessor.requests[0].current_message == "Hello" + assert len(model.requests[0].messages) == 3 + capability_ids = { + item["identifier"] for item in model.requests[0].capabilities + } + assert capability_ids == { + "echo_capability", + "clarify_capability", + "failing_capability", + } + assert "conversation_relevance" not in capability_ids + assert "policy_information" in model.requests[0].system + assert "household_analysis" in model.requests[0].system + assert "society_analysis" in model.requests[0].system + assert "Do not conduct a" in model.requests[0].system + assert "separate informal household intake" in model.requests[0].system + assert "sole authority for" in model.requests[0].system + assert "which household details require clarification" in model.requests[0].system + assert "Assumptions used" in model.requests[0].system + assert "Markdown bullet" in model.requests[0].system + assert "referenced_household_id" in model.requests[0].system + assert "same household" in model.requests[0].system + assert "every category in required_output_ids" in model.requests[0].system + assert "Do not infer policy mechanisms" in model.requests[0].system + assert events[-1].usage.input_tokens == 7 + assert events[-1].usage.output_tokens == 5 + + +def test_clearly_out_of_scope_turn_is_local_and_skips_conversation_model(): + assessor = FakeRelevanceAssessor(RelevanceResult.CLEARLY_OUT_OF_SCOPE) + model = FakeConversationModel([]) + _composition, service, context = _runtime(model, assessor) + + events = _collect(service, _turn("Calculate Canadian tax"), context) + + assert [type(event) for event in events] == [TextChunk, TurnCompleted] + assert events[-1].outcome == "out_of_scope" + assert model.requests == [] + assert len(assessor.requests) == 1 + + +def test_multiple_capability_outcomes_return_to_same_model_without_global_waiting(): + assessor = FakeRelevanceAssessor() + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-1", + capability_id="echo_capability", + input={"text": "complete"}, + ), + ModelCapabilityCall( + call_id="call-2", + capability_id="clarify_capability", + input={"text": "missing"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="I completed one part. Which amount should I use for the other?", + model="fake-model", + stop_reason="end_turn", + ), + ] + ) + _composition, service, context = _runtime(model, assessor) + + events = _collect(service, _turn(), context) + + assert [type(event) for event in events] == [ + InvocationActivity, + InvocationActivity, + InvocationActivity, + InvocationActivity, + TextChunk, + TurnCompleted, + ] + statuses = [ + event.record.status.value + for event in events + if isinstance(event, InvocationActivity) and event.phase == "finished" + ] + assert statuses == ["completed", "needs_input"] + result_blocks = model.requests[1].messages[-1]["content"] + assert '"status": "completed"' in result_blocks[0]["content"] + assert '"status": "needs_input"' in result_blocks[1]["content"] + + +def test_debug_projection_changes_activity_visibility_but_not_capability_result(): + responses = [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-1", + capability_id="echo_capability", + input={"text": "same calculation"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="Same grounded result.", + model="fake-model", + stop_reason="end_turn", + ), + ] + normal_model = FakeConversationModel(responses) + debug_model = FakeConversationModel(responses) + normal_composition, normal_service, normal_context = _runtime( + normal_model, + FakeRelevanceAssessor(), + ) + debug_composition, debug_service, debug_context = _runtime( + debug_model, + FakeRelevanceAssessor(), + ) + + normal_events = _collect(normal_service, _turn(debug=False), normal_context) + debug_events = _collect(debug_service, _turn(debug=True), debug_context) + + assert next(event.content for event in normal_events if isinstance(event, TextChunk)) == ( + "Same grounded result." + ) + assert next(event.content for event in debug_events if isinstance(event, TextChunk)) == ( + "Same grounded result." + ) + normal_ids = { + event.record.identifier + for event in normal_events + if isinstance(event, InvocationActivity) + } + debug_ids = { + event.record.identifier + for event in debug_events + if isinstance(event, InvocationActivity) + } + assert normal_ids == {"echo_capability"} + assert debug_ids == { + "conversation_relevance", + "assess_relevance", + "echo_capability", + } + assert all( + event.record.debug_input is None and event.record.debug_output is None + for event in normal_events + if isinstance(event, InvocationActivity) + ) + assert all( + event.record.debug_input is not None + for event in debug_events + if isinstance(event, InvocationActivity) + ) + assert all( + event.record.debug_output is not None + for event in debug_events + if isinstance(event, InvocationActivity) and event.phase == "finished" + ) + assert { + record.identifier + for record in normal_composition.tracer.records( + "conversation-1", + include_private=True, + ) + } == debug_ids + assert { + record.identifier + for record in debug_composition.tracer.records( + "conversation-1", + include_private=True, + ) + } == debug_ids + + +def test_model_capability_trace_matches_the_json_exchanged_with_the_model(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-1", + capability_id="clarify_capability", + input={"text": "missing"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="Which amount?", + model="fake-model", + stop_reason="end_turn", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(debug=True), context) + + finished = next( + event.record + for event in events + if isinstance(event, InvocationActivity) + and event.phase == "finished" + and event.record.identifier == "clarify_capability" + ) + model_result = json.loads( + model.requests[1].messages[-1]["content"][0]["content"] + ) + assert finished.debug_input == {"text": "missing"} + assert "optional_note" not in finished.debug_input + assert finished.debug_output == model_result + assert finished.debug_output["response_guidance"].startswith( + "Ask the supplied prompt" + ) + + +def test_completed_capability_tells_model_to_report_results_before_follow_up(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-completed", + capability_id="echo_capability", + input={"text": "completed result"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="Here is the completed result.", + model="fake-model", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + _collect(service, _turn(), context) + + result = json.loads(model.requests[1].messages[-1]["content"][0]["content"]) + assert result["status"] == "completed" + assert "Answer the current request now" in result["response_guidance"] + assert "Do not ask for input or an output choice" in ( + result["response_guidance"] + ) + + +def test_model_result_omits_verifier_only_narration_fields(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-numeric", + capability_id="echo_capability", + input={"text": "numeric"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="Net cost is £1.2 billion.", + model="fake-model", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(debug=True), context) + + model_result = json.loads( + model.requests[1].messages[-1]["content"][0]["content"] + ) + assert "narration_facts" not in model_result["value"] + assert "narration_fallback" not in model_result["value"] + assert model_result["value"]["text"] == "numeric" + finished = next( + event.record + for event in events + if isinstance(event, InvocationActivity) + and event.phase == "finished" + and event.record.identifier == "echo_capability" + ) + assert finished.debug_output == model_result + + +def test_failed_capability_is_returned_to_model_without_blocking_sibling_call(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-failed", + capability_id="failing_capability", + input={"text": "fail"}, + ), + ModelCapabilityCall( + call_id="call-complete", + capability_id="echo_capability", + input={"text": "continue"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="One operation failed, while the other completed.", + model="fake-model", + stop_reason="end_turn", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(debug=True), context) + + result_blocks = model.requests[1].messages[-1]["content"] + assert '"status": "failed"' in result_blocks[0]["content"] + assert "private provider failure" not in result_blocks[0]["content"] + assert '"status": "completed"' in result_blocks[1]["content"] + failed_trace = next( + event.record + for event in events + if isinstance(event, InvocationActivity) + and event.phase == "finished" + and event.record.identifier == "failing_capability" + ) + assert failed_trace.debug_output == json.loads(result_blocks[0]["content"]) + assert any( + isinstance(event, TextChunk) + and event.content == "One operation failed, while the other completed." + for event in events + ) + + +def test_needs_input_cannot_be_retried_or_replaced_with_estimated_numbers(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-needs-input", + capability_id="clarify_capability", + input={"text": "missing"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-retry", + capability_id="clarify_capability", + input={"text": "invented"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="The estimated answer is £15,000.", + model="fake-model", + stop_reason="end_turn", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(), context) + + finished_clarifications = [ + event + for event in events + if isinstance(event, InvocationActivity) + and event.phase == "finished" + and event.record.identifier == "clarify_capability" + ] + assert len(finished_clarifications) == 1 + assert any( + definition["identifier"] == "clarify_capability" + for definition in model.requests[0].capabilities + ) + assert all( + definition["identifier"] != "clarify_capability" + for request in model.requests[1:] + for definition in request.capabilities + ) + assert len(model.requests) == 2 + repeated_result = json.loads( + model.requests[1].messages[-1]["content"][0]["content"] + ) + assert repeated_result["status"] == "needs_input" + assert repeated_result["prompt"] == "Which amount?" + assert "safe_message" not in repeated_result + assert "capability_repeated_without_new_user_evidence" not in str( + repeated_result + ) + assert events[-2].content == "Which amount?" + assert "15,000" not in events[-2].content + assert model.redraft_calls == [] + + +def test_natural_numbered_clarification_skips_calculation_verification(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-needs-input", + capability_id="clarify_capability", + input={"text": "missing"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text=( + "I need a little more information:\n\n" + "1. What is your age?\n" + "2. Which amount should I use?" + ), + model="fake-model", + stop_reason="end_turn", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(debug=True), context) + + assert events[-2].content == ( + "I need a little more information:\n\n" + "1. What is your age?\n" + "2. Which amount should I use?" + ) + assert not any( + isinstance(event, InvocationActivity) + and event.record.identifier == "verify_numerical_response" + for event in events + ) + + +def test_quantitative_response_gets_one_free_form_correction_and_usage_accounting(): + assessor = FakeRelevanceAssessor() + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-1", + capability_id="echo_capability", + input={"text": "numeric"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="The reform costs £2 billion.", + model="fake-model", + usage=ModelUsage(input_tokens=4, output_tokens=3), + ), + ], + redraft="The reform costs £1.2 billion.", + ) + _composition, service, context = _runtime(model, assessor) + + events = _collect(service, _turn(), context) + + assert events[-2].content == "The reform costs £1.2 billion." + assert len(model.redraft_calls) == 1 + assert events[-1].usage.input_tokens == 9 + assert events[-1].usage.output_tokens == 6 + + +def test_remaining_unsupported_sentence_is_removed_before_fact_list_fallback(): + draft = ( + "The verified net cost is £1.2 billion.\n\n" + "An unsupported estimate is £2 billion.\n\n" + "The policy direction is unchanged." + ) + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-sanitize", + capability_id="echo_capability", + input={"text": "numeric"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse(text=draft, model="fake-model"), + ], + redraft=draft, + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(), context) + + assert events[-2].content == ( + "The verified net cost is £1.2 billion.\n\n" + "The policy direction is unchanged." + ) + assert "£2 billion" not in events[-2].content + assert not events[-2].content.startswith("### Results") + assert len(model.redraft_calls) == 1 + + +def test_missing_assumptions_are_appended_as_a_markdown_list(): + model = FakeConversationModel( + [ + ConversationModelResponse( + capability_calls=( + ModelCapabilityCall( + call_id="call-requirements", + capability_id="echo_capability", + input={"text": "requirements"}, + ), + ), + model="fake-model", + ), + ConversationModelResponse( + text="Here is the naturally written result.", + model="fake-model", + ), + ] + ) + _composition, service, context = _runtime(model, FakeRelevanceAssessor()) + + events = _collect(service, _turn(), context) + + assert events[-2].content == ( + "Here is the naturally written result.\n\n" + "### Assumptions used\n\n" + "- The household has no childcare expenses." + ) + + +def test_missing_assumption_is_added_to_an_existing_markdown_list(): + response = ( + "Here is the result.\n\n" + "### Assumptions used\n\n" + "- Policy year: 2026.\n\n" + "### Next steps\n\n" + "You can revise the household." + ) + + completed = ChatTurnService._ensure_assumption_list( + response, + ["Policy year: 2026.", "The household lives in England."], + ) + + assert completed.count("### Assumptions used") == 1 + assert "- Policy year: 2026." in completed + assert "- The household lives in England.\n### Next steps" in completed + + +def test_bold_assumption_heading_is_not_duplicated(): + response = ( + "Here is the result.\n\n" + "**Assumptions used**\n\n" + "- Policy year: 2026.\n" + "- The household lives in England." + ) + + completed = ChatTurnService._ensure_assumption_list( + response, + ["Policy year: 2026.", "The household lives in England."], + ) + + assert completed == response + assert completed.casefold().count("assumptions used") == 1 + + +def test_request_cancellation_stops_before_relevance_or_model(): + assessor = FakeRelevanceAssessor() + model = FakeConversationModel([]) + _composition, service, context = _runtime(model, assessor) + + async def cancelled() -> bool: + return True + + events = _collect(service, _turn(), context, cancelled) + + assert len(events) == 1 + assert isinstance(events[0], TurnCancelled) + assert assessor.requests == [] + assert model.requests == [] + + +def test_turn_idempotency_replays_and_rejects_conflicting_input(tmp_path): + engine = create_engine( + f"sqlite:///{tmp_path / 'idempotency.sqlite'}", + connect_args={"check_same_thread": False}, + ) + SQLModel.metadata.create_all(engine) + idempotency = SQLIdempotencyRepository(engine=engine) + assessor = FakeRelevanceAssessor() + first_model = FakeConversationModel( + [ConversationModelResponse(text="First answer", model="fake-model")] + ) + _composition, service, context = _runtime( + first_model, + assessor, + idempotency=idempotency, + ) + first = _collect(service, _turn(), context) + + replay_model = FakeConversationModel([]) + _composition, replay_service, replay_context = _runtime( + replay_model, + assessor, + idempotency=idempotency, + ) + replay = _collect(replay_service, _turn(), replay_context) + conflict = _collect( + replay_service, + _turn("Different input"), + replay_context, + ) + + assert first[-1].outcome == "completed" + assert replay[-1].outcome == "replay" + assert replay[-1].content == "First answer" + assert isinstance(conflict[-1], TurnFailed) + assert conflict[-1].stop_reason == "idempotency_conflict" + assert replay_model.requests == [] + assert len(assessor.requests) == 1 diff --git a/backend/tests/test_capability_composition.py b/backend/tests/test_capability_composition.py new file mode 100644 index 00000000..7ba77d1e --- /dev/null +++ b/backend/tests/test_capability_composition.py @@ -0,0 +1,402 @@ +"""Unit and architecture tests for typed capability composition.""" + +from __future__ import annotations + +import asyncio + +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from capabilities.composition import compose_runtime +from capabilities.application import build_capability_chat_application +from capabilities.contracts import ( + ArtifactContract, + Capability, + CapabilityDependency, + CapabilitySpec, + Completed, + NeedsInput, +) +from capabilities.executor import InvocationCancelled +from capabilities.registry import CapabilityRegistry +from capabilities.tracing import InvocationKind, InvocationStatus +from tools.contracts import CallerType, Tool, ToolSpec, Visibility + + +def test_concrete_application_composes_every_registered_operation(tmp_path): + from sqlmodel import SQLModel, create_engine + + engine = create_engine(f"sqlite:///{tmp_path / 'application.sqlite'}") + SQLModel.metadata.create_all(engine) + + application = build_capability_chat_application(engine=engine) + + assert {spec.identifier for spec in application.composition.capabilities.specs()} == { + "conversation_relevance", + "policy_information", + "policy_reform", + "household_analysis", + "society_analysis", + "analysis_follow_up", + "society_chart", + } + assert len(application.composition.tools.specs()) == 32 + operation_ids = { + spec.identifier for spec in application.composition.tools.specs() + } + assert { + "propose_context_change", + "validate_context_change", + "resolve_context_change", + "apply_context_change", + } <= operation_ids +from tools.registry import ToolRegistry + + +class StrictModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class NumberInput(StrictModel): + value: int + + +class NumberOutput(StrictModel): + value: int + + +class AddOneTool(Tool[NumberInput, NumberOutput]): + spec = ToolSpec( + identifier="add_one", + version="1", + description="Add one to a number.", + visibility=Visibility.PRIVATE, + allowed_callers=frozenset({CallerType.CAPABILITY}), + input_model=NumberInput, + output_model=NumberOutput, + ) + + async def run(self, tool_input: NumberInput, context) -> NumberOutput: + del context + return NumberOutput(value=tool_input.value + 1) + + +class PublicRestrictedTool(AddOneTool): + spec = AddOneTool.spec.model_copy( + update={ + "identifier": "public_restricted", + "visibility": Visibility.PUBLIC, + } + ) + + +class BrokenOutputTool(AddOneTool): + spec = AddOneTool.spec.model_copy(update={"identifier": "broken"}) + + async def run(self, tool_input: NumberInput, context) -> NumberOutput: + del tool_input, context + return {"wrong": 1} # type: ignore[return-value] + + +class RuntimeBrokenOutputTool(BrokenOutputTool): + spec = BrokenOutputTool.spec.model_copy( + update={ + "identifier": "runtime_broken", + "allowed_callers": frozenset({CallerType.RUNTIME}), + } + ) + + +class AddOneCapability(Capability[NumberInput, NumberOutput]): + spec = CapabilitySpec( + identifier="add_one_capability", + version="1", + description="Return the input plus one.", + required_use="Use only in this composition test.", + visibility=Visibility.PUBLIC, + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=NumberInput, + output_model=NumberOutput, + tool_dependencies=("add_one",), + ) + + async def run(self, capability_input: NumberInput, context): + output = await context.invoke_tool("add_one", capability_input) + return Completed(value=output) + + +class ClarificationCapability(AddOneCapability): + spec = AddOneCapability.spec.model_copy( + update={"identifier": "clarification", "tool_dependencies": ()} + ) + + async def run(self, capability_input: NumberInput, context): + del capability_input, context + return NeedsInput( + prompt="What value should I use?", + missing_fields=("value",), + partial_input={}, + ) + + +async def not_cancelled() -> bool: + return False + + +def _context(composition, cancellation=not_cancelled): + return composition.executor.context( + request_id="request-1", + conversation_id="conversation-1", + turn_id="turn-1", + is_cancelled=cancellation, + ) + + +def test_tool_and_capability_specs_require_explicit_visibility_and_are_immutable(): + with pytest.raises(ValidationError): + ToolSpec( + identifier="missing_visibility", + version="1", + description="Invalid.", + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=NumberInput, + output_model=NumberOutput, + ) + + with pytest.raises(ValidationError): + CapabilitySpec( + identifier="missing_visibility", + version="1", + description="Invalid.", + required_use="Never.", + allowed_callers=frozenset({CallerType.RUNTIME}), + input_model=NumberInput, + output_model=NumberOutput, + ) + + with pytest.raises(ValidationError): + AddOneTool.spec.identifier = "changed" # type: ignore[misc] + + +def test_tool_registry_rejects_duplicates_and_filters_by_caller_and_visibility(): + registry = ToolRegistry() + registry.register(AddOneTool()) + registry.register(PublicRestrictedTool()) + + assert registry.definitions_for(CallerType.MODEL) == [] + assert registry.definitions_for(CallerType.CAPABILITY) == [ + { + "name": "public_restricted", + "description": "Add one to a number.", + "input_schema": NumberInput.model_json_schema(), + } + ] + assert [ + definition["name"] + for definition in registry.definitions_for( + CallerType.CAPABILITY, + include_private=True, + ) + ] == ["add_one", "public_restricted"] + with pytest.raises(ValueError, match="Duplicate typed tool"): + registry.register(AddOneTool()) + with pytest.raises(PermissionError): + registry.get("public_restricted", caller=CallerType.MODEL) + + +def _dependency_capability( + identifier: str, + *, + dependencies=(), + accepted_artifacts=(), + produced_artifacts=(), +): + class DependencyCapability(AddOneCapability): + spec = AddOneCapability.spec.model_copy( + update={ + "identifier": identifier, + "tool_dependencies": (), + "dependencies": dependencies, + "accepted_artifacts": accepted_artifacts, + "produced_artifacts": produced_artifacts, + } + ) + + async def run(self, capability_input, context): + del context + return Completed(value=NumberOutput(value=capability_input.value)) + + return DependencyCapability() + + +def test_capability_registry_rejects_unknown_dependencies_cycles_and_artifact_mismatch(): + unknown = CapabilityRegistry() + unknown.register( + _dependency_capability( + "consumer", + dependencies=(CapabilityDependency(capability_id="missing"),), + ) + ) + with pytest.raises(ValueError, match="unknown capability missing"): + unknown.validate() + + cyclic = CapabilityRegistry() + cyclic.register( + _dependency_capability( + "first", + dependencies=(CapabilityDependency(capability_id="second"),), + ) + ) + cyclic.register( + _dependency_capability( + "second", + dependencies=(CapabilityDependency(capability_id="first"),), + ) + ) + with pytest.raises(ValueError, match="first -> second -> first"): + cyclic.validate() + + scenario_v1 = ArtifactContract( + artifact_type="policy_scenario", + schema_version="1", + ) + scenario_v2 = scenario_v1.model_copy(update={"schema_version": "2"}) + incompatible = CapabilityRegistry() + incompatible.register( + _dependency_capability( + "provider", + produced_artifacts=(scenario_v1,), + ) + ) + incompatible.register( + _dependency_capability( + "consumer", + dependencies=( + CapabilityDependency( + capability_id="provider", + artifact=scenario_v2, + ), + ), + accepted_artifacts=(scenario_v2,), + ) + ) + with pytest.raises(ValueError, match="incompatible artifact"): + incompatible.validate() + + +def test_executor_validates_nested_calls_and_records_parent_aware_trace(): + composition = compose_runtime( + tools=[AddOneTool()], + capabilities=[AddOneCapability()], + ) + + outcome = asyncio.run( + composition.executor.invoke_capability( + "add_one_capability", + {"value": 4}, + caller=CallerType.RUNTIME, + context=_context(composition), + ) + ) + + assert isinstance(outcome, Completed) + assert outcome.value == NumberOutput(value=5) + records = composition.tracer.records( + "conversation-1", + include_private=True, + ) + assert [record.kind for record in records] == [ + InvocationKind.CAPABILITY, + InvocationKind.TOOL, + ] + assert records[1].parent_invocation_id == records[0].invocation_id + assert [record.status for record in records] == [ + InvocationStatus.COMPLETED, + InvocationStatus.COMPLETED, + ] + assert composition.tracer.records( + "conversation-1", + include_private=False, + ) == (records[0],) + + +def test_executor_rejects_invalid_input_output_and_undeclared_nested_calls(): + composition = compose_runtime( + tools=[BrokenOutputTool(), RuntimeBrokenOutputTool()], + capabilities=[ClarificationCapability()], + ) + context = _context(composition) + + with pytest.raises(TypeError, match="Invalid input"): + asyncio.run( + composition.executor.invoke_tool( + "runtime_broken", + {"value": 1, "extra": True}, + caller=CallerType.RUNTIME, + context=context, + ) + ) + + with pytest.raises(PermissionError, match="did not declare tool dependency"): + asyncio.run( + composition.executor.invoke_tool( + "broken", + {"value": 1}, + caller=CallerType.CAPABILITY, + context=context.for_capability("clarification"), + ) + ) + + with pytest.raises(TypeError, match="Invalid output"): + asyncio.run( + composition.executor.invoke_tool( + "runtime_broken", + {"value": 1}, + caller=CallerType.RUNTIME, + context=context, + ) + ) + + +def test_needs_input_is_typed_and_cancellation_is_checked_before_dispatch(): + composition = compose_runtime( + tools=[], + capabilities=[ClarificationCapability()], + ) + outcome = asyncio.run( + composition.executor.invoke_capability( + "clarification", + {"value": 1}, + caller=CallerType.RUNTIME, + context=_context(composition), + ) + ) + assert isinstance(outcome, NeedsInput) + assert outcome.missing_fields == ("value",) + + async def cancelled() -> bool: + return True + + with pytest.raises(InvocationCancelled): + asyncio.run( + composition.executor.invoke_capability( + "clarification", + {"value": 1}, + caller=CallerType.RUNTIME, + context=_context(composition, cancelled), + ) + ) + + +def test_startup_composition_rejects_unknown_tool_dependencies(): + with pytest.raises(ValueError, match="requires unknown tools"): + compose_runtime(tools=[], capabilities=[AddOneCapability()]) + + +def test_executor_exposes_no_selection_or_input_resolution_operations(): + composition = compose_runtime(tools=[], capabilities=[]) + executor = composition.executor + + assert not hasattr(executor, "select_capability") + assert not hasattr(executor, "infer_intent") + assert not hasattr(executor, "resolve_input") diff --git a/backend/tests/test_capability_persistence.py b/backend/tests/test_capability_persistence.py new file mode 100644 index 00000000..8da45997 --- /dev/null +++ b/backend/tests/test_capability_persistence.py @@ -0,0 +1,485 @@ +from __future__ import annotations + +import asyncio +import json +import os +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from pydantic import BaseModel, ConfigDict +from sqlalchemy import event +from sqlmodel import Session, SQLModel, create_engine, select + +from capabilities.artifacts import ( + ArtifactProvenance, + PolicyScenarioRef, +) +from capabilities.repository import WaitingCapabilityInvocation +from capabilities.tracing import ( + InvocationKind, + InvocationRecord, + InvocationStatus, +) +from conversations.models import ChatConversation +from conversation_context.models import ConversationContext +from persistence.capability_repository import ( + InvalidPersistedRecord, + PartialInputRegistry, + RepositoryArtifactAccess, + SQLConversationCapabilityRepository, +) +from persistence.deletion import delete_capability_records +from persistence.context_repository import SQLConversationContextRepository +from persistence.idempotency import ( + IdempotencyDecision, + ReceiptStatus, + SQLIdempotencyRepository, + request_fingerprint, +) +from persistence.rows import ( + CapabilityArtifactRow, + CapabilityCallReceiptRow, + ConversationContextRow, + InvocationTraceRow, + TurnReceiptRow, + WaitingCapabilityInvocationRow, +) +from persistence.trace_repository import SQLInvocationTraceRepository +from tools.contracts import Visibility + + +class PartialSocietyInput(BaseModel): + model_config = ConfigDict(extra="forbid") + + instruction: str | None = None + year: int | None = None + + +def _engine(tmp_path): + engine = create_engine( + f"sqlite:///{tmp_path / 'capability.sqlite'}", + connect_args={"check_same_thread": False}, + ) + + @event.listens_for(engine, "connect") + def enable_foreign_keys(connection, _record): + cursor = connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + SQLModel.metadata.create_all(engine) + with Session(engine) as session: + session.add( + ChatConversation( + session_id="conversation-1", + title="Test", + messages="[]", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + ) + session.commit() + return engine + + +def _repository(engine): + partials = PartialInputRegistry() + partials.register( + "society_analysis", + schema_version="1", + model=PartialSocietyInput, + ) + return SQLConversationCapabilityRepository( + engine=engine, + partial_inputs=partials, + ) + + +def _scenario() -> PolicyScenarioRef: + return PolicyScenarioRef( + artifact_id="scenario-1", + provenance=ArtifactProvenance( + conversation_id="conversation-1", + turn_id="turn-1", + capability_id="policy_reform", + capability_version="1", + invocation_id="invocation-1", + ), + year=2026, + scenario_revision="revision-1", + catalogue_version="catalogue-1", + calculation_engine_version="engine-1", + baseline=True, + ) + + +def test_sqlite_artifact_round_trip_and_invalid_payload_rejection(tmp_path): + engine = _engine(tmp_path) + repository = _repository(engine) + scenario = _scenario() + + repository.save_artifact("conversation-1", scenario) + + assert repository.get_artifact("conversation-1", "scenario-1") == scenario + assert repository.find_artifacts( + "conversation-1", + PolicyScenarioRef, + ) == (scenario,) + with Session(engine) as session: + row = session.get(CapabilityArtifactRow, "scenario-1") + row.payload_json = json.dumps( + { + "artifact_type": "policy_scenario", + "schema_version": "1", + "artifact_id": "scenario-1", + } + ) + session.add(row) + session.commit() + + with pytest.raises(InvalidPersistedRecord, match="Invalid persisted artifact"): + repository.get_artifact("conversation-1", "scenario-1") + + +def test_waiting_invocations_resume_branch_and_remain_independent(tmp_path): + repository = _repository(_engine(tmp_path)) + first = WaitingCapabilityInvocation( + invocation_id="waiting-1", + conversation_id="conversation-1", + capability_id="society_analysis", + capability_version="1", + input_schema_version="1", + partial_input=PartialSocietyInput(instruction="raise UC"), + source_turn_id="turn-1", + ) + repository.create_waiting(first) + + resumed = repository.resume_waiting("waiting-1", {"year": 2026}) + branch = repository.branch_waiting("waiting-1", "waiting-2", "turn-2") + repository.resume_waiting("waiting-2", {"year": 2025}) + + assert resumed.partial_input == PartialSocietyInput( + instruction="raise UC", + year=2026, + ) + assert branch.invocation_id == "waiting-2" + assert repository.get_waiting("waiting-1").partial_input.year == 2026 + assert repository.get_waiting("waiting-2").partial_input.year == 2025 + assert len(repository.list_waiting("conversation-1")) == 2 + + repository.remove_waiting("waiting-1") + assert [ + item.invocation_id for item in repository.list_waiting("conversation-1") + ] == ["waiting-2"] + + +def test_async_artifact_access_lists_updates_and_removes_waiting_input(tmp_path): + repository = _repository(_engine(tmp_path)) + access = RepositoryArtifactAccess(repository) + invocation = WaitingCapabilityInvocation( + invocation_id="waiting-async", + conversation_id="conversation-1", + capability_id="society_analysis", + capability_version="1", + input_schema_version="1", + partial_input=PartialSocietyInput(instruction="raise UC"), + source_turn_id="turn-1", + ) + + async def exercise_access(): + await access.save_waiting(invocation) + listed = await access.list_waiting( + conversation_id="conversation-1", + capability_id="society_analysis", + ) + updated = await access.update_waiting( + invocation_id="waiting-async", + partial_input=PartialSocietyInput(instruction="raise UC", year=2026), + ) + await access.remove_waiting(invocation_id="waiting-async") + remaining = await access.list_waiting( + conversation_id="conversation-1", + capability_id="society_analysis", + ) + return listed, updated, remaining + + listed, updated, remaining = asyncio.run(exercise_access()) + + assert [item.invocation_id for item in listed] == ["waiting-async"] + assert updated.partial_input.year == 2026 + assert remaining == () + + +def test_trace_repository_round_trips_typed_records_and_filters_private(tmp_path): + engine = _engine(tmp_path) + repository = SQLInvocationTraceRepository(engine=engine) + now = datetime.now(timezone.utc) + public = InvocationRecord( + conversation_id="conversation-1", + turn_id="turn-1", + invocation_id="trace-1", + sequence=1, + kind=InvocationKind.CAPABILITY, + identifier="society_analysis", + version="1", + visibility=Visibility.PUBLIC, + started_at=now, + completed_at=now, + duration_ms=0, + status=InvocationStatus.COMPLETED, + summary="society analysis completed", + ) + private = public.model_copy( + update={ + "invocation_id": "trace-2", + "sequence": 2, + "kind": InvocationKind.TOOL, + "identifier": "validate_reform", + "visibility": Visibility.PRIVATE, + } + ) + repository.save(public) + repository.save(private) + + assert repository.list_for_conversation( + "conversation-1", include_private=False + ) == (public,) + assert repository.list_for_conversation( + "conversation-1", include_private=True + ) == (public, private) + + +def test_turn_and_call_idempotency_conflicts_replay_and_billing_claim(tmp_path): + repository = SQLIdempotencyRepository(engine=_engine(tmp_path)) + fingerprint = request_fingerprint({"messages": ["hello"]}) + + started = repository.begin_turn( + conversation_id="conversation-1", + turn_id="turn-1", + fingerprint=fingerprint, + ) + in_progress = repository.begin_turn( + conversation_id="conversation-1", + turn_id="turn-1", + fingerprint=fingerprint, + ) + conflict = repository.begin_turn( + conversation_id="conversation-1", + turn_id="turn-1", + fingerprint=request_fingerprint({"messages": ["different"]}), + ) + repository.complete_turn( + turn_id="turn-1", + fingerprint=fingerprint, + outcome={"content": "answer"}, + ) + replay = repository.begin_turn( + conversation_id="conversation-1", + turn_id="turn-1", + fingerprint=fingerprint, + ) + + assert started.decision is IdempotencyDecision.STARTED + assert in_progress.decision is IdempotencyDecision.IN_PROGRESS + assert conflict.decision is IdempotencyDecision.CONFLICT + assert replay.decision is IdempotencyDecision.REPLAY + assert replay.outcome == {"content": "answer"} + assert repository.claim_billing("turn-1") is True + assert repository.claim_billing("turn-1") is False + + call = repository.begin_call( + conversation_id="conversation-1", + turn_id="turn-1", + call_id="call-1", + operation_id="run_society_simulation", + fingerprint=fingerprint, + ) + repository.complete_call( + call_id="call-1", + fingerprint=fingerprint, + outcome={"artifact_id": "result-1"}, + ) + call_replay = repository.begin_call( + conversation_id="conversation-1", + turn_id="turn-1", + call_id="call-1", + operation_id="run_society_simulation", + fingerprint=fingerprint, + ) + assert call.decision is IdempotencyDecision.STARTED + assert call_replay.status is ReceiptStatus.COMPLETED + assert call_replay.outcome == {"artifact_id": "result-1"} + + +@pytest.mark.parametrize( + "conflicting_update", + [ + {"conversation_id": "conversation-2"}, + {"turn_id": "turn-2"}, + {"operation_id": "different_operation"}, + {"fingerprint": "different-fingerprint"}, + ], +) +def test_each_capability_call_identity_dimension_rejects_reuse( + tmp_path, + conflicting_update, +): + repository = SQLIdempotencyRepository(engine=_engine(tmp_path)) + original = { + "conversation_id": "conversation-1", + "turn_id": "turn-1", + "call_id": "call-property", + "operation_id": "society_analysis", + "fingerprint": "fingerprint-1", + } + assert repository.begin_call(**original).decision is IdempotencyDecision.STARTED + + conflicting = {**original, **conflicting_update} + assert repository.begin_call(**conflicting).decision is IdempotencyDecision.CONFLICT + + +def test_competing_turn_retries_start_only_once(tmp_path): + repository = SQLIdempotencyRepository(engine=_engine(tmp_path)) + fingerprint = request_fingerprint({"turn": 1}) + + def begin(): + return repository.begin_turn( + conversation_id="conversation-1", + turn_id="turn-concurrent", + fingerprint=fingerprint, + ).decision + + with ThreadPoolExecutor(max_workers=4) as pool: + decisions = list(pool.map(lambda _: begin(), range(4))) + + assert decisions.count(IdempotencyDecision.STARTED) == 1 + assert decisions.count(IdempotencyDecision.IN_PROGRESS) == 3 + + +def test_deletion_removes_only_conversation_owned_capability_rows(tmp_path): + engine = _engine(tmp_path) + repository = _repository(engine) + repository.save_artifact("conversation-1", _scenario()) + repository.create_waiting( + WaitingCapabilityInvocation( + invocation_id="waiting-1", + conversation_id="conversation-1", + capability_id="society_analysis", + capability_version="1", + input_schema_version="1", + partial_input=PartialSocietyInput(), + source_turn_id="turn-1", + ) + ) + idempotency = SQLIdempotencyRepository(engine=engine) + idempotency.begin_turn( + conversation_id="conversation-1", + turn_id="turn-1", + fingerprint="fingerprint", + ) + SQLConversationContextRepository(engine=engine).save( + ConversationContext.initial("conversation-1"), + expected_revision=0, + ) + + with Session(engine) as session: + delete_capability_records(session, "conversation-1") + session.commit() + + with Session(engine) as session: + assert session.exec(select(CapabilityArtifactRow)).all() == [] + assert session.exec(select(WaitingCapabilityInvocationRow)).all() == [] + assert session.exec(select(InvocationTraceRow)).all() == [] + assert session.exec(select(TurnReceiptRow)).all() == [] + assert session.exec(select(CapabilityCallReceiptRow)).all() == [] + assert session.exec(select(ConversationContextRow)).all() == [] + assert session.exec(select(ChatConversation)).one().session_id == "conversation-1" + + +def test_additive_alembic_revision_uses_conversation_identity_without_foreign_key(): + migration = ( + Path(__file__).parents[2] + / "backend" + / "migrations" + / "versions" + / "0002_add_capability_persistence_tables.py" + ).read_text() + + assert migration.count("sa.Column('conversation_id'") == 5 + assert "ForeignKeyConstraint" not in migration + assert "chat_conversations" not in migration + assert "capability_artifacts" in migration + assert "waiting_capability_invocations" in migration + assert "capability_invocation_traces" in migration + assert "capability_turn_receipts" in migration + assert "capability_call_receipts" in migration + assert "analysis_workflows" not in migration + assert "analysis_plans" not in migration + assert "analysis_request_revisions" not in migration + + +def test_context_alembic_revision_was_cli_autogenerated_and_is_additive(): + migration = ( + Path(__file__).parents[2] + / "backend" + / "migrations" + / "versions" + / "9526d8c80914_add_conversation_context_persistence.py" + ).read_text() + + assert "commands auto generated by Alembic" in migration + assert "op.create_table('conversation_contexts'" in migration + assert "op.create_index('idx_conversation_context_updated'" in migration + assert "ForeignKeyConstraint" not in migration + assert "analysis_" not in migration + + +POSTGRES_URL = os.environ.get("CAPABILITY_TEST_POSTGRES_URL") + + +@pytest.mark.skipif( + not POSTGRES_URL, + reason="CAPABILITY_TEST_POSTGRES_URL is not configured for disposable PostgreSQL tests", +) +def test_disposable_postgres_uses_the_same_typed_repository_contract(tmp_path): + del tmp_path + engine = create_engine(POSTGRES_URL) + SQLModel.metadata.create_all(engine) + conversation_id = f"capability-postgres-{os.getpid()}" + try: + with Session(engine) as session: + session.add( + ChatConversation( + session_id=conversation_id, + title="Disposable test", + messages="[]", + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + ) + session.commit() + scenario = _scenario().model_copy( + update={ + "artifact_id": f"scenario-{os.getpid()}", + "provenance": _scenario().provenance.model_copy( + update={"conversation_id": conversation_id} + ), + } + ) + repository = SQLConversationCapabilityRepository(engine=engine) + repository.save_artifact(conversation_id, scenario) + assert repository.get_artifact(conversation_id, scenario.artifact_id) == scenario + finally: + with Session(engine) as session: + delete_capability_records(session, conversation_id) + row = session.exec( + select(ChatConversation).where( + ChatConversation.session_id == conversation_id + ) + ).first() + if row is not None: + session.delete(row) + session.commit() + engine.dispose() diff --git a/backend/tests/test_capability_public_api.py b/backend/tests/test_capability_public_api.py new file mode 100644 index 00000000..11baa056 --- /dev/null +++ b/backend/tests/test_capability_public_api.py @@ -0,0 +1,275 @@ +from __future__ import annotations + +import json +from datetime import datetime, timezone + +from fastapi.testclient import TestClient +from sqlmodel import Session, select + +from api.main import app +from capabilities.tracing import ( + InvocationKind, + InvocationRecord, + InvocationStatus, + InvocationTracer, +) +from chat.events import ( + ChatUsage, + InvocationActivity, + TextChunk, + TurnCompleted, +) +from conversations.models import get_engine +from persistence.deletion import delete_capability_records +from persistence.rows import InvocationTraceRow +from persistence.trace_repository import SQLInvocationTraceRepository +from tools.contracts import Visibility + + +client = TestClient(app) + + +def _events(response) -> list[dict[str, object]]: + return [ + json.loads(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") + ] + + +def _record( + identifier: str, + sequence: int, + *, + visibility: Visibility = Visibility.PUBLIC, + status: InvocationStatus = InvocationStatus.COMPLETED, + debug_details: bool = False, +) -> InvocationRecord: + now = datetime.now(timezone.utc) + return InvocationRecord( + conversation_id="api-session", + turn_id="api-turn", + invocation_id=f"invocation-{sequence}", + sequence=sequence, + kind=( + InvocationKind.CAPABILITY + if identifier not in {"assess_relevance", "validate_reform"} + else InvocationKind.TOOL + ), + identifier=identifier, + version="1", + visibility=visibility, + started_at=now, + completed_at=now, + duration_ms=1, + status=status, + summary=f"operation {identifier} {status.value}", + debug_input={"identifier": identifier} if debug_details else None, + debug_output={"status": status.value} if debug_details else None, + ) + + +def test_public_api_projects_direct_mandatory_mixed_and_localized_outcomes(monkeypatch): + from chat import public_service + + monkeypatch.setenv("BILLING_ENABLED", "false") + + async def fake_capability_turn(turn, *, is_cancelled): + del is_cancelled + prompt = str(turn.messages[-1]["content"]) + identifiers: list[str] = [] + if "mixed" in prompt: + identifiers = [ + "policy_information", + "household_analysis", + "society_analysis", + ] + elif "clarify" in prompt: + identifiers = ["household_analysis"] + if turn.debug: + yield InvocationActivity( + phase="finished", + record=_record( + "assess_relevance", + 1, + visibility=Visibility.PRIVATE, + debug_details=True, + ), + ) + fact_record = _record( + "reduce_context_patch", + 2, + visibility=Visibility.PRIVATE, + debug_details=True, + ).model_copy( + update={ + "kind": InvocationKind.TOOL, + "debug_input": { + "patch": { + "expected_revision": 0, + "operations": [{"definition_key": "person.age"}], + } + }, + "debug_output": { + "context": {"revision": 1}, + "decisions": [ + { + "status": "accepted", + "definition_key": "person.age", + "subject_entity_id": "person:self", + } + ], + }, + } + ) + yield InvocationActivity(phase="finished", record=fact_record) + for offset, identifier in enumerate( + identifiers, + start=3 if turn.debug else 2, + ): + status = ( + InvocationStatus.NEEDS_INPUT + if "clarify" in prompt + else InvocationStatus.COMPLETED + ) + yield InvocationActivity( + phase="finished", + record=_record( + identifier, + offset, + status=status, + debug_details=turn.debug, + ), + ) + answer = ( + "Which household ages and relationships should I use?" + if "clarify" in prompt + else "Natural response." + ) + yield TextChunk(answer) + yield TurnCompleted( + content=answer, + session_id="api-session", + model="fake-model", + route="capability", + outcome="completed", + stop_reason="end_turn", + usage=ChatUsage(), + turn_id=turn.turn_id, + ) + + monkeypatch.setattr( + public_service, + "run_capability_chat_turn", + fake_capability_turn, + ) + + direct = client.post( + "/chat/message", + json={"messages": [{"role": "user", "content": "direct"}]}, + ) + mixed = client.post( + "/chat/message", + json={"messages": [{"role": "user", "content": "mixed"}]}, + ) + clarification = client.post( + "/chat/message", + json={"messages": [{"role": "user", "content": "clarify"}]}, + ) + debug = client.post( + "/chat/message", + json={ + "messages": [{"role": "user", "content": "mixed"}], + "debug": True, + }, + ) + + assert [event["type"] for event in _events(direct)] == ["chunk", "done"] + mixed_activity = [ + event["invocation"]["identifier"] + for event in _events(mixed) + if event["type"] == "invocation_activity" + ] + assert mixed_activity == [ + "policy_information", + "household_analysis", + "society_analysis", + ] + assert all( + "debug_input" not in event["invocation"] + and "debug_output" not in event["invocation"] + for event in _events(mixed) + if event["type"] == "invocation_activity" + ) + clarification_events = _events(clarification) + assert next( + event["invocation"]["status"] + for event in clarification_events + if event["type"] == "invocation_activity" + ) == "needs_input" + assert "Which household" in clarification_events[-1]["content"] + debug_activity = [ + event["invocation"] + for event in _events(debug) + if event["type"] == "invocation_activity" + ] + assert debug_activity[0]["identifier"] == "assess_relevance" + assert debug_activity[0]["visibility"] == "private" + assert all("conversation_id" not in item for item in debug_activity) + assert all("debug_input" in item and "debug_output" in item for item in debug_activity) + fact_activity = next( + item for item in debug_activity if item["identifier"] == "reduce_context_patch" + ) + assert fact_activity["debug_output"]["decisions"][0] == { + "status": "accepted", + "definition_key": "person.age", + "subject_entity_id": "person:self", + } + + +def test_conversation_delete_removes_trace_endpoint_records( + isolated_conversations_table, +): + del isolated_conversations_table + saved = client.post( + "/conversations", + json={ + "session_id": "delete-capability-session", + "title": "Delete capability state", + "messages": [{"role": "user", "content": "hello"}], + "user_id": "owner-1", + }, + ) + assert saved.status_code == 200 + tracer = InvocationTracer(sink=SQLInvocationTraceRepository()) + record = tracer.start( + conversation_id="delete-capability-session", + turn_id="turn-1", + parent_invocation_id=None, + kind=InvocationKind.CAPABILITY, + identifier="policy_information", + version="1", + visibility=Visibility.PUBLIC, + summary="started", + ) + tracer.finish( + record.invocation_id, + status=InvocationStatus.COMPLETED, + summary="completed", + ) + activity = client.get( + "/chat/delete-capability-session/activity", + params={"user_id": "owner-1", "debug": "true"}, + ) + assert activity.status_code == 200 + assert len(activity.json()["invocations"]) == 1 + + deleted = client.delete(f"/conversations/{saved.json()['id']}") + + assert deleted.status_code == 204 + assert client.get( + "/chat/delete-capability-session/activity", + params={"user_id": "owner-1", "debug": "true"}, + ).status_code == 404 + with Session(get_engine()) as session: + assert session.exec(select(InvocationTraceRow)).all() == [] diff --git a/backend/tests/test_chat_events.py b/backend/tests/test_chat_events.py index af8e0f7b..c0c96b63 100644 --- a/backend/tests/test_chat_events.py +++ b/backend/tests/test_chat_events.py @@ -1,4 +1,8 @@ -from chat.events import ChatUsage, ToolCompleted, TurnCompleted +from datetime import datetime, timezone + +from capabilities.tracing import InvocationKind, InvocationRecord, InvocationStatus +from chat.events import ChatUsage, InvocationActivity, TurnCompleted +from tools.contracts import Visibility def test_chat_usage_exposes_the_existing_public_shape(): @@ -17,21 +21,32 @@ def test_chat_usage_exposes_the_existing_public_shape(): } -def test_tool_completed_retains_the_complete_structured_result(): +def test_invocation_activity_retains_the_sanitized_structured_projection(): output = { "status": "success", "rows": [{"income": 25_000, "nested": {"values": list(range(30))}}], } - event = ToolCompleted( - tool_name="run_society_simulation", - tool_id="tool-1", - status="success", - output=output, + record = InvocationRecord( + conversation_id="session-1", + turn_id="turn-1", + invocation_id="invocation-1", + sequence=1, + kind=InvocationKind.TOOL, + identifier="run_society_simulation", + version="1", + visibility=Visibility.PRIVATE, + started_at=datetime.now(timezone.utc), + completed_at=datetime.now(timezone.utc), + duration_ms=1, + status=InvocationStatus.COMPLETED, + summary="Society simulation completed.", + debug_output=output, ) + event = InvocationActivity("finished", record) - assert event.output is output - assert event.output["rows"][0]["nested"]["values"][-1] == 29 + assert event.record.debug_output == output + assert event.record.debug_output["rows"][0]["nested"]["values"][-1] == 29 def test_turn_completed_carries_execution_metadata_without_http_fields(): diff --git a/backend/tests/test_chat_orchestrator.py b/backend/tests/test_chat_orchestrator.py deleted file mode 100644 index 14d53a7b..00000000 --- a/backend/tests/test_chat_orchestrator.py +++ /dev/null @@ -1,462 +0,0 @@ -import asyncio -from types import SimpleNamespace - -from chat.events import ( - TextChunk, - ToolCompleted, - ToolUsed, - TurnCancelled, - TurnCompleted, - TurnFailed, -) -from chat.turn_input import ChatTurnInput -from gateway.assessment import ( - GatewayCatalogueUnavailable, - ReformAlternative, - ReformAssessment, - ValidatedParameterBinding, -) -from gateway.intent import ReformIntent -from gateway.policy import GatingReason, SlotFact -from gateway.runtime import GatewayVerdict - - -def _event(name: str, **attrs): - value = type(name, (), {})() - for key, item in attrs.items(): - setattr(value, key, item) - return value - - -class FakeStream: - def __init__(self, *, chunks=None, final_content=None, stop_reason="end_turn"): - self._events = [ - _event( - "RawContentBlockDeltaEvent", - delta=SimpleNamespace(type="text_delta", text=chunk), - ) - for chunk in (chunks or []) - ] - self._final = SimpleNamespace( - content=final_content or [], - stop_reason=stop_reason, - ) - - async def __aenter__(self): - return self - - async def __aexit__(self, *_args): - return None - - def __aiter__(self): - return self._iterate() - - async def _iterate(self): - for event in self._events: - yield event - - async def get_final_message(self): - return self._final - - -class FakeMessages: - def __init__(self, streams): - self.streams = list(streams) - self.calls = [] - - def stream(self, **kwargs): - self.calls.append(kwargs) - return self.streams.pop(0) - - -class FakeClient: - def __init__(self, streams): - self.messages = FakeMessages(streams) - - -def _tool_use(name, tool_input, tool_id="tool-1"): - return SimpleNamespace(type="tool_use", id=tool_id, name=name, input=tool_input) - - -async def _connected(): - return False - - -def test_run_chat_turn_emits_complete_tool_trace_and_completion(monkeypatch): - import chat.orchestrator as orchestrator - - tool_input = {"year": 2026, "reform": {"income_tax": {"basic_rate": 0.19}}} - tool_output = { - "status": "success", - "result_id": "result-1", - "rows": [{"budgetary_impact": -1_234_567_890}], - } - client = FakeClient( - [ - FakeStream( - chunks=["I will calculate."], - final_content=[_tool_use("run_society_simulation", tool_input)], - stop_reason="tool_use", - ), - FakeStream(chunks=["The result is £1.23bn."], final_content=[]), - ] - ) - - monkeypatch.setattr(orchestrator, "get_async_client", lambda: client) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: True) - monkeypatch.setattr(orchestrator, "execute_tool", lambda *_args, **_kwargs: tool_output) - - async def no_suggestions(*_args, **_kwargs): - return [] - - monkeypatch.setattr(orchestrator, "generate_followup_suggestions", no_suggestions) - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[{"role": "user", "content": "Calculate it."}], - session_id="eval-session", - ), - is_cancelled=_connected, - ) - ] - - events = asyncio.run(collect()) - - assert [ - type(event) - for event in events - if isinstance(event, (ToolUsed, ToolCompleted, TurnCompleted)) - ] == [ToolUsed, ToolCompleted, TurnCompleted] - completed_tool = next(event for event in events if isinstance(event, ToolCompleted)) - assert completed_tool.output is tool_output - assert completed_tool.output["rows"][0]["budgetary_impact"] == -1_234_567_890 - done = next(event for event in events if isinstance(event, TurnCompleted)) - assert done.content == "The result is £1.23bn." - assert done.session_id == "eval-session" - assert any( - isinstance(event, TextChunk) and event.content == "I will calculate." - for event in events - ) - assert client.messages.calls[1]["messages"][-1]["content"][0]["type"] == "tool_result" - - -def test_run_chat_turn_cancels_without_calling_the_model(monkeypatch): - import chat.orchestrator as orchestrator - - client = FakeClient([]) - monkeypatch.setattr(orchestrator, "get_async_client", lambda: client) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: True) - - async def cancelled(): - return True - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[{"role": "user", "content": "Calculate it."}], - session_id="cancel-session", - ), - is_cancelled=cancelled, - ) - ] - - events = asyncio.run(collect()) - - assert len(events) == 1 - assert isinstance(events[0], TurnCancelled) - assert events[0].session_id == "cancel-session" - assert client.messages.calls == [] - - -def test_needs_plan_is_rendered_without_async_writer_model(monkeypatch): - import chat.orchestrator as orchestrator - - verdict = GatewayVerdict( - outcome="needs_plan", - route="lightweight", - gating_reasons=[GatingReason("missing_output", "output")], - ) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: False) - monkeypatch.setattr(orchestrator, "run_gateway", lambda _prompt: verdict) - - def no_client(): - raise AssertionError("deterministic clarification must not create a client") - - monkeypatch.setattr(orchestrator, "get_async_client", no_client) - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[{"role": "user", "content": "Model a reform."}], - session_id="clarify-session", - ), - is_cancelled=_connected, - ) - ] - - events = asyncio.run(collect()) - - assert [type(event) for event in events] == [TextChunk, TurnCompleted] - assert events[0].content.startswith("What result") - done = events[1] - assert done.model is None - assert done.route == "lightweight" - assert done.outcome == "needs_plan" - assert done.stop_reason == "gateway_clarification" - assert done.gateway_trace.gating_reasons[0].code == "missing_output" - - -def test_cancellation_before_deterministic_clarification_emits_cancelled(monkeypatch): - import chat.orchestrator as orchestrator - - verdict = GatewayVerdict( - outcome="needs_plan", - route="lightweight", - gating_reasons=[GatingReason("missing_output", "output")], - ) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: False) - monkeypatch.setattr(orchestrator, "run_gateway", lambda _prompt: verdict) - probes = iter([False, True]) - - async def cancelled_after_gateway(): - return next(probes) - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[{"role": "user", "content": "Model a reform."}], - session_id="cancel-clarification", - ), - is_cancelled=cancelled_after_gateway, - ) - ] - - events = asyncio.run(collect()) - - assert len(events) == 1 - assert isinstance(events[0], TurnCancelled) - assert events[0].gateway_trace.gating_reasons[0].code == "missing_output" - - -def test_unrenderable_gateway_reason_fails_open_to_compute(monkeypatch): - import chat.orchestrator as orchestrator - - verdict = GatewayVerdict( - outcome="needs_plan", - route="lightweight", - gating_reasons=[GatingReason("internal_slot", "internal")], - ) - client = FakeClient([FakeStream(chunks=["Computed response."], final_content=[])]) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: False) - monkeypatch.setattr(orchestrator, "run_gateway", lambda _prompt: verdict) - monkeypatch.setattr(orchestrator, "get_async_client", lambda: client) - - async def no_suggestions(*_args, **_kwargs): - return [] - - monkeypatch.setattr(orchestrator, "generate_followup_suggestions", no_suggestions) - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[{"role": "user", "content": "Calculate it."}], - session_id="fail-open-session", - ), - is_cancelled=_connected, - ) - ] - - events = asyncio.run(collect()) - - done = next(event for event in events if isinstance(event, TurnCompleted)) - assert done.route == "compute" - assert done.outcome == "ready" - assert client.messages.calls[0]["tools"] - - -def test_low_confidence_clarification_contains_signed_resume_marker( - monkeypatch, -): - import chat.orchestrator as orchestrator - - assessment = ReformAssessment( - reform={"path.best": 0.21}, - summary="Best proposal", - confidence=72, - parameter_bindings=( - ValidatedParameterBinding("path.best", "Best label", "best"), - ), - alternatives=( - ReformAlternative( - "Other proposal", - ( - ValidatedParameterBinding( - "path.other", - "Other label", - "other", - ), - ), - {"path.other": 0.22}, - ), - ), - search_queries=("basic rate",), - catalogue_version="test-version", - ) - verdict = GatewayVerdict( - outcome="needs_plan", - route="lightweight", - tool="run_society_simulation", - slots=[ - SlotFact( - "output", - "prompt", - kind="output", - value="budgetary_impact", - ) - ], - gating_reasons=[GatingReason("confirm_reform", "reform")], - reform_intent=ReformIntent( - policy_phrase="basic rate", - action="increase", - amount="one percentage point", - scope="unspecified", - evidence="increasing the basic rate by one percentage point", - ), - reform_assessment=assessment, - ) - - monkeypatch.setenv( - "GATEWAY_PROPOSAL_SIGNING_KEY", - "orchestrator-test-signing-key-at-least-32-bytes", - ) - monkeypatch.setattr(orchestrator, "is_followup", lambda _messages: False) - monkeypatch.setattr( - orchestrator, - "run_gateway", - lambda _prompt: verdict, - ) - monkeypatch.setattr( - orchestrator, - "get_async_client", - lambda: (_ for _ in ()).throw(AssertionError("writer client created")), - ) - - async def collect(): - return [ - event - async for event in orchestrator.run_chat_turn( - ChatTurnInput( - messages=[ - { - "role": "user", - "content": ( - "What is the cost of increasing the basic rate " - "by one percentage point?" - ), - } - ], - session_id="proposal-session", - ), - is_cancelled=_connected, - ) - ] - - events = asyncio.run(collect()) - - assert [type(event) for event in events] == [TextChunk, TurnCompleted] - assert "Best label" in events[0].content - assert "path.best" not in events[0].content - assert "