Feat add context window detection

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"]

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

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
         )

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"

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

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