chore: init monorepo snapshot
This commit is contained in:
liaibo
@@ -0,0 +1,35 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EmbeddingGenerator:
|
||||
"""简单的 OpenAI 兼容嵌入生成器"""
|
||||
|
||||
def __init__(self, api_base: str, api_key: str, model: str, timeout: float = 30.0) -> None:
|
||||
self.api_base = api_base.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
|
||||
async def embed_texts(self, texts: List[str]) -> List[Optional[List[float]]]:
|
||||
if not texts:
|
||||
return []
|
||||
payload = {"input": texts, "model": self.model}
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
resp = await client.post(f"{self.api_base}/embeddings", json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
embeddings = []
|
||||
for item in data.get("data", []):
|
||||
embeddings.append(item.get("embedding"))
|
||||
if len(embeddings) != len(texts):
|
||||
logger.warning("嵌入结果数量与输入不一致：%s != %s", len(embeddings), len(texts))
|
||||
return embeddings
|
||||
@@ -0,0 +1,24 @@
|
||||
import asyncio
|
||||
from typing import Dict, Optional
|
||||
|
||||
from app.models.job import IngestJob
|
||||
|
||||
|
||||
class InMemoryJobStore:
|
||||
"""简单的内存 JobStore，后续可替换为 Supabase 持久化"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._jobs: Dict[str, IngestJob] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def save(self, job: IngestJob) -> None:
|
||||
async with self._lock:
|
||||
self._jobs[job.id] = job
|
||||
|
||||
async def get(self, job_id: str) -> Optional[IngestJob]:
|
||||
async with self._lock:
|
||||
return self._jobs.get(job_id)
|
||||
|
||||
async def all(self) -> Dict[str, IngestJob]:
|
||||
async with self._lock:
|
||||
return dict(self._jobs)
|
||||
@@ -0,0 +1,42 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.models.job import IngestJob
|
||||
from app.services.embeddings import EmbeddingGenerator
|
||||
from app.services.job_store import InMemoryJobStore
|
||||
from app.services.scheduler import SchedulerManager
|
||||
from app.services.tasks import IngestWorker
|
||||
from app.services.webhooks import LightRAGWebhook
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
sys.path.append(str(ROOT / "packages" / "siyuan_ingest" / "src"))
|
||||
from siyuan_ingest import SiYuanClient, SupabaseSync, SupabaseWriter # type: ignore # noqa: E402
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
job_store = InMemoryJobStore()
|
||||
scheduler = SchedulerManager()
|
||||
siyuan_client = SiYuanClient()
|
||||
webhook = LightRAGWebhook()
|
||||
embedding_generator = EmbeddingGenerator(
|
||||
api_base=settings.ollama_base_url,
|
||||
api_key=settings.ollama_api_key,
|
||||
model=settings.embeddings_model,
|
||||
)
|
||||
supabase_sync = SupabaseSync(
|
||||
supabase_url=settings.supabase_url + "/rest/v1",
|
||||
supabase_key=settings.supabase_key,
|
||||
)
|
||||
supabase_writer = SupabaseWriter(
|
||||
supabase_url=settings.supabase_url + "/rest/v1",
|
||||
supabase_key=settings.supabase_key,
|
||||
)
|
||||
worker = IngestWorker(
|
||||
job_store=job_store,
|
||||
siyuan_client=siyuan_client,
|
||||
webhook=webhook,
|
||||
supabase_sync=supabase_sync,
|
||||
supabase_writer=supabase_writer,
|
||||
embedding_generator=embedding_generator,
|
||||
)
|
||||
@@ -0,0 +1,41 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Awaitable, Callable, Dict, Optional
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from apscheduler.triggers.interval import IntervalTrigger
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SchedulerManager:
|
||||
"""封装 APScheduler，提供周期性增量任务"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.scheduler = AsyncIOScheduler()
|
||||
self._jobs: Dict[str, str] = {}
|
||||
|
||||
def start(self) -> None:
|
||||
if not self.scheduler.running:
|
||||
self.scheduler.start()
|
||||
logger.info("APScheduler 启动完成")
|
||||
|
||||
def shutdown(self) -> None:
|
||||
if self.scheduler.running:
|
||||
self.scheduler.shutdown(wait=False)
|
||||
logger.info("APScheduler 已关闭")
|
||||
|
||||
def add_interval_job(
|
||||
self,
|
||||
job_id: str,
|
||||
func: Callable[..., Awaitable[None]],
|
||||
seconds: int,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
if job_id in self._jobs:
|
||||
logger.info("任务 %s 已存在，跳过重复注册", job_id)
|
||||
return
|
||||
trigger = IntervalTrigger(seconds=seconds)
|
||||
job = self.scheduler.add_job(func, trigger=trigger, id=job_id, kwargs=kwargs, max_instances=1)
|
||||
self._jobs[job_id] = job.id
|
||||
logger.info("已注册周期任务 %s，每隔 %s 秒执行", job_id, seconds)
|
||||
@@ -0,0 +1,125 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from app.models.job import IngestJob
|
||||
from app.services.embeddings import EmbeddingGenerator
|
||||
from app.services.job_store import InMemoryJobStore
|
||||
from app.services.webhooks import LightRAGWebhook
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
sys.path.append(str(ROOT / "packages" / "siyuan_ingest" / "src"))
|
||||
from siyuan_ingest import ( # type: ignore # noqa: E402
|
||||
IngestState,
|
||||
SiYuanClient,
|
||||
SiYuanContentExtractor,
|
||||
SupabaseSync,
|
||||
SupabaseWriter,
|
||||
chunk_blocks,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class IngestWorker:
|
||||
"""处理 ingest 队列，串联 SiYuan → Supabase → LightRAG"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
job_store: InMemoryJobStore,
|
||||
siyuan_client: SiYuanClient,
|
||||
webhook: LightRAGWebhook,
|
||||
supabase_sync: Optional[SupabaseSync] = None,
|
||||
supabase_writer: Optional[SupabaseWriter] = None,
|
||||
embedding_generator: Optional[EmbeddingGenerator] = None,
|
||||
) -> None:
|
||||
self.job_store = job_store
|
||||
self.siyuan_client = siyuan_client
|
||||
self.extractor = SiYuanContentExtractor(siyuan_client)
|
||||
self.webhook = webhook
|
||||
self.supabase_sync = supabase_sync
|
||||
self.supabase_writer = supabase_writer
|
||||
self.embedding_generator = embedding_generator
|
||||
self.queue: asyncio.Queue[IngestJob] = asyncio.Queue()
|
||||
self._consumer_task: Optional[asyncio.Task[None]] = None
|
||||
|
||||
async def start(self) -> None:
|
||||
if not self._consumer_task:
|
||||
self._consumer_task = asyncio.create_task(self._consume())
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._consumer_task:
|
||||
self._consumer_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self._consumer_task
|
||||
|
||||
async def enqueue(self, notebook_id: Optional[str], doc_id: Optional[str], source: str, blocks: List[str]) -> IngestJob:
|
||||
job = IngestJob(
|
||||
id=str(uuid4()),
|
||||
notebook_id=notebook_id,
|
||||
doc_id=doc_id,
|
||||
source=source,
|
||||
blocks=blocks,
|
||||
)
|
||||
await self.job_store.save(job)
|
||||
await self.queue.put(job)
|
||||
logger.info("排队 ingest 任务 %s", job.id)
|
||||
return job
|
||||
|
||||
async def _consume(self) -> None:
|
||||
while True:
|
||||
job = await self.queue.get()
|
||||
await self._handle_job(job)
|
||||
self.queue.task_done()
|
||||
|
||||
async def _handle_job(self, job: IngestJob) -> None:
|
||||
try:
|
||||
job.mark_running()
|
||||
await self.job_store.save(job)
|
||||
blocks = await self.extractor.list_blocks(job.notebook_id or "")
|
||||
chunks = chunk_blocks(blocks)
|
||||
logger.info("任务 %s 产生 %s 个 chunk", job.id, len(chunks))
|
||||
await self._maybe_embed(chunks)
|
||||
await self._persist_to_supabase(job, chunks)
|
||||
await self.webhook.notify_ingest(job.doc_id, [c["block_id"] for c in chunks if c.get("block_id")])
|
||||
job.mark_done()
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.exception("任务 %s 失败: %s", job.id, exc)
|
||||
job.mark_failed(str(exc))
|
||||
finally:
|
||||
await self.job_store.save(job)
|
||||
|
||||
async def _persist_to_supabase(self, job: IngestJob, chunks: List[Dict[str, str]]) -> None:
|
||||
await self._maybe_sync_state(job, chunks)
|
||||
if not self.supabase_writer:
|
||||
logger.debug("未配置 SupabaseWriter，跳过落库")
|
||||
return
|
||||
await self.supabase_writer.ingest_chunks(job.notebook_id, chunks)
|
||||
|
||||
async def _maybe_sync_state(self, job: IngestJob, chunks: List[Dict[str, str]]) -> None:
|
||||
if not self.supabase_sync:
|
||||
return
|
||||
if not chunks:
|
||||
return
|
||||
merged_checksum = chunks[-1]["checksum"]
|
||||
state = await self.supabase_sync.fetch_state(job.notebook_id, job.doc_id)
|
||||
if state and not state.should_ingest(merged_checksum):
|
||||
logger.info("任务 %s 未发生变化，跳过写入 Supabase", job.id)
|
||||
return
|
||||
state = state or IngestState(notebook_id=job.notebook_id, doc_id=job.doc_id)
|
||||
state.update(merged_checksum)
|
||||
await self.supabase_sync.upsert_state(state)
|
||||
|
||||
async def _maybe_embed(self, chunks: List[Dict[str, str]]) -> None:
|
||||
if not self.embedding_generator:
|
||||
return
|
||||
texts = [c.get("text") or "" for c in chunks]
|
||||
if not texts:
|
||||
return
|
||||
embeddings = await self.embedding_generator.embed_texts(texts)
|
||||
for chunk, embedding in zip(chunks, embeddings):
|
||||
chunk["embedding"] = embedding
|
||||
@@ -0,0 +1,26 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from app.core.config import get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LightRAGWebhook:
|
||||
"""封装与 LightRAG 的交互，默认调用 /ingest"""
|
||||
|
||||
def __init__(self, base_url: Optional[str] = None):
|
||||
settings = get_settings()
|
||||
self.base_url = base_url or settings.lightrag_url.rstrip("/")
|
||||
|
||||
async def notify_ingest(self, doc_id: Optional[str], block_ids: List[str]) -> None:
|
||||
if not block_ids:
|
||||
logger.info("无可用 block 推送给 LightRAG")
|
||||
return
|
||||
payload = {"doc_id": doc_id, "block_ids": block_ids}
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
resp = await client.post(f"{self.base_url}/ingest", json=payload)
|
||||
resp.raise_for_status()
|
||||
logger.info("LightRAG ingest 完成，返回 %s", resp.status_code)
|
||||
Reference in New Issue
Block a user