Andrey Lutsenko
94 lines
2.6 KiB
Python
94 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
import os
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, PointStruct, VectorParams
|
|
|
|
DOCS_DIR = Path(os.getenv("RAG_DOCS_DIR", "./docs"))
|
|
QDRANT_URL = os.getenv("QDRANT_URL", "http://127.0.0.1:6333")
|
|
COLLECTION = os.getenv("QDRANT_COLLECTION", "docs")
|
|
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
|
|
CHUNK_SIZE = int(os.getenv("CHUNK_SIZE", "1000"))
|
|
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "120"))
|
|
|
|
|
|
def read_text(path: Path) -> str:
|
|
if path.suffix.lower() == ".pdf":
|
|
from pypdf import PdfReader
|
|
|
|
reader = PdfReader(str(path))
|
|
return "\n".join((p.extract_text() or "") for p in reader.pages)
|
|
return path.read_text(encoding="utf-8", errors="ignore")
|
|
|
|
|
|
def chunk_text(text: str, size: int, overlap: int) -> list[str]:
|
|
chunks = []
|
|
i = 0
|
|
step = max(1, size - overlap)
|
|
while i < len(text):
|
|
chunks.append(text[i : i + size])
|
|
i += step
|
|
return chunks
|
|
|
|
|
|
def embed(text: str) -> list[float]:
|
|
r = requests.post(
|
|
"http://127.0.0.1:11434/api/embeddings",
|
|
json={"model": EMBED_MODEL, "prompt": text},
|
|
timeout=120,
|
|
)
|
|
r.raise_for_status()
|
|
return r.json()["embedding"]
|
|
|
|
|
|
def main() -> None:
|
|
if not DOCS_DIR.exists():
|
|
print(f"Docs directory not found: {DOCS_DIR}")
|
|
return
|
|
|
|
client = QdrantClient(url=QDRANT_URL)
|
|
test_vec = embed("ping")
|
|
dim = len(test_vec)
|
|
|
|
if not client.collection_exists(COLLECTION):
|
|
client.create_collection(
|
|
collection_name=COLLECTION,
|
|
vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
|
|
)
|
|
|
|
points: list[PointStruct] = []
|
|
supported = {".md", ".txt", ".pdf"}
|
|
|
|
for path in DOCS_DIR.rglob("*"):
|
|
if not path.is_file() or path.suffix.lower() not in supported:
|
|
continue
|
|
raw = read_text(path).strip()
|
|
if not raw:
|
|
continue
|
|
for idx, chunk in enumerate(chunk_text(raw, CHUNK_SIZE, CHUNK_OVERLAP)):
|
|
vec = embed(chunk)
|
|
points.append(
|
|
PointStruct(
|
|
id=str(uuid.uuid4()),
|
|
vector=vec,
|
|
payload={
|
|
"source": str(path),
|
|
"chunk_id": idx,
|
|
"text": chunk[:4000],
|
|
},
|
|
)
|
|
)
|
|
|
|
if points:
|
|
client.upsert(collection_name=COLLECTION, points=points)
|
|
print(f"Indexed chunks: {len(points)}")
|
|
else:
|
|
print("No documents indexed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|