update knowledge module;

* Knowledge.storage to storages list;
* optimize Knowledge & Storage save and load method;
* optimize Knowledge query prompt;

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

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))

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: