fix: fix embed paths

This commit is contained in:
George Panchuk
2026-07-23 00:05:33 +07:00
parent f326784879
commit e1f2789338
4 changed files with 34 additions and 1 deletions

View File

@@ -235,7 +235,11 @@ class ModelEmbedder:
current_model = getattr(item, path.current, None)
if current_model is None:
continue
if path.tail:
# a node with a tail can still hold an inference object itself,
# e.g. `prefetch.query` can be a Document, while `prefetch.query.nearest`
# can be a Document as well. If an inference object is found at the current
# node, treat it as a leaf, otherwise continue traversing the tail.
if path.tail and not isinstance(current_model, get_args(INFERENCE_OBJECT_TYPES)):
self._process_model(
current_model,
path.tail,

View File

@@ -7,11 +7,16 @@ from pydantic import BaseModel, Field
class FieldPath(BaseModel):
current: str
tail: list["FieldPath"] | None = Field(default=None)
# marks the node as a valid endpoint, e.g. in ["a.b", "a.b.c"] both "b" and "c" are endpoints.
# Nodes without a tail are always endpoints, regardless of the flag value.
leaf: bool = Field(default=False)
def as_str_list(self) -> list[str]:
"""
>>> FieldPath(current='a', tail=[FieldPath(current='b', tail=[FieldPath(current='c'), FieldPath(current='d')])]).as_str_list()
['a.b.c', 'a.b.d']
>>> FieldPath(current='a', tail=[FieldPath(current='b', leaf=True, tail=[FieldPath(current='c')])]).as_str_list()
['a.b', 'a.b.c']
"""
# Recursive function to collect all paths
@@ -21,6 +26,8 @@ class FieldPath(BaseModel):
return [current_path]
else:
paths = []
if path.leaf:
paths.append(current_path)
for sub_path in path.tail:
paths.extend(collect_paths(sub_path, current_path + "."))
return paths
@@ -64,6 +71,7 @@ def convert_paths(paths: list[str]) -> list[FieldPath]:
assert current.tail is not None
current.tail.append(new_tail)
current = new_tail
current.leaf = True # the last node of each path is a valid endpoint
return converted_paths

View File

@@ -391,6 +391,22 @@ def test_inspect_query_requests():
paths = inspector_embed.inspect([document_only_prefetch_request])
assert len(paths) == 1 and paths[0].as_str_list() == ["prefetch.query"]
# mixed direct Document and nested NearestQuery — both prefix paths must be kept
mixed_prefix_prefetch_request = models.QueryRequest(
query=doc,
prefetch=[
models.Prefetch(query=doc),
models.Prefetch(query=models.NearestQuery(nearest=doc)),
],
)
assert inspector.inspect(mixed_prefix_prefetch_request)
paths = inspector_embed.inspect(mixed_prefix_prefetch_request)
assert {path for field_path in paths for path in field_path.as_str_list()} == {
"query",
"prefetch.query",
"prefetch.query.nearest",
}
assert inspector.inspect([query_request_vector, document_only_query_request])
paths = inspector_embed.inspect([query_request_vector, document_only_query_request])
assert len(paths) == 1 and paths[0].as_str_list() == ["query"]

View File

@@ -51,3 +51,8 @@ def test_parser(model):
count += check_path_recursive(plain_path.split("."), paths)
assert count == len(plain_paths)
# convert_paths must keep every plain path, including ones that are
# prefixes of longer paths (e.g. "query" and "query.nearest")
flattened = {path for field_path in paths for path in field_path.as_str_list()}
assert set(plain_paths) == flattened