Merge pull request #389 from namin/ollama-autoconf

Support provider Ollama for automatic configuration.

Author: jlowin

docs/guides/configure-llms.mdx

@@ -47,6 +47,7 @@ At this time, supported providers for automatic configuration include:
 | Anthropic | `anthropic` | (included) |
 | Google   | `google`       | `langchain_google_genai` |
 | Groq      | `groq`       | `langchain_groq` |
+| Ollama    | `ollama`    | `langchain-ollama` |
 
 If the required dependencies are not installed, ControlFlow will be unable to load the model and will raise an error.
 

pyproject.toml

@@ -51,6 +51,7 @@ tests = [
     "langchain_community",
     "langchain_google_genai",
     "langchain_groq",
+    "langchain-ollama',
     "pytest-asyncio>=0.18.2,!=0.22.0,<0.23.0",
     "pytest-env>=0.8,<2.0",
     "pytest-rerunfailures>=10,<14",

src/controlflow/llm/models.py

@@ -60,6 +60,14 @@ def get_model(
                 "To use Groq as an LLM provider, please install the `langchain_groq` package."
             )
         cls = ChatGroq
+    elif provider == "ollama":
+        try:
+            from langchain_ollama import ChatOllama
+        except ImportError:
+            raise ImportError(
+                "To use Ollama as an LLM provider, please install the `langchain-ollama` package."
+            )
+        cls = ChatOllama
     else:
         raise ValueError(
             f"Could not load provider `{provider}` automatically. Please provide the LLM class manually."

tests/llm/test_models.py

@@ -2,6 +2,7 @@
 from langchain_anthropic import ChatAnthropic
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 
 from controlflow.llm.models import get_model
@@ -45,6 +46,10 @@ def test_get_groq_model(monkeypatch):
     assert isinstance(model, ChatGroq)
     assert model.model_name == "mixtral-8x7b-32768"
 
+def test_get_ollama_model(monkeypatch):
+    model = get_model("ollama/qwen2.5")
+    assert isinstance(model, ChatOllama)
+    assert model.model == "qwen2.5"
 
 def test_get_model_with_invalid_format():
     with pytest.raises(ValueError, match="The model `gpt-4o` is not valid."):