nl2sql-engine 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (192) hide show
  1. nl2sql/__init__.py +38 -0
  2. nl2sql/adapters/__init__.py +0 -0
  3. nl2sql/adapters/duckdb/__init__.py +0 -0
  4. nl2sql/adapters/duckdb/adapter.py +71 -0
  5. nl2sql/adapters/mssql/__init__.py +0 -0
  6. nl2sql/adapters/mssql/adapter.py +122 -0
  7. nl2sql/adapters/mysql/__init__.py +0 -0
  8. nl2sql/adapters/mysql/adapter.py +123 -0
  9. nl2sql/adapters/postgres/__init__.py +0 -0
  10. nl2sql/adapters/postgres/adapter.py +115 -0
  11. nl2sql/adapters/sqlalchemy_base/__init__.py +17 -0
  12. nl2sql/adapters/sqlalchemy_base/adapter.py +476 -0
  13. nl2sql/adapters/sqlalchemy_base/models.py +36 -0
  14. nl2sql/adapters/sqlite/__init__.py +0 -0
  15. nl2sql/adapters/sqlite/adapter.py +88 -0
  16. nl2sql/aggregation/__init__.py +3 -0
  17. nl2sql/aggregation/aggregator.py +98 -0
  18. nl2sql/aggregation/engines/__init__.py +3 -0
  19. nl2sql/aggregation/engines/polars_duckdb.py +125 -0
  20. nl2sql/api/__init__.py +0 -0
  21. nl2sql/api/auth_api.py +60 -0
  22. nl2sql/api/benchmark_api.py +114 -0
  23. nl2sql/api/datasource_api.py +132 -0
  24. nl2sql/api/indexing_api.py +59 -0
  25. nl2sql/api/llm_api.py +82 -0
  26. nl2sql/api/policy_api.py +135 -0
  27. nl2sql/api/query_api.py +138 -0
  28. nl2sql/api/result_api.py +24 -0
  29. nl2sql/api/settings_api.py +65 -0
  30. nl2sql/auth/__init__.py +8 -0
  31. nl2sql/auth/models.py +36 -0
  32. nl2sql/auth/rbac.py +25 -0
  33. nl2sql/cli/__init__.py +0 -0
  34. nl2sql/cli/checks.py +53 -0
  35. nl2sql/cli/commands/__init__.py +0 -0
  36. nl2sql/cli/commands/benchmark.py +34 -0
  37. nl2sql/cli/commands/doctor.py +49 -0
  38. nl2sql/cli/commands/indexing.py +126 -0
  39. nl2sql/cli/commands/info.py +25 -0
  40. nl2sql/cli/commands/install.py +27 -0
  41. nl2sql/cli/commands/policy.py +57 -0
  42. nl2sql/cli/commands/run.py +166 -0
  43. nl2sql/cli/commands/setup.py +415 -0
  44. nl2sql/cli/commands/visualize.py +34 -0
  45. nl2sql/cli/common/decorators.py +34 -0
  46. nl2sql/cli/config.py +24 -0
  47. nl2sql/cli/console.py +52 -0
  48. nl2sql/cli/demo/__init__.py +1 -0
  49. nl2sql/cli/demo/data.py +87 -0
  50. nl2sql/cli/demo/defaults.py +122 -0
  51. nl2sql/cli/demo/factory.py +289 -0
  52. nl2sql/cli/demo/manager.py +230 -0
  53. nl2sql/cli/demo/schemas.py +336 -0
  54. nl2sql/cli/demo/writers/__init__.py +0 -0
  55. nl2sql/cli/demo/writers/docker.py +182 -0
  56. nl2sql/cli/demo/writers/sqlite.py +88 -0
  57. nl2sql/cli/generators/datasources/__init__.py +3 -0
  58. nl2sql/cli/generators/datasources/generator.py +24 -0
  59. nl2sql/cli/generators/datasources/templates.py +7 -0
  60. nl2sql/cli/generators/env/__init__.py +3 -0
  61. nl2sql/cli/generators/env/generator.py +46 -0
  62. nl2sql/cli/generators/env/templates.py +25 -0
  63. nl2sql/cli/generators/llm/__init__.py +3 -0
  64. nl2sql/cli/generators/llm/generator.py +24 -0
  65. nl2sql/cli/generators/llm/templates.py +4 -0
  66. nl2sql/cli/generators/policies/__init__.py +3 -0
  67. nl2sql/cli/generators/policies/generator.py +20 -0
  68. nl2sql/cli/generators/policies/templates.py +2 -0
  69. nl2sql/cli/main.py +195 -0
  70. nl2sql/cli/reporting.py +878 -0
  71. nl2sql/cli/types.py +13 -0
  72. nl2sql/common/__init__.py +1 -0
  73. nl2sql/common/cancellation.py +25 -0
  74. nl2sql/common/context.py +5 -0
  75. nl2sql/common/errors.py +109 -0
  76. nl2sql/common/event_logger.py +88 -0
  77. nl2sql/common/exceptions.py +3 -0
  78. nl2sql/common/logger.py +119 -0
  79. nl2sql/common/metrics.py +50 -0
  80. nl2sql/common/resilience.py +59 -0
  81. nl2sql/common/settings.py +195 -0
  82. nl2sql/configs/__init__.py +6 -0
  83. nl2sql/configs/datasources.py +10 -0
  84. nl2sql/configs/llm.py +36 -0
  85. nl2sql/configs/manager.py +176 -0
  86. nl2sql/configs/policies.py +14 -0
  87. nl2sql/configs/sample_questions.py +11 -0
  88. nl2sql/configs/secrets.py +11 -0
  89. nl2sql/context.py +106 -0
  90. nl2sql/datasources/__init__.py +21 -0
  91. nl2sql/datasources/discovery.py +28 -0
  92. nl2sql/datasources/models.py +21 -0
  93. nl2sql/datasources/protocols.py +3 -0
  94. nl2sql/datasources/registry.py +172 -0
  95. nl2sql/evaluation/__init__.py +6 -0
  96. nl2sql/evaluation/benchmark_runner.py +320 -0
  97. nl2sql/evaluation/evaluator.py +134 -0
  98. nl2sql/evaluation/types.py +22 -0
  99. nl2sql/execution/__init__.py +4 -0
  100. nl2sql/execution/artifacts/__init__.py +3 -0
  101. nl2sql/execution/artifacts/parquet.py +41 -0
  102. nl2sql/execution/artifacts/store.py +165 -0
  103. nl2sql/execution/contracts.py +57 -0
  104. nl2sql/execution/execution_store.py +25 -0
  105. nl2sql/execution/executor/__init__.py +3 -0
  106. nl2sql/execution/executor/sql_executor.py +116 -0
  107. nl2sql/indexing/__init__.py +7 -0
  108. nl2sql/indexing/chunk_builder.py +227 -0
  109. nl2sql/indexing/embeddings.py +180 -0
  110. nl2sql/indexing/enrichment_service.py +316 -0
  111. nl2sql/indexing/models.py +209 -0
  112. nl2sql/indexing/orchestrator.py +90 -0
  113. nl2sql/indexing/vector_store.py +422 -0
  114. nl2sql/llm/__init__.py +8 -0
  115. nl2sql/llm/models.py +10 -0
  116. nl2sql/llm/registry.py +214 -0
  117. nl2sql/pipeline/__init__.py +1 -0
  118. nl2sql/pipeline/graph.py +73 -0
  119. nl2sql/pipeline/graph_utils.py +141 -0
  120. nl2sql/pipeline/nodes/__init__.py +25 -0
  121. nl2sql/pipeline/nodes/aggregator/__init__.py +4 -0
  122. nl2sql/pipeline/nodes/aggregator/node.py +55 -0
  123. nl2sql/pipeline/nodes/aggregator/prompts.py +20 -0
  124. nl2sql/pipeline/nodes/aggregator/schemas.py +28 -0
  125. nl2sql/pipeline/nodes/answer_synthesizer/__init__.py +4 -0
  126. nl2sql/pipeline/nodes/answer_synthesizer/node.py +98 -0
  127. nl2sql/pipeline/nodes/answer_synthesizer/prompts.py +19 -0
  128. nl2sql/pipeline/nodes/answer_synthesizer/schemas.py +24 -0
  129. nl2sql/pipeline/nodes/ast_planner/__init__.py +4 -0
  130. nl2sql/pipeline/nodes/ast_planner/node.py +104 -0
  131. nl2sql/pipeline/nodes/ast_planner/prompts.py +138 -0
  132. nl2sql/pipeline/nodes/ast_planner/schemas.py +236 -0
  133. nl2sql/pipeline/nodes/datasource_resolver/__init__.py +4 -0
  134. nl2sql/pipeline/nodes/datasource_resolver/node.py +253 -0
  135. nl2sql/pipeline/nodes/datasource_resolver/schemas.py +21 -0
  136. nl2sql/pipeline/nodes/decomposer/__init__.py +3 -0
  137. nl2sql/pipeline/nodes/decomposer/node.py +219 -0
  138. nl2sql/pipeline/nodes/decomposer/prompts.py +96 -0
  139. nl2sql/pipeline/nodes/decomposer/schemas.py +143 -0
  140. nl2sql/pipeline/nodes/executor/__init__.py +3 -0
  141. nl2sql/pipeline/nodes/executor/node.py +107 -0
  142. nl2sql/pipeline/nodes/generator/__init__.py +4 -0
  143. nl2sql/pipeline/nodes/generator/node.py +267 -0
  144. nl2sql/pipeline/nodes/generator/schemas.py +13 -0
  145. nl2sql/pipeline/nodes/global_planner/__init__.py +4 -0
  146. nl2sql/pipeline/nodes/global_planner/node.py +186 -0
  147. nl2sql/pipeline/nodes/global_planner/schemas.py +101 -0
  148. nl2sql/pipeline/nodes/refiner/__init__.py +4 -0
  149. nl2sql/pipeline/nodes/refiner/node.py +132 -0
  150. nl2sql/pipeline/nodes/refiner/prompts.py +28 -0
  151. nl2sql/pipeline/nodes/refiner/schemas.py +13 -0
  152. nl2sql/pipeline/nodes/schema_retriever/__init__.py +3 -0
  153. nl2sql/pipeline/nodes/schema_retriever/node.py +252 -0
  154. nl2sql/pipeline/nodes/schema_retriever/schema.py +27 -0
  155. nl2sql/pipeline/nodes/validator/__init__.py +7 -0
  156. nl2sql/pipeline/nodes/validator/node.py +839 -0
  157. nl2sql/pipeline/nodes/validator/schemas.py +12 -0
  158. nl2sql/pipeline/pipeline_runner.py +72 -0
  159. nl2sql/pipeline/routes.py +72 -0
  160. nl2sql/pipeline/runtime.py +153 -0
  161. nl2sql/pipeline/state.py +92 -0
  162. nl2sql/pipeline/subgraphs/__init__.py +5 -0
  163. nl2sql/pipeline/subgraphs/schemas.py +23 -0
  164. nl2sql/pipeline/subgraphs/sql_agent.py +167 -0
  165. nl2sql/public_api.py +199 -0
  166. nl2sql/schema/__init__.py +37 -0
  167. nl2sql/schema/in_memory_store.py +173 -0
  168. nl2sql/schema/protocol.py +88 -0
  169. nl2sql/schema/sqlite_store.py +233 -0
  170. nl2sql/schema/store.py +29 -0
  171. nl2sql/secrets/__init__.py +14 -0
  172. nl2sql/secrets/factory.py +85 -0
  173. nl2sql/secrets/interfaces.py +16 -0
  174. nl2sql/secrets/manager.py +139 -0
  175. nl2sql/secrets/models.py +56 -0
  176. nl2sql/secrets/providers/aws.py +30 -0
  177. nl2sql/secrets/providers/azure.py +49 -0
  178. nl2sql/secrets/providers/env.py +8 -0
  179. nl2sql/secrets/providers/hashi.py +46 -0
  180. nl2sql/services/__init__.py +0 -0
  181. nl2sql/services/callbacks/__init__.py +0 -0
  182. nl2sql/services/callbacks/monitor.py +84 -0
  183. nl2sql/services/callbacks/node_context.py +7 -0
  184. nl2sql/services/callbacks/node_handlers.py +187 -0
  185. nl2sql/services/callbacks/node_metrics.py +14 -0
  186. nl2sql/services/callbacks/presenter.py +12 -0
  187. nl2sql/services/callbacks/token_handler.py +56 -0
  188. nl2sql_engine-0.1.0.dist-info/METADATA +295 -0
  189. nl2sql_engine-0.1.0.dist-info/RECORD +192 -0
  190. nl2sql_engine-0.1.0.dist-info/WHEEL +5 -0
  191. nl2sql_engine-0.1.0.dist-info/entry_points.txt +9 -0
  192. nl2sql_engine-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,122 @@
