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

flowllm/context/flow_context.py

@@ -2,27 +2,15 @@ import uuid
 
 from flowllm.context.base_context import BaseContext
 from flowllm.schema.flow_response import FlowResponse
-from flowllm.schema.service_config import ServiceConfig
 
 
 class FlowContext(BaseContext):
 
-    def __init__(self, flow_id: str = uuid.uuid4().hex, **kwargs):
+    def __init__(self,
+                 flow_id: str = uuid.uuid4().hex,
+                 response: FlowResponse = None,
+                 **kwargs):
         super().__init__(**kwargs)
-        self.flow_id: str = flow_id
-
-    @property
-    def service_config(self) -> ServiceConfig:
-        return self._data.get("service_config")
-
-    @service_config.setter
-    def service_config(self, service_config: ServiceConfig):
-        self._data["service_config"] = service_config
 
-    @property
-    def response(self) -> FlowResponse:
-        return self._data.get("response")
-
-    @response.setter
-    def response(self, response: FlowResponse):
-        self._data["response"] = response
+        self.flow_id: str = flow_id
+        self.response: FlowResponse = FlowResponse() if response is None else response

flowllm/context/registry.py

@@ -6,13 +6,17 @@ from flowllm.utils.common_utils import camel_to_snake
 
 class Registry(BaseContext):
 
-    def __init__(self, registry_name: str, enable_log: bool = True, **kwargs):
+    def __init__(self, registry_name: str, enable_log: bool = True, register_flow_module: bool = True, **kwargs):
         super().__init__(**kwargs)
         self.registry_name: str = registry_name
         self.enable_log: bool = enable_log
+        self.register_flow_module: bool = register_flow_module
 
     def register(self, name: str = ""):
         def decorator(cls):
+            if not self.register_flow_module and cls.__module__.startswith("flowllm"):
+                return cls
+
             class_name = name if name else camel_to_snake(cls.__name__)
             if self.enable_log:
                 if class_name in self._data:

flowllm/context/service_context.py

@@ -1,9 +1,15 @@
+import os
 import uuid
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict
+from inspect import isclass
+from typing import Dict, List
+
+import ray
+from loguru import logger
 
 from flowllm.context.base_context import BaseContext
 from flowllm.context.registry import Registry
+from flowllm.schema.service_config import ServiceConfig, EmbeddingModelConfig
 from flowllm.utils.singleton import singleton
 
 
@@ -12,41 +18,80 @@ class ServiceContext(BaseContext):
 
     def __init__(self, service_id: str = uuid.uuid4().hex, **kwargs):
         super().__init__(**kwargs)
-        self.service_id: str = service_id
-        self.registry_dict: Dict[str, Registry] = \
-            {k: Registry(k) for k in ["embedding_model", "llm", "vector_store", "op", "flow_engine", "service"]}
-
-    @property
-    def language(self) -> str:
-        return self._data.get("language", "")
-
-    @language.setter
-    def language(self, value: str):
-        self._data["language"] = value
-
-    @property
-    def thread_pool(self) -> ThreadPoolExecutor:
-        return self._data["thread_pool"]
 
