File size: 7,660 Bytes
4c67792 f98ae7d 4c67792 f98ae7d 4c67792 | 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 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 | """
Fact reranker following the original DSPy filter logic.
Uses QAFD-RAG's async LLM functions (wrapped synchronously) to call the
same prompt structure that the original DSPyFilter uses.
"""
import ast
import asyncio
import difflib
import json
import logging
import re
from copy import deepcopy
from typing import Callable, Dict, List, Tuple, Any
from .prompts import make_reranker_messages
logger = logging.getLogger(__name__)
def _run_sync(coro):
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None and loop.is_running():
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
return pool.submit(asyncio.run, coro).result()
else:
return asyncio.run(coro)
class FactReranker:
"""Rerank candidate fact triples using an LLM (DSPy-style filtering).
Parameters
----------
llm_func : callable
Async LLM function from ``src/llm.py``.
dspy_file_path : str or None
Path to a DSPy-saved JSON file with custom demos/system prompt.
If ``None``, uses built-in demos from ``prompts.py``.
"""
def __init__(self, llm_func: Callable, dspy_file_path: str = None):
self.llm_func = llm_func
self.dspy_file_path = dspy_file_path
if dspy_file_path is not None:
self._custom_template = self._load_dspy_template(dspy_file_path)
else:
self._custom_template = None
# ------------------------------------------------------------------
@staticmethod
def _load_dspy_template(path: str) -> List[Dict[str, str]]:
"""Load a DSPy-saved JSON and convert to chat messages."""
data = json.load(open(path))
system_prompt = data["prog"]["system"]
demos = data["prog"]["demos"]
one_in = (
"[[ ## question ## ]]\n{question}\n\n"
"[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
"Respond with the corresponding output fields, starting with the field "
"`[[ ## fact_after_filter ## ]]` (must be formatted as a valid Python Fact), "
"and then ending with the marker for `[[ ## completed ## ]]`."
)
one_out = (
"[[ ## fact_after_filter ## ]]\n{fact_after_filter}\n\n"
"[[ ## completed ## ]]"
)
msgs = [{"role": "system", "content": system_prompt}]
for demo in demos:
msgs.append({
"role": "user",
"content": one_in.format(
question=demo["question"],
fact_before_filter=demo["fact_before_filter"],
),
})
if "fact_after_filter" in demo:
msgs.append({
"role": "assistant",
"content": one_out.format(
fact_after_filter=demo["fact_after_filter"],
),
})
return msgs
# ------------------------------------------------------------------
def _build_messages(
self, question: str, fact_before_filter_json: str
) -> List[Dict[str, str]]:
if self._custom_template is not None:
msgs = deepcopy(self._custom_template)
one_in = (
"[[ ## question ## ]]\n{question}\n\n"
"[[ ## fact_before_filter ## ]]\n{fact_before_filter}\n\n"
"Respond with the corresponding output fields, starting with the field "
"`[[ ## fact_after_filter ## ]]` (must be formatted as a valid Python Fact), "
"and then ending with the marker for `[[ ## completed ## ]]`."
)
msgs.append({
"role": "user",
"content": one_in.format(
question=question,
fact_before_filter=fact_before_filter_json,
),
})
return msgs
else:
return make_reranker_messages(question, fact_before_filter_json)
# ------------------------------------------------------------------
def _call_llm(self, messages: List[Dict[str, str]]) -> str:
# Pass messages directly to OpenAI — the old decompose/recompose
# loop scrambled demo order (assistant before user in each pair).
return _run_sync(
self.llm_func(
prompt=messages[-1]["content"],
system_prompt=messages[0]["content"] if messages[0]["role"] == "system" else None,
history_messages=messages[1:-1],
max_tokens=512,
)
)
# ------------------------------------------------------------------
@staticmethod
def _parse_filter(response: str) -> List[List[str]]:
"""Extract fact_after_filter from the DSPy-style response."""
sections = [(None, [])]
header_re = re.compile(r'\[\[ ## (\w+) ## \]\]')
for line in response.splitlines():
m = header_re.match(line.strip())
if m:
sections.append((m.group(1), []))
else:
sections[-1][1].append(line)
sections = [(k, "\n".join(v).strip()) for k, v in sections]
parsed: List[List[str]] = []
for k, value in sections:
if k == "fact_after_filter":
try:
try:
pv = json.loads(value)
except json.JSONDecodeError:
try:
pv = ast.literal_eval(value)
except (ValueError, SyntaxError):
pv = value
if isinstance(pv, dict) and "fact" in pv:
parsed = pv["fact"]
except Exception as e:
logger.warning(f"Error parsing reranker output: {e}")
return parsed
# ------------------------------------------------------------------
def rerank(
self,
query: str,
candidate_items: List[Tuple],
candidate_indices: List[int],
len_after_rerank: int = None,
) -> Tuple[List[int], List[Tuple], dict]:
"""Rerank candidate facts by LLM-based filtering.
Returns
-------
(sorted_indices, sorted_facts, metadata_dict)
"""
fact_json = json.dumps({"fact": [list(c) for c in candidate_items]})
try:
msgs = self._build_messages(query, fact_json)
response = self._call_llm(msgs)
generated_facts = self._parse_filter(response)
except Exception as e:
logger.warning(f"Reranker exception: {e}")
generated_facts = []
result_indices = []
for gf in generated_facts:
matches = difflib.get_close_matches(
str(gf), [str(i) for i in candidate_items], n=1, cutoff=0.0
)
if matches:
try:
result_indices.append(candidate_items.index(eval(matches[0])))
except Exception as e:
logger.warning(f"Index matching error: {e}")
sorted_indices = [candidate_indices[i] for i in result_indices]
sorted_items = [candidate_items[i] for i in result_indices]
if len_after_rerank is not None:
sorted_indices = sorted_indices[:len_after_rerank]
sorted_items = sorted_items[:len_after_rerank]
return sorted_indices, sorted_items, {"confidence": None}
def __call__(self, *args, **kwargs):
return self.rerank(*args, **kwargs)
|