Spaces:
Running
Running
| """Query transformation strategies via LiteLLM, as used by Notebook 6. | |
| `QueryTransformCache` wraps `rewrite_query` / `expand_queries` / `hyde_generate` | |
| with a disk-backed cache (keyed by `(transform_type, question)`), so repeated | |
| demo runs for the same question + strategy don't re-call the LLM. | |
| Sync methods are used by the single-query path (called from a thread via | |
| `asyncio.to_thread`). A `threading.Lock` guards the cache file writes so | |
| concurrent benchmark threads do not corrupt the JSON on disk. | |
| """ | |
| import hashlib | |
| import json | |
| import re | |
| import threading | |
| from pathlib import Path | |
| from src import settings | |
| from src.generation.llm_client import DEFAULT_MODEL, generate | |
| REWRITE_PROMPT = """Rewrite the following question to make it more effective for retrieving \ | |
| relevant passages from a document corpus. Preserve the original meaning. Output ONLY \ | |
| the rewritten question, nothing else. | |
| Question: {question} | |
| Rewritten question:""" | |
| MULTIQUERY_PROMPT = """Generate {n} alternative phrasings of the following question. Each \ | |
| alternative should preserve the original meaning but use different wording or \ | |
| emphasis, to help retrieve relevant passages from a document corpus. | |
| Question: {question} | |
| Respond with ONLY a JSON array of {n} strings, e.g. ["...", "..."]""" | |
| HYDE_PROMPT = """Write a short hypothetical passage (2-4 sentences) that would directly \ | |
| answer the following question. Write it as if it were an excerpt from a reference \ | |
| document, not as a conversational response. Output ONLY the passage, nothing else. | |
| Question: {question} | |
| Hypothetical passage:""" | |
| _write_lock = threading.Lock() | |
| class QueryTransformCache: | |
| def __init__(self, path: Path = None): | |
| self.path = path or (settings.APP_CACHE_DIR / "query_transforms.json") | |
| if self.path.exists(): | |
| self._cache = json.loads(self.path.read_text()) | |
| else: | |
| self._cache = {} | |
| def _get_or_compute(self, question: str, kind: str, compute_fn): | |
| key = hashlib.md5(f"{kind}|{question}".encode("utf-8")).hexdigest() | |
| if key in self._cache: | |
| return self._cache[key] | |
| result = compute_fn(question) | |
| with _write_lock: | |
| self._cache[key] = result | |
| self.path.write_text(json.dumps(self._cache)) | |
| return result | |
| def rewrite_query(self, question: str, model: str = DEFAULT_MODEL) -> str: | |
| def _compute(q): | |
| return generate(REWRITE_PROMPT.format(question=q), model=model).strip() or q | |
| return self._get_or_compute(question, "rewrite", _compute) | |
| def expand_queries(self, question: str, n: int = 2, model: str = DEFAULT_MODEL) -> list: | |
| def _compute(q): | |
| raw = generate(MULTIQUERY_PROMPT.format(question=q, n=n), model=model) | |
| match = re.search(r"\[.*\]", raw, re.DOTALL) | |
| try: | |
| variants = json.loads(match.group(0)) if match else [] | |
| except Exception: | |
| variants = [] | |
| variants = [v.strip() for v in variants if isinstance(v, str) and v.strip()] | |
| return variants[:n] or [q] | |
| return self._get_or_compute(question, f"multiquery_{n}", _compute) | |
| def hyde_generate(self, question: str, model: str = DEFAULT_MODEL) -> str: | |
| def _compute(q): | |
| return generate(HYDE_PROMPT.format(question=q), model=model).strip() or q | |
| return self._get_or_compute(question, "hyde", _compute) | |