Skip to content

Commit

Permalink
fix: Support Chinese characters in prompt generation stage (#1168)
Browse files Browse the repository at this point in the history
Update dataframe_serializer.py
add test case to #1168
  • Loading branch information
cFireworks authored and chuanqing committed Aug 8, 2024
1 parent 456ff73 commit cd3bf42
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions tests/unit_tests/helpers/test_dataframe_serializer.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
import unittest

import pandas as pd

from pandasai.connectors import PandasConnector
from pandasai.helpers.dataframe_serializer import (
DataframeSerializer,
DataframeSerializerType,
)


class TestDataframeSerializer(unittest.TestCase):
def setUp(self):
self.serializer = DataframeSerializer()

def test_convert_df_to_yml(self):
# Test convert df to yml
data = {"name": ["en_name", "中文_名称"]}
connector = PandasConnector(
{"original_df": pd.DataFrame(data)},
name="en_table_name",
description="中文_描述",
field_descriptions={k: k for k in data},
)
result = self.serializer.serialize(
connector,
type_=DataframeSerializerType.YML,
extras={"index": 0, "type": "pd.Dataframe"},
)
self.assertIn("中文_描述", result)

0 comments on commit cd3bf42

Please sign in to comment.