toolchemy 0.2.192__py3-none-any.whl → 0.2.194__py3-none-any.whl

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/ai/trackers/common.py

@@ -11,11 +11,21 @@ class ITracker(ABC):
     def experiment_name(self) -> str:
         pass
 
+    @property
+    @abstractmethod
+    def experiment_id(self) -> str:
+        pass
+
     @property
     @abstractmethod
     def run_name(self) -> str:
         pass
 
+    @property
+    @abstractmethod
+    def run_id(self) -> str:
+        pass
+
     @abstractmethod
     def start_run(
             self, run_id: str = None,
@@ -101,6 +111,10 @@ class TrackerBase(ITracker, ABC):
     def experiment_name(self) -> str:
         return self._experiment_name
 
+    @property
+    def experiment_id(self) -> str:
+        return self._experiment_name
+
     def get_max_metric_value(self, name: str) -> float:
         return max(self._metrics[name], key=lambda el: el['value'])
 
@@ -166,6 +180,10 @@ class InMemoryTracker(TrackerBase):
     def run_name(self) -> str:
         return self._run_name
 
+    @property
+    def run_id(self) -> str:
+        return self._run_name
+
     def start_run(
             self, run_id: str = None,
             run_name: str = None,

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")
@@ -41,6 +39,16 @@ class MLFlowTracker(TrackerBase):
             raise RuntimeError("There is no active run!")
         return self._active_run.info.run_name
 
+    @property
+    def run_id(self) -> str:
+        if not self._active_run:
+            raise RuntimeError("There is no active run!")
+        return self._active_run.info.run_id
+
+    @property
+    def experiment_id(self) -> str:
+        return self._experiment_id
+
     def start_run(
             self, run_id: str = None,
             run_name: str = None,
@@ -67,25 +75,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 +106,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 +114,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 +123,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 +138,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 +156,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 +173,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

toolchemy/ai/trackers/neptune_tracker.py

@@ -20,6 +20,10 @@ class NeptuneAITracker(TrackerBase):
             raise RuntimeError("There is no active run!")
         return self._active_run_id
 
+    @property
+    def run_id(self) -> str:
+        return self.run_name
+
 
     def start_run(
             self, run_id: str = None,

{toolchemy-0.2.192.dist-info → toolchemy-0.2.194.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: toolchemy
-Version: 0.2.192
+Version: 0.2.194
 Summary: a set of auxiliary programming tools
 License-File: LICENSE
 Author: Cyprian Nosek

{toolchemy-0.2.192.dist-info → toolchemy-0.2.194.dist-info}/RECORD

@@ -8,11 +8,11 @@ toolchemy/ai/clients/ollama_client.py,sha256=_PYI_udyTBmtKZlvRECzeWKrrR-NwTlaGJS
 toolchemy/ai/clients/openai_client.py,sha256=20Yt8azldRzbd8ylKqoD5frIEj8LhScjSYhGCt0WcWk,3611
 toolchemy/ai/clients/pricing.py,sha256=FuAuarQIaLl-e6KLzLbiWawWpqjY_kV4poO_qiTuHfA,1953
 toolchemy/ai/clients/whisper_client.py,sha256=AQcY2FiGLu0JTYw0zzjjH-7WyZ6NbQWrtxFa6iYhR4M,5252
-toolchemy/ai/prompter.py,sha256=29s1vWOL-ROacqprAmRCLMd_MlEOfUWjT48ZdVoZIAQ,4997
+toolchemy/ai/prompter.py,sha256=2Vf00d3pHUC1giG39REAn0YhuFUqxHXAQVNr5PCQQJ0,5367
 toolchemy/ai/trackers/__init__.py,sha256=FUv2NuGSmCpDrbE7rmeG0VmQr9rJRC-Y_D0MuQa0e-M,214
-toolchemy/ai/trackers/common.py,sha256=Y8_QHD1BAsJcBEHAekO-0Txe2MeBwGqh9TM53Ta49qk,6202
-toolchemy/ai/trackers/mlflow_tracker.py,sha256=fVX3YdNEDs9aS5UhlY4boaAsz2ArKnawrhLKli0AS2s,8627
-toolchemy/ai/trackers/neptune_tracker.py,sha256=vNvPzfrqqi7tRCfBadB-4VWobE5JR1i6ae23SQQC7I0,4639
+toolchemy/ai/trackers/common.py,sha256=tEYQv1OqcyqMVEfBP0tRse3V9f-gypOrfNBj1U9IcsE,6525
+toolchemy/ai/trackers/mlflow_tracker.py,sha256=pzxIBRdODtA4qp5CezSZHY1b1oMSH_Jpo8xjbA1PDO4,8522
+toolchemy/ai/trackers/neptune_tracker.py,sha256=GIbkLmHAzzO8Q7FuCmGf--y2NTEfcjDF9CVqqgmWzVg,4712
 toolchemy/db/lightdb.py,sha256=Bp7AimIuOfp2jIORFvBndk8kuKnEV12Ua1hoQXR5brY,9377
 toolchemy/utils/__init__.py,sha256=6Mp65G6Ce_zY86AxvjyzkMjU_e8cbeij7TTHYqLwO0U,574
 toolchemy/utils/at_exit_collector.py,sha256=YXupFjnujFK86n7dxH1Qol_i7FjNXJyo2q5f4SA3OEU,3682
@@ -29,8 +29,8 @@ toolchemy/utils/utils.py,sha256=JCRNRBJ4vRWzl-GtbetJs8GM72ImVAbDkPIXrvp3IE0,5167
 toolchemy/vision/__init__.py,sha256=p-sQaVd1FUQ1LMGdF51y28Z9IJvMVWaCdvrL0-jKDgc,141
 toolchemy/vision/caption_overlay.py,sha256=fLxHMD7IPSh5XJC5pvdPf1wZvHidlOvnXt_wgW7UJjk,2825
 toolchemy/vision/image.py,sha256=5q7MRK2SEyqWPB62duuzJqGmqBfQP1Tgphu01nP7O8s,2392
-toolchemy-0.2.192.dist-info/METADATA,sha256=L2cRA56cq-YT_jcPVabOvmIXyzGomuAsMoCUYS5S414,980
-toolchemy-0.2.192.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-toolchemy-0.2.192.dist-info/entry_points.txt,sha256=6l1k81pMAx7Ap1FaX8qW8SDwsONnNLOb28QQD9t6gto,66
-toolchemy-0.2.192.dist-info/licenses/LICENSE,sha256=JY5EPk3yvwZqQHk8JrQ-TFARdyfVTK9uUKV6CuKo_Bs,1062
-toolchemy-0.2.192.dist-info/RECORD,,
+toolchemy-0.2.194.dist-info/METADATA,sha256=nE_yD7yEedvdxxOhPfaWvmV62S1mC4vbx-nH7S1cTQs,980
+toolchemy-0.2.194.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
+toolchemy-0.2.194.dist-info/entry_points.txt,sha256=6l1k81pMAx7Ap1FaX8qW8SDwsONnNLOb28QQD9t6gto,66
+toolchemy-0.2.194.dist-info/licenses/LICENSE,sha256=JY5EPk3yvwZqQHk8JrQ-TFARdyfVTK9uUKV6CuKo_Bs,1062
+toolchemy-0.2.194.dist-info/RECORD,,