feat: add max_continue_runs (#685)

Co-authored-by: luyan <suluyan.sly@aliabab-inc.com>

Author: suluyana

.github/workflows/citest.yaml

@@ -66,6 +66,8 @@ jobs:
     - name: Run tests
       env:
         OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+        DASHSCOPE_API_KEY: ${{ secrets.DASHSCOPE_API_KEY }}
+        MODELSCOPE_API_KEY: ${{ secrets.MODELSCOPE_API_KEY }}
 
       shell: bash
       run: bash .dev_scripts/dockerci.sh

ms_agent/llm/openai_llm.py

@@ -12,6 +12,8 @@
 
 logger = get_logger()
 
+MAX_CONTINUE_RUNS = 3
+
 
 class OpenAI(LLM):
     """Base Class for OpenAI SDK LLMs.
@@ -28,14 +30,18 @@ class OpenAI(LLM):
         'role', 'content', 'tool_calls', 'partial', 'prefix', 'tool_call_id'
     }
 
-    def __init__(self,
-                 config: DictConfig,
-                 base_url: Optional[str] = None,
-                 api_key: Optional[str] = None):
+    def __init__(
+        self,
+        config: DictConfig,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+    ):
         super().__init__(config)
         assert_package_exist('openai')
         import openai
         self.model: str = config.llm.model
+        self.max_continue_runs = getattr(config.llm, 'max_continue_runs',
+                                         None) or MAX_CONTINUE_RUNS
         base_url = base_url or config.llm.openai_base_url
         api_key = api_key or config.llm.openai_api_key
 
@@ -76,6 +82,7 @@ def format_tools(self,
     def generate(self,
                  messages: List[Message],
                  tools: Optional[List[Tool]] = None,
+                 max_continue_runs: Optional[int] = None,
                  **kwargs) -> Message | Generator[Message, None, None]:
         """Generates a response based on the given conversation history and optional tools.
 
@@ -99,11 +106,14 @@ def generate(self,
 
         # Complex task may produce long response
         # Call continue_generate to keep generating if the finish_reason is `length`
+        max_continue_runs = max_continue_runs or self.max_continue_runs
         if stream:
             return self._stream_continue_generate(messages, completion, tools,
+                                                  max_continue_runs - 1,
                                                   **args)
         else:
-            return self._continue_generate(messages, completion, tools, **args)
+            return self._continue_generate(messages, completion, tools,
+                                           max_continue_runs - 1, **args)
 
     def _call_llm(self,
                   messages: List[Message],
@@ -180,6 +190,7 @@ def _stream_continue_generate(self,
                                   messages: List[Message],
                                   completion: Iterable,
                                   tools: Optional[List[Tool]] = None,
+                                  max_runs: Optional[int] = None,
                                   **kwargs) -> Generator[Message, None, None]:
         """Recursively continues generating until the model finishes naturally in streaming mode.
 
@@ -193,7 +204,6 @@ def _stream_continue_generate(self,
             Message: Incremental chunks of the generated message.
         """
         message = None
-
         for chunk in completion:
             message_chunk = self._stream_format_output_message(chunk)
             message = self._merge_stream_message(message, message_chunk)
@@ -208,15 +218,18 @@ def _stream_continue_generate(self,
                     # The stream may end without a final usage chunk, which is acceptable.
                     pass
                 first_run = not messages[-1].to_dict().get('partial', False)
-                if chunk.choices[0].finish_reason in ['length', 'null']:
-                    print(
-                        f'finish_reason: {chunk.choices[0].finish_reason}， continue generate.'
+                if chunk.choices[0].finish_reason in [
+                        'length', 'null'
+                ] and (max_runs is None or max_runs != 0):
+                    logger.info(
+                        f'finish_reason: {chunk.choices[0].finish_reason}, continue generate.'
                     )
-
                     completion = self._call_llm_for_continue_gen(
                         messages, message, tools, **kwargs)
                     for chunk in self._stream_continue_generate(
-                            messages, completion, tools, **kwargs):
+                            messages, completion, tools,
+                            max_runs - 1 if max_runs is not None else None,
+                            **kwargs):
                         if first_run:
                             yield self._merge_stream_message(
                                 messages[-1], chunk)
@@ -265,8 +278,9 @@ def _stream_format_output_message(completion_chunk) -> Message:
             reasoning_content=reasoning_content,
             tool_calls=tool_calls,
             id=completion_chunk.id,
-            prompt_tokens=completion_chunk.usage.prompt_tokens,
-            completion_tokens=completion_chunk.usage.completion_tokens)
+            prompt_tokens=getattr(completion_chunk.usage, 'prompt_tokens', 0),
+            completion_tokens=getattr(completion_chunk.usage,
+                                      'completion_tokens', 0))
 
     @staticmethod
     def _format_output_message(completion) -> Message:
@@ -359,6 +373,7 @@ def _continue_generate(self,
                            messages: List[Message],
                            completion,
                            tools: List[Tool] = None,
+                           max_runs: Optional[int] = None,
                            **kwargs) -> Message:
         """Recursively continues generating until the model finishes naturally.
 
@@ -375,14 +390,17 @@ def _continue_generate(self,
             Message: A fully formed Message object containing the complete response.
         """
         new_message = self._format_output_message(completion)
-        if completion.choices[0].finish_reason in ['length', 'null']:
+        if completion.choices[0].finish_reason in [
+                'length', 'null'
+        ] and (max_runs is None or max_runs != 0):
             logger.info(
                 f'finish_reason: {completion.choices[0].finish_reason}， continue generate.'
             )
             completion = self._call_llm_for_continue_gen(
                 messages, new_message, tools, **kwargs)
-            return self._continue_generate(messages, completion, tools,
-                                           **kwargs)
+            return self._continue_generate(
+                messages, completion, tools,
+                max_runs - 1 if max_runs is not None else None, **kwargs)
         elif messages[-1].to_dict().get('partial', False):
             self._merge_partial_message(messages, new_message)
             messages[-1].partial = False

tests/llm/test_openai.py

@@ -3,19 +3,22 @@
 import os
 import unittest
 
-from ms_agent.llm.openai_llm import OpenAI
-from ms_agent.llm.utils import Message, Tool, ToolCall
+from ms_agent.llm.openai_llm import MAX_CONTINUE_RUNS, OpenAI
+from ms_agent.llm.utils import Message, Tool
 from omegaconf import DictConfig, OmegaConf
 
+from modelscope.utils.test_utils import test_level
+
 API_CALL_MAX_TOKEN = 50
 
 
 class OpenaiLLM(unittest.TestCase):
     conf: DictConfig = OmegaConf.create({
         'llm': {
-            'model': 'Qwen/Qwen3-235B-A22B',
-            'openai_base_url': 'https://api-inference.modelscope.cn/v1',
-            'openai_api_key': os.getenv('MODELSCOPE_API_KEY'),
+            'model': 'qwen3-235b-a22b',
+            'openai_base_url':
+            'https://dashscope.aliyuncs.com/compatible-mode/v1',
+            'openai_api_key': os.getenv('DASHSCOPE_API_KEY'),
         },
         'generation_config': {
             'stream': False,
@@ -68,62 +71,73 @@ class OpenaiLLM(unittest.TestCase):
             })
     ]
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_no_stream(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.messages, tools=None)
         print(res)
         assert (res.content)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_stream(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.messages, tools=None, stream=True)
         for chunk in res:
             print(chunk)
         assert (len(chunk.content))
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_thinking(self):
         llm = OpenAI(self.conf)
         res = llm.generate(
             messages=self.messages,
             tools=None,
+            max_continue_runs=1,
             stream=True,
             extra_body={'enable_thinking': True})
         for chunk in res:
             print(chunk)
         assert (chunk.reasoning_content)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_continue_run(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.continue_messages, tools=None)
         print(res)
         assert (res.completion_tokens > 100)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_tool(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.tool_messages, tools=self.tools)
         print(res)
         assert (len(res.tool_calls))
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_apis_count(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.messages, tools=None)
         print(res)
         assert res.api_calls == 1
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_apis_count_stream(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.messages, stream=True, tools=None)
         for chunk in res:
             print(chunk)
         assert chunk.api_calls == 1
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_apis_count_continue(self):
         llm = OpenAI(self.conf)
         res = llm.generate(messages=self.continue_messages, tools=None)
         print(res)
         assert math.ceil(res.completion_tokens
                          / API_CALL_MAX_TOKEN) == res.api_calls
+        assert res.api_calls <= MAX_CONTINUE_RUNS
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_apis_count_continue_stream(self):
         llm = OpenAI(self.conf)
         res = llm.generate(
@@ -132,7 +146,9 @@ def test_call_apis_count_continue_stream(self):
             print(chunk)
         assert math.ceil(chunk.completion_tokens
                          / API_CALL_MAX_TOKEN) == chunk.api_calls
+        assert chunk.api_calls <= MAX_CONTINUE_RUNS
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_call_tool_stream(self):
         llm = OpenAI(self.conf)
         res = llm.generate(