Merge branch 'add-json-schema' of github.com:wkukka1/ai-server into add-json-schema

Author: Will Kukkamalla

test/test_cli_mode.py

@@ -108,7 +108,9 @@ def test_cli_mode_uses_llamacpp_when_available(self):
 
         assert result == "CLI response from DeepSeek V3"
         self.mock_available.assert_called_once_with(TEST_LLAMACPP_MODEL)
-        self.mock_chat_llamacpp.assert_called_once_with(TEST_LLAMACPP_MODEL, 'Write a function', system_prompt=None, image_files=None, json_schema=None)
+        self.mock_chat_llamacpp.assert_called_once_with(
+            TEST_LLAMACPP_MODEL, 'Write a function', system_prompt=None, image_files=None, json_schema=None
+        )
 
     def test_cli_mode_fallback_to_ollama_when_unavailable(self):
         """Test CLI mode falls back to ollama when model not available in llama.cpp."""
@@ -119,7 +121,9 @@ def test_cli_mode_fallback_to_ollama_when_unavailable(self):
 
         assert result == "Ollama response from DeepSeek Coder"
         self.mock_available.assert_called_once_with(TEST_OLLAMA_MODEL)
-        self.mock_chat_ollama.assert_called_once_with(TEST_OLLAMA_MODEL, 'Help with coding', system_prompt=None, image_files=None, json_schema=None)
+        self.mock_chat_ollama.assert_called_once_with(
+            TEST_OLLAMA_MODEL, 'Help with coding', system_prompt=None, image_files=None, json_schema=None
+        )
 
     def test_default_mode_is_cli(self):
         """Test that default mode is CLI when no llama_mode specified."""
@@ -130,7 +134,9 @@ def test_default_mode_is_cli(self):
 
         assert result == "Default CLI mode response"
         self.mock_available.assert_called_once_with(TEST_LLAMACPP_MODEL)
-        self.mock_chat_llamacpp.assert_called_once_with(TEST_LLAMACPP_MODEL, 'Help me', system_prompt=None, image_files=None, json_schema=None)
+        self.mock_chat_llamacpp.assert_called_once_with(
+            TEST_LLAMACPP_MODEL, 'Help me', system_prompt=None, image_files=None, json_schema=None
+        )
 
 
 class TestCLIModeIntegration:
