Saibalaji Namburi commited on
Commit
3cfefe7
·
0 Parent(s):

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)

This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .env.example +35 -0
  2. .gitattributes +23 -0
  3. .gitignore +49 -0
  4. BUILD_GUIDE.md +0 -0
  5. CustomerCore_Progress_Summary.md +0 -0
  6. docker-compose.yml +158 -0
  7. infra/otel-collector.yaml +49 -0
  8. pyproject.toml +27 -0
  9. requirements.txt +0 -0
  10. src/__init__.py +0 -0
  11. src/agent/memory.py +107 -0
  12. src/agent/nodes/churn_agent.py +161 -0
  13. src/agent/nodes/classify_agent.py +143 -0
  14. src/agent/nodes/hitl_agent.py +95 -0
  15. src/agent/nodes/incident_agent.py +69 -0
  16. src/agent/nodes/memory_agent.py +21 -0
  17. src/agent/nodes/rag_agent.py +198 -0
  18. src/agent/schemas.py +131 -0
  19. src/agent/state.py +38 -0
  20. src/agent/supervisor.py +199 -0
  21. src/dbt/.user.yml +1 -0
  22. src/dbt/dbt_project.yml +22 -0
  23. src/dbt/models/gold/billing_failure_summary.sql +12 -0
  24. src/dbt/models/gold/customer_health_daily.sql +58 -0
  25. src/dbt/models/gold/incident_severity_hourly.sql +11 -0
  26. src/dbt/models/gold/product_adoption_features.sql +11 -0
  27. src/dbt/models/gold/retention_cohort_metrics.sql +33 -0
  28. src/dbt/models/gold/stg_silver_events.sql +36 -0
  29. src/dbt/models/gold/support_agent_performance.sql +13 -0
  30. src/dbt/models/gold/ticket_funnel_daily.sql +13 -0
  31. src/dbt/models/schema.yml +101 -0
  32. src/dbt/models/sources.yml +10 -0
  33. src/dbt/package-lock.yml +5 -0
  34. src/dbt/packages.yml +3 -0
  35. src/dbt/profiles.yml +16 -0
  36. src/rag/__init__.py +14 -0
  37. src/rag/graph_rag.py +617 -0
  38. src/rag/hybrid_retriever.py +470 -0
  39. src/rag/llm_client.py +263 -0
  40. src/rag/multilingual.py +208 -0
  41. src/rag/router.py +377 -0
  42. src/responsible_ai/__init__.py +10 -0
  43. src/responsible_ai/audit_log.py +124 -0
  44. src/responsible_ai/key_manager.py +101 -0
  45. src/responsible_ai/privacy_vault.py +442 -0
  46. src/streaming/__init__.py +0 -0
  47. src/streaming/bronze_consumer.py +177 -0
  48. src/streaming/bronze_to_silver.py +328 -0
  49. src/streaming/data_loader.py +378 -0
  50. src/streaming/minio_setup.py +66 -0
.env.example ADDED
@@ -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
.gitattributes ADDED
@@ -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
.gitignore ADDED
@@ -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
BUILD_GUIDE.md ADDED
The diff for this file is too large to render. See raw diff
 
CustomerCore_Progress_Summary.md ADDED
The diff for this file is too large to render. See raw diff
 
docker-compose.yml ADDED
@@ -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
infra/otel-collector.yaml ADDED
@@ -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]
pyproject.toml ADDED
@@ -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
requirements.txt ADDED
Binary file (14.8 kB). View file
 
src/__init__.py ADDED
File without changes
src/agent/memory.py ADDED
@@ -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, [])
src/agent/nodes/churn_agent.py ADDED
@@ -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
+ }
src/agent/nodes/classify_agent.py ADDED
@@ -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
+ }
src/agent/nodes/hitl_agent.py ADDED
@@ -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
+ }
src/agent/nodes/incident_agent.py ADDED
@@ -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
+ }
src/agent/nodes/memory_agent.py ADDED
@@ -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
+ }
src/agent/nodes/rag_agent.py ADDED
@@ -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
+ }
src/agent/schemas.py ADDED
@@ -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
+ }
src/agent/state.py ADDED
@@ -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]
src/agent/supervisor.py ADDED
@@ -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")
src/dbt/.user.yml ADDED
@@ -0,0 +1 @@
 
 
1
+ id: e384ffcc-f3ea-4709-bda5-25a26b0c069b
src/dbt/dbt_project.yml ADDED
@@ -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
src/dbt/models/gold/billing_failure_summary.sql ADDED
@@ -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
src/dbt/models/gold/customer_health_daily.sql ADDED
@@ -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)
src/dbt/models/gold/incident_severity_hourly.sql ADDED
@@ -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
src/dbt/models/gold/product_adoption_features.sql ADDED
@@ -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
src/dbt/models/gold/retention_cohort_metrics.sql ADDED
@@ -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
src/dbt/models/gold/stg_silver_events.sql ADDED
@@ -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') }}
src/dbt/models/gold/support_agent_performance.sql ADDED
@@ -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
src/dbt/models/gold/ticket_funnel_daily.sql ADDED
@@ -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
src/dbt/models/schema.yml ADDED
@@ -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
src/dbt/models/sources.yml ADDED
@@ -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)"
src/dbt/package-lock.yml ADDED
@@ -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
src/dbt/packages.yml ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ packages:
2
+ - package: dbt-labs/dbt_utils
3
+ version: 1.3.0
src/dbt/profiles.yml ADDED
@@ -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"
src/rag/__init__.py ADDED
@@ -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
+ """
src/rag/graph_rag.py ADDED
@@ -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)
src/rag/hybrid_retriever.py ADDED
@@ -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)
src/rag/llm_client.py ADDED
@@ -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
src/rag/multilingual.py ADDED
@@ -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)
src/rag/router.py ADDED
@@ -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)
src/responsible_ai/__init__.py ADDED
@@ -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
+ """
src/responsible_ai/audit_log.py ADDED
@@ -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))}")
src/responsible_ai/key_manager.py ADDED
@@ -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}")
src/responsible_ai/privacy_vault.py ADDED
@@ -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)
src/streaming/__init__.py ADDED
File without changes
src/streaming/bronze_consumer.py ADDED
@@ -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)
src/streaming/bronze_to_silver.py ADDED
@@ -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)
src/streaming/data_loader.py ADDED
@@ -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
+
src/streaming/minio_setup.py ADDED
@@ -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()