flowllm 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

flowllm/op/akshare/get_ak_a_code_op.py

@@ -0,0 +1,116 @@
+import json
+import time
+from typing import List
+
+import akshare as ak
+import pandas as pd
+from loguru import logger
+
+from flowllm.config.pydantic_config_parser import get_default_config
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.enumeration.role import Role
+from flowllm.op.llm_base_op import BaseLLMOp
+from flowllm.schema.message import Message
+from flowllm.utils.dataframe_cache import DataFrameCache
+from flowllm.utils.timer import timer
+
+
+@C.register_op()
+class GetAkACodeOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def __init__(self, language: str = "zh", llm="qwen3_30b_instruct", **kwargs):
+        super().__init__(language=language, llm=llm, **kwargs)
+
+    @staticmethod
+    def download_a_stock_df():
+        df_cache = DataFrameCache()
+        save_df_key: str = "all_a_stock_name_code"
+        if not df_cache.exists(save_df_key):
+            stock_sh_a_spot_em_df = ak.stock_sh_a_spot_em()
+            stock_sz_a_spot_em_df = ak.stock_sz_a_spot_em()
+            stock_bj_a_spot_em_df = ak.stock_bj_a_spot_em()
+
+            df: pd.DataFrame = pd.concat([stock_sh_a_spot_em_df, stock_sz_a_spot_em_df, stock_bj_a_spot_em_df], axis=0)
+            df = df.drop(columns=["序号"])
+            df = df.reset_index(drop=True)
+            df = df.sort_values(by="代码")
+            df_cache.save(save_df_key, df, expire_hours=0.25)
+
+        df = df_cache.load(save_df_key, dtype={"代码": str})
+        return df
+
+    def get_name_code_dict(self) -> dict:
+        df = self.download_a_stock_df()
+
+        name_code_dict = {}
+        for line in df.to_dict(orient="records"):
+            name = line["名称"].replace(" ", "")
+            code = line["代码"]
+            name_code_dict[name] = code
+        logger.info(f"name_code_dict.size={len(name_code_dict)} content={str(name_code_dict)[:50]}...")
+        return name_code_dict
+
+    @staticmethod
+    def split_list(array_list: list, n: int):
+        if n <= 0:
+            raise ValueError
+
+        length = len(array_list)
+        base_size = length // n
+        remainder = length % n
+
+        start = 0
+        for i in range(n):
+            size = base_size + (1 if i < remainder else 0)
+            end = start + size
+            yield array_list[start:end]
+            start = end
+
+    @timer()
+    def find_stock_codes(self, stock_names: List[str]):
+        stock_names = "\n".join([x.strip() for x in stock_names if x])
+        prompt = self.prompt_format(prompt_name="find_stock_name",
+                                    stock_names=stock_names,
+                                    query=self.flow_context.query)
+        logger.info(f"prompt={prompt}")
+
+        def callback_fn(msg: Message):
+            content = msg.content
+            if "```" in content:
+                content = content.split("```")[1]
+                content = content.strip("json")
+            content = json.loads(content.strip())
+            return content
+
+        codes: List[str] = self.llm.chat(messages=[Message(role=Role.USER, content=prompt)],
+                                         enable_stream_print=False,
+                                         callback_fn=callback_fn)
+        return codes
+
+    def execute(self):
+        name_code_dict = self.get_name_code_dict()
+        stock_names = list(name_code_dict.keys())
+        for p_stock_names in self.split_list(stock_names, n=2):
+            self.submit_task(self.find_stock_codes, stock_names=p_stock_names)
+            time.sleep(1)
+
+        stock_names = sorted(set(self.join_task()))
+        self.flow_context.code_infos = {name_code_dict[n]: {} for n in stock_names}
+        logger.info(f"code_infos={self.flow_context.code_infos}")
+
+
+if __name__ == "__main__":
+    from concurrent.futures import ThreadPoolExecutor
+
+    C.thread_pool = ThreadPoolExecutor(max_workers=10)
+    flow_context = FlowContext()
+    service_config = get_default_config()
+    flow_context.query = "茅台和五粮现在价格多少？"
+    flow_context.service_config = service_config
+
+    op = GetAkACodeOp(flow_context=flow_context)
+    # for x in op.split_list(list(range(10)), 3):
+    #     print(x)
+    op.execute()

flowllm/op/akshare/get_ak_a_code_prompt.yaml

@@ -0,0 +1,21 @@
+find_stock_name_zh: |
+  # 股票标准名称
+  {stock_names}
+
+  # 用户问题
+  {query}
+
+  # 任务
+  请你提取出用户问题中提到的股票名称，并通过**股票标准名称**找到对应的标准名称，并以json格式返回。
+  如果用户问题中没有提及股票名称或者**股票标准名称**中没有找到，请返回空"[]"。
+  如果用户只是提到了某个行业，并且没有明确提及股票名称，请返回空"[]"。
+  请思考后输出你的答案。
+  
+  # 答案格式
+  ```json
+  [
+      "股票标准名称1",
+      "股票标准名称2",
+      ...
+  ]
+  ```

flowllm/op/akshare/get_ak_a_info_op.py

@@ -0,0 +1,143 @@
+import json
+import time
+
+import akshare as ak
+import pandas as pd
+from loguru import logger
+
+from flowllm.config.pydantic_config_parser import get_default_config
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.op.base_op import BaseOp
+from flowllm.utils.fetch_url import fetch_webpage_text
+
+
+@C.register_op()
+class GetAkAInfoOp(BaseOp):
+
+    def execute_code(self, code: str) -> dict:
+        df = ak.stock_individual_info_em(symbol=code)
+        result = {}
+        for line in df.to_dict(orient="records"):
+            result[line["item"].strip()] = line["value"]
+        return {"基本信息": result}
+
+    def execute(self):
+        max_retries: int = self.op_params.get("max_retries", 3)
+        for code, info_dict in self.flow_context.code_infos.items():
+            result = {}
+            for i in range(max_retries):
+                try:
+                    result = self.execute_code(code)
+                    break
+
+                except Exception as _:
+                    if i != max_retries - 1:
+                        time.sleep(i * 2 + 1)
+
+            if result:
+                info_dict.update(result)
+
+        time.sleep(1)
+        logger.info(f"code_infos={json.dumps(self.flow_context.code_infos, ensure_ascii=False, indent=2)}")
+
+
+@C.register_op()
+class GetAkASpotOp(GetAkAInfoOp):
+
+    def execute_code(self, code: str) -> dict:
+        from flowllm.op import GetAkACodeOp
+
+        df: pd.DataFrame = GetAkACodeOp.download_a_stock_df()
+        df = df.loc[df["代码"] == code, :]
+        result = {}
+        if len(df) > 0:
+            result["实时行情"] = df.to_dict(orient="records")[-1]
+
+        return result
+
+
+@C.register_op()
+class GetAkAMoneyFlowOp(GetAkAInfoOp):
+
+    def execute_code(self, code: str) -> dict:
+        df = ak.stock_individual_fund_flow(stock=code)
+        result = {}
+        if len(df) > 0:
+            result["资金流入流出"] = {k: str(v) for k, v in df.to_dict(orient="records")[-1].items()}
+        return result
+
+
+@C.register_op()
+class GetAkAFinancialInfoOp(GetAkAInfoOp):
+
+    def execute_code(self, code: str) -> dict:
+        df = ak.stock_financial_abstract_ths(symbol=code, indicator="按报告期")
+        result = {}
+        if len(df) > 0:
+            result["财务信息"] = {k: str(v) for k, v in df.to_dict(orient="records")[-1].items()}
+        return result
+
+
+@C.register_op()
+class GetAkANewsOp(GetAkAInfoOp):
+
+    def execute_code(self, code: str) -> dict:
+        stock_news_em_df = ak.stock_news_em(symbol=code)
+        top_n_news: int = self.op_params.get("top_n_news", 1)
+
+        news_content_list = []
+        for i, line in enumerate(stock_news_em_df.to_dict(orient="records")[:top_n_news]):
+            url = line["新闻链接"]
+            # http://finance.eastmoney.com/a/202508133482756869.html
+            ts = url.split("/")[-1].split(".")[0]
+            date = ts[:8]
+            content = fetch_webpage_text(url).strip()
+            content = f"新闻{i}\n时间{date}\n{content}"
+            news_content_list.append(content)
+
+        return {"新闻": "\n\n".join(news_content_list)}
+
+
+@C.register_op()
+class MergeAkAInfoOp(BaseOp):
+
+    def execute(self):
+        code_content = {}
+        for code, info_dict in self.flow_context.code_infos.items():
+            content_list = [f"\n\n### {code}"]
+            for key, value in info_dict.items():
+                content_list.append(f"\n#### {code}-{key}")
+                if isinstance(value, str):
+                    content_list.append(value)
+                elif isinstance(value, dict):
+                    for attr_name, attr_value in value.items():
+                        content_list.append(f"{attr_name}: {attr_value}")
+                elif isinstance(value, list):
+                    content_list.extend([x.strip() for x in value if x])
+
+            code_content[code] = "\n".join(content_list)
+
+        answer = "\n".join(code_content.values())
+        logger.info(f"answer=\n{answer}")
+        self.flow_context.response.answer = answer.strip()
+
+
+if __name__ == "__main__":
+    from flowllm.schema.flow_response import FlowResponse
+
+    code_infos = {"000858": {}, "600519": {}}
+    flow_context = FlowContext(code_infos=code_infos, response=FlowResponse())
+    service_config = get_default_config()
+    flow_context.query = "茅台和五粮现在价格多少？"
+    flow_context.service_config = service_config
+
+    op1 = GetAkAInfoOp(flow_context=flow_context)
+    op2 = GetAkASpotOp(flow_context=flow_context)
+    op3 = GetAkAMoneyFlowOp(flow_context=flow_context)
+    op4 = GetAkAFinancialInfoOp(flow_context=flow_context)
+    op5 = GetAkANewsOp(flow_context=flow_context)
+    op6 = MergeAkAInfoOp(flow_context=flow_context)
+
+    op = op1 >> op2 >> op3 >> op4 >> op5 >> op6
+    op.execute()

