kgmodule-utils 0.4.3__tar.gz → 0.4.5__tar.gz

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kgmodule-utils
-Version: 0.4.3
+Version: 0.4.5
 Summary: Shared types, graph store, semantic index, and pipeline base for the KGModule SDK
 License: Elastic-2.0
 License-File: LICENSE
@@ -35,7 +35,7 @@ Description-Content-Type: text/markdown
 
 [![Python](https://img.shields.io/badge/python-3.12%20%7C%203.13-blue.svg)](https://www.python.org/)
 [![License: Elastic-2.0](https://img.shields.io/badge/License-Elastic%202.0-blue.svg)](https://www.elastic.co/licensing/elastic-license)
-[![Version](https://img.shields.io/badge/version-0.4.0-blue.svg)](https://github.com/Flux-Frontiers/KG_utils/releases)
+[![Version](https://img.shields.io/badge/version-0.4.4-blue.svg)](https://github.com/Flux-Frontiers/KG_utils/releases)
 [![CI](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml/badge.svg)](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml)
 [![Poetry](https://img.shields.io/endpoint?url=https://python-poetry.org/badge/v0.json)](https://python-poetry.org/)
 

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/README.md

@@ -1,7 +1,7 @@
 
 [![Python](https://img.shields.io/badge/python-3.12%20%7C%203.13-blue.svg)](https://www.python.org/)
 [![License: Elastic-2.0](https://img.shields.io/badge/License-Elastic%202.0-blue.svg)](https://www.elastic.co/licensing/elastic-license)
-[![Version](https://img.shields.io/badge/version-0.4.0-blue.svg)](https://github.com/Flux-Frontiers/KG_utils/releases)
+[![Version](https://img.shields.io/badge/version-0.4.4-blue.svg)](https://github.com/Flux-Frontiers/KG_utils/releases)
 [![CI](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml/badge.svg)](https://github.com/Flux-Frontiers/KG_utils/actions/workflows/ci.yml)
 [![Poetry](https://img.shields.io/endpoint?url=https://python-poetry.org/badge/v0.json)](https://python-poetry.org/)
 

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "kgmodule-utils"
-version = "0.4.3"
+version = "0.4.5"
 description = "Shared types, graph store, semantic index, and pipeline base for the KGModule SDK"
 readme = "README.md"
 license = { text = "Elastic-2.0" }

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/src/kg_utils/__init__.py

@@ -24,4 +24,4 @@ Optional extras
     pip install 'kgmodule-utils[synthesis-mflux]'  # + mflux (Apple Silicon local gen)
 """
 
-__version__ = "0.4.3"
+__version__ = "0.4.5"

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/src/kg_utils/embedder.py

@@ -34,12 +34,12 @@ License: Elastic 2.0
 
 from __future__ import annotations
 
+import importlib
 import os
 from typing import Any
 
 from kg_utils.embed import DEFAULT_MODEL, KNOWN_MODELS, resolve_model_path
 
-
 # ---------------------------------------------------------------------------
 # Abstract base
 # ---------------------------------------------------------------------------
@@ -75,7 +75,7 @@ class Embedder:
 # ---------------------------------------------------------------------------
 
 
-def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
+def load_sentence_transformer(model_name: str = DEFAULT_MODEL, device: str | None = None) -> Any:
     """Load a ``SentenceTransformer`` with the canonical safe-load sequence.
 
     Resolution order:
@@ -90,13 +90,22 @@ def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
     retry loops leave stale thread state that causes SIGBUS on the first
     ``encode()`` call.
 
+    Device precedence: explicit *device* arg > ``KG_EMBED_DEVICE`` env >
+    auto-detect.  The env var exists because ``spawn``-based embedding workers
+    inherit ``os.environ`` but can't easily receive a Python arg — without a way
+    to pin the device, each worker auto-selects MPS and N parallel workers stack
+    N GPU allocations into an OOM.  So CPU multiprocessing embedding on Apple
+    Silicon is only safe with this knob.
+
     :param model_name: HuggingFace model ID or KNOWN_MODELS alias.
+    :param device: Explicit device (``"cpu"``/``"mps"``/``"cuda"``).  ``None``
+        falls back to ``KG_EMBED_DEVICE`` then CUDA→MPS→CPU auto-detect.
     :return: Loaded ``SentenceTransformer`` instance.
     """
-    from sentence_transformers import SentenceTransformer  # pylint: disable=import-outside-toplevel
+    SentenceTransformer = importlib.import_module("sentence_transformers").SentenceTransformer
 
     try:
-        from transformers import logging as hf_logging  # pylint: disable=import-outside-toplevel
+        hf_logging = importlib.import_module("transformers.logging")
 
         hf_logging.set_verbosity_error()
         # TQDM_DISABLE alone misses transformers' _tqdm_active gate
@@ -106,9 +115,18 @@ def load_sentence_transformer(model_name: str = DEFAULT_MODEL) -> Any:
 
     os.environ["TQDM_DISABLE"] = "1"
 
-    import torch  # pylint: disable=import-outside-toplevel
-
-    if torch.cuda.is_available():
+    torch = importlib.import_module("torch")
+
+    # Device precedence: explicit arg > KG_EMBED_DEVICE env > auto-detect.
+    # The env var lets spawn-based embedding workers (which inherit os.environ
+    # but can't easily receive a Python arg) pin to e.g. CPU — without it each
+    # worker auto-selects MPS and N parallel workers stack N GPU allocations →
+    # MPS OOM. The override is why CPU multiprocessing embedding is safe on
+    # Apple Silicon.
+    sel = (device or os.environ.get("KG_EMBED_DEVICE", "")).strip().lower()
+    if sel:
+        device = sel
+    elif torch.cuda.is_available():
         device = "cuda"
     else:
         try:
@@ -158,7 +176,7 @@ class SentenceTransformerEmbedder(Embedder):
 
     def __init__(self, model_name: str = DEFAULT_MODEL) -> None:
         try:
-            from transformers import logging as hf_logging  # pylint: disable=import-outside-toplevel
+            hf_logging = importlib.import_module("transformers.logging")
 
             hf_logging.set_verbosity_error()
             hf_logging.disable_progress_bar()
@@ -188,7 +206,7 @@ class SentenceTransformerEmbedder(Embedder):
         :param texts: Input strings.
         :param encode_batch_size: Passed to ``model.encode()`` — tune down if OOM on MPS.
         """
-        import numpy as np  # pylint: disable=import-outside-toplevel
+        np = importlib.import_module("numpy")
 
         vecs = self.model.encode(
             texts,
@@ -200,7 +218,7 @@ class SentenceTransformerEmbedder(Embedder):
 
     def embed_query(self, query: str) -> list[float]:
         """Embed a single query string into a float32 vector."""
-        import numpy as np  # pylint: disable=import-outside-toplevel
+        np = importlib.import_module("numpy")
 
         vec = self.model.encode([query], normalize_embeddings=True)[0]
         return list(np.asarray(vec, dtype="float32").tolist())
@@ -233,7 +251,7 @@ def wrap_embedder(st_model: Any, model_name: str = DEFAULT_MODEL) -> Embedder:
     :param model_name: Model name stored as metadata on the wrapper.
     :return: An :class:`Embedder` that delegates all calls to *st_model*.
     """
-    import numpy as np  # pylint: disable=import-outside-toplevel
+    np = importlib.import_module("numpy")
 
     resolved = KNOWN_MODELS.get(model_name, model_name)
     _dim_fn = getattr(st_model, "get_embedding_dimension", None) or getattr(

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/src/kg_utils/synthesis/_image.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import base64
+import importlib
 import random
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
@@ -102,7 +103,10 @@ class ImageSynthesizer:
     def _load_mflux(self, model_name: str) -> Any:
         if self._mflux_model is not None and self._mflux_model_name == model_name:
             return self._mflux_model
-        from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein  # type: ignore
+
+        Flux2Klein = importlib.import_module(
+            "mflux.models.flux2.variants.txt2img.flux2_klein"
+        ).Flux2Klein
 
         self._mflux_model = Flux2Klein(model_path=model_name)
         self._mflux_model_name = model_name

{kgmodule_utils-0.4.3 → kgmodule_utils-0.4.5}/src/kg_utils/synthesis/_text.py

@@ -98,6 +98,31 @@ class TextSynthesizer:
     # Public API
     # ------------------------------------------------------------------
 
+    def complete(
+        self,
+        messages: list[dict],
+        *,
+        model: str | None = None,
+        max_tokens: int | None = None,
+        temperature: float = 0.7,
+    ) -> str | None:
+        """Run a raw chat completion against the configured backend.
+
+        General-purpose entry point for callers that build their own message
+        list — summarization, classification, arbitrary prompting — rather than
+        using :meth:`synthesize_rag` or :meth:`rewrite_for_image`. Applies the
+        same oMLX thinking suppression and ``<think>`` stripping as the other
+        public methods.
+
+        :param messages: OpenAI-style chat messages (``[{"role": ..., "content": ...}]``).
+        :param model: Override the configured model for this single call.
+        :param max_tokens: Override the configured ``max_tokens`` for this call.
+        :param temperature: Sampling temperature (lower = more deterministic).
+        :returns: The completion text with thinking stripped, or ``None`` on
+                  failure or empty output.
+        """
+        return self._complete(messages, model=model, max_tokens=max_tokens, temperature=temperature)
+
     def list_models(self) -> list[str]:
         """Return the model IDs available at the configured endpoint.