From 37d83fd74731d479081407c962854385b58f116c Mon Sep 17 00:00:00 2001 From: Cadenza-Li <362237642@qq.com> Date: Thu, 13 Jul 2023 17:20:22 +0800 Subject: [PATCH] update knowledge module; * Knowledge.storage to storages list; * optimize Knowledge & Storage save and load method; * optimize Knowledge query prompt; --- qlib/finco/knowledge.py | 205 +++++++++++++++++++++++++++------------- qlib/finco/utils.py | 7 ++ qlib/finco/workflow.py | 6 +- 3 files changed, 150 insertions(+), 68 deletions(-) diff --git a/qlib/finco/knowledge.py b/qlib/finco/knowledge.py index 9a781bdd8..6920b6295 100644 --- a/qlib/finco/knowledge.py +++ b/qlib/finco/knowledge.py @@ -3,12 +3,13 @@ from jinja2 import Template from typing import List, Union import pickle import yaml -import inspect from qlib.workflow import R from qlib.finco.log import FinCoLog from qlib.finco.llm import APIBackend -from qlib.finco.utils import similarity +from qlib.finco.utils import similarity, random_string + +logger = FinCoLog() class Storage: @@ -17,9 +18,9 @@ class Storage: """ - def __init__(self, path: Union[str, Path]): + def __init__(self, path: Union[str, Path], name: str = None): self.path = path if isinstance(path, Path) else Path(path) - self.logger = FinCoLog() + self.name = name if name else self.path.name self.source = None # todo: get document by key @@ -54,6 +55,7 @@ class PickleStorage(Storage): def save(self, **kwargs): """use pickle as the default save method""" + Path.mkdir(self.path.parent, exist_ok=True) with open(self.path, "wb") as f: pickle.dump(self, f) @@ -63,9 +65,11 @@ class YamlStorage(Storage): This class is responsible for storage and loading of Knowledge related data in yaml format. """ + DEFAULT_NAME = "storage.yml" def __init__(self, path: Union[str, Path]): super().__init__(path) + assert self.path.name, "Yaml storage should specify file name." self.load() def load(self): @@ -73,10 +77,11 @@ class YamlStorage(Storage): try: self.documents = yaml.load(open(self.path, "r"), Loader=yaml.FullLoader) except FileNotFoundError: - self.logger.warning(f"YamlStorage: file {self.path} doesn't exist.") + logger.warning(f"YamlStorage: file {self.path} doesn't exist.") def save(self, **kwargs): """use pickle as the default save method""" + Path.mkdir(self.path.parent, exist_ok=True) with open(self.path, 'w') as f: yaml.dump(self.documents, f) @@ -109,11 +114,21 @@ class Knowledge: Use to handle knowledge in finCo such as experiment and outside domain information """ - def __init__(self, storage: Storage): - self.logger = FinCoLog() - self.storage = storage + def __init__(self, storages: Union[List[Storage], Storage], name: str = None): + self.name = name if name else random_string() + self.workdir = Path.cwd().joinpath("knowledge") + self.storages = [storages] if isinstance(storages, Storage) else storages self.knowledge = [] + def get_storage(self, name: str): + """ + return first storage matched given name, else return None + """ + for storage in self.storages: + if storage.name == name: + return storage + return None + def summarize(self, **kwargs): """ summarize storage data to knowledge, default knowledge is storage.documents @@ -124,20 +139,24 @@ class Knowledge: Return ------ """ + for storage in self.storages: + self.knowledge.extend(storage.documents) - self.knowledge = self.storage.documents - - def load(self, **kwargs): + @classmethod + def load(cls, path: Union[str, Path]): """ Load knowledge in memory - + use pickle as the default file type Parameters ---------- Return ------ """ - raise NotImplementedError(f"Please implement the `load` method.") + """""" + path = path if isinstance(path, Path) else Path(path) + with open(path, "rb") as f: + return pickle.load(f) def brief(self, **kwargs): """ @@ -153,7 +172,10 @@ class Knowledge: def save(self, **kwargs): """save knowledge persistently""" - self.storage.save(**kwargs) + # todo: storages save index only + Path.mkdir(self.workdir.joinpath(self.name), exist_ok=True) + with open(self.workdir.joinpath(self.name).joinpath("knowledge.pkl"), "wb") as f: + pickle.dump(self, f) class ExperimentKnowledge(Knowledge): @@ -161,9 +183,9 @@ class ExperimentKnowledge(Knowledge): Handle knowledge from experiments """ - def __init__(self, storage: ExperimentStorage): - super().__init__(storage=storage) - self.storage = storage + def __init__(self, storages: Union[List[ExperimentStorage], ExperimentStorage]): + super().__init__(storages=storages) + self.storage = storages def brief(self): docs = [] @@ -179,13 +201,15 @@ class PracticeKnowledge(Knowledge): some template sentence for now """ - def __init__(self, storage: YamlStorage): - super().__init__(storage=storage) + def __init__(self, storages: Union[List[YamlStorage], YamlStorage]): + super().__init__(storages=storages, name="practice") self.summarize() def add(self, docs: List): - self.storage.add(docs) + storage = YamlStorage(path=self.workdir.joinpath(self.name).joinpath(YamlStorage.DEFAULT_NAME)) + storage.add(documents=docs) + self.storages.append(storage) self.summarize() self.save() @@ -196,15 +220,17 @@ class FinanceKnowledge(Knowledge): Knowledge from articles """ - def __init__(self, storage: YamlStorage): - super().__init__(storage=storage) - if len(self.storage.documents) == 0: - docs = self.read_files_in_directory(self.storage.path) - self.add(docs) + def __init__(self, storages: Union[List[YamlStorage], YamlStorage]): + super().__init__(storages=storages, name="finance") + + docs = self.read_files_in_directory(self.workdir.joinpath(self.name)) + self.add(docs) self.summarize() def add(self, docs: List): - self.storage.add(docs) + storage = YamlStorage(path=self.workdir.joinpath(self.name).joinpath(YamlStorage.DEFAULT_NAME)) + storage.add(documents=docs) + self.storages.append(storage) self.summarize() self.save() @@ -228,12 +254,14 @@ class ExecuteKnowledge(Knowledge): Config and associate execution result(pass or error message). We can regard the example in prompt as pass execution """ - def __init__(self, storage: YamlStorage): - super().__init__(storage=storage) + def __init__(self, storages: Union[List[YamlStorage], YamlStorage]): + super().__init__(storages=storages, name="execute") self.summarize() def add(self, docs: List): - self.storage.add(docs) + storage = YamlStorage(path=self.workdir.joinpath(YamlStorage.DEFAULT_NAME)) + storage.add(documents=docs) + self.storages.append(storage) self.summarize() self.save() @@ -244,34 +272,65 @@ class InfrastructureKnowledge(Knowledge): Knowledge from sentences, docstring, and code """ - def __init__(self, storage: YamlStorage): - super().__init__(storage=storage) + def __init__(self, storages: Union[List[YamlStorage], YamlStorage]): + super().__init__(storages=storages, name="infrastructure") - if len(self.storage.documents) == 0: - # todo: change the path to qlib root path - docs = self.get_functions_and_docstrings(Path.cwd().parent) + storage = self.get_storage(YamlStorage.DEFAULT_NAME) + if len(storage.documents) == 0: + docs = self.get_functions_and_docstrings(Path(__file__).parent.parent.parent) self.add(docs) def add(self, docs: List): - self.storage.add(docs) + storage = YamlStorage(path=self.workdir.joinpath(self.name).joinpath(YamlStorage.DEFAULT_NAME)) + storage.add(documents=docs) + self.storages.append(storage) self.summarize() self.save() - @staticmethod - def get_functions_and_docstrings(directory): + def get_functions_and_docstrings(self, directory): """ get all method and docstring in .py files under directory + """ functions = [] - for file_path in Path(directory).rglob("*.py"): - with file_path.open("r", encoding="utf-8") as f: - lines = f.readlines() - for line in lines: - if line.startswith("def "): - function_name = line.split("(")[0][4:].strip() - function_docstring = inspect.getdoc(eval(function_name)) - functions.append({"function_name": function_name, "docstring": function_docstring}) + for py_file_path in Path(directory).rglob('*.py'): + for _functions in self.get_functions_with_docstrings(py_file_path): + functions.append(_functions) + + return functions + + @staticmethod + def get_functions_with_docstrings(file_path): + """ + Extract method name and docstring using string matching method + """ + with open(file_path, "r", encoding="utf8") as f: + lines = f.readlines() + + functions = [] + current_func = None + docstring = None + for line in lines: + if line.strip().startswith("def ") or line.strip().startswith("class "): + func = line.strip().split(' ')[1].split('(')[0] + if func.startswith("__"): + continue + if current_func is not None: + docstring = docstring.replace('"""', "") if docstring else docstring + functions.append({"function": current_func, "docstring": docstring}) + current_func = f"{file_path.name.split('.')[0]}.{func}" + docstring = None + elif current_func is not None and docstring is None and line.strip().startswith('"""'): + docstring = line + elif current_func is not None and docstring is not None: + docstring += line.strip() + if line.strip().endswith('"""'): + docstring = docstring.replace('"""', "") if docstring else docstring + functions.append({"function": current_func, "docstring": docstring}) + current_func = None + docstring = None + return functions @@ -303,18 +362,20 @@ class KnowledgeBase: KT_EXECUTE = "execute" KT_PRACTICE = "practice" KT_FINANCE = "finance" + KT_INFRASTRUCTURE = "infrastructure" - def __init__(self, init_path=None): + def __init__(self, workdir=None): self.logger = FinCoLog() - self.init_path = Path(init_path) if init_path else Path.cwd() + self.workdir = Path(workdir) if workdir else Path.cwd() - if not self.init_path.exists(): - self.logger.warning(f"{self.init_path} not exist, create empty directory.") - Path.mkdir(self.init_path) + if not self.workdir.exists(): + self.logger.warning(f"{self.workdir} not exist, create empty directory.") + Path.mkdir(self.workdir) - self.practice_knowledge = self.load_practice_knowledge(self.init_path) - self.execute_knowledge = self.load_execute_knowledge(self.init_path) - self.finance_knowledge = self.load_finance_knowledge(self.init_path) + self.practice_knowledge = self.load_practice_knowledge(self.workdir) + self.execute_knowledge = self.load_execute_knowledge(self.workdir) + self.finance_knowledge = self.load_finance_knowledge(self.workdir) + self.infrastructure_knowledge = self.load_infrastructure_knowledge(self.workdir) def load_experiment_knowledge(self, path) -> List: # similar to practice knowledge, not use for now @@ -326,33 +387,43 @@ class KnowledgeBase: # todo: check the influence of set uri R.set_uri(path.as_uri()) for exp_name in R.list_experiments(): - knowledge.append(ExperimentKnowledge(storage=ExperimentStorage(exp_name=exp_name))) + knowledge.append(ExperimentKnowledge(storages=ExperimentStorage(exp_name=exp_name))) self.logger.plain_info(f"Load knowledge from: {path} finished.") return knowledge def load_practice_knowledge(self, path: Path) -> PracticeKnowledge: - self.practice_knowledge = PracticeKnowledge(YamlStorage(path.joinpath("practice_knowledge.yaml"))) + self.practice_knowledge = PracticeKnowledge( + YamlStorage(path.joinpath(f"{self.KT_PRACTICE}/{YamlStorage.DEFAULT_NAME}"))) return self.practice_knowledge def load_execute_knowledge(self, path: Path) -> ExecuteKnowledge: - self.execute_knowledge = ExecuteKnowledge(YamlStorage(path.joinpath("execute_knowledge.yaml"))) + self.execute_knowledge = ExecuteKnowledge( + YamlStorage(path.joinpath(f"{self.KT_EXECUTE}/{YamlStorage.DEFAULT_NAME}"))) return self.execute_knowledge def load_finance_knowledge(self, path: Path) -> FinanceKnowledge: - self.finance_knowledge = FinanceKnowledge(YamlStorage(path.joinpath("finance_knowledge.yaml"))) + self.finance_knowledge = FinanceKnowledge( + YamlStorage(path.joinpath(f"{self.KT_FINANCE}/{YamlStorage.DEFAULT_NAME}"))) return self.finance_knowledge - def knowledge(self, knowledge_type: str = None): + def load_infrastructure_knowledge(self, path: Path) -> InfrastructureKnowledge: + self.infrastructure_knowledge = InfrastructureKnowledge( + YamlStorage(path.joinpath(f"{self.KT_INFRASTRUCTURE}/{YamlStorage.DEFAULT_NAME}"))) + return self.infrastructure_knowledge + + def get_knowledge(self, knowledge_type: str = None): if knowledge_type == self.KT_EXECUTE: - knowledge = self.execute_knowledge + knowledge = self.execute_knowledge.knowledge elif knowledge_type == self.KT_PRACTICE: - knowledge = self.practice_knowledge + knowledge = self.practice_knowledge.knowledge elif knowledge_type == self.KT_FINANCE: - knowledge = self.finance_knowledge + knowledge = self.finance_knowledge.knowledge + elif knowledge_type == self.KT_INFRASTRUCTURE: + knowledge = self.infrastructure_knowledge.knowledge else: knowledge = self.execute_knowledge.knowledge + self.practice_knowledge.knowledge \ - + self.finance_knowledge.knowledge + + self.finance_knowledge.knowledge + self.infrastructure_knowledge.knowledge return knowledge def query(self, knowledge_type: str = None, content: str = None, n: int = 5): @@ -366,7 +437,7 @@ class KnowledgeBase: # todo: replace list with persistent storage strategy such as ES/pinecone to enable # literal search/semantic search - knowledge = self.knowledge(knowledge_type=knowledge_type) + knowledge = self.get_knowledge(knowledge_type=knowledge_type) scores = [] for k in knowledge: scores.append(similarity(str(k), content)) @@ -374,10 +445,12 @@ class KnowledgeBase: similar_n_indexes = sorted_indexes[:n] similar_n_docs = [knowledge[i] for i in similar_n_indexes] - prompt = Template("""summarize this information: '{{docs}}'""") - prompt_workflow_selection = prompt.render(docs=similar_n_docs) + prompt = Template( + """find the most relevant doc with this query: '{{content}}' from docs='{{docs}}'. + Just return the most relevant item I provided, no more explain. For example: {'function': 'config.resolve_path', 'docstring': None}""") + prompt_workflow_selection = prompt.render(content=content, docs=similar_n_docs) response = APIBackend().build_messages_and_create_chat_completion( - user_prompt=prompt_workflow_selection + user_prompt=prompt_workflow_selection, system_prompt="You are an excellent assistant." ) return response diff --git a/qlib/finco/utils.py b/qlib/finco/utils.py index aa4107021..379ef6375 100644 --- a/qlib/finco/utils.py +++ b/qlib/finco/utils.py @@ -1,4 +1,6 @@ import json +import string +import random from fuzzywuzzy import fuzz @@ -36,3 +38,8 @@ def similarity(text1, text2): # Maybe we can use other similarity algorithm such as tfidf return fuzz.ratio(text1, text2) + + +def random_string(length=10): + letters = string.ascii_letters + string.digits + return ''.join(random.choice(letters) for i in range(length)) diff --git a/qlib/finco/workflow.py b/qlib/finco/workflow.py index b6d00f966..3084a2012 100644 --- a/qlib/finco/workflow.py +++ b/qlib/finco/workflow.py @@ -176,7 +176,9 @@ class LearnManager: self.topics = [Topic(name=topic, describe=self.wm.prompt_template.get(f"Topic_{topic}")) for topic in self.__DEFAULT_TOPICS] - self.knowledge_base = KnowledgeBase(init_path=Path.cwd().joinpath('knowledge')) + self.knowledge_base = KnowledgeBase(workdir=Path.cwd().joinpath('knowledge')) + self.knowledge_base.execute_knowledge.add([]) + self.knowledge_base.query(knowledge_type="infrastructure", content="resolve_path") def run(self, prompt): # todo: add early stop condition @@ -202,7 +204,7 @@ class LearnManager: user_prompt = self.wm.context.get_context("user_prompt") summary = self.wm.context.get_context("summary") - [topic.summarize(self.knowledge_base.knowledge()) for topic in self.topics] + [topic.summarize(self.knowledge_base.get_knowledge()) for topic in self.topics] knowledge_of_topics = [{topic.name: topic.knowledge} for topic in self.topics] for task in task_finished: