diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 64d4276ac..1614ceb8d 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -1,3 +1,4 @@ +import json import logging from typing import TYPE_CHECKING @@ -109,9 +110,20 @@ class AIClient: ) from llama_index.core.llms import ChatMessage - from llama_index.core.program.function_program import get_function_tool user_msg = ChatMessage(role="user", content=prompt) + if self.settings.llm_backend == LLMBackend.OLLAMA: + result = self.llm.chat( + [user_msg], + format=DocumentClassifierSchema.model_json_schema(), + think=False, + ) + logger.debug("LLM query result: %s", result) + parsed = DocumentClassifierSchema(**json.loads(result.message.content)) + return parsed.model_dump() + + from llama_index.core.program.function_program import get_function_tool + tool = get_function_tool(DocumentClassifierSchema) result = self.llm.chat_with_tools( tools=[tool], diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index b2fbf68ff..be660274f 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -1,3 +1,4 @@ +import json from unittest.mock import ANY from unittest.mock import MagicMock from unittest.mock import patch @@ -90,12 +91,42 @@ def test_get_llm_unsupported_backend(mock_ai_config): AIClient() -def test_run_llm_query(mock_ai_config, mock_ollama_llm): +def test_run_llm_query_ollama_uses_structured_json(mock_ai_config, mock_ollama_llm): mock_ai_config.llm_backend = "ollama" mock_ai_config.llm_model = "test_model" mock_ai_config.llm_endpoint = "http://test-url" mock_llm_instance = mock_ollama_llm.return_value + mock_llm_instance.chat.return_value = MagicMock() + mock_llm_instance.chat.return_value.message.content = json.dumps( + { + "title": "Test Title", + "tags": ["test", "document"], + "correspondents": ["John Doe"], + "document_types": ["report"], + "storage_paths": ["Reports"], + "dates": ["2023-01-01"], + }, + ) + + client = AIClient() + result = client.run_llm_query("test_prompt") + + assert result["title"] == "Test Title" + mock_llm_instance.chat.assert_called_once_with( + [ANY], + format=ANY, + think=False, + ) + + +def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm): + mock_ai_config.llm_backend = "openai-like" + mock_ai_config.llm_model = "test_model" + mock_ai_config.llm_api_key = "test_api_key" + mock_ai_config.llm_endpoint = "http://test-url" + + mock_llm_instance = mock_openai_llm.return_value tool_selection = ToolSelection( tool_id="call_test", @@ -117,6 +148,7 @@ def test_run_llm_query(mock_ai_config, mock_ollama_llm): result = client.run_llm_query("test_prompt") assert result["title"] == "Test Title" + mock_llm_instance.chat_with_tools.assert_called_once() def test_run_chat(mock_ai_config, mock_ollama_llm):