diff --git a/prompt_orchestrator/config/settings.py b/prompt_orchestrator/config/settings.py index d98c45e..39f5edc 100644 --- a/prompt_orchestrator/config/settings.py +++ b/prompt_orchestrator/config/settings.py @@ -22,6 +22,10 @@ class OrchestratorSettings(BaseModel): ), ) token_chars_ratio: float = 4.0 + openai_context_probe_enabled: bool = False + openai_context_probe_start_size: int = 20000 + openai_context_probe_step: int = 2000 + openai_context_probe_max_attempts: int = 50 section_priority: list[str] = Field( default_factory=lambda: ["rag", "recent", "summary"] diff --git a/prompt_orchestrator/llm/openai_client.py b/prompt_orchestrator/llm/openai_client.py index 412d257..f847ec4 100644 --- a/prompt_orchestrator/llm/openai_client.py +++ b/prompt_orchestrator/llm/openai_client.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Mapping + from pydantic import BaseModel from .base_client import SummaryLLMClient @@ -51,3 +53,135 @@ def generate(self, prompt: str, model: str, max_tokens: int, temperature: float) max_tokens=max_tokens, ) return chat.choices[0].message.content or "" + + +def _extract_context_window(payload: object) -> int | None: + if payload is None: + return None + + for field in ("context_window", "input_token_limit", "max_input_tokens"): + value = getattr(payload, field, None) + if isinstance(value, int) and value > 0: + return value + + if isinstance(payload, Mapping): + for field in ("context_window", "input_token_limit", "max_input_tokens"): + value = payload.get(field) + if isinstance(value, int) and value > 0: + return value + + return None + + +def discover_openai_context_window(config: OpenAIConfig, model: str) -> int | None: + client = _build_openai_client(config) + if client is None: + return None + + try: + model_payload = client.models.retrieve(model) + value = _extract_context_window(model_payload) + if value is not None: + return value + except Exception: + pass + + try: + models = client.models.list() + for item in getattr(models, "data", []): + if getattr(item, "id", None) != model: + continue + value = _extract_context_window(item) + if value is not None: + return value + except Exception: + pass + + return None + + +def discover_openai_context_window_by_probe( + config: OpenAIConfig, + model: str, + start_size: int = 20000, + step: int = 2000, + max_attempts: int = 50, +) -> int | None: + if start_size <= 0 or step <= 0 or max_attempts <= 0: + return None + + client = _build_openai_client(config) + if client is None: + return None + + attempts_left = max_attempts + + if not _probe_input_size(client=client, model=model, size=start_size): + return None + + attempts_left -= 1 + best_ok = start_size + low_ok = start_size + high_fail: int | None = None + + growth = step + while attempts_left > 0: + candidate = low_ok + growth + if _probe_input_size(client=client, model=model, size=candidate): + best_ok = candidate + low_ok = candidate + growth *= 2 + attempts_left -= 1 + continue + + high_fail = candidate + attempts_left -= 1 + break + + if high_fail is None: + return best_ok + + left = low_ok + right = high_fail - 1 + + while attempts_left > 0 and left <= right: + mid = (left + right) // 2 + if _probe_input_size(client=client, model=model, size=mid): + best_ok = mid + left = mid + 1 + attempts_left -= 1 + continue + + right = mid - 1 + attempts_left -= 1 + + return best_ok + + +def _probe_input_size(client: object, model: str, size: int) -> bool: + try: + prompt = "A" * size + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + max_tokens=1, + ) + return True + except Exception: + return False + + +def _build_openai_client(config: OpenAIConfig): + try: + from openai import OpenAI + except ImportError: + return None + + try: + return OpenAI( + api_key=config.api_key, + base_url=config.base_url, + organization=config.organization, + ) + except Exception: + return None diff --git a/prompt_orchestrator/orchestrator/factory.py b/prompt_orchestrator/orchestrator/factory.py index 5631c14..dd90c17 100644 --- a/prompt_orchestrator/orchestrator/factory.py +++ b/prompt_orchestrator/orchestrator/factory.py @@ -5,6 +5,10 @@ from ..config.config_store import ConfigStore from ..context.manager import PromptContextManager from ..llm.base_client import SummaryLLMClient +from ..llm.openai_client import ( + discover_openai_context_window, + discover_openai_context_window_by_probe, +) from ..llm.summary_llm import SummaryLLM from ..rag.base import RAGProvider from ..rag.no_rag import NoRAGProvider @@ -24,6 +28,22 @@ def from_config_store( prompt_config = config_store.get_prompt() summary_llm_config = config_store.get_summary_llm() + if summary_llm_config.provider == "openai": + discovered_window = discover_openai_context_window( + config=summary_llm_config.openai, + model=settings.token_model, + ) + if discovered_window is None and settings.openai_context_probe_enabled: + discovered_window = discover_openai_context_window_by_probe( + config=summary_llm_config.openai, + model=settings.token_model, + start_size=settings.openai_context_probe_start_size, + step=settings.openai_context_probe_step, + max_attempts=settings.openai_context_probe_max_attempts, + ) + if discovered_window is not None: + settings.max_prompt_tokens = discovered_window + cache = cache_backend or LocalTTLCacheBackend( default_ttl_seconds=settings.cache_ttl_seconds ) diff --git a/pyproject.toml b/pyproject.toml index d30bc9d..3dc31fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "prompt-orchestrator" -version = "0.1.6" +version = "0.1.7" description = "Structured prompt orchestration with cache, safety, and analyzer layers" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.10" diff --git a/tests/test_config_store.py b/tests/test_config_store.py index 38cce5f..bb3e219 100644 --- a/tests/test_config_store.py +++ b/tests/test_config_store.py @@ -91,3 +91,96 @@ def test_orchestrator_settings_accepts_legacy_safety_auto_rewrite_alias() -> Non settings = OrchestratorSettings.model_validate({"safety_auto_rewrite": False}) assert settings.security_checks_auto_rewrite is False + + +def test_factory_replaces_max_prompt_tokens_with_openai_context_window(monkeypatch) -> None: + cfg = _module_config() + cfg.summary_llm.provider = "openai" + cfg.settings.token_model = "gpt-4o-mini" + cfg.settings.max_prompt_tokens = 1000 + store = ConfigStore(cfg) + + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window", + lambda config, model: 128000, + ) + + PromptOrchestratorFactory.from_config_store( + store, + summary_llm=SummaryLLM(config=SummaryLLMConfig(provider="none")), + ) + + assert store.get_settings().max_prompt_tokens == 128000 + + +def test_factory_keeps_config_max_prompt_tokens_when_openai_context_window_missing(monkeypatch) -> None: + cfg = _module_config() + cfg.summary_llm.provider = "openai" + cfg.settings.token_model = "gpt-4o-mini" + cfg.settings.max_prompt_tokens = 1000 + store = ConfigStore(cfg) + + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window", + lambda config, model: None, + ) + + PromptOrchestratorFactory.from_config_store( + store, + summary_llm=SummaryLLM(config=SummaryLLMConfig(provider="none")), + ) + + assert store.get_settings().max_prompt_tokens == 1000 + + +def test_factory_uses_probe_fallback_when_enabled(monkeypatch) -> None: + cfg = _module_config() + cfg.summary_llm.provider = "openai" + cfg.settings.token_model = "qwen3-32b" + cfg.settings.max_prompt_tokens = 1000 + cfg.settings.openai_context_probe_enabled = True + cfg.settings.openai_context_probe_start_size = 20000 + cfg.settings.openai_context_probe_step = 2000 + cfg.settings.openai_context_probe_max_attempts = 10 + store = ConfigStore(cfg) + + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window", + lambda config, model: None, + ) + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window_by_probe", + lambda config, model, start_size, step, max_attempts: 36000, + ) + + PromptOrchestratorFactory.from_config_store( + store, + summary_llm=SummaryLLM(config=SummaryLLMConfig(provider="none")), + ) + + assert store.get_settings().max_prompt_tokens == 36000 + + +def test_factory_skips_probe_fallback_when_disabled(monkeypatch) -> None: + cfg = _module_config() + cfg.summary_llm.provider = "openai" + cfg.settings.token_model = "qwen3-32b" + cfg.settings.max_prompt_tokens = 1000 + cfg.settings.openai_context_probe_enabled = False + store = ConfigStore(cfg) + + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window", + lambda config, model: None, + ) + monkeypatch.setattr( + "prompt_orchestrator.orchestrator.factory.discover_openai_context_window_by_probe", + lambda config, model, start_size, step, max_attempts: 36000, + ) + + PromptOrchestratorFactory.from_config_store( + store, + summary_llm=SummaryLLM(config=SummaryLLMConfig(provider="none")), + ) + + assert store.get_settings().max_prompt_tokens == 1000 diff --git a/tests/test_openai_context_probe.py b/tests/test_openai_context_probe.py new file mode 100644 index 0000000..c17ba81 --- /dev/null +++ b/tests/test_openai_context_probe.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +from prompt_orchestrator.llm.openai_client import ( + OpenAIConfig, + discover_openai_context_window_by_probe, +) + + +class _FakeCompletions: + def __init__(self, limit: int) -> None: + self.limit = limit + + def create(self, model: str, messages: list[dict[str, str]], max_tokens: int) -> None: + _ = model + _ = max_tokens + content = messages[0]["content"] + if len(content) > self.limit: + raise RuntimeError("context limit exceeded") + + +class _FakeChat: + def __init__(self, limit: int) -> None: + self.completions = _FakeCompletions(limit=limit) + + +class _FakeClient: + def __init__(self, limit: int) -> None: + self.chat = _FakeChat(limit=limit) + + +def test_probe_returns_none_on_invalid_params() -> None: + result = discover_openai_context_window_by_probe( + config=OpenAIConfig(api_key="x"), + model="test-model", + start_size=0, + step=2000, + max_attempts=10, + ) + + assert result is None + + +def test_probe_returns_none_when_first_probe_fails(monkeypatch) -> None: + monkeypatch.setattr( + "prompt_orchestrator.llm.openai_client._build_openai_client", + lambda config: _FakeClient(limit=10000), + ) + + result = discover_openai_context_window_by_probe( + config=OpenAIConfig(api_key="x"), + model="test-model", + start_size=20000, + step=2000, + max_attempts=10, + ) + + assert result is None + + +def test_probe_uses_exponential_and_binary_search(monkeypatch) -> None: + monkeypatch.setattr( + "prompt_orchestrator.llm.openai_client._build_openai_client", + lambda config: _FakeClient(limit=23500), + ) + + result = discover_openai_context_window_by_probe( + config=OpenAIConfig(api_key="x"), + model="qwen3-32b", + start_size=20000, + step=2000, + max_attempts=20, + ) + + assert result == 23500 + + +def test_probe_returns_best_known_value_when_attempt_budget_is_exhausted(monkeypatch) -> None: + monkeypatch.setattr( + "prompt_orchestrator.llm.openai_client._build_openai_client", + lambda config: _FakeClient(limit=500000), + ) + + result = discover_openai_context_window_by_probe( + config=OpenAIConfig(api_key="x"), + model="qwen3-32b", + start_size=20000, + step=2000, + max_attempts=3, + ) + + assert result == 26000