Spaces:
Sleeping
Sleeping
File size: 3,801 Bytes
d8d14f1 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
from __future__ import annotations
from abc import abstractmethod
from typing import Sequence
class Message:
"""
The base abstract Message class.
Messages are the inputs and outputs of ChatModels.
"""
def __init__(
self, content: str, role: str, additional_kwargs: dict = None
):
self.content = content
self.role = role
self.additional_kwargs = (
additional_kwargs if additional_kwargs else {}
)
@abstractmethod
def get_type(self) -> str:
pass
class HumanMessage(Message):
"""
A Message from a human.
"""
def __init__(
self,
content: str,
role: str = "Human",
additional_kwargs: dict = None,
example: bool = False,
):
super().__init__(content, role, additional_kwargs)
self.example = example
def get_type(self) -> str:
return "human"
class AIMessage(Message):
"""
A Message from an AI.
"""
def __init__(
self,
content: str,
role: str = "AI",
additional_kwargs: dict = None,
example: bool = False,
):
super().__init__(content, role, additional_kwargs)
self.example = example
def get_type(self) -> str:
return "ai"
class SystemMessage(Message):
"""
A Message for priming AI behavior, usually passed in as the first of a sequence
of input messages.
"""
def __init__(
self,
content: str,
role: str = "System",
additional_kwargs: dict = None,
):
super().__init__(content, role, additional_kwargs)
def get_type(self) -> str:
return "system"
class FunctionMessage(Message):
"""
A Message for passing the result of executing a function back to a model.
"""
def __init__(
self,
content: str,
role: str = "Function",
name: str = None,
additional_kwargs: dict = None,
):
super().__init__(content, role, additional_kwargs)
self.name = name
def get_type(self) -> str:
return "function"
class ChatMessage(Message):
"""
A Message that can be assigned an arbitrary speaker (i.e. role).
"""
def __init__(
self, content: str, role: str, additional_kwargs: dict = None
):
super().__init__(content, role, additional_kwargs)
def get_type(self) -> str:
return "chat"
def get_buffer_string(
messages: Sequence[Message],
human_prefix: str = "Human",
ai_prefix: str = "AI",
) -> str:
string_messages = []
for m in messages:
message = f"{m.role}: {m.content}"
if (
isinstance(m, AIMessage)
and "function_call" in m.additional_kwargs
):
message += f"{m.additional_kwargs['function_call']}"
string_messages.append(message)
return "\n".join(string_messages)
def message_to_dict(message: Message) -> dict:
return {"type": message.get_type(), "data": message.__dict__}
def messages_to_dict(messages: Sequence[Message]) -> list[dict]:
return [message_to_dict(m) for m in messages]
def message_from_dict(message: dict) -> Message:
_type = message["type"]
if _type == "human":
return HumanMessage(**message["data"])
elif _type == "ai":
return AIMessage(**message["data"])
elif _type == "system":
return SystemMessage(**message["data"])
elif _type == "chat":
return ChatMessage(**message["data"])
elif _type == "function":
return FunctionMessage(**message["data"])
else:
raise ValueError(f"Got unexpected message type: {_type}")
def messages_from_dict(messages: list[dict]) -> list[Message]:
return [message_from_dict(m) for m in messages]
|