diff --git a/apps/settings/workspace.html b/apps/settings/workspace.html index cb7185671..cdcebde44 100644 --- a/apps/settings/workspace.html +++ b/apps/settings/workspace.html @@ -3786,6 +3786,14 @@ function updateTypeProviderKeyWarning(type, provider, apiKeys, auth) { if (!warning || !link) return; const message = warning.querySelectorAll('span')[1]; + // Providers with no env_key (e.g., Ollama local) don't need API keys + const providerMeta = providersData?.providers?.find(p => p.name === provider); + if (providerMeta && !providerMeta.env_key) { + warning.style.display = 'none'; + warning.title = ''; + return; + } + // For cogitate, hide key warning when using platform auth const authMode = (auth && type === 'cogitate') ? (auth[provider] || 'platform') : 'api_key'; const validation = provider ? keyValidationData[provider] : null; diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index bea8fcaae..e567e85a7 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -315,6 +315,34 @@ client = OpenAI( This allows reusing much of the OpenAI provider's patterns for request/response handling. +The Ollama provider (`think/providers/ollama.py`) takes a different approach — +it uses Ollama's native ``/api/chat`` endpoint directly via ``httpx`` for +reliable thinking control. See the Ollama section below. + +## Ollama (Local) Provider + +The ``ollama`` provider connects to a local Ollama instance via the native +``/api/chat`` endpoint (not the OpenAI-compatible endpoint, which silently +ignores the ``think`` parameter on models like Qwen3.5). Key differences +from cloud providers: + +- **No API key required.** ``validate_key()`` checks Ollama reachability + instead of key validity. +- **Model prefix convention:** Models use the ``ollama-local/`` prefix + (e.g., ``ollama-local/qwen3.5:9b``). The prefix is stripped before + sending requests to the Ollama API. +- **Thinking support:** Controlled via Ollama's ``think`` parameter, + mapped from ``thinking_budget``. Budget > 0 enables thinking; + None or 0 disables it. +- **Cogitate via OpenCode CLI.** ``run_cogitate()`` uses the OpenCode CLI + (``opencode run --format json``) as a subprocess, following the same + CLIRunner pattern as the other providers. Requires OpenCode CLI installed + and configured with a user-level ``.opencode/opencode.json`` that registers + the local Ollama instance as a provider. Do not place this config in the + project root — it belongs in the user's config directory. +- **Base URL:** Reads ``OLLAMA_BASE_URL`` env var, defaults to + ``http://localhost:11434``. + ## Checklist for New Providers **Core implementation:** diff --git a/docs/THINK.md b/docs/THINK.md index 4f6e2ed48..ae5effbb4 100644 --- a/docs/THINK.md +++ b/docs/THINK.md @@ -160,7 +160,7 @@ The `sol agents` command is primarily used internally by Cortex. For testing pur sol agents [TASK_FILE] [--provider PROVIDER] [--model MODEL] [--max-tokens N] [-o OUT_FILE] ``` -The provider can be ``openai`` (default), ``google`` or ``anthropic``. Configure the corresponding API key in the ``env`` section of ``journal/config/journal.json`` (e.g., ``OPENAI_API_KEY``, ``GOOGLE_API_KEY``, or ``ANTHROPIC_API_KEY``). Keys are loaded into ``os.environ`` by ``setup_cli()`` at process startup. +The provider can be ``openai`` (default), ``google``, ``anthropic``, or ``ollama``. Configure the corresponding API key in the ``env`` section of ``journal/config/journal.json`` (e.g., ``OPENAI_API_KEY``, ``GOOGLE_API_KEY``, or ``ANTHROPIC_API_KEY``). The ``ollama`` provider requires no API key — it connects to a local Ollama instance. Keys are loaded into ``os.environ`` by ``setup_cli()`` at process startup. ### Provider modules @@ -228,7 +228,7 @@ Cortex (orchestrator) ├── Tool execution via `sol call` └── Agent subprocess management ↓ - Providers (openai, google, anthropic) + Providers (openai, google, anthropic, ollama) ``` ## Providers @@ -238,6 +238,7 @@ Cortex (orchestrator) | OpenAI | `think/providers/openai.py` | GPT models via Agents SDK | | Google | `think/providers/google.py` | Gemini models | | Anthropic | `think/providers/anthropic.py` | Claude via Anthropic SDK | +| Ollama | `think/providers/ollama.py` | Local models via Ollama | Providers implement `run_generate()`, `run_agenerate()`, and `run_cogitate()` functions. See [PROVIDERS.md](PROVIDERS.md) for implementation details. diff --git a/pyproject.toml b/pyproject.toml index 509edfbff..66f1e8700 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ dependencies = [ "openai>=1.2.0", "openai-agents>=0.1.0", "anthropic", + "httpx", "genai-prices", "pypdf", "pdf2image", diff --git a/tests/baselines/api/settings/config.json b/tests/baselines/api/settings/config.json index 3c1e68da5..9c66f04e8 100644 --- a/tests/baselines/api/settings/config.json +++ b/tests/baselines/api/settings/config.json @@ -55,6 +55,10 @@ "provider": "google", "tier": 2 }, + "test.ollama": { + "model": "ollama-local/qwen3.5:2b", + "provider": "ollama" + }, "test.openai": { "model": "gpt-5-mini", "provider": "openai" diff --git a/tests/baselines/api/settings/providers.json b/tests/baselines/api/settings/providers.json index 60bbf0bb5..b393ba403 100644 --- a/tests/baselines/api/settings/providers.json +++ b/tests/baselines/api/settings/providers.json @@ -2,11 +2,13 @@ "api_keys": { "anthropic": false, "google": false, + "ollama": false, "openai": false }, "auth": { "anthropic": "platform", "google": "platform", + "ollama": "platform", "openai": "platform" }, "cogitate": { @@ -408,6 +410,10 @@ "provider": "google", "tier": 2 }, + "test.ollama": { + "model": "ollama-local/qwen3.5:2b", + "provider": "ollama" + }, "test.openai": { "model": "gpt-5-mini", "provider": "openai" @@ -432,6 +438,11 @@ "google_backend": "auto", "key_validation": {}, "providers": [ + { + "env_key": "", + "label": "Ollama (Local)", + "name": "ollama" + }, { "env_key": "ANTHROPIC_API_KEY", "label": "Anthropic (Claude)", diff --git a/tests/fixtures/journal/config/journal.json b/tests/fixtures/journal/config/journal.json index ab235c418..dd27340dd 100644 --- a/tests/fixtures/journal/config/journal.json +++ b/tests/fixtures/journal/config/journal.json @@ -52,7 +52,8 @@ "test.tier.inherit": {"tier": 3}, "test.tier.override": {"provider": "openai", "tier": 2}, "test.config.override": {"provider": "google", "tier": 2}, - "observe.*": {"provider": "google", "tier": 3} + "observe.*": {"provider": "google", "tier": 3}, + "test.ollama": {"provider": "ollama", "model": "ollama-local/qwen3.5:2b"} }, "models": { "google": { diff --git a/tests/integration/test_ollama_provider.py b/tests/integration/test_ollama_provider.py new file mode 100644 index 000000000..74caabed8 --- /dev/null +++ b/tests/integration/test_ollama_provider.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +"""Integration tests for Ollama provider with a live local Ollama instance.""" + +import asyncio +import shutil + +import pytest + +from think.models import OLLAMA_LITE, OLLAMA_PRO + +# Use the smallest available model for fast integration tests +_TEST_MODEL = OLLAMA_LITE + + +def _ollama_reachable() -> bool: + """Check if the local Ollama instance is reachable.""" + try: + from think.providers.ollama import validate_key + + return validate_key("")["valid"] + except Exception: + return False + + +# Skip all tests in this module if Ollama is not running +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + not _ollama_reachable(), + reason="Local Ollama instance not reachable", + ), +] + + +class TestOllamaGenerate: + def test_basic_generation(self): + from think.providers.ollama import run_generate + + result = run_generate( + "What is 2 + 2? Reply with just the number.", + model=_TEST_MODEL, + max_output_tokens=64, + thinking_budget=0, + ) + + assert result["text"] + assert "4" in result["text"] + assert result["usage"] is not None + assert result["usage"]["input_tokens"] > 0 + assert result["usage"]["output_tokens"] > 0 + assert result["finish_reason"] == "stop" + # With think=False via native API, thinking should be absent + assert result["thinking"] is None + + def test_system_instruction(self): + from think.providers.ollama import run_generate + + result = run_generate( + "What color is the sky?", + model=_TEST_MODEL, + max_output_tokens=64, + system_instruction="Always respond in exactly one word.", + thinking_budget=0, + ) + + assert result["text"] + + def test_json_output(self): + from think.providers.ollama import run_generate + + result = run_generate( + 'Return a JSON object with key "answer" and value 42.', + model=_TEST_MODEL, + max_output_tokens=256, + json_output=True, + thinking_budget=0, + ) + + # The response should contain JSON content. Small models may wrap + # it in markdown fences, so we check for the key rather than strict + # parsing. (JSON validation is handled centrally by think/models.py.) + assert result["text"] + assert "answer" in result["text"] + + def test_thinking_enabled(self): + from think.providers.ollama import run_generate + + result = run_generate( + "What is 15 * 17?", + model=_TEST_MODEL, + max_output_tokens=512, + thinking_budget=4096, + ) + + assert result["text"] + assert result["usage"] is not None + # With think=True, thinking content should be present + assert result["thinking"] is not None + assert len(result["thinking"]) > 0 + assert result["thinking"][0]["summary"] + + def test_thinking_disabled_no_reasoning(self): + """Verify that think=False actually suppresses reasoning on the native API.""" + from think.providers.ollama import run_generate + + result = run_generate( + "What is 2 + 2? Reply with just the number.", + model=_TEST_MODEL, + max_output_tokens=64, + thinking_budget=0, + ) + + assert result["text"] + assert result["thinking"] is None + + +class TestOllamaAgenerate: + def test_async_generation(self): + from think.providers.ollama import run_agenerate + + result = asyncio.run( + run_agenerate( + "What is 3 + 5? Reply with just the number.", + model=_TEST_MODEL, + max_output_tokens=64, + thinking_budget=0, + ) + ) + + assert result["text"] + assert "8" in result["text"] + assert result["usage"] is not None + + +class TestOllamaListModels: + def test_list_models(self): + from think.providers.ollama import list_models + + models = list_models() + assert isinstance(models, list) + assert len(models) > 0 + # Native API returns models with "name" field + assert "name" in models[0] + + +class TestOllamaValidateKey: + def test_reachable(self): + from think.providers.ollama import validate_key + + result = validate_key("") + assert result["valid"] is True + + +def _opencode_available() -> bool: + """Check if the OpenCode CLI is installed.""" + return shutil.which("opencode") is not None + + +@pytest.mark.skipif( + not _opencode_available(), + reason="OpenCode CLI not installed", +) +class TestOllamaCogitate: + def test_basic_cogitate(self): + """Test cogitate with a simple prompt that doesn't require tool use.""" + from think.providers.ollama import run_cogitate + + events = [] + result = asyncio.run( + run_cogitate( + { + "prompt": "Say the word 'solstone' and nothing else.", + "model": OLLAMA_PRO, + }, + on_event=lambda e: events.append(e), + ) + ) + + assert result + + # Should have emitted a finish event + finish_events = [e for e in events if e.get("event") == "finish"] + assert len(finish_events) == 1 + assert finish_events[0]["result"] == result + + def test_cogitate_with_tool_use(self): + """Test cogitate with a prompt that triggers bash tool use.""" + from think.providers.ollama import run_cogitate + + events = [] + result = asyncio.run( + run_cogitate( + { + "prompt": "Use the bash tool to run 'echo solstone_test_ok' and tell me exactly what it output.", + "model": OLLAMA_PRO, + }, + on_event=lambda e: events.append(e), + ) + ) + + assert result + assert "solstone_test_ok" in result + + # Should have tool_start and tool_end events + tool_starts = [e for e in events if e.get("event") == "tool_start"] + tool_ends = [e for e in events if e.get("event") == "tool_end"] + assert len(tool_starts) >= 1 + assert len(tool_ends) >= 1 diff --git a/tests/test_ollama.py b/tests/test_ollama.py new file mode 100644 index 000000000..77d9af51b --- /dev/null +++ b/tests/test_ollama.py @@ -0,0 +1,956 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +"""Unit tests for the Ollama (Local) provider.""" + +import asyncio +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from think.models import OLLAMA_FLASH, OLLAMA_LITE, OLLAMA_PRO + + +def _ollama_provider(): + import importlib + + return importlib.reload(importlib.import_module("think.providers.ollama")) + + +# --------------------------------------------------------------------------- +# _strip_model_prefix +# --------------------------------------------------------------------------- + + +class TestStripModelPrefix: + def test_strips_ollama_local_prefix(self): + provider = _ollama_provider() + assert provider._strip_model_prefix("ollama-local/qwen3.5:9b") == "qwen3.5:9b" + + def test_strips_prefix_from_complex_name(self): + provider = _ollama_provider() + assert ( + provider._strip_model_prefix("ollama-local/qwen3.5:35b-a3b-bf16") + == "qwen3.5:35b-a3b-bf16" + ) + + def test_no_prefix_passthrough(self): + provider = _ollama_provider() + assert provider._strip_model_prefix("llama3.1:8b") == "llama3.1:8b" + + +# --------------------------------------------------------------------------- +# _build_messages +# --------------------------------------------------------------------------- + + +class TestBuildMessages: + def test_string_contents(self): + provider = _ollama_provider() + msgs = provider._build_messages("hello") + assert msgs == [{"role": "user", "content": "hello"}] + + def test_string_with_system(self): + provider = _ollama_provider() + msgs = provider._build_messages("hello", system_instruction="be helpful") + assert msgs == [ + {"role": "system", "content": "be helpful"}, + {"role": "user", "content": "hello"}, + ] + + def test_list_of_strings(self): + provider = _ollama_provider() + msgs = provider._build_messages(["line1", "line2"]) + assert msgs == [{"role": "user", "content": "line1\nline2"}] + + def test_list_of_dicts_passthrough(self): + provider = _ollama_provider() + input_msgs = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + msgs = provider._build_messages(input_msgs) + assert msgs == input_msgs + + def test_list_of_dicts_with_system(self): + provider = _ollama_provider() + input_msgs = [{"role": "user", "content": "hi"}] + msgs = provider._build_messages(input_msgs, system_instruction="sys") + assert msgs == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + + def test_non_string_contents(self): + provider = _ollama_provider() + msgs = provider._build_messages(42) + assert msgs == [{"role": "user", "content": "42"}] + + +# --------------------------------------------------------------------------- +# _build_request_body +# --------------------------------------------------------------------------- + + +class TestBuildRequestBody: + def test_basic_body(self): + provider = _ollama_provider() + body = provider._build_request_body( + model="qwen3.5:9b", + messages=[{"role": "user", "content": "hi"}], + temperature=0.3, + max_output_tokens=1024, + json_output=False, + thinking_budget=None, + ) + assert body["model"] == "qwen3.5:9b" + assert body["stream"] is False + assert body["options"]["temperature"] == 0.3 + assert body["options"]["num_predict"] == 1024 + assert body["think"] is False + + def test_thinking_enabled(self): + provider = _ollama_provider() + body = provider._build_request_body( + "m", + [{"role": "user", "content": "hi"}], + 0.3, + 1024, + False, + thinking_budget=4096, + ) + assert body["think"] is True + + def test_thinking_disabled_zero(self): + provider = _ollama_provider() + body = provider._build_request_body( + "m", + [{"role": "user", "content": "hi"}], + 0.3, + 1024, + False, + thinking_budget=0, + ) + assert body["think"] is False + + def test_json_output(self): + provider = _ollama_provider() + body = provider._build_request_body( + "m", [{"role": "user", "content": "hi"}], 0.3, 1024, True, None + ) + assert body["format"] == "json" + + def test_no_json_output(self): + provider = _ollama_provider() + body = provider._build_request_body( + "m", [{"role": "user", "content": "hi"}], 0.3, 1024, False, None + ) + assert "format" not in body + + +# --------------------------------------------------------------------------- +# _normalize_finish_reason +# --------------------------------------------------------------------------- + + +class TestNormalizeFinishReason: + def test_stop(self): + provider = _ollama_provider() + assert ( + provider._normalize_finish_reason({"done": True, "done_reason": "stop"}) + == "stop" + ) + + def test_length_to_max_tokens(self): + provider = _ollama_provider() + assert ( + provider._normalize_finish_reason({"done": True, "done_reason": "length"}) + == "max_tokens" + ) + + def test_done_no_reason(self): + provider = _ollama_provider() + assert provider._normalize_finish_reason({"done": True}) == "stop" + + def test_not_done(self): + provider = _ollama_provider() + assert provider._normalize_finish_reason({"done": False}) is None + + def test_unknown_passthrough(self): + provider = _ollama_provider() + assert ( + provider._normalize_finish_reason({"done": True, "done_reason": "other"}) + == "other" + ) + + +# --------------------------------------------------------------------------- +# _extract_usage +# --------------------------------------------------------------------------- + + +class TestExtractUsage: + def test_normal_usage(self): + provider = _ollama_provider() + result = provider._extract_usage( + { + "prompt_eval_count": 100, + "eval_count": 50, + } + ) + assert result == { + "input_tokens": 100, + "output_tokens": 50, + "total_tokens": 150, + } + + def test_missing_fields(self): + provider = _ollama_provider() + result = provider._extract_usage({}) + assert result == { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + } + + +# --------------------------------------------------------------------------- +# _extract_thinking +# --------------------------------------------------------------------------- + + +class TestExtractThinking: + def test_with_thinking(self): + provider = _ollama_provider() + result = provider._extract_thinking( + {"message": {"content": "4", "thinking": "Let me calculate..."}} + ) + assert result == [{"summary": "Let me calculate..."}] + + def test_no_thinking(self): + provider = _ollama_provider() + result = provider._extract_thinking({"message": {"content": "4"}}) + assert result is None + + def test_empty_thinking(self): + provider = _ollama_provider() + result = provider._extract_thinking( + {"message": {"content": "4", "thinking": " "}} + ) + assert result is None + + def test_no_message(self): + provider = _ollama_provider() + result = provider._extract_thinking({}) + assert result is None + + +# --------------------------------------------------------------------------- +# _parse_response +# --------------------------------------------------------------------------- + + +class TestParseResponse: + def test_full_response(self): + provider = _ollama_provider() + data = { + "message": { + "role": "assistant", + "content": "Hello!", + "thinking": "Reasoning...", + }, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 25, + "eval_count": 10, + } + result = provider._parse_response(data) + assert result["text"] == "Hello!" + assert result["finish_reason"] == "stop" + assert result["usage"]["input_tokens"] == 25 + assert result["usage"]["output_tokens"] == 10 + assert result["thinking"] == [{"summary": "Reasoning..."}] + + def test_no_thinking(self): + provider = _ollama_provider() + data = { + "message": {"role": "assistant", "content": "4"}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 10, + "eval_count": 2, + } + result = provider._parse_response(data) + assert result["text"] == "4" + assert result["thinking"] is None + + +# --------------------------------------------------------------------------- +# run_generate +# --------------------------------------------------------------------------- + + +def _make_ollama_response( + content="Hello!", + thinking=None, + done=True, + done_reason="stop", + prompt_eval_count=10, + eval_count=5, +): + """Build a mock native Ollama /api/chat response dict.""" + message = {"role": "assistant", "content": content} + if thinking is not None: + message["thinking"] = thinking + return { + "model": "qwen3.5:9b", + "message": message, + "done": done, + "done_reason": done_reason, + "prompt_eval_count": prompt_eval_count, + "eval_count": eval_count, + } + + +class TestRunGenerate: + def test_basic_generation(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + result = provider.run_generate("hello", model=OLLAMA_FLASH) + + assert result["text"] == "Hello!" + assert result["finish_reason"] == "stop" + assert result["usage"]["input_tokens"] == 10 + + def test_model_prefix_stripped(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate("hello", model="ollama-local/qwen3.5:9b") + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + assert body["model"] == "qwen3.5:9b" + + def test_thinking_enabled(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response(thinking="Reasoning...") + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + result = provider.run_generate( + "hello", model=OLLAMA_FLASH, thinking_budget=4096 + ) + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + assert body["think"] is True + assert result["thinking"] == [{"summary": "Reasoning..."}] + + def test_thinking_disabled_when_none(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate("hello", model=OLLAMA_FLASH, thinking_budget=None) + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + assert body["think"] is False + + def test_thinking_disabled_when_zero(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate("hello", model=OLLAMA_FLASH, thinking_budget=0) + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + assert body["think"] is False + + def test_json_output(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response( + content='{"key": "value"}' + ) + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate("hello", model=OLLAMA_FLASH, json_output=True) + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + assert body["format"] == "json" + + def test_system_instruction(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate( + "hello", model=OLLAMA_FLASH, system_instruction="be concise" + ) + + call_kwargs = mock_client.post.call_args + body = call_kwargs.kwargs["json"] + messages = body["messages"] + assert messages[0] == {"role": "system", "content": "be concise"} + assert messages[1] == {"role": "user", "content": "hello"} + + def test_timeout(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.post.return_value = mock_response + mock_get.return_value = mock_client + + provider.run_generate("hello", model=OLLAMA_FLASH, timeout_s=30.0) + + call_kwargs = mock_client.post.call_args + assert call_kwargs.kwargs["timeout"] == 30.0 + + +# --------------------------------------------------------------------------- +# run_agenerate +# --------------------------------------------------------------------------- + + +class TestRunAgenerate: + def test_async_generation(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = _make_ollama_response() + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_async_client") as mock_get: + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_get.return_value = mock_client + + result = asyncio.run(provider.run_agenerate("hello", model=OLLAMA_FLASH)) + + assert result["text"] == "Hello!" + assert result["finish_reason"] == "stop" + + +# --------------------------------------------------------------------------- +# _translate_opencode +# --------------------------------------------------------------------------- + + +def _make_test_harness(): + """Create a callback/aggregator pair for testing _translate_opencode.""" + from think.providers.cli import ThinkingAggregator + from think.providers.shared import JSONEventCallback + + events = [] + cb = JSONEventCallback(lambda e: events.append(e)) + aggregator = ThinkingAggregator(cb, "qwen3.5:9b") + return events, cb, aggregator + + +class TestTranslateOpencode: + def test_step_start_returns_session_id(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + event = { + "type": "step_start", + "sessionID": "ses_abc123", + "part": {"type": "step-start"}, + } + usage = {} + + result = provider._translate_opencode(event, aggregator, cb, usage) + + assert result == "ses_abc123" + assert events == [] + + def test_text_accumulates(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + event = { + "type": "text", + "part": {"type": "text", "text": "Hello world"}, + } + usage = {} + + result = provider._translate_opencode(event, aggregator, cb, usage) + + assert result is None + assert aggregator.has_content + assert aggregator.flush_as_result() == "Hello world" + + def test_tool_use_emits_start_and_end(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + event = { + "type": "tool_use", + "part": { + "type": "tool", + "tool": "bash", + "callID": "call_xyz", + "state": { + "status": "completed", + "input": {"command": "echo hello"}, + "output": "hello\n", + }, + }, + } + usage = {} + + result = provider._translate_opencode(event, aggregator, cb, usage) + + assert result is None + assert len(events) == 2 + assert events[0]["event"] == "tool_start" + assert events[0]["tool"] == "bash" + assert events[0]["args"] == {"command": "echo hello"} + assert events[0]["call_id"] == "call_xyz" + assert events[1]["event"] == "tool_end" + assert events[1]["tool"] == "bash" + assert events[1]["result"] == "hello\n" + assert events[1]["call_id"] == "call_xyz" + + def test_tool_use_flushes_thinking(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + usage = {} + + # Accumulate some text first + aggregator.accumulate("Let me run a command...") + + # Then a tool use event + event = { + "type": "tool_use", + "part": { + "type": "tool", + "tool": "bash", + "callID": "call_1", + "state": { + "status": "completed", + "input": {"command": "ls"}, + "output": "file.txt\n", + }, + }, + } + provider._translate_opencode(event, aggregator, cb, usage) + + # First event should be thinking (flushed), then tool_start, tool_end + assert len(events) == 3 + assert events[0]["event"] == "thinking" + assert events[0]["summary"] == "Let me run a command..." + assert events[1]["event"] == "tool_start" + assert events[2]["event"] == "tool_end" + + def test_step_finish_captures_usage(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + event = { + "type": "step_finish", + "part": { + "type": "step-finish", + "reason": "stop", + "tokens": { + "total": 14681, + "input": 14649, + "output": 32, + "reasoning": 0, + "cache": {"write": 0, "read": 0}, + }, + }, + } + usage = {} + + result = provider._translate_opencode(event, aggregator, cb, usage) + + assert result is None + assert usage["input_tokens"] == 14649 + assert usage["output_tokens"] == 32 + assert usage["total_tokens"] == 14681 + + def test_step_finish_accumulates_usage_across_steps(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + usage = {} + + # First step + event1 = { + "type": "step_finish", + "part": { + "type": "step-finish", + "reason": "tool-calls", + "tokens": {"total": 100, "input": 80, "output": 20, "reasoning": 0}, + }, + } + provider._translate_opencode(event1, aggregator, cb, usage) + + # Second step + event2 = { + "type": "step_finish", + "part": { + "type": "step-finish", + "reason": "stop", + "tokens": {"total": 150, "input": 120, "output": 30, "reasoning": 0}, + }, + } + provider._translate_opencode(event2, aggregator, cb, usage) + + assert usage["input_tokens"] == 200 + assert usage["output_tokens"] == 50 + assert usage["total_tokens"] == 250 + + def test_unknown_event_ignored(self): + provider = _ollama_provider() + events, cb, aggregator = _make_test_harness() + event = {"type": "unknown_type", "part": {}} + usage = {} + + result = provider._translate_opencode(event, aggregator, cb, usage) + + assert result is None + assert events == [] + + +# --------------------------------------------------------------------------- +# run_cogitate +# --------------------------------------------------------------------------- + + +class TestRunCogitate: + def test_basic_cogitate(self): + provider = _ollama_provider() + + class MockCLIRunner: + last_instance = None + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.cmd = kwargs["cmd"] + self.prompt_text = kwargs["prompt_text"] + self.cli_session_id = "ses_test123" + self.run = AsyncMock(return_value="test result") + MockCLIRunner.last_instance = self + + with patch("shutil.which", return_value="/usr/bin/opencode"), \ + patch("think.providers.ollama.CLIRunner", MockCLIRunner): + events = [] + asyncio.run( + provider.run_cogitate( + {"prompt": "hello", "model": OLLAMA_FLASH}, + lambda e: events.append(e), + ) + ) + + instance = MockCLIRunner.last_instance + assert "opencode" in instance.cmd + assert "--format" in instance.cmd + assert "json" in instance.cmd + assert "-m" in instance.cmd + m_idx = instance.cmd.index("-m") + assert instance.cmd[m_idx + 1] == "ollama/qwen3.5:9b" + + # Should emit finish event + finish_events = [e for e in events if e.get("event") == "finish"] + assert len(finish_events) == 1 + assert finish_events[0]["result"] == "test result" + assert finish_events[0]["cli_session_id"] == "ses_test123" + + def test_cogitate_strips_model_prefix(self): + provider = _ollama_provider() + + class MockCLIRunner: + last_instance = None + + def __init__(self, **kwargs): + self.cmd = kwargs["cmd"] + self.prompt_text = kwargs["prompt_text"] + self.cli_session_id = None + self.run = AsyncMock(return_value="ok") + MockCLIRunner.last_instance = self + + with patch("shutil.which", return_value="/usr/bin/opencode"), \ + patch("think.providers.ollama.CLIRunner", MockCLIRunner): + asyncio.run( + provider.run_cogitate( + {"prompt": "test", "model": "ollama-local/qwen3.5:35b-a3b-bf16"}, + lambda e: None, + ) + ) + + cmd = MockCLIRunner.last_instance.cmd + m_idx = cmd.index("-m") + assert cmd[m_idx + 1] == "ollama/qwen3.5:35b-a3b-bf16" + + def test_cogitate_session_resume(self): + provider = _ollama_provider() + + class MockCLIRunner: + last_instance = None + + def __init__(self, **kwargs): + self.cmd = kwargs["cmd"] + self.prompt_text = kwargs["prompt_text"] + self.cli_session_id = None + self.run = AsyncMock(return_value="ok") + MockCLIRunner.last_instance = self + + with patch("shutil.which", return_value="/usr/bin/opencode"), \ + patch("think.providers.ollama.CLIRunner", MockCLIRunner): + asyncio.run( + provider.run_cogitate( + { + "prompt": "continue", + "model": OLLAMA_FLASH, + "session_id": "ses_previous", + }, + lambda e: None, + ) + ) + + cmd = MockCLIRunner.last_instance.cmd + assert "--session" in cmd + s_idx = cmd.index("--session") + assert cmd[s_idx + 1] == "ses_previous" + + def test_cogitate_prepends_system_instruction(self): + provider = _ollama_provider() + + class MockCLIRunner: + last_instance = None + + def __init__(self, **kwargs): + self.prompt_text = kwargs["prompt_text"] + self.cmd = kwargs["cmd"] + self.cli_session_id = None + self.run = AsyncMock(return_value="ok") + MockCLIRunner.last_instance = self + + with patch("shutil.which", return_value="/usr/bin/opencode"), \ + patch("think.providers.ollama.CLIRunner", MockCLIRunner): + asyncio.run( + provider.run_cogitate( + { + "prompt": "user prompt", + "system_instruction": "be helpful", + "model": OLLAMA_FLASH, + }, + lambda e: None, + ) + ) + + prompt = MockCLIRunner.last_instance.prompt_text + assert prompt.startswith("be helpful") + assert "user prompt" in prompt + + def test_cogitate_emits_error_on_failure(self): + provider = _ollama_provider() + + class MockCLIRunner: + def __init__(self, **kwargs): + self.cmd = kwargs["cmd"] + self.prompt_text = kwargs["prompt_text"] + self.cli_session_id = None + self.run = AsyncMock(side_effect=RuntimeError("CLI not found")) + + events = [] + with patch("shutil.which", return_value="/usr/bin/opencode"), \ + patch("think.providers.ollama.CLIRunner", MockCLIRunner): + with pytest.raises(RuntimeError, match="CLI not found"): + asyncio.run( + provider.run_cogitate( + {"prompt": "test", "model": OLLAMA_FLASH}, + lambda e: events.append(e), + ) + ) + + error_events = [e for e in events if e.get("event") == "error"] + assert len(error_events) == 1 + assert "CLI not found" in error_events[0]["error"] + + def test_cogitate_raises_when_opencode_not_installed(self): + provider = _ollama_provider() + + events = [] + with patch("shutil.which", return_value=None): + with pytest.raises(RuntimeError, match="Cogitate requires OpenCode CLI"): + asyncio.run( + provider.run_cogitate( + {"prompt": "test", "model": OLLAMA_FLASH}, + lambda e: events.append(e), + ) + ) + + error_events = [e for e in events if e.get("event") == "error"] + assert len(error_events) == 1 + assert "OpenCode CLI" in error_events[0]["error"] + + +# --------------------------------------------------------------------------- +# _build_opencode_env +# --------------------------------------------------------------------------- + + +class TestBuildOpencodeEnv: + def test_sets_api_key_placeholder(self): + provider = _ollama_provider() + with patch.dict(os.environ, {}, clear=True): + env = provider._build_opencode_env() + assert env.get("OPENAI_API_KEY") == "ollama" + + def test_preserves_existing_api_key(self): + provider = _ollama_provider() + with patch.dict(os.environ, {"OPENAI_API_KEY": "real-key"}, clear=False): + env = provider._build_opencode_env() + assert env["OPENAI_API_KEY"] == "real-key" + + +# --------------------------------------------------------------------------- +# list_models / validate_key +# --------------------------------------------------------------------------- + + +class TestListModels: + def test_returns_model_list(self): + provider = _ollama_provider() + mock_response = MagicMock() + mock_response.json.return_value = { + "models": [ + {"name": "qwen3.5:9b", "size": 6600000000}, + {"name": "llama3.1:8b", "size": 4900000000}, + ] + } + mock_response.raise_for_status = MagicMock() + + with patch.object(provider, "_get_client") as mock_get: + mock_client = MagicMock() + mock_client.get.return_value = mock_response + mock_get.return_value = mock_client + + result = provider.list_models() + + assert len(result) == 2 + assert result[0]["name"] == "qwen3.5:9b" + + +class TestValidateKey: + def test_reachable(self): + provider = _ollama_provider() + + with patch("httpx.get") as mock_get: + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = {"version": "0.18.3"} + mock_get.return_value = mock_response + + result = provider.validate_key("ignored") + + assert result == {"valid": True} + + def test_unreachable(self): + provider = _ollama_provider() + + with patch("httpx.get") as mock_get: + mock_get.side_effect = httpx.ConnectError("Connection refused") + + result = provider.validate_key("ignored") + + assert result["valid"] is False + assert "Connection refused" in result["error"] + + +# --------------------------------------------------------------------------- +# Model constants +# --------------------------------------------------------------------------- + +import httpx + + +class TestModelConstants: + def test_default_models_have_prefix(self): + assert OLLAMA_PRO.startswith("ollama-local/") + assert OLLAMA_FLASH.startswith("ollama-local/") + assert OLLAMA_LITE.startswith("ollama-local/") + + def test_get_model_provider(self): + from think.models import get_model_provider + + assert get_model_provider(OLLAMA_PRO) == "ollama" + assert get_model_provider(OLLAMA_FLASH) == "ollama" + assert get_model_provider(OLLAMA_LITE) == "ollama" + + def test_provider_defaults_exist(self): + from think.models import PROVIDER_DEFAULTS + + assert "ollama" in PROVIDER_DEFAULTS + assert 1 in PROVIDER_DEFAULTS["ollama"] + assert 2 in PROVIDER_DEFAULTS["ollama"] + assert 3 in PROVIDER_DEFAULTS["ollama"] + + def test_calc_token_cost_zero(self): + from think.models import calc_token_cost + + result = calc_token_cost( + { + "model": OLLAMA_FLASH, + "usage": {"input_tokens": 100, "output_tokens": 50}, + } + ) + assert result is not None + assert result["total_cost"] == 0.0 + + def test_provider_registry(self): + from think.providers import PROVIDER_METADATA, PROVIDER_REGISTRY + + assert "ollama" in PROVIDER_REGISTRY + assert "ollama" in PROVIDER_METADATA + assert PROVIDER_METADATA["ollama"]["label"] == "Ollama (Local)" + assert PROVIDER_METADATA["ollama"]["env_key"] == "" diff --git a/think/agents.py b/think/agents.py index fc6bf36d2..494406cbc 100644 --- a/think/agents.py +++ b/think/agents.py @@ -506,7 +506,7 @@ def prepare_config(request: dict) -> dict: backup = get_backup_provider(agent_type) if backup and backup != provider: env_key = PROVIDER_METADATA.get(backup, {}).get("env_key") - if env_key and os.getenv(env_key): + if not env_key or os.getenv(env_key): config["fallback_from"] = provider config["provider"] = backup config["model"] = resolve_model_for_provider( @@ -820,7 +820,7 @@ async def _execute_with_tools( if not backup or backup == provider: raise env_key = PROVIDER_METADATA.get(backup, {}).get("env_key") - if not env_key or not os.getenv(env_key): + if env_key and not os.getenv(env_key): raise context = config.get("context") @@ -941,7 +941,7 @@ async def _execute_generate( if not backup or backup == provider: raise env_key = PROVIDER_METADATA.get(backup, {}).get("env_key") - if not env_key or not os.getenv(env_key): + if env_key and not os.getenv(env_key): raise backup_model = resolve_model_for_provider(context, backup, "generate") @@ -1171,9 +1171,17 @@ def _check_generate(provider_name: str, tier: int, timeout: int) -> tuple[bool, from think.providers import PROVIDER_METADATA, get_provider_module env_key = PROVIDER_METADATA[provider_name]["env_key"] - if not os.getenv(env_key): + if env_key and not os.getenv(env_key): return False, f"FAIL: {env_key} not set" + # For keyless providers (e.g., Ollama), check reachability instead + if not env_key: + from think.providers import validate_key + + result = validate_key(provider_name, "") + if not result.get("valid"): + return False, f"FAIL: {result.get('error', 'unreachable')}" + try: module = get_provider_module(provider_name) model = PROVIDER_DEFAULTS[provider_name][tier] diff --git a/think/models.py b/think/models.py index cc012ec7c..b5cc38052 100644 --- a/think/models.py +++ b/think/models.py @@ -58,6 +58,10 @@ CLAUDE_OPUS_4 = "claude-opus-4-5" CLAUDE_SONNET_4 = "claude-sonnet-4-5" CLAUDE_HAIKU_4 = "claude-haiku-4-5" +OLLAMA_PRO = "ollama-local/qwen3.5:35b-a3b-bf16" +OLLAMA_FLASH = "ollama-local/qwen3.5:9b" +OLLAMA_LITE = "ollama-local/qwen3.5:2b" + # --------------------------------------------------------------------------- # System defaults: provider -> tier -> model # --------------------------------------------------------------------------- @@ -78,6 +82,11 @@ PROVIDER_DEFAULTS: Dict[str, Dict[int, str]] = { TIER_FLASH: CLAUDE_SONNET_4, TIER_LITE: CLAUDE_HAIKU_4, }, + "ollama": { + TIER_PRO: OLLAMA_PRO, + TIER_FLASH: OLLAMA_FLASH, + TIER_LITE: OLLAMA_LITE, + }, } TYPE_DEFAULTS: Dict[str, Dict[str, Any]] = { @@ -637,11 +646,13 @@ def get_model_provider(model: str) -> str: Returns ------- str - Provider name: "openai", "google", "anthropic", or "unknown" + Provider name: "openai", "google", "anthropic", "ollama", or "unknown" """ model_lower = model.lower() - if model_lower.startswith("gpt"): + if model_lower.startswith("ollama-local/"): + return "ollama" + elif model_lower.startswith("gpt"): return "openai" elif model_lower.startswith("gemini"): return "google" @@ -701,6 +712,15 @@ def calc_token_cost(token_data: Dict[str, Any]) -> Optional[Dict[str, Any]]: if provider_id == "unknown": return None + # Ollama models are local — no cost + if provider_id == "ollama": + return { + "total_cost": 0.0, + "input_cost": 0.0, + "output_cost": 0.0, + "currency": "USD", + } + # Apply price aliases for models genai-prices doesn't recognize yet model = MODEL_PRICE_ALIASES.get(model, model) diff --git a/think/providers/__init__.py b/think/providers/__init__.py index 8174693e6..b1fb20986 100644 --- a/think/providers/__init__.py +++ b/think/providers/__init__.py @@ -17,6 +17,7 @@ Available providers: - google: Google Gemini models - openai: OpenAI GPT models - anthropic: Anthropic Claude models +- ollama: Ollama local models """ from importlib import import_module @@ -37,6 +38,7 @@ PROVIDER_REGISTRY: Dict[str, str] = { "google": "think.providers.google", "openai": "think.providers.openai", "anthropic": "think.providers.anthropic", + "ollama": "think.providers.ollama", } # --------------------------------------------------------------------------- @@ -57,6 +59,7 @@ PROVIDER_METADATA: Dict[str, Dict[str, Any]] = { }, "openai": {"label": "OpenAI (GPT)", "env_key": "OPENAI_API_KEY"}, "anthropic": {"label": "Anthropic (Claude)", "env_key": "ANTHROPIC_API_KEY"}, + "ollama": {"label": "Ollama (Local)", "env_key": ""}, } diff --git a/think/providers/ollama.py b/think/providers/ollama.py new file mode 100644 index 000000000..a28076914 --- /dev/null +++ b/think/providers/ollama.py @@ -0,0 +1,607 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright (c) 2026 sol pbc + +"""Ollama (Local) provider for LLM generation and tool-calling agents. + +This module provides the Ollama provider for run_generate/run_agenerate +(text generation) and run_cogitate (tool-calling agents). + +**Generation** uses Ollama's native ``/api/chat`` endpoint via ``httpx`` +for reliable control over the ``think`` parameter, which the OpenAI-compatible +endpoint silently ignores on models like Qwen3.5. + +**Cogitate** uses the OpenCode CLI (``opencode run --format json``) as a +subprocess, following the same CLIRunner + translate pattern as the Google, +OpenAI, and Anthropic providers. OpenCode connects to local Ollama via its +OpenAI-compatible endpoint and handles tool execution internally. + +Common Parameters +----------------- +contents : str or list + The content to send to the model. +model : str + Model name with ``ollama-local/`` prefix (e.g., ``ollama-local/qwen3.5:9b``). + The prefix is stripped before sending to the Ollama API. +temperature : float + Temperature for generation (default: 0.3). +max_output_tokens : int + Maximum tokens for the model's response output. +system_instruction : str, optional + System instruction for the model. +json_output : bool + Whether to request JSON response format. +thinking_budget : int, optional + Token budget for model thinking. When > 0, enables Ollama's ``think`` + parameter. When None or 0, thinking is explicitly disabled. +timeout_s : float, optional + Request timeout in seconds. +**kwargs + Additional provider-specific options (absorbed for forward compatibility). + +Environment Variables +--------------------- +OLLAMA_BASE_URL : str + Base URL for the Ollama server (default: ``http://localhost:11434``). +""" + +from __future__ import annotations + +import logging +import os +import traceback +from typing import Any, Callable + +import httpx + +from think.models import OLLAMA_FLASH +from think.utils import now_ms + +from .cli import CLIRunner, ThinkingAggregator, assemble_prompt +from .shared import GenerateResult, JSONEventCallback, safe_raw + +LOG = logging.getLogger("think.providers.ollama") + +_OLLAMA_LOCAL_PREFIX = "ollama-local/" +_DEFAULT_BASE_URL = "http://localhost:11434" +_DEFAULT_TIMEOUT = 120.0 + +# --------------------------------------------------------------------------- +# Client management +# --------------------------------------------------------------------------- + +_sync_client: httpx.Client | None = None +_async_client: httpx.AsyncClient | None = None + + +def _get_base_url() -> str: + """Get Ollama base URL from environment or default.""" + return os.getenv("OLLAMA_BASE_URL", _DEFAULT_BASE_URL).rstrip("/") + + +def _get_client() -> httpx.Client: + """Get or create cached sync httpx client.""" + global _sync_client + if _sync_client is None: + _sync_client = httpx.Client( + base_url=_get_base_url(), + timeout=_DEFAULT_TIMEOUT, + ) + return _sync_client + + +def _get_async_client() -> httpx.AsyncClient: + """Get or create cached async httpx client.""" + global _async_client + if _async_client is None: + _async_client = httpx.AsyncClient( + base_url=_get_base_url(), + timeout=_DEFAULT_TIMEOUT, + ) + return _async_client + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _strip_model_prefix(model: str) -> str: + """Strip the ``ollama-local/`` prefix for the Ollama API. + + The Ollama API expects bare model names like ``qwen3.5:9b``, but + Solstone uses the ``ollama-local/`` prefix for provider routing. + """ + if model.startswith(_OLLAMA_LOCAL_PREFIX): + return model[len(_OLLAMA_LOCAL_PREFIX) :] + return model + + +def _build_messages( + contents: Any, + system_instruction: str | None = None, +) -> list[dict[str, str]]: + """Convert contents and system instruction to chat messages. + + Parameters + ---------- + contents + String, list of strings, or list of message dicts with ``role`` keys. + system_instruction + Optional system prompt, prepended as a system message. + + Returns + ------- + list[dict[str, str]] + Messages in ``[{role, content}, ...]`` format. + """ + messages: list[dict[str, str]] = [] + + if system_instruction: + messages.append({"role": "system", "content": system_instruction}) + + if isinstance(contents, str): + messages.append({"role": "user", "content": contents}) + elif isinstance(contents, list): + if contents and isinstance(contents[0], dict) and "role" in contents[0]: + messages.extend(contents) + else: + messages.append( + {"role": "user", "content": "\n".join(str(c) for c in contents)} + ) + else: + messages.append({"role": "user", "content": str(contents)}) + + return messages + + +def _build_request_body( + model: str, + messages: list[dict[str, str]], + temperature: float, + max_output_tokens: int, + json_output: bool, + thinking_budget: int | None, +) -> dict[str, Any]: + """Build the native Ollama /api/chat request body. + + Parameters + ---------- + model + Bare model name (prefix already stripped). + messages + Chat messages list. + temperature + Sampling temperature. + max_output_tokens + Maximum response tokens (``num_predict`` in Ollama). + json_output + Whether to request JSON response format. + thinking_budget + Thinking token budget; > 0 enables, None/0 disables. + + Returns + ------- + dict + Request body for ``POST /api/chat``. + """ + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": False, + "options": { + "temperature": temperature, + "num_predict": max_output_tokens, + }, + } + + # Thinking control: this is the reason we use the native API. + # The OpenAI-compat endpoint ignores this parameter. + if thinking_budget is not None and thinking_budget > 0: + body["think"] = True + else: + body["think"] = False + + if json_output: + body["format"] = "json" + + return body + + +def _normalize_finish_reason(data: dict[str, Any]) -> str | None: + """Normalize Ollama's done_reason to standard values. + + Returns ``"stop"``, ``"max_tokens"``, or None. + """ + if not data.get("done"): + return None + + reason = data.get("done_reason", "") + if reason == "stop": + return "stop" + elif reason == "length": + return "max_tokens" + elif reason: + return reason + return "stop" # done=True with no reason implies normal completion + + +def _extract_usage(data: dict[str, Any]) -> dict[str, int]: + """Extract normalized usage dict from native Ollama response. + + Ollama uses ``prompt_eval_count`` and ``eval_count`` instead of the + OpenAI-style ``prompt_tokens`` / ``completion_tokens``. + """ + input_tokens = data.get("prompt_eval_count", 0) or 0 + output_tokens = data.get("eval_count", 0) or 0 + return { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + + +def _extract_thinking(data: dict[str, Any]) -> list | None: + """Extract thinking content from native Ollama response. + + The native API returns a ``thinking`` field on the message when + thinking is enabled. + """ + message = data.get("message", {}) + thinking = message.get("thinking") + if thinking and isinstance(thinking, str) and thinking.strip(): + return [{"summary": thinking.strip()}] + return None + + +def _parse_response(data: dict[str, Any]) -> GenerateResult: + """Parse the native Ollama /api/chat response into GenerateResult.""" + message = data.get("message", {}) + text = message.get("content", "") + + return GenerateResult( + text=text, + usage=_extract_usage(data), + finish_reason=_normalize_finish_reason(data), + thinking=_extract_thinking(data), + ) + + +# --------------------------------------------------------------------------- +# run_generate / run_agenerate +# --------------------------------------------------------------------------- + + +def run_generate( + contents: str | list[Any], + model: str, + temperature: float = 0.3, + max_output_tokens: int = 8192 * 2, + system_instruction: str | None = None, + json_output: bool = False, + thinking_budget: int | None = None, + timeout_s: float | None = None, + **kwargs: Any, +) -> GenerateResult: + """Generate text synchronously via local Ollama. + + Returns GenerateResult with text, usage, finish_reason, and thinking. + See module docstring for parameter details. + """ + client = _get_client() + api_model = _strip_model_prefix(model) + messages = _build_messages(contents, system_instruction) + body = _build_request_body( + api_model, + messages, + temperature, + max_output_tokens, + json_output, + thinking_budget, + ) + + response = client.post( + "/api/chat", + json=body, + timeout=timeout_s or _DEFAULT_TIMEOUT, + ) + response.raise_for_status() + return _parse_response(response.json()) + + +async def run_agenerate( + contents: str | list[Any], + model: str, + temperature: float = 0.3, + max_output_tokens: int = 8192 * 2, + system_instruction: str | None = None, + json_output: bool = False, + thinking_budget: int | None = None, + timeout_s: float | None = None, + **kwargs: Any, +) -> GenerateResult: + """Generate text asynchronously via local Ollama. + + Returns GenerateResult with text, usage, finish_reason, and thinking. + See module docstring for parameter details. + """ + client = _get_async_client() + api_model = _strip_model_prefix(model) + messages = _build_messages(contents, system_instruction) + body = _build_request_body( + api_model, + messages, + temperature, + max_output_tokens, + json_output, + thinking_budget, + ) + + response = await client.post( + "/api/chat", + json=body, + timeout=timeout_s or _DEFAULT_TIMEOUT, + ) + response.raise_for_status() + return _parse_response(response.json()) + + +# --------------------------------------------------------------------------- +# run_cogitate via OpenCode CLI +# --------------------------------------------------------------------------- + + +def _translate_opencode( + event: dict[str, Any], + aggregator: ThinkingAggregator, + callback: JSONEventCallback, + usage_out: dict[str, Any], +) -> str | None: + """Translate an OpenCode JSONL event into our standard Event types. + + Args: + event: Raw JSONL event dict from ``opencode run --format json``. + aggregator: ThinkingAggregator for buffering text. + callback: JSONEventCallback for emitting events. + usage_out: Mutable dict to receive usage stats from step_finish events. + + Returns: + The CLI session ID from step_start events, or None. + """ + event_type = event.get("type") + part = event.get("part", {}) + + # -- step_start: capture session ID ------------------------------------ + if event_type == "step_start": + return event.get("sessionID") + + # -- text: accumulate assistant text ----------------------------------- + if event_type == "text": + text = part.get("text", "") + if text: + aggregator.accumulate(text) + return None + + # -- tool_use: emit tool_start + tool_end ------------------------------ + # OpenCode reports tools as already completed, so we emit both events + # back-to-back from a single JSONL line. + if event_type == "tool_use": + aggregator.flush_as_thinking(raw_events=[event]) + + tool_name = part.get("tool", "") + call_id = part.get("callID", "") + state = part.get("state", {}) + tool_input = state.get("input", {}) + tool_output = state.get("output", "") + + callback.emit( + { + "event": "tool_start", + "tool": tool_name, + "args": tool_input, + "call_id": call_id, + "raw": safe_raw([event]), + "ts": now_ms(), + } + ) + callback.emit( + { + "event": "tool_end", + "tool": tool_name, + "args": tool_input, + "result": tool_output, + "call_id": call_id, + "raw": safe_raw([event]), + "ts": now_ms(), + } + ) + return None + + # -- step_finish: capture usage ---------------------------------------- + if event_type == "step_finish": + tokens = part.get("tokens") + if tokens and usage_out is not None: + input_tokens = tokens.get("input", 0) + output_tokens = tokens.get("output", 0) + total_tokens = tokens.get("total", 0) + # Accumulate across steps (OpenCode emits one per turn) + usage_out["input_tokens"] = usage_out.get("input_tokens", 0) + input_tokens + usage_out["output_tokens"] = ( + usage_out.get("output_tokens", 0) + output_tokens + ) + usage_out["total_tokens"] = usage_out.get("total_tokens", 0) + total_tokens + reasoning = tokens.get("reasoning", 0) + if reasoning: + usage_out["reasoning_tokens"] = ( + usage_out.get("reasoning_tokens", 0) + reasoning + ) + cache = tokens.get("cache", {}) + cached_read = cache.get("read", 0) + if cached_read: + usage_out["cached_tokens"] = ( + usage_out.get("cached_tokens", 0) + cached_read + ) + return None + + # Unknown event type — log and skip + LOG.debug("Unknown OpenCode CLI event type: %s", event_type) + return None + + +def _build_opencode_env() -> dict[str, str]: + """Build environment dict for the OpenCode subprocess. + + Sets ``OPENAI_API_KEY`` to a placeholder if not already set, since + OpenCode's OpenAI-compatible provider requires it even for local Ollama. + """ + env = os.environ.copy() + if not env.get("OPENAI_API_KEY"): + env["OPENAI_API_KEY"] = "ollama" + return env + + +async def run_cogitate( + config: dict[str, Any], + on_event: Callable[[dict], None] | None = None, +) -> str: + """Run a prompt with tool-calling support via OpenCode CLI + local Ollama. + + Uses the OpenCode CLI as a subprocess agent, which connects to the local + Ollama instance and provides built-in tools (bash, read, glob, grep, etc.). + + Args: + config: Complete configuration dictionary including prompt, system_instruction, + user_instruction, extra_context, model, etc. + on_event: Optional event callback + """ + model = _strip_model_prefix(config.get("model", OLLAMA_FLASH)) + session_id = config.get("session_id") + callback = JSONEventCallback(on_event) + + try: + # Check that OpenCode CLI is available + import shutil + + if not shutil.which("opencode"): + raise RuntimeError( + "Cogitate requires OpenCode CLI (opencode). " + "Install from https://opencode.ai and configure it with a local " + "Ollama provider. Generate works without it." + ) + + # Assemble prompt from config fields + prompt_body, system_instruction = assemble_prompt(config) + + # OpenCode has no --system-prompt flag; prepend to prompt body + if system_instruction: + prompt_body = system_instruction + "\n\n" + prompt_body + + # Build CLI command. + # --title skips OpenCode's title-generation LLM call (avoids delays). + agent_name = config.get("name", "sol-agent") + cmd = [ + "opencode", + "run", + "--format", + "json", + "--title", + agent_name, + "-m", + f"ollama/{model}", + ] + + # Resume from previous session if continuing + if session_id: + cmd.extend(["--session", session_id]) + + # Mutable container for usage accumulation + usage: dict[str, Any] = {} + + def translate( + event: dict[str, Any], agg: ThinkingAggregator, cb: JSONEventCallback + ) -> str | None: + return _translate_opencode(event, agg, cb, usage) + + aggregator = ThinkingAggregator(callback, model=model) + runner = CLIRunner( + cmd=cmd, + prompt_text=prompt_body, + translate=translate, + callback=callback, + aggregator=aggregator, + env=_build_opencode_env(), + # Local models are slower than cloud APIs; allow more time for + # the first event (model loading + initial inference). + first_event_timeout=120, + ) + + result = await runner.run() + + # Emit finish event (CLIRunner does not emit one) + finish_event: dict[str, Any] = { + "event": "finish", + "result": result, + "ts": now_ms(), + } + if usage: + finish_event["usage"] = usage + if runner.cli_session_id: + finish_event["cli_session_id"] = runner.cli_session_id + callback.emit(finish_event) + return result + except Exception as exc: + callback.emit( + { + "event": "error", + "error": str(exc), + "trace": traceback.format_exc(), + } + ) + setattr(exc, "_evented", True) + raise + + +# --------------------------------------------------------------------------- +# list_models / validate_key +# --------------------------------------------------------------------------- + + +def list_models() -> list[dict]: + """List available models from the local Ollama instance. + + Returns + ------- + list[dict] + List of model info dicts from the Ollama ``/api/tags`` endpoint. + """ + client = _get_client() + response = client.get("/api/tags") + response.raise_for_status() + return response.json().get("models", []) + + +def validate_key(api_key: str) -> dict: + """Check that the local Ollama instance is reachable. + + The ``api_key`` parameter is ignored — Ollama requires no authentication. + Connectivity is validated by hitting the version endpoint. + + Returns ``{"valid": True}`` if reachable, ``{"valid": False, "error": "..."}`` + if not. + """ + try: + base_url = _get_base_url() + response = httpx.get(f"{base_url}/api/version", timeout=5) + response.raise_for_status() + return {"valid": True} + except Exception as e: + return {"valid": False, "error": str(e)} + + +__all__ = [ + "run_generate", + "run_agenerate", + "run_cogitate", + "list_models", + "validate_key", +] diff --git a/uv.lock b/uv.lock index 4245e8133..a31618fda 100644 --- a/uv.lock +++ b/uv.lock @@ -3512,6 +3512,7 @@ dependencies = [ { name = "flask-sock" }, { name = "genai-prices" }, { name = "google-genai" }, + { name = "httpx" }, { name = "icalendar" }, { name = "markdown" }, { name = "mistune" }, @@ -3557,6 +3558,7 @@ requires-dist = [ { name = "flask-sock" }, { name = "genai-prices" }, { name = "google-genai" }, + { name = "httpx" }, { name = "icalendar" }, { name = "markdown" }, { name = "mistune" },