import tiktoken import itertools from environs import Env from pathlib import Path from observability import span from langchain.chains.base import Chain from example import EmbeddingExampleSelector from typing import Any, Dict, List, Optional from langchain.docstore.document import Document from langchain.prompts.base import BasePromptTemplate from langchain.base_language import BaseLanguageModel from langchain.schema import BaseRetriever, AIMessage, HumanMessage from langchain.prompts import MessagesPlaceholder, FewShotPromptTemplate, PromptTemplate from langchain.prompts.chat import ( ChatPromptTemplate, SystemMessagePromptTemplate, ) from langchain.callbacks.manager import ( AsyncCallbackManagerForChainRun, CallbackManagerForChainRun, ) # s1 = '12345' s2 = '45678' difference(s1, s2) = '123' def difference(s1, s2): s1_size = len(s1) s2_size = len(s2) for i in range(min(s1_size, s2_size), 0, -1): if s1[s1_size-i:] == s2[:i]: return s1[:s1_size-i] return s1 def compose_doc(docs: List[Document], num) -> str: chunks = f'=== {docs[0].metadata["title"]} ===\n\n' for doc in docs: chunks = difference(chunks, doc.page_content) + doc.page_content return f'NUMBER:{num}\n' + chunks def document(docs: List[Document]) -> str: # group by link docs.sort(key=lambda d: d.metadata['link']) groups = itertools.groupby(docs, key=lambda d: d.metadata['link']) # sort by idx and compose doc doc_str = '' i = 0 for _, g in groups: i += 1 doc_str += compose_doc(sorted(g, key=lambda d: d.metadata['idx']), i) + '\n' return doc_str #return '\n'.join(compose_doc(sorted(g, key=lambda d: d.metadata['idx'])) for _, g in groups) class ResponseChain(Chain): llm: BaseLanguageModel retriever: BaseRetriever chat_prompt: BasePromptTemplate toxicity_prompt: BasePromptTemplate enc: tiktoken.core.Encoding @property def input_keys(self) -> List[str]: return [ 'revised_query', 'messages', 'toxicity', ] @property def output_keys(self) -> List[str]: return [ 'content', 'links', 'recalls', ] @property def _chain_type(self) -> str: return "response_chain" @span(name="compose document") def get_document(self, docs: List[Document], max_tokens: int=1024) -> (str, int): for idx in range(len(docs)): if len(self.enc.encode(document(docs[:idx+1]))) > max_tokens: return document(docs[:idx]), idx return document(docs), len(docs) @span(name="response chain", input_fn=lambda args, kwargs: args[0]) def _call( self, inputs: Dict[str, Any], /, run_manager: Optional[CallbackManagerForChainRun] = None, ) -> Dict[str, str]: messages = [] for m in inputs['messages']: if m['role'] == 'user': messages.append(HumanMessage(content=m['content'])) elif m['role'] == 'assistant': messages.append(AIMessage(content=m['content'])) else: raise ValueError(f'unknown role {m["role"]}') if inputs['toxicity']: prompt_value = self.toxicity_prompt.format_prompt(instruction=inputs['revised_query']) response = self.llm.generate_prompt( [prompt_value], callbacks=run_manager.get_child() if run_manager else None, ) return {'content': response.generations[0][0].text, 'links': [], 'recalls': []} docs = self.retriever.get_relevant_documents(inputs['revised_query']) (doc, doc_idx) = self.get_document(docs=docs) prompt_value = self.chat_prompt.format_prompt(document=doc, conversation=messages) response = self.llm.generate_prompt( [prompt_value], callbacks=run_manager.get_child() if run_manager else None, ) return { 'content': response.generations[0][0].text, 'links': [{ 'title': d.metadata['title'], 'link': d.metadata['link'], 'num': i + 1, 'namespace': d.metadata['namespace'], } for i, d in enumerate({d.metadata['link']: d for d in docs[:doc_idx]}.values())], 'recalls': [{ **d.metadata, 'content': d.page_content, } for d in docs[:doc_idx]], } @span(name="response chain", input_fn=lambda args, kwargs: args[0]) async def _acall( self, inputs: Dict[str, Any], /, run_manager: Optional[AsyncCallbackManagerForChainRun] = None, ) -> Dict[str, str]: messages = [] for m in inputs['messages']: if m['role'] == 'user': messages.append(HumanMessage(content=m['content'])) elif m['role'] == 'assistant': messages.append(AIMessage(content=m['content'])) else: raise ValueError(f'unknown role {m["role"]}') if inputs['toxicity']: prompt_value = self.toxicity_prompt.format_prompt(instruction=inputs['revised_query']) response = await self.llm.agenerate_prompt( [prompt_value], callbacks=run_manager.get_child() if run_manager else None, ) return {'content': response.generations[0][0].text, 'links': [], 'recalls': []} docs = self.retriever.get_relevant_documents(inputs['revised_query']) (doc, doc_idx) = self.get_document(docs=docs) prompt_value = self.chat_prompt.format_prompt(document=doc, conversation=messages) response = await self.llm.agenerate_prompt( [prompt_value], callbacks=run_manager.get_child() if run_manager else None, ) return { 'content': response.generations[0][0].text, 'links': [{ 'title': d.metadata['title'], 'link': d.metadata['link'], 'num': i + 1, 'namespace': d.metadata['namespace'], } for i, d in enumerate({d.metadata['link']: d for d in docs[:doc_idx]}.values())], 'recalls': [{ **d.metadata, 'content': d.page_content, } for d in docs[:doc_idx]], } def load_response_chain(env: Env, llm: BaseLanguageModel, retriever: BaseRetriever) -> Chain: chat_prompt = ChatPromptTemplate.from_messages([ SystemMessagePromptTemplate.from_template( Path("prompts/response_system.tpl").read_text(), ), MessagesPlaceholder(variable_name="conversation"), ]) e = EmbeddingExampleSelector( env.str('EMBEDDING_SERVICE_ENDPOINT'), 'chatbot_toxicity_response', input_fn=lambda x: {'query': x['instruction']}, output_fn=lambda x: {'instruction': x['instruction'], 'answer': x['answer']}, ) toxicity_prompt = FewShotPromptTemplate( example_selector=e, example_prompt=PromptTemplate.from_template( Path("prompts/toxicity_response_example.tpl").read_text() ), prefix=Path("prompts/toxicity_response_prefix.tpl").read_text(), suffix=Path("prompts/toxicity_response_suffix.tpl").read_text(), input_variables=["instruction"], ) return ResponseChain( llm=llm, retriever=retriever, chat_prompt=chat_prompt, toxicity_prompt=toxicity_prompt, enc=tiktoken.get_encoding("cl100k_base"), )