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

flowllm/llm/openai_compatible_llm.py

@@ -2,7 +2,7 @@ import os
 from typing import List, Dict
 
 from loguru import logger
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI
 from openai.types import CompletionUsage
 from pydantic import Field, PrivateAttr, model_validator
 
@@ -33,19 +33,21 @@ class OpenAICompatibleBaseLLM(BaseLLM):
     base_url: str = Field(default_factory=lambda: os.getenv("FLOW_LLM_BASE_URL"),
                           description="Base URL for the API endpoint")
     _client: OpenAI = PrivateAttr()
+    _aclient: 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 validator runs after all field validation is complete,
-        ensuring we have valid API credentials before creating the client.
+        ensuring we have valid API credentials before creating the clients.
         
         Returns:
             Self for method chaining
         """
         self._client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+        self._aclient = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
         return self
 
     def stream_chat(self, messages: List[Message], tools: List[ToolCall] = None, **kwargs):
@@ -151,6 +153,109 @@ class OpenAICompatibleBaseLLM(BaseLLM):
                 else:
                     yield e.args, ChunkEnum.ERROR
 
+    async def astream_chat(self, messages: List[Message], tools: List[ToolCall] = None, **kwargs):
+        """
+        Async stream chat completions from OpenAI-compatible API.
+        
+        This method handles async streaming responses and categorizes chunks into different types:
+        - THINK: Reasoning/thinking content from the model
+        - ANSWER: Regular response content
+        - TOOL: Tool calls that need to be executed
+        - USAGE: Token usage statistics
+        - ERROR: Error information
+        
+        Args:
+            messages: List of conversation messages
+            tools: Optional list of tools available to the model
+            **kwargs: Additional parameters
+            
+        Yields:
+            Tuple of (chunk_content, ChunkEnum) for each streaming piece
+        """
+        for i in range(self.max_retries):
+            try:
+                extra_body = {}
+                if self.enable_thinking:
+                    extra_body["enable_thinking"] = True  # qwen3 params
+
+                completion = await self._aclient.chat.completions.create(
+                    model=self.model_name,
+                    messages=[x.simple_dump() for x in messages],
+                    seed=self.seed,
+                    top_p=self.top_p,
+                    stream=True,
+                    stream_options=self.stream_options,
+                    temperature=self.temperature,
+                    extra_body=extra_body,
+                    tools=[x.simple_input_dump() for x in tools] if tools else None,
+                    parallel_tool_calls=self.parallel_tool_calls)
+
+                # Initialize tool call tracking
+                ret_tools: List[ToolCall] = []  # Accumulate tool calls across chunks
+                is_answering: bool = False  # Track when model starts answering
+
+                # Process each chunk in the streaming response
+                async for chunk in completion:
+                    # Handle chunks without choices (usually usage info)
+                    if not chunk.choices:
+                        yield chunk.usage, ChunkEnum.USAGE
+
+                    else:
+                        delta = chunk.choices[0].delta
+
+                        # Handle reasoning/thinking content (model's internal thoughts)
+                        if hasattr(delta, 'reasoning_content') and delta.reasoning_content is not None:
+                            yield delta.reasoning_content, ChunkEnum.THINK
+
+                        else:
+                            # Mark transition from thinking to answering
+                            if not is_answering:
+                                is_answering = True
+
+                            # Handle regular response content
+                            if delta.content is not None:
+                                yield delta.content, ChunkEnum.ANSWER
+
+                            # Handle tool calls (function calling)
+                            if delta.tool_calls is not None:
+                                for tool_call in delta.tool_calls:
+                                    index = tool_call.index
+
+                                    # Ensure we have enough tool call slots
+                                    while len(ret_tools) <= index:
+                                        ret_tools.append(ToolCall(index=index))
+
+                                    # Accumulate tool call information across chunks
+                                    if tool_call.id:
+                                        ret_tools[index].id += tool_call.id
+
+                                    if tool_call.function and tool_call.function.name:
+                                        ret_tools[index].name += tool_call.function.name
+
+                                    if tool_call.function and tool_call.function.arguments:
+                                        ret_tools[index].arguments += tool_call.function.arguments
+
+                # Yield completed tool calls after streaming finishes
+                if ret_tools:
+                    tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {}
+                    for tool in ret_tools:
+                        # Only yield tool calls that correspond to available tools
+                        if tool.name not in tool_dict:
+                            continue
+
+                        yield tool, ChunkEnum.TOOL
+
+                return
+
+            except Exception as e:
+                logger.exception(f"async stream chat with model={self.model_name} encounter error with e={e.args}")
+
+                # Handle retry logic
+                if i == self.max_retries - 1 and self.raise_exception:
+                    raise e
+                else:
+                    yield e.args, ChunkEnum.ERROR
+
     def _chat(self, messages: List[Message], tools: List[ToolCall] = None, enable_stream_print: bool = False,
               **kwargs) -> Message:
         """