-    @thread_pool.setter
-    def thread_pool(self, thread_pool: ThreadPoolExecutor):
-        self._data["thread_pool"] = thread_pool
+        self.service_id: str = service_id
+        self.service_config: ServiceConfig | None = None
+        self.language: str = ""
+        self.thread_pool: ThreadPoolExecutor | None = None
+        self.vector_store_dict: dict = {}
+
+        self.registry_dict: Dict[str, Registry] = {}
+        use_framework: bool = os.environ.get("FLOW_USE_FRAMEWORK", "").lower() == "true"
+        for key in ["embedding_model", "llm", "vector_store", "op", "tool_flow", "service"]:
+            enable_log = True
+            register_flow_module = True
+
+            if use_framework:
+                enable_log = False
+                if key in ["op", "tool_flow"]:
+                    register_flow_module = False
+            self.registry_dict[key] = Registry(key, enable_log=enable_log, register_flow_module=register_flow_module)
+
+        self.tool_flow_dict: dict = {}
+
+    def set_default_service_config(self):
+        from flowllm.config.pydantic_config_parser import PydanticConfigParser
+
+        config_parser = PydanticConfigParser(ServiceConfig)
+        self.service_config = config_parser.parse_args("config=default")
+        return self
+
+    def init_by_service_config(self, service_config: ServiceConfig = None):
+        if service_config:
+            self.service_config = service_config
+
+        self.language = self.service_config.language
+        self.thread_pool = ThreadPoolExecutor(max_workers=self.service_config.thread_pool_max_workers)
+        if self.service_config.ray_max_workers > 1:
+            ray.init(num_cpus=self.service_config.ray_max_workers)
+
+        # add vector store
+        for name, config in self.service_config.vector_store.items():
+            vector_store_cls = self.resolve_vector_store(config.backend)
+            embedding_model_config: EmbeddingModelConfig = self.service_config.embedding_model[config.embedding_model]
+            embedding_model_cls = self.resolve_embedding_model(embedding_model_config.backend)
+            embedding_model = embedding_model_cls(model_name=embedding_model_config.model_name,
+                                                  **embedding_model_config.params)
+            self.vector_store_dict[name] = vector_store_cls(embedding_model=embedding_model, **config.params)
+
+        from flowllm.flow.base_tool_flow import BaseToolFlow
+        from flowllm.flow.gallery import ExpressionToolFlow
+
+        # add tool flow cls
+        for name, tool_flow_cls in self.registry_dict["tool_flow"].items():
+            if not isclass(tool_flow_cls):
+                continue
+
+            tool_flow: BaseToolFlow = tool_flow_cls()
+            self.tool_flow_dict[tool_flow.name] = tool_flow
+            logger.info(f"add diy tool_flow: {tool_flow.name}")
+
+        # add tool flow config
+        for name, flow_config in self.service_config.flow.items():
+            flow_config.name = name
+            tool_flow: BaseToolFlow = ExpressionToolFlow(flow_config=flow_config)
+            self.tool_flow_dict[tool_flow.name] = tool_flow
+            logger.info(f"add expression tool_flow:{tool_flow.name}")
 
     def get_vector_store(self, name: str = "default"):
-        vector_store_dict: dict = self._data["vector_store_dict"]
-        if name not in vector_store_dict:
-            raise KeyError(f"vector store {name} not found")
-
-        return vector_store_dict[name]
+        return self.vector_store_dict[name]
 
-    def set_vector_store(self, name: str, vector_store):
-        if "vector_store_dict" not in self._data:
-            self.set_vector_stores({})
+    def get_tool_flow(self, name: str = "default"):
+        return self.tool_flow_dict[name]
 
-        self._data["vector_store_dict"][name] = vector_store
-
-    def set_vector_stores(self, vector_store_dict: dict):
-        self._data["vector_store_dict"] = vector_store_dict
+    @property
+    def tool_flow_names(self) -> List[str]:
+        return sorted(self.tool_flow_dict.keys())
 
     """
     register models
@@ -64,8 +109,8 @@ class ServiceContext(BaseContext):
     def register_op(self, name: str = ""):
         return self.registry_dict["op"].register(name=name)
 
-    def register_flow_engine(self, name: str = ""):
-        return self.registry_dict["flow_engine"].register(name=name)
+    def register_tool_flow(self, name: str = ""):
+        return self.registry_dict["tool_flow"].register(name=name)
 
     def register_service(self, name: str = ""):
         return self.registry_dict["service"].register(name=name)
@@ -90,14 +135,13 @@ class ServiceContext(BaseContext):
         assert name in self.registry_dict["op"], f"op={name} not found!"
         return self.registry_dict["op"][name]
 
-    def resolve_flow_engine(self, name: str):
-        assert name in self.registry_dict["flow_engine"], f"flow={name} not found!"
-        return self.registry_dict["flow_engine"][name]
+    def resolve_tool_flow(self, name: str):
+        assert name in self.registry_dict["tool_flow"], f"tool_flow={name} not found!"
+        return self.registry_dict["tool_flow"][name]
 
     def resolve_service(self, name: str):
         assert name in self.registry_dict["service"], f"service={name} not found!"
         return self.registry_dict["service"][name]
 
 
-
 C = ServiceContext()

flowllm/embedding_model/__init__.py

@@ -1 +1 @@
-from flowllm.embedding_model.openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel
+from .openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel

flowllm/embedding_model/base_embedding_model.py

@@ -1,3 +1,4 @@
+import asyncio
 from abc import ABC
 from typing import List
 
@@ -37,6 +38,21 @@ class BaseEmbeddingModel(BaseModel, ABC):
         """
         raise NotImplementedError
 
