edsl 0.1.39.dev2__py3-none-any.whl → 0.1.39.dev3__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 (334) hide show
  1. edsl/Base.py +332 -385
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +49 -57
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +867 -1079
  7. edsl/agents/AgentList.py +413 -551
  8. edsl/agents/Invigilator.py +233 -285
  9. edsl/agents/InvigilatorBase.py +270 -254
  10. edsl/agents/PromptConstructor.py +354 -252
  11. edsl/agents/__init__.py +3 -2
  12. edsl/agents/descriptors.py +99 -99
  13. edsl/agents/prompt_helpers.py +129 -129
  14. edsl/auto/AutoStudy.py +117 -117
  15. edsl/auto/StageBase.py +230 -230
  16. edsl/auto/StageGenerateSurvey.py +178 -178
  17. edsl/auto/StageLabelQuestions.py +125 -125
  18. edsl/auto/StagePersona.py +61 -61
  19. edsl/auto/StagePersonaDimensionValueRanges.py +88 -88
  20. edsl/auto/StagePersonaDimensionValues.py +74 -74
  21. edsl/auto/StagePersonaDimensions.py +69 -69
  22. edsl/auto/StageQuestions.py +73 -73
  23. edsl/auto/SurveyCreatorPipeline.py +21 -21
  24. edsl/auto/utilities.py +224 -224
  25. edsl/base/Base.py +279 -279
  26. edsl/config.py +157 -177
  27. edsl/conversation/Conversation.py +290 -290
  28. edsl/conversation/car_buying.py +58 -59
  29. edsl/conversation/chips.py +95 -95
  30. edsl/conversation/mug_negotiation.py +81 -81
  31. edsl/conversation/next_speaker_utilities.py +93 -93
  32. edsl/coop/PriceFetcher.py +54 -54
  33. edsl/coop/__init__.py +2 -2
  34. edsl/coop/coop.py +1028 -1090
  35. edsl/coop/utils.py +131 -131
  36. edsl/data/Cache.py +555 -562
  37. edsl/data/CacheEntry.py +233 -230
  38. edsl/data/CacheHandler.py +149 -170
  39. edsl/data/RemoteCacheSync.py +78 -78
  40. edsl/data/SQLiteDict.py +292 -292
  41. edsl/data/__init__.py +4 -5
  42. edsl/data/orm.py +10 -10
  43. edsl/data_transfer_models.py +73 -74
  44. edsl/enums.py +175 -195
  45. edsl/exceptions/BaseException.py +21 -21
  46. edsl/exceptions/__init__.py +54 -54
  47. edsl/exceptions/agents.py +42 -54
  48. edsl/exceptions/cache.py +5 -5
  49. edsl/exceptions/configuration.py +16 -16
  50. edsl/exceptions/coop.py +10 -10
  51. edsl/exceptions/data.py +14 -14
  52. edsl/exceptions/general.py +34 -34
  53. edsl/exceptions/jobs.py +33 -33
  54. edsl/exceptions/language_models.py +63 -63
  55. edsl/exceptions/prompts.py +15 -15
  56. edsl/exceptions/questions.py +91 -109
  57. edsl/exceptions/results.py +29 -29
  58. edsl/exceptions/scenarios.py +22 -29
  59. edsl/exceptions/surveys.py +37 -37
  60. edsl/inference_services/AnthropicService.py +87 -84
  61. edsl/inference_services/AwsBedrock.py +120 -118
  62. edsl/inference_services/AzureAI.py +217 -215
  63. edsl/inference_services/DeepInfraService.py +18 -18
  64. edsl/inference_services/GoogleService.py +148 -139
  65. edsl/inference_services/GroqService.py +20 -20
  66. edsl/inference_services/InferenceServiceABC.py +147 -80
  67. edsl/inference_services/InferenceServicesCollection.py +97 -122
  68. edsl/inference_services/MistralAIService.py +123 -120
  69. edsl/inference_services/OllamaService.py +18 -18
  70. edsl/inference_services/OpenAIService.py +224 -221
  71. edsl/inference_services/PerplexityService.py +163 -160
  72. edsl/inference_services/TestService.py +89 -92
  73. edsl/inference_services/TogetherAIService.py +170 -170
  74. edsl/inference_services/models_available_cache.py +118 -118
  75. edsl/inference_services/rate_limits_cache.py +25 -25
  76. edsl/inference_services/registry.py +41 -41
  77. edsl/inference_services/write_available.py +10 -10
  78. edsl/jobs/Answers.py +56 -43
  79. edsl/jobs/Jobs.py +898 -757
  80. edsl/jobs/JobsChecks.py +147 -172
  81. edsl/jobs/JobsPrompts.py +268 -270
  82. edsl/jobs/JobsRemoteInferenceHandler.py +239 -287
  83. edsl/jobs/__init__.py +1 -1
  84. edsl/jobs/buckets/BucketCollection.py +63 -104
  85. edsl/jobs/buckets/ModelBuckets.py +65 -65
  86. edsl/jobs/buckets/TokenBucket.py +251 -283
  87. edsl/jobs/interviews/Interview.py +661 -358
  88. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  89. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  90. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  91. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  92. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  93. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  94. edsl/jobs/interviews/ReportErrors.py +66 -66
  95. edsl/jobs/interviews/interview_status_enum.py +9 -9
  96. edsl/jobs/runners/JobsRunnerAsyncio.py +466 -421
  97. edsl/jobs/runners/JobsRunnerStatus.py +330 -330
  98. edsl/jobs/tasks/QuestionTaskCreator.py +242 -244
  99. edsl/jobs/tasks/TaskCreators.py +64 -64
  100. edsl/jobs/tasks/TaskHistory.py +450 -449
  101. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  102. edsl/jobs/tasks/task_status_enum.py +163 -161
  103. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  104. edsl/jobs/tokens/TokenUsage.py +34 -34
  105. edsl/language_models/KeyLookup.py +30 -0
  106. edsl/language_models/LanguageModel.py +668 -571
  107. edsl/language_models/ModelList.py +155 -153
  108. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  109. edsl/language_models/__init__.py +3 -2
  110. edsl/language_models/fake_openai_call.py +15 -15
  111. edsl/language_models/fake_openai_service.py +61 -61
  112. edsl/language_models/registry.py +190 -180
  113. edsl/language_models/repair.py +156 -156
  114. edsl/language_models/unused/ReplicateBase.py +83 -0
  115. edsl/language_models/utilities.py +64 -65
  116. edsl/notebooks/Notebook.py +258 -263
  117. edsl/notebooks/__init__.py +1 -1
  118. edsl/prompts/Prompt.py +362 -352
  119. edsl/prompts/__init__.py +2 -2
  120. edsl/questions/AnswerValidatorMixin.py +289 -334
  121. edsl/questions/QuestionBase.py +664 -509
  122. edsl/questions/QuestionBaseGenMixin.py +161 -165
  123. edsl/questions/QuestionBasePromptsMixin.py +217 -221
  124. edsl/questions/QuestionBudget.py +227 -227
  125. edsl/questions/QuestionCheckBox.py +359 -359
  126. edsl/questions/QuestionExtract.py +182 -182
  127. edsl/questions/QuestionFreeText.py +114 -113
  128. edsl/questions/QuestionFunctional.py +166 -166
  129. edsl/questions/QuestionList.py +231 -229
  130. edsl/questions/QuestionMultipleChoice.py +286 -330
  131. edsl/questions/QuestionNumerical.py +153 -151
  132. edsl/questions/QuestionRank.py +324 -314
  133. edsl/questions/Quick.py +41 -41
  134. edsl/questions/RegisterQuestionsMeta.py +71 -71
  135. edsl/questions/ResponseValidatorABC.py +174 -200
  136. edsl/questions/SimpleAskMixin.py +73 -74
  137. edsl/questions/__init__.py +26 -27
  138. edsl/questions/compose_questions.py +98 -98
  139. edsl/questions/decorators.py +21 -21
  140. edsl/questions/derived/QuestionLikertFive.py +76 -76
  141. edsl/questions/derived/QuestionLinearScale.py +87 -90
  142. edsl/questions/derived/QuestionTopK.py +93 -93
  143. edsl/questions/derived/QuestionYesNo.py +82 -82
  144. edsl/questions/descriptors.py +413 -427
  145. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  146. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  147. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  148. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  149. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  150. edsl/questions/prompt_templates/question_list.jinja +17 -17
  151. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  152. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  153. edsl/questions/question_registry.py +177 -177
  154. edsl/questions/settings.py +12 -12
  155. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  156. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  157. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  158. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  159. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  160. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  161. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  162. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  163. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  164. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  165. edsl/questions/templates/list/question_presentation.jinja +5 -5
  166. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  167. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  168. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  169. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  170. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  171. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  172. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  173. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  174. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  175. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  176. edsl/results/CSSParameterizer.py +108 -108
  177. edsl/results/Dataset.py +424 -587
  178. edsl/results/DatasetExportMixin.py +731 -653
  179. edsl/results/DatasetTree.py +275 -295
  180. edsl/results/Result.py +465 -451
  181. edsl/results/Results.py +1165 -1172
  182. edsl/results/ResultsDBMixin.py +238 -0
  183. edsl/results/ResultsExportMixin.py +43 -45
  184. edsl/results/ResultsFetchMixin.py +33 -33
  185. edsl/results/ResultsGGMixin.py +121 -121
  186. edsl/results/ResultsToolsMixin.py +98 -98
  187. edsl/results/Selector.py +135 -145
  188. edsl/results/TableDisplay.py +198 -125
  189. edsl/results/__init__.py +2 -2
  190. edsl/results/table_display.css +77 -77
  191. edsl/results/tree_explore.py +115 -115
  192. edsl/scenarios/FileStore.py +632 -511
  193. edsl/scenarios/Scenario.py +601 -498
  194. edsl/scenarios/ScenarioHtmlMixin.py +64 -65
  195. edsl/scenarios/ScenarioJoin.py +127 -131
  196. edsl/scenarios/ScenarioList.py +1287 -1430
  197. edsl/scenarios/ScenarioListExportMixin.py +52 -45
  198. edsl/scenarios/ScenarioListPdfMixin.py +261 -239
  199. edsl/scenarios/__init__.py +4 -3
  200. edsl/shared.py +1 -1
  201. edsl/study/ObjectEntry.py +173 -173
  202. edsl/study/ProofOfWork.py +113 -113
  203. edsl/study/SnapShot.py +80 -80
  204. edsl/study/Study.py +528 -521
  205. edsl/study/__init__.py +4 -4
  206. edsl/surveys/DAG.py +148 -148
  207. edsl/surveys/Memory.py +31 -31
  208. edsl/surveys/MemoryPlan.py +244 -244
  209. edsl/surveys/Rule.py +326 -327
  210. edsl/surveys/RuleCollection.py +387 -385
  211. edsl/surveys/Survey.py +1801 -1229
  212. edsl/surveys/SurveyCSS.py +261 -273
  213. edsl/surveys/SurveyExportMixin.py +259 -259
  214. edsl/surveys/{SurveyFlowVisualization.py → SurveyFlowVisualizationMixin.py} +179 -181
  215. edsl/surveys/SurveyQualtricsImport.py +284 -284
  216. edsl/surveys/__init__.py +3 -5
  217. edsl/surveys/base.py +53 -53
  218. edsl/surveys/descriptors.py +56 -60
  219. edsl/surveys/instructions/ChangeInstruction.py +49 -48
  220. edsl/surveys/instructions/Instruction.py +65 -56
  221. edsl/surveys/instructions/InstructionCollection.py +77 -82
  222. edsl/templates/error_reporting/base.html +23 -23
  223. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  224. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  225. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  226. edsl/templates/error_reporting/interview_details.html +115 -115
  227. edsl/templates/error_reporting/interviews.html +19 -19
  228. edsl/templates/error_reporting/overview.html +4 -4
  229. edsl/templates/error_reporting/performance_plot.html +1 -1
  230. edsl/templates/error_reporting/report.css +73 -73
  231. edsl/templates/error_reporting/report.html +117 -117
  232. edsl/templates/error_reporting/report.js +25 -25
  233. edsl/tools/__init__.py +1 -1
  234. edsl/tools/clusters.py +192 -192
  235. edsl/tools/embeddings.py +27 -27
  236. edsl/tools/embeddings_plotting.py +118 -118
  237. edsl/tools/plotting.py +112 -112
  238. edsl/tools/summarize.py +18 -18
  239. edsl/utilities/SystemInfo.py +28 -28
  240. edsl/utilities/__init__.py +22 -22
  241. edsl/utilities/ast_utilities.py +25 -25
  242. edsl/utilities/data/Registry.py +6 -6
  243. edsl/utilities/data/__init__.py +1 -1
  244. edsl/utilities/data/scooter_results.json +1 -1
  245. edsl/utilities/decorators.py +77 -77
  246. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  247. edsl/utilities/interface.py +627 -627
  248. edsl/utilities/naming_utilities.py +263 -263
  249. edsl/utilities/repair_functions.py +28 -28
  250. edsl/utilities/restricted_python.py +70 -70
  251. edsl/utilities/utilities.py +424 -436
  252. {edsl-0.1.39.dev2.dist-info → edsl-0.1.39.dev3.dist-info}/LICENSE +21 -21
  253. {edsl-0.1.39.dev2.dist-info → edsl-0.1.39.dev3.dist-info}/METADATA +10 -12
  254. edsl-0.1.39.dev3.dist-info/RECORD +277 -0
  255. edsl/agents/QuestionInstructionPromptBuilder.py +0 -128
  256. edsl/agents/QuestionOptionProcessor.py +0 -172
  257. edsl/agents/QuestionTemplateReplacementsBuilder.py +0 -137
  258. edsl/coop/CoopFunctionsMixin.py +0 -15
  259. edsl/coop/ExpectedParrotKeyHandler.py +0 -125
  260. edsl/exceptions/inference_services.py +0 -5
  261. edsl/inference_services/AvailableModelCacheHandler.py +0 -184
  262. edsl/inference_services/AvailableModelFetcher.py +0 -209
  263. edsl/inference_services/ServiceAvailability.py +0 -135
  264. edsl/inference_services/data_structures.py +0 -62
  265. edsl/jobs/AnswerQuestionFunctionConstructor.py +0 -188
  266. edsl/jobs/FetchInvigilator.py +0 -40
  267. edsl/jobs/InterviewTaskManager.py +0 -98
  268. edsl/jobs/InterviewsConstructor.py +0 -48
  269. edsl/jobs/JobsComponentConstructor.py +0 -189
  270. edsl/jobs/JobsRemoteInferenceLogger.py +0 -239
  271. edsl/jobs/RequestTokenEstimator.py +0 -30
  272. edsl/jobs/buckets/TokenBucketAPI.py +0 -211
  273. edsl/jobs/buckets/TokenBucketClient.py +0 -191
  274. edsl/jobs/decorators.py +0 -35
  275. edsl/jobs/jobs_status_enums.py +0 -9
  276. edsl/jobs/loggers/HTMLTableJobLogger.py +0 -304
  277. edsl/language_models/ComputeCost.py +0 -63
  278. edsl/language_models/PriceManager.py +0 -127
  279. edsl/language_models/RawResponseHandler.py +0 -106
  280. edsl/language_models/ServiceDataSources.py +0 -0
  281. edsl/language_models/key_management/KeyLookup.py +0 -63
  282. edsl/language_models/key_management/KeyLookupBuilder.py +0 -273
  283. edsl/language_models/key_management/KeyLookupCollection.py +0 -38
  284. edsl/language_models/key_management/__init__.py +0 -0
  285. edsl/language_models/key_management/models.py +0 -131
  286. edsl/notebooks/NotebookToLaTeX.py +0 -142
  287. edsl/questions/ExceptionExplainer.py +0 -77
  288. edsl/questions/HTMLQuestion.py +0 -103
  289. edsl/questions/LoopProcessor.py +0 -149
  290. edsl/questions/QuestionMatrix.py +0 -265
  291. edsl/questions/ResponseValidatorFactory.py +0 -28
  292. edsl/questions/templates/matrix/__init__.py +0 -1
  293. edsl/questions/templates/matrix/answering_instructions.jinja +0 -5
  294. edsl/questions/templates/matrix/question_presentation.jinja +0 -20
  295. edsl/results/MarkdownToDocx.py +0 -122
  296. edsl/results/MarkdownToPDF.py +0 -111
  297. edsl/results/TextEditor.py +0 -50
  298. edsl/results/smart_objects.py +0 -96
  299. edsl/results/table_data_class.py +0 -12
  300. edsl/results/table_renderers.py +0 -118
  301. edsl/scenarios/ConstructDownloadLink.py +0 -109
  302. edsl/scenarios/DirectoryScanner.py +0 -96
  303. edsl/scenarios/DocumentChunker.py +0 -102
  304. edsl/scenarios/DocxScenario.py +0 -16
  305. edsl/scenarios/PdfExtractor.py +0 -40
  306. edsl/scenarios/ScenarioSelector.py +0 -156
  307. edsl/scenarios/file_methods.py +0 -85
  308. edsl/scenarios/handlers/__init__.py +0 -13
  309. edsl/scenarios/handlers/csv.py +0 -38
  310. edsl/scenarios/handlers/docx.py +0 -76
  311. edsl/scenarios/handlers/html.py +0 -37
  312. edsl/scenarios/handlers/json.py +0 -111
  313. edsl/scenarios/handlers/latex.py +0 -5
  314. edsl/scenarios/handlers/md.py +0 -51
  315. edsl/scenarios/handlers/pdf.py +0 -68
  316. edsl/scenarios/handlers/png.py +0 -39
  317. edsl/scenarios/handlers/pptx.py +0 -105
  318. edsl/scenarios/handlers/py.py +0 -294
  319. edsl/scenarios/handlers/sql.py +0 -313
  320. edsl/scenarios/handlers/sqlite.py +0 -149
  321. edsl/scenarios/handlers/txt.py +0 -33
  322. edsl/surveys/ConstructDAG.py +0 -92
  323. edsl/surveys/EditSurvey.py +0 -221
  324. edsl/surveys/InstructionHandler.py +0 -100
  325. edsl/surveys/MemoryManagement.py +0 -72
  326. edsl/surveys/RuleManager.py +0 -172
  327. edsl/surveys/Simulator.py +0 -75
  328. edsl/surveys/SurveyToApp.py +0 -141
  329. edsl/utilities/PrettyList.py +0 -56
  330. edsl/utilities/is_notebook.py +0 -18
  331. edsl/utilities/is_valid_variable_name.py +0 -11
  332. edsl/utilities/remove_edsl_version.py +0 -24
  333. edsl-0.1.39.dev2.dist-info/RECORD +0 -352
  334. {edsl-0.1.39.dev2.dist-info → edsl-0.1.39.dev3.dist-info}/WHEEL +0 -0
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,5 +1,4 @@
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
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
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)