toolchemy 0.2.185__py3-none-any.whl

toolchemy/ai/clients/openai_client.py

@@ -0,0 +1,76 @@
+import time
+from abc import ABC, abstractmethod
+from openai import OpenAI, AzureOpenAI, NOT_GIVEN
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
+from typing import Iterable
+
+from toolchemy.ai.clients.common import LLMClientBase, ModelConfig, Usage, prepare_chat_messages
+
+
+class BaseOpenAIClient(LLMClientBase, ABC):
+    @property
+    @abstractmethod
+    def _client(self) -> OpenAI:
+        pass
+
+    def embeddings(self, text: str) -> list[float]:
+        response = self._client.embeddings.create(
+            input=text,
+            model=self.embedding_name
+        )
+        return response.data[0].embedding
+
+    def _completion(self, prompt: str, system_prompt: str | None, model_config: ModelConfig | None = None,
+                    images_base64: list[str] | None = None) -> tuple[str, Usage]:
+        messages = self._prepare_chat_messages(prompt=prompt, system_prompt=system_prompt, images_base64=images_base64)
+
+        duration_time_start = time.time()
+        response = self._client.chat.completions.create(
+            model=model_config.model_name,
+            messages=messages,
+            top_p=model_config.top_p,
+        )
+        duration = time.time() - duration_time_start
+
+        usage = Usage(input_tokens=response.usage.prompt_tokens, output_tokens=response.usage.completion_tokens,
+                      duration=duration)
+
+        return response.choices[0].message.content, usage
+
+    def _prepare_chat_messages(self, prompt: str, system_prompt: str | None = None, images_base64: list[str] | None = None) -> list[ChatCompletionMessageParam]:
+        messages_all = []
+        system_prompt = system_prompt or self._system_prompt or NOT_GIVEN
+
+        if self._keep_chat_session:
+            messages_all.extend(self._session_messages)
+
+        return prepare_chat_messages(prompt=prompt, system_prompt=system_prompt, images_base64=images_base64, messages_history=messages_all)
+
+
+class OpenAIClient(BaseOpenAIClient):
+    def __init__(self, api_key: str, model_name: str = "gpt-3.5-turbo",
+                 embedding_model_name: str = "text-embedding-3-large", default_model_config: ModelConfig | None = None,
+                 system_prompt: str | None = None, keep_chat_session: bool = False, no_cache: bool = False):
+        super().__init__(default_model_name=model_name, default_embedding_model_name=embedding_model_name,
+                         default_model_config=default_model_config,
+                         system_prompt=system_prompt, keep_chat_session=keep_chat_session, disable_cache=no_cache)
+        self._openai_client = OpenAI(api_key=api_key)
+
+    @property
+    def _client(self) -> OpenAI:
+        return self._openai_client
+
+
+class AzureOpenAIClient(BaseOpenAIClient):
+    def __init__(self, api_key: str, api_endpoint: str, api_version: str, model_name: str = "gpt-3.5-turbo",
+                 embedding_model_name: str = "text-embedding-3-large", default_model_config: ModelConfig | None = None,
+                 system_prompt: str | None = None, keep_chat_session: bool = False):
+        super().__init__(default_model_name=model_name, default_embedding_model_name=embedding_model_name,
+                         default_model_config=default_model_config,
+                         system_prompt=system_prompt, keep_chat_session=keep_chat_session)
+        self._openai_client = AzureOpenAI(api_key=api_key, azure_endpoint=api_endpoint, api_version=api_version,
+                                          azure_deployment=model_name)
+
+    @property
+    def _client(self) -> OpenAI:
+        return self._openai_client

toolchemy/ai/clients/pricing.py

