MVP stage 1: extract -> Postgres + atom_versions (web) + Qdrant
This commit is contained in:
neo committed
1 parent
ecd3c2e7b3
commit
ffb19d6000
3 files changed
+428
No files matched your search
@@ -0,0 +1,319 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
MVP stage 1: extract atoms from Markdown -> Postgres -> atom_versions (web).
|
||||
|
||||
Разбирает .md-статью на типизированные атомы, сохраняет в Postgres (atom-db),
|
||||
генерирует atom_versions под web-канал.
|
||||
|
||||
bge-m3 (dim=1024) — единственная embed-модель.
|
||||
Qdrant — эмбеддинги + дедупликация.
|
||||
Postgres (atom-db) — атомы как источник правды.
|
||||
|
||||
Запуск:
|
||||
python3 stage1/extract_to_postgres.py --article file.md --project rabota-ufa
|
||||
python3 stage1/extract_to_postgres.py --article a.md --article b.md --project ufa-kvartira
|
||||
|
||||
.env:
|
||||
POSTGRES_HOST=127.0.0.1
|
||||
POSTGRES_PORT=5434
|
||||
POSTGRES_USER=atom
|
||||
POSTGRES_DB=content_atom
|
||||
POSTGRES_PASSWORD=atompass123
|
||||
OLLAMA_URL=http://127.0.0.1:11434
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import RealDictCursor
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.models import Distance, VectorParams, PointStruct
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Конфиг
|
||||
# ---------------------------------------------------------------------------
|
||||
for var in ("POSTGRES_HOST", "POSTGRES_PORT", "POSTGRES_USER",
|
||||
"POSTGRES_DB", "POSTGRES_PASSWORD"):
|
||||
if not os.environ.get(var):
|
||||
print(f"ERROR: {var} not set (see .env)")
|
||||
sys.exit(1)
|
||||
|
||||
POSTGRES = {
|
||||
"host": os.environ.get("POSTGRES_HOST", "127.0.0.1"),
|
||||
"port": int(os.environ.get("POSTGRES_PORT", "5434")),
|
||||
"user": os.environ.get("POSTGRES_USER", "atom"),
|
||||
"dbname": os.environ.get("POSTGRES_DB", "content_atom"),
|
||||
"password": os.environ.get("POSTGRES_PASSWORD", "atompass123"),
|
||||
}
|
||||
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434")
|
||||
EMBED_MODEL = "bge-m3"
|
||||
COLLECTION = "content_atoms"
|
||||
DIM = 1024
|
||||
WEB_CHANNEL = "web_article"
|
||||
WEB_FORMAT = "markdown"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Типизированные атомы
|
||||
# ---------------------------------------------------------------------------
|
||||
class Atom:
|
||||
def __init__(self, project, atom_type, title, summary, body, plain_text, word_count):
|
||||
self.atom_id = str(uuid.uuid4())
|
||||
self.project = project
|
||||
self.type = atom_type
|
||||
self.title = title
|
||||
self.summary = summary
|
||||
self.body = body
|
||||
self.plain_text = plain_text
|
||||
self.word_count = word_count
|
||||
self.status = "draft"
|
||||
self.entities = []
|
||||
self.topics = []
|
||||
self.confidence = 0.8
|
||||
self.quality_score = 7.0
|
||||
self.pipeline_stage = "extract"
|
||||
self.created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
|
||||
def to_row(self):
|
||||
return {
|
||||
"atom_id": self.atom_id,
|
||||
"project": self.project,
|
||||
"type": self.type,
|
||||
"status": self.status,
|
||||
"title": self.title,
|
||||
"summary": self.summary,
|
||||
"body": json.dumps({"statement": self.body}, ensure_ascii=False),
|
||||
"plain_text": self.plain_text,
|
||||
"language": "ru",
|
||||
"entities": json.dumps(self.entities),
|
||||
"topics": self.topics,
|
||||
"confidence": self.confidence,
|
||||
"quality_score": self.quality_score,
|
||||
"word_count": self.word_count,
|
||||
"pipeline_stage": self.pipeline_stage,
|
||||
"created_at": self.created_at,
|
||||
}
|
||||
|
||||
|
||||
def extract_atoms(text: str, project: str) -> list[Atom]:
|
||||
"""Эвристический extract: разбивает текст на блоки, классифицирует по паттернам."""
|
||||
lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
|
||||
atoms: list[Atom] = []
|
||||
seq = 0
|
||||
|
||||
for ln in lines:
|
||||
if ln.startswith("#") or ln.startswith(">") or ln.startswith("http"):
|
||||
continue
|
||||
if re.match(r"^\|.*\|$", ln) or ln.startswith("---") or ln.startswith("```"):
|
||||
continue
|
||||
if len(ln) < 20:
|
||||
continue
|
||||
seq += 1
|
||||
atoms.append(Atom(
|
||||
project=project,
|
||||
atom_type=classify(ln),
|
||||
title=ln[:80],
|
||||
summary=ln[:160],
|
||||
body=ln,
|
||||
plain_text=ln,
|
||||
word_count=len(re.findall(r"\S+", ln)),
|
||||
))
|
||||
return atoms
|
||||
|
||||
|
||||
def classify(line: str) -> str:
|
||||
low = line.lower()
|
||||
if re.search(r"\d", line) and re.search(r"[%₽$€]", line):
|
||||
return "stat"
|
||||
if "?" in line:
|
||||
return "qa"
|
||||
if re.search(r"(подпишись|подписка|перейди|заказывай|оставь|свяжись|позвони)", low):
|
||||
return "cta"
|
||||
if re.search(r"(правильно|неправильно|ошибаеш|ты не|вот почему|секрет)", low):
|
||||
return "hook"
|
||||
if re.match(r"\d+[\.\)]", line):
|
||||
return "instruction"
|
||||
if line.startswith("«") or line.startswith("«"):
|
||||
return "quote"
|
||||
if re.search(r"(vs|против|по сравнению|лучше|дороже|дешевле)", low):
|
||||
return "comparison"
|
||||
return "claim"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Эмбеддинг bge-m3
|
||||
# ---------------------------------------------------------------------------
|
||||
def embed(text: str) -> list[float]:
|
||||
import urllib.request
|
||||
body = json.dumps({"model": EMBED_MODEL, "prompt": text}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{OLLAMA_URL}/api/embeddings", data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=60) as resp:
|
||||
data = json.loads(resp.read().decode())
|
||||
emb = data["embedding"]
|
||||
if len(emb) != DIM:
|
||||
raise ValueError(f"bge-m3 dim={len(emb)}, expected {DIM}")
|
||||
return emb
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Postgres
|
||||
# ---------------------------------------------------------------------------
|
||||
def get_conn():
|
||||
return psycopg2.connect(**POSTGRES)
|
||||
|
||||
|
||||
def ensure_project(conn, project: str) -> str:
|
||||
"""Создать project, если нет. Возвращает project_id."""
|
||||
with conn.cursor() as cur:
|
||||
cur.execute("SELECT id FROM projects WHERE slug=%s", (project,))
|
||||
row = cur.fetchone()
|
||||
if row:
|
||||
return row[0]
|
||||
cur.execute(
|
||||
"INSERT INTO projects (slug, name, kind, role, status) VALUES (%s, %s, %s, 'satellite', 'active') "
|
||||
"RETURNING id",
|
||||
(project, project, "site"),
|
||||
)
|
||||
return cur.fetchone()[0]
|
||||
|
||||
|
||||
def insert_atoms(conn, project_id: str, atoms: list[Atom]) -> int:
|
||||
"""Вставить атомы в atoms. Возвращает количество."""
|
||||
n = 0
|
||||
with conn.cursor() as cur:
|
||||
for atom in atoms:
|
||||
row = atom.to_row()
|
||||
cur.execute("""
|
||||
INSERT INTO atoms
|
||||
(id, project_id, type, status, title, summary, body, plain_text,
|
||||
language, entities, topics, confidence, quality_score, word_count, pipeline_stage, created_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (id) DO NOTHING
|
||||
""", (
|
||||
atom.atom_id, project_id, atom.type, atom.status, atom.title,
|
||||
atom.summary, row["body"], atom.plain_text, row["language"],
|
||||
row["entities"], row["topics"], row["confidence"],
|
||||
row["quality_score"], atom.word_count, row["pipeline_stage"],
|
||||
row["created_at"],
|
||||
))
|
||||
n += cur.rowcount
|
||||
conn.commit()
|
||||
return n
|
||||
|
||||
|
||||
def insert_atom_versions(conn, atoms: list[Atom]) -> int:
|
||||
"""Генерировать atom_versions под web-канал для каждого атома."""
|
||||
n = 0
|
||||
with conn.cursor() as cur:
|
||||
for atom in atoms:
|
||||
content = generate_web_version(atom)
|
||||
cur.execute("""
|
||||
INSERT INTO atom_versions
|
||||
(atom_id, channel, format, content, quality_score)
|
||||
VALUES (%s, %s, %s, %s, %s)
|
||||
ON CONFLICT (atom_id, channel, format) DO NOTHING
|
||||
""", (atom.atom_id, WEB_CHANNEL, WEB_FORMAT, content, 7.0))
|
||||
n += cur.rowcount
|
||||
conn.commit()
|
||||
return n
|
||||
|
||||
|
||||
def generate_web_version(atom: Atom) -> str:
|
||||
"""Переупаковка атома под web-канал (markdown)."""
|
||||
if atom.type in ("hook", "cta"):
|
||||
return f"> **{atom.title}**\n\n{atom.body}"
|
||||
if atom.type == "stat":
|
||||
return f"**{atom.summary}** — источник: {atom.atom_id}"
|
||||
if atom.type == "qa":
|
||||
return f"**{atom.title}**\n\n{atom.body}"
|
||||
return f"{atom.body}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Qdrant
|
||||
# ---------------------------------------------------------------------------
|
||||
def ensure_collection(client: QdrantClient) -> None:
|
||||
if client.collection_exists(COLLECTION):
|
||||
return
|
||||
client.create_collection(
|
||||
collection_name=COLLECTION,
|
||||
vectors_config=VectorParams(size=DIM, distance=Distance.COSINE),
|
||||
)
|
||||
|
||||
|
||||
def upsert_qdrant(client: QdrantClient, atoms: list[Atom]) -> int:
|
||||
n = 0
|
||||
for atom in atoms:
|
||||
emb = embed(atom.body)
|
||||
client.upsert(
|
||||
collection_name=COLLECTION,
|
||||
points=[PointStruct(
|
||||
id=str(uuid.uuid4()),
|
||||
vector=emb,
|
||||
payload=atom.to_row() | {"atom_id": atom.atom_id},
|
||||
)],
|
||||
)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description="MVP stage 1: extract -> Postgres -> atom_versions")
|
||||
ap.add_argument("--article", action="append", help=".md файл (можно несколько)")
|
||||
ap.add_argument("--project", default="spike", help="project_id (домен)")
|
||||
args = ap.parse_args()
|
||||
|
||||
if not args.article:
|
||||
print("Использование: --article file.md [--article file2.md] --project domain")
|
||||
return 1
|
||||
|
||||
# 1. Построить атомы
|
||||
all_atoms: list[Atom] = []
|
||||
for path in args.article:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
text = f.read()
|
||||
atoms = extract_atoms(text, args.project)
|
||||
all_atoms.extend(atoms)
|
||||
print(f" • {os.path.basename(path)}: {len(atoms)} атомов")
|
||||
|
||||
if not all_atoms:
|
||||
print("Нет атомов")
|
||||
return 1
|
||||
|
||||
# 2. Postgres
|
||||
start = time.time()
|
||||
conn = get_conn()
|
||||
project_id = ensure_project(conn, args.project)
|
||||
n_atoms = insert_atoms(conn, project_id, all_atoms)
|
||||
n_versions = insert_atom_versions(conn, all_atoms)
|
||||
conn.close()
|
||||
print(f" ✓ Postgres: {n_atoms} атомов, {n_versions} atom_versions за {time.time()-start:.1f}с")
|
||||
|
||||
# 3. Qdrant (эмбеддинги + дедуп)
|
||||
start = time.time()
|
||||
client = QdrantClient(host="127.0.0.1", port=6333)
|
||||
ensure_collection(client)
|
||||
n_emb = upsert_qdrant(client, all_atoms)
|
||||
print(f" ✓ Qdrant: {n_emb} эмбеддингов за {time.time()-start:.1f}с (~{ (time.time()-start)/n_emb*1000:.0f}ms/atom)")
|
||||
|
||||
# 4. Итог
|
||||
print(f"\nMVP stage 1 OK: {len(all_atoms)} атомов → Postgres + Qdrant")
|
||||
print(f" projects.slug={args.project}, atom_versions channel={WEB_CHANNEL}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in new issue
Block a user