Skip to content

Commit f37e3a0

Browse files
danmichaeljonesDan Jones
danmichaeljones
and
Dan Jones
authored
Allow additional filters to be added deterministically (#66)
* Add search-only API models * Initial search-only client * Add async search-only client * Use old union syntax for old python * Bump version * Lint * Ignore user filters for now * Check collections configured somewhere * Remove user_filters * Export the SearchModeResponse model * Add searchers to query namespace * Add tests for search only client * Lint * Add additional connection headers * Add default values for limit / offset * Better docstrings for search-only methods * Rename methods * Lint * Ruff format * Move serialisation utils upwards to share * Add additional_filter option to collection config * Fix typing * Ruff format * Move tests --------- Co-authored-by: Dan Jones <[email protected]>
1 parent 7b5b2ae commit f37e3a0

File tree

5 files changed

+119
-95
lines changed

5 files changed

+119
-95
lines changed

‎test/personalization_agent/classes/__init__.py‎

Whitespace-only changes.

‎test/personalization_agent/classes/test_query.py‎ renamed to ‎test/test_serialise.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
Move,
1010
)
1111

12-
from weaviate_agents.personalization.classes.query import (
12+
from weaviate_agents.serialise import (
1313
serialise_filter,
1414
serialise_hybrid_vector_type,
1515
serialise_move,

‎weaviate_agents/personalization/classes/query.py‎

Lines changed: 8 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -3,26 +3,23 @@
33
from typing import Annotated, List, Literal, Optional, Union
44
from uuid import UUID
55

6-
from pydantic import BaseModel, ConfigDict, Field, PlainSerializer
7-
from typing_extensions import TypedDict
6+
from pydantic import BaseModel, ConfigDict, Field
87
from weaviate.classes.query import Move, Rerank
9-
from weaviate.collections.classes.filters import (
10-
_FilterAnd,
11-
_FilterOr,
12-
_Filters,
13-
_FilterValue,
14-
)
8+
from weaviate.collections.classes.filters import _Filters
159
from weaviate.collections.classes.grpc import (
1610
METADATA,
1711
HybridFusion,
1812
HybridVectorType,
19-
NearVectorInputType,
2013
TargetVectorJoinType,
21-
_HybridNearText,
22-
_HybridNearVector,
2314
)
2415
from weaviate.collections.classes.internal import ReturnProperties, ReturnReferences
2516

17+
from weaviate_agents.serialise import (
18+
serialise_filter,
19+
serialise_hybrid_vector_type,
20+
serialise_move,
21+
)
22+
2623

2724
class NearTextQueryParameters(BaseModel):
2825
model_config = ConfigDict(arbitrary_types_allowed=True)
@@ -101,84 +98,3 @@ class QueryRequest(BaseModel):
10198
decay_rate: float
10299
overfetch_factor: float
103100
query_parameters: QueryParameters = Field(discriminator="query_method")
104-
105-
106-
class _MoveSerialise(TypedDict):
107-
force: float
108-
objects: Optional[List[str]]
109-
concepts: Optional[List[str]]
110-
111-
112-
@PlainSerializer
113-
def serialise_move(move: Optional[Move]) -> Optional[_MoveSerialise]:
114-
if move is None:
115-
return None
116-
return _MoveSerialise(
117-
force=move.force, objects=move._objects_list, concepts=move._concepts_list
118-
)
119-
120-
121-
class _FilterAndOrSerialise(BaseModel):
122-
model_config = ConfigDict(arbitrary_types_allowed=True)
123-
124-
combine: Literal["and", "or"]
125-
filters: list[_Filters]
126-
127-
128-
@PlainSerializer
129-
def serialise_filter(
130-
filter_and_or_value: Union[_FilterAnd, _FilterOr, _FilterValue],
131-
) -> Union[_FilterValue, _FilterAndOrSerialise]:
132-
if isinstance(filter_and_or_value, _FilterValue):
133-
return filter_and_or_value
134-
135-
if isinstance(filter_and_or_value, _FilterAnd):
136-
combine: Literal["and", "or"] = "and"
137-
elif isinstance(filter_and_or_value, _FilterOr):
138-
combine = "or"
139-
else:
140-
raise TypeError(f"Unknown filter type {type(filter_and_or_value)}")
141-
return _FilterAndOrSerialise(combine=combine, filters=filter_and_or_value.filters)
142-
143-
144-
# A set of models to help serialise HybridVectorType, which is a union of
145-
# NearVectorInputType (itself a union over std types), _HybridNearText and _HybridNearVector.
146-
class _HybridNearTextSerialise(_HybridNearText):
147-
serialised_class: Literal["_HybridNearText"] = "_HybridNearText"
148-
149-
move_to: Annotated[Optional[Move], serialise_move] = None
150-
move_away: Annotated[Optional[Move], serialise_move] = None
151-
152-
153-
def _serialise_hybrid_near_text(model: _HybridNearText) -> _HybridNearTextSerialise:
154-
return _HybridNearTextSerialise.model_validate(model.model_dump())
155-
156-
157-
class _HybridNearVectorSerialise(BaseModel):
158-
serialised_class: Literal["_HybridNearVector"] = "_HybridNearVector"
159-
160-
vector: NearVectorInputType
161-
distance: Optional[float]
162-
certainty: Optional[float]
163-
164-
165-
def _serialise_hybrid_near_vector(
166-
model: _HybridNearVector,
167-
) -> _HybridNearVectorSerialise:
168-
return _HybridNearVectorSerialise(
169-
vector=model.vector,
170-
distance=model.distance,
171-
certainty=model.certainty,
172-
)
173-
174-
175-
@PlainSerializer
176-
def serialise_hybrid_vector_type(
177-
vector: HybridVectorType,
178-
) -> Union[_HybridNearTextSerialise, _HybridNearVectorSerialise, NearVectorInputType]:
179-
if isinstance(vector, _HybridNearText):
180-
return _serialise_hybrid_near_text(vector)
181-
elif isinstance(vector, _HybridNearVector):
182-
return _serialise_hybrid_near_vector(vector)
183-
else:
184-
return vector

‎weaviate_agents/query/classes/collection.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1-
from typing import Union
1+
from typing import Annotated, Union
22

3-
from pydantic import BaseModel
3+
from pydantic import BaseModel, ConfigDict
4+
from weaviate.collections.classes.filters import _Filters
5+
6+
from weaviate_agents.serialise import serialise_filter
47

58

69
class QueryAgentCollectionConfig(BaseModel):
@@ -13,9 +16,14 @@ class QueryAgentCollectionConfig(BaseModel):
1316
for this specific collection.
1417
target_vector: Optional target vector name(s) for collections with named vectors.
1518
Can be a single vector name or a list of vector names.
19+
additional_filters: Optional filters to apply when the query is executed, in addition
20+
to filters selected by the Query Agent (i.e., there are AND combined).
1621
"""
1722

23+
model_config = ConfigDict(arbitrary_types_allowed=True)
24+
1825
name: str
1926
tenant: Union[str, None] = None
2027
view_properties: Union[list[str], None] = None
2128
target_vector: Union[str, list[str], None] = None
29+
additional_filters: Union[Annotated[_Filters, serialise_filter], None] = None

‎weaviate_agents/serialise.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
from __future__ import annotations
2+
3+
from typing import Annotated, List, Literal, Optional, Union
4+
5+
from pydantic import BaseModel, ConfigDict, PlainSerializer
6+
from typing_extensions import TypedDict
7+
from weaviate.classes.query import Move
8+
from weaviate.collections.classes.filters import (
9+
_FilterAnd,
10+
_FilterOr,
11+
_Filters,
12+
_FilterValue,
13+
)
14+
from weaviate.collections.classes.grpc import (
15+
HybridVectorType,
16+
NearVectorInputType,
17+
_HybridNearText,
18+
_HybridNearVector,
19+
)
20+
21+
22+
class _MoveSerialise(TypedDict):
23+
force: float
24+
objects: Optional[List[str]]
25+
concepts: Optional[List[str]]
26+
27+
28+
@PlainSerializer
29+
def serialise_move(move: Optional[Move]) -> Optional[_MoveSerialise]:
30+
if move is None:
31+
return None
32+
return _MoveSerialise(
33+
force=move.force, objects=move._objects_list, concepts=move._concepts_list
34+
)
35+
36+
37+
class _FilterAndOrSerialise(BaseModel):
38+
model_config = ConfigDict(arbitrary_types_allowed=True)
39+
40+
combine: Literal["and", "or"]
41+
filters: list[_Filters]
42+
43+
44+
@PlainSerializer
45+
def serialise_filter(
46+
filter_and_or_value: Union[_FilterAnd, _FilterOr, _FilterValue],
47+
) -> Union[_FilterValue, _FilterAndOrSerialise]:
48+
if isinstance(filter_and_or_value, _FilterValue):
49+
return filter_and_or_value
50+
51+
if isinstance(filter_and_or_value, _FilterAnd):
52+
combine: Literal["and", "or"] = "and"
53+
elif isinstance(filter_and_or_value, _FilterOr):
54+
combine = "or"
55+
else:
56+
raise TypeError(f"Unknown filter type {type(filter_and_or_value)}")
57+
return _FilterAndOrSerialise(combine=combine, filters=filter_and_or_value.filters)
58+
59+
60+
# A set of models to help serialise HybridVectorType, which is a union of
61+
# NearVectorInputType (itself a union over std types), _HybridNearText and _HybridNearVector.
62+
class _HybridNearTextSerialise(_HybridNearText):
63+
serialised_class: Literal["_HybridNearText"] = "_HybridNearText"
64+
65+
move_to: Annotated[Optional[Move], serialise_move] = None
66+
move_away: Annotated[Optional[Move], serialise_move] = None
67+
68+
69+
def _serialise_hybrid_near_text(model: _HybridNearText) -> _HybridNearTextSerialise:
70+
return _HybridNearTextSerialise.model_validate(model.model_dump())
71+
72+
73+
class _HybridNearVectorSerialise(BaseModel):
74+
serialised_class: Literal["_HybridNearVector"] = "_HybridNearVector"
75+
76+
vector: NearVectorInputType
77+
distance: Optional[float]
78+
certainty: Optional[float]
79+
80+
81+
def _serialise_hybrid_near_vector(
82+
model: _HybridNearVector,
83+
) -> _HybridNearVectorSerialise:
84+
return _HybridNearVectorSerialise(
85+
vector=model.vector,
86+
distance=model.distance,
87+
certainty=model.certainty,
88+
)
89+
90+
91+
@PlainSerializer
92+
def serialise_hybrid_vector_type(
93+
vector: HybridVectorType,
94+
) -> Union[_HybridNearTextSerialise, _HybridNearVectorSerialise, NearVectorInputType]:
95+
if isinstance(vector, _HybridNearText):
96+
return _serialise_hybrid_near_text(vector)
97+
elif isinstance(vector, _HybridNearVector):
98+
return _serialise_hybrid_near_vector(vector)
99+
else:
100+
return vector

0 commit comments

Comments
 (0)