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

flowllm/llm/openai_compatible_llm.py

@@ -0,0 +1,439 @@
+import os
+from typing import List, Dict
+
+from loguru import logger
+from openai import OpenAI, AsyncOpenAI
+from openai.types import CompletionUsage
+from pydantic import Field, PrivateAttr, model_validator
+
+from flowllm.context.service_context import C
+from flowllm.enumeration.chunk_enum import ChunkEnum
+from flowllm.enumeration.role import Role
+from flowllm.llm.base_llm import BaseLLM
+from flowllm.schema.message import Message
+from flowllm.schema.tool_call import ToolCall
+
+
+@C.register_llm("openai_compatible")
+class OpenAICompatibleBaseLLM(BaseLLM):
+    """
+    OpenAI-compatible LLM implementation supporting streaming and tool calls.
+    
+    This class implements the BaseLLM interface for OpenAI-compatible APIs,
+    including support for:
+    - Streaming responses with different chunk types (thinking, answer, tools)
+    - Tool calling with parallel execution
+    - Reasoning/thinking content from supported models
+    - Robust error handling and retries
+    """
+
+    # API configuration
+    api_key: str = Field(default_factory=lambda: os.getenv("FLOW_LLM_API_KEY"),
+                         description="API key for authentication")
+    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 clients after model validation.
+        
+        This validator runs after all field validation is complete,
+        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):
+        """
+        Stream chat completions from OpenAI-compatible API.
+        
+        This method handles 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 = self._client.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
+                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"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
+
+    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:
+        """
+        Perform a complete chat completion by aggregating streaming chunks.
+        
+        This method consumes the entire 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 streaming response and aggregate chunks by type
+        for chunk, chunk_enum in self.stream_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 _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
+
+    load_env()
+
+    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("Sync result:", message)
+
+
+if __name__ == "__main__":
+    # main()
+
+    import asyncio
+    asyncio.run(async_main())

flowllm/op/__init__.py

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

llmflow/op/react/react_v1_op.py → flowllm/op/agent/react_op.py

@@ -4,36 +4,30 @@ from typing import List, Dict
 
 from loguru import logger
 
-from llmflow.enumeration.role import Role
-from llmflow.op import OP_REGISTRY
-from llmflow.op.base_op import BaseOp
-from llmflow.schema.message import Message
-from llmflow.schema.request import AgentRequest
-from llmflow.schema.response import AgentResponse
-from llmflow.tool import TOOL_REGISTRY
-from llmflow.tool.base_tool import BaseTool
+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
 
 
-@OP_REGISTRY.register()
-class ReactV1Op(BaseOp):
-    current_path: str = __file__
+@C.register_op()
+class ReactOp(BaseLLMOp):
+    # TODO: test react op
+    file_path: str = __file__
 
     def execute(self):
-        request: AgentRequest = self.context.request
-        response: AgentResponse = self.context.response
+        query: str = self.context.query
 
         max_steps: int = int(self.op_params.get("max_steps", 10))
-        # dashscope_search_tool tavily_search_tool
-        tool_names = self.op_params.get("tool_names", "code_tool,tavily_search_tool,terminate_tool")
-        tools: List[BaseTool] = [TOOL_REGISTRY[x.strip()]() for x in tool_names.split(",") if x]
-        tool_dict: Dict[str, BaseTool] = {x.name: x for x in tools}
+        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=request.query)
+                                         query=query)
         messages: List[Message] = [Message(role=Role.USER, content=user_prompt)]
         logger.info(f"step.0 user_prompt={user_prompt}")
 
@@ -41,7 +35,7 @@ class ReactV1Op(BaseOp):
             if has_terminate_tool:
                 assistant_message: Message = self.llm.chat(messages)
             else:
-                assistant_message: Message = self.llm.chat(messages, tools=tools)
+                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"
@@ -67,7 +61,7 @@ class ReactV1Op(BaseOp):
                 if tool_call.name not in tool_dict:
                     continue
 
-                self.submit_task(tool_dict[tool_call.name].execute, **tool_call.argument_dict)
+                self.submit_task(tool_dict[tool_call.name].__call__, **tool_call.argument_dict)
                 time.sleep(1)
 
             if not has_terminate_tool:
@@ -84,5 +78,6 @@ class ReactV1Op(BaseOp):
                 assistant_message.tool_calls.clear()
                 messages.append(Message(role=Role.USER, content=self.prompt_format(prompt_name="final_prompt")))
 
