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

flowllm/op/akshare/get_ak_a_info_op.py

@@ -0,0 +1,140 @@
+import json
+import time
+
+import akshare as ak
+import pandas as pd
+from loguru import logger
+from tqdm import tqdm
+
+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.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.context.code_infos, ensure_ascii=False, indent=2)}")
+
+
+@C.register_op()
+class GetAkASpotOp(GetAkAInfoOp):
+
+    def execute_code(self, code: str) -> dict:
+        from flowllm.op.akshare 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(tqdm(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.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.context.response.answer = answer.strip()
+
+
+if __name__ == "__main__":
+    C.set_default_service_config().init_by_service_config()
+
+    code_infos = {"000858": {}, "600519": {}}
+    context = FlowContext(code_infos=code_infos, query="茅台和五粮现在价格多少？")
+
+    op1 = GetAkAInfoOp()
+    op2 = GetAkASpotOp()
+    op3 = GetAkAMoneyFlowOp()
+    op4 = GetAkAFinancialInfoOp()
+    op5 = GetAkANewsOp()
+    op6 = MergeAkAInfoOp()
+
+    op = op1 >> op2 >> op3 >> op4 >> op5 >> op6
+    op(context=context)

flowllm/op/base_llm_op.py

@@ -0,0 +1,64 @@
+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,
+                 language: str = "",
+                 prompt_path: str = "",
+                 llm: str = "default",
+                 embedding_model: str = "default",
+                 vector_store: str = "default",
+                 **kwargs):
+        super().__init__(**kwargs)
+
+        self.language: str = language or C.language
+        self.prompt_path: Path = Path(prompt_path) if prompt_path else \
+            Path(self.file_path).parent / self.name.replace("_op", "_prompt.yaml")
+        self._llm: BaseLLM | str = llm
+        self._embedding_model: BaseEmbeddingModel | str = embedding_model
+        self._vector_store: BaseVectorStore | str = vector_store
+
+        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 = C.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 = \
+                C.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/base_op.py

@@ -0,0 +1,148 @@
+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 = "",
+                 raise_exception: bool = True,
+                 **kwargs):
+        super().__init__()
+
+        self.name: str = name or camel_to_snake(self.__class__.__name__)
+        self.raise_exception: bool = raise_exception
+        self.op_params: dict = kwargs
+
+        self.task_list: List[Future] = []
+        self.ray_task_list: List = []  # Ray ObjectRef list
+        self.timer = Timer(name=self.name)
+        self.context: FlowContext | None = None
+
+    @abstractmethod
+    def execute(self):
+        ...
+
+    def __call__(self, context: FlowContext = None):
+        self.context = context
+        with self.timer:
+            if self.raise_exception:
+                self.execute()
+
+            else:
+
+                try:
+                    self.execute()
+                except Exception as e:
+                    logger.exception(f"op={self.name} execute failed, error={e.args}")
+
+        return self.context.response if self.context else None
+
+    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])
+
+        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])
+
+        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):
+            time.sleep(0.1)
+            op_result = f"{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/base_ray_op.py

