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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kgmodule-utils
-Version: 0.4.1
+Version: 0.4.3
 Summary: Shared types, graph store, semantic index, and pipeline base for the KGModule SDK
 License: Elastic-2.0
 License-File: LICENSE

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

@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "kgmodule-utils"
-version = "0.4.1"
+version = "0.4.3"
 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.1 → kgmodule_utils-0.4.3}/src/kg_utils/__init__.py

@@ -14,6 +14,8 @@ Sub-packages / modules:
     kg_utils.synthesis  — Unified text + image synthesis: TextSynthesizer, ImageSynthesizer.
                           Backends: omlx | ollama | openai (text);
                                     mflux-local | mflux-serve | openai (image).
+    kg_utils.worker     — RunPod worker protocol helpers and WorkerClient for /runsync calls.
+    kg_utils.retrieval  — Shared retrieval helpers: hit_to_dict, attach_content_by_sqlite.
 
 Optional extras
 ---------------
@@ -22,4 +24,4 @@ Optional extras
     pip install 'kgmodule-utils[synthesis-mflux]'  # + mflux (Apple Silicon local gen)
 """
 
-__version__ = "0.4.1"
+__version__ = "0.4.3"

kgmodule_utils-0.4.3/src/kg_utils/retrieval/__init__.py

@@ -0,0 +1,5 @@
+"""Shared retrieval helpers for serializing and enriching KG hits."""
+
+from kg_utils.retrieval.hits import attach_content_by_sqlite, hit_to_dict
+
+__all__ = ["hit_to_dict", "attach_content_by_sqlite"]

kgmodule_utils-0.4.3/src/kg_utils/retrieval/hits.py

@@ -0,0 +1,75 @@
+# © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
+"""Hit serialization and content hydration helpers for KG retrieval responses."""
+
+from __future__ import annotations
+
+import sqlite3
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+__all__ = ["hit_to_dict", "attach_content_by_sqlite"]
+
+
+def _is_diary_kind(kind_value: Any) -> bool:
+    kind_str = str(kind_value)
+    return kind_str == "KGKind.DIARY" or kind_str.lower().endswith("diary")
+
+
+def hit_to_dict(hit: Any, include_diary_timestamp: bool = False) -> dict:
+    """Serialize a KGRAG hit object into a plain dictionary.
+
+    :param hit: Hit-like object with standard retrieval attributes.
+    :param include_diary_timestamp: Include ``timestamp`` field for diary hits.
+    :returns: Serialized hit dictionary.
+    """
+    out = {
+        "kg_name": hit.kg_name,
+        "kg_kind": str(hit.kg_kind),
+        "node_id": hit.node_id,
+        "name": hit.name,
+        "kind": hit.kind,
+        "score": round(float(hit.score), 4),
+        "summary": hit.summary,
+        "source_path": hit.source_path,
+    }
+    if include_diary_timestamp:
+        out["timestamp"] = hit.name if _is_diary_kind(hit.kg_kind) else None
+    return out
+
+
+def attach_content_by_sqlite(hits: list[dict], kg_sqlite_map: dict[str, Path]) -> None:
+    """Attach full node text under ``content`` via batched SQLite lookups.
+
+    Missing or unreadable databases are ignored to preserve permissive behavior.
+
+    :param hits: Mutable hit dictionaries. Each hit should include ``kg_name`` and ``node_id``.
+    :param kg_sqlite_map: Mapping of KG name to sqlite database path.
+    """
+    by_kg: dict[str, list[dict]] = defaultdict(list)
+    for hit in hits:
+        by_kg[hit.get("kg_name", "")].append(hit)
+
+    for kg_name, kg_hits in by_kg.items():
+        db_path = kg_sqlite_map.get(kg_name)
+        if not db_path or not Path(db_path).exists():
+            continue
+
+        ids = [h.get("node_id") for h in kg_hits if h.get("node_id")]
+        if not ids:
+            continue
+
+        text_by_id: dict[str, str] = {}
+        try:
+            with sqlite3.connect(str(db_path)) as con:
+                placeholders = ",".join("?" * len(ids))
+                query = f"SELECT id, text FROM nodes WHERE id IN ({placeholders})"
+                for node_id, text in con.execute(query, ids):
+                    text_by_id[node_id] = text or ""
+        except Exception:  # noqa: BLE001  # pylint: disable=broad-exception-caught
+            continue
+
+        for hit in kg_hits:
+            node_id = hit.get("node_id")
+            if node_id:
+                hit["content"] = text_by_id.get(node_id, "")

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

@@ -53,6 +53,11 @@ from kg_utils.synthesis._config import (
 )
 from kg_utils.synthesis._image import ImageSynthesizer
 from kg_utils.synthesis._text import TextSynthesizer
+from kg_utils.synthesis.factory import (
+    image_synth_for_backend,
+    normalize_openai_base_url,
+    text_synth_for_backend,
+)
 
 
 def text_synthesizer_from_env() -> TextSynthesizer:
@@ -76,4 +81,7 @@ __all__ = [
     "image_synthesizer_from_env",
     "text_config_from_env",
     "image_config_from_env",
+    "normalize_openai_base_url",
+    "text_synth_for_backend",
+    "image_synth_for_backend",
 ]

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

@@ -51,6 +51,29 @@ _GPT_IMAGE_SIZES: dict[str, str] = {
 }
 
 
+# ---------------------------------------------------------------------------
+# Size parsing
+# ---------------------------------------------------------------------------
+
+
+def _parse_size(size: str | None) -> tuple[int, int] | None:
+    """Parse an explicit ``"WIDTHxHEIGHT"`` string into an ``(width, height)`` pair.
+
+    :param size: Size string such as ``"768x512"`` (case-insensitive ``x``), or None.
+    :returns: ``(width, height)`` when *size* parses to two positive ints, else None.
+    """
+    if not size:
+        return None
+    try:
+        w_str, h_str = size.lower().split("x", 1)
+        width, height = int(w_str), int(h_str)
+    except (ValueError, AttributeError):
+        return None
+    if width <= 0 or height <= 0:
+        return None
+    return width, height
+
+
 # ---------------------------------------------------------------------------
 # Synthesizer
 # ---------------------------------------------------------------------------
@@ -97,6 +120,7 @@ class ImageSynthesizer:
         seed: int | None = None,
         steps: int | None = None,
         model: str | None = None,
+        size: str | None = None,
     ) -> PILImage:
         """Generate an image and return a PIL Image.
 