@@ -0,0 +1,66 @@
+KEY_INPUT_TOKENS = "input_tokens_cost"
+KEY_OUTPUT_TOKENS = "output_tokens_cost"
+
+
+
+class Pricing:
+    pricing_per_1_mln = {
+        "gpt-5.2": {
+            KEY_INPUT_TOKENS: 1.75,
+            KEY_OUTPUT_TOKENS: 14.00,
+        },
+        "gpt-5.2-pro": {
+            KEY_INPUT_TOKENS: 21.00,
+            KEY_OUTPUT_TOKENS: 168.00,
+        },
+        "gpt-5-mini": {
+            KEY_INPUT_TOKENS: 0.25,
+            KEY_OUTPUT_TOKENS: 2.00,
+        },
+        "gpt-4.1": {
+            KEY_INPUT_TOKENS: 3.00,
+            KEY_OUTPUT_TOKENS: 12.00,
+        },
+        "gpt-4.1-mini": {
+            KEY_INPUT_TOKENS: 0.8,
+            KEY_OUTPUT_TOKENS: 3.2,
+        },
+        "gpt-4.1-nano": {
+            KEY_INPUT_TOKENS: 0.2,
+            KEY_OUTPUT_TOKENS: 0.8,
+        },
+        "o4-mini": {
+            KEY_INPUT_TOKENS: 4.00,
+            KEY_OUTPUT_TOKENS: 16.00,
+        },
+        "mistral-small3.2:24b": {
+            KEY_INPUT_TOKENS: 0.01,
+            KEY_OUTPUT_TOKENS: 0.03,
+        },
+        "gpt-oss:120b": {
+            KEY_INPUT_TOKENS: 0.01,
+            KEY_OUTPUT_TOKENS: 0.03,
+        },
+        "gemma3:27b": {
+            KEY_INPUT_TOKENS: 0.01,
+            KEY_OUTPUT_TOKENS: 0.03,
+        },
+        "qwen3:32b-q8_0": {
+            KEY_INPUT_TOKENS: 0.01,
+            KEY_OUTPUT_TOKENS: 0.03,
+        },
+        "dummy-model": {
+            KEY_INPUT_TOKENS: 0.1,
+            KEY_OUTPUT_TOKENS: 1.0,
+        }
+    }
+
+    @classmethod
+    def estimate(cls, model_name: str, input_tokens: int, output_tokens: int) -> float:
+        if model_name not in cls.pricing_per_1_mln:
+            raise ValueError(f"Model '{model_name}' not supported for pricing estimation")
+
+        model_pricing = cls.pricing_per_1_mln[model_name]
+        input_cost = model_pricing[KEY_INPUT_TOKENS] * (input_tokens / 1_000_000)
+        output_cost = model_pricing[KEY_OUTPUT_TOKENS] * (output_tokens / 1_000_000)
+        return input_cost + output_cost

toolchemy/ai/clients/whisper_client.py

