add ctx detection for ollama and vllm

prompt_orchestrator/llm/ollama_client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import re
 from urllib import request
 
 from pydantic import BaseModel
@@ -40,3 +41,48 @@ def generate(self, prompt: str, model: str, max_tokens: int, temperature: float)
             body = response.read().decode("utf-8")
             parsed = json.loads(body)
             return str(parsed.get("response", "")).strip()
+
+
+def discover_ollama_context_window(config: OllamaConfig, model: str) -> int | None:
+    endpoint = f"{config.base_url.rstrip('/')}/api/show"
+    payload = {"name": model}
+    data = json.dumps(payload).encode("utf-8")
+    req = request.Request(
+        endpoint,
+        data=data,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+
+    try:
+        with request.urlopen(req, timeout=config.timeout_seconds) as response:
+            body = response.read().decode("utf-8")
+            parsed = json.loads(body)
+    except Exception:
+        return None
+
+    if not isinstance(parsed, dict):
+        return None
+
+    parameters = parsed.get("parameters")
+    if isinstance(parameters, dict):
+        value = parameters.get("num_ctx")
+        if isinstance(value, int) and value > 0:
+            return value
+        if isinstance(value, str) and value.strip().isdigit():
+            number = int(value.strip())
+            return number if number > 0 else None
+
+    if isinstance(parameters, str):
+        match = re.search(r"num_ctx\s+(\d+)", parameters)
+        if match:
+            return int(match.group(1))
+
+    num_ctx = parsed.get("num_ctx")
+    if isinstance(num_ctx, int) and num_ctx > 0:
+        return num_ctx
+    if isinstance(num_ctx, str) and num_ctx.strip().isdigit():
+        number = int(num_ctx.strip())
+        return number if number > 0 else None
+
+    return None

prompt_orchestrator/llm/openai_client.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import json
+import re
 from collections.abc import Mapping
+from urllib import request
 
 from pydantic import BaseModel
 
@@ -97,6 +100,17 @@ def discover_openai_context_window(config: OpenAIConfig, model: str) -> int | No
     except Exception:
         pass
 
+    vllm_window = _discover_vllm_context_window(config.base_url)
+    if vllm_window is not None:
+        return vllm_window
+
+    ollama_window = _discover_ollama_context_window_from_base_url(
+        base_url=config.base_url,
+        model=model,
+    )
+    if ollama_window is not None:
+        return ollama_window
+
     return None
 
 
@@ -185,3 +199,113 @@ def _build_openai_client(config: OpenAIConfig):
         )
     except Exception:
         return None
+
+
+def _discover_vllm_context_window(base_url: str | None) -> int | None:
+    for endpoint in _candidate_vllm_info_endpoints(base_url):
+        payload = _http_json(endpoint=endpoint, method="GET", payload=None, timeout=10)
+        value = _parse_positive_int(payload.get("max_model_len")) if payload else None
+        if value is not None:
+            return value
+    return None
+
+
+def _discover_ollama_context_window_from_base_url(base_url: str | None, model: str) -> int | None:
+    for endpoint in _candidate_ollama_show_endpoints(base_url):
+        payload = _http_json(
+            endpoint=endpoint,
+            method="POST",
+            payload={"name": model},
+            timeout=10,
+        )
+        if not payload:
+            continue
+
+        parameters = payload.get("parameters")
+        if isinstance(parameters, Mapping):
+            value = _parse_positive_int(parameters.get("num_ctx"))
+            if value is not None:
+                return value
+
+        if isinstance(parameters, str):
+            match = re.search(r"num_ctx\s+(\d+)", parameters)
+            if match:
+                return int(match.group(1))
+
+        value = _parse_positive_int(payload.get("num_ctx"))
+        if value is not None:
+            return value
+
+    return None
+
+
+def _candidate_vllm_info_endpoints(base_url: str | None) -> list[str]:
+    if not base_url:
+        return []
+
+    clean = base_url.rstrip("/")
+    endpoints = [f"{clean}/v1/internal/model/info"]
+    if clean.endswith("/v1"):
+        endpoints.append(f"{clean}/internal/model/info")
+
+    return _unique(endpoints)
+
+
+def _candidate_ollama_show_endpoints(base_url: str | None) -> list[str]:
+    if not base_url:
+        return []
+
+    clean = base_url.rstrip("/")
+    endpoints = [f"{clean}/api/show"]
+    if clean.endswith("/v1"):
+        endpoints.append(f"{clean[:-3].rstrip('/')}/api/show")
+
+    return _unique(endpoints)
+
+
+def _http_json(
+    endpoint: str,
+    method: str,
+    payload: dict[str, str] | None,
+    timeout: int,
+) -> dict[str, object] | None:
+    data = None
+    headers: dict[str, str] = {}
+    if payload is not None:
+        data = json.dumps(payload).encode("utf-8")
+        headers["Content-Type"] = "application/json"
+
+    req = request.Request(endpoint, data=data, headers=headers, method=method)
+    try:
+        with request.urlopen(req, timeout=timeout) as response:
+            body = response.read().decode("utf-8")
+            parsed = json.loads(body)
+            if isinstance(parsed, dict):
+                return parsed
+            return None
+    except Exception:
+        return None
+
+
+def _parse_positive_int(value: object) -> int | None:
+    if isinstance(value, int) and value > 0:
+        return value
+
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped.isdigit():
+            number = int(stripped)
+            return number if number > 0 else None
+
+    return None
+
+
+def _unique(items: list[str]) -> list[str]:
+    seen: set[str] = set()
+    result: list[str] = []
+    for item in items:
+        if item in seen:
+            continue
+        seen.add(item)
+        result.append(item)
+    return result

