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

flowllm/op/{agent/react_v2_op.py → llm/react_llm_op.py}

@@ -1,3 +1,4 @@
+import asyncio
 import datetime
 import json
 import time
@@ -7,36 +8,46 @@ 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.op import BaseToolOp
 from flowllm.schema.flow_response import FlowResponse
 from flowllm.schema.message import Message, Role
+from flowllm.schema.tool_call import ToolCall
 
 
-@C.register_op()
-class ReactV2Op(BaseLLMOp):
+@C.register_op(name="react_llm_op")
+class ReactLLMOp(BaseToolOp):
     file_path: str = __file__
 
-    def execute(self):
+    def __init__(self, llm="qwen3_30b_instruct", **kwargs):
+        super().__init__(llm=llm, **kwargs)
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "query_llm",
+            "description": "use this query to query an LLM",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "search keyword",
+                    "required": True
+                }
+            }
+        })
+
+    async def async_execute(self):
         query: str = self.context.query
 
         max_steps: int = int(self.op_params.get("max_steps", 10))
-        from flowllm.flow.base_tool_flow import BaseToolFlow
-        from flowllm.flow.gallery import DashscopeSearchToolFlow, CodeToolFlow
+        from flowllm.op import BaseToolOp
+        from flowllm.op.search import DashscopeSearchOp
 
-        tools: List[BaseToolFlow] = [DashscopeSearchToolFlow(), CodeToolFlow()]
-
-        """
-        NOTE : x.tool_call.name != x.name
-        `x.tool_call.name` is tool's namex.name is flow's name(unique service name)
-        """
-        tool_dict: Dict[str, BaseToolFlow] = {x.tool_call.name: x for x in tools}
+        tools: List[BaseToolOp] = [DashscopeSearchOp(save_answer=True)]
+        tool_dict: Dict[str, BaseToolOp] = {x.tool_call.name: x for x in tools}
         for name, tool_call in tool_dict.items():
             logger.info(f"name={name} "
                         f"tool_call={json.dumps(tool_call.tool_call.simple_input_dump(), ensure_ascii=False)}")
 
         now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-        has_terminate_tool = False
-
         user_prompt = self.prompt_format(prompt_name="role_prompt",
                                          time=now_time,
                                          tools=",".join(list(tool_dict.keys())),
@@ -45,7 +56,7 @@ class ReactV2Op(BaseLLMOp):
         logger.info(f"step.0 user_prompt={user_prompt}")
 
         for i in range(max_steps):
-            assistant_message: Message = self.llm.chat(messages, tools=[x.tool_call for x in tools])
+            assistant_message: Message = await self.llm.achat(messages, tools=[x.tool_call for x in tools])
             messages.append(assistant_message)
             logger.info(f"assistant.round{i}.reasoning_content={assistant_message.reasoning_content}\n"
                         f"content={assistant_message.content}\n"
@@ -61,11 +72,15 @@ class ReactV2Op(BaseLLMOp):
                     logger.warning(f"step={i} no tool_call.name={tool_call.name}")
                     continue
 
-                self.submit_task(tool_dict[tool_call.name].__call__, **tool_call.argument_dict)
+                self.submit_async_task(tool_dict[tool_call.name].copy().async_call,
+                                       context=self.context.copy(**tool_call.argument_dict))
                 time.sleep(1)
 
-            for i, (tool_result, tool_call) in enumerate(zip(self.join_task(), assistant_message.tool_calls)):
-                logger.info(f"submit step={i} tool_calls.name={tool_call.name} tool_result={tool_result}")
+            task_results = await self.join_async_task()
+
+            for j, tool_result in enumerate(task_results):
+                tool_call = assistant_message.tool_calls[j]
+                logger.info(f"submit step.index={i}.{j} tool_result={tool_result}")
                 if isinstance(tool_result, FlowResponse):
                     tool_result = tool_result.answer
                 else:
@@ -73,14 +88,18 @@ class ReactV2Op(BaseLLMOp):
                 tool_message = Message(role=Role.TOOL, content=tool_result, tool_call_id=tool_call.id)
                 messages.append(tool_message)
 
-        # Store results in context instead of response
         self.context.response.messages = messages
         self.context.response.answer = messages[-1].content
 
 
-if __name__ == "__main__":
-    C.set_default_service_config().init_by_service_config()
+async def main():
+    C.set_service_config().init_by_service_config()
     context = FlowContext(query="茅台和五粮现在股价多少？")
 
-    op = ReactV2Op()
-    op(context=context)
+    op = ReactLLMOp()
+    result = await op.async_call(context=context)
+    print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

flowllm/op/llm/simple_llm_op.py

@@ -0,0 +1,48 @@
+import asyncio
+from typing import List
+
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.op import BaseToolOp
+from flowllm.schema.message import Message, Role
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_op(name="simple_llm_op")
+class SimpleLLMOp(BaseToolOp):
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "query_llm",
+            "description": "use this query to query an LLM",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "search keyword",
+                    "required": True
+                }
+            }
+        })
+
+    async def async_execute(self):
+        query: str = self.input_dict["query"]
+        logger.info(f"query={query}")
+        messages: List[Message] = [Message(role=Role.USER, content=query)]
+
+        assistant_message: Message = await self.llm.achat(messages)
+        self.context.response.answer = assistant_message.content
+
+
+async def main():
+    C.set_service_config().init_by_service_config()
+    context = FlowContext(query="hello", stream_queue=asyncio.Queue())
+
+    op = SimpleLLMOp()
+    result = await op.async_call(context=context)
+    print(result)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

