From c5055310483d7b9d5122e1fae2d4caf3f8040c74 Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Wed, 30 Oct 2024 17:08:20 -0700 Subject: [PATCH] fix: Make imports more explicit for `BaseModel` v1 or v2 (#1959) --- letta/functions/schema_generator.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/letta/functions/schema_generator.py b/letta/functions/schema_generator.py index ace26cf2..2cf79631 100644 --- a/letta/functions/schema_generator.py +++ b/letta/functions/schema_generator.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, Type, Union, get_args, get_origin from docstring_parser import parse from pydantic import BaseModel +from pydantic.v1 import BaseModel as V1BaseModel def is_optional(annotation): @@ -139,7 +140,7 @@ def generate_schema(function, name: Optional[str] = None, description: Optional[ def generate_schema_from_args_schema( - args_schema: Type[BaseModel], name: Optional[str] = None, description: Optional[str] = None, append_heartbeat: bool = True + args_schema: Type[V1BaseModel], name: Optional[str] = None, description: Optional[str] = None, append_heartbeat: bool = True ) -> Dict[str, Any]: properties = {} required = []