reformat

app.py

@@ -28,7 +28,9 @@ with st.sidebar:
 
 def main() -> None:
     st.header("X-Pipe Wiki 机器人 💬")
-    robot_manager = XPipeRobotManagerFactory.get_or_create(XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0)
+    robot_manager = XPipeRobotManagerFactory.get_or_create(
+        XPipeRobotRevision.SIMPLE_OPENAI_VERSION_0
+    )
     robot = robot_manager.get_robot()
     query = st.text_input("X-Pipe Wiki 问题:")
     if query:

core/helper.py

@@ -2,23 +2,30 @@ from core.lifecycle import Lifecycle
 
 
 class LifecycleHelper:
-
     @classmethod
     def initialize_if_possible(cls, ls: Lifecycle) -> None:
-        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_initialize(ls.lifecycle_state.phase):
+        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_initialize(
+            ls.lifecycle_state.phase
+        ):
             ls.initialize()
 
     @classmethod
     def start_if_possible(cls, ls: Lifecycle) -> None:
-        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_start(ls.lifecycle_state.phase):
+        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_start(
+            ls.lifecycle_state.phase
+        ):
             ls.start()
 
     @classmethod
     def stop_if_possible(cls, ls: Lifecycle) -> None:
-        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_stop(ls.lifecycle_state.phase):
+        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_stop(
+            ls.lifecycle_state.phase
+        ):
             ls.stop()
 
     @classmethod
     def dispose_if_possible(cls, ls: Lifecycle) -> None:
-        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_dispose(ls.lifecycle_state.phase):
+        if isinstance(ls, Lifecycle) and ls.get_lifecycle_state().can_dispose(
+            ls.lifecycle_state.phase
+        ):
             ls.dispose()

core/lifecycle.py

@@ -112,7 +112,7 @@ class LifecycleController(ABC):
 
     def can_start(self, phase: Optional[LifecyclePhase]) -> bool:
         return phase is not None and (
-                phase == LifecyclePhase.INITIALIZED or phase == LifecyclePhase.STOPPED
+            phase == LifecyclePhase.INITIALIZED or phase == LifecyclePhase.STOPPED
         )
 
     def can_stop(self, phase: Optional[LifecyclePhase]) -> bool:
@@ -120,7 +120,7 @@ class LifecycleController(ABC):
 
     def can_dispose(self, phase: Optional[LifecyclePhase]) -> bool:
         return phase is not None and (
-                phase == LifecyclePhase.INITIALIZED or phase == LifecyclePhase.STOPPED
+            phase == LifecyclePhase.INITIALIZED or phase == LifecyclePhase.STOPPED
         )
 
 

langchain_manager/manager.py

@@ -8,7 +8,7 @@ from langchain_manager.llms import AzureOpenAI
 from core.lifecycle import Lifecycle
 
 
-class BaseLangChainManager(Lifecycle, ABC):
+class BaseLangChainManager(ABC):
     def __init__(self) -> None:
         super().__init__()
 

llama/context.py

@@ -5,10 +5,10 @@ from llama_index import StorageContext
 
 from core.lifecycle import Lifecycle
 from langchain_manager.manager import BaseLangChainManager
+from typing import List
 
 
 class ServiceContextManager(Lifecycle, ABC):
-
     @abstractmethod
     def get_service_context(self) -> ServiceContext:
         pass
@@ -44,23 +44,33 @@ class AzureServiceContextManager(ServiceContextManager):
         )
 
     def do_start(self) -> None:
