diff --git a/llmengine/rerank.py b/llmengine/rerank.py index 645cc8f..fc7c3c6 100644 --- a/llmengine/rerank.py +++ b/llmengine/rerank.py @@ -60,17 +60,22 @@ async def docs(request, params_kw, *params, **kw): return helptext async def rerank(request, params_kw, *params, **kw): - debug(f'{params_kw.input=}') + debug(f'{params_kw.query=}, {params_kw.documents=}, {params_kw.top_n=}') se = ServerEnv() engine = se.engine f = awaitify(engine.rerank) query = params_kw.query if query is None: - e = exception(f'query is None') + e = Exception(f'query is None') raise e - if isinstance(query, str): - input = [input] - arr = await f(input) + documents = params_kw.documents + if documents is None: + e = Exception(f'documents is None') + raise e + if isinstance(documents, str): + documents = [documents] + top_n = parmas_kw.top_n + arr = await f(query, params_kw.documents, top_n = top_n) debug(f'{arr=}, type(arr)') return arr