diff --git a/aisuite/framework/chat_completion_response.py b/aisuite/framework/chat_completion_response.py index ef13fc2c..5e985ea1 100644 --- a/aisuite/framework/chat_completion_response.py +++ b/aisuite/framework/chat_completion_response.py @@ -6,3 +6,20 @@ class ChatCompletionResponse: def __init__(self): self.choices = [Choice()] # Adjust the range as needed for more choices + + def __repr_name__(self) -> str: + """Name of the instance's class, used in __repr__.""" + return self.__class__.__name__ + + def __repr_args__(self): + attrs_names = self.__slots__ if hasattr(self, '__slots__') else None + if not attrs_names and hasattr(self, '__dict__'): + attrs_names = self.__dict__.keys() + attrs = ((s, getattr(self, s)) for s in attrs_names) + return [(a, v) for a, v in attrs] + + def __repr_str__(self, join_str: str) -> str: + return join_str.join(repr(v) if a is None else f'{a}={v!r}' for a, v in self.__repr_args__()) + + def __repr__(self): + return f'{self.__repr_name__()}({self.__repr_str__(", ")})' diff --git a/aisuite/framework/choice.py b/aisuite/framework/choice.py index 38d128b3..92a0a510 100644 --- a/aisuite/framework/choice.py +++ b/aisuite/framework/choice.py @@ -2,7 +2,8 @@ from typing import Literal, Optional, List -class Choice: +class Choice(): + def __init__(self): self.finish_reason: Optional[Literal["stop", "tool_calls"]] = None self.message = Message( @@ -13,3 +14,20 @@ def __init__(self): reasoning_content=None, ) self.intermediate_messages: List[Message] = [] + + def __repr_name__(self) -> str: + """Name of the instance's class, used in __repr__.""" + return self.__class__.__name__ + + def __repr_args__(self): + attrs_names = self.__slots__ if hasattr(self, '__slots__') else None + if not attrs_names and hasattr(self, '__dict__'): + attrs_names = self.__dict__.keys() + attrs = ((s, getattr(self, s)) for s in attrs_names) + return [(a, v) for a, v in attrs] + + def __repr_str__(self, join_str: str) -> str: + return join_str.join(repr(v) if a is None else f'{a}={v!r}' for a, v in self.__repr_args__()) + + def __repr__(self): + return f'{self.__repr_name__()}({self.__repr_str__(", ")})'