51 lines
2.0 KiB
Python
51 lines
2.0 KiB
Python
from __future__ import annotations
|
|
import os
|
|
from typing import List, Tuple
|
|
from qdrant_client.http import models as rest
|
|
|
|
|
|
def _names(prefix: str) -> Tuple[str, str, str]:
|
|
return f"{prefix}_notes", f"{prefix}_chunks", f"{prefix}_edges"
|
|
|
|
|
|
def points_for_note(prefix: str, note_payload: dict, note_vec: List[float] | None, dim: int) -> Tuple[str, List[rest.PointStruct]]:
|
|
"""
|
|
Liefert (collection_name, [PointStruct]) für die Notes-Collection.
|
|
Falls kein Note-Embedding übergeben wurde, wird ein Nullvektor der Länge `dim` verwendet.
|
|
Hintergrund: Die Notes-Collection ist in ensure_collections mit Vektor-Dimension angelegt.
|
|
"""
|
|
notes_col, _, _ = _names(prefix)
|
|
vector = note_vec if note_vec is not None else [0.0] * int(dim)
|
|
pt = rest.PointStruct(id=note_payload["note_id"], vector=vector, payload=note_payload)
|
|
return notes_col, [pt]
|
|
|
|
|
|
def points_for_chunks(prefix: str, chunk_payloads: List[dict], vectors: List[List[float]]) -> Tuple[str, List[rest.PointStruct]]:
|
|
"""
|
|
Liefert (collection_name, [PointStruct]) für die Chunks-Collection.
|
|
Erwartet für jeden Chunk einen Embedding-Vektor (oder Nullvektor, wenn --skip-embed).
|
|
"""
|
|
_, chunks_col, _ = _names(prefix)
|
|
points: List[rest.PointStruct] = []
|
|
for pl, vec in zip(chunk_payloads, vectors):
|
|
points.append(rest.PointStruct(id=pl["chunk_id"], vector=vec, payload=pl))
|
|
return chunks_col, points
|
|
|
|
|
|
def points_for_edges(prefix: str, edge_payloads: List[dict]) -> Tuple[str, List[rest.PointStruct]]:
|
|
"""
|
|
Liefert (collection_name, [PointStruct]) für die Edges-Collection.
|
|
Edges-Collection ist VEKTORENLOS angelegt → nur Payload.
|
|
"""
|
|
_, _, edges_col = _names(prefix)
|
|
points: List[rest.PointStruct] = []
|
|
for pl in edge_payloads:
|
|
points.append(rest.PointStruct(id=pl["edge_id"], payload=pl))
|
|
return edges_col, points
|
|
|
|
|
|
def upsert_batch(client, collection: str, points: List[rest.PointStruct]) -> None:
|
|
if not points:
|
|
return
|
|
client.upsert(collection_name=collection, points=points, wait=True)
|