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

flowllm/op/akshare/get_ak_a_info_op.py

@@ -4,8 +4,8 @@ import time
 import akshare as ak
 import pandas as pd
 from loguru import logger
+from tqdm import tqdm
 
-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
@@ -24,7 +24,7 @@ class GetAkAInfoOp(BaseOp):
 
     def execute(self):
         max_retries: int = self.op_params.get("max_retries", 3)
-        for code, info_dict in self.flow_context.code_infos.items():
+        for code, info_dict in self.context.code_infos.items():
             result = {}
             for i in range(max_retries):
                 try:
@@ -39,14 +39,14 @@ class GetAkAInfoOp(BaseOp):
                 info_dict.update(result)
 
         time.sleep(1)
-        logger.info(f"code_infos={json.dumps(self.flow_context.code_infos, ensure_ascii=False, indent=2)}")
+        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 import GetAkACodeOp
+        from flowllm.op.akshare import GetAkACodeOp
 
         df: pd.DataFrame = GetAkACodeOp.download_a_stock_df()
         df = df.loc[df["代码"] == code, :]
@@ -87,7 +87,7 @@ class GetAkANewsOp(GetAkAInfoOp):
         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]):
+        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]
@@ -104,7 +104,7 @@ class MergeAkAInfoOp(BaseOp):
 
     def execute(self):
         code_content = {}
-        for code, info_dict in self.flow_context.code_infos.items():
+        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}")
@@ -120,24 +120,21 @@ class MergeAkAInfoOp(BaseOp):
 
         answer = "\n".join(code_content.values())
         logger.info(f"answer=\n{answer}")
-        self.flow_context.response.answer = answer.strip()
+        self.context.response.answer = answer.strip()
 
 
 if __name__ == "__main__":
-    from flowllm.schema.flow_response import FlowResponse
+    C.set_default_service_config().init_by_service_config()
 
     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)
+    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.execute()
+    op(context=context)

flowllm/op/{llm_base_op.py → base_llm_op.py}

@@ -14,26 +14,28 @@ 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
+        default_prompt_path = self.file_path.replace("op.py", "prompt.yaml")
+        self.prompt_path: Path = Path(prompt_path) if prompt_path else default_prompt_path
+
         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_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)
 
@@ -43,7 +45,7 @@ class BaseLLMOp(BaseOp, ABC):
     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]
+                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)

flowllm/op/base_op.py

@@ -1,23 +1,3 @@
-"""
-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
@@ -35,60 +15,76 @@ class BaseOp(ABC):
 
     def __init__(self,
                  name: str = "",
-                 language: str = "",
                  raise_exception: bool = True,
-                 flow_context: FlowContext | None = None,
+                 enable_multithread: bool = True,
                  **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.enable_multithread: bool = enable_multithread
         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, *args, **kwargs):
+    def __call__(self, context: FlowContext = None):
+        self.context = context
         with self.timer:
             if self.raise_exception:
-                return self.execute()
+                self.execute()
 
             else:
-                try:
-                    return self.execute()
 
+                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)
+        if self.enable_multithread:
+            task = C.thread_pool.submit(fn, *args, **kwargs)
+            self.task_list.append(task)
+
+        else:
+            result = fn(*args, **kwargs)
+            if result:
+                if isinstance(result, list):
+                    result.extend(result)
+                else:
+                    result.append(result)
+
         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)
+        if self.enable_multithread:
+            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)
+
+        else:
+            result.extend(self.task_list)
+
         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)
+        sequential_op = SequentialOp(ops=[self])
 
         if isinstance(op, SequentialOp):
             sequential_op.ops.extend(op.ops)
@@ -99,7 +95,7 @@ class BaseOp(ABC):
     def __or__(self, op: "BaseOp"):
         from flowllm.op.parallel_op import ParallelOp
 
-        parallel_op = ParallelOp(ops=[self], flow_context=self.flow_context)
+        parallel_op = ParallelOp(ops=[self])
 
         if isinstance(op, ParallelOp):
             parallel_op.ops.extend(op.ops)
@@ -127,9 +123,9 @@ def run2():
 
     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
+        def execute(self):
+            time.sleep(0.1)
+            op_result = f"{self.name}"
             logger.info(f"Executing {op_result}")
             return op_result
 

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

flowllm/op/code/execute_code_op.py

@@ -0,0 +1,42 @@
+import sys
+from io import StringIO
+
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.op.base_op import BaseOp
+
+
+@C.register_op()
+class ExecuteCodeOp(BaseOp):
+
+    def execute(self):
+        old_stdout = sys.stdout
+        redirected_output = sys.stdout = StringIO()
+
+        try:
+            code_key: str = self.op_params.get("code_key", "code")
+            code_str: str = self.context[code_key]
+            exec(code_str)
+            code_result = redirected_output.getvalue()
+
+        except Exception as e:
+            logger.info(f"{self.name} encounter exception! error={e.args}")
+            code_result = str(e)
+
+        sys.stdout = old_stdout
+        self.context.code_result = code_result
+
+
+if __name__ == "__main__":
+    C.set_default_service_config().init_by_service_config()
+    op = ExecuteCodeOp()
+
+    context = FlowContext(code="print('Hello World')")
+    op(context=context)
+    print(context.code_result)
+
+    context.code = "print('Hello World!'"
+    op(context=context)
+    print(context.code_result)

flowllm/op/gallery/__init__.py

@@ -0,0 +1,2 @@
+from .mock_op import Mock1Op, Mock2Op, Mock3Op, Mock4Op, Mock5Op, Mock6Op
+from .terminate_op import TerminateOp

flowllm/op/{mock_op.py → gallery/mock_op.py}

@@ -3,18 +3,18 @@ import time
 from loguru import logger
 
 from flowllm.context.service_context import C
-from flowllm.op.llm_base_op import BaseLLMOp
+from flowllm.op.base_llm_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
+        a = self.context.a
+        b = self.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"
+        self.context.response.answer = f"{self.name} {a} {b} answer=47"
 
 
 @C.register_op()

flowllm/op/gallery/terminate_op.py

@@ -0,0 +1,29 @@
+from flowllm.context.service_context import C
+from flowllm.op.base_op import BaseOp
+
+
+@C.register_op()
+class TerminateOp(BaseOp):
+
+    def execute(self):
+        # Get status from context
+        status = self.context.status
+        assert status in ["success", "failure"], f"Invalid status: {status}"
+        self.context.terminate_answer = f"The interaction has been completed with status: {status}"
+
+
+if __name__ == "__main__":
+    from flowllm.context.flow_context import FlowContext
+
+    C.set_default_service_config().init_by_service_config()
+
+    # Test success termination
+    op = TerminateOp()
+    context = FlowContext(status="success")
+    result = op(context=context)
+    print(f"Result: {context.terminate_answer}")
+
+    # Test failure termination
+    context.status = "failure"
+    op(context=context)
+    print(f"Result: {context.terminate_answer}")