# BaseLanguageModel 是一个抽象基类，是所有语言模型的基类
class BaseLanguageModel(...):
    # 基于用户输入生成prompt
    @abstractmethod
    def generate_prompt(self,prompts: List[PromptValue],stop: Optional[List[str]] = None,...) -> LLMResult:
    @abstractmethod
    def predict(self, text: str, *, stop: Optional[Sequence[str]] = None, **kwargs: Any ) -> str:
    @abstractmethod
    def predict_messages(self, messages: List[BaseMessage],) -> BaseMessage:
# BaseLLM 增加了缓存选项, 回调选项, 持有各种参数
class BaseLLM(BaseLanguageModel[str], ABC):
    1. 覆盖了 __call__ 实现  ==> generate ==> _generate_helper ==> prompt 处理 +  _generate 留给子类实现。 
    2. predict ==> __call__
# LLM类期望它的子类可以更加简单，将大模型的调用方法完全封装，不需要用户实现完整的_generate方法，只需要对外提供一个非常简单的call方法就可以操作LLMs
class LLM(BaseLLM):
    1. _generate ==> _call 留给子类实现。 输入文本格式提示，返回文本格式的答案

class Chain(...,ABC):
    memory: Optional[BaseMemory] = None
    callbacks: Callbacks = Field(default=None, exclude=True)
    @abstractmethod
    def _call(self,inputs: Dict[str, Any],run_manager: Optional[CallbackManagerForChainRun] = None,) -> Dict[str, Any]:
    1. 覆盖了 __call__ 实现 ==> 输入处理 + _call + 输出处理  
    def generate(self,input_list: List[Dict[str, Any]],run_manager)) -> LLMResult:
        prompts, stop = self.prep_prompts(input_list, run_manager=run_manager)
        return self.llm.generate_prompt(prompts,stop,callbacks=callbacks,**self.llm_kwargs,) 

from langchain import PromptTemplate, OpenAI, LLMChain
 
llm = OpenAI(temperature=0.9)
prompt = PromptTemplate.from_template("将下面的句子翻译成英文：{sentence}")
llm_chain = LLMChain(
    llm = llm, 
    prompt = prompt
)
result = llm_chain("今天的天气真不错")
print(result['text'])

class Chain(Serializable, Runnable[Dict[str, Any], Dict[str, Any]], ABC):
    @property
    @abstractmethod
    def input_keys(self) -> List[str]:
    @property
    @abstractmethod
    def output_keys(self) -> List[str]:
    # # Chain 的本质其实就是根据一个 Dict 输入，得到一个 Dict 输出
    def __call__(self,inputs: Union[Dict[str, Any], Any],...) -> Dict[str, Any]:
        inputs = self.prep_inputs(inputs)
        outputs = (self._call(inputs, run_manager=run_manager) if new_arg_supported else self._call(inputs))
class LLMChain(Chain):
    @property
    def input_keys(self) -> List[str]:
        return self.prompt.input_variables
    def _call(self,inputs: Dict[str, Any],run_manager: Optional[CallbackManagerForChainRun] = None,) -> Dict[str, str]:
        response = self.generate([inputs], run_manager=run_manager)
            prompts, stop = self.prep_prompts(input_list, run_manager=run_manager)
                for inputs in input_list:
                    selected_inputs = {k: inputs[k] for k in self.prompt.input_variables}
                    prompt = self.prompt.format_prompt(**selected_inputs)
                    prompts.append(prompt)
            return self.llm.generate_prompt(prompts,stop)  # 
        return self.create_outputs(response)[0]

memory = ConversationBufferMemory()
conversation = ConversationChain(llm=llm, memory=memory, verbose=True)
print(conversation.prompt)
print(conversation.predict(input="1+1=?"))

class Chain(Serializable, Runnable[Dict[str, Any], Dict[str, Any]], ABC):
    memory: Optional[BaseMemory] = None
    callbacks: Callbacks = Field(default=None, exclude=True)
    def invoke( self,input: Dict[str, Any],...) -> Dict[str, Any]:
        return self(input,callbacks=config.get("callbacks"),...)

    @property
    @abstractmethod
    def input_keys(self) -> List[str]:
        """Keys expected to be in the chain input."""
    @property
    @abstractmethod
    def output_keys(self) -> List[str]:
        """Keys expected to be in the chain output."""

    def __call__(self,inputs: Union[Dict[str, Any], Any],callbacks,...)-> Dict[str, Any]:
        inputs = self.prep_inputs(inputs)
        callback_manager = CallbackManager.configure(callbacks,...)
        run_manager = callback_manager.on_chain_start(inputs,...)
        outputs = self._call(inputs, run_manager=run_manager)
        run_manager.on_chain_end(outputs)
        final_outputs = self.prep_outputs(inputs, outputs, return_only_outputs)
        return final_outputs  
        
    @abstractmethod
    def _call(self,inputs: Dict[str, Any],...) -> Dict[str, Any]:   
        """Execute the chain.This is a private method that is not user-facing. It is only called within
            `Chain.__call__`, which is the user-facing wrapper method that handles
            callbacks configuration and some input/output processing.""" 

    def prep_inputs(self, inputs: Union[Dict[str, Any], Any]) -> Dict[str, str]:
        ...
        if self.memory is not None:
            external_context = self.memory.load_memory_variables(inputs)
            inputs = dict(inputs, **external_context)
        return inputs
     def prep_outputs(self,inputs: Dict[str, str],outputs: Dict[str, str],...)-> Dict[str, str]:
        ...
        if self.memory is not None:
            self.memory.save_context(inputs, outputs)     

