119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Entrypoint при старте стека: проверяет, что в Qdrant есть коллекция chapter_analyses,
|
||
и создаёт её при отсутствии. Запускается отдельным контейнером после старта Qdrant.
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import urllib.error
|
||
import urllib.request
|
||
from typing import Any
|
||
|
||
|
||
def env(name: str, default: str) -> str:
|
||
"""Читает переменную окружения или default."""
|
||
return os.environ.get(name, default).strip()
|
||
|
||
|
||
def wait_for_qdrant(
|
||
base_url: str,
|
||
max_attempts: int = 45,
|
||
interval: float = 2.0,
|
||
initial_delay: float = 15.0,
|
||
) -> bool:
|
||
"""Ждёт, пока Qdrant станет доступен (сначала пауза, затем повторные попытки)."""
|
||
if initial_delay > 0:
|
||
print(f"Ожидание {initial_delay:.0f} с, пока Qdrant запустится...", flush=True)
|
||
time.sleep(initial_delay)
|
||
health = f"{base_url.rstrip('/')}/readyz"
|
||
last_error: Exception | None = None
|
||
for i in range(max_attempts):
|
||
try:
|
||
req = urllib.request.Request(health, method="GET")
|
||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||
if resp.status == 200:
|
||
return True
|
||
except Exception as e:
|
||
last_error = e
|
||
time.sleep(interval)
|
||
if last_error is not None:
|
||
print(f"Последняя ошибка при подключении: {last_error}", file=sys.stderr)
|
||
return False
|
||
|
||
|
||
def collection_exists(base_url: str, name: str) -> bool:
|
||
"""Проверяет, существует ли коллекция."""
|
||
url = f"{base_url.rstrip('/')}/collections/{name}"
|
||
try:
|
||
req = urllib.request.Request(url, method="GET")
|
||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||
return resp.status == 200
|
||
except urllib.error.HTTPError as e:
|
||
if e.code == 404:
|
||
return False
|
||
raise
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def create_collection(base_url: str, name: str, size: int) -> dict[str, Any]:
|
||
"""Создаёт коллекцию с заданной размерностью вектора (Cosine)."""
|
||
url = f"{base_url.rstrip('/')}/collections/{name}"
|
||
payload = {
|
||
"vectors": {
|
||
"size": size,
|
||
"distance": "Cosine",
|
||
},
|
||
}
|
||
body = json.dumps(payload).encode("utf-8")
|
||
req = urllib.request.Request(
|
||
url,
|
||
data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
method="PUT",
|
||
)
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
return json.loads(resp.read().decode("utf-8"))
|
||
|
||
|
||
def main() -> int:
|
||
"""Ждёт Qdrant, при отсутствии коллекции создаёт её, выходит 0."""
|
||
base_url = env("QDRANT_URL", "http://qdrant:6333")
|
||
collection_name = env("QDRANT_COLLECTION_CHAPTER_ANALYSES", "chapter_analyses")
|
||
vector_size = int(env("QDRANT_VECTOR_SIZE", "1024"))
|
||
initial_delay = float(env("QDRANT_INIT_DELAY", "15"))
|
||
|
||
print(f"Qdrant init: URL={base_url}, collection={collection_name}, size={vector_size}", flush=True)
|
||
|
||
if not wait_for_qdrant(base_url, initial_delay=initial_delay):
|
||
print("Ошибка: Qdrant недоступен.", file=sys.stderr)
|
||
return 1
|
||
|
||
if collection_exists(base_url, collection_name):
|
||
print(f"Коллекция '{collection_name}' уже существует.", flush=True)
|
||
return 0
|
||
|
||
try:
|
||
create_collection(base_url, collection_name, vector_size)
|
||
print(f"Коллекция '{collection_name}' создана.", flush=True)
|
||
except urllib.error.HTTPError as e:
|
||
print(f"Ошибка HTTP {e.code}: {e.reason}", file=sys.stderr)
|
||
if e.fp:
|
||
try:
|
||
print(e.fp.read().decode("utf-8")[:500], file=sys.stderr)
|
||
except Exception:
|
||
pass
|
||
return 1
|
||
except Exception as e:
|
||
print(f"Ошибка: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|