Dont lose the complete name suggestions, do it as a second step

This commit is contained in:
shamoon
2026-08-30 10:28:55 -07:00
parent b20d707f2a
commit 699d856286
9 changed files with 270 additions and 156 deletions
+10 -15
View File
@@ -67,7 +67,6 @@ def build_prompt_without_rag(
document: Document,
config: AIConfig,
candidates: TaxonomyCandidates | None = None,
assigned: AssignedMetadata | None = None,
) -> str:
filename = document.filename or ""
content = truncate_content(
@@ -77,9 +76,7 @@ def build_prompt_without_rag(
)
taxonomy_block = (
format_taxonomy_for_prompt(candidates, assigned)
if candidates is not None and assigned is not None
else ""
format_taxonomy_for_prompt(candidates) if candidates is not None else ""
)
has_candidates = candidates is not None and any(candidates.values())
@@ -97,14 +94,12 @@ def build_prompt_with_rag(
document: Document,
config: AIConfig,
candidates: TaxonomyCandidates | None = None,
assigned: AssignedMetadata | None = None,
context: str = "",
) -> str:
base_prompt = build_prompt_without_rag(
document,
config,
candidates=candidates,
assigned=assigned,
)
truncated_context = truncate_content(
context,
@@ -277,28 +272,28 @@ def get_ai_document_classification(
ai_config = AIConfig()
if ai_config.llm_embedding_backend:
candidates, assigned, context = get_taxonomy_context(document, user)
candidates, _assigned, context = get_taxonomy_context(document, user)
prompt = build_prompt_with_rag(
document,
ai_config,
candidates=candidates,
assigned=assigned,
context=context,
)
else:
candidates = empty_taxonomy_candidates()
prompt = build_prompt_without_rag(
document,
ai_config,
candidates=candidates,
assigned=get_assigned_metadata(document, user),
)
prompt = build_prompt_without_rag(document, ai_config, candidates=candidates)
client = AIClient()
# Hand the pooled DB connection back while the (slow) LLM query runs so it
# is not pinned for the call's duration; see paperless_ai.db and #12976.
with db_connection_released():
result = client.run_llm_query(prompt)
result = client.run_llm_query(
prompt,
allowed_candidate_ids={
category: {candidate["id"] for candidate in values}
for category, values in candidates.items()
},
)
suggestions = _restrict_to_shown_candidates(
parse_ai_response(result),
candidates,
+120 -41
View File
@@ -48,75 +48,107 @@ class DocumentClassifierSchema(BaseModel):
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names of topic labels describing what this document is about, "
"e.g. 'Insurance', 'Car', 'Warranty'. When an available tag "
"represents the same label, use its ID in tag_ids instead."
"All topic labels you would suggest from the document itself, e.g. "
"'Insurance', 'Car', 'Warranty'. Always include every suggested "
"name here, even when it matches an available tag."
),
)
matched_tags: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names copied exactly from tags that mean the same thing as an "
"available tag. Align each name by position with tag_ids."
),
)
tag_ids: list[int] = Field(
default_factory=list,
max_length=MAX_EXISTING_IDS,
description=(
"IDs of available tags that clearly apply to this document. Only "
"use IDs shown in the prompt; never invent one or choose a weak match."
"Available tag IDs matching matched_tags, in the same order. "
"Only use IDs shown in the prompt."
),
)
correspondents: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names of people, institutions or companies this document is from "
"or was sent to, not every party merely mentioned. When an "
"available correspondent is the same entity, use its ID in "
"correspondent_ids instead."
"All people, institutions or companies you would suggest as who "
"this document is from or was sent to, not every party merely "
"mentioned. Always include every suggested name here, even when it "
"matches an available correspondent."
),
)
matched_correspondents: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names copied exactly from correspondents that identify the same "
"entity as an available correspondent. Align each name by position "
"with correspondent_ids."
),
)
correspondent_ids: list[int] = Field(
default_factory=list,
max_length=MAX_EXISTING_IDS,
description=(
"IDs of available correspondents that clearly apply to this "
"document. Only use IDs shown in the prompt; never invent one or "
"choose a weak match."
"Available correspondent IDs matching matched_correspondents, in "
"the same order. Only use IDs shown in the prompt."
),
)
document_types: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names describing what kind of document this is, e.g. 'Invoice', "
"All names describing what kind of document this is, e.g. 'Invoice', "
"'Contract', 'Bank Statement', 'Letter'. Never use its subject or "
"sender as a document type. When an available document type is the "
"same kind, use its ID in document_type_ids instead."
"sender as a document type. Always include every suggested name "
"here, even when it matches an available document type."
),
)
matched_document_types: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names copied exactly from document_types that mean the same thing "
"as an available document type. Align each name by position with "
"document_type_ids."
),
)
document_type_ids: list[int] = Field(
default_factory=list,
max_length=MAX_EXISTING_IDS,
description=(
"IDs of available document types that clearly apply to this "
"document. Only use IDs shown in the prompt; never invent one or "
"choose a weak match."
"Available document type IDs matching matched_document_types, in "
"the same order. Only use IDs shown in the prompt."
),
)
storage_paths: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names of folder-style filing locations, e.g. "
"All folder-style filing locations you would suggest, e.g. "
"'Finance/Invoices'. Leave empty unless a filing location is "
"clearly implied - never put tags, document types or "
"correspondents here. When an available storage path is the same "
"location, use its ID in storage_path_ids instead."
"correspondents here. Always include every suggested name here, "
"even when it matches an available storage path."
),
)
matched_storage_paths: list[str] = Field(
default_factory=list,
max_length=MAX_NEW_NAMES,
description=(
"Names copied exactly from storage_paths that mean the same filing "
"location as an available storage path. Align each name by position "
"with storage_path_ids."
),
)
storage_path_ids: list[int] = Field(
default_factory=list,
max_length=MAX_EXISTING_IDS,
description=(
"IDs of available storage paths that clearly apply to this "
"document. Only use IDs shown in the prompt; never invent one or "
"choose a weak match."
"Available storage path IDs matching matched_storage_paths, in the "
"same order. Only use IDs shown in the prompt."
),
)
dates: list[str] = Field(
@@ -132,12 +164,16 @@ class DocumentClassifierSchema(BaseModel):
@field_validator(
"title",
"tags",
"matched_tags",
"tag_ids",
"correspondents",
"matched_correspondents",
"correspondent_ids",
"document_types",
"matched_document_types",
"document_type_ids",
"storage_paths",
"matched_storage_paths",
"storage_path_ids",
"dates",
mode="before",
@@ -167,25 +203,64 @@ class ClassificationSuggestions(TypedDict):
def model_to_classification_suggestions(
model: DocumentClassifierSchema,
allowed_candidate_ids: dict[str, set[int]] | None = None,
) -> ClassificationSuggestions:
"""Convert the flat, model-friendly response to the internal shape."""
"""Validate optional candidate mappings and convert to the internal shape.
A mapping is accepted only when its name is copied from the model's own
complete suggestion list and its ID was actually shown for that category.
Invalid or unpaired mappings leave the original name untouched.
"""
allowed_candidate_ids = allowed_candidate_ids or {}
def _choice(
names: list[str],
matched_names: list[str],
ids: list[int],
category: str,
) -> TaxonomyChoiceDict:
remaining_names = list(names)
existing_ids: list[int] = []
allowed_ids = allowed_candidate_ids.get(category, set())
for name, object_id in zip(matched_names, ids, strict=False):
if (
name not in remaining_names
or object_id not in allowed_ids
or object_id in existing_ids
):
continue
remaining_names.remove(name)
existing_ids.append(object_id)
return TaxonomyChoiceDict(
existing_ids=existing_ids,
new_names=remaining_names,
)
return ClassificationSuggestions(
title=model.title,
tags=TaxonomyChoiceDict(
existing_ids=model.tag_ids,
new_names=model.tags,
tags=_choice(
model.tags,
model.matched_tags,
model.tag_ids,
"tags",
),
correspondents=TaxonomyChoiceDict(
existing_ids=model.correspondent_ids,
new_names=model.correspondents,
correspondents=_choice(
model.correspondents,
model.matched_correspondents,
model.correspondent_ids,
"correspondents",
),
document_types=TaxonomyChoiceDict(
existing_ids=model.document_type_ids,
new_names=model.document_types,
document_types=_choice(
model.document_types,
model.matched_document_types,
model.document_type_ids,
"document_types",
),
storage_paths=TaxonomyChoiceDict(
existing_ids=model.storage_path_ids,
new_names=model.storage_paths,
storage_paths=_choice(
model.storage_paths,
model.matched_storage_paths,
model.storage_path_ids,
"storage_paths",
),
dates=model.dates,
)
@@ -198,12 +273,16 @@ def classification_suggestions_to_model(
return DocumentClassifierSchema(
title=suggestions["title"],
tags=suggestions["tags"]["new_names"],
tag_ids=suggestions["tags"]["existing_ids"],
matched_tags=[],
tag_ids=[],
correspondents=suggestions["correspondents"]["new_names"],
correspondent_ids=suggestions["correspondents"]["existing_ids"],
matched_correspondents=[],
correspondent_ids=[],
document_types=suggestions["document_types"]["new_names"],
document_type_ids=suggestions["document_types"]["existing_ids"],
matched_document_types=[],
document_type_ids=[],
storage_paths=suggestions["storage_paths"]["new_names"],
storage_path_ids=suggestions["storage_paths"]["existing_ids"],
matched_storage_paths=[],
storage_path_ids=[],
dates=suggestions["dates"],
)
+14 -3
View File
@@ -117,7 +117,12 @@ class AIClient:
else:
raise ValueError(f"Unsupported LLM backend: {self.settings.llm_backend}")
def run_llm_query(self, prompt: str) -> ClassificationSuggestions:
def run_llm_query(
self,
prompt: str,
*,
allowed_candidate_ids: dict[str, set[int]] | None = None,
) -> ClassificationSuggestions:
logger.debug(
"Running LLM query against %s with model %s",
self.settings.llm_backend,
@@ -136,7 +141,10 @@ class AIClient:
)
logger.debug("LLM query result: %s", result)
parsed = DocumentClassifierSchema(**json.loads(result.message.content))
return model_to_classification_suggestions(parsed)
return model_to_classification_suggestions(
parsed,
allowed_candidate_ids,
)
from llama_index.core.program.function_program import get_function_tool
@@ -155,7 +163,10 @@ class AIClient:
)
logger.debug("LLM query result: %s", tool_calls)
parsed = DocumentClassifierSchema(**tool_calls[0].tool_kwargs)
return model_to_classification_suggestions(parsed)
return model_to_classification_suggestions(
parsed,
allowed_candidate_ids,
)
@contextmanager
def _normalize_timeouts(self) -> Iterator[None]:
+1 -4
View File
@@ -13,10 +13,7 @@ Analyze the following document and fill in these fields:
- dates: up to 3 relevant dates in YYYY-MM-DD format
{% if has_candidates %}
For tags, correspondents, document types, and storage paths: first decide whether there is a useful, well-supported suggestion. If an available candidate clearly represents that suggestion, put its id in the matching tag_ids, correspondent_ids, document_type_ids, or storage_path_ids field instead of repeating its name. If no candidate represents the suggestion, put its name in tags, correspondents, document_types, or storage_paths. Do not choose a weak candidate merely because it exists.
{% else %}
No candidates are shown for this document, so leave every field ending in "_ids" empty and put suggestions in the corresponding name fields.
First produce the complete name suggestions from the document itself in tags, correspondents, document_types, and storage_paths. Always include every suggested name in those fields, even when an available candidate represents the same value. Then, as a separate reconciliation step, copy each name that means the same thing as an available candidate into the matching matched_* field and put that candidate's ID at the same position in the corresponding *_ids field. Candidates must not create, replace, or suppress suggestions. Do not match a candidate that is merely related.
{% endif %}
Filename:
+1 -1
View File
@@ -1,6 +1,6 @@
You are localizing document classification suggestions for display in Paperless-ngx.
Rewrite only the "title", "tags", "document_types", and "storage_paths" fields in {{ language_name }}. Leave every field ending in "_ids" exactly as given - these are database identifiers, not text, and are not used from your response even if changed.
Rewrite only the "title", "tags", "document_types", and "storage_paths" fields in {{ language_name }}.
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.
+7 -3
View File
@@ -245,7 +245,7 @@ def _assigned_block(assigned: AssignedMetadata) -> str:
def format_taxonomy_for_prompt(
candidates: TaxonomyCandidates,
assigned: AssignedMetadata,
assigned: AssignedMetadata | None = None,
) -> str:
"""Render assigned metadata and ranked candidates as labelled prompt
blocks. Candidate names are untrusted, user-controlled data, so they are
@@ -255,7 +255,7 @@ def format_taxonomy_for_prompt(
is nothing to say (no assigned metadata and no candidates), so callers can
treat the result the same as no hints at all.
"""
has_assigned = any(
has_assigned = assigned is not None and any(
[
assigned["tags"],
assigned["document_type"],
@@ -271,7 +271,11 @@ def format_taxonomy_for_prompt(
return render_prompt(
TaxonomyBlockPromptContext(
assigned_block=_assigned_block(assigned) if has_assigned else "",
assigned_block=(
_assigned_block(assigned)
if assigned is not None and has_assigned
else ""
),
candidate_payload_json=(
json.dumps(candidate_payload, ensure_ascii=False)
if candidate_payload
+23 -61
View File
@@ -606,10 +606,11 @@ def test_build_prompt_without_rag_includes_taxonomy_block():
GIVEN:
- Non-empty taxonomy candidates
WHEN:
- build_prompt_without_rag() is called with candidates and assigned metadata
- build_prompt_without_rag() is called with candidates
THEN:
- The candidate's id and the flat name/ID instructions appear
- Candidates are presented as deduplication options, not requirements
- The candidate and single-call reconciliation instructions appear
- Complete name suggestions remain mandatory
- Assigned metadata is not included
"""
document = DocumentFactory.create(content="Some content")
config = AIConfig()
@@ -619,40 +620,31 @@ def test_build_prompt_without_rag_includes_taxonomy_block():
"correspondents": [],
"storage_paths": [],
}
assigned = {
"tags": [],
"document_type": None,
"correspondent": None,
"storage_path": None,
}
prompt = build_prompt_without_rag(
document,
config,
candidates=candidates,
assigned=assigned,
)
assert '"id": 12' in prompt
assert "tag_ids" in prompt
assert "correspondent_ids" in prompt
assert "not requirements" in prompt
assert "weak candidate" in prompt
assert "Always include every suggested name" in prompt
assert "matched_*" in prompt
assert "corresponding *_ids" in prompt
assert "Candidates must not create, replace, or suppress suggestions" in prompt
assert "already assigned" not in prompt
@pytest.mark.django_db
def test_build_prompt_without_rag_identical_when_no_hints():
def test_build_prompt_without_rag_identical_when_no_candidates():
"""
GIVEN:
- Empty taxonomy candidates and empty assigned metadata
- Empty taxonomy candidates
WHEN:
- build_prompt_without_rag() is called with those empty values, and
separately with no candidates/assigned at all
separately with no candidates at all
THEN:
- Both prompts are identical
- Neither carries the "Available ..." candidate block or the
id-vs-name routing instruction
- Both still tell the model to leave every ID field empty
- Neither carries candidate reconciliation instructions
"""
document = DocumentFactory.create(content="Some content")
config = AIConfig()
@@ -662,67 +654,37 @@ def test_build_prompt_without_rag_identical_when_no_hints():
"correspondents": [],
"storage_paths": [],
}
empty_assigned = {
"tags": [],
"document_type": None,
"correspondent": None,
"storage_path": None,
}
with_empty_hints = build_prompt_without_rag(
document,
config,
candidates=empty_candidates,
assigned=empty_assigned,
)
with_no_hints = build_prompt_without_rag(document, config)
assert with_empty_hints == with_no_hints
assert "Available " not in with_no_hints
assert "put its id in the matching" not in with_no_hints
assert 'leave every field ending in "_ids" empty' in with_no_hints
assert "matched_*" not in with_no_hints
@pytest.mark.django_db
def test_build_prompt_without_rag_tells_model_to_skip_ids_when_no_candidates():
def test_build_prompt_without_rag_never_includes_assigned_metadata():
"""
GIVEN:
- Assigned metadata but empty taxonomy candidates
- A document with assigned taxonomy metadata
WHEN:
- build_prompt_without_rag() is called with candidates and assigned metadata
- build_prompt_without_rag() is called
THEN:
- The assigned-metadata block appears (taxonomy_block is non-empty)
- The prompt tells the model to leave every ID field empty
Staying silent about ID fields here is not enough: the response schema
advertises the field whatever the prompt says, and models fill it with
placeholder ids that resolve to real but unrelated objects (#13831).
- Assigned metadata is absent so it cannot anchor classification
"""
document = DocumentFactory.create(content="Some content")
config = AIConfig()
empty_candidates = {
"tags": [],
"document_types": [],
"correspondents": [],
"storage_paths": [],
}
assigned = {
"tags": ["Bloodwork"],
"document_type": None,
"correspondent": None,
"storage_path": None,
}
assigned_tag = TagFactory.create(name="Bloodwork")
document.tags.add(assigned_tag)
prompt = build_prompt_without_rag(
document,
config,
candidates=empty_candidates,
assigned=assigned,
)
prompt = build_prompt_without_rag(document, config)
assert "already assigned" in prompt
assert "No candidates are shown" in prompt
assert 'leave every field ending in "_ids" empty' in prompt
assert "Bloodwork" not in prompt
assert "already assigned" not in prompt
@pytest.mark.django_db
+82 -18
View File
@@ -1,7 +1,6 @@
import json
from paperless_ai.base_model import MAX_DATES
from paperless_ai.base_model import MAX_EXISTING_IDS
from paperless_ai.base_model import MAX_NEW_NAMES
from paperless_ai.base_model import MAX_TITLE_LENGTH
from paperless_ai.base_model import ClassificationSuggestions
@@ -19,7 +18,7 @@ def test_document_classifier_schema_declared_defaults():
WHEN:
- The schema is dumped to a dict via model_dump()
THEN:
- Every name and ID field, and dates, dump as empty lists
- Every optional name field, and dates, dump as empty lists
The model may omit optional fields, so the schema must provide the complete
empty shape expected by the conversion and matching pipeline.
@@ -31,57 +30,106 @@ def test_document_classifier_schema_declared_defaults():
assert dumped == {
"title": "Test Title",
"tags": [],
"matched_tags": [],
"tag_ids": [],
"correspondents": [],
"matched_correspondents": [],
"correspondent_ids": [],
"document_types": [],
"matched_document_types": [],
"document_type_ids": [],
"storage_paths": [],
"matched_storage_paths": [],
"storage_path_ids": [],
"dates": [],
}
def test_flat_model_response_converts_to_internal_taxonomy_choices():
def test_model_response_converts_names_to_internal_taxonomy_choices():
"""
GIVEN:
- A flat model response with separate name and candidate-ID fields
- A model response containing taxonomy names
WHEN:
- It is converted to Paperless' internal suggestion representation
THEN:
- Names and IDs are paired under their taxonomy category
- Names enter the internal taxonomy representation as new names
- Existing IDs remain empty for deterministic application-side matching
"""
parsed = DocumentClassifierSchema(
title="Electricity Bill",
tags=["Utilities", "Electricity"],
tag_ids=[12],
correspondents=["Power Company"],
correspondent_ids=[23],
document_types=["Utility Bill"],
document_type_ids=[34],
storage_paths=["Finance/Utilities"],
storage_path_ids=[45],
)
suggestions = model_to_classification_suggestions(parsed)
assert suggestions["tags"] == {
"existing_ids": [12],
"existing_ids": [],
"new_names": ["Utilities", "Electricity"],
}
assert suggestions["correspondents"] == {
"existing_ids": [23],
"existing_ids": [],
"new_names": ["Power Company"],
}
assert suggestions["document_types"] == {
"existing_ids": [34],
"existing_ids": [],
"new_names": ["Utility Bill"],
}
assert suggestions["storage_paths"] == {
"existing_ids": [45],
"existing_ids": [],
"new_names": ["Finance/Utilities"],
}
def test_valid_candidate_mappings_replace_only_the_matched_names():
parsed = DocumentClassifierSchema(
title="Electricity Bill",
tags=["Utilities", "Electricity"],
matched_tags=["Utilities"],
tag_ids=[12],
correspondents=["Power Company"],
matched_correspondents=["Power Company"],
correspondent_ids=[23],
)
suggestions = model_to_classification_suggestions(
parsed,
{
"tags": {12},
"correspondents": {23},
},
)
assert suggestions["tags"] == {
"existing_ids": [12],
"new_names": ["Electricity"],
}
assert suggestions["correspondents"] == {
"existing_ids": [23],
"new_names": [],
}
def test_invalid_or_unpaired_candidate_mappings_do_not_remove_names():
parsed = DocumentClassifierSchema(
title="Electricity Bill",
tags=["Utilities", "Electricity", "Energy"],
matched_tags=["Invented", "Utilities", "Electricity"],
tag_ids=[12, 999],
)
suggestions = model_to_classification_suggestions(
parsed,
{"tags": {12}},
)
assert suggestions["tags"] == {
"existing_ids": [],
"new_names": ["Utilities", "Electricity", "Energy"],
}
def test_document_classifier_schema_json_schema_is_self_contained():
"""
GIVEN:
@@ -175,13 +223,11 @@ def test_over_long_response_is_truncated_rather_than_rejected():
parsed = DocumentClassifierSchema(
title="T" * (MAX_TITLE_LENGTH + 50),
tags=["n"] * (MAX_NEW_NAMES + 20),
tag_ids=list(range(MAX_EXISTING_IDS + 20)),
dates=[f"2016-{month:02d}-01" for month in range(1, 13)],
)
assert len(parsed.title) == MAX_TITLE_LENGTH
assert len(parsed.dates) == MAX_DATES
assert len(parsed.tag_ids) == MAX_EXISTING_IDS
assert len(parsed.tags) == MAX_NEW_NAMES
@@ -214,7 +260,7 @@ def test_model_conversion_matches_internal_typed_dict_keys():
- The converted tags dict's keys exactly match TaxonomyChoiceDict's
declared keys
"""
schema = DocumentClassifierSchema(title="T", tags=["Tag"], tag_ids=[1])
schema = DocumentClassifierSchema(title="T", tags=["Tag"])
suggestions = model_to_classification_suggestions(schema)
assert set(suggestions.keys()) == set(
@@ -225,7 +271,7 @@ def test_model_conversion_matches_internal_typed_dict_keys():
)
def test_internal_suggestions_round_trip_through_flat_model():
def test_internal_suggestions_convert_to_names_only_model():
suggestions = ClassificationSuggestions(
title="Electricity Bill",
tags=TaxonomyChoiceDict(existing_ids=[1], new_names=["Utilities"]),
@@ -246,4 +292,22 @@ def test_internal_suggestions_round_trip_through_flat_model():
model = classification_suggestions_to_model(suggestions)
assert model_to_classification_suggestions(model) == suggestions
converted = model_to_classification_suggestions(model)
assert converted == ClassificationSuggestions(
title="Electricity Bill",
tags=TaxonomyChoiceDict(existing_ids=[], new_names=["Utilities"]),
correspondents=TaxonomyChoiceDict(
existing_ids=[],
new_names=["Power Company"],
),
document_types=TaxonomyChoiceDict(
existing_ids=[],
new_names=["Utility Bill"],
),
storage_paths=TaxonomyChoiceDict(
existing_ids=[],
new_names=["Finance/Utilities"],
),
dates=["2026-08-30"],
)
+12 -10
View File
@@ -124,22 +124,23 @@ def test_run_llm_query_ollama_uses_structured_json(mock_ai_config, mock_ollama_l
{
"title": "Test Title",
"tags": ["document"],
"matched_tags": ["document"],
"tag_ids": [1],
"correspondents": ["John Doe"],
"correspondent_ids": [],
"document_types": ["report"],
"document_type_ids": [],
"storage_paths": ["Reports"],
"storage_path_ids": [],
"dates": ["2023-01-01"],
},
)
client = AIClient()
result = client.run_llm_query("test_prompt")
result = client.run_llm_query(
"test_prompt",
allowed_candidate_ids={"tags": {1}},
)
assert result["title"] == "Test Title"
assert result["tags"] == {"existing_ids": [1], "new_names": ["document"]}
assert result["tags"] == {"existing_ids": [1], "new_names": []}
mock_llm_instance.chat.assert_called_once_with(
[ANY],
format=ANY,
@@ -161,13 +162,11 @@ def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm):
tool_kwargs={
"title": "Test Title",
"tags": ["document"],
"matched_tags": ["document"],
"tag_ids": [1],
"correspondents": ["John Doe"],
"correspondent_ids": [],
"document_types": ["report"],
"document_type_ids": [],
"storage_paths": ["Reports"],
"storage_path_ids": [],
"dates": ["2023-01-01"],
},
)
@@ -176,10 +175,13 @@ def test_run_llm_query_openai_uses_tools(mock_ai_config, mock_openai_llm):
mock_llm_instance.get_tool_calls_from_response.return_value = [tool_selection]
client = AIClient()
result = client.run_llm_query("test_prompt")
result = client.run_llm_query(
"test_prompt",
allowed_candidate_ids={"tags": {1}},
)
assert result["title"] == "Test Title"
assert result["tags"] == {"existing_ids": [1], "new_names": ["document"]}
assert result["tags"] == {"existing_ids": [1], "new_names": []}
mock_llm_instance.chat_with_tools.assert_called_once()