1
+
2
+ # Default Configurations for Demo
3
+
4
+ SAMPLE_QUESTIONS = {
5
+ "manufacturing_ref": [
6
+ "List all factories in the US",
7
+ "Show me the capacity of Berlin Plant",
8
+ "What shifts are available?",
9
+ "List all machine types produced by TechCorp",
10
+ "Which factories have capacity greater than 4000?",
11
+ "Show all employee roles in the Maintenance department",
12
+ "List customer segments available for reporting"
13
+ ],
14
+ "manufacturing_ops": [
15
+ "Show me active employees in the Austin Gigafactory",
16
+ "Which machines have error logs in the last 7 days?",
17
+ "Who is the operator for machine 5?",
18
+ "Count the number of active machines per factory",
19
+ "List maintenance logs for Vibration sensor alerts",
20
+ "Which machines are overdue for maintenance based on last_maintenance_date?",
21
+ "Show employees hired in the last 12 months",
22
+ "Compare machine status counts by factory",
23
+ "List maintenance technicians with the most downtime hours logged"
24
+ ],
25
+ "manufacturing_supply": [
26
+ "Total sales amount for 'Industrial Controller'",
27
+ "Find suppliers for high value components",
28
+ "Check inventory levels for 'Bolt M5' in Berlin",
29
+ "List products with base cost greater than 500",
30
+ "Show me suppliers from Germany",
31
+ "List products with low inventory across all factories",
32
+ "Which suppliers provide EV Battery Pack Long Range?",
33
+ "Show inventory last updated more than 7 days ago",
34
+ "Compare inventory levels for Hardware category products"
35
+ ],
36
+ "manufacturing_history": [
37
+ "Show total sales orders in Q4",
38
+ "Calculate average production output per run",
39
+ "Summarize sales by customer for last year",
40
+ "List the top 5 largest orders",
41
+ "Show sales orders by status for the last 30 days",
42
+ "Compare production output by factory over the last quarter",
43
+ "List highest revenue products by month",
44
+ "Show average discount percentage by customer segment"
45
+ ]
46
+ }
47
+
48
+ DEMO_POLICIES = {
49
+ "admin": {
50
+ "description": "Demo Admin",
51
+ "role": "admin",
52
+ "allowed_datasources": ["*"],
53
+ "allowed_tables": ["*"]
54
+ }
55
+ }
56
+
57
+ DEMO_LLM_CONFIG = {
58
+ "default": {
59
+ "provider": "openai",
60
+ "model": "gpt-4o",
61
+ "api_key": "${env:OPENAI_API_KEY}"
62
+ }
63
+ }
64
+
65
+ DEMO_LITE_DATASOURCES = [
66
+ {"id": "manufacturing_ref", "connection": {"type": "sqlite", "database": "data/demo_lite/manufacturing_ref.db"}, "description": "Master Data (Factories)"},
67
+ {"id": "manufacturing_ops", "connection": {"type": "sqlite", "database": "data/demo_lite/manufacturing_ops.db"}, "description": "Operational Data (Employees, Machines)"},
68
+ {"id": "manufacturing_supply", "connection": {"type": "sqlite", "database": "data/demo_lite/manufacturing_supply.db"}, "description": "Supply Chain (Inventory)"},
69
+ {"id": "manufacturing_history", "connection": {"type": "sqlite", "database": "data/demo_lite/manufacturing_history.db"}, "description": "Historical Data (Sales)"},
70
+ ]
71
+
72
+ DEMO_DOCKER_DATASOURCES = [
73
+ {
74
+ "id": "manufacturing_ref",
75
+ "connection": {
76
+ "type": "postgres",
77
+ "host": "localhost",
78
+ "port": 5433,
79
+ "user": "${env:DEMO_REF_USER}",
80
+ "password": "${env:DEMO_REF_PASSWORD}",
81
+ "database": "manufacturing_ref"
82
+ },
83
+ "description": "Master Data (Factories)"
84
+ },
85
+ {
86
+ "id": "manufacturing_ops",
87
+ "connection": {
88
+ "type": "postgres",
89
+ "host": "localhost",
90
+ "port": 5434,
91
+ "user": "${env:DEMO_OPS_USER}",
92
+ "password": "${env:DEMO_OPS_PASSWORD}",
93
+ "database": "manufacturing_ops"
94
+ },
95
+ "description": "Operational Data (Employees, Machines)"
96
+ },
97
+ {
98
+ "id": "manufacturing_supply",
99
+ "connection": {
100
+ "type": "mysql",
101
+ "host": "localhost",
102
+ "port": 3307,
103
+ "user": "${env:DEMO_SUPPLY_USER}",
104
+ "password": "${env:DEMO_SUPPLY_PASSWORD}",
105
+ "database": "manufacturing_supply"
106
+ },
107
+ "description": "Supply Chain (Inventory)"
108
+ },
109
+ {
110
+ "id": "manufacturing_history",
111
+ "connection": {
112
+ "type": "mssql",
113
+ "host": "localhost",
114
+ "port": 1434,
115
+ "user": "${env:DEMO_HISTORY_USER}",
116
+ "password": "${env:DEMO_HISTORY_PASSWORD}",
117
+ "database": "manufacturing_history",
118
+ "driver": "ODBC Driver 17 for SQL Server"
119
+ },
120
+ "description": "Historical Data (Sales)"
121
+ },
122
+ ]
@@ -0,0 +1,289 @@
1
+ import random
2
+ import datetime
3
+ from typing import List, Dict, Any, Tuple
4
+ from .data import (
5
+ FACTORIES,
6
+ SHIFTS,
7
+ MACHINE_TYPES,
8
+ PRODUCTS,
9
+ SUPPLIERS,
10
+ DEPARTMENTS,
11
+ CUSTOMER_SEGMENTS,
12
+ EMPLOYEE_ROLES,
13
+ SUPPLIER_PRODUCTS,
14
+ )
15
+
16
+ class DemoDataFactory:
17
+ """
18
+ Generates pure Python data structures (Lists of Dicts) for the Demo Environment.
19
+ Ensures consistency between Lite (SQLite) and Docker (SQL Scripts) by being the single source of truth.
20
+ """
21
+
22
+ def __init__(self, seed: int = 42):
23
+ self.seed = seed
24
+ self._reset_random()
25
+
26
+ def _reset_random(self):
27
+ random.seed(self.seed)
28
+
29
+ def _weighted_choice(self, values: List[int], weights: List[float]) -> int:
30
+ total = sum(weights)
31
+ pick = random.random() * total
32
+ cumulative = 0.0
33
+ for value, weight in zip(values, weights):
34
+ cumulative += weight
35
+ if pick <= cumulative:
36
+ return value
37
+ return values[-1]
38
+
39
+ def generate_secrets(self) -> Dict[str, str]:
40
+ """Generates random passwords for the demo environment."""
41
+ self._reset_random() # Consistent secrets if called deterministically
42
+
43
+ def _gen(length=24):
44
+ chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
45
+ return "".join(random.choice(chars) for _ in range(length))
46
+
47
+ return {
48
+ "DEMO_REF_PASSWORD": _gen(),
49
+ "DEMO_OPS_PASSWORD": _gen(),
50
+ "DEMO_POSTGRES_PASSWORD": _gen(),
51
+ "DEMO_SUPPLY_PASSWORD": _gen(),
52
+ "DEMO_MYSQL_ROOT_PASSWORD": _gen(),
53
+ "DEMO_HISTORY_PASSWORD": f"StrongP@ss{_gen(8)}!",
54
+ "DEMO_MSSQL_SA_PASSWORD": f"StrongP@ss{_gen(8)}!",
55
+ }
56
+
57
+ def get_ref_data(self) -> Tuple[List[Dict], List[Dict], List[Dict], List[Dict], List[Dict], List[Dict]]:
58
+ """Returns reference data for shared dimensions."""
59
+ return FACTORIES, MACHINE_TYPES, SHIFTS, DEPARTMENTS, EMPLOYEE_ROLES, CUSTOMER_SEGMENTS
60
+
61
+ def get_ops_data(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
62
+ """Generates (employees, machines, maintenance_logs)."""
63
+ self._reset_random()
64
+
65
+ factory_ids = [f["id"] for f in FACTORIES]
66
+ factory_weights = [0.45, 0.2, 0.15, 0.12, 0.08]
67
+
68
+ # 1. Employees
69
+ employees = []
70
+ role_ids = [r["id"] for r in EMPLOYEE_ROLES]
71
+ role_weights = [0.45, 0.12, 0.12, 0.1, 0.16, 0.05]
72
+ shift_ids = [s["id"] for s in SHIFTS]
73
+ shift_weights = [0.5, 0.35, 0.15]
74
+ last_names = ["Smith", "Garcia", "Kim", "Muller", "Chen", "Patel", "Jones", "Brown", "Nguyen", "Singh"]
75
+ first_names = ["John", "Maria", "Wei", "Hans", "Rahul", "Sarah", "Aisha", "Luis", "Mina", "Kenji"]
76
+
77
+ for i in range(1, 501):
78
+ fid = self._weighted_choice(factory_ids, factory_weights)
79
+ sid = self._weighted_choice(shift_ids, shift_weights)
80
+ role_id = self._weighted_choice(role_ids, role_weights)
81
+ role = next(r for r in EMPLOYEE_ROLES if r["id"] == role_id)
82
+ hire_date = (datetime.date.today() - datetime.timedelta(days=random.randint(30, 365 * 8))).isoformat()
83
+ status = "Active"
84
+ roll = random.random()
85
+ if roll < 0.07:
86
+ status = "Leave"
87
+ elif roll < 0.1:
88
+ status = "Contractor"
89
+
90
+ employees.append({
91
+ "id": i,
92
+ "name": f"{random.choice(first_names)} {random.choice(last_names)}",
93
+ "factory_id": fid,
94
+ "shift_id": sid,
95
+ "hire_date": hire_date,
96
+ "role_id": role_id,
97
+ "department_id": role["department_id"],
98
+ "status": status
99
+ })
100
+
101
+ # 2. Machines
102
+ machines = []
103
+ mach_id = 1
104
+ total_capacity = sum(f["capacity"] for f in FACTORIES)
105
+ for factory in FACTORIES:
106
+ base = max(10, int((factory["capacity"] / total_capacity) * 150))
107
+ count = base + random.randint(-2, 3)
108
+ for _ in range(count):
109
+ mtype = random.choice(MACHINE_TYPES)
110
+ install_date = datetime.date.today() - datetime.timedelta(days=random.randint(120, 2000))
111
+ last_maint = install_date + datetime.timedelta(days=random.randint(30, 900))
112
+ if last_maint > datetime.date.today():
113
+ last_maint = datetime.date.today() - datetime.timedelta(days=random.randint(7, 90))
114
+ status = "Active"
115
+ if (datetime.date.today() - last_maint).days > int(mtype["maintenance_interval_days"] * 1.2):
116
+ status = "Maintenance"
117
+ elif random.random() < 0.02:
118
+ status = "Error"
119
+
120
+ machines.append({
121
+ "id": mach_id,
122
+ "factory_id": factory["id"],
123
+ "type_id": mtype["id"],
124
+ "status": status,
125
+ "installation_date": install_date.isoformat(),
126
+ "last_maintenance_date": last_maint.isoformat()
127
+ })
128
+ mach_id += 1
129
+
130
+ # 3. Logs
131
+ logs = []
132
+ log_id = 1
133
+ technician_ids = [e["id"] for e in employees if e["role_id"] == 5]
134
+ severities = ["Low", "Medium", "High", "Critical"]
135
+ descriptions = [
136
+ "Vibration sensor alert triggered",
137
+ "Unusual noise reported",
138
+ "Temperature threshold exceeded",
139
+ "Calibration drift detected",
140
+ "Unexpected shutdown",
141
+ "Hydraulic pressure drop",
142
+ ]
143
+ for _ in range(250):
144
+ machine = random.choice(machines)
145
+ tech_id = random.choice(technician_ids) if technician_ids else random.choice(employees)["id"]
146
+ sev = random.choices(severities, weights=[0.45, 0.3, 0.2, 0.05], k=1)[0]
147
+ downtime = random.randint(1, 4) if sev in ["Low", "Medium"] else random.randint(4, 12)
148
+ log_date = datetime.date.today() - datetime.timedelta(days=random.randint(0, 120))
149
+ logs.append({
150
+ "id": log_id,
151
+ "machine_id": machine["id"],
152
+ "date": log_date.isoformat(),
153
+ "description": random.choice(descriptions),
154
+ "technician_id": tech_id,
155
+ "severity": sev,
156
+ "downtime_hours": downtime
157
+ })
158
+ log_id += 1
159
+
160
+ return employees, machines, logs
161
+
162
+ def get_supply_data(self) -> Tuple[List[Dict], List[Dict], List[Dict], List[Dict]]:
163
+ """Generates (products, suppliers, inventory, supplier_products)."""
164
+ self._reset_random()
165
+
166
+ # Products & Suppliers are static
167
+ products = PRODUCTS
168
+ suppliers = SUPPLIERS
169
+ supplier_products = SUPPLIER_PRODUCTS
170
+
171
+ # Inventory
172
+ inventory = []
173
+ for p in PRODUCTS:
174
+ for f in FACTORIES:
175
+ if p["id"] in [1, 2, 3]:
176
+ qty = random.randint(20, 200)
177
+ else:
178
+ qty = random.randint(200, 2000)
179
+ # Scenario: Low Stock
180
+ if p["id"] == 4 and f["id"] == 2:
181
+ qty = 5
182
+ if p["id"] == 2 and f["id"] == 4:
183
+ qty = 12
184
+ last_updated = (datetime.date.today() - datetime.timedelta(days=random.randint(0, 14))).isoformat()
185
+
186
+ inventory.append({
187
+ "product_id": p["id"],
188
+ "factory_id": f["id"],
189
+ "quantity": qty,
190
+ "last_updated": last_updated
191
+ })
192
+
193
+ return products, suppliers, inventory, supplier_products
194
+
195
+ def get_history_data(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
196
+ """Generates (sales_orders, sales_items, production_runs)."""
197
+ self._reset_random()
198
+
199
+ start_date = datetime.date.today() - datetime.timedelta(days=365)
200
+ customers = ["Acme Inc", "Cyberdyne", "Wayne Ent", "Stark Ind", "Massive Dynamic"]
201
+ segment_ids = [s["id"] for s in CUSTOMER_SEGMENTS]
202
+ factory_ids = [f["id"] for f in FACTORIES]
203
+ factory_weights = [0.45, 0.2, 0.15, 0.12, 0.08]
204
+ order_statuses = ["Pending", "Shipped", "Delivered", "Cancelled"]
205
+
206
+ orders = []
207
+ items = []
208
+ item_id_counter = 1
209
+ monthly_volume = {m: 0 for m in range(1, 13)}
210
+
211
+ # 5000 Orders
212
+ for i in range(1, 5001):
213
+ cust = random.choice(customers)
214
+ if random.random() < 0.6:
215
+ delta = random.randint(270, 365)
216
+ else:
217
+ delta = random.randint(0, 270)
218
+
219
+ order_date = start_date + datetime.timedelta(days=delta)
220
+ date_str = order_date.isoformat()
221
+ fid = self._weighted_choice(factory_ids, factory_weights)
222
+ status = random.choices(order_statuses, weights=[0.12, 0.25, 0.58, 0.05], k=1)[0]
223
+ segment_id = random.choice(segment_ids)
224
+
225
+ order_total = 0.0
226
+ num_items = random.randint(1, 6)
227
+
228
+ # Generate Items
229
+ for _ in range(num_items):
230
+ prod = random.choices(PRODUCTS, weights=[1, 1, 2, 6, 6, 4, 3, 4, 3], k=1)[0]
231
+ qty = random.randint(1, 120)
232
+ markup = random.uniform(1.3, 1.8)
233
+ discount = 0.0
234
+ if qty >= 50:
235
+ discount = random.choice([0.02, 0.05, 0.08])
236
+ price = prod["base_cost"] * markup
237
+ line_total = qty * price * (1 - discount)
238
+ order_total += line_total
239
+
240
+ items.append({
241
+ "id": item_id_counter,
242
+ "order_id": i,
243
+ "product_id": prod["id"],
244
+ "quantity": qty,
245
+ "unit_price": round(price, 2),
246
+ "discount_pct": round(discount * 100, 2)
247
+ })
248
+ item_id_counter += 1
249
+ monthly_volume[order_date.month] += qty
250
+
251
+ orders.append({
252
+ "id": i,
253
+ "customer_name": cust,
254
+ "order_date": date_str,
255
+ "total_amount": round(order_total, 2),
256
+ "status": status,
257
+ "customer_segment_id": segment_id,
258
+ "factory_id": fid
259
+ })
260
+
261
+ # Production Runs
262
+ runs = []
263
+ run_id = 1
264
+ for day_offset in range(0, 365):
265
+ run_date = start_date + datetime.timedelta(days=day_offset)
266
+ month = run_date.month
267
+ season_multiplier = 1.0
268
+ if month in [10, 11, 12]:
269
+ season_multiplier = 1.35
270
+ elif month in [6, 7, 8]:
271
+ season_multiplier = 1.15
272
+
273
+ for factory in FACTORIES:
274
+ base = int((factory["capacity"] / 10) * season_multiplier)
275
+ qty = max(250, base + random.randint(-120, 150))
276
+ run_shift = self._weighted_choice([s["id"] for s in SHIFTS], [0.5, 0.35, 0.15])
277
+ run_status = "Complete" if random.random() > 0.03 else "Delayed"
278
+
279
+ runs.append({
280
+ "id": run_id,
281
+ "factory_id": factory["id"],
282
+ "date": run_date.isoformat(),
283
+ "output_quantity": qty,
284
+ "shift_id": run_shift,
285
+ "status": run_status
286
+ })
287
+ run_id += 1
288
+
289
+ return orders, items, runs
@@ -0,0 +1,230 @@
1
+ from typing import Dict, Any, Optional
2
+ import pathlib
3
+ import subprocess
4
+ import yaml
5
+ from rich.console import Console
6
+ from rich.markup import escape
7
+ from nl2sql.configs import (
8
+ ConfigManager,
9
+ LLMFileConfig,
10
+ DatasourceConfig,
11
+ DatasourceFileConfig,
12
+ PolicyFileConfig
13
+ )
14
+
15
+ from nl2sql.cli.generators.env import EnvFileGenerator
16
+ from nl2sql.cli.generators.datasources import DatasourceGenerator
17
+ from nl2sql.cli.generators.llm import LLMGenerator
18
+ from nl2sql.cli.generators.policies import PolicyGenerator
19
+
20
+ from .factory import DemoDataFactory
21
+ from .writers.sqlite import SQLiteWriter
22
+ from .writers.docker import DockerWriter
23
+ from .defaults import (
24
+ SAMPLE_QUESTIONS,
25
+ DEMO_POLICIES,
26
+ DEMO_LLM_CONFIG,
27
+ DEMO_LITE_DATASOURCES,
28
+ DEMO_DOCKER_DATASOURCES
29
+ )
30
+
31
+ class DemoManager:
32
+ """
33
+ Manages the creation of the Demo Environment (Lite or Docker).
34
+ Orchestrates Data Factory, Writers, and Generators.
35
+ """
36
+
37
+ def __init__(self, console: Console, project_root: pathlib.Path):
38
+ self.console = console
39
+ self.project_root = project_root
40
+ self.config_manager = ConfigManager(project_root)
41
+ self.factory = DemoDataFactory(seed=42)
42
+
43
+ def print_step(self, msg: str):
44
+ self.console.print(f"[dim]{escape(str(msg))}[/dim]")
45
+
46
+ def print_success(self, msg: str):
47
+ self.console.print(f"[green][OK][/green] {escape(str(msg))}")
48
+
49
+ def print_error(self, msg: str):
50
+ self.console.print(f"[red][ERROR] {escape(str(msg))}[/red]")
51
+
52
+ def setup_lite(self, api_key: Optional[str] = None):
53
+ """Sets up the SQLite-based demo environment."""
54
+ data_dir = self.project_root / "data" / "demo_lite"
55
+ self.print_step(f"Generating SQLite Databases in {data_dir}...")
56
+
57
+ # 1. Generate Data
58
+ ref = self.factory.get_ref_data()
59
+ ops = self.factory.get_ops_data()
60
+ supply = self.factory.get_supply_data()
61
+ history = self.factory.get_history_data()
62
+
63
+ # 2. Write DBs
64
+ SQLiteWriter.write_lite(data_dir, ref, ops, supply, history)
65
+
66
+ configs = DEMO_LITE_DATASOURCES
67
+
68
+ # Write using Generator (Strict Typed)
69
+ ds_configs = [DatasourceConfig(**c) for c in configs]
70
+ file_config = DatasourceFileConfig(datasources=ds_configs)
71
+ content = DatasourceGenerator.generate(file_config)
72
+ ds_path = self.project_root / "configs" / "datasources.demo.yaml"
73
+ ds_path.parent.mkdir(parents=True, exist_ok=True)
74
+ with open(ds_path, "w", encoding="utf-8") as f:
75
+ f.write(content)
76
+
77
+ self._write_common_artifacts()
78
+
79
+ self.print_step("Writing .env.demo configuration...")
80
+
81
+ secrets = {}
82
+ if api_key:
83
+ secrets["OPENAI_API_KEY"] = api_key
84
+
85
+ env_content = EnvFileGenerator.generate("demo", secrets=secrets)
86
+ env_path = self.project_root / ".env.demo"
87
+ with open(env_path, "w", encoding="utf-8") as f:
88
+ f.write(env_content)
89
+
90
+ self.print_success("Lite Demo Setup Complete")
91
+
92
+
93
+ def setup_docker(self, api_key: Optional[str] = None):
94
+ """Sets up the Docker-based demo environment."""
95
+ docker_dir = self.project_root / "demo_docker"
96
+ self.print_step(f"Generating Docker Configuration in {docker_dir}...")
97
+
98
+ # 1. Generate Data & Secrets
99
+ secrets = self.factory.generate_secrets()
100
+ if api_key:
101
+ secrets["OPENAI_API_KEY"] = api_key
102
+
103
+ ref = self.factory.get_ref_data()
104
+ ops = self.factory.get_ops_data()
105
+ supply = self.factory.get_supply_data()
106
+ history = self.factory.get_history_data()
107
+
108
+ # 2. Write Artifacts
109
+ DockerWriter.write_docker(docker_dir, secrets, ref, ops, supply, history)
110
+
111
+ self.print_step("Writing datasources config...")
112
+ # Use defaults from configuration
113
+ configs = DEMO_DOCKER_DATASOURCES
114
+
115
+
116
+ ds_configs = [DatasourceConfig(**c) for c in configs]
117
+ file_config = DatasourceFileConfig(datasources=ds_configs)
118
+ content = DatasourceGenerator.generate(file_config)
119
+ ds_path = self.project_root / "configs" / "datasources.demo.yaml"
120
+ ds_path.parent.mkdir(parents=True, exist_ok=True)
121
+ with open(ds_path, "w", encoding="utf-8") as f:
122
+ f.write(content)
123
+
124
+ self._write_common_artifacts()
125
+
126
+ # The app container reads the same .env.demo the lite path writes, so it
127
+ # has to exist before `docker compose up` (and before `docker compose config`).
128
+ self.print_step("Writing .env.demo configuration...")
129
+ self.copy_docker_env_to_root(docker_dir)
130
+
131
+ self.print_success("Docker Configuration Generated")
132
+ return docker_dir
133
+
134
+ def _write_common_artifacts(self):
135
+ """Writes policies and sample questions."""
136
+ self.print_step("Writing policies...")
137
+ policy_config = PolicyFileConfig(roles=DEMO_POLICIES)
138
+ content = PolicyGenerator.generate(policy_config)
139
+ policy_path = self.project_root / "configs" / "policies.demo.json"
140
+
141
+ if not policy_path.parent.exists():
142
+ policy_path.parent.mkdir(parents=True, exist_ok=True)
143
+
144
+ with open(policy_path, "w", encoding="utf-8") as f:
145
+ f.write(content)
146
+
147
+ self.print_step("Writing sample questions...")
148
+ samples_path = self.project_root / "configs" / "sample_questions.demo.yaml"
149
+ with open(samples_path, "w") as f:
150
+ yaml.dump(SAMPLE_QUESTIONS, f, sort_keys=False)
151
+
152
+ self.print_step("Writing LLM config...")
153
+ llm_config = LLMFileConfig(**DEMO_LLM_CONFIG)
154
+ content = LLMGenerator.generate(llm_config)
155
+ llm_path = self.project_root / "configs" / "llm.demo.yaml"
156
+ with open(llm_path, "w", encoding="utf-8") as f:
157
+ f.write(content)
158
+
159
+ def start_docker_containers(self, docker_dir: pathlib.Path) -> bool:
160
+ """Starts the docker containers using subprocess."""
161
+ try:
162
+ subprocess.run(
163
+ ["docker", "compose", "-f", "docker-compose.demo.yml", "up", "-d"],
164
+ cwd=docker_dir,
165
+ check=True
166
+ )
167
+ return True
168
+ except Exception as e:
169
+ self.print_error(f"Failed to start Docker: {e}")
170
+ return False
171
+
172
+ def copy_docker_env_to_root(self, docker_dir: pathlib.Path) -> bool:
173
+ """Copies the generated .env from docker dir to project root as .env.demo, using Protocol."""
174
+ try:
175
+ src = docker_dir / ".env"
176
+ dest = self.project_root / ".env.demo"
177
+
178
+ # Read secrets from source env
179
+ secrets = {}
180
+ if src.exists():
181
+ with open(src, "r", encoding="utf-8") as f:
182
+ for line in f:
183
+ line = line.strip()
184
+ if line and not line.startswith("#") and "=" in line:
185
+ k, v = line.split("=", 1)
186
+ secrets[k.strip()] = v.strip()
187
+
188
+ # Generate Standard Env Content with secrets injected
189
+ content = EnvFileGenerator.generate("demo", secrets=secrets)
190
+
191
+ # Write content
192
+ with open(dest, "w", encoding="utf-8") as f:
193
+ f.write(content)
194
+
195
+ return True
196
+ except Exception:
197
+ return False
198
+
199
+ def index_demo_data(self):
200
+ """Triggers the indexing process for the demo."""
201
+
202
+ from dotenv import load_dotenv
203
+ from nl2sql.common.settings import settings, reload_settings
204
+
205
+ env_path = self.project_root / ".env.demo"
206
+ if not env_path.exists():
207
+ self.print_error(f"Could not find {env_path}")
208
+ return False
209
+
210
+ # The demo environment must be active before the context is built:
211
+ # NL2SQLContext validates vector store settings during construction.
212
+ load_dotenv(env_path, override=True)
213
+ reload_settings()
214
+
215
+ from nl2sql.context import NL2SQLContext
216
+ from nl2sql.cli.commands.indexing import run_indexing
217
+
218
+ try:
219
+ ctx = NL2SQLContext(
220
+ ds_config_path=self.project_root / settings.datasource_config_path,
221
+ secrets_config_path=self.project_root / settings.secrets_config_path,
222
+ llm_config_path=self.project_root / settings.llm_config_path,
223
+ vector_store_path=self.project_root / settings.vector_store_path,
224
+ policies_config_path=self.project_root / settings.policies_config_path,
225
+ )
226
+ run_indexing(ctx)
227
+ return True
228
+ except Exception as e:
229
+ self.print_error(f"Indexing Failed: {e}")
230
+ return False