flowllm/op/llm/stream_llm_op.py

@@ -0,0 +1,61 @@
+import asyncio
+from typing import List
+
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.enumeration.chunk_enum import ChunkEnum
+from flowllm.op import BaseToolOp
+from flowllm.schema.message import Message, Role
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_op(name="stream_llm_op")
+class StreamLLMOp(BaseToolOp):
+
+    def build_tool_call(self) -> ToolCall:
+        return ToolCall(**{
+            "name": "query_llm",
+            "description": "use this query to query an LLM",
+            "input_schema": {
+                "query": {
+                    "type": "str",
+                    "description": "search keyword",
+                    "required": True
+                }
+            }
+        })
+
+    async def async_execute(self):
+        query: str = self.input_dict["query"]
+        logger.info(f"query={query}")
+        messages: List[Message] = [Message(role=Role.USER, content=query)]
+
+        async for chunk, chunk_type in self.llm.astream_chat(messages):  # noqa
+            if chunk_type == ChunkEnum.ANSWER:
+                await self.context.add_stream_answer(chunk)
+
+        await self.context.add_stream_done()
+
+
+async def main():
+    C.set_service_config().init_by_service_config()
+    context = FlowContext(query="hello, introduce yourself.", stream_queue=asyncio.Queue())
+
+    op = StreamLLMOp()
+    task = asyncio.create_task(op.async_call(context=context))
+
+    while True:
+        stream_chunk = await context.stream_queue.get()
+        if stream_chunk.done:
+            print("\nend")
+            break
+        else:
+            print(stream_chunk.chunk, end="")
+
+    await task
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

flowllm/op/mcp/__init__.py

@@ -0,0 +1,2 @@
+from .ant_op import AntSearchOp, AntInvestmentOp
+from .base_sse_mcp_op import BaseSSEMcpOp

flowllm/op/mcp/ant_op.py

