from loguru import logger from .helper import ErrorCode, LogManager from .retriever import CacheRetriever from .inferencer import LLMInference from .feature_database import DocumentProcessor, FeatureDataBase class ChatAgent: def __init__(self, config, tensor_parallel_size) -> None: self.work_dir = config['default']['work_dir'] self.embedding_model_path = config['feature_database']['embedding_model_path'] self.reranker_model_path = config['feature_database']['reranker_model_path'] reject_throttle = float(config['feature_database']['reject_throttle']) local_llm_path = config['llm']['local_llm_path'] accelerate = config.getboolean('llm', 'accelerate') self.retriever = CacheRetriever(self.embedding_model_path, self.reranker_model_path).get(reject_throttle=reject_throttle, work_dir=self.work_dir) self.llm_server = LLMInference(local_llm_path, tensor_parallel_size, accelerate=accelerate) def generate_prompt(self, history_pair, instruction: str, template: str, context: str = ''): if context is not None and len(context) > 0: instruction = template.format(context, instruction) real_history = [] for pair in history_pair: if pair[0] is None or pair[1] is None: continue if len(pair[0]) < 1 or len(pair[1]) < 1: continue real_history.append(pair) return instruction, real_history def call_rag_retrieve(self, query): return self.retriever.query(query) def call_llm_response(self, prompt, history=None): text, error = self.llm_server.generate_response(prompt=prompt, history=history) return text def parse_file_and_merge(self, file_dir): file_opr = DocumentProcessor() files = file_opr.scan_directory(repo_dir=file_dir) file_handler = FeatureDataBase(embeddings=self.retriever.embeddings, reranker=self.retriever.reranker) file_handler.preprocess(files=files, work_dir=self.work_dir, file_opr=file_opr) file_handler.merge_db_response(self.retriever.vector_store, files=files, work_dir=self.work_dir, file_opr=file_opr) file_opr.summarize(files) self.retriever = CacheRetriever(self.embedding_model_path, self.reranker_model_path).get(work_dir=self.work_dir) class Worker: def __init__(self, config, tensor_parallel_size): self.agent = ChatAgent(config, tensor_parallel_size) self.TOPIC_TEMPLATE = '告诉我这句话的主题，直接说主题不要解释：“{}”' self.SCORING_RELAVANCE_TEMPLATE = '问题：“{}”\n材料：“{}”\n请仔细阅读以上内容，材料里为一个列表，列表里面有若干子列表，请判断每个子列表的内容和问题的相关度，不要解释直接给出相关度得分列表并以空格分隔，用0～10表示。判断标准：非常相关得 10 分；完全没关联得 0 分。\n' # noqa E501 self.KEYWORDS_TEMPLATE = '谷歌搜索是一个通用搜索引擎，可用于访问互联网、查询百科知识、了解时事新闻等。搜索参数类型 string， 内容是短语或关键字，以空格分隔。\n你现在是搜搜小助手，用户提问“{}”，你打算通过谷歌搜索查询相关资料，请提供用于搜索的关键字或短语，不要解释直接给出关键字或短语。' # noqa E501 self.SECURITY_TEMAPLTE = '判断以下句子是否涉及政治、辱骂、色情、恐暴、宗教、网络暴力、种族歧视等违禁内容，结果用 0～10 表示，不要解释直接给出得分。判断标准：涉其中任一问题直接得 10 分；完全不涉及得 0 分。直接给得分不要解释：“{}”' # noqa E501 self.PERPLESITY_TEMPLATE = '“question:{} answer:{}”\n阅读以上对话，answer 是否在表达自己不知道，回答越全面得分越少，用0～10表示，不要解释直接给出得分。\n判断标准：准确回答问题得 0 分；答案详尽得 1 分；知道部分答案但有不确定信息得 8 分；知道小部分答案但推荐求助其他人得 9 分；不知道任何答案直接推荐求助别人得 10 分。直接打分不要解释。' # noqa E501 self.SUMMARIZE_TEMPLATE = '{} \n 仔细阅读以上内容，总结得简短有力点' # noqa E501 self.GENERATE_TEMPLATE = '“{}” \n问题：“{}” \n请仔细阅读上述文字, 并使用markdown格式回答问题，直接给出回答不做任何解释。' # noqa E501 self.MARKDOWN_TEMPLATE = '问题：“{}” \n请使用markdown格式回答此问题' def judgment_results(self, query, chunks, throttle): relation_score = self.agent.call_llm_response( prompt=self.SCORING_RELAVANCE_TEMPLATE.format(query, chunks)) logger.info('score: %s' % [relation_score, throttle]) # 过滤操作 filtered_chunks = [] for chunk, score in zip(chunks, relation_score.split()): if float(score) >= float(throttle): filtered_chunks.append(chunk) return filtered_chunks def extract_topic(self, query): topic = self.agent.call_llm_response(self.TOPIC_TEMPLATE.format(query)) return topic def response_direct_by_llm(self, query): # Compliant check prompt = self.SECURITY_TEMAPLTE.format(query) score = self.agent.call_llm_response(prompt=prompt) logger.debug("score:{}, prompt:{}".format(score, prompt)) if int(score) > 5: return ErrorCode.NON_COMPLIANCE_QUESTION, "您的问题中涉及敏感话题，请重新提问。", None logger.info('LLM direct response and prompt is: {}'.format(query)) prompt = self.MARKDOWN_TEMPLATE.format(query) response_direct = self.agent.call_llm_response(prompt=prompt) return ErrorCode.NOT_FIND_RELATED_DOCS, response_direct, None def produce_response(self, query, history, judgment, topic=False, rag=True): response = '' references = [] if query is None: return ErrorCode.NOT_A_QUESTION, response, references logger.info('input: %s' % [query, history]) if rag: if topic: query = self.extract_topic(query) logger.info('topic: %s' % query) if len(query) <= 0: return ErrorCode.NO_TOPIC, response, references chunks, references = self.agent.call_rag_retrieve(query) if len(chunks) == 0: return self.response_direct_by_llm(query) if judgment: chunks = self.judgment_results( query, chunks, throttle=5, ) # 如果DataBase检索到了，就用检索到的块去回答 if len(chunks) > 0: prompt, history = self.agent.generate_prompt( instruction=query, context=chunks, history_pair=history, template=self.GENERATE_TEMPLATE) logger.debug('prompt: {}'.format(prompt)) response = self.agent.call_llm_response(prompt=prompt, history=history) return ErrorCode.SUCCESS, response, references else: return self.response_direct_by_llm(query)