edsl 0.1.39.dev3__py3-none-any.whl → 0.1.39.dev5__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 (341) hide show
  1. edsl/Base.py +413 -332
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +57 -49
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +1071 -867
  7. edsl/agents/AgentList.py +551 -413
  8. edsl/agents/Invigilator.py +284 -233
  9. edsl/agents/InvigilatorBase.py +257 -270
  10. edsl/agents/PromptConstructor.py +272 -354
  11. edsl/agents/QuestionInstructionPromptBuilder.py +128 -0
  12. edsl/agents/QuestionTemplateReplacementsBuilder.py +137 -0
  13. edsl/agents/__init__.py +2 -3
  14. edsl/agents/descriptors.py +99 -99
  15. edsl/agents/prompt_helpers.py +129 -129
  16. edsl/agents/question_option_processor.py +172 -0
  17. edsl/auto/AutoStudy.py +130 -117
  18. edsl/auto/StageBase.py +243 -230
  19. edsl/auto/StageGenerateSurvey.py +178 -178
  20. edsl/auto/StageLabelQuestions.py +125 -125
  21. edsl/auto/StagePersona.py +61 -61
  22. edsl/auto/StagePersonaDimensionValueRanges.py +88 -88
  23. edsl/auto/StagePersonaDimensionValues.py +74 -74
  24. edsl/auto/StagePersonaDimensions.py +69 -69
  25. edsl/auto/StageQuestions.py +74 -73
  26. edsl/auto/SurveyCreatorPipeline.py +21 -21
  27. edsl/auto/utilities.py +218 -224
  28. edsl/base/Base.py +279 -279
  29. edsl/config.py +177 -157
  30. edsl/conversation/Conversation.py +290 -290
  31. edsl/conversation/car_buying.py +59 -58
  32. edsl/conversation/chips.py +95 -95
  33. edsl/conversation/mug_negotiation.py +81 -81
  34. edsl/conversation/next_speaker_utilities.py +93 -93
  35. edsl/coop/CoopFunctionsMixin.py +15 -0
  36. edsl/coop/ExpectedParrotKeyHandler.py +125 -0
  37. edsl/coop/PriceFetcher.py +54 -54
  38. edsl/coop/__init__.py +2 -2
  39. edsl/coop/coop.py +1106 -1028
  40. edsl/coop/utils.py +131 -131
  41. edsl/data/Cache.py +573 -555
  42. edsl/data/CacheEntry.py +230 -233
  43. edsl/data/CacheHandler.py +168 -149
  44. edsl/data/RemoteCacheSync.py +186 -78
  45. edsl/data/SQLiteDict.py +292 -292
  46. edsl/data/__init__.py +5 -4
  47. edsl/data/orm.py +10 -10
  48. edsl/data_transfer_models.py +74 -73
  49. edsl/enums.py +202 -175
  50. edsl/exceptions/BaseException.py +21 -21
  51. edsl/exceptions/__init__.py +54 -54
  52. edsl/exceptions/agents.py +54 -42
  53. edsl/exceptions/cache.py +5 -5
  54. edsl/exceptions/configuration.py +16 -16
  55. edsl/exceptions/coop.py +10 -10
  56. edsl/exceptions/data.py +14 -14
  57. edsl/exceptions/general.py +34 -34
  58. edsl/exceptions/inference_services.py +5 -0
  59. edsl/exceptions/jobs.py +33 -33
  60. edsl/exceptions/language_models.py +63 -63
  61. edsl/exceptions/prompts.py +15 -15
  62. edsl/exceptions/questions.py +109 -91
  63. edsl/exceptions/results.py +29 -29
  64. edsl/exceptions/scenarios.py +29 -22
  65. edsl/exceptions/surveys.py +37 -37
  66. edsl/inference_services/AnthropicService.py +106 -87
  67. edsl/inference_services/AvailableModelCacheHandler.py +184 -0
  68. edsl/inference_services/AvailableModelFetcher.py +215 -0
  69. edsl/inference_services/AwsBedrock.py +118 -120
  70. edsl/inference_services/AzureAI.py +215 -217
  71. edsl/inference_services/DeepInfraService.py +18 -18
  72. edsl/inference_services/GoogleService.py +143 -148
  73. edsl/inference_services/GroqService.py +20 -20
  74. edsl/inference_services/InferenceServiceABC.py +80 -147
  75. edsl/inference_services/InferenceServicesCollection.py +138 -97
  76. edsl/inference_services/MistralAIService.py +120 -123
  77. edsl/inference_services/OllamaService.py +18 -18
  78. edsl/inference_services/OpenAIService.py +236 -224
  79. edsl/inference_services/PerplexityService.py +160 -163
  80. edsl/inference_services/ServiceAvailability.py +135 -0
  81. edsl/inference_services/TestService.py +90 -89
  82. edsl/inference_services/TogetherAIService.py +172 -170
  83. edsl/inference_services/data_structures.py +134 -0
  84. edsl/inference_services/models_available_cache.py +118 -118
  85. edsl/inference_services/rate_limits_cache.py +25 -25
  86. edsl/inference_services/registry.py +41 -41
  87. edsl/inference_services/write_available.py +10 -10
  88. edsl/jobs/AnswerQuestionFunctionConstructor.py +223 -0
  89. edsl/jobs/Answers.py +43 -56
  90. edsl/jobs/FetchInvigilator.py +47 -0
  91. edsl/jobs/InterviewTaskManager.py +98 -0
  92. edsl/jobs/InterviewsConstructor.py +50 -0
  93. edsl/jobs/Jobs.py +823 -898
  94. edsl/jobs/JobsChecks.py +172 -147
  95. edsl/jobs/JobsComponentConstructor.py +189 -0
  96. edsl/jobs/JobsPrompts.py +270 -268
  97. edsl/jobs/JobsRemoteInferenceHandler.py +311 -239
  98. edsl/jobs/JobsRemoteInferenceLogger.py +239 -0
  99. edsl/jobs/RequestTokenEstimator.py +30 -0
  100. edsl/jobs/__init__.py +1 -1
  101. edsl/jobs/async_interview_runner.py +138 -0
  102. edsl/jobs/buckets/BucketCollection.py +104 -63
  103. edsl/jobs/buckets/ModelBuckets.py +65 -65
  104. edsl/jobs/buckets/TokenBucket.py +283 -251
  105. edsl/jobs/buckets/TokenBucketAPI.py +211 -0
  106. edsl/jobs/buckets/TokenBucketClient.py +191 -0
  107. edsl/jobs/check_survey_scenario_compatibility.py +85 -0
  108. edsl/jobs/data_structures.py +120 -0
  109. edsl/jobs/decorators.py +35 -0
  110. edsl/jobs/interviews/Interview.py +396 -661
  111. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  112. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  113. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  114. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  115. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  116. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  117. edsl/jobs/interviews/ReportErrors.py +66 -66
  118. edsl/jobs/interviews/interview_status_enum.py +9 -9
  119. edsl/jobs/jobs_status_enums.py +9 -0
  120. edsl/jobs/loggers/HTMLTableJobLogger.py +304 -0
  121. edsl/jobs/results_exceptions_handler.py +98 -0
  122. edsl/jobs/runners/JobsRunnerAsyncio.py +151 -466
  123. edsl/jobs/runners/JobsRunnerStatus.py +297 -330
  124. edsl/jobs/tasks/QuestionTaskCreator.py +244 -242
  125. edsl/jobs/tasks/TaskCreators.py +64 -64
  126. edsl/jobs/tasks/TaskHistory.py +470 -450
  127. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  128. edsl/jobs/tasks/task_status_enum.py +161 -163
  129. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  130. edsl/jobs/tokens/TokenUsage.py +34 -34
  131. edsl/language_models/ComputeCost.py +63 -0
  132. edsl/language_models/LanguageModel.py +626 -668
  133. edsl/language_models/ModelList.py +164 -155
  134. edsl/language_models/PriceManager.py +127 -0
  135. edsl/language_models/RawResponseHandler.py +106 -0
  136. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  137. edsl/language_models/ServiceDataSources.py +0 -0
  138. edsl/language_models/__init__.py +2 -3
  139. edsl/language_models/fake_openai_call.py +15 -15
  140. edsl/language_models/fake_openai_service.py +61 -61
  141. edsl/language_models/key_management/KeyLookup.py +63 -0
  142. edsl/language_models/key_management/KeyLookupBuilder.py +273 -0
  143. edsl/language_models/key_management/KeyLookupCollection.py +38 -0
  144. edsl/language_models/key_management/__init__.py +0 -0
  145. edsl/language_models/key_management/models.py +131 -0
  146. edsl/language_models/model.py +256 -0
  147. edsl/language_models/repair.py +156 -156
  148. edsl/language_models/utilities.py +65 -64
  149. edsl/notebooks/Notebook.py +263 -258
  150. edsl/notebooks/NotebookToLaTeX.py +142 -0
  151. edsl/notebooks/__init__.py +1 -1
  152. edsl/prompts/Prompt.py +352 -362
  153. edsl/prompts/__init__.py +2 -2
  154. edsl/questions/ExceptionExplainer.py +77 -0
  155. edsl/questions/HTMLQuestion.py +103 -0
  156. edsl/questions/QuestionBase.py +518 -664
  157. edsl/questions/QuestionBasePromptsMixin.py +221 -217
  158. edsl/questions/QuestionBudget.py +227 -227
  159. edsl/questions/QuestionCheckBox.py +359 -359
  160. edsl/questions/QuestionExtract.py +180 -182
  161. edsl/questions/QuestionFreeText.py +113 -114
  162. edsl/questions/QuestionFunctional.py +166 -166
  163. edsl/questions/QuestionList.py +223 -231
  164. edsl/questions/QuestionMatrix.py +265 -0
  165. edsl/questions/QuestionMultipleChoice.py +330 -286
  166. edsl/questions/QuestionNumerical.py +151 -153
  167. edsl/questions/QuestionRank.py +314 -324
  168. edsl/questions/Quick.py +41 -41
  169. edsl/questions/SimpleAskMixin.py +74 -73
  170. edsl/questions/__init__.py +27 -26
  171. edsl/questions/{AnswerValidatorMixin.py → answer_validator_mixin.py} +334 -289
  172. edsl/questions/compose_questions.py +98 -98
  173. edsl/questions/data_structures.py +20 -0
  174. edsl/questions/decorators.py +21 -21
  175. edsl/questions/derived/QuestionLikertFive.py +76 -76
  176. edsl/questions/derived/QuestionLinearScale.py +90 -87
  177. edsl/questions/derived/QuestionTopK.py +93 -93
  178. edsl/questions/derived/QuestionYesNo.py +82 -82
  179. edsl/questions/descriptors.py +427 -413
  180. edsl/questions/loop_processor.py +149 -0
  181. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  182. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  183. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  184. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  185. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  186. edsl/questions/prompt_templates/question_list.jinja +17 -17
  187. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  188. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  189. edsl/questions/{QuestionBaseGenMixin.py → question_base_gen_mixin.py} +168 -161
  190. edsl/questions/question_registry.py +177 -177
  191. edsl/questions/{RegisterQuestionsMeta.py → register_questions_meta.py} +71 -71
  192. edsl/questions/{ResponseValidatorABC.py → response_validator_abc.py} +188 -174
  193. edsl/questions/response_validator_factory.py +34 -0
  194. edsl/questions/settings.py +12 -12
  195. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  196. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  197. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  198. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  199. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  200. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  201. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  202. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  203. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  204. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  205. edsl/questions/templates/list/question_presentation.jinja +5 -5
  206. edsl/questions/templates/matrix/__init__.py +1 -0
  207. edsl/questions/templates/matrix/answering_instructions.jinja +5 -0
  208. edsl/questions/templates/matrix/question_presentation.jinja +20 -0
  209. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  210. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  211. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  212. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  213. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  214. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  215. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  216. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  217. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  218. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  219. edsl/results/CSSParameterizer.py +108 -108
  220. edsl/results/Dataset.py +587 -424
  221. edsl/results/DatasetExportMixin.py +594 -731
  222. edsl/results/DatasetTree.py +295 -275
  223. edsl/results/MarkdownToDocx.py +122 -0
  224. edsl/results/MarkdownToPDF.py +111 -0
  225. edsl/results/Result.py +557 -465
  226. edsl/results/Results.py +1183 -1165
  227. edsl/results/ResultsExportMixin.py +45 -43
  228. edsl/results/ResultsGGMixin.py +121 -121
  229. edsl/results/TableDisplay.py +125 -198
  230. edsl/results/TextEditor.py +50 -0
  231. edsl/results/__init__.py +2 -2
  232. edsl/results/file_exports.py +252 -0
  233. edsl/results/{ResultsFetchMixin.py → results_fetch_mixin.py} +33 -33
  234. edsl/results/{Selector.py → results_selector.py} +145 -135
  235. edsl/results/{ResultsToolsMixin.py → results_tools_mixin.py} +98 -98
  236. edsl/results/smart_objects.py +96 -0
  237. edsl/results/table_data_class.py +12 -0
  238. edsl/results/table_display.css +77 -77
  239. edsl/results/table_renderers.py +118 -0
  240. edsl/results/tree_explore.py +115 -115
  241. edsl/scenarios/ConstructDownloadLink.py +109 -0
  242. edsl/scenarios/DocumentChunker.py +102 -0
  243. edsl/scenarios/DocxScenario.py +16 -0
  244. edsl/scenarios/FileStore.py +511 -632
  245. edsl/scenarios/PdfExtractor.py +40 -0
  246. edsl/scenarios/Scenario.py +498 -601
  247. edsl/scenarios/ScenarioHtmlMixin.py +65 -64
  248. edsl/scenarios/ScenarioList.py +1458 -1287
  249. edsl/scenarios/ScenarioListExportMixin.py +45 -52
  250. edsl/scenarios/ScenarioListPdfMixin.py +239 -261
  251. edsl/scenarios/__init__.py +3 -4
  252. edsl/scenarios/directory_scanner.py +96 -0
  253. edsl/scenarios/file_methods.py +85 -0
  254. edsl/scenarios/handlers/__init__.py +13 -0
  255. edsl/scenarios/handlers/csv.py +38 -0
  256. edsl/scenarios/handlers/docx.py +76 -0
  257. edsl/scenarios/handlers/html.py +37 -0
  258. edsl/scenarios/handlers/json.py +111 -0
  259. edsl/scenarios/handlers/latex.py +5 -0
  260. edsl/scenarios/handlers/md.py +51 -0
  261. edsl/scenarios/handlers/pdf.py +68 -0
  262. edsl/scenarios/handlers/png.py +39 -0
  263. edsl/scenarios/handlers/pptx.py +105 -0
  264. edsl/scenarios/handlers/py.py +294 -0
  265. edsl/scenarios/handlers/sql.py +313 -0
  266. edsl/scenarios/handlers/sqlite.py +149 -0
  267. edsl/scenarios/handlers/txt.py +33 -0
  268. edsl/scenarios/{ScenarioJoin.py → scenario_join.py} +131 -127
  269. edsl/scenarios/scenario_selector.py +156 -0
  270. edsl/shared.py +1 -1
  271. edsl/study/ObjectEntry.py +173 -173
  272. edsl/study/ProofOfWork.py +113 -113
  273. edsl/study/SnapShot.py +80 -80
  274. edsl/study/Study.py +521 -528
  275. edsl/study/__init__.py +4 -4
  276. edsl/surveys/ConstructDAG.py +92 -0
  277. edsl/surveys/DAG.py +148 -148
  278. edsl/surveys/EditSurvey.py +221 -0
  279. edsl/surveys/InstructionHandler.py +100 -0
  280. edsl/surveys/Memory.py +31 -31
  281. edsl/surveys/MemoryManagement.py +72 -0
  282. edsl/surveys/MemoryPlan.py +244 -244
  283. edsl/surveys/Rule.py +327 -326
  284. edsl/surveys/RuleCollection.py +385 -387
  285. edsl/surveys/RuleManager.py +172 -0
  286. edsl/surveys/Simulator.py +75 -0
  287. edsl/surveys/Survey.py +1280 -1801
  288. edsl/surveys/SurveyCSS.py +273 -261
  289. edsl/surveys/SurveyExportMixin.py +259 -259
  290. edsl/surveys/{SurveyFlowVisualizationMixin.py → SurveyFlowVisualization.py} +181 -179
  291. edsl/surveys/SurveyQualtricsImport.py +284 -284
  292. edsl/surveys/SurveyToApp.py +141 -0
  293. edsl/surveys/__init__.py +5 -3
  294. edsl/surveys/base.py +53 -53
  295. edsl/surveys/descriptors.py +60 -56
  296. edsl/surveys/instructions/ChangeInstruction.py +48 -49
  297. edsl/surveys/instructions/Instruction.py +56 -65
  298. edsl/surveys/instructions/InstructionCollection.py +82 -77
  299. edsl/templates/error_reporting/base.html +23 -23
  300. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  301. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  302. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  303. edsl/templates/error_reporting/interview_details.html +115 -115
  304. edsl/templates/error_reporting/interviews.html +19 -19
  305. edsl/templates/error_reporting/overview.html +4 -4
  306. edsl/templates/error_reporting/performance_plot.html +1 -1
  307. edsl/templates/error_reporting/report.css +73 -73
  308. edsl/templates/error_reporting/report.html +117 -117
  309. edsl/templates/error_reporting/report.js +25 -25
  310. edsl/tools/__init__.py +1 -1
  311. edsl/tools/clusters.py +192 -192
  312. edsl/tools/embeddings.py +27 -27
  313. edsl/tools/embeddings_plotting.py +118 -118
  314. edsl/tools/plotting.py +112 -112
  315. edsl/tools/summarize.py +18 -18
  316. edsl/utilities/PrettyList.py +56 -0
  317. edsl/utilities/SystemInfo.py +28 -28
  318. edsl/utilities/__init__.py +22 -22
  319. edsl/utilities/ast_utilities.py +25 -25
  320. edsl/utilities/data/Registry.py +6 -6
  321. edsl/utilities/data/__init__.py +1 -1
  322. edsl/utilities/data/scooter_results.json +1 -1
  323. edsl/utilities/decorators.py +77 -77
  324. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  325. edsl/utilities/interface.py +627 -627
  326. edsl/utilities/is_notebook.py +18 -0
  327. edsl/utilities/is_valid_variable_name.py +11 -0
  328. edsl/utilities/naming_utilities.py +263 -263
  329. edsl/utilities/remove_edsl_version.py +24 -0
  330. edsl/utilities/repair_functions.py +28 -28
  331. edsl/utilities/restricted_python.py +70 -70
  332. edsl/utilities/utilities.py +436 -424
  333. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/LICENSE +21 -21
  334. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/METADATA +13 -11
  335. edsl-0.1.39.dev5.dist-info/RECORD +358 -0
  336. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/WHEEL +1 -1
  337. edsl/language_models/KeyLookup.py +0 -30
  338. edsl/language_models/registry.py +0 -190
  339. edsl/language_models/unused/ReplicateBase.py +0 -83
  340. edsl/results/ResultsDBMixin.py +0 -238
  341. edsl-0.1.39.dev3.dist-info/RECORD +0 -277