+    async def _get_embeddings_async(self, input_text: str | List[str]):
+        """
+        Abstract async method to get embeddings from the model.
+        
+        This method must be implemented by concrete subclasses to provide
+        the actual async embedding functionality.
+        
+        Args:
+            input_text: Single text string or list of text strings to embed
+            
+        Returns:
+            Embedding vector(s) corresponding to the input text(s)
+        """
+        raise NotImplementedError
+
     def get_embeddings(self, input_text: str | List[str]):
         """
         Get embeddings with retry logic and error handling.
@@ -64,6 +80,33 @@ class BaseEmbeddingModel(BaseModel, ABC):
         # Return None if all retries failed and raise_exception is False
         return None
 
+    async def get_embeddings_async(self, input_text: str | List[str]):
+        """
+        Get embeddings asynchronously with retry logic and error handling.
+        
+        This method wraps the _get_embeddings_async method with automatic retry
+        functionality in case of failures.
+        
+        Args:
+            input_text: Single text string or list of text strings to embed
+            
+        Returns:
+            Embedding vector(s) or None if all retries failed and raise_exception is False
+        """
+        # Retry loop with exponential backoff potential
+        for i in range(self.max_retries):
+            try:
+                return await self._get_embeddings_async(input_text)
+
+            except Exception as e:
+                logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}")
+                # If this is the last retry and raise_exception is True, re-raise the exception
+                if i == self.max_retries - 1 and self.raise_exception:
+                    raise e
+
+        # Return None if all retries failed and raise_exception is False
+        return None
+
     def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]):
         """
         Generate embeddings for VectorNode objects and update their vector fields.
@@ -102,3 +145,51 @@ class BaseEmbeddingModel(BaseModel, ABC):
 
         else:
             raise TypeError(f"unsupported type={type(nodes)}")
+
+    async def get_node_embeddings_async(self, nodes: VectorNode | List[VectorNode]):
+        """
+        Generate embeddings asynchronously for VectorNode objects and update their vector fields.
+        
+        This method handles both single nodes and lists of nodes, with automatic
+        batching for efficient processing of large node lists.
+        
+        Args:
+            nodes: Single VectorNode or list of VectorNode objects to embed
+            
+        Returns:
+            The same node(s) with updated vector fields containing embeddings
+            
+        Raises:
+            RuntimeError: If unsupported node type is provided
+        """
+        # Handle single VectorNode
+        if isinstance(nodes, VectorNode):
+            nodes.vector = await self.get_embeddings_async(nodes.content)
+            return nodes
+
+        # Handle list of VectorNodes with batch processing
+        elif isinstance(nodes, list):
+            # Process nodes in batches to respect max_batch_size limits
+            batch_tasks = []
+            for i in range(0, len(nodes), self.max_batch_size):
+                batch_nodes = nodes[i:i + self.max_batch_size]
+                batch_content = [node.content for node in batch_nodes]
+                batch_tasks.append(self.get_embeddings_async(batch_content))
+
+            # Execute all batch tasks concurrently
+            batch_results = await asyncio.gather(*batch_tasks)
+
+            # Flatten the results
+            embeddings = [emb for batch_result in batch_results for emb in batch_result]
+
+            # Validate that we got the expected number of embeddings
+            if len(embeddings) != len(nodes):
+                logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}")
+            else:
+                # Assign embeddings to corresponding nodes
+                for node, embedding in zip(nodes, embeddings):
+                    node.vector = embedding
+            return nodes
+
+        else:
+            raise TypeError(f"unsupported type={type(nodes)}")

flowllm/embedding_model/openai_compatible_embedding_model.py

@@ -1,7 +1,7 @@
 import os
 from typing import Literal, List
 
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI
 from pydantic import Field, PrivateAttr, model_validator
 
 from flowllm.context.service_context import C
@@ -26,21 +26,23 @@ class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
     dimensions: int = Field(default=1024, description="Dimensionality of the embedding vectors")
     encoding_format: Literal["float", "base64"] = Field(default="float", description="Encoding format for embeddings")
 
-    # Private OpenAI client instance
+    # Private OpenAI client instances
     _client: OpenAI = PrivateAttr()
+    _async_client: AsyncOpenAI = PrivateAttr()
 
     @model_validator(mode="after")
     def init_client(self):
         """
-        Initialize the OpenAI client after model validation.
+        Initialize the OpenAI clients after model validation.
         
         This method is called automatically after Pydantic model validation
-        to set up the OpenAI client with the provided API key and base URL.
+        to set up both sync and async OpenAI clients with the provided API key and base URL.
         
         Returns:
             self: The model instance for method chaining
         """
         self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+        self._async_client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
         return self
 
     def _get_embeddings(self, input_text: str | List[str]):