@@ -0,0 +1,313 @@
+from abc import ABC
+
+import pandas as pd
+import ray
+from loguru import logger
+from tqdm import tqdm
+
+from flowllm.context.service_context import C
+from flowllm.op.base_op import BaseOp
+
+
+class BaseRayOp(BaseOp, ABC):
+    """
+    Base class for Ray-based operations that provides parallel task execution capabilities.
+    Inherits from BaseOp and provides methods for submitting and joining Ray tasks.
+    """
+
+    def submit_and_join_ray_task(self, fn, parallel_key: str = "", task_desc: str = "",
+                                 enable_test: bool = False, **kwargs):
+        """
+        Submit multiple Ray tasks in parallel and wait for all results.
+        
+        This method automatically detects a list parameter to parallelize over, distributes
+        the work across multiple Ray workers, and returns the combined results.
+        
+        Args:
+            fn: Function to execute in parallel
+            parallel_key: Key of the parameter to parallelize over (auto-detected if empty)
+            task_desc: Description for logging and progress bars
+            enable_test: Enable test mode (prints results instead of executing)
+            **kwargs: Arguments to pass to the function, including the list to parallelize over
+        
+        Returns:
+            List of results from all parallel tasks
+        """
+        max_workers = C.service_config.ray_max_workers
+        self.ray_task_list.clear()
+
+        # Auto-detect parallel key if not provided
+        if not parallel_key:
+            for key, value in kwargs.items():
+                if isinstance(value, list):
+                    parallel_key = key
+                    logger.info(f"using first list parallel_key={parallel_key}")
+                    break
+
+        # Extract the list to parallelize over
+        parallel_list = kwargs.pop(parallel_key)
+        assert isinstance(parallel_list, list)
+
+        # Convert pandas DataFrames to Ray objects for efficient sharing
+        for key in sorted(kwargs.keys()):
+            value = kwargs[key]
+            if isinstance(value, pd.DataFrame):
+                kwargs[key] = ray.put(value)
+
+        if enable_test:
+            test_result_list = []
+            for value in parallel_list:
+                kwargs.update({"actor_index": 0, parallel_key: value})
+                t_result = fn(**kwargs)
+                if t_result:
+                    if isinstance(t_result, list):
+                        test_result_list.extend(t_result)
+                    else:
+                        test_result_list.append(t_result)
+            return test_result_list
+
+        # Create and submit tasks for each worker
+        for i in range(max_workers):
+            def fn_wrapper():
+                result_list = []
+                # Distribute work using stride: worker i-th processes items [i, i+max_workers, i+2*max_workers, ...]
+                for parallel_value in parallel_list[i::max_workers]:
+                    kwargs.update({
+                        "actor_index": i,
+                        parallel_key: parallel_value,
+                    })
+                    part_result = fn(**kwargs)
+                    if part_result:
+                        if isinstance(part_result, list):
+                            result_list.extend(part_result)
+                        else:
+                            result_list.append(part_result)
+                return result_list
+
+            self.submit_ray_task(fn=fn_wrapper)
+            logger.info(f"ray.submit task_desc={task_desc} id={i}")
+
+        # Wait for all tasks to complete and collect results
+        result = self.join_ray_task(task_desc=task_desc)
+        logger.info(f"{task_desc} complete. result_size={len(result)} resources={ray.available_resources()}")
+        return result
+
+    def submit_ray_task(self, fn, *args, **kwargs):
+        """
+        Submit a single Ray task for asynchronous execution.
+        
+        Args:
+            fn: Function to execute remotely
+            *args: Positional arguments for the function
+            **kwargs: Keyword arguments for the function
+        
+        Returns:
+            Self for method chaining
+        
+        Raises:
+            RuntimeError: If Ray is not configured (ray_max_workers <= 1)
+        """
+        if C.service_config.ray_max_workers <= 1:
+            raise RuntimeError("Ray is not configured. Please set ray_max_workers > 1 in service config.")
+
+        # Initialize Ray if not already done
+        if not ray.is_initialized():
+            logger.warning(f"Ray is not initialized. Initializing Ray with {C.service_config.ray_max_workers} workers.")
+            ray.init(num_cpus=C.service_config.ray_max_workers)
+
+        # Create remote function and submit task
+        remote_fn = ray.remote(fn)
+        task = remote_fn.remote(*args, **kwargs)
+        self.ray_task_list.append(task)
+        return self
+
+    def join_ray_task(self, task_desc: str = None) -> list:
+        """
+        Wait for all submitted Ray tasks to complete and collect their results.
+        
+        Args:
+            task_desc: Description for the progress bar
+            
+        Returns:
+            Combined list of results from all completed tasks
+        """
+        result = []
+        # Process each task and collect results with progress bar
+        for task in tqdm(self.ray_task_list, desc=task_desc or f"{self.name}_ray"):
+            t_result = ray.get(task)
+            if t_result:
+                if isinstance(t_result, list):
+                    result.extend(t_result)
+                else:
+                    result.append(t_result)
+        self.ray_task_list.clear()
+        return result
+
+
+def run():
+    """Test Ray multiprocessing functionality"""
+    import time
+    import math
+
+    # CPU intensive task for testing
+    def cpu_intensive_task(n: int, task_id: str):
+        """CPU intensive task: calculate prime numbers"""
+        start_t = time.time()
+
+        def is_prime(num):
+            if num < 2:
+                return False
+            for j in range(2, int(math.sqrt(num)) + 1):
+                if num % j == 0:
+                    return False
+            return True
+
+        primes = [x for x in range(2, n) if is_prime(x)]
+        end_t = time.time()
+
+        result = {
+            'task_id': task_id,
+            'prime_count': len(primes),
+            'max_prime': max(primes) if primes else 0,
+            'execution_time': end_t - start_t
+        }
+        logger.info(f"Task {task_id} completed: found {len(primes)} primes, time: {result['execution_time']:.2f}s")
+        return result
+
+    class TestRayOp(BaseRayOp):
+        def execute(self):
+            logger.info(f"Executing {self.name}")
+            return f"Result from {self.name}"
+
+    # Initialize service config for Ray
+    from flowllm.schema.service_config import ServiceConfig
+
+    # Create a test service config with Ray enabled
+    test_config = ServiceConfig()
+    test_config.ray_max_workers = 4  # Enable Ray with 4 workers
+    test_config.thread_pool_max_workers = 4
+
+    # Set the service config
+    C.init_by_service_config(test_config)
+
+    logger.info("=== Testing Ray multiprocessing ===")
+
+    # Create test operation
+    ray_op = TestRayOp("ray_test_op")
+
+    logger.info("--- Testing submit_ray_task and join_ray_task ---")
+
+    # Test 1: Basic Ray task submission
+    task_size = 50000  # Find primes up to 50000 (more CPU intensive)
+    num_tasks = 4
+
+    try:
+        # Submit multiple CPU-intensive tasks
+
+        logger.info(f"Submitting {num_tasks} Ray tasks (finding primes up to {task_size})")
+        start_time = time.time()
+
+        for i in range(num_tasks):
+            ray_op.submit_ray_task(cpu_intensive_task, task_size, f"ray_task_{i}")
+
+        # Wait for all tasks to complete
+        results = ray_op.join_ray_task("Processing Ray tasks")
+        end_time = time.time()
+
+        logger.info(f"Ray tasks completed in {end_time - start_time:.2f}s")
+        logger.info(f"Ray results: {results}")
+
+    except Exception as e:
+        logger.error(f"Ray task execution failed: {e}")
+
+    # Test 2: Compare Ray vs ThreadPool performance
+    logger.info("\n--- Performance Comparison: Ray vs ThreadPool ---")
+
+    try:
+        # Test with ThreadPool
+        thread_op = TestRayOp("thread_test_op")
+
+        logger.info(f"Testing ThreadPool with {num_tasks} tasks")
+        start_time = time.time()
+
+        for i in range(num_tasks):
+            thread_op.submit_task(cpu_intensive_task, task_size, f"thread_task_{i}")
+
+        thread_results = thread_op.join_task("Processing ThreadPool tasks")
+        print(thread_results)
+        thread_time = time.time() - start_time
+
+        logger.info(f"ThreadPool completed in {thread_time:.2f}s")
+
+        # Test with Ray again for comparison
+        ray_op2 = TestRayOp("ray_test_op2")
+
+        logger.info(f"Testing Ray with {num_tasks} tasks")
+        start_time = time.time()
+
+        for i in range(num_tasks):
+            ray_op2.submit_ray_task(cpu_intensive_task, task_size, f"ray_task2_{i}")
+
+        ray_results2 = ray_op2.join_ray_task("Processing Ray tasks (comparison)")
+        print(ray_results2)
+        ray_time = time.time() - start_time
+
+        logger.info(f"Ray completed in {ray_time:.2f}s")
+
+        # Performance comparison
+        speedup = thread_time / ray_time if ray_time > 0 else 0
+        logger.info(f"\n=== Performance Summary ===")
+        logger.info(f"ThreadPool time: {thread_time:.2f}s")
+        logger.info(f"Ray time: {ray_time:.2f}s")
+        logger.info(f"Ray speedup: {speedup:.2f}x")
+
+    except Exception as e:
+        logger.error(f"Performance comparison failed: {e}")
+
+    # Test 3: Error handling
+    logger.info("\n--- Testing Error Handling ---")
+
+    def failing_task(task_id: str):
+        if task_id == "fail_task":
+            raise ValueError(f"Intentional error in {task_id}")
+        return f"Success: {task_id}"
+
+    try:
+        error_op = TestRayOp("error_test_op")
+
+        # Submit mix of successful and failing tasks
+        error_op.submit_ray_task(failing_task, "success_task_1")
+        error_op.submit_ray_task(failing_task, "fail_task")
+        error_op.submit_ray_task(failing_task, "success_task_2")
+
+        error_results = error_op.join_ray_task("Testing error handling")
+        logger.info(f"Error handling results: {error_results}")
+
+    except Exception as e:
+        logger.error(f"Expected error occurred: {e}")
+
+    # Test 4: Ray without proper configuration (should fail)
+    logger.info("\n--- Testing Ray Configuration Validation ---")
+
+    original_workers = C.service_config.ray_max_workers
+    try:
+        # Temporarily disable Ray in config
+        C.service_config.ray_max_workers = 1  # Disable Ray
+
+        config_test_op = TestRayOp("config_test_op")
+        config_test_op.submit_ray_task(cpu_intensive_task, 100, "config_test")
+
+        logger.error("This should not be reached - Ray should be disabled")
+
+    except RuntimeError as e:
+        logger.info(f"✓ Correctly caught configuration error: {e}")
+
+    finally:
+        # Restore original configuration
+        C.service_config.ray_max_workers = original_workers
+
+    logger.info("\n=== Ray testing completed ===")
+
+
+if __name__ == "__main__":
+    run()

flowllm/op/code/__init__.py

@@ -0,0 +1 @@
+from .execute_code_op import ExecuteCodeOp