001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.store;
017
018import com.agentsflex.core.store.condition.*;
019
020import java.util.*;
021import java.util.function.Consumer;
022
023public class SearchWrapper extends VectorData {
024
025    /**
026     * the default value of search data count
027     */
028    public static final int DEFAULT_MAX_RESULTS = 4;
029
030    /**
031     * search text, Vector store will convert the text to vector data
032     */
033    private String text;
034
035    /**
036     * search max result, like the sql "limit" in mysql
037     */
038    private Integer maxResults = DEFAULT_MAX_RESULTS;
039
040    /**
041     * The lowest correlation score, ranging from 0 to 1 (including 0 and 1). Only embeddings with a score of this value or higher will be returned.
042     * 0.0 indicates accepting any similarity or disabling similarity threshold filtering. A threshold of 1.0 indicates the need for a perfect match.
043     */
044    private Double minScore;
045
046    /**
047     * The flag of include vector data queries. If the current value is true and the vector content is null,
048     * the query text will be automatically converted into vector data through the vector store.
049     */
050    private boolean withVector = true;
051
052    /**
053     * query condition
054     */
055    private Condition condition;
056
057    /**
058     * query fields
059     */
060    private List<String> outputFields;
061
062    /**
063     * whether to output vector data
064     */
065    private boolean outputVector = false;
066
067
068    public String getText() {
069        return text;
070    }
071
072    public void setText(String text) {
073        this.text = text;
074    }
075
076    public SearchWrapper text(String text) {
077        setText(text);
078        return this;
079    }
080
081    public Integer getMaxResults() {
082        return maxResults;
083    }
084
085    public void setMaxResults(Integer maxResults) {
086        this.maxResults = maxResults;
087    }
088
089    public SearchWrapper maxResults(Integer maxResults) {
090        setMaxResults(maxResults);
091        return this;
092    }
093
094    public Double getMinScore() {
095        return minScore;
096    }
097
098    public void setMinScore(Double minScore) {
099        this.minScore = minScore;
100    }
101
102    public SearchWrapper minScore(Double minScore) {
103        setMinScore(minScore);
104        return this;
105    }
106
107    public boolean isWithVector() {
108        return withVector;
109    }
110
111    public void setWithVector(boolean withVector) {
112        this.withVector = withVector;
113    }
114
115    public SearchWrapper withVector(Boolean withVector) {
116        setWithVector(withVector);
117        return this;
118    }
119
120    public Condition getCondition() {
121        return condition;
122    }
123
124    public void setCondition(Condition condition) {
125        this.condition = condition;
126    }
127
128    public List<String> getOutputFields() {
129        return outputFields;
130    }
131
132    public void setOutputFields(List<String> outputFields) {
133        this.outputFields = outputFields;
134    }
135
136    public SearchWrapper outputFields(Collection<String> outputFields) {
137        setOutputFields(new ArrayList<>(outputFields));
138        return this;
139    }
140
141    public SearchWrapper outputFields(String... outputFields) {
142        setOutputFields(Arrays.asList(outputFields));
143        return this;
144    }
145
146    public boolean isOutputVector() {
147        return outputVector;
148    }
149
150    public void setOutputVector(boolean outputVector) {
151        this.outputVector = outputVector;
152    }
153
154    public SearchWrapper outputVector(boolean outputVector) {
155        setOutputVector(outputVector);
156        return this;
157    }
158
159
160    public SearchWrapper eq(String key, Object value) {
161        return eq(Connector.AND, key, value);
162    }
163
164    public SearchWrapper eq(Connector connector, String key, Object value) {
165        if (this.condition == null) {
166            this.condition = new Condition(ConditionType.EQ, new Key(key), new Value(value));
167        } else {
168            this.condition.connect(new Condition(ConditionType.EQ, new Key(key), new Value(value)), connector);
169        }
170        return this;
171    }
172
173    public SearchWrapper ne(String key, Object value) {
174        return ne(Connector.AND, key, value);
175    }
176
177    public SearchWrapper ne(Connector connector, String key, Object value) {
178        if (this.condition == null) {
179            this.condition = new Condition(ConditionType.NE, new Key(key), new Value(value));
180        } else {
181            this.condition.connect(new Condition(ConditionType.NE, new Key(key), new Value(value)), connector);
182        }
183        return this;
184    }
185
186    public SearchWrapper gt(String key, Object value) {
187        return gt(Connector.AND, key, value);
188    }
189
190    public SearchWrapper gt(Connector connector, String key, Object value) {
191        if (this.condition == null) {
192            this.condition = new Condition(ConditionType.GT, new Key(key), new Value(value));
193        } else {
194            this.condition.connect(new Condition(ConditionType.GT, new Key(key), new Value(value)), connector);
195        }
196        return this;
197    }
198
199
200    public SearchWrapper ge(String key, Object value) {
201        return ge(Connector.AND, key, value);
202    }
203
204    public SearchWrapper ge(Connector connector, String key, Object value) {
205        if (this.condition == null) {
206            this.condition = new Condition(ConditionType.GE, new Key(key), new Value(value));
207        } else {
208            this.condition.connect(new Condition(ConditionType.GE, new Key(key), new Value(value)), connector);
209        }
210        return this;
211    }
212
213
214    public SearchWrapper lt(String key, Object value) {
215        return lt(Connector.AND, key, value);
216    }
217
218    public SearchWrapper lt(Connector connector, String key, Object value) {
219        if (this.condition == null) {
220            this.condition = new Condition(ConditionType.LT, new Key(key), new Value(value));
221        } else {
222            this.condition.connect(new Condition(ConditionType.LT, new Key(key), new Value(value)), connector);
223        }
224        return this;
225    }
226
227
228    public SearchWrapper le(String key, Object value) {
229        return le(Connector.AND, key, value);
230    }
231
232    public SearchWrapper le(Connector connector, String key, Object value) {
233        if (this.condition == null) {
234            this.condition = new Condition(ConditionType.LE, new Key(key), new Value(value));
235        } else {
236            this.condition.connect(new Condition(ConditionType.LE, new Key(key), new Value(value)), connector);
237        }
238        return this;
239    }
240
241
242    public SearchWrapper in(String key, Collection<?> values) {
243        return in(Connector.AND, key, values);
244    }
245
246    public SearchWrapper in(Connector connector, String key, Collection<?> values) {
247        if (this.condition == null) {
248            this.condition = new Condition(ConditionType.IN, new Key(key), new Value(values.toArray()));
249        } else {
250            this.condition.connect(new Condition(ConditionType.IN, new Key(key), new Value(values.toArray())), connector);
251        }
252        return this;
253    }
254
255    public SearchWrapper min(String key, Object value) {
256        return min(Connector.AND, key, value);
257    }
258
259    public SearchWrapper min(Connector connector, String key, Object value) {
260        if (this.condition == null) {
261            this.condition = new Condition(ConditionType.NIN, new Key(key), new Value(value));
262        } else {
263            this.condition.connect(new Condition(ConditionType.NIN, new Key(key), new Value(value)), connector);
264        }
265        return this;
266    }
267
268    public SearchWrapper between(String key, Object start, Object end) {
269        return between(Connector.AND, key, start, end);
270    }
271
272    public SearchWrapper between(Connector connector, String key, Object start, Object end) {
273        if (this.condition == null) {
274            this.condition = new Condition(ConditionType.BETWEEN, new Key(key), new Value(start, end));
275        } else {
276            this.condition.connect(new Condition(ConditionType.BETWEEN, new Key(key), new Value(start, end)), connector);
277        }
278        return this;
279    }
280
281
282    public SearchWrapper group(SearchWrapper wrapper) {
283        return group(wrapper.condition);
284    }
285
286    public SearchWrapper group(Condition condition) {
287        if (this.condition == null) {
288            this.condition = new Group(condition);
289        } else {
290            this.condition.connect(new Group(condition), Connector.AND);
291        }
292        return this;
293    }
294
295    public SearchWrapper group(Consumer<SearchWrapper> consumer) {
296        SearchWrapper newWrapper = new SearchWrapper();
297        consumer.accept(newWrapper);
298        Condition condition = newWrapper.condition;
299        if (condition != null) {
300            group(condition);
301        }
302        return this;
303    }
304
305    public SearchWrapper andCriteria(Consumer<SearchWrapper> consumer) {
306        return group(consumer);
307    }
308
309    public SearchWrapper orCriteria(Consumer<SearchWrapper> consumer) {
310        SearchWrapper newWrapper = new SearchWrapper();
311        consumer.accept(newWrapper);
312        Condition condition = newWrapper.condition;
313        if (condition != null) {
314            if (this.condition == null) {
315                this.condition = new Group(condition);
316            } else {
317                this.condition.connect(new Group(condition), Connector.OR);
318            }
319        }
320        return this;
321    }
322
323    /**
324     * Convert to expressions for filtering conditions, with different expression requirements for each vendor.
325     * Customized adaptor can be achieved through ExpressionAdaptor
326     */
327    public String toFilterExpression() {
328        return toFilterExpression(ExpressionAdaptor.DEFAULT);
329    }
330
331    public String toFilterExpression(ExpressionAdaptor adaptor) {
332        if (this.condition == null) {
333            return null;
334        } else {
335            Objects.requireNonNull(adaptor, "adaptor must not be null");
336            return this.condition.toExpression(adaptor);
337        }
338    }
339
340}