init project

metagpt/schema.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/8 22:12
+@Author  : alexanderwu
+@File    : schema.py
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Optional, Type, TypedDict
+
+from metagpt.logs import logger
+# from pydantic import BaseModel
+
+# from metagpt.actions import Action
+
+
+class RawMessage(TypedDict):
+    content: str
+    role: str
+
+
+@dataclass
+class Message:
+    """list[<role>: <content>]"""
+    content: str
+    role: str = field(default='user') # system / user / assistant
+    cause_by: Type["Action"] = field(default="")
+
+    def __str__(self):
+        # prefix = '-'.join([self.role, str(self.cause_by)])
+        return f"{self.role}: {self.content}"
+
+    def __repr__(self):
+        return self.__str__()
+
+    def to_dict(self) -> dict:
+        return {
+            "role": self.role,
+            "content": self.content
+        }
+
+
+@dataclass
+class UserMessage(Message):
+    """便于支持OpenAI的消息"""
+    def __init__(self, content: str):
+        super().__init__(content, 'user')
+
+
+@dataclass
+class SystemMessage(Message):
+    """便于支持OpenAI的消息"""
+    def __init__(self, content: str):
+        super().__init__(content, 'system')
+
+
+@dataclass
+class AIMessage(Message):
+    """便于支持OpenAI的消息"""
+    def __init__(self, content: str):
+        super().__init__(content, 'assistant')
+
+
+if __name__ == '__main__':
+    test_content = 'test_message'
+    msgs = [
+        UserMessage(test_content),
+        SystemMessage(test_content),
+        AIMessage(test_content),
+        Message(test_content, role='QA')
+    ]
+    logger.info(msgs)