084d3a8c63
- scripts/pkm_utils.py: 공통 유틸 (로거, dotenv, osascript 래퍼) - scripts/prompts/classify_document.txt: Ollama 분류 프롬프트 - applescript/auto_classify.scpt: Inbox → AI 분류 → DB 이동 - applescript/omnifocus_sync.scpt: Projects → OmniFocus 작업 생성 - scripts/law_monitor.py: 법령 변경 모니터링 + DEVONthink 임포트 - scripts/mailplus_archive.py: MailPlus IMAP → Archive DB - scripts/pkm_daily_digest.py: 일일 다이제스트 + OmniFocus 액션 - scripts/embed_to_chroma.py: GPU 서버 벡터 임베딩 → ChromaDB - launchd/*.plist: 3개 스케줄 (07:00, 07:00+18:00, 20:00) - docs/deploy.md: Mac mini 배포 가이드 - docs/devonagent-setup.md: 검색 세트 9종 설정 가이드 - tests/test_classify.py: 5종 문서 분류 테스트 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
105 lines
3.1 KiB
Python
105 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
벡터 임베딩 스크립트
|
|
- DEVONthink 문서 UUID로 텍스트 추출
|
|
- GPU 서버(nomic-embed-text)로 임베딩 생성
|
|
- ChromaDB에 저장
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import requests
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
from pkm_utils import setup_logger, load_credentials, run_applescript_inline
|
|
|
|
logger = setup_logger("embed")
|
|
|
|
# ChromaDB 저장 경로
|
|
CHROMA_DIR = Path.home() / ".local" / "share" / "pkm" / "chromadb"
|
|
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def get_document_text(uuid: str) -> tuple[str, str]:
|
|
"""DEVONthink에서 UUID로 문서 텍스트 + 제목 추출"""
|
|
script = f'''
|
|
tell application id "DNtp"
|
|
set theRecord to get record with uuid "{uuid}"
|
|
set docText to plain text of theRecord
|
|
set docTitle to name of theRecord
|
|
return docTitle & "|||" & docText
|
|
end tell
|
|
'''
|
|
result = run_applescript_inline(script)
|
|
parts = result.split("|||", 1)
|
|
title = parts[0] if len(parts) > 0 else ""
|
|
text = parts[1] if len(parts) > 1 else ""
|
|
return title, text
|
|
|
|
|
|
def get_embedding(text: str, gpu_server_ip: str) -> list[float] | None:
|
|
"""GPU 서버의 nomic-embed-text로 임베딩 생성"""
|
|
url = f"http://{gpu_server_ip}:11434/api/embeddings"
|
|
try:
|
|
resp = requests.post(url, json={
|
|
"model": "nomic-embed-text",
|
|
"prompt": text[:8000] # 토큰 제한
|
|
}, timeout=60)
|
|
resp.raise_for_status()
|
|
return resp.json().get("embedding")
|
|
except Exception as e:
|
|
logger.error(f"임베딩 생성 실패: {e}")
|
|
return None
|
|
|
|
|
|
def store_in_chromadb(doc_id: str, title: str, text: str, embedding: list[float]):
|
|
"""ChromaDB에 저장"""
|
|
import chromadb
|
|
client = chromadb.PersistentClient(path=str(CHROMA_DIR))
|
|
collection = client.get_or_create_collection(
|
|
name="pkm_documents",
|
|
metadata={"hnsw:space": "cosine"}
|
|
)
|
|
collection.upsert(
|
|
ids=[doc_id],
|
|
embeddings=[embedding],
|
|
documents=[text[:2000]],
|
|
metadatas=[{"title": title, "source": "devonthink"}]
|
|
)
|
|
logger.info(f"ChromaDB 저장: {doc_id} ({title[:30]})")
|
|
|
|
|
|
def run(uuid: str):
|
|
"""단일 문서 임베딩 처리"""
|
|
logger.info(f"임베딩 처리 시작: {uuid}")
|
|
|
|
creds = load_credentials()
|
|
gpu_ip = creds.get("GPU_SERVER_IP")
|
|
if not gpu_ip:
|
|
logger.warning("GPU_SERVER_IP 미설정 — 임베딩 건너뜀")
|
|
return
|
|
|
|
try:
|
|
title, text = get_document_text(uuid)
|
|
if not text or len(text) < 10:
|
|
logger.warning(f"텍스트 부족 [{uuid}]: {len(text)}자")
|
|
return
|
|
|
|
embedding = get_embedding(text, gpu_ip)
|
|
if embedding:
|
|
store_in_chromadb(uuid, title, text, embedding)
|
|
logger.info(f"임베딩 완료: {uuid}")
|
|
else:
|
|
logger.error(f"임베딩 실패: {uuid}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"임베딩 처리 에러 [{uuid}]: {e}", exc_info=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("사용법: python3 embed_to_chroma.py <DEVONthink_UUID>")
|
|
sys.exit(1)
|
|
run(sys.argv[1])
|