flowllm/op/base_op.py

@@ -0,0 +1,169 @@
+"""
+BaseOp operator overloading implementation
+
+Supported operators:
+- op1 >> op2: Sequential execution, output of op1 becomes input of op2
+- op1 | op2: Parallel execution, both operations use the same input, returns list of results
+- Mixed calls: op1 >> (op2 | op3) >> op4
+
+Usage examples:
+    # Sequential execution
+    result = op1 >> op2 >> op3
+    
+    # Parallel execution
+    results = op1 | op2 | op3
+    
+    # Mixed calls
+    result = op1 >> (op2 | op3) >> op4
+    result = op1 >> (op1 | (op2 >> op3)) >> op4
+"""
+
+from abc import abstractmethod, ABC
+from concurrent.futures import Future
+from typing import List
+
+from loguru import logger
+from tqdm import tqdm
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.utils.common_utils import camel_to_snake
+from flowllm.utils.timer import Timer
+
+
+class BaseOp(ABC):
+
+    def __init__(self,
+                 name: str = "",
+                 language: str = "",
+                 raise_exception: bool = True,
+                 flow_context: FlowContext | None = None,
+                 **kwargs):
+
+        super().__init__()
+
+        self.name: str = name or camel_to_snake(self.__class__.__name__)
+        self.language: str = language or C.language
+        self.raise_exception: bool = raise_exception
+
+        self.flow_context: FlowContext | None = flow_context
+        self.op_params: dict = kwargs
+
+        self.task_list: List[Future] = []
+        self.timer = Timer(name=self.name)
+
+    @abstractmethod
+    def execute(self):
+        ...
+
+    def __call__(self, *args, **kwargs):
+        with self.timer:
+            if self.raise_exception:
+                return self.execute()
+
+            else:
+                try:
+                    return self.execute()
+
+                except Exception as e:
+                    logger.exception(f"op={self.name} execute failed, error={e.args}")
+
+    def submit_task(self, fn, *args, **kwargs):
+        task = C.thread_pool.submit(fn, *args, **kwargs)
+        self.task_list.append(task)
+        return self
+
+    def join_task(self, task_desc: str = None) -> list:
+        result = []
+        for task in tqdm(self.task_list, desc=task_desc or self.name):
+            t_result = task.result()
+            if t_result:
+                if isinstance(t_result, list):
+                    result.extend(t_result)
+                else:
+                    result.append(t_result)
+        self.task_list.clear()
+        return result
+
+    def __rshift__(self, op: "BaseOp"):
+        from flowllm.op.sequential_op import SequentialOp
+
+        sequential_op = SequentialOp(ops=[self], flow_context=self.flow_context)
+
+        if isinstance(op, SequentialOp):
+            sequential_op.ops.extend(op.ops)
+        else:
+            sequential_op.ops.append(op)
+        return sequential_op
+
+    def __or__(self, op: "BaseOp"):
+        from flowllm.op.parallel_op import ParallelOp
+
+        parallel_op = ParallelOp(ops=[self], flow_context=self.flow_context)
+
+        if isinstance(op, ParallelOp):
+            parallel_op.ops.extend(op.ops)
+        else:
+            parallel_op.ops.append(op)
+
+        return parallel_op
+
+
+def run1():
+    """Basic test"""
+
+    class MockOp(BaseOp):
+        def execute(self):
+            logger.info(f"op={self.name} execute")
+
+    mock_op = MockOp()
+    mock_op()
+
+
+def run2():
+    """Test operator overloading functionality"""
+    from concurrent.futures import ThreadPoolExecutor
+    import time
+
+    class TestOp(BaseOp):
+
+        def execute(self, data=None):
+            time.sleep(0.1)  # Simulate execution time
+            op_result = f"{self.name}({data})" if data else self.name
+            logger.info(f"Executing {op_result}")
+            return op_result
+
+    # Create service_context for parallel execution
+    C["thread_pool"] = ThreadPoolExecutor(max_workers=4)
+
+    # Create test operations
+    op1 = TestOp("op1")
+    op2 = TestOp("op2")
+    op3 = TestOp("op3")
+    op4 = TestOp("op4")
+
+    logger.info("=== Testing sequential execution op1 >> op2 ===")
+    sequential = op1 >> op2
+    result = sequential()
+    logger.info(f"Sequential result: {result}")
+
+    logger.info("=== Testing parallel execution op1 | op2 ===")
+    parallel = op1 | op2
+    result = parallel()
+    logger.info(f"Parallel result: {result}")
+
+    logger.info("=== Testing mixed calls op1 >> (op2 | op3) >> op4 ===")
+    mixed = op1 >> (op2 | op3) >> op4
+    result = mixed()
+    logger.info(f"Mixed result: {result}")
+
+    logger.info("=== Testing complex mixed calls op1 >> (op1 | (op2 >> op3)) >> op4 ===")
+    complex_mixed = op1 >> (op1 | (op2 >> op3)) >> op4
+    result = complex_mixed()
+    logger.info(f"Complex mixed result: {result}")
+
+
+if __name__ == "__main__":
+    run1()
+    print("\n" + "=" * 50 + "\n")
+    run2()