-        self.logger.info("[do_start][embedding] last used usage: %d",
-                         self.service_context.embed_model.total_tokens_used)
-        self.logger.info("[do_start][predict] last used usage: %d",
-                         self.service_context.llm_predictor.total_tokens_used)
+        self.logger.info(
+            "[do_start][embedding] last used usage: %d",
+            self.service_context.embed_model.total_tokens_used,
+        )
+        self.logger.info(
+            "[do_start][predict] last used usage: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
 
     def do_stop(self) -> None:
-        self.logger.info("[do_stop][embedding] last used usage: %d",
-                         self.service_context.embed_model.total_tokens_used)
-        self.logger.info("[do_stop][predict] last used usage: %d",
-                         self.service_context.llm_predictor.total_tokens_used)
+        self.logger.info(
+            "[do_stop][embedding] last used usage: %d",
+            self.service_context.embed_model.total_tokens_used,
+        )
+        self.logger.info(
+            "[do_stop][predict] last used usage: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
 
     def do_dispose(self) -> None:
-        self.logger.info("[do_dispose] total used token: %d", self.service_context.llm_predictor.total_tokens_used)
+        self.logger.info(
+            "[do_dispose] total used token: %d",
+            self.service_context.llm_predictor.total_tokens_used,
+        )
 
 
 class StorageContextManager(Lifecycle, ABC):
-
     @abstractmethod
     def get_storage_context(self) -> StorageContext:
         pass
@@ -69,9 +79,11 @@ class StorageContextManager(Lifecycle, ABC):
 class LocalStorageContextManager(StorageContextManager):
     storage_context: StorageContext
 
-    def __init__(self,
-                 dataset_path: str = "./dataset",
-                 service_context_manager: ServiceContextManager = None) -> None:
+    def __init__(
+        self,
+        service_context_manager: ServiceContextManager,
+        dataset_path: str = "./dataset",
+    ) -> None:
         super().__init__()
         self.dataset_path = dataset_path
         self.service_context_manager = service_context_manager
@@ -81,8 +93,11 @@ class LocalStorageContextManager(StorageContextManager):
 
     def do_init(self) -> None:
         from llama.utils import is_local_storage_files_ready
+
         if is_local_storage_files_ready(self.dataset_path):
-            self.storage_context = StorageContext.from_defaults(persist_dir=self.dataset_path)
+            self.storage_context = StorageContext.from_defaults(
+                persist_dir=self.dataset_path
+            )
         else:
             docs = self._download()
             self._indexing(docs)
@@ -96,14 +111,17 @@ class LocalStorageContextManager(StorageContextManager):
     def do_dispose(self) -> None:
         self.storage_context.persist(self.dataset_path)
 
-    def _download(self) -> [Document]:
+    def _download(self) -> List[Document]:
         from llama.data_loader import GithubLoader
+
         loader = GithubLoader()
         return loader.load()
 
-    def _indexing(self, docs: [Document]):
+    def _indexing(self, docs: List[Document]) -> None:
         from llama_index import GPTVectorStoreIndex
-        index = GPTVectorStoreIndex.from_documents(docs,
-                                                   service_context=self.service_context_manager.get_service_context())
+
+        index = GPTVectorStoreIndex.from_documents(
+            docs, service_context=self.service_context_manager.get_service_context()
+        )
         index.storage_context.persist(persist_dir=self.dataset_path)
         self.storage_context = index.storage_context

llama/data_loader.py

@@ -1,7 +1,7 @@
 import os
 import pickle
 from abc import abstractmethod, ABC
-from typing import Optional, Sequence
+from typing import Optional, Sequence, List
 
 from llama_hub.github_repo import GithubRepositoryReader, GithubClient
 from llama_index import download_loader
@@ -10,16 +10,16 @@ from llama_index.readers.schema.base import Document
 
 class WikiLoader(ABC):
     @abstractmethod
-    def load(self) -> [Document]:
+    def load(self) -> List[Document]:
         pass
 
 
 class GithubLoader(WikiLoader):
     def __init__(
-            self,
-            github_owner: Optional[str] = None,
-            repo: Optional[str] = None,
-            dirs: Optional[Sequence[str]] = None,
+        self,
+        github_owner: Optional[str] = None,
+        repo: Optional[str] = None,
+        dirs: Optional[Sequence[str]] = None,
     ):
         super().__init__()
         self.owner = (
@@ -28,7 +28,7 @@ class GithubLoader(WikiLoader):
         self.repo = repo if repo is not None else os.environ["GITHUB_REPO"]
         self.dirs = dirs if dirs is not None else [".", "doc"]
 
-    def load(self) -> [Document]:
+    def load(self) -> List[Document]:
         download_loader("GithubRepositoryReader")
         docs = None
         if os.path.exists("docs/docs.pkl"):

xpipe_wiki/manager_factory.py

@@ -9,11 +9,10 @@ class XPipeRobotRevision(enum.Enum):
     SIMPLE_OPENAI_VERSION_0 = 1
 
 
-CAPABLE = dict[XPipeRobotRevision, XPipeWikiRobotManager]
+CAPABLE = dict()
 
 
 class XPipeRobotManagerFactory:
-
     @classmethod
     def get_or_create(cls, revision: XPipeRobotRevision) -> XPipeWikiRobotManager:
         if CAPABLE.get(revision) is not None:
@@ -25,18 +24,24 @@ class XPipeRobotManagerFactory:
 
     @classmethod
     def create_simple_openai_version_0(cls) -> AzureXPipeWikiRobotManager:
-
         from llama.context import AzureServiceContextManager
         from langchain_manager.manager import LangChainAzureManager
-        service_context_manager = AzureServiceContextManager(lc_manager=LangChainAzureManager())
+
+        service_context_manager = AzureServiceContextManager(
+            lc_manager=LangChainAzureManager()
+        )
 
         from llama.context import LocalStorageContextManager
-        dataset_path = os.getenv("XPIPE_WIKI_DATASET_PATH", "./dataset")
-        storage_context_manager = LocalStorageContextManager(dataset_path=dataset_path,
-                                                             service_context_manager=service_context_manager)
 
-        robot_manager = AzureXPipeWikiRobotManager(service_context_manager=service_context_manager,
-                                                   storage_context_manager=storage_context_manager)
+        dataset_path = os.getenv("XPIPE_WIKI_DATASET_PATH", "./dataset")
+        storage_context_manager = LocalStorageContextManager(
+            dataset_path=dataset_path, service_context_manager=service_context_manager
+        )
+
+        robot_manager = AzureXPipeWikiRobotManager(
+            service_context_manager=service_context_manager,
+            storage_context_manager=storage_context_manager,
+        )
         LifecycleHelper.initialize_if_possible(robot_manager)
         LifecycleHelper.start_if_possible(robot_manager)
         return robot_manager

xpipe_wiki/robot_manager.py

@@ -16,29 +16,43 @@ class XPipeWikiRobot(ABC):
         pass
 
 
-@dataclasses
 class AzureOpenAIXPipeWikiRobot(XPipeWikiRobot):
     query_engine: BaseQueryEngine
 
+    def __init__(self, query_engine: BaseQueryEngine) -> None:
+        super().__init__()
+        self.query_engine = query_engine
+
     def ask(self, question: str) -> Any:
         return self.query_engine.query(question)
 
 
 class XPipeWikiRobotManager(Lifecycle):
-
     @abstractmethod
     def get_robot(self) -> XPipeWikiRobot:
         pass
 
 
-@dataclasses
 class AzureXPipeWikiRobotManager(XPipeWikiRobotManager):
     service_context_manager: ServiceContextManager
     storage_context_manager: StorageContextManager
 
+    def __init__(
+        self,
+        service_context_manager: ServiceContextManager,
+        storage_context_manager: StorageContextManager,
+    ) -> None:
+        super().__init__()
+        self.service_context_manager = service_context_manager
+        self.storage_context_manager = storage_context_manager
+
     def get_robot(self) -> XPipeWikiRobot:
-        index = load_index_from_storage(storage_context=self.storage_context_manager.get_storage_context())
-        query_engine = index.as_query_engine(service_context=self.service_context_manager.get_service_context())
+        index = load_index_from_storage(
+            storage_context=self.storage_context_manager.get_storage_context()
+        )
+        query_engine = index.as_query_engine(
+            service_context=self.service_context_manager.get_service_context()
+        )
         return AzureOpenAIXPipeWikiRobot(query_engine)
 
     def do_init(self) -> None: