Skip to content

Commit

Permalink
use custom JSON type for database for more generic support (#389)
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier authored Apr 25, 2024
1 parent 3a5b82d commit 7963453
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 3 deletions.
2 changes: 1 addition & 1 deletion ragna/deploy/_api/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ def _orm_to_schema_chat(chat: orm.Chat) -> schemas.Chat:
documents=documents,
source_storage=chat.source_storage,
assistant=chat.assistant,
params=chat.params, # type: ignore[arg-type]
params=chat.params,
),
messages=messages,
prepared=chat.prepared,
Expand Down
30 changes: 28 additions & 2 deletions ragna/deploy/_api/orm.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,35 @@
import json
from typing import Any

from sqlalchemy import Column, ForeignKey, Table, types
from sqlalchemy.engine import Dialect
from sqlalchemy.orm import DeclarativeBase, relationship # type: ignore[attr-defined]

from ragna.core import MessageRole


class Json(types.TypeDecorator):
"""Universal JSON type which stores values as strings.
This is needed because sqlalchemy.types.JSON only works for a limited subset of
databases.
"""

impl = types.String

cache_ok = True

def process_bind_param(self, value: Any, dialect: Dialect) -> str:
return json.dumps(value)

def process_result_value(
self,
value: str, # type: ignore[override]
dialect: Dialect,
) -> Any:
return json.loads(value)


class Base(DeclarativeBase):
pass

Expand Down Expand Up @@ -34,7 +60,7 @@ class Document(Base):
name = Column(types.String)
# Mind the trailing underscore here. Unfortunately, this is necessary, because
# metadata without the underscore is reserved by SQLAlchemy
metadata_ = Column(types.JSON)
metadata_ = Column(Json)
chats = relationship(
"Chat",
secondary=document_chat_association_table,
Expand All @@ -59,7 +85,7 @@ class Chat(Base):
)
source_storage = Column(types.String)
assistant = Column(types.String)
params = Column(types.JSON)
params = Column(Json)
messages = relationship("Message", cascade="all, delete")
prepared = Column(types.Boolean)

Expand Down

0 comments on commit 7963453

Please sign in to comment.