兼容 pydantic v1

lmclient/parsers/__init__.py

@@ -1,6 +1,6 @@
 from lmclient.parsers.base import ModelResponseParser
 from lmclient.parsers.minimax import MinimaxTextParser
-from lmclient.parsers.openai import OpenAIContentParser, OpenAIFunctionCallParser, OpenAIParser, OpenAISchema
+from lmclient.parsers.openai import Field, OpenAIContentParser, OpenAIFunctionCallParser, OpenAIParser, OpenAISchema
 from lmclient.parsers.zhipu import ZhiPuParser
 
 __all__ = [
@@ -11,4 +11,5 @@
     'OpenAIFunctionCallParser',
     'OpenAIParser',
     'OpenAISchema',
+    'Field',
 ]

lmclient/parsers/openai.py

@@ -3,7 +3,12 @@
 import json
 from typing import TypeVar
 
-from pydantic import BaseModel
+try:
+    from pydantic.v1 import BaseModel
+    from pydantic.v1 import Field as Field
+except ImportError:
+    from pydantic import BaseModel
+    from pydantic import Field as Field
 
 from lmclient.exceptions import ParserError
 from lmclient.parsers.base import ModelResponseParser
@@ -59,7 +64,7 @@ def _remove_a_key(d, remove_key) -> None:
                 _remove_a_key(d[key], remove_key)
 
 
-class OpenAISchema(BaseModel):
+class OpenAISchema(BaseModel):  # type: ignore
     @classmethod
     def openai_schema(cls):
         """
@@ -71,7 +76,7 @@ def openai_schema(cls):
         Returns:
             model_json_schema (dict): A dictionary in the format of OpenAI's schema as jsonschema
         """
-        schema = cls.model_json_schema()
+        schema = cls.schema()
         parameters = {k: v for k, v in schema.items() if k not in ('title', 'description')}
         parameters['required'] = sorted(parameters['properties'])
         _remove_a_key(parameters, 'title')

lmclient/types.py

@@ -2,7 +2,11 @@
 
 from typing import Any, Dict, Generic, Optional, Sequence, TypedDict, TypeVar, Union
 
-from pydantic import BaseModel, Field
+try:
+    from pydantic.v1 import BaseModel, Field
+except ImportError:
+    from pydantic import BaseModel, Field
+
 from typing_extensions import NotRequired
 
 T = TypeVar('T')
@@ -22,7 +26,7 @@ class Message(TypedDict):
 Prompt = Union[str, Sequence[dict]]
 
 
-class TaskResult(BaseModel, Generic[T]):
+class TaskResult(BaseModel, Generic[T]):  # type: ignore
     output: Optional[T] = None
     response: ModelResponse = Field(default_factory=dict)
     error_message: Optional[str] = None

lmclient/version.py

@@ -1,2 +1,2 @@
-__version__ = '0.5.0'
+__version__ = '0.5.1'
 __cache_version__ = '3'

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lmclient-core"
-version = "0.5.0"
+version = "0.5.1"
 description = "LM Async Client, openai client, azure openai client ..."
 authors = ["wangyuxin <[email protected]>"]
 readme = "README.md"
@@ -16,11 +16,11 @@ typing-extensions = "^4.6.2"
 diskcache = "^5.6.1"
 httpx = "^0.24.1"
 tenacity = "^8.2.2"
-pydantic = "^2.1.1"
 websocket-client = "^1.6.1"
 pyjwt = "^2.8.0"
 cachetools = "^5.3.1"
 tqdm = "^4.66.1"
+pydantic = ">1.0"
 
 [tool.ruff]
 line-length = 128

scripts/ner.py

@@ -6,9 +6,8 @@
 from typing import List
 
 import typer
-from pydantic import Field
 
-from lmclient import AzureChat, LMClientForStructuredData, OpenAIChat, OpenAISchema
+from lmclient import AzureChat, Field, LMClientForStructuredData, OpenAIChat, OpenAISchema
 from lmclient.client import ErrorMode
 
 
@@ -18,6 +17,8 @@ class ModelType(str, Enum):
 
 
 class NerInfo(OpenAISchema):
+    """命名实体信息，包括人名，地名和组织名"""
+
     person: List[str] = Field(default_factory=list)
     location: List[str] = Field(default_factory=list)
     organization: List[str] = Field(default_factory=list)
@@ -67,7 +68,10 @@ def main(
             if result.output is None:
                 output = None
             else:
-                output = result.output.model_dump()
+                try:
+                    output = result.output.model_dump()
+                except AttributeError:
+                    output = result.output.dict()
             output_dict = {'text': text, 'output': output}
             f.write(json.dumps(output_dict, ensure_ascii=False) + '\n')