diff --git a/v1/em_service.py b/v1/em_service.py new file mode 100644 index 00000000..91d5c419 --- /dev/null +++ b/v1/em_service.py @@ -0,0 +1,62 @@ +from concurrent.futures import ThreadPoolExecutor + +from v1.pipeline.pipeline import Pipeline +from v1.pipeline.pipeline_context import PipelineContext +from v1.schema.app_config import AppConfig +from v1.schema.request import SummarizerRequest, RetrieverRequest, VectorStoreRequest, AgentRequest +from v1.schema.response import SummarizerResponse, RetrieverResponse, VectorStoreResponse, AgentResponse + + +class EMService(object): + def __init__(self, app_config: AppConfig, thread_pool: ThreadPoolExecutor): + self.context: PipelineContext = PipelineContext(app_config=app_config, thread_pool=thread_pool) + + def __call__(self, service: str, **kwargs) -> dict: + if service == "retriever": + response = self.call_retriever(RetrieverRequest(**kwargs)) + + elif service == "summarizer": + response = self.call_summarizer(SummarizerRequest(**kwargs)) + + elif service == "vector_store": + response = self.call_vector_store(VectorStoreRequest(**kwargs)) + + elif service == "agent": + response = self.call_agent(AgentRequest(**kwargs)) + + else: + raise Exception(f"Invalid service={service}") + + return response.model_dump() + + def call_retriever(self, request: RetrieverRequest) -> RetrieverResponse: + self.context.set_context("request", request) + response = RetrieverResponse() + self.context.set_context("response", response) + pipeline = Pipeline(pipeline=self.context.app_config.api.retriever, context=self.context) + pipeline.execute_pipeline() + return response + + def call_summarizer(self, request: SummarizerRequest) -> SummarizerResponse: + self.context.set_context("request", request) + response = SummarizerResponse() + self.context.set_context("request", response) + pipeline = Pipeline(pipeline=self.context.app_config.api.summarizer, context=self.context) + pipeline.execute_pipeline() + return response + + def call_vector_store(self, request: VectorStoreRequest) -> VectorStoreResponse: + self.context.set_context("request", request) + response = VectorStoreResponse() + self.context.set_context("request", response) + pipeline = Pipeline(pipeline=self.context.app_config.api.vector_store, context=self.context) + pipeline.execute_pipeline() + return response + + def call_agent(self, request: AgentRequest) -> AgentResponse: + self.context.set_context("request", request) + response = AgentResponse() + self.context.set_context("request", response) + pipeline = Pipeline(pipeline=self.context.app_config.api.agent, context=self.context) + pipeline.execute_pipeline() + return response diff --git a/v1/main.py b/v1/main.py index 161a752d..2011ece8 100644 --- a/v1/main.py +++ b/v1/main.py @@ -1,29 +1,51 @@ -# main.py -# config.py -import json -from dataclasses import dataclass, field, asdict -from typing import List +import sys +from concurrent.futures.thread import ThreadPoolExecutor -from omegaconf import OmegaConf +import uvicorn +from fastapi import FastAPI + +from v1.em_service import EMService +from v1.schema.request import RetrieverRequest, SummarizerRequest, VectorStoreRequest, AgentRequest +from v1.schema.response import RetrieverResponse, SummarizerResponse, VectorStoreResponse, AgentResponse +from v1.utils.config_parser import ConfigParser + +app = FastAPI() +config_parser = ConfigParser(sys.argv[1:]) +global_app_config = config_parser.get_app_config() +thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=global_app_config.thread_pool.max_workers) -main() +@app.post('/retriever', response_model=RetrieverResponse) +def call_retriever(request: RetrieverRequest): + app_config = config_parser.get_app_config(**request.params) + ems = EMService(app_config=app_config, thread_pool=thread_pool) + return ems.call_retriever(request) + + +@app.post('/summarizer', response_model=SummarizerResponse) +def call_summarizer(request: SummarizerRequest): + app_config = config_parser.get_app_config(**request.params) + ems = EMService(app_config=app_config, thread_pool=thread_pool) + return ems.call_summarizer(request) + + +@app.post('/vector_store', response_model=VectorStoreResponse) +def call_vector_store(request: VectorStoreRequest): + app_config = config_parser.get_app_config(**request.params) + ems = EMService(app_config=app_config, thread_pool=thread_pool) + return ems.call_vector_store(request) + + +@app.post('/agent', response_model=AgentResponse) +def call_agent(request: AgentRequest): + app_config = config_parser.get_app_config(**request.params) + ems = EMService(app_config=app_config, thread_pool=thread_pool) + return ems.call_agent(request) + if __name__ == "__main__": - default_cfg = OmegaConf.structured(AppConfig) - - # 2. 从 YAML 文件加载 - yaml_cfg = OmegaConf.load("config.yaml") - - # 3. 合并（YAML 覆盖默认） - cfg = OmegaConf.merge(default_cfg, yaml_cfg) - - # 4. 再合并命令行（命令行优先级最高） - cli_cfg = OmegaConf.from_cli() - cfg = OmegaConf.merge(cfg, cli_cfg) - - # 5. 转为 dataclass 实例 - cfg_obj: AppConfig = OmegaConf.to_object(cfg) # 递归转为 dataclass - - print(json.dumps(asdict(cfg_obj), indent=2)) - # print(cfg_obj) # 如果你需要 dataclass 实例 + uvicorn.run(app=app, + host=global_app_config.http_service.host, + port=global_app_config.http_service.port, + timeout_keep_alive=global_app_config.http_service.timeout_keep_alive, + limit_concurrency=global_app_config.http_service.limit_concurrency) diff --git a/v1/op/__init__.py b/v1/op/__init__.py index a09751c6..b8c42b57 100644 --- a/v1/op/__init__.py +++ b/v1/op/__init__.py @@ -1,3 +1,3 @@ from v1.utils.registry import Registry -OPERATION_REGISTRY = Registry() +OP_REGISTRY = Registry() diff --git a/v1/op/base_op.py b/v1/op/base_op.py index d54854d7..ce08a31a 100644 --- a/v1/op/base_op.py +++ b/v1/op/base_op.py @@ -8,18 +8,23 @@ from v1.utils.timer import Timer class BaseOp(ABC): - def __init__(self, **kwargs): + def __init__(self, context: PipelineContext, **kwargs): super().__init__(**kwargs) - self.timer = Timer(name=self.__class__.__name__) + self.context: PipelineContext = context + self.timer = Timer(name=self.simple_name) + + @property + def simple_name(self) -> str: + return self.__class__.__name__.lower().replace("op", "") @abstractmethod - def execute(self, context: PipelineContext): + def execute(self): ... - def execute_wrap(self, context: PipelineContext): + def execute_wrap(self): try: with self.timer: - return self.execute(context) + return self.execute() except Exception as e: - logger.exception(f"OP.{self.__class__.__name__} execute failed, error={e.args}") + logger.exception(f"op={self.simple_name} execute failed, error={e.args}") diff --git a/v1/op/mock_op.py b/v1/op/mock_op.py index ea2f4a83..7449011c 100644 --- a/v1/op/mock_op.py +++ b/v1/op/mock_op.py @@ -2,12 +2,11 @@ import time from loguru import logger -from v1.op import OPERATION_REGISTRY +from v1.op import OP_REGISTRY from v1.op.base_op import BaseOp -from v1.pipeline.pipeline_context import PipelineContext -@OPERATION_REGISTRY.register("mock1") +@OP_REGISTRY.register("mock1_op") class MockOp1(BaseOp): def __init__(self, a: int, b: str, **kwargs): @@ -15,31 +14,31 @@ class MockOp1(BaseOp): self.a = a self.b = b - def execute(self, context: PipelineContext): + def execute(self): time.sleep(3) logger.info(f"enter class={self.__class__.__name__}. a={self.a} b={self.b}") -@OPERATION_REGISTRY.register("mock2") +@OP_REGISTRY.register("mock2_op") class MockOp2(MockOp1): ... -@OPERATION_REGISTRY.register("mock3") +@OP_REGISTRY.register("mock3_op") class MockOp3(MockOp1): ... -@OPERATION_REGISTRY.register("mock4") +@OP_REGISTRY.register("mock4_op") class MockOp4(MockOp1): ... -@OPERATION_REGISTRY.register("mock5") +@OP_REGISTRY.register("mock5_op") class MockOp5(MockOp1): ... -@OPERATION_REGISTRY.register("mock6") +@OP_REGISTRY.register("mock6_op") class MockOp6(MockOp1): ... diff --git a/v1/pipeline/pipeline.py b/v1/pipeline/pipeline.py index f984055b..f8147885 100644 --- a/v1/pipeline/pipeline.py +++ b/v1/pipeline/pipeline.py @@ -1,9 +1,10 @@ -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import as_completed from itertools import zip_longest -from typing import Dict, List +from typing import List from loguru import logger +from v1.op import OP_REGISTRY from v1.op.base_op import BaseOp from v1.pipeline.pipeline_context import PipelineContext from v1.utils.timer import Timer, timer @@ -13,20 +14,9 @@ class Pipeline(object): seq_symbol: str = "->" parallel_symbol: str = "|" - def __init__(self, - name: str, - pipeline: str, - op_config_dict: Dict[str, dict], - op_registry: Dict[str, type[BaseOp]], - context: PipelineContext, - thread_pool: ThreadPoolExecutor): - - self.name: str = name + def __init__(self, pipeline: str, context: PipelineContext): self.pipeline_list: List[str | List[str]] = self._parse_pipline(pipeline) - self.op_config_dict: Dict[str, dict] = op_config_dict - self.op_registry: Dict[str, type[BaseOp]] = op_registry self.context: PipelineContext = context - self.thread_pool: ThreadPoolExecutor = thread_pool def _parse_pipline(self, pipeline: str) -> List[str | List[str]]: pipeline_list: List[str | List[str]] = [] @@ -44,21 +34,22 @@ class Pipeline(object): return pipeline_list - def execute_sub_pipeline(self, pipeline: str, context: PipelineContext): + def _execute_sub_pipeline(self, pipeline: str): + op_config_dict = self.context.app_config.op for op in pipeline.split(self.seq_symbol): op = op.strip() if not op: continue - assert op in self.op_config_dict, f"op({op}).config is missing!" - backend = self.op_config_dict.pop("backend", "") - assert backend in self.op_registry, f"op({op}).backend({backend}) is not registered!" + assert op in op_config_dict, f"op={op} config is missing!" + backend = op_config_dict[op].backend + assert backend in OP_REGISTRY, f"op={op} backend={backend} is not registered!" - op_cls = self.op_registry[backend] - op_obj: BaseOp = op_cls(**self.op_config_dict[op]) - op_obj.execute_wrap(context) + op_cls = OP_REGISTRY[backend] + op_obj: BaseOp = op_cls(context=self.context, **op_config_dict[op].params) + op_obj.execute_wrap() - def parse_sub_pipeline(self, pipeline: str): + def _parse_sub_pipeline(self, pipeline: str): for op in pipeline.split(self.seq_symbol): op = op.strip() if not op: @@ -66,17 +57,16 @@ class Pipeline(object): yield op - @timer() def print_pipeline(self): i: int = 0 for pipeline in self.pipeline_list: if isinstance(pipeline, str): - for op in self.parse_sub_pipeline(pipeline): + for op in self._parse_sub_pipeline(pipeline): i += 1 logger.info(f"stage_{i}: {op}") elif isinstance(pipeline, list): - for op_list in zip_longest(*[self.parse_sub_pipeline(x) for x in pipeline], fillvalue="-"): + for op_list in zip_longest(*[self._parse_sub_pipeline(x) for x in pipeline], fillvalue="-"): i += 1 logger.info(f"stage{i}: {' | '.join(op_list)}") @@ -84,18 +74,19 @@ class Pipeline(object): raise ValueError(f"unknown pipeline.type={type(pipeline)}") @timer() - def execute_pipeline(self): + def execute_pipeline(self, enable_print: bool = True): + if enable_print: + self.print_pipeline() + for i, pipeline in enumerate(self.pipeline_list): with Timer(f"step_{i}"): if isinstance(pipeline, str): - self.execute_sub_pipeline(pipeline, self.context) + self._execute_sub_pipeline(pipeline) else: future_list = [] for sub_pipeline in pipeline: - future = self.thread_pool.submit(self.execute_sub_pipeline, - pipeline=sub_pipeline, - context=self.context) + future = self.context.thread_pool.submit(self._execute_sub_pipeline, pipeline=sub_pipeline) future_list.append(future) for future in as_completed(future_list): diff --git a/v1/pipeline/pipeline_context.py b/v1/pipeline/pipeline_context.py index c13937b2..9a705178 100644 --- a/v1/pipeline/pipeline_context.py +++ b/v1/pipeline/pipeline_context.py @@ -1,7 +1,18 @@ from concurrent.futures import ThreadPoolExecutor +from typing import Any + +from v1.schema.app_config import AppConfig class PipelineContext(object): - def __init__(self): - self.thread_pool: ThreadPoolExecutor = ThreadPoolExecutor(max_workers=10) \ No newline at end of file + def __init__(self, app_config: AppConfig, thread_pool: ThreadPoolExecutor): + self.app_config: AppConfig = app_config + self.thread_pool: ThreadPoolExecutor = thread_pool + self.context: dict = {} + + def set_context(self, key: str, value: Any): + self.context[key] = value + + def get_context(self, key: str): + return self.context.get(key) diff --git a/v1/schema/app_config.py b/v1/schema/app_config.py index 7d1a2bed..bcc05937 100644 --- a/v1/schema/app_config.py +++ b/v1/schema/app_config.py @@ -17,10 +17,10 @@ class ThreadPoolConfig: @dataclass class APIConfig: - retriever_pipeline: str = field(default="") - summarizer_pipeline: str = field(default="") - vector_store_pipeline: str = field(default="") - agent_pipeline: str = field(default="") + retriever: str = field(default="") + summarizer: str = field(default="") + vector_store: str = field(default="") + agent: str = field(default="") @dataclass diff --git a/v1/schema/experience.py b/v1/schema/experience.py index 0fc1ae64..4e8382cf 100644 --- a/v1/schema/experience.py +++ b/v1/schema/experience.py @@ -5,7 +5,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field -from experiencemaker.schema.vector_store_node import VectorStoreNode +from v1.schema.vector_store_node import VectorStoreNode class ExperienceMeta(BaseModel): diff --git a/v1/schema/request.py b/v1/schema/request.py index 67e49a3d..939fc0e5 100644 --- a/v1/schema/request.py +++ b/v1/schema/request.py @@ -7,14 +7,14 @@ from v1.schema.message import Message, Trajectory class BaseRequest(BaseModel, ABC): - metadata: dict = Field(default_factory=dict) workspace_id: str = Field(default="") + params: dict = Field(default_factory=dict) + metadata: dict | None = Field(default=None) class RetrieverRequest(BaseRequest): query: str = Field(default="") messages: List[Message] = Field(default_factory=list) - top_k: int = Field(default=3) class SummarizerRequest(BaseRequest): @@ -23,7 +23,6 @@ class SummarizerRequest(BaseRequest): class VectorStoreRequest(BaseRequest): action: str = Field(default="") - params: dict = Field(default_factory=dict) class AgentRequest(BaseRequest): diff --git a/v1/schema/response.py b/v1/schema/response.py index a3d04d94..c7c369e2 100644 --- a/v1/schema/response.py +++ b/v1/schema/response.py @@ -1,4 +1,3 @@ -from abc import ABC from typing import List from pydantic import BaseModel, Field @@ -7,7 +6,7 @@ from v1.schema.experience import BaseExperienceNode from v1.schema.message import Message -class BaseResponse(BaseModel, ABC): +class BaseResponse(BaseModel): success: bool = Field(default=True) metadata: dict = Field(default_factory=dict) @@ -22,9 +21,7 @@ class SummarizerResponse(BaseResponse): class VectorStoreResponse(BaseResponse): - action: str = Field(default="") - params: dict = Field(default_factory=dict) - + ... class AgentResponse(BaseResponse): answer: str = Field(default="") diff --git a/v1/service/service.py b/v1/service/service.py deleted file mode 100644 index d6a78903..00000000 --- a/v1/service/service.py +++ /dev/null @@ -1,2 +0,0 @@ - -class Service \ No newline at end of file diff --git a/v1/utils/config_parser.py b/v1/utils/config_parser.py new file mode 100644 index 00000000..6f03b610 --- /dev/null +++ b/v1/utils/config_parser.py @@ -0,0 +1,35 @@ +import json + +from loguru import logger +from omegaconf import OmegaConf, DictConfig + +from v1.schema.app_config import AppConfig + + +class ConfigParser(object): + + def __init__(self, args: list): + assert len(args) >= 1, "The `args` require at least one argument." + assert "config_path=" in args[0], f"The 0th args must be config_path, for example: `config_path=XXX`." + config_path: str = args[0].replace("config_path=", "") + yaml_config = OmegaConf.load(config_path) + + self.app_config: DictConfig = OmegaConf.structured(AppConfig) + self.app_config = OmegaConf.merge(self.app_config, yaml_config) + + if args[1:]: + cli_config = OmegaConf.from_dotlist(args[1:]) + self.app_config = OmegaConf.merge(self.app_config, cli_config) + + app_config_dict = OmegaConf.to_container(self.app_config, resolve=True) + app_config_str = json.dumps(app_config_dict, indent=2, ensure_ascii=False) + logger.info(f"app_config_str={app_config_str}") + + def get_app_config(self, **kwargs) -> AppConfig: + app_config = self.app_config.copy() + if kwargs: + kwargs_list = [f"{k}={v}" for k, v in kwargs.items()] + update_config = OmegaConf.from_dotlist(kwargs_list) + app_config = OmegaConf.merge(app_config, update_config) + + return OmegaConf.to_object(app_config) diff --git a/v1/singleton.py b/v1/utils/singleton.py similarity index 100% rename from v1/singleton.py rename to v1/utils/singleton.py