def test_default(self): query = Query() self.assertDictEqual( query.create_body(query=self.query), { "yql": 'select * from sources * where (userInput("this is a test"));', "ranking": { "profile": "default", "listFeatures": "false" }, }, )
def test_match_and_rank(self): query = Query( match_phase=ANN( doc_vector="doc_vector", query_vector="query_vector", embedding_model=lambda x: [1, 2, 3], hits=10, label="label", ), rank_profile=RankProfile(name="bm25", list_features=True), ) self.assertDictEqual( query.create_body(query=self.query), { "yql": 'select * from sources * where ([{"targetNumHits": 10, "label": "label"}]nearestNeighbor(doc_vector, query_vector));', "ranking": { "profile": "bm25", "listFeatures": "true" }, "ranking.features.query(query_vector)": "[1, 2, 3]", }, )