@@ -105,6 +129,9 @@ class ImageSynthesizer:
         :param seed: Random seed for reproducibility (random int if omitted).
         :param steps: Override inference steps (mflux backends only; ignored for OpenAI).
         :param model: Override the configured model for this single call.
+        :param size: Explicit ``"WIDTHxHEIGHT"`` override (mflux backends only). When given
+                     and valid, it takes priority over the *aspect_ratio* size table. OpenAI
+                     backends ignore it because they accept only a fixed set of sizes.
         :returns: PIL Image.
         """
         cfg = self._cfg
@@ -118,6 +145,7 @@ class ImageSynthesizer:
                 seed=seed,
                 steps=effective_steps,
                 aspect_ratio=aspect_ratio,
+                size=size,
             )
         elif cfg.backend == ImageBackend.MFLUX_SERVE:
             return self._generate_via_server(
@@ -126,6 +154,7 @@ class ImageSynthesizer:
                 seed=seed,
                 steps=effective_steps,
                 aspect_ratio=aspect_ratio,
+                size=size,
             )
         else:
             return self._generate_openai(
@@ -142,6 +171,7 @@ class ImageSynthesizer:
         seed: int | None = None,
         steps: int | None = None,
         model: str | None = None,
+        size: str | None = None,
     ) -> str:
         """Generate an image and return it as a base64-encoded PNG string.
 
@@ -152,9 +182,12 @@ class ImageSynthesizer:
         :param seed: Random seed for reproducibility.
         :param steps: Override inference steps (mflux backends only).
         :param model: Override the configured model for this single call.
+        :param size: Explicit ``"WIDTHxHEIGHT"`` override (mflux backends only).
         :returns: Base64-encoded PNG string.
         """
-        img = self.generate(prompt, aspect_ratio=aspect_ratio, seed=seed, steps=steps, model=model)
+        img = self.generate(
+            prompt, aspect_ratio=aspect_ratio, seed=seed, steps=steps, model=model, size=size
+        )
         buf = BytesIO()
         img.save(buf, format="PNG")
         return base64.b64encode(buf.getvalue()).decode()
@@ -171,8 +204,9 @@ class ImageSynthesizer:
         seed: int | None,
         steps: int,
         aspect_ratio: str,
+        size: str | None = None,
     ) -> PILImage:
-        width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
+        width, height = _parse_size(size) or _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
         effective_seed = seed if seed is not None else random.randint(0, 2**31 - 1)
         flux = self._load_mflux(model)
         result = flux.generate_image(
@@ -194,11 +228,12 @@ class ImageSynthesizer:
         seed: int | None,
         steps: int,
         aspect_ratio: str,
+        size: str | None = None,
     ) -> PILImage:
         import httpx
         from PIL import Image  # type: ignore[import-unresolved]
 
-        width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
+        width, height = _parse_size(size) or _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
         payload: dict[str, Any] = {
             "prompt": prompt,
             "n": 1,

kgmodule_utils-0.4.3/src/kg_utils/synthesis/factory.py

@@ -0,0 +1,97 @@
+# © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
+"""Synthesis backend factory helpers for per-request backend overrides."""
+
+from __future__ import annotations
+
+import os
+
+from kg_utils.synthesis._config import (
+    ImageBackend,
+    ImageConfig,
+    TextBackend,
+    TextConfig,
+)
+from kg_utils.synthesis._image import ImageSynthesizer
+from kg_utils.synthesis._text import TextSynthesizer
+
+__all__ = [
+    "normalize_openai_base_url",
+    "text_synth_for_backend",
+    "image_synth_for_backend",
+]
+
+
+def normalize_openai_base_url(endpoint: str) -> str:
+    """Normalize an OpenAI-wire endpoint so it ends with /v1.
+
+    Returns an empty string when endpoint is empty.
+    """
+    ep = (endpoint or "").strip().rstrip("/")
+    if not ep:
+        return ""
+    if ep.endswith("/v1"):
+        return ep
+    return f"{ep}/v1"
+
+
+def text_synth_for_backend(backend: str, fallback: TextSynthesizer) -> TextSynthesizer:
+    """Return a TextSynthesizer configured for a specific backend override.
+
+    Unknown or empty backend strings return ``fallback``.
+    """
+    backend_str = (backend or "").strip().lower()
+    if not backend_str:
+        return fallback
+
+    try:
+        selected = TextBackend(backend_str)
+    except ValueError:
+        return fallback
+
+    if selected == TextBackend.OMLX:
+        endpoint = os.environ.get("SYNTH_ENDPOINT") or os.environ.get("VLLM_ENDPOINT_URL") or ""
+        endpoint = normalize_openai_base_url(endpoint)
+        api_key = os.environ.get("SYNTH_API_KEY") or os.environ.get("VLLM_API_KEY") or ""
+        model = os.environ.get("SYNTH_MODEL") or os.environ.get("VLLM_MODEL") or ""
+        return TextSynthesizer(
+            TextConfig(backend=selected, endpoint=endpoint, api_key=api_key, model=model)
+        )
+
+    if selected == TextBackend.OLLAMA:
+        endpoint = os.environ.get("OLLAMA_ENDPOINT") or ""
+        return TextSynthesizer(TextConfig(backend=selected, endpoint=endpoint))
+
+    if selected == TextBackend.OPENAI:
+        api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("SYNTH_API_KEY") or ""
+        return TextSynthesizer(TextConfig(backend=selected, api_key=api_key))
+
+    return fallback
+
+
+def image_synth_for_backend(backend: str, fallback: ImageSynthesizer) -> ImageSynthesizer:
+    """Return an ImageSynthesizer configured for a specific backend override.
+
+    Unknown or empty backend strings return ``fallback``.
+    """
+    backend_str = (backend or "").strip().lower()
+    if not backend_str:
+        return fallback
+
+    try:
+        selected = ImageBackend(backend_str)
+    except ValueError:
+        return fallback
+
+    if selected == ImageBackend.OPENAI:
+        api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("IMAGE_API_KEY") or ""
+        return ImageSynthesizer(ImageConfig(backend=selected, api_key=api_key))
+
+    if selected == ImageBackend.MFLUX_SERVE:
+        server_url = os.environ.get("IMAGE_ENDPOINT") or ""
+        return ImageSynthesizer(ImageConfig(backend=selected, server_url=server_url))
+
+    if selected == ImageBackend.MFLUX_LOCAL:
+        model = os.environ.get("IMAGE_MODEL") or os.environ.get("GUTENKG_IMAGE_MODEL") or ""
+        return ImageSynthesizer(ImageConfig(backend=selected, model=model))
+
+    return fallback

kgmodule_utils-0.4.3/src/kg_utils/worker/__init__.py

@@ -0,0 +1,17 @@
+"""Worker protocol helpers and client for RunPod ``/runsync`` endpoints."""
+
+from kg_utils.worker.client import (
+    WorkerClient,
+    WorkerError,
+    decode_worker_response,
+    extract_worker_error,
+)
+from kg_utils.worker.ops import handle_aux_ops
+
+__all__ = [
+    "WorkerClient",
+    "WorkerError",
+    "decode_worker_response",
+    "extract_worker_error",
+    "handle_aux_ops",
+]

kgmodule_utils-0.4.3/src/kg_utils/worker/client.py

@@ -0,0 +1,204 @@
+# © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
+"""RunPod worker client utilities for chat and handler front-ends.
+
+This module centralizes payload construction and response/error decoding for
+``/runsync`` worker calls used by Streamlit clients.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, cast
+
+import httpx
+
+
+class WorkerError(Exception):
+    """Raised when a worker response contains a structured application-level error."""
+
+
+def _format_error_data(error_data: object) -> str:
+    if isinstance(error_data, str):
+        try:
+            decoded = json.loads(error_data)
+        except (ValueError, TypeError):
+            return error_data
+        if isinstance(decoded, dict):
+            err_type = decoded.get("error_type", "Unknown")
+            err_msg = decoded.get("error_message", str(decoded))
+            return f"{err_type}: {err_msg}"
+        return str(decoded)
+
+    if isinstance(error_data, dict):
+        d = cast("dict[str, Any]", error_data)
+        err_type = d.get("error_type", "Unknown")
+        err_msg = d.get("error_message", str(error_data))
+        return f"{err_type}: {err_msg}"
+
+    return str(error_data)
+
+
+def extract_worker_error(data: object) -> str | None:
+    """Extract a readable worker error from a raw RunPod response payload."""
+    if not isinstance(data, dict):
+        return str(data)
+
+    d = cast("dict[str, Any]", data)
+    if d.get("status") == "FAILED" or "error_type" in d:
+        return _format_error_data(d.get("error", d))
+
+    out = d.get("output")
+    if isinstance(out, dict) and isinstance(out.get("error"), str):
+        return out["error"]
+
+    return None
+
+
+def decode_worker_response(data: object) -> dict:
+    """Decode a worker response payload and raise WorkerError on application errors."""
+    error = extract_worker_error(data)
+    if error:
+        raise WorkerError(error)
+
+    if not isinstance(data, dict):
+        raise WorkerError(f"unexpected worker response type: {type(data).__name__}")
+
+    d2 = cast("dict[str, Any]", data)
+    out = d2.get("output", d2)
+    if not isinstance(out, dict):
+        raise WorkerError(f"unexpected worker output type: {type(out).__name__}")
+    return out
+
+
+class WorkerClient:
+    """Small client for RunPod ``/runsync`` worker endpoints."""
+
+    def __init__(self, base_url: str, secret: str = "") -> None:
+        self._base_url = base_url.rstrip("/")
+        self._secret = secret
+
+    def _post(self, payload: dict, timeout: httpx.Timeout) -> dict:
+        resp = httpx.post(f"{self._base_url}/runsync", json=payload, timeout=timeout)
+        resp.raise_for_status()
+        return resp.json()
+
+    def list_models(self, backend: str = "") -> tuple[list[str], str]:
+        payload: dict = {"input": {"op": "models"}}
+        if backend:
+            payload["input"]["backend"] = backend
+        if self._secret:
+            payload["input"]["secret"] = self._secret
+
+        try:
+            data = self._post(
+                payload,
+                timeout=httpx.Timeout(connect=5.0, read=20.0, write=5.0, pool=5.0),
+            )
+            out = data.get("output", {}) if isinstance(data, dict) else {}
+            if not isinstance(out, dict):
+                return [], ""
+            return out.get("models", []), out.get("default", "")
+        except Exception:  # noqa: BLE001
+            return [], ""
+
+    def rewrite(
+        self,
+        text: str,
+        backend: str = "",
+        model: str = "",
+    ) -> tuple[str, str | None]:
+        payload: dict = {"input": {"op": "rewrite", "text": text}}
+        if backend:
+            payload["input"]["backend"] = backend
+        if model:
+            payload["input"]["model"] = model
+        if self._secret:
+            payload["input"]["secret"] = self._secret
+
+        try:
+            data = self._post(
+                payload,
+                timeout=httpx.Timeout(connect=5.0, read=60.0, write=10.0, pool=5.0),
+            )
+            err = extract_worker_error(data)
+            if err:
+                return text, err
+            out = data.get("output", {}) if isinstance(data, dict) else {}
+            if not isinstance(out, dict):
+                return text, "unexpected worker output"
+            return out.get("prompt", text), out.get("error")
+        except Exception as exc:  # noqa: BLE001
+            return text, str(exc)
+
+    def imagine(
+        self,
+        prompt: str,
+        *,
+        image_backend: str = "",
+        aspect_ratio: str = "3:2",
+        steps: int | None = None,
+        size: str | None = None,
+    ) -> tuple[str | None, str | None, str | None, str | None]:
+        inp: dict[str, Any] = {"op": "imagine", "prompt": prompt, "aspect_ratio": aspect_ratio}
+        if image_backend:
+            inp["image_backend"] = image_backend
+        if steps is not None:
+            inp["steps"] = steps
+        if size:
+            inp["size"] = size
+        if self._secret:
+            inp["secret"] = self._secret
+        payload: dict[str, Any] = {"input": inp}
+
+        try:
+            data = self._post(
+                payload,
+                timeout=httpx.Timeout(connect=5.0, read=300.0, write=10.0, pool=5.0),
+            )
+            err = extract_worker_error(data)
+            if err:
+                return None, None, None, err
+
+            out = data.get("output", {}) if isinstance(data, dict) else {}
+            if not isinstance(out, dict):
+                return None, None, None, "unexpected worker output"
+            if "error" in out:
+                return None, None, None, str(out["error"])
+            return out.get("image_b64"), out.get("image_model"), out.get("image_backend"), None
+        except Exception as exc:  # noqa: BLE001
+            return None, None, None, str(exc)
+
+    def query(
+        self,
+        query: str,
+        *,
+        corpus: str = "all",
+        k: int = 8,
+        min_score: float = 0.0,
+        semantic_floor: float = 0.0,
+        synthesize: bool = False,
+        model: str = "",
+        backend: str = "",
+    ) -> dict:
+        payload: dict = {
+            "input": {
+                "query": query,
+                "corpus": corpus,
+                "k": k,
+                "min_score": min_score,
+                "semantic_floor": semantic_floor,
+                "synthesize": synthesize,
+            }
+        }
+        if model:
+            payload["input"]["model"] = model
+        if backend:
+            payload["input"]["backend"] = backend
+        if self._secret:
+            payload["input"]["secret"] = self._secret
+
+        data = self._post(
+            payload,
+            timeout=httpx.Timeout(connect=5.0, read=600.0, write=30.0, pool=5.0),
+        )
+        return decode_worker_response(data)

kgmodule_utils-0.4.3/src/kg_utils/worker/ops.py

@@ -0,0 +1,87 @@
+# © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
+"""Shared handler operation dispatch for models, rewrite, and imagine."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from kg_utils.synthesis._image import ImageSynthesizer
+from kg_utils.synthesis._text import TextSynthesizer
+
+__all__ = ["handle_aux_ops"]
+
+
+def handle_aux_ops(
+    inp: dict,
+    text_synth_factory: Callable[[str], TextSynthesizer],
+    image_synth_factory: Callable[[str], ImageSynthesizer],
+) -> dict | None:
+    """Handle shared non-query worker operations.
+
+    Returns:
+    - operation payload dict when op is recognized
+    - ``None`` when input has no recognized operation
+    """
+    op = inp.get("op")
+
+    if op == "models":
+        synth = text_synth_factory(inp.get("backend", ""))
+        # Existing handlers expose the active model via synthesizer config internals.
+        return {
+            "models": synth.list_models(),
+            "default": synth._cfg.resolved_model(),  # pylint: disable=protected-access
+        }
+
+    if op == "rewrite":
+        text = (inp.get("text") or "").strip()
+        if not text:
+            return {"error": "rewrite requires a non-empty 'text'"}
+
+        synth = text_synth_factory(inp.get("backend", ""))
+        model_override = (inp.get("model") or "").strip() or None
+        prompt, error = synth.rewrite_for_image(text, model=model_override)
+        return {"prompt": prompt, "error": error}
+
+    if op == "imagine":
+        prompt = (inp.get("prompt") or "").strip()
+        if not prompt:
+            return {"error": "imagine requires a non-empty 'prompt'"}
+
+        aspect = inp.get("aspect_ratio", "3:2")
+        size = inp.get("size")
+        seed = inp.get("seed")
+        steps = inp.get("steps")
+        img_synth = image_synth_factory(inp.get("image_backend", ""))
+        seed_int = int(seed) if seed is not None else None
+        steps_int = int(steps) if steps is not None else None
+
+        try:
+            if size is not None:
+                b64 = img_synth.generate_b64(
+                    prompt,
+                    aspect_ratio=aspect,
+                    seed=seed_int,
+                    steps=steps_int,
+                    size=size,
+                )
+            else:
+                b64 = img_synth.generate_b64(
+                    prompt,
+                    aspect_ratio=aspect,
+                    seed=seed_int,
+                    steps=steps_int,
+                )
+            result = {
+                "image_b64": b64,
+                "prompt": prompt,
+                "aspect_ratio": aspect,
+                "image_model": img_synth._cfg.resolved_model(),  # pylint: disable=protected-access
+                "image_backend": img_synth._cfg.backend.value,  # pylint: disable=protected-access
+            }
+            if size is not None:
+                result["size"] = size
+            return result
+        except Exception as exc:  # noqa: BLE001  # pylint: disable=broad-exception-caught
+            return {"error": f"image generation failed: {exc}"}
+
+    return None