@@ -224,6 +329,95 @@ class OpenAICompatibleBaseLLM(BaseLLM):
                        content=answer_content,
                        tool_calls=tool_calls)
 
+    async def _achat(self, messages: List[Message], tools: List[ToolCall] = None, enable_stream_print: bool = False,
+                     **kwargs) -> Message:
+        """
+        Perform an async complete chat completion by aggregating streaming chunks.
+        
+        This method consumes the entire async streaming response and combines all
+        chunks into a single Message object. It separates reasoning content,
+        regular answer content, and tool calls.
+        
+        Args:
+            messages: List of conversation messages
+            tools: Optional list of tools available to the model
+            enable_stream_print: Whether to print streaming response to console
+            **kwargs: Additional parameters
+            
+        Returns:
+            Complete Message with all content aggregated
+        """
+
+        enter_think = False  # Whether we've started printing thinking content
+        enter_answer = False  # Whether we've started printing answer content
+        reasoning_content = ""  # Model's internal reasoning
+        answer_content = ""  # Final response content
+        tool_calls = []  # List of tool calls to execute
+
+        # Consume async streaming response and aggregate chunks by type
+        async for chunk, chunk_enum in self.astream_chat(messages, tools, **kwargs):
+            if chunk_enum is ChunkEnum.USAGE:
+                # Display token usage statistics
+                if enable_stream_print:
+                    if isinstance(chunk, CompletionUsage):
+                        print(f"\n<usage>{chunk.model_dump_json(indent=2)}</usage>")
+                    else:
+                        print(f"\n<usage>{chunk}</usage>")
+
+            elif chunk_enum is ChunkEnum.THINK:
+                if enable_stream_print:
+                    # Format thinking/reasoning content
+                    if not enter_think:
+                        enter_think = True
+                        print("<think>\n", end="")
+                    print(chunk, end="")
+
+                reasoning_content += chunk
+
+            elif chunk_enum is ChunkEnum.ANSWER:
+                if enable_stream_print:
+                    if not enter_answer:
+                        enter_answer = True
+                        # Close thinking section if we were in it
+                        if enter_think:
+                            print("\n</think>")
+                    print(chunk, end="")
+
+                answer_content += chunk
+
+            elif chunk_enum is ChunkEnum.TOOL:
+                if enable_stream_print:
+                    print(f"\n<tool>{chunk.model_dump_json()}</tool>", end="")
+
+                tool_calls.append(chunk)
+
+            elif chunk_enum is ChunkEnum.ERROR:
+                if enable_stream_print:
+                    # Display error information
+                    print(f"\n<error>{chunk}</error>", end="")
+
+        # Construct complete response message
+        return Message(role=Role.ASSISTANT,
+                       reasoning_content=reasoning_content,
+                       content=answer_content,
+                       tool_calls=tool_calls)
+
+
+async def async_main():
+    from flowllm.utils.common_utils import load_env
+
+    load_env()
+
+    # model_name = "qwen-max-2025-01-25"
+    model_name = "qwen3-30b-a3b-thinking-2507"
+    llm = OpenAICompatibleBaseLLM(model_name=model_name)
+
+    # Test async chat
+    message: Message = await llm.achat([Message(role=Role.USER, content="hello")], [],
+                                       enable_stream_print=True)
+    print("Async result:", message)
+
+
 def main():
     from flowllm.utils.common_utils import load_env
 
@@ -231,9 +425,15 @@ def main():
 
     model_name = "qwen-max-2025-01-25"
     llm = OpenAICompatibleBaseLLM(model_name=model_name)
+
+    # Test sync chat
     message: Message = llm.chat([Message(role=Role.USER, content="hello")], [],
                                 enable_stream_print=False)
-    print(message)
+    print("Sync result:", message)
+
 
 if __name__ == "__main__":
-    main()
+    # main()
+
+    import asyncio
+    asyncio.run(async_main())

flowllm/op/__init__.py

