toolchemy 0.2.191__tar.gz → 0.2.193__tar.gz

{toolchemy-0.2.191 → toolchemy-0.2.193}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: toolchemy
-Version: 0.2.191
+Version: 0.2.193
 Summary: a set of auxiliary programming tools
 License-File: LICENSE
 Author: Cyprian Nosek
@@ -8,7 +8,7 @@ Author-email: cyprian.nosek@protonmail.com
 Requires-Python: >=3.12,<3.13
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: click (>=8.2.1,<9.0.0)
+Requires-Dist: click (>=8.2.1,<8.3.0)
 Requires-Dist: colorlog (>=6.10.1,<7.0.0)
 Requires-Dist: diskcache (>=5.6.3,<6.0.0)
 Requires-Dist: google-genai (>=1.55.0,<2.0.0)

{toolchemy-0.2.191 → toolchemy-0.2.193}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "toolchemy"
-version = "0.2.191"
+version = "0.2.193"
 description = "a set of auxiliary programming tools"
 authors = ["Cyprian Nosek <cyprian.nosek@protonmail.com>"]
 
@@ -21,7 +21,7 @@ diskcache = "^5.6.3"
 google-genai = "^1.55.0"
 colorlog = "^6.10.1"
 jsonschema = "^4.26.0"
-click = "^8.2.1"
+click = "~8.2.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.5"

{toolchemy-0.2.191 → toolchemy-0.2.193}/toolchemy/ai/prompter.py

@@ -67,8 +67,8 @@ class PrompterMLflow(PrompterBase):
 
         self._registry_store_uri = f"sqlite:///{registry_store_dir}/registry.db"
         self._tracking_uri = f"sqlite:///{registry_store_dir}/tracking.db"
-        mlflow.set_tracking_uri(self._tracking_uri)
-        mlflow.set_registry_uri(self._registry_store_uri)
+
+        self._client = mlflow.tracking.MlflowClient(tracking_uri=self._tracking_uri, registry_uri=self._registry_store_uri)
 
         self._logger.info(f"Prompter-MLflow initialized")
         self._logger.info(f"> tracking store uri: {self._tracking_uri}")
@@ -79,11 +79,12 @@ class PrompterMLflow(PrompterBase):
 
         cache_key = self._cacher.create_cache_key(["prompt_render", prompt_uri], [variables])
         if self._cacher.exists(cache_key):
+            self._logger.debug(f"Retrieving from the cache")
             return self._cacher.get(cache_key)
 
         self._logger.debug(f"Rendering prompt: '{name}' (version: '{version}') -> prompt uri: '{prompt_uri}'")
 
-        prompt_template = mlflow.genai.load_prompt(prompt_uri)
+        prompt_template = self._client.load_prompt(prompt_uri)
         prompt = prompt_template.format(**variables)
 
         self._cacher.set(cache_key, prompt)
@@ -99,16 +100,16 @@ class PrompterMLflow(PrompterBase):
 
         self._logger.debug(f"Getting prompt template: '{name}' (version: '{version}') -> prompt uri: '{prompt_uri}'")
 
-        prompt_template = mlflow.genai.load_prompt(prompt_uri)
+        prompt_template = self._client.load_prompt(prompt_uri)
 
         self._cacher.set(cache_key, prompt_template.template)
 
         return prompt_template.template
 
     def create_template(self, name: str, template: str, overwrite: bool = False):
-        if mlflow.genai.load_prompt(name_or_uri=name, allow_missing=True) and not overwrite:
+        if self._client.load_prompt(name_or_uri=name, allow_missing=True) and not overwrite:
             return
-        mlflow.genai.register_prompt(name=name, template=template)
+        self._client.register_prompt(name=name, template=template)
 
     def _build_prompt_uri(self, name: str, version: str | int | None = None) -> str:
         prompt_version = self._prompt_version(name, version_mapping=version)
@@ -122,3 +123,15 @@ class PrompterMLflow(PrompterBase):
 
     def _latest_value(self) -> str | None:
         return "latest"
+
+
+def testing():
+    import logging
+    locations = Locations()
+    prompter = PrompterMLflow(registry_store_dir=locations.in_resources("tests/prompts_mlflow"), log_level=logging.DEBUG)
+    p = prompter.render("create_model_task_prompt")
+    print(p)
+
+
+if __name__ == "__main__":
+    testing()

{toolchemy-0.2.191 → toolchemy-0.2.193}/toolchemy/ai/trackers/mlflow_tracker.py

@@ -14,21 +14,19 @@ from toolchemy.utils.logger import get_logger
 
 
 class MLFlowTracker(TrackerBase):
