mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-08-31 06:57:16 +00:00
Dont lose the complete name suggestions, do it as a second step
This commit is contained in:
@@ -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
@@ -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"],
|
||||
)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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,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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"],
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user