Spaces:
Sleeping
Sleeping
File size: 2,804 Bytes
257bc0d bd59653 257bc0d bd59653 257bc0d b404642 257bc0d bd59653 257bc0d 9ecd254 257bc0d b404642 257bc0d 9ecd254 257bc0d 9ecd254 5ea412d 9ecd254 bd59653 d092e87 9ecd254 257bc0d 01d6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | from abc import ABC, abstractmethod
from typing import Any
from llama_index import load_index_from_storage
from llama_index.indices.query.base import BaseQueryEngine
from llama_index.indices.response import ResponseMode
from core.helper import LifecycleHelper
from core.lifecycle import Lifecycle
from llama.service_context import ServiceContextManager
from llama.storage_context import StorageContextManager
class XPipeWikiRobot(ABC):
@abstractmethod
def ask(self, question: str) -> Any:
pass
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:
print("question: ", question)
response = self.query_engine.query(question)
print("response type: ", type(response))
return response.__str__()
class XPipeWikiRobotManager(Lifecycle):
@abstractmethod
def get_robot(self) -> XPipeWikiRobot:
pass
class AzureXPipeWikiRobotManager(XPipeWikiRobotManager):
service_context_manager: ServiceContextManager
storage_context_manager: StorageContextManager
query_engine: BaseQueryEngine
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:
return AzureOpenAIXPipeWikiRobot(self.query_engine)
def do_init(self) -> None:
LifecycleHelper.initialize_if_possible(self.service_context_manager)
LifecycleHelper.initialize_if_possible(self.storage_context_manager)
def do_start(self) -> None:
LifecycleHelper.start_if_possible(self.service_context_manager)
LifecycleHelper.start_if_possible(self.storage_context_manager)
index = load_index_from_storage(
storage_context=self.storage_context_manager.get_storage_context(),
service_context=self.service_context_manager.get_service_context(),
)
self.query_engine = index.as_query_engine(
service_context=self.service_context_manager.get_service_context(),
response_mode=ResponseMode.TREE_SUMMARIZE,
)
def do_stop(self) -> None:
LifecycleHelper.stop_if_possible(self.storage_context_manager)
LifecycleHelper.stop_if_possible(self.service_context_manager)
def do_dispose(self) -> None:
LifecycleHelper.dispose_if_possible(self.storage_context_manager)
LifecycleHelper.dispose_if_possible(self.service_context_manager)
|