232 lines
8.2 KiB
Python
232 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
CLI для получения эмбеддинга по JSON шага 5 (мерж анализа и тегов).
|
||
|
||
На вход — путь к JSON-файлу, сформированному на этапе 5 (merged_with_tags.json),
|
||
и имя модели. Текст для эмбеддинга собирается из блоков framework, insights, application
|
||
по спецификации embed_input_spec.txt. На выход — вектор.
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import sys
|
||
import urllib.error
|
||
import urllib.request
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
|
||
def _serialize_framework(fw: dict[str, Any]) -> str:
|
||
"""Собирает текст секции FRAMEWORK из JSON блока framework."""
|
||
parts: list[str] = []
|
||
principles = fw.get("principles") or []
|
||
for p in principles:
|
||
line = f"Принцип: {p.get('title', '')}. Описание: {p.get('description', '')}."
|
||
if p.get("example"):
|
||
line += f" Пример: {p['example']}."
|
||
parts.append(line)
|
||
for chain in p.get("chains") or []:
|
||
c_line = "Причина: {} Механизм: {} Результат: {}.".format(
|
||
chain.get("cause", ""),
|
||
chain.get("mechanism", ""),
|
||
chain.get("result", ""),
|
||
)
|
||
parts.append(c_line)
|
||
terms = fw.get("terms") or {}
|
||
if terms:
|
||
terms_str = " ".join(f"{k} — {v}" for k, v in terms.items())
|
||
parts.append(f"Термины: {terms_str}.")
|
||
return "\n".join(parts)
|
||
|
||
|
||
def _serialize_insights(insights: list[Any]) -> str:
|
||
"""Собирает текст секции INSIGHTS."""
|
||
parts: list[str] = []
|
||
for i in insights:
|
||
line = f"Инсайт: {i.get('title', '')}. {i.get('description', '')}."
|
||
if i.get("example"):
|
||
line += f" Пример: {i['example']}."
|
||
parts.append(line)
|
||
return "\n".join(parts)
|
||
|
||
|
||
def _serialize_application(app: dict[str, Any]) -> str:
|
||
"""Собирает текст секции APPLICATION из блока application."""
|
||
parts: list[str] = []
|
||
techniques = app.get("techniques") or []
|
||
for t in techniques:
|
||
line = f"Техника: {t.get('name', '')}. Цель: {t.get('goal', '')}."
|
||
if t.get("steps"):
|
||
steps = t["steps"] if isinstance(t["steps"], list) else [t["steps"]]
|
||
line += f" Шаги: {'; '.join(steps)}."
|
||
if t.get("context_example"):
|
||
line += f" Контекст: {t['context_example']}."
|
||
parts.append(line)
|
||
return "\n".join(parts)
|
||
|
||
|
||
def merged_json_to_embed_text(merged: dict[str, Any]) -> str:
|
||
"""
|
||
Собирает текст для эмбеддинга из JSON шага 5 по embed_input_spec.txt.
|
||
|
||
Входят только блоки framework, insights, application. Порядок секций:
|
||
FRAMEWORK, INSIGHTS, APPLICATION.
|
||
Limitations и теги не включаются: теги хранятся в payload Qdrant для фильтрации,
|
||
вектор строится по смыслу (принципы, инсайты, техники), не по тегам.
|
||
|
||
Args:
|
||
merged: Документ из шага 5 (merge.json + выход_valid_tag.json).
|
||
|
||
Returns:
|
||
Нормализованная строка для подачи в модель эмбеддингов.
|
||
"""
|
||
fw = merged.get("framework") or {}
|
||
insights = merged.get("insights") or []
|
||
app = merged.get("application") or {}
|
||
|
||
framework_text = _serialize_framework(fw)
|
||
insights_text = _serialize_insights(insights)
|
||
application_text = _serialize_application(app)
|
||
|
||
sections = [
|
||
"FRAMEWORK",
|
||
framework_text,
|
||
"INSIGHTS",
|
||
insights_text,
|
||
"APPLICATION",
|
||
application_text,
|
||
]
|
||
text = "\n\n".join(sections)
|
||
# Нормализация: схлопнуть множественные пробелы и переносы, обрезать по краям
|
||
text = re.sub(r"[ \t]+", " ", text)
|
||
text = re.sub(r"\n\s*\n", "\n\n", text)
|
||
return text.strip()
|
||
|
||
|
||
def get_embedding(
|
||
base_url: str,
|
||
model: str,
|
||
text: str,
|
||
) -> list[float]:
|
||
"""
|
||
Запрашивает эмбеддинг текста у OpenAI-совместимого API.
|
||
|
||
Args:
|
||
base_url: Базовый URL API (например http://localhost:1234/v1).
|
||
model: Имя модели эмбеддингов.
|
||
text: Текст для эмбеддинга.
|
||
|
||
Returns:
|
||
Вектор эмбеддинга (список float).
|
||
|
||
Raises:
|
||
urllib.error.HTTPError: При ошибке HTTP.
|
||
ValueError: Если в ответе нет ожидаемой структуры.
|
||
"""
|
||
url = f"{base_url.rstrip('/')}/embeddings"
|
||
payload = {
|
||
"model": model,
|
||
"input": text,
|
||
}
|
||
body = json.dumps(payload).encode("utf-8")
|
||
req = urllib.request.Request(
|
||
url,
|
||
data=body,
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(req, timeout=120) as resp:
|
||
data: dict[str, Any] = json.loads(resp.read().decode("utf-8"))
|
||
if "data" not in data or not data["data"]:
|
||
raise ValueError("В ответе API нет поля data с эмбеддингом")
|
||
embedding = data["data"][0].get("embedding")
|
||
if not embedding:
|
||
raise ValueError("В ответе API нет поля embedding")
|
||
return list(embedding)
|
||
|
||
|
||
def main() -> int:
|
||
"""Точка входа CLI."""
|
||
parser = argparse.ArgumentParser(
|
||
description="Эмбеддинг по JSON шага 5 (мерж анализа и тегов). На вход — JSON-файл, на выход — вектор.",
|
||
)
|
||
parser.add_argument(
|
||
"input_json",
|
||
type=Path,
|
||
help="Путь к JSON-файлу шага 5 (merged_with_tags.json).",
|
||
)
|
||
parser.add_argument(
|
||
"model",
|
||
help="Имя модели эмбеддингов (например text-embedding-bge-m3).",
|
||
)
|
||
parser.add_argument(
|
||
"--base-url",
|
||
default="http://localhost:1234/v1",
|
||
help="Базовый URL API (по умолчанию LM Studio).",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--output",
|
||
type=Path,
|
||
default=None,
|
||
help="Путь к файлу для записи вектора (JSON). По умолчанию — stdout.",
|
||
)
|
||
parser.add_argument(
|
||
"--format",
|
||
choices=("json", "compact"),
|
||
default="json",
|
||
help="Формат вывода: json — массив, compact — числа через пробел.",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
if not args.input_json.exists():
|
||
print(f"Ошибка: файл не найден: {args.input_json}", file=sys.stderr)
|
||
return 1
|
||
|
||
try:
|
||
with open(args.input_json, encoding="utf-8") as f:
|
||
merged = json.load(f)
|
||
except json.JSONDecodeError as e:
|
||
print(f"Ошибка разбора JSON: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
text = merged_json_to_embed_text(merged)
|
||
if not text:
|
||
print("Ошибка: текст для эмбеддинга пуст (нет framework/insights/application).", file=sys.stderr)
|
||
return 1
|
||
|
||
try:
|
||
vector = get_embedding(args.base_url, args.model, text)
|
||
except urllib.error.HTTPError as e:
|
||
print(f"Ошибка HTTP {e.code}: {e.reason}", file=sys.stderr)
|
||
if e.fp:
|
||
try:
|
||
body = e.fp.read().decode("utf-8")
|
||
print(body[:500], file=sys.stderr)
|
||
except Exception:
|
||
pass
|
||
return 1
|
||
except urllib.error.URLError as e:
|
||
print(f"Ошибка запроса: {e.reason}", file=sys.stderr)
|
||
return 1
|
||
except ValueError as e:
|
||
print(f"Ошибка: {e}", file=sys.stderr)
|
||
return 1
|
||
|
||
if args.format == "json":
|
||
payload = json.dumps(vector, ensure_ascii=False)
|
||
else:
|
||
payload = " ".join(str(x) for x in vector)
|
||
|
||
if args.output is None:
|
||
print(payload)
|
||
else:
|
||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
args.output.write_text(payload, encoding="utf-8")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|