add doc and more tests

modelscope · May 9, 2024 · bebfba7 · bebfba7
1 parent dafc60e
commit bebfba7
Show file tree

Hide file tree

Showing 7 changed files with 96 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -87,6 +87,7 @@ services and third-party model APIs.
 | ollama                 | Chat            | [`OllamaChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_chat_template.json)       | llama3, llama2, Mistral, ...                  |
 |                        | Embedding       | [`OllamaEmbeddingWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)            | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_embedding_template.json)       | llama2, Mistral, ...                          |
 |                        | Generation      | [`OllamaGenerationWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)           | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_generate_template.json)       | llama2, Mistral, ...                          |
+| LiteLLM API             | Chat            | [`LiteLLMChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#litellm-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/litellm_chat_template.json)       | ...                               |
 | Post Request based API | -               | [`PostAPIModelWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#post-request-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/postapi_model_config_template.json) | -                                             |
 
 **Supported Local Model Deployment**

diff --git a/README_ZH.md b/README_ZH.md
@@ -76,6 +76,7 @@ AgentScope提供了一系列`ModelWrapper`来支持本地模型服务和第三
 | ollama                 | Chat            | [`OllamaChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_chat_template.json)       | llama3, llama2, Mistral, ...                  |
 |                        | Embedding       | [`OllamaEmbeddingWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)            | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_embedding_template.json)       | llama2, Mistral, ...                          |
 |                        | Generation      | [`OllamaGenerationWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)           | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#ollama-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/ollama_generate_template.json)       | llama2, Mistral, ...                          |
+| LiteLLM API             | Chat            | [`LiteLLMChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#litellm-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/litellm_chat_template.json)       | ...                               |
 | Post Request based API | -               | [`PostAPIModelWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py)                 | [guidance](https://modelscope.github.io/agentscope/en/tutorial/203-model.html#post-request-api) <br> [template](https://github.com/modelscope/agentscope/blob/main/examples/model_configs_template/postapi_model_config_template.json) | -                                             |
 
 **支持的本地模型部署**

diff --git a/docs/sphinx_doc/en/source/tutorial/203-model.md b/docs/sphinx_doc/en/source/tutorial/203-model.md
@@ -16,6 +16,7 @@ Currently, AgentScope supports the following model service APIs:
 - Gemini API, including chat and embedding.
 - ZhipuAI API, including chat and embedding.
 - Ollama API, including chat, embedding and generation.
+- LiteLLM API, including chat, with various model APIs.
 - Post Request API, model inference services based on Post
   requests, including Huggingface/ModelScope Inference API and various
   post request based model APIs.
@@ -87,6 +88,7 @@ In the current AgentScope, the supported `model_type` types, the corresponding
 | ollama                 | Chat            | [`OllamaChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)                 | `"ollama_chat"`               | llama2, ...                                      |
 |                        | Embedding       | [`OllamaEmbeddingWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)            | `"ollama_embedding"`          | llama2, ...                                      |
 |                        | Generation      | [`OllamaGenerationWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)           | `"ollama_generate"`           | llama2, ...                                      |
+| LiteLLM API | Chat               | [`LiteLLMModelWrapperChat`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py)             | `"litellm_chat"`                  | -                                                |
 | Post Request based API | -               | [`PostAPIModelWrapperBase`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py)             | `"post_api"`                  | -                                                |
 |  | Chat | [`PostAPIChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py) | `"post_api_chat"` | meta-llama/Meta-Llama-3-8B-Instruct, ... |
 
@@ -440,6 +442,28 @@ Here we provide example configurations for different model wrappers.
 
 <br/>
 
+
+#### LiteLLM Chat API
+
+<details>
+<summary>LiteLLM Chat API (<code><a href="https://github.
+com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py">agentscope.models.LiteLLMChatModelWrapper</a></code>)</summary>
+
+```python
+{
+    "config_name": "lite_llm_openai_chat_gpt-3.5-turbo",
+    "model_type": "litellm_chat",
+    "model_name": "gpt-3.5-turbo",
+    "api_key": "{your_api_key}",
+    "api_key_name": "OPENAI_API_KEY"
+},
+```
+
+</details>
+
+<br/>
+
+
 #### Post Request Chat API
 
 <details>

diff --git a/docs/sphinx_doc/zh_CN/source/tutorial/203-model.md b/docs/sphinx_doc/zh_CN/source/tutorial/203-model.md
@@ -13,6 +13,7 @@ AgentScope中，模型的部署和调用是通过`ModelWrapper`来解耦开的
 - Gemini API，包括对话（Chat）和嵌入（Embedding）。
 - ZhipuAi API，包括对话（Chat）和嵌入（Embedding）。
 - Ollama API，包括对话（Chat），嵌入（Embedding）和生成（Generation）。
+- LiteLLM API, 包括对话（Chat）, 支持各种模型的API.
 - Post请求API，基于Post请求实现的模型推理服务，包括Huggingface/ModelScope
   Inference API和各种符合Post请求格式的API。
 
@@ -107,6 +108,7 @@ API如下：
 | ollama                 | Chat            | [`OllamaChatWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)                 | `"ollama_chat"`               | llama2, ...                                      |
 |                        | Embedding       | [`OllamaEmbeddingWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)            | `"ollama_embedding"`          | llama2, ...                                      |
 |                        | Generation      | [`OllamaGenerationWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/ollama_model.py)           | `"ollama_generate"`           | llama2, ...                                      |
+| LiteLLM API | Chat               | [`LiteLLMModelWrapperChat`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py)             | `"litellm_chat"`                  | -                                                |
 | Post Request based API | -               | [`PostAPIModelWrapperBase`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py)             | `"post_api"`                  | -                                                |
 |  | Chat | [`PostAPIChatModelWrapper`](https://github.com/modelscope/agentscope/blob/main/src/agentscope/models/post_model.py) | `"post_api_chat"` | meta-llama/Meta-Llama-3-8B-Instruct, ... |
 
@@ -435,6 +437,28 @@ API如下：
 
 <br/>
 
+
+#### LiteLLM Chat API
+
+<details>
+<summary>LiteLLM Chat API (<code><a href="https://github.
+com/modelscope/agentscope/blob/main/src/agentscope/models/litellm_model.py">agentscope.models.LiteLLMChatModelWrapper</a></code>)</summary>
+
+```python
+{
+    "config_name": "lite_llm_openai_chat_gpt-3.5-turbo",
+    "model_type": "litellm_chat",
+    "model_name": "gpt-3.5-turbo",
+    "api_key": "{your_api_key}",
+    "api_key_name": "OPENAI_API_KEY"
+},
+```
+
+</details>
+
+<br/>
+
+
 #### Post Request API
 
 <details>

diff --git a/examples/model_configs_template/litellm_chat_template.json b/examples/model_configs_template/litellm_chat_template.json
@@ -0,0 +1,15 @@
+[{
+    "config_name": "lite_llm_openai_chat_gpt-3.5-turbo",
+    "model_type": "litellm_chat",
+    "model_name": "gpt-3.5-turbo",
+    "api_key": "{your_api_key}",
+    "api_key_name": "OPENAI_API_KEY"
+},
+{
+    "config_name": "lite_llm_claude3",
+    "model_type": "litellm_chat",
+    "model_name": "claude-3-opus-20240229",
+    "api_key": "{your_api_key}",
+    "api_key_name": "ANTHROPIC_API_KEY"
+}
+]
diff --git a/src/agentscope/models/__init__.py b/src/agentscope/models/__init__.py
@@ -37,6 +37,9 @@
     ZhipuAIChatWrapper,
     ZhipuAIEmbeddingWrapper,
 )
+from .litellm_model import (
+    LiteLLMChatWrapper,
+)
 
 
 __all__ = [
@@ -59,6 +62,7 @@
     "GeminiEmbeddingWrapper",
     "ZhipuAIChatWrapper",
     "ZhipuAIEmbeddingWrapper",
+    "LiteLLMChatWrapper",
     "load_model_by_config_name",
     "read_model_configs",
     "clear_model_configs",

diff --git a/tests/format_test.py b/tests/format_test.py
@@ -12,6 +12,7 @@
     ZhipuAIChatWrapper,
     DashScopeChatWrapper,
     DashScopeMultiModalWrapper,
+    LiteLLMChatWrapper,
 )
 
 
@@ -204,6 +205,32 @@ def test_zhipuai_chat(self) -> None:
         with self.assertRaises(TypeError):
             model.format(*self.wrong_inputs)  # type: ignore[arg-type]
 
+    def test_litellm_chat(self) -> None:
+        """Unit test for the format function in litellm chat api wrapper."""
+        model = LiteLLMChatWrapper(
+            config_name="",
+            model_name="gpt-3.5-turbo",
+            api_key="xxx",
+        )
+
+        ground_truth = [
+            {
+                "role": "user",
+                "content": (
+                    "You are a helpful assistant\n\n"
+                    "## Dialogue History\nuser: What is the weather today?\n"
+                    "assistant: It is sunny today"
+                ),
+            },
+        ]
+
+        prompt = model.format(*self.inputs)
+        self.assertListEqual(prompt, ground_truth)
+
+        # wrong format
+        with self.assertRaises(TypeError):
+            model.format(*self.wrong_inputs)  # type: ignore[arg-type]
+
     def test_dashscope_multimodal_image(self) -> None:
         """Unit test for the format function in dashscope multimodal
         conversation api wrapper for image."""