-        response.messages = messages
-        response.answer = response.messages[-1].content
+        # Store results in context instead of response
+        self.context.messages = messages
+        self.context.answer = messages[-1].content

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

@@ -0,0 +1,108 @@
+import json
+import time
+from typing import List
+
+import akshare as ak
+import pandas as pd
+from loguru import logger
+
+from flowllm.context.flow_context import FlowContext
+from flowllm.context.service_context import C
+from flowllm.enumeration.role import Role
+from flowllm.op.base_llm_op import BaseLLMOp
+from flowllm.schema.message import Message
+from flowllm.storage.cache.data_cache import DataCache
+from flowllm.utils.timer import timer
+
+
+@C.register_op()
+class GetAkACodeOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def __init__(self, language: str = "zh", llm="qwen3_30b_instruct", **kwargs):
+        super().__init__(language=language, llm=llm, **kwargs)
+
+    @staticmethod
+    def download_a_stock_df():
+        cache = DataCache()
+        save_df_key: str = "all_a_stock_name_code"
+        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()
+
+            df: pd.DataFrame = pd.concat([stock_sh_a_spot_em_df, stock_sz_a_spot_em_df, stock_bj_a_spot_em_df], axis=0)
+            df = df.drop(columns=["序号"])
+            df = df.reset_index(drop=True)
+            df = df.sort_values(by="代码")
+            cache.save(save_df_key, df, expire_hours=0.25)
+
+        df = cache.load(save_df_key, dtype={"代码": str})
+        return df
+
+    def get_name_code_dict(self) -> dict:
+        df = self.download_a_stock_df()
+
+        name_code_dict = {}
+        for line in df.to_dict(orient="records"):
+            name = line["名称"].replace(" ", "")
+            code = line["代码"]
+            name_code_dict[name] = code
+        logger.info(f"name_code_dict.size={len(name_code_dict)} content={str(name_code_dict)[:50]}...")
+        return name_code_dict
+
+    @staticmethod
+    def split_list(array_list: list, n: int):
+        if n <= 0:
+            raise ValueError
+
+        length = len(array_list)
+        base_size = length // n
+        remainder = length % n
+
+        start = 0
+        for i in range(n):
+            size = base_size + (1 if i < remainder else 0)
+            end = start + size
+            yield array_list[start:end]
+            start = end
+
+    @timer()
+    def find_stock_codes(self, stock_names: List[str]):
+        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.context.query)
+        logger.info(f"prompt={prompt}")
+
+        def callback_fn(msg: Message):
+            content = msg.content
+            if "```" in content:
+                content = content.split("```")[1]
+                content = content.strip("json")
+            content = json.loads(content.strip())
+            return content
+
+        codes: List[str] = self.llm.chat(messages=[Message(role=Role.USER, content=prompt)],
+                                         enable_stream_print=False,
+                                         callback_fn=callback_fn)
+        return codes
+
+    def execute(self):
+        name_code_dict = self.get_name_code_dict()
+        stock_names = list(name_code_dict.keys())
+        for p_stock_names in self.split_list(stock_names, n=2):
+            self.submit_task(self.find_stock_codes, stock_names=p_stock_names)
+            time.sleep(1)
+
+        stock_names = sorted(set(self.join_task()))
+        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__":
+    C.set_default_service_config().init_by_service_config()
+    context = FlowContext(query="茅台和五粮现在价格多少？")
+
+    op = GetAkACodeOp()
+    op(context=context)

flowllm/op/akshare/get_ak_a_code_prompt.yaml

@@ -0,0 +1,21 @@
+find_stock_name_zh: |
+  # 股票标准名称
+  {stock_names}
+
+  # 用户问题
+  {query}
+
+  # 任务
+  请你提取出用户问题中提到的股票名称，并通过**股票标准名称**找到对应的标准名称，并以json格式返回。
+  如果用户问题中没有提及股票名称或者**股票标准名称**中没有找到，请返回空"[]"。
+  如果用户只是提到了某个行业，并且没有明确提及股票名称，请返回空"[]"。
+  请思考后输出你的答案。
+  
+  # 答案格式
+  ```json
+  [
+      "股票标准名称1",
+      "股票标准名称2",
+      ...
+  ]
+  ```