diff --git a/README.md b/README.md index 7cfac3a..af9a61b 100644 --- a/README.md +++ b/README.md @@ -202,6 +202,37 @@ if result: --- +## Vector Search And Smart Search + +Use `db.records.vector_search()` for direct semantic/vector retrieval over an +embedding index: + +```python +results = db.records.vector_search({ + 'labels': ['MEMORY'], + 'propertyName': 'content', + 'query': 'how agents remember things', + 'where': {'agent_id': 'agent-42'}, + 'limit': 5, +}) + +for record in results: + print(record.score, record.get('content')) +``` + +Use `db.ai.search()` when you want RushDB to turn a natural-language request +into a SearchQuery and execute it: + +```python +results = db.ai.search('Find active memories about Q4 results for agent-42') +print(results.search_query) +``` + +`db.ai.search({...})` still works as a deprecated vector-search alias, but new +code should use `db.records.vector_search({...})`. + +--- + ## Record API ```python diff --git a/pyproject.toml b/pyproject.toml index 77303eb..243412e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "rushdb" -version = "2.9.0" +version = "2.10.0" description = "RushDB Python SDK — memory layer for AI agents and modern apps" authors = [ {name = "RushDB Team", email = "hi@rushdb.com"} diff --git a/src/rushdb/api/ai.py b/src/rushdb/api/ai.py index d524f0c..3d998a0 100644 --- a/src/rushdb/api/ai.py +++ b/src/rushdb/api/ai.py @@ -1,13 +1,16 @@ """AI API for RushDB Python SDK. -Provides methods for graph schema exploration, semantic vector search, +Provides methods for AI-assisted graph exploration, smart search, and embedding index management. """ -from typing import TYPE_CHECKING, Any, Dict, Optional, Union +import warnings +from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast from ..models.api_response import ApiResponse from ..models.record import Record +from ..models.result import RecordSearchResult +from ..models.search_query import SearchQuery from ..models.transaction import Transaction from .base import BaseAPI @@ -124,7 +127,12 @@ class AIAPI(BaseAPI): >>> db = RushDB(api_key="...") >>> schema = db.ai.get_schema() - >>> results = db.ai.search({"query": "fast cars", "propertyName": "description"}) + >>> results = db.ai.search("books about fast cars") + >>> vectors = db.records.vector_search({ + ... "query": "fast cars", + ... "propertyName": "description", + ... "labels": ["Book"], + ... }) >>> db.ai.indexes.create({"propertyName": "description", "label": "Book"}) """ @@ -145,7 +153,7 @@ def get_schema( indexes exist for that property — each entry exposes ``id``, ``sourceType``, ``similarityFunction``, ``dimensions``, ``status``, and ``modelKey``. A non-empty ``vectorIndexes`` list means the property is - queryable with ``db.ai.search()``. + queryable with ``db.records.vector_search()``. Args: params: Optional filter. Pass ``{"labels": ["Label1"]}`` to scope @@ -214,35 +222,60 @@ def get_schema_markdown( total=response.get("total"), ) - def search(self, params: Dict[str, Any]) -> ApiResponse: - """Perform semantic (vector) search over indexed record properties. + def search( + self, + prompt: Union[str, Dict[str, Any]], + current_query: Optional[SearchQuery] = None, + transaction: Optional[Union[Transaction, str]] = None, + ) -> RecordSearchResult: + """Perform AI-assisted smart search from natural language. - **Direct vector-index mode** (default, fast): used when no ``where`` - filter and at most one ``labels`` entry. Queries the shared global - vector index directly. + RushDB converts ``prompt`` into a SearchQuery using the project schema, + executes it, and returns matching records with the generated query + attached to ``result.search_query``. - **Prefilter mode** (exact, slower): activated when a ``where`` - filter is supplied or ``labels`` contains more than one value. - Candidates are first narrowed by MATCH/WHERE, then ranked by exact - cosine similarity. + Use ``db.records.vector_search({...})`` for direct vector similarity + over embedding indexes. Args: - params: Search parameters. Expected keys: - - - ``query`` (str): The natural-language query text. - - ``propertyName`` (str): Property that has been embedded. - - ``labels`` (list[str], optional): Scope to specific labels. - - ``where`` (dict, optional): Additional property filters. - - ``limit`` (int, optional): Maximum number of results. + prompt: Natural-language search request. Passing a dict is + deprecated and delegates to ``db.records.vector_search``. + current_query: Optional current SearchQuery context from a + dashboard/query-builder session. + transaction: Optional transaction context. Returns: - ApiResponse: Response whose ``data`` is a list of semantic search - result objects (each includes the matched record and a score). + RecordSearchResult: Matching records. The generated SearchQuery is + available as ``result.search_query`` and server warnings as + ``result.warnings``. """ - response = self.client._make_request("POST", "/ai/search", params) + if isinstance(prompt, dict): + warnings.warn( + "db.ai.search({...}) is deprecated for vector search; " + "use db.records.vector_search({...}) instead.", + DeprecationWarning, + stacklevel=2, + ) + return self.client.records.vector_search(prompt, transaction=transaction) + + headers = Transaction._build_transaction_header(transaction) + generated = self.client._make_request( + "POST", + "/ai/search-query", + {"prompt": prompt, "currentQuery": current_query}, + headers, + ) + generated_data = generated.get("data") or {} + search_query = generated_data.get("searchQuery") or {} + + response = self.client._make_request( + "POST", "/records/search", search_query, headers + ) records = [Record(self.client, item) for item in response.get("data", [])] - return ApiResponse( + return RecordSearchResult( data=records, - success=response.get("success", True), - total=response.get("total"), + total=response.get("total", len(records)), + search_query=cast(SearchQuery, search_query), + client=self.client, + warnings=generated_data.get("warnings") or [], ) diff --git a/src/rushdb/api/records.py b/src/rushdb/api/records.py index 4710b20..137329e 100644 --- a/src/rushdb/api/records.py +++ b/src/rushdb/api/records.py @@ -688,6 +688,46 @@ def find( except Exception: return RecordSearchResult(data=[], total=0, client=self.client) + def vector_search( + self, + params: Dict[str, Any], + transaction: Optional[Union[Transaction, str]] = None, + ) -> RecordSearchResult: + """Perform vector similarity search over indexed record properties. + + This is the Python SDK counterpart to the TypeScript + ``db.records.vectorSearch({...})`` method. RushDB narrows candidates by + ``labels`` and optional ``where`` filters first, then ranks them by + vector similarity. Pass ``query`` for managed indexes or ``queryVector`` + for external/custom vectors. + + Args: + params: Vector search parameters. Expected keys: + + - ``propertyName`` (str): Property that has an embedding index. + - ``labels`` (list[str]): Labels to scope the search. + - ``query`` (str, optional): Text query for managed indexes. + - ``queryVector`` (list[float], optional): Pre-computed vector. + - ``where`` (dict, optional): Additional property filters. + - ``limit`` (int, optional): Maximum number of results. + - ``skip`` (int, optional): Number of results to skip. + - ``topK`` (int, optional): Candidate count in direct vector mode. + transaction: Optional transaction context. + + Returns: + RecordSearchResult: Matching records ranked by similarity. Each + record may include ``__score``. + """ + headers = Transaction._build_transaction_header(transaction) + response = self.client._make_request("POST", "/ai/search", params, headers) + records = [Record(self.client, item) for item in response.get("data", [])] + return RecordSearchResult( + data=records, + total=response.get("total", len(records)), + search_query=typing.cast(SearchQuery, params), + client=self.client, + ) + def find_one( self, search_query: Optional[SearchQuery] = None, diff --git a/src/rushdb/models/result.py b/src/rushdb/models/result.py index 4f7ef80..2f712ba 100644 --- a/src/rushdb/models/result.py +++ b/src/rushdb/models/result.py @@ -30,6 +30,7 @@ def __init__( total: Optional[int] = None, search_query: Optional[SearchQuery] = None, client: Optional["RushDB"] = None, + warnings: Optional[List[str]] = None, ): """ Initialize search result. @@ -39,11 +40,13 @@ def __init__( total: Total number of matching records (may be larger than len(data)) search_query: The search query used to generate this result client: Optional RushDB client instance (required for delete_all, next, set_properties) + warnings: Optional warnings returned by AI-assisted query generation """ self._data = data self._total = total or len(data) self._search_query = search_query or {} self._client = client + self._warnings = warnings or [] @property def data(self) -> List[T]: @@ -60,6 +63,11 @@ def search_query(self) -> SearchQuery: """Get the search query used to generate this result.""" return self._search_query + @property + def warnings(self) -> List[str]: + """Get warnings returned by AI-assisted query generation.""" + return self._warnings + @property def has_more(self) -> bool: """Check if there are more records available beyond this result set.""" @@ -106,6 +114,7 @@ def to_dict(self) -> dict: "total": self.total, "data": self.data, "search_query": self.search_query, + "warnings": self.warnings, } def get_page_info(self) -> dict: diff --git a/tests/test_ai_vector_search.py b/tests/test_ai_vector_search.py new file mode 100644 index 0000000..f32271e --- /dev/null +++ b/tests/test_ai_vector_search.py @@ -0,0 +1,117 @@ +import unittest +import warnings +from unittest.mock import Mock + +from src.rushdb.api.ai import AIAPI +from src.rushdb.api.records import RecordsAPI +from src.rushdb.models.result import RecordSearchResult + + +class TestVectorAndSmartSearch(unittest.TestCase): + def test_records_vector_search_posts_to_ai_search(self): + client = Mock() + client._make_request.return_value = { + "success": True, + "total": 1, + "data": [ + {"__id": "doc_1", "__label": "Doc", "title": "Alpha", "__score": 0.91} + ], + } + + result = RecordsAPI(client).vector_search( + { + "labels": ["Doc"], + "propertyName": "description", + "queryVector": [1, 0, 0], + "limit": 5, + } + ) + + client._make_request.assert_called_once_with( + "POST", + "/ai/search", + { + "labels": ["Doc"], + "propertyName": "description", + "queryVector": [1, 0, 0], + "limit": 5, + }, + None, + ) + self.assertIsInstance(result, RecordSearchResult) + self.assertEqual(result.total, 1) + self.assertEqual(result[0].get("title"), "Alpha") + self.assertEqual(result[0].score, 0.91) + + def test_ai_search_prompt_generates_and_executes_search_query(self): + client = Mock() + client._make_request.side_effect = [ + { + "success": True, + "data": { + "searchQuery": {"labels": ["Pilot"], "where": {"ship": "Falcon"}}, + "warnings": ["ambiguous ship name"], + }, + }, + { + "success": True, + "total": 1, + "data": [{"__id": "pilot_1", "__label": "Pilot", "name": "Han"}], + }, + ] + + result = AIAPI(client).search( + "Who are piloting Falcon?", + current_query={"labels": ["Pilot"]}, + transaction="tx_123", + ) + + self.assertEqual( + client._make_request.call_args_list[0].args, + ( + "POST", + "/ai/search-query", + { + "prompt": "Who are piloting Falcon?", + "currentQuery": {"labels": ["Pilot"]}, + }, + {"X-Transaction-Id": "tx_123"}, + ), + ) + self.assertEqual( + client._make_request.call_args_list[1].args, + ( + "POST", + "/records/search", + {"labels": ["Pilot"], "where": {"ship": "Falcon"}}, + {"X-Transaction-Id": "tx_123"}, + ), + ) + self.assertEqual( + result.search_query, {"labels": ["Pilot"], "where": {"ship": "Falcon"}} + ) + self.assertEqual(result.warnings, ["ambiguous ship name"]) + self.assertEqual(result[0].get("name"), "Han") + + def test_ai_search_dict_is_deprecated_vector_search_alias(self): + client = Mock() + client.records = Mock() + client.records.vector_search.return_value = RecordSearchResult([], total=0) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = AIAPI(client).search( + {"labels": ["Doc"], "propertyName": "body", "query": "graph"} + ) + + client.records.vector_search.assert_called_once_with( + {"labels": ["Doc"], "propertyName": "body", "query": "graph"}, + transaction=None, + ) + self.assertIsInstance(result, RecordSearchResult) + self.assertEqual(len(caught), 1) + self.assertIs(caught[0].category, DeprecationWarning) + + +if __name__ == "__main__": + unittest.main() diff --git a/uv.lock b/uv.lock index 8187e37..7e2189a 100644 --- a/uv.lock +++ b/uv.lock @@ -785,7 +785,7 @@ wheels = [ [[package]] name = "rushdb" -version = "2.9.0" +version = "2.10.0" source = { editable = "." } dependencies = [ { name = "python-dotenv", version = "1.0.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" },