class BaseMemory(Serializable, ABC):
    @property
    @abstractmethod
    def memory_variables(self) -> List[str]:
        """The string keys this memory class will add to chain inputs."""
    @abstractmethod
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Return key-value pairs given the text input to the chain."""
    @abstractmethod
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save the context of this chain run to memory."""
class BaseChatMemory(BaseMemory, ABC):
    """Abstract base class for chat memory."""
    chat_memory: BaseChatMessageHistory = Field(default_factory=ChatMessageHistory) # 这就是所谓的存在内存里了
    output_key: Optional[str] = None
    input_key: Optional[str] = None
    return_messages: bool = False
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save context from this conversation to buffer."""
        input_str, output_str = self._get_input_output(inputs, outputs)
        self.chat_memory.add_user_message(input_str)
        self.chat_memory.add_ai_message(output_str)
class ConversationBufferMemory(BaseChatMemory):
    memory_key: str = "history"
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]: 
        """Return history buffer."""
        return {self.memory_key: self.buffer} # self.buffer = self.chat_memory.messages

# 消息在内存中的形态
class BaseChatMessageHistory(ABC):
    messages: List[BaseMessage] 
    def add_user_message(self, message: str) -> None:   # 带有Chat字样的类是为聊天设计的，主要保存聊天的历史信息，实现了add_xx_message方法
    def add_ai_message(self, message: str) -> None:
# 消息格式
class BaseMessage(Serializable):
    content: str
    additional_kwargs: dict = Field(default_factory=dict)

loader = TextLoader('./test.txt', encoding='utf8')
docs = loader.load()
print(docs)
# [Document(page_content='ChatGPT是OpenAI开发的一个大型语言模型，...', metadata={'source': './test.txt'})]
text_splitter = CharacterTextSplitter(separator = "\n\n",chunk_size = 1000,chunk_overlap  = 200,length_function = len,is_separator_regex = False,)
texts = text_splitter.create_documents([d.page_content for d in docs])
print(texts)
# [
#	Document(page_content='ChatGPT是OpenAI开发的一个大型语言模型，...', metadata={}), 
#	Document(page_content='我们将探讨如何使用不同的提示工程技术来实现不同的目标。...', metadata={}), 
#	Document(page_content='无论您是普通人、研究人员、开发人员，...', metadata={}), 
#	Document(page_content='在整本书中，...', metadata={})
#]
embeddings = HuggingFaceEmbeddings(model_name='BAAI/bge-large-en',multi_process=True)
db = Chroma.from_documents(texts, embeddings)
query = "ChatGPT是什么？"
docs = db.similarity_search(query)
print(docs[0].page_content)
# ChatGPT是OpenAI开发的一个大型语言模型，可以提供各种主题的信息
retriever = db.as_retriever()
retrieved_docs = retriever.invoke(query)
print(retrieved_docs[0].page_content)

class BaseRetriever(ABC):
    def invoke(self, input: str, config: Optional[RunnableConfig] = None) -> List[Document]:
        config = config or {}
        return self.get_relevant_documents(input,...)     
    def get_relevant_documents(self, query: str, *, callbacks: Callbacks = None, **kwargs: Any) -> List[Document]:
    @abstractmethod
    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
    async def aget_relevant_documents(self, query: str, *, callbacks: Callbacks = None, **kwargs: Any) -> List[Document]:
class VectorStoreRetriever(BaseRetriever):
    def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
        return docs

class MultiQueryRetriever(BaseRetriever):
    retriever: BaseRetriever
    llm_chain: LLMChain
    def _get_relevant_documents(self,query: str,*,run_manager: CallbackManagerForRetrieverRun,) -> List[Document]:
        queries = self.generate_queries(query, run_manager)
        documents = self.retrieve_documents(queries, run_manager)
        return self.unique_union(documents) 
class MultiVectorRetriever(BaseRetriever):
    id_key: str = "doc_id"
    vectorstore: VectorStore
    docstore: BaseStore[str, Document]
    def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerForRetrieverRun) -> List[Document]:
        sub_docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
        ids = [] # We do this to maintain the order of the ids that are returned
        for d in sub_docs:
            if d.metadata[self.id_key] not in ids:
                ids.append(d.metadata[self.id_key])
        docs = self.docstore.mget(ids)
        return [d for d in docs if d is not None]
# 将文档拆分为较小的块，同时每块关联其父文档的id，小块用于提高检索准确度，大块父文档用于返回上下文
class ParentDocumentRetriever(MultiVectorRetriever):
    child_splitter: TextSplitter
    parent_splitter: Optional[TextSplitter] = None
    def add_documents(self,documents: List[Document],ids: Optional[List[str]] = None,add_to_docstore: bool = True,) -> None:
        documents = self.parent_splitter.split_documents(documents)
        docs = []
        full_docs = []
        for i, doc in enumerate(documents):
            sub_docs = self.child_splitter.split_documents([doc])
            for _doc in sub_docs:
                _doc.metadata[self.id_key] = _id
            docs.extend(sub_docs)
            full_docs.append((_id, doc))
        self.vectorstore.add_documents(docs)
        self.docstore.mset(full_docs)

qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff",retriever=self.vector_db.as_retriever())
answer = qa(query)                                          

# langchain/chains/retrieval_qa/base.py
@classmethod
def from_chain_type(cls,llm: BaseLanguageModel,chain_type: str = "stuff",chain_type_kwargs: Optional[dict] = None,**kwargs: Any,) -> BaseRetrievalQA:
    """Load chain from chain type."""
    _chain_type_kwargs = chain_type_kwargs or {}
    combine_documents_chain = load_qa_chain(llm, chain_type=chain_type, **_chain_type_kwargs)
    return cls(combine_documents_chain=combine_documents_chain, **kwargs)
# langchain/chains/question_answering/__init__.py
def load_qa_chain(llm: BaseLanguageModel,chain_type: str = "stuff",verbose: Optional[bool] = None,callback_manager: Optional[BaseCallbackManager] = None,**kwargs: Any,) -> BaseCombineDocumentsChain:
    loader_mapping: Mapping[str, LoadingCallable] = {
        "stuff": _load_stuff_chain,
        "map_reduce": _load_map_reduce_chain,
        "refine": _load_refine_chain,
        "map_rerank": _load_map_rerank_chain,
    }
    return loader_mapping[chain_type](
        llm, verbose=verbose, callback_manager=callback_manager, **kwargs
    )
def _load_stuff_chain(...)-> StuffDocumentsChain:
    _prompt = prompt or stuff_prompt.PROMPT_SELECTOR.get_prompt(llm)
    llm_chain = LLMChain(llm=llm,prompt=_prompt,verbose=verbose,callback_manager=callback_manager,callbacks=callbacks,)
    return StuffDocumentsChain(llm_chain=llm_chain,...) 
class RetrievalQA(BaseRetrievalQA):
    retriever: BaseRetriever = Field(exclude=True)

class AsyncCallbackHandler(BaseCallbackHandler):
    async def on_llm_start(self,prompts: List[str],...)
        """Run when LLM starts running."""
    async def on_llm_end(self,response: LLMResult,...):
        """Run when LLM ends running.""" 

    async def on_chat_model_start(self,serialized: Dict[str, Any],messages: List[List[BaseMessage]],...):
        """Run when a chat model starts running."""
        
    async def on_llm_new_token(self,token: str,...):
        """Run on new LLM token. Only available when streaming is enabled."""
    
    async def on_chain_start(self,inputs: Dict[str, Any],...):
        """Run when chain starts running."""
    async def on_chain_end(self,outputs: Dict[str, Any],...):
        """Run when chain ends running."""
    async def on_chain_error(self,error: BaseException,...):
        """Run when chain errors."""

    async def on_tool_start(self,serialized: Dict[str, Any],input_str: str,...):
        """Run when tool starts running."""
    async def on_tool_end(self,output: str,...):
        """Run when tool ends running."""
    async def on_tool_error(self,error: BaseException, ...):
        """Run when tool errors."""
        

    async def on_agent_action(self,action: AgentAction,...):
        """Run on agent action."""
    async def on_agent_finish(self,finish: AgentFinish,...):
        """Run on agent end."""
        
    async def on_retriever_start(self,serialized: Dict[str, Any],query: str,...):
        """Run on retriever start."""
    async def on_retriever_end(self,documents: Sequence[Document],...):
        """Run on retriever end."""
    async def on_retriever_error(self,error: BaseException,...):
        """Run on retriever error."""

# 创建Chain
chain = prompt | llm | output_parser
# 调用Chain
answer = chain.invoke({'question': '什么是图计算？'})
print(answer)