@@ -78,6 +80,41 @@ class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
         else:
             raise RuntimeError(f"unsupported type={type(input_text)}")
 
+    async def _get_embeddings_async(self, input_text: str | List[str]):
+        """
+        Get embeddings asynchronously from the OpenAI-compatible API.
+        
+        This method implements the abstract _get_embeddings_async method from BaseEmbeddingModel
+        by calling the OpenAI-compatible embeddings API asynchronously.
+        
+        Args:
+            input_text: Single text string or list of text strings to embed
+            
+        Returns:
+            Embedding vector(s) corresponding to the input text(s)
+            
+        Raises:
+            RuntimeError: If unsupported input type is provided
+        """
+        completion = await self._async_client.embeddings.create(
+            model=self.model_name,
+            input=input_text,
+            dimensions=self.dimensions,
+            encoding_format=self.encoding_format
+        )
+
+        if isinstance(input_text, str):
+            return completion.data[0].embedding
+
+        elif isinstance(input_text, list):
+            result_emb = [[] for _ in range(len(input_text))]
+            for emb in completion.data:
+                result_emb[emb.index] = emb.embedding
+            return result_emb
+
+        else:
+            raise RuntimeError(f"unsupported type={type(input_text)}")
+
 
 def main():
     from flowllm.utils.common_utils import load_env
@@ -91,5 +128,26 @@ def main():
     print(res2)
 
 
+async def async_main():
+    from flowllm.utils.common_utils import load_env
+
+    load_env()
+    model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
+
+    # Test async single text embedding
+    res1 = await model.get_embeddings_async(
+        "The clothes are of good quality and look good, definitely worth the wait. I love them.")
+
+    # Test async batch text embedding
+    res2 = await model.get_embeddings_async(["aa", "bb"])
+
+    print("Async results:")
+    print(res1)
+    print(res2)
+
+
 if __name__ == "__main__":
-    main()
+    # main()
+    import asyncio
+
+    asyncio.run(async_main())

flowllm/flow/__init__.py

@@ -0,0 +1 @@
+from . import gallery

flowllm/flow/base_flow.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.op.base_op import BaseOp
+from flowllm.op.parallel_op import ParallelOp
+from flowllm.op.sequential_op import SequentialOp
+from flowllm.schema.flow_response import FlowResponse
+from flowllm.utils.common_utils import camel_to_snake
+
+
+class BaseFlow(ABC):
+
+    def __init__(self, name: str = "", **kwargs):
+        self.name: str = name or camel_to_snake(self.__class__.__name__)
+        self.flow_params: dict = kwargs
+
+        self.flow_op: Optional[BaseOp] = self.build_flow()
+        self.print_flow()
+
+    @abstractmethod
+    def build_flow(self):
+        ...
+
+    def print_flow(self):
+        assert self.flow_op is not None, "flow_content is not parsed!"
+        logger.info(f"---------- start print flow={self.name} ----------")
+        self._print_operation_tree(self.flow_op, indent=0)
+        logger.info(f"---------- end print flow={self.name} ----------")
+
+    def _print_operation_tree(self, op: BaseOp, indent: int):
+        """
+        Recursively print the operation tree structure.
+
+        Args:
+            op: The operation to print
+            indent: Current indentation level
+        """
+        prefix = "  " * indent
+        if isinstance(op, SequentialOp):
+            logger.info(f"{prefix}Sequential Execution:")
+            for i, sub_op in enumerate(op.ops):
+                logger.info(f"{prefix}  Step {i + 1}:")
+                self._print_operation_tree(sub_op, indent + 2)
+
+        elif isinstance(op, ParallelOp):
+            logger.info(f"{prefix}Parallel Execution:")
+            for i, sub_op in enumerate(op.ops):
+                logger.info(f"{prefix}  Branch {i + 1}:")
+                self._print_operation_tree(sub_op, indent + 2)
+
+        else:
+            logger.info(f"{prefix}Operation: {op.name}")
+
+    def return_callback(self, context: FlowContext):
+        return context.response
+
+    def __call__(self, **kwargs) -> FlowResponse:
+        context = FlowContext(**kwargs)
+        logger.info(f"request.params={kwargs}")
+
+        try:
+            self.flow_op(context=context)
+
+        except Exception as e:
+            logger.exception(f"flow_name={self.name} encounter error={e.args}")
+            context.response.success = False
+            context.response.answer = str(e.args)
+
+        return self.return_callback(context=context)

