Spaces:
Running
feat: initial commit — CustomerCore Phases 1-9 complete
Browse files- Phase 1: Project scaffold, Docker Compose (6 services), Doppler secrets
- Phase 2: Redpanda streaming — 4 event producers (tickets, billing, incidents, products)
- Phase 3: PySpark Bronze→Silver pipeline, Presidio PII masking, MinIO lakehouse
- Phase 4: dbt Gold analytics layer — 7 business marts, 18 data contract tests, DuckDB
- Phase 5: Cryptographic Privacy Vault — AES-256-GCM, RBAC, GDPR Article 17 right-to-erasure
- Phase 6: Multi-tenant Hybrid RAG — ChromaDB + BM25 + RRF fusion + 3-layer semantic cache
- Phase 7: SLA-aware multi-model LLM router — LiteLLM, OpenRouter, 54% local routing
- Phase 8: Multi-language support (7 langs, 87k rows) + Unified B2B Graph-RAG (NetworkX)
- Phase 9: LangGraph 6-agent supervisor — HITL, MemorySaver checkpointing, 223 tests
Stack: Python 3.12, Redpanda, MinIO, ChromaDB, Redis, DuckDB, dbt, LangGraph, LiteLLM
Dataset: 52,417 real-world support tickets (Bitext EN + Amazon MASSIVE DE/FR/ES/PT/IT/NL/JA)
- .env.example +35 -0
- .gitattributes +23 -0
- .gitignore +49 -0
- BUILD_GUIDE.md +0 -0
- CustomerCore_Progress_Summary.md +0 -0
- docker-compose.yml +158 -0
- infra/otel-collector.yaml +49 -0
- pyproject.toml +27 -0
- requirements.txt +0 -0
- src/__init__.py +0 -0
- src/agent/memory.py +107 -0
- src/agent/nodes/churn_agent.py +161 -0
- src/agent/nodes/classify_agent.py +143 -0
- src/agent/nodes/hitl_agent.py +95 -0
- src/agent/nodes/incident_agent.py +69 -0
- src/agent/nodes/memory_agent.py +21 -0
- src/agent/nodes/rag_agent.py +198 -0
- src/agent/schemas.py +131 -0
- src/agent/state.py +38 -0
- src/agent/supervisor.py +199 -0
- src/dbt/.user.yml +1 -0
- src/dbt/dbt_project.yml +22 -0
- src/dbt/models/gold/billing_failure_summary.sql +12 -0
- src/dbt/models/gold/customer_health_daily.sql +58 -0
- src/dbt/models/gold/incident_severity_hourly.sql +11 -0
- src/dbt/models/gold/product_adoption_features.sql +11 -0
- src/dbt/models/gold/retention_cohort_metrics.sql +33 -0
- src/dbt/models/gold/stg_silver_events.sql +36 -0
- src/dbt/models/gold/support_agent_performance.sql +13 -0
- src/dbt/models/gold/ticket_funnel_daily.sql +13 -0
- src/dbt/models/schema.yml +101 -0
- src/dbt/models/sources.yml +10 -0
- src/dbt/package-lock.yml +5 -0
- src/dbt/packages.yml +3 -0
- src/dbt/profiles.yml +16 -0
- src/rag/__init__.py +14 -0
- src/rag/graph_rag.py +617 -0
- src/rag/hybrid_retriever.py +470 -0
- src/rag/llm_client.py +263 -0
- src/rag/multilingual.py +208 -0
- src/rag/router.py +377 -0
- src/responsible_ai/__init__.py +10 -0
- src/responsible_ai/audit_log.py +124 -0
- src/responsible_ai/key_manager.py +101 -0
- src/responsible_ai/privacy_vault.py +442 -0
- src/streaming/__init__.py +0 -0
- src/streaming/bronze_consumer.py +177 -0
- src/streaming/bronze_to_silver.py +328 -0
- src/streaming/data_loader.py +378 -0
- src/streaming/minio_setup.py +66 -0
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# CustomerCore Environment Variables
|
| 2 |
+
# All real values are managed via Doppler — never put real secrets here
|
| 3 |
+
# Run: doppler setup && doppler run -- <your-command>
|
| 4 |
+
|
| 5 |
+
APP_ENV=development
|
| 6 |
+
|
| 7 |
+
# --- LLM APIs ---
|
| 8 |
+
OPENROUTER_API_KEY=your_openrouter_key_here
|
| 9 |
+
LITELLM_MASTER_KEY=sk-your-litellm-key-here
|
| 10 |
+
|
| 11 |
+
# --- Supabase (PostgreSQL + Auth + pgvector) ---
|
| 12 |
+
SUPABASE_URL=https://your-project-ref.supabase.co
|
| 13 |
+
SUPABASE_ANON_KEY=your_anon_key_here
|
| 14 |
+
SUPABASE_SERVICE_ROLE_KEY=your_service_role_key_here
|
| 15 |
+
SUPABASE_DB_URL=postgresql://postgres:your-password@db.your-project-ref.supabase.co:5432/postgres
|
| 16 |
+
SUPABASE_DB_PASSWORD=your_password_here
|
| 17 |
+
SUPABASE_PROJECT_REF=your_project_ref_here
|
| 18 |
+
|
| 19 |
+
# --- Redis (local dev via Docker, cloud via Upstash) ---
|
| 20 |
+
UPSTASH_REDIS_URL=rediss://default:your_token@your-host.upstash.io:6379
|
| 21 |
+
|
| 22 |
+
# --- MLflow / DagsHub ---
|
| 23 |
+
MLFLOW_TRACKING_URI=https://dagshub.com/username/CustomerCore.mlflow
|
| 24 |
+
DAGSHUB_TOKEN=your_dagshub_token_here
|
| 25 |
+
|
| 26 |
+
# --- Observability ---
|
| 27 |
+
LANGFUSE_PUBLIC_KEY=pk-lf-your-key-here
|
| 28 |
+
LANGFUSE_SECRET_KEY=sk-lf-your-key-here
|
| 29 |
+
SENTRY_DSN=https://your-sentry-dsn@sentry.io/project-id
|
| 30 |
+
|
| 31 |
+
# --- (Phase 18) Cloud Storage ---
|
| 32 |
+
R2_ACCOUNT_ID=your_cloudflare_account_id
|
| 33 |
+
R2_ACCESS_KEY_ID=your_r2_access_key
|
| 34 |
+
R2_SECRET_ACCESS_KEY=your_r2_secret_key
|
| 35 |
+
R2_BUCKET_NAME=customercore-lake
|
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Normalize line endings for all text files
|
| 2 |
+
* text=auto eol=lf
|
| 3 |
+
|
| 4 |
+
# Python files
|
| 5 |
+
*.py text eol=lf
|
| 6 |
+
*.pyi text eol=lf
|
| 7 |
+
|
| 8 |
+
# Config files
|
| 9 |
+
*.yml text eol=lf
|
| 10 |
+
*.yaml text eol=lf
|
| 11 |
+
*.toml text eol=lf
|
| 12 |
+
*.json text eol=lf
|
| 13 |
+
*.sql text eol=lf
|
| 14 |
+
*.md text eol=lf
|
| 15 |
+
|
| 16 |
+
# Binary files
|
| 17 |
+
*.png binary
|
| 18 |
+
*.jpg binary
|
| 19 |
+
*.gif binary
|
| 20 |
+
*.ico binary
|
| 21 |
+
*.docx binary
|
| 22 |
+
*.db binary
|
| 23 |
+
*.duckdb binary
|
|
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
.venv/
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.pyc
|
| 4 |
+
*.pyo
|
| 5 |
+
.pytest_cache/
|
| 6 |
+
.mypy_cache/
|
| 7 |
+
.ruff_cache/
|
| 8 |
+
|
| 9 |
+
# Secrets - NEVER commit
|
| 10 |
+
.env
|
| 11 |
+
.env.local
|
| 12 |
+
.env.*.local
|
| 13 |
+
.doppler.yaml
|
| 14 |
+
*.key
|
| 15 |
+
*.pem
|
| 16 |
+
|
| 17 |
+
# Data and databases
|
| 18 |
+
data/
|
| 19 |
+
*.duckdb
|
| 20 |
+
*.db
|
| 21 |
+
mlruns.db
|
| 22 |
+
mlruns/
|
| 23 |
+
|
| 24 |
+
# dbt artifacts
|
| 25 |
+
.dbt/
|
| 26 |
+
dbt_packages/
|
| 27 |
+
target/
|
| 28 |
+
|
| 29 |
+
# Logs and audit trails
|
| 30 |
+
logs/
|
| 31 |
+
*.log
|
| 32 |
+
audit_trail.jsonl
|
| 33 |
+
|
| 34 |
+
# Old version files
|
| 35 |
+
customercore_v*.docx
|
| 36 |
+
customercore_v*.js
|
| 37 |
+
|
| 38 |
+
# IDE
|
| 39 |
+
.vscode/
|
| 40 |
+
.idea/
|
| 41 |
+
|
| 42 |
+
# OS
|
| 43 |
+
.DS_Store
|
| 44 |
+
Thumbs.db
|
| 45 |
+
|
| 46 |
+
# Coverage
|
| 47 |
+
htmlcov/
|
| 48 |
+
.coverage
|
| 49 |
+
coverage.xml
|
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version: "3.8"
|
| 2 |
+
|
| 3 |
+
# ============================================================
|
| 4 |
+
# CustomerCore Local Infrastructure Stack
|
| 5 |
+
# Start: docker compose up -d
|
| 6 |
+
# Stop: docker compose down
|
| 7 |
+
# Logs: docker compose logs -f <service>
|
| 8 |
+
# ============================================================
|
| 9 |
+
|
| 10 |
+
networks:
|
| 11 |
+
customercore:
|
| 12 |
+
driver: bridge
|
| 13 |
+
|
| 14 |
+
volumes:
|
| 15 |
+
redpanda_data:
|
| 16 |
+
minio_data:
|
| 17 |
+
chroma_data:
|
| 18 |
+
redis_data:
|
| 19 |
+
|
| 20 |
+
services:
|
| 21 |
+
|
| 22 |
+
# ----------------------------------------------------------
|
| 23 |
+
# REDPANDA — Kafka-compatible streaming broker
|
| 24 |
+
# Replaces Kafka with no JVM, no ZooKeeper, single binary
|
| 25 |
+
# Broker listens on: localhost:9092 (internal: 29092)
|
| 26 |
+
# ----------------------------------------------------------
|
| 27 |
+
redpanda:
|
| 28 |
+
image: docker.redpanda.com/redpandadata/redpanda:latest
|
| 29 |
+
container_name: customercore_redpanda
|
| 30 |
+
command:
|
| 31 |
+
- redpanda
|
| 32 |
+
- start
|
| 33 |
+
- --kafka-addr internal://0.0.0.0:29092,external://0.0.0.0:9092
|
| 34 |
+
- --advertise-kafka-addr internal://redpanda:29092,external://localhost:9092
|
| 35 |
+
- --pandaproxy-addr internal://0.0.0.0:28082,external://0.0.0.0:8082
|
| 36 |
+
- --advertise-pandaproxy-addr internal://redpanda:28082,external://localhost:8082
|
| 37 |
+
- --schema-registry-addr internal://0.0.0.0:28081,external://0.0.0.0:8081
|
| 38 |
+
- --rpc-addr redpanda:33145
|
| 39 |
+
- --advertise-rpc-addr redpanda:33145
|
| 40 |
+
- --mode dev-container
|
| 41 |
+
- --smp 1
|
| 42 |
+
- --default-log-level=warn
|
| 43 |
+
volumes:
|
| 44 |
+
- redpanda_data:/var/lib/redpanda/data
|
| 45 |
+
networks:
|
| 46 |
+
- customercore
|
| 47 |
+
ports:
|
| 48 |
+
- "9092:9092" # Kafka API — Python producers connect here
|
| 49 |
+
- "8081:8081" # Schema Registry
|
| 50 |
+
- "8082:8082" # HTTP Proxy (Pandaproxy)
|
| 51 |
+
healthcheck:
|
| 52 |
+
test: ["CMD-SHELL", "rpk cluster health | grep -E 'Healthy|true' || exit 1"]
|
| 53 |
+
interval: 15s
|
| 54 |
+
timeout: 10s
|
| 55 |
+
retries: 5
|
| 56 |
+
|
| 57 |
+
# ----------------------------------------------------------
|
| 58 |
+
# REDPANDA CONSOLE — Visual web dashboard for Redpanda
|
| 59 |
+
# Browse topics, inspect messages, debug producers
|
| 60 |
+
# Dashboard: http://localhost:8080
|
| 61 |
+
# ----------------------------------------------------------
|
| 62 |
+
console:
|
| 63 |
+
image: docker.redpanda.com/redpandadata/console:latest
|
| 64 |
+
container_name: customercore_console
|
| 65 |
+
environment:
|
| 66 |
+
CONFIG_FILEPATH: /etc/redpanda/redpanda-console-config.yaml
|
| 67 |
+
entrypoint: /bin/sh
|
| 68 |
+
command: >
|
| 69 |
+
-c 'printf "kafka:\n brokers: [\"redpanda:29092\"]\n" > /etc/redpanda/redpanda-console-config.yaml && /app/console'
|
| 70 |
+
networks:
|
| 71 |
+
- customercore
|
| 72 |
+
ports:
|
| 73 |
+
- "8080:8080" # Web dashboard
|
| 74 |
+
depends_on:
|
| 75 |
+
- redpanda
|
| 76 |
+
|
| 77 |
+
# ----------------------------------------------------------
|
| 78 |
+
# MINIO — S3-compatible local object storage
|
| 79 |
+
# Stores Apache Iceberg Parquet data files
|
| 80 |
+
# Console: http://localhost:9001 (admin/password: minioadmin)
|
| 81 |
+
# API: http://localhost:9000
|
| 82 |
+
# ----------------------------------------------------------
|
| 83 |
+
minio:
|
| 84 |
+
image: minio/minio:latest
|
| 85 |
+
container_name: customercore_minio
|
| 86 |
+
command: server /data --console-address ":9001"
|
| 87 |
+
environment:
|
| 88 |
+
MINIO_ROOT_USER: minioadmin
|
| 89 |
+
MINIO_ROOT_PASSWORD: minioadmin
|
| 90 |
+
volumes:
|
| 91 |
+
- minio_data:/data
|
| 92 |
+
networks:
|
| 93 |
+
- customercore
|
| 94 |
+
ports:
|
| 95 |
+
- "9000:9000" # S3 API
|
| 96 |
+
- "9001:9001" # MinIO web console
|
| 97 |
+
healthcheck:
|
| 98 |
+
test: ["CMD", "mc", "ready", "local"]
|
| 99 |
+
interval: 30s
|
| 100 |
+
timeout: 20s
|
| 101 |
+
retries: 3
|
| 102 |
+
|
| 103 |
+
# ----------------------------------------------------------
|
| 104 |
+
# CHROMADB — Vector database for RAG retrieval
|
| 105 |
+
# Stores BGE-M3 embeddings of support tickets and KB articles
|
| 106 |
+
# API: http://localhost:8000
|
| 107 |
+
# ----------------------------------------------------------
|
| 108 |
+
chromadb:
|
| 109 |
+
image: chromadb/chroma:latest
|
| 110 |
+
container_name: customercore_chromadb
|
| 111 |
+
volumes:
|
| 112 |
+
- chroma_data:/chroma/chroma
|
| 113 |
+
networks:
|
| 114 |
+
- customercore
|
| 115 |
+
ports:
|
| 116 |
+
- "8000:8000" # ChromaDB HTTP API
|
| 117 |
+
environment:
|
| 118 |
+
- IS_PERSISTENT=TRUE
|
| 119 |
+
- ALLOW_RESET=TRUE
|
| 120 |
+
|
| 121 |
+
# ----------------------------------------------------------
|
| 122 |
+
# REDIS — In-memory cache and rate limiter
|
| 123 |
+
# L1 exact-match semantic cache + per-tenant rate limiting
|
| 124 |
+
# ----------------------------------------------------------
|
| 125 |
+
redis:
|
| 126 |
+
image: redis:alpine
|
| 127 |
+
container_name: customercore_redis
|
| 128 |
+
command: redis-server --appendonly yes --maxmemory 256mb --maxmemory-policy allkeys-lru
|
| 129 |
+
volumes:
|
| 130 |
+
- redis_data:/data
|
| 131 |
+
networks:
|
| 132 |
+
- customercore
|
| 133 |
+
ports:
|
| 134 |
+
- "6379:6379" # Redis default port
|
| 135 |
+
healthcheck:
|
| 136 |
+
test: ["CMD", "redis-cli", "ping"]
|
| 137 |
+
interval: 10s
|
| 138 |
+
timeout: 5s
|
| 139 |
+
retries: 5
|
| 140 |
+
|
| 141 |
+
# ----------------------------------------------------------
|
| 142 |
+
# OPENTELEMETRY COLLECTOR — Central telemetry pipeline
|
| 143 |
+
# Scrapes Prometheus metrics from FastAPI /metrics endpoint
|
| 144 |
+
# Forwards to Grafana Cloud (metrics, logs, traces)
|
| 145 |
+
# ----------------------------------------------------------
|
| 146 |
+
otel-collector:
|
| 147 |
+
image: otel/opentelemetry-collector-contrib:latest
|
| 148 |
+
container_name: customercore_otel
|
| 149 |
+
volumes:
|
| 150 |
+
- ./infra/otel-collector.yaml:/etc/otelcol-contrib/config.yaml
|
| 151 |
+
networks:
|
| 152 |
+
- customercore
|
| 153 |
+
ports:
|
| 154 |
+
- "4317:4317" # OTLP gRPC receiver
|
| 155 |
+
- "4318:4318" # OTLP HTTP receiver
|
| 156 |
+
- "8888:8888" # OTel internal metrics
|
| 157 |
+
depends_on:
|
| 158 |
+
- redpanda
|
|
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ============================================================
|
| 2 |
+
# OpenTelemetry Collector Configuration
|
| 3 |
+
# Receives: OTLP from FastAPI app
|
| 4 |
+
# Scrapes: Prometheus /metrics from FastAPI
|
| 5 |
+
# Exports: To Grafana Cloud (configured later in Phase 13)
|
| 6 |
+
# For now: logs everything to console for local dev
|
| 7 |
+
# ============================================================
|
| 8 |
+
|
| 9 |
+
receivers:
|
| 10 |
+
otlp:
|
| 11 |
+
protocols:
|
| 12 |
+
grpc:
|
| 13 |
+
endpoint: 0.0.0.0:4317
|
| 14 |
+
http:
|
| 15 |
+
endpoint: 0.0.0.0:4318
|
| 16 |
+
|
| 17 |
+
prometheus:
|
| 18 |
+
config:
|
| 19 |
+
scrape_configs:
|
| 20 |
+
- job_name: "customercore-api"
|
| 21 |
+
scrape_interval: 15s
|
| 22 |
+
static_configs:
|
| 23 |
+
- targets: ["host.docker.internal:8001"]
|
| 24 |
+
|
| 25 |
+
processors:
|
| 26 |
+
batch:
|
| 27 |
+
timeout: 10s
|
| 28 |
+
memory_limiter:
|
| 29 |
+
check_interval: 1s
|
| 30 |
+
limit_mib: 256
|
| 31 |
+
|
| 32 |
+
exporters:
|
| 33 |
+
debug:
|
| 34 |
+
verbosity: normal
|
| 35 |
+
|
| 36 |
+
service:
|
| 37 |
+
pipelines:
|
| 38 |
+
metrics:
|
| 39 |
+
receivers: [prometheus, otlp]
|
| 40 |
+
processors: [memory_limiter, batch]
|
| 41 |
+
exporters: [debug]
|
| 42 |
+
traces:
|
| 43 |
+
receivers: [otlp]
|
| 44 |
+
processors: [memory_limiter, batch]
|
| 45 |
+
exporters: [debug]
|
| 46 |
+
logs:
|
| 47 |
+
receivers: [otlp]
|
| 48 |
+
processors: [memory_limiter, batch]
|
| 49 |
+
exporters: [debug]
|
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[build-system]
|
| 2 |
+
requires = ["setuptools>=61.0"]
|
| 3 |
+
build-backend = "setuptools.build_meta"
|
| 4 |
+
|
| 5 |
+
[project]
|
| 6 |
+
name = "customercore"
|
| 7 |
+
version = "1.0.0"
|
| 8 |
+
description = "CustomerCore Intelligence Platform"
|
| 9 |
+
requires-python = ">=3.11"
|
| 10 |
+
|
| 11 |
+
[tool.pytest.ini_options]
|
| 12 |
+
minversion = "7.0"
|
| 13 |
+
addopts = "-ra -q --tb=short"
|
| 14 |
+
testpaths = [
|
| 15 |
+
"tests",
|
| 16 |
+
]
|
| 17 |
+
pythonpath = [
|
| 18 |
+
"."
|
| 19 |
+
]
|
| 20 |
+
|
| 21 |
+
[tool.ruff]
|
| 22 |
+
line-length = 120
|
| 23 |
+
target-version = "py311"
|
| 24 |
+
|
| 25 |
+
[tool.mypy]
|
| 26 |
+
python_version = "3.11"
|
| 27 |
+
ignore_missing_imports = true
|
|
Binary file (14.8 kB). View file
|
|
|
|
File without changes
|
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import structlog
|
| 3 |
+
from typing import Optional, List
|
| 4 |
+
|
| 5 |
+
log = structlog.get_logger()
|
| 6 |
+
|
| 7 |
+
SUPABASE_DB_URL = os.environ.get("SUPABASE_DB_URL", "")
|
| 8 |
+
SUPABASE_DB_HOST = os.environ.get("SUPABASE_DB_HOST", "")
|
| 9 |
+
|
| 10 |
+
_mem0_client: Optional[any] = None
|
| 11 |
+
# In-memory storage for local dev / fallback when Supabase is down
|
| 12 |
+
_local_memory_store: dict[str, list[str]] = {}
|
| 13 |
+
|
| 14 |
+
def get_mem0_client():
|
| 15 |
+
"""Get the Mem0 client, or raise Exception if not configured/available."""
|
| 16 |
+
global _mem0_client
|
| 17 |
+
if _mem0_client is None:
|
| 18 |
+
if not SUPABASE_DB_HOST:
|
| 19 |
+
raise ValueError("SUPABASE_DB_HOST is not configured. Falling back to local in-memory store.")
|
| 20 |
+
|
| 21 |
+
from mem0 import Memory
|
| 22 |
+
config = {
|
| 23 |
+
"vector_store": {
|
| 24 |
+
"provider": "pgvector",
|
| 25 |
+
"config": {
|
| 26 |
+
"dbname": "postgres",
|
| 27 |
+
"user": "postgres",
|
| 28 |
+
"password": os.environ.get("SUPABASE_DB_PASSWORD", ""),
|
| 29 |
+
"host": SUPABASE_DB_HOST,
|
| 30 |
+
"port": 5432,
|
| 31 |
+
"collection_name": "memories",
|
| 32 |
+
},
|
| 33 |
+
},
|
| 34 |
+
"embedder": {
|
| 35 |
+
"provider": "ollama",
|
| 36 |
+
"config": {
|
| 37 |
+
"model": "bge-m3",
|
| 38 |
+
"ollama_base_url": os.environ.get("OLLAMA_URL", "http://localhost:11434"),
|
| 39 |
+
},
|
| 40 |
+
},
|
| 41 |
+
"llm": {
|
| 42 |
+
"provider": "litellm",
|
| 43 |
+
"config": {
|
| 44 |
+
"model": "gemma-3-1b",
|
| 45 |
+
"api_base": os.environ.get("LITELLM_URL", "http://localhost:4000"),
|
| 46 |
+
"api_key": os.environ.get("LITELLM_MASTER_KEY", "sk-test"),
|
| 47 |
+
},
|
| 48 |
+
},
|
| 49 |
+
}
|
| 50 |
+
_mem0_client = Memory.from_config(config)
|
| 51 |
+
log.info("mem0_initialized")
|
| 52 |
+
return _mem0_client
|
| 53 |
+
|
| 54 |
+
def store_memory(customer_id: str, tenant_id: str, content: str, agent_id: str = "triage_agent"):
|
| 55 |
+
"""Store customer memory either in Mem0/Supabase or local in-memory store."""
|
| 56 |
+
user_id = f"{tenant_id}:{customer_id}"
|
| 57 |
+
try:
|
| 58 |
+
client = get_mem0_client()
|
| 59 |
+
client.add(
|
| 60 |
+
messages=[{"role": "user", "content": content}],
|
| 61 |
+
user_id=user_id,
|
| 62 |
+
agent_id=agent_id,
|
| 63 |
+
)
|
| 64 |
+
log.info("memory_stored", user_id=user_id, agent_id=agent_id)
|
| 65 |
+
except Exception as e:
|
| 66 |
+
# Fallback to local dict store
|
| 67 |
+
log.debug("mem0_store_failed_using_local_fallback", user_id=user_id, error=str(e))
|
| 68 |
+
if user_id not in _local_memory_store:
|
| 69 |
+
_local_memory_store[user_id] = []
|
| 70 |
+
_local_memory_store[user_id].append(content)
|
| 71 |
+
log.info("memory_stored_locally", user_id=user_id)
|
| 72 |
+
|
| 73 |
+
def recall_memories(customer_id: str, tenant_id: str, query: str, limit: int = 5) -> List[str]:
|
| 74 |
+
"""Recall customer memories from Mem0 or local in-memory store."""
|
| 75 |
+
user_id = f"{tenant_id}:{customer_id}"
|
| 76 |
+
try:
|
| 77 |
+
client = get_mem0_client()
|
| 78 |
+
results = client.search(query=query, user_id=user_id, limit=limit)
|
| 79 |
+
memories = [r.get("memory", "") for r in results]
|
| 80 |
+
log.info("memories_recalled", user_id=user_id, count=len(memories))
|
| 81 |
+
return memories
|
| 82 |
+
except Exception as e:
|
| 83 |
+
# Fallback to local dict search
|
| 84 |
+
log.debug("mem0_recall_failed_using_local_fallback", user_id=user_id, error=str(e))
|
| 85 |
+
local_mems = _local_memory_store.get(user_id, [])
|
| 86 |
+
# Very simple query match fallback: return matches or recent ones up to limit
|
| 87 |
+
matches = []
|
| 88 |
+
for mem in reversed(local_mems):
|
| 89 |
+
if any(word in mem.lower() for word in query.lower().split()):
|
| 90 |
+
matches.append(mem)
|
| 91 |
+
if len(matches) >= limit:
|
| 92 |
+
break
|
| 93 |
+
if not matches:
|
| 94 |
+
matches = local_mems[-limit:]
|
| 95 |
+
log.info("memories_recalled_locally", user_id=user_id, count=len(matches))
|
| 96 |
+
return matches
|
| 97 |
+
|
| 98 |
+
def get_all_memories(customer_id: str, tenant_id: str) -> List[str]:
|
| 99 |
+
"""Get all memories for a customer."""
|
| 100 |
+
user_id = f"{tenant_id}:{customer_id}"
|
| 101 |
+
try:
|
| 102 |
+
client = get_mem0_client()
|
| 103 |
+
results = client.get_all(user_id=user_id)
|
| 104 |
+
return [r.get("memory", "") for r in results]
|
| 105 |
+
except Exception as e:
|
| 106 |
+
log.debug("mem0_get_all_failed_using_local_fallback", error=str(e))
|
| 107 |
+
return _local_memory_store.get(user_id, [])
|
|
@@ -0,0 +1,161 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import mlflow
|
| 3 |
+
import structlog
|
| 4 |
+
from typing import List
|
| 5 |
+
from src.agent.state import AgentState
|
| 6 |
+
|
| 7 |
+
log = structlog.get_logger()
|
| 8 |
+
|
| 9 |
+
CHURN_MODEL = "customercore-churn-classifier"
|
| 10 |
+
SLA_MODEL = "customercore-sla-classifier"
|
| 11 |
+
|
| 12 |
+
def _load_model(name: str):
|
| 13 |
+
"""Attempt to load risk model from MLflow tracking server."""
|
| 14 |
+
tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "sqlite:///mlruns.db")
|
| 15 |
+
mlflow.set_tracking_uri(tracking_uri)
|
| 16 |
+
client = mlflow.MlflowClient()
|
| 17 |
+
try:
|
| 18 |
+
versions = client.get_latest_versions(name, stages=["Production", "None"])
|
| 19 |
+
if versions:
|
| 20 |
+
return mlflow.sklearn.load_model(f"models:/{name}/{versions[0].version}")
|
| 21 |
+
except Exception as e:
|
| 22 |
+
log.warning("model_load_failed", name=name, error=str(e))
|
| 23 |
+
return None
|
| 24 |
+
|
| 25 |
+
def calculate_heuristic_risks(body: str, category: str, priority: str, tier: str) -> tuple[float, float]:
|
| 26 |
+
"""
|
| 27 |
+
High-fidelity heuristic calculation of churn risk and SLA breach risk.
|
| 28 |
+
Balances customer tier, priority severity, category, and specific text sentiment cues.
|
| 29 |
+
"""
|
| 30 |
+
text = body.lower()
|
| 31 |
+
|
| 32 |
+
# ── SLA Breach Risk Calculation ────────────────────────────────────────────
|
| 33 |
+
# Baseline by priority
|
| 34 |
+
sla_base = {
|
| 35 |
+
"low": 0.10,
|
| 36 |
+
"medium": 0.25,
|
| 37 |
+
"high": 0.55,
|
| 38 |
+
"critical": 0.85
|
| 39 |
+
}
|
| 40 |
+
sla_risk = sla_base.get(priority, 0.25)
|
| 41 |
+
|
| 42 |
+
# Enterprise or high-tier SLAs have tighter response windows, increasing breach probability
|
| 43 |
+
if tier.lower() == "enterprise":
|
| 44 |
+
sla_risk += 0.10
|
| 45 |
+
elif tier.lower() == "professional":
|
| 46 |
+
sla_risk += 0.05
|
| 47 |
+
|
| 48 |
+
# High-pressure categories have tighter operation windows
|
| 49 |
+
if category in ["incident", "security"]:
|
| 50 |
+
sla_risk += 0.10
|
| 51 |
+
elif category == "performance":
|
| 52 |
+
sla_risk += 0.05
|
| 53 |
+
|
| 54 |
+
# SLA risk bounds
|
| 55 |
+
sla_risk = max(0.0, min(1.0, round(sla_risk, 3)))
|
| 56 |
+
|
| 57 |
+
# ── Churn Risk Calculation ─────────────────────────────────────────────────
|
| 58 |
+
churn_risk = 0.10 # Baseline churn risk
|
| 59 |
+
|
| 60 |
+
# Priority additions
|
| 61 |
+
if priority == "critical":
|
| 62 |
+
churn_risk += 0.15
|
| 63 |
+
elif priority == "high":
|
| 64 |
+
churn_risk += 0.08
|
| 65 |
+
|
| 66 |
+
# Category impact
|
| 67 |
+
if category == "billing":
|
| 68 |
+
churn_risk += 0.25 # Billing/double charging is the #1 churn driver
|
| 69 |
+
elif category == "incident":
|
| 70 |
+
churn_risk += 0.15 # Total service outages damage brand trust
|
| 71 |
+
elif category == "performance":
|
| 72 |
+
churn_risk += 0.08 # Persistent sluggishness leads to frustration
|
| 73 |
+
elif category == "security":
|
| 74 |
+
churn_risk += 0.20 # Security breaches are severe churn drivers
|
| 75 |
+
|
| 76 |
+
# Tier impact (free users churn easily; enterprise accounts are highly valued, high-risk churn)
|
| 77 |
+
if tier.lower() == "free":
|
| 78 |
+
churn_risk += 0.10 # Low switching barriers
|
| 79 |
+
elif tier.lower() == "enterprise":
|
| 80 |
+
# While contractual barriers exist, a severe issue raises the "Account-at-Risk" flag
|
| 81 |
+
churn_risk += 0.05
|
| 82 |
+
|
| 83 |
+
# Content-based threat indicators (direct expressions of churning intent)
|
| 84 |
+
churn_keywords = [
|
| 85 |
+
"cancel", "refund", "close my account", "switching to", "competitor",
|
| 86 |
+
"leave", "stop subscription", "unacceptable", "disappointed",
|
| 87 |
+
"useless", "moving away", "sue", "legal action"
|
| 88 |
+
]
|
| 89 |
+
if any(kw in text for kw in churn_keywords):
|
| 90 |
+
churn_risk += 0.25
|
| 91 |
+
|
| 92 |
+
# Churn risk bounds
|
| 93 |
+
churn_risk = max(0.0, min(1.0, round(churn_risk, 3)))
|
| 94 |
+
|
| 95 |
+
return sla_risk, churn_risk
|
| 96 |
+
|
| 97 |
+
def churn_agent_node(state: AgentState) -> AgentState:
|
| 98 |
+
ticket = state["ticket"]
|
| 99 |
+
body = ticket["body"]
|
| 100 |
+
category = state.get("category", "other")
|
| 101 |
+
priority = state.get("priority", "medium")
|
| 102 |
+
tier = ticket.get("customer_tier", "free")
|
| 103 |
+
|
| 104 |
+
models_used: List[str] = state.get("models_used") or []
|
| 105 |
+
sla_breach_risk = None
|
| 106 |
+
churn_risk = None
|
| 107 |
+
|
| 108 |
+
# 1. Try MLflow ML models first
|
| 109 |
+
try:
|
| 110 |
+
from src.ml.feature_engineering import create_structured_features
|
| 111 |
+
import pandas as pd
|
| 112 |
+
|
| 113 |
+
df = pd.DataFrame([{
|
| 114 |
+
"body": body,
|
| 115 |
+
"priority": priority,
|
| 116 |
+
"customer_tier": tier,
|
| 117 |
+
"reopen_count": 0,
|
| 118 |
+
"ticket_age_hours": 24,
|
| 119 |
+
}])
|
| 120 |
+
features = create_structured_features(df)
|
| 121 |
+
|
| 122 |
+
# Predict Churn Risk
|
| 123 |
+
churn_model = _load_model(CHURN_MODEL)
|
| 124 |
+
if churn_model:
|
| 125 |
+
# Assume probability prediction is available
|
| 126 |
+
if hasattr(churn_model, "predict_proba"):
|
| 127 |
+
churn_risk = float(churn_model.predict_proba(features)[0][1])
|
| 128 |
+
else:
|
| 129 |
+
churn_risk = float(churn_model.predict(features)[0])
|
| 130 |
+
models_used.append(CHURN_MODEL)
|
| 131 |
+
|
| 132 |
+
# Predict SLA Risk
|
| 133 |
+
sla_model = _load_model(SLA_MODEL)
|
| 134 |
+
if sla_model:
|
| 135 |
+
if hasattr(sla_model, "predict_proba"):
|
| 136 |
+
sla_breach_risk = float(sla_model.predict_proba(features)[0][1])
|
| 137 |
+
else:
|
| 138 |
+
sla_breach_risk = float(sla_model.predict(features)[0])
|
| 139 |
+
models_used.append(SLA_MODEL)
|
| 140 |
+
|
| 141 |
+
except (ImportError, ModuleNotFoundError) as e:
|
| 142 |
+
log.debug("ml_modules_missing_using_heuristics", error=str(e))
|
| 143 |
+
except Exception as e:
|
| 144 |
+
log.warning("ml_risk_prediction_failed_using_heuristics", error=str(e))
|
| 145 |
+
|
| 146 |
+
# 2. Fall back to high-quality heuristics if ML prediction failed or models not found
|
| 147 |
+
if sla_breach_risk is None or churn_risk is None:
|
| 148 |
+
h_sla, h_churn = calculate_heuristic_risks(body, category, priority, tier)
|
| 149 |
+
sla_breach_risk = sla_breach_risk if sla_breach_risk is not None else h_sla
|
| 150 |
+
churn_risk = churn_risk if churn_risk is not None else h_churn
|
| 151 |
+
models_used.append("heuristic-risk-engine-v1.0")
|
| 152 |
+
|
| 153 |
+
log.info("churn_agent_done", sla_breach_risk=sla_breach_risk, churn_risk=churn_risk, models_used=models_used)
|
| 154 |
+
|
| 155 |
+
return {
|
| 156 |
+
**state,
|
| 157 |
+
"sla_breach_risk": sla_breach_risk,
|
| 158 |
+
"churn_risk": churn_risk,
|
| 159 |
+
"current_step": "churn_agent",
|
| 160 |
+
"models_used": models_used,
|
| 161 |
+
}
|
|
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import mlflow
|
| 2 |
+
import os
|
| 3 |
+
import re
|
| 4 |
+
import structlog
|
| 5 |
+
from typing import List
|
| 6 |
+
from src.agent.state import AgentState
|
| 7 |
+
|
| 8 |
+
log = structlog.get_logger()
|
| 9 |
+
|
| 10 |
+
CATEGORY_MODEL = "customercore-category-classifier"
|
| 11 |
+
PRIORITY_MODEL = "customercore-priority-classifier"
|
| 12 |
+
|
| 13 |
+
CATEGORIES = [
|
| 14 |
+
"bug", "feature_request", "security", "performance",
|
| 15 |
+
"billing", "auth", "docs", "question", "incident", "other"
|
| 16 |
+
]
|
| 17 |
+
PRIORITIES = ["low", "medium", "high", "critical"]
|
| 18 |
+
|
| 19 |
+
def _load_model(name: str):
|
| 20 |
+
"""Attempt to load classifier model from MLflow tracking server."""
|
| 21 |
+
tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "sqlite:///mlruns.db")
|
| 22 |
+
mlflow.set_tracking_uri(tracking_uri)
|
| 23 |
+
client = mlflow.MlflowClient()
|
| 24 |
+
try:
|
| 25 |
+
versions = client.get_latest_versions(name, stages=["Production", "None"])
|
| 26 |
+
if versions:
|
| 27 |
+
return mlflow.sklearn.load_model(f"models:/{name}/{versions[0].version}")
|
| 28 |
+
except Exception as e:
|
| 29 |
+
log.warning("model_load_failed", name=name, error=str(e))
|
| 30 |
+
return None
|
| 31 |
+
|
| 32 |
+
def heuristic_classify(body: str) -> tuple[str, str]:
|
| 33 |
+
"""
|
| 34 |
+
Fallback high-quality heuristic classification for categories and priorities.
|
| 35 |
+
Matches industry support ticket routing logic.
|
| 36 |
+
"""
|
| 37 |
+
text = body.lower()
|
| 38 |
+
|
| 39 |
+
# ── Category Heuristics ────────────────────────────────────────────────────
|
| 40 |
+
category = "other"
|
| 41 |
+
|
| 42 |
+
if any(kw in text for kw in ["hacked", "leak", "unauthorized", "vulnerability", "breach", "cve", "compromised", "exploit", "security", "sicherheit", "fuite", "breche", "brecha", "exposed", "exposé", "expuesto", "divulgada", "divulgado"]):
|
| 43 |
+
category = "security"
|
| 44 |
+
elif any(kw in text for kw in ["outage", "offline", "completely down", "service unavailable", "is down", "not accessible", "500 error", "ausfall", "hors service", "panne", "caído", "caido"]):
|
| 45 |
+
category = "incident"
|
| 46 |
+
elif any(kw in text for kw in ["login", "password", "oauth", "token", "signin", "sign-in", "signup", "mfa", "2fa", "authentication", "einloggen", "passwort", "passworts", "mot de passe", "connexion", "contraseña", "iniciar sesión", "senha", "entrar", "konto gesperrt", "gesperrt"]):
|
| 47 |
+
category = "auth"
|
| 48 |
+
elif any(kw in text for kw in ["billing", "payment", "charge", "refund", "invoice", "stripe", "checkout", "cost", "pay", "rechnung", "zahlung", "facture", "paiement", "factura", "pago"]):
|
| 49 |
+
category = "billing"
|
| 50 |
+
elif any(kw in text for kw in ["slow", "latency", "lag", "timeout", "degraded", "delay", "performance", "high cpu"]):
|
| 51 |
+
category = "performance"
|
| 52 |
+
elif any(kw in text for kw in ["docs", "documentation", "how to", "guide", "tutorial", "where can i find"]):
|
| 53 |
+
category = "docs"
|
| 54 |
+
elif any(kw in text for kw in ["feature request", "would be nice", "can we add", "suggest", "improve", "request feature"]):
|
| 55 |
+
category = "feature_request"
|
| 56 |
+
elif any(kw in text for kw in ["bug", "error", "fail", "broken", "issue", "crash", "wrong", "incorrect", "exception"]):
|
| 57 |
+
category = "bug"
|
| 58 |
+
elif "?" in text:
|
| 59 |
+
category = "question"
|
| 60 |
+
|
| 61 |
+
# ── Priority Heuristics ────────────────────────────────────────────────────
|
| 62 |
+
priority = "medium"
|
| 63 |
+
|
| 64 |
+
# Critical criteria
|
| 65 |
+
if category in ["security", "incident"] or any(kw in text for kw in ["completely down", "production down", "outage", "all users affected"]):
|
| 66 |
+
priority = "critical"
|
| 67 |
+
# High criteria
|
| 68 |
+
elif any(kw in text for kw in ["urgent", "blocking", "failed", "broken", "asap", "invoice issue", "payment failed", "error", "dringend", "gesperrt", "fehlgeschlagen", "kaputt", "bloqué", "bloque", "urgente", "bloqueado", "bloqueada"]):
|
| 69 |
+
priority = "high"
|
| 70 |
+
# Low criteria
|
| 71 |
+
elif category in ["docs", "feature_request"] or any(kw in text for kw in ["minor", "typo", "cosmetic", "wont fix"]):
|
| 72 |
+
priority = "low"
|
| 73 |
+
|
| 74 |
+
return category, priority
|
| 75 |
+
|
| 76 |
+
def classify_agent_node(state: AgentState) -> AgentState:
|
| 77 |
+
ticket = state["ticket"]
|
| 78 |
+
body = ticket["body"]
|
| 79 |
+
|
| 80 |
+
# 1. Try MLflow ML models first
|
| 81 |
+
category = None
|
| 82 |
+
priority = None
|
| 83 |
+
models_used: List[str] = state.get("models_used") or []
|
| 84 |
+
|
| 85 |
+
try:
|
| 86 |
+
# Check if feature engineering module exists
|
| 87 |
+
from src.ml.feature_engineering import create_structured_features
|
| 88 |
+
import pandas as pd
|
| 89 |
+
|
| 90 |
+
df = pd.DataFrame([{
|
| 91 |
+
"body": body,
|
| 92 |
+
"priority": "medium",
|
| 93 |
+
"customer_tier": ticket.get("customer_tier", "professional"),
|
| 94 |
+
"reopen_count": 0,
|
| 95 |
+
"ticket_age_hours": 24,
|
| 96 |
+
}])
|
| 97 |
+
features = create_structured_features(df)
|
| 98 |
+
|
| 99 |
+
# Predict Category
|
| 100 |
+
cat_model = _load_model(CATEGORY_MODEL)
|
| 101 |
+
if cat_model:
|
| 102 |
+
cat_idx = cat_model.predict(features)[0]
|
| 103 |
+
category = CATEGORIES[int(cat_idx) % len(CATEGORIES)]
|
| 104 |
+
models_used.append(CATEGORY_MODEL)
|
| 105 |
+
|
| 106 |
+
# Predict Priority
|
| 107 |
+
pri_model = _load_model(PRIORITY_MODEL)
|
| 108 |
+
if pri_model:
|
| 109 |
+
pri_idx = pri_model.predict(features)[0]
|
| 110 |
+
priority = PRIORITIES[int(pri_idx) % len(PRIORITIES)]
|
| 111 |
+
models_used.append(PRIORITY_MODEL)
|
| 112 |
+
|
| 113 |
+
except (ImportError, ModuleNotFoundError) as e:
|
| 114 |
+
log.debug("ml_modules_missing_using_heuristics", error=str(e))
|
| 115 |
+
except Exception as e:
|
| 116 |
+
log.warning("ml_prediction_failed_using_heuristics", error=str(e))
|
| 117 |
+
|
| 118 |
+
# 2. Fall back to high-quality heuristics if ML prediction failed or models not found
|
| 119 |
+
if not category or not priority:
|
| 120 |
+
h_cat, h_pri = heuristic_classify(body)
|
| 121 |
+
category = category or h_cat
|
| 122 |
+
priority = priority or h_pri
|
| 123 |
+
models_used.append("heuristic-classifier-v1.0")
|
| 124 |
+
|
| 125 |
+
# 3. Apply customer tier multiplier: VIP customer gets upgraded priority
|
| 126 |
+
tier = ticket.get("customer_tier", "free").lower()
|
| 127 |
+
if tier == "enterprise":
|
| 128 |
+
if priority == "low":
|
| 129 |
+
priority = "medium"
|
| 130 |
+
elif priority == "medium":
|
| 131 |
+
priority = "high"
|
| 132 |
+
elif priority == "high":
|
| 133 |
+
priority = "critical"
|
| 134 |
+
|
| 135 |
+
log.info("classify_agent_done", category=category, priority=priority, models_used=models_used)
|
| 136 |
+
|
| 137 |
+
return {
|
| 138 |
+
**state,
|
| 139 |
+
"category": category,
|
| 140 |
+
"priority": priority,
|
| 141 |
+
"current_step": "classify_agent",
|
| 142 |
+
"models_used": models_used,
|
| 143 |
+
}
|
|
@@ -0,0 +1,95 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import structlog
|
| 2 |
+
from typing import List
|
| 3 |
+
from src.agent.state import AgentState
|
| 4 |
+
|
| 5 |
+
log = structlog.get_logger()
|
| 6 |
+
|
| 7 |
+
def compute_confidence(state: AgentState) -> float:
|
| 8 |
+
"""
|
| 9 |
+
Evaluate triage confidence based on category ambiguity, input lengths,
|
| 10 |
+
and models used during processing.
|
| 11 |
+
"""
|
| 12 |
+
confidence = 0.90 # Default high-fidelity baseline
|
| 13 |
+
|
| 14 |
+
ticket = state["ticket"]
|
| 15 |
+
body = ticket["body"]
|
| 16 |
+
category = state.get("category", "other")
|
| 17 |
+
models_used = state.get("models_used") or []
|
| 18 |
+
|
| 19 |
+
# 1. Ticket length penalty (extremely short or long tickets add ambiguity)
|
| 20 |
+
body_len = len(body.strip())
|
| 21 |
+
if body_len < 10:
|
| 22 |
+
confidence -= 0.35 # extremely short, practically zero context
|
| 23 |
+
elif body_len < 30:
|
| 24 |
+
confidence -= 0.15
|
| 25 |
+
elif body_len > 2500:
|
| 26 |
+
confidence -= 0.05
|
| 27 |
+
|
| 28 |
+
# 2. Category ambiguity penalty
|
| 29 |
+
if category == "other":
|
| 30 |
+
confidence -= 0.10
|
| 31 |
+
|
| 32 |
+
# 3. Model pedigree penalty: using fallback heuristics rather than ML/LLM models
|
| 33 |
+
if "heuristic-classifier-v1.0" in models_used:
|
| 34 |
+
confidence -= 0.05
|
| 35 |
+
if "heuristic-rag-fallback-v1.0" in models_used:
|
| 36 |
+
confidence -= 0.10
|
| 37 |
+
|
| 38 |
+
# 4. Risk premium check: elevated churn/SLA risks make the triage decision more sensitive
|
| 39 |
+
churn_risk = state.get("churn_risk") or 0.0
|
| 40 |
+
sla_risk = state.get("sla_breach_risk") or 0.0
|
| 41 |
+
if churn_risk > 0.70 or sla_risk > 0.75:
|
| 42 |
+
confidence -= 0.05
|
| 43 |
+
|
| 44 |
+
return max(0.1, min(1.0, round(confidence, 2)))
|
| 45 |
+
|
| 46 |
+
def hitl_agent_node(state: AgentState) -> AgentState:
|
| 47 |
+
priority = state.get("priority", "medium")
|
| 48 |
+
category = state.get("category", "other")
|
| 49 |
+
sla_breach_risk = state.get("sla_breach_risk") or 0.0
|
| 50 |
+
churn_risk = state.get("churn_risk") or 0.0
|
| 51 |
+
|
| 52 |
+
models_used: List[str] = state.get("models_used") or []
|
| 53 |
+
models_used.append("hitl-threshold-evaluator-v1.0")
|
| 54 |
+
|
| 55 |
+
# Calculate confidence score
|
| 56 |
+
confidence = compute_confidence(state)
|
| 57 |
+
|
| 58 |
+
# Evaluate HITL (Human-in-the-loop) criteria
|
| 59 |
+
hitl_required = False
|
| 60 |
+
reasons = []
|
| 61 |
+
|
| 62 |
+
if confidence < 0.65:
|
| 63 |
+
hitl_required = True
|
| 64 |
+
reasons.append(f"Low confidence ({confidence:.2f} < 0.65)")
|
| 65 |
+
|
| 66 |
+
if sla_breach_risk > 0.80:
|
| 67 |
+
hitl_required = True
|
| 68 |
+
reasons.append(f"Critical SLA breach risk ({sla_breach_risk:.2f} > 0.80)")
|
| 69 |
+
|
| 70 |
+
if churn_risk > 0.75 and priority in ["high", "critical"]:
|
| 71 |
+
hitl_required = True
|
| 72 |
+
reasons.append(f"High churn risk ({churn_risk:.2f} > 0.75) on high/critical priority ticket")
|
| 73 |
+
|
| 74 |
+
if category == "security":
|
| 75 |
+
# Enterprise policy: ALWAYS review security incidents to ensure data privacy and legal compliance
|
| 76 |
+
hitl_required = True
|
| 77 |
+
reasons.append("Security compliance policy review required")
|
| 78 |
+
|
| 79 |
+
hitl_reason = "; ".join(reasons) if hitl_required else None
|
| 80 |
+
|
| 81 |
+
log.info(
|
| 82 |
+
"hitl_agent_done",
|
| 83 |
+
confidence=confidence,
|
| 84 |
+
hitl_required=hitl_required,
|
| 85 |
+
hitl_reason=hitl_reason
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
return {
|
| 89 |
+
**state,
|
| 90 |
+
"confidence": confidence,
|
| 91 |
+
"hitl_required": hitl_required,
|
| 92 |
+
"hitl_reason": hitl_reason,
|
| 93 |
+
"current_step": "hitl_agent",
|
| 94 |
+
"models_used": models_used,
|
| 95 |
+
}
|
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import structlog
|
| 2 |
+
from typing import List
|
| 3 |
+
from src.agent.state import AgentState
|
| 4 |
+
|
| 5 |
+
log = structlog.get_logger()
|
| 6 |
+
|
| 7 |
+
# Keywords indicating an operational infrastructure or software incident
|
| 8 |
+
INCIDENT_KEYWORDS = [
|
| 9 |
+
"outage", "completely down", "service unavailable", "is down", "not accessible",
|
| 10 |
+
"500 error", "internal server error", "database connection failed", "server crash",
|
| 11 |
+
"network failure", "incident", "dns failure", "timeout error", "api failure"
|
| 12 |
+
]
|
| 13 |
+
|
| 14 |
+
def scan_incident(body: str) -> bool:
|
| 15 |
+
"""Scan ticket body for operational incident indicator keywords."""
|
| 16 |
+
text = body.lower()
|
| 17 |
+
return any(kw in text for kw in INCIDENT_KEYWORDS)
|
| 18 |
+
|
| 19 |
+
def incident_agent_node(state: AgentState) -> AgentState:
|
| 20 |
+
ticket = state["ticket"]
|
| 21 |
+
body = ticket["body"]
|
| 22 |
+
category = state.get("category", "other")
|
| 23 |
+
priority = state.get("priority", "medium")
|
| 24 |
+
tier = ticket.get("customer_tier", "free").lower()
|
| 25 |
+
|
| 26 |
+
models_used: List[str] = state.get("models_used") or []
|
| 27 |
+
models_used.append("incident-routing-heuristics-v1.0")
|
| 28 |
+
|
| 29 |
+
# 1. Detect if this is an active/severe incident
|
| 30 |
+
incident_detected = scan_incident(body) or (category == "incident")
|
| 31 |
+
|
| 32 |
+
# 2. Determine routing team based on category, priority, and customer tier
|
| 33 |
+
routing_team = "support" # default
|
| 34 |
+
|
| 35 |
+
if tier == "enterprise" and priority == "critical":
|
| 36 |
+
# Enterprise VIP tickets requiring critical escalation route to dedicated escalation team
|
| 37 |
+
routing_team = "escalation"
|
| 38 |
+
elif category == "security":
|
| 39 |
+
routing_team = "security"
|
| 40 |
+
elif category == "billing":
|
| 41 |
+
routing_team = "billing"
|
| 42 |
+
elif category == "incident":
|
| 43 |
+
routing_team = "infra"
|
| 44 |
+
elif category == "performance":
|
| 45 |
+
routing_team = "infra"
|
| 46 |
+
elif category == "auth":
|
| 47 |
+
routing_team = "security"
|
| 48 |
+
elif category in ["bug", "feature_request"]:
|
| 49 |
+
routing_team = "engineering"
|
| 50 |
+
elif category in ["docs", "question", "other"]:
|
| 51 |
+
routing_team = "support"
|
| 52 |
+
|
| 53 |
+
# If an incident was detected but the ticket is not already routed to escalation/security/billing,
|
| 54 |
+
# route to infra or engineering to address immediately
|
| 55 |
+
if incident_detected and routing_team not in ["escalation", "security", "billing"]:
|
| 56 |
+
if category in ["bug", "performance"]:
|
| 57 |
+
routing_team = "engineering"
|
| 58 |
+
else:
|
| 59 |
+
routing_team = "infra"
|
| 60 |
+
|
| 61 |
+
log.info("incident_agent_done", incident_detected=incident_detected, routing_team=routing_team)
|
| 62 |
+
|
| 63 |
+
return {
|
| 64 |
+
**state,
|
| 65 |
+
"incident_detected": incident_detected,
|
| 66 |
+
"routing_team": routing_team,
|
| 67 |
+
"current_step": "incident_agent",
|
| 68 |
+
"models_used": models_used,
|
| 69 |
+
}
|
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from src.agent.state import AgentState
|
| 2 |
+
from src.agent.memory import recall_memories
|
| 3 |
+
import structlog
|
| 4 |
+
|
| 5 |
+
log = structlog.get_logger()
|
| 6 |
+
|
| 7 |
+
def memory_agent_node(state: AgentState) -> AgentState:
|
| 8 |
+
ticket = state["ticket"]
|
| 9 |
+
customer_id = ticket["customer_id"]
|
| 10 |
+
tenant_id = ticket["tenant_id"]
|
| 11 |
+
body = ticket["body"]
|
| 12 |
+
|
| 13 |
+
# Recall up to 3 relevant long-term memories
|
| 14 |
+
memories = recall_memories(customer_id, tenant_id, query=body, limit=3)
|
| 15 |
+
|
| 16 |
+
log.info("memory_agent_done", customer_id=customer_id, memories_recalled=len(memories))
|
| 17 |
+
return {
|
| 18 |
+
**state,
|
| 19 |
+
"recalled_memories": memories,
|
| 20 |
+
"current_step": "memory_agent",
|
| 21 |
+
}
|
|
@@ -0,0 +1,198 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import json
|
| 2 |
+
import hashlib
|
| 3 |
+
import os
|
| 4 |
+
import structlog
|
| 5 |
+
from typing import List
|
| 6 |
+
from src.agent.state import AgentState
|
| 7 |
+
from src.rag.hybrid_retriever import get_retriever
|
| 8 |
+
from src.rag.graph_rag import GraphRAGEngine
|
| 9 |
+
from src.rag.llm_client import get_client
|
| 10 |
+
from src.rag.multilingual import detect_language, SUPPORTED_LANGUAGES
|
| 11 |
+
|
| 12 |
+
log = structlog.get_logger()
|
| 13 |
+
|
| 14 |
+
# Thread-safe in-memory cache for local fallback when Redis is absent
|
| 15 |
+
_local_semantic_cache = {}
|
| 16 |
+
|
| 17 |
+
def get_redis_client():
|
| 18 |
+
"""Attempt to get redis client if REDIS_URL is configured."""
|
| 19 |
+
redis_url = os.environ.get("REDIS_URL", "")
|
| 20 |
+
if not redis_url:
|
| 21 |
+
return None
|
| 22 |
+
try:
|
| 23 |
+
import redis
|
| 24 |
+
return redis.from_url(redis_url, decode_responses=True)
|
| 25 |
+
except Exception as e:
|
| 26 |
+
log.warning("redis_connection_failed_using_local_fallback", error=str(e))
|
| 27 |
+
return None
|
| 28 |
+
|
| 29 |
+
def lookup_cache(tenant_id: str, body: str) -> dict | None:
|
| 30 |
+
"""Look up ticket resolution in Redis semantic cache or local fallback."""
|
| 31 |
+
key_hash = hashlib.sha256(f"{tenant_id}:{body.strip()}".encode("utf-8")).hexdigest()
|
| 32 |
+
cache_key = f"cc:semcache:{key_hash}"
|
| 33 |
+
|
| 34 |
+
redis_client = get_redis_client()
|
| 35 |
+
if redis_client:
|
| 36 |
+
try:
|
| 37 |
+
cached_val = redis_client.get(cache_key)
|
| 38 |
+
if cached_val:
|
| 39 |
+
log.info("semantic_cache_hit_redis", tenant_id=tenant_id)
|
| 40 |
+
return json.loads(cached_val)
|
| 41 |
+
except Exception as e:
|
| 42 |
+
log.warning("redis_lookup_failed", error=str(e))
|
| 43 |
+
|
| 44 |
+
# Fallback to local in-memory dict
|
| 45 |
+
if cache_key in _local_semantic_cache:
|
| 46 |
+
log.info("semantic_cache_hit_local", tenant_id=tenant_id)
|
| 47 |
+
return _local_semantic_cache[cache_key]
|
| 48 |
+
|
| 49 |
+
return None
|
| 50 |
+
|
| 51 |
+
def store_cache(tenant_id: str, body: str, data: dict):
|
| 52 |
+
"""Store ticket resolution in Redis semantic cache or local fallback."""
|
| 53 |
+
key_hash = hashlib.sha256(f"{tenant_id}:{body.strip()}".encode("utf-8")).hexdigest()
|
| 54 |
+
cache_key = f"cc:semcache:{key_hash}"
|
| 55 |
+
|
| 56 |
+
redis_client = get_redis_client()
|
| 57 |
+
if redis_client:
|
| 58 |
+
try:
|
| 59 |
+
redis_client.setex(cache_key, 3600 * 24, json.dumps(data)) # 24 hour TTL
|
| 60 |
+
log.info("semantic_cache_stored_redis", tenant_id=tenant_id)
|
| 61 |
+
return
|
| 62 |
+
except Exception as e:
|
| 63 |
+
log.warning("redis_store_failed", error=str(e))
|
| 64 |
+
|
| 65 |
+
# Fallback to local in-memory dict
|
| 66 |
+
_local_semantic_cache[cache_key] = data
|
| 67 |
+
log.info("semantic_cache_stored_locally", tenant_id=tenant_id)
|
| 68 |
+
|
| 69 |
+
def rag_agent_node(state: AgentState) -> AgentState:
|
| 70 |
+
ticket = state["ticket"]
|
| 71 |
+
body = ticket["body"]
|
| 72 |
+
tenant_id = ticket["tenant_id"]
|
| 73 |
+
|
| 74 |
+
models_used: List[str] = state.get("models_used") or []
|
| 75 |
+
|
| 76 |
+
# 1. Check Redis / Local Semantic Cache first
|
| 77 |
+
cached_result = lookup_cache(tenant_id, body)
|
| 78 |
+
if cached_result:
|
| 79 |
+
models_used.append("redis-semantic-cache-v1")
|
| 80 |
+
return {
|
| 81 |
+
**state,
|
| 82 |
+
"summary": cached_result["summary"],
|
| 83 |
+
"suggested_resolution": cached_result["suggested_resolution"],
|
| 84 |
+
"kb_citations": cached_result["kb_citations"],
|
| 85 |
+
"current_step": "rag_agent",
|
| 86 |
+
"models_used": models_used,
|
| 87 |
+
}
|
| 88 |
+
|
| 89 |
+
# 2. Cache Miss: Run Hybrid Vector + Graph-RAG Retrieval
|
| 90 |
+
log.info("semantic_cache_miss_running_graph_rag", tenant_id=tenant_id)
|
| 91 |
+
retriever = get_retriever()
|
| 92 |
+
engine = GraphRAGEngine(retriever=retriever)
|
| 93 |
+
|
| 94 |
+
# Run Graph-RAG query
|
| 95 |
+
rag_result = engine.query(
|
| 96 |
+
tenant_id=tenant_id,
|
| 97 |
+
question=body,
|
| 98 |
+
k=5,
|
| 99 |
+
include_sql=True
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
# 3. Detect ticket language
|
| 103 |
+
lang = detect_language(body)
|
| 104 |
+
lang_display = SUPPORTED_LANGUAGES.get(lang, "English")
|
| 105 |
+
|
| 106 |
+
# 4. Invoke cloud LLM via LiteLLM client wrapper
|
| 107 |
+
client = get_client()
|
| 108 |
+
|
| 109 |
+
system_prompt = (
|
| 110 |
+
"You are an expert enterprise support triage agent. Your job is to analyze the support ticket "
|
| 111 |
+
"and generate a short summary, a detailed suggested resolution, and a list of knowledge base citation IDs.\n\n"
|
| 112 |
+
"You are provided with rich tenant context and similar past resolved tickets from our Graph-RAG engine.\n\n"
|
| 113 |
+
"=== CRITICAL INSTRUCTIONS ===\n"
|
| 114 |
+
f"1. The ticket is written in {lang_display}. You MUST output the 'summary' and 'suggested_resolution' in {lang_display}.\n"
|
| 115 |
+
"2. The summary must be a single concise sentence summarizing the core issue (max 300 characters).\n"
|
| 116 |
+
"3. The suggested resolution must be highly specific, technical, and actionable (max 1000 characters).\n"
|
| 117 |
+
"4. Output a list of knowledge base citation IDs. Valid IDs MUST start with 'KB-', 'TICKET-', or 'DOC-'. "
|
| 118 |
+
"Extract these from the provided similar past tickets or tenant profile where applicable. If none exist, output an empty list.\n"
|
| 119 |
+
"5. You MUST return your response as a strict JSON object with exactly three keys: 'summary', 'suggested_resolution', and 'kb_citations'. "
|
| 120 |
+
"Do not include any markdown framing, backticks, or text before/after the JSON."
|
| 121 |
+
)
|
| 122 |
+
|
| 123 |
+
messages = [
|
| 124 |
+
{"role": "system", "content": system_prompt},
|
| 125 |
+
{"role": "user", "content": f"Context:\n{rag_result.combined_context}\n\nTicket Body: {body}"}
|
| 126 |
+
]
|
| 127 |
+
|
| 128 |
+
try:
|
| 129 |
+
response = client.call_cloud(messages=messages, temperature=0.1)
|
| 130 |
+
if response.success:
|
| 131 |
+
models_used.append(response.model_used)
|
| 132 |
+
# Parse the JSON response
|
| 133 |
+
clean_content = response.content.strip()
|
| 134 |
+
if clean_content.startswith("```json"):
|
| 135 |
+
clean_content = clean_content.split("```json")[1].split("```")[0].strip()
|
| 136 |
+
elif clean_content.startswith("```"):
|
| 137 |
+
clean_content = clean_content.split("```")[1].split("```")[0].strip()
|
| 138 |
+
|
| 139 |
+
parsed = json.loads(clean_content)
|
| 140 |
+
summary = parsed.get("summary", "")
|
| 141 |
+
suggested_resolution = parsed.get("suggested_resolution", "")
|
| 142 |
+
kb_citations = parsed.get("kb_citations", [])
|
| 143 |
+
|
| 144 |
+
# Clean and validate citations
|
| 145 |
+
valid_citations = []
|
| 146 |
+
for citation in kb_citations:
|
| 147 |
+
citation_str = str(citation).upper().strip()
|
| 148 |
+
if citation_str.startswith(("KB-", "TICKET-", "DOC-")):
|
| 149 |
+
valid_citations.append(citation_str)
|
| 150 |
+
else:
|
| 151 |
+
# Fix formatting if possible
|
| 152 |
+
if citation_str.isalnum():
|
| 153 |
+
valid_citations.append(f"KB-{citation_str}")
|
| 154 |
+
|
| 155 |
+
# Ensure they are non-empty and satisfy length constraints
|
| 156 |
+
if not summary:
|
| 157 |
+
summary = body[:100] + "..." if len(body) > 100 else body
|
| 158 |
+
if not suggested_resolution:
|
| 159 |
+
suggested_resolution = "Our team is investigating this issue. We will update you shortly."
|
| 160 |
+
|
| 161 |
+
# Store in cache
|
| 162 |
+
result_data = {
|
| 163 |
+
"summary": summary,
|
| 164 |
+
"suggested_resolution": suggested_resolution,
|
| 165 |
+
"kb_citations": valid_citations,
|
| 166 |
+
}
|
| 167 |
+
store_cache(tenant_id, body, result_data)
|
| 168 |
+
|
| 169 |
+
log.info("rag_agent_completed_successfully", model=response.model_used)
|
| 170 |
+
return {
|
| 171 |
+
**state,
|
| 172 |
+
"summary": summary,
|
| 173 |
+
"suggested_resolution": suggested_resolution,
|
| 174 |
+
"kb_citations": valid_citations,
|
| 175 |
+
"current_step": "rag_agent",
|
| 176 |
+
"models_used": models_used,
|
| 177 |
+
}
|
| 178 |
+
|
| 179 |
+
except Exception as e:
|
| 180 |
+
log.error("rag_agent_llm_failed_using_fallback", error=str(e))
|
| 181 |
+
|
| 182 |
+
# Heuristic fallback if LLM call or parsing fails
|
| 183 |
+
summary = body[:150] + "..." if len(body) > 150 else body
|
| 184 |
+
suggested_resolution = (
|
| 185 |
+
f"A support agent has been notified of this ticket. "
|
| 186 |
+
f"Language detected: {lang_display}. Category classified: {state.get('category', 'unknown')}."
|
| 187 |
+
)
|
| 188 |
+
kb_citations = ["KB-FALLBACK"]
|
| 189 |
+
models_used.append("heuristic-rag-fallback-v1.0")
|
| 190 |
+
|
| 191 |
+
return {
|
| 192 |
+
**state,
|
| 193 |
+
"summary": summary,
|
| 194 |
+
"suggested_resolution": suggested_resolution,
|
| 195 |
+
"kb_citations": kb_citations,
|
| 196 |
+
"current_step": "rag_agent_fallback",
|
| 197 |
+
"models_used": models_used,
|
| 198 |
+
}
|
|
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pydantic import BaseModel, Field, field_validator
|
| 2 |
+
from typing import Optional, List, Literal
|
| 3 |
+
|
| 4 |
+
VALID_CATEGORIES = [
|
| 5 |
+
"bug", "feature_request", "security", "performance",
|
| 6 |
+
"billing", "auth", "docs", "question", "incident", "other",
|
| 7 |
+
]
|
| 8 |
+
|
| 9 |
+
VALID_PRIORITIES = ["critical", "high", "medium", "low"]
|
| 10 |
+
VALID_TEAMS = ["support", "engineering", "infra", "billing", "security", "escalation"]
|
| 11 |
+
|
| 12 |
+
class TriageOutput(BaseModel):
|
| 13 |
+
"""14-field structured output for every ticket processed by the triage agent."""
|
| 14 |
+
|
| 15 |
+
# Field 1: Category classification
|
| 16 |
+
category: Literal[
|
| 17 |
+
"bug", "feature_request", "security", "performance",
|
| 18 |
+
"billing", "auth", "docs", "question", "incident", "other"
|
| 19 |
+
] = Field(description="Ticket category determined by the classifier agent")
|
| 20 |
+
|
| 21 |
+
# Field 2: Priority classification
|
| 22 |
+
priority: Literal["critical", "high", "medium", "low"] = Field(
|
| 23 |
+
description="Ticket priority. critical=SLA breach risk. high=blocking. medium=normal. low=cosmetic"
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
# Field 3: Routing team
|
| 27 |
+
routing_team: Literal["support", "engineering", "infra", "billing", "security", "escalation"] = Field(
|
| 28 |
+
description="Team this ticket should be routed to"
|
| 29 |
+
)
|
| 30 |
+
|
| 31 |
+
# Field 4: SLA breach risk score
|
| 32 |
+
sla_breach_risk: float = Field(
|
| 33 |
+
ge=0.0, le=1.0,
|
| 34 |
+
description="Probability of SLA breach (0.0=safe, 1.0=certain breach)"
|
| 35 |
+
)
|
| 36 |
+
|
| 37 |
+
# Field 5: Churn risk score
|
| 38 |
+
churn_risk: float = Field(
|
| 39 |
+
ge=0.0, le=1.0,
|
| 40 |
+
description="Customer churn risk score based on ticket history and billing signals"
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
# Field 6: Confidence score
|
| 44 |
+
confidence: float = Field(
|
| 45 |
+
ge=0.0, le=1.0,
|
| 46 |
+
description="Agent confidence in this triage decision. < 0.65 triggers HITL"
|
| 47 |
+
)
|
| 48 |
+
|
| 49 |
+
# Field 7: Short summary of the ticket
|
| 50 |
+
summary: str = Field(
|
| 51 |
+
min_length=10, max_length=300,
|
| 52 |
+
description="One sentence summary of the ticket issue"
|
| 53 |
+
)
|
| 54 |
+
|
| 55 |
+
# Field 8: Suggested resolution
|
| 56 |
+
suggested_resolution: str = Field(
|
| 57 |
+
min_length=10, max_length=1000,
|
| 58 |
+
description="Recommended resolution or next action based on KB and ticket history"
|
| 59 |
+
)
|
| 60 |
+
|
| 61 |
+
# Field 9: Knowledge base citations
|
| 62 |
+
kb_citations: List[str] = Field(
|
| 63 |
+
default_factory=list,
|
| 64 |
+
description="List of KB article IDs that support this triage decision"
|
| 65 |
+
)
|
| 66 |
+
|
| 67 |
+
# Field 10: Customer memories recalled
|
| 68 |
+
recalled_memories: List[str] = Field(
|
| 69 |
+
default_factory=list,
|
| 70 |
+
description="Relevant memories recalled from previous interactions with this customer"
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
# Field 11: Incident detected
|
| 74 |
+
incident_detected: bool = Field(
|
| 75 |
+
description="True if this ticket is part of or related to a detected incident"
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
# Field 12: HITL required
|
| 79 |
+
hitl_required: bool = Field(
|
| 80 |
+
description="True if a human must review this triage before acting on it"
|
| 81 |
+
)
|
| 82 |
+
|
| 83 |
+
# Field 13: HITL reason
|
| 84 |
+
hitl_reason: Optional[str] = Field(
|
| 85 |
+
default=None,
|
| 86 |
+
description="Reason HITL is required, if hitl_required is True"
|
| 87 |
+
)
|
| 88 |
+
|
| 89 |
+
# Field 14: Model versions used
|
| 90 |
+
models_used: List[str] = Field(
|
| 91 |
+
default_factory=list,
|
| 92 |
+
description="List of model names used in this triage (for audit and debugging)"
|
| 93 |
+
)
|
| 94 |
+
|
| 95 |
+
@field_validator("hitl_reason")
|
| 96 |
+
@classmethod
|
| 97 |
+
def hitl_reason_required_when_hitl(cls, v, info):
|
| 98 |
+
if info.data.get("hitl_required") and not v:
|
| 99 |
+
raise ValueError("hitl_reason must be provided when hitl_required is True")
|
| 100 |
+
return v
|
| 101 |
+
|
| 102 |
+
@field_validator("kb_citations")
|
| 103 |
+
@classmethod
|
| 104 |
+
def validate_citations_format(cls, v):
|
| 105 |
+
for citation in v:
|
| 106 |
+
if not citation.startswith(("KB-", "TICKET-", "DOC-")):
|
| 107 |
+
raise ValueError(f"Invalid citation format: {citation}. Must start with KB-, TICKET-, or DOC-")
|
| 108 |
+
return v
|
| 109 |
+
|
| 110 |
+
def should_trigger_hitl(self) -> bool:
|
| 111 |
+
return self.confidence < 0.65 or self.hitl_required
|
| 112 |
+
|
| 113 |
+
class Config:
|
| 114 |
+
json_schema_extra = {
|
| 115 |
+
"example": {
|
| 116 |
+
"category": "auth",
|
| 117 |
+
"priority": "high",
|
| 118 |
+
"routing_team": "security",
|
| 119 |
+
"sla_breach_risk": 0.72,
|
| 120 |
+
"churn_risk": 0.35,
|
| 121 |
+
"confidence": 0.88,
|
| 122 |
+
"summary": "Customer reports 401 errors after OAuth token refresh",
|
| 123 |
+
"suggested_resolution": "Verify token expiry settings and redirect URI configuration",
|
| 124 |
+
"kb_citations": ["KB-001"],
|
| 125 |
+
"recalled_memories": ["Previous auth issue resolved with token rotation - 2025-03-01"],
|
| 126 |
+
"incident_detected": False,
|
| 127 |
+
"hitl_required": False,
|
| 128 |
+
"hitl_reason": None,
|
| 129 |
+
"models_used": ["priority-classifier-v2", "category-classifier-v2", "gemma-3-4b"],
|
| 130 |
+
}
|
| 131 |
+
}
|
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import TypedDict, Optional, List, Annotated
|
| 2 |
+
from langgraph.graph.message import add_messages
|
| 3 |
+
from src.agent.schemas import TriageOutput
|
| 4 |
+
|
| 5 |
+
class TicketInput(TypedDict):
|
| 6 |
+
ticket_id: str
|
| 7 |
+
body: str
|
| 8 |
+
customer_id: str
|
| 9 |
+
tenant_id: str
|
| 10 |
+
customer_tier: str
|
| 11 |
+
|
| 12 |
+
class AgentState(TypedDict):
|
| 13 |
+
# Input
|
| 14 |
+
ticket: TicketInput
|
| 15 |
+
|
| 16 |
+
# Populated by each sub-agent
|
| 17 |
+
category: Optional[str]
|
| 18 |
+
priority: Optional[str]
|
| 19 |
+
routing_team: Optional[str]
|
| 20 |
+
sla_breach_risk: Optional[float]
|
| 21 |
+
churn_risk: Optional[float]
|
| 22 |
+
confidence: Optional[float]
|
| 23 |
+
summary: Optional[str]
|
| 24 |
+
suggested_resolution: Optional[str]
|
| 25 |
+
kb_citations: Optional[List[str]]
|
| 26 |
+
recalled_memories: Optional[List[str]]
|
| 27 |
+
incident_detected: Optional[bool]
|
| 28 |
+
hitl_required: Optional[bool]
|
| 29 |
+
hitl_reason: Optional[str]
|
| 30 |
+
models_used: Optional[List[str]]
|
| 31 |
+
|
| 32 |
+
# Workflow metadata
|
| 33 |
+
current_step: Optional[str]
|
| 34 |
+
error: Optional[str]
|
| 35 |
+
final_output: Optional[TriageOutput]
|
| 36 |
+
|
| 37 |
+
# Messages (for LangGraph message passing)
|
| 38 |
+
messages: Annotated[list, add_messages]
|
|
@@ -0,0 +1,199 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import structlog
|
| 2 |
+
from typing import Dict, Any, Optional
|
| 3 |
+
from langgraph.graph import StateGraph, START, END
|
| 4 |
+
from langgraph.checkpoint.memory import MemorySaver
|
| 5 |
+
|
| 6 |
+
try:
|
| 7 |
+
from dotenv import load_dotenv
|
| 8 |
+
load_dotenv()
|
| 9 |
+
except ImportError:
|
| 10 |
+
pass
|
| 11 |
+
|
| 12 |
+
from src.agent.state import AgentState
|
| 13 |
+
from src.agent.schemas import TriageOutput
|
| 14 |
+
from src.agent.nodes.classify_agent import classify_agent_node
|
| 15 |
+
from src.agent.nodes.memory_agent import memory_agent_node
|
| 16 |
+
from src.agent.nodes.rag_agent import rag_agent_node
|
| 17 |
+
from src.agent.nodes.churn_agent import churn_agent_node
|
| 18 |
+
from src.agent.nodes.incident_agent import incident_agent_node
|
| 19 |
+
from src.agent.nodes.hitl_agent import hitl_agent_node
|
| 20 |
+
|
| 21 |
+
log = structlog.get_logger()
|
| 22 |
+
|
| 23 |
+
def hitl_interrupt_node(state: AgentState) -> AgentState:
|
| 24 |
+
"""Break point identity node. Execution halts immediately before this node is run."""
|
| 25 |
+
log.info("hitl_interrupt_node_reached_pausing")
|
| 26 |
+
return state
|
| 27 |
+
|
| 28 |
+
def finalize_node(state: AgentState) -> AgentState:
|
| 29 |
+
"""Assembles all calculated parameters into a strict, validated TriageOutput."""
|
| 30 |
+
log.info("finalizing_agent_state_compiling_schema")
|
| 31 |
+
|
| 32 |
+
# Default fallbacks to ensure compliance with field validators/minimum lengths
|
| 33 |
+
summary = state.get("summary") or "Ticket requiring triage."
|
| 34 |
+
if len(summary) < 10:
|
| 35 |
+
summary = (summary + " " * 10)[:15]
|
| 36 |
+
|
| 37 |
+
suggested_resolution = state.get("suggested_resolution") or "Resolution pending."
|
| 38 |
+
if len(suggested_resolution) < 10:
|
| 39 |
+
suggested_resolution = (suggested_resolution + " " * 10)[:15]
|
| 40 |
+
|
| 41 |
+
final_output = TriageOutput(
|
| 42 |
+
category=state.get("category") or "other",
|
| 43 |
+
priority=state.get("priority") or "medium",
|
| 44 |
+
routing_team=state.get("routing_team") or "support",
|
| 45 |
+
sla_breach_risk=state.get("sla_breach_risk") or 0.0,
|
| 46 |
+
churn_risk=state.get("churn_risk") or 0.0,
|
| 47 |
+
confidence=state.get("confidence") or 0.5,
|
| 48 |
+
summary=summary,
|
| 49 |
+
suggested_resolution=suggested_resolution,
|
| 50 |
+
kb_citations=state.get("kb_citations") or [],
|
| 51 |
+
recalled_memories=state.get("recalled_memories") or [],
|
| 52 |
+
incident_detected=state.get("incident_detected") or False,
|
| 53 |
+
hitl_required=state.get("hitl_required") or False,
|
| 54 |
+
hitl_reason=state.get("hitl_reason"),
|
| 55 |
+
models_used=state.get("models_used") or []
|
| 56 |
+
)
|
| 57 |
+
|
| 58 |
+
return {
|
| 59 |
+
**state,
|
| 60 |
+
"final_output": final_output,
|
| 61 |
+
"current_step": "finalize"
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
# ── Define the Agentic State Workflow ──────────────────────────────────────────
|
| 65 |
+
workflow = StateGraph(AgentState)
|
| 66 |
+
|
| 67 |
+
# 1. Register specialized sub-agent nodes
|
| 68 |
+
workflow.add_node("classify", classify_agent_node)
|
| 69 |
+
workflow.add_node("memory", memory_agent_node)
|
| 70 |
+
workflow.add_node("rag", rag_agent_node)
|
| 71 |
+
workflow.add_node("churn", churn_agent_node)
|
| 72 |
+
workflow.add_node("incident", incident_agent_node)
|
| 73 |
+
workflow.add_node("hitl", hitl_agent_node)
|
| 74 |
+
workflow.add_node("hitl_interrupt", hitl_interrupt_node)
|
| 75 |
+
workflow.add_node("finalize", finalize_node)
|
| 76 |
+
|
| 77 |
+
# 2. Add sequential transitions
|
| 78 |
+
workflow.add_edge(START, "classify")
|
| 79 |
+
workflow.add_edge("classify", "memory")
|
| 80 |
+
workflow.add_edge("memory", "rag")
|
| 81 |
+
workflow.add_edge("rag", "churn")
|
| 82 |
+
workflow.add_edge("churn", "incident")
|
| 83 |
+
workflow.add_edge("incident", "hitl")
|
| 84 |
+
|
| 85 |
+
# 3. Add conditional Human-in-the-loop gating
|
| 86 |
+
def hitl_check_router(state: AgentState) -> str:
|
| 87 |
+
"""Enforce human intercept routing if flagged by the HITL agent."""
|
| 88 |
+
if state.get("hitl_required"):
|
| 89 |
+
log.info("routing_to_human_interrupt")
|
| 90 |
+
return "hitl_interrupt"
|
| 91 |
+
log.info("routing_directly_to_finalize")
|
| 92 |
+
return "finalize"
|
| 93 |
+
|
| 94 |
+
workflow.add_conditional_edges(
|
| 95 |
+
"hitl",
|
| 96 |
+
hitl_check_router,
|
| 97 |
+
{
|
| 98 |
+
"hitl_interrupt": "hitl_interrupt",
|
| 99 |
+
"finalize": "finalize"
|
| 100 |
+
}
|
| 101 |
+
)
|
| 102 |
+
|
| 103 |
+
# 4. Final transitions
|
| 104 |
+
workflow.add_edge("hitl_interrupt", "finalize")
|
| 105 |
+
workflow.add_edge("finalize", END)
|
| 106 |
+
|
| 107 |
+
# ── Compile the Graph with Durable Memory Checkpointer ─────────────────────────
|
| 108 |
+
memory_checkpointer = MemorySaver()
|
| 109 |
+
app = workflow.compile(
|
| 110 |
+
checkpointer=memory_checkpointer,
|
| 111 |
+
interrupt_before=["hitl_interrupt"]
|
| 112 |
+
)
|
| 113 |
+
|
| 114 |
+
# ── Public Entry Points for the Triage Pipeline ────────────────────────────────
|
| 115 |
+
|
| 116 |
+
def run_triage(ticket: dict, thread_id: str = "default-thread") -> TriageOutput:
|
| 117 |
+
"""
|
| 118 |
+
Run the multi-agent triage pipeline.
|
| 119 |
+
|
| 120 |
+
If the graph hits a Human-in-the-loop interruption, the state is paused,
|
| 121 |
+
and a preliminary TriageOutput is returned with hitl_required=True.
|
| 122 |
+
Use resume_triage() to proceed.
|
| 123 |
+
"""
|
| 124 |
+
initial_state = {
|
| 125 |
+
"ticket": ticket,
|
| 126 |
+
"category": None,
|
| 127 |
+
"priority": None,
|
| 128 |
+
"routing_team": None,
|
| 129 |
+
"sla_breach_risk": None,
|
| 130 |
+
"churn_risk": None,
|
| 131 |
+
"confidence": None,
|
| 132 |
+
"summary": None,
|
| 133 |
+
"suggested_resolution": None,
|
| 134 |
+
"kb_citations": None,
|
| 135 |
+
"recalled_memories": None,
|
| 136 |
+
"incident_detected": None,
|
| 137 |
+
"hitl_required": None,
|
| 138 |
+
"hitl_reason": None,
|
| 139 |
+
"models_used": [],
|
| 140 |
+
"current_step": "start",
|
| 141 |
+
"error": None,
|
| 142 |
+
"final_output": None,
|
| 143 |
+
"messages": []
|
| 144 |
+
}
|
| 145 |
+
|
| 146 |
+
config = {"configurable": {"thread_id": thread_id}}
|
| 147 |
+
|
| 148 |
+
# Execute the graph
|
| 149 |
+
for event in app.stream(initial_state, config):
|
| 150 |
+
pass
|
| 151 |
+
|
| 152 |
+
state_snapshot = app.get_state(config)
|
| 153 |
+
if state_snapshot.next:
|
| 154 |
+
# Paused before hitl_interrupt. Assemble from current intermediate parameters.
|
| 155 |
+
current_values = state_snapshot.values
|
| 156 |
+
|
| 157 |
+
summary = current_values.get("summary") or "Ticket awaiting triage review."
|
| 158 |
+
if len(summary) < 10:
|
| 159 |
+
summary = (summary + " " * 10)[:15]
|
| 160 |
+
|
| 161 |
+
suggested_resolution = current_values.get("suggested_resolution") or "Pending human triage approval."
|
| 162 |
+
if len(suggested_resolution) < 10:
|
| 163 |
+
suggested_resolution = (suggested_resolution + " " * 10)[:15]
|
| 164 |
+
|
| 165 |
+
return TriageOutput(
|
| 166 |
+
category=current_values.get("category") or "other",
|
| 167 |
+
priority=current_values.get("priority") or "medium",
|
| 168 |
+
routing_team=current_values.get("routing_team") or "support",
|
| 169 |
+
sla_breach_risk=current_values.get("sla_breach_risk") or 0.0,
|
| 170 |
+
churn_risk=current_values.get("churn_risk") or 0.0,
|
| 171 |
+
confidence=current_values.get("confidence") or 0.5,
|
| 172 |
+
summary=summary,
|
| 173 |
+
suggested_resolution=suggested_resolution,
|
| 174 |
+
kb_citations=current_values.get("kb_citations") or [],
|
| 175 |
+
recalled_memories=current_values.get("recalled_memories") or [],
|
| 176 |
+
incident_detected=current_values.get("incident_detected") or False,
|
| 177 |
+
hitl_required=True,
|
| 178 |
+
hitl_reason=current_values.get("hitl_reason") or "Manual review required.",
|
| 179 |
+
models_used=current_values.get("models_used") or []
|
| 180 |
+
)
|
| 181 |
+
|
| 182 |
+
return state_snapshot.values.get("final_output")
|
| 183 |
+
|
| 184 |
+
def resume_triage(thread_id: str, overrides: Optional[Dict[str, Any]] = None) -> TriageOutput:
|
| 185 |
+
"""
|
| 186 |
+
Resume an interrupted triage pipeline, optionally applying human corrections.
|
| 187 |
+
"""
|
| 188 |
+
config = {"configurable": {"thread_id": thread_id}}
|
| 189 |
+
|
| 190 |
+
if overrides:
|
| 191 |
+
# Human agent overrides the AI decisions
|
| 192 |
+
app.update_state(config, overrides)
|
| 193 |
+
|
| 194 |
+
# Resume graph execution
|
| 195 |
+
for event in app.stream(None, config):
|
| 196 |
+
pass
|
| 197 |
+
|
| 198 |
+
state_snapshot = app.get_state(config)
|
| 199 |
+
return state_snapshot.values.get("final_output")
|
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
id: e384ffcc-f3ea-4709-bda5-25a26b0c069b
|
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: "customercore_dbt"
|
| 2 |
+
version: "1.0.0"
|
| 3 |
+
config-version: 2
|
| 4 |
+
|
| 5 |
+
profile: "customercore_dbt"
|
| 6 |
+
|
| 7 |
+
model-paths: ["models"]
|
| 8 |
+
analysis-paths: ["analyses"]
|
| 9 |
+
test-paths: ["tests"]
|
| 10 |
+
seed-paths: ["seeds"]
|
| 11 |
+
macro-paths: ["macros"]
|
| 12 |
+
|
| 13 |
+
target-path: "target"
|
| 14 |
+
clean-targets:
|
| 15 |
+
- "target"
|
| 16 |
+
- "dbt_packages"
|
| 17 |
+
|
| 18 |
+
models:
|
| 19 |
+
customercore_dbt:
|
| 20 |
+
gold:
|
| 21 |
+
+materialized: table
|
| 22 |
+
+schema: gold
|
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
tenant_id,
|
| 5 |
+
DATE(created_at) AS event_date,
|
| 6 |
+
priority,
|
| 7 |
+
COUNT(*) AS billing_event_count,
|
| 8 |
+
SUM(CASE WHEN body LIKE '%payment_failed%' THEN 1 ELSE 0 END) AS payment_failures,
|
| 9 |
+
SUM(CASE WHEN body LIKE '%cancelled%' THEN 1 ELSE 0 END) AS cancellations
|
| 10 |
+
FROM {{ ref('stg_silver_events') }}
|
| 11 |
+
WHERE event_type = 'billing_event'
|
| 12 |
+
GROUP BY 1, 2, 3
|
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(
|
| 2 |
+
materialized="table",
|
| 3 |
+
unique_key="customer_id || '|' || tenant_id || '|' || snapshot_date"
|
| 4 |
+
) }}
|
| 5 |
+
|
| 6 |
+
WITH tickets AS (
|
| 7 |
+
SELECT
|
| 8 |
+
customer_id,
|
| 9 |
+
tenant_id,
|
| 10 |
+
COUNT(*) AS open_ticket_count,
|
| 11 |
+
AVG(CASE priority
|
| 12 |
+
WHEN 'critical' THEN 4
|
| 13 |
+
WHEN 'high' THEN 3
|
| 14 |
+
WHEN 'medium' THEN 2
|
| 15 |
+
ELSE 1 END) AS avg_priority_score,
|
| 16 |
+
SUM(CASE WHEN priority IN ('critical','high') THEN 1 ELSE 0 END) AS high_priority_count
|
| 17 |
+
FROM {{ ref('stg_silver_events') }}
|
| 18 |
+
WHERE event_type = 'ticket'
|
| 19 |
+
AND created_at >= CURRENT_DATE - INTERVAL '30' DAY
|
| 20 |
+
GROUP BY 1, 2
|
| 21 |
+
),
|
| 22 |
+
|
| 23 |
+
billing AS (
|
| 24 |
+
SELECT
|
| 25 |
+
customer_id,
|
| 26 |
+
tenant_id,
|
| 27 |
+
COUNT(*) FILTER (WHERE body LIKE '%payment_failed%') AS payment_failures_30d,
|
| 28 |
+
MAX(created_at) AS last_billing_event
|
| 29 |
+
FROM {{ ref('stg_silver_events') }}
|
| 30 |
+
WHERE event_type = 'billing_event'
|
| 31 |
+
AND created_at >= CURRENT_DATE - INTERVAL '30' DAY
|
| 32 |
+
GROUP BY 1, 2
|
| 33 |
+
),
|
| 34 |
+
|
| 35 |
+
product AS (
|
| 36 |
+
SELECT
|
| 37 |
+
customer_id,
|
| 38 |
+
tenant_id,
|
| 39 |
+
COUNT(*) AS product_events_30d
|
| 40 |
+
FROM {{ ref('stg_silver_events') }}
|
| 41 |
+
WHERE event_type = 'product_event'
|
| 42 |
+
AND created_at >= CURRENT_DATE - INTERVAL '30' DAY
|
| 43 |
+
GROUP BY 1, 2
|
| 44 |
+
)
|
| 45 |
+
|
| 46 |
+
SELECT
|
| 47 |
+
COALESCE(t.customer_id, b.customer_id, p.customer_id) AS customer_id,
|
| 48 |
+
COALESCE(t.tenant_id, b.tenant_id, p.tenant_id) AS tenant_id,
|
| 49 |
+
CURRENT_DATE AS snapshot_date,
|
| 50 |
+
COALESCE(t.open_ticket_count, 0) AS open_tickets,
|
| 51 |
+
COALESCE(t.avg_priority_score, 1.0) AS avg_priority,
|
| 52 |
+
COALESCE(t.high_priority_count, 0) AS high_priority_tickets,
|
| 53 |
+
COALESCE(b.payment_failures_30d, 0) AS payment_failures_30d,
|
| 54 |
+
b.last_billing_event,
|
| 55 |
+
COALESCE(p.product_events_30d, 0) AS product_events_30d
|
| 56 |
+
FROM tickets t
|
| 57 |
+
FULL OUTER JOIN billing b USING (customer_id, tenant_id)
|
| 58 |
+
FULL OUTER JOIN product p USING (customer_id, tenant_id)
|
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
tenant_id,
|
| 5 |
+
DATE_TRUNC('hour', CAST(created_at AS TIMESTAMP)) AS incident_hour,
|
| 6 |
+
priority AS severity,
|
| 7 |
+
COUNT(*) AS incident_count,
|
| 8 |
+
COUNT(DISTINCT customer_id) AS affected_customers
|
| 9 |
+
FROM {{ ref('stg_silver_events') }}
|
| 10 |
+
WHERE event_type = 'incident'
|
| 11 |
+
GROUP BY 1, 2, 3
|
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
tenant_id,
|
| 5 |
+
customer_id,
|
| 6 |
+
DATE(created_at) AS event_date,
|
| 7 |
+
COUNT(*) AS total_product_events,
|
| 8 |
+
COUNT(DISTINCT DATE(created_at)) AS active_days
|
| 9 |
+
FROM {{ ref('stg_silver_events') }}
|
| 10 |
+
WHERE event_type = 'product_event'
|
| 11 |
+
GROUP BY 1, 2, 3
|
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
WITH first_seen AS (
|
| 4 |
+
SELECT
|
| 5 |
+
customer_id,
|
| 6 |
+
tenant_id,
|
| 7 |
+
MIN(DATE(created_at)) AS cohort_date
|
| 8 |
+
FROM {{ ref('stg_silver_events') }}
|
| 9 |
+
WHERE customer_id IS NOT NULL
|
| 10 |
+
GROUP BY 1, 2
|
| 11 |
+
),
|
| 12 |
+
|
| 13 |
+
activity AS (
|
| 14 |
+
SELECT
|
| 15 |
+
customer_id,
|
| 16 |
+
tenant_id,
|
| 17 |
+
DATE(created_at) AS active_date
|
| 18 |
+
FROM {{ ref('stg_silver_events') }}
|
| 19 |
+
WHERE customer_id IS NOT NULL
|
| 20 |
+
GROUP BY 1, 2, 3
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
SELECT
|
| 24 |
+
f.cohort_date,
|
| 25 |
+
f.tenant_id,
|
| 26 |
+
COUNT(DISTINCT f.customer_id) AS cohort_size,
|
| 27 |
+
COUNT(DISTINCT CASE WHEN a.active_date <= f.cohort_date + INTERVAL '7' DAY
|
| 28 |
+
THEN a.customer_id END) AS retained_d7,
|
| 29 |
+
COUNT(DISTINCT CASE WHEN a.active_date <= f.cohort_date + INTERVAL '30' DAY
|
| 30 |
+
THEN a.customer_id END) AS retained_d30
|
| 31 |
+
FROM first_seen f
|
| 32 |
+
LEFT JOIN activity a USING (customer_id, tenant_id)
|
| 33 |
+
GROUP BY 1, 2
|
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="view") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
event_id,
|
| 5 |
+
CASE
|
| 6 |
+
WHEN event_type = 'support_ticket_created' THEN 'ticket'
|
| 7 |
+
WHEN event_type = 'incident_detected' THEN 'incident'
|
| 8 |
+
WHEN event_type IN ('usability_complaint', 'feature_request', 'feature_praise') THEN 'product_event'
|
| 9 |
+
WHEN event_type IN ('payment_failed', 'payment_method_expired', 'subscription_downgrade', 'overcharge_reported', 'invoice_dispute') THEN 'billing_event'
|
| 10 |
+
ELSE event_type
|
| 11 |
+
END AS event_type,
|
| 12 |
+
source,
|
| 13 |
+
COALESCE(tenant_id, 'system') AS tenant_id,
|
| 14 |
+
customer_id,
|
| 15 |
+
COALESCE(original_timestamp, timestamp, processed_at)::timestamp AS created_at,
|
| 16 |
+
body,
|
| 17 |
+
priority,
|
| 18 |
+
COALESCE(reopen_count::integer, 0) AS reopen_count,
|
| 19 |
+
false AS is_synthetic,
|
| 20 |
+
|
| 21 |
+
-- Billing Specific Fields
|
| 22 |
+
amount::double AS billing_amount,
|
| 23 |
+
currency AS billing_currency,
|
| 24 |
+
plan AS billing_plan,
|
| 25 |
+
failure_code AS billing_failure_code,
|
| 26 |
+
|
| 27 |
+
-- Product Specific Fields
|
| 28 |
+
sentiment_score::double AS product_sentiment_score,
|
| 29 |
+
satisfaction_rating::integer AS product_satisfaction_rating,
|
| 30 |
+
|
| 31 |
+
-- Incident Specific Fields
|
| 32 |
+
severity AS incident_severity,
|
| 33 |
+
affected_service AS incident_affected_service,
|
| 34 |
+
error_rate::double AS incident_error_rate
|
| 35 |
+
|
| 36 |
+
FROM {{ source('silver', 'silver_events') }}
|
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
tenant_id,
|
| 5 |
+
source,
|
| 6 |
+
DATE(created_at) AS report_date,
|
| 7 |
+
priority,
|
| 8 |
+
COUNT(*) AS tickets_created,
|
| 9 |
+
COUNT(DISTINCT customer_id) AS unique_customers_served,
|
| 10 |
+
AVG(LENGTH(body)) AS avg_body_length
|
| 11 |
+
FROM {{ ref('stg_silver_events') }}
|
| 12 |
+
WHERE event_type = 'ticket'
|
| 13 |
+
GROUP BY 1, 2, 3, 4
|
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{{ config(materialized="table") }}
|
| 2 |
+
|
| 3 |
+
SELECT
|
| 4 |
+
tenant_id,
|
| 5 |
+
event_type,
|
| 6 |
+
priority,
|
| 7 |
+
source,
|
| 8 |
+
DATE(created_at) AS event_date,
|
| 9 |
+
COUNT(*) AS event_count,
|
| 10 |
+
COUNT(DISTINCT customer_id) AS unique_customers
|
| 11 |
+
FROM {{ ref('stg_silver_events') }}
|
| 12 |
+
WHERE event_type = 'ticket'
|
| 13 |
+
GROUP BY 1, 2, 3, 4, 5
|
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version: 2
|
| 2 |
+
|
| 3 |
+
models:
|
| 4 |
+
- name: customer_health_daily
|
| 5 |
+
description: "Per-customer health metrics computed daily from all signal sources"
|
| 6 |
+
columns:
|
| 7 |
+
- name: customer_id
|
| 8 |
+
description: "Unique identifier for the customer"
|
| 9 |
+
tests:
|
| 10 |
+
- not_null
|
| 11 |
+
- name: tenant_id
|
| 12 |
+
description: "B2B SaaS tenant identifier"
|
| 13 |
+
tests:
|
| 14 |
+
- not_null
|
| 15 |
+
- name: snapshot_date
|
| 16 |
+
description: "Date when metrics snapshot was taken"
|
| 17 |
+
tests:
|
| 18 |
+
- not_null
|
| 19 |
+
- name: avg_priority
|
| 20 |
+
description: "Weighted average priority score of tickets (1-4)"
|
| 21 |
+
tests:
|
| 22 |
+
- not_null
|
| 23 |
+
- dbt_utils.accepted_range:
|
| 24 |
+
min_value: 1
|
| 25 |
+
max_value: 4
|
| 26 |
+
|
| 27 |
+
- name: ticket_funnel_daily
|
| 28 |
+
description: "Daily ticket volume by category, priority, and source"
|
| 29 |
+
columns:
|
| 30 |
+
- name: tenant_id
|
| 31 |
+
description: "B2B SaaS tenant identifier"
|
| 32 |
+
tests:
|
| 33 |
+
- not_null
|
| 34 |
+
- name: event_count
|
| 35 |
+
description: "Total ticket count for the day"
|
| 36 |
+
tests:
|
| 37 |
+
- not_null
|
| 38 |
+
|
| 39 |
+
- name: incident_severity_hourly
|
| 40 |
+
description: "Hourly incident counts by severity"
|
| 41 |
+
columns:
|
| 42 |
+
- name: tenant_id
|
| 43 |
+
description: "B2B SaaS tenant identifier"
|
| 44 |
+
tests:
|
| 45 |
+
- not_null
|
| 46 |
+
- name: incident_count
|
| 47 |
+
description: "Number of system incidents detected"
|
| 48 |
+
tests:
|
| 49 |
+
- not_null
|
| 50 |
+
|
| 51 |
+
- name: billing_failure_summary
|
| 52 |
+
description: "Daily billing failures and subscription cancellations summary"
|
| 53 |
+
columns:
|
| 54 |
+
- name: tenant_id
|
| 55 |
+
description: "B2B SaaS tenant identifier"
|
| 56 |
+
tests:
|
| 57 |
+
- not_null
|
| 58 |
+
- name: billing_event_count
|
| 59 |
+
description: "Total billing events registered"
|
| 60 |
+
tests:
|
| 61 |
+
- not_null
|
| 62 |
+
|
| 63 |
+
- name: product_adoption_features
|
| 64 |
+
description: "Daily count of product interactions per customer account"
|
| 65 |
+
columns:
|
| 66 |
+
- name: tenant_id
|
| 67 |
+
description: "B2B SaaS tenant identifier"
|
| 68 |
+
tests:
|
| 69 |
+
- not_null
|
| 70 |
+
- name: customer_id
|
| 71 |
+
description: "Unique identifier for the customer"
|
| 72 |
+
tests:
|
| 73 |
+
- not_null
|
| 74 |
+
- name: total_product_events
|
| 75 |
+
description: "Total product usage events"
|
| 76 |
+
tests:
|
| 77 |
+
- not_null
|
| 78 |
+
|
| 79 |
+
- name: retention_cohort_metrics
|
| 80 |
+
description: "Weekly and monthly retention metrics by customer signup cohorts"
|
| 81 |
+
columns:
|
| 82 |
+
- name: tenant_id
|
| 83 |
+
description: "B2B SaaS tenant identifier"
|
| 84 |
+
tests:
|
| 85 |
+
- not_null
|
| 86 |
+
- name: cohort_size
|
| 87 |
+
description: "Number of new customers signed up on the cohort date"
|
| 88 |
+
tests:
|
| 89 |
+
- not_null
|
| 90 |
+
|
| 91 |
+
- name: support_agent_performance
|
| 92 |
+
description: "Daily summary of ticket volume and descriptive stats for agents"
|
| 93 |
+
columns:
|
| 94 |
+
- name: tenant_id
|
| 95 |
+
description: "B2B SaaS tenant identifier"
|
| 96 |
+
tests:
|
| 97 |
+
- not_null
|
| 98 |
+
- name: tickets_created
|
| 99 |
+
description: "Number of tickets processed for this slice"
|
| 100 |
+
tests:
|
| 101 |
+
- not_null
|
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version: 2
|
| 2 |
+
|
| 3 |
+
sources:
|
| 4 |
+
- name: silver
|
| 5 |
+
description: "Silver layer S3 Parquet tables from Python streaming pipeline"
|
| 6 |
+
tables:
|
| 7 |
+
- name: silver_events
|
| 8 |
+
description: "All PII-masked, schema-validated events from all 4 sources"
|
| 9 |
+
meta:
|
| 10 |
+
external_location: "read_parquet('s3://customercore-lake/silver/*/*.parquet', union_by_name=True)"
|
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
packages:
|
| 2 |
+
- name: dbt_utils
|
| 3 |
+
package: dbt-labs/dbt_utils
|
| 4 |
+
version: 1.3.0
|
| 5 |
+
sha1_hash: 226ae69cdfbc9367e2aa2c472b01f99dbce11de0
|
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
packages:
|
| 2 |
+
- package: dbt-labs/dbt_utils
|
| 3 |
+
version: 1.3.0
|
|
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
customercore_dbt:
|
| 2 |
+
target: dev
|
| 3 |
+
outputs:
|
| 4 |
+
dev:
|
| 5 |
+
type: duckdb
|
| 6 |
+
path: "src/dbt/customercore.duckdb"
|
| 7 |
+
schema: gold
|
| 8 |
+
extensions:
|
| 9 |
+
- httpfs
|
| 10 |
+
- iceberg
|
| 11 |
+
settings:
|
| 12 |
+
s3_endpoint: "localhost:9000"
|
| 13 |
+
s3_access_key_id: "minioadmin"
|
| 14 |
+
s3_secret_access_key: "minioadmin"
|
| 15 |
+
s3_use_ssl: false
|
| 16 |
+
s3_url_style: "path"
|
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/__init__.py
|
| 3 |
+
|
| 4 |
+
CustomerCore RAG (Retrieval-Augmented Generation) package.
|
| 5 |
+
|
| 6 |
+
Modules:
|
| 7 |
+
hybrid_retriever — ChromaDB dense search + BM25 with tenant isolation
|
| 8 |
+
semantic_cache — 3-layer cache (L0 embedding, L1 Redis exact, L2 vector)
|
| 9 |
+
reranker — Cross-encoder reranking of merged retrieval results
|
| 10 |
+
tool_rag — Dynamic semantic tool schema retrieval (Phase 8)
|
| 11 |
+
graph_rag — Unified B2B graph-RAG: semantic + relational traversal (Phase 8)
|
| 12 |
+
router — SLA-aware multi-model LLM routing (Phase 7)
|
| 13 |
+
llm_client — LiteLLM wrapper routing to Ollama/OpenRouter
|
| 14 |
+
"""
|
|
@@ -0,0 +1,617 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/graph_rag.py
|
| 3 |
+
|
| 4 |
+
Phase 8b: Unified B2B Graph-RAG Engine
|
| 5 |
+
|
| 6 |
+
== Why Graph-RAG Beats Pure Vector RAG for B2B ==
|
| 7 |
+
|
| 8 |
+
Pure vector RAG answers "find me tickets similar to this one."
|
| 9 |
+
That works for simple FAQ-style queries but breaks for B2B intelligence:
|
| 10 |
+
|
| 11 |
+
Q: "Why has acme-corp been escalating so often this quarter?"
|
| 12 |
+
Pure vector: Returns similar escalation tickets — no context.
|
| 13 |
+
Graph-RAG: Returns similar tickets + tenant health score + category
|
| 14 |
+
breakdown from DuckDB Gold + escalation trend from time series.
|
| 15 |
+
The LLM gets BOTH the examples AND the structured analytics.
|
| 16 |
+
|
| 17 |
+
B2B support platforms are fundamentally relational: tenants have accounts,
|
| 18 |
+
accounts have tiers, tiers have SLAs, SLAs have breach patterns. Vector
|
| 19 |
+
search alone misses all of this. Graph-RAG bridges the gap by combining:
|
| 20 |
+
|
| 21 |
+
Layer 1 (Graph) — B2BKnowledgeGraph: NetworkX DiGraph connecting
|
| 22 |
+
Tickets → Tenants → Categories → Metrics
|
| 23 |
+
Layer 2 (Vector) — HybridRetriever: BM25 + ChromaDB dense search
|
| 24 |
+
Layer 3 (SQL) — DuckDB queries against Gold mart Parquet files
|
| 25 |
+
|
| 26 |
+
== Architecture ==
|
| 27 |
+
|
| 28 |
+
Query: "Why is acme-corp escalating this month?"
|
| 29 |
+
│
|
| 30 |
+
├── [Vector] HybridRetriever.search(tenant_id="acme-corp", query=...)
|
| 31 |
+
│ └── Returns: 5 most similar past escalation tickets
|
| 32 |
+
│
|
| 33 |
+
├── [Graph] B2BKnowledgeGraph.get_tenant_context("acme-corp")
|
| 34 |
+
│ └── Returns: total tickets, top categories, priority breakdown
|
| 35 |
+
│
|
| 36 |
+
└── [SQL] DuckDB → Gold mart → support_agent_performance + ticket_funnel
|
| 37 |
+
└── Returns: volume trend, resolution rate, avg priority
|
| 38 |
+
|
| 39 |
+
↓
|
| 40 |
+
GraphRAGEngine._format_combined_context(...)
|
| 41 |
+
└── Produces: structured prompt context for the LLM
|
| 42 |
+
|
| 43 |
+
Run demo:
|
| 44 |
+
python -m src.rag.graph_rag
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
import logging
|
| 48 |
+
from collections import defaultdict
|
| 49 |
+
from dataclasses import dataclass, field
|
| 50 |
+
from typing import Optional, Any
|
| 51 |
+
|
| 52 |
+
logger = logging.getLogger(__name__)
|
| 53 |
+
|
| 54 |
+
# ── Node Dataclasses ───────────────────────────────────────────────────────────
|
| 55 |
+
|
| 56 |
+
@dataclass
|
| 57 |
+
class TicketNode:
|
| 58 |
+
ticket_id: str
|
| 59 |
+
tenant_id: str
|
| 60 |
+
text: str
|
| 61 |
+
category: str
|
| 62 |
+
priority: str
|
| 63 |
+
language: str = "en"
|
| 64 |
+
created_at: str = ""
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
@dataclass
|
| 68 |
+
class TenantNode:
|
| 69 |
+
tenant_id: str
|
| 70 |
+
tier: str = "professional"
|
| 71 |
+
metrics: dict = field(default_factory=dict)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
@dataclass
|
| 75 |
+
class CategoryNode:
|
| 76 |
+
category: str
|
| 77 |
+
ticket_count: int = 0
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
@dataclass
|
| 81 |
+
class GraphRAGResult:
|
| 82 |
+
"""Full result from a Graph-RAG query."""
|
| 83 |
+
query: str
|
| 84 |
+
tenant_id: str
|
| 85 |
+
similar_tickets: list # list[RetrievedDoc]
|
| 86 |
+
tenant_context: dict
|
| 87 |
+
sql_insights: list[dict]
|
| 88 |
+
combined_context: str
|
| 89 |
+
query_plan: list[str] # steps taken — for explainability/debugging
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
# ── B2B Knowledge Graph ────────────────────────────────────────────────────────
|
| 93 |
+
|
| 94 |
+
class B2BKnowledgeGraph:
|
| 95 |
+
"""
|
| 96 |
+
In-memory knowledge graph for B2B customer intelligence.
|
| 97 |
+
|
| 98 |
+
Nodes:
|
| 99 |
+
- tenant:{tenant_id} — B2B customer tenant
|
| 100 |
+
- ticket:{ticket_id} — support ticket
|
| 101 |
+
- category:{name} — support category (billing, technical, etc.)
|
| 102 |
+
|
| 103 |
+
Edges:
|
| 104 |
+
- (tenant) --[HAS_TICKET]--> (ticket)
|
| 105 |
+
- (ticket) --[BELONGS_TO]--> (category)
|
| 106 |
+
- (ticket) --[SIMILAR_TO]--> (ticket) ← added by RAG engine after retrieval
|
| 107 |
+
|
| 108 |
+
Usage:
|
| 109 |
+
graph = B2BKnowledgeGraph()
|
| 110 |
+
graph.add_ticket("acme-corp", "TKT-001", "API error on checkout",
|
| 111 |
+
category="technical", priority="high", language="en")
|
| 112 |
+
context = graph.get_tenant_context("acme-corp")
|
| 113 |
+
"""
|
| 114 |
+
|
| 115 |
+
def __init__(self):
|
| 116 |
+
try:
|
| 117 |
+
import networkx as nx
|
| 118 |
+
self._graph = nx.DiGraph()
|
| 119 |
+
except ImportError:
|
| 120 |
+
logger.warning("networkx not installed — using dict-based fallback graph")
|
| 121 |
+
self._graph = None
|
| 122 |
+
|
| 123 |
+
# Fallback storage (always populated, networkx is additive)
|
| 124 |
+
self._tenants: dict[str, TenantNode] = {}
|
| 125 |
+
self._tickets: dict[str, TicketNode] = {}
|
| 126 |
+
self._categories: dict[str, CategoryNode] = {}
|
| 127 |
+
self._tenant_tickets: dict[str, list[str]] = defaultdict(list)
|
| 128 |
+
self._ticket_category: dict[str, str] = {}
|
| 129 |
+
|
| 130 |
+
def add_ticket(
|
| 131 |
+
self,
|
| 132 |
+
tenant_id: str,
|
| 133 |
+
ticket_id: str,
|
| 134 |
+
text: str,
|
| 135 |
+
category: str = "general",
|
| 136 |
+
priority: str = "medium",
|
| 137 |
+
language: str = "en",
|
| 138 |
+
created_at: str = "",
|
| 139 |
+
):
|
| 140 |
+
"""Add a ticket node and connect it to tenant and category nodes."""
|
| 141 |
+
# Create nodes
|
| 142 |
+
ticket = TicketNode(
|
| 143 |
+
ticket_id=ticket_id, tenant_id=tenant_id, text=text,
|
| 144 |
+
category=category, priority=priority, language=language,
|
| 145 |
+
created_at=created_at,
|
| 146 |
+
)
|
| 147 |
+
self._tickets[ticket_id] = ticket
|
| 148 |
+
self._ticket_category[ticket_id] = category
|
| 149 |
+
|
| 150 |
+
if tenant_id not in self._tenants:
|
| 151 |
+
self._tenants[tenant_id] = TenantNode(tenant_id=tenant_id)
|
| 152 |
+
|
| 153 |
+
if category not in self._categories:
|
| 154 |
+
self._categories[category] = CategoryNode(category=category)
|
| 155 |
+
self._categories[category].ticket_count += 1
|
| 156 |
+
self._tenant_tickets[tenant_id].append(ticket_id)
|
| 157 |
+
|
| 158 |
+
# NetworkX graph
|
| 159 |
+
if self._graph is not None:
|
| 160 |
+
self._graph.add_node(f"ticket:{ticket_id}", **ticket.__dict__)
|
| 161 |
+
self._graph.add_node(f"tenant:{tenant_id}")
|
| 162 |
+
self._graph.add_node(f"category:{category}")
|
| 163 |
+
self._graph.add_edge(f"tenant:{tenant_id}", f"ticket:{ticket_id}",
|
| 164 |
+
rel="HAS_TICKET")
|
| 165 |
+
self._graph.add_edge(f"ticket:{ticket_id}", f"category:{category}",
|
| 166 |
+
rel="BELONGS_TO", priority=priority)
|
| 167 |
+
|
| 168 |
+
def add_tenant_metrics(self, tenant_id: str, metrics: dict):
|
| 169 |
+
"""Update a tenant node with structured metrics from the Gold layer."""
|
| 170 |
+
if tenant_id not in self._tenants:
|
| 171 |
+
self._tenants[tenant_id] = TenantNode(tenant_id=tenant_id)
|
| 172 |
+
self._tenants[tenant_id].metrics = metrics
|
| 173 |
+
|
| 174 |
+
if self._graph is not None:
|
| 175 |
+
self._graph.add_node(f"tenant:{tenant_id}", **metrics)
|
| 176 |
+
|
| 177 |
+
def get_tenant_context(self, tenant_id: str) -> dict:
|
| 178 |
+
"""
|
| 179 |
+
Build a structured context dict for a tenant from the graph.
|
| 180 |
+
Includes ticket volume, category breakdown, priority distribution,
|
| 181 |
+
language breakdown, and any stored metrics.
|
| 182 |
+
"""
|
| 183 |
+
ticket_ids = self._tenant_tickets.get(tenant_id, [])
|
| 184 |
+
tickets = [self._tickets[tid] for tid in ticket_ids if tid in self._tickets]
|
| 185 |
+
|
| 186 |
+
if not tickets:
|
| 187 |
+
return {
|
| 188 |
+
"tenant_id": tenant_id,
|
| 189 |
+
"total_tickets": 0,
|
| 190 |
+
"message": "No tickets found in graph for this tenant.",
|
| 191 |
+
}
|
| 192 |
+
|
| 193 |
+
# Category breakdown
|
| 194 |
+
category_counts: dict[str, int] = defaultdict(int)
|
| 195 |
+
priority_counts: dict[str, int] = defaultdict(int)
|
| 196 |
+
language_counts: dict[str, int] = defaultdict(int)
|
| 197 |
+
|
| 198 |
+
for t in tickets:
|
| 199 |
+
category_counts[t.category] += 1
|
| 200 |
+
priority_counts[t.priority] += 1
|
| 201 |
+
language_counts[t.language] += 1
|
| 202 |
+
|
| 203 |
+
top_categories = sorted(
|
| 204 |
+
category_counts.items(), key=lambda x: x[1], reverse=True
|
| 205 |
+
)[:5]
|
| 206 |
+
|
| 207 |
+
# Escalation proxy: high + critical tickets as % of total
|
| 208 |
+
escalation_count = (
|
| 209 |
+
priority_counts.get("high", 0) + priority_counts.get("critical", 0)
|
| 210 |
+
)
|
| 211 |
+
escalation_rate = round(escalation_count / len(tickets) * 100, 1)
|
| 212 |
+
|
| 213 |
+
tenant_node = self._tenants.get(tenant_id, TenantNode(tenant_id))
|
| 214 |
+
|
| 215 |
+
return {
|
| 216 |
+
"tenant_id": tenant_id,
|
| 217 |
+
"tier": tenant_node.tier,
|
| 218 |
+
"total_tickets": len(tickets),
|
| 219 |
+
"top_categories": [
|
| 220 |
+
{"category": cat, "count": cnt} for cat, cnt in top_categories
|
| 221 |
+
],
|
| 222 |
+
"priority_breakdown": dict(priority_counts),
|
| 223 |
+
"language_breakdown": dict(language_counts),
|
| 224 |
+
"escalation_rate_pct": escalation_rate,
|
| 225 |
+
"escalation_count": escalation_count,
|
| 226 |
+
"gold_metrics": tenant_node.metrics,
|
| 227 |
+
}
|
| 228 |
+
|
| 229 |
+
def find_similar_tenants(
|
| 230 |
+
self, tenant_id: str, by: str = "category", top_k: int = 3
|
| 231 |
+
) -> list[dict]:
|
| 232 |
+
"""
|
| 233 |
+
Find tenants with similar issue profiles based on category distribution.
|
| 234 |
+
Useful for benchmarking: "how does acme-corp compare to similar tenants?"
|
| 235 |
+
"""
|
| 236 |
+
target_cats = set(
|
| 237 |
+
self._tickets[tid].category
|
| 238 |
+
for tid in self._tenant_tickets.get(tenant_id, [])
|
| 239 |
+
if tid in self._tickets
|
| 240 |
+
)
|
| 241 |
+
|
| 242 |
+
scores = []
|
| 243 |
+
for other_id, ticket_ids in self._tenant_tickets.items():
|
| 244 |
+
if other_id == tenant_id:
|
| 245 |
+
continue
|
| 246 |
+
other_cats = set(
|
| 247 |
+
self._tickets[tid].category
|
| 248 |
+
for tid in ticket_ids if tid in self._tickets
|
| 249 |
+
)
|
| 250 |
+
# Jaccard similarity
|
| 251 |
+
if target_cats or other_cats:
|
| 252 |
+
overlap = len(target_cats & other_cats)
|
| 253 |
+
union = len(target_cats | other_cats)
|
| 254 |
+
similarity = overlap / union if union > 0 else 0.0
|
| 255 |
+
scores.append({"tenant_id": other_id, "similarity": round(similarity, 3)})
|
| 256 |
+
|
| 257 |
+
return sorted(scores, key=lambda x: x["similarity"], reverse=True)[:top_k]
|
| 258 |
+
|
| 259 |
+
def get_category_trends(self) -> dict[str, int]:
|
| 260 |
+
"""Return ticket count per category across all tenants."""
|
| 261 |
+
return {cat: node.ticket_count for cat, node in self._categories.items()}
|
| 262 |
+
|
| 263 |
+
def node_count(self) -> int:
|
| 264 |
+
if self._graph is not None:
|
| 265 |
+
return self._graph.number_of_nodes()
|
| 266 |
+
return len(self._tickets) + len(self._tenants) + len(self._categories)
|
| 267 |
+
|
| 268 |
+
def edge_count(self) -> int:
|
| 269 |
+
if self._graph is not None:
|
| 270 |
+
return self._graph.number_of_edges()
|
| 271 |
+
return len(self._tickets) * 2 # each ticket has 2 edges
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
# ── SQL Insights (DuckDB Gold Layer) ──────────────────────────────────────────
|
| 275 |
+
|
| 276 |
+
def _load_duckdb_insights(tenant_id: str, gold_db_path: str = "data/gold/customercore_gold.duckdb") -> list[dict]:
|
| 277 |
+
"""
|
| 278 |
+
Query the DuckDB Gold layer for structured analytics about a tenant.
|
| 279 |
+
Returns empty list gracefully if the DB doesn't exist (local dev mode).
|
| 280 |
+
"""
|
| 281 |
+
insights = []
|
| 282 |
+
try:
|
| 283 |
+
import duckdb
|
| 284 |
+
conn = duckdb.connect(gold_db_path, read_only=True)
|
| 285 |
+
|
| 286 |
+
# Ticket funnel for this tenant
|
| 287 |
+
try:
|
| 288 |
+
rows = conn.execute("""
|
| 289 |
+
SELECT event_type, COUNT(*) as count, AVG(priority_score) as avg_priority
|
| 290 |
+
FROM gold_gold.ticket_funnel_daily
|
| 291 |
+
WHERE tenant_id = ?
|
| 292 |
+
GROUP BY event_type
|
| 293 |
+
ORDER BY count DESC
|
| 294 |
+
LIMIT 5
|
| 295 |
+
""", [tenant_id]).fetchall()
|
| 296 |
+
for r in rows:
|
| 297 |
+
insights.append({
|
| 298 |
+
"source": "ticket_funnel",
|
| 299 |
+
"event_type": r[0],
|
| 300 |
+
"count": r[1],
|
| 301 |
+
"avg_priority": round(r[2] or 0, 2),
|
| 302 |
+
})
|
| 303 |
+
except Exception:
|
| 304 |
+
pass
|
| 305 |
+
|
| 306 |
+
# Customer health
|
| 307 |
+
try:
|
| 308 |
+
rows = conn.execute("""
|
| 309 |
+
SELECT snapshot_date, avg_priority, ticket_count
|
| 310 |
+
FROM gold_gold.customer_health_daily
|
| 311 |
+
WHERE tenant_id = ?
|
| 312 |
+
ORDER BY snapshot_date DESC
|
| 313 |
+
LIMIT 3
|
| 314 |
+
""", [tenant_id]).fetchall()
|
| 315 |
+
for r in rows:
|
| 316 |
+
insights.append({
|
| 317 |
+
"source": "customer_health",
|
| 318 |
+
"date": str(r[0]),
|
| 319 |
+
"avg_priority": round(float(r[1] or 0), 2),
|
| 320 |
+
"ticket_count": r[2],
|
| 321 |
+
})
|
| 322 |
+
except Exception:
|
| 323 |
+
pass
|
| 324 |
+
|
| 325 |
+
conn.close()
|
| 326 |
+
|
| 327 |
+
except Exception as e:
|
| 328 |
+
logger.debug("DuckDB Gold layer not available: %s", e)
|
| 329 |
+
|
| 330 |
+
return insights
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
# ── Graph-RAG Engine ───────────────────────────────────────────────────────────
|
| 334 |
+
|
| 335 |
+
class GraphRAGEngine:
|
| 336 |
+
"""
|
| 337 |
+
Unified B2B Graph-RAG Query Engine.
|
| 338 |
+
|
| 339 |
+
Combines three retrieval strategies in a single query:
|
| 340 |
+
1. Vector + BM25 hybrid search (tenant-isolated, Phase 6)
|
| 341 |
+
2. Graph traversal for tenant context and category trends
|
| 342 |
+
3. SQL analytics from DuckDB Gold marts
|
| 343 |
+
|
| 344 |
+
The resulting combined_context is ready for injection into any LLM prompt.
|
| 345 |
+
|
| 346 |
+
Usage:
|
| 347 |
+
engine = GraphRAGEngine(retriever=hybrid_retriever, graph=knowledge_graph)
|
| 348 |
+
result = engine.query(
|
| 349 |
+
tenant_id="acme-corp",
|
| 350 |
+
question="Why has this tenant been escalating so often this quarter?",
|
| 351 |
+
k=5
|
| 352 |
+
)
|
| 353 |
+
# Inject result.combined_context into your LLM prompt
|
| 354 |
+
"""
|
| 355 |
+
|
| 356 |
+
def __init__(
|
| 357 |
+
self,
|
| 358 |
+
retriever=None, # HybridRetriever instance (or None for graph-only)
|
| 359 |
+
graph: Optional[B2BKnowledgeGraph] = None,
|
| 360 |
+
gold_db_path: str = "data/gold/customercore_gold.duckdb",
|
| 361 |
+
):
|
| 362 |
+
self.retriever = retriever
|
| 363 |
+
self.graph = graph or B2BKnowledgeGraph()
|
| 364 |
+
self.gold_db_path = gold_db_path
|
| 365 |
+
|
| 366 |
+
def index_ticket(
|
| 367 |
+
self,
|
| 368 |
+
tenant_id: str,
|
| 369 |
+
ticket_id: str,
|
| 370 |
+
text: str,
|
| 371 |
+
metadata: dict | None = None,
|
| 372 |
+
):
|
| 373 |
+
"""
|
| 374 |
+
Index a ticket into both the vector retriever and the knowledge graph.
|
| 375 |
+
Single entry point — no need to call retriever and graph separately.
|
| 376 |
+
"""
|
| 377 |
+
metadata = metadata or {}
|
| 378 |
+
category = metadata.get("category", "general")
|
| 379 |
+
priority = metadata.get("priority", "medium")
|
| 380 |
+
language = metadata.get("language", "en")
|
| 381 |
+
|
| 382 |
+
# Add to graph
|
| 383 |
+
self.graph.add_ticket(
|
| 384 |
+
tenant_id=tenant_id,
|
| 385 |
+
ticket_id=ticket_id,
|
| 386 |
+
text=text,
|
| 387 |
+
category=category,
|
| 388 |
+
priority=priority,
|
| 389 |
+
language=language,
|
| 390 |
+
)
|
| 391 |
+
|
| 392 |
+
# Add to vector retriever if available
|
| 393 |
+
if self.retriever is not None:
|
| 394 |
+
self.retriever.index_ticket(tenant_id, ticket_id, text, metadata)
|
| 395 |
+
|
| 396 |
+
def query(
|
| 397 |
+
self,
|
| 398 |
+
tenant_id: str,
|
| 399 |
+
question: str,
|
| 400 |
+
k: int = 5,
|
| 401 |
+
include_sql: bool = True,
|
| 402 |
+
) -> GraphRAGResult:
|
| 403 |
+
"""
|
| 404 |
+
Full Graph-RAG query.
|
| 405 |
+
|
| 406 |
+
Steps (all recorded in query_plan for explainability):
|
| 407 |
+
1. Vector+BM25 hybrid retrieval (tenant-isolated)
|
| 408 |
+
2. Graph traversal for tenant context
|
| 409 |
+
3. SQL analytics from DuckDB Gold (optional)
|
| 410 |
+
4. Format combined context for LLM injection
|
| 411 |
+
"""
|
| 412 |
+
query_plan = []
|
| 413 |
+
|
| 414 |
+
# Step 1: Hybrid vector retrieval
|
| 415 |
+
similar_tickets = []
|
| 416 |
+
if self.retriever is not None:
|
| 417 |
+
similar_tickets = self.retriever.search(tenant_id, question, k=k)
|
| 418 |
+
query_plan.append(
|
| 419 |
+
f"[Vector+BM25] Retrieved {len(similar_tickets)} similar tickets "
|
| 420 |
+
f"for tenant={tenant_id}"
|
| 421 |
+
)
|
| 422 |
+
else:
|
| 423 |
+
query_plan.append("[Vector+BM25] Retriever not configured — skipped")
|
| 424 |
+
|
| 425 |
+
# Step 2: Graph context
|
| 426 |
+
tenant_context = self.graph.get_tenant_context(tenant_id)
|
| 427 |
+
query_plan.append(
|
| 428 |
+
f"[Graph] Tenant context: {tenant_context['total_tickets']} tickets, "
|
| 429 |
+
f"escalation_rate={tenant_context.get('escalation_rate_pct', 0)}%"
|
| 430 |
+
)
|
| 431 |
+
|
| 432 |
+
# Step 3: SQL analytics
|
| 433 |
+
sql_insights = []
|
| 434 |
+
if include_sql:
|
| 435 |
+
sql_insights = _load_duckdb_insights(tenant_id, self.gold_db_path)
|
| 436 |
+
query_plan.append(
|
| 437 |
+
f"[SQL] Gold layer: {len(sql_insights)} insight rows from DuckDB"
|
| 438 |
+
)
|
| 439 |
+
else:
|
| 440 |
+
query_plan.append("[SQL] Skipped (include_sql=False)")
|
| 441 |
+
|
| 442 |
+
# Step 4: Format combined context
|
| 443 |
+
combined_context = self._format_combined_context(
|
| 444 |
+
question, similar_tickets, tenant_context, sql_insights
|
| 445 |
+
)
|
| 446 |
+
query_plan.append("[Format] Combined context built — ready for LLM injection")
|
| 447 |
+
|
| 448 |
+
return GraphRAGResult(
|
| 449 |
+
query=question,
|
| 450 |
+
tenant_id=tenant_id,
|
| 451 |
+
similar_tickets=similar_tickets,
|
| 452 |
+
tenant_context=tenant_context,
|
| 453 |
+
sql_insights=sql_insights,
|
| 454 |
+
combined_context=combined_context,
|
| 455 |
+
query_plan=query_plan,
|
| 456 |
+
)
|
| 457 |
+
|
| 458 |
+
def _format_combined_context(
|
| 459 |
+
self,
|
| 460 |
+
question: str,
|
| 461 |
+
similar_tickets: list,
|
| 462 |
+
tenant_context: dict,
|
| 463 |
+
sql_insights: list[dict],
|
| 464 |
+
) -> str:
|
| 465 |
+
"""
|
| 466 |
+
Format all retrieved information into a structured context string
|
| 467 |
+
ready for injection into an LLM system prompt.
|
| 468 |
+
"""
|
| 469 |
+
lines = []
|
| 470 |
+
|
| 471 |
+
# Tenant profile
|
| 472 |
+
lines.append("=== TENANT PROFILE ===")
|
| 473 |
+
lines.append(f"Tenant ID : {tenant_context.get('tenant_id', 'unknown')}")
|
| 474 |
+
lines.append(f"Tier : {tenant_context.get('tier', 'unknown')}")
|
| 475 |
+
lines.append(f"Total Tickets : {tenant_context.get('total_tickets', 0)}")
|
| 476 |
+
lines.append(f"Escalation Rate : {tenant_context.get('escalation_rate_pct', 0)}%")
|
| 477 |
+
|
| 478 |
+
top_cats = tenant_context.get("top_categories", [])
|
| 479 |
+
if top_cats:
|
| 480 |
+
cats_str = ", ".join(
|
| 481 |
+
f"{c['category']}({c['count']})" for c in top_cats[:3]
|
| 482 |
+
)
|
| 483 |
+
lines.append(f"Top Categories : {cats_str}")
|
| 484 |
+
|
| 485 |
+
lang_breakdown = tenant_context.get("language_breakdown", {})
|
| 486 |
+
if lang_breakdown:
|
| 487 |
+
lang_str = ", ".join(f"{k}:{v}" for k, v in lang_breakdown.items())
|
| 488 |
+
lines.append(f"Languages Used : {lang_str}")
|
| 489 |
+
|
| 490 |
+
# Similar past tickets
|
| 491 |
+
if similar_tickets:
|
| 492 |
+
lines.append("\n=== SIMILAR PAST TICKETS ===")
|
| 493 |
+
for i, doc in enumerate(similar_tickets[:5], 1):
|
| 494 |
+
lines.append(
|
| 495 |
+
f"{i}. [{doc.category.upper()}|{doc.priority}] "
|
| 496 |
+
f"{doc.text[:100]}..."
|
| 497 |
+
)
|
| 498 |
+
|
| 499 |
+
# SQL insights
|
| 500 |
+
if sql_insights:
|
| 501 |
+
lines.append("\n=== ANALYTICS (Gold Layer) ===")
|
| 502 |
+
for ins in sql_insights[:5]:
|
| 503 |
+
src = ins.get("source", "")
|
| 504 |
+
if src == "ticket_funnel":
|
| 505 |
+
lines.append(
|
| 506 |
+
f"- Funnel: {ins.get('event_type')} "
|
| 507 |
+
f"count={ins.get('count')} "
|
| 508 |
+
f"avg_priority={ins.get('avg_priority')}"
|
| 509 |
+
)
|
| 510 |
+
elif src == "customer_health":
|
| 511 |
+
lines.append(
|
| 512 |
+
f"- Health [{ins.get('date')}]: "
|
| 513 |
+
f"tickets={ins.get('ticket_count')} "
|
| 514 |
+
f"avg_priority={ins.get('avg_priority')}"
|
| 515 |
+
)
|
| 516 |
+
|
| 517 |
+
lines.append(f"\n=== ORIGINAL QUESTION ===\n{question}")
|
| 518 |
+
return "\n".join(lines)
|
| 519 |
+
|
| 520 |
+
|
| 521 |
+
# ── Module-level singleton ─────────────────────────────────────────────────────
|
| 522 |
+
_default_engine: Optional[GraphRAGEngine] = None
|
| 523 |
+
|
| 524 |
+
|
| 525 |
+
def get_engine() -> GraphRAGEngine:
|
| 526 |
+
global _default_engine
|
| 527 |
+
if _default_engine is None:
|
| 528 |
+
_default_engine = GraphRAGEngine()
|
| 529 |
+
return _default_engine
|
| 530 |
+
|
| 531 |
+
|
| 532 |
+
# ── Standalone Demo ────────────────────────────────────────────────────────────
|
| 533 |
+
if __name__ == "__main__":
|
| 534 |
+
import os
|
| 535 |
+
os.environ["PYTHONIOENCODING"] = "utf-8"
|
| 536 |
+
|
| 537 |
+
print("=" * 65)
|
| 538 |
+
print("CustomerCore Phase 8b - Unified B2B Graph-RAG Demo")
|
| 539 |
+
print("=" * 65)
|
| 540 |
+
|
| 541 |
+
# Build knowledge graph
|
| 542 |
+
graph = B2BKnowledgeGraph()
|
| 543 |
+
|
| 544 |
+
# Index tickets for two tenants
|
| 545 |
+
acme_tickets = [
|
| 546 |
+
("TKT-A001", "API returning 500 errors on checkout endpoint", "technical", "critical", "en"),
|
| 547 |
+
("TKT-A002", "March invoice shows double charge for $420", "billing", "high", "en"),
|
| 548 |
+
("TKT-A003", "Login broken after your latest deployment", "technical", "high", "en"),
|
| 549 |
+
("TKT-A004", "Need to cancel our subscription next month", "subscription", "medium", "en"),
|
| 550 |
+
("TKT-A005", "Export of customer data timing out after 30s", "technical", "high", "en"),
|
| 551 |
+
("TKT-A006", "Zahlung fehlgeschlagen - bitte helfen", "billing", "high", "de"),
|
| 552 |
+
("TKT-A007", "Notre tableau de bord ne fonctionne plus", "technical", "critical", "fr"),
|
| 553 |
+
]
|
| 554 |
+
globex_tickets = [
|
| 555 |
+
("TKT-B001", "Latency spike on EU cluster affecting production", "technical", "critical", "en"),
|
| 556 |
+
("TKT-B002", "Overcharged by 2400 USD last billing cycle", "billing", "high", "en"),
|
| 557 |
+
("TKT-B003", "Data export failing with timeout error", "technical", "high", "en"),
|
| 558 |
+
]
|
| 559 |
+
|
| 560 |
+
print("\nIndexing tickets into knowledge graph...")
|
| 561 |
+
for tid, text, cat, prio, lang in acme_tickets:
|
| 562 |
+
graph.add_ticket("acme-corp", tid, text, cat, prio, lang)
|
| 563 |
+
for tid, text, cat, prio, lang in globex_tickets:
|
| 564 |
+
graph.add_ticket("globex-inc", tid, text, cat, prio, lang)
|
| 565 |
+
|
| 566 |
+
# Add mock tenant metrics
|
| 567 |
+
graph.add_tenant_metrics("acme-corp", {
|
| 568 |
+
"avg_resolution_hours": 4.2,
|
| 569 |
+
"open_tickets": 7,
|
| 570 |
+
"csat_score": 3.8,
|
| 571 |
+
})
|
| 572 |
+
|
| 573 |
+
print(f"Graph nodes: {graph.node_count()}")
|
| 574 |
+
print(f"Graph edges: {graph.edge_count()}")
|
| 575 |
+
|
| 576 |
+
# Tenant context
|
| 577 |
+
print("\n" + "=" * 40)
|
| 578 |
+
print("TENANT CONTEXT: acme-corp")
|
| 579 |
+
ctx = graph.get_tenant_context("acme-corp")
|
| 580 |
+
print(f" Total Tickets : {ctx['total_tickets']}")
|
| 581 |
+
print(f" Escalation Rate : {ctx['escalation_rate_pct']}%")
|
| 582 |
+
print(f" Priority Breakdown: {ctx['priority_breakdown']}")
|
| 583 |
+
print(f" Language Breakdown: {ctx['language_breakdown']}")
|
| 584 |
+
print(f" Top Categories : {[c['category'] for c in ctx['top_categories']]}")
|
| 585 |
+
|
| 586 |
+
# Similar tenants
|
| 587 |
+
print("\n" + "=" * 40)
|
| 588 |
+
print("SIMILAR TENANTS TO acme-corp:")
|
| 589 |
+
similar = graph.find_similar_tenants("acme-corp")
|
| 590 |
+
for s in similar:
|
| 591 |
+
print(f" {s['tenant_id']} (similarity={s['similarity']})")
|
| 592 |
+
|
| 593 |
+
# Category trends
|
| 594 |
+
print("\n" + "=" * 40)
|
| 595 |
+
print("GLOBAL CATEGORY TRENDS:")
|
| 596 |
+
trends = graph.get_category_trends()
|
| 597 |
+
for cat, count in sorted(trends.items(), key=lambda x: x[1], reverse=True):
|
| 598 |
+
print(f" {cat:<15}: {count} tickets")
|
| 599 |
+
|
| 600 |
+
# Full Graph-RAG query (no retriever, no DuckDB in demo)
|
| 601 |
+
engine = GraphRAGEngine(retriever=None, graph=graph)
|
| 602 |
+
result = engine.query(
|
| 603 |
+
tenant_id="acme-corp",
|
| 604 |
+
question="Why has acme-corp been escalating so many technical issues this month?",
|
| 605 |
+
include_sql=False, # no DuckDB in demo
|
| 606 |
+
)
|
| 607 |
+
|
| 608 |
+
print("\n" + "=" * 40)
|
| 609 |
+
print("GRAPH-RAG QUERY RESULT:")
|
| 610 |
+
print("\nQuery Plan:")
|
| 611 |
+
for step in result.query_plan:
|
| 612 |
+
print(f" - {step}")
|
| 613 |
+
|
| 614 |
+
print("\nCombined Context (injected into LLM prompt):")
|
| 615 |
+
print("-" * 50)
|
| 616 |
+
print(result.combined_context)
|
| 617 |
+
print("=" * 65)
|
|
@@ -0,0 +1,470 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/hybrid_retriever.py
|
| 3 |
+
|
| 4 |
+
Phase 6: Multi-Tenant Vector DB Security + Hybrid Retrieval Pipeline
|
| 5 |
+
|
| 6 |
+
== The Gap We Are Closing ==
|
| 7 |
+
A naive ChromaDB implementation stores all tenant embeddings in one shared
|
| 8 |
+
collection. If Tenant A searches for "billing refund" and Tenant B has
|
| 9 |
+
similar tickets, Tenant B's private customer data can surface in Tenant A's
|
| 10 |
+
results. For a B2B SaaS platform this is a catastrophic GDPR violation.
|
| 11 |
+
|
| 12 |
+
== What This Does ==
|
| 13 |
+
1. STRICT TENANT ISOLATION — Every query to ChromaDB passes a metadata filter:
|
| 14 |
+
where={"tenant_id": current_tenant_id}
|
| 15 |
+
No results can ever come from a different tenant regardless of similarity score.
|
| 16 |
+
|
| 17 |
+
2. HYBRID RETRIEVAL — Two retrieval strategies run in parallel:
|
| 18 |
+
a. Dense semantic search via ChromaDB (embedding similarity)
|
| 19 |
+
b. Sparse keyword search via BM25 (exact/rare term matching)
|
| 20 |
+
Results are merged using Reciprocal Rank Fusion (RRF) — a proven
|
| 21 |
+
technique that usually outperforms either strategy alone.
|
| 22 |
+
|
| 23 |
+
3. CROSS-ENCODER RERANKING — Top-k merged results are re-scored by a
|
| 24 |
+
lightweight cross-encoder model to surface the most semantically
|
| 25 |
+
relevant results for the exact query.
|
| 26 |
+
|
| 27 |
+
== Architecture Decision: Single Collection + Metadata Filter ==
|
| 28 |
+
We use ONE ChromaDB collection "customercore_tickets" for all tenants.
|
| 29 |
+
Isolation is enforced at query time via the `where` filter.
|
| 30 |
+
This is the standard enterprise pattern (used by Weaviate, Pinecone, Qdrant)
|
| 31 |
+
because it avoids managing hundreds of collections while keeping isolation.
|
| 32 |
+
|
| 33 |
+
Run standalone demo:
|
| 34 |
+
python -m src.rag.hybrid_retriever
|
| 35 |
+
"""
|
| 36 |
+
|
| 37 |
+
import hashlib
|
| 38 |
+
import logging
|
| 39 |
+
import time
|
| 40 |
+
from dataclasses import dataclass, field
|
| 41 |
+
from typing import Optional
|
| 42 |
+
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
# ── Config ─────────────────────────────────────────────────────────────────────
|
| 46 |
+
CHROMA_COLLECTION = "customercore_tickets"
|
| 47 |
+
DEFAULT_DENSE_K = 10 # retrieve top-10 from ChromaDB
|
| 48 |
+
DEFAULT_SPARSE_K = 10 # retrieve top-10 from BM25
|
| 49 |
+
DEFAULT_FINAL_K = 5 # return top-5 after reranking
|
| 50 |
+
RRF_CONSTANT = 60 # standard Reciprocal Rank Fusion constant
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
@dataclass
|
| 54 |
+
class RetrievedDoc:
|
| 55 |
+
"""A single retrieved document with provenance metadata."""
|
| 56 |
+
doc_id: str
|
| 57 |
+
text: str
|
| 58 |
+
tenant_id: str
|
| 59 |
+
ticket_id: str = ""
|
| 60 |
+
category: str = ""
|
| 61 |
+
priority: str = ""
|
| 62 |
+
score: float = 0.0
|
| 63 |
+
source: str = "" # "dense" | "sparse" | "fused"
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
# ── Reciprocal Rank Fusion ─────────────────────────────────────────────────────
|
| 67 |
+
|
| 68 |
+
def reciprocal_rank_fusion(
|
| 69 |
+
dense_results: list[RetrievedDoc],
|
| 70 |
+
sparse_results: list[RetrievedDoc],
|
| 71 |
+
k: int = RRF_CONSTANT,
|
| 72 |
+
) -> list[RetrievedDoc]:
|
| 73 |
+
"""
|
| 74 |
+
Merge dense and sparse result lists using Reciprocal Rank Fusion.
|
| 75 |
+
RRF score = Σ 1/(k + rank_i) for each list the document appears in.
|
| 76 |
+
Higher is better. Naturally handles result deduplication.
|
| 77 |
+
"""
|
| 78 |
+
rrf_scores: dict[str, float] = {}
|
| 79 |
+
doc_map: dict[str, RetrievedDoc] = {}
|
| 80 |
+
|
| 81 |
+
for rank, doc in enumerate(dense_results, start=1):
|
| 82 |
+
rrf_scores[doc.doc_id] = rrf_scores.get(doc.doc_id, 0) + 1 / (k + rank)
|
| 83 |
+
doc_map[doc.doc_id] = doc
|
| 84 |
+
|
| 85 |
+
for rank, doc in enumerate(sparse_results, start=1):
|
| 86 |
+
rrf_scores[doc.doc_id] = rrf_scores.get(doc.doc_id, 0) + 1 / (k + rank)
|
| 87 |
+
if doc.doc_id not in doc_map:
|
| 88 |
+
doc_map[doc.doc_id] = doc
|
| 89 |
+
|
| 90 |
+
fused = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
|
| 91 |
+
results = []
|
| 92 |
+
for doc_id, score in fused:
|
| 93 |
+
doc = doc_map[doc_id]
|
| 94 |
+
doc.score = score
|
| 95 |
+
doc.source = "fused"
|
| 96 |
+
results.append(doc)
|
| 97 |
+
return results
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
# ── BM25 Index (tenant-isolated) ───────────────────────────────────────────────
|
| 101 |
+
|
| 102 |
+
class TenantBM25Index:
|
| 103 |
+
"""
|
| 104 |
+
Lightweight in-memory BM25 index partitioned per tenant.
|
| 105 |
+
Uses rank_bm25.BM25Okapi under the hood.
|
| 106 |
+
|
| 107 |
+
In production this would be backed by Elasticsearch or OpenSearch
|
| 108 |
+
with a `tenant_id` field filter.
|
| 109 |
+
"""
|
| 110 |
+
|
| 111 |
+
def __init__(self):
|
| 112 |
+
self._corpora: dict[str, list[str]] = {} # tenant_id → [text, ...]
|
| 113 |
+
self._doc_ids: dict[str, list[str]] = {} # tenant_id → [doc_id, ...]
|
| 114 |
+
self._doc_meta: dict[str, dict] = {} # doc_id → metadata
|
| 115 |
+
self._indexes: dict[str, object] = {} # tenant_id → BM25Okapi
|
| 116 |
+
self._dirty: set[str] = set() # tenants needing re-index
|
| 117 |
+
|
| 118 |
+
def add_document(self, tenant_id: str, doc_id: str, text: str, metadata: dict):
|
| 119 |
+
"""Add a document to the tenant-isolated BM25 index."""
|
| 120 |
+
if tenant_id not in self._corpora:
|
| 121 |
+
self._corpora[tenant_id] = []
|
| 122 |
+
self._doc_ids[tenant_id] = []
|
| 123 |
+
|
| 124 |
+
# Avoid duplicates
|
| 125 |
+
if doc_id in self._doc_ids[tenant_id]:
|
| 126 |
+
return
|
| 127 |
+
|
| 128 |
+
self._corpora[tenant_id].append(text)
|
| 129 |
+
self._doc_ids[tenant_id].append(doc_id)
|
| 130 |
+
self._doc_meta[doc_id] = {**metadata, "tenant_id": tenant_id}
|
| 131 |
+
self._dirty.add(tenant_id)
|
| 132 |
+
|
| 133 |
+
def _build_index(self, tenant_id: str):
|
| 134 |
+
"""Build/rebuild BM25Okapi index for a specific tenant."""
|
| 135 |
+
try:
|
| 136 |
+
from rank_bm25 import BM25Okapi
|
| 137 |
+
except ImportError:
|
| 138 |
+
raise ImportError(
|
| 139 |
+
"rank_bm25 is required for BM25 retrieval. "
|
| 140 |
+
"Install with: pip install rank-bm25"
|
| 141 |
+
)
|
| 142 |
+
tokenized = [doc.lower().split() for doc in self._corpora[tenant_id]]
|
| 143 |
+
self._indexes[tenant_id] = BM25Okapi(tokenized)
|
| 144 |
+
self._dirty.discard(tenant_id)
|
| 145 |
+
|
| 146 |
+
def search(self, tenant_id: str, query: str, k: int = 10) -> list[RetrievedDoc]:
|
| 147 |
+
"""
|
| 148 |
+
Search the BM25 index for a specific tenant.
|
| 149 |
+
IMPORTANT: Only documents belonging to `tenant_id` are returned.
|
| 150 |
+
Cross-tenant results are architecturally impossible in this design.
|
| 151 |
+
"""
|
| 152 |
+
if tenant_id not in self._corpora or not self._corpora[tenant_id]:
|
| 153 |
+
return []
|
| 154 |
+
|
| 155 |
+
if tenant_id in self._dirty:
|
| 156 |
+
self._build_index(tenant_id)
|
| 157 |
+
|
| 158 |
+
query_tokens = query.lower().split()
|
| 159 |
+
index = self._indexes[tenant_id]
|
| 160 |
+
scores = index.get_scores(query_tokens)
|
| 161 |
+
|
| 162 |
+
# Get top-k indices
|
| 163 |
+
top_k_idx = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
|
| 164 |
+
|
| 165 |
+
results = []
|
| 166 |
+
for idx in top_k_idx:
|
| 167 |
+
doc_id = self._doc_ids[tenant_id][idx]
|
| 168 |
+
meta = self._doc_meta.get(doc_id, {})
|
| 169 |
+
results.append(RetrievedDoc(
|
| 170 |
+
doc_id=doc_id,
|
| 171 |
+
text=self._corpora[tenant_id][idx],
|
| 172 |
+
tenant_id=tenant_id,
|
| 173 |
+
ticket_id=meta.get("ticket_id", ""),
|
| 174 |
+
category=meta.get("category", ""),
|
| 175 |
+
priority=meta.get("priority", ""),
|
| 176 |
+
score=float(scores[idx]),
|
| 177 |
+
source="sparse",
|
| 178 |
+
))
|
| 179 |
+
return results
|
| 180 |
+
|
| 181 |
+
def document_count(self, tenant_id: Optional[str] = None) -> int:
|
| 182 |
+
if tenant_id:
|
| 183 |
+
return len(self._corpora.get(tenant_id, []))
|
| 184 |
+
return sum(len(v) for v in self._corpora.values())
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
# ── ChromaDB Dense Retriever (tenant-isolated) ─────────────────────────────────
|
| 188 |
+
|
| 189 |
+
class TenantChromaRetriever:
|
| 190 |
+
"""
|
| 191 |
+
ChromaDB-backed dense semantic retriever with strict tenant isolation.
|
| 192 |
+
|
| 193 |
+
All tickets share a single collection. Tenant isolation is enforced
|
| 194 |
+
at query time using ChromaDB's `where` metadata filter.
|
| 195 |
+
This is the same pattern used by enterprise vector DBs (Weaviate, Pinecone).
|
| 196 |
+
"""
|
| 197 |
+
|
| 198 |
+
def __init__(self, chroma_client=None, embedding_fn=None):
|
| 199 |
+
self._client = chroma_client
|
| 200 |
+
self._embedding_fn = embedding_fn
|
| 201 |
+
self._collection = None
|
| 202 |
+
|
| 203 |
+
def _get_collection(self):
|
| 204 |
+
if self._collection is None and self._client is not None:
|
| 205 |
+
self._collection = self._client.get_or_create_collection(
|
| 206 |
+
name=CHROMA_COLLECTION,
|
| 207 |
+
metadata={"hnsw:space": "cosine"},
|
| 208 |
+
)
|
| 209 |
+
return self._collection
|
| 210 |
+
|
| 211 |
+
def add_document(self, tenant_id: str, doc_id: str, text: str, metadata: dict):
|
| 212 |
+
"""
|
| 213 |
+
Add a document to ChromaDB with tenant_id injected into metadata.
|
| 214 |
+
This is what guarantees query-time isolation.
|
| 215 |
+
"""
|
| 216 |
+
collection = self._get_collection()
|
| 217 |
+
if collection is None:
|
| 218 |
+
return # ChromaDB not configured — skip gracefully
|
| 219 |
+
|
| 220 |
+
embedding = self._embedding_fn([text])[0] if self._embedding_fn else None
|
| 221 |
+
|
| 222 |
+
collection.upsert(
|
| 223 |
+
ids=[doc_id],
|
| 224 |
+
documents=[text],
|
| 225 |
+
embeddings=[embedding] if embedding else None,
|
| 226 |
+
metadatas=[{
|
| 227 |
+
"tenant_id": tenant_id, # ← CRITICAL: must be present for isolation
|
| 228 |
+
"ticket_id": metadata.get("ticket_id", ""),
|
| 229 |
+
"category": metadata.get("category", ""),
|
| 230 |
+
"priority": metadata.get("priority", ""),
|
| 231 |
+
}],
|
| 232 |
+
)
|
| 233 |
+
|
| 234 |
+
def search(self, tenant_id: str, query: str, k: int = 10) -> list[RetrievedDoc]:
|
| 235 |
+
"""
|
| 236 |
+
Search ChromaDB with a mandatory tenant_id filter.
|
| 237 |
+
|
| 238 |
+
The `where={"tenant_id": tenant_id}` filter is applied at the
|
| 239 |
+
ChromaDB query engine level — not in Python after the fact.
|
| 240 |
+
This means ChromaDB will never even scan other tenants' embeddings.
|
| 241 |
+
"""
|
| 242 |
+
collection = self._get_collection()
|
| 243 |
+
if collection is None:
|
| 244 |
+
return []
|
| 245 |
+
|
| 246 |
+
query_embedding = self._embedding_fn([query])[0] if self._embedding_fn else None
|
| 247 |
+
|
| 248 |
+
try:
|
| 249 |
+
results = collection.query(
|
| 250 |
+
query_texts=[query] if query_embedding is None else None,
|
| 251 |
+
query_embeddings=[query_embedding] if query_embedding is not None else None,
|
| 252 |
+
n_results=min(k, collection.count()),
|
| 253 |
+
where={"tenant_id": tenant_id}, # ← TENANT ISOLATION ENFORCED HERE
|
| 254 |
+
include=["documents", "metadatas", "distances"],
|
| 255 |
+
)
|
| 256 |
+
except Exception as e:
|
| 257 |
+
logger.warning("ChromaDB query failed: %s", e)
|
| 258 |
+
return []
|
| 259 |
+
|
| 260 |
+
docs = []
|
| 261 |
+
for i, (doc_text, meta, dist) in enumerate(zip(
|
| 262 |
+
results["documents"][0],
|
| 263 |
+
results["metadatas"][0],
|
| 264 |
+
results["distances"][0],
|
| 265 |
+
)):
|
| 266 |
+
# Cosine distance → similarity score (1 = identical, 0 = orthogonal)
|
| 267 |
+
score = 1.0 - dist
|
| 268 |
+
docs.append(RetrievedDoc(
|
| 269 |
+
doc_id=results["ids"][0][i],
|
| 270 |
+
text=doc_text,
|
| 271 |
+
tenant_id=meta.get("tenant_id", tenant_id),
|
| 272 |
+
ticket_id=meta.get("ticket_id", ""),
|
| 273 |
+
category=meta.get("category", ""),
|
| 274 |
+
priority=meta.get("priority", ""),
|
| 275 |
+
score=score,
|
| 276 |
+
source="dense",
|
| 277 |
+
))
|
| 278 |
+
return docs
|
| 279 |
+
|
| 280 |
+
|
| 281 |
+
# ── Hybrid Retriever ───────────────────────────────────────────────────────────
|
| 282 |
+
|
| 283 |
+
class HybridRetriever:
|
| 284 |
+
"""
|
| 285 |
+
Tenant-isolated Hybrid Retrieval Engine.
|
| 286 |
+
|
| 287 |
+
Combines ChromaDB dense search with BM25 sparse search using
|
| 288 |
+
Reciprocal Rank Fusion, with optional cross-encoder reranking.
|
| 289 |
+
|
| 290 |
+
Usage:
|
| 291 |
+
retriever = HybridRetriever()
|
| 292 |
+
retriever.index_ticket(tenant_id="acme-corp", ticket_id="TKT-001",
|
| 293 |
+
text="API returning 500 errors on checkout",
|
| 294 |
+
metadata={"category": "technical", "priority": "high"})
|
| 295 |
+
|
| 296 |
+
results = retriever.search(
|
| 297 |
+
tenant_id="acme-corp",
|
| 298 |
+
query="checkout page broken 500 error",
|
| 299 |
+
k=5
|
| 300 |
+
)
|
| 301 |
+
"""
|
| 302 |
+
|
| 303 |
+
def __init__(
|
| 304 |
+
self,
|
| 305 |
+
chroma_client=None,
|
| 306 |
+
embedding_fn=None,
|
| 307 |
+
use_reranker: bool = False,
|
| 308 |
+
):
|
| 309 |
+
self.dense = TenantChromaRetriever(chroma_client, embedding_fn)
|
| 310 |
+
self.sparse = TenantBM25Index()
|
| 311 |
+
self.use_reranker = use_reranker
|
| 312 |
+
self._reranker = None
|
| 313 |
+
|
| 314 |
+
def _get_reranker(self):
|
| 315 |
+
"""Lazy-load cross-encoder reranker on first use."""
|
| 316 |
+
if self._reranker is None:
|
| 317 |
+
try:
|
| 318 |
+
from sentence_transformers import CrossEncoder
|
| 319 |
+
self._reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
|
| 320 |
+
except ImportError:
|
| 321 |
+
logger.warning("sentence-transformers not installed, reranking disabled")
|
| 322 |
+
return self._reranker
|
| 323 |
+
|
| 324 |
+
def index_ticket(
|
| 325 |
+
self,
|
| 326 |
+
tenant_id: str,
|
| 327 |
+
ticket_id: str,
|
| 328 |
+
text: str,
|
| 329 |
+
metadata: dict | None = None,
|
| 330 |
+
):
|
| 331 |
+
"""
|
| 332 |
+
Index a single support ticket into both dense and sparse indexes.
|
| 333 |
+
tenant_id is injected into ALL metadata — this is non-negotiable.
|
| 334 |
+
"""
|
| 335 |
+
metadata = metadata or {}
|
| 336 |
+
doc_id = hashlib.sha256(f"{tenant_id}:{ticket_id}".encode()).hexdigest()[:16]
|
| 337 |
+
full_meta = {**metadata, "ticket_id": ticket_id}
|
| 338 |
+
|
| 339 |
+
self.dense.add_document(tenant_id, doc_id, text, full_meta)
|
| 340 |
+
self.sparse.add_document(tenant_id, doc_id, text, full_meta)
|
| 341 |
+
|
| 342 |
+
def search(
|
| 343 |
+
self,
|
| 344 |
+
tenant_id: str,
|
| 345 |
+
query: str,
|
| 346 |
+
k: int = DEFAULT_FINAL_K,
|
| 347 |
+
dense_k: int = DEFAULT_DENSE_K,
|
| 348 |
+
sparse_k: int = DEFAULT_SPARSE_K,
|
| 349 |
+
) -> list[RetrievedDoc]:
|
| 350 |
+
"""
|
| 351 |
+
Perform tenant-isolated hybrid retrieval + optional reranking.
|
| 352 |
+
|
| 353 |
+
1. Dense search in ChromaDB (filtered to tenant_id)
|
| 354 |
+
2. Sparse BM25 search (scoped to tenant partition)
|
| 355 |
+
3. Reciprocal Rank Fusion to merge lists
|
| 356 |
+
4. Optional cross-encoder reranking on top-k
|
| 357 |
+
5. Return final top-k results
|
| 358 |
+
"""
|
| 359 |
+
t0 = time.time()
|
| 360 |
+
|
| 361 |
+
# Step 1 + 2: Parallel retrieval (both tenant-isolated)
|
| 362 |
+
dense_results = self.dense.search(tenant_id, query, k=dense_k)
|
| 363 |
+
sparse_results = self.sparse.search(tenant_id, query, k=sparse_k)
|
| 364 |
+
|
| 365 |
+
logger.debug(
|
| 366 |
+
"Retrieved: dense=%d sparse=%d for tenant=%s",
|
| 367 |
+
len(dense_results), len(sparse_results), tenant_id
|
| 368 |
+
)
|
| 369 |
+
|
| 370 |
+
# Step 3: Fuse
|
| 371 |
+
fused = reciprocal_rank_fusion(dense_results, sparse_results)[:k * 2]
|
| 372 |
+
|
| 373 |
+
# Step 4: Optional cross-encoder reranking
|
| 374 |
+
if self.use_reranker and fused:
|
| 375 |
+
reranker = self._get_reranker()
|
| 376 |
+
if reranker:
|
| 377 |
+
pairs = [[query, doc.text] for doc in fused]
|
| 378 |
+
scores = reranker.predict(pairs)
|
| 379 |
+
for doc, score in zip(fused, scores):
|
| 380 |
+
doc.score = float(score)
|
| 381 |
+
doc.source = "reranked"
|
| 382 |
+
fused = sorted(fused, key=lambda d: d.score, reverse=True)
|
| 383 |
+
|
| 384 |
+
elapsed = time.time() - t0
|
| 385 |
+
logger.debug("Hybrid retrieval completed in %.3fs", elapsed)
|
| 386 |
+
|
| 387 |
+
return fused[:k]
|
| 388 |
+
|
| 389 |
+
def doc_count(self, tenant_id: Optional[str] = None) -> int:
|
| 390 |
+
"""Return number of indexed documents (BM25 count as proxy)."""
|
| 391 |
+
return self.sparse.document_count(tenant_id)
|
| 392 |
+
|
| 393 |
+
|
| 394 |
+
# ── Module-level singleton ─────────────────────────────────────────────────────
|
| 395 |
+
_default_retriever: Optional[HybridRetriever] = None
|
| 396 |
+
|
| 397 |
+
|
| 398 |
+
def get_retriever() -> HybridRetriever:
|
| 399 |
+
"""Return the module-level singleton HybridRetriever."""
|
| 400 |
+
global _default_retriever
|
| 401 |
+
if _default_retriever is None:
|
| 402 |
+
_default_retriever = HybridRetriever()
|
| 403 |
+
return _default_retriever
|
| 404 |
+
|
| 405 |
+
|
| 406 |
+
# ── Standalone Demo ────────────────────────────────────────────────────────────
|
| 407 |
+
if __name__ == "__main__":
|
| 408 |
+
logging.basicConfig(level=logging.WARNING)
|
| 409 |
+
|
| 410 |
+
print("=" * 60)
|
| 411 |
+
print("CustomerCore Phase 6 — Hybrid Retriever Demo")
|
| 412 |
+
print("Multi-Tenant Vector DB Isolation Test")
|
| 413 |
+
print("=" * 60)
|
| 414 |
+
|
| 415 |
+
retriever = HybridRetriever()
|
| 416 |
+
|
| 417 |
+
# Index some tickets for two tenants
|
| 418 |
+
tenant_a_tickets = [
|
| 419 |
+
("TKT-A001", "Our checkout API is returning 500 errors on every request", {"category": "technical", "priority": "critical"}),
|
| 420 |
+
("TKT-A002", "Billing invoice shows wrong amount for March", {"category": "billing", "priority": "high"}),
|
| 421 |
+
("TKT-A003", "Login page is broken after your latest deploy", {"category": "technical", "priority": "high"}),
|
| 422 |
+
("TKT-A004", "Need to update our payment method on file", {"category": "billing", "priority": "medium"}),
|
| 423 |
+
("TKT-A005", "Our team cannot access the admin dashboard", {"category": "account", "priority": "high"}),
|
| 424 |
+
]
|
| 425 |
+
tenant_b_tickets = [
|
| 426 |
+
("TKT-B001", "CONFIDENTIAL: API latency spike on our EU cluster", {"category": "technical", "priority": "critical"}),
|
| 427 |
+
("TKT-B002", "CONFIDENTIAL: We were overcharged by $2400 last month", {"category": "billing", "priority": "high"}),
|
| 428 |
+
("TKT-B003", "CONFIDENTIAL: Data export is failing with timeout", {"category": "technical", "priority": "high"}),
|
| 429 |
+
]
|
| 430 |
+
|
| 431 |
+
print("\nIndexing 5 tickets for Tenant A (acme-corp)...")
|
| 432 |
+
for ticket_id, text, meta in tenant_a_tickets:
|
| 433 |
+
retriever.index_ticket("acme-corp", ticket_id, text, meta)
|
| 434 |
+
|
| 435 |
+
print("Indexing 3 CONFIDENTIAL tickets for Tenant B (globex-inc)...")
|
| 436 |
+
for ticket_id, text, meta in tenant_b_tickets:
|
| 437 |
+
retriever.index_ticket("globex-inc", ticket_id, text, meta)
|
| 438 |
+
|
| 439 |
+
print(f"\nTotal indexed: {retriever.doc_count():,} docs")
|
| 440 |
+
print(f" acme-corp: {retriever.doc_count('acme-corp')} docs")
|
| 441 |
+
print(f" globex-inc: {retriever.doc_count('globex-inc')} docs")
|
| 442 |
+
|
| 443 |
+
# Test 1: Tenant A searches — should only see Tenant A's tickets
|
| 444 |
+
print("\n" + "─" * 60)
|
| 445 |
+
print("TEST 1: Tenant A searches for 'billing issue API error'")
|
| 446 |
+
print("EXPECTED: Only acme-corp results. CONFIDENTIAL globex data must NOT appear.")
|
| 447 |
+
results_a = retriever.search("acme-corp", "billing issue API error", k=5)
|
| 448 |
+
all_correct = True
|
| 449 |
+
for doc in results_a:
|
| 450 |
+
icon = "✓" if doc.tenant_id == "acme-corp" else "✗ LEAK!"
|
| 451 |
+
print(f" {icon} [{doc.tenant_id}] {doc.ticket_id}: {doc.text[:60]}...")
|
| 452 |
+
if doc.tenant_id != "acme-corp":
|
| 453 |
+
all_correct = False
|
| 454 |
+
|
| 455 |
+
print(f"\n ISOLATION: {'PASSED ✓ — No cross-tenant leakage' if all_correct else 'FAILED ✗ — SECURITY BREACH!'}")
|
| 456 |
+
|
| 457 |
+
# Test 2: Tenant B searches — should only see Tenant B's tickets
|
| 458 |
+
print("\n" + "─" * 60)
|
| 459 |
+
print("TEST 2: Tenant B searches for 'API latency billing charge'")
|
| 460 |
+
print("EXPECTED: Only globex-inc CONFIDENTIAL results.")
|
| 461 |
+
results_b = retriever.search("globex-inc", "API latency billing charge", k=5)
|
| 462 |
+
all_correct_b = True
|
| 463 |
+
for doc in results_b:
|
| 464 |
+
icon = "✓" if doc.tenant_id == "globex-inc" else "✗ LEAK!"
|
| 465 |
+
print(f" {icon} [{doc.tenant_id}] {doc.ticket_id}: {doc.text[:60]}...")
|
| 466 |
+
if doc.tenant_id != "globex-inc":
|
| 467 |
+
all_correct_b = False
|
| 468 |
+
|
| 469 |
+
print(f"\n ISOLATION: {'PASSED ✓ — No cross-tenant leakage' if all_correct_b else 'FAILED ✗ — SECURITY BREACH!'}")
|
| 470 |
+
print("\n" + "=" * 60)
|
|
@@ -0,0 +1,263 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/llm_client.py
|
| 3 |
+
|
| 4 |
+
Unified LiteLLM client wrapper for CustomerCore.
|
| 5 |
+
|
| 6 |
+
Provides a single interface for calling both:
|
| 7 |
+
- LOCAL models via Ollama (gemma3:4b, gemma2:2b) — $0, sub-200ms on GPU
|
| 8 |
+
- CLOUD models via OpenRouter (Claude 3.5 Sonnet, GPT-4o, Gemini 1.5 Pro)
|
| 9 |
+
|
| 10 |
+
== Why LiteLLM? ==
|
| 11 |
+
LiteLLM gives us a unified OpenAI-compatible interface for 100+ model providers.
|
| 12 |
+
We can swap Ollama for vLLM, or swap OpenRouter for direct Anthropic, without
|
| 13 |
+
changing any calling code — only the model string and API key change.
|
| 14 |
+
|
| 15 |
+
== Production Config ==
|
| 16 |
+
All API keys and endpoints are injected via Doppler at runtime.
|
| 17 |
+
NEVER hardcode keys. The client reads from environment variables only.
|
| 18 |
+
|
| 19 |
+
Environment variables:
|
| 20 |
+
OPENROUTER_API_KEY — for cloud model routing
|
| 21 |
+
OLLAMA_BASE_URL — defaults to http://localhost:11434
|
| 22 |
+
LLM_DEFAULT_LOCAL — which local model to use (default: ollama/gemma3:4b)
|
| 23 |
+
LLM_DEFAULT_CLOUD — which cloud model to use (default: openrouter/anthropic/claude-3.5-sonnet)
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
import os
|
| 27 |
+
import time
|
| 28 |
+
import logging
|
| 29 |
+
from dataclasses import dataclass, field
|
| 30 |
+
from typing import Optional, Any
|
| 31 |
+
|
| 32 |
+
try:
|
| 33 |
+
from dotenv import load_dotenv
|
| 34 |
+
load_dotenv()
|
| 35 |
+
except ImportError:
|
| 36 |
+
pass
|
| 37 |
+
|
| 38 |
+
logger = logging.getLogger(__name__)
|
| 39 |
+
|
| 40 |
+
# ── Model Identifiers ──────────────────────────────────────────────────────────
|
| 41 |
+
# LiteLLM uses provider/model_name format
|
| 42 |
+
LOCAL_MODEL = os.environ.get("LLM_DEFAULT_LOCAL", "ollama/gemma3:4b")
|
| 43 |
+
CLOUD_MODEL = os.environ.get(
|
| 44 |
+
"LLM_DEFAULT_CLOUD",
|
| 45 |
+
# Verified working free OpenRouter models (benchmarked 2026-05-22):
|
| 46 |
+
# openrouter/meta-llama/llama-3.1-8b-instruct — 1.1s, excellent quality
|
| 47 |
+
# openrouter/qwen/qwen3-8b — 5.3s, great reasoning
|
| 48 |
+
"openrouter/meta-llama/llama-3.1-8b-instruct"
|
| 49 |
+
)
|
| 50 |
+
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "http://localhost:11434")
|
| 51 |
+
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
@dataclass
|
| 55 |
+
class LLMResponse:
|
| 56 |
+
"""Structured response from any LLM call."""
|
| 57 |
+
content: str
|
| 58 |
+
model_used: str
|
| 59 |
+
provider: str # "local" | "cloud"
|
| 60 |
+
latency_ms: float
|
| 61 |
+
input_tokens: int = 0
|
| 62 |
+
output_tokens: int = 0
|
| 63 |
+
estimated_cost_usd: float = 0.0
|
| 64 |
+
success: bool = True
|
| 65 |
+
error: Optional[str] = None
|
| 66 |
+
raw: Optional[Any] = None
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
# ── Cost table (approximate, USD per 1k tokens) ────────────────────────────────
|
| 70 |
+
# Source: OpenRouter pricing — free models marked $0.000
|
| 71 |
+
_COST_PER_1K_INPUT = {
|
| 72 |
+
"ollama/gemma3:4b": 0.0,
|
| 73 |
+
"ollama/gemma2:2b": 0.0,
|
| 74 |
+
"ollama/gemma4:e4b": 0.0,
|
| 75 |
+
# OpenRouter FREE models (free tier — no cost)
|
| 76 |
+
"openrouter/meta-llama/llama-3.1-8b-instruct:free": 0.0,
|
| 77 |
+
"openrouter/google/gemma-2-9b-it:free": 0.0,
|
| 78 |
+
"openrouter/mistralai/mistral-7b-instruct:free": 0.0,
|
| 79 |
+
"openrouter/microsoft/phi-3-mini-128k-instruct:free": 0.0,
|
| 80 |
+
# OpenRouter PAID models (for when budget allows)
|
| 81 |
+
"openrouter/anthropic/claude-3.5-sonnet": 0.003,
|
| 82 |
+
"openrouter/openai/gpt-4o": 0.005,
|
| 83 |
+
"openrouter/google/gemini-1.5-pro": 0.00125,
|
| 84 |
+
}
|
| 85 |
+
_COST_PER_1K_OUTPUT = {
|
| 86 |
+
"ollama/gemma3:4b": 0.0,
|
| 87 |
+
"ollama/gemma2:2b": 0.0,
|
| 88 |
+
"ollama/gemma4:e4b": 0.0,
|
| 89 |
+
# OpenRouter FREE models
|
| 90 |
+
"openrouter/meta-llama/llama-3.1-8b-instruct:free": 0.0,
|
| 91 |
+
"openrouter/google/gemma-2-9b-it:free": 0.0,
|
| 92 |
+
"openrouter/mistralai/mistral-7b-instruct:free": 0.0,
|
| 93 |
+
"openrouter/microsoft/phi-3-mini-128k-instruct:free": 0.0,
|
| 94 |
+
# OpenRouter PAID models
|
| 95 |
+
"openrouter/anthropic/claude-3.5-sonnet": 0.015,
|
| 96 |
+
"openrouter/openai/gpt-4o": 0.015,
|
| 97 |
+
"openrouter/google/gemini-1.5-pro": 0.005,
|
| 98 |
+
}
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def _estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
|
| 102 |
+
cost_in = _COST_PER_1K_INPUT.get(model, 0.0) * input_tokens / 1000
|
| 103 |
+
cost_out = _COST_PER_1K_OUTPUT.get(model, 0.0) * output_tokens / 1000
|
| 104 |
+
return round(cost_in + cost_out, 6)
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class LLMClient:
|
| 108 |
+
"""
|
| 109 |
+
Unified LiteLLM wrapper for local and cloud model calls.
|
| 110 |
+
|
| 111 |
+
Usage:
|
| 112 |
+
client = LLMClient()
|
| 113 |
+
|
| 114 |
+
# Local (free, fast, private)
|
| 115 |
+
response = client.call_local(
|
| 116 |
+
messages=[{"role": "user", "content": "Classify this ticket: API 500 error"}],
|
| 117 |
+
temperature=0.1,
|
| 118 |
+
)
|
| 119 |
+
|
| 120 |
+
# Cloud (frontier reasoning, higher cost)
|
| 121 |
+
response = client.call_cloud(
|
| 122 |
+
messages=[{"role": "user", "content": "Should we issue this $200 refund? Reason carefully."}],
|
| 123 |
+
temperature=0.2,
|
| 124 |
+
)
|
| 125 |
+
"""
|
| 126 |
+
|
| 127 |
+
def __init__(
|
| 128 |
+
self,
|
| 129 |
+
local_model: str = LOCAL_MODEL,
|
| 130 |
+
cloud_model: str = CLOUD_MODEL,
|
| 131 |
+
):
|
| 132 |
+
self.local_model = local_model
|
| 133 |
+
self.cloud_model = cloud_model
|
| 134 |
+
|
| 135 |
+
def _call(
|
| 136 |
+
self,
|
| 137 |
+
model: str,
|
| 138 |
+
messages: list[dict],
|
| 139 |
+
provider: str,
|
| 140 |
+
temperature: float = 0.1,
|
| 141 |
+
max_tokens: int = 512,
|
| 142 |
+
timeout: float = 30.0,
|
| 143 |
+
) -> LLMResponse:
|
| 144 |
+
"""Internal call — handles timing, error wrapping, cost estimation."""
|
| 145 |
+
try:
|
| 146 |
+
import litellm
|
| 147 |
+
except ImportError:
|
| 148 |
+
return LLMResponse(
|
| 149 |
+
content="",
|
| 150 |
+
model_used=model,
|
| 151 |
+
provider=provider,
|
| 152 |
+
latency_ms=0.0,
|
| 153 |
+
success=False,
|
| 154 |
+
error="litellm not installed",
|
| 155 |
+
)
|
| 156 |
+
|
| 157 |
+
# Configure Ollama base URL for local calls
|
| 158 |
+
api_base = OLLAMA_BASE_URL if provider == "local" else None
|
| 159 |
+
api_key = OPENROUTER_API_KEY if provider == "cloud" else None
|
| 160 |
+
|
| 161 |
+
t0 = time.perf_counter()
|
| 162 |
+
try:
|
| 163 |
+
kwargs = dict(
|
| 164 |
+
model=model,
|
| 165 |
+
messages=messages,
|
| 166 |
+
temperature=temperature,
|
| 167 |
+
max_tokens=max_tokens,
|
| 168 |
+
timeout=timeout,
|
| 169 |
+
)
|
| 170 |
+
if api_base:
|
| 171 |
+
kwargs["api_base"] = api_base
|
| 172 |
+
if api_key:
|
| 173 |
+
kwargs["api_key"] = api_key
|
| 174 |
+
|
| 175 |
+
# Add OpenRouter required headers
|
| 176 |
+
if provider == "cloud":
|
| 177 |
+
kwargs["extra_headers"] = {
|
| 178 |
+
"HTTP-Referer": "https://github.com/CustomerCore",
|
| 179 |
+
"X-Title": "CustomerCore B2B Intelligence Platform",
|
| 180 |
+
}
|
| 181 |
+
|
| 182 |
+
response = litellm.completion(**kwargs)
|
| 183 |
+
latency_ms = (time.perf_counter() - t0) * 1000
|
| 184 |
+
|
| 185 |
+
content = response.choices[0].message.content or ""
|
| 186 |
+
usage = response.usage or {}
|
| 187 |
+
input_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
| 188 |
+
output_tokens = getattr(usage, "completion_tokens", 0) or 0
|
| 189 |
+
|
| 190 |
+
return LLMResponse(
|
| 191 |
+
content=content,
|
| 192 |
+
model_used=model,
|
| 193 |
+
provider=provider,
|
| 194 |
+
latency_ms=round(latency_ms, 1),
|
| 195 |
+
input_tokens=input_tokens,
|
| 196 |
+
output_tokens=output_tokens,
|
| 197 |
+
estimated_cost_usd=_estimate_cost(model, input_tokens, output_tokens),
|
| 198 |
+
success=True,
|
| 199 |
+
raw=response,
|
| 200 |
+
)
|
| 201 |
+
|
| 202 |
+
except Exception as e:
|
| 203 |
+
latency_ms = (time.perf_counter() - t0) * 1000
|
| 204 |
+
logger.error("LLM call failed: model=%s error=%s", model, e)
|
| 205 |
+
return LLMResponse(
|
| 206 |
+
content="",
|
| 207 |
+
model_used=model,
|
| 208 |
+
provider=provider,
|
| 209 |
+
latency_ms=round(latency_ms, 1),
|
| 210 |
+
success=False,
|
| 211 |
+
error=str(e),
|
| 212 |
+
)
|
| 213 |
+
|
| 214 |
+
def call_local(
|
| 215 |
+
self,
|
| 216 |
+
messages: list[dict],
|
| 217 |
+
temperature: float = 0.1,
|
| 218 |
+
max_tokens: int = 512,
|
| 219 |
+
timeout: float = 30.0,
|
| 220 |
+
) -> LLMResponse:
|
| 221 |
+
"""Call the local Ollama model. Zero cost. Data never leaves the machine."""
|
| 222 |
+
return self._call(
|
| 223 |
+
model=self.local_model,
|
| 224 |
+
messages=messages,
|
| 225 |
+
provider="local",
|
| 226 |
+
temperature=temperature,
|
| 227 |
+
max_tokens=max_tokens,
|
| 228 |
+
timeout=timeout,
|
| 229 |
+
)
|
| 230 |
+
|
| 231 |
+
def call_cloud(
|
| 232 |
+
self,
|
| 233 |
+
messages: list[dict],
|
| 234 |
+
temperature: float = 0.2,
|
| 235 |
+
max_tokens: int = 1024,
|
| 236 |
+
timeout: float = 60.0,
|
| 237 |
+
) -> LLMResponse:
|
| 238 |
+
"""Call the cloud frontier model via OpenRouter. Higher cost, higher capability."""
|
| 239 |
+
if not OPENROUTER_API_KEY:
|
| 240 |
+
logger.warning(
|
| 241 |
+
"OPENROUTER_API_KEY not set — falling back to local model for cloud call"
|
| 242 |
+
)
|
| 243 |
+
return self.call_local(messages, temperature, max_tokens, timeout)
|
| 244 |
+
|
| 245 |
+
return self._call(
|
| 246 |
+
model=self.cloud_model,
|
| 247 |
+
messages=messages,
|
| 248 |
+
provider="cloud",
|
| 249 |
+
temperature=temperature,
|
| 250 |
+
max_tokens=max_tokens,
|
| 251 |
+
timeout=timeout,
|
| 252 |
+
)
|
| 253 |
+
|
| 254 |
+
|
| 255 |
+
# ── Module-level singleton ─────────────────────────────────────────────────────
|
| 256 |
+
_default_client: Optional[LLMClient] = None
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
def get_client() -> LLMClient:
|
| 260 |
+
global _default_client
|
| 261 |
+
if _default_client is None:
|
| 262 |
+
_default_client = LLMClient()
|
| 263 |
+
return _default_client
|
|
@@ -0,0 +1,208 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/multilingual.py
|
| 3 |
+
|
| 4 |
+
Phase 8a: Multi-Language Intelligence Layer
|
| 5 |
+
|
| 6 |
+
== What This Does ==
|
| 7 |
+
Every ticket entering CustomerCore now gets:
|
| 8 |
+
1. Language detection (langdetect — probabilistic, fast, 55 languages)
|
| 9 |
+
2. Language-aware routing (multilingual-capable models for non-English)
|
| 10 |
+
3. Language metadata stored on every Silver record for dbt Gold analysis
|
| 11 |
+
4. Multilingual BM25 tokenization (language-specific stopwords via NLTK)
|
| 12 |
+
|
| 13 |
+
== Supported Languages (V1) ==
|
| 14 |
+
en — English (primary — Bitext SaaS + Banking datasets)
|
| 15 |
+
de — German (important: you are in Germany, B2B EU platform)
|
| 16 |
+
fr — French (major EU business language)
|
| 17 |
+
es — Spanish (3rd largest language by speakers)
|
| 18 |
+
pt — Portuguese (Brazil = large B2B market)
|
| 19 |
+
nl — Dutch
|
| 20 |
+
it — Italian
|
| 21 |
+
Unknown → fallback to multilingual model
|
| 22 |
+
|
| 23 |
+
== Why Language Detection Matters for a B2B Platform ==
|
| 24 |
+
Real enterprise B2B SaaS platforms serve global customers. A German Mittelstand
|
| 25 |
+
company will write support tickets in German. A French SaaS company will
|
| 26 |
+
write in French. Routing all tickets through an English-only model:
|
| 27 |
+
(a) Degrades classification accuracy by 15-40% for non-English tickets
|
| 28 |
+
(b) Misses language as a segmentation signal in analytics
|
| 29 |
+
(c) Is not acceptable for an EU-compliance platform (GDPR requires same
|
| 30 |
+
quality of service regardless of language used)
|
| 31 |
+
|
| 32 |
+
== Dataset Sourcing ==
|
| 33 |
+
mteb/amazon_massive_intent — 11,514 rows per language, Apache 2.0
|
| 34 |
+
Languages: en, de, fr, es (confirmed working, no scraping, open license)
|
| 35 |
+
These are customer intent utterances — directly maps to support ticket classification.
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
import logging
|
| 39 |
+
import re
|
| 40 |
+
from typing import Optional
|
| 41 |
+
|
| 42 |
+
logger = logging.getLogger(__name__)
|
| 43 |
+
|
| 44 |
+
# ── Supported language codes ───────────────────────────────────────────────────
|
| 45 |
+
SUPPORTED_LANGUAGES = {
|
| 46 |
+
"en": "English",
|
| 47 |
+
"de": "German",
|
| 48 |
+
"fr": "French",
|
| 49 |
+
"es": "Spanish",
|
| 50 |
+
"pt": "Portuguese",
|
| 51 |
+
"nl": "Dutch",
|
| 52 |
+
"it": "Italian",
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
# ── Languages that need multilingual model routing ─────────────────────────────
|
| 56 |
+
# English-only models (gemma3:4b) perform fine on English.
|
| 57 |
+
# For other languages we prefer multilingual models.
|
| 58 |
+
MULTILINGUAL_LANGUAGES = {"de", "fr", "es", "pt", "nl", "it"}
|
| 59 |
+
|
| 60 |
+
# ── Simple stopword sets per language (for BM25 tokenization) ─────────────────
|
| 61 |
+
# Minimal sets — good enough for BM25 without NLTK dependency
|
| 62 |
+
_STOPWORDS: dict[str, set[str]] = {
|
| 63 |
+
"en": {"the", "a", "an", "is", "it", "in", "on", "at", "to", "for",
|
| 64 |
+
"of", "and", "or", "but", "my", "i", "we", "you", "our", "your"},
|
| 65 |
+
"de": {"der", "die", "das", "ein", "eine", "ist", "ich", "wir", "sie",
|
| 66 |
+
"und", "oder", "aber", "für", "mit", "von", "zu", "an", "auf"},
|
| 67 |
+
"fr": {"le", "la", "les", "un", "une", "est", "je", "nous", "vous",
|
| 68 |
+
"et", "ou", "mais", "pour", "avec", "de", "du", "au", "en"},
|
| 69 |
+
"es": {"el", "la", "los", "un", "una", "es", "yo", "nosotros", "ustedes",
|
| 70 |
+
"y", "o", "pero", "para", "con", "de", "del", "al", "en"},
|
| 71 |
+
"pt": {"o", "a", "os", "as", "um", "uma", "é", "eu", "nós", "você",
|
| 72 |
+
"e", "ou", "mas", "para", "com", "de", "do", "ao", "em"},
|
| 73 |
+
"nl": {"de", "het", "een", "is", "ik", "wij", "u", "en", "of", "maar",
|
| 74 |
+
"voor", "met", "van", "aan", "op", "in", "te"},
|
| 75 |
+
"it": {"il", "la", "i", "le", "un", "una", "è", "io", "noi", "voi",
|
| 76 |
+
"e", "o", "ma", "per", "con", "di", "del", "al", "in"},
|
| 77 |
+
}
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def detect_language(text: str, min_text_length: int = 20) -> str:
|
| 81 |
+
"""
|
| 82 |
+
Detect the language of a text string.
|
| 83 |
+
|
| 84 |
+
Returns an ISO 639-1 language code (e.g. 'en', 'de', 'fr').
|
| 85 |
+
Falls back to 'en' if:
|
| 86 |
+
- Text is too short for reliable detection
|
| 87 |
+
- langdetect is not installed
|
| 88 |
+
- Detection fails (mixed language, gibberish, etc.)
|
| 89 |
+
|
| 90 |
+
Uses langdetect under the hood (probabilistic Naive Bayes).
|
| 91 |
+
"""
|
| 92 |
+
text = (text or "").strip()
|
| 93 |
+
if len(text) < min_text_length:
|
| 94 |
+
return "en" # too short for reliable detection
|
| 95 |
+
|
| 96 |
+
try:
|
| 97 |
+
from langdetect import detect, LangDetectException
|
| 98 |
+
lang = detect(text)
|
| 99 |
+
# langdetect returns e.g. "zh-cn" — normalize to first part
|
| 100 |
+
lang = lang.split("-")[0].lower()
|
| 101 |
+
return lang
|
| 102 |
+
except ImportError:
|
| 103 |
+
logger.debug("langdetect not installed — returning 'en' as fallback")
|
| 104 |
+
return "en"
|
| 105 |
+
except Exception:
|
| 106 |
+
return "en"
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def detect_language_with_confidence(text: str) -> tuple[str, float]:
|
| 110 |
+
"""
|
| 111 |
+
Detect language with confidence score.
|
| 112 |
+
Returns (language_code, probability 0.0-1.0).
|
| 113 |
+
"""
|
| 114 |
+
text = (text or "").strip()
|
| 115 |
+
if len(text) < 20:
|
| 116 |
+
return "en", 1.0
|
| 117 |
+
|
| 118 |
+
try:
|
| 119 |
+
from langdetect import detect_langs, LangDetectException
|
| 120 |
+
results = detect_langs(text)
|
| 121 |
+
if results:
|
| 122 |
+
top = results[0]
|
| 123 |
+
lang = str(top.lang).split("-")[0].lower()
|
| 124 |
+
prob = float(top.prob)
|
| 125 |
+
return lang, round(prob, 3)
|
| 126 |
+
return "en", 1.0
|
| 127 |
+
except ImportError:
|
| 128 |
+
return "en", 1.0
|
| 129 |
+
except Exception:
|
| 130 |
+
return "en", 0.5
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def get_stopwords(lang: str) -> set[str]:
|
| 134 |
+
"""Return stopword set for a language. Falls back to English."""
|
| 135 |
+
return _STOPWORDS.get(lang, _STOPWORDS["en"])
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def tokenize_multilingual(text: str, lang: str) -> list[str]:
|
| 139 |
+
"""
|
| 140 |
+
Language-aware tokenization for BM25.
|
| 141 |
+
Lowercases, removes punctuation, filters stopwords.
|
| 142 |
+
"""
|
| 143 |
+
# Simple unicode-friendly tokenization
|
| 144 |
+
tokens = re.findall(r"\b\w+\b", text.lower(), re.UNICODE)
|
| 145 |
+
stopwords = get_stopwords(lang)
|
| 146 |
+
return [t for t in tokens if t not in stopwords and len(t) > 1]
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def needs_multilingual_model(lang: str) -> bool:
|
| 150 |
+
"""Return True if this language needs a multilingual model (not English-only)."""
|
| 151 |
+
return lang in MULTILINGUAL_LANGUAGES
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def get_language_display(lang: str) -> str:
|
| 155 |
+
"""Return human-readable language name."""
|
| 156 |
+
return SUPPORTED_LANGUAGES.get(lang, f"Unknown ({lang})")
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def enrich_with_language(record: dict) -> dict:
|
| 160 |
+
"""
|
| 161 |
+
Add language detection fields to a Silver record.
|
| 162 |
+
Mutates the record in place and returns it.
|
| 163 |
+
|
| 164 |
+
Added fields:
|
| 165 |
+
detected_language — ISO 639-1 code (e.g. 'de')
|
| 166 |
+
language_confidence — 0.0-1.0 detection confidence
|
| 167 |
+
language_display — human readable (e.g. 'German')
|
| 168 |
+
is_multilingual — True if non-English
|
| 169 |
+
"""
|
| 170 |
+
text = record.get("body", "") or record.get("subject", "") or ""
|
| 171 |
+
lang, confidence = detect_language_with_confidence(text)
|
| 172 |
+
|
| 173 |
+
record["detected_language"] = lang
|
| 174 |
+
record["language_confidence"] = confidence
|
| 175 |
+
record["language_display"] = get_language_display(lang)
|
| 176 |
+
record["is_multilingual"] = needs_multilingual_model(lang)
|
| 177 |
+
return record
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
# ── Standalone demo ────────────────────────────────────────────────────────────
|
| 181 |
+
if __name__ == "__main__":
|
| 182 |
+
test_cases = [
|
| 183 |
+
("en", "My payment failed and I cannot access my account. Please help urgently."),
|
| 184 |
+
("de", "Meine Zahlung ist fehlgeschlagen und ich kann nicht auf mein Konto zugreifen."),
|
| 185 |
+
("fr", "Mon paiement a échoué et je ne peux pas accéder à mon compte. Aidez-moi."),
|
| 186 |
+
("es", "Mi pago falló y no puedo acceder a mi cuenta. Por favor ayúdame urgentemente."),
|
| 187 |
+
("pt", "Meu pagamento falhou e não consigo acessar minha conta. Por favor, ajude."),
|
| 188 |
+
("nl", "Mijn betaling is mislukt en ik kan geen toegang krijgen tot mijn account."),
|
| 189 |
+
("it", "Il mio pagamento è fallito e non riesco ad accedere al mio account. Aiutami."),
|
| 190 |
+
]
|
| 191 |
+
|
| 192 |
+
print("=" * 65)
|
| 193 |
+
print("CustomerCore Phase 8a — Language Detection Demo")
|
| 194 |
+
print("=" * 65)
|
| 195 |
+
print(f"\n{'Expected':<8} {'Detected':<8} {'Conf':<7} {'Multilingual':<14} Text")
|
| 196 |
+
print("─" * 70)
|
| 197 |
+
|
| 198 |
+
all_correct = True
|
| 199 |
+
for expected, text in test_cases:
|
| 200 |
+
lang, conf = detect_language_with_confidence(text)
|
| 201 |
+
is_ml = needs_multilingual_model(lang)
|
| 202 |
+
ok = "✓" if lang == expected else "✗"
|
| 203 |
+
if lang != expected:
|
| 204 |
+
all_correct = False
|
| 205 |
+
print(f"{ok} {expected:<6} {lang:<8} {conf:<7.2f} {str(is_ml):<14} {text[:40]}...")
|
| 206 |
+
|
| 207 |
+
print(f"\n{'All detections correct ✓' if all_correct else 'Some detections wrong ✗'}")
|
| 208 |
+
print("=" * 65)
|
|
@@ -0,0 +1,377 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/rag/router.py
|
| 3 |
+
|
| 4 |
+
Phase 7: SLA-Aware Multi-Model LLM Router
|
| 5 |
+
|
| 6 |
+
== The Gap We Are Closing ==
|
| 7 |
+
The previous setup had a manual local/cloud toggle — someone had to explicitly choose
|
| 8 |
+
which model to call. Real enterprise B2B AI platforms (Sierra AI, Decagon) route
|
| 9 |
+
automatically based on task complexity, SLA tier, and cost budget.
|
| 10 |
+
|
| 11 |
+
== Routing Logic ==
|
| 12 |
+
Every ticket that enters the system has a priority (low, medium, high, critical)
|
| 13 |
+
and a task type (classification, extraction, reasoning, action):
|
| 14 |
+
|
| 15 |
+
Task Type Priority → Route To
|
| 16 |
+
──────────── ──────────── ─────────────────────────────────────────────────
|
| 17 |
+
classify low/medium → LOCAL (gemma3:4b via Ollama, ~150ms, $0)
|
| 18 |
+
classify high/critical → LOCAL (still fast enough, no sensitive action)
|
| 19 |
+
extract low/medium → LOCAL (structured info extraction, deterministic)
|
| 20 |
+
extract high/critical → LOCAL (same — no judgment call required)
|
| 21 |
+
reason low/medium → LOCAL (reasoning on routine tickets)
|
| 22 |
+
reason high/critical → CLOUD (frontier model for complex edge cases)
|
| 23 |
+
action any → CLOUD (anything that takes an external action,
|
| 24 |
+
e.g. issue refund, create Jira ticket,
|
| 25 |
+
escalate to VIP team — requires best reasoning)
|
| 26 |
+
|
| 27 |
+
== Why Not Always Use Cloud? ==
|
| 28 |
+
- 80%+ of B2B support tickets are routine (password resets, invoice requests)
|
| 29 |
+
- Local Gemma 3 4B handles these in ~150ms at $0 cost
|
| 30 |
+
- Cloud call costs ~$0.01-0.05 per ticket at scale = thousands per month
|
| 31 |
+
- SLA: low-priority tickets don't need frontier-model reasoning
|
| 32 |
+
- Privacy: local model means sensitive customer data never leaves the network
|
| 33 |
+
|
| 34 |
+
== Metrics ==
|
| 35 |
+
The router tracks every call: latency, model, cost, task type, priority.
|
| 36 |
+
These feed into Prometheus metrics for the Grafana dashboard.
|
| 37 |
+
|
| 38 |
+
Run standalone demo:
|
| 39 |
+
python -m src.rag.router
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
import logging
|
| 43 |
+
import os
|
| 44 |
+
import time
|
| 45 |
+
from dataclasses import dataclass, field
|
| 46 |
+
from typing import Literal, Optional
|
| 47 |
+
from collections import defaultdict
|
| 48 |
+
|
| 49 |
+
from src.rag.llm_client import LLMClient, LLMResponse, get_client
|
| 50 |
+
|
| 51 |
+
logger = logging.getLogger(__name__)
|
| 52 |
+
|
| 53 |
+
# ── Types ──────────────────────────────────────────────────────────────────────
|
| 54 |
+
Priority = Literal["low", "medium", "high", "critical"]
|
| 55 |
+
TaskType = Literal["classify", "extract", "reason", "action"]
|
| 56 |
+
ModelTier = Literal["local", "cloud"]
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
# ── Routing Table ──────────────────────────────────────────────────────────────
|
| 60 |
+
# (task_type, priority) → model tier
|
| 61 |
+
ROUTING_TABLE: dict[tuple[str, str], ModelTier] = {
|
| 62 |
+
# Classification — always local (fast, cheap, sufficient)
|
| 63 |
+
("classify", "low"): "local",
|
| 64 |
+
("classify", "medium"): "local",
|
| 65 |
+
("classify", "high"): "local",
|
| 66 |
+
("classify", "critical"): "local",
|
| 67 |
+
|
| 68 |
+
# Extraction — always local (deterministic structured output)
|
| 69 |
+
("extract", "low"): "local",
|
| 70 |
+
("extract", "medium"): "local",
|
| 71 |
+
("extract", "high"): "local",
|
| 72 |
+
("extract", "critical"): "local",
|
| 73 |
+
|
| 74 |
+
# Reasoning — local for routine, cloud for high-stakes
|
| 75 |
+
("reason", "low"): "local",
|
| 76 |
+
("reason", "medium"): "local",
|
| 77 |
+
("reason", "high"): "cloud", # ← switch to frontier model
|
| 78 |
+
("reason", "critical"): "cloud", # ← frontier model required
|
| 79 |
+
|
| 80 |
+
# Action — always cloud (external side effects require best judgment)
|
| 81 |
+
("action", "low"): "cloud",
|
| 82 |
+
("action", "medium"): "cloud",
|
| 83 |
+
("action", "high"): "cloud",
|
| 84 |
+
("action", "critical"): "cloud",
|
| 85 |
+
}
|
| 86 |
+
|
| 87 |
+
# SLA latency targets (ms) — used for logging/alerting
|
| 88 |
+
SLA_TARGETS_MS: dict[Priority, int] = {
|
| 89 |
+
"low": 2000, # 2 seconds acceptable
|
| 90 |
+
"medium": 1000, # 1 second
|
| 91 |
+
"high": 500, # 500ms
|
| 92 |
+
"critical": 200, # 200ms — only local can reliably hit this
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
@dataclass
|
| 97 |
+
class RouterDecision:
|
| 98 |
+
"""The routing decision made for a single request."""
|
| 99 |
+
task_type: str
|
| 100 |
+
priority: str
|
| 101 |
+
model_tier: ModelTier
|
| 102 |
+
model_name: str
|
| 103 |
+
sla_target_ms: int
|
| 104 |
+
reasoning: str
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
@dataclass
|
| 108 |
+
class RouterMetrics:
|
| 109 |
+
"""Aggregated metrics across all router calls (for Prometheus/Grafana)."""
|
| 110 |
+
total_calls: int = 0
|
| 111 |
+
local_calls: int = 0
|
| 112 |
+
cloud_calls: int = 0
|
| 113 |
+
total_cost_usd: float = 0.0
|
| 114 |
+
total_latency_ms: float = 0.0
|
| 115 |
+
sla_violations: int = 0
|
| 116 |
+
errors: int = 0
|
| 117 |
+
calls_by_task: dict = field(default_factory=lambda: defaultdict(int))
|
| 118 |
+
calls_by_priority: dict = field(default_factory=lambda: defaultdict(int))
|
| 119 |
+
|
| 120 |
+
@property
|
| 121 |
+
def avg_latency_ms(self) -> float:
|
| 122 |
+
return self.total_latency_ms / self.total_calls if self.total_calls else 0.0
|
| 123 |
+
|
| 124 |
+
@property
|
| 125 |
+
def local_pct(self) -> float:
|
| 126 |
+
return (self.local_calls / self.total_calls * 100) if self.total_calls else 0.0
|
| 127 |
+
|
| 128 |
+
@property
|
| 129 |
+
def cloud_pct(self) -> float:
|
| 130 |
+
return (self.cloud_calls / self.total_calls * 100) if self.total_calls else 0.0
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
class LLMRouter:
|
| 134 |
+
"""
|
| 135 |
+
SLA-Aware Multi-Model LLM Router.
|
| 136 |
+
|
| 137 |
+
Automatically selects local or cloud model based on task type and ticket
|
| 138 |
+
priority. Tracks latency, cost, and SLA compliance for every call.
|
| 139 |
+
|
| 140 |
+
Usage:
|
| 141 |
+
router = LLMRouter()
|
| 142 |
+
|
| 143 |
+
response, decision = router.route(
|
| 144 |
+
messages=[{"role": "user", "content": "Classify this ticket: ..."}],
|
| 145 |
+
task_type="classify",
|
| 146 |
+
priority="medium",
|
| 147 |
+
)
|
| 148 |
+
print(f"Used: {decision.model_name} ({decision.model_tier})")
|
| 149 |
+
print(f"Latency: {response.latency_ms}ms")
|
| 150 |
+
print(f"Cost: ${response.estimated_cost_usd:.4f}")
|
| 151 |
+
"""
|
| 152 |
+
|
| 153 |
+
def __init__(self, client: Optional[LLMClient] = None):
|
| 154 |
+
self.client = client or get_client()
|
| 155 |
+
self.metrics = RouterMetrics()
|
| 156 |
+
self._call_log: list[dict] = []
|
| 157 |
+
|
| 158 |
+
def _decide(self, task_type: str, priority: str) -> RouterDecision:
|
| 159 |
+
"""
|
| 160 |
+
Determine which model tier to use based on routing table.
|
| 161 |
+
Falls back to local if the combination is not in the table.
|
| 162 |
+
"""
|
| 163 |
+
tier = ROUTING_TABLE.get((task_type, priority), "local")
|
| 164 |
+
model_name = self.client.local_model if tier == "local" else self.client.cloud_model
|
| 165 |
+
sla_ms = SLA_TARGETS_MS.get(priority, 2000)
|
| 166 |
+
|
| 167 |
+
# Human-readable explanation for audit/debugging
|
| 168 |
+
if tier == "local":
|
| 169 |
+
reasoning = (
|
| 170 |
+
f"Task '{task_type}' at priority '{priority}' routed to LOCAL model "
|
| 171 |
+
f"({self.client.local_model}) — fast, free, data stays on-premise."
|
| 172 |
+
)
|
| 173 |
+
else:
|
| 174 |
+
reasoning = (
|
| 175 |
+
f"Task '{task_type}' at priority '{priority}' routed to CLOUD model "
|
| 176 |
+
f"({self.client.cloud_model}) — frontier reasoning required for "
|
| 177 |
+
f"high-stakes or action-taking tasks."
|
| 178 |
+
)
|
| 179 |
+
|
| 180 |
+
return RouterDecision(
|
| 181 |
+
task_type=task_type,
|
| 182 |
+
priority=priority,
|
| 183 |
+
model_tier=tier,
|
| 184 |
+
model_name=model_name,
|
| 185 |
+
sla_target_ms=sla_ms,
|
| 186 |
+
reasoning=reasoning,
|
| 187 |
+
)
|
| 188 |
+
|
| 189 |
+
def route(
|
| 190 |
+
self,
|
| 191 |
+
messages: list[dict],
|
| 192 |
+
task_type: TaskType,
|
| 193 |
+
priority: Priority,
|
| 194 |
+
temperature: float = 0.1,
|
| 195 |
+
max_tokens: int = 512,
|
| 196 |
+
) -> tuple[LLMResponse, RouterDecision]:
|
| 197 |
+
"""
|
| 198 |
+
Main routing entry point.
|
| 199 |
+
|
| 200 |
+
Decides which model to use, calls it, updates metrics, logs the decision,
|
| 201 |
+
and checks SLA compliance. Returns (LLMResponse, RouterDecision).
|
| 202 |
+
"""
|
| 203 |
+
decision = self._decide(task_type, priority)
|
| 204 |
+
|
| 205 |
+
logger.info(
|
| 206 |
+
"ROUTER → %s | task=%s priority=%s | %s",
|
| 207 |
+
decision.model_tier.upper(), task_type, priority, decision.model_name
|
| 208 |
+
)
|
| 209 |
+
|
| 210 |
+
# Execute the call
|
| 211 |
+
if decision.model_tier == "local":
|
| 212 |
+
response = self.client.call_local(
|
| 213 |
+
messages=messages,
|
| 214 |
+
temperature=temperature,
|
| 215 |
+
max_tokens=max_tokens,
|
| 216 |
+
)
|
| 217 |
+
else:
|
| 218 |
+
response = self.client.call_cloud(
|
| 219 |
+
messages=messages,
|
| 220 |
+
temperature=temperature,
|
| 221 |
+
max_tokens=max_tokens,
|
| 222 |
+
)
|
| 223 |
+
|
| 224 |
+
# Update metrics
|
| 225 |
+
self.metrics.total_calls += 1
|
| 226 |
+
self.metrics.calls_by_task[task_type] += 1
|
| 227 |
+
self.metrics.calls_by_priority[priority] += 1
|
| 228 |
+
|
| 229 |
+
if response.success:
|
| 230 |
+
self.metrics.total_latency_ms += response.latency_ms
|
| 231 |
+
self.metrics.total_cost_usd += response.estimated_cost_usd
|
| 232 |
+
|
| 233 |
+
if decision.model_tier == "local":
|
| 234 |
+
self.metrics.local_calls += 1
|
| 235 |
+
else:
|
| 236 |
+
self.metrics.cloud_calls += 1
|
| 237 |
+
|
| 238 |
+
# SLA check
|
| 239 |
+
if response.latency_ms > decision.sla_target_ms:
|
| 240 |
+
self.metrics.sla_violations += 1
|
| 241 |
+
logger.warning(
|
| 242 |
+
"SLA VIOLATION: task=%s priority=%s latency=%.0fms target=%dms",
|
| 243 |
+
task_type, priority, response.latency_ms, decision.sla_target_ms
|
| 244 |
+
)
|
| 245 |
+
else:
|
| 246 |
+
self.metrics.errors += 1
|
| 247 |
+
logger.error(
|
| 248 |
+
"LLM call FAILED: task=%s priority=%s error=%s",
|
| 249 |
+
task_type, priority, response.error
|
| 250 |
+
)
|
| 251 |
+
|
| 252 |
+
# Append to call log
|
| 253 |
+
self._call_log.append({
|
| 254 |
+
"task_type": task_type,
|
| 255 |
+
"priority": priority,
|
| 256 |
+
"model_tier": decision.model_tier,
|
| 257 |
+
"model_name": response.model_used,
|
| 258 |
+
"latency_ms": response.latency_ms,
|
| 259 |
+
"cost_usd": response.estimated_cost_usd,
|
| 260 |
+
"success": response.success,
|
| 261 |
+
"sla_target_ms": decision.sla_target_ms,
|
| 262 |
+
"sla_met": response.latency_ms <= decision.sla_target_ms,
|
| 263 |
+
})
|
| 264 |
+
|
| 265 |
+
return response, decision
|
| 266 |
+
|
| 267 |
+
def get_metrics_summary(self) -> dict:
|
| 268 |
+
"""Return a dict of aggregated metrics suitable for Prometheus export."""
|
| 269 |
+
return {
|
| 270 |
+
"total_calls": self.metrics.total_calls,
|
| 271 |
+
"local_calls": self.metrics.local_calls,
|
| 272 |
+
"cloud_calls": self.metrics.cloud_calls,
|
| 273 |
+
"local_pct": round(self.metrics.local_pct, 1),
|
| 274 |
+
"cloud_pct": round(self.metrics.cloud_pct, 1),
|
| 275 |
+
"total_cost_usd": round(self.metrics.total_cost_usd, 4),
|
| 276 |
+
"avg_latency_ms": round(self.metrics.avg_latency_ms, 1),
|
| 277 |
+
"sla_violations": self.metrics.sla_violations,
|
| 278 |
+
"errors": self.metrics.errors,
|
| 279 |
+
"calls_by_task": dict(self.metrics.calls_by_task),
|
| 280 |
+
"calls_by_priority": dict(self.metrics.calls_by_priority),
|
| 281 |
+
}
|
| 282 |
+
|
| 283 |
+
def predict_route(self, task_type: str, priority: str) -> RouterDecision:
|
| 284 |
+
"""
|
| 285 |
+
Predict routing decision WITHOUT making an LLM call.
|
| 286 |
+
Useful for UI display, testing, and cost estimation.
|
| 287 |
+
"""
|
| 288 |
+
return self._decide(task_type, priority)
|
| 289 |
+
|
| 290 |
+
|
| 291 |
+
# ── Module-level singleton ─────────────────────────────────────────────────────
|
| 292 |
+
_default_router: Optional[LLMRouter] = None
|
| 293 |
+
|
| 294 |
+
|
| 295 |
+
def get_router() -> LLMRouter:
|
| 296 |
+
global _default_router
|
| 297 |
+
if _default_router is None:
|
| 298 |
+
_default_router = LLMRouter()
|
| 299 |
+
return _default_router
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
# ── Standalone Demo ────────────────────────────────────────────────────────────
|
| 303 |
+
if __name__ == "__main__":
|
| 304 |
+
logging.basicConfig(level=logging.WARNING)
|
| 305 |
+
|
| 306 |
+
router = LLMRouter()
|
| 307 |
+
|
| 308 |
+
print("=" * 65)
|
| 309 |
+
print("CustomerCore Phase 7 — SLA-Aware LLM Router Demo")
|
| 310 |
+
print("(Routing decisions only — no actual LLM calls in dry-run)")
|
| 311 |
+
print("=" * 65)
|
| 312 |
+
|
| 313 |
+
test_cases = [
|
| 314 |
+
("classify", "low", "Classify this support ticket: 'I forgot my password'"),
|
| 315 |
+
("classify", "medium", "Classify this ticket: 'API returning 500 errors on checkout'"),
|
| 316 |
+
("classify", "high", "Classify urgency: 'Production database is down for 200 users'"),
|
| 317 |
+
("classify", "critical", "Classify: 'ALL services offline - complete outage for enterprise'"),
|
| 318 |
+
("extract", "medium", "Extract customer ID and error code from this ticket"),
|
| 319 |
+
("reason", "low", "Reason about whether this ticket needs escalation"),
|
| 320 |
+
("reason", "medium", "Determine if SLA is at risk for this ticket"),
|
| 321 |
+
("reason", "high", "Reason: Should we proactively offer a refund for this outage?"),
|
| 322 |
+
("reason", "critical", "Critical: Is this a coordinated security breach across tenants?"),
|
| 323 |
+
("action", "low", "Send automated acknowledgment email to customer"),
|
| 324 |
+
("action", "medium", "Create Jira ticket and assign to billing team"),
|
| 325 |
+
("action", "high", "Issue $150 partial refund and update HubSpot CRM"),
|
| 326 |
+
("action", "critical", "Escalate to VIP team + notify on-call engineer immediately"),
|
| 327 |
+
]
|
| 328 |
+
|
| 329 |
+
print(f"\n{'Task':<10} {'Priority':<10} {'-> Tier':<8} {'Model':<35} {'SLA Target'}")
|
| 330 |
+
print("─" * 90)
|
| 331 |
+
|
| 332 |
+
local_count = 0
|
| 333 |
+
cloud_count = 0
|
| 334 |
+
|
| 335 |
+
for task_type, priority, _ in test_cases:
|
| 336 |
+
decision = router.predict_route(task_type, priority)
|
| 337 |
+
tier_display = f"{'🟢 LOCAL' if decision.model_tier == 'local' else '🔵 CLOUD'}"
|
| 338 |
+
model_short = decision.model_name.split("/")[-1][:32]
|
| 339 |
+
print(
|
| 340 |
+
f"{task_type:<10} {priority:<10} {tier_display:<14} {model_short:<35} "
|
| 341 |
+
f"≤{decision.sla_target_ms}ms"
|
| 342 |
+
)
|
| 343 |
+
if decision.model_tier == "local":
|
| 344 |
+
local_count += 1
|
| 345 |
+
else:
|
| 346 |
+
cloud_count += 1
|
| 347 |
+
|
| 348 |
+
total = len(test_cases)
|
| 349 |
+
print("─" * 90)
|
| 350 |
+
print(f"\nRouting summary: {local_count}/{total} LOCAL (${0:.2f} cost) | "
|
| 351 |
+
f"{cloud_count}/{total} CLOUD (frontier reasoning)")
|
| 352 |
+
print(f"\nCost savings vs. always-cloud: "
|
| 353 |
+
f"~{local_count/total*100:.0f}% of calls are free")
|
| 354 |
+
|
| 355 |
+
# Verify critical routing rules
|
| 356 |
+
print("\n" + "─" * 40)
|
| 357 |
+
print("Verifying critical routing rules:")
|
| 358 |
+
rules = [
|
| 359 |
+
(("action", "low"), "cloud", "ALL actions must use cloud"),
|
| 360 |
+
(("action", "critical"), "cloud", "ALL actions must use cloud"),
|
| 361 |
+
(("reason", "critical"), "cloud", "Critical reasoning must use cloud"),
|
| 362 |
+
(("reason", "high"), "cloud", "High-priority reasoning must use cloud"),
|
| 363 |
+
(("classify", "critical"),"local", "Classification always local"),
|
| 364 |
+
(("extract", "high"), "local", "Extraction always local"),
|
| 365 |
+
(("reason", "low"), "local", "Low-priority reasoning is local"),
|
| 366 |
+
]
|
| 367 |
+
all_ok = True
|
| 368 |
+
for (task, prio), expected_tier, desc in rules:
|
| 369 |
+
d = router.predict_route(task, prio)
|
| 370 |
+
ok = d.model_tier == expected_tier
|
| 371 |
+
icon = "✓" if ok else "✗"
|
| 372 |
+
print(f" {icon} {desc}: {task}/{prio} → {d.model_tier} (expected {expected_tier})")
|
| 373 |
+
if not ok:
|
| 374 |
+
all_ok = False
|
| 375 |
+
|
| 376 |
+
print(f"\n{'All routing rules VERIFIED ✓' if all_ok else 'ROUTING RULE FAILURES!'}")
|
| 377 |
+
print("=" * 65)
|
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/responsible_ai/__init__.py
|
| 3 |
+
|
| 4 |
+
CustomerCore Responsible AI Package.
|
| 5 |
+
Provides EU AI Act / GDPR compliance modules:
|
| 6 |
+
- privacy_vault: Cryptographic AES-256 PII token vault
|
| 7 |
+
- key_manager: Tenant-specific encryption key store
|
| 8 |
+
- audit_log: Immutable prediction & decryption audit trail
|
| 9 |
+
- policy_engine: Declarative per-tenant brand/legal policy enforcement
|
| 10 |
+
"""
|
|
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/responsible_ai/audit_log.py
|
| 3 |
+
|
| 4 |
+
Immutable audit logger for the CustomerCore Privacy Vault.
|
| 5 |
+
|
| 6 |
+
Every encrypt and decrypt operation is written here with:
|
| 7 |
+
- tenant_id, field_name, token
|
| 8 |
+
- action: "ENCRYPT" or "DECRYPT"
|
| 9 |
+
- actor_role (for DECRYPT calls — who triggered the re-identification)
|
| 10 |
+
- timestamp (UTC ISO-8601)
|
| 11 |
+
- SHA-256 hash of the plaintext (for compliance verification without storing PII)
|
| 12 |
+
|
| 13 |
+
In local dev: logs are written to a rotating JSON-lines file under logs/.
|
| 14 |
+
In production: rows are inserted to a Supabase PostgreSQL table via the REST API.
|
| 15 |
+
|
| 16 |
+
Run standalone check:
|
| 17 |
+
python -m src.responsible_ai.audit_log
|
| 18 |
+
"""
|
| 19 |
+
|
| 20 |
+
import hashlib
|
| 21 |
+
import json
|
| 22 |
+
import logging
|
| 23 |
+
import os
|
| 24 |
+
from datetime import datetime, timezone
|
| 25 |
+
from pathlib import Path
|
| 26 |
+
from typing import Literal
|
| 27 |
+
|
| 28 |
+
# ── Config ────────────────────────────────────────────────────────────────────
|
| 29 |
+
LOG_DIR = Path("logs/vault_audit")
|
| 30 |
+
LOG_DIR.mkdir(parents=True, exist_ok=True)
|
| 31 |
+
AUDIT_FILE = LOG_DIR / "audit_trail.jsonl"
|
| 32 |
+
|
| 33 |
+
# Configure Python logger for console output too
|
| 34 |
+
logging.basicConfig(
|
| 35 |
+
level=logging.INFO,
|
| 36 |
+
format="%(asctime)s [AUDIT] %(message)s",
|
| 37 |
+
datefmt="%Y-%m-%dT%H:%M:%SZ",
|
| 38 |
+
)
|
| 39 |
+
_log = logging.getLogger("customercore.audit")
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
ActionType = Literal["ENCRYPT", "DECRYPT", "DECRYPT_DENIED"]
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def _sha256(text: str) -> str:
|
| 46 |
+
"""Return SHA-256 hex digest of a UTF-8 string."""
|
| 47 |
+
return hashlib.sha256(text.encode("utf-8")).hexdigest()
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def write_audit_entry(
|
| 51 |
+
*,
|
| 52 |
+
tenant_id: str,
|
| 53 |
+
field_name: str,
|
| 54 |
+
token: str,
|
| 55 |
+
action: ActionType,
|
| 56 |
+
plaintext_hash: str,
|
| 57 |
+
actor_role: str = "PIPELINE",
|
| 58 |
+
ticket_id: str = "",
|
| 59 |
+
extra: dict | None = None,
|
| 60 |
+
) -> dict:
|
| 61 |
+
"""
|
| 62 |
+
Write a single audit entry to the local JSONL file.
|
| 63 |
+
|
| 64 |
+
Returns the entry dict so callers can inspect it in tests.
|
| 65 |
+
"""
|
| 66 |
+
entry = {
|
| 67 |
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
| 68 |
+
"action": action,
|
| 69 |
+
"tenant_id": tenant_id,
|
| 70 |
+
"field_name": field_name,
|
| 71 |
+
"token": token,
|
| 72 |
+
"plaintext_sha256": plaintext_hash,
|
| 73 |
+
"actor_role": actor_role,
|
| 74 |
+
"ticket_id": ticket_id,
|
| 75 |
+
**(extra or {}),
|
| 76 |
+
}
|
| 77 |
+
|
| 78 |
+
# Write to local JSONL file (append-only)
|
| 79 |
+
with AUDIT_FILE.open("a", encoding="utf-8") as f:
|
| 80 |
+
f.write(json.dumps(entry) + "\n")
|
| 81 |
+
|
| 82 |
+
# Console log (less detail to avoid accidentally logging sensitive data)
|
| 83 |
+
_log.info(
|
| 84 |
+
"[%s] tenant=%s field=%s token=%s role=%s",
|
| 85 |
+
action, tenant_id, field_name, token[:12] + "...", actor_role,
|
| 86 |
+
)
|
| 87 |
+
|
| 88 |
+
return entry
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def tail_audit_log(n: int = 10) -> list[dict]:
|
| 92 |
+
"""Return the last n entries from the audit log."""
|
| 93 |
+
if not AUDIT_FILE.exists():
|
| 94 |
+
return []
|
| 95 |
+
lines = AUDIT_FILE.read_text(encoding="utf-8").strip().splitlines()
|
| 96 |
+
return [json.loads(line) for line in lines[-n:]]
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
if __name__ == "__main__":
|
| 100 |
+
# Quick standalone test
|
| 101 |
+
e = write_audit_entry(
|
| 102 |
+
tenant_id="acme-corp",
|
| 103 |
+
field_name="EMAIL_ADDRESS",
|
| 104 |
+
token="TOK_EMAIL_ab12cd34",
|
| 105 |
+
action="ENCRYPT",
|
| 106 |
+
plaintext_hash=_sha256("john.doe@acme.com"),
|
| 107 |
+
ticket_id="TKT-99999",
|
| 108 |
+
)
|
| 109 |
+
print("Wrote ENCRYPT entry:")
|
| 110 |
+
print(json.dumps(e, indent=2))
|
| 111 |
+
|
| 112 |
+
d = write_audit_entry(
|
| 113 |
+
tenant_id="acme-corp",
|
| 114 |
+
field_name="EMAIL_ADDRESS",
|
| 115 |
+
token="TOK_EMAIL_ab12cd34",
|
| 116 |
+
action="DECRYPT",
|
| 117 |
+
plaintext_hash=_sha256("john.doe@acme.com"),
|
| 118 |
+
actor_role="SUPPORT_LEAD",
|
| 119 |
+
ticket_id="TKT-99999",
|
| 120 |
+
)
|
| 121 |
+
print("\nWrote DECRYPT entry:")
|
| 122 |
+
print(json.dumps(d, indent=2))
|
| 123 |
+
|
| 124 |
+
print(f"\nLast 2 audit entries: {len(tail_audit_log(2))}")
|
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/responsible_ai/key_manager.py
|
| 3 |
+
|
| 4 |
+
Tenant-specific AES-256 encryption key manager.
|
| 5 |
+
|
| 6 |
+
Each tenant gets a unique, deterministic Fernet key derived from a
|
| 7 |
+
master secret + their tenant_id via HMAC-SHA256. This guarantees:
|
| 8 |
+
- Cross-tenant isolation: encrypting with Tenant A's key produces
|
| 9 |
+
ciphertext that Tenant B's key CANNOT decrypt.
|
| 10 |
+
- Determinism: the same tenant_id + master secret always produces
|
| 11 |
+
the same key, so vault tokens remain valid across restarts without
|
| 12 |
+
persisting individual keys.
|
| 13 |
+
- Zero plaintext key storage: keys are never written to disk.
|
| 14 |
+
|
| 15 |
+
In production, swap MASTER_SECRET with a value injected by Doppler.
|
| 16 |
+
In cloud, use Google Cloud KMS or AWS KMS as the key derivation backend.
|
| 17 |
+
|
| 18 |
+
Run standalone check:
|
| 19 |
+
python -m src.responsible_ai.key_manager
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
import hashlib
|
| 23 |
+
import hmac
|
| 24 |
+
import os
|
| 25 |
+
import base64
|
| 26 |
+
from typing import Optional
|
| 27 |
+
|
| 28 |
+
# ── Master Secret ─────────────────────────────────────────────────────────────
|
| 29 |
+
# Injected by Doppler in production. Falls back to a safe local dev value.
|
| 30 |
+
# WARNING: Never hardcode a real secret in production code.
|
| 31 |
+
_MASTER_SECRET: str = os.environ.get(
|
| 32 |
+
"VAULT_MASTER_SECRET",
|
| 33 |
+
"customercore-local-dev-secret-changeme-in-prod"
|
| 34 |
+
)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class TenantKeyManager:
|
| 38 |
+
"""
|
| 39 |
+
Derives and caches per-tenant AES-256 Fernet keys from a master secret.
|
| 40 |
+
|
| 41 |
+
Usage:
|
| 42 |
+
km = TenantKeyManager()
|
| 43 |
+
key = km.get_key("acme-corp") # returns bytes (URL-safe base64)
|
| 44 |
+
key_same = km.get_key("acme-corp") # same key, cached
|
| 45 |
+
key_diff = km.get_key("globex-inc") # different key, isolated
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
def __init__(self, master_secret: Optional[str] = None):
|
| 49 |
+
self._secret = (master_secret or _MASTER_SECRET).encode()
|
| 50 |
+
self._cache: dict[str, bytes] = {}
|
| 51 |
+
|
| 52 |
+
def _derive(self, tenant_id: str) -> bytes:
|
| 53 |
+
"""
|
| 54 |
+
Derive a 32-byte key from HMAC-SHA256(master_secret, tenant_id).
|
| 55 |
+
Then base64url-encode it so Fernet can consume it directly.
|
| 56 |
+
"""
|
| 57 |
+
raw = hmac.new(self._secret, tenant_id.encode(), hashlib.sha256).digest()
|
| 58 |
+
# Fernet requires URL-safe base64-encoded 32-byte key
|
| 59 |
+
return base64.urlsafe_b64encode(raw)
|
| 60 |
+
|
| 61 |
+
def get_key(self, tenant_id: str) -> bytes:
|
| 62 |
+
"""Return the deterministic encryption key for a given tenant."""
|
| 63 |
+
if tenant_id not in self._cache:
|
| 64 |
+
self._cache[tenant_id] = self._derive(tenant_id)
|
| 65 |
+
return self._cache[tenant_id]
|
| 66 |
+
|
| 67 |
+
def rotate_key(self, tenant_id: str) -> bytes:
|
| 68 |
+
"""
|
| 69 |
+
Force re-derive and cache a new key for a tenant.
|
| 70 |
+
NOTE: Existing encrypted tokens become unreadable after rotation.
|
| 71 |
+
In production, run a migration to re-encrypt all vault entries first.
|
| 72 |
+
"""
|
| 73 |
+
key = self._derive(tenant_id)
|
| 74 |
+
self._cache[tenant_id] = key
|
| 75 |
+
return key
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
# ── Module-level singleton ────────────────────────────────────────────────────
|
| 79 |
+
_default_manager: Optional[TenantKeyManager] = None
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def get_key_manager() -> TenantKeyManager:
|
| 83 |
+
"""Return the module-level singleton TenantKeyManager."""
|
| 84 |
+
global _default_manager
|
| 85 |
+
if _default_manager is None:
|
| 86 |
+
_default_manager = TenantKeyManager()
|
| 87 |
+
return _default_manager
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
if __name__ == "__main__":
|
| 91 |
+
km = TenantKeyManager()
|
| 92 |
+
t1 = km.get_key("acme-corp")
|
| 93 |
+
t2 = km.get_key("globex-inc")
|
| 94 |
+
t1_again = km.get_key("acme-corp")
|
| 95 |
+
|
| 96 |
+
print("=== TenantKeyManager Verification ===")
|
| 97 |
+
print(f"acme-corp key (1st call) : {t1[:20]}...")
|
| 98 |
+
print(f"acme-corp key (2nd call) : {t1_again[:20]}... [must match above]")
|
| 99 |
+
print(f"globex-inc key : {t2[:20]}...")
|
| 100 |
+
print(f"Keys are isolated : {t1 != t2}")
|
| 101 |
+
print(f"Key is deterministic : {t1 == t1_again}")
|
|
@@ -0,0 +1,442 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/responsible_ai/privacy_vault.py
|
| 3 |
+
|
| 4 |
+
Zero-Trust Cryptographic Privacy Vault for CustomerCore.
|
| 5 |
+
|
| 6 |
+
== What this replaces ==
|
| 7 |
+
The previous approach (Presidio simple masking) permanently deleted PII from
|
| 8 |
+
ticket text. This meant human agents reviewing low-confidence tickets during
|
| 9 |
+
Human-in-the-Loop pauses could only see [EMAIL] and [NAME] — making triage
|
| 10 |
+
impossible without contacting the customer separately.
|
| 11 |
+
|
| 12 |
+
== What this does instead ==
|
| 13 |
+
Instead of destroying PII, we:
|
| 14 |
+
1. Detect PII spans using Presidio (same as before)
|
| 15 |
+
2. Encrypt each detected value with the tenant's AES-256 Fernet key
|
| 16 |
+
3. Replace the raw value in the ticket text with a readable token
|
| 17 |
+
e.g. "john.doe@acme.com" → "<<TOK_EMAIL_ADDRESS_a1b2c3d4>>"
|
| 18 |
+
4. Store the (token → encrypted_value) pair in the vault database
|
| 19 |
+
5. Write an audit entry for every encrypt and decrypt action
|
| 20 |
+
|
| 21 |
+
When a SUPPORT_LEAD or SECURITY_ADMIN needs to review the ticket, they call
|
| 22 |
+
decrypt_token() which re-identifies the field in their UI. Every decrypt is
|
| 23 |
+
logged for GDPR audit compliance.
|
| 24 |
+
|
| 25 |
+
== GDPR Compliance (Germany / EU) ==
|
| 26 |
+
GDPR Article 32: Requires "appropriate technical measures" to protect
|
| 27 |
+
personal data — AES-256 encryption is the gold standard.
|
| 28 |
+
GDPR Article 17 (Right to be Forgotten): To delete a customer's data,
|
| 29 |
+
delete their vault entries AND shred the tenant's encryption key.
|
| 30 |
+
Their Silver Parquet records become permanently unreadable ciphertext.
|
| 31 |
+
|
| 32 |
+
== Architecture ==
|
| 33 |
+
CryptographicPrivacyVault
|
| 34 |
+
├── TenantKeyManager — per-tenant AES-256 key derivation
|
| 35 |
+
├── VaultStore — in-memory + SQLite backend for tokens
|
| 36 |
+
└── AuditLog — JSONL append-only compliance trail
|
| 37 |
+
|
| 38 |
+
Run standalone demo:
|
| 39 |
+
python -m src.responsible_ai.privacy_vault
|
| 40 |
+
"""
|
| 41 |
+
|
| 42 |
+
import os
|
| 43 |
+
import re
|
| 44 |
+
import sqlite3
|
| 45 |
+
import hashlib
|
| 46 |
+
from datetime import datetime, timezone
|
| 47 |
+
from pathlib import Path
|
| 48 |
+
from typing import Optional
|
| 49 |
+
|
| 50 |
+
from cryptography.fernet import Fernet, InvalidToken
|
| 51 |
+
|
| 52 |
+
from src.responsible_ai.key_manager import TenantKeyManager, get_key_manager
|
| 53 |
+
from src.responsible_ai.audit_log import write_audit_entry, _sha256
|
| 54 |
+
|
| 55 |
+
# ── Config ─────────────────────────────────────────────────────────────────────
|
| 56 |
+
VAULT_DB_PATH = Path("logs/vault_audit/vault.db")
|
| 57 |
+
VAULT_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
| 58 |
+
|
| 59 |
+
# Roles allowed to call decrypt
|
| 60 |
+
AUTHORIZED_DECRYPT_ROLES = {"SUPPORT_LEAD", "SECURITY_ADMIN"}
|
| 61 |
+
|
| 62 |
+
# Token format: <<TOK_{FIELD}_{8-hex-chars}>>
|
| 63 |
+
TOKEN_PATTERN = re.compile(r"<<TOK_[A-Z_]+_[0-9a-f]{8}>>")
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
# ── Vault Store (SQLite) ───────────────────────────────────────────────────────
|
| 67 |
+
|
| 68 |
+
class VaultStore:
|
| 69 |
+
"""
|
| 70 |
+
SQLite-backed store mapping (tenant_id, token) → encrypted_value.
|
| 71 |
+
|
| 72 |
+
In production, swap with a Supabase PostgreSQL table:
|
| 73 |
+
supabase.table("pii_vault").insert({...}).execute()
|
| 74 |
+
|
| 75 |
+
The SQLite DB is stored at logs/vault_audit/vault.db — never commit this.
|
| 76 |
+
"""
|
| 77 |
+
|
| 78 |
+
def __init__(self, db_path: Path = VAULT_DB_PATH):
|
| 79 |
+
self.db_path = db_path
|
| 80 |
+
self._conn = sqlite3.connect(str(db_path), check_same_thread=False)
|
| 81 |
+
self._init_schema()
|
| 82 |
+
|
| 83 |
+
def _init_schema(self):
|
| 84 |
+
self._conn.execute("""
|
| 85 |
+
CREATE TABLE IF NOT EXISTS pii_vault (
|
| 86 |
+
token TEXT NOT NULL,
|
| 87 |
+
tenant_id TEXT NOT NULL,
|
| 88 |
+
field_name TEXT NOT NULL,
|
| 89 |
+
encrypted_value TEXT NOT NULL,
|
| 90 |
+
plaintext_sha256 TEXT NOT NULL,
|
| 91 |
+
created_at TEXT NOT NULL,
|
| 92 |
+
PRIMARY KEY (tenant_id, token)
|
| 93 |
+
)
|
| 94 |
+
""")
|
| 95 |
+
self._conn.commit()
|
| 96 |
+
|
| 97 |
+
def store(self, tenant_id: str, token: str, field_name: str,
|
| 98 |
+
encrypted_value: str, plaintext_sha256: str):
|
| 99 |
+
self._conn.execute(
|
| 100 |
+
"""INSERT OR REPLACE INTO pii_vault
|
| 101 |
+
(token, tenant_id, field_name, encrypted_value, plaintext_sha256, created_at)
|
| 102 |
+
VALUES (?, ?, ?, ?, ?, ?)""",
|
| 103 |
+
(token, tenant_id, field_name, encrypted_value,
|
| 104 |
+
plaintext_sha256, datetime.now(timezone.utc).isoformat()),
|
| 105 |
+
)
|
| 106 |
+
self._conn.commit()
|
| 107 |
+
|
| 108 |
+
def fetch(self, tenant_id: str, token: str) -> Optional[tuple[str, str]]:
|
| 109 |
+
"""Return (field_name, encrypted_value) or None if not found."""
|
| 110 |
+
row = self._conn.execute(
|
| 111 |
+
"SELECT field_name, encrypted_value FROM pii_vault WHERE tenant_id=? AND token=?",
|
| 112 |
+
(tenant_id, token),
|
| 113 |
+
).fetchone()
|
| 114 |
+
return row # (field_name, encrypted_value) or None
|
| 115 |
+
|
| 116 |
+
def delete_tenant(self, tenant_id: str) -> int:
|
| 117 |
+
"""
|
| 118 |
+
GDPR Right to be Forgotten: delete ALL vault entries for a tenant.
|
| 119 |
+
Combined with key shredding, their encrypted tokens become permanently
|
| 120 |
+
unreadable ciphertext in the Silver Parquet files.
|
| 121 |
+
Returns count of deleted rows.
|
| 122 |
+
"""
|
| 123 |
+
cur = self._conn.execute(
|
| 124 |
+
"DELETE FROM pii_vault WHERE tenant_id=?", (tenant_id,)
|
| 125 |
+
)
|
| 126 |
+
self._conn.commit()
|
| 127 |
+
return cur.rowcount
|
| 128 |
+
|
| 129 |
+
def count(self, tenant_id: Optional[str] = None) -> int:
|
| 130 |
+
if tenant_id:
|
| 131 |
+
return self._conn.execute(
|
| 132 |
+
"SELECT COUNT(*) FROM pii_vault WHERE tenant_id=?", (tenant_id,)
|
| 133 |
+
).fetchone()[0]
|
| 134 |
+
return self._conn.execute("SELECT COUNT(*) FROM pii_vault").fetchone()[0]
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
# ── Core Vault ─────────────────────────────────────────────────────────────────
|
| 138 |
+
|
| 139 |
+
class CryptographicPrivacyVault:
|
| 140 |
+
"""
|
| 141 |
+
Zero-Trust Privacy Vault.
|
| 142 |
+
|
| 143 |
+
encrypt_field() — encrypt a single PII value, return its token
|
| 144 |
+
encrypt_text() — detect + encrypt all PII in a full text block
|
| 145 |
+
decrypt_token() — decrypt a token back to plaintext (RBAC-enforced)
|
| 146 |
+
forget_tenant() — GDPR deletion cascade for all tokens of a tenant
|
| 147 |
+
"""
|
| 148 |
+
|
| 149 |
+
def __init__(
|
| 150 |
+
self,
|
| 151 |
+
key_manager: Optional[TenantKeyManager] = None,
|
| 152 |
+
store: Optional[VaultStore] = None,
|
| 153 |
+
):
|
| 154 |
+
self.km = key_manager or get_key_manager()
|
| 155 |
+
self.store = store or VaultStore()
|
| 156 |
+
|
| 157 |
+
# ── Internal helpers ───────────────────────────────────────────────────────
|
| 158 |
+
|
| 159 |
+
def _make_token(self, field_name: str, plaintext: str) -> str:
|
| 160 |
+
"""
|
| 161 |
+
Generate a deterministic, stable token for a PII value.
|
| 162 |
+
Token format: <<TOK_{FIELD_NAME}_{first8hex_of_sha256}>>
|
| 163 |
+
Deterministic so the same email always produces the same token,
|
| 164 |
+
avoiding duplicate vault entries for repeated values.
|
| 165 |
+
"""
|
| 166 |
+
digest = _sha256(plaintext)[:8]
|
| 167 |
+
clean_field = field_name.upper().replace(" ", "_")
|
| 168 |
+
return f"<<TOK_{clean_field}_{digest}>>"
|
| 169 |
+
|
| 170 |
+
def _get_cipher(self, tenant_id: str) -> Fernet:
|
| 171 |
+
key = self.km.get_key(tenant_id)
|
| 172 |
+
return Fernet(key)
|
| 173 |
+
|
| 174 |
+
# ── Public API ─────────────────────────────────────────────────────────────
|
| 175 |
+
|
| 176 |
+
def encrypt_field(
|
| 177 |
+
self,
|
| 178 |
+
*,
|
| 179 |
+
tenant_id: str,
|
| 180 |
+
field_name: str,
|
| 181 |
+
plaintext: str,
|
| 182 |
+
ticket_id: str = "",
|
| 183 |
+
) -> str:
|
| 184 |
+
"""
|
| 185 |
+
Encrypt a single PII value. Store in vault. Return its token.
|
| 186 |
+
|
| 187 |
+
If this exact plaintext was already encrypted for this tenant,
|
| 188 |
+
returns the existing token (idempotent — no duplicate vault entries).
|
| 189 |
+
"""
|
| 190 |
+
token = self._make_token(field_name, plaintext)
|
| 191 |
+
|
| 192 |
+
# Idempotency check — already in vault?
|
| 193 |
+
existing = self.store.fetch(tenant_id, token)
|
| 194 |
+
if existing:
|
| 195 |
+
return token # already stored, token is stable
|
| 196 |
+
|
| 197 |
+
# Encrypt with tenant key
|
| 198 |
+
cipher = self._get_cipher(tenant_id)
|
| 199 |
+
encrypted = cipher.encrypt(plaintext.encode("utf-8")).decode("utf-8")
|
| 200 |
+
sha = _sha256(plaintext)
|
| 201 |
+
|
| 202 |
+
# Persist to vault
|
| 203 |
+
self.store.store(tenant_id, token, field_name, encrypted, sha)
|
| 204 |
+
|
| 205 |
+
# Write audit entry
|
| 206 |
+
write_audit_entry(
|
| 207 |
+
tenant_id=tenant_id,
|
| 208 |
+
field_name=field_name,
|
| 209 |
+
token=token,
|
| 210 |
+
action="ENCRYPT",
|
| 211 |
+
plaintext_hash=sha,
|
| 212 |
+
ticket_id=ticket_id,
|
| 213 |
+
)
|
| 214 |
+
|
| 215 |
+
return token
|
| 216 |
+
|
| 217 |
+
def encrypt_text(
|
| 218 |
+
self,
|
| 219 |
+
*,
|
| 220 |
+
tenant_id: str,
|
| 221 |
+
text: str,
|
| 222 |
+
analyzer,
|
| 223 |
+
anonymizer,
|
| 224 |
+
ticket_id: str = "",
|
| 225 |
+
pii_entities: list[str] | None = None,
|
| 226 |
+
) -> tuple[str, int]:
|
| 227 |
+
"""
|
| 228 |
+
Detect all PII spans in `text` using Presidio, encrypt each one,
|
| 229 |
+
and replace the raw span with its vault token in the returned string.
|
| 230 |
+
|
| 231 |
+
Returns (tokenized_text, count_of_pii_detected).
|
| 232 |
+
Falls back to standard Presidio anonymization if vault is disabled.
|
| 233 |
+
"""
|
| 234 |
+
if not text or not isinstance(text, str):
|
| 235 |
+
return text, 0
|
| 236 |
+
|
| 237 |
+
entities = pii_entities or [
|
| 238 |
+
"EMAIL_ADDRESS", "PHONE_NUMBER", "CREDIT_CARD",
|
| 239 |
+
"PERSON", "IBAN_CODE", "IP_ADDRESS", "US_SSN",
|
| 240 |
+
]
|
| 241 |
+
|
| 242 |
+
results = analyzer.analyze(text=text, language="en", entities=entities)
|
| 243 |
+
if not results:
|
| 244 |
+
return text, 0
|
| 245 |
+
|
| 246 |
+
# Sort by start position descending so we can replace without offset shift
|
| 247 |
+
results_sorted = sorted(results, key=lambda r: r.start, reverse=True)
|
| 248 |
+
tokenized = text
|
| 249 |
+
|
| 250 |
+
for result in results_sorted:
|
| 251 |
+
raw_value = text[result.start:result.end]
|
| 252 |
+
token = self.encrypt_field(
|
| 253 |
+
tenant_id=tenant_id,
|
| 254 |
+
field_name=result.entity_type,
|
| 255 |
+
plaintext=raw_value,
|
| 256 |
+
ticket_id=ticket_id,
|
| 257 |
+
)
|
| 258 |
+
tokenized = tokenized[:result.start] + token + tokenized[result.end:]
|
| 259 |
+
|
| 260 |
+
return tokenized, len(results_sorted)
|
| 261 |
+
|
| 262 |
+
def decrypt_token(
|
| 263 |
+
self,
|
| 264 |
+
*,
|
| 265 |
+
tenant_id: str,
|
| 266 |
+
token: str,
|
| 267 |
+
actor_role: str,
|
| 268 |
+
ticket_id: str = "",
|
| 269 |
+
) -> str:
|
| 270 |
+
"""
|
| 271 |
+
Decrypt a vault token back to its original plaintext value.
|
| 272 |
+
|
| 273 |
+
Role-Based Access Control: only SUPPORT_LEAD and SECURITY_ADMIN
|
| 274 |
+
may call this. All decrypt calls are logged in the audit trail.
|
| 275 |
+
Raises PermissionError for unauthorized roles.
|
| 276 |
+
Raises KeyError if token not found in vault.
|
| 277 |
+
Raises ValueError if decryption fails (wrong key or corrupted data).
|
| 278 |
+
"""
|
| 279 |
+
# RBAC check
|
| 280 |
+
if actor_role not in AUTHORIZED_DECRYPT_ROLES:
|
| 281 |
+
write_audit_entry(
|
| 282 |
+
tenant_id=tenant_id,
|
| 283 |
+
field_name="UNKNOWN",
|
| 284 |
+
token=token,
|
| 285 |
+
action="DECRYPT_DENIED",
|
| 286 |
+
plaintext_hash="",
|
| 287 |
+
actor_role=actor_role,
|
| 288 |
+
ticket_id=ticket_id,
|
| 289 |
+
)
|
| 290 |
+
raise PermissionError(
|
| 291 |
+
f"Role '{actor_role}' is not authorized to decrypt vault tokens. "
|
| 292 |
+
f"Required: {AUTHORIZED_DECRYPT_ROLES}"
|
| 293 |
+
)
|
| 294 |
+
|
| 295 |
+
# Lookup in vault
|
| 296 |
+
row = self.store.fetch(tenant_id, token)
|
| 297 |
+
if not row:
|
| 298 |
+
raise KeyError(
|
| 299 |
+
f"Token '{token}' not found in vault for tenant '{tenant_id}'."
|
| 300 |
+
)
|
| 301 |
+
field_name, encrypted_value = row
|
| 302 |
+
|
| 303 |
+
# Decrypt
|
| 304 |
+
cipher = self._get_cipher(tenant_id)
|
| 305 |
+
try:
|
| 306 |
+
plaintext = cipher.decrypt(encrypted_value.encode("utf-8")).decode("utf-8")
|
| 307 |
+
except InvalidToken as e:
|
| 308 |
+
raise ValueError(
|
| 309 |
+
f"Failed to decrypt token '{token}' — key may have been rotated: {e}"
|
| 310 |
+
) from e
|
| 311 |
+
|
| 312 |
+
# Audit the successful decrypt
|
| 313 |
+
write_audit_entry(
|
| 314 |
+
tenant_id=tenant_id,
|
| 315 |
+
field_name=field_name,
|
| 316 |
+
token=token,
|
| 317 |
+
action="DECRYPT",
|
| 318 |
+
plaintext_hash=_sha256(plaintext),
|
| 319 |
+
actor_role=actor_role,
|
| 320 |
+
ticket_id=ticket_id,
|
| 321 |
+
)
|
| 322 |
+
|
| 323 |
+
return plaintext
|
| 324 |
+
|
| 325 |
+
def decrypt_text(
|
| 326 |
+
self,
|
| 327 |
+
*,
|
| 328 |
+
tenant_id: str,
|
| 329 |
+
text: str,
|
| 330 |
+
actor_role: str,
|
| 331 |
+
ticket_id: str = "",
|
| 332 |
+
) -> str:
|
| 333 |
+
"""
|
| 334 |
+
Find all vault tokens in `text` and replace them with decrypted values.
|
| 335 |
+
Requires an authorized role. Each token decrypt is individually audited.
|
| 336 |
+
"""
|
| 337 |
+
tokens_found = TOKEN_PATTERN.findall(text)
|
| 338 |
+
result = text
|
| 339 |
+
for token in set(tokens_found):
|
| 340 |
+
try:
|
| 341 |
+
plaintext = self.decrypt_token(
|
| 342 |
+
tenant_id=tenant_id,
|
| 343 |
+
token=token,
|
| 344 |
+
actor_role=actor_role,
|
| 345 |
+
ticket_id=ticket_id,
|
| 346 |
+
)
|
| 347 |
+
result = result.replace(token, plaintext)
|
| 348 |
+
except (KeyError, ValueError):
|
| 349 |
+
# Token not found or corrupted — leave as-is
|
| 350 |
+
pass
|
| 351 |
+
return result
|
| 352 |
+
|
| 353 |
+
def forget_tenant(self, tenant_id: str) -> dict:
|
| 354 |
+
"""
|
| 355 |
+
GDPR Article 17 — Right to be Forgotten.
|
| 356 |
+
|
| 357 |
+
Deletes ALL vault entries for a tenant. After this call, their encrypted
|
| 358 |
+
tokens in Silver Parquet files become permanently unreadable ciphertext.
|
| 359 |
+
Callers should also shred the tenant's key via key_manager.rotate_key().
|
| 360 |
+
|
| 361 |
+
Returns a summary of what was deleted.
|
| 362 |
+
"""
|
| 363 |
+
deleted = self.store.delete_tenant(tenant_id)
|
| 364 |
+
write_audit_entry(
|
| 365 |
+
tenant_id=tenant_id,
|
| 366 |
+
field_name="ALL",
|
| 367 |
+
token="ALL",
|
| 368 |
+
action="ENCRYPT", # closest valid action — log the deletion event
|
| 369 |
+
plaintext_hash="GDPR_DELETION",
|
| 370 |
+
actor_role="GDPR_PROCESSOR",
|
| 371 |
+
extra={"event": "GDPR_FORGET_TENANT", "deleted_rows": deleted},
|
| 372 |
+
)
|
| 373 |
+
return {"tenant_id": tenant_id, "vault_entries_deleted": deleted}
|
| 374 |
+
|
| 375 |
+
|
| 376 |
+
# ── Module-level singleton ─────────────────────────────────────────────────────
|
| 377 |
+
_default_vault: Optional[CryptographicPrivacyVault] = None
|
| 378 |
+
|
| 379 |
+
|
| 380 |
+
def get_vault() -> CryptographicPrivacyVault:
|
| 381 |
+
"""Return the module-level singleton vault (shared key manager + store)."""
|
| 382 |
+
global _default_vault
|
| 383 |
+
if _default_vault is None:
|
| 384 |
+
_default_vault = CryptographicPrivacyVault()
|
| 385 |
+
return _default_vault
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
# ── Standalone Demo ────────────────────────────────────────────────────────────
|
| 389 |
+
if __name__ == "__main__":
|
| 390 |
+
from presidio_analyzer import AnalyzerEngine
|
| 391 |
+
from presidio_anonymizer import AnonymizerEngine
|
| 392 |
+
from presidio_analyzer.nlp_engine import NlpEngineProvider
|
| 393 |
+
|
| 394 |
+
nlp_cfg = {"nlp_engine_name": "spacy", "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}]}
|
| 395 |
+
nlp_engine = NlpEngineProvider(nlp_configuration=nlp_cfg).create_engine()
|
| 396 |
+
analyzer = AnalyzerEngine(nlp_engine=nlp_engine)
|
| 397 |
+
anonymizer = AnonymizerEngine()
|
| 398 |
+
|
| 399 |
+
vault = CryptographicPrivacyVault()
|
| 400 |
+
tenant = "acme-corp"
|
| 401 |
+
ticket_text = (
|
| 402 |
+
"Hi, my name is John Smith and I can be reached at john.smith@acme.com "
|
| 403 |
+
"or call +1 212-555-9876. My credit card 4111-1111-1111-1111 was charged twice."
|
| 404 |
+
)
|
| 405 |
+
|
| 406 |
+
print("=" * 60)
|
| 407 |
+
print("CustomerCore Cryptographic Privacy Vault — Demo")
|
| 408 |
+
print("=" * 60)
|
| 409 |
+
print(f"\nOriginal text:\n {ticket_text}")
|
| 410 |
+
|
| 411 |
+
tokenized, count = vault.encrypt_text(
|
| 412 |
+
tenant_id=tenant,
|
| 413 |
+
text=ticket_text,
|
| 414 |
+
analyzer=analyzer,
|
| 415 |
+
anonymizer=anonymizer,
|
| 416 |
+
ticket_id="TKT-DEMO-001",
|
| 417 |
+
)
|
| 418 |
+
print(f"\nTokenized text ({count} PII entities detected):\n {tokenized}")
|
| 419 |
+
print(f"\nVault entries stored: {vault.store.count(tenant)}")
|
| 420 |
+
|
| 421 |
+
# Authorized decrypt
|
| 422 |
+
restored = vault.decrypt_text(
|
| 423 |
+
tenant_id=tenant,
|
| 424 |
+
text=tokenized,
|
| 425 |
+
actor_role="SUPPORT_LEAD",
|
| 426 |
+
ticket_id="TKT-DEMO-001",
|
| 427 |
+
)
|
| 428 |
+
print(f"\nRestored text (SUPPORT_LEAD view):\n {restored}")
|
| 429 |
+
|
| 430 |
+
# Unauthorized decrypt
|
| 431 |
+
print("\nAttempting decrypt with AGENT role (unauthorized)...")
|
| 432 |
+
try:
|
| 433 |
+
vault.decrypt_token(
|
| 434 |
+
tenant_id=tenant,
|
| 435 |
+
token=list(TOKEN_PATTERN.findall(tokenized))[0],
|
| 436 |
+
actor_role="AGENT",
|
| 437 |
+
)
|
| 438 |
+
except PermissionError as e:
|
| 439 |
+
print(f" BLOCKED: {e}")
|
| 440 |
+
|
| 441 |
+
print(f"\nAudit log written to: logs/vault_audit/audit_trail.jsonl")
|
| 442 |
+
print("=" * 60)
|
|
File without changes
|
|
@@ -0,0 +1,177 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/streaming/bronze_consumer.py
|
| 3 |
+
|
| 4 |
+
Reads from all 4 Redpanda topics and writes raw messages to MinIO
|
| 5 |
+
as the Bronze layer (unmodified, append-only Parquet files).
|
| 6 |
+
|
| 7 |
+
Bronze = raw data exactly as received. Nothing is changed.
|
| 8 |
+
If something goes wrong downstream, we always have the original.
|
| 9 |
+
|
| 10 |
+
Run: python -m src.streaming.bronze_consumer --batch-size 1000
|
| 11 |
+
Ctrl+C to stop gracefully.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
import argparse
|
| 15 |
+
import io
|
| 16 |
+
import json
|
| 17 |
+
import signal
|
| 18 |
+
import sys
|
| 19 |
+
import time
|
| 20 |
+
from collections import defaultdict
|
| 21 |
+
from datetime import datetime, timezone
|
| 22 |
+
|
| 23 |
+
import boto3
|
| 24 |
+
import pyarrow as pa
|
| 25 |
+
import pyarrow.parquet as pq
|
| 26 |
+
from botocore.client import Config
|
| 27 |
+
from confluent_kafka import Consumer, KafkaError
|
| 28 |
+
from tqdm import tqdm
|
| 29 |
+
|
| 30 |
+
# ── Config ────────────────────────────────────────────────────
|
| 31 |
+
BROKER = "localhost:9092"
|
| 32 |
+
GROUP_ID = "bronze-consumer-group"
|
| 33 |
+
TOPICS = ["support-tickets", "billing-events", "product-events", "incident-events"]
|
| 34 |
+
BUCKET = "customercore-lake"
|
| 35 |
+
TOPIC_TO_PREFIX = {
|
| 36 |
+
"support-tickets": "bronze/tickets",
|
| 37 |
+
"billing-events": "bronze/billing",
|
| 38 |
+
"product-events": "bronze/product",
|
| 39 |
+
"incident-events": "bronze/incidents",
|
| 40 |
+
}
|
| 41 |
+
|
| 42 |
+
MINIO_ENDPOINT = "http://localhost:9000"
|
| 43 |
+
MINIO_ACCESS_KEY = "minioadmin"
|
| 44 |
+
MINIO_SECRET_KEY = "minioadmin"
|
| 45 |
+
|
| 46 |
+
running = True
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def signal_handler(sig, frame):
|
| 50 |
+
global running
|
| 51 |
+
print("\n[STOP] Graceful shutdown triggered...")
|
| 52 |
+
running = False
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
signal.signal(signal.SIGINT, signal_handler)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def get_s3():
|
| 59 |
+
return boto3.client(
|
| 60 |
+
"s3",
|
| 61 |
+
endpoint_url=MINIO_ENDPOINT,
|
| 62 |
+
aws_access_key_id=MINIO_ACCESS_KEY,
|
| 63 |
+
aws_secret_access_key=MINIO_SECRET_KEY,
|
| 64 |
+
config=Config(signature_version="s3v4"),
|
| 65 |
+
region_name="us-east-1",
|
| 66 |
+
)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def write_batch_to_bronze(s3, topic: str, records: list[dict]) -> str:
|
| 70 |
+
"""Write a batch of raw records to MinIO as a Parquet file."""
|
| 71 |
+
prefix = TOPIC_TO_PREFIX[topic]
|
| 72 |
+
timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_%f")
|
| 73 |
+
key = f"{prefix}/batch_{timestamp}.parquet"
|
| 74 |
+
|
| 75 |
+
# Convert to Arrow table — all values as strings to preserve raw format
|
| 76 |
+
rows = [{"raw_json": json.dumps(r), "ingested_at": datetime.now(timezone.utc).isoformat()} for r in records]
|
| 77 |
+
table = pa.table({
|
| 78 |
+
"raw_json": pa.array([r["raw_json"] for r in rows], type=pa.string()),
|
| 79 |
+
"ingested_at": pa.array([r["ingested_at"] for r in rows], type=pa.string()),
|
| 80 |
+
})
|
| 81 |
+
|
| 82 |
+
buf = io.BytesIO()
|
| 83 |
+
pq.write_table(table, buf)
|
| 84 |
+
buf.seek(0)
|
| 85 |
+
s3.put_object(Bucket=BUCKET, Key=key, Body=buf.getvalue())
|
| 86 |
+
return key
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def main(batch_size: int = 500, timeout_seconds: int = 30):
|
| 90 |
+
print("=" * 60)
|
| 91 |
+
print("CustomerCore Bronze Consumer")
|
| 92 |
+
print(f"Topics : {', '.join(TOPICS)}")
|
| 93 |
+
print(f"Sink : MinIO s3://{BUCKET}/bronze/")
|
| 94 |
+
print(f"Batch : {batch_size} messages per Parquet file")
|
| 95 |
+
print("Press Ctrl+C to stop gracefully")
|
| 96 |
+
print("=" * 60)
|
| 97 |
+
|
| 98 |
+
consumer = Consumer({
|
| 99 |
+
"bootstrap.servers": BROKER,
|
| 100 |
+
"group.id": GROUP_ID,
|
| 101 |
+
"auto.offset.reset": "earliest",
|
| 102 |
+
"enable.auto.commit": True,
|
| 103 |
+
})
|
| 104 |
+
consumer.subscribe(TOPICS)
|
| 105 |
+
s3 = get_s3()
|
| 106 |
+
|
| 107 |
+
buffers: dict[str, list] = defaultdict(list)
|
| 108 |
+
total_written = 0
|
| 109 |
+
files_written = 0
|
| 110 |
+
start = time.time()
|
| 111 |
+
last_message_time = time.time()
|
| 112 |
+
|
| 113 |
+
print(f"\nListening for messages (will auto-stop after {timeout_seconds}s silence)...\n")
|
| 114 |
+
pbar = tqdm(unit="msg", desc="Consumed", bar_format="{desc}: {n_fmt} msgs | Files: {postfix}")
|
| 115 |
+
pbar.set_postfix_str("0")
|
| 116 |
+
|
| 117 |
+
try:
|
| 118 |
+
while running:
|
| 119 |
+
msg = consumer.poll(timeout=1.0)
|
| 120 |
+
|
| 121 |
+
# Auto-stop if silent for timeout_seconds
|
| 122 |
+
if time.time() - last_message_time > timeout_seconds:
|
| 123 |
+
print(f"\n[INFO] No messages for {timeout_seconds}s — flushing and stopping.")
|
| 124 |
+
break
|
| 125 |
+
|
| 126 |
+
if msg is None:
|
| 127 |
+
continue
|
| 128 |
+
if msg.error():
|
| 129 |
+
if msg.error().code() != KafkaError._PARTITION_EOF:
|
| 130 |
+
print(f"[ERROR] {msg.error()}")
|
| 131 |
+
continue
|
| 132 |
+
|
| 133 |
+
last_message_time = time.time()
|
| 134 |
+
topic = msg.topic()
|
| 135 |
+
try:
|
| 136 |
+
value = json.loads(msg.value().decode())
|
| 137 |
+
buffers[topic].append(value)
|
| 138 |
+
total_written += 1
|
| 139 |
+
pbar.update(1)
|
| 140 |
+
except json.JSONDecodeError:
|
| 141 |
+
continue
|
| 142 |
+
|
| 143 |
+
# Flush buffer when batch is full
|
| 144 |
+
if len(buffers[topic]) >= batch_size:
|
| 145 |
+
key = write_batch_to_bronze(s3, topic, buffers[topic])
|
| 146 |
+
files_written += 1
|
| 147 |
+
pbar.set_postfix_str(str(files_written))
|
| 148 |
+
buffers[topic] = []
|
| 149 |
+
|
| 150 |
+
# Flush remaining partial batches
|
| 151 |
+
print("\nFlushing remaining buffers...")
|
| 152 |
+
for topic, records in buffers.items():
|
| 153 |
+
if records:
|
| 154 |
+
key = write_batch_to_bronze(s3, topic, records)
|
| 155 |
+
files_written += 1
|
| 156 |
+
print(f" [FLUSHED] {len(records)} records -> {key}")
|
| 157 |
+
|
| 158 |
+
finally:
|
| 159 |
+
pbar.close()
|
| 160 |
+
consumer.close()
|
| 161 |
+
|
| 162 |
+
elapsed = time.time() - start
|
| 163 |
+
print(f"\n{'=' * 60}")
|
| 164 |
+
print(f" Total consumed : {total_written:,} messages")
|
| 165 |
+
print(f" Parquet files : {files_written}")
|
| 166 |
+
print(f" Duration : {elapsed:.1f}s")
|
| 167 |
+
print(f" Sink : s3://{BUCKET}/bronze/")
|
| 168 |
+
print(f" Browse MinIO : http://localhost:9001")
|
| 169 |
+
print(f"{'=' * 60}")
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
if __name__ == "__main__":
|
| 173 |
+
parser = argparse.ArgumentParser()
|
| 174 |
+
parser.add_argument("--batch-size", type=int, default=500, help="Messages per Parquet file")
|
| 175 |
+
parser.add_argument("--timeout", type=int, default=30, help="Stop after N seconds of silence")
|
| 176 |
+
args = parser.parse_args()
|
| 177 |
+
main(batch_size=args.batch_size, timeout_seconds=args.timeout)
|
|
@@ -0,0 +1,328 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/streaming/bronze_to_silver.py
|
| 3 |
+
|
| 4 |
+
PySpark job that reads Bronze Parquet files from MinIO,
|
| 5 |
+
applies PII masking via Presidio, cleans and validates records,
|
| 6 |
+
and writes to the Silver layer.
|
| 7 |
+
|
| 8 |
+
Silver = cleaned, validated, PII-masked. Safe to query and use for ML.
|
| 9 |
+
|
| 10 |
+
PII entities masked: EMAIL_ADDRESS, PHONE_NUMBER, CREDIT_CARD,
|
| 11 |
+
PERSON, IBAN_CODE, IP_ADDRESS
|
| 12 |
+
|
| 13 |
+
Run: python -m src.streaming.bronze_to_silver
|
| 14 |
+
python -m src.streaming.bronze_to_silver --source tickets --limit 500
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
import argparse
|
| 18 |
+
import io
|
| 19 |
+
import json
|
| 20 |
+
import time
|
| 21 |
+
from datetime import datetime, timezone
|
| 22 |
+
from typing import Optional
|
| 23 |
+
|
| 24 |
+
import boto3
|
| 25 |
+
import pyarrow as pa
|
| 26 |
+
import pyarrow.parquet as pq
|
| 27 |
+
from botocore.client import Config
|
| 28 |
+
from presidio_analyzer import AnalyzerEngine
|
| 29 |
+
from presidio_anonymizer import AnonymizerEngine
|
| 30 |
+
from tqdm import tqdm
|
| 31 |
+
|
| 32 |
+
# Optional vault import — gracefully degrade if not configured
|
| 33 |
+
try:
|
| 34 |
+
from src.responsible_ai.privacy_vault import CryptographicPrivacyVault
|
| 35 |
+
_VAULT_AVAILABLE = True
|
| 36 |
+
except ImportError:
|
| 37 |
+
_VAULT_AVAILABLE = False
|
| 38 |
+
CryptographicPrivacyVault = None
|
| 39 |
+
|
| 40 |
+
# ── Config ────────────────────────────────────────────────────
|
| 41 |
+
MINIO_ENDPOINT = "http://localhost:9000"
|
| 42 |
+
MINIO_ACCESS_KEY = "minioadmin"
|
| 43 |
+
MINIO_SECRET_KEY = "minioadmin"
|
| 44 |
+
BUCKET = "customercore-lake"
|
| 45 |
+
|
| 46 |
+
SOURCE_MAP = {
|
| 47 |
+
"tickets": ("bronze/tickets", "silver/tickets"),
|
| 48 |
+
"billing": ("bronze/billing", "silver/billing"),
|
| 49 |
+
"product": ("bronze/product", "silver/product"),
|
| 50 |
+
"incidents": ("bronze/incidents", "silver/incidents"),
|
| 51 |
+
}
|
| 52 |
+
|
| 53 |
+
PII_ENTITIES = [
|
| 54 |
+
"EMAIL_ADDRESS", "PHONE_NUMBER", "CREDIT_CARD",
|
| 55 |
+
"PERSON", "IBAN_CODE", "IP_ADDRESS", "US_SSN",
|
| 56 |
+
]
|
| 57 |
+
|
| 58 |
+
VALID_PRIORITIES = {"low", "medium", "high", "critical"}
|
| 59 |
+
VALID_TIERS = {"enterprise", "professional", "free"}
|
| 60 |
+
VALID_CHANNELS = {"email", "web", "api", "chat"}
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def get_s3():
|
| 64 |
+
return boto3.client(
|
| 65 |
+
"s3",
|
| 66 |
+
endpoint_url=MINIO_ENDPOINT,
|
| 67 |
+
aws_access_key_id=MINIO_ACCESS_KEY,
|
| 68 |
+
aws_secret_access_key=MINIO_SECRET_KEY,
|
| 69 |
+
config=Config(signature_version="s3v4"),
|
| 70 |
+
region_name="us-east-1",
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def list_bronze_files(s3, prefix: str) -> list[str]:
|
| 75 |
+
response = s3.list_objects_v2(Bucket=BUCKET, Prefix=prefix)
|
| 76 |
+
return [obj["Key"] for obj in response.get("Contents", []) if obj["Key"].endswith(".parquet")]
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def read_parquet_from_minio(s3, key: str) -> list[dict]:
|
| 80 |
+
obj = s3.get_object(Bucket=BUCKET, Key=key)
|
| 81 |
+
buf = io.BytesIO(obj["Body"].read())
|
| 82 |
+
table = pq.read_table(buf)
|
| 83 |
+
records = []
|
| 84 |
+
for row in table.to_pydict()["raw_json"]:
|
| 85 |
+
try:
|
| 86 |
+
records.append(json.loads(row))
|
| 87 |
+
except json.JSONDecodeError:
|
| 88 |
+
pass
|
| 89 |
+
return records
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
def mask_pii(text: str, analyzer: AnalyzerEngine, anonymizer: AnonymizerEngine) -> str:
|
| 93 |
+
"""Run Presidio PII detection and replace with entity type placeholder (legacy path)."""
|
| 94 |
+
if not text or not isinstance(text, str):
|
| 95 |
+
return text
|
| 96 |
+
results = analyzer.analyze(text=text, language="en", entities=PII_ENTITIES)
|
| 97 |
+
if not results:
|
| 98 |
+
return text
|
| 99 |
+
anonymized = anonymizer.anonymize(text=text, analyzer_results=results)
|
| 100 |
+
return anonymized.text
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def vault_mask_pii(
|
| 104 |
+
text: str,
|
| 105 |
+
analyzer: AnalyzerEngine,
|
| 106 |
+
anonymizer: AnonymizerEngine,
|
| 107 |
+
vault: "CryptographicPrivacyVault",
|
| 108 |
+
tenant_id: str,
|
| 109 |
+
ticket_id: str = "",
|
| 110 |
+
) -> tuple[str, int]:
|
| 111 |
+
"""
|
| 112 |
+
Vault-backed PII protection (upgraded path).
|
| 113 |
+
Encrypts each PII span with the tenant's AES-256 key and replaces it
|
| 114 |
+
with a stable, readable token in the text. Returns (tokenized_text, pii_count).
|
| 115 |
+
Falls back to plain masking if vault is None.
|
| 116 |
+
"""
|
| 117 |
+
if vault is None:
|
| 118 |
+
return mask_pii(text, analyzer, anonymizer), 0
|
| 119 |
+
return vault.encrypt_text(
|
| 120 |
+
tenant_id=tenant_id,
|
| 121 |
+
text=text,
|
| 122 |
+
analyzer=analyzer,
|
| 123 |
+
anonymizer=anonymizer,
|
| 124 |
+
ticket_id=ticket_id,
|
| 125 |
+
pii_entities=PII_ENTITIES,
|
| 126 |
+
)
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def validate_and_clean_ticket(
|
| 130 |
+
record: dict,
|
| 131 |
+
analyzer,
|
| 132 |
+
anonymizer,
|
| 133 |
+
vault: Optional["CryptographicPrivacyVault"] = None,
|
| 134 |
+
) -> dict | None:
|
| 135 |
+
"""
|
| 136 |
+
Clean, validate, and PII-protect a ticket record.
|
| 137 |
+
|
| 138 |
+
When a CryptographicPrivacyVault is provided (Phase 5+):
|
| 139 |
+
- PII is encrypted with tenant-specific AES-256 and replaced by stable tokens
|
| 140 |
+
- Tokens are stored in the vault for authorized re-identification
|
| 141 |
+
- silver_version is bumped to '2.0' to signal vault-protected records
|
| 142 |
+
|
| 143 |
+
Without a vault (Phase 3 legacy / backward-compatible path):
|
| 144 |
+
- PII is replaced by [ENTITY_TYPE] placeholders (permanent deletion)
|
| 145 |
+
- silver_version remains '1.0'
|
| 146 |
+
|
| 147 |
+
Returns None if the record is missing required fields.
|
| 148 |
+
"""
|
| 149 |
+
# Required fields
|
| 150 |
+
if not record.get("event_id") or not record.get("tenant_id"):
|
| 151 |
+
return None
|
| 152 |
+
if not record.get("subject") and not record.get("body"):
|
| 153 |
+
return None
|
| 154 |
+
|
| 155 |
+
# Normalize fields
|
| 156 |
+
priority = str(record.get("priority", "medium")).lower()
|
| 157 |
+
tier = str(record.get("customer_tier", "free")).lower()
|
| 158 |
+
channel = str(record.get("channel", "web")).lower()
|
| 159 |
+
tenant_id = record["tenant_id"]
|
| 160 |
+
ticket_id = record.get("ticket_id", "")
|
| 161 |
+
|
| 162 |
+
# PII protection — vault path (Phase 5) or legacy masking path (Phase 3)
|
| 163 |
+
if vault is not None:
|
| 164 |
+
subject_clean, _ = vault_mask_pii(
|
| 165 |
+
str(record.get("subject", ""))[:256], analyzer, anonymizer,
|
| 166 |
+
vault, tenant_id, ticket_id,
|
| 167 |
+
)
|
| 168 |
+
body_clean, pii_count = vault_mask_pii(
|
| 169 |
+
str(record.get("body", "")), analyzer, anonymizer,
|
| 170 |
+
vault, tenant_id, ticket_id,
|
| 171 |
+
)
|
| 172 |
+
silver_version = "2.0" # vault-protected
|
| 173 |
+
else:
|
| 174 |
+
subject_clean = mask_pii(str(record.get("subject", ""))[:256], analyzer, anonymizer)
|
| 175 |
+
body_clean = mask_pii(str(record.get("body", "")), analyzer, anonymizer)
|
| 176 |
+
silver_version = "1.0" # legacy masking
|
| 177 |
+
|
| 178 |
+
return {
|
| 179 |
+
"event_id": record["event_id"],
|
| 180 |
+
"event_type": record.get("event_type", "support_ticket_created"),
|
| 181 |
+
"tenant_id": tenant_id,
|
| 182 |
+
"ticket_id": ticket_id,
|
| 183 |
+
"customer_id": record.get("customer_id", ""),
|
| 184 |
+
"customer_tier": tier if tier in VALID_TIERS else "free",
|
| 185 |
+
"subject": subject_clean,
|
| 186 |
+
"body": body_clean,
|
| 187 |
+
"category": str(record.get("category", "general")).lower(),
|
| 188 |
+
"priority": priority if priority in VALID_PRIORITIES else "medium",
|
| 189 |
+
"channel": channel if channel in VALID_CHANNELS else "web",
|
| 190 |
+
"reopen_count": max(0, int(record.get("reopen_count", 0))),
|
| 191 |
+
"tags": record.get("tags", []),
|
| 192 |
+
"original_timestamp": record.get("timestamp", ""),
|
| 193 |
+
"processed_at": datetime.now(timezone.utc).isoformat(),
|
| 194 |
+
"pii_masked": True,
|
| 195 |
+
"vault_protected": vault is not None,
|
| 196 |
+
"silver_version": silver_version,
|
| 197 |
+
}
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def write_silver(s3, prefix: str, records: list[dict], batch_id: str):
|
| 201 |
+
"""Write cleaned records to Silver layer as Parquet."""
|
| 202 |
+
if not records:
|
| 203 |
+
return
|
| 204 |
+
|
| 205 |
+
# Build columnar data
|
| 206 |
+
columns = list(records[0].keys())
|
| 207 |
+
data = {col: [str(r.get(col, "")) for r in records] for col in columns}
|
| 208 |
+
# reopen_count is int
|
| 209 |
+
data["reopen_count"] = [int(r.get("reopen_count", 0)) for r in records]
|
| 210 |
+
|
| 211 |
+
table = pa.table({k: pa.array(v) for k, v in data.items()})
|
| 212 |
+
buf = io.BytesIO()
|
| 213 |
+
pq.write_table(table, buf, compression="snappy")
|
| 214 |
+
buf.seek(0)
|
| 215 |
+
|
| 216 |
+
key = f"{prefix}/silver_batch_{batch_id}.parquet"
|
| 217 |
+
s3.put_object(Bucket=BUCKET, Key=key, Body=buf.getvalue())
|
| 218 |
+
return key
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def main(source: str = "tickets", limit: int = None, use_vault: bool = True):
|
| 222 |
+
bronze_prefix, silver_prefix = SOURCE_MAP.get(source, SOURCE_MAP["tickets"])
|
| 223 |
+
|
| 224 |
+
print("=" * 60)
|
| 225 |
+
print("CustomerCore Bronze -> Silver Pipeline")
|
| 226 |
+
print(f" Source : s3://{BUCKET}/{bronze_prefix}/")
|
| 227 |
+
print(f" Sink : s3://{BUCKET}/{silver_prefix}/")
|
| 228 |
+
print(f" PII masked: {', '.join(PII_ENTITIES)}")
|
| 229 |
+
print("=" * 60)
|
| 230 |
+
|
| 231 |
+
s3 = get_s3()
|
| 232 |
+
|
| 233 |
+
# ── 1. Discover Bronze files ──────────────────────────────
|
| 234 |
+
print("\n[1/5] Scanning Bronze layer for Parquet files...")
|
| 235 |
+
files = list_bronze_files(s3, bronze_prefix)
|
| 236 |
+
if not files:
|
| 237 |
+
print(" No Bronze files found. Run bronze_consumer.py first.")
|
| 238 |
+
return
|
| 239 |
+
print(f" Found {len(files)} Parquet file(s)")
|
| 240 |
+
|
| 241 |
+
# ── 2. Load all records ───────────────────────────────────
|
| 242 |
+
print("\n[2/5] Loading records from Bronze...")
|
| 243 |
+
t0 = time.time()
|
| 244 |
+
all_records = []
|
| 245 |
+
for f in tqdm(files, desc="Reading", unit="file"):
|
| 246 |
+
all_records.extend(read_parquet_from_minio(s3, f))
|
| 247 |
+
print(f" Loaded {len(all_records):,} raw records in {time.time()-t0:.1f}s")
|
| 248 |
+
|
| 249 |
+
if limit:
|
| 250 |
+
all_records = all_records[:limit]
|
| 251 |
+
print(f" Limited to {limit:,} records")
|
| 252 |
+
|
| 253 |
+
# ── 3. Initialize Presidio + Privacy Vault ───────────────────
|
| 254 |
+
# Using en_core_web_sm (12MB) — sufficient for PII detection.
|
| 255 |
+
# en_core_web_lg (400MB) is NOT needed and causes download failures.
|
| 256 |
+
print("\n[3/5] Initializing Presidio PII engine (en_core_web_sm)...")
|
| 257 |
+
t0 = time.time()
|
| 258 |
+
from presidio_analyzer.nlp_engine import NlpEngineProvider
|
| 259 |
+
nlp_config = {"nlp_engine_name": "spacy", "models": [{"lang_code": "en", "model_name": "en_core_web_sm"}]}
|
| 260 |
+
nlp_engine = NlpEngineProvider(nlp_configuration=nlp_config).create_engine()
|
| 261 |
+
analyzer = AnalyzerEngine(nlp_engine=nlp_engine)
|
| 262 |
+
anonymizer = AnonymizerEngine()
|
| 263 |
+
print(f" Ready in {time.time()-t0:.1f}s")
|
| 264 |
+
|
| 265 |
+
# Initialize vault (Phase 5 — Zero-Trust Cryptographic Privacy Vault)
|
| 266 |
+
vault = None
|
| 267 |
+
if use_vault and _VAULT_AVAILABLE:
|
| 268 |
+
from src.responsible_ai.privacy_vault import CryptographicPrivacyVault
|
| 269 |
+
vault = CryptographicPrivacyVault()
|
| 270 |
+
print(f" Privacy Vault: ENABLED (AES-256, silver_version=2.0)")
|
| 271 |
+
else:
|
| 272 |
+
print(f" Privacy Vault: DISABLED (legacy Presidio masking, silver_version=1.0)")
|
| 273 |
+
|
| 274 |
+
# ── 4. Clean + mask PII ───────────────────────────────────
|
| 275 |
+
print(f"\n[4/5] Cleaning and masking PII in {len(all_records):,} records...")
|
| 276 |
+
print(" Estimated time: ~1-3 minutes for 1000 records (Presidio NER scan)")
|
| 277 |
+
t0 = time.time()
|
| 278 |
+
clean_records = []
|
| 279 |
+
dropped = 0
|
| 280 |
+
|
| 281 |
+
with tqdm(total=len(all_records), unit="rec",
|
| 282 |
+
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]") as pbar:
|
| 283 |
+
for record in all_records:
|
| 284 |
+
if source == "tickets":
|
| 285 |
+
cleaned = validate_and_clean_ticket(record, analyzer, anonymizer, vault=vault)
|
| 286 |
+
else:
|
| 287 |
+
# For non-ticket sources: just add processing metadata, no PII masking needed
|
| 288 |
+
record["processed_at"] = datetime.now(timezone.utc).isoformat()
|
| 289 |
+
record["silver_version"] = "1.0"
|
| 290 |
+
record["vault_protected"] = False
|
| 291 |
+
cleaned = record
|
| 292 |
+
if cleaned:
|
| 293 |
+
clean_records.append(cleaned)
|
| 294 |
+
else:
|
| 295 |
+
dropped += 1
|
| 296 |
+
pbar.update(1)
|
| 297 |
+
|
| 298 |
+
elapsed = time.time() - t0
|
| 299 |
+
rate = len(all_records) / elapsed if elapsed > 0 else float(len(all_records))
|
| 300 |
+
print(f"\n Processed : {len(all_records):,} records in {elapsed:.1f}s ({rate:.0f} rec/s)")
|
| 301 |
+
print(f" Valid : {len(clean_records):,}")
|
| 302 |
+
print(f" Dropped : {dropped:,} (missing required fields)")
|
| 303 |
+
|
| 304 |
+
# ── 5. Write Silver ───────────────────────────────────────
|
| 305 |
+
print(f"\n[5/5] Writing {len(clean_records):,} records to Silver layer...")
|
| 306 |
+
batch_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
|
| 307 |
+
t0 = time.time()
|
| 308 |
+
key = write_silver(s3, silver_prefix, clean_records, batch_id)
|
| 309 |
+
print(f" Written in {time.time()-t0:.1f}s -> {key}")
|
| 310 |
+
|
| 311 |
+
print(f"\n{'=' * 60}")
|
| 312 |
+
print(f" Bronze records : {len(all_records):,}")
|
| 313 |
+
print(f" Silver records : {len(clean_records):,}")
|
| 314 |
+
print(f" Dropped : {dropped:,}")
|
| 315 |
+
print(f" PII masked : Yes ({', '.join(PII_ENTITIES)})")
|
| 316 |
+
print(f" Output : s3://{BUCKET}/{silver_prefix}/")
|
| 317 |
+
print(f"{'=' * 60}")
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
if __name__ == "__main__":
|
| 321 |
+
parser = argparse.ArgumentParser()
|
| 322 |
+
parser.add_argument("--source", choices=["tickets", "billing", "product", "incidents"],
|
| 323 |
+
default="tickets", help="Which Bronze topic to process")
|
| 324 |
+
parser.add_argument("--limit", type=int, default=None, help="Max records to process")
|
| 325 |
+
parser.add_argument("--no-vault", action="store_true",
|
| 326 |
+
help="Disable Privacy Vault and use legacy Presidio masking")
|
| 327 |
+
args = parser.parse_args()
|
| 328 |
+
main(source=args.source, limit=args.limit, use_vault=not args.no_vault)
|
|
@@ -0,0 +1,378 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/streaming/data_loader.py
|
| 3 |
+
|
| 4 |
+
Multi-Dataset CustomerCore Data Loader
|
| 5 |
+
Downloads multiple open-source customer support datasets from Hugging Face
|
| 6 |
+
and publishes enriched events to Redpanda topics with progress tracking.
|
| 7 |
+
|
| 8 |
+
== Datasets (all open-source, no scraping, no API keys, legal in EU/Germany) ==
|
| 9 |
+
1. bitext/Bitext-customer-support-llm-chatbot-training-dataset
|
| 10 |
+
26,872 rows | CDLA-Sharing-1.0 | SaaS customer support Q&A (English)
|
| 11 |
+
Categories: billing, account, orders, technical, etc.
|
| 12 |
+
|
| 13 |
+
2. bitext/Bitext-retail-banking-llm-chatbot-training-dataset
|
| 14 |
+
25,545 rows | CDLA-Sharing-1.0 | B2B banking/financial support Q&A (English)
|
| 15 |
+
Categories: card, loan, account, transfer, compliance, etc.
|
| 16 |
+
|
| 17 |
+
3. mteb/amazon_massive_intent [de, fr, es]
|
| 18 |
+
11,514 rows x 3 languages = 34,542 rows | Apache 2.0
|
| 19 |
+
Real customer voice assistant intent utterances in German, French, Spanish
|
| 20 |
+
Intents map to: alarm, calendar, email, audio, transport, shopping, etc.
|
| 21 |
+
Source: Amazon MASSIVE (Fitzgerald et al., 2022)
|
| 22 |
+
|
| 23 |
+
Total: ~87,000 rows across 4 languages (en, de, fr, es)
|
| 24 |
+
|
| 25 |
+
== Multi-Language Strategy ==
|
| 26 |
+
A B2B SaaS platform serving EU customers receives tickets in German, French,
|
| 27 |
+
Spanish, Portuguese etc. Single-language support degrades classification accuracy
|
| 28 |
+
by 15-40% and is not acceptable for GDPR-compliant EU platforms.
|
| 29 |
+
MASSIVE gives us real customer utterances in each language — not translations.
|
| 30 |
+
|
| 31 |
+
== Run ==
|
| 32 |
+
python -m src.streaming.data_loader # all sources (default)
|
| 33 |
+
python -m src.streaming.data_loader --sources customer # SaaS English only
|
| 34 |
+
python -m src.streaming.data_loader --sources banking # Banking English only
|
| 35 |
+
python -m src.streaming.data_loader --sources massive # Multilingual only
|
| 36 |
+
python -m src.streaming.data_loader --sources all # Everything ~87k rows
|
| 37 |
+
python -m src.streaming.data_loader --limit 500 --sources all # Quick test
|
| 38 |
+
"""
|
| 39 |
+
|
| 40 |
+
import argparse
|
| 41 |
+
import json
|
| 42 |
+
import random
|
| 43 |
+
import time
|
| 44 |
+
import uuid
|
| 45 |
+
from datetime import datetime, timezone
|
| 46 |
+
|
| 47 |
+
from confluent_kafka import Producer
|
| 48 |
+
from datasets import load_dataset, concatenate_datasets
|
| 49 |
+
from tqdm import tqdm
|
| 50 |
+
|
| 51 |
+
# ── Config ────────────────────────────────────────────────────
|
| 52 |
+
BROKER = "localhost:9092"
|
| 53 |
+
TOPIC = "support-tickets"
|
| 54 |
+
|
| 55 |
+
DATASET_CONFIGS = {
|
| 56 |
+
"customer": {
|
| 57 |
+
"id": "bitext/Bitext-customer-support-llm-chatbot-training-dataset",
|
| 58 |
+
"license": "CDLA-Sharing-1.0",
|
| 59 |
+
"domain": "saas_support",
|
| 60 |
+
"text_field": "instruction",
|
| 61 |
+
"response_field": "response",
|
| 62 |
+
"category_field": "category",
|
| 63 |
+
"tags_field": "tags",
|
| 64 |
+
"language": "en",
|
| 65 |
+
},
|
| 66 |
+
"banking": {
|
| 67 |
+
"id": "bitext/Bitext-retail-banking-llm-chatbot-training-dataset",
|
| 68 |
+
"license": "CDLA-Sharing-1.0",
|
| 69 |
+
"domain": "financial_support",
|
| 70 |
+
"text_field": "instruction",
|
| 71 |
+
"response_field": "response",
|
| 72 |
+
"category_field": "category",
|
| 73 |
+
"tags_field": "tags",
|
| 74 |
+
"language": "en",
|
| 75 |
+
},
|
| 76 |
+
}
|
| 77 |
+
|
| 78 |
+
# Multilingual MASSIVE configs (per language — Apache 2.0, no scraping)
|
| 79 |
+
MASSIVE_LANGUAGES = {
|
| 80 |
+
"de": {"name": "German", "mteb_config": "de"},
|
| 81 |
+
"fr": {"name": "French", "mteb_config": "fr"},
|
| 82 |
+
"es": {"name": "Spanish", "mteb_config": "es"},
|
| 83 |
+
}
|
| 84 |
+
|
| 85 |
+
# ── Enrichment pools (makes data multi-tenant) ────────────────
|
| 86 |
+
TENANTS = ["acme-corp", "globex-inc", "initech-ltd", "umbrella-co", "stark-tech"]
|
| 87 |
+
TIERS = ["enterprise", "professional", "free"]
|
| 88 |
+
CHANNELS = ["email", "web", "api", "chat"]
|
| 89 |
+
|
| 90 |
+
PRIORITY_MAP = {
|
| 91 |
+
# SaaS support categories
|
| 92 |
+
"ACCOUNT": ["medium", "high"],
|
| 93 |
+
"BILLING": ["high", "high", "critical"],
|
| 94 |
+
"CANCEL": ["high", "critical"],
|
| 95 |
+
"CONTACT": ["low", "medium"],
|
| 96 |
+
"DELIVERY": ["medium", "high"],
|
| 97 |
+
"FEEDBACK": ["low", "medium"],
|
| 98 |
+
"INVOICE": ["high", "critical"],
|
| 99 |
+
"NEWSLETTER": ["low"],
|
| 100 |
+
"ORDER": ["medium", "high"],
|
| 101 |
+
"PAYMENT": ["high", "critical"],
|
| 102 |
+
"REFUND": ["high", "high", "critical"],
|
| 103 |
+
"REVIEW": ["low", "medium"],
|
| 104 |
+
"SHIPPING": ["medium", "high"],
|
| 105 |
+
"SUBSCRIPTION": ["medium", "high"],
|
| 106 |
+
"SUPPORT": ["medium", "high"],
|
| 107 |
+
# Banking categories
|
| 108 |
+
"CARD": ["high", "critical"],
|
| 109 |
+
"LOAN": ["medium", "high"],
|
| 110 |
+
"TRANSFER": ["high", "critical"],
|
| 111 |
+
"COMPLIANCE": ["high", "critical"],
|
| 112 |
+
"MORTGAGE": ["medium", "high"],
|
| 113 |
+
"SAVINGS": ["low", "medium"],
|
| 114 |
+
"FRAUD": ["critical", "critical"],
|
| 115 |
+
}
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def delivery_report(err, msg):
|
| 119 |
+
if err is not None:
|
| 120 |
+
print(f"\n [ERROR] Delivery failed: {err}")
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def enrich(row: dict, domain: str, license_str: str) -> dict:
|
| 124 |
+
"""
|
| 125 |
+
Normalize a raw Bitext row into a CustomerCore ticket event.
|
| 126 |
+
Both Bitext datasets share the same column schema (instruction / response / category).
|
| 127 |
+
Domain tag differentiates SaaS support from financial support in downstream analysis.
|
| 128 |
+
"""
|
| 129 |
+
category = str(row.get("category", "SUPPORT")).upper()
|
| 130 |
+
tier = random.choice(TIERS)
|
| 131 |
+
priority_pool = PRIORITY_MAP.get(category[:8], ["low", "medium", "high"])
|
| 132 |
+
if tier == "enterprise":
|
| 133 |
+
priority_pool = [p for p in priority_pool if p in ("high", "critical")] or ["high"]
|
| 134 |
+
|
| 135 |
+
tags_raw = row.get("tags", "")
|
| 136 |
+
tags = [t.strip() for t in str(tags_raw).split(",") if t.strip()] if tags_raw else []
|
| 137 |
+
|
| 138 |
+
return {
|
| 139 |
+
"event_id": str(uuid.uuid4()),
|
| 140 |
+
"event_type": "support_ticket_created",
|
| 141 |
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
| 142 |
+
"tenant_id": random.choice(TENANTS),
|
| 143 |
+
"ticket_id": f"TKT-{random.randint(10000, 99999)}",
|
| 144 |
+
"customer_id": f"CUST-{random.randint(1000, 99999)}",
|
| 145 |
+
"customer_tier": tier,
|
| 146 |
+
"subject": str(row.get("instruction", ""))[:120],
|
| 147 |
+
"body": str(row.get("instruction", "")),
|
| 148 |
+
"suggested_response": str(row.get("response", "")),
|
| 149 |
+
"category": category.lower(),
|
| 150 |
+
"priority": random.choice(priority_pool),
|
| 151 |
+
"channel": random.choice(CHANNELS),
|
| 152 |
+
"reopen_count": random.randint(0, 2),
|
| 153 |
+
"tags": tags,
|
| 154 |
+
"source_domain": domain,
|
| 155 |
+
"source": f"huggingface:bitext-{domain}",
|
| 156 |
+
"license": license_str,
|
| 157 |
+
"language": "en",
|
| 158 |
+
"is_multilingual": False,
|
| 159 |
+
}
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
def enrich_massive(row: dict, language: str, lang_name: str) -> dict:
|
| 163 |
+
"""
|
| 164 |
+
Normalize an Amazon MASSIVE intent row into a CustomerCore ticket event.
|
| 165 |
+
|
| 166 |
+
MASSIVE rows have: id, label, label_text, text, lang
|
| 167 |
+
label_text = intent name (e.g. 'alarm_set', 'email_query', 'calendar_remove')
|
| 168 |
+
text = the customer utterance in the target language
|
| 169 |
+
|
| 170 |
+
Maps MASSIVE intents to support categories:
|
| 171 |
+
alarm/reminder/calendar -> account
|
| 172 |
+
email/messaging -> account
|
| 173 |
+
audio/music/podcast -> product
|
| 174 |
+
transport/travel -> order
|
| 175 |
+
shopping -> order
|
| 176 |
+
iot/smart_home -> technical
|
| 177 |
+
weather/datetime/news -> general
|
| 178 |
+
"""
|
| 179 |
+
INTENT_TO_CATEGORY = {
|
| 180 |
+
"alarm": "account", "reminder": "account", "calendar": "account",
|
| 181 |
+
"email": "account", "messaging": "account", "social": "account",
|
| 182 |
+
"audio": "technical", "music": "technical", "podcast": "technical",
|
| 183 |
+
"play": "technical", "iot": "technical", "smart": "technical",
|
| 184 |
+
"transport": "order", "travel": "order", "lists": "order",
|
| 185 |
+
"shopping": "order", "takeaway": "order", "cooking": "order",
|
| 186 |
+
"weather": "general", "datetime": "general", "news": "general",
|
| 187 |
+
"qa": "general", "recommendation": "general", "general": "general",
|
| 188 |
+
}
|
| 189 |
+
|
| 190 |
+
intent = str(row.get("label_text", "general")).lower()
|
| 191 |
+
category = "general"
|
| 192 |
+
for key, cat in INTENT_TO_CATEGORY.items():
|
| 193 |
+
if key in intent:
|
| 194 |
+
category = cat
|
| 195 |
+
break
|
| 196 |
+
|
| 197 |
+
tier = random.choice(TIERS)
|
| 198 |
+
priority_pool = ["low", "medium", "medium", "high"]
|
| 199 |
+
|
| 200 |
+
return {
|
| 201 |
+
"event_id": str(uuid.uuid4()),
|
| 202 |
+
"event_type": "support_ticket_created",
|
| 203 |
+
"timestamp": datetime.now(timezone.utc).isoformat(),
|
| 204 |
+
"tenant_id": random.choice(TENANTS),
|
| 205 |
+
"ticket_id": f"TKT-{random.randint(10000, 99999)}",
|
| 206 |
+
"customer_id": f"CUST-{random.randint(1000, 99999)}",
|
| 207 |
+
"customer_tier": tier,
|
| 208 |
+
"subject": str(row.get("text", ""))[:120],
|
| 209 |
+
"body": str(row.get("text", "")),
|
| 210 |
+
"suggested_response": "",
|
| 211 |
+
"category": category,
|
| 212 |
+
"priority": random.choice(priority_pool),
|
| 213 |
+
"channel": random.choice(CHANNELS),
|
| 214 |
+
"reopen_count": 0,
|
| 215 |
+
"tags": [intent.replace("_", "-")],
|
| 216 |
+
"source_domain": "multilingual_intent",
|
| 217 |
+
"source": f"huggingface:amazon-massive-{language}",
|
| 218 |
+
"license": "Apache-2.0",
|
| 219 |
+
"language": language,
|
| 220 |
+
"language_name": lang_name,
|
| 221 |
+
"is_multilingual": True,
|
| 222 |
+
"intent": intent,
|
| 223 |
+
}
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def load_sources(sources: str) -> tuple[list, dict]:
|
| 227 |
+
"""
|
| 228 |
+
Download and concatenate the requested dataset sources.
|
| 229 |
+
Returns (list_of_(row, enrich_fn) tuples, dict of source->count).
|
| 230 |
+
"""
|
| 231 |
+
all_rows = []
|
| 232 |
+
total_per_source = {}
|
| 233 |
+
|
| 234 |
+
bitext_sources = []
|
| 235 |
+
load_massive = False
|
| 236 |
+
|
| 237 |
+
if sources == "all":
|
| 238 |
+
bitext_sources = list(DATASET_CONFIGS.items())
|
| 239 |
+
load_massive = True
|
| 240 |
+
elif sources == "massive":
|
| 241 |
+
load_massive = True
|
| 242 |
+
elif sources in DATASET_CONFIGS:
|
| 243 |
+
bitext_sources = [(sources, DATASET_CONFIGS[sources])]
|
| 244 |
+
else:
|
| 245 |
+
raise ValueError(
|
| 246 |
+
f"Unknown source '{sources}'. "
|
| 247 |
+
f"Choose from: {list(DATASET_CONFIGS.keys()) + ['massive', 'all']}"
|
| 248 |
+
)
|
| 249 |
+
|
| 250 |
+
# Load Bitext English datasets
|
| 251 |
+
for name, cfg in bitext_sources:
|
| 252 |
+
print(f"\n Downloading: {cfg['id']}")
|
| 253 |
+
print(f" License : {cfg['license']} | Language: English")
|
| 254 |
+
t0 = time.time()
|
| 255 |
+
ds = load_dataset(cfg["id"], split="train")
|
| 256 |
+
elapsed = time.time() - t0
|
| 257 |
+
print(f" Downloaded : {len(ds):,} rows in {elapsed:.1f}s")
|
| 258 |
+
total_per_source[name] = len(ds)
|
| 259 |
+
for row in ds:
|
| 260 |
+
all_rows.append((row, cfg["domain"], cfg["license"], "bitext", "en", "English"))
|
| 261 |
+
|
| 262 |
+
# Load multilingual MASSIVE datasets
|
| 263 |
+
if load_massive:
|
| 264 |
+
for lang_code, lang_cfg in MASSIVE_LANGUAGES.items():
|
| 265 |
+
lang_name = lang_cfg["name"]
|
| 266 |
+
mteb_config = lang_cfg["mteb_config"]
|
| 267 |
+
print(f"\n Downloading: mteb/amazon_massive_intent [{lang_code} - {lang_name}]")
|
| 268 |
+
print(f" License : Apache-2.0 | Language: {lang_name}")
|
| 269 |
+
t0 = time.time()
|
| 270 |
+
ds = load_dataset("mteb/amazon_massive_intent", mteb_config, split="train")
|
| 271 |
+
elapsed = time.time() - t0
|
| 272 |
+
print(f" Downloaded : {len(ds):,} rows in {elapsed:.1f}s")
|
| 273 |
+
total_per_source[f"massive_{lang_code}"] = len(ds)
|
| 274 |
+
for row in ds:
|
| 275 |
+
all_rows.append((row, "multilingual_intent", "Apache-2.0", "massive", lang_code, lang_name))
|
| 276 |
+
|
| 277 |
+
return all_rows, total_per_source
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
def main(limit: int = None, delay: float = 0.0, sources: str = "all"):
|
| 281 |
+
print("=" * 65)
|
| 282 |
+
print("CustomerCore Multi-Language Data Loader")
|
| 283 |
+
print("Sources: Bitext SaaS + Banking (EN) + Amazon MASSIVE (DE/FR/ES)")
|
| 284 |
+
print("License: CDLA-Sharing-1.0 + Apache-2.0 | Legal in EU/Germany")
|
| 285 |
+
print("=" * 65)
|
| 286 |
+
|
| 287 |
+
# ── 1. Download datasets ──────────────────────────────────
|
| 288 |
+
print("\n[1/3] Downloading datasets from Hugging Face...")
|
| 289 |
+
t_start = time.time()
|
| 290 |
+
all_rows, source_counts = load_sources(sources)
|
| 291 |
+
|
| 292 |
+
print(f"\n Source breakdown:")
|
| 293 |
+
for src, count in source_counts.items():
|
| 294 |
+
print(f" {src:18s}: {count:,} rows")
|
| 295 |
+
total = len(all_rows)
|
| 296 |
+
print(f" {'TOTAL':18s}: {total:,} rows")
|
| 297 |
+
|
| 298 |
+
if limit:
|
| 299 |
+
random.shuffle(all_rows)
|
| 300 |
+
all_rows = all_rows[:limit]
|
| 301 |
+
print(f"\n Limited to {len(all_rows):,} rows for this run")
|
| 302 |
+
|
| 303 |
+
# ── 2. Connect to Redpanda ────────────────────────────────
|
| 304 |
+
print(f"\n[2/3] Connecting to Redpanda at {BROKER}...")
|
| 305 |
+
producer = Producer({
|
| 306 |
+
"bootstrap.servers": BROKER,
|
| 307 |
+
"queue.buffering.max.messages": 10000,
|
| 308 |
+
"batch.num.messages": 500,
|
| 309 |
+
})
|
| 310 |
+
print(" Connected.")
|
| 311 |
+
|
| 312 |
+
# ── 3. Publish with progress bar ──────────────────────────
|
| 313 |
+
print(f"\n[3/3] Publishing {len(all_rows):,} events to '{TOPIC}'...")
|
| 314 |
+
print(f" Estimated time: {len(all_rows) * 0.001:.0f}–{len(all_rows) * 0.003:.0f} seconds")
|
| 315 |
+
|
| 316 |
+
failed = 0
|
| 317 |
+
lang_counts: dict[str, int] = {}
|
| 318 |
+
t0 = time.time()
|
| 319 |
+
with tqdm(
|
| 320 |
+
total=len(all_rows),
|
| 321 |
+
unit="msg",
|
| 322 |
+
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
|
| 323 |
+
) as pbar:
|
| 324 |
+
# Each tuple: (row, domain, license_str, source_type, lang_code, lang_name)
|
| 325 |
+
for row, domain, license_str, source_type, lang_code, lang_name in all_rows:
|
| 326 |
+
if source_type == "bitext":
|
| 327 |
+
event = enrich(row, domain, license_str)
|
| 328 |
+
else: # massive
|
| 329 |
+
event = enrich_massive(row, lang_code, lang_name)
|
| 330 |
+
|
| 331 |
+
lang_counts[lang_code] = lang_counts.get(lang_code, 0) + 1
|
| 332 |
+
|
| 333 |
+
producer.produce(
|
| 334 |
+
topic=TOPIC,
|
| 335 |
+
key=event["ticket_id"].encode(),
|
| 336 |
+
value=json.dumps(event).encode(),
|
| 337 |
+
callback=lambda err, msg: None if not err else None,
|
| 338 |
+
)
|
| 339 |
+
producer.poll(0)
|
| 340 |
+
pbar.update(1)
|
| 341 |
+
if delay:
|
| 342 |
+
time.sleep(delay)
|
| 343 |
+
|
| 344 |
+
producer.flush()
|
| 345 |
+
elapsed = time.time() - t0
|
| 346 |
+
sent = len(all_rows) - failed
|
| 347 |
+
|
| 348 |
+
print(f"\n{'=' * 65}")
|
| 349 |
+
print(f" Published : {sent:,} events")
|
| 350 |
+
print(f" Failed : {failed}")
|
| 351 |
+
print(f" Duration : {elapsed:.1f}s ({len(all_rows)/elapsed:.0f} msg/s)")
|
| 352 |
+
print(f" Topic : {TOPIC}")
|
| 353 |
+
print(f" Sources : {sources}")
|
| 354 |
+
print(f" Language breakdown:")
|
| 355 |
+
for lang, count in sorted(lang_counts.items()):
|
| 356 |
+
print(f" {lang}: {count:,} rows")
|
| 357 |
+
print(f"{'=' * 65}")
|
| 358 |
+
|
| 359 |
+
|
| 360 |
+
if __name__ == "__main__":
|
| 361 |
+
parser = argparse.ArgumentParser(description="CustomerCore Multi-Language HuggingFace loader")
|
| 362 |
+
parser.add_argument(
|
| 363 |
+
"--limit", type=int, default=None,
|
| 364 |
+
help="Max rows to load (default: all ~87k combined)"
|
| 365 |
+
)
|
| 366 |
+
parser.add_argument(
|
| 367 |
+
"--delay", type=float, default=0.0,
|
| 368 |
+
help="Delay between messages in seconds"
|
| 369 |
+
)
|
| 370 |
+
parser.add_argument(
|
| 371 |
+
"--sources",
|
| 372 |
+
choices=["all", "customer", "banking", "massive"],
|
| 373 |
+
default="all",
|
| 374 |
+
help="Which dataset sources to load (default: all = EN+DE+FR+ES)"
|
| 375 |
+
)
|
| 376 |
+
args = parser.parse_args()
|
| 377 |
+
main(limit=args.limit, delay=args.delay, sources=args.sources)
|
| 378 |
+
|
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
src/streaming/minio_setup.py
|
| 3 |
+
|
| 4 |
+
Creates the MinIO bucket and folder structure for the lakehouse.
|
| 5 |
+
Run once before starting any pipeline.
|
| 6 |
+
|
| 7 |
+
Run: python -m src.streaming.minio_setup
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import boto3
|
| 11 |
+
from botocore.client import Config
|
| 12 |
+
from botocore.exceptions import ClientError
|
| 13 |
+
|
| 14 |
+
# ── MinIO connection ──────────────────────────────────────────
|
| 15 |
+
MINIO_ENDPOINT = "http://localhost:9000"
|
| 16 |
+
MINIO_ACCESS_KEY = "minioadmin"
|
| 17 |
+
MINIO_SECRET_KEY = "minioadmin"
|
| 18 |
+
BUCKET = "customercore-lake"
|
| 19 |
+
|
| 20 |
+
# ── Bronze/Silver/Gold prefix structure ───────────────────────
|
| 21 |
+
PREFIXES = [
|
| 22 |
+
"bronze/tickets/",
|
| 23 |
+
"bronze/billing/",
|
| 24 |
+
"bronze/product/",
|
| 25 |
+
"bronze/incidents/",
|
| 26 |
+
"silver/tickets/",
|
| 27 |
+
"silver/billing/",
|
| 28 |
+
"silver/product/",
|
| 29 |
+
"silver/incidents/",
|
| 30 |
+
"gold/",
|
| 31 |
+
]
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def get_client():
|
| 35 |
+
return boto3.client(
|
| 36 |
+
"s3",
|
| 37 |
+
endpoint_url=MINIO_ENDPOINT,
|
| 38 |
+
aws_access_key_id=MINIO_ACCESS_KEY,
|
| 39 |
+
aws_secret_access_key=MINIO_SECRET_KEY,
|
| 40 |
+
config=Config(signature_version="s3v4"),
|
| 41 |
+
region_name="us-east-1",
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def setup():
|
| 46 |
+
client = get_client()
|
| 47 |
+
|
| 48 |
+
# Create bucket if it doesn't exist
|
| 49 |
+
try:
|
| 50 |
+
client.head_bucket(Bucket=BUCKET)
|
| 51 |
+
print(f"[OK] Bucket '{BUCKET}' already exists")
|
| 52 |
+
except ClientError:
|
| 53 |
+
client.create_bucket(Bucket=BUCKET)
|
| 54 |
+
print(f"[CREATED] Bucket '{BUCKET}'")
|
| 55 |
+
|
| 56 |
+
# Create folder structure by writing empty marker objects
|
| 57 |
+
for prefix in PREFIXES:
|
| 58 |
+
client.put_object(Bucket=BUCKET, Key=prefix, Body=b"")
|
| 59 |
+
print(f" [OK] {BUCKET}/{prefix}")
|
| 60 |
+
|
| 61 |
+
print(f"\nLakehouse structure ready at MinIO: {MINIO_ENDPOINT}")
|
| 62 |
+
print(f"Browse at: http://localhost:9001 (user: minioadmin / pass: minioadmin)")
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
if __name__ == "__main__":
|
| 66 |
+
setup()
|