add tests

config/examples/huoshan_ark.yaml

@@ -1,8 +1,5 @@
-# Full Example: https://github.com/geekan/MetaGPT/blob/main/config/config2.example.yaml
-# Reflected Code: https://github.com/geekan/MetaGPT/blob/main/metagpt/config2.py
-# Config Docs: https://docs.deepwisdom.ai/main/en/guide/get_started/configuration.html
 llm:
-  api_type: "ark"  # or azure / ollama / groq etc.
+  api_type: "ark"  
   model: ""  # your model endpoint like ep-xxx
-  base_url: "https://ark.cn-beijing.volces.com/api/v3"  # or forward url / other llm url
+  base_url: "https://ark.cn-beijing.volces.com/api/v3"  
   api_key: "" # your api-key like ey……

metagpt/provider/ark_api.py

@@ -20,7 +20,7 @@ class ArkLLM(OpenAILLM):
         response: AsyncStream[ChatCompletionChunk] = await self.aclient.chat.completions.create(
             **self._cons_kwargs(messages, timeout=self.get_timeout(timeout)),
             stream=True,
-            extra_body={"stream_options": {"include_usage": True}}
+            extra_body={"stream_options": {"include_usage": True}}  # 只有增加这个参数才会在流式时最后返回usage
         )
         usage = None
         collected_messages = []
@@ -29,7 +29,7 @@ class ArkLLM(OpenAILLM):
             log_llm_stream(chunk_message)
             collected_messages.append(chunk_message)
             if chunk.usage:
-                # the usage of ark when streaming is in the last chunk while others are null
+                # 火山方舟的流式调用会在最后一个chunk中返回usage,最后一个chunk的choices为[]
                 usage = CompletionUsage(**chunk.usage)
 
         log_llm_stream("\n")

metagpt/utils/token_counter.py

@@ -68,13 +68,15 @@ TOKEN_COSTS = {
     "openai/gpt-4-turbo-preview": {"prompt": 0.01, "completion": 0.03},
     "deepseek-chat": {"prompt": 0.00014, "completion": 0.00028},
     "deepseek-coder": {"prompt": 0.00014, "completion": 0.00028},
-    # For ark model
+    # For ark model https://www.volcengine.com/docs/82379/1099320
     "doubao-lite-4k-240515": {"prompt": 0.000042, "completion": 0.000084},
     "doubao-lite-32k-240515": {"prompt": 0.000042, "completion": 0.000084},
     "doubao-lite-128k-240515": {"prompt": 0.00011, "completion": 0.00013},
     "doubao-pro-4k-240515": {"prompt": 0.00011, "completion": 0.00028},
     "doubao-pro-32k-240515": {"prompt": 0.00011, "completion": 0.00028},
     "doubao-pro-128k-240515": {"prompt": 0.0007, "completion": 0.0012},
+    "llama3-70b-llama3-70b-instruct": {"prompt": 0.0, "completion": 0.0},
+    "llama3-8b-llama3-8b-instruct": {"prompt": 0.0, "completion": 0.0},
 }
 
 

tests/metagpt/provider/mock_llm_config.py

@@ -69,3 +69,5 @@ mock_llm_config_bedrock = LLMConfig(
     secret_key="123abc",
     max_token=10000,
 )
+
+mock_llm_config_ark = LLMConfig(api_type="ark", api_key="eyxxx", base_url="xxx", model="ep-xxx")

tests/metagpt/provider/test_ark.py

@@ -0,0 +1,85 @@
+"""
+用于火山方舟Python SDK V3的测试用例
+API文档：https://www.volcengine.com/docs/82379/1263482
+"""
+
+from typing import AsyncIterator, List, Union
+
+import pytest
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta
+
+from metagpt.provider.ark_api import ArkLLM
+from tests.metagpt.provider.mock_llm_config import mock_llm_config
+from tests.metagpt.provider.req_resp_const import (
+    get_openai_chat_completion,
+    llm_general_chat_funcs_test,
+    messages,
+    prompt,
+    resp_cont_tmpl,
+)
+
+name = "AI assistant"
+resp_cont = resp_cont_tmpl.format(name=name)
+USAGE = {"completion_tokens": 1000, "prompt_tokens": 1000, "total_tokens": 2000}
+default_resp = get_openai_chat_completion(name)
+default_resp.model = "doubao-pro-32k-240515"
+default_resp.usage = USAGE
+
+
+def create_chat_completion_chunk(
+    content: str, finish_reason: str = None, choices: List[Choice] = None
+) -> ChatCompletionChunk:
+    if choices is None:
+        choices = [
+            Choice(
+                delta=ChoiceDelta(content=content, function_call=None, role="assistant", tool_calls=None),
+                finish_reason=finish_reason,
+                index=0,
+                logprobs=None,
+            )
+        ]
+
+    return ChatCompletionChunk(
+        id="012",
+        choices=choices,
+        created=1716278586,
+        model="doubao-pro-32k-240515",
+        object="chat.completion.chunk",
+        system_fingerprint=None,
+        usage=None if choices else USAGE,
+    )
+
+
+ark_resp_chunk = create_chat_completion_chunk(content="")
+ark_resp_chunk_finish = create_chat_completion_chunk(content=resp_cont, finish_reason="stop")
+ark_resp_chunk_last = create_chat_completion_chunk(content="", choices=[])
+
+
+async def chunk_iterator(chunks: List[ChatCompletionChunk]) -> AsyncIterator[ChatCompletionChunk]:
+    for chunk in chunks:
+        yield chunk
+
+
+async def mock_ark_acompletions_create(
+    self, stream: bool = False, **kwargs
+) -> Union[ChatCompletionChunk, ChatCompletion]:
+    if stream:
+        chunks = [ark_resp_chunk, ark_resp_chunk_finish, ark_resp_chunk_last]
+        return chunk_iterator(chunks)
+    else:
+        return default_resp
+
+
+@pytest.mark.asyncio
+async def test_ark_acompletion(mocker):
+    mocker.patch("openai.resources.chat.completions.AsyncCompletions.create", mock_ark_acompletions_create)
+
+    llm = ArkLLM(mock_llm_config)
+
+    resp = await llm.acompletion(messages)
+    assert resp.choices[0].finish_reason == "stop"
+    assert resp.choices[0].message.content == resp_cont
+    assert resp.usage == USAGE
+
+    await llm_general_chat_funcs_test(llm, prompt, messages, resp_cont)