Add support for schemas

handyman5 · web-flow · commit 041b3f5f1a31 · 2025-03-02T16:32:33.000+01:00
diff --git a/llm_ollama.py b/llm_ollama.py
@@ -52,6 +52,7 @@ def register_embedding_models(register):
 
 class _SharedOllama:
     can_stream: bool = True
+    supports_schema: bool = True
     attachment_types = {
         "image/png",
         "image/jpeg",
@@ -198,6 +199,8 @@ def execute(
         usage = None
         if json_object:
             kwargs["format"] = "json"
+        elif prompt.schema:
+            kwargs["format"] = prompt.schema
 
         if stream:
             response_stream = ollama.Client().chat(
@@ -256,6 +259,8 @@ async def execute(
         usage = None
         if json_object:
             kwargs["format"] = "json"
+        elif prompt.schema:
+            kwargs["format"] = prompt.schema
 
         try:
             if stream:
diff --git a/pyproject.toml b/pyproject.toml
@@ -6,7 +6,7 @@ readme = "README.md"
 authors = [{ name = "Sergey Alexandrov" }]
 license = { text = "Apache-2.0" }
 classifiers = ["License :: OSI Approved :: Apache Software License"]
-dependencies = ["llm>=0.19", "ollama>=0.4", "pydantic>=2"]
+dependencies = ["llm>=0.23", "ollama>=0.4", "pydantic>=2"]
 requires-python = ">=3.9"
 
 [project.urls]
diff --git a/tests/test_ollama_integration.py b/tests/test_ollama_integration.py
@@ -1,6 +1,8 @@
+import json
+
 import pytest
 
-from llm import get_async_model
+from llm import get_async_model, schema_dsl
 
 
 @pytest.mark.integration
@@ -11,3 +13,21 @@ async def test_async_model_prompt():
     response = model.prompt("a short poem about tea")
     response_text = await response.text()
     assert len(response_text) > 0
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_async_model_prompt_with_schema():
+    """Tests actual run. Needs llama3.2"""
+    model = get_async_model("llama3.2:latest")
+    response = model.prompt(
+        "Describe a nice dog with a surprising name",
+        schema=schema_dsl("name, age int, bio"),
+    )
+    response_text = await response.text()
+    assert len(response_text) > 0
+    json_response = json.loads(response_text)
+    assert "name" in json_response
+    assert "bio" in json_response
+    assert "age" in json_response
+    assert isinstance(json_response["age"], int)
