diff --git a/metagpt/provider/metagpt_llm_api.py b/metagpt/provider/metagpt_llm_api.py new file mode 100644 index 000000000..bfd003fff --- /dev/null +++ b/metagpt/provider/metagpt_llm_api.py @@ -0,0 +1,33 @@ +# -*- coding: utf-8 -*- +""" +@Time : 2023/8/30 +@Author : mashenquan +@File : metagpt_llm_api.py +@Desc : MetaGPT LLM related APIs +""" + +import openai + +from metagpt.config import CONFIG +from metagpt.provider import OpenAIGPTAPI +from metagpt.provider.openai_api import RateLimiter + + +class MetaGPTLLMAPI(OpenAIGPTAPI): + """MetaGPT LLM api""" + + def __init__(self): + self.__init_openai(CONFIG) + self.llm = openai + self.model = CONFIG.METAGPT_API_MODEL + self.auto_max_tokens = False + RateLimiter.__init__(self, rpm=self.rpm) + + def __init_openai(self, config): + openai.api_key = CONFIG.METAGPT_API_KEY + if config.openai_api_base: + openai.api_base = CONFIG.METAGPT_API_BASE + if config.openai_api_type: + openai.api_type = CONFIG.METAGPT_API_TYPE + openai.api_version = CONFIG.METAGPT_API_VERSION + self.rpm = int(config.get("RPM", 10)) diff --git a/tests/metagpt/provider/test_metagpt_llm_api.py b/tests/metagpt/provider/test_metagpt_llm_api.py new file mode 100644 index 000000000..9c8356ca6 --- /dev/null +++ b/tests/metagpt/provider/test_metagpt_llm_api.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +@Time : 2023/8/30 +@Author : mashenquan +@File : test_metagpt_llm_api.py +""" +from metagpt.provider.metagpt_llm_api import MetaGPTLLMAPI + + +def test_metagpt(): + llm = MetaGPTLLMAPI() + assert llm + + +if __name__ == "__main__": + test_metagpt()