Spaces:
Sleeping
Sleeping
File size: 2,245 Bytes
8428312 82915e5 8428312 82915e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import json
import csv
from io import StringIO
from pydantic import BaseModel, validator, ValidationError
class Card(BaseModel):
question: str
answer: str
class Message(BaseModel):
role: str
content: list[Card]
@validator('content', pre=True)
def parse_content(cls, v: str) -> 'Message':
if isinstance(v, str):
try:
content_list = json.loads(v)
return content_list
except json.JSONDecodeError as e:
raise ValueError(f"Error decoding 'content' JSON: {e}") from e
return v
def content_to_json(self) -> str:
return json.dumps([card.dict() for card in self.content], indent=2)
def content_to_csv(self) -> str:
"""
Converts the content of the Message instance into a CSV string.
"""
output = StringIO()
# Step 2: Create a CSV writer - windows style is the default, so set Unix-style line endings
writer = csv.writer(output, lineterminator='\n')
writer.writerow(["Question", "Answer"])
for card in self.content:
writer.writerow([card.question, card.answer])
csv_content = output.getvalue()
output.close()
return csv_content
class PydanticEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, BaseModel):
return obj.dict()
return super().default(obj)
def parse_message(input_dict: dict[str, any]) -> Message:
try:
# Extract the role
role: str = input_dict['role']
# Parse the content
content: str = input_dict['content']
# If content is a string, try to parse it as JSON
if isinstance(content, str):
content = json.loads(content)
# Create Card objects from the content
cards = [Card(**item) for item in content]
# Create and return the Message object
return Message(role=role, content=cards)
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON in content: {str(e)}")
except ValidationError as e:
raise ValueError(f"Validation error: {str(e)}")
except KeyError as e:
raise ValueError(f"Missing required key: {str(e)}")
|