# Copyright (c) OpenMMLab. All rights reserved. """Pipeline.""" import argparse import datetime import re import time import os import json import pytoml from loguru import logger from .helper import ErrorCode, QueryTracker from .llm_client import ChatClient from .retriever import CacheRetriever from .sg_search import SourceGraphProxy from .web_search import WebSearch def convertid2url(text): # Regular expression to find all PMC references pattern = r"\[PMC(\d+)\]" # Function to replace each match with a URL link replacement = lambda match: f"[PMC{match.group(1)}](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC{match.group(1)}/)" # Replace all occurrences in the text formatted_text = re.sub(pattern, replacement, text) return formatted_text class Worker: """The Worker class orchestrates the logic of handling user queries, generating responses and managing several aspects of a chat assistant. It enables feature storage, language model client setup, time scheduling and much more. Attributes: llm: A ChatClient instance that communicates with the language model. fs: An instance of FeatureStore for loading and querying features. config_path: A string indicating the path of the configuration file. config: A dictionary holding the configuration settings. language: A string indicating the language of the chat, default is 'zh' (Chinese). # noqa E501 context_max_length: An integer representing the maximum length of the context used by the language model. # noqa E501 Several template strings for various prompts are also defined. """ def __init__(self, work_dir: str, config_path: str, language: str = 'zh'): """Constructs all the necessary attributes for the worker object. Args: work_dir (str): The working directory where feature files are located. config_path (str): The location of the configuration file. language (str, optional): Specifies the language to be used. Defaults to 'zh' (Chinese). # noqa E501 """ self.llm = ChatClient(config_path=config_path) #每次实例化worker都会重新实例化chatclient并读取本地最新的配置文件 self.retriever = CacheRetriever(config_path=config_path).get(work_dir=work_dir) self.config_path = config_path self.config = None self.language = language with open(config_path, encoding='utf8') as f: self.config = pytoml.load(f) if self.config is None: raise Exception('worker config can not be None') llm_config = self.config['llm'] self.context_max_length = llm_config['server']['local_llm_max_text_length'] self.logsavepath = os.path.join(self.config['feature_store']['work_dir'],self.config['worker']['save_path']) if not os.path.exists(self.logsavepath.replace(self.logsavepath.split('/')[-1],'')): os.makedirs(self.logsavepath.replace(self.logsavepath.split('/')[-1],'')) if llm_config['enable_remote']: self.context_max_length = llm_config['server'][ 'remote_llm_max_text_length'] # Switch languages according to the scenario. if self.language == 'zh': self.TOPIC_TEMPLATE = '告诉我这句话的主题，直接说主题不要解释：“{}”' self.SCORING_QUESTION_TEMPLTE = '“{}”\n请仔细阅读以上内容，判断句子是否是个有主题的疑问句，结果用 0～10 表示。直接提供得分不要解释。\n判断标准：有主语谓语宾语并且是疑问句得 10 分；缺少主谓宾扣分；陈述句直接得 0 分；不是疑问句直接得 0 分。直接提供得分不要解释。' # noqa E501 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 请仔细阅读参考材料回答问题，材料可能和问题无关。如果材料和问题无关，尝试用你自己的理解来回答问题。如果无法确定答案，直接回答不知道。' # noqa E501 self.GENERATE_TEMPLATE = '材料：“{}”\n 问题：“{}” \n 请仔细阅读参考材料回答问题，回答中附上对应内容的参考id，例如：治疗方法的主要决定因素是年龄、合并症和诊断分子特征[PMC9958586]' # yyj self.ANNOTATE_CLUSTER = '这是关于{}的不同论文的分块句子，它们通过相似性进行了聚类，以下是其中一个聚类的10个样本：“{}”\n 请用一句话标注这个聚类。' # noqa E501 self.INSPIRATION_TEMPLATE = '以下是一些有关{0}的文章内容的总结 {1}，请提出一个关于{0}的综述子问题，一个问题即可。' else: self.TOPIC_TEMPLATE = 'Tell me the theme of this sentence, just state the theme without explanation: "{}"' # noqa E501 self.SCORING_QUESTION_TEMPLTE = '"{}"\nPlease read the content above carefully and judge whether the sentence is a thematic question. Rate it on a scale of 0-10. Only provide the score, no explanation.\nThe criteria are as follows: a sentence gets 10 points if it has a subject, predicate, object and is a question; points are deducted for missing subject, predicate or object; declarative sentences get 0 points; sentences that are not questions also get 0 points. Just give the score, no explanation.' # noqa E501 self.SCORING_RELAVANCE_TEMPLATE = 'Question: "{}", Background Information: "{}"\nPlease read the content above carefully and assess the relevance between the question and the material on a scale of 0-10. The scoring standard is as follows: extremely relevant gets 10 points; completely irrelevant gets 0 points. Only provide the score, no explanation needed.' # noqa E501 self.KEYWORDS_TEMPLATE = 'Google search is a general-purpose search engine that can be used to access the internet, look up encyclopedic knowledge, keep abreast of current affairs and more. Search parameters type: string, content consists of phrases or keywords separated by spaces.\nYou are now the assistant in the "{}" communication group. A user asked "{}", you plan to use Google search to find related information, please provide the keywords or phrases for the search, no explanation, just give the keywords or phrases.' # noqa E501 self.SECURITY_TEMAPLTE = 'Evaluate whether the following sentence involves prohibited content such as politics, insult, pornography, terror, religion, cyber violence, racial discrimination, etc., rate it on a scale of 0-10, do not explain, just give the score. The scoring standard is as follows: any violation directly gets 10 points; completely unrelated gets 0 points. Give the score, no explanation: "{}"' # noqa E501 self.PERPLESITY_TEMPLATE = 'Question: {} Answer: {}\nRead the dialogue above, does the answer express that they don\'t know? The more comprehensive the answer, the lower the score. Rate it on a scale of 0-10, no explanation, just give the score.\nThe scoring standard is as follows: an accurate answer to the question gets 0 points; a detailed answer gets 1 point; knowing some answers but having uncertain information gets 8 points; knowing a small part of the answer but recommends seeking help from others gets 9 points; not knowing any of the answers and directly recommending asking others for help gets 10 points. Just give the score, no explanation.' # noqa E501 self.SUMMARIZE_TEMPLATE = '"{}" \n Read the content above carefully, summarize it in a short and powerful way.' # noqa E501 self.GENERATE_TEMPLATE = 'Background Information: "{}"\n Question: "{}"\n Please read the reference material carefully and answer the question with reference id at the end of the corresponding content for example: Primary determinants of the therapeutic approach are age, comorbidities, and diagnostic molecular profile [PMC9958586]' # yyj self.ANNOTATE_CLUSTER = 'these are chunklized sentences from different papers about{}, they are clustered by similarity, the following is 10 samples from one of the cluster: "{}"\n Please tag the cluster in one breif sentence.' self.INSPIRATION_TEMPLATE = 'Given the following summary of the articles content about {0} {1}, give some idea or sub-questions of the review about {0}, one question is sufficient.' # noqa E501 def single_judge(self, prompt, tracker, throttle: int, default: int,backend): """Generates a score based on the prompt, and then compares it to threshold. Args: prompt (str): The prompt for the language model. tracker (obj): An instance of QueryTracker logs the operations. throttle (int): Threshold value to compare the score against. default (int): Default score to be assigned in case of failure in score calculation. # noqa E501 Returns: bool: True if the score surpasses the throttle, otherwise False. """ if prompt is None or len(prompt) == 0: return False score = default relation = self.llm.generate_response(prompt=prompt, backend=backend) tracker.log('score', [relation, throttle, default]) filtered_relation = ''.join([c for c in relation if c.isdigit()]) try: score_str = re.sub(r'[^\d]', ' ', filtered_relation).strip() score = int(score_str.split(' ')[0]) except Exception as e: logger.error(str(e)) if score >= throttle: return True return False def work_time(self): """Determines if the current time falls within the scheduled working hours of the chat assistant. Returns: bool: True if the current time is within working hours, otherwise False. # noqa E501 """ time_config = self.config['worker']['time'] beginWork = datetime.datetime.now().strftime( '%Y-%m-%d') + ' ' + time_config['start'] endWork = datetime.datetime.now().strftime( '%Y-%m-%d') + ' ' + time_config['end'] beginWorkSeconds = time.time() - time.mktime( time.strptime(beginWork, '%Y-%m-%d %H:%M:%S')) endWorkSeconds = time.time() - time.mktime( time.strptime(endWork, '%Y-%m-%d %H:%M:%S')) if int(beginWorkSeconds) > 0 and int(endWorkSeconds) < 0: if not time_config['has_weekday']: return True if int(datetime.datetime.now().weekday()) in range(7): return True return False def generate(self, query, history, groupname,backend): """Processes user queries and generates appropriate responses. It involves several steps including checking for valid questions, extracting topics, querying the feature store, searching the web, and generating responses from the language model. Args: query (str): User's query. history (str): Chat history. groupname (str): The group name in which user asked the query. backend (str): The backend to use for generating the response. Returns: ErrorCode: An error code indicating the status of response generation. # noqa E501 str: Generated response to the user query. references: List for referenced filename or web url """ response = '' references = [] if not self.work_time(): return ErrorCode.NOT_WORK_TIME, response, references if query is None or len(query) < 6: return ErrorCode.NOT_A_QUESTION, response, references reborn_code = ErrorCode.SUCCESS tracker = QueryTracker(self.logsavepath) tracker.log('input', [query, history, groupname]) if not self.single_judge( prompt=self.SCORING_QUESTION_TEMPLTE.format(query), tracker=tracker, throttle=6, default=3,backend=backend): return ErrorCode.NOT_A_QUESTION, response, references topic = self.llm.generate_response(self.TOPIC_TEMPLATE.format(query),backend=backend) tracker.log('topic', topic) if len(topic) <= 2: return ErrorCode.NO_TOPIC, response, references # chunk, db_context, references = self.retriever.query( # topic, # context_max_length=self.context_max_length - # 2 * len(self.GENERATE_TEMPLATE)) chunks, references = self.retriever.query( topic, context_max_length=self.context_max_length - 2 * len(self.GENERATE_TEMPLATE), tracker=tracker) # if db_context is None: # tracker.log('feature store reject') # return ErrorCode.UNRELATED, response, references context = '' refs = [] for chunk,ref in zip(chunks,references): if self.single_judge(self.SCORING_RELAVANCE_TEMPLATE.format(query, chunk), tracker=tracker, throttle=5, default=10,backend=backend): context += f"reference: {ref} content: {chunk}" context += '\n\n' refs.append(ref) refs = list(set(refs)) if len(context) > 0: prompt, history = self.llm.build_prompt( instruction=query, context=context, history_pair=history, template=self.GENERATE_TEMPLATE) response = self.llm.generate_response(prompt=prompt, backend=backend, history=history) tracker.log('feature store doc', [chunk, response]) response = convertid2url(response) return ErrorCode.SUCCESS, response, refs else: return ErrorCode.NO_SEARCH_RESULT, response, references # try: # references = [] # web_context = '' # web_search = WebSearch(config_path=self.config_path) # articles, error = web_search.get(query=topic, max_article=2) # if error is not None: # return ErrorCode.SEARCH_FAIL, response, references # tracker.log('search returned') # web_context_max_length = self.context_max_length - 2 * len( # self.SCORING_RELAVANCE_TEMPLATE) # for article in articles: # if len(article) > 0: # article.cut(0, web_context_max_length) # if self.single_judge( # self.SCORING_RELAVANCE_TEMPLATE.format( # query, article.content), # tracker=tracker, # throttle=5, # default=10): # web_context += '\n\n' # web_context += article.content # references.append(article.source) # web_context = web_context[0:web_context_max_length] # web_context = web_context.strip() # if len(web_context) > 0: # prompt, history = self.llm.build_prompt( # instruction=query, # context=web_context, # history_pair=history, # template=self.GENERATE_TEMPLATE) # response = self.llm.generate_response(prompt=prompt, # history=history) # else: # reborn_code = ErrorCode.NO_SEARCH_RESULT # tracker.log('web response', [web_context, response, reborn_code]) # except Exception as e: # logger.error(e) # if response is not None and len(response) > 0: # prompt = self.PERPLESITY_TEMPLATE.format(query, response) # if self.single_judge(prompt=prompt, # tracker=tracker, # throttle=10, # default=0,backend = backend): # reborn_code = ErrorCode.BAD_ANSWER # if self.config['worker']['enable_sg_search']: # if reborn_code == ErrorCode.BAD_ANSWER or reborn_code == ErrorCode.NO_SEARCH_RESULT: # noqa E501 # # reborn # sg = SourceGraphProxy(config_path=self.config_path, # language=self.language) # sg_context = sg.search(llm_client=self.llm, # question=query, # groupname=groupname) # if sg_context is not None and len(sg_context) > 2: # prompt, history = self.llm.build_prompt( # instruction=query, # context=sg_context, # history_pair=history, # template=self.GENERATE_TEMPLATE) # response = self.llm.generate_response(prompt=prompt, # history=history, # backend='remote') # tracker.log('source graph', [sg_context, response]) # prompt = self.PERPLESITY_TEMPLATE.format(query, response) # if self.single_judge(prompt=prompt, # tracker=tracker, # throttle=9, # default=0): # return ErrorCode.BAD_ANSWER, response, references # if response is not None and len(response) >= 800: # # reply too long, summarize it # response = self.llm.generate_response( # prompt=self.SUMMARIZE_TEMPLATE.format(response), # backend=backend) # if len(response) > 0 and self.single_judge( # self.SECURITY_TEMAPLTE.format(response), # tracker=tracker, # throttle=3, # default=0): # return ErrorCode.SECURITY, response, references # if reborn_code != ErrorCode.SUCCESS: # return reborn_code, response, references # return ErrorCode.SUCCESS, response, references def annotate_cluster(self,theme, cluster_no, chunk, history, groupname,backend): """Annotates a cluster of questions based on the user query and generates appropriate responses. Args: cluster_no (str): The cluster number of the questions. chunks (str): A cluster of questions. history (str): Chat history. groupname (str): The group name in which user asked the query. backend (str): The backend to use for generating the response. Returns: ErrorCode: An error code indicating the status of response generation. # noqa E501 str: Generated response to the user query. references: List for referenced filename or web url """ response = '' references = [] if not self.work_time(): return ErrorCode.NOT_WORK_TIME, response, references tracker = QueryTracker(self.logsavepath) tracker.log('input', [chunk[:1000], history, groupname]) response = self.llm.generate_response(prompt=self.ANNOTATE_CLUSTER.format(theme,chunk), backend=backend, history=history) if response is not None and len(response) >= 800: # reply too long, summarize it response = self.llm.generate_response( prompt=self.SUMMARIZE_TEMPLATE.format(response), backend=backend, history=history) tracker.log('annotate cluster', [cluster_no,chunk, response]) return ErrorCode.SUCCESS, response, cluster_no def getinspiration(self,theme,annotations,history,groupname,backend): """ give some idea of the review given the summary of the articals content """ response = '' references = [] if not self.work_time(): return ErrorCode.NOT_WORK_TIME, response, references tracker = QueryTracker(self.logsavepath) tracker.log('input', [annotations[:1000], history, groupname]) response = self.llm.generate_response(prompt=self.INSPIRATION_TEMPLATE.format(theme,annotations), backend=backend, history=history) tracker.log('get inspiration', [theme,annotations, response]) return ErrorCode.SUCCESS, response def parse_args(): """Parses command-line arguments.""" parser = argparse.ArgumentParser(description='Worker.') parser.add_argument('work_dir', type=str, help='Working directory.') parser.add_argument( '--config_path', default='config.ini', help='Worker configuration path. Default value is config.ini') return parser.parse_args() if __name__ == '__main__': args = parse_args() bot = Worker(work_dir=args.work_dir, config_path=args.config_path) queries = ['茴香豆是怎么做的'] for example in queries: print(bot.generate(query=example, history=[], groupname=''))