Source code for langchain.chains.flare.base
from __future__ import annotations
import re
from abc import abstractmethod
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
from langchain.callbacks.manager import (
    CallbackManagerForChainRun,
)
from langchain.chains.base import Chain
from langchain.chains.flare.prompts import (
    PROMPT,
    QUESTION_GENERATOR_PROMPT,
    FinishedOutputParser,
)
from langchain.chains.llm import LLMChain
from langchain.llms import OpenAI
from langchain.pydantic_v1 import Field
from langchain.schema import BasePromptTemplate, BaseRetriever, Generation
from langchain.schema.language_model import BaseLanguageModel
class _ResponseChain(LLMChain):
    """Base class for chains that generate responses."""
    prompt: BasePromptTemplate = PROMPT
    @property
    def input_keys(self) -> List[str]:
        return self.prompt.input_variables
    def generate_tokens_and_log_probs(
        self,
        _input: Dict[str, Any],
        *,
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Tuple[Sequence[str], Sequence[float]]:
        llm_result = self.generate([_input], run_manager=run_manager)
        return self._extract_tokens_and_log_probs(llm_result.generations[0])
    @abstractmethod
    def _extract_tokens_and_log_probs(
        self, generations: List[Generation]
    ) -> Tuple[Sequence[str], Sequence[float]]:
        """Extract tokens and log probs from response."""
class _OpenAIResponseChain(_ResponseChain):
    """Chain that generates responses from user input and context."""
    llm: OpenAI = Field(
        default_factory=lambda: OpenAI(