@@ -0,0 +1,141 @@
+import asyncio
+import requests
+import sys
+import os
+import subprocess
+import tempfile
+from wyoming.client import AsyncClient
+from wyoming.audio import AudioChunk, AudioStart, AudioStop, AudioChunkConverter
+from wyoming.asr import Transcript
+from wyoming.ping import Ping
+from wyoming.info import Describe
+import wave
+
+from toolchemy.utils.logger import get_logger
+
+
+class WhisperClient:
+    def __init__(self, url: str):
+        self._logger = get_logger()
+        self._endpoint = url
+        self._whisper_client_wyoming = None
+        if self._endpoint.startswith("http"):
+            if not self._endpoint.endswith("transcribe"):
+                if not self._endpoint.endswith("/"):
+                    self._endpoint += "/"
+                self._endpoint += "transcribe"
+        elif self._endpoint.startswith("tcp"):
+            self._whisper_client_wyoming = AsyncClient.from_uri(self._endpoint)
+        else:
+            raise ValueError(f"Unknown protocol for the whisper server endpoint: '{self._endpoint}'")
+        self._logger.info(f"Whisper client initialized (endpoint: '{self._endpoint}')")
+
+    def transcribe(self, audio_path: str) -> str:
+        transcription = None
+
+        if self._endpoint.startswith("tcp"):
+            transcription = asyncio.run(self._transcribe_wyoming(audio_path))
+
+        if self._endpoint.startswith("http"):
+            transcription = self._transcribe_http(audio_path)
+
+        if transcription is None:
+            raise RuntimeError(f"Transcription failed...")
+
+        return transcription.strip()
+
+    def _transcribe_http(self, audio_path: str) -> str:
+        if not os.path.exists(audio_path):
+            raise ValueError(f"Error: File '{audio_path}' not found.")
+
+        with open(audio_path, "rb") as audio_file:
+            files = {"file": audio_file}
+
+            self._logger.info(f"Sending '{audio_path}' to Whisper server...")
+            response = requests.post(self._endpoint, files=files)
+
+        if response.status_code == 200:
+            result = response.json()
+            result_transcription = result.get("text")
+            self._logger.info(f"Transcription: '{result_transcription}'")
+            return result_transcription
+
+        err_msg = f"Error: Failed to transcribe. Status Code: {response.status_code}"
+        self._logger.error(err_msg)
+        raise RuntimeError(err_msg)
+
+    async def _transcribe_wyoming(self, audio_path: str, audio_rate: int = 16000, audio_width: int = 2,
+                                  audio_channels: int = 1, chunk_size: int = 1024) -> str:
+        wav_path = self._convert_to_wav(audio_path, audio_rate=audio_rate, audio_channels=audio_channels)
+
+        await self._whisper_client_wyoming.connect()
+        await self._whisper_client_wyoming.write_event(Ping(text="test").event())
+
+        await self._whisper_client_wyoming.write_event(Describe().event())
+
+        info_event = await self._whisper_client_wyoming.read_event()
+        self._logger.info(f"info event: {info_event}")
+
+        with wave.open(wav_path, 'rb') as wav:
+            assert wav.getframerate() == audio_rate
+            assert wav.getsampwidth() == audio_width
+            assert wav.getnchannels() == audio_channels
+
+            await self._whisper_client_wyoming.write_event(AudioStart(audio_rate, audio_width, audio_channels).event())
+
+            audio_bytes = wav.readframes(chunk_size)
+
+            converter = AudioChunkConverter(
+                rate=audio_rate, width=audio_width, channels=audio_channels,
+            )
+
+            while audio_bytes:
+                chunk = converter.convert(AudioChunk(audio_rate, audio_width, audio_channels, audio_bytes))
+                await self._whisper_client_wyoming.write_event(chunk.event())
+                audio_bytes = wav.readframes(chunk_size)
+
+            await self._whisper_client_wyoming.write_event(AudioStop().event())
+
+            while True:
+                event = await asyncio.wait_for(self._whisper_client_wyoming.read_event(), timeout=30)
+                if event is None:
+                    break
+                transcript = Transcript.from_event(event)
+                if transcript.text:
+                    self._logger.info(f"Transcription: {transcript.text}")
+                    break
+        await self._whisper_client_wyoming.disconnect()
+
+        return transcript.text
+
+    def _convert_to_wav(self, input_path: str, audio_rate: int, audio_channels: int) -> str:
+        temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+        temp_wav.close()
+
+        cmd = [
+            "ffmpeg",
+            "-y",
+            "-i", input_path,
+            "-ar", str(audio_rate),
+            "-ac", str(audio_channels),
+            "-f", "wav",
+            "-sample_fmt", "s16",
+            temp_wav.name
+        ]
+
+        subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        return temp_wav.name
+
+
+def main(argv: list):
+    if len(argv) < 2:
+        raise ValueError("Usage: python transcribe_audio.py <audio_file.mp3/wav>")
+
+    file_path = argv[1]
+    client = WhisperClient(url="tcp://hal:10300")
+    transcription = client.transcribe(file_path)
+    print(f"> '{transcription}'")
+
+
+if __name__ == "__main__":
+    main(sys.argv)

toolchemy/ai/prompter.py