flowllm/flow/base_tool_flow.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+
+from flowllm.flow.base_flow import BaseFlow
+from flowllm.schema.tool_call import ToolCall
+
+
+class BaseToolFlow(BaseFlow, ABC):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.tool_call: ToolCall = self.build_tool_call()
+
+    @abstractmethod
+    def build_tool_call(self) -> ToolCall:
+        ...

flowllm/flow/gallery/__init__.py

@@ -0,0 +1,8 @@
+from .cmd_flow import CmdFlow
+from .code_tool_flow import CodeToolFlow
+from .dashscope_search_tool_flow import DashscopeSearchToolFlow
+from .deepsearch_tool_flow import DeepSearchToolFlow
+from .expression_tool_flow import ExpressionToolFlow
+from .mock_tool_flow import MockToolFlow
+from .tavily_search_tool_flow import TavilySearchToolFlow
+from .terminate_tool_flow import TerminateToolFlow

flowllm/flow/gallery/cmd_flow.py

@@ -0,0 +1,11 @@
+from flowllm.flow.base_flow import BaseFlow
+from flowllm.flow.parser.expression_parser import ExpressionParser
+
+
+class CmdFlow(BaseFlow):
+
+    def build_flow(self):
+        flow: str = self.flow_params["flow"]
+        assert flow, "add `flow=<op_flow>` in cmd!"
+        parser = ExpressionParser(flow)
+        return parser.parse_flow()

flowllm/flow/gallery/code_tool_flow.py

@@ -0,0 +1,30 @@
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.op.code.execute_code_op import ExecuteCodeOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class CodeToolFlow(BaseToolFlow):
+
+    def build_flow(self):
+        return ExecuteCodeOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "python_execute",
+            "description": "Execute python code can be used in scenarios such as analysis or calculation, and the final result can be printed using the `print` function.",
+            "input_schema": {
+                "code": {
+                    "type": "str",
+                    "description": "code to be executed. Please do not execute any matplotlib code here.",
+                    "required": True
+                }
+            }
+        })
+
+    def return_callback(self, context: FlowContext):
+        context.response.answer = context.code_result
+        return context.response
+

flowllm/flow/gallery/dashscope_search_tool_flow.py

@@ -0,0 +1,34 @@
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.op.search import DashscopeSearchOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class DashscopeSearchToolFlow(BaseToolFlow):
+
+    def build_flow(self):
+        return DashscopeSearchOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "web_search",
+            "description": "Use search keywords to retrieve relevant information from the internet. If there are multiple search keywords, please use each keyword separately to call this tool.",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "search keyword",
+                    "required": True
+                }
+            }
+        })
+
+    def return_callback(self, context: FlowContext):
+        context.response.answer = context.dashscope_search_result
+        return context.response
+
+
+if __name__ == "__main__":
+    flow = DashscopeSearchToolFlow()
+    flow(query="what is AI?")

flowllm/flow/gallery/deepsearch_tool_flow.py

@@ -0,0 +1,39 @@
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.op.search.dashscope_deep_research_op import DashscopeDeepResearchOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class DeepSearchToolFlow(BaseToolFlow):
+
+    def build_flow(self):
+        return DashscopeDeepResearchOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "deep_search",
+            "description": "Perform deep research on a topic using Dashscope's qwen-deep-research model. This tool will conduct multi-phase research including model questioning, web research, and result generation.",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "Research topic or question",
+                    "required": True
+                }
+            }
+        })
+
+    def return_callback(self, context: FlowContext):
+        context.response.answer = context.dashscope_deep_research_result
+        return context.response
+
+
+if __name__ == "__main__":
+    from flowllm.utils.common_utils import load_env
+
+    load_env()
+
+    flow = DeepSearchToolFlow()
+    result = flow(query="中国电解铝行业值得投资吗，有哪些值得投资的标的，各个标的之间需要对比优劣势")
+    print(result.answer)

flowllm/flow/gallery/expression_tool_flow.py

@@ -0,0 +1,18 @@
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.flow.parser.expression_parser import ExpressionParser
+from flowllm.schema.service_config import FlowConfig
+from flowllm.schema.tool_call import ToolCall
+
+
+class ExpressionToolFlow(BaseToolFlow):
+
+    def __init__(self, flow_config: FlowConfig = None, **kwargs):
+        self.flow_config: FlowConfig = flow_config
+        super().__init__(name=flow_config.name, **kwargs)
+
+    def build_flow(self):
+        parser = ExpressionParser(self.flow_config.flow_content)
+        return parser.parse_flow()
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**self.flow_config.model_dump())