flowllm/op/llm_base_op.py

@@ -0,0 +1,63 @@
+from abc import ABC
+from pathlib import Path
+
+from flowllm.context.prompt_handler import PromptHandler
+from flowllm.context.service_context import C
+from flowllm.embedding_model.base_embedding_model import BaseEmbeddingModel
+from flowllm.llm.base_llm import BaseLLM
+from flowllm.op.base_op import BaseOp
+from flowllm.schema.service_config import LLMConfig, EmbeddingModelConfig
+from flowllm.storage.vector_store.base_vector_store import BaseVectorStore
+
+
+class BaseLLMOp(BaseOp, ABC):
+    file_path: str = __file__
+
+    def __init__(self,
+                 prompt_path: str = "",
+                 llm: str = "default",
+                 embedding_model: str = "default",
+                 vector_store: str = "default",
+                 **kwargs):
+
+        super().__init__(**kwargs)
+
+        self._llm: BaseLLM | str = llm
+        self._embedding_model: BaseEmbeddingModel | str = embedding_model
+        self._vector_store: BaseVectorStore | str = vector_store
+
+        default_prompt_path: Path = Path(self.file_path).parent / self.name.replace("_op", "_prompt.yaml")
+        self.prompt_path: Path = Path(prompt_path) if prompt_path else default_prompt_path
+        self.prompt = PromptHandler(language=self.language).load_prompt_by_file(self.prompt_path)
+
+    @property
+    def llm(self) -> BaseLLM:
+        if isinstance(self._llm, str):
+            llm_config: LLMConfig = self.flow_context.service_config.llm[self._llm]
+            llm_cls = C.resolve_llm(llm_config.backend)
+            self._llm = llm_cls(model_name=llm_config.model_name, **llm_config.params)
+
+        return self._llm
+
+    @property
+    def embedding_model(self) -> BaseEmbeddingModel:
+        if isinstance(self._embedding_model, str):
+            embedding_model_config: EmbeddingModelConfig = \
+                self.flow_context.service_config.embedding_model[self._embedding_model]
+            embedding_model_cls = C.resolve_embedding_model(embedding_model_config.backend)
+            self._embedding_model = embedding_model_cls(model_name=embedding_model_config.model_name,
+                                                        **embedding_model_config.params)
+
+        return self._embedding_model
+
+    @property
+    def vector_store(self) -> BaseVectorStore:
+        if isinstance(self._vector_store, str):
+            self._vector_store = C.get_vector_store(self._vector_store)
+        return self._vector_store
+
+    def prompt_format(self, prompt_name: str, **kwargs) -> str:
+        return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs)
+
+    def get_prompt(self, prompt_name: str) -> str:
+        return self.prompt.get_prompt(prompt_name=prompt_name)