@@ -0,0 +1,124 @@
+import logging
+import mlflow
+import sys
+import subprocess
+from abc import ABC, abstractmethod
+
+from toolchemy.utils.cacher import Cacher, ICacher
+from toolchemy.utils.logger import get_logger
+from toolchemy.utils.locations import Locations
+
+
+class IPrompter(ABC):
+    @abstractmethod
+    def render(self, name: str, version: str | dict[str, str | list[str]] | None = None, **variables) -> str:
+        pass
+
+    @abstractmethod
+    def template(self, name: str, version: str | dict[str, str | list[str]] | None = None) -> str:
+        pass
+
+    @abstractmethod
+    def create_template(self, name: str, template: str, overwrite: bool = False):
+        pass
+
+    @abstractmethod
+    def run_studio(self):
+        pass
+
+
+def run_studio():
+    prompter = PrompterMLflow()
+    prompter.run_studio()
+
+
+class PrompterBase(IPrompter, ABC):
+    def __init__(self, cacher: ICacher | None = None, no_cache: bool = False, log_level: int = logging.INFO):
+        self._logger = get_logger(level=log_level)
+        self._cacher = cacher or Cacher(disabled=no_cache)
+
+    def _prompt_version(self, name: str, version_mapping: str | dict[str, str | list[str]] | None = None) -> str | None:
+        if version_mapping is None:
+            return self._latest_value()
+        if isinstance(version_mapping, int):
+            return str(version_mapping)
+        if isinstance(version_mapping, str):
+            return version_mapping
+        if name not in version_mapping:
+            return None
+        if isinstance(version_mapping[name], str):
+            return version_mapping[name]
+        return version_mapping[name][0]
+
+    @abstractmethod
+    def _latest_value(self) -> str |  None:
+        pass
+
+
+class PrompterMLflow(PrompterBase):
+    DEFAULT_PROMPT_REGISTRY_NAME = "prompts_mlflow"
+
+    def __init__(self, registry_store_dir: str | None = None, cacher: ICacher | None = None, no_cache: bool = False, log_level: int = logging.INFO):
+        super().__init__(cacher=cacher, no_cache=no_cache, log_level=log_level)
+        locations = Locations()
+        if registry_store_dir is None:
+            registry_store_dir = locations.in_root(self.DEFAULT_PROMPT_REGISTRY_NAME)
+        registry_store_dir = locations.abs(registry_store_dir).rstrip("/")
+
+        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._logger.info(f"Prompter-MLflow initialized")
+        self._logger.info(f"> tracking store uri: {self._tracking_uri}")
+        self._logger.info(f"> registry store uri: {self._registry_store_uri}")
+
+    def render(self, name: str, version: str | dict[str, str | list[str]] | None = None, **variables) -> str:
+        prompt_uri = self._build_prompt_uri(name=name, version=version)
+
+        cache_key = self._cacher.create_cache_key(["prompt_render", prompt_uri], [variables])
+        if self._cacher.exists(cache_key):
+            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 = prompt_template.format(**variables)
+
+        self._cacher.set(cache_key, prompt)
+
+        return prompt
+
+    def template(self, name: str, version: str | dict[str, str | list[str]] | None = None) -> str:
+        prompt_uri = self._build_prompt_uri(name=name, version=version)
+
+        cache_key = self._cacher.create_cache_key(["prompt_template", prompt_uri])
+        if self._cacher.exists(cache_key):
+            return self._cacher.get(cache_key)
+
+        self._logger.debug(f"Getting prompt template: '{name}' (version: '{version}') -> prompt uri: '{prompt_uri}'")
+
+        prompt_template = mlflow.genai.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:
+            return
+        mlflow.genai.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)
+        prompt_uri = f"prompts:/{name}"
+        prompt_uri += f"@{prompt_version}" if isinstance(prompt_version, str) and not prompt_version.isdigit() else f"/{prompt_version}"
+        return prompt_uri
+
+    def run_studio(self):
+        command = [sys.executable, "-m", "mlflow", "ui", "--registry-store-uri", self._registry_store_uri, "--backend-store-uri", self._tracking_uri]
+        sys.exit(subprocess.call(command))
+
+    def _latest_value(self) -> str | None:
+        return "latest"

toolchemy/ai/trackers/__init__.py

@@ -0,0 +1,5 @@
+from .common import ITracker, InMemoryTracker
+from .mlflow_tracker import MLFlowTracker
+from .neptune_tracker import NeptuneAITracker
+
+__all__ = ["ITracker", "InMemoryTracker", "MLFlowTracker", "NeptuneAITracker"]

toolchemy/ai/trackers/common.py

