Spaces:
Sleeping
Sleeping
File size: 3,755 Bytes
9002555 |
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 |
# Experimental
from sqlalchemy import Column, String, Enum, ForeignKey, DateTime
from sqlalchemy.dialects.postgresql import UUID, ENUM, JSONB
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from enum import Enum
from sqlalchemy.ext.declarative import as_declarative, declared_attr
from llama_index.core.callbacks.schema import CBEventType
# Model
@as_declarative()
class Base:
id = Column(UUID, primary_key=True, index=True, default=func.uuid_generate_v4())
created_at = Column(DateTime, server_default=func.now(), nullable=False)
updated_at = Column(
DateTime, server_default=func.now(), onupdate=func.now(), nullable=False
)
__name__: str
# Generate __tablename__ automatically
@declared_attr
def __tablename__(cls) -> str:
return cls.__name__.lower()
# DB
class MessageRoleEnum(str, Enum):
user = "user"
assistant = "assistant"
class MessageStatusEnum(str, Enum):
PENDING = "PENDING"
SUCCESS = "SUCCESS"
ERROR = "ERROR"
class MessageSubProcessStatusEnum(str, Enum):
PENDING = "PENDING"
FINISHED = "FINISHED"
# python doesn't allow enums to be extended, so we have to do this
additional_message_subprocess_fields = {
"CONSTRUCTED_QUERY_ENGINE": "constructed_query_engine",
"SUB_QUESTIONS": "sub_questions",
}
MessageSubProcessSourceEnum = Enum(
"MessageSubProcessSourceEnum",
[(event_type.name, event_type.value) for event_type in CBEventType]
+ list(additional_message_subprocess_fields.items()),
)
def to_pg_enum(enum_class) -> ENUM:
return ENUM(enum_class, name=enum_class.__name__)
class Document(Base):
"""
A document along with its metadata
"""
# URL to the actual document (e.g. a PDF)
url = Column(String, nullable=False, unique=True)
metadata_map = Column(JSONB, nullable=True)
conversations = relationship("ConversationDocument", back_populates="document")
class Conversation(Base):
"""
A conversation with messages and linked documents
"""
messages = relationship("Message", back_populates="conversation")
conversation_documents = relationship(
"ConversationDocument", back_populates="conversation"
)
class ConversationDocument(Base):
"""
A many-to-many relationship between a conversation and a document
"""
conversation_id = Column(
UUID(as_uuid=True), ForeignKey("conversation.id"), index=True
)
document_id = Column(UUID(as_uuid=True), ForeignKey("document.id"), index=True)
conversation = relationship("Conversation", back_populates="conversation_documents")
document = relationship("Document", back_populates="conversations")
class Message(Base):
"""
A message in a conversation
"""
conversation_id = Column(
UUID(as_uuid=True), ForeignKey("conversation.id"), index=True
)
content = Column(String)
role = Column(to_pg_enum(MessageRoleEnum))
status = Column(to_pg_enum(MessageStatusEnum), default=MessageStatusEnum.PENDING)
conversation = relationship("Conversation", back_populates="messages")
sub_processes = relationship("MessageSubProcess", back_populates="message")
class MessageSubProcess(Base):
"""
A record of a sub-process that occurred as part of the generation of a message from an AI assistant
"""
message_id = Column(UUID(as_uuid=True), ForeignKey("message.id"), index=True)
source = Column(to_pg_enum(MessageSubProcessSourceEnum))
message = relationship("Message", back_populates="sub_processes")
status = Column(
to_pg_enum(MessageSubProcessStatusEnum),
default=MessageSubProcessStatusEnum.FINISHED,
nullable=False,
)
metadata_map = Column(JSONB, nullable=True) |