You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
31 lines
845 B
Python
31 lines
845 B
Python
# backend/app/rag/store.py
|
|
import math
|
|
|
|
from sqlmodel import Session, select
|
|
|
|
from ..models import RagChunk
|
|
|
|
|
|
def cosine(a: list[float], b: list[float]) -> float:
|
|
if not a or not b:
|
|
return 0.0
|
|
dot = sum(x * y for x, y in zip(a, b, strict=False))
|
|
na = math.sqrt(sum(x * x for x in a)) or 1.0
|
|
nb = math.sqrt(sum(y * y for y in b)) or 1.0
|
|
return dot / (na * nb)
|
|
|
|
|
|
def add_chunks(s: Session, chunks: list[RagChunk]) -> None:
|
|
for c in chunks:
|
|
s.add(c)
|
|
s.commit()
|
|
|
|
|
|
def query_store(s: Session, qvec: list[float], k: int, source_ids: list[str] | None):
|
|
rows = s.exec(select(RagChunk)).all()
|
|
if source_ids:
|
|
rows = [r for r in rows if r.source_id in source_ids]
|
|
scored = [(cosine(qvec, r.embedding), r) for r in rows]
|
|
scored.sort(key=lambda x: x[0], reverse=True)
|
|
return scored[:k]
|