feat: Phase 1.1 数据模型 + 持久化队列
- Episode 模型 (src/models/episode.py) - Distilled 模型 (src/models/distilled.py) - SQLite 持久化队列 (src/distill/queue.py) - 模型测试 (tests/test_models.py) - 队列测试 (tests/test_queue.py)
This commit is contained in:
xiaowei
parent
2e503d82ae
commit
1802ffd40b
Binary file not shown.
Binary file not shown.
|
|
@ -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()
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .episode import Episode
|
||||
from .distilled import Distilled
|
||||
|
||||
__all__ = ["Episode", "Distilled"]
|
||||
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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!')
|
||||
|
|
@ -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!')
|
||||
Loading…
Reference in New Issue