kanshan/ollama-python
1
0
Fork 0
mirror of https://github.com/ollama/ollama-python.git synced 2026-01-13 21:57:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Pydantic Fixes and Tests (#311)

Browse Source 
* 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
ollama/_types.py
View File

@ -97,6 +97,7 @@ class BaseGenerateRequest(BaseStreamableRequest):
class Image(BaseModel):
value: Union[FilePath, Base64Str, bytes]
# This overloads the `model_dump` method and returns values depending on the type of the `value` field
@model_serializer
def serialize_model(self):
if isinstance(self.value, Path):
@ -333,7 +334,7 @@ class ModelDetails(SubscriptableBaseModel):
class ListResponse(SubscriptableBaseModel):
class Model(BaseModel):
class Model(SubscriptableBaseModel):
modified_at: Optional[datetime] = None
digest: Optional[str] = None
size: Optional[ByteSize] = None
@ -394,7 +395,7 @@ class ShowResponse(SubscriptableBaseModel):
class ProcessResponse(SubscriptableBaseModel):
class Model(BaseModel):
class Model(SubscriptableBaseModel):
model: Optional[str] = None
name: Optional[str] = None
digest: Optional[str] = None

15
tests/test_type_serialization.py Normal file
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()