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

flowllm/flow/gallery/mock_tool_flow.py

@@ -7,61 +7,74 @@ from flowllm.schema.tool_call import ToolCall, ParamAttrs
 @C.register_tool_flow()
 class MockToolFlow(BaseToolFlow):
 
+    def __init__(self,
+                 use_async: bool = False,
+                 stream: bool = False,
+                 service_type: str = "http+mcp",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
     def build_flow(self):
         mock1_op = Mock1Op()
         mock2_op = Mock2Op()
         mock3_op = Mock3Op()
-        mock4_op = Mock4Op()
-        mock5_op = Mock5Op()
-        mock6_op = Mock6Op()
 
-        op = mock1_op >> ((mock4_op >> mock2_op) | mock5_op) >> (mock3_op | mock6_op)
+        op = mock1_op >> ((mock2_op >> mock3_op) | mock1_op) >> (mock2_op | mock3_op)
         return op
 
     def build_tool_call(self) -> ToolCall:
         return ToolCall(**{
-            "index": 0,
-            "id": "call_mock_tool_12345",
             "type": "function",
-            "name": "mock_data_processor",
+            "name": "mock_data",
             "description": "A mock tool that processes data through multiple operations and returns structured results",
-            "arguments": {
-                "input_data": "sample_data",
-                "processing_mode": "advanced",
-                "output_format": "json"
-            },
             "input_schema": {
-                "input_data": ParamAttrs(
-                    type="string",
+                "a": ParamAttrs(
+                    type="str",
                     description="The input data to be processed",
                     required=True
                 ),
-                "processing_mode": ParamAttrs(
+                "b": ParamAttrs(
                     type="string",
                     description="Processing mode: basic, advanced, or expert",
                     required=False
                 ),
-                "output_format": ParamAttrs(
-                    type="string",
-                    description="Output format: json, xml, or plain",
-                    required=False
-                )
-            },
-            "output_schema": {
-                "result": ParamAttrs(
-                    type="object",
-                    description="Processed result data",
+            }
+        })
+
+
+@C.register_tool_flow()
+class MockAsyncToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = False,
+                 service_type: str = "http+mcp",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        mock4_op = Mock4Op()
+        mock5_op = Mock5Op()
+        mock6_op = Mock6Op()
+
+        op = mock4_op >> ((mock5_op >> mock6_op) | mock4_op) >> (mock5_op | mock6_op)
+        return op
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "type": "function",
+            "name": "mock_data",
+            "description": "A mock tool that processes data through multiple operations and returns structured results",
+            "input_schema": {
+                "a": ParamAttrs(
+                    type="str",
+                    description="The input data to be processed",
                     required=True
                 ),
-                "status": ParamAttrs(
+                "b": ParamAttrs(
                     type="string",
-                    description="Processing status: success, warning, or error",
-                    required=True
-                ),
-                "metadata": ParamAttrs(
-                    type="object",
-                    description="Additional metadata about the processing",
+                    description="Processing mode: basic, advanced, or expert",
                     required=False
-                )
+                ),
             }
         })

flowllm/flow/tool_op_flow.py

@@ -0,0 +1,97 @@
+from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.op.llm import SimpleLLMOp, ReactLLMOp, StreamLLMOp
+from flowllm.op.search import DashscopeSearchOp
+from flowllm.op.search import TavilySearchOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class TavilySearchToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = False,
+                 service_type: str = "http+mcp",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        return TavilySearchOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return self.flow_op.tool_call
+
+    def after_flow(self, context):
+        context.response.answer = context.tavily_search_result
+
+
+@C.register_tool_flow()
+class DashscopeSearchToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = False,
+                 service_type: str = "http+mcp",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        return DashscopeSearchOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return self.flow_op.tool_call
+
+    def after_flow(self, context):
+        context.response.answer = context.dashscope_search_result
+
+
+@C.register_tool_flow(name="stream_llm_tool_flow")
+class StreamLLMToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = True,
+                 service_type: str = "http",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        return StreamLLMOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return self.flow_op.tool_call
+
+
+@C.register_tool_flow(name="simple_llm_tool_flow")
+class SimpleLLMToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = False,
+                 service_type: str = "http",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        return SimpleLLMOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return self.flow_op.tool_call
+
+
+@C.register_tool_flow(name="react_llm_tool_flow")
+class ReactLLMToolFlow(BaseToolFlow):
+
+    def __init__(self,
+                 use_async: bool = True,
+                 stream: bool = False,
+                 service_type: str = "http",
+                 **kwargs):
+        super().__init__(use_async=use_async, stream=stream, service_type=service_type, **kwargs)
+
+    def build_flow(self):
+        return ReactLLMOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return self.flow_op.tool_call

flowllm/llm/base_llm.py

@@ -155,10 +155,8 @@ class BaseLLM(BaseModel, ABC):
             except Exception as e:
                 logger.exception(f"chat with model={self.model_name} encounter error with e={e.args}")
 
-                # Exponential backoff: wait longer after each failure
                 time.sleep(1 + i)
 
-                # Handle final retry failure
                 if i == self.max_retries - 1:
                     if self.raise_exception:
                         raise e

flowllm/llm/litellm_llm.py

@@ -2,7 +2,6 @@ import asyncio
 import os
 from typing import List, Dict
 
-from litellm import completion, acompletion
 from loguru import logger
 from pydantic import Field, PrivateAttr, model_validator
 
@@ -98,6 +97,7 @@ class LiteLLMBaseLLM(BaseLLM):
         Yields:
             Tuple of (chunk_content, ChunkEnum) for each streaming piece
         """
+        from litellm import completion
         for i in range(self.max_retries):
             try:
                 # Prepare parameters for LiteLLM
@@ -200,6 +200,7 @@ class LiteLLMBaseLLM(BaseLLM):
         Yields:
             Tuple of (chunk_content, ChunkEnum) for each streaming piece
         """
+        from litellm import acompletion
         for i in range(self.max_retries):
             try:
                 # Prepare parameters for LiteLLM

flowllm/op/__init__.py

@@ -1,11 +1,11 @@
 from .base_llm_op import BaseLLMOp
 from .base_op import BaseOp
-from .base_ray_op import BaseRayOp
+from .base_tool_op import BaseToolOp
 
 """
-op folder
 """
 from . import akshare
-from . import code
 from . import gallery
+from . import llm
 from . import search
+from . import mcp

flowllm/op/akshare/get_ak_a_code_op.py

@@ -101,7 +101,7 @@ class GetAkACodeOp(BaseLLMOp):
 
 
 if __name__ == "__main__":
-    C.set_default_service_config().init_by_service_config()
+    C.set_service_config().init_by_service_config()
     context = FlowContext(query="茅台和五粮现在价格多少？")
 
     op = GetAkACodeOp()

flowllm/op/akshare/get_ak_a_info_op.py

@@ -124,7 +124,7 @@ class MergeAkAInfoOp(BaseOp):
 
 
 if __name__ == "__main__":
-    C.set_default_service_config().init_by_service_config()
+    C.set_service_config().init_by_service_config()
 
     code_infos = {"000858": {}, "600519": {}}
     context = FlowContext(code_infos=code_infos, query="茅台和五粮现在价格多少？")

flowllm/op/base_llm_op.py

@@ -23,8 +23,9 @@ class BaseLLMOp(BaseOp, ABC):
         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")
+        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

flowllm/op/base_op.py

@@ -1,6 +1,8 @@
-from abc import abstractmethod, ABC
+import asyncio
+import copy
+from abc import ABC
 from concurrent.futures import Future
-from typing import List
+from typing import List, Any, Optional, Callable
 
 from loguru import logger
 from tqdm import tqdm
@@ -13,54 +15,153 @@ from flowllm.utils.timer import Timer
 
 class BaseOp(ABC):
 
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+        instance._init_args = copy.copy(args)
+        instance._init_kwargs = copy.copy(kwargs)
+        return instance
+
     def __init__(self,
                  name: str = "",
+                 max_retries: int = 1,
                  raise_exception: bool = True,
+                 enable_multithread: bool = True,
                  **kwargs):
         super().__init__()
 
         self.name: str = name or camel_to_snake(self.__class__.__name__)
+        self.max_retries: int = max_retries
         self.raise_exception: bool = raise_exception
+        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.async_task_list: List = []
+        self.ray_task_list: List = []
         self.timer = Timer(name=self.name)
         self.context: FlowContext | None = None
+        self.sub_op: Optional["BaseOp"] = None
+
+    def before_execute(self):
+        ...
+
+    def after_execute(self):
+        ...
 
-    @abstractmethod
     def execute(self):
         ...
 
+    def default_execute(self):
+        ...
+
+    async def async_execute(self):
+        ...
+
     def __call__(self, context: FlowContext = None):
         self.context = context
         with self.timer:
-            if self.raise_exception:
+            if self.max_retries == 1 and self.raise_exception:
+                self.before_execute()
                 self.execute()
+                self.after_execute()
 
             else:
+                for i in range(self.max_retries):
+                    try:
+                        self.before_execute()
+                        self.execute()
+                        self.after_execute()
+
+                    except Exception as e:
+                        logger.exception(f"op={self.name} execute failed, error={e.args}")
+
+                        if i == self.max_retries - 1:
+                            if self.raise_exception:
+                                raise e
+                            else:
+                                self.default_execute()
+
+        if self.context is not None and self.context.response is not None:
+            return self.context.response
+        return None
+
+    async def async_call(self, context: FlowContext = None) -> Any:
+        self.context = context
+        with self.timer:
+            if self.max_retries == 1 and self.raise_exception:
+                self.before_execute()
+                await self.async_execute()
+                self.after_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)
-        return self
+            else:
+                for i in range(self.max_retries):
+                    try:
+                        self.before_execute()
+                        await self.async_execute()
+                        self.after_execute()
+
+                    except Exception as e:
+                        logger.exception(f"op={self.name} async execute failed, error={e.args}")
+
+                        if i == self.max_retries:
+                            if self.raise_exception:
+                                raise e
+                            else:
+                                self.default_execute()
+
+        if self.context is not None and self.context.response is not None:
+            return self.context.response
+        return None
+
+    def submit_async_task(self, fn: Callable, *args, **kwargs):
+        if asyncio.iscoroutinefunction(fn):
+            task = asyncio.create_task(fn(*args, **kwargs))
+            self.async_task_list.append(task)
+        else:
+            logger.warning("submit_async_task failed, fn is not a coroutine function!")
 
-    def join_task(self, task_desc: str = None) -> list:
+    async def join_async_task(self):
         result = []
-        for task in tqdm(self.task_list, desc=task_desc or self.name):
-            t_result = task.result()
+        for task in self.async_task_list:
+            t_result = await task
             if t_result:
                 if isinstance(t_result, list):
                     result.extend(t_result)
                 else:
                     result.append(t_result)
+
+        self.async_task_list.clear()
+        return result
+
+    def submit_task(self, fn, *args, **kwargs):
+        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 = []
+        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
 
@@ -75,6 +176,9 @@ class BaseOp(ABC):
             sequential_op.ops.append(op)
         return sequential_op
 
+    def __lshift__(self, op: "BaseOp"):
+        self.sub_op = op
+
     def __or__(self, op: "BaseOp"):
         from flowllm.op.parallel_op import ParallelOp
 
@@ -87,6 +191,9 @@ class BaseOp(ABC):
 
         return parallel_op
 
+    def copy(self) -> "BaseOp":
+        return self.__class__(*self._init_args, **self._init_kwargs)
+
 
 def run1():
     """Basic test"""
@@ -95,6 +202,9 @@ def run1():
         def execute(self):
             logger.info(f"op={self.name} execute")
 
+        async def async_execute(self):
+            logger.info(f"op={self.name} async_execute")
+
     mock_op = MockOp()
     mock_op()
 
@@ -112,6 +222,12 @@ def run2():
             logger.info(f"Executing {op_result}")
             return op_result
 
+        async def async_execute(self):
+            await asyncio.sleep(0.1)
+            op_result = f"{self.name}"
+            logger.info(f"Async executing {op_result}")
+            return op_result
+
     # Create service_context for parallel execution
     C["thread_pool"] = ThreadPoolExecutor(max_workers=4)
 
@@ -131,18 +247,135 @@ def run2():
     result = parallel()
     logger.info(f"Parallel result: {result}")
 
-    logger.info("=== Testing mixed calls op1 >> (op2 | op3) >> op4 ===")
+
+async def async_run1():
+    """Basic async test"""
+
+    class MockOp(BaseOp):
+        def execute(self):
+            logger.info(f"op={self.name} execute")
+
+        async def async_execute(self):
+            logger.info(f"op={self.name} async_execute")
+
+    mock_op = MockOp()
+    await mock_op.async_call()
+
+
+async def async_run2():
+    """Test async 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
+
+        async def async_execute(self):
+            await asyncio.sleep(0.1)
+            op_result = f"{self.name}"
+            logger.info(f"Async 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 async sequential execution op1 >> op2 ===")
+    sequential = op1 >> op2
+    result = await sequential.async_call()
+    logger.info(f"Async sequential result: {result}")
+
+    logger.info("=== Testing async parallel execution op1 | op2 ===")
+    parallel = op1 | op2
+    result = await parallel.async_call()
+    logger.info(f"Async parallel result: {result}")
+
+    logger.info("=== Testing async mixed calls op1 >> (op2 | op3) >> op4 ===")
     mixed = op1 >> (op2 | op3) >> op4
-    result = mixed()
-    logger.info(f"Mixed result: {result}")
+    result = await mixed.async_call()
+    logger.info(f"Async mixed result: {result}")
 
-    logger.info("=== Testing complex mixed calls op1 >> (op1 | (op2 >> op3)) >> op4 ===")
+    logger.info("=== Testing async 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}")
+    result = await complex_mixed.async_call()
+    logger.info(f"Async complex mixed result: {result}")
+
+
+def test_copy():
+    class TestOp(BaseOp):
+        def __init__(self, name="", custom_param="default", **kwargs):
+            super().__init__(name=name, **kwargs)
+            self.custom_param = custom_param
+
+        def execute(self):
+            logger.info(f"TestOp {self.name} executing with custom_param={self.custom_param}")
+
+    class AdvancedOp(TestOp):
+        def __init__(self, name="", custom_param="default", advanced_param=42, **kwargs):
+            super().__init__(name=name, custom_param=custom_param, **kwargs)
+            self.advanced_param = advanced_param
+
+        def execute(self):
+            logger.info(
+                f"AdvancedOp {self.name} executing with custom_param={self.custom_param}, advanced_param={self.advanced_param}")
+
+    logger.info("=== Testing copy functionality ===")
+
+    original_op = TestOp(name="test_op", custom_param="custom_value", max_retries=3, enable_multithread=False)
+    copied_op = original_op.copy()
+
+    logger.info(
+        f"Original: name={original_op.name}, custom_param={original_op.custom_param}, max_retries={original_op.max_retries}")
+    logger.info(
+        f"Copied: name={copied_op.name}, custom_param={copied_op.custom_param}, max_retries={copied_op.max_retries}")
+    logger.info(f"Same object? {original_op is copied_op}")
+    logger.info(f"Same class? {type(original_op) == type(copied_op)}")
+
+    original_advanced = AdvancedOp(
+        name="advanced_op",
+        custom_param="advanced_custom",
+        advanced_param=100,
+        max_retries=5,
+        raise_exception=False
+    )
+    copied_advanced = original_advanced.copy()
+
+    logger.info(
+        f"Advanced Original: name={original_advanced.name}, custom_param={original_advanced.custom_param}, advanced_param={original_advanced.advanced_param}")
+    logger.info(
+        f"Advanced Copied: name={copied_advanced.name}, custom_param={copied_advanced.custom_param}, advanced_param={copied_advanced.advanced_param}")
+    logger.info(f"Advanced Same object? {original_advanced is copied_advanced}")
+    logger.info(f"Advanced Same class? {type(original_advanced) == type(copied_advanced)}")
+
+    copied_op.name = "modified_name"
+    logger.info(f"After modifying copy - Original name: {original_op.name}, Copied name: {copied_op.name}")
 
 
 if __name__ == "__main__":
     run1()
     print("\n" + "=" * 50 + "\n")
     run2()
+    print("\n" + "=" * 50 + "\n")
+    test_copy()
+
+    print("\n" + "=" * 50 + "\n")
+    print("Running async tests:")
+
+
+    async def main():
+        await async_run1()
+        print("\n" + "=" * 50 + "\n")
+        await async_run2()
+
+
+    asyncio.run(main())