xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-06-14 17:56:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add plugin rerank invocation action

Browse Source
This commit is contained in:
huanghuoguoguo
2026-06-13 16:47:56 +08:00
parent 7965d333ac
commit 7d94a3e8dd
2 changed files with 89 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/langbot/pkg/plugin/handler.py
View File

@@ -514,6 +514,35 @@ class RuntimeConnectionHandler(handler.Handler):
except Exception as e:
return _make_rag_error_response(e, 'EmbeddingError', embedding_model_uuid=embedding_model_uuid)
@self.action(PluginToRuntimeAction.INVOKE_RERANK)
async def invoke_rerank(data: dict[str, Any]) -> handler.ActionResponse:
rerank_model_uuid = data['rerank_model_uuid']
query = data['query']
documents = data['documents']
top_k = data.get('top_k')
extra_args = data.get('extra_args', {})
try:
rerank_model = await self.ap.model_mgr.get_rerank_model_by_uuid(rerank_model_uuid)
except ValueError:
return handler.ActionResponse.error(
message=f'Rerank model with rerank_model_uuid {rerank_model_uuid} not found',
)
try:
scores = await rerank_model.provider.invoke_rerank(
model=rerank_model,
query=query,
documents=documents[:64],
extra_args=extra_args,
)
scored = sorted(scores, key=lambda x: x.get('relevance_score', 0), reverse=True)
if top_k is not None:
scored = scored[: int(top_k)]
return handler.ActionResponse.success(data={'results': scored})
except Exception as e:
return _make_rag_error_response(e, 'RerankError', rerank_model_uuid=rerank_model_uuid)
@self.action(PluginToRuntimeAction.VECTOR_UPSERT)
async def vector_upsert(data: dict[str, Any]) -> handler.ActionResponse:
collection_id = data['collection_id']