@@ -173,26 +179,18 @@ def test_cli_mode_passes_json_schema_to_ollama(self, tmp_path):
         When json_schema is supplied, chat_with_model should forward the parsed
         schema (as a dict) to chat_with_ollama.
         """
-        test_schema = {"schema": {"type":"object","properties":{"answer":{"type":"string"}}}}
+        test_schema = {"schema": {"type": "object", "properties": {"answer": {"type": "string"}}}}
 
         #  Prepare mocks
-        with patch('ai_server.server.is_llamacpp_available', return_value=False), \
-             patch('ai_server.server.chat_with_ollama') as mock_ollama:
+        with patch('ai_server.server.is_llamacpp_available', return_value=False), patch(
+            'ai_server.server.chat_with_ollama'
+        ) as mock_ollama:
             mock_ollama.return_value = "schema-aware response"
 
-            result = chat_with_model(
-                TEST_OLLAMA_MODEL,
-                "Give me an answer",
-                llama_mode='cli',
-                json_schema=test_schema
-            )
+            result = chat_with_model(TEST_OLLAMA_MODEL, "Give me an answer", llama_mode='cli', json_schema=test_schema)
 
             assert result == "schema-aware response"
 
             mock_ollama.assert_called_once_with(
-                TEST_OLLAMA_MODEL,
-                "Give me an answer",
-                system_prompt=None,
-                image_files=None,
-                json_schema=test_schema
+                TEST_OLLAMA_MODEL, "Give me an answer", system_prompt=None, image_files=None, json_schema=test_schema
             )

test/test_core.py

@@ -83,7 +83,7 @@ def test_chat_with_ollama_success(self, mock_ollama):
             model=TEST_OLLAMA_MODEL,
             messages=[{'role': 'user', 'content': 'Help me write a Python function'}],
             stream=False,
-            format=None
+            format=None,
         )
 
     def test_chat_with_ollama_service_unavailable(self, mock_ollama):
@@ -103,18 +103,14 @@ def test_chat_with_ollama_model_not_found(self, mock_ollama):
     def test_chat_with_ollama_with_json_schema(self, mock_ollama, tmp_path):
         """Ollama chat should forward the JSON schema (format=…) when provided."""
         # Fake schema file
-        test_schema = {"schema": {"type":"object","properties":{"answer":{"type":"string"}}}}
+        test_schema = {"schema": {"type": "object", "properties": {"answer": {"type": "string"}}}}
 
         # Mock ollama response
         mock_response = MagicMock()
         mock_response.message.content = "42"
         mock_ollama.return_value = mock_response
 
-        result = chat_with_ollama(
-            TEST_OLLAMA_MODEL,
-            "What is the meaning of life?",
-            json_schema=test_schema
-        )
+        result = chat_with_ollama(TEST_OLLAMA_MODEL, "What is the meaning of life?", json_schema=test_schema)
 
         assert result == "42"
         mock_ollama.assert_called_once_with(

test/test_server_mode.py

@@ -111,7 +111,9 @@ def test_server_mode_uses_llamacpp_when_available(self):
 
         assert result == "Server response from DeepSeek V3"
         self.mock_available.assert_called_once_with(TEST_LLAMACPP_MODEL)
-        self.mock_chat_server.assert_called_once_with(TEST_LLAMACPP_MODEL, 'Explain code', system_prompt=None, image_files=None, json_schema=None)
+        self.mock_chat_server.assert_called_once_with(
+            TEST_LLAMACPP_MODEL, 'Explain code', system_prompt=None, image_files=None, json_schema=None
+        )
 
     def test_server_mode_fallback_to_ollama_when_unavailable(self):
         """Test server mode falls back to ollama when model not available in llama.cpp."""
@@ -122,7 +124,9 @@ def test_server_mode_fallback_to_ollama_when_unavailable(self):
 
         assert result == "Ollama fallback response"
         self.mock_available.assert_called_once_with(TEST_OLLAMA_MODEL)
-        self.mock_chat_ollama.assert_called_once_with(TEST_OLLAMA_MODEL, 'Debug code', system_prompt=None, image_files=None, json_schema=None)
+        self.mock_chat_ollama.assert_called_once_with(
+            TEST_OLLAMA_MODEL, 'Debug code', system_prompt=None, image_files=None, json_schema=None
+        )
 
     def test_server_mode_requires_server_url(self):
         """Test server mode requires LLAMA_SERVER_URL to be set."""
@@ -184,21 +188,16 @@ def test_server_mode_passes_json_schema_to_llama_server(self, tmp_path, mock_req
         chat_with_model (server mode) should forward a json_schema file path
         and llama-server should receive the parsed schema in its JSON body.
         """
-        test_schema = {"schema": {"type":"object","properties":{"answer":{"type":"string"}}}}
+        test_schema = {"schema": {"type": "object", "properties": {"answer": {"type": "string"}}}}
 
         with patch('ai_server.server.is_llamacpp_available', return_value=True):
             mock_response = MagicMock()
             mock_response.status_code = 200
-            mock_response.json.return_value = {
-                "choices": [{"message": {"content": "Schema-aware server reply"}}]
-            }
+            mock_response.json.return_value = {"choices": [{"message": {"content": "Schema-aware server reply"}}]}
             mock_requests_post.return_value = mock_response
 
             result = chat_with_model(
-                TEST_LLAMACPP_MODEL,
-                "Give me an answer",
-                llama_mode="server",
-                json_schema=test_schema
+                TEST_LLAMACPP_MODEL, "Give me an answer", llama_mode="server", json_schema=test_schema
             )
 
             assert result == "Schema-aware server reply"
@@ -213,4 +212,3 @@ def test_server_mode_passes_json_schema_to_llama_server(self, tmp_path, mock_req
             assert body["model"] == TEST_LLAMACPP_MODEL
             assert body["messages"][0]["content"] == "Give me an answer"
             assert body["json_schema"] == {"type": "object", "properties": {"answer": {"type": "string"}}}
-

test/test_system_prompt.py

@@ -77,4 +77,6 @@ def test_chat_with_model_routing(self, mock_available, mock_chat):
         mock_chat.return_value = "result"
 
         chat_with_model(TEST_MODEL, TEST_USER_CONTENT, 'cli', TEST_SYSTEM_PROMPT)
-        mock_chat.assert_called_once_with(TEST_MODEL, TEST_USER_CONTENT, system_prompt=TEST_SYSTEM_PROMPT, image_files=None, json_schema=None)
+        mock_chat.assert_called_once_with(
+            TEST_MODEL, TEST_USER_CONTENT, system_prompt=TEST_SYSTEM_PROMPT, image_files=None, json_schema=None
+        )