多场景对架构一期

pilot/scene/message.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+from datetime import datetime, timedelta
+from pydantic import BaseModel, Field, root_validator, validator
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+)
+
+from pilot.scene.base_message import BaseMessage, AIMessage, HumanMessage, SystemMessage, messages_to_dict, \
+    messages_from_dict
+
+
+class OnceConversation:
+    """
+       All the information of a conversation, the current single service in memory, can expand cache and database support distributed services
+    """
+
+    def __init__(self):
+        self.messages: List[BaseMessage] = []
+        self.start_date: str = ""
+        self.chat_order: int = 0
+        self.cost: int = 0
+        self.tokens: int = 0
+
+    def add_user_message(self, message: str) -> None:
+        """Add a user message to the store"""
+        self.messages.append(HumanMessage(content=message))
+
+    def add_ai_message(self, message: str) -> None:
+        """Add an AI message to the store"""
+        self.messages.append(AIMessage(content=message))
+        """  """
+
+    def add_system_message(self, message: str) -> None:
+        """Add an AI message to the store"""
+        self.messages.append(SystemMessage(content=message))
+
+    def set_start_time(self, datatime: datetime):
+        dt_str = datatime.strftime("%Y-%m-%d %H:%M:%S")
+        self.start_date = dt_str;
+
+    def clear(self) -> None:
+        """Remove all messages from the store"""
+        self.messages.clear()
+        self.session_id = None
+
+
+def _conversation_to_dic(once: OnceConversation) -> dict:
+    start_str: str = ""
+    if once.start_date:
+        if isinstance(once.start_date, datetime):
+            start_str = once.start_date.strftime("%Y-%m-%d %H:%M:%S")
+        else:
+            start_str = once.start_date
+
+    return {
+        "chat_order": once.chat_order,
+        "start_date": start_str,
+        "cost": once.cost if once.cost else 0,
+        "tokens": once.tokens if once.tokens else 0,
+        "messages": messages_to_dict(once.messages)
+    }
+
+
+def conversations_to_dict(conversations: List[OnceConversation]) -> List[dict]:
+    return [_conversation_to_dic(m) for m in conversations]
+
+
+def conversation_from_dict(once: dict) -> OnceConversation:
+    conversation = OnceConversation()
+    conversation.cost = once.get('cost', 0)
+    conversation.tokens = once.get('tokens', 0)
+    conversation.start_date = once.get('start_date', '')
+    conversation.chat_order = int(once.get('chat_order'))
+    print(once.get('messages'))
+    conversation.messages = messages_from_dict(once.get('messages', []))
+    return conversation