diff --git a/.opencode/opencode.db-shm b/.opencode/opencode.db-shm new file mode 100644 index 0000000..a8e275b Binary files /dev/null and b/.opencode/opencode.db-shm differ diff --git a/.opencode/opencode.db-wal b/.opencode/opencode.db-wal new file mode 100644 index 0000000..ed64eef Binary files /dev/null and b/.opencode/opencode.db-wal differ diff --git a/src/distill/queue.py b/src/distill/queue.py new file mode 100644 index 0000000..18c4315 --- /dev/null +++ b/src/distill/queue.py @@ -0,0 +1,89 @@ +import sqlite3 +import json +import threading +from datetime import datetime +from pathlib import Path +from typing import Optional +import sys +sys.path.insert(0, str(Path(__file__).parent.parent)) +from models.episode import Episode + +class PersistenceQueue: + def __init__(self, db_path: str = 'zhiyi.db'): + self.db_path = db_path + self.lock = threading.Lock() + self.conn = sqlite3.connect(db_path, check_same_thread=False) + self.conn.row_factory = sqlite3.Row + self._init_db() + + def _init_db(self): + self.conn.execute(''' + CREATE TABLE IF NOT EXISTS episode_queue ( + id TEXT PRIMARY KEY, + data TEXT NOT NULL, + enqueued_at TEXT NOT NULL, + dequeued_at TEXT, + status TEXT DEFAULT 'pending' + ) + ''') + self.conn.execute('CREATE INDEX IF NOT EXISTS idx_status ON episode_queue(status)') + self.conn.commit() + + def enqueue(self, episode: Episode) -> bool: + with self.lock: + try: + data = json.dumps(episode.to_dict()) + self.conn.execute( + 'INSERT INTO episode_queue (id, data, enqueued_at, status) VALUES (?, ?, ?, ?)', + (episode.id, data, datetime.now().isoformat(), 'pending') + ) + self.conn.commit() + return True + except sqlite3.IntegrityError: + return False + + def dequeue(self) -> Optional[dict]: + with self.lock: + row = self.conn.execute( + 'SELECT * FROM episode_queue WHERE status = ? ORDER BY enqueued_at ASC LIMIT 1', + ('pending',) + ).fetchone() + if row is None: + return None + self.conn.execute( + 'UPDATE episode_queue SET status = ?, dequeued_at = ? WHERE id = ?', + ('dequeued', datetime.now().isoformat(), row['id']) + ) + self.conn.commit() + return json.loads(row['data']) + + def peek(self) -> Optional[dict]: + with self.lock: + row = self.conn.execute( + 'SELECT data FROM episode_queue WHERE status = ? ORDER BY enqueued_at ASC LIMIT 1', + ('pending',) + ).fetchone() + return json.loads(row['data']) if row else None + + def size(self) -> int: + with self.lock: + row = self.conn.execute( + 'SELECT COUNT(*) as cnt FROM episode_queue WHERE status = ?', + ('pending',) + ).fetchone() + return row['cnt'] if row else 0 + + def is_empty(self) -> bool: + return self.size() == 0 + + def requeue(self, episode: Episode) -> bool: + with self.lock: + self.conn.execute( + 'UPDATE episode_queue SET status = ?, dequeued_at = ? WHERE id = ?', + ('pending', None, episode.id) + ) + self.conn.commit() + return True + + def close(self): + self.conn.close() \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000..4da3e42 --- /dev/null +++ b/src/models/__init__.py @@ -0,0 +1,4 @@ +from .episode import Episode +from .distilled import Distilled + +__all__ = ["Episode", "Distilled"] \ No newline at end of file diff --git a/src/models/distilled.py b/src/models/distilled.py new file mode 100644 index 0000000..a8e14ff --- /dev/null +++ b/src/models/distilled.py @@ -0,0 +1,47 @@ +from dataclasses import dataclass, field +from datetime import datetime + +TYPE_DECISION = 'decision' +TYPE_REQUEST = 'request' +TYPE_FACT = 'fact' +TYPE_PATTERN = 'pattern' + +STATUS_PENDING = 'pending' +STATUS_VALIDATED = 'validated' +STATUS_DEPRECATED = 'deprecated' + +@dataclass +class Distilled: + id: str + episode_id: str + type: str + summary: str + entities: list[str] = field(default_factory=list) + facts: list[str] = field(default_factory=list) + confidence: float = 0.5 + status: str = 'pending' + importance: int = 0 + created_at: datetime = field(default_factory=datetime.now) + updated_at: datetime = field(default_factory=datetime.now) + + def to_dict(self) -> dict: + return { + 'id': self.id, + 'episode_id': self.episode_id, + 'type': self.type, + 'summary': self.summary, + 'entities': self.entities, + 'facts': self.facts, + 'confidence': self.confidence, + 'status': self.status, + 'importance': self.importance, + 'created_at': self.created_at.isoformat(), + 'updated_at': self.updated_at.isoformat(), + } + + @classmethod + def from_dict(cls, d: dict): + d = d.copy() + d['created_at'] = datetime.fromisoformat(d['created_at']) + d['updated_at'] = datetime.fromisoformat(d['updated_at']) + return cls(**d) \ No newline at end of file diff --git a/src/models/episode.py b/src/models/episode.py new file mode 100644 index 0000000..2e6f54a --- /dev/null +++ b/src/models/episode.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass, field +from datetime import datetime +from typing import Optional +import uuid + +@dataclass +class Episode: + id: str + timestamp: datetime + content: str + entities: list[str] = field(default_factory=list) + facts: list[str] = field(default_factory=list) + metadata: dict = field(default_factory=dict) + source: str = 'hermes' + + def to_dict(self) -> dict: + return { + 'id': self.id, + 'timestamp': self.timestamp.isoformat(), + 'content': self.content, + 'entities': self.entities, + 'facts': self.facts, + 'metadata': self.metadata, + 'source': self.source, + } + + @classmethod + def from_dict(cls, d: dict): + d = d.copy() + d['timestamp'] = datetime.fromisoformat(d['timestamp']) + return cls(**d) + + @classmethod + def create(cls, content: str, source: str = 'hermes', entities: list = None, facts: list = None, metadata: dict = None): + return cls( + id=str(uuid.uuid4()), + timestamp=datetime.now(), + content=content, + entities=entities or [], + facts=facts or [], + metadata=metadata or {}, + source=source, + ) \ No newline at end of file diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..65ab79b --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,48 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) +from datetime import datetime +from models.episode import Episode +from models.distilled import Distilled, TYPE_FACT, STATUS_PENDING, STATUS_VALIDATED + +def test_episode_to_dict_from_dict(): + ep = Episode.create(content='test content', source='test', entities=['entity1'], facts=['fact1']) + d = ep.to_dict() + assert d['content'] == 'test content' + assert d['source'] == 'test' + ep2 = Episode.from_dict(d) + assert ep2.content == ep.content + assert ep2.id == ep.id + +def test_episode_create(): + ep = Episode.create(content='hello') + assert ep.id is not None + assert ep.content == 'hello' + assert ep.source == 'hermes' + assert isinstance(ep.timestamp, datetime) + +def test_distilled_to_dict_from_dict(): + ep = Episode.create(content='test') + d = Distilled( + id='d1', episode_id=ep.id, type=TYPE_FACT, + summary='a summary', confidence=0.8, status=STATUS_PENDING + ) + dd = d.to_dict() + assert dd['type'] == TYPE_FACT + assert dd['confidence'] == 0.8 + d2 = Distilled.from_dict(dd) + assert d2.id == d.id + assert d2.type == d.type + +def test_distilled_defaults(): + d = Distilled(id='d1', episode_id='e1', type=TYPE_FACT, summary='s') + assert d.status == STATUS_PENDING + assert d.confidence == 0.5 + assert d.importance == 0 + +if __name__ == '__main__': + test_episode_to_dict_from_dict() + test_episode_create() + test_distilled_to_dict_from_dict() + test_distilled_defaults() + print('All model tests passed!') \ No newline at end of file diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 0000000..4dfd2f2 --- /dev/null +++ b/tests/test_queue.py @@ -0,0 +1,67 @@ +import sys +import os +import threading +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent / 'src')) +from models.episode import Episode +from distill.queue import PersistenceQueue + +def test_queue_basic(): + q = PersistenceQueue(':memory:') + ep = Episode.create(content='test1') + assert q.enqueue(ep) == True + assert q.size() == 1 + assert q.is_empty() == False + dequeued = q.dequeue() + assert dequeued is not None + assert dequeued['content'] == 'test1' + assert q.size() == 0 + assert q.is_empty() == True + q.close() + +def test_queue_requeue(): + q = PersistenceQueue(':memory:') + ep = Episode.create(content='test2') + q.enqueue(ep) + deq = q.dequeue() + d = Episode.from_dict(deq) + assert q.requeue(d) == True + assert q.size() == 1 + q.close() + +def test_queue_peek(): + q = PersistenceQueue(':memory:') + ep1 = Episode.create(content='first') + ep2 = Episode.create(content='second') + q.enqueue(ep1) + q.enqueue(ep2) + peeked = q.peek() + assert peeked['content'] == 'first' + assert q.size() == 2 # peek does not remove + q.close() + +def test_queue_concurrent(): + q = PersistenceQueue(':memory:') + errors = [] + def enqueue_many(n): + try: + for i in range(n): + ep = Episode.create(content=f'thread-{i}') + q.enqueue(ep) + except Exception as e: + errors.append(str(e)) + threads = [threading.Thread(target=enqueue_many, args=(20,)) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + assert len(errors) == 0, errors + assert q.size() == 100 + q.close() + +if __name__ == '__main__': + test_queue_basic() + test_queue_requeue() + test_queue_peek() + test_queue_concurrent() + print('All queue tests passed!') \ No newline at end of file