mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-31 04:45:23 +08:00
261 lines
9.1 KiB
Python
261 lines
9.1 KiB
Python
"""QA Index strategy — generates question-answer pairs from document chunks.
|
||
|
||
Instead of embedding raw chunk text, this strategy uses an LLM to generate
|
||
question-answer pairs from each chunk. The *questions* are embedded and stored
|
||
as vectors, while the *answers* are stored in metadata. At retrieval time the
|
||
user's query (naturally phrased as a question) matches against generated
|
||
questions, yielding better semantic alignment than raw-chunk matching.
|
||
|
||
This strategy yields Q&A pairs incrementally (per chunk) so that the ingest
|
||
pipeline can embed and upsert concurrently with ongoing LLM generation.
|
||
|
||
When *sections* are provided, section-aware chunks are used instead of flat
|
||
text chunks, preserving heading/page metadata in each Q&A pair's metadata.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import re
|
||
from collections.abc import AsyncGenerator
|
||
|
||
from langbot_plugin.api.entities.builtin.provider.message import Message
|
||
|
||
from .base import IndexStrategy
|
||
from components.offload import BoundedOffload
|
||
from ..chunker import (
|
||
chunk_text,
|
||
chunk_sections,
|
||
DEFAULT_CHUNK_SIZE,
|
||
DEFAULT_CHUNK_OVERLAP,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_QUESTIONS_PER_CHUNK = 1
|
||
|
||
QA_GENERATION_PROMPT = """\
|
||
Based on the following text, generate {n} question-answer pairs.
|
||
Each question should be directly answerable from the text.
|
||
Each answer should be concise and grounded in the text.
|
||
|
||
Text:
|
||
{chunk_text}
|
||
|
||
Output exactly {n} pairs in this JSON format (no other text):
|
||
[
|
||
{{"q": "question", "a": "answer"}},
|
||
...
|
||
]"""
|
||
|
||
|
||
def _parse_qa_pairs(text: str) -> list[tuple[str, str]]:
|
||
"""Parse Q&A pairs from LLM response.
|
||
|
||
Tries JSON first, falls back to ``Q:/A:`` text format.
|
||
"""
|
||
# --- Try JSON ---
|
||
try:
|
||
# Extract JSON array even if surrounded by markdown fences
|
||
json_match = re.search(r"\[.*?\]", text, re.DOTALL)
|
||
if json_match:
|
||
items = json.loads(json_match.group())
|
||
pairs = []
|
||
for item in items:
|
||
q = str(item.get("q", "")).strip()
|
||
a = str(item.get("a", "")).strip()
|
||
if q and a:
|
||
pairs.append((q, a))
|
||
if pairs:
|
||
return pairs
|
||
except (json.JSONDecodeError, AttributeError):
|
||
pass
|
||
|
||
# --- Fallback: Q:/A: text format ---
|
||
pairs: list[tuple[str, str]] = []
|
||
blocks = re.split(r"\nQ:\s*", "\n" + text)
|
||
for block in blocks:
|
||
block = block.strip()
|
||
if not block:
|
||
continue
|
||
parts = re.split(r"\nA:\s*", block, maxsplit=1)
|
||
if len(parts) == 2:
|
||
q, a = parts[0].strip(), parts[1].strip()
|
||
if q and a:
|
||
pairs.append((q, a))
|
||
return pairs
|
||
|
||
|
||
def _extract_text(msg: Message) -> str:
|
||
"""Extract plain text from an LLM response Message."""
|
||
if isinstance(msg.content, str):
|
||
return msg.content.strip()
|
||
if isinstance(msg.content, list):
|
||
return "".join(e.text for e in msg.content if e.type == "text").strip()
|
||
return ""
|
||
|
||
|
||
class QAStrategy(IndexStrategy):
|
||
"""QA Index strategy.
|
||
|
||
During ingestion each chunk is sent to an LLM which generates N
|
||
question-answer pairs. The **questions** become the texts that get
|
||
embedded. The **answers** are stored in the ``text`` metadata field so
|
||
they are returned as retrieval results. The original chunk is kept in
|
||
``source_chunk`` for additional context.
|
||
|
||
Yields Q&A pairs per chunk so the ingest pipeline can start embedding
|
||
while the next chunk's LLM call is in flight.
|
||
|
||
When *sections* are provided, section-aware chunks are used, and each
|
||
Q&A pair inherits the chunk's heading/page/heading_path metadata.
|
||
"""
|
||
|
||
async def build_chunks_and_metadata(
|
||
self,
|
||
text: str,
|
||
doc_id: str,
|
||
filename: str,
|
||
creation_settings: dict,
|
||
plugin=None,
|
||
*,
|
||
sections: list | None = None,
|
||
doc_metadata: dict | None = None,
|
||
) -> AsyncGenerator[tuple[list[str], list[str], list[dict]], None]:
|
||
if plugin is None:
|
||
raise RuntimeError("QAStrategy requires a plugin reference for LLM calls")
|
||
|
||
chunk_size = creation_settings.get("chunk_size") or DEFAULT_CHUNK_SIZE
|
||
overlap = creation_settings.get("overlap") or DEFAULT_CHUNK_OVERLAP
|
||
qa_llm_uuid = creation_settings.get("qa_llm_model_uuid", "")
|
||
n_questions = (
|
||
creation_settings.get("questions_per_chunk") or DEFAULT_QUESTIONS_PER_CHUNK
|
||
)
|
||
doc_fields = self._build_doc_meta_fields(doc_metadata)
|
||
offload = getattr(plugin, "offload", None) or BoundedOffload()
|
||
|
||
if not qa_llm_uuid:
|
||
raise ValueError(
|
||
"QA Index requires 'qa_llm_model_uuid' in creation settings"
|
||
)
|
||
|
||
# Build chunk list — section-aware or flat
|
||
if sections:
|
||
s_chunks = await offload.run(chunk_sections, sections, chunk_size, overlap)
|
||
# Wrap into a uniform structure for the loop below
|
||
chunk_items = [
|
||
{
|
||
"text": sc.text,
|
||
"heading": sc.heading,
|
||
"level": sc.level,
|
||
"page": sc.page,
|
||
"heading_path": sc.heading_path,
|
||
"has_table": sc.has_table,
|
||
"section_index": sc.section_index,
|
||
}
|
||
for sc in s_chunks
|
||
]
|
||
else:
|
||
flat_chunks = await offload.run(chunk_text, text, chunk_size, overlap)
|
||
chunk_items = [{"text": c} for c in flat_chunks]
|
||
|
||
logger.info(
|
||
f"[QA] Generating Q&A pairs: {len(chunk_items)} chunks × "
|
||
f"{n_questions} questions, LLM={qa_llm_uuid}"
|
||
)
|
||
|
||
total_qa = 0
|
||
|
||
for chunk_idx, chunk_info in enumerate(chunk_items):
|
||
chunk = chunk_info["text"]
|
||
prompt = QA_GENERATION_PROMPT.format(n=n_questions, chunk_text=chunk)
|
||
messages = [Message(role="user", content=prompt)]
|
||
|
||
try:
|
||
response = await plugin.invoke_llm(qa_llm_uuid, messages)
|
||
response_text = _extract_text(response)
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"[QA] LLM call failed for chunk {chunk_idx}: {e}, skipping"
|
||
)
|
||
continue
|
||
|
||
qa_pairs = _parse_qa_pairs(response_text)
|
||
if not qa_pairs:
|
||
logger.warning(
|
||
f"[QA] No Q&A pairs parsed from chunk {chunk_idx}, "
|
||
f"raw response: {response_text[:200]!r}"
|
||
)
|
||
continue
|
||
|
||
logger.info(
|
||
f"[QA] Chunk {chunk_idx}/{len(chunk_items) - 1}: "
|
||
f"generated {len(qa_pairs)} Q&A pairs"
|
||
)
|
||
|
||
questions: list[str] = []
|
||
ids: list[str] = []
|
||
metadatas: list[dict] = []
|
||
|
||
for qa_idx, (question, answer) in enumerate(qa_pairs):
|
||
vec_id = f"{doc_id}_{chunk_idx}_qa{qa_idx}"
|
||
questions.append(question)
|
||
ids.append(vec_id)
|
||
meta = {
|
||
"file_id": doc_id,
|
||
"document_id": doc_id,
|
||
"document_name": filename,
|
||
"chunk_index": chunk_idx,
|
||
"qa_index": qa_idx,
|
||
"question": question,
|
||
"answer": answer,
|
||
"text": answer,
|
||
"source_chunk": chunk,
|
||
"index_type": "qa",
|
||
}
|
||
# Add section metadata if available
|
||
if "heading" in chunk_info:
|
||
meta["heading"] = chunk_info["heading"]
|
||
meta["level"] = chunk_info["level"]
|
||
meta["page"] = chunk_info["page"]
|
||
meta["heading_path"] = chunk_info["heading_path"]
|
||
meta["has_table"] = chunk_info["has_table"]
|
||
meta["section_index"] = chunk_info["section_index"]
|
||
meta.update(doc_fields)
|
||
metadatas.append(meta)
|
||
|
||
total_qa += len(questions)
|
||
yield questions, ids, metadatas
|
||
|
||
logger.info(f"[QA] Total: {total_qa} Q&A pairs from {len(chunk_items)} chunks")
|
||
|
||
def postprocess_results(self, results: list[dict], top_k: int) -> list[dict]:
|
||
"""Deduplicate by source chunk, keeping the highest-scoring QA per chunk."""
|
||
|
||
def _metric(res: dict) -> tuple[int, float]:
|
||
distance = res.get("distance")
|
||
if isinstance(distance, (int, float)):
|
||
return (0, float(distance))
|
||
|
||
score = res.get("score")
|
||
if isinstance(score, (int, float)):
|
||
return (1, float(score))
|
||
|
||
return (2, float("inf"))
|
||
|
||
seen: dict[str, dict] = {}
|
||
for res in results:
|
||
meta = res.get("metadata", {})
|
||
doc_id = meta.get("document_id", "")
|
||
chunk_idx = meta.get("chunk_index", "")
|
||
key = f"{doc_id}:{chunk_idx}"
|
||
|
||
if key not in seen:
|
||
seen[key] = res
|
||
else:
|
||
if _metric(res) < _metric(seen[key]):
|
||
seen[key] = res
|
||
|
||
return list(seen.values())[:top_k]
|