kiln-ai 0.5.3__tar.gz → 0.5.5__tar.gz

{kiln_ai-0.5.3 → kiln_ai-0.5.5}/.gitignore

@@ -7,6 +7,7 @@ __pycache__/
 *$py.class
 .coverage
 **/.venv
+**/*.egg-info
 
 libs/core/docs
 libs/core/build

kiln_ai-0.5.5/LICENSE.txt

@@ -0,0 +1,13 @@
+
+
+This license applies only to the software in the libs/core directory.
+
+=======================================================
+
+Copyright 2024 - Chesterfield Laboratories Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

{kiln_ai-0.5.3 → kiln_ai-0.5.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: kiln-ai
-Version: 0.5.3
+Version: 0.5.5
 Summary: Kiln AI
 Project-URL: Homepage, https://getkiln.ai
 Project-URL: Repository, https://github.com/Kiln-AI/kiln
@@ -9,9 +9,11 @@ Project-URL: Issues, https://github.com/Kiln-AI/kiln/issues
 Author-email: "Steve Cosman, Chesterfield Laboratories Inc" <scosman@users.noreply.github.com>
 License-File: LICENSE.txt
 Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
 Requires-Dist: coverage>=7.6.4
 Requires-Dist: jsonschema>=4.23.0
@@ -29,6 +31,12 @@ Description-Content-Type: text/markdown
 
 # kiln_ai
 
+<p align="center">
+    <picture>
+        <img width="205" alt="Kiln AI Logo" src="https://github.com/user-attachments/assets/5fbcbdf7-1feb-45c9-bd73-99a46dd0a47f">
+    </picture>
+</p>
+
 [![PyPI - Version](https://img.shields.io/pypi/v/kiln-ai.svg?logo=pypi&label=PyPI&logoColor=gold)](https://pypi.org/project/kiln-ai)
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/kiln-ai.svg)](https://pypi.org/project/kiln-ai)
 [![Docs](https://img.shields.io/badge/docs-pdoc-blue)](https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html)
@@ -41,12 +49,41 @@ Description-Content-Type: text/markdown
 pip install kiln_ai
 ```
 
-## About Kiln AI
+## About
+
+This package is the Kiln AI core library. There is also a separate desktop application and server package. Learn more about Kiln AI at [getkiln.ai](https://getkiln.ai)
+
+- Github: [github.com/Kiln-AI/kiln](https://github.com/Kiln-AI/kiln)
+- Core Library Docs: [https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html](https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html)
+
+## Quick Start
 
-Learn more about Kiln AI at [getkiln.ai](https://getkiln.ai)
+```python
+from kiln_ai.datamodel import Project
 
-This package is the Kiln AI core library. There is also a separate desktop application and server package.
+print("Reading Kiln project")
+project = Project.load_from_file("path/to/project.kiln")
+print("Project: ", project.name, " - ", project.description)
 
-Github: [github.com/Kiln-AI/kiln](https://github.com/Kiln-AI/kiln)
+task = project.tasks()[0]
+print("Task: ", task.name, " - ", task.description)
+print("Total dataset size:", len(task.runs()))
 
-Docs: [https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html](https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html)
+# ... app specific code using the typed kiln datamodel
+
+# Alternatively, load data into pandas or a similar tool:
+import glob
+import json
+import pandas as pd
+from pathlib import Path
+
+dataitem_glob = str(task.path.parent) + "/runs/*/task_run.kiln"
+
+dfs = []
+for file in glob.glob(dataitem_glob):
+    js = json.loads(Path(file).read_text())
+    df = pd.json_normalize(js)
+    dfs.append(df)
+final_df = pd.concat(dfs, ignore_index=True)
+print(final_df)
+```

kiln_ai-0.5.5/README.md

@@ -0,0 +1,58 @@
+# kiln_ai
+
+<p align="center">
+    <picture>
+        <img width="205" alt="Kiln AI Logo" src="https://github.com/user-attachments/assets/5fbcbdf7-1feb-45c9-bd73-99a46dd0a47f">
+    </picture>
+</p>
+
+[![PyPI - Version](https://img.shields.io/pypi/v/kiln-ai.svg?logo=pypi&label=PyPI&logoColor=gold)](https://pypi.org/project/kiln-ai)
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/kiln-ai.svg)](https://pypi.org/project/kiln-ai)
+[![Docs](https://img.shields.io/badge/docs-pdoc-blue)](https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html)
+
+---
+
+## Installation
+
+```console
+pip install kiln_ai
+```
+
+## About
+
+This package is the Kiln AI core library. There is also a separate desktop application and server package. Learn more about Kiln AI at [getkiln.ai](https://getkiln.ai)
+
+- Github: [github.com/Kiln-AI/kiln](https://github.com/Kiln-AI/kiln)
+- Core Library Docs: [https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html](https://kiln-ai.github.io/Kiln/kiln_core_docs/index.html)
+
+## Quick Start
+
+```python
+from kiln_ai.datamodel import Project
+
+print("Reading Kiln project")
+project = Project.load_from_file("path/to/project.kiln")
+print("Project: ", project.name, " - ", project.description)
+
+task = project.tasks()[0]
+print("Task: ", task.name, " - ", task.description)
+print("Total dataset size:", len(task.runs()))
+
+# ... app specific code using the typed kiln datamodel
+
+# Alternatively, load data into pandas or a similar tool:
+import glob
+import json
+import pandas as pd
+from pathlib import Path
+
+dataitem_glob = str(task.path.parent) + "/runs/*/task_run.kiln"
+
+dfs = []
+for file in glob.glob(dataitem_glob):
+    js = json.loads(Path(file).read_text())
+    df = pd.json_normalize(js)
+    dfs.append(df)
+final_df = pd.concat(dfs, ignore_index=True)
+print(final_df)
+```

{kiln_ai-0.5.3 → kiln_ai-0.5.5}/kiln_ai/adapters/langchain_adapters.py

@@ -1,16 +1,23 @@
 from typing import Dict
 
+from langchain_core.language_models import LanguageModelInput
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
+from langchain_core.runnables import Runnable
+from pydantic import BaseModel
 
 import kiln_ai.datamodel as datamodel
 
 from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder
 from .ml_model_list import langchain_model_from
 
+LangChainModelType = BaseChatModel | Runnable[LanguageModelInput, Dict | BaseModel]
+
 
 class LangChainPromptAdapter(BaseAdapter):
+    _model: LangChainModelType | None = None
+
     def __init__(
         self,
         kiln_task: datamodel.Task,
@@ -21,7 +28,7 @@ class LangChainPromptAdapter(BaseAdapter):
     ):
         super().__init__(kiln_task, prompt_builder=prompt_builder)
         if custom_model is not None:
-            self.model = custom_model
+            self._model = custom_model
 
             # Attempt to infer model provider and name from custom model
             self.model_provider = "custom.langchain:" + custom_model.__class__.__name__
@@ -37,19 +44,32 @@ class LangChainPromptAdapter(BaseAdapter):
             ):
                 self.model_name = "custom.langchain:" + getattr(custom_model, "model")
         elif model_name is not None:
-            self.model = langchain_model_from(model_name, provider)
             self.model_name = model_name
             self.model_provider = provider or "custom.langchain.default_provider"
         else:
             raise ValueError(
                 "model_name and provider must be provided if custom_model is not provided"
             )
+
+    def adapter_specific_instructions(self) -> str | None:
+        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
         if self.has_structured_output():
-            if not hasattr(self.model, "with_structured_output") or not callable(
-                getattr(self.model, "with_structured_output")
+            return "Always respond with a tool call. Never respond with a human readable message."
+        return None
+
+    async def model(self) -> LangChainModelType:
+        # cached model
+        if self._model:
+            return self._model
+
+        self._model = await langchain_model_from(self.model_name, self.model_provider)
+
+        if self.has_structured_output():
+            if not hasattr(self._model, "with_structured_output") or not callable(
+                getattr(self._model, "with_structured_output")
             ):
                 raise ValueError(
-                    f"model {self.model} does not support structured output, cannot use output_json_schema"
+                    f"model {self._model} does not support structured output, cannot use output_json_schema"
                 )
             # Langchain expects title/description to be at top level, on top of json schema
             output_schema = self.kiln_task.output_schema()
@@ -59,15 +79,10 @@ class LangChainPromptAdapter(BaseAdapter):
                 )
             output_schema["title"] = "task_response"
             output_schema["description"] = "A response from the task"
-            self.model = self.model.with_structured_output(
+            self._model = self._model.with_structured_output(
                 output_schema, include_raw=True
             )
-
-    def adapter_specific_instructions(self) -> str | None:
-        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
-        if self.has_structured_output():
-            return "Always respond with a tool call. Never respond with a human readable message."
-        return None
+        return self._model
 
     async def _run(self, input: Dict | str) -> Dict | str:
         prompt = self.build_prompt()
@@ -76,7 +91,8 @@ class LangChainPromptAdapter(BaseAdapter):
             SystemMessage(content=prompt),
             HumanMessage(content=user_msg),
         ]
-        response = self.model.invoke(messages)
+        model = await self.model()
+        response = model.invoke(messages)
 
         if self.has_structured_output():
             if (

{kiln_ai-0.5.3 → kiln_ai-0.5.5}/kiln_ai/adapters/ml_model_list.py

@@ -2,9 +2,10 @@ import os
 from dataclasses import dataclass
 from enum import Enum
 from os import getenv
-from typing import Dict, List, NoReturn
+from typing import Any, Dict, List, NoReturn
 
 import httpx
+import requests
 from langchain_aws import ChatBedrockConverse
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_groq import ChatGroq
@@ -43,6 +44,8 @@ class ModelFamily(str, Enum):
     phi = "phi"
     mistral = "mistral"
     gemma = "gemma"
+    gemini = "gemini"
+    claude = "claude"
 
 
 # Where models have instruct and raw versions, instruct is default and raw is specified
@@ -55,6 +58,9 @@ class ModelName(str, Enum):
     llama_3_1_8b = "llama_3_1_8b"
     llama_3_1_70b = "llama_3_1_70b"
     llama_3_1_405b = "llama_3_1_405b"
+    llama_3_2_3b = "llama_3_2_3b"
+    llama_3_2_11b = "llama_3_2_11b"
+    llama_3_2_90b = "llama_3_2_90b"
     gpt_4o_mini = "gpt_4o_mini"
     gpt_4o = "gpt_4o"
     phi_3_5 = "phi_3_5"
@@ -63,6 +69,12 @@ class ModelName(str, Enum):
     gemma_2_2b = "gemma_2_2b"
     gemma_2_9b = "gemma_2_9b"
     gemma_2_27b = "gemma_2_27b"
+    claude_3_5_haiku = "claude_3_5_haiku"
+    claude_3_5_sonnet = "claude_3_5_sonnet"
+    gemini_1_5_flash = "gemini_1_5_flash"
+    gemini_1_5_flash_8b = "gemini_1_5_flash_8b"
+    gemini_1_5_pro = "gemini_1_5_pro"
+    nemotron_70b = "nemotron_70b"
 
 
 class KilnModelProvider(BaseModel):
@@ -132,6 +144,79 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Claude 3.5 Haiku
+    KilnModel(
+        family=ModelFamily.claude,
+        name=ModelName.claude_3_5_haiku,
+        friendly_name="Claude 3.5 Haiku",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "anthropic/claude-3-5-haiku"},
+            ),
+        ],
+    ),
+    # Claude 3.5 Sonnet
+    KilnModel(
+        family=ModelFamily.claude,
+        name=ModelName.claude_3_5_sonnet,
+        friendly_name="Claude 3.5 Sonnet",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "anthropic/claude-3.5-sonnet"},
+            ),
+        ],
+    ),
+    # Gemini 1.5 Pro
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_1_5_pro,
+        friendly_name="Gemini 1.5 Pro",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "google/gemini-pro-1.5"},
+            ),
+        ],
+    ),
+    # Gemini 1.5 Flash
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_1_5_flash,
+        friendly_name="Gemini 1.5 Flash",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "google/gemini-flash-1.5"},
+            ),
+        ],
+    ),
+    # Gemini 1.5 Flash 8B
+    KilnModel(
+        family=ModelFamily.gemini,
+        name=ModelName.gemini_1_5_flash_8b,
+        friendly_name="Gemini 1.5 Flash 8B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "google/gemini-flash-1.5-8b"},
+            ),
+        ],
+    ),
+    # Nemotron 70B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.nemotron_70b,
+        friendly_name="Nemotron 70B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                provider_options={"model": "nvidia/llama-3.1-nemotron-70b-instruct"},
+            ),
+        ],
+    ),
     # Llama 3.1-8b
     KilnModel(
         family=ModelFamily.llama,
@@ -144,6 +229,7 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
+                supports_structured_output=False,
                 provider_options={
                     "model": "meta.llama3-1-8b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -151,10 +237,14 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
-                provider_options={"model": "llama3.1"},  # 8b is default
+                provider_options={
+                    "model": "llama3.1:8b",
+                    "model_aliases": ["llama3.1"],  # 8b is default
+                },
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,
                 provider_options={"model": "meta-llama/llama-3.1-8b-instruct"},
             ),
         ],
@@ -171,7 +261,6 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
-                # TODO: this should work but a bug in the bedrock response schema
                 supports_structured_output=False,
                 provider_options={
                     "model": "meta.llama3-1-70b-instruct-v1:0",
@@ -182,11 +271,10 @@ built_in_models: List[KilnModel] = [
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "meta-llama/llama-3.1-70b-instruct"},
             ),
-            # TODO: enable once tests update to check if model is available
-            # KilnModelProvider(
-            #     provider=ModelProviders.ollama,
-            #     provider_options={"model": "llama3.1:70b"},
-            # ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "llama3.1:70b"},
+            ),
         ],
     ),
     # Llama 3.1 405b
@@ -195,11 +283,6 @@ built_in_models: List[KilnModel] = [
         name=ModelName.llama_3_1_405b,
         friendly_name="Llama 3.1 405B",
         providers=[
-            # TODO: bring back when groq does: https://console.groq.com/docs/models
-            # KilnModelProvider(
-            #     name=ModelProviderName.groq,
-            #     provider_options={"model": "llama-3.1-405b-instruct-v1:0"},
-            # ),
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
                 provider_options={
@@ -207,11 +290,10 @@ built_in_models: List[KilnModel] = [
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
                 },
             ),
-            # TODO: enable once tests update to check if model is available
-            # KilnModelProvider(
-            #     name=ModelProviderName.ollama,
-            #     provider_options={"model": "llama3.1:405b"},
-            # ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "llama3.1:405b"},
+            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "meta-llama/llama-3.1-405b-instruct"},
@@ -247,11 +329,49 @@ built_in_models: List[KilnModel] = [
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "mistralai/mistral-large"},
             ),
-            # TODO: enable once tests update to check if model is available
-            # KilnModelProvider(
-            #     provider=ModelProviders.ollama,
-            #     provider_options={"model": "mistral-large"},
-            # ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={"model": "mistral-large"},
+            ),
+        ],
+    ),
+    # Llama 3.2 3B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_2_3b,
+        friendly_name="Llama 3.2 3B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                provider_options={"model": "meta-llama/llama-3.2-3b-instruct"},
+            ),
+        ],
+    ),
+    # Llama 3.2 11B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_2_11b,
+        friendly_name="Llama 3.2 11B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                provider_options={"model": "meta-llama/llama-3.2-11b-vision-instruct"},
+            ),
+        ],
+    ),
+    # Llama 3.2 90B
+    KilnModel(
+        family=ModelFamily.llama,
+        name=ModelName.llama_3_2_90b,
+        friendly_name="Llama 3.2 90B",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                provider_options={"model": "meta-llama/llama-3.2-90b-vision-instruct"},
+            ),
         ],
     ),
     # Phi 3.5
@@ -263,6 +383,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_structured_output=False,
                 provider_options={"model": "phi3.5"},
             ),
             KilnModelProvider(
@@ -280,6 +401,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_structured_output=False,
                 provider_options={
                     "model": "gemma2:2b",
                 },
@@ -293,13 +415,12 @@ built_in_models: List[KilnModel] = [
         friendly_name="Gemma 2 9B",
         supports_structured_output=False,
         providers=[
-            # TODO: enable once tests update to check if model is available
-            # KilnModelProvider(
-            #     name=ModelProviderName.ollama,
-            #     provider_options={
-            #         "model": "gemma2:9b",
-            #     },
-            # ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={
+                    "model": "gemma2:9b",
+                },
+            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "google/gemma-2-9b-it"},
@@ -313,13 +434,12 @@ built_in_models: List[KilnModel] = [
         friendly_name="Gemma 2 27B",
         supports_structured_output=False,
         providers=[
-            # TODO: enable once tests update to check if model is available
-            # KilnModelProvider(
-            #     name=ModelProviderName.ollama,
-            #     provider_options={
-            #         "model": "gemma2:27b",
-            #     },
-            # ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                provider_options={
+                    "model": "gemma2:27b",
+                },
+            ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 provider_options={"model": "google/gemma-2-27b-it"},
@@ -417,7 +537,9 @@ def check_provider_warnings(provider_name: ModelProviderName):
             raise ValueError(warning_check.message)
 
 
-def langchain_model_from(name: str, provider_name: str | None = None) -> BaseChatModel:
+async def langchain_model_from(
+    name: str, provider_name: str | None = None
+) -> BaseChatModel:
     """
     Creates a LangChain chat model instance for the specified model and provider.
 
@@ -476,7 +598,23 @@ def langchain_model_from(name: str, provider_name: str | None = None) -> BaseCha
             **provider.provider_options,
         )
     elif provider.name == ModelProviderName.ollama:
-        return ChatOllama(**provider.provider_options, base_url=ollama_base_url())
+        # Ollama model naming is pretty flexible. We try a few versions of the model name
+        potential_model_names = []
+        if "model" in provider.provider_options:
+            potential_model_names.append(provider.provider_options["model"])
+        if "model_aliases" in provider.provider_options:
+            potential_model_names.extend(provider.provider_options["model_aliases"])
+
+        # Get the list of models Ollama supports
+        ollama_connection = await get_ollama_connection()
+        if ollama_connection is None:
+            raise ValueError("Failed to connect to Ollama. Ensure Ollama is running.")
+
+        for model_name in potential_model_names:
+            if ollama_model_supported(ollama_connection, model_name):
+                return ChatOllama(model=model_name, base_url=ollama_base_url())
+
+        raise ValueError(f"Model {name} not installed on Ollama")
     elif provider.name == ModelProviderName.openrouter:
         api_key = Config.shared().open_router_api_key
         base_url = getenv("OPENROUTER_BASE_URL") or "https://openrouter.ai/api/v1"
@@ -519,3 +657,67 @@ async def ollama_online() -> bool:
     except httpx.RequestError:
         return False
     return True
+
+
+class OllamaConnection(BaseModel):
+    message: str
+    models: List[str]
+
+
+# Parse the Ollama /api/tags response
+def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
+    # Build a list of models we support for Ollama from the built-in model list
+    supported_ollama_models = [
+        provider.provider_options["model"]
+        for model in built_in_models
+        for provider in model.providers
+        if provider.name == ModelProviderName.ollama
+    ]
+    # Append model_aliases to supported_ollama_models
+    supported_ollama_models.extend(
+        [
+            alias
+            for model in built_in_models
+            for provider in model.providers
+            for alias in provider.provider_options.get("model_aliases", [])
+        ]
+    )
+
+    if "models" in tags:
+        models = tags["models"]
+        if isinstance(models, list):
+            model_names = [model["model"] for model in models]
+            print(f"model_names: {model_names}")
+            available_supported_models = [
+                model
+                for model in model_names
+                if model in supported_ollama_models
+                or model in [f"{m}:latest" for m in supported_ollama_models]
+            ]
+            if available_supported_models:
+                return OllamaConnection(
+                    message="Ollama connected",
+                    models=available_supported_models,
+                )
+
+    return OllamaConnection(
+        message="Ollama is running, but no supported models are installed. Install one or more supported model, like 'ollama pull phi3.5'.",
+        models=[],
+    )
+
+
+async def get_ollama_connection() -> OllamaConnection | None:
+    """
+    Gets the connection status for Ollama.
+    """
+    try:
+        tags = requests.get(ollama_base_url() + "/api/tags", timeout=5).json()
+
+    except Exception:
+        return None
+
+    return parse_ollama_tags(tags)
+
+
+def ollama_model_supported(conn: OllamaConnection, model_name: str) -> bool:
+    return model_name in conn.models or f"{model_name}:latest" in conn.models