rushdb 3.0.1__tar.gz → 3.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {rushdb-3.0.1 → rushdb-3.1.0}/PKG-INFO +13 -4
- {rushdb-3.0.1 → rushdb-3.1.0}/README.md +12 -3
- {rushdb-3.0.1 → rushdb-3.1.0}/pyproject.toml +1 -1
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/ai.py +43 -21
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/query.py +1 -1
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/direct/capabilities.py +2 -1
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/direct/router.py +2 -1
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/result.py +43 -3
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_ai_vector_search.py +28 -24
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_direct_mode.py +1 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/uv.lock +1 -1
- {rushdb-3.0.1 → rushdb-3.1.0}/.env.example +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/.github/workflows/ci.yml +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/.gitignore +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/CODE_OF_CONDUCT.md +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/CONTRIBUTING.md +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/LICENSE +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/rushdb-logo.svg +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/base.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/labels.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/properties.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/records.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/relationship_patterns.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/relationships.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/text_indexes.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/api/transactions.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/client.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/common.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/direct/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/direct/engine.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/direct/schema.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/api_response.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/property.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/record.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/relationship.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/search_query.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/models/transaction.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/src/rushdb/utils/token_prefix.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/__init__.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_base_setup.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_create_import.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_direct_mode_live.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_relationship_patterns.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_search_query.py +0 -0
- {rushdb-3.0.1 → rushdb-3.1.0}/tests/test_search_result.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: rushdb
|
|
3
|
-
Version: 3.0
|
|
3
|
+
Version: 3.1.0
|
|
4
4
|
Summary: RushDB SDK for Python: the client for RushDB, the database for AI (RushDB Cloud or self-hosted)
|
|
5
5
|
Project-URL: Homepage, https://github.com/rush-db/rushdb-python
|
|
6
6
|
Project-URL: Repository, https://github.com/rush-db/rushdb-python
|
|
@@ -102,7 +102,7 @@ What each offers:
|
|
|
102
102
|
| Keyword (BM25) search (`records.text_search`), text indexes (`text_indexes.*`), hybrid search | ✓ | ✓ |
|
|
103
103
|
| Transactions | ✓ | record creates, updates, sets, deletes by id, attaches and imports are staged; other writes refuse a transaction |
|
|
104
104
|
| Managed embeddings: managed indexes, `records.vector_search` by text `query` | ✓ | — |
|
|
105
|
-
| Natural-language search (`ai.search`) | ✓ | — |
|
|
105
|
+
| Natural-language search on SQL (`ai.search`, `ai.run_sql`) | ✓ | — |
|
|
106
106
|
| Relationship-pattern suggestions (`relationships.patterns.*`) | ✓ | — |
|
|
107
107
|
| Project settings (`settings.get`), dashboard, team access | ✓ | — |
|
|
108
108
|
|
|
@@ -280,6 +280,7 @@ if result:
|
|
|
280
280
|
| `total` | `int` | Total matching records in the database |
|
|
281
281
|
| `has_more` | `bool` | Whether more records exist beyond this page |
|
|
282
282
|
| `search_query` | `dict` | The query that produced this result |
|
|
283
|
+
| `sql` | `str` | For Smart search, the SQL statement it ran |
|
|
283
284
|
|
|
284
285
|
---
|
|
285
286
|
|
|
@@ -302,13 +303,21 @@ for record in results:
|
|
|
302
303
|
```
|
|
303
304
|
|
|
304
305
|
Use `db.ai.search()` when you want RushDB to turn a natural-language request
|
|
305
|
-
into
|
|
306
|
+
into one read-only SQL statement over your schema and run it (RushDB Cloud):
|
|
306
307
|
|
|
307
308
|
```python
|
|
308
309
|
results = db.ai.search('Find active memories about Q4 results for agent-42')
|
|
309
|
-
print(results.
|
|
310
|
+
print(results.name) # a short title, e.g. to save the query under
|
|
311
|
+
print(results.sql) # the statement it ran
|
|
312
|
+
more = results.next() # the next page of the statement's rows
|
|
313
|
+
|
|
314
|
+
# Run a statement directly, e.g. a saved one:
|
|
315
|
+
rows = db.ai.run_sql('SELECT agent_id, count(*) AS memories FROM "MEMORY" GROUP BY agent_id')
|
|
310
316
|
```
|
|
311
317
|
|
|
318
|
+
Rows that select a record's `__id` come back as full records; aggregates come
|
|
319
|
+
back as plain rows.
|
|
320
|
+
|
|
312
321
|
`db.ai.search({...})` still works as a deprecated vector-search alias, but new
|
|
313
322
|
code should use `db.records.vector_search({...})`.
|
|
314
323
|
|
|
@@ -88,7 +88,7 @@ What each offers:
|
|
|
88
88
|
| Keyword (BM25) search (`records.text_search`), text indexes (`text_indexes.*`), hybrid search | ✓ | ✓ |
|
|
89
89
|
| Transactions | ✓ | record creates, updates, sets, deletes by id, attaches and imports are staged; other writes refuse a transaction |
|
|
90
90
|
| Managed embeddings: managed indexes, `records.vector_search` by text `query` | ✓ | — |
|
|
91
|
-
| Natural-language search (`ai.search`) | ✓ | — |
|
|
91
|
+
| Natural-language search on SQL (`ai.search`, `ai.run_sql`) | ✓ | — |
|
|
92
92
|
| Relationship-pattern suggestions (`relationships.patterns.*`) | ✓ | — |
|
|
93
93
|
| Project settings (`settings.get`), dashboard, team access | ✓ | — |
|
|
94
94
|
|
|
@@ -266,6 +266,7 @@ if result:
|
|
|
266
266
|
| `total` | `int` | Total matching records in the database |
|
|
267
267
|
| `has_more` | `bool` | Whether more records exist beyond this page |
|
|
268
268
|
| `search_query` | `dict` | The query that produced this result |
|
|
269
|
+
| `sql` | `str` | For Smart search, the SQL statement it ran |
|
|
269
270
|
|
|
270
271
|
---
|
|
271
272
|
|
|
@@ -288,13 +289,21 @@ for record in results:
|
|
|
288
289
|
```
|
|
289
290
|
|
|
290
291
|
Use `db.ai.search()` when you want RushDB to turn a natural-language request
|
|
291
|
-
into
|
|
292
|
+
into one read-only SQL statement over your schema and run it (RushDB Cloud):
|
|
292
293
|
|
|
293
294
|
```python
|
|
294
295
|
results = db.ai.search('Find active memories about Q4 results for agent-42')
|
|
295
|
-
print(results.
|
|
296
|
+
print(results.name) # a short title, e.g. to save the query under
|
|
297
|
+
print(results.sql) # the statement it ran
|
|
298
|
+
more = results.next() # the next page of the statement's rows
|
|
299
|
+
|
|
300
|
+
# Run a statement directly, e.g. a saved one:
|
|
301
|
+
rows = db.ai.run_sql('SELECT agent_id, count(*) AS memories FROM "MEMORY" GROUP BY agent_id')
|
|
296
302
|
```
|
|
297
303
|
|
|
304
|
+
Rows that select a record's `__id` come back as full records; aggregates come
|
|
305
|
+
back as plain rows.
|
|
306
|
+
|
|
298
307
|
`db.ai.search({...})` still works as a deprecated vector-search alias, but new
|
|
299
308
|
code should use `db.records.vector_search({...})`.
|
|
300
309
|
|
|
@@ -5,12 +5,11 @@ and embedding index management.
|
|
|
5
5
|
"""
|
|
6
6
|
|
|
7
7
|
import warnings
|
|
8
|
-
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
|
8
|
+
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
|
|
9
9
|
|
|
10
10
|
from ..models.api_response import ApiResponse
|
|
11
11
|
from ..models.record import Record
|
|
12
12
|
from ..models.result import RecordSearchResult
|
|
13
|
-
from ..models.search_query import SearchQuery
|
|
14
13
|
from ..models.transaction import Transaction
|
|
15
14
|
from .base import BaseAPI
|
|
16
15
|
|
|
@@ -225,14 +224,18 @@ class AIAPI(BaseAPI):
|
|
|
225
224
|
def search(
|
|
226
225
|
self,
|
|
227
226
|
prompt: Union[str, Dict[str, Any]],
|
|
228
|
-
|
|
227
|
+
skip: Optional[int] = None,
|
|
228
|
+
limit: Optional[int] = None,
|
|
229
229
|
transaction: Optional[Union[Transaction, str]] = None,
|
|
230
230
|
) -> RecordSearchResult:
|
|
231
231
|
"""Perform AI-assisted smart search from natural language.
|
|
232
232
|
|
|
233
|
-
RushDB
|
|
234
|
-
|
|
235
|
-
|
|
233
|
+
RushDB turns ``prompt`` into one read-only SQL statement over the
|
|
234
|
+
project's schema, runs it, and returns the result with the statement
|
|
235
|
+
as ``result.sql``, a short title as ``result.name`` and any generation
|
|
236
|
+
warnings as ``result.warnings``. Rows that select a record's ``__id``
|
|
237
|
+
come back as full records; aggregates come back as plain rows.
|
|
238
|
+
``result.next()`` pages through the statement.
|
|
236
239
|
|
|
237
240
|
Use ``db.records.vector_search({...})`` for direct vector similarity
|
|
238
241
|
over embedding indexes.
|
|
@@ -240,14 +243,12 @@ class AIAPI(BaseAPI):
|
|
|
240
243
|
Args:
|
|
241
244
|
prompt: Natural-language search request. Passing a dict is
|
|
242
245
|
deprecated and delegates to ``db.records.vector_search``.
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
transaction:
|
|
246
|
+
skip: Rows to skip (default 0).
|
|
247
|
+
limit: Rows to return, at most 1000 (default 100).
|
|
248
|
+
transaction: Only used by the deprecated dict form.
|
|
246
249
|
|
|
247
250
|
Returns:
|
|
248
|
-
RecordSearchResult:
|
|
249
|
-
available as ``result.search_query`` and server warnings as
|
|
250
|
-
``result.warnings``.
|
|
251
|
+
RecordSearchResult: The statement's first page of rows.
|
|
251
252
|
"""
|
|
252
253
|
if isinstance(prompt, dict):
|
|
253
254
|
warnings.warn(
|
|
@@ -258,24 +259,45 @@ class AIAPI(BaseAPI):
|
|
|
258
259
|
)
|
|
259
260
|
return self.client.records.vector_search(prompt, transaction=transaction)
|
|
260
261
|
|
|
261
|
-
headers = Transaction._build_transaction_header(transaction)
|
|
262
262
|
generated = self.client._make_request(
|
|
263
|
-
"POST",
|
|
264
|
-
"/ai/search-query",
|
|
265
|
-
{"prompt": prompt, "currentQuery": current_query},
|
|
266
|
-
headers,
|
|
263
|
+
"POST", "/ai/sql-search", {"prompt": prompt}
|
|
267
264
|
)
|
|
268
265
|
generated_data = generated.get("data") or {}
|
|
269
|
-
|
|
266
|
+
result = self.run_sql(generated_data.get("sql", ""), skip=skip, limit=limit)
|
|
267
|
+
result._name = generated_data.get("name")
|
|
268
|
+
result._warnings = generated_data.get("warnings") or []
|
|
269
|
+
return result
|
|
270
|
+
|
|
271
|
+
def run_sql(
|
|
272
|
+
self,
|
|
273
|
+
sql: str,
|
|
274
|
+
skip: Optional[int] = None,
|
|
275
|
+
limit: Optional[int] = None,
|
|
276
|
+
) -> RecordSearchResult:
|
|
277
|
+
"""Run one read-only SQL statement over the project's schema, as Smart
|
|
278
|
+
search does: a statement from ``ai.search`` (its ``sql``), or one
|
|
279
|
+
written by hand. Labels are tables; relationship types are
|
|
280
|
+
``rushdb_edges.<type>`` tables.
|
|
270
281
|
|
|
282
|
+
Args:
|
|
283
|
+
sql: One read-only SELECT statement.
|
|
284
|
+
skip: Rows to skip (default 0).
|
|
285
|
+
limit: Rows to return, at most 1000 (default 100).
|
|
286
|
+
|
|
287
|
+
Returns:
|
|
288
|
+
RecordSearchResult: Full records for rows that select a record's
|
|
289
|
+
``__id``, otherwise the rows themselves; ``result.sql`` holds the
|
|
290
|
+
statement and ``result.next()`` fetches the following page.
|
|
291
|
+
"""
|
|
292
|
+
page = {"skip": skip or 0, "limit": limit or 100}
|
|
271
293
|
response = self.client._make_request(
|
|
272
|
-
"POST", "/records/search",
|
|
294
|
+
"POST", "/records/sql-search", {"sql": sql, **page}
|
|
273
295
|
)
|
|
274
296
|
records = [Record(self.client, item) for item in response.get("data", [])]
|
|
275
297
|
return RecordSearchResult(
|
|
276
298
|
data=records,
|
|
277
299
|
total=response.get("total", len(records)),
|
|
278
|
-
search_query=cast(SearchQuery, search_query),
|
|
279
300
|
client=self.client,
|
|
280
|
-
|
|
301
|
+
sql=sql,
|
|
302
|
+
sql_page=page,
|
|
281
303
|
)
|
|
@@ -8,7 +8,7 @@ class QueryAPI(BaseAPI):
|
|
|
8
8
|
"""API client for executing raw openCypher queries.
|
|
9
9
|
|
|
10
10
|
Works on RushDB Cloud and on a self-hosted RushDB engine in direct mode
|
|
11
|
-
(the engine must be started with ``--languages cypher``).
|
|
11
|
+
(the engine must be started with ``--languages cypher,sql``).
|
|
12
12
|
|
|
13
13
|
Example:
|
|
14
14
|
>>> from rushdb import RushDB
|
|
@@ -10,7 +10,8 @@ from typing import Dict
|
|
|
10
10
|
_PATTERNS = "relationship-pattern suggestions are produced and stored by the platform"
|
|
11
11
|
|
|
12
12
|
CLOUD_ONLY_METHODS: Dict[str, str] = {
|
|
13
|
-
"ai.search": "natural-language search turns the prompt into a
|
|
13
|
+
"ai.search": "natural-language search turns the prompt into a SQL statement with an LLM",
|
|
14
|
+
"ai.run_sql": "Smart search SQL runs on the platform, which loads the full records the statement selects",
|
|
14
15
|
"relationships.patterns.list": _PATTERNS,
|
|
15
16
|
"relationships.patterns.analyze": _PATTERNS,
|
|
16
17
|
"relationships.patterns.approve": _PATTERNS,
|
|
@@ -1372,7 +1372,8 @@ _ROUTES: List[Tuple[str, "re.Pattern[str]", Callable[[_Request], Body]]] = [
|
|
|
1372
1372
|
),
|
|
1373
1373
|
("DELETE", re.compile(rf"^/text-indexes/{_ID}$"), _drop_text_index),
|
|
1374
1374
|
("POST", re.compile(r"^/search/text$"), _text_search),
|
|
1375
|
-
("POST", re.compile(r"^/ai/search
|
|
1375
|
+
("POST", re.compile(r"^/ai/sql-search$"), _cloud_only("ai.search")),
|
|
1376
|
+
("POST", re.compile(r"^/records/sql-search$"), _cloud_only("ai.run_sql")),
|
|
1376
1377
|
("POST", re.compile(r"^/ai/search$"), _vector_search),
|
|
1377
1378
|
(
|
|
1378
1379
|
"GET",
|
|
@@ -32,6 +32,9 @@ class SearchResult(Generic[T]):
|
|
|
32
32
|
client: Optional["RushDB"] = None,
|
|
33
33
|
warnings: Optional[List[str]] = None,
|
|
34
34
|
meta: Optional[Dict[str, Any]] = None,
|
|
35
|
+
sql: Optional[str] = None,
|
|
36
|
+
name: Optional[str] = None,
|
|
37
|
+
sql_page: Optional[Dict[str, int]] = None,
|
|
35
38
|
):
|
|
36
39
|
"""
|
|
37
40
|
Initialize search result.
|
|
@@ -44,6 +47,10 @@ class SearchResult(Generic[T]):
|
|
|
44
47
|
warnings: Optional warnings returned by AI-assisted query generation
|
|
45
48
|
meta: How a vector, keyword or hybrid search ran (``vectorSearch``,
|
|
46
49
|
``textSearch``, ``fusion``)
|
|
50
|
+
sql: For Smart search, the SQL statement that produced this result
|
|
51
|
+
name: For Smart search, a short title of the query
|
|
52
|
+
sql_page: For Smart search, the page (``skip``, ``limit``) of the
|
|
53
|
+
statement's rows this result holds
|
|
47
54
|
"""
|
|
48
55
|
self._data = data
|
|
49
56
|
self._total = total or len(data)
|
|
@@ -51,6 +58,9 @@ class SearchResult(Generic[T]):
|
|
|
51
58
|
self._client = client
|
|
52
59
|
self._warnings = warnings or []
|
|
53
60
|
self._meta = meta or {}
|
|
61
|
+
self._sql = sql
|
|
62
|
+
self._name = name
|
|
63
|
+
self._sql_page = sql_page or {"skip": 0, "limit": 100}
|
|
54
64
|
|
|
55
65
|
@property
|
|
56
66
|
def data(self) -> List[T]:
|
|
@@ -67,6 +77,16 @@ class SearchResult(Generic[T]):
|
|
|
67
77
|
"""Get the search query used to generate this result."""
|
|
68
78
|
return self._search_query
|
|
69
79
|
|
|
80
|
+
@property
|
|
81
|
+
def sql(self) -> Optional[str]:
|
|
82
|
+
"""For Smart search: the SQL statement that produced this result."""
|
|
83
|
+
return self._sql
|
|
84
|
+
|
|
85
|
+
@property
|
|
86
|
+
def name(self) -> Optional[str]:
|
|
87
|
+
"""For Smart search: a short title of the query."""
|
|
88
|
+
return self._name
|
|
89
|
+
|
|
70
90
|
@property
|
|
71
91
|
def warnings(self) -> List[str]:
|
|
72
92
|
"""Get warnings returned by AI-assisted query generation."""
|
|
@@ -87,11 +107,15 @@ class SearchResult(Generic[T]):
|
|
|
87
107
|
@property
|
|
88
108
|
def skip(self) -> int:
|
|
89
109
|
"""Get the number of records that were skipped."""
|
|
110
|
+
if self._sql:
|
|
111
|
+
return self._sql_page["skip"]
|
|
90
112
|
return self._search_query.get("skip") or 0
|
|
91
113
|
|
|
92
114
|
@property
|
|
93
115
|
def limit(self) -> Optional[int]:
|
|
94
116
|
"""Get the limit that was applied to the search."""
|
|
117
|
+
if self._sql:
|
|
118
|
+
return self._sql_page["limit"]
|
|
95
119
|
return self._search_query.get("limit")
|
|
96
120
|
|
|
97
121
|
def __len__(self) -> int:
|
|
@@ -119,12 +143,14 @@ class SearchResult(Generic[T]):
|
|
|
119
143
|
Return the result in a standardized dictionary format.
|
|
120
144
|
|
|
121
145
|
Returns:
|
|
122
|
-
Dict with keys: total, data, search_query
|
|
146
|
+
Dict with keys: total, data, search_query, sql, name, warnings
|
|
123
147
|
"""
|
|
124
148
|
return {
|
|
125
149
|
"total": self.total,
|
|
126
150
|
"data": self.data,
|
|
127
151
|
"search_query": self.search_query,
|
|
152
|
+
"sql": self.sql,
|
|
153
|
+
"name": self.name,
|
|
128
154
|
"warnings": self.warnings,
|
|
129
155
|
}
|
|
130
156
|
|
|
@@ -164,7 +190,8 @@ class SearchResult(Generic[T]):
|
|
|
164
190
|
return self._get_client().records.delete_by_id(ids, transaction)
|
|
165
191
|
|
|
166
192
|
def next(self, preserve_data: bool = False) -> "SearchResult":
|
|
167
|
-
"""Fetch the next page of results
|
|
193
|
+
"""Fetch the next page of results: of the original search query, or of
|
|
194
|
+
the SQL statement for a Smart search result.
|
|
168
195
|
|
|
169
196
|
Args:
|
|
170
197
|
preserve_data: If True, appends new records to this instance's data
|
|
@@ -174,8 +201,21 @@ class SearchResult(Generic[T]):
|
|
|
174
201
|
SearchResult with the next page (or self when preserve_data=True)
|
|
175
202
|
|
|
176
203
|
Raises:
|
|
177
|
-
RuntimeError: If
|
|
204
|
+
RuntimeError: If neither a search_query nor a sql statement was provided
|
|
178
205
|
"""
|
|
206
|
+
if self._sql:
|
|
207
|
+
next_page = {
|
|
208
|
+
"skip": self._sql_page["skip"] + self._sql_page["limit"],
|
|
209
|
+
"limit": self._sql_page["limit"],
|
|
210
|
+
}
|
|
211
|
+
result = self._get_client().ai.run_sql(self._sql, **next_page)
|
|
212
|
+
if preserve_data:
|
|
213
|
+
self._data = self._data + result._data # type: ignore[operator]
|
|
214
|
+
self._total = result._total
|
|
215
|
+
self._sql_page = next_page
|
|
216
|
+
return self
|
|
217
|
+
result._name = self._name
|
|
218
|
+
return result
|
|
179
219
|
if not self._search_query:
|
|
180
220
|
raise RuntimeError(
|
|
181
221
|
"SearchResult: cannot paginate — no search_query was provided."
|
|
@@ -43,13 +43,14 @@ class TestVectorAndSmartSearch(unittest.TestCase):
|
|
|
43
43
|
self.assertEqual(result[0].get("title"), "Alpha")
|
|
44
44
|
self.assertEqual(result[0].score, 0.91)
|
|
45
45
|
|
|
46
|
-
def
|
|
46
|
+
def test_ai_search_prompt_generates_and_runs_sql(self):
|
|
47
47
|
client = Mock()
|
|
48
48
|
client._make_request.side_effect = [
|
|
49
49
|
{
|
|
50
50
|
"success": True,
|
|
51
51
|
"data": {
|
|
52
|
-
"
|
|
52
|
+
"sql": 'SELECT __id FROM "Pilot" WHERE "ship" = \'Falcon\'',
|
|
53
|
+
"name": "Falcon pilots",
|
|
53
54
|
"warnings": ["ambiguous ship name"],
|
|
54
55
|
},
|
|
55
56
|
},
|
|
@@ -60,38 +61,41 @@ class TestVectorAndSmartSearch(unittest.TestCase):
|
|
|
60
61
|
},
|
|
61
62
|
]
|
|
62
63
|
|
|
63
|
-
result = AIAPI(client).search(
|
|
64
|
-
"Who are piloting Falcon?",
|
|
65
|
-
current_query={"labels": ["Pilot"]},
|
|
66
|
-
transaction="tx_123",
|
|
67
|
-
)
|
|
64
|
+
result = AIAPI(client).search("Who are piloting Falcon?", limit=10)
|
|
68
65
|
|
|
69
66
|
self.assertEqual(
|
|
70
67
|
client._make_request.call_args_list[0].args,
|
|
71
|
-
(
|
|
72
|
-
"POST",
|
|
73
|
-
"/ai/search-query",
|
|
74
|
-
{
|
|
75
|
-
"prompt": "Who are piloting Falcon?",
|
|
76
|
-
"currentQuery": {"labels": ["Pilot"]},
|
|
77
|
-
},
|
|
78
|
-
{"X-Transaction-Id": "tx_123"},
|
|
79
|
-
),
|
|
68
|
+
("POST", "/ai/sql-search", {"prompt": "Who are piloting Falcon?"}),
|
|
80
69
|
)
|
|
70
|
+
sql = 'SELECT __id FROM "Pilot" WHERE "ship" = \'Falcon\''
|
|
71
|
+
self.assertEqual(
|
|
72
|
+
client._make_request.call_args_list[1].args,
|
|
73
|
+
("POST", "/records/sql-search", {"sql": sql, "skip": 0, "limit": 10}),
|
|
74
|
+
)
|
|
75
|
+
self.assertEqual(result.sql, sql)
|
|
76
|
+
self.assertEqual(result.name, "Falcon pilots")
|
|
77
|
+
self.assertEqual(result.warnings, ["ambiguous ship name"])
|
|
78
|
+
self.assertEqual(result.search_query, {})
|
|
79
|
+
self.assertEqual(result[0].get("name"), "Han")
|
|
80
|
+
|
|
81
|
+
def test_smart_search_result_pages_through_its_sql(self):
|
|
82
|
+
client = Mock()
|
|
83
|
+
client._make_request.return_value = {"success": True, "total": 3, "data": []}
|
|
84
|
+
client.ai = AIAPI(client)
|
|
85
|
+
|
|
86
|
+
first = client.ai.run_sql('SELECT __id FROM "Pilot"', limit=2)
|
|
87
|
+
self.assertEqual((first.skip, first.limit), (0, 2))
|
|
88
|
+
second = first.next()
|
|
89
|
+
|
|
81
90
|
self.assertEqual(
|
|
82
91
|
client._make_request.call_args_list[1].args,
|
|
83
92
|
(
|
|
84
93
|
"POST",
|
|
85
|
-
"/records/search",
|
|
86
|
-
{"
|
|
87
|
-
{"X-Transaction-Id": "tx_123"},
|
|
94
|
+
"/records/sql-search",
|
|
95
|
+
{"sql": 'SELECT __id FROM "Pilot"', "skip": 2, "limit": 2},
|
|
88
96
|
),
|
|
89
97
|
)
|
|
90
|
-
self.assertEqual(
|
|
91
|
-
result.search_query, {"labels": ["Pilot"], "where": {"ship": "Falcon"}}
|
|
92
|
-
)
|
|
93
|
-
self.assertEqual(result.warnings, ["ambiguous ship name"])
|
|
94
|
-
self.assertEqual(result[0].get("name"), "Han")
|
|
98
|
+
self.assertEqual((second.skip, second.limit, second.sql), (2, 2, first.sql))
|
|
95
99
|
|
|
96
100
|
def test_ai_search_dict_is_deprecated_vector_search_alias(self):
|
|
97
101
|
client = Mock()
|
|
@@ -271,6 +271,7 @@ EVERY_METHOD = {
|
|
|
271
271
|
"book-vec", {"items": [{"recordId": RECORD["id"], "vector": [1, 0, 0]}]}
|
|
272
272
|
),
|
|
273
273
|
"ai.search": lambda db: db.ai.search("books about deserts"),
|
|
274
|
+
"ai.run_sql": lambda db: db.ai.run_sql('SELECT __id FROM "Book"'),
|
|
274
275
|
"settings.get": lambda db: db.settings.get(),
|
|
275
276
|
}
|
|
276
277
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|