-    def __init__(self, tracking_uri: str, experiment_name: str, with_artifact_logging=True,
+    def __init__(self, tracking_uri: str, experiment_name: str, with_artifact_logging=True, registry_uri: str | None = None,
                  tracking_client: MlflowClient | None = None):
         super().__init__(experiment_name, with_artifact_logging)
-        self._tracking_client = None
+        self._client = None
         self._active_run = None
         self._active_run_id = None
         self._experiment_id = None
         self._reset_run()
 
         if tracking_client:
-            self._tracking_client = tracking_client
+            self._client = tracking_client
         else:
-            self._tracking_client = self._build_tracking_client(tracking_uri)
-
-        mlflow.set_tracking_uri(tracking_uri)
+            self._client = mlflow.tracking.MlflowClient(tracking_uri=tracking_uri, registry_uri=registry_uri)
 
         logger = get_logger()
         logger.info(f"Mlflow tracker created")
@@ -67,25 +65,25 @@ class MLFlowTracker(TrackerBase):
 
         tags = resolve_tags(user_specified_tags)
 
-        experiment = self._tracking_client.get_experiment_by_name(self._experiment_name)
+        experiment = self._client.get_experiment_by_name(self._experiment_name)
         if experiment:
             experiment_comment_msg = "(already exists)"
             self._logger.debug(f"Experiment '{self._experiment_name}' already exists")
             self._experiment_id = experiment.experiment_id
             if experiment.lifecycle_stage == "deleted":
                 self._logger.info(f"Restoring deleted experiment")
-                self._tracking_client.restore_experiment(self._experiment_id)
+                self._client.restore_experiment(self._experiment_id)
         else:
             experiment_comment_msg = "(does not exist, creating a new one)"
             self._logger.debug(f"Experiment '{self._experiment_name}' does not exist, creating a new one")
-            self._experiment_id = self._tracking_client.create_experiment(self._experiment_name)
+            self._experiment_id = self._client.create_experiment(self._experiment_name)
 
         self._logger.info(f"Starting the experiment tracking")
         self._logger.info(f"> experiment name: {self._experiment_name} {experiment_comment_msg}")
         self._logger.info(f"> experiment id: {self._experiment_id}")
         self._logger.info(f"> run name: {run_name}")
 
-        self._active_run = self._tracking_client.create_run(
+        self._active_run = self._client.create_run(
             experiment_id=self._experiment_id,
             start_time=None,
             run_name=run_name,
@@ -98,7 +96,7 @@ class MLFlowTracker(TrackerBase):
             return
 
         status = RunStatus.to_string(RunStatus.FINISHED)
-        self._tracking_client.set_terminated(self._active_run_id, status)
+        self._client.set_terminated(self._active_run_id, status)
         self._reset_run()
 
     def log(self, name: str, value: Any):
@@ -106,7 +104,7 @@ class MLFlowTracker(TrackerBase):
             return
 
         if isinstance(value, dict):
-            self._tracking_client.log_dict(self._active_run_id, value, f"{name}.json")
+            self._client.log_dict(self._active_run_id, value, f"{name}.json")
 
         raise ValueError(f"Unsupported logged object type: {type(value)}")
 
@@ -115,7 +113,7 @@ class MLFlowTracker(TrackerBase):
             return
 
         self._store_param(name, value)
-        self._tracking_client.log_param(self._active_run_id, name, value)
+        self._client.log_param(self._active_run_id, name, value)
 
     def log_params(self, params: Dict[str, Any]):
         if self._disabled:
@@ -130,16 +128,16 @@ class MLFlowTracker(TrackerBase):
                 params_to_store.append(Param(key, str(value)))
             self._store_param(key, value)
 
-        self._tracking_client.log_batch(self._active_run_id, [], params_to_store)
+        self._client.log_batch(self._active_run_id, [], params_to_store)
 
     def log_text(self, name: str, value: str):
         if self._disabled:
             return
         try:
-            self._tracking_client.log_text(run_id=self._active_run_id, text=value, artifact_file=name)
+            self._client.log_text(run_id=self._active_run_id, text=value, artifact_file=name)
         except MlflowException as e:
             self._logger.error(f"An error occurred during text logging: {e}")
-            self._logger.error(f"> tracking uri: {self._tracking_client.tracking_uri}")
+            self._logger.error(f"> tracking uri: {self._client.tracking_uri}")
             self._logger.error(f"> artifact uri: {self._active_run.info.artifact_uri}")
             raise e
 
@@ -148,7 +146,7 @@ class MLFlowTracker(TrackerBase):
             return
 
         metric_value = self._store_metric(name, value, metric_metadata)
-        self._tracking_client.log_metric(self._active_run_id, name, metric_value, step)
+        self._client.log_metric(self._active_run_id, name, metric_value, step)
 
     def log_metrics(self, metrics: Dict[str, float | list], step: Optional[int] = None):
         if self._disabled:
@@ -165,32 +163,27 @@ class MLFlowTracker(TrackerBase):
                 metric_value = self._store_metric(k, value)
                 metrics_to_store.append(Metric(k, metric_value, timestamp, step or 0))
 
-        self._tracking_client.log_batch(self._active_run_id, metrics_to_store)
+        self._client.log_batch(self._active_run_id, metrics_to_store)
 
     def log_artifact(self, artifact_path: str, save_dir: str = None):
         if self._disabled:
             return
 
         if self._artifact_logging:
-            self._tracking_client.log_artifact(self._active_run_id, artifact_path, save_dir)
+            self._client.log_artifact(self._active_run_id, artifact_path, save_dir)
 
     def log_figure(self, figure, save_path: str):
         if self._disabled:
             return
-        self._tracking_client.log_figure(self._active_run_id, figure, save_path)
+        self._client.log_figure(self._active_run_id, figure, save_path)
 
     def set_run_tag(self, name: str, value: str | int | float):
         self._store_tag(name, value, run_name=self.run_name)
-        self._tracking_client.set_tag(self._active_run_id, name, value)
+        self._client.set_tag(self._active_run_id, name, value)
 
     def set_experiment_tag(self, name: str, value: str | int | float):
         self._store_tag(name, value)
-        self._tracking_client.set_experiment_tag(self._experiment_id, name, value)
-
-    @staticmethod
-    def _build_tracking_client(tracking_uri: str) -> MlflowClient:
-        tracking_client = mlflow.tracking.MlflowClient(tracking_uri=tracking_uri, registry_uri=tracking_uri)
-        return tracking_client
+        self._client.set_experiment_tag(self._experiment_id, name, value)
 
     def _reset_run(self):
         self._active_run = None