kgmodule-utils 0.3.1__tar.gz → 0.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (21) hide show
  1. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/PKG-INFO +11 -1
  2. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/pyproject.toml +14 -22
  3. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/__init__.py +7 -2
  4. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/pipeline.py +2 -2
  5. kgmodule_utils-0.4.0/src/kg_utils/synthesis/__init__.py +79 -0
  6. kgmodule_utils-0.4.0/src/kg_utils/synthesis/_config.py +155 -0
  7. kgmodule_utils-0.4.0/src/kg_utils/synthesis/_image.py +235 -0
  8. kgmodule_utils-0.4.0/src/kg_utils/synthesis/_text.py +178 -0
  9. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/LICENSE +0 -0
  10. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/README.md +0 -0
  11. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/embed.py +0 -0
  12. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/embedder.py +0 -0
  13. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/extractor.py +0 -0
  14. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/module.py +0 -0
  15. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/py.typed +0 -0
  16. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/semantic.py +0 -0
  17. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/snapshots/__init__.py +0 -0
  18. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/snapshots/manager.py +0 -0
  19. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/snapshots/models.py +0 -0
  20. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/specs.py +0 -0
  21. {kgmodule_utils-0.3.1 → kgmodule_utils-0.4.0}/src/kg_utils/store.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: kgmodule-utils
3
- Version: 0.3.1
3
+ Version: 0.4.0
4
4
  Summary: Shared types, graph store, semantic index, and pipeline base for the KGModule SDK
5
5
  License: Elastic-2.0
6
6
  License-File: LICENSE
@@ -14,11 +14,21 @@ Classifier: Programming Language :: Python :: 3
14
14
  Classifier: Programming Language :: Python :: 3.12
15
15
  Classifier: Programming Language :: Python :: 3.13
16
16
  Provides-Extra: semantic
17
+ Provides-Extra: synthesis
18
+ Provides-Extra: synthesis-mflux
19
+ Requires-Dist: httpx (>=0.27.0) ; extra == "synthesis"
20
+ Requires-Dist: httpx (>=0.27.0) ; extra == "synthesis-mflux"
17
21
  Requires-Dist: lancedb (>=0.19.0) ; extra == "semantic"
22
+ Requires-Dist: mflux (>=0.9.0) ; extra == "synthesis-mflux"
18
23
  Requires-Dist: numpy (>=1.24.0) ; extra == "semantic"
24
+ Requires-Dist: openai (>=1.30.0) ; extra == "synthesis"
25
+ Requires-Dist: openai (>=1.30.0) ; extra == "synthesis-mflux"
26
+ Requires-Dist: pillow (>=10.0.0) ; extra == "synthesis"
27
+ Requires-Dist: pillow (>=10.0.0) ; extra == "synthesis-mflux"
19
28
  Requires-Dist: sentence-transformers (>=5.4.1) ; extra == "semantic"
20
29
  Requires-Dist: torch (>=2.5.1) ; extra == "semantic"
21
30
  Requires-Dist: transformers (>=4.40.0,<4.57) ; extra == "semantic"
31
+ Requires-Dist: ty (>=0.0.44,<0.0.45)
22
32
  Project-URL: Repository, https://github.com/Flux-Frontiers/kg_utils
23
33
  Description-Content-Type: text/markdown
24
34
 
@@ -10,7 +10,7 @@ build-backend = "poetry.core.masonry.api"
10
10
 
11
11
  [project]
12
12
  name = "kgmodule-utils"
13
- version = "0.3.1"
13
+ version = "0.4.0"
14
14
  description = "Shared types, graph store, semantic index, and pipeline base for the KGModule SDK"
15
15
  readme = "README.md"
16
16
  license = { text = "Elastic-2.0" }
@@ -26,7 +26,7 @@ classifiers = [
26
26
  "Programming Language :: Python :: 3.13",
27
27
  ]
28
28
  requires-python = ">=3.12,<3.14"
29
- dependencies = []
29
+ dependencies = ["ty (>=0.0.44,<0.0.45)"]
30
30
 
31
31
  [project.optional-dependencies]
