Pydantic Fixes and Tests (#311)

* Added SubscriptableBaseModel to the Model classes and added Image codec test

---------

Co-authored-by: Parth Sareen <parth@Parths-MacBook-Pro.local>
This commit is contained in:
Parth Sareen
2024-11-06 14:04:56 -08:00
committed by Michael Yang
parent 0bbc246007
commit f25834217b
2 changed files with 18 additions and 2 deletions
+15
View File
@@ -0,0 +1,15 @@
from base64 import b64decode, b64encode
from ollama._types import Image
def test_image_serialization():
# Test bytes serialization
image_bytes = b'test image bytes'
img = Image(value=image_bytes)
assert img.model_dump() == b64encode(image_bytes).decode()
# Test base64 string serialization
b64_str = 'dGVzdCBiYXNlNjQgc3RyaW5n'
img = Image(value=b64_str)
assert img.model_dump() == b64decode(b64_str).decode()