kgmodule-utils 0.4.4__tar.gz → 0.4.6__tar.gz

{kgmodule_utils-0.4.4 → kgmodule_utils-0.4.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kgmodule-utils
-Version: 0.4.4
+Version: 0.4.6
 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.4 → kgmodule_utils-0.4.6}/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.4 → kgmodule_utils-0.4.6}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "kgmodule-utils"
-version = "0.4.4"
+version = "0.4.6"
 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.4 → kgmodule_utils-0.4.6}/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.4"
+__version__ = "0.4.6"

{kgmodule_utils-0.4.4 → kgmodule_utils-0.4.6}/src/kg_utils/embedder.py

@@ -40,6 +40,14 @@ from typing import Any
 
 from kg_utils.embed import DEFAULT_MODEL, KNOWN_MODELS, resolve_model_path
 
+#: Default per-call encode batch fed to ``model.encode(batch_size=...)``.
+#: Transformer attention memory scales with ``batch x seq^2``, so a large batch
+#: on long (near-max-sequence) chunks can allocate many GB per call and OOM /
+#: stall MPS.  Throughput is flat above ~128 on both CPU and MPS for the models
+#: in use, so 128 is the safe default; raise it only for a large-VRAM CUDA GPU
+#: with short sequences.
+DEFAULT_ENCODE_BATCH: int = 128
+
 # ---------------------------------------------------------------------------
 # Abstract base
 # ---------------------------------------------------------------------------
@@ -53,10 +61,14 @@ class Embedder:
 
     dim: int
 
-    def embed_texts(self, texts: list[str]) -> list[list[float]]:
+    def embed_texts(
+        self, texts: list[str], encode_batch_size: int = DEFAULT_ENCODE_BATCH
+    ) -> list[list[float]]:
         """Embed a list of strings into float32 vectors.
 
         :param texts: Input strings.
+        :param encode_batch_size: Per-call ``model.encode`` batch (default
+            :data:`DEFAULT_ENCODE_BATCH`); memory scales with ``batch x seq^2``.
         :return: One float32 vector per input.
         """
         raise NotImplementedError
@@ -200,11 +212,15 @@ class SentenceTransformerEmbedder(Embedder):
         )
         self.dim: int = (_dim_fn() if _dim_fn is not None else None) or 384
 
-    def embed_texts(self, texts: list[str], encode_batch_size: int = 512) -> list[list[float]]:
+    def embed_texts(
+        self, texts: list[str], encode_batch_size: int = DEFAULT_ENCODE_BATCH
+    ) -> list[list[float]]:
         """Embed a list of strings into float32 vectors.
 
         :param texts: Input strings.
-        :param encode_batch_size: Passed to ``model.encode()`` — tune down if OOM on MPS.
+        :param encode_batch_size: Per-call ``model.encode`` batch (default
+            :data:`DEFAULT_ENCODE_BATCH`).  Memory scales with ``batch x seq^2``;
+            tune down further if OOM on MPS, up only for large-VRAM CUDA.
         """
         np = importlib.import_module("numpy")
 
@@ -263,10 +279,12 @@ def wrap_embedder(st_model: Any, model_name: str = DEFAULT_MODEL) -> Embedder:
         model_name: str = resolved
         dim: int = _dim
 
-        def embed_texts(self, texts: list[str]) -> list[list[float]]:
+        def embed_texts(
+            self, texts: list[str], encode_batch_size: int = DEFAULT_ENCODE_BATCH
+        ) -> list[list[float]]:
             vecs = st_model.encode(
                 texts,
-                batch_size=512,
+                batch_size=encode_batch_size,
                 normalize_embeddings=True,
                 show_progress_bar=False,
             )

{kgmodule_utils-0.4.4 → kgmodule_utils-0.4.6}/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.