from typing import Optional, List, Dict, Any, TYPE_CHECKING, Union

from pydantic import BaseModel, validator

from ..math.ndarray import to_list

if TYPE_CHECKING:
    from ..types import ArrayType

_ProtoValueType = Optional[Union[str, bool, float]]
_StructValueType = Union[
    _ProtoValueType, List[_ProtoValueType], Dict[str, _ProtoValueType]
]


def _convert_ndarray_to_list(v: 'ArrayType'):
    if v is not None:
        return to_list(v)


class PydanticDocument(BaseModel):
    id: str
    parent_id: Optional[str]
    granularity: Optional[int]
    adjacency: Optional[int]
    blob: Optional[bytes]
    tensor: Optional[Any]
    mime_type: Optional[str]
    text: Optional[str]
    weight: Optional[float]
    uri: Optional[str]
    tags: Optional[Dict[str, '_StructValueType']]
    offset: Optional[float]
    location: Optional[List[float]]
    embedding: Optional[Any]
    modality: Optional[str]
    evaluations: Optional[Dict[str, Dict[str, '_StructValueType']]]
    scores: Optional[Dict[str, Dict[str, '_StructValueType']]]
    chunks: Optional[List['PydanticDocument']]
    matches: Optional[List['PydanticDocument']]

    _tensor2list = validator('tensor', allow_reuse=True)(_convert_ndarray_to_list)
    _embedding2list = validator('embedding', allow_reuse=True)(_convert_ndarray_to_list)


PydanticDocument.update_forward_refs()

PydanticDocumentArray = List[PydanticDocument]

docarray/document/pydantic_model.py
from typing import Optional, List, Dict, Any, TYPE_CHECKING, Union

from pydantic import BaseModel, validator

from ..math.ndarray import to_list

if TYPE_CHECKING:
    from ..types import ArrayType

# this order must be preserved: https://pydantic-docs.helpmanual.io/usage/types/#unions
_ProtoValueType = Optional[Union[bool, float, str]]
_StructValueType = Union[
    _ProtoValueType, List[_ProtoValueType], Dict[str, _ProtoValueType]
]


def _convert_ndarray_to_list(v: 'ArrayType'):
    if v is not None:
        return to_list(v)


class PydanticDocument(BaseModel):
    id: str
    parent_id: Optional[str]
    granularity: Optional[int]
    adjacency: Optional[int]
    blob: Optional[bytes]
    tensor: Optional[Any]
    mime_type: Optional[str]
    text: Optional[str]
    weight: Optional[float]
    uri: Optional[str]
    tags: Optional[Dict[str, '_StructValueType']]
    offset: Optional[float]
    location: Optional[List[float]]
    embedding: Optional[Any]
    modality: Optional[str]
    evaluations: Optional[Dict[str, Dict[str, '_StructValueType']]]
    scores: Optional[Dict[str, Dict[str, '_StructValueType']]]
    chunks: Optional[List['PydanticDocument']]
    matches: Optional[List['PydanticDocument']]

    _tensor2list = validator('tensor', allow_reuse=True)(_convert_ndarray_to_list)
    _embedding2list = validator('embedding', allow_reuse=True)(_convert_ndarray_to_list)


PydanticDocument.update_forward_refs()

PydanticDocumentArray = List[PydanticDocument]

docarray/document/pydantic_model.py
diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py
index bc4cb2a9aa7..2be077c75eb 100644
--- a/docarray/document/pydantic_model.py
+++ b/docarray/document/pydantic_model.py
@@ -7,7 +7,8 @@
 if TYPE_CHECKING:
     from ..types import ArrayType
 
-_ProtoValueType = Optional[Union[str, bool, float]]
+# this order must be preserved: https://pydantic-docs.helpmanual.io/usage/types/#unions
+_ProtoValueType = Optional[Union[bool, float, str]]
 _StructValueType = Union[
     _ProtoValueType, List[_ProtoValueType], Dict[str, _ProtoValueType]
 ]
diff --git a/tests/__init__.py b/tests/__init__.py
index 25ab626a64f..f358cb479bc 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -20,7 +20,7 @@ def random_docs(
 
         d = Document(id=doc_id)
         d.text = text
-        d.tags['id'] = doc_id
+        d.tags['id'] = f'myself id is: {doc_id}'
         if embedding:
             if sparse_embedding:
                 from scipy.sparse import coo_matrix
@@ -42,8 +42,8 @@ def random_docs(
                 c.embedding = np.random.random(
                     [embed_dim + np.random.randint(0, jitter)]
                 )
-            c.tags['parent_id'] = doc_id
-            c.tags['id'] = chunk_doc_id
+            c.tags['parent_id'] = f'my parent is: {id}'
+            c.tags['id'] = f'myself id is: {doc_id}'
             d.chunks.append(c)
             next_chunk_doc_id += 1
 
diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py
index b6b43216db6..9f9c3e1c22f 100644
--- a/tests/unit/test_pydantic.py
+++ b/tests/unit/test_pydantic.py
@@ -2,6 +2,7 @@
 from typing import List, Optional
 
 import numpy as np
+import pytest
 from fastapi import FastAPI
 from pydantic import BaseModel
 from starlette.testclient import TestClient
@@ -116,3 +117,28 @@ def test_match_to_from_pydantic():
 def test_with_embedding_no_tensor():
     d = Document(embedding=np.random.rand(2, 2))
     PydanticDocument.parse_obj(d.to_pydantic_model().dict())
+
+
+@pytest.mark.parametrize(
+    'tag_value, tag_type',
+    [(3, float), (3.4, float), ('hello', str), (True, bool), (False, bool)],
+)
+@pytest.mark.parametrize('protocol', ['protobuf', 'jsonschema'])
+def test_tags_int_float_str_bool(tag_type, tag_value, protocol):
+    d = Document(tags={'hello': tag_value})
+    dd = d.to_dict(protocol=protocol)['tags']['hello']
+    assert dd == tag_value
+    assert isinstance(dd, tag_type)
+
+    # now nested tags in dict
+
+    d = Document(tags={'hello': {'world': tag_value}})
+    dd = d.to_dict(protocol=protocol)['tags']['hello']['world']
+    assert dd == tag_value
+    assert isinstance(dd, tag_type)
+
+    # now nested in list
+    d = Document(tags={'hello': [tag_value] * 10})
+    dd = d.to_dict(protocol=protocol)['tags']['hello'][-1]
+    assert dd == tag_value
+    assert isinstance(dd, tag_type)

fix: fix tags type after pydantic model

