From 0c7e853f1f29d49798a85f13acf524f2b093c934 Mon Sep 17 00:00:00 2001 From: Somebody Else Date: Wed, 17 Jun 2026 04:58:04 +0300 Subject: [PATCH] add ctx detection for ollama and vllm --- prompt_orchestrator/llm/ollama_client.py | 46 ++++++++ prompt_orchestrator/llm/openai_client.py | 124 ++++++++++++++++++++ prompt_orchestrator/orchestrator/factory.py | 9 ++ tests/test_config_store.py | 40 +++++++ tests/test_openai_context_probe.py | 72 ++++++++++++ 5 files changed, 291 insertions(+) diff --git a/prompt_orchestrator/llm/ollama_client.py b/prompt_orchestrator/llm/ollama_client.py index 79337ee..6a61455 100644 --- a/prompt_orchestrator/llm/ollama_client.py +++ b/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 diff --git a/prompt_orchestrator/llm/openai_client.py b/prompt_orchestrator/llm/openai_client.py index f847ec4..28c9ddb 100644 --- a/prompt_orchestrator/llm/openai_client.py +++ b/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 diff --git a/prompt_orchestrator/orchestrator/factory.py b/prompt_orchestrator/orchestrator/factory.py index dd90c17..3c0ce18 100644 --- a/prompt_orchestrator/orchestrator/factory.py +++ b/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 ) diff --git a/tests/test_config_store.py b/tests/test_config_store.py index bb3e219..dbd5257 100644 --- a/tests/test_config_store.py +++ b/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 diff --git a/tests/test_openai_context_probe.py b/tests/test_openai_context_probe.py index c17ba81..bd976bb 100644 --- a/tests/test_openai_context_probe.py +++ b/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