flowllm/op/mock_op.py

@@ -0,0 +1,42 @@
+import time
+
+from loguru import logger
+
+from flowllm.context.service_context import C
+from flowllm.op.llm_base_op import BaseLLMOp
+
+
+@C.register_op()
+class Mock1Op(BaseLLMOp):
+    def execute(self):
+        time.sleep(1)
+        a = self.flow_context.a
+        b = self.flow_context.b
+        logger.info(f"enter class={self.name}. a={a} b={b}")
+
+        self.flow_context.response.answer = f"{self.name} {a} {b} answer=47"
+
+
+@C.register_op()
+class Mock2Op(Mock1Op):
+    ...
+
+
+@C.register_op()
+class Mock3Op(Mock1Op):
+    ...
+
+
+@C.register_op()
+class Mock4Op(Mock1Op):
+    ...
+
+
+@C.register_op()
+class Mock5Op(Mock1Op):
+    ...
+
+
+@C.register_op()
+class Mock6Op(Mock1Op):
+    ...

flowllm/op/parallel_op.py

@@ -0,0 +1,30 @@
+from typing import List
+
+from flowllm.op.base_op import BaseOp
+
+
+class ParallelOp(BaseOp):
+    """Container class for parallel operation execution
+
+    Executes multiple operations in parallel, all operations use the same input,
+    returns a list of results from all operations.
+    Supports parallel calls: op1 | op2 | op3
+    Falls back to sequential execution if no thread pool is available.
+    """
+
+    def __init__(self, ops: List[BaseOp], **kwargs):
+        super().__init__(**kwargs)
+        self.ops = ops
+
+    def execute(self):
+        for op in self.ops:
+            self.submit_task(op.__call__)
+
+        return self.join_task(task_desc="Parallel execution")
+
+    def __or__(self, op: BaseOp):
+        if isinstance(op, ParallelOp):
+            self.ops.extend(op.ops)
+        else:
+            self.ops.append(op)
+        return self

flowllm/op/sequential_op.py

@@ -0,0 +1,29 @@
+from typing import List
+
+from flowllm.op.base_op import BaseOp
+
+
+class SequentialOp(BaseOp):
+    """Container class for sequential operation execution
+
+    Executes multiple operations in sequence, where the output of the previous operation
+    becomes the input of the next operation.
+    Supports chaining: op1 >> op2 >> op3
+    """
+
+    def __init__(self, ops: List[BaseOp], **kwargs):
+        super().__init__(**kwargs)
+        self.ops = ops
+
+    def execute(self):
+        result = None
+        for op in self.ops:
+            result = op.execute()
+        return result
+
+    def __rshift__(self, op: BaseOp):
+        if isinstance(op, SequentialOp):
+            self.ops.extend(op.ops)
+        else:
+            self.ops.append(op)
+        return self

flowllm/schema/flow_response.py

@@ -0,0 +1,12 @@
+from typing import List
+
+from pydantic import Field, BaseModel
+
+from flowllm.schema.message import Message
+
+
+class FlowResponse(BaseModel):
+    answer: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    success: bool = Field(default=True)
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/message.py

@@ -0,0 +1,35 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from flowllm.enumeration.role import Role
+from flowllm.schema.tool_call import ToolCall
+
+
+class Message(BaseModel):
+    role: Role = Field(default=Role.USER)
+    content: str | bytes = Field(default="")
+    reasoning_content: str = Field(default="")
+    tool_calls: List[ToolCall] = Field(default_factory=list)
+    tool_call_id: str = Field(default="")
+    metadata: dict = Field(default_factory=dict)
+
+    def simple_dump(self, add_reason_content: bool = True) -> dict:
+        result: dict
+        if self.content:
+            result = {"role": self.role.value, "content": self.content}
+        elif add_reason_content and self.reasoning_content:
+            result = {"role": self.role.value, "content": self.reasoning_content}
+        else:
+            result = {"role": self.role.value, "content": ""}
+
+        if self.tool_calls:
+            result["tool_calls"] = [x.simple_output_dump() for x in self.tool_calls]
+        return result
+
+
+class Trajectory(BaseModel):
+    task_id: str = Field(default="")
+    messages: List[Message] = Field(default_factory=list)
+    score: float = Field(default=0.0)
+    metadata: dict = Field(default_factory=dict)