flowllm 0.1.1__py3-none-any.whl → 0.1.2__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,11 @@
-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

flowllm/op/agent/react_op.py

@@ -0,0 +1,83 @@
+import datetime
+import time
+from typing import List, Dict
+
+from loguru import logger
+
+from flowllm import C, BaseLLMOp
+from flowllm.flow.base_tool_flow import BaseToolFlow
+from flowllm.flow.gallery import DashscopeSearchToolFlow, CodeToolFlow, TerminateToolFlow
+from flowllm.schema.message import Message, Role
+
+
+@C.register_op()
+class ReactOp(BaseLLMOp):
+    # TODO: test react op
+    file_path: str = __file__
+
+    def execute(self):
+        query: str = self.context.query
+
+        max_steps: int = int(self.op_params.get("max_steps", 10))
+        tools: List[BaseToolFlow] = [DashscopeSearchToolFlow(), CodeToolFlow(), TerminateToolFlow()]
+        tool_dict: Dict[str, BaseToolFlow] = {x.name: x for x in tools}
+        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([x.name for x in tools]),
+                                         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:
+                    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}")
+                    assert isinstance(tool_result, str)
+                    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()
+                messages.append(Message(role=Role.USER, content=self.prompt_format(prompt_name="final_prompt")))
+
+        # Store results in context instead of response
+        self.context.messages = messages
+        self.context.answer = messages[-1].content

flowllm/op/agent/react_prompt.yaml

@@ -0,0 +1,28 @@
+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}
+

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)

flowllm/op/akshare/get_ak_a_info_op.py

@@ -4,8 +4,8 @@ import time
 import akshare as ak
 import pandas as pd
 from loguru import logger
+from tqdm import tqdm
 
-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.op.base_op import BaseOp
@@ -24,7 +24,7 @@ class GetAkAInfoOp(BaseOp):
 
     def execute(self):
         max_retries: int = self.op_params.get("max_retries", 3)
-        for code, info_dict in self.flow_context.code_infos.items():
+        for code, info_dict in self.context.code_infos.items():
             result = {}
             for i in range(max_retries):
                 try:
@@ -39,14 +39,14 @@ class GetAkAInfoOp(BaseOp):
                 info_dict.update(result)
 
         time.sleep(1)
-        logger.info(f"code_infos={json.dumps(self.flow_context.code_infos, ensure_ascii=False, indent=2)}")
+        logger.info(f"code_infos={json.dumps(self.context.code_infos, ensure_ascii=False, indent=2)}")
 
 
 @C.register_op()
 class GetAkASpotOp(GetAkAInfoOp):
 
     def execute_code(self, code: str) -> dict:
-        from flowllm.op import GetAkACodeOp
+        from flowllm.op.akshare import GetAkACodeOp
 
         df: pd.DataFrame = GetAkACodeOp.download_a_stock_df()
         df = df.loc[df["代码"] == code, :]
@@ -87,7 +87,7 @@ class GetAkANewsOp(GetAkAInfoOp):
         top_n_news: int = self.op_params.get("top_n_news", 1)
 
         news_content_list = []
-        for i, line in enumerate(stock_news_em_df.to_dict(orient="records")[:top_n_news]):
+        for i, line in enumerate(tqdm(stock_news_em_df.to_dict(orient="records")[:top_n_news])):
             url = line["新闻链接"]
             # http://finance.eastmoney.com/a/202508133482756869.html
             ts = url.split("/")[-1].split(".")[0]
@@ -104,7 +104,7 @@ class MergeAkAInfoOp(BaseOp):
 
     def execute(self):
         code_content = {}
-        for code, info_dict in self.flow_context.code_infos.items():
+        for code, info_dict in self.context.code_infos.items():
             content_list = [f"\n\n### {code}"]
             for key, value in info_dict.items():
                 content_list.append(f"\n#### {code}-{key}")
@@ -120,24 +120,21 @@ class MergeAkAInfoOp(BaseOp):
 
         answer = "\n".join(code_content.values())
         logger.info(f"answer=\n{answer}")
-        self.flow_context.response.answer = answer.strip()
+        self.context.response.answer = answer.strip()
 
 
 if __name__ == "__main__":
-    from flowllm.schema.flow_response import FlowResponse
+    C.set_default_service_config().init_by_service_config()
 
     code_infos = {"000858": {}, "600519": {}}
-    flow_context = FlowContext(code_infos=code_infos, response=FlowResponse())
-    service_config = get_default_config()
-    flow_context.query = "茅台和五粮现在价格多少？"
-    flow_context.service_config = service_config
-
-    op1 = GetAkAInfoOp(flow_context=flow_context)
-    op2 = GetAkASpotOp(flow_context=flow_context)
-    op3 = GetAkAMoneyFlowOp(flow_context=flow_context)
-    op4 = GetAkAFinancialInfoOp(flow_context=flow_context)
-    op5 = GetAkANewsOp(flow_context=flow_context)
-    op6 = MergeAkAInfoOp(flow_context=flow_context)
+    context = FlowContext(code_infos=code_infos, query="茅台和五粮现在价格多少？")
+
+    op1 = GetAkAInfoOp()
+    op2 = GetAkASpotOp()
+    op3 = GetAkAMoneyFlowOp()
+    op4 = GetAkAFinancialInfoOp()
+    op5 = GetAkANewsOp()
+    op6 = MergeAkAInfoOp()
 
     op = op1 >> op2 >> op3 >> op4 >> op5 >> op6
-    op.execute()
+    op(context=context)

flowllm/op/{llm_base_op.py → base_llm_op.py}

@@ -14,26 +14,27 @@ class BaseLLMOp(BaseOp, ABC):
     file_path: str = __file__
 
     def __init__(self,
+                 language: str = "",
                  prompt_path: str = "",
                  llm: str = "default",
                  embedding_model: str = "default",
                  vector_store: str = "default",
                  **kwargs):
-
         super().__init__(**kwargs)
 
+        self.language: str = language or C.language
+        self.prompt_path: Path = Path(prompt_path) if prompt_path else \
+            Path(self.file_path).parent / self.name.replace("_op", "_prompt.yaml")
         self._llm: BaseLLM | str = llm
         self._embedding_model: BaseEmbeddingModel | str = embedding_model
         self._vector_store: BaseVectorStore | str = vector_store
 
-        default_prompt_path: Path = Path(self.file_path).parent / self.name.replace("_op", "_prompt.yaml")
-        self.prompt_path: Path = Path(prompt_path) if prompt_path else default_prompt_path
         self.prompt = PromptHandler(language=self.language).load_prompt_by_file(self.prompt_path)
 
     @property
     def llm(self) -> BaseLLM:
         if isinstance(self._llm, str):
-            llm_config: LLMConfig = self.flow_context.service_config.llm[self._llm]
+            llm_config: LLMConfig = C.service_config.llm[self._llm]
             llm_cls = C.resolve_llm(llm_config.backend)
             self._llm = llm_cls(model_name=llm_config.model_name, **llm_config.params)
 
@@ -43,7 +44,7 @@ class BaseLLMOp(BaseOp, ABC):
     def embedding_model(self) -> BaseEmbeddingModel:
         if isinstance(self._embedding_model, str):
             embedding_model_config: EmbeddingModelConfig = \
-                self.flow_context.service_config.embedding_model[self._embedding_model]
+                C.service_config.embedding_model[self._embedding_model]
             embedding_model_cls = C.resolve_embedding_model(embedding_model_config.backend)
             self._embedding_model = embedding_model_cls(model_name=embedding_model_config.model_name,
                                                         **embedding_model_config.params)