Skip to content

Commit 3cf5840

Browse files
JqRrtJqRrt
authored
trace support pydantic v1 and v2 (#24)
* support pydantic v1 and v2
1 parent 78a34b3 commit 3cf5840

File tree

11 files changed

+122
-52
lines changed

11 files changed

+122
-52
lines changed

‎cozeloop/integration/langchain/trace_callback.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import traceback
88
from typing import List, Dict, Union, Any, Optional
99

10+
import pydantic
1011
from pydantic import Field, BaseModel
1112
from langchain.callbacks.base import BaseCallbackHandler
1213
from langchain.schema import AgentFinish, AgentAction, LLMResult
@@ -445,7 +446,10 @@ def _convert_inputs(inputs: Any) -> Any:
445446
if isinstance(inputs, PromptValue):
446447
return _convert_inputs(inputs.to_messages())
447448
if isinstance(inputs, BaseModel):
448-
return inputs.model_dump_json()
449+
if pydantic.VERSION.startswith('1'):
450+
return inputs.json()
451+
else:
452+
return inputs.model_dump_json()
449453
if inputs is None:
450454
return 'None'
451-
return 'type of inputs is not supported'
455+
return str(inputs)

‎cozeloop/internal/httpclient/auth_client.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from urllib.parse import urlparse, quote_plus
88

99
import httpx
10+
import pydantic
1011
from authlib.jose import jwt
1112
from pydantic import BaseModel
1213

@@ -199,9 +200,15 @@ def get_access_token(
199200
jwt_token = self._gen_jwt(self._public_key_id, self._private_key, 3600, session_name)
200201
url = f"{self._base_url}/api/permission/oauth2/token"
201202
headers = {"Authorization": f"Bearer {jwt_token}"}
203+
scope_str = None
204+
if scope:
205+
if pydantic.VERSION.startswith('1'):
206+
scope_str = scope.dict()
207+
else:
208+
scope_str = scope.model_dump()
202209
body = {
203210
"duration_seconds": ttl,
204211
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
205-
"scope": scope.model_dump() if scope else None,
212+
"scope": scope_str,
206213
}
207214
return self._do_request(url, "POST", OAuthToken, headers=headers, json=body)

‎cozeloop/internal/httpclient/client.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from typing import Optional, Dict, Union, IO, Type, Tuple
77

88
import httpx
9+
import pydantic
910
from pydantic import BaseModel
1011

1112
from cozeloop.internal import consts
@@ -71,7 +72,10 @@ def request(
7172
_timeout = timeout if timeout is not None else self.timeout
7273

7374
if isinstance(json, BaseModel):
74-
json = json.model_dump(by_alias=True)
75+
if pydantic.VERSION.startswith('1'):
76+
json = json.dict(by_alias=True)
77+
else:
78+
json = json.model_dump(by_alias=True)
7579

7680
try:
7781
response = self.http_client.request(

‎cozeloop/internal/httpclient/http_client.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from typing import Dict, Type, TypeVar
66

77
import httpx
8+
import pydantic
89
from pydantic import ValidationError
910

1011
from cozeloop.internal import consts
@@ -50,7 +51,14 @@ def parse_response(url: str, response: httpx.Response, response_model: Type[T])
5051
raise e
5152

5253
try:
53-
res = response_model.model_validate(data) if data is not None else response_model()
54+
res = None
55+
if data is not None:
56+
if pydantic.VERSION.startswith('1'):
57+
res = response_model.parse_obj(data)
58+
else:
59+
res = response_model.model_validate(data)
60+
else:
61+
res = response_model()
5462
except ValidationError as e:
5563
logger.error(f"Failed to parse response. Path: {url}, http code: {http_code}, log id: {log_id}, error: {e}.")
5664
raise consts.InternalError from e

‎cozeloop/internal/prompt/openapi.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from enum import Enum
55
from typing import List, Optional
66

7+
import pydantic
78
from pydantic import BaseModel
89

910
from cozeloop.internal.httpclient import Client, BaseResponse
@@ -165,6 +166,10 @@ def _do_mpull_prompt(self, workspace_id: str, queries: List[PromptQuery]) -> Opt
165166
return None
166167
request = MPullPromptRequest(workspace_id=workspace_id, queries=queries)
167168
response = self.http_client.post(MPULL_PROMPT_PATH, MPullPromptResponse, request)
168-
real_resp = MPullPromptResponse.model_validate(response)
169+
real_resp = None
170+
if pydantic.VERSION.startswith('1'):
171+
real_resp = MPullPromptResponse.parse_obj(response)
172+
else:
173+
real_resp = MPullPromptResponse.model_validate(response)
169174
if real_resp.data is not None:
170175
return real_resp.data.items

‎cozeloop/internal/prompt/prompt.py‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import json
55
from typing import Dict, Any, List, Optional
6+
import pydantic
67

78
from jinja2 import BaseLoader, Undefined
89
from jinja2.sandbox import SandboxedEnvironment
@@ -52,9 +53,15 @@ def get_prompt(self, prompt_key: str, version: str = '', label: str = '') -> Opt
5253
try:
5354
prompt = self._get_prompt(prompt_key, version, label)
5455
if prompt is not None:
56+
57+
output = None
58+
if pydantic.VERSION.startswith('1'):
59+
output = prompt.json()
60+
else:
61+
output = prompt.model_dump_json(exclude_none=True)
5562
prompt_hub_pan.set_tags({
5663
PROMPT_VERSION: prompt.version,
57-
consts.OUTPUT: prompt.model_dump_json(exclude_none=True),
64+
consts.OUTPUT: output,
5865
})
5966
return prompt
6067
except RemoteServiceError as e:
@@ -91,15 +98,25 @@ def prompt_format(
9198
with self.trace_provider.start_span(consts.TRACE_PROMPT_TEMPLATE_SPAN_NAME,
9299
consts.TRACE_PROMPT_TEMPLATE_SPAN_TYPE,
93100
scene=V_SCENE_PROMPT_TEMPLATE) as prompt_template_span:
101+
input = None
102+
if pydantic.VERSION.startswith('1'):
103+
input = _to_span_prompt_input(prompt.prompt_template.messages, variables).json()
104+
else:
105+
input = _to_span_prompt_input(prompt.prompt_template.messages, variables).model_dump_json(exclude_none=True)
94106
prompt_template_span.set_tags({
95107
PROMPT_KEY: prompt.prompt_key,
96108
PROMPT_VERSION: prompt.version,
97-
consts.INPUT: _to_span_prompt_input(prompt.prompt_template.messages, variables).model_dump_json(exclude_none=True)
109+
consts.INPUT: input
98110
})
99111
try:
100112
results = self._prompt_format(prompt, variables)
113+
output = None
114+
if pydantic.VERSION.startswith('1'):
115+
output = _to_span_prompt_output(results).json()
116+
else:
117+
output = _to_span_prompt_output(results).model_dump_json(exclude_none=True)
101118
prompt_template_span.set_tags({
102-
consts.OUTPUT: _to_span_prompt_output(results).model_dump_json(exclude_none=True),
119+
consts.OUTPUT: output,
103120
})
104121
return results
105122
except RemoteServiceError as e:

‎cozeloop/internal/trace/exporter.py‎

Lines changed: 47 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import time
77
from typing import Dict, List, Optional, Tuple, Callable, Any
88

9+
import pydantic
10+
911
from cozeloop.spec.tracespec import ModelInput, ModelMessagePart, ModelMessagePartType, ModelImageURL, ModelFileURL, ModelOutput
1012
from cozeloop.internal.consts import *
1113
from cozeloop.internal.httpclient import Client, BaseResponse
@@ -17,6 +19,35 @@
1719

1820
logger = logging.getLogger(__name__)
1921

22+
23+
class UploadSpan(BaseModel):
24+
started_at_micros: int
25+
log_id: str
26+
span_id: str
27+
parent_id: str
28+
trace_id: str
29+
duration_micros: int
30+
service_name: str
31+
workspace_id: str
32+
span_name: str
33+
span_type: str
34+
status_code: int
35+
input: str
36+
output: str
37+
object_storage: str
38+
system_tags_string: Dict[str, str]
39+
system_tags_long: Dict[str, int]
40+
system_tags_double: Dict[str, float]
41+
tags_string: Dict[str, str]
42+
tags_long: Dict[str, int]
43+
tags_double: Dict[str, float]
44+
tags_bool: Dict[str, bool]
45+
46+
47+
class UploadSpanData(BaseModel):
48+
spans: List['UploadSpan']
49+
50+
2051
class Exporter:
2152
def export_spans(self, ctx: dict, spans: List['UploadSpan']):
2253
raise NotImplementedError
@@ -92,34 +123,6 @@ def export_spans(self, ctx: dict, spans: List['UploadSpan']):
92123
raise Exception(f"export spans fail, err:[{e}]")
93124

94125

95-
class UploadSpanData(BaseModel):
96-
spans: List['UploadSpan']
97-
98-
99-
class UploadSpan(BaseModel):
100-
started_at_micros: int
101-
log_id: str
102-
span_id: str
103-
parent_id: str
104-
trace_id: str
105-
duration_micros: int
106-
service_name: str
107-
workspace_id: str
108-
span_name: str
109-
span_type: str
110-
status_code: int
111-
input: str
112-
output: str
113-
object_storage: str
114-
system_tags_string: Dict[str, str]
115-
system_tags_long: Dict[str, int]
116-
system_tags_double: Dict[str, float]
117-
tags_string: Dict[str, str]
118-
tags_long: Dict[str, int]
119-
tags_double: Dict[str, float]
120-
tags_bool: Dict[str, bool]
121-
122-
123126
class UploadFile(BaseModel):
124127
class Config:
125128
arbitrary_types_allowed = True
@@ -216,7 +219,10 @@ def convert_input(span_key: str, span: Span) -> (str, List[UploadFile]):
216219
model_input = ModelInput()
217220
if isinstance(value, str):
218221
try:
219-
model_input = ModelInput.model_validate_json(value)
222+
if pydantic.VERSION.startswith('1'):
223+
model_input = ModelInput.parse_raw(value)
224+
else:
225+
model_input = ModelInput.model_validate_json(value)
220226
except Exception as e:
221227
logger.error(f"unmarshal ModelInput failed, err: {e}")
222228
return "", []
@@ -226,7 +232,10 @@ def convert_input(span_key: str, span: Span) -> (str, List[UploadFile]):
226232
files = transfer_message_part(part, span, span_key)
227233
upload_files.extend(files)
228234

229-
value_res = model_input.model_dump_json()
235+
if pydantic.VERSION.startswith('1'):
236+
value_res = model_input.json()
237+
else:
238+
value_res = model_input.model_dump_json()
230239

231240
if len(value_res) > MAX_BYTES_OF_ONE_TAG_VALUE_OF_INPUT_OUTPUT:
232241
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]):
251260
model_output = ModelOutput()
252261
if isinstance(value, str):
253262
try:
254-
model_output = ModelOutput.model_validate_json(value)
263+
if pydantic.VERSION.startswith('1'):
264+
model_output = ModelOutput.parse_raw(value)
265+
else:
266+
model_output = ModelOutput.model_validate_json(value)
255267
except Exception as e:
256268
logger.error(f"unmarshal ModelOutput failed, err: {e}")
257269
return "", []
@@ -330,8 +342,10 @@ def transfer_object_storage(span_upload_files: List[UploadFile]) -> str:
330342
if not is_exist:
331343
return ""
332344

333-
334-
return object_storage.model_dump_json()
345+
if pydantic.VERSION.startswith('1'):
346+
return object_storage.json()
347+
else:
348+
return object_storage.model_dump_json()
335349

336350

337351
def transfer_message_part(src: ModelMessagePart, span: 'Span', tag_key: str) -> List[UploadFile]:

‎cozeloop/internal/trace/model/model.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,19 @@
88
from pydantic.dataclasses import dataclass
99

1010

11-
class ObjectStorage(BaseModel):
12-
input_tos_key: Optional[str] = None # The key for reporting long input data
13-
output_tos_key: Optional[str] = None # The key for reporting long output data
14-
attachments: List['Attachment'] = None # attachments in input or output
15-
16-
1711
class Attachment(BaseModel):
1812
field: Optional[str] = None
1913
name: Optional[str] = None
2014
type: Optional[str] = None # text, image, file
2115
tos_key: Optional[str] = None
2216

2317

18+
class ObjectStorage(BaseModel):
19+
input_tos_key: Optional[str] = None # The key for reporting long input data
20+
output_tos_key: Optional[str] = None # The key for reporting long output data
21+
attachments: List['Attachment'] = None # attachments in input or output
22+
23+
2424
class UploadType(str, Enum):
2525
LONG = 1
2626
MULTI_MODALITY = 2

‎cozeloop/internal/trace/span.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import json
1010
import urllib.parse
1111

12+
import pydantic
13+
1214
from cozeloop import span
1315
from cozeloop.internal.trace.model.model import TagTruncateConf
1416
from cozeloop.spec.tracespec import (ModelInput, ModelOutput, ModelMessagePartType, ModelMessage, ModelMessagePart,
@@ -218,7 +220,11 @@ def get_model_input_bytes_size(self, m_content):
218220
part.file_url.url = ""
219221

220222
try:
221-
m_content_json = m_content.model_dump_json()
223+
m_content_json = ""
224+
if pydantic.VERSION.startswith('1'):
225+
m_content_json = m_content.json()
226+
else:
227+
m_content_json = m_content.model_dump_json()
222228
return len(m_content_json)
223229
except Exception as e:
224230
logger.error(f"Failed to get model input size, m_content model_dump_json err: {e}")

‎cozeloop/internal/utils/convert.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
import string
77
from typing import Any, Dict, List, Optional, TypeVar, Sequence
88
from functools import singledispatch
9+
10+
import pydantic
911
from pydantic import BaseModel
1012

1113
T = TypeVar('T')
@@ -75,7 +77,10 @@ def to_json(param: Any) -> str:
7577
return param
7678
try:
7779
if isinstance(param, BaseModel):
78-
return param.model_dump_json()
80+
if pydantic.VERSION.startswith('1'):
81+
return param.json()
82+
else:
83+
return param.model_dump_json()
7984
return json.dumps(param, ensure_ascii=False)
8085
except json.JSONDecodeError:
8186
return param.__str__()

0 commit comments

Comments
 (0)