@@ -0,0 +1,42 @@
+import asyncio
+import json
+import os
+
+from flowllm.context import FlowContext, C
+from flowllm.op.mcp.base_sse_mcp_op import BaseSSEMcpOp
+
+
+@C.register_op()
+class AntSearchOp(BaseSSEMcpOp):
+
+    def __init__(self, **kwargs):
+        host = os.getenv("FLOW_MCP_HOSTS", "").split(",")[0]
+        super().__init__(host=host, tool_name="search", **kwargs)
+
+
+@C.register_op()
+class AntInvestmentOp(BaseSSEMcpOp):
+
+    def __init__(self, **kwargs):
+        host = os.getenv("FLOW_MCP_HOSTS", "").split(",")[0]
+        super().__init__(host=host, tool_name="investment_analysis", **kwargs)
+
+
+async def async_main():
+    op = AntSearchOp()
+    context = FlowContext(query="阿里巴巴怎么样？", entity="阿里巴巴")
+    await op.async_call(context=context)
+    print(json.dumps(op.tool_call.simple_input_dump(), ensure_ascii=False))
+    print(context.response.answer)
+
+    op = AntInvestmentOp()
+    context = FlowContext(entity="阿里巴巴", analysis_category="股票")
+    await op.async_call(context=context)
+    print(json.dumps(op.tool_call.simple_input_dump(), ensure_ascii=False))
+    print(context.response.answer)
+
+
+if __name__ == "__main__":
+    C.prepare_sse_mcp().set_service_config().init_by_service_config()
+
+    asyncio.run(async_main())

flowllm/op/mcp/base_sse_mcp_op.py

@@ -0,0 +1,28 @@
+from fastmcp import Client
+from mcp.types import CallToolResult
+
+from flowllm.context import C
+from flowllm.op import BaseToolOp
+from flowllm.schema.tool_call import ToolCall
+
+
+class BaseSSEMcpOp(BaseToolOp):
+
+    def __init__(self, host: str = "", tool_name: str = "", **kwargs):
+        self.host: str = host
+        self.tool_name: str = tool_name
+        super().__init__(**kwargs)
+
+    def build_tool_call(self) -> ToolCall:
+        key = f"{self.host}/{self.tool_name}"
+        assert key in C.sse_mcp_dict, \
+            f"host={self.host} tool_name={self.tool_name} not found in mcp_tool_call_dict"
+        return C.sse_mcp_dict[key]
+
+    def default_execute(self):
+        self.context.response.answer = self.output_dict[f"{self.name}_result"] = f"{self.name} execute failed!"
+
+    async def async_execute(self):
+        async with Client(f"{self.host}/sse/") as client:
+            result: CallToolResult = await client.call_tool(self.tool_name, arguments=self.input_dict)
+        self.context.response.answer = self.output_dict[f"{self.name}_result"] = result.content[0].text

flowllm/op/parallel_op.py

@@ -12,9 +12,13 @@ class ParallelOp(BaseOp):
     def execute(self):
         for op in self.ops:
             self.submit_task(op.__call__, context=self.context)
-
         self.join_task(task_desc="parallel execution")
 
+    async def async_execute(self):
+        for op in self.ops:
+            self.submit_async_task(op.async_call, context=self.context)
+        return await self.join_async_task()
+
     def __or__(self, op: BaseOp):
         if isinstance(op, ParallelOp):
             self.ops.extend(op.ops)

flowllm/op/search/__init__.py

@@ -1,3 +1,2 @@
-from .dashscope_deep_research_op import DashscopeDeepResearchOp
 from .dashscope_search_op import DashscopeSearchOp
-from .tavily_search_op import TavilySearchOp
+from .tavily_search_op import TavilySearchOp

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,36 +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.environ["FLOW_DASHSCOPE_API_KEY"]
-        self.cache_path: str = cache_path
-        self._cache: DataCache | None = None
-
-    @property
-    def cache(self):
-        if self.enable_cache and self._cache is None:
-            self._cache = DataCache(self.cache_path)
-        return self._cache
+        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:
@@ -65,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"])
+    def default_execute(self):
+        self.output_dict["dashscope_search_result"] = "dashscope search failed!"
 
-        return response_data
+    async def async_execute(self):
+        query: str = self.input_dict["query"]
 
-    def execute(self):
-        """Execute the Dashscope search operation"""
-        # Get query from context - support multiple parameter names
-        query = self.context.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)
-
-                # 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
+        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
+        }
 
-                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())