File size: 4,212 Bytes
9abace2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """
Reranking Tools
Part of SOVEREIGN PYTHON LLM ENGINE
"""
from typing import Any
from ..registry import tool, RiskClass, ApprovalPolicy
@tool(
tool_id="rerank.documents",
version="1.0.0",
title="Rerank Documents",
description="Rerank documents by relevance to query using cross-encoder model",
input_schema={
"type": "object",
"properties": {
"query": {"type": "string"},
"documents": {"type": "array", "items": {"type": "string"}},
"top_k": {"type": "integer", "default": 5},
"model": {"type": "string", "default": "rerank-english-v3.0"},
"provider": {"type": "string", "enum": ["cohere"], "default": "cohere"}
},
"required": ["query", "documents"]
},
output_schema={
"type": "object",
"properties": {
"ranked_documents": {
"type": "array",
"items": {
"type": "object",
"properties": {
"document": {"type": "string"},
"score": {"type": "number"},
"index": {"type": "integer"}
}
}
},
"model": {"type": "string"}
}
},
risk_class=RiskClass.READ_ONLY_REMOTE,
approval_policy=ApprovalPolicy.AUTOMATIC,
tags=["rerank", "rag", "retrieval"]
)
async def rerank_documents(params: dict[str, Any]) -> dict[str, Any]:
"""
Rerank documents by relevance to query.
Args:
params: {query, documents, top_k, model, provider}
Returns:
{ranked_documents, model}
"""
query = params['query']
documents = params['documents']
top_k = params.get('top_k', 5)
model = params.get('model', 'rerank-english-v3.0')
provider = params.get('provider', 'cohere')
# Get provider implementation
if provider == 'cohere':
from .providers.cohere import CohereReranker
provider_impl = CohereReranker()
else:
raise ValueError(f"Unknown provider: {provider}")
# Rerank
ranked = await provider_impl.rerank(
query=query,
documents=documents,
top_k=top_k,
model=model
)
return {
'ranked_documents': ranked,
'model': model
}
@tool(
tool_id="rerank.score",
version="1.0.0",
title="Score Document Relevance",
description="Compute relevance score for single document given query",
input_schema={
"type": "object",
"properties": {
"query": {"type": "string"},
"document": {"type": "string"},
"model": {"type": "string", "default": "rerank-english-v3.0"},
"provider": {"type": "string", "enum": ["cohere"], "default": "cohere"}
},
"required": ["query", "document"]
},
output_schema={
"type": "object",
"properties": {
"score": {"type": "number"},
"model": {"type": "string"}
}
},
risk_class=RiskClass.READ_ONLY_REMOTE,
approval_policy=ApprovalPolicy.AUTOMATIC,
tags=["rerank", "rag", "relevance"]
)
async def score_document(params: dict[str, Any]) -> dict[str, Any]:
"""
Compute relevance score for document.
Args:
params: {query, document, model, provider}
Returns:
{score, model}
"""
query = params['query']
document = params['document']
model = params.get('model', 'rerank-english-v3.0')
provider = params.get('provider', 'cohere')
# Get provider implementation
if provider == 'cohere':
from .providers.cohere import CohereReranker
provider_impl = CohereReranker()
else:
raise ValueError(f"Unknown provider: {provider}")
# Score single document
ranked = await provider_impl.rerank(
query=query,
documents=[document],
top_k=1,
model=model
)
score = ranked[0]['score'] if ranked else 0.0
return {
'score': score,
'model': model
}
|