@@ -1,3 +1,12 @@
-from flowllm.op.akshare.get_ak_a_code_op import GetAkACodeOp
-from flowllm.op.akshare.get_ak_a_info_op import GetAkAInfoOp, GetAkASpotOp, GetAkAMoneyFlowOp, GetAkAFinancialInfoOp, GetAkANewsOp, MergeAkAInfoOp
-from flowllm.op.mock_op import Mock1Op, Mock2Op, Mock3Op, Mock4Op, Mock5Op, Mock6Op
+from .base_llm_op import BaseLLMOp
+from .base_op import BaseOp
+from .base_ray_op import BaseRayOp
+
+"""
+op folder
+"""
+from . import akshare
+from . import code
+from . import gallery
+from . import search
+from . import agent

flowllm/op/agent/__init__.py

@@ -0,0 +1 @@
+from .react_v1_op import ReactV1Op

flowllm/op/agent/react_v1_op.py

@@ -0,0 +1,109 @@
+import datetime
+import json
+import time
+from typing import List, Dict
+
+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.schema.flow_response import FlowResponse
+from flowllm.schema.message import Message, Role
+
+
+@C.register_op()
+class ReactV1Op(BaseLLMOp):
+    file_path: str = __file__
+
+    def 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, TerminateToolFlow
+
+        tools: List[BaseToolFlow] = [DashscopeSearchToolFlow(), CodeToolFlow(), TerminateToolFlow()]
+
+        """
+        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}
+        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())),
+                                         query=query)
+        messages: List[Message] = [Message(role=Role.USER, content=user_prompt)]
+        logger.info(f"step.0 user_prompt={user_prompt}")
+
+        for i in range(max_steps):
+            if has_terminate_tool:
+                assistant_message: Message = self.llm.chat(messages)
+            else:
+                assistant_message: Message = self.llm.chat(messages, tools=[x.tool_call for x in tools])
+
+            messages.append(assistant_message)
+            logger.info(f"assistant.{i}.reasoning_content={assistant_message.reasoning_content}\n"
+                        f"content={assistant_message.content}\n"
+                        f"tool.size={len(assistant_message.tool_calls)}")
+
+            if has_terminate_tool:
+                break
+
+            for tool in assistant_message.tool_calls:
+                if tool.name == "terminate":
+                    has_terminate_tool = True
+                    logger.info(f"step={i} find terminate tool, break.")
+                    break
+
+            if not has_terminate_tool and not assistant_message.tool_calls:
+                logger.warning(f"【bugfix】step={i} no tools, break.")
+                has_terminate_tool = True
+
+            for j, tool_call in enumerate(assistant_message.tool_calls):
+                logger.info(f"submit step={i} tool_calls.name={tool_call.name} argument_dict={tool_call.argument_dict}")
+
+                if tool_call.name not in tool_dict:
+                    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)
+                time.sleep(1)
+
+            if not has_terminate_tool:
+                user_content_list = []
+                for tool_result, tool_call in zip(self.join_task(), assistant_message.tool_calls):
+                    logger.info(f"submit step={i} tool_calls.name={tool_call.name} tool_result={tool_result}")
+                    if isinstance(tool_result, FlowResponse):
+                        tool_result = tool_result.answer
+                    else:
+                        tool_result = str(tool_result)
+                    user_content_list.append(f"<tool_response>\n{tool_result}\n</tool_response>")
+                user_content_list.append(self.prompt_format(prompt_name="next_prompt"))
+                assistant_message.tool_calls.clear()
+                messages.append(Message(role=Role.USER, content="\n".join(user_content_list)))
+
+            else:
+                assistant_message.tool_calls.clear()
+                query = self.prompt_format(prompt_name="final_prompt", query=query)
+                messages.append(Message(role=Role.USER, content=query))
+
+        # 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()
+    context = FlowContext(query="茅台和五粮现在股价多少？")
+
+    op = ReactV1Op()
+    op(context=context)

flowllm/op/agent/react_v1_prompt.yaml

@@ -0,0 +1,54 @@
+role_prompt: |
+  You are a helpful assistant.
+  The current time is {time}.
+  Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc.
+  Please first think about how to break down the problem into subtasks, what tools and parameters should be used for each subtask, and finally provide the tool call name and parameters.
+  Try calling the same tool multiple times with different parameters to obtain information from various perspectives.
+  Please determine the response language based on the language of the user's question.
+  
+  {query}
+
+# write a complete and rigorous report to answer user's questions based on the context.
+next_prompt: |
+  Think based on the current content and the user's question: Is the current context sufficient to answer the user's question?  
+  
+  - If the current context is not sufficient to answer the user's question, consider what information is missing. 
+    Re-plan and think about how to break down the missing information into subtasks. 
+    For each subtask, determine what tools and parameters should be used for the query.
+    Please first provide the reasoning process, then give the tool call name and parameters.
+  
+  - If the current context is sufficient to answer the user's question, use the **terminate** tool.
+
+#   Please determine the response language based on the language of the user's question.
+final_prompt: |
+  Please integrate the context and provide a complete answer to the user's question.
+
+  # User's Question
+  {query}
+
+
+
+
+role_prompt_zh: |
+  你是一个有用的助手。
+  当前时间是 {time}。
+  请根据用户的问题，主动选择最合适的工具或工具组合，包括 {tools} 等。
+  请先思考如何将问题分解为子任务，每个子任务应使用哪些工具和参数，最后提供工具调用名称和参数。
+  尝试多次使用相同的工具，但使用不同的参数，从多个角度获取信息。
+  请根据用户问题的语言来确定回复的语言。
+  
+  {query}
+
+next_prompt_zh: |
+  根据当前内容和用户的问题进行思考：当前上下文是否足以回答用户的问题？
+  - 如果当前上下文不足以回答用户的问题，请考虑缺少哪些信息。
+    重新规划并思考如何将缺失的信息分解为子任务。
+    对于每个子任务，确定应使用哪些工具和参数进行查询。
+    请先提供推理过程，然后给出工具调用名称和参数。
+  - 如果当前上下文足以回答用户的问题，请使用 **terminate** 工具。
+
+final_prompt_zh: |
+  请整合上下文，为用户的问题提供一个完整的答案。
+  
+  # 用户的问题
+  {query}

flowllm/op/agent/react_v2_op.py

@@ -0,0 +1,86 @@
+import datetime
+import json
+import time
+from typing import List, Dict
+
+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.schema.flow_response import FlowResponse
+from flowllm.schema.message import Message, Role
+
+
+@C.register_op()
+class ReactV2Op(BaseLLMOp):
+    file_path: str = __file__
+
+    def 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
+
+        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}
+        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())),
+                                         query=query)
+        messages: List[Message] = [Message(role=Role.USER, content=user_prompt)]
+        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])
+            messages.append(assistant_message)
+            logger.info(f"assistant.round{i}.reasoning_content={assistant_message.reasoning_content}\n"
+                        f"content={assistant_message.content}\n"
+                        f"tool.size={len(assistant_message.tool_calls)}")
+
+            if not assistant_message.tool_calls:
+                break
+
+            for j, tool_call in enumerate(assistant_message.tool_calls):
+                logger.info(f"submit step={i} tool_calls.name={tool_call.name} argument_dict={tool_call.argument_dict}")
+
+                if tool_call.name not in tool_dict:
+                    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)
+                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}")
+                if isinstance(tool_result, FlowResponse):
+                    tool_result = tool_result.answer
+                else:
+                    tool_result = str(tool_result)
+                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()
+    context = FlowContext(query="茅台和五粮现在股价多少？")
+
+    op = ReactV2Op()
+    op(context=context)

flowllm/op/agent/react_v2_prompt.yaml

@@ -0,0 +1,35 @@
+role_prompt: |
+  You are a helpful assistant.
+  The current time is {time}.
+  Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc.
+  Please first think about how to break down the problem into subtasks, what tools and parameters should be used for each subtask, and finally provide the tool call name and parameters.
+  Try calling the same tool multiple times with different parameters to obtain information from various perspectives.
+  Please determine the response language based on the language of the user's question.
+  
+  {query}
+
+next_prompt: |
+  Think based on the current content and the user's question: Is the current context sufficient to answer the user's question?    
+  - If the current context is not sufficient to answer the user's question, consider what information is missing. 
+    Re-plan and think about how to break down the missing information into subtasks. 
+    For each subtask, determine what tools and parameters should be used for the query.
+    Please first provide the reasoning process, then give the tool call name and parameters.
+  - If the current context is sufficient to answer the user's question, please integrate the context and provide a complete answer to the user's question.
+
+role_prompt_zh: |
+  你是一个有用的助手。
+  当前时间是 {time}。
+  请根据用户的问题，主动选择最合适的工具或工具组合，包括 {tools} 等。
+  请先思考如何将问题分解为子任务，每个子任务应使用哪些工具和参数，最后提供工具调用名称和参数。
+  尝试多次使用相同的工具，但使用不同的参数，从多个角度获取信息。
+  请根据用户问题的语言来确定回复的语言。
+  
+  {query}
+
+next_prompt_zh: |
+  根据当前内容和用户的问题进行思考：当前上下文是否足以回答用户的问题？
+  - 如果当前上下文不足以回答用户的问题，请考虑缺少哪些信息。
+    重新规划并思考如何将缺失的信息分解为子任务。
+    对于每个子任务，确定应使用哪些工具和参数进行查询。
+    请先提供推理过程，然后给出工具调用名称和参数。
+  - 如果当前上下文足以回答用户的问题，请整合上下文，为用户的问题提供一个完整的答案。

flowllm/op/akshare/__init__.py

@@ -0,0 +1,3 @@
+from .get_ak_a_code_op import GetAkACodeOp
+from .get_ak_a_info_op import GetAkAInfoOp, GetAkASpotOp, GetAkAMoneyFlowOp, GetAkAFinancialInfoOp, \
+    GetAkANewsOp, MergeAkAInfoOp

flowllm/op/akshare/get_ak_a_code_op.py

@@ -6,13 +6,12 @@ import akshare as ak
 import pandas as pd
 from loguru import logger
 
-from flowllm.config.pydantic_config_parser import get_default_config
 from flowllm.context.flow_context import FlowContext
 from flowllm.context.service_context import C
 from flowllm.enumeration.role import Role
-from flowllm.op.llm_base_op import BaseLLMOp
+from flowllm.op.base_llm_op import BaseLLMOp
 from flowllm.schema.message import Message
-from flowllm.utils.dataframe_cache import DataFrameCache
+from flowllm.storage.cache.data_cache import DataCache
 from flowllm.utils.timer import timer
 
 
@@ -25,9 +24,9 @@ class GetAkACodeOp(BaseLLMOp):
 
     @staticmethod
     def download_a_stock_df():
-        df_cache = DataFrameCache()
+        cache = DataCache()
         save_df_key: str = "all_a_stock_name_code"
-        if not df_cache.exists(save_df_key):
+        if not cache.exists(save_df_key):
             stock_sh_a_spot_em_df = ak.stock_sh_a_spot_em()
             stock_sz_a_spot_em_df = ak.stock_sz_a_spot_em()
             stock_bj_a_spot_em_df = ak.stock_bj_a_spot_em()
@@ -36,9 +35,9 @@ class GetAkACodeOp(BaseLLMOp):
             df = df.drop(columns=["序号"])
             df = df.reset_index(drop=True)
             df = df.sort_values(by="代码")
-            df_cache.save(save_df_key, df, expire_hours=0.25)
+            cache.save(save_df_key, df, expire_hours=0.25)
 
-        df = df_cache.load(save_df_key, dtype={"代码": str})
+        df = cache.load(save_df_key, dtype={"代码": str})
         return df
 
     def get_name_code_dict(self) -> dict:
@@ -73,7 +72,7 @@ class GetAkACodeOp(BaseLLMOp):
         stock_names = "\n".join([x.strip() for x in stock_names if x])
         prompt = self.prompt_format(prompt_name="find_stock_name",
                                     stock_names=stock_names,
-                                    query=self.flow_context.query)
+                                    query=self.context.query)
         logger.info(f"prompt={prompt}")
 
         def callback_fn(msg: Message):
@@ -97,20 +96,13 @@ class GetAkACodeOp(BaseLLMOp):
             time.sleep(1)
 
         stock_names = sorted(set(self.join_task()))
-        self.flow_context.code_infos = {name_code_dict[n]: {} for n in stock_names}
-        logger.info(f"code_infos={self.flow_context.code_infos}")
+        self.context.code_infos = {name_code_dict[n]: {"股票名称": n} for n in stock_names}
+        logger.info(f"code_infos={self.context.code_infos}")
 
 
 if __name__ == "__main__":
-    from concurrent.futures import ThreadPoolExecutor
-
-    C.thread_pool = ThreadPoolExecutor(max_workers=10)
-    flow_context = FlowContext()
-    service_config = get_default_config()
-    flow_context.query = "茅台和五粮现在价格多少？"
-    flow_context.service_config = service_config
-
-    op = GetAkACodeOp(flow_context=flow_context)
-    # for x in op.split_list(list(range(10)), 3):
-    #     print(x)
-    op.execute()
+    C.set_default_service_config().init_by_service_config()
+    context = FlowContext(query="茅台和五粮现在价格多少？")
+
+    op = GetAkACodeOp()
+    op(context=context)