mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from paperforge.embedding._chroma import get_collection
|
|
from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def embed_paper(vault: Path, zotero_key: str, chunks: list[dict]) -> int:
|
|
"""Embed chunks for one paper using API and insert into ChromaDB. Returns count."""
|
|
collection = get_collection(vault)
|
|
provider = OpenAICompatibleProvider(vault)
|
|
|
|
texts = [c["text"] for c in chunks]
|
|
ids = [f"{zotero_key}_{c['chunk_index']}" for c in chunks]
|
|
metadatas = [
|
|
{
|
|
"paper_id": zotero_key,
|
|
"section": c["section"],
|
|
"page_number": c["page_number"],
|
|
"chunk_index": c["chunk_index"],
|
|
"token_estimate": c["token_estimate"],
|
|
}
|
|
for c in chunks
|
|
]
|
|
|
|
embeddings = provider.encode(texts)
|
|
try:
|
|
collection.add(
|
|
ids=ids,
|
|
embeddings=embeddings,
|
|
documents=texts,
|
|
metadatas=metadatas,
|
|
)
|
|
except Exception as exc:
|
|
err = str(exc).lower()
|
|
if "hnsw" in err or "compaction" in err or "segment" in err:
|
|
raise RuntimeError(
|
|
f"ChromaDB index error (possibly corrupted). "
|
|
f"Run 'paperforge embed build --force' to rebuild from scratch. "
|
|
f"Original error: {exc}"
|
|
) from exc
|
|
raise
|
|
return len(chunks)
|