@@ -1,275 +1,295 @@
1
- from typing import Dict, List, Any, Optional, List
2
- from docx import Document
3
- from docx.shared import Inches, Pt
4
- from docx.enum.text import WD_ALIGN_PARAGRAPH
5
- from docx.enum.style import WD_STYLE_TYPE
6
-
7
-
8
- def is_hashable(v):
9
- try:
10
- hash(v)
11
- return True
12
- except TypeError:
13
- return False
14
-
15
-
16
- class TreeNode:
17
- def __init__(self, key=None, value=None):
18
- self.key = key
19
- self.value = value
20
- self.children = {}
21
-
22
-
23
- class Tree:
24
- def __init__(self, data: "Dataset", node_order: Optional[List[str]] = None):
25
- d = {}
26
- for entry in data:
27
- d.update(entry)
28
- self.data = d
29
- self.root = None
30
-
31
- self.node_order = node_order
32
-
33
- self.construct_tree(node_order)
34
-
35
- def unique_values_by_keys(self) -> dict:
36
- unique_values = {}
37
- for key, raw_values in self.data.items():
38
- values = [v if is_hashable(v) else str(v) for v in raw_values]
39
- unique_values[key] = list(set(values))
40
- return unique_values
41
-
42
- def construct_tree(self, node_order: Optional[List[str]] = None):
43
- # Validate node_order
44
- if node_order is None:
45
- unique_values = self.unique_values_by_keys()
46
- # Sort keys by number of unique values
47
- node_order = sorted(
48
- unique_values, key=lambda k: len(unique_values[k]), reverse=True
49
- )
50
- else:
51
- if not set(node_order).issubset(set(self.data.keys())):
52
- invalid_keys = set(node_order) - set(self.data.keys())
53
- raise ValueError(f"Invalid keys in node_order: {invalid_keys}")
54
-
55
- self.root = TreeNode()
56
-
57
- for i in range(len(self.data[node_order[0]])):
58
- current = self.root
59
- for level in node_order[:-1]:
60
- value = self.data[level][i]
61
- if not is_hashable(value):
62
- value = str(value)
63
- if value not in current.children:
64
- current.children[value] = TreeNode(key=level, value=value)
65
- current = current.children[value]
66
-
67
- leaf_key = node_order[-1]
68
- leaf_value = self.data[leaf_key][i]
69
- if not is_hashable(leaf_value):
70
- leaf_value = str(leaf_value)
71
- if leaf_value not in current.children:
72
- current.children[leaf_value] = TreeNode(key=leaf_key, value=leaf_value)
73
-
74
- def __repr__(self):
75
- if self.node_order is not None:
76
- return f"Tree(Dataset({self.data}), node_order={self.node_order})"
77
- else:
78
- return f"Tree(Dataset({self.data}))"
79
-
80
- def print_tree(
81
- self, node: Optional[TreeNode] = None, level: int = 0, print_keys: bool = False
82
- ):
83
- if node is None:
84
- node = self.root
85
- if node is None:
86
- print("Tree has not been constructed yet.")
87
- return
88
-
89
- if node.value is not None:
90
- if print_keys and node.key is not None:
91
- print(" " * level + f"{node.key}: {node.value}")
92
- else:
93
- print(" " * level + str(node.value))
94
- for child in node.children.values():
95
- self.print_tree(child, level + 1, print_keys)
96
-
97
- def to_docx(self, filename: Optional[str] = None):
98
- if filename is None:
99
- filename = "tree_structure.docx"
100
-
101
- doc = Document()
102
-
103
- # Create styles for headings
104
- for i in range(1, 10): # Up to 9 levels of headings
105
- style_name = f"Heading {i}"
106
- if style_name not in doc.styles:
107
- doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
108
-
109
- # Get or create the 'Body Text' style
110
- if "Body Text" not in doc.styles:
111
- body_style = doc.styles.add_style("Body Text", WD_STYLE_TYPE.PARAGRAPH)
112
- else:
113
- body_style = doc.styles["Body Text"]
114
-
115
- body_style.font.size = Pt(11)
116
-
117
- self._add_to_docx(doc, self.root, 0)
118
- doc.save(filename)
119
- from edsl.utilities.utilities import file_notice
120
-
121
- file_notice(filename)
122
-
123
- def _repr_html_(self):
124
- """Returns an interactive HTML representation of the tree with collapsible sections."""
125
-
126
- # Generate a unique ID for this tree instance
127
- import uuid
128
-
129
- tree_id = f"tree_{uuid.uuid4().hex[:8]}"
130
-
131
- styles = f"""
132
- <div class="{tree_id}">
133
- <style>
134
- .{tree_id} details {{
135
- margin-left: 20px;
136
- }}
137
- .{tree_id} summary {{
138
- cursor: pointer;
139
- margin: 2px 0;
140
- }}
141
- .{tree_id} .value {{
142
- font-family: monospace;
143
- background: #f5f5f5;
144
- padding: 2px 6px;
145
- border-radius: 3px;
146
- margin: 1px 0;
147
- }}
148
- .{tree_id} .key {{
149
- color: #666;
150
- font-style: italic;
151
- }}
152
- </style>
153
- """
154
-
155
- def node_to_html(node, level=0, print_keys=True):
156
- if node is None:
157
- return "Tree has not been constructed yet."
158
-
159
- html = []
160
-
161
- if node.value is not None:
162
- # Create the node content
163
- content = []
164
- if print_keys and node.key is not None:
165
- content.append(f'<span class="key">{node.key}: </span>')
166
- content.append(f'<span class="value">{node.value}</span>')
167
- content_html = "".join(content)
168
-
169
- if node.children:
170
- # Node with children
171
- html.append(f'<details {"open" if level < 1 else ""}>')
172
- html.append(f"<summary>{content_html}</summary>")
173
- for child in node.children.values():
174
- html.append(node_to_html(child, level + 1, print_keys))
175
- html.append("</details>")
176
- else:
177
- # Leaf node
178
- html.append(f"<div>{content_html}</div>")
179
- else:
180
- # Root node with no value
181
- if node.children:
182
- for child in node.children.values():
183
- html.append(node_to_html(child, level, print_keys))
184
-
185
- return "\n".join(html)
186
-
187
- tree_html = node_to_html(self.root)
188
- return f"{styles}{tree_html}</div>"
189
-
190
- # def _repr_html_(self):
191
- # """Returns an HTML representation of the tree, following the same logic as print_tree."""
192
- # styles = """
193
- # <style>
194
- # .tree-container {
195
- # font-family: monospace;
196
- # white-space: pre;
197
- # margin: 10px;
198
- # }
199
- # </style>
200
- # """
201
-
202
- # def node_to_html(node, level=0, print_keys=False):
203
- # if node is None:
204
- # node = self.root
205
- # if node is None:
206
- # return "Tree has not been constructed yet."
207
-
208
- # html = []
209
- # if node.value is not None:
210
- # indent = "&nbsp;" * 2 * level # Using &nbsp; for HTML spaces
211
- # if print_keys and node.key is not None:
212
- # html.append(f"{indent}{node.key}: {node.value}<br>")
213
- # else:
214
- # html.append(f"{indent}{node.value}<br>")
215
-
216
- # for child in node.children.values():
217
- # html.append(node_to_html(child, level + 1, print_keys))
218
-
219
- # return "".join(html)
220
-
221
- # tree_html = node_to_html(self.root)
222
- # return f'<div class="tree-container">{tree_html}</div>{styles}'
223
-
224
- def _add_to_docx(self, doc, node: TreeNode, level: int):
225
- if node.value is not None:
226
- if level == 0:
227
- doc.add_heading(str(node.value), level=level + 1)
228
- elif node.children: # If the node has children, it's not the last level
229
- para = doc.add_paragraph(str(node.value))
230
- para.style = f"Heading {level+1}"
231
- else: # If the node has no children, it's the last level (body text)
232
- para = doc.add_paragraph(str(node.value))
233
- para.style = "Body Text"
234
-
235
- # Process child nodes (moved outside the if block)
236
- for child in node.children.values():
237
- self._add_to_docx(doc, child, level + 1)
238
-
239
-
240
- # Example usage (commented out)
241
- """
242
- from edsl.results.Dataset import Dataset
243
-
244
- data = Dataset(
245
- [
246
- {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
247
- {"country": ["US", "China", "France", "Canada", "Japan"]},
248
- {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
249
- {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
250
- ]
251
- )
252
-
253
- tree = Tree(data)
254
-
255
- try:
256
- tree.construct_tree(["continent", "country", "city", "population"])
257
- print("Tree without key names:")
258
- tree.print_tree()
259
- print("\nTree with key names:")
260
- tree.print_tree(print_keys=True)
261
- except ValueError as e:
262
- print(f"Error: {e}")
263
-
264
- # Demonstrating validation
265
- try:
266
- tree.construct_tree(["continent", "country", "invalid_key"])
267
- except ValueError as e:
268
- print(f"\nValidation Error: {e}")
269
-
270
- tree = Tree(data)
271
- tree.construct_tree(["continent", "country", "city", "population"])
272
- tree.print_tree(print_keys=True)
273
- tree.to_docx("tree_structure.docx")
274
- print("DocX file 'tree_structure.docx' has been created.")
275
- """
1
+ from typing import Dict, List, Any, Optional, List
2
+
3
+
4
+ def is_hashable(v):
5
+ try:
6
+ hash(v)
7
+ return True
8
+ except TypeError:
9
+ return False
10
+
11
+
12
+ class TreeNode:
13
+ def __init__(self, key=None, value=None):
14
+ self.key = key
15
+ self.value = value
16
+ self.children = {}
17
+
18
+
19
+ class Tree:
20
+ def __init__(self, data: "Dataset", node_order: Optional[List[str]] = None):
21
+ d = {}
22
+ for entry in data:
23
+ d.update(entry)
24
+ self.data = d
25
+ self.root = None
26
+
27
+ self.node_order = node_order
28
+
29
+ self.construct_tree(node_order)
30
+
31
+ def unique_values_by_keys(self) -> dict:
32
+ unique_values = {}
33
+ for key, raw_values in self.data.items():
34
+ values = [v if is_hashable(v) else str(v) for v in raw_values]
35
+ unique_values[key] = list(set(values))
36
+ return unique_values
37
+
38
+ def construct_tree(self, node_order: Optional[List[str]] = None):
39
+ # Validate node_order
40
+ if node_order is None:
41
+ unique_values = self.unique_values_by_keys()
42
+ # Sort keys by number of unique values
43
+ node_order = sorted(
44
+ unique_values, key=lambda k: len(unique_values[k]), reverse=True
45
+ )
46
+ else:
47
+ if not set(node_order).issubset(set(self.data.keys())):
48
+ invalid_keys = set(node_order) - set(self.data.keys())
49
+ raise ValueError(f"Invalid keys in node_order: {invalid_keys}")
50
+
51
+ self.root = TreeNode()
52
+
53
+ for i in range(len(self.data[node_order[0]])):
54
+ current = self.root
55
+ for level in node_order[:-1]:
56
+ value = self.data[level][i]
57
+ if not is_hashable(value):
58
+ value = str(value)
59
+ if value not in current.children:
60
+ current.children[value] = TreeNode(key=level, value=value)
61
+ current = current.children[value]
62
+
63
+ leaf_key = node_order[-1]
64
+ leaf_value = self.data[leaf_key][i]
65
+ if not is_hashable(leaf_value):
66
+ leaf_value = str(leaf_value)
67
+ if leaf_value not in current.children:
68
+ current.children[leaf_value] = TreeNode(key=leaf_key, value=leaf_value)
69
+
70
+ def __repr__(self):
71
+ if self.node_order is not None:
72
+ return f"Tree(Dataset({self.data}), node_order={self.node_order})"
73
+ else:
74
+ return f"Tree(Dataset({self.data}))"
75
+
76
+ def print_tree(
77
+ self, node: Optional[TreeNode] = None, level: int = 0, print_keys: bool = False
78
+ ):
79
+ if node is None:
80
+ node = self.root
81
+ if node is None:
82
+ print("Tree has not been constructed yet.")
83
+ return
84
+
85
+ if node.value is not None:
86
+ if print_keys and node.key is not None:
87
+ print(" " * level + f"{node.key}: {node.value}")
88
+ else:
89
+ print(" " * level + str(node.value))
90
+ for child in node.children.values():
91
+ self.print_tree(child, level + 1, print_keys)
92
+
93
+ def to_docx(self, filename: Optional[str] = None):
94
+ if filename is None:
95
+ filename = "tree_structure.docx"
96
+
97
+ from docx import Document
98
+ from docx.shared import Inches, Pt
99
+ from docx.enum.text import WD_ALIGN_PARAGRAPH
100
+ from docx.enum.style import WD_STYLE_TYPE
101
+
102
+ doc = Document()
103
+
104
+ # Create styles for headings
105
+ for i in range(1, 10): # Up to 9 levels of headings
106
+ style_name = f"Heading {i}"
107
+ if style_name not in doc.styles:
108
+ doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
109
+
110
+ # Get or create the 'Body Text' style
111
+ if "Body Text" not in doc.styles:
112
+ body_style = doc.styles.add_style("Body Text", WD_STYLE_TYPE.PARAGRAPH)
113
+ else:
114
+ body_style = doc.styles["Body Text"]
115
+
116
+ body_style.font.size = Pt(11)
117
+
118
+ self._add_to_docx(doc, self.root, 0)
119
+ import base64
120
+ from io import BytesIO
121
+ import base64
122
+
123
+ # Save document to bytes buffer
124
+ doc_buffer = BytesIO()
125
+ doc.save(doc_buffer)
126
+ doc_buffer.seek(0)
127
+
128
+ base64_string = base64.b64encode(doc_buffer.getvalue()).decode("utf-8")
129
+ from edsl.scenarios.FileStore import FileStore
130
+
131
+ # Create and return FileStore instance
132
+ return FileStore(
133
+ path="tree_structure.docx", # Default name
134
+ mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
135
+ binary=True,
136
+ suffix="docx",
137
+ base64_string=base64_string,
138
+ )
139
+ # doc.save(filename)
140
+ # from edsl.utilities.utilities import file_notice
141
+ # file_notice(filename)
142
+
143
+ def _repr_html_(self):
144
+ """Returns an interactive HTML representation of the tree with collapsible sections."""
145
+
146
+ # Generate a unique ID for this tree instance
147
+ import uuid
148
+
149
+ tree_id = f"tree_{uuid.uuid4().hex[:8]}"
150
+
151
+ styles = f"""
152
+ <div class="{tree_id}">
153
+ <style>
154
+ .{tree_id} details {{
155
+ margin-left: 20px;
156
+ }}
157
+ .{tree_id} summary {{
158
+ cursor: pointer;
159
+ margin: 2px 0;
160
+ }}
161
+ .{tree_id} .value {{
162
+ font-family: monospace;
163
+ background: #f5f5f5;
164
+ padding: 2px 6px;
165
+ border-radius: 3px;
166
+ margin: 1px 0;
167
+ }}
168
+ .{tree_id} .key {{
169
+ color: #666;
170
+ font-style: italic;
171
+ }}
172
+ </style>
173
+ """
174
+
175
+ def node_to_html(node, level=0, print_keys=True):
176
+ if node is None:
177
+ return "Tree has not been constructed yet."
178
+
179
+ html = []
180
+
181
+ if node.value is not None:
182
+ # Create the node content
183
+ content = []
184
+ if print_keys and node.key is not None:
185
+ content.append(f'<span class="key">{node.key}: </span>')
186
+ content.append(f'<span class="value">{node.value}</span>')
187
+ content_html = "".join(content)
188
+
189
+ if node.children:
190
+ # Node with children
191
+ html.append(f'<details {"open" if level < 1 else ""}>')
192
+ html.append(f"<summary>{content_html}</summary>")
193
+ for child in node.children.values():
194
+ html.append(node_to_html(child, level + 1, print_keys))
195
+ html.append("</details>")
196
+ else:
197
+ # Leaf node
198
+ html.append(f"<div>{content_html}</div>")
199
+ else:
200
+ # Root node with no value
201
+ if node.children:
202
+ for child in node.children.values():
203
+ html.append(node_to_html(child, level, print_keys))
204
+
205
+ return "\n".join(html)
206
+
207
+ tree_html = node_to_html(self.root)
208
+ return f"{styles}{tree_html}</div>"
209
+
210
+ # def _repr_html_(self):
211
+ # """Returns an HTML representation of the tree, following the same logic as print_tree."""
212
+ # styles = """
213
+ # <style>
214
+ # .tree-container {
215
+ # font-family: monospace;
216
+ # white-space: pre;
217
+ # margin: 10px;
218
+ # }
219
+ # </style>
220
+ # """
221
+
222
+ # def node_to_html(node, level=0, print_keys=False):
223
+ # if node is None:
224
+ # node = self.root
225
+ # if node is None:
226
+ # return "Tree has not been constructed yet."
227
+
228
+ # html = []
229
+ # if node.value is not None:
230
+ # indent = "&nbsp;" * 2 * level # Using &nbsp; for HTML spaces
231
+ # if print_keys and node.key is not None:
232
+ # html.append(f"{indent}{node.key}: {node.value}<br>")
233
+ # else:
234
+ # html.append(f"{indent}{node.value}<br>")
235
+
236
+ # for child in node.children.values():
237
+ # html.append(node_to_html(child, level + 1, print_keys))
238
+
239
+ # return "".join(html)
240
+
241
+ # tree_html = node_to_html(self.root)
242
+ # return f'<div class="tree-container">{tree_html}</div>{styles}'
243
+
244
+ def _add_to_docx(self, doc, node: TreeNode, level: int):
245
+ if node.value is not None:
246
+ if level == 0:
247
+ doc.add_heading(str(node.value), level=level + 1)
248
+ elif node.children: # If the node has children, it's not the last level
249
+ para = doc.add_paragraph(str(node.value))
250
+ para.style = f"Heading {level+1}"
251
+ else: # If the node has no children, it's the last level (body text)
252
+ para = doc.add_paragraph(str(node.value))
253
+ para.style = "Body Text"
254
+
255
+ # Process child nodes (moved outside the if block)
256
+ for child in node.children.values():
257
+ self._add_to_docx(doc, child, level + 1)
258
+
259
+
260
+ # Example usage (commented out)
261
+ """
262
+ from edsl.results.Dataset import Dataset
263
+
264
+ data = Dataset(
265
+ [
266
+ {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
267
+ {"country": ["US", "China", "France", "Canada", "Japan"]},
268
+ {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
269
+ {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
270
+ ]
271
+ )
272
+
273
+ tree = Tree(data)
274
+
275
+ try:
276
+ tree.construct_tree(["continent", "country", "city", "population"])
277
+ print("Tree without key names:")
278
+ tree.print_tree()
279
+ print("\nTree with key names:")
280
+ tree.print_tree(print_keys=True)
281
+ except ValueError as e:
282
+ print(f"Error: {e}")
283
+
284
+ # Demonstrating validation
285
+ try:
286
+ tree.construct_tree(["continent", "country", "invalid_key"])
287
+ except ValueError as e:
288
+ print(f"\nValidation Error: {e}")
289
+
290
+ tree = Tree(data)
291
+ tree.construct_tree(["continent", "country", "city", "population"])
292
+ tree.print_tree(print_keys=True)
293
+ tree.to_docx("tree_structure.docx")
294
+ print("DocX file 'tree_structure.docx' has been created.")
295
+ """