Merge pull request #603 from better629/feat_ollama

feat: add ollama api support

Author: geekan

config/config.yaml

@@ -48,6 +48,10 @@ RPM: 10
 #FIREWORKS_API_BASE: "https://api.fireworks.ai/inference/v1"
 #FIREWORKS_API_MODEL: "YOUR_LLM_MODEL"  # example, accounts/fireworks/models/llama-v2-13b-chat
 
+#### if use self-host open llm model by ollama
+# OLLAMA_API_BASE: http://127.0.0.1:11434/api
+# OLLAMA_API_MODEL: llama2
+
 #### for Search
 
 ## Supported values: serpapi/google/serper/ddg

metagpt/config.py

@@ -42,6 +42,7 @@ class LLMProviderEnum(Enum):
     FIREWORKS = "fireworks"
     OPEN_LLM = "open_llm"
     GEMINI = "gemini"
+    OLLAMA = "ollama"
 
 
 class Config(metaclass=Singleton):
@@ -78,7 +79,8 @@ def get_default_llm_provider_enum(self) -> LLMProviderEnum:
             (self.zhipuai_api_key, LLMProviderEnum.ZHIPUAI),
             (self.fireworks_api_key, LLMProviderEnum.FIREWORKS),
             (self.open_llm_api_base, LLMProviderEnum.OPEN_LLM),
-            (self.gemini_api_key, LLMProviderEnum.GEMINI),  # reuse logic. but not a key
+            (self.gemini_api_key, LLMProviderEnum.GEMINI),
+            (self.ollama_api_base, LLMProviderEnum.OLLAMA),  # reuse logic. but not a key
         ]:
             if self._is_valid_llm_key(k):
                 # logger.debug(f"Use LLMProvider: {v.value}")
@@ -103,6 +105,8 @@ def _update(self):
         self.open_llm_api_model = self._get("OPEN_LLM_API_MODEL")
         self.fireworks_api_key = self._get("FIREWORKS_API_KEY")
         self.gemini_api_key = self._get("GEMINI_API_KEY")
+        self.ollama_api_base = self._get("OLLAMA_API_BASE")
+        self.ollama_api_model = self._get("OLLAMA_API_MODEL")
         _ = self.get_default_llm_provider_enum()
 
         self.openai_base_url = self._get("OPENAI_BASE_URL")

metagpt/const.py

@@ -102,3 +102,5 @@ def get_metagpt_root():
 CODE_SUMMARIES_PDF_FILE_REPO = "resources/code_summaries"
 
 YAPI_URL = "http://yapi.deepwisdomai.com/"
+
+LLM_API_TIMEOUT = 300

metagpt/provider/__init__.py

@@ -8,8 +8,9 @@
 
 from metagpt.provider.fireworks_api import FireWorksGPTAPI
 from metagpt.provider.google_gemini_api import GeminiGPTAPI
+from metagpt.provider.ollama_api import OllamaGPTAPI
 from metagpt.provider.open_llm_api import OpenLLMGPTAPI
 from metagpt.provider.openai_api import OpenAIGPTAPI
 from metagpt.provider.zhipuai_api import ZhiPuAIGPTAPI
 
-__all__ = ["FireWorksGPTAPI", "GeminiGPTAPI", "OpenLLMGPTAPI", "OpenAIGPTAPI", "ZhiPuAIGPTAPI"]
+__all__ = ["FireWorksGPTAPI", "GeminiGPTAPI", "OpenLLMGPTAPI", "OpenAIGPTAPI", "ZhiPuAIGPTAPI", "OllamaGPTAPI"]

metagpt/provider/general_api_base.py

@@ -1,3 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Desc   : refs to openai 0.x sdk
+
 import asyncio
 import json
 import os
@@ -43,8 +47,8 @@
 # Has one attribute per thread, 'session'.
 _thread_context = threading.local()
 
-OPENAI_LOG = os.environ.get("OPENAI_LOG")
-OPENAI_LOG = "debug"
+LLM_LOG = os.environ.get("LLM_LOG")
+LLM_LOG = "debug"
 
 
 class ApiType(Enum):
@@ -74,8 +78,8 @@ def from_str(label):
 
 
 def _console_log_level():
