Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| """Search enhancement proxy.""" | |
| import argparse | |
| import json | |
| import os | |
| import pytoml | |
| from loguru import logger | |
| from .llm_client import ChatClient | |
| class SourceGraphProxy: | |
| """A class to serve as a proxy for interacting with the Source Graph. | |
| Args: | |
| config_path (dict): Path to the configuration file. | |
| topk (int, optional): Top K results to consider from the search. Defaults to 1. # noqa E501 | |
| language (str, optional): Language for the system prompts - 'zh' for Chinese and 'en' for English. Defaults to 'zh'. # noqa E501 | |
| Attributes: | |
| config_path (str): The path of the configuration file. | |
| sg_config (dict): Configuration settings for sourcegraph search. | |
| topk (int): Top K results to consider from the search. | |
| language (str): Language for the system prompts. | |
| CHOICE_TEMPLATE (str): Template string for generating choice based on selected language. # noqa E501 | |
| KEYWORDS_TEMPLATE (str): Template string for generating keywords based on selected language. # noqa E501 | |
| """ | |
| def __init__(self, | |
| config_path: dict, | |
| topk=1, | |
| language: str = 'zh') -> None: | |
| """Init searcher with config.""" | |
| self.config_path = config_path | |
| self.sg_config = None | |
| with open(self.config_path, encoding='utf8') as f: | |
| config = pytoml.load(f) | |
| self.sg_config = config['sg_search'] | |
| self.topk = topk | |
| self.language = language | |
| if self.language == 'zh': | |
| self.CHOICE_TEMPLATE = '“{}”\n请仔细阅读以上问题，请问应该查询以下哪个开源项目：\n' # noqa E501 | |
| self.KEYWORDS_TEMPLATE = '“{}”\n请仔细阅读以上问题，提取其中可用作搜索引擎的关键字，关键字之间,分隔，不要解释。' # noqa E501 | |
| else: | |
| self.CHOICE_TEMPLATE = '"{}"\nPlease read the above question carefully, which of the following open-source projects should this question refer to: \n' # noqa E501 | |
| self.KEYWORDS_TEMPLATE = '"{}"\nPlease read the above questions carefully, extract the keywords which can be used as search engines, between keywords, separate, do not explain.' # noqa E501 | |
| def command(self, txt: str): | |
| """Executes a shell command and returns its output. | |
| Args: | |
| txt (str): Command to be executed in the shell. | |
| Returns: | |
| str: Output of the shell command execution. | |
| """ | |
| logger.debug('cmd: {}'.format(txt)) | |
| cmd = os.popen(txt) | |
| return cmd.read().rstrip().lstrip() | |
| def extract_sg_result(self, jsonstr): | |
| """Extracts the desired data from the source graph result. | |
| Args: | |
| jsonstr (str): JSON string containing source graph search result. | |
| Returns: | |
| list: List of dictionaries each contains 'filepath' and 'content' of the files returned by source graph. # noqa E501 | |
| """ | |
| ret = [] | |
| try: | |
| root = json.loads(jsonstr) | |
| results = root['Results'] | |
| for result in results: | |
| if 'FileMatch' != result['__typename']: | |
| continue | |
| content = result['file']['content'] | |
| path = result['file']['path'] | |
| ret.append({'filepath': path, 'content': content}) | |
| if len(ret) >= self.topk: | |
| break | |
| except Exception as e: | |
| logger.warning('{} when source graph parse {}'.format( | |
| str(e), jsonstr)) | |
| return ret | |
| def choose_repo(self, llm_client, question, groupname): | |
| """Interactively assists user to select a repository for search based | |
| on user's question. | |
| Args: | |
| llm_client: Client instance for LLM. | |
| question (str): User's question. | |
| groupname (str): Name of the user's group. | |
| Returns: | |
| str: The ID of selected repository. | |
| """ | |
| prompt = self.CHOICE_TEMPLATE.format(question) | |
| keys = self.sg_config.keys() | |
| skip = ['binary_src_path', 'src_access_token'] | |
| repos = {} | |
| for key in keys: | |
| if key in skip: | |
| continue | |
| introduction = self.sg_config[key]['introduction'] | |
| prompt += f'* {key} {introduction}\n' | |
| repos[key] = self.sg_config[key] | |
| prompt += '* none ' | |
| choice = llm_client.generate_response(prompt=prompt, | |
| backend='remote').strip() | |
| target_repo_id = None | |
| for key in repos.keys(): | |
| if key in choice: | |
| target_repo_id = repos[key]['github_repo_id'] | |
| break | |
| return target_repo_id | |
| def search(self, llm_client, question, groupname): | |
| """Performs a search operation in the selected repository based on the | |
| user's question. | |
| Args: | |
| llm_client: Client instance for LLM. | |
| question (str): User's question. | |
| groupname (str): Name of the user's group. | |
| Returns: | |
| str: Search result from source graph in JSON format. | |
| """ | |
| repo_id = self.choose_repo(llm_client, question, groupname) | |
| if repo_id is None: | |
| logger.warning('cannot choose repo_id') | |
| return '' | |
| ENV = 'export SRC_ACCESS_TOKEN="{}" && '.format( | |
| self.sg_config['src_access_token']) | |
| BINARY = self.sg_config['binary_src_path'] | |
| prompt = self.KEYWORDS_TEMPLATE.format(question) | |
| entities = [] | |
| entity_str = '' | |
| try: | |
| entity_str = llm_client.generate_response(prompt=prompt) | |
| entities = [item for item in entity_str.split(',') if item.strip()] | |
| except Exception as e: | |
| logger.error('parse {} failed {}.'.format(entity_str, str(e))) | |
| # return '' | |
| entities = [] | |
| search_items = [] | |
| for entity in entities: | |
| # search doc and source code based on entities | |
| # search -json 'repo:open-compass/opencompass summarizers' | |
| cmd_doc = '''{} search -json 'repo:{} lang:MarkDown {}' '''.format( | |
| BINARY, repo_id, entity) | |
| cmd_return = self.command(ENV + cmd_doc) | |
| search_items += self.extract_sg_result(cmd_return) | |
| cmd_python = '''{} search -json 'repo:{} lang:Python {}' '''.format( # noqa E501 | |
| BINARY, repo_id, entity) | |
| cmd_return = self.command(ENV + cmd_python) | |
| search_items += self.extract_sg_result(cmd_return) | |
| search_text = json.dumps(search_items, ensure_ascii=False, indent=2) | |
| return search_text | |
| def parse_args(): | |
| """Parses command line arguments.""" | |
| parser = argparse.ArgumentParser(description='Source graph proxy search') | |
| parser.add_argument( | |
| '--config_path', | |
| default='config.ini', | |
| help= # noqa E251 | |
| 'Source graph proxy configuration path. Default value is config.ini') | |
| args = parser.parse_args() | |
| return args | |
| if __name__ == '__main__': | |
| """Test search.""" | |
| logger.add('logs/sg_search.log', rotation='4MB') | |
| args = parse_args() | |
| llm = ChatClient(config_path=args.config_path) | |
| sg = SourceGraphProxy(config_path=args.config_path) | |
| context = sg.search(llm, | |
| question='请问triviaqa 5shot结果怎么在summarizer里输出呢', | |
| groupname='opencompass') | |
| print(context) | |