prompt_orchestrator/orchestrator/factory.py

@@ -5,6 +5,7 @@
 from ..config.config_store import ConfigStore
 from ..context.manager import PromptContextManager
 from ..llm.base_client import SummaryLLMClient
+from ..llm.ollama_client import discover_ollama_context_window
 from ..llm.openai_client import (
     discover_openai_context_window,
     discover_openai_context_window_by_probe,
@@ -44,6 +45,14 @@ def from_config_store(
             if discovered_window is not None:
                 settings.max_prompt_tokens = discovered_window
 
+        if summary_llm_config.provider == "ollama":
+            discovered_window = discover_ollama_context_window(
+                config=summary_llm_config.ollama,
+                model=settings.token_model,
+            )
+            if discovered_window is not None:
+                settings.max_prompt_tokens = discovered_window
+
         cache = cache_backend or LocalTTLCacheBackend(
             default_ttl_seconds=settings.cache_ttl_seconds
         )

tests/test_config_store.py

@@ -184,3 +184,43 @@ def test_factory_skips_probe_fallback_when_disabled(monkeypatch) -> None:
     )
 
     assert store.get_settings().max_prompt_tokens == 1000
+
+
+def test_factory_replaces_max_prompt_tokens_with_ollama_context_window(monkeypatch) -> None:
+    cfg = _module_config()
+    cfg.summary_llm.provider = "ollama"
+    cfg.settings.token_model = "qwen3-32b"
+    cfg.settings.max_prompt_tokens = 1000
+    store = ConfigStore(cfg)
+
+    monkeypatch.setattr(
+        "prompt_orchestrator.orchestrator.factory.discover_ollama_context_window",
+        lambda config, model: 32768,
+    )
+
+    PromptOrchestratorFactory.from_config_store(
+        store,
+        summary_llm=SummaryLLM(config=SummaryLLMConfig(provider="none")),
+    )
+
+    assert store.get_settings().max_prompt_tokens == 32768
+
+
+def test_factory_keeps_config_max_prompt_tokens_when_ollama_window_missing(monkeypatch) -> None:
+    cfg = _module_config()
+    cfg.summary_llm.provider = "ollama"
+    cfg.settings.token_model = "qwen3-32b"
+    cfg.settings.max_prompt_tokens = 1000
+    store = ConfigStore(cfg)
+
+    monkeypatch.setattr(
+        "prompt_orchestrator.orchestrator.factory.discover_ollama_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

tests/test_openai_context_probe.py

@@ -2,6 +2,7 @@
 
 from prompt_orchestrator.llm.openai_client import (
     OpenAIConfig,
+    discover_openai_context_window,
     discover_openai_context_window_by_probe,
 )
 
@@ -28,6 +29,28 @@ def __init__(self, limit: int) -> None:
         self.chat = _FakeChat(limit=limit)
 
 
+class _FakeModelMeta:
+    def __init__(self, model_id: str) -> None:
+        self.id = model_id
+
+
+class _FakeModelsApi:
+    def retrieve(self, model: str):
+        _ = model
+        return _FakeModelMeta(model_id="without-window")
+
+    def list(self):
+        class _List:
+            data = [_FakeModelMeta(model_id="qwen3-32b")]
+
+        return _List()
+
+
+class _FakeOpenAIClientNoWindow:
+    def __init__(self) -> None:
+        self.models = _FakeModelsApi()
+
+
 def test_probe_returns_none_on_invalid_params() -> None:
     result = discover_openai_context_window_by_probe(
         config=OpenAIConfig(api_key="x"),
@@ -89,3 +112,52 @@ def test_probe_returns_best_known_value_when_attempt_budget_is_exhausted(monkeyp
     )
 
     assert result == 26000
+
+
+def test_discover_openai_context_window_falls_back_to_vllm_endpoint(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "prompt_orchestrator.llm.openai_client._build_openai_client",
+        lambda config: _FakeOpenAIClientNoWindow(),
+    )
+
+    def _fake_http_json(endpoint, method, payload, timeout):
+        _ = method
+        _ = payload
+        _ = timeout
+        if endpoint.endswith("/v1/internal/model/info"):
+            return {"max_model_len": 32768}
+        return None
+
+    monkeypatch.setattr("prompt_orchestrator.llm.openai_client._http_json", _fake_http_json)
+
+    result = discover_openai_context_window(
+        config=OpenAIConfig(api_key="x", base_url="http://localhost:8000/v1"),
+        model="qwen3-32b",
+    )
+
+    assert result == 32768
+
+
+def test_discover_openai_context_window_falls_back_to_ollama_endpoint(monkeypatch) -> None:
+    monkeypatch.setattr(
+        "prompt_orchestrator.llm.openai_client._build_openai_client",
+        lambda config: _FakeOpenAIClientNoWindow(),
+    )
+
+    def _fake_http_json(endpoint, method, payload, timeout):
+        _ = method
+        _ = timeout
+        if endpoint.endswith("/v1/internal/model/info"):
+            return None
+        if endpoint.endswith("/api/show") and isinstance(payload, dict) and payload.get("name") == "qwen3-32b":
+            return {"parameters": {"num_ctx": 65536}}
+        return None
+
+    monkeypatch.setattr("prompt_orchestrator.llm.openai_client._http_json", _fake_http_json)
+
+    result = discover_openai_context_window(
+        config=OpenAIConfig(api_key="x", base_url="http://localhost:11434/v1"),
+        model="qwen3-32b",
+    )
+
+    assert result == 65536