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

{llmflow → flowllm}/embedding_model/base_embedding_model.py

@@ -1,10 +1,11 @@
+import asyncio
 from abc import ABC
 from typing import List
 
 from loguru import logger
 from pydantic import BaseModel, Field
 
-from llmflow.schema.vector_node import VectorNode
+from flowllm.schema.vector_node import VectorNode
 
 
 class BaseEmbeddingModel(BaseModel, ABC):
@@ -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.
@@ -101,4 +144,52 @@ class BaseEmbeddingModel(BaseModel, ABC):
             return nodes
 
         else:
-            raise RuntimeError(f"unsupported type={type(nodes)}")
+            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)}")

{llmflow → flowllm}/embedding_model/openai_compatible_embedding_model.py

@@ -1,15 +1,14 @@
 import os
 from typing import Literal, List
 
-from dotenv import load_dotenv
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI
 from pydantic import Field, PrivateAttr, model_validator
 
-from llmflow.embedding_model import EMBEDDING_MODEL_REGISTRY
-from llmflow.embedding_model.base_embedding_model import BaseEmbeddingModel
+from flowllm.context.service_context import C
+from flowllm.embedding_model.base_embedding_model import BaseEmbeddingModel
 
 
-@EMBEDDING_MODEL_REGISTRY.register("openai_compatible")
+@C.register_embedding_model("openai_compatible")
 class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
     """
     OpenAI-compatible embedding model implementation.
@@ -19,29 +18,31 @@ class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel):
     other services that follow the same interface.
     """
     # API configuration fields
-    api_key: str = Field(default_factory=lambda: os.getenv("EMBEDDING_API_KEY"),
+    api_key: str = Field(default_factory=lambda: os.getenv("FLOW_EMBEDDING_API_KEY"),
                          description="API key for authentication")
-    base_url: str = Field(default_factory=lambda: os.getenv("EMBEDDING_BASE_URL"),
+    base_url: str = Field(default_factory=lambda: os.getenv("FLOW_EMBEDDING_BASE_URL"),
                           description="Base URL for the API endpoint")
     model_name: str = Field(default="", description="Name of the embedding model to use")
     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]):
@@ -79,9 +80,46 @@ 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():
-    load_dotenv()
+    from flowllm.utils.common_utils import load_env
+
+    load_env()
     model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4")
     res1 = model.get_embeddings(
         "The clothes are of good quality and look good, definitely worth the wait. I love them.")
@@ -90,6 +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()
-    # launch with: python -m llmflow.model.openai_compatible_embedding_model
+    # 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())

flowllm/flow/gallery/mock_tool_flow.py

@@ -0,0 +1,67 @@
+from flowllm.context.service_context import C
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.op.gallery import Mock1Op, Mock2Op, Mock3Op, Mock4Op, Mock5Op, Mock6Op
+from flowllm.schema.tool_call import ToolCall, ParamAttrs
+
+
+@C.register_tool_flow()
+class MockToolFlow(BaseToolFlow):
+
+    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)
+        return op
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "index": 0,
+            "id": "call_mock_tool_12345",
+            "type": "function",
+            "name": "mock_data_processor",
+            "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",
+                    description="The input data to be processed",
+                    required=True
+                ),
+                "processing_mode": 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",
+                    required=True
+                ),
+                "status": ParamAttrs(
+                    type="string",
+                    description="Processing status: success, warning, or error",
+                    required=True
+                ),
+                "metadata": ParamAttrs(
+                    type="object",
+                    description="Additional metadata about the processing",
+                    required=False
+                )
+            }
+        })

flowllm/flow/gallery/tavily_search_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.search import TavilySearchOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class TavilySearchToolFlow(BaseToolFlow):
+
+    def build_flow(self):
+        return TavilySearchOp()
+
+    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.tavily_search_result
+        return context.response
+

flowllm/flow/gallery/terminate_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.gallery.terminate_op import TerminateOp
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_tool_flow()
+class TerminateToolFlow(BaseToolFlow):
+
+    def build_flow(self):
+        return TerminateOp()
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "terminate",
+            "description": "If you can answer the user's question based on the context, be sure to use the **terminate** tool.",
+            "input_schema": {
+                "status": {
+                    "type": "str",
+                    "description": "Please determine whether the user's question has been completed. (success / failure)",
+                    "required": True,
+                    "enum": ["success", "failure"],
+                }
+            }
+        })
+
+    def return_callback(self, context: FlowContext):
+        context.response.answer = context.terminate_answer
+        return context.response