-    if OPENAI_LOG in ["debug", "info"]:
-        return OPENAI_LOG
+    if LLM_LOG in ["debug", "info"]:
+        return LLM_LOG
     else:
         return None
 
@@ -140,7 +144,7 @@ def operation_location(self) -> Optional[str]:
 
     @property
     def organization(self) -> Optional[str]:
-        return self._headers.get("OpenAI-Organization")
+        return self._headers.get("LLM-Organization")
 
     @property
     def response_ms(self) -> Optional[int]:
@@ -478,7 +482,7 @@ def handle_error_response(self, rbody, rcode, resp, rheaders, stream_error=False
             error_data["message"] += "\n\n" + error_data["internal_message"]
 
         log_info(
-            "OpenAI API error received",
+            "LLM API error received",
             error_code=error_data.get("code"),
             error_type=error_data.get("type"),
             error_message=error_data.get("message"),
@@ -516,7 +520,7 @@ def handle_error_response(self, rbody, rcode, resp, rheaders, stream_error=False
             )
 
     def request_headers(self, method: str, extra, request_id: Optional[str]) -> Dict[str, str]:
-        user_agent = "OpenAI/v1 PythonBindings/%s" % (version.VERSION,)
+        user_agent = "LLM/v1 PythonBindings/%s" % (version.VERSION,)
 
         uname_without_node = " ".join(v for k, v in platform.uname()._asdict().items() if k != "node")
         ua = {
@@ -530,17 +534,17 @@ def request_headers(self, method: str, extra, request_id: Optional[str]) -> Dict
         }
 
         headers = {
-            "X-OpenAI-Client-User-Agent": json.dumps(ua),
+            "X-LLM-Client-User-Agent": json.dumps(ua),
             "User-Agent": user_agent,
         }
 
         headers.update(api_key_to_header(self.api_type, self.api_key))
 
         if self.organization:
-            headers["OpenAI-Organization"] = self.organization
+            headers["LLM-Organization"] = self.organization
 
         if self.api_version is not None and self.api_type == ApiType.OPEN_AI:
-            headers["OpenAI-Version"] = self.api_version
+            headers["LLM-Version"] = self.api_version
         if request_id is not None:
             headers["X-Request-Id"] = request_id
         headers.update(extra)
@@ -592,15 +596,14 @@ def _prepare_request_raw(
                 headers["Content-Type"] = "application/json"
         else:
             raise openai.APIConnectionError(
-                "Unrecognized HTTP method %r. This may indicate a bug in the "
-                "OpenAI bindings. Please contact us through our help center at help.openai.com for "
-                "assistance." % (method,)
+                message=f"Unrecognized HTTP method {method}. This may indicate a bug in the LLM bindings.",
+                request=None,
             )
 
         headers = self.request_headers(method, headers, request_id)
 
-        log_debug("Request to OpenAI API", method=method, path=abs_url)
-        log_debug("Post details", data=data, api_version=self.api_version)
+        # log_debug("Request to LLM API", method=method, path=abs_url)
+        # log_debug("Post details", data=data, api_version=self.api_version)
 
         return abs_url, headers, data
 
@@ -639,14 +642,14 @@ def request_raw(
         except requests.exceptions.Timeout as e:
             raise openai.APITimeoutError("Request timed out: {}".format(e)) from e
         except requests.exceptions.RequestException as e:
-            raise openai.APIConnectionError("Error communicating with OpenAI: {}".format(e)) from e
-        log_debug(
-            "OpenAI API response",
-            path=abs_url,
-            response_code=result.status_code,
-            processing_ms=result.headers.get("OpenAI-Processing-Ms"),
-            request_id=result.headers.get("X-Request-Id"),
-        )
+            raise openai.APIConnectionError(message="Error communicating with LLM: {}".format(e), request=None) from e
+        # log_debug(
+        #     "LLM API response",
+        #     path=abs_url,
+        #     response_code=result.status_code,
+        #     processing_ms=result.headers.get("LLM-Processing-Ms"),
+        #     request_id=result.headers.get("X-Request-Id"),
+        # )
         return result
 
     async def arequest_raw(
@@ -685,18 +688,18 @@ async def arequest_raw(
         }
         try:
             result = await session.request(**request_kwargs)
-            log_info(
-                "OpenAI API response",
-                path=abs_url,
-                response_code=result.status,
-                processing_ms=result.headers.get("OpenAI-Processing-Ms"),
-                request_id=result.headers.get("X-Request-Id"),
-            )
+            # log_info(
+            #     "LLM API response",
+            #     path=abs_url,
+            #     response_code=result.status,
+            #     processing_ms=result.headers.get("LLM-Processing-Ms"),
+            #     request_id=result.headers.get("X-Request-Id"),
+            # )
             return result
         except (aiohttp.ServerTimeoutError, asyncio.TimeoutError) as e:
             raise openai.APITimeoutError("Request timed out") from e
         except aiohttp.ClientError as e:
-            raise openai.APIConnectionError("Error communicating with OpenAI") from e
+            raise openai.APIConnectionError(message="Error communicating with LLM", request=None) from e
 
     def _interpret_response(
         self, result: requests.Response, stream: bool

metagpt/provider/general_api_requestor.py

@@ -3,14 +3,38 @@
 # @Desc   : General Async API for http-based LLM model
 
 import asyncio
-from typing import AsyncGenerator, Tuple, Union
+from typing import AsyncGenerator, Generator, Iterator, Tuple, Union
 
 import aiohttp
+import requests
 
 from metagpt.logs import logger
 from metagpt.provider.general_api_base import APIRequestor
 
 
+def parse_stream_helper(line: bytes) -> Union[bytes, None]:
+    if line and line.startswith(b"data:"):
+        if line.startswith(b"data: "):
+            # SSE event may be valid when it contain whitespace
+            line = line[len(b"data: ") :]
+        else:
+            line = line[len(b"data:") :]
+        if line.strip() == b"[DONE]":
+            # return here will cause GeneratorExit exception in urllib3
+            # and it will close http connection with TCP Reset
+            return None
+        else:
+            return line
+    return None
+
+
+def parse_stream(rbody: Iterator[bytes]) -> Iterator[bytes]:
+    for line in rbody:
+        _line = parse_stream_helper(line)
+        if _line is not None:
+            yield _line
+
+
 class GeneralAPIRequestor(APIRequestor):
     """
     usage
@@ -26,16 +50,40 @@ class GeneralAPIRequestor(APIRequestor):
         )
     """
 
-    def _interpret_response_line(self, rbody: str, rcode: int, rheaders, stream: bool) -> str:
+    def _interpret_response_line(self, rbody: bytes, rcode: int, rheaders, stream: bool) -> bytes:
         # just do nothing to meet the APIRequestor process and return the raw data
         # due to the openai sdk will convert the data into OpenAIResponse which we don't need in general cases.
 
         return rbody
 
+    def _interpret_response(
+        self, result: requests.Response, stream: bool
+    ) -> Tuple[Union[bytes, Iterator[Generator]], bytes]:
+        """Returns the response(s) and a bool indicating whether it is a stream."""
+        if stream and "text/event-stream" in result.headers.get("Content-Type", ""):
+            return (
+                self._interpret_response_line(line, result.status_code, result.headers, stream=True)
+                for line in parse_stream(result.iter_lines())
+            ), True
+        else:
+            return (
+                self._interpret_response_line(
+                    result.content,  # let the caller to decode the msg
+                    result.status_code,
+                    result.headers,
+                    stream=False,
+                ),
+                False,
+            )
+
     async def _interpret_async_response(
         self, result: aiohttp.ClientResponse, stream: bool
-    ) -> Tuple[Union[str, AsyncGenerator[str, None]], bool]:
-        if stream and "text/event-stream" in result.headers.get("Content-Type", ""):
+    ) -> Tuple[Union[bytes, AsyncGenerator[bytes, None]], bool]:
+        if stream and (
+            "text/event-stream" in result.headers.get("Content-Type", "")
+            or "application/x-ndjson" in result.headers.get("Content-Type", "")
+        ):
+            # the `Content-Type` of ollama stream resp is "application/x-ndjson"
             return (
                 self._interpret_response_line(line, result.status, result.headers, stream=True)
                 async for line in result.content