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

flowllm/op/search/dashscope_search_op.py

@@ -1,18 +1,17 @@
+import asyncio
 import os
-import time
 from typing import Dict, Any, List
 
 import dashscope
 from loguru import logger
 
-from flowllm.context.flow_context import FlowContext
-from flowllm.context.service_context import C
-from flowllm.op.base_llm_op import BaseLLMOp
-from flowllm.storage.cache.data_cache import DataCache
+from flowllm.context import FlowContext, C
+from flowllm.op.base_tool_op import BaseToolOp
+from flowllm.schema.tool_call import ToolCall
 
 
 @C.register_op()
-class DashscopeSearchOp(BaseLLMOp):
+class DashscopeSearchOp(BaseToolOp):
     file_path: str = __file__
 
     """
@@ -24,29 +23,31 @@ class DashscopeSearchOp(BaseLLMOp):
 
     def __init__(self,
                  model: str = "qwen-plus",
-                 enable_print: bool = True,
-                 enable_cache: bool = False,
-                 cache_path: str = "./dashscope_search_cache",
-                 cache_expire_hours: float = 0.1,
-                 max_retries: int = 3,
                  search_strategy: str = "max",
-                 return_only_content: bool = True,
                  enable_role_prompt: bool = True,
+                 save_answer: bool = False,
                  **kwargs):
         super().__init__(**kwargs)
 
-        self.model = model
-        self.enable_print = enable_print
-        self.enable_cache = enable_cache
-        self.cache_expire_hours = cache_expire_hours
-        self.max_retries = max_retries
-        self.search_strategy = search_strategy
-        self.return_only_content = return_only_content
-        self.enable_role_prompt = enable_role_prompt
-
-        # Ensure API key is available
-        self.api_key = os.getenv("FLOW_DASHSCOPE_API_KEY")
-        self.cache = DataCache(cache_path) if self.enable_cache else None
+        self.model: str = model
+        self.search_strategy: str = search_strategy
+        self.enable_role_prompt: bool = enable_role_prompt
+        self.save_answer: bool = save_answer
+
+        self.api_key = os.getenv("FLOW_DASHSCOPE_API_KEY", "")
+
+    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
+                }
+            }
+        })
 
     @staticmethod
     def format_search_results(search_results: List[Dict[str, Any]]) -> str:
@@ -58,122 +59,73 @@ class DashscopeSearchOp(BaseLLMOp):
 
         return "\n".join(formatted_results)
 
-    def post_process(self, response_data: dict) -> dict:
-        """Post-process the response and optionally print results"""
-        if self.enable_print:
-            # Print search information
-            if "search_results" in response_data:
-                search_info = self.format_search_results(response_data["search_results"])
-                logger.info(f"Search Information:\n{search_info}")
-
-            # Print response content
-            if "response_content" in response_data:
-                logger.info("=" * 20 + " Response Content " + "=" * 20)
-                logger.info(response_data["response_content"])
-
-        return response_data
+    def default_execute(self):
+        self.output_dict["dashscope_search_result"] = "dashscope search failed!"
 
-    def execute(self):
-        """Execute the Dashscope search operation"""
-        # Get query from context - support multiple parameter names
-        query = self.context.query
+    async def async_execute(self):
+        query: str = self.input_dict["query"]
 
-        # Check cache first
         if self.enable_cache and self.cache:
             cached_result = self.cache.load(query)
             if cached_result:
-                result = self.post_process(cached_result)
-                if self.return_only_content:
-                    self.context.dashscope_search_result = result["response_content"]
-                else:
-                    self.context.dashscope_search_result = result
-
+                self.output_dict["dashscope_search_result"] = cached_result["response_content"]
                 return
 
         if self.enable_role_prompt:
             user_query = self.prompt_format(prompt_name="role_prompt", query=query)
         else:
             user_query = query
+        logger.info(f"user_query={user_query}")
         messages: list = [{"role": "user", "content": user_query}]
 
-        # Retry logic for API calls
-        for attempt in range(self.max_retries):
-            try:
-                # Call Dashscope Generation API with search enabled
-                response = dashscope.Generation.call(
-                    api_key=self.api_key,
-                    model=self.model,
-                    messages=messages,
-                    enable_search=True,  # Enable web search
-                    search_options={
-                        "forced_search": True,  # Force web search
-                        "enable_source": True,  # Include search source information
-                        "enable_citation": False,  # Enable citation markers
-                        "search_strategy": self.search_strategy,  # Search strategy
-                    },
-                    result_format="message",
-                )
-
-                # Extract search results and response content
-                search_results = []
-                response_content = ""
-
-                if hasattr(response, 'output') and response.output:
-                    # Extract search information
-                    if hasattr(response.output, 'search_info') and response.output.search_info:
-                        search_results = response.output.search_info.get("search_results", [])
-
-                    # Extract response content
-                    if (hasattr(response.output, 'choices') and
-                            response.output.choices and
-                            len(response.output.choices) > 0):
-                        response_content = response.output.choices[0].message.content
-
-                # Prepare final result
-                final_result = {
-                    "query": query,
-                    "search_results": search_results,
-                    "response_content": response_content,
-                    "model": self.model,
-                    "search_strategy": self.search_strategy
-                }
-
-                # Cache the result if enabled
-                if self.enable_cache and self.cache:
-                    self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
+        response = await dashscope.AioGeneration.call(
+            api_key=self.api_key,
+            model=self.model,
+            messages=messages,
+            enable_search=True,  # Enable web search
+            search_options={
+                "forced_search": True,  # Force web search
+                "enable_source": True,  # Include search source information
+                "enable_citation": False,  # Enable citation markers
+                "search_strategy": self.search_strategy,  # Search strategy
+            },
+            result_format="message",
+        )
+
+        search_results = []
+        response_content = ""
+
+        if hasattr(response, "output") and response.output:
+            if hasattr(response.output, "search_info") and response.output.search_info:
+                search_results = response.output.search_info.get("search_results", [])
+
+            if hasattr(response.output, "choices") and response.output.choices and len(response.output.choices) > 0:
+                response_content = response.output.choices[0].message.content
+
+        final_result = {
+            "query": query,
+            "search_results": search_results,
+            "response_content": response_content,
+            "model": self.model,
+            "search_strategy": self.search_strategy
+        }
 
-                # Post-process and set context
-                result = self.post_process(final_result)
-                if self.return_only_content:
-                    self.context.dashscope_search_result = result["response_content"]
-                else:
-                    self.context.dashscope_search_result = result
-
-                return
-
-            except Exception as e:
-                logger.warning(f"Dashscope search attempt {attempt + 1} failed for query='{query}': {e}")
-                if attempt < self.max_retries - 1:
-                    time.sleep(attempt + 1)  # Exponential backoff
-                else:
-                    logger.error(f"All {self.max_retries} attempts failed for Dashscope search")
-
-        self.context.dashscope_search_result = "dashscope_search failed"
-
-
-def main():
-    from flowllm.utils.common_utils import load_env
+        if self.enable_cache and self.cache:
+            self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
 
-    load_env()
+        self.output_dict["dashscope_search_result"] = final_result["response_content"]
+        if self.save_answer:
+            self.context.response.answer = final_result["response_content"]
 
-    C.set_default_service_config().init_by_service_config()
 
-    op = DashscopeSearchOp(enable_print=True, enable_cache=False)
+async def async_main():
+    C.set_service_config().init_by_service_config()
 
-    context = FlowContext(query="杭州明天天气")
-    op(context=context)
+    op = DashscopeSearchOp()
+    context = FlowContext(query="what is AI?")
+    await op.async_call(context=context)
     print(context.dashscope_search_result)
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(async_main())

flowllm/op/search/tavily_search_op.py

@@ -1,102 +1,91 @@
+import asyncio
 import json
 import os
-import time
-from typing import Literal
+from functools import partial
+from typing import Union, List
 
-from loguru import logger
 from tavily import TavilyClient
 
-from flowllm.context.flow_context import FlowContext
-from flowllm.context.service_context import C
-from flowllm.op.base_op import BaseOp
-from flowllm.storage.cache.data_cache import DataCache
+from flowllm.context import FlowContext, C
+from flowllm.op.base_tool_op import BaseToolOp
+from flowllm.schema.tool_call import ToolCall
 
 
 @C.register_op()
-class TavilySearchOp(BaseOp):
-    def __init__(self,
-                 enable_print: bool = True,
-                 enable_cache: bool = False,
-                 cache_path: str = "./web_search_cache",
-                 cache_expire_hours: float = 0.1,
-                 topic: Literal["general", "news", "finance"] = "general",
-                 max_retries: int = 3,
-                 return_only_content: bool = True,
-                 **kwargs):
+class TavilySearchOp(BaseToolOp):
+    def __init__(self, save_answer: bool = False, **kwargs):
         super().__init__(**kwargs)
+        self.save_answer = save_answer
+        self._client: TavilyClient | None = None
+
+    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
+                }
+            }
+        })
+
+    @property
+    def client(self):
+        if self._client is None:
+            self._client = TavilyClient(api_key=os.environ["FLOW_TAVILY_API_KEY"])
+        return self._client
+
+    def default_execute(self):
+        self.output_dict["tavily_search_result"] = "tavily search failed!"
+
+    async def search(self, query: str):
+        loop = asyncio.get_event_loop()
+        func = partial(self.client.search, query=query)
+        task = loop.run_in_executor(executor=C.thread_pool, func=func)  # noqa
+        return await task
+
+    async def extract(self, urls: Union[List[str], str]):
+        loop = asyncio.get_event_loop()
+        func = partial(self.client.extract, urls=urls, format="text")
+        task = loop.run_in_executor(executor=C.thread_pool, func=func)  # noqa
+        return await task
+
+    async def async_execute(self):
+        query: str = self.input_dict["query"]
 
-        self.enable_print = enable_print
-        self.enable_cache = enable_cache
-        self.cache_expire_hours = cache_expire_hours
-        self.topic = topic
-        self.max_retries = max_retries
-        self.return_only_content = return_only_content
-
-        # Initialize DataCache if caching is enabled
-        self.cache = DataCache(cache_path) if self.enable_cache else None
-        self._client = TavilyClient(api_key=os.getenv("FLOW_TAVILY_API_KEY", ""))
-
-    def post_process(self, response):
-        if self.enable_print:
-            logger.info("response=\n" + json.dumps(response, indent=2, ensure_ascii=False))
-
-        return response
-
-    def execute(self):
-        # Get query from context
-        query: str = self.context.query
-
-        # Check cache first
         if self.enable_cache and self.cache:
             cached_result = self.cache.load(query)
             if cached_result:
-                final_result = self.post_process(cached_result)
-                if self.return_only_content:
-                    self.context.tavily_search_result = json.dumps(final_result, ensure_ascii=False, indent=2)
-                else:
-                    self.context.tavily_search_result = final_result
+                self.output_dict["tavily_search_result"] = json.dumps(cached_result, ensure_ascii=False, indent=2)
                 return
 
-        for i in range(self.max_retries):
-            try:
-                response = self._client.search(query=query, topic=self.topic)
-                url_info_dict = {item["url"]: item for item in response["results"]}
-                response_extract = self._client.extract(urls=[item["url"] for item in response["results"]],
-                                                        format="text")
-
-                final_result = {}
-                for item in response_extract["results"]:
-                    url = item["url"]
-                    final_result[url] = url_info_dict[url]
-                    final_result[url]["raw_content"] = item["raw_content"]
-
-                # Cache the result if enabled
-                if self.enable_cache and self.cache:
-                    self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
-
-                final_result = self.post_process(final_result)
-
-                if self.return_only_content:
-                    self.context.tavily_search_result = json.dumps(final_result, ensure_ascii=False, indent=2)
-                else:
-                    self.context.tavily_search_result = final_result
-                return
+        response = await self.search(query=query)
+        url_info_dict = {item["url"]: item for item in response["results"]}
+        response_extract = await self.extract(urls=[item["url"] for item in response["results"]])
 
-            except Exception as e:
-                logger.exception(f"tavily search with query={query} encounter error with e={e.args}")
-                time.sleep(i + 1)
+        final_result = {}
+        for item in response_extract["results"]:
+            url = item["url"]
+            final_result[url] = url_info_dict[url]
+            final_result[url]["raw_content"] = item["raw_content"]
 
-        self.context.tavily_search_result = "tavily search failed!"
+        if self.enable_cache and self.cache is not None:
+            self.cache.save(query, final_result, expire_hours=self.cache_expire_hours)
 
+        self.output_dict["tavily_search_result"] = json.dumps(final_result, ensure_ascii=False, indent=2)
+        if self.save_answer:
+            self.context.response.answer = self.output_dict["tavily_search_result"]
 
-if __name__ == "__main__":
-    from flowllm.utils.common_utils import load_env
+async def async_main():
+    C.set_service_config().init_by_service_config()
 
-    load_env()
+    op = TavilySearchOp()
+    context = FlowContext(query="what is AI?")
+    await op.async_call(context=context)
+    print(context.tavily_search_result)
 
-    C.set_default_service_config().init_by_service_config()
 
-    op = TavilySearchOp(enable_cache=True)
-    context = FlowContext(query="A股医药为什么一直涨")
-    op(context=context)
-    print(context.tavily_search_result)
+if __name__ == "__main__":
+    asyncio.run(async_main())

flowllm/op/sequential_op.py

@@ -13,6 +13,10 @@ class SequentialOp(BaseOp):
         for op in self.ops:
             op.__call__(self.context)
 
+    async def async_execute(self):
+        for op in self.ops:
+            await op.async_call(self.context)
+
     def __rshift__(self, op: BaseOp):
         if isinstance(op, SequentialOp):
             self.ops.extend(op.ops)

flowllm/schema/flow_stream_chunk.py

@@ -0,0 +1,11 @@
+from pydantic import Field, BaseModel
+
+from flowllm.enumeration.chunk_enum import ChunkEnum
+
+
+class FlowStreamChunk(BaseModel):
+    flow_id: str = Field(default="")
+    chunk_type: ChunkEnum = Field(default=ChunkEnum.ANSWER)
+    chunk: str | bytes = Field(default="")
+    done: bool = Field(default=False)
+    metadata: dict = Field(default_factory=dict)

flowllm/schema/message.py

@@ -1,3 +1,4 @@
+import datetime
 from typing import List
 
 from pydantic import BaseModel, Field
@@ -12,6 +13,7 @@ class Message(BaseModel):
     reasoning_content: str = Field(default="")
     tool_calls: List[ToolCall] = Field(default_factory=list)
     tool_call_id: str = Field(default="")
+    time_created: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
     metadata: dict = Field(default_factory=dict)
 
     def simple_dump(self, add_reason_content: bool = True) -> dict:

flowllm/schema/service_config.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, List
 
 from pydantic import BaseModel, Field
 
@@ -6,7 +6,7 @@ from flowllm.schema.tool_call import ToolCall
 
 
 class MCPConfig(BaseModel):
-    transport: str = Field(default="", description="stdio/http/sse/streamable-http")
+    transport: str = Field(default="", description="stdio/http/sse")
     host: str = Field(default="0.0.0.0")
     port: int = Field(default=8001)
 
@@ -25,11 +25,14 @@ class CmdConfig(BaseModel):
 
 class FlowConfig(ToolCall):
     flow_content: str = Field(default="")
-    service_type: str = Field(default="all", description="all/http/mcp/cmd")
+    use_async: bool = Field(default=True)
+    service_type: str = Field(default="http+mcp", description="http+mcp/http/mcp")
+    stream: bool = Field(default=False)
 
 class OpConfig(BaseModel):
     backend: str = Field(default="")
     language: str = Field(default="")
+    max_retries: int = Field(default=1)
     raise_exception: bool = Field(default=True)
     prompt_path: str = Field(default="")
     llm: str = Field(default="")
@@ -61,6 +64,8 @@ class ServiceConfig(BaseModel):
     language: str = Field(default="")
     thread_pool_max_workers: int = Field(default=16)
     ray_max_workers: int = Field(default=1)
+    import_config: str = Field(default="", description="Import the configuration in the same path as the base")
+    disabled_flows: List[str] = Field(default_factory=list)
 
     cmd: CmdConfig = Field(default_factory=CmdConfig)
     mcp: MCPConfig = Field(default_factory=MCPConfig)

flowllm/schema/tool_call.py

@@ -1,7 +1,8 @@
 import json
 from typing import Dict, List
 
-from pydantic import BaseModel, Field
+from mcp.types import Tool
+from pydantic import BaseModel, Field, model_validator
 
 
 class ParamAttrs(BaseModel):
@@ -58,12 +59,27 @@ class ToolCall(BaseModel):
     type: str = Field(default="function")
     name: str = Field(default="")
 
-    arguments: dict = Field(default_factory=dict, description="tool execution arguments")
+    arguments: str = Field(default="", description="tool execution arguments")
 
     description: str = Field(default="")
     input_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
     output_schema: Dict[str, ParamAttrs] = Field(default_factory=dict)
 
+    @model_validator(mode="before")
+    @classmethod
+    def init_tool_call(cls, data: dict):
+        tool_type = data.get("type", "")
+        tool_type_dict = data.get(tool_type, {})
+
+        for key in ["name", "arguments"]:
+            if key not in data:
+                data[key] = tool_type_dict.get(key, "")
+        return data
+
+    @property
+    def argument_dict(self) -> dict:
+        return json.loads(self.arguments)
+
     def simple_input_dump(self, version: str = "v1") -> dict:
         if version == "v1":
             required_list = [name for name, tool_param in self.input_schema.items() if tool_param.required]
@@ -91,7 +107,7 @@ class ToolCall(BaseModel):
                 "index": self.index,
                 "id": self.id,
                 self.type: {
-                    "arguments": json.dumps(self.arguments, ensure_ascii=False),
+                    "arguments": self.arguments,
                     "name": self.name
                 },
                 "type": self.type,
@@ -111,8 +127,41 @@ class ToolCall(BaseModel):
                 if name:
                     self.name = name
                 if arguments:
-                    self.arguments = json.loads(arguments)
+                    self.arguments = arguments
         else:
             raise NotImplementedError(f"version {version} not supported")
 
         return self
+
+    @classmethod
+    def from_mcp_tool(cls, tool: Tool) -> "ToolCall":
+        input_schema = {}
+        properties = tool.inputSchema["properties"]
+        required = tool.inputSchema["required"]
+        for name, attr_dict in properties.items():
+            param_attrs = ParamAttrs()
+
+            if name in required:
+                param_attrs.required = True
+            param_attrs.type = attr_dict.get("type", "str")
+            param_attrs.description = attr_dict.get("description", "")
+            if "enum" in attr_dict:
+                param_attrs.enum = attr_dict["enum"]
+            input_schema[name] = param_attrs
+
+        return cls(name=tool.name,
+                   description=tool.description,
+                   input_schema=input_schema)
+
+
+if __name__ == "__main__":
+    tool_call = ToolCall(**{
+        "id": "call_0fb6077ad56f4647b0b04a",
+        "function": {
+            "arguments": "{\"symbol\": \"ZETA\"}",
+            "name": "get_stock_info"
+        },
+        "type": "function",
+        "index": 0
+    })
+    print(tool_call.simple_output_dump())

flowllm/service/__init__.py

@@ -1,3 +1,2 @@
-from .cmd_service import CmdService
 from .http_service import HttpService
 from .mcp_service import MCPService