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
edsl/data/SQLiteDict.py CHANGED
@@ -1,292 +1,292 @@
1
- from __future__ import annotations
2
- import json
3
- from typing import Any, Generator, Optional, Union
4
-
5
- from edsl.config import CONFIG
6
- from edsl.data.CacheEntry import CacheEntry
7
- from edsl.data.orm import Base, Data
8
-
9
-
10
- class SQLiteDict:
11
- """
12
- A dictionary-like object that is an interface for an local database.
13
- - You can use SQLiteDict as a regular dictionary.
14
- - Supports only SQLite for now.
15
- """
16
-
17
- def __init__(self, db_path: Optional[str] = None):
18
- """
19
-
20
- >>> temp_db_path = SQLiteDict._get_temp_path()
21
- >>> SQLiteDict(f"sqlite:///{temp_db_path}") # Use the temp file for SQLite
22
- SQLiteDict(db_path='...')
23
- >>> import os; os.unlink(temp_db_path) # Clean up the temp file after the test
24
-
25
- """
26
- from sqlalchemy.exc import SQLAlchemyError
27
- from sqlalchemy.orm import sessionmaker
28
- from sqlalchemy import create_engine
29
-
30
- self.db_path = db_path or CONFIG.get("EDSL_DATABASE_PATH")
31
- if not self.db_path.startswith("sqlite:///"):
32
- self.db_path = f"sqlite:///{self.db_path}"
33
- try:
34
- from edsl.data.orm import Base, Data
35
-
36
- self.engine = create_engine(self.db_path, echo=False, future=True)
37
- Base.metadata.create_all(self.engine)
38
- self.Session = sessionmaker(bind=self.engine)
39
- except SQLAlchemyError as e:
40
- raise Exception(
41
- f"""Database initialization error: {e}. The attempted DB path was {db_path}"""
42
- ) from e
43
-
44
- @classmethod
45
- def _get_temp_path(self):
46
- import tempfile
47
- import os
48
-
49
- _, temp_db_path = tempfile.mkstemp(suffix=".db")
50
- return temp_db_path
51
-
52
- def __setitem__(self, key: str, value: CacheEntry) -> None:
53
- """
54
- Stores a key-value pair.
55
-
56
- >>> d = SQLiteDict.example()
57
- >>> d["foo"] = CacheEntry.example()
58
- """
59
- if not isinstance(value, CacheEntry):
60
- raise ValueError(f"Value must be a CacheEntry object (got {type(value)}).")
61
- with self.Session() as db:
62
- from edsl.data.orm import Base, Data
63
-
64
- db.merge(Data(key=key, value=json.dumps(value.to_dict())))
65
- db.commit()
66
-
67
- def __getitem__(self, key: str) -> CacheEntry:
68
- """
69
- Gets a value for a given key.
70
- - Raises a KeyError if the key is not found.
71
-
72
- >>> d = SQLiteDict.example()
73
- >>> d["foo"] = CacheEntry.example()
74
- >>> d["foo"] == CacheEntry.example()
75
- True
76
- """
77
- with self.Session() as db:
78
- from edsl.data.orm import Base, Data
79
-
80
- value = db.query(Data).filter_by(key=key).first()
81
- if not value:
82
- raise KeyError(f"Key '{key}' not found.")
83
- return CacheEntry.from_dict(json.loads(value.value))
84
-
85
- def get(self, key: str, default: Optional[Any] = None) -> Union[CacheEntry, Any]:
86
- """
87
- Gets the value for a given key
88
- - Returns the `default` value if the key is not found.
89
-
90
- >>> d = SQLiteDict.example()
91
- >>> d.get("foo", "bar")
92
- 'bar'
93
- """
94
- try:
95
- return self[key]
96
- except KeyError:
97
- return default
98
-
99
- def __bool__(self) -> bool:
100
- """This is so likes like
101
- self.data = data or {} 'work' as expected
102
- """
103
- return True
104
-
105
- def update(
106
- self,
107
- new_d: Union[dict, SQLiteDict],
108
- overwrite: Optional[bool] = False,
109
- max_batch_size: Optional[int] = 100,
110
- ) -> None:
111
- """
112
- Update the dictionary with the values from another dictionary.
113
-
114
- :param new_d: The dictionary to update the current dictionary with.
115
- :param overwrite: If `overwrite` is False, existing values will not be overwritten.
116
- :param max_batch_size: The maximum number of items to update in a single transaction.
117
-
118
- - If `overwrite` is True, existing values will be overwritten.
119
-
120
- >>> d = SQLiteDict.example()
121
- >>> d.update({"foo": CacheEntry.example()})
122
- >>> d["foo"] == CacheEntry.example()
123
- True
124
- """
125
- if not (isinstance(new_d, dict) or isinstance(new_d, SQLiteDict)):
126
- raise ValueError(
127
- f"new_d must be a dict or SQLiteDict object (got {type(new_d)})"
128
- )
129
- current_batch = 0
130
- with self.Session() as db:
131
- for key, value in new_d.items():
132
- if current_batch == max_batch_size:
133
- db.commit()
134
- current_batch = 0
135
- current_batch += 1
136
- if key in self:
137
- if overwrite:
138
- db.merge(Data(key=key, value=json.dumps(value.to_dict())))
139
- else:
140
- db.merge(Data(key=key, value=json.dumps(value.to_dict())))
141
- db.commit()
142
-
143
- def values(self) -> Generator[CacheEntry, None, None]:
144
- """
145
- Returns a generator that yields the values in the cache.
146
-
147
- >>> d = SQLiteDict.example()
148
- >>> d["foo"] = CacheEntry.example()
149
- >>> list(d.values()) == [CacheEntry.example()]
150
- True
151
- """
152
- with self.Session() as db:
153
- for instance in db.query(Data).all():
154
- yield CacheEntry.from_dict(json.loads(instance.value))
155
-
156
- def items(self) -> Generator[tuple[str, CacheEntry], None, None]:
157
- """
158
- Returns a generator that yields the items in the cache.
159
-
160
- >>> d = SQLiteDict.example()
161
- >>> d["foo"] = CacheEntry.example()
162
- >>> list(d.items()) == [("foo", CacheEntry.example())]
163
- True
164
- """
165
- with self.Session() as db:
166
- for instance in db.query(Data).all():
167
- yield (instance.key, CacheEntry.from_dict(json.loads(instance.value)))
168
-
169
- def to_dict(self):
170
- """
171
- Returns the cache as a dictionary.
172
-
173
- >>> d = SQLiteDict.example()
174
- >>> d["foo"] = CacheEntry.example()
175
- >>> d.to_dict() == {"foo": CacheEntry.example()}
176
- True
177
- """
178
- return dict(self.items())
179
-
180
- def __delitem__(self, key: str) -> None:
181
- """
182
- Deletes the value for a given key.
183
-
184
- >>> d = SQLiteDict.example()
185
- >>> d["foo"] = CacheEntry.example()
186
- >>> del d["foo"]
187
- >>> d.get("foo", "missing")
188
- 'missing'
189
- """
190
- with self.Session() as db:
191
- instance = db.query(Data).filter_by(key=key).one_or_none()
192
- if instance:
193
- db.delete(instance)
194
- db.commit()
195
- else:
196
- raise KeyError(f"Key '{key}' not found.")
197
-
198
- def __contains__(self, key: str) -> bool:
199
- """
200
- Checks if the dict contains the given key.
201
-
202
- >>> d = SQLiteDict.example()
203
- >>> d["foo"] = CacheEntry.example()
204
- >>> "foo" in d
205
- True
206
- >>> "bar" in d
207
- False
208
- """
209
- with self.Session() as db:
210
- return db.query(Data).filter_by(key=key).first() is not None
211
-
212
- def __iter__(self) -> Generator[str, None, None]:
213
- """
214
- Returns a generator that yields the keys in the dict.
215
-
216
- >>> d = SQLiteDict.example()
217
- >>> d["foo"] = CacheEntry.example()
218
- >>> list(iter(d)) == ["foo"]
219
- True
220
- """
221
- with self.Session() as db:
222
- for instance in db.query(Data).all():
223
- yield instance.key
224
-
225
- def __len__(self) -> int:
226
- """
227
- Returns the number of items in the cache.
228
-
229
- >>> d = SQLiteDict.example()
230
- >>> len(d)
231
- 0
232
- >>> d["foo"] = CacheEntry.example()
233
- >>> len(d)
234
- 1
235
- """
236
- with self.Session() as db:
237
- return db.query(Data).count()
238
-
239
- def keys(self) -> Generator[str, None, None]:
240
- """
241
- Returns a generator that yields the keys in the cache.
242
-
243
- >>> d = SQLiteDict.example()
244
- >>> d["foo"] = CacheEntry.example()
245
- >>> list(d.keys()) == ["foo"]
246
- True
247
- """
248
- return self.__iter__()
249
-
250
- def __repr__(self) -> str:
251
- return f"{self.__class__.__name__}(db_path={self.db_path!r})"
252
-
253
- @classmethod
254
- def example(cls) -> SQLiteDict:
255
- """
256
- Returns an example SQLiteDict object.
257
- - The example SQLiteDict is empty and stored in memory.
258
-
259
- >>> SQLiteDict.example()
260
- SQLiteDict(db_path='sqlite:///:memory:')
261
- """
262
- return cls(db_path="sqlite:///:memory:")
263
-
264
-
265
- def main():
266
- from edsl.data.CacheEntry import CacheEntry
267
- from edsl.data.SQLiteDict import SQLiteDict
268
-
269
- d = SQLiteDict.example()
270
- d["foo"] = CacheEntry.example()
271
- d["foo"]
272
- d.get("foo")
273
- d.get("poo")
274
- d.get("poo", "not found")
275
- d.update({"poo": CacheEntry.example()})
276
- d["poo"]
277
- len(d)
278
- list(d.keys())
279
- list(d.values())
280
- list(d.items())
281
- assert "poo" in d
282
- assert "loo" not in d
283
- del d["poo"]
284
- assert len(d) == 1
285
- repr(d)
286
- d
287
-
288
-
289
- if __name__ == "__main__":
290
- import doctest
291
-
292
- doctest.testmod()
1
+ from __future__ import annotations
2
+ import json
3
+ from typing import Any, Generator, Optional, Union
4
+
5
+ from edsl.config import CONFIG
6
+ from edsl.data.CacheEntry import CacheEntry
7
+ from edsl.data.orm import Base, Data
8
+
9
+
10
+ class SQLiteDict:
11
+ """
12
+ A dictionary-like object that is an interface for an local database.
13
+ - You can use SQLiteDict as a regular dictionary.
14
+ - Supports only SQLite for now.
15
+ """
16
+
17
+ def __init__(self, db_path: Optional[str] = None):
18
+ """
19
+
20
+ >>> temp_db_path = SQLiteDict._get_temp_path()
21
+ >>> SQLiteDict(f"sqlite:///{temp_db_path}") # Use the temp file for SQLite
22
+ SQLiteDict(db_path='...')
23
+ >>> import os; os.unlink(temp_db_path) # Clean up the temp file after the test
24
+
25
+ """
26
+ from sqlalchemy.exc import SQLAlchemyError
27
+ from sqlalchemy.orm import sessionmaker
28
+ from sqlalchemy import create_engine
29
+
30
+ self.db_path = db_path or CONFIG.get("EDSL_DATABASE_PATH")
31
+ if not self.db_path.startswith("sqlite:///"):
32
+ self.db_path = f"sqlite:///{self.db_path}"
33
+ try:
34
+ from edsl.data.orm import Base, Data
35
+
36
+ self.engine = create_engine(self.db_path, echo=False, future=True)
37
+ Base.metadata.create_all(self.engine)
38
+ self.Session = sessionmaker(bind=self.engine)
39
+ except SQLAlchemyError as e:
40
+ raise Exception(
41
+ f"""Database initialization error: {e}. The attempted DB path was {db_path}"""
42
+ ) from e
43
+
44
+ @classmethod
45
+ def _get_temp_path(self):
46
+ import tempfile
47
+ import os
48
+
49
+ _, temp_db_path = tempfile.mkstemp(suffix=".db")
50
+ return temp_db_path
51
+
52
+ def __setitem__(self, key: str, value: CacheEntry) -> None:
53
+ """
54
+ Stores a key-value pair.
55
+
56
+ >>> d = SQLiteDict.example()
57
+ >>> d["foo"] = CacheEntry.example()
58
+ """
59
+ if not isinstance(value, CacheEntry):
60
+ raise ValueError(f"Value must be a CacheEntry object (got {type(value)}).")
61
+ with self.Session() as db:
62
+ from edsl.data.orm import Base, Data
63
+
64
+ db.merge(Data(key=key, value=json.dumps(value.to_dict())))
65
+ db.commit()
66
+
67
+ def __getitem__(self, key: str) -> CacheEntry:
68
+ """
69
+ Gets a value for a given key.
70
+ - Raises a KeyError if the key is not found.
71
+
72
+ >>> d = SQLiteDict.example()
73
+ >>> d["foo"] = CacheEntry.example()
74
+ >>> d["foo"] == CacheEntry.example()
75
+ True
76
+ """
77
+ with self.Session() as db:
78
+ from edsl.data.orm import Base, Data
79
+
80
+ value = db.query(Data).filter_by(key=key).first()
81
+ if not value:
82
+ raise KeyError(f"Key '{key}' not found.")
83
+ return CacheEntry.from_dict(json.loads(value.value))
84
+
85
+ def get(self, key: str, default: Optional[Any] = None) -> Union[CacheEntry, Any]:
86
+ """
87
+ Gets the value for a given key
88
+ - Returns the `default` value if the key is not found.
89
+
90
+ >>> d = SQLiteDict.example()
91
+ >>> d.get("foo", "bar")
92
+ 'bar'
93
+ """
94
+ try:
95
+ return self[key]
96
+ except KeyError:
97
+ return default
98
+
99
+ def __bool__(self) -> bool:
100
+ """This is so likes like
101
+ self.data = data or {} 'work' as expected
102
+ """
103
+ return True
104
+
105
+ def update(
106
+ self,
107
+ new_d: Union[dict, SQLiteDict],
108
+ overwrite: Optional[bool] = False,
109
+ max_batch_size: Optional[int] = 100,
110
+ ) -> None:
111
+ """
112
+ Update the dictionary with the values from another dictionary.
113
+
114
+ :param new_d: The dictionary to update the current dictionary with.
115
+ :param overwrite: If `overwrite` is False, existing values will not be overwritten.
116
+ :param max_batch_size: The maximum number of items to update in a single transaction.
117
+
118
+ - If `overwrite` is True, existing values will be overwritten.
119
+
120
+ >>> d = SQLiteDict.example()
121
+ >>> d.update({"foo": CacheEntry.example()})
122
+ >>> d["foo"] == CacheEntry.example()
123
+ True
124
+ """
125
+ if not (isinstance(new_d, dict) or isinstance(new_d, SQLiteDict)):
126
+ raise ValueError(
127
+ f"new_d must be a dict or SQLiteDict object (got {type(new_d)})"
128
+ )
129
+ current_batch = 0
130
+ with self.Session() as db:
131
+ for key, value in new_d.items():
132
+ if current_batch == max_batch_size:
133
+ db.commit()
134
+ current_batch = 0
135
+ current_batch += 1
136
+ if key in self:
137
+ if overwrite:
138
+ db.merge(Data(key=key, value=json.dumps(value.to_dict())))
139
+ else:
140
+ db.merge(Data(key=key, value=json.dumps(value.to_dict())))
141
+ db.commit()
142
+
143
+ def values(self) -> Generator[CacheEntry, None, None]:
144
+ """
145
+ Returns a generator that yields the values in the cache.
146
+
147
+ >>> d = SQLiteDict.example()
148
+ >>> d["foo"] = CacheEntry.example()
149
+ >>> list(d.values()) == [CacheEntry.example()]
150
+ True
151
+ """
152
+ with self.Session() as db:
153
+ for instance in db.query(Data).all():
154
+ yield CacheEntry.from_dict(json.loads(instance.value))
155
+
156
+ def items(self) -> Generator[tuple[str, CacheEntry], None, None]:
157
+ """
158
+ Returns a generator that yields the items in the cache.
159
+
160
+ >>> d = SQLiteDict.example()
161
+ >>> d["foo"] = CacheEntry.example()
162
+ >>> list(d.items()) == [("foo", CacheEntry.example())]
163
+ True
164
+ """
165
+ with self.Session() as db:
166
+ for instance in db.query(Data).all():
167
+ yield (instance.key, CacheEntry.from_dict(json.loads(instance.value)))
168
+
169
+ def to_dict(self):
170
+ """
171
+ Returns the cache as a dictionary.
172
+
173
+ >>> d = SQLiteDict.example()
174
+ >>> d["foo"] = CacheEntry.example()
175
+ >>> d.to_dict() == {"foo": CacheEntry.example()}
176
+ True
177
+ """
178
+ return dict(self.items())
179
+
180
+ def __delitem__(self, key: str) -> None:
181
+ """
182
+ Deletes the value for a given key.
183
+
184
+ >>> d = SQLiteDict.example()
185
+ >>> d["foo"] = CacheEntry.example()
186
+ >>> del d["foo"]
187
+ >>> d.get("foo", "missing")
188
+ 'missing'
189
+ """
190
+ with self.Session() as db:
191
+ instance = db.query(Data).filter_by(key=key).one_or_none()
192
+ if instance:
193
+ db.delete(instance)
194
+ db.commit()
195
+ else:
196
+ raise KeyError(f"Key '{key}' not found.")
197
+
198
+ def __contains__(self, key: str) -> bool:
199
+ """
200
+ Checks if the dict contains the given key.
201
+
202
+ >>> d = SQLiteDict.example()
203
+ >>> d["foo"] = CacheEntry.example()
204
+ >>> "foo" in d
205
+ True
206
+ >>> "bar" in d
207
+ False
208
+ """
209
+ with self.Session() as db:
210
+ return db.query(Data).filter_by(key=key).first() is not None
211
+
212
+ def __iter__(self) -> Generator[str, None, None]:
213
+ """
214
+ Returns a generator that yields the keys in the dict.
215
+
216
+ >>> d = SQLiteDict.example()
217
+ >>> d["foo"] = CacheEntry.example()
218
+ >>> list(iter(d)) == ["foo"]
219
+ True
220
+ """
221
+ with self.Session() as db:
222
+ for instance in db.query(Data).all():
223
+ yield instance.key
224
+
225
+ def __len__(self) -> int:
226
+ """
227
+ Returns the number of items in the cache.
228
+
229
+ >>> d = SQLiteDict.example()
230
+ >>> len(d)
231
+ 0
232
+ >>> d["foo"] = CacheEntry.example()
233
+ >>> len(d)
234
+ 1
235
+ """
236
+ with self.Session() as db:
237
+ return db.query(Data).count()
238
+
239
+ def keys(self) -> Generator[str, None, None]:
240
+ """
241
+ Returns a generator that yields the keys in the cache.
242
+
243
+ >>> d = SQLiteDict.example()
244
+ >>> d["foo"] = CacheEntry.example()
245
+ >>> list(d.keys()) == ["foo"]
246
+ True
247
+ """
248
+ return self.__iter__()
249
+
250
+ def __repr__(self) -> str:
251
+ return f"{self.__class__.__name__}(db_path={self.db_path!r})"
252
+
253
+ @classmethod
254
+ def example(cls) -> SQLiteDict:
255
+ """
256
+ Returns an example SQLiteDict object.
257
+ - The example SQLiteDict is empty and stored in memory.
258
+
259
+ >>> SQLiteDict.example()
260
+ SQLiteDict(db_path='sqlite:///:memory:')
261
+ """
262
+ return cls(db_path="sqlite:///:memory:")
263
+
264
+
265
+ def main():
266
+ from edsl.data.CacheEntry import CacheEntry
267
+ from edsl.data.SQLiteDict import SQLiteDict
268
+
269
+ d = SQLiteDict.example()
270
+ d["foo"] = CacheEntry.example()
271
+ d["foo"]
272
+ d.get("foo")
273
+ d.get("poo")
274
+ d.get("poo", "not found")
275
+ d.update({"poo": CacheEntry.example()})
276
+ d["poo"]
277
+ len(d)
278
+ list(d.keys())
279
+ list(d.values())
280
+ list(d.items())
281
+ assert "poo" in d
282
+ assert "loo" not in d
283
+ del d["poo"]
284
+ assert len(d) == 1
285
+ repr(d)
286
+ d
287
+
288
+
289
+ if __name__ == "__main__":
290
+ import doctest
291
+
292
+ doctest.testmod()
edsl/data/__init__.py CHANGED
@@ -1,4 +1,5 @@
1
- from edsl.data.CacheEntry import CacheEntry
2
- from edsl.data.SQLiteDict import SQLiteDict
3
- from edsl.data.Cache import Cache
4
- from edsl.data.CacheHandler import CacheHandler
1
+ # from edsl.data.CacheEntry import CacheEntry
2
+ # from edsl.data.SQLiteDict import SQLiteDict
3
+ from edsl.data.Cache import Cache
4
+
5
+ # from edsl.data.CacheHandler import CacheHandler
edsl/data/orm.py CHANGED
@@ -1,10 +1,10 @@
1
- from sqlalchemy import Column, String
2
- from sqlalchemy.ext.declarative import declarative_base
3
-
4
- Base = declarative_base()
5
-
6
-
7
- class Data(Base):
8
- __tablename__ = "data"
9
- key = Column(String, primary_key=True)
10
- value = Column(String)
1
+ from sqlalchemy import Column, String
2
+ from sqlalchemy.ext.declarative import declarative_base
3
+
4
+ Base = declarative_base()
5
+
6
+
7
+ class Data(Base):
8
+ __tablename__ = "data"
9
+ key = Column(String, primary_key=True)
10
+ value = Column(String)