[Model] support Claude (#283)

* [Model] support Claude * fix openai
InternLM · Dec 20, 2024 · c4523f0 · c4523f0
1 parent bcc62ed
commit c4523f0
Show file tree

Hide file tree

Showing 8 changed files with 562 additions and 230 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,20 +1,22 @@
 exclude: ^(tests/data|scripts|ftdp/protocols|ftdp/template_configs|ftdp/tool_dicts)/
 repos:
   - repo: https://github.com/PyCQA/flake8
-    rev: 7.0.0
+    rev: 7.1.1
     hooks:
       - id: flake8
   - repo: https://github.com/PyCQA/isort
     rev: 5.13.2
     hooks:
       - id: isort
+        args: ["--profile", "black", "--filter-files", "--line-width", "119"]
   - repo: https://github.com/psf/black
-    rev: 22.8.0
+    rev: 24.10.0
     hooks:
       - id: black
         args: ["--line-length", "119", "--skip-string-normalization"]
+
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v5.0.0
     hooks:
       - id: trailing-whitespace
       - id: check-yaml
@@ -27,7 +29,7 @@ repos:
       - id: mixed-line-ending
         args: ["--fix=lf"]
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.17
+    rev: 0.7.21
     hooks:
       - id: mdformat
         args: ["--number"]
@@ -36,11 +38,11 @@ repos:
           - mdformat_frontmatter
           - linkify-it-py
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.2.6
+    rev: v2.3.0
     hooks:
       - id: codespell
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.15.0
+    rev: v3.19.1
     hooks:
       - id: pyupgrade
         args: ["--py36-plus"]
diff --git a/lagent/actions/__init__.py b/lagent/actions/__init__.py
@@ -1,6 +1,6 @@
 from .action_executor import ActionExecutor, AsyncActionExecutor
 from .arxiv_search import ArxivSearch, AsyncArxivSearch
-from .base_action import BaseAction, tool_api
+from .base_action import AsyncActionMixin, BaseAction, tool_api
 from .bing_map import AsyncBINGMap, BINGMap
 from .builtin_actions import FinishAction, InvalidAction, NoAction
 from .google_scholar_search import AsyncGoogleScholar, GoogleScholar
@@ -14,12 +14,34 @@
 from .web_browser import AsyncWebBrowser, WebBrowser
 
 __all__ = [
-    'BaseAction', 'ActionExecutor', 'AsyncActionExecutor', 'InvalidAction',
-    'FinishAction', 'NoAction', 'BINGMap', 'AsyncBINGMap', 'ArxivSearch',
-    'AsyncArxivSearch', 'GoogleSearch', 'AsyncGoogleSearch', 'GoogleScholar',
-    'AsyncGoogleScholar', 'IPythonInterpreter', 'AsyncIPythonInterpreter',
-    'IPythonInteractive', 'AsyncIPythonInteractive',
-    'IPythonInteractiveManager', 'PythonInterpreter', 'AsyncPythonInterpreter',
-    'PPT', 'AsyncPPT', 'WebBrowser', 'AsyncWebBrowser', 'BaseParser',
-    'JsonParser', 'TupleParser', 'tool_api'
+    'BaseAction',
+    'ActionExecutor',
+    'AsyncActionExecutor',
+    'InvalidAction',
+    'FinishAction',
+    'NoAction',
+    'BINGMap',
+    'AsyncBINGMap',
+    'ArxivSearch',
+    'AsyncArxivSearch',
+    'GoogleSearch',
+    'AsyncGoogleSearch',
+    'GoogleScholar',
+    'AsyncGoogleScholar',
+    'IPythonInterpreter',
+    'AsyncIPythonInterpreter',
+    'IPythonInteractive',
+    'AsyncIPythonInteractive',
+    'IPythonInteractiveManager',
+    'PythonInterpreter',
+    'AsyncPythonInterpreter',
+    'PPT',
+    'AsyncPPT',
+    'WebBrowser',
+    'AsyncWebBrowser',
+    'BaseParser',
+    'JsonParser',
+    'TupleParser',
+    'tool_api',
+    'AsyncActionMixin',
 ]
diff --git a/lagent/hooks/logger.py b/lagent/hooks/logger.py
@@ -1,5 +1,4 @@
 import random
-from typing import Optional
 
 from termcolor import COLORS, colored
 
@@ -8,10 +7,10 @@
 
 
 class MessageLogger(Hook):
-
-    def __init__(self, name: str = 'lagent'):
+    def __init__(self, name: str = 'lagent', add_file_handler: bool = False):
         self.logger = get_logger(
-            name, 'info', '%(asctime)s %(levelname)8s %(name)8s - %(message)s')
+            name, 'info', '%(asctime)s %(levelname)8s %(name)8s - %(message)s', add_file_handler=add_file_handler
+        )
         self.sender2color = {}
 
     def before_agent(self, agent, messages, session_id):
@@ -29,9 +28,5 @@ def after_action(self, executor, message, session_id):
 
     def _process_message(self, message, session_id):
         sender = message.sender
-        color = self.sender2color.setdefault(sender,
-                                             random.choice(list(COLORS)))
-        self.logger.info(
-            colored(
-                f'session id: {session_id}, message sender: {sender}\n'
-                f'{message.content}', color))
+        color = self.sender2color.setdefault(sender, random.choice(list(COLORS)))
+        self.logger.info(colored(f'session id: {session_id}, message sender: {sender}\n' f'{message.content}', color))
diff --git a/lagent/llms/__init__.py b/lagent/llms/__init__.py
@@ -1,9 +1,15 @@
+from .anthropic_llm import AsyncClaudeAPI, ClaudeAPI
 from .base_api import AsyncBaseAPILLM, BaseAPILLM
 from .base_llm import AsyncBaseLLM, BaseLLM
 from .huggingface import HFTransformer, HFTransformerCasualLM, HFTransformerChat
-from .lmdeploy_wrapper import (AsyncLMDeployClient, AsyncLMDeployPipeline,
-                               AsyncLMDeployServer, LMDeployClient,
-                               LMDeployPipeline, LMDeployServer)
+from .lmdeploy_wrapper import (
+    AsyncLMDeployClient,
+    AsyncLMDeployPipeline,
+    AsyncLMDeployServer,
+    LMDeployClient,
+    LMDeployPipeline,
+    LMDeployServer,
+)
 from .meta_template import INTERNLM2_META
 from .openai import GPTAPI, AsyncGPTAPI
 from .sensenova import SensenovaAPI
@@ -29,4 +35,6 @@
     'VllmModel',
     'AsyncVllmModel',
     'SensenovaAPI',
+    'AsyncClaudeAPI',
+    'ClaudeAPI',
 ]