diff --git a/src/paperless_ai/client.py b/src/paperless_ai/client.py index 2ca5f0eea..b14f4a726 100644 --- a/src/paperless_ai/client.py +++ b/src/paperless_ai/client.py @@ -131,11 +131,10 @@ class AIClient: from llama_index.core.llms import ChatMessage - user_msg = ChatMessage(role="user", content=prompt) if self.settings.llm_backend == LLMBackend.OLLAMA: with self._normalize_timeouts(): result = self.llm.chat( - [user_msg], + [ChatMessage(role="user", content=prompt)], format=DocumentClassifierSchema.model_json_schema(), think=False, ) @@ -149,6 +148,11 @@ class AIClient: from llama_index.core.program.function_program import get_function_tool tool = get_function_tool(DocumentClassifierSchema) + user_msg = ChatMessage( + role="user", + content=f"{prompt}\n\n" + f"Answer by calling the {tool.metadata.name} tool. Do not write the answer as text.", + ) with self._normalize_timeouts(): result = self.llm.chat_with_tools( tools=[tool], diff --git a/src/paperless_ai/prompts/localization.j2 b/src/paperless_ai/prompts/localization.j2 index 72bc27344..be311f807 100644 --- a/src/paperless_ai/prompts/localization.j2 +++ b/src/paperless_ai/prompts/localization.j2 @@ -4,7 +4,7 @@ Rewrite only the "title", "tags", "document_types", and "storage_paths" fields i Do not translate correspondents or dates. Preserve proper nouns, organization names, product names, and exact official document names. Translate generic category words when a {{ language_name }} equivalent exists. -Return the same JSON schema with all fields present. +Keep every entry you were given in those four fields, in the same order, using the original wording where no translation applies. Suggestions: {{ suggestions_json }} diff --git a/src/paperless_ai/tests/test_client.py b/src/paperless_ai/tests/test_client.py index 7ab07c1ea..7568e90b1 100644 --- a/src/paperless_ai/tests/test_client.py +++ b/src/paperless_ai/tests/test_client.py @@ -146,6 +146,8 @@ def test_run_llm_query_ollama_uses_structured_json(mock_ai_config, mock_ollama_l format=ANY, think=False, ) + messages = mock_llm_instance.chat.call_args.args[0] + assert messages[0].content == "test_prompt" def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm): @@ -183,6 +185,13 @@ def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm): assert result["title"] == "Test Title" assert result["tags"] == {"existing_ids": [1], "new_names": []} mock_llm_instance.chat_with_tools.assert_called_once() + kwargs = mock_llm_instance.chat_with_tools.call_args.kwargs + offered_tool_name = kwargs["tools"][0].metadata.name + assert kwargs["user_msg"].content == ( + "test_prompt\n\n" + f"Answer by calling the {offered_tool_name} tool. " + "Do not write the answer as text." + ) def test_run_llm_query_openai_timeout_raises_local_error(