| """Default LLM-backed RAG operations (the evaluated v2 prompts and parsers).""" |
| from __future__ import annotations |
|
|
| import asyncio |
| import inspect |
| from typing import Any, Callable, List, Optional, Sequence |
|
|
| from . import prompts as P |
| from .documents import Document |
| from .llms import acomplete, as_llm |
|
|
|
|
| class LLMGenerator: |
| """Extractive answer generation; returns '' to abstain.""" |
|
|
| def __init__(self, llm: Any, prompt: Callable[[str, Sequence[Document]], str] = P.generate_prompt, |
| parse: Callable[[str], str] = P.parse_answer) -> None: |
| self.llm, self.prompt, self.parse = as_llm(llm), prompt, parse |
|
|
| def __call__(self, question: str, docs: Sequence[Document]) -> str: |
| return self.parse(self.llm.complete(self.prompt(question, docs))) |
|
|
| async def acall(self, question: str, docs: Sequence[Document]) -> str: |
| return self.parse(await acomplete(self.llm, self.prompt(question, docs))) |
|
|
|
|
| class LLMGrader: |
| """Relevance grading: one score in {0, 1, 2} per passage; passages scoring >= 1 are kept.""" |
|
|
| def __init__(self, llm: Any, prompt: Callable[[str, Sequence[Document]], str] = P.grade_prompt, |
| parse: Callable[[str, int], List[int]] = P.parse_grades) -> None: |
| self.llm, self.prompt, self.parse = as_llm(llm), prompt, parse |
|
|
| def __call__(self, question: str, docs: Sequence[Document]) -> List[int]: |
| return self.parse(self.llm.complete(self.prompt(question, docs)), len(docs)) |
|
|
| async def acall(self, question: str, docs: Sequence[Document]) -> List[int]: |
| return self.parse(await acomplete(self.llm, self.prompt(question, docs)), len(docs)) |
|
|
|
|
| class LLMRewriter: |
| """Query rewriting for retrieval; '' keeps the original question.""" |
|
|
| def __init__(self, llm: Any, prompt: Callable[[str], str] = P.rewrite_prompt, |
| parse: Callable[[str], str] = P.parse_query) -> None: |
| self.llm, self.prompt, self.parse = as_llm(llm), prompt, parse |
|
|
| def __call__(self, question: str) -> str: |
| return self.parse(self.llm.complete(self.prompt(question))) |
|
|
| async def acall(self, question: str) -> str: |
| return self.parse(await acomplete(self.llm, self.prompt(question))) |
|
|
|
|
| class LLMSupportChecker: |
| """Support check: does the passage that contains the answer support it?""" |
|
|
| def __init__(self, llm: Any, prompt: Callable[[str, Document, str], str] = P.verify_prompt, |
| parse: Callable[[str], bool] = P.parse_yes) -> None: |
| self.llm, self.prompt, self.parse = as_llm(llm), prompt, parse |
|
|
| def __call__(self, question: str, passage: Document, answer: str) -> bool: |
| return bool(self.parse(self.llm.complete(self.prompt(question, passage, answer)))) |
|
|
| async def acall(self, question: str, passage: Document, answer: str) -> bool: |
| return bool(self.parse(await acomplete(self.llm, self.prompt(question, passage, answer)))) |
|
|
|
|
| async def acall(fn: Any, *args: Any) -> Any: |
| """Call an operation natively async when it supports it, otherwise in a worker thread.""" |
| a = getattr(fn, "acall", None) |
| if a is not None: |
| return await a(*args) |
| if inspect.iscoroutinefunction(fn) or inspect.iscoroutinefunction(getattr(fn, "__call__", None)): |
| return await fn(*args) |
| return await asyncio.to_thread(fn, *args) |
|
|
