from __future__ import annotations import re from typing import Any, Dict, List, Optional, Sequence, Tuple import numpy as np from langchain_core.callbacks import ( CallbackManagerForChainRun, ) from langchain_core.language_models import BaseLanguageModel from langchain_core.messages import AIMessage from langchain_core.output_parsers import StrOutputParser from langchain_core.prompts import BasePromptTemplate from langchain_core.retrievers import BaseRetriever from langchain_core.runnables import Runnable from pydantic import Field from langchain.chains.base import Chain from langchain.chains.flare.prompts import ( PROMPT, QUESTION_GENERATOR_PROMPT, FinishedOutputParser, ) from langchain.chains.llm import LLMChain def _extract_tokens_and_log_probs(response: AIMessage) -> Tuple[List[str], List[float]]: """Extract tokens and log probabilities from chat model response.""" tokens = [] log_probs = [] for token in response.response_metadata["logprobs"]["content"]: tokens.append(token["token"]) log_probs.append(token["logprob"]) return tokens, log_probs class QuestionGeneratorChain(LLMChain): """Chain that generates questions from uncertain spans.""" prompt: BasePromptTemplate = QUESTION_GENERATOR_PROMPT """Prompt template for the chain.""" @classmethod def is_lc_serializable(cls) -> bool: return False @property def input_keys(self) -> List[str]: """Input keys for the chain.""" return ["user_input", "context", "response"] def _low_confidence_spans( tokens: Sequence[str], log_probs: Sequence[float], min_prob: float, min_token_gap: int, num_pad_tokens: int, ) -> List[str]: _low_idx = np.where(np.exp(log_probs) < min_prob)[0] low_idx = [i for i in _low_idx if re.search(r"\w", tokens[i])] if len(low_idx) == 0: return [] spans = [[low_idx[0], low_idx[0] + num_pad_tokens + 1]] for i, idx in enumerate(low_idx[1:]): end = idx + num_pad_tokens + 1 if idx - low_idx[i] < min_token_gap: spans[-1][1] = end else: spans.append([idx, end]) return ["".join(tokens[start:end]) for start, end in spans] class FlareChain(Chain): """Chain that combines a retriever, a question generator, and a response generator. See [Active Retrieval Augmented Generation](https://arxiv.org/abs/2305.06983) paper. """ question_generator_chain: Runnable """Chain that generates questions from uncertain spans.""" response_chain: Runnable """Chain that generates responses from user input and context.""" output_parser: FinishedOutputParser = Field(default_factory=FinishedOutputParser) """Parser that determines whether the chain is finished.""" retriever: BaseRetriever """Retriever that retrieves relevant documents from a user input.""" min_prob: float = 0.2 """Minimum probability for a token to be considered low confidence.""" min_token_gap: int = 5 """Minimum number of tokens between two low confidence spans.""" num_pad_tokens: int = 2 """Number of tokens to pad around a low confidence span.""" max_iter: int = 10 """Maximum number of iterations.""" start_with_retrieval: bool = True """Whether to start with retrieval.""" @property def input_keys(self) -> List[str]: """Input keys for the chain.""" return ["user_input"] @property def output_keys(self) -> List[str]: """Output keys for the chain.""" return ["response"] def _do_generation( self, questions: List[str], user_input: str, response: str, _run_manager: CallbackManagerForChainRun, ) -> Tuple[str, bool]: callbacks = _run_manager.get_child() docs = [] for question in questions: docs.extend(self.retriever.invoke(question)) context = "\n\n".join(d.page_content for d in docs) result = self.response_chain.invoke( { "user_input": user_input, "context": context, "response": response, }, {"callbacks": callbacks}, ) if isinstance(result, AIMessage): result = result.content marginal, finished = self.output_parser.parse(result) return marginal, finished def _do_retrieval( self, low_confidence_spans: List[str], _run_manager: CallbackManagerForChainRun, user_input: str, response: str, initial_response: str, ) -> Tuple[str, bool]: question_gen_inputs = [ { "user_input": user_input, "current_response": initial_response, "uncertain_span": span, } for span in low_confidence_spans ] callbacks = _run_manager.get_child() if isinstance(self.question_generator_chain, LLMChain): question_gen_outputs = self.question_generator_chain.apply( question_gen_inputs, callbacks=callbacks ) questions = [ output[self.question_generator_chain.output_keys[0]] for output in question_gen_outputs ] else: questions = self.question_generator_chain.batch( question_gen_inputs, config={"callbacks": callbacks} ) _run_manager.on_text( f"Generated Questions: {questions}", color="yellow", end="\n" ) return self._do_generation(questions, user_input, response, _run_manager) def _call( self, inputs: Dict[str, Any], run_manager: Optional[CallbackManagerForChainRun] = None, ) -> Dict[str, Any]: _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager() user_input = inputs[self.input_keys[0]] response = "" for i in range(self.max_iter): _run_manager.on_text( f"Current Response: {response}", color="blue", end="\n" ) _input = {"user_input": user_input, "context": "", "response": response} tokens, log_probs = _extract_tokens_and_log_probs( self.response_chain.invoke( _input, {"callbacks": _run_manager.get_child()} ) ) low_confidence_spans = _low_confidence_spans( tokens, log_probs, self.min_prob, self.min_token_gap, self.num_pad_tokens, ) initial_response = response.strip() + " " + "".join(tokens) if not low_confidence_spans: response = initial_response final_response, finished = self.output_parser.parse(response) if finished: return {self.output_keys[0]: final_response} continue marginal, finished = self._do_retrieval( low_confidence_spans, _run_manager, user_input, response, initial_response, ) response = response.strip() + " " + marginal if finished: break return {self.output_keys[0]: response} @classmethod def from_llm( cls, llm: BaseLanguageModel, max_generation_len: int = 32, **kwargs: Any ) -> FlareChain: """Creates a FlareChain from a language model. Args: llm: Language model to use. max_generation_len: Maximum length of the generated response. kwargs: Additional arguments to pass to the constructor. Returns: FlareChain class with the given language model. """ try: from langchain_openai import ChatOpenAI except ImportError: raise ImportError( "OpenAI is required for FlareChain. " "Please install langchain-openai." "pip install langchain-openai" ) llm = ChatOpenAI(max_tokens=max_generation_len, logprobs=True, temperature=0) response_chain = PROMPT | llm question_gen_chain = QUESTION_GENERATOR_PROMPT | llm | StrOutputParser() return cls( question_generator_chain=question_gen_chain, response_chain=response_chain, **kwargs, )