@@ -0,0 +1,216 @@
+from abc import ABC, abstractmethod
+import statistics
+from typing import Dict, Any
+
+from toolchemy.utils.logger import get_logger
+
+
+class ITracker(ABC):
+    @property
+    @abstractmethod
+    def experiment_name(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def run_name(self) -> str:
+        pass
+
+    @abstractmethod
+    def start_run(
+            self, run_id: str = None,
+            run_name: str = None,
+            parent_run_id: str = None,
+            user_specified_tags: Dict[str, str] = None
+    ):
+        pass
+
+    @abstractmethod
+    def end_run(self):
+        pass
+
+    @abstractmethod
+    def log(self, name: str, value: Any):
+        pass
+
+    @abstractmethod
+    def log_param(self, name: str, value: Any):
+        pass
+
+    @abstractmethod
+    def log_params(self, params: Dict[str, Any]):
+        pass
+
+    @abstractmethod
+    def log_metric(self, name: str, value: float, step: int | None = None, metric_metadata: dict | None = None):
+        pass
+
+    @abstractmethod
+    def log_text(self, name: str, value: str):
+        pass
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, float | list], step: int | None = None):
+        pass
+
+    @abstractmethod
+    def log_artifact(self, artifact_path: str, save_dir: str = None):
+        pass
+
+    @abstractmethod
+    def log_figure(self, figure, save_path: str):
+        """
+        Parameters:
+             figure (matplotlib.figure.Figure, _matplotlib.figure.Figure, _plotly.graph_objects.Figure): plot figure
+             save_path (str): the run-relative artifact file path in posixpath format to which
+                              the figure is saved (e.g. "dir/file.png").
+        """
+        pass
+
+    @abstractmethod
+    def set_run_tag(self, name: str, value: str | int | float):
+        pass
+
+    @abstractmethod
+    def set_experiment_tag(self, name: str, value: str | int | float):
+        pass
+
+    @abstractmethod
+    def disable(self):
+        pass
+
+    @abstractmethod
+    def get_data(self) -> dict:
+        pass
+
+
+class TrackerBase(ITracker, ABC):
+    def __init__(self, experiment_name: str, with_artifact_logging: bool = True, disabled: bool = False):
+        self._disabled = disabled
+        self._logger = get_logger()
+        self._experiment_name = experiment_name
+        self._artifact_logging = with_artifact_logging
+        self._metrics = {}
+        self._params = {}
+        self._tags = {
+            "experiment": {},
+            "runs": {},
+        }
+
+    @property
+    def experiment_name(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'])
+
+    def get_min_metric_value(self, name: str) -> float:
+        return min(self._metrics[name], key=lambda el: el['value'])
+
+    def get_avg_metric_value(self, name: str) -> float:
+        metric_values = [m['value'] for m in self._metrics[name]]
+        return statistics.mean(metric_values)
+
+    def get_data(self) -> dict:
+        return {
+            "metrics": self._metrics.copy(),
+            "params": self._params.copy(),
+            "tags": self._tags.copy(),
+        }
+
+    def _store_param(self, name: str, value: Any):
+        if self._disabled:
+            raise RuntimeError(f"Disabled trackers cannot store params!")
+
+        self._params[name] = value
+
+    def _store_tag(self, name: str, value: str | int | float, run_name: str | None = None):
+        if run_name is None:
+            self._tags["experiment"][name] = value
+            return
+        if run_name not in self._tags["runs"]:
+            self._tags["runs"][run_name] = {}
+        self._tags["runs"][run_name][name] = value
+
+    def _store_metric(self, name: str, value: float, metric_metadata: dict | None = None) -> float:
+        if self._disabled:
+            raise RuntimeError(f"Disabled trackers cannot store metrics!")
+        if name not in self._metrics:
+            self._metrics[name] = []
+
+        if isinstance(value, dict):
+            new_entry = value
+        else:
+            new_entry = {
+                'value': value
+            }
+
+        if metric_metadata:
+            new_entry.update(metric_metadata)
+
+        self._metrics[name] += [new_entry]
+
+        return new_entry['value']
+
+    def disable(self):
+        self._disabled = True
+
+
+class InMemoryTracker(TrackerBase):
+    def __init__(self, experiment_name: str = "dummy", disabled: bool = False):
+        super().__init__(experiment_name=experiment_name, disabled=disabled)
+        self._run_name = None
+        self._reset()
+
+    @property
+    def run_name(self) -> str:
+        return self._run_name
+
+    def start_run(
+            self, run_id: str = None,
+            run_name: str = None,
+            parent_run_id: str = None,
+            user_specified_tags: Dict[str, str] = None
+    ):
+        self._run_name = run_name
+
+    def _reset(self):
+        self._run_name = None
+        self._data = {}
+        self._params = {}
+        self._metrics = {}
+
+    def end_run(self):
+        self._reset()
+
+    def log(self, name: str, value: Any):
+        self._data[name] = value
+
+    def log_param(self, name: str, value: Any):
+        self._params[name] = value
+
+    def log_params(self, params: Dict[str, Any]):
+        for name, value in params.items():
+            self.log_param(name, value)
+
+    def log_text(self, name: str, value: str):
+        self.log(name, value)
+
+    def log_metric(self, name: str, value: float, step: int | None = None, metric_metadata: dict | None = None):
+        self._metrics[name] = value
+
+    def log_metrics(self, metrics: Dict[str, float | list], step: int | None = None):
+        for name, value in metrics.items():
+            self.log_metric(name, value, step)
+
+    def log_artifact(self, artifact_path: str, save_dir: str = None):
+        raise NotImplementedError()
+
+    def log_figure(self, figure, save_path: str):
+        raise NotImplementedError()
+
+    def set_run_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value, run_name=self.run_name)
+
+    def set_experiment_tag(self, name: str, value: str | int | float):
+        self._store_tag(name, value)