Spaces:
Sleeping
Sleeping
File size: 6,672 Bytes
de6fb09 | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, List
from chatkit.store import NotFoundError, Store
from chatkit.types import Attachment, Page, Thread, ThreadItem, ThreadMetadata
@dataclass
class _ThreadState:
thread: ThreadMetadata
items: List[ThreadItem]
class MemoryStore(Store[dict[str, Any]]):
"""Simple in-memory store compatible with the ChatKit server interface."""
def __init__(self) -> None:
self._threads: Dict[str, _ThreadState] = {}
# Attachments intentionally unsupported; use a real store that enforces auth.
@staticmethod
def _coerce_thread_metadata(thread: ThreadMetadata | Thread) -> ThreadMetadata:
"""Return thread metadata without any embedded items (openai-chatkit>=1.0)."""
has_items = isinstance(thread, Thread) or "items" in getattr(
thread, "model_fields_set", set()
)
if not has_items:
return thread.model_copy(deep=True)
data = thread.model_dump()
data.pop("items", None)
return ThreadMetadata(**data).model_copy(deep=True)
# -- Thread metadata -------------------------------------------------
async def load_thread(self, thread_id: str, context: dict[str, Any]) -> ThreadMetadata:
state = self._threads.get(thread_id)
if not state:
raise NotFoundError(f"Thread {thread_id} not found")
return self._coerce_thread_metadata(state.thread)
async def save_thread(self, thread: ThreadMetadata, context: dict[str, Any]) -> None:
metadata = self._coerce_thread_metadata(thread)
state = self._threads.get(thread.id)
if state:
state.thread = metadata
else:
self._threads[thread.id] = _ThreadState(
thread=metadata,
items=[],
)
async def load_threads(
self,
limit: int,
after: str | None,
order: str,
context: dict[str, Any],
) -> Page[ThreadMetadata]:
threads = sorted(
(self._coerce_thread_metadata(state.thread) for state in self._threads.values()),
key=lambda t: t.created_at or datetime.min,
reverse=(order == "desc"),
)
if after:
index_map = {thread.id: idx for idx, thread in enumerate(threads)}
start = index_map.get(after, -1) + 1
else:
start = 0
slice_threads = threads[start : start + limit + 1]
has_more = len(slice_threads) > limit
slice_threads = slice_threads[:limit]
next_after = slice_threads[-1].id if has_more and slice_threads else None
return Page(
data=slice_threads,
has_more=has_more,
after=next_after,
)
async def delete_thread(self, thread_id: str, context: dict[str, Any]) -> None:
self._threads.pop(thread_id, None)
# -- Thread items ----------------------------------------------------
def _items(self, thread_id: str) -> List[ThreadItem]:
state = self._threads.get(thread_id)
if state is None:
state = _ThreadState(
thread=ThreadMetadata(id=thread_id, created_at=datetime.utcnow()),
items=[],
)
self._threads[thread_id] = state
return state.items
async def load_thread_items(
self,
thread_id: str,
after: str | None,
limit: int,
order: str,
context: dict[str, Any],
) -> Page[ThreadItem]:
items = [item.model_copy(deep=True) for item in self._items(thread_id)]
items.sort(
key=lambda item: getattr(item, "created_at", datetime.utcnow()),
reverse=(order == "desc"),
)
if after:
index_map = {item.id: idx for idx, item in enumerate(items)}
start = index_map.get(after, -1) + 1
else:
start = 0
slice_items = items[start : start + limit + 1]
has_more = len(slice_items) > limit
slice_items = slice_items[:limit]
next_after = slice_items[-1].id if has_more and slice_items else None
return Page(data=slice_items, has_more=has_more, after=next_after)
async def add_thread_item(
self, thread_id: str, item: ThreadItem, context: dict[str, Any]
) -> None:
print("Adding item", item)
self._items(thread_id).append(item.model_copy(deep=True))
async def save_item(self, thread_id: str, item: ThreadItem, context: dict[str, Any]) -> None:
items = self._items(thread_id)
for idx, existing in enumerate(items):
if existing.id == item.id:
items[idx] = item.model_copy(deep=True)
return
items.append(item.model_copy(deep=True))
async def load_item(self, thread_id: str, item_id: str, context: dict[str, Any]) -> ThreadItem:
for item in self._items(thread_id):
if item.id == item_id:
return item.model_copy(deep=True)
raise NotFoundError(f"Item {item_id} not found")
async def delete_thread_item(
self, thread_id: str, item_id: str, context: dict[str, Any]
) -> None:
items = self._items(thread_id)
self._threads[thread_id].items = [item for item in items if item.id != item_id]
# -- Files -----------------------------------------------------------
# These methods are not currently used but required to be compatible with the Store interface.
async def save_attachment(
self,
attachment: Attachment,
context: dict[str, Any],
) -> None:
raise NotImplementedError(
"MemoryStore does not persist attachments. Provide a Store implementation "
"that enforces authentication and authorization before enabling uploads."
)
async def load_attachment(
self,
attachment_id: str,
context: dict[str, Any],
) -> Attachment:
raise NotImplementedError(
"MemoryStore does not load attachments. Provide a Store implementation "
"that enforces authentication and authorization before enabling uploads."
)
async def delete_attachment(self, attachment_id: str, context: dict[str, Any]) -> None:
raise NotImplementedError(
"MemoryStore does not delete attachments because they are never stored."
)
|