trace support pydantic v1 and v2

cozeloop/integration/langchain/trace_callback.py

@@ -7,6 +7,7 @@
 import traceback
 from typing import List, Dict, Union, Any, Optional
 
+import pydantic
 from pydantic import Field, BaseModel
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentFinish, AgentAction, LLMResult
@@ -445,7 +446,10 @@ def _convert_inputs(inputs: Any) -> Any:
     if isinstance(inputs, PromptValue):
         return _convert_inputs(inputs.to_messages())
     if isinstance(inputs, BaseModel):
-        return inputs.model_dump_json()
+        if pydantic.VERSION.startswith('1'):
+            return inputs.json()
+        else:
+            return inputs.model_dump_json()
     if inputs is None:
         return 'None'
-    return 'type of inputs is not supported'
+    return str(inputs)

cozeloop/internal/httpclient/auth_client.py

@@ -7,6 +7,7 @@
 from urllib.parse import urlparse, quote_plus
 
 import httpx
+import pydantic
 from authlib.jose import jwt
 from pydantic import BaseModel
 
@@ -199,9 +200,15 @@ def get_access_token(
         jwt_token = self._gen_jwt(self._public_key_id, self._private_key, 3600, session_name)
         url = f"{self._base_url}/api/permission/oauth2/token"
         headers = {"Authorization": f"Bearer {jwt_token}"}
+        scope_str = None
+        if scope:
+            if pydantic.VERSION.startswith('1'):
+                scope_str = scope.dict()
+            else:
+                scope_str = scope.model_dump()
         body = {
             "duration_seconds": ttl,
             "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
-            "scope": scope.model_dump() if scope else None,
+            "scope": scope_str,
         }
         return self._do_request(url, "POST", OAuthToken, headers=headers, json=body)

cozeloop/internal/httpclient/client.py

@@ -6,6 +6,7 @@
 from typing import Optional, Dict, Union, IO, Type, Tuple
 
 import httpx
+import pydantic
 from pydantic import BaseModel
 
 from cozeloop.internal import consts
@@ -71,7 +72,10 @@ def request(
         _timeout = timeout if timeout is not None else self.timeout
 
         if isinstance(json, BaseModel):
-            json = json.model_dump(by_alias=True)
+            if pydantic.VERSION.startswith('1'):
+                json = json.dict(by_alias=True)
+            else:
+                json = json.model_dump(by_alias=True)
 
         try:
             response = self.http_client.request(

cozeloop/internal/httpclient/http_client.py

@@ -5,6 +5,7 @@
 from typing import Dict, Type, TypeVar
 
 import httpx
+import pydantic
 from pydantic import ValidationError
 
 from cozeloop.internal import consts
@@ -50,7 +51,14 @@ def parse_response(url: str, response: httpx.Response, response_model: Type[T])
         raise e
 
     try:
-        res = response_model.model_validate(data) if data is not None else response_model()
+        res = None
+        if data is not None:
+            if pydantic.VERSION.startswith('1'):
+                res = response_model.parse_obj(data)
+            else:
+                res = response_model.model_validate(data)
+        else:
+            res = response_model()
     except ValidationError as e:
         logger.error(f"Failed to parse response. Path: {url}, http code: {http_code}, log id: {log_id}, error: {e}.")
         raise consts.InternalError from e

cozeloop/internal/prompt/openapi.py

@@ -4,6 +4,7 @@
 from enum import Enum
 from typing import List, Optional
 
+import pydantic
 from pydantic import BaseModel
 
 from cozeloop.internal.httpclient import Client, BaseResponse
@@ -153,6 +154,10 @@ def _do_mpull_prompt(self, workspace_id: str, queries: List[PromptQuery]) -> Opt
             return None
         request = MPullPromptRequest(workspace_id=workspace_id, queries=queries)
         response = self.http_client.post(MPULL_PROMPT_PATH, MPullPromptResponse, request)
-        real_resp = MPullPromptResponse.model_validate(response)
+        real_resp = None
+        if pydantic.VERSION.startswith('1'):
+            real_resp = MPullPromptResponse.parse_obj(response)
+        else:
+            real_resp = MPullPromptResponse.model_validate(response)
         if real_resp.data is not None:
             return real_resp.data.items

cozeloop/internal/prompt/prompt.py

@@ -3,6 +3,7 @@
 
 import json
 from typing import Dict, Any, List, Optional
+import pydantic
 
 from jinja2 import BaseLoader, Undefined
 from jinja2.sandbox import SandboxedEnvironment
@@ -52,9 +53,15 @@ def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Opt
                 try:
                     prompt = self._get_prompt(prompt_key, version, label)
                     if prompt is not None:
+
+                        output = None
+                        if pydantic.VERSION.startswith('1'):
+                            output = prompt.json()
+                        else:
+                            output = prompt.model_dump_json(exclude_none=True)
                         prompt_hub_pan.set_tags({
                             PROMPT_VERSION: prompt.version,
-                            consts.OUTPUT: prompt.model_dump_json(exclude_none=True),
+                            consts.OUTPUT: output,
                         })
                     return prompt
                 except RemoteServiceError as e:
@@ -91,15 +98,25 @@ def prompt_format(
             with self.trace_provider.start_span(consts.TRACE_PROMPT_TEMPLATE_SPAN_NAME,
                                                 consts.TRACE_PROMPT_TEMPLATE_SPAN_TYPE,
                                                 scene=V_SCENE_PROMPT_TEMPLATE) as prompt_template_span:
+                input = None
+                if pydantic.VERSION.startswith('1'):
+                    input = _to_span_prompt_input(prompt.prompt_template.messages, variables).json()
+                else:
+                    input = _to_span_prompt_input(prompt.prompt_template.messages, variables).model_dump_json(exclude_none=True)
                 prompt_template_span.set_tags({
                     PROMPT_KEY: prompt.prompt_key,
                     PROMPT_VERSION: prompt.version,
-                    consts.INPUT: _to_span_prompt_input(prompt.prompt_template.messages, variables).model_dump_json(exclude_none=True)
+                    consts.INPUT: input
                 })
                 try:
                     results = self._prompt_format(prompt, variables)
+                    output = None
+                    if pydantic.VERSION.startswith('1'):
+                        output = _to_span_prompt_output(results).json()
+                    else:
+                        output = _to_span_prompt_output(results).model_dump_json(exclude_none=True)
                     prompt_template_span.set_tags({
-                        consts.OUTPUT: _to_span_prompt_output(results).model_dump_json(exclude_none=True),
+                        consts.OUTPUT: output,
                     })
                     return results
                 except RemoteServiceError as e:

cozeloop/internal/trace/exporter.py

@@ -6,6 +6,8 @@
 import time
 from typing import Dict, List, Optional, Tuple, Callable, Any
 
+import pydantic
+
 from cozeloop.spec.tracespec import ModelInput, ModelMessagePart, ModelMessagePartType, ModelImageURL, ModelFileURL, ModelOutput
 from cozeloop.internal.consts import *
 from cozeloop.internal.httpclient import Client, BaseResponse
@@ -17,6 +19,35 @@
 
 logger = logging.getLogger(__name__)
 
+
+class UploadSpan(BaseModel):
+    started_at_micros: int
+    log_id: str
+    span_id: str
+    parent_id: str
+    trace_id: str
+    duration_micros: int
+    service_name: str
+    workspace_id: str
+    span_name: str
+    span_type: str
+    status_code: int
+    input: str
+    output: str
+    object_storage: str
+    system_tags_string: Dict[str, str]
+    system_tags_long: Dict[str, int]
+    system_tags_double: Dict[str, float]
+    tags_string: Dict[str, str]
+    tags_long: Dict[str, int]
+    tags_double: Dict[str, float]
+    tags_bool: Dict[str, bool]
+
+
+class UploadSpanData(BaseModel):
+    spans: List['UploadSpan']
+
+
 class Exporter:
     def export_spans(self, ctx: dict, spans: List['UploadSpan']):
         raise NotImplementedError
@@ -92,34 +123,6 @@ def export_spans(self, ctx: dict, spans: List['UploadSpan']):
             raise Exception(f"export spans fail, err:[{e}]")
 
 
-class UploadSpanData(BaseModel):
-    spans: List['UploadSpan']
-
-
-class UploadSpan(BaseModel):
-    started_at_micros: int
-    log_id: str
-    span_id: str
-    parent_id: str
-    trace_id: str
-    duration_micros: int
-    service_name: str
-    workspace_id: str
-    span_name: str
-    span_type: str
-    status_code: int
-    input: str
-    output: str
-    object_storage: str
-    system_tags_string: Dict[str, str]
-    system_tags_long: Dict[str, int]
-    system_tags_double: Dict[str, float]
-    tags_string: Dict[str, str]
-    tags_long: Dict[str, int]
-    tags_double: Dict[str, float]
-    tags_bool: Dict[str, bool]
-
-
 class UploadFile(BaseModel):
     class Config:
         arbitrary_types_allowed = True
@@ -216,7 +219,10 @@ def convert_input(span_key: str, span: Span) -> (str, List[UploadFile]):
         model_input = ModelInput()
         if isinstance(value, str):
             try:
-                model_input = ModelInput.model_validate_json(value)
+                if pydantic.VERSION.startswith('1'):
+                    model_input = ModelInput.parse_raw(value)
+                else:
+                    model_input = ModelInput.model_validate_json(value)
             except Exception as e:
                 logger.error(f"unmarshal ModelInput failed, err: {e}")
                 return "", []
@@ -226,7 +232,10 @@ def convert_input(span_key: str, span: Span) -> (str, List[UploadFile]):
                 files = transfer_message_part(part, span, span_key)
                 upload_files.extend(files)
 
-        value_res = model_input.model_dump_json()
+        if pydantic.VERSION.startswith('1'):
+            value_res = model_input.json()
+        else:
+            value_res = model_input.model_dump_json()
 
         if len(value_res) > MAX_BYTES_OF_ONE_TAG_VALUE_OF_INPUT_OUTPUT:
             value_res, f = transfer_text(value_res, span, span_key)
@@ -251,7 +260,10 @@ def convert_output(span_key: str, span: Span) -> (str, List[UploadFile]):
         model_output = ModelOutput()
         if isinstance(value, str):
             try:
-                model_output = ModelOutput.model_validate_json(value)
+                if pydantic.VERSION.startswith('1'):
+                    model_output = ModelOutput.parse_raw(value)
+                else:
+                    model_output = ModelOutput.model_validate_json(value)
             except Exception as e:
                 logger.error(f"unmarshal ModelOutput failed, err: {e}")
                 return "", []
@@ -330,8 +342,10 @@ def transfer_object_storage(span_upload_files: List[UploadFile]) -> str:
     if not is_exist:
         return ""
 
-
-    return object_storage.model_dump_json()
+    if pydantic.VERSION.startswith('1'):
+        return object_storage.json()
+    else:
+        return object_storage.model_dump_json()
 
 
 def transfer_message_part(src: ModelMessagePart, span: 'Span', tag_key: str) -> List[UploadFile]:

cozeloop/internal/trace/model/model.py

@@ -8,19 +8,19 @@
 from pydantic.dataclasses import dataclass
 
 
-class ObjectStorage(BaseModel):
-    input_tos_key: Optional[str] = None  # The key for reporting long input data
-    output_tos_key: Optional[str] = None  # The key for reporting long output data
-    attachments: List['Attachment'] = None  # attachments in input or output
-
-
 class Attachment(BaseModel):
     field: Optional[str] = None
     name: Optional[str] = None
     type: Optional[str] = None  # text, image, file
     tos_key: Optional[str] = None
 
 
+class ObjectStorage(BaseModel):
+    input_tos_key: Optional[str] = None  # The key for reporting long input data
+    output_tos_key: Optional[str] = None  # The key for reporting long output data
+    attachments: List['Attachment'] = None  # attachments in input or output
+
+
 class UploadType(str, Enum):
     LONG = 1
     MULTI_MODALITY = 2

cozeloop/internal/trace/span.py

@@ -9,6 +9,8 @@
 import json
 import urllib.parse
 
+import pydantic
+
 from cozeloop import span
 from cozeloop.internal.trace.model.model import TagTruncateConf
 from cozeloop.spec.tracespec import (ModelInput, ModelOutput, ModelMessagePartType, ModelMessage, ModelMessagePart,
@@ -218,7 +220,11 @@ def get_model_input_bytes_size(self, m_content):
                     part.file_url.url = ""
 
         try:
-            m_content_json = m_content.model_dump_json()
+            m_content_json = ""
+            if pydantic.VERSION.startswith('1'):
+                m_content_json = m_content.json()
+            else:
+                m_content_json = m_content.model_dump_json()
             return len(m_content_json)
         except Exception as e:
             logger.error(f"Failed to get model input size, m_content model_dump_json err: {e}")

cozeloop/internal/utils/convert.py

@@ -6,6 +6,8 @@
 import string
 from typing import Any, Dict, List, Optional, TypeVar, Sequence
 from functools import singledispatch
+
+import pydantic
 from pydantic import BaseModel
 
 T = TypeVar('T')
@@ -75,7 +77,10 @@ def to_json(param: Any) -> str:
         return param
     try:
         if isinstance(param, BaseModel):
-            return param.model_dump_json()
+            if pydantic.VERSION.startswith('1'):
+                return param.json()
+            else:
+                return param.model_dump_json()
         return json.dumps(param, ensure_ascii=False)
     except json.JSONDecodeError:
         return param.__str__()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cozeloop"
-version = "0.1.11"
+version = "0.1.12"
 description = "coze loop sdk"
 authors = ["JiangQi715 <jiangqi.rrt@bytedance.com>"]
 license = "MIT"
@@ -9,7 +9,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = ">=3.8,<4.0"
 httpx = ">=0.23.0,<1.0.0"
-pydantic = ">=2.10.6,<3.0.0"
+pydantic = ">=1.10.12,<3.0.0"
 cachetools = "^5.5.2"
 apscheduler = "^3.11.0"
 jinja2 = "^3.1.6"