32
32
  semantic = [
@@ -36,6 +36,17 @@ semantic = [
36
36
  "torch>=2.5.1",
37
37
  "transformers>=4.40.0,<4.57",
38
38
  ]
39
+ synthesis = [
40
+ "httpx>=0.27.0",
41
+ "openai>=1.30.0",
42
+ "pillow>=10.0.0",
43
+ ]
44
+ synthesis-mflux = [
45
+ "httpx>=0.27.0",
46
+ "openai>=1.30.0",
47
+ "pillow>=10.0.0",
48
+ "mflux>=0.9.0",
49
+ ]
39
50
 
40
51
  [project.urls]
41
52
  Repository = "https://github.com/Flux-Frontiers/kg_utils"
@@ -58,7 +69,7 @@ pytest = "^8.0.0"
58
69
  pytest-cov = "^5.0.0"
59
70
  black = ">=22.0"
60
71
  ruff = ">=0.4.0"
61
- mypy = ">=1.10.0"
72
+ ty = ">=0.0.41"
62
73
 
63
74
  [tool.black]
64
75
  line-length = 100
@@ -76,22 +87,3 @@ init-hook = "import sys; sys.path.insert(0, 'src')"
76
87
  disable = [
77
88
  "missing-module-docstring",
78
89
  ]
79
-
80
- [tool.mypy]
81
- python_version = "3.12"
82
- strict = true
83
- warn_unused_ignores = true
84
- disallow_untyped_defs = true
85
-
86
- [[tool.mypy.overrides]]
87
- module = [
88
- "sentence_transformers.*",
89
- "transformers.*",
90
- "numpy.*",
91
- "lancedb",
92
- ]
93
- ignore_missing_imports = true
94
-
95
- [[tool.mypy.overrides]]
96
- module = "kg_utils.embedder"
97
- disallow_untyped_calls = false
@@ -11,10 +11,15 @@ Sub-packages / modules:
11
11
  kg_model_cache_dir(), resolve_model_path().
12
12
  kg_utils.embedder — Concrete SentenceTransformerEmbedder, get_embedder(),
13
13
  wrap_embedder(), load_sentence_transformer().
14
+ kg_utils.synthesis — Unified text + image synthesis: TextSynthesizer, ImageSynthesizer.
15
+ Backends: omlx | ollama | openai (text);
16
+ mflux-local | mflux-serve | openai (image).
14
17
 
15
18
  Optional extras
16
19
  ---------------
17
- pip install 'kgmodule-utils[semantic]' # lancedb + sentence-transformers
20
+ pip install 'kgmodule-utils[semantic]' # lancedb + sentence-transformers
21
+ pip install 'kgmodule-utils[synthesis]' # httpx + openai + pillow
22
+ pip install 'kgmodule-utils[synthesis-mflux]' # + mflux (Apple Silicon local gen)
18
23
  """
19
24
 
20
- __version__ = "0.3.1"
25
+ __version__ = "0.4.0"
@@ -38,7 +38,7 @@ from __future__ import annotations
38
38
  import re
39
39
  from abc import ABC, abstractmethod
40
40
  from pathlib import Path
41
- from typing import Any
41
+ from typing import Any, cast
42
42
 
43
43
  from kg_utils.semantic import (
44
44
  DEFAULT_MODEL,
@@ -715,7 +715,7 @@ class KGModule(ABC):
715
715
  file_nlines=len(lines),
716
716
  )
717
717
  if n.get("qualname") and n.get("_span"):
718
- spans_by_qualname[(mp, n["qualname"])] = n["_span"]
718
+ spans_by_qualname[(mp, n["qualname"])] = cast(tuple[int, int], n["_span"])
719
719
 
720
720
  raw_nodes.sort(key=lambda x: x["_rank_key"])
721
721
 
@@ -0,0 +1,79 @@
1
+ """kg_utils.synthesis — Unified text and image synthesis backends.
2
+
3
+ Text backends (TextBackend):
4
+ omlx — local oMLX / vLLM (OpenAI wire protocol; MLX chain-of-thought suppressed)
5
+ ollama — local Ollama (OpenAI wire protocol; no api_key required)
6
+ openai — OpenAI cloud API (requires OPENAI_API_KEY or SYNTH_API_KEY)
7
+
8
+ Image backends (ImageBackend):
9
+ mflux-local — local Flux2Klein via mflux (Apple Silicon only)
10
+ mflux-serve — HTTP proxy to mflux-serve (runs anywhere)
11
+ openai — OpenAI DALL-E (requires OPENAI_API_KEY or IMAGE_API_KEY)
12
+
13
+ Env vars
14
+ --------
15
+ Text synthesis:
16
+ SYNTH_BACKEND omlx | ollama | openai (default: omlx)
17
+ SYNTH_ENDPOINT override base URL
18
+ SYNTH_API_KEY bearer token / OpenAI key
19
+ SYNTH_MODEL model-id override
20
+ Legacy aliases: VLLM_ENDPOINT_URL → SYNTH_ENDPOINT
21
+ VLLM_API_KEY → SYNTH_API_KEY
22
+ VLLM_MODEL → SYNTH_MODEL
23
+
24
+ Image synthesis:
25
+ IMAGE_BACKEND mflux-local | mflux-serve | openai (default: mflux-serve)
26
+ IMAGE_ENDPOINT mflux-serve base URL
27
+ IMAGE_API_KEY OpenAI key (also reads OPENAI_API_KEY)
28
+ IMAGE_MODEL model-id override
29
+ IMAGE_STEPS inference steps (default: 4, mflux backends only)
30
+ Legacy alias: GUTENKG_IMAGE_MODEL → IMAGE_MODEL
31
+
32
+ Quick start
33
+ -----------
34
+ from kg_utils.synthesis import text_synthesizer_from_env, image_synthesizer_from_env
35
+
36
+ text = text_synthesizer_from_env()
37
+ image = image_synthesizer_from_env()
38
+
39
+ answer = text.synthesize_rag(query, hits)
40
+ b64 = image.generate_b64(prompt)
41
+ rewritten_prompt, err = text.rewrite_for_image(corpus_passage)
42
+ """
43
+
44
+ from __future__ import annotations
45
+
46
+ from kg_utils.synthesis._config import (
47
+ ImageBackend,
48
+ ImageConfig,
49
+ TextBackend,
50
+ TextConfig,
51
+ image_config_from_env,
52
+ text_config_from_env,
53
+ )
54
+ from kg_utils.synthesis._image import ImageSynthesizer
55
+ from kg_utils.synthesis._text import TextSynthesizer
56
+
57
+
58
+ def text_synthesizer_from_env() -> TextSynthesizer:
59
+ """Build a TextSynthesizer configured from the current process environment."""
60
+ return TextSynthesizer(text_config_from_env())
61
+
62
+
63
+ def image_synthesizer_from_env() -> ImageSynthesizer:
64
+ """Build an ImageSynthesizer configured from the current process environment."""
65
+ return ImageSynthesizer(image_config_from_env())
66
+
67
+
68
+ __all__ = [
69
+ "TextBackend",
70
+ "ImageBackend",
71
+ "TextConfig",
72
+ "ImageConfig",
73
+ "TextSynthesizer",
74
+ "ImageSynthesizer",
75
+ "text_synthesizer_from_env",
76
+ "image_synthesizer_from_env",
77
+ "text_config_from_env",
78
+ "image_config_from_env",
79
+ ]
@@ -0,0 +1,155 @@
1
+ # © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
2
+ """Synthesis backend configuration — enums, dataclasses, and env-var factories."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import os
7
+ from dataclasses import dataclass
8
+ from enum import Enum
9
+
10
+
11
+ class TextBackend(str, Enum):
12
+ """LLM backend for text synthesis."""
13
+
14
+ OMLX = "omlx" # local oMLX / vLLM — OpenAI wire protocol with MLX thinking suppression
15
+ OLLAMA = "ollama" # local Ollama — OpenAI wire protocol, no api_key required
16
+ OPENAI = "openai" # OpenAI cloud API — requires OPENAI_API_KEY / SYNTH_API_KEY
17
+
18
+
19
+ class ImageBackend(str, Enum):
20
+ """Image generation backend."""
21
+
22
+ MFLUX_LOCAL = "mflux-local" # local Flux2Klein via mflux (Apple Silicon only)
23
+ MFLUX_SERVE = "mflux-serve" # HTTP proxy to a running mflux-serve instance
24
+ OPENAI = "openai" # OpenAI DALL-E — requires OPENAI_API_KEY / IMAGE_API_KEY
25
+
26
+
27
+ # Per-backend defaults — filled in when the user has not provided an override.
28
+ _TEXT_DEFAULTS: dict[TextBackend, dict[str, str]] = {
29
+ TextBackend.OMLX: {
30
+ "endpoint": "http://localhost:8080/v1",
31
+ "model": "Qwen3-4B-Instruct-2507-MLX-8bit",
32
+ },
33
+ TextBackend.OLLAMA: {
34
+ "endpoint": "http://localhost:11434/v1",
35
+ "model": "hf.co/unsloth/Qwen3-4B-Instruct-2507-GGUF:Q8_0",
36
+ },
37
+ TextBackend.OPENAI: {
38
+ "endpoint": "https://api.openai.com/v1",
39
+ "model": "gpt-4o-mini",
40
+ },
41
+ }
42
+
43
+ _IMAGE_DEFAULTS: dict[ImageBackend, dict[str, str]] = {
44
+ ImageBackend.MFLUX_LOCAL: {
45
+ "model": "mlx-community/flux2-klein-4b-4bit",
46
+ },
47
+ ImageBackend.MFLUX_SERVE: {
48
+ "server_url": "http://localhost:8090",
49
+ "model": "flux2-klein-4b",
50
+ },
51
+ ImageBackend.OPENAI: {
52
+ "model": "dall-e-3",
53
+ },
54
+ }
55
+
56
+
57
+ @dataclass
58
+ class TextConfig:
59
+ """Configuration for a text synthesis backend.
60
+
61
+ :param backend: Which LLM provider to use.
62
+ :param endpoint: Base URL override (empty = use backend default).
63
+ :param api_key: Bearer token / OpenAI key (empty = not required for omlx/ollama).
64
+ :param model: Model-id override (empty = use backend default).
65
+ :param max_tokens: Maximum tokens in the generated response.
66
+ :param suppress_thinking: Strip ``<think>`` blocks and pass ``extra_body`` to disable
67
+ chain-of-thought for oMLX backends.
68
+ """
69
+
70
+ backend: TextBackend = TextBackend.OMLX
71
+ endpoint: str = ""
72
+ api_key: str = ""
73
+ model: str = ""
74
+ max_tokens: int = 2048
75
+ suppress_thinking: bool = True
76
+
77
+ def resolved_endpoint(self) -> str:
78
+ return self.endpoint or _TEXT_DEFAULTS[self.backend]["endpoint"]
79
+
80
+ def resolved_model(self) -> str:
81
+ return self.model or _TEXT_DEFAULTS[self.backend]["model"]
82
+
83
+
84
+ @dataclass
85
+ class ImageConfig:
86
+ """Configuration for an image synthesis backend.
87
+
88
+ :param backend: Which image provider to use.
89
+ :param server_url: mflux-serve base URL override (empty = use backend default).
90
+ :param api_key: OpenAI API key (empty = read from OPENAI_API_KEY at call time).
91
+ :param model: Model-id override (empty = use backend default).
92
+ :param steps: Inference steps for mflux backends (ignored for OpenAI).
93
+ """
94
+
95
+ backend: ImageBackend = ImageBackend.MFLUX_SERVE
96
+ server_url: str = ""
97
+ api_key: str = ""
98
+ model: str = ""
99
+ steps: int = 4
100
+
101
+ def resolved_server_url(self) -> str:
102
+ return self.server_url or _IMAGE_DEFAULTS[self.backend].get("server_url", "")
103
+
104
+ def resolved_model(self) -> str:
105
+ return self.model or _IMAGE_DEFAULTS[self.backend].get("model", "")
106
+
107
+
108
+ # ---------------------------------------------------------------------------
109
+ # Env-var factories
110
+ # ---------------------------------------------------------------------------
111
+
112
+
113
+ def text_config_from_env() -> TextConfig:
114
+ """Build a TextConfig from environment variables.
115
+
116
+ Reads ``SYNTH_BACKEND`` (default ``omlx``) plus endpoint/key/model overrides.
117
+ Legacy vars ``VLLM_ENDPOINT_URL``, ``VLLM_API_KEY``, and ``VLLM_MODEL`` are
118
+ honoured as fallbacks so existing deployments need no changes.
119
+ """
120
+ backend_str = os.environ.get("SYNTH_BACKEND", "omlx")
121
+ try:
122
+ backend = TextBackend(backend_str)
123
+ except ValueError:
124
+ backend = TextBackend.OMLX
125
+
126
+ endpoint = os.environ.get("SYNTH_ENDPOINT") or os.environ.get("VLLM_ENDPOINT_URL") or ""
127
+ api_key = (
128
+ os.environ.get("SYNTH_API_KEY")
129
+ or os.environ.get("VLLM_API_KEY")
130
+ or os.environ.get("OPENAI_API_KEY")
131
+ or ""
132
+ )
133
+ model = os.environ.get("SYNTH_MODEL") or os.environ.get("VLLM_MODEL") or ""
134
+ return TextConfig(backend=backend, endpoint=endpoint, api_key=api_key, model=model)
135
+
136
+
137
+ def image_config_from_env() -> ImageConfig:
138
+ """Build an ImageConfig from environment variables.
139
+
140
+ Reads ``IMAGE_BACKEND`` (default ``mflux-serve``) plus server-url/key/model/steps overrides.
141
+ Legacy var ``GUTENKG_IMAGE_MODEL`` is honoured as a fallback for ``IMAGE_MODEL``.
142
+ """
143
+ backend_str = os.environ.get("IMAGE_BACKEND", "mflux-serve")
144
+ try:
145
+ backend = ImageBackend(backend_str)
146
+ except ValueError:
147
+ backend = ImageBackend.MFLUX_SERVE
148
+
149
+ server_url = os.environ.get("IMAGE_ENDPOINT") or ""
150
+ api_key = os.environ.get("IMAGE_API_KEY") or os.environ.get("OPENAI_API_KEY") or ""
151
+ model = os.environ.get("IMAGE_MODEL") or os.environ.get("GUTENKG_IMAGE_MODEL") or ""
152
+ steps = int(os.environ.get("IMAGE_STEPS", "4"))
153
+ return ImageConfig(
154
+ backend=backend, server_url=server_url, api_key=api_key, model=model, steps=steps
155
+ )
@@ -0,0 +1,235 @@
1
+ # © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
2
+ """ImageSynthesizer — unified image generation across mflux-local, mflux-serve, and DALL-E."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import base64
7
+ import random
8
+ from io import BytesIO
9
+ from typing import TYPE_CHECKING, Any
10
+
11
+ if TYPE_CHECKING:
12
+ from PIL.Image import Image as PILImage # type: ignore[import-unresolved]
13
+
14
+ from kg_utils.synthesis._config import ImageBackend, ImageConfig
15
+
16
+ # ---------------------------------------------------------------------------
17
+ # Aspect-ratio → pixel dimension tables
18
+ # ---------------------------------------------------------------------------
19
+
20
+ # Used by mflux backends (arbitrary resolution).
21
+ _MFLUX_SIZES: dict[str, tuple[int, int]] = {
22
+ "1:1": (1024, 1024),
23
+ "3:2": (1536, 1024),
24
+ "2:3": (1024, 1536),
25
+ "16:9": (1536, 864),
26
+ "9:16": (864, 1536),
27
+ "4:3": (1365, 1024),
28
+ "3:4": (1024, 1365),
29
+ }
30
+
31
+ # DALL-E 3 only supports three sizes; map aspect ratios to the nearest valid option.
32
+ _DALLE3_SIZES: dict[str, str] = {
33
+ "1:1": "1024x1024",
34
+ "3:2": "1792x1024",
35
+ "2:3": "1024x1792",
36
+ "16:9": "1792x1024",
37
+ "9:16": "1024x1792",
38
+ "4:3": "1792x1024",
39
+ "3:4": "1024x1792",
40
+ }
41
+
42
+
43
+ # ---------------------------------------------------------------------------
44
+ # Synthesizer
45
+ # ---------------------------------------------------------------------------
46
+
47
+
48
+ class ImageSynthesizer:
49
+ """Unified image generation across mflux-local, mflux-serve, and OpenAI DALL-E.
50
+
51
+ All ``generate*`` methods return a PIL ``Image`` — callers that need a base64
52
+ string (e.g. HTTP wire responses) should use ``generate_b64()``. A per-instance
53
+ mflux model cache avoids reloading across calls to the same synthesizer.
54
+
55
+ :param config: Backend configuration produced by ``image_config_from_env()`` or
56
+ built directly as an ``ImageConfig`` dataclass.
57
+ """
58
+
59
+ def __init__(self, config: ImageConfig) -> None:
60
+ self._cfg = config
61
+ self._mflux_model: Any = None
62
+ self._mflux_model_name: str | None = None
63
+
64
+ # ------------------------------------------------------------------
65
+ # Internal helpers
66
+ # ------------------------------------------------------------------
67
+
68
+ def _load_mflux(self, model_name: str) -> Any:
69
+ if self._mflux_model is not None and self._mflux_model_name == model_name:
70
+ return self._mflux_model
71
+ from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein # type: ignore
72
+
73
+ self._mflux_model = Flux2Klein(model_path=model_name)
74
+ self._mflux_model_name = model_name
75
+ return self._mflux_model
76
+
77
+ # ------------------------------------------------------------------
78
+ # Public API
79
+ # ------------------------------------------------------------------
80
+
81
+ def generate(
82
+ self,
83
+ prompt: str,
84
+ *,
85
+ aspect_ratio: str = "3:2",
86
+ seed: int | None = None,
87
+ steps: int | None = None,
88
+ model: str | None = None,
89
+ ) -> PILImage:
90
+ """Generate an image and return a PIL Image.
91
+
92
+ :param prompt: Text description of the scene to generate.
93
+ :param aspect_ratio: One of 1:1, 3:2, 2:3, 16:9, 9:16, 4:3, 3:4.
94
+ :param seed: Random seed for reproducibility (random int if omitted).
95
+ :param steps: Override inference steps (mflux backends only; ignored for OpenAI).
96
+ :param model: Override the configured model for this single call.
97
+ :returns: PIL Image.
98
+ """
99
+ cfg = self._cfg
100
+ effective_steps = steps if steps is not None else cfg.steps
101
+ effective_model = model or cfg.resolved_model()
102
+
103
+ if cfg.backend == ImageBackend.MFLUX_LOCAL:
104
+ return self._generate_local(
105
+ prompt,
106
+ model=effective_model,
107
+ seed=seed,
108
+ steps=effective_steps,
109
+ aspect_ratio=aspect_ratio,
110
+ )
111
+ elif cfg.backend == ImageBackend.MFLUX_SERVE:
112
+ return self._generate_via_server(
113
+ prompt,
114
+ server_url=cfg.resolved_server_url(),
115
+ seed=seed,
116
+ steps=effective_steps,
117
+ aspect_ratio=aspect_ratio,
118
+ )
119
+ else:
120
+ return self._generate_openai(
121
+ prompt,
122
+ model=effective_model,
123
+ aspect_ratio=aspect_ratio,
124
+ )
125
+
126
+ def generate_b64(
127
+ self,
128
+ prompt: str,
129
+ *,
130
+ aspect_ratio: str = "3:2",
131
+ seed: int | None = None,
132
+ steps: int | None = None,
133
+ model: str | None = None,
134
+ ) -> str:
135
+ """Generate an image and return it as a base64-encoded PNG string.
136
+
137
+ Convenience wrapper around ``generate()`` for HTTP wire responses.
138
+
139
+ :param prompt: Text description of the scene to generate.
140
+ :param aspect_ratio: One of 1:1, 3:2, 2:3, 16:9, 9:16, 4:3, 3:4.
141
+ :param seed: Random seed for reproducibility.
142
+ :param steps: Override inference steps (mflux backends only).
143
+ :param model: Override the configured model for this single call.
144
+ :returns: Base64-encoded PNG string.
145
+ """
146
+ img = self.generate(prompt, aspect_ratio=aspect_ratio, seed=seed, steps=steps, model=model)
147
+ buf = BytesIO()
148
+ img.save(buf, format="PNG")
149
+ return base64.b64encode(buf.getvalue()).decode()
150
+
151
+ # ------------------------------------------------------------------
152
+ # Backend implementations
153
+ # ------------------------------------------------------------------
154
+
155
+ def _generate_local(
156
+ self,
157
+ prompt: str,
158
+ *,
159
+ model: str,
160
+ seed: int | None,
161
+ steps: int,
162
+ aspect_ratio: str,
163
+ ) -> PILImage:
164
+ width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
165
+ effective_seed = seed if seed is not None else random.randint(0, 2**31 - 1)
166
+ flux = self._load_mflux(model)
167
+ result = flux.generate_image(
168
+ seed=effective_seed,
169
+ prompt=prompt,
170
+ width=width,
171
+ height=height,
172
+ guidance=1.0,
173
+ num_inference_steps=steps,
174
+ scheduler="flow_match_euler_discrete",
175
+ )
176
+ return result.image
177
+
178
+ def _generate_via_server(
179
+ self,
180
+ prompt: str,
181
+ *,
182
+ server_url: str,
183
+ seed: int | None,
184
+ steps: int,
185
+ aspect_ratio: str,
186
+ ) -> PILImage:
187
+ import httpx
188
+ from PIL import Image # type: ignore[import-unresolved]
189
+
190
+ width, height = _MFLUX_SIZES.get(aspect_ratio, _MFLUX_SIZES["3:2"])
191
+ payload: dict[str, Any] = {
192
+ "prompt": prompt,
193
+ "n": 1,
194
+ "size": f"{width}x{height}",
195
+ "num_inference_steps": steps,
196
+ "response_format": "b64_json",
197
+ }
198
+ if seed is not None:
199
+ payload["seed"] = seed
200
+
201
+ resp = httpx.post(
202
+ server_url.rstrip("/") + "/v1/images/generations",
203
+ json=payload,
204
+ timeout=httpx.Timeout(connect=5.0, read=300.0, write=10.0, pool=5.0),
205
+ )
206
+ resp.raise_for_status()
207
+ b64 = resp.json()["data"][0]["b64_json"]
208
+ return Image.open(BytesIO(base64.b64decode(b64)))
209
+
210
+ def _generate_openai(
211
+ self,
212
+ prompt: str,
213
+ *,
214
+ model: str,
215
+ aspect_ratio: str,
216
+ ) -> PILImage:
217
+ import os
218
+
219
+ from openai import OpenAI # type: ignore[import-unresolved]
220
+ from PIL import Image # type: ignore[import-unresolved]
221
+
222
+ cfg = self._cfg
223
+ api_key = cfg.api_key or os.environ.get("OPENAI_API_KEY", "")
224
+ size = _DALLE3_SIZES.get(aspect_ratio, _DALLE3_SIZES["3:2"])
225
+
226
+ client = OpenAI(api_key=api_key)
227
+ resp = client.images.generate(
228
+ model=model,
229
+ prompt=prompt,
230
+ n=1,
231
+ size=size, # type: ignore[arg-type]
232
+ response_format="b64_json",
233
+ )
234
+ b64 = (resp.data[0].b64_json or "") if resp.data else ""
235
+ return Image.open(BytesIO(base64.b64decode(b64)))
@@ -0,0 +1,178 @@
1
+ # © 2026 Eric G. Suchanek, PhD — Flux-Frontiers · SPDX-License-Identifier: Elastic-2.0
2
+ """TextSynthesizer — unified LLM completion across oMLX, Ollama, and OpenAI backends."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import re
7
+
8
+ from kg_utils.synthesis._config import TextBackend, TextConfig
9
+
10
+ # ---------------------------------------------------------------------------
11
+ # System prompts
12
+ # ---------------------------------------------------------------------------
13
+
14
+ _RAG_SYSTEM = (
15
+ "You are a literary guide to the Project Gutenberg corpus. "
16
+ "Answer the question using ONLY the provided source passages. "
17
+ "Do NOT use any prior knowledge — if something is in the passages, "
18
+ "report it; if it is not in the passages, say so. "
19
+ "Never contradict or override what the passages say based on what "
20
+ "you believe to be true. Be concise and specific. "
21
+ "Cite the author and work when relevant."
22
+ )
23
+
24
+ _IMAGE_REWRITE_SYSTEM = (
25
+ "You are an expert art director. Given a passage of historical text, write a single concise "
26
+ "image generation prompt (one paragraph, no bullet points, no quotation marks) that vividly "
27
+ "describes the scene for a text-to-image model. Focus on visual elements: setting, lighting, "
28
+ "figures, mood, and artistic style. Do NOT include any text, labels, captions, or words in "
29
+ "the scene description. Output ONLY the prompt, nothing else."
30
+ )
31
+
32
+
33
+ # ---------------------------------------------------------------------------
34
+ # Synthesizer
35
+ # ---------------------------------------------------------------------------
36
+
37
+
38
+ class TextSynthesizer:
39
+ """Unified text completion client across oMLX, Ollama, and OpenAI backends.
40
+
41
+ Uses the ``openai`` Python package for all three — each backend is just a
42
+ different ``base_url`` and optionally an ``extra_body`` for thinking suppression.
43
+
44
+ :param config: Backend configuration produced by ``text_config_from_env()`` or
45
+ built directly as a ``TextConfig`` dataclass.
46
+ """
47
+
48
+ def __init__(self, config: TextConfig) -> None:
49
+ self._cfg = config
50
+
51
+ # ------------------------------------------------------------------
52
+ # Internal helpers
53
+ # ------------------------------------------------------------------
54
+
55
+ def _client(self): # type: ignore[no-untyped-def]
56
+ from openai import OpenAI # type: ignore[import-unresolved]
57
+
58
+ cfg = self._cfg
59
+ api_key = cfg.api_key or "not-needed"
60
+ return OpenAI(base_url=cfg.resolved_endpoint(), api_key=api_key)
61
+
62
+ def _extra_body(self) -> dict | None:
63
+ if self._cfg.suppress_thinking and self._cfg.backend == TextBackend.OMLX:
64
+ return {"think": False, "chat_template_kwargs": {"enable_thinking": False}}
65
+ return None
66
+
67
+ @staticmethod
68
+ def _strip_thinking(text: str) -> str:
69
+ return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()
70
+
71
+ def _complete(
72
+ self,
73
+ messages: list[dict],
74
+ *,
75
+ model: str | None = None,
76
+ max_tokens: int | None = None,
77
+ temperature: float = 0.7,
78
+ ) -> str | None:
79
+ try:
80
+ client = self._client()
81
+ mdl = model or self._cfg.resolved_model()
82
+ extra = self._extra_body()
83
+ kwargs: dict = {
84
+ "model": mdl,
85
+ "messages": messages,
86
+ "max_tokens": max_tokens or self._cfg.max_tokens,
87
+ "temperature": temperature,
88
+ }
89
+ if extra:
90
+ kwargs["extra_body"] = extra
91
+ resp = client.chat.completions.create(**kwargs)
92
+ raw = (resp.choices[0].message.content or "").strip()
93
+ return self._strip_thinking(raw) or None
94
+ except Exception: # noqa: BLE001
95
+ return None
96
+
97
+ # ------------------------------------------------------------------
98
+ # Public API
99
+ # ------------------------------------------------------------------
100
+
101
+ def list_models(self) -> list[str]:
102
+ """Return the model IDs available at the configured endpoint.
103
+
104
+ :returns: List of model ID strings; empty list on failure or if the
105
+ endpoint does not support the ``/v1/models`` route.
106
+ """
107
+ try:
108
+ return [m.id for m in self._client().models.list().data if m.id]
109
+ except Exception: # noqa: BLE001
110
+ return []
111
+
112
+ def synthesize_rag(
113
+ self,
114
+ query: str,
115
+ snippets: list[dict],
116
+ *,
117
+ model: str | None = None,
118
+ max_k: int = 12,
119
+ system: str | None = None,
120
+ ) -> str | None:
121
+ """Synthesize a grounded answer from corpus retrieval snippets.
122
+
123
+ Formats snippet metadata (genre · author · title) as a structured context
124
+ block, then calls the configured LLM to answer *only* from those passages.
125
+
126
+ :param query: Natural-language question.
127
+ :param snippets: Hit dicts from KGRAG — must contain a ``content`` key.
128
+ :param model: Override the configured model for this single call.
129
+ :param max_k: Maximum number of snippets included in the context window.
130
+ :param system: Override the default RAG system prompt.
131
+ :returns: Synthesized answer, or ``None`` on failure or when no snippets
132
+ carry content.
133
+ """
134
+ items = [s for s in snippets[:max_k] if (s.get("content") or "").strip()]
135
+ if not items:
136
+ return None
137
+
138
+ ctx_parts: list[str] = []
139
+ for s in items:
140
+ genre = s.get("genre", s.get("kg_kind", ""))
141
+ author = s.get("author") or ""
142
+ title = s.get("title") or s.get("name") or ""
143
+ header = " · ".join(x for x in [genre, author, title] if x)
144
+ ctx_parts.append(f"[{header}]\n{s['content'].strip()}")
145
+
146
+ ctx = "\n\n".join(ctx_parts)
147
+ messages = [
148
+ {"role": "system", "content": system or _RAG_SYSTEM},
149
+ {"role": "user", "content": f"Source passages:\n{ctx}\n\nQuestion: {query}"},
150
+ ]
151
+ return self._complete(messages, model=model, temperature=0.3)
152
+
153
+ def rewrite_for_image(
154
+ self,
155
+ corpus_text: str,
156
+ *,
157
+ model: str | None = None,
158
+ ) -> tuple[str, str | None]:
159
+ """Rewrite a historical corpus passage into an image-generation prompt.
160
+
161
+ Uses the same backend and model as RAG synthesis so no second model load
162
+ is required. Pass ``model=`` to use a lighter model for this step if
163
+ synthesis and rewrite are called from the same hot path.
164
+
165
+ :param corpus_text: Raw historical passage to convert into a visual scene.
166
+ :param model: Override the configured model for this single call.
167
+ :returns: ``(prompt, error)`` — prompt is the rewritten scene description
168
+ (falls back to the original text on failure); error is ``None``
169
+ on success or a short message otherwise.
170
+ """
171
+ messages = [
172
+ {"role": "system", "content": _IMAGE_REWRITE_SYSTEM},
173
+ {"role": "user", "content": corpus_text},
174
+ ]
175
+ result = self._complete(messages, model=model, max_tokens=300, temperature=0.7)
176
+ if result is None:
177
+ return corpus_text, "synthesis returned no content"
178
+ return result, None
File without changes
File without changes