wehub-resource-sync
228 lines
8.0 KiB
Python
228 lines
8.0 KiB
Python
"""Node Re-ranker class for async execution"""
|
|
|
|
from collections.abc import Callable
|
|
import logging
|
|
|
|
import google.auth
|
|
import google.auth.transport.requests
|
|
from llama_index.core import QueryBundle, Settings
|
|
from llama_index.core.bridge.pydantic import Field, PrivateAttr
|
|
from llama_index.core.indices.utils import (
|
|
default_format_node_batch_fn,
|
|
default_parse_choice_select_answer_fn,
|
|
)
|
|
from llama_index.core.llms.llm import LLM
|
|
from llama_index.core.postprocessor.types import BaseNodePostprocessor
|
|
from llama_index.core.prompts import BasePromptTemplate
|
|
from llama_index.core.prompts.default_prompts import DEFAULT_CHOICE_SELECT_PROMPT
|
|
from llama_index.core.prompts.mixin import PromptDictType
|
|
from llama_index.core.schema import NodeWithScore, TextNode
|
|
from llama_index.core.service_context import ServiceContext
|
|
from llama_index.core.settings import llm_from_settings_or_context
|
|
from llama_index.llms.vertex import Vertex
|
|
import requests
|
|
|
|
logging.basicConfig(level=logging.INFO) # Set the desired logging level
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Initialize the LLM and set it in the Settings
|
|
llm = Vertex(model="gemini-2.0-flash", temperature=0.0)
|
|
Settings.llm = llm
|
|
|
|
|
|
def authenticate_google():
|
|
"""Authenticate using Google credentials and return the access token."""
|
|
credentials, project_id = google.auth.default(
|
|
quota_project_id="pr-sbx-vertex-genai"
|
|
)
|
|
auth_req = google.auth.transport.requests.Request()
|
|
credentials.refresh(auth_req)
|
|
return credentials.token
|
|
|
|
|
|
def call_reranker(query, records, google_token):
|
|
"""Calls the reranker API with the given query and records.
|
|
|
|
Args:
|
|
query: The search query.
|
|
records: A list of dictionaries, where each dictionary represents a record
|
|
with "id", "title", and "content" fields.
|
|
|
|
Returns:
|
|
The API response as a dictionary.
|
|
"""
|
|
# Replace 'your-project-id' with your actual Google Cloud project ID
|
|
project_id = "pr-sbx-vertex-genai"
|
|
model_name = "semantic-ranker-512@latest"
|
|
url = f"https://discoveryengine.googleapis.com/v1alpha/projects/{project_id}/locations/global/rankingConfigs/default_ranking_config:rank"
|
|
|
|
headers = {
|
|
"Authorization": "Bearer " + google_token,
|
|
"Content-Type": "application/json",
|
|
"X-Goog-User-Project": project_id,
|
|
}
|
|
|
|
data = {
|
|
"model": model_name,
|
|
"query": query,
|
|
"records": records,
|
|
}
|
|
|
|
response = requests.post(url, headers=headers, json=data)
|
|
print(response)
|
|
response.raise_for_status() # Raise an error if the request failed
|
|
return response.json()
|
|
|
|
|
|
class GoogleReRankerSecretSauce(BaseNodePostprocessor):
|
|
def _postprocess_nodes(
|
|
self, nodes: list[NodeWithScore], query_bundle: QueryBundle | None
|
|
) -> list[NodeWithScore]:
|
|
google_token = authenticate_google()
|
|
|
|
records = []
|
|
for node_wscore in nodes:
|
|
records.append(
|
|
{
|
|
"id": node_wscore.node.id_,
|
|
"title": node_wscore.node.metadata["title"],
|
|
"content": node_wscore.node.text,
|
|
}
|
|
)
|
|
response_json = call_reranker(query_bundle.query_str, records, google_token)
|
|
|
|
records = response_json["records"]
|
|
new_nodes_wscores = []
|
|
for r in records:
|
|
node = TextNode(id_=r["id"], text=r["content"])
|
|
node_wscore = NodeWithScore(node=node, score=r["score"])
|
|
new_nodes_wscores.append(node_wscore)
|
|
|
|
return sorted(new_nodes_wscores, key=lambda x: x.score or 0.0, reverse=True)
|
|
|
|
|
|
class CustomLLMRerank(BaseNodePostprocessor):
|
|
"""LLM-based reranker."""
|
|
|
|
top_n: int = Field(description="Top N nodes to return.")
|
|
choice_select_prompt: BasePromptTemplate = Field(
|
|
description="Choice select prompt."
|
|
)
|
|
choice_batch_size: int = Field(description="Batch size for choice select.")
|
|
llm: LLM = Field(description="The LLM to rerank with.")
|
|
|
|
_format_node_batch_fn: Callable = PrivateAttr()
|
|
_parse_choice_select_answer_fn: Callable = PrivateAttr()
|
|
|
|
def __init__(
|
|
self,
|
|
llm: LLM | None = None,
|
|
choice_select_prompt: BasePromptTemplate | None = None,
|
|
choice_batch_size: int = 10,
|
|
format_node_batch_fn: Callable | None = None,
|
|
parse_choice_select_answer_fn: Callable | None = None,
|
|
service_context: ServiceContext | None = None,
|
|
top_n: int = 10,
|
|
) -> None:
|
|
choice_select_prompt = choice_select_prompt or DEFAULT_CHOICE_SELECT_PROMPT
|
|
|
|
llm = llm or llm_from_settings_or_context(Settings, service_context)
|
|
|
|
self._format_node_batch_fn = (
|
|
format_node_batch_fn or default_format_node_batch_fn
|
|
)
|
|
self._parse_choice_select_answer_fn = (
|
|
parse_choice_select_answer_fn or default_parse_choice_select_answer_fn
|
|
)
|
|
|
|
super().__init__(
|
|
llm=llm,
|
|
choice_select_prompt=choice_select_prompt,
|
|
choice_batch_size=choice_batch_size,
|
|
service_context=service_context,
|
|
top_n=top_n,
|
|
)
|
|
|
|
def _get_prompts(self) -> PromptDictType:
|
|
"""Get prompts."""
|
|
return {"choice_select_prompt": self.choice_select_prompt}
|
|
|
|
def _update_prompts(self, prompts: PromptDictType) -> None:
|
|
"""Update prompts."""
|
|
if "choice_select_prompt" in prompts:
|
|
self.choice_select_prompt = prompts["choice_select_prompt"]
|
|
|
|
@classmethod
|
|
def class_name(cls) -> str:
|
|
return "LLMRerank"
|
|
|
|
async def postprocess_nodes(
|
|
self,
|
|
nodes: list[NodeWithScore],
|
|
query_bundle: QueryBundle | None = None,
|
|
query_str: str | None = None,
|
|
) -> list[NodeWithScore]:
|
|
"""Postprocess nodes."""
|
|
if query_str is not None and query_bundle is not None:
|
|
raise ValueError("Cannot specify both query_str and query_bundle")
|
|
elif query_str is not None:
|
|
query_bundle = QueryBundle(query_str)
|
|
else:
|
|
pass
|
|
return await self._postprocess_nodes(nodes, query_bundle)
|
|
|
|
async def _postprocess_nodes(
|
|
self,
|
|
nodes: list[NodeWithScore],
|
|
query_bundle: QueryBundle | None = None,
|
|
) -> list[NodeWithScore]:
|
|
if query_bundle is None:
|
|
raise ValueError("Query bundle must be provided.")
|
|
if len(nodes) == 0:
|
|
return []
|
|
|
|
initial_results: list[NodeWithScore] = []
|
|
for idx in range(0, len(nodes), self.choice_batch_size):
|
|
nodes_batch = [
|
|
node.node for node in nodes[idx : idx + self.choice_batch_size]
|
|
]
|
|
|
|
query_str = query_bundle.query_str
|
|
fmt_batch_str = self._format_node_batch_fn(nodes_batch)
|
|
# call each batch independently
|
|
raw_response = await self.llm.apredict(
|
|
self.choice_select_prompt,
|
|
context_str=fmt_batch_str,
|
|
query_str=query_str,
|
|
)
|
|
|
|
logging.info(raw_response)
|
|
try:
|
|
raw_choices, relevances = self._parse_choice_select_answer_fn(
|
|
raw_response, len(nodes_batch)
|
|
)
|
|
# Try again
|
|
except IndexError:
|
|
raw_response = await self.llm.apredict(
|
|
self.choice_select_prompt,
|
|
context_str=fmt_batch_str,
|
|
query_str=query_str,
|
|
)
|
|
raw_choices, relevances = self._parse_choice_select_answer_fn(
|
|
raw_response, len(nodes_batch)
|
|
)
|
|
choice_idxs = [int(choice) - 1 for choice in raw_choices]
|
|
choice_nodes = [nodes_batch[idx] for idx in choice_idxs]
|
|
relevances = relevances or [1.0 for _ in choice_nodes]
|
|
initial_results.extend(
|
|
[
|
|
NodeWithScore(node=node, score=relevance)
|
|
for node, relevance in zip(choice_nodes, relevances)
|
|
]
|
|
)
|
|
|
|
return sorted(initial_results, key=lambda x: x.score or 0.0, reverse=True)[
|
|
: self.top_n
|
|
]
|