mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
feat(embedding): add API-only retrieve_chunks() semantic search
This commit is contained in:
parent
22bf7fb1f5
commit
711be5e5bb
1 changed files with 39 additions and 0 deletions
39
paperforge/embedding/search.py
Normal file
39
paperforge/embedding/search.py
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
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 retrieve_chunks(vault: Path, query: str, limit: int = 5, expand: bool = True) -> list[dict]:
|
||||
"""Search chunks via API embedding. Returns list with metadata and similarity scores."""
|
||||
collection = get_collection(vault)
|
||||
provider = OpenAICompatibleProvider(vault)
|
||||
query_embedding = provider.encode_single(query)
|
||||
|
||||
results = collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=limit * 3 if expand else limit,
|
||||
include=["documents", "metadatas", "distances"],
|
||||
)
|
||||
|
||||
chunks = []
|
||||
for i, (doc, meta, dist) in enumerate(zip(
|
||||
results["documents"][0],
|
||||
results["metadatas"][0],
|
||||
results["distances"][0],
|
||||
)):
|
||||
chunks.append({
|
||||
"paper_id": meta["paper_id"],
|
||||
"section": meta.get("section", "Text"),
|
||||
"page_number": meta.get("page_number", 1),
|
||||
"chunk_index": meta.get("chunk_index", 0),
|
||||
"chunk_text": doc,
|
||||
"score": round(1.0 - dist, 4),
|
||||
})
|
||||
|
||||
return chunks
|
||||
Loading…
Reference in a new issue