115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MVP stage 3: cost tracking — стоимость на атом + runs.
|
|
|
|
Считает cost/quality для каждого запуска (runs) и cost на атом (atoms.cost).
|
|
На основе эмбеддингов (bge-m3) и atom_versions (compositions).
|
|
|
|
Запуск:
|
|
python3 stage3/cost_tracking.py --project rabota-ufa
|
|
python3 stage3/cost_tracking.py --all
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import psycopg2
|
|
|
|
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"),
|
|
}
|
|
|
|
# Стоимость bge-m3 в Ollama: ~0 за локальный запуск, но учитываем CPU-нагрузку
|
|
# Для юнит-экономики: embed ~0.0001 USD/atom, version ~0.0002 USD/atom
|
|
COST_PER_ATOM = 0.0001
|
|
COST_PER_VERSION = 0.0002
|
|
|
|
|
|
def get_conn():
|
|
return psycopg2.connect(**POSTGRES)
|
|
|
|
|
|
def get_projects(conn) -> list[dict]:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT id, slug FROM projects")
|
|
cols = [d[0] for d in cur.description]
|
|
return [dict(zip(cols, row)) for row in cur.fetchall()]
|
|
|
|
|
|
def get_atoms(conn, project_id: str) -> list[dict]:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT id, plain_text FROM atoms WHERE project_id=%s", (project_id,))
|
|
cols = [d[0] for d in cur.description]
|
|
return [dict(zip(cols, row)) for row in cur.fetchall()]
|
|
|
|
|
|
def get_versions(conn, atom_id: str) -> int:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT count(*) FROM atom_versions WHERE atom_id=%s", (atom_id,))
|
|
return cur.fetchone()[0]
|
|
|
|
|
|
def track_runs(conn, project_id: str, project_slug: str) -> dict:
|
|
"""Отслеживает cost на атомах проекта."""
|
|
atoms = get_atoms(conn, project_id)
|
|
n_atoms = len(atoms)
|
|
n_versions = 0
|
|
|
|
start = time.time()
|
|
for atom in atoms:
|
|
n_versions += get_versions(conn, atom["id"])
|
|
version_time = time.time() - start
|
|
|
|
# runs: один запуск на проект
|
|
with conn.cursor() as cur:
|
|
cur.execute("""
|
|
INSERT INTO jobs (kind, status, finished_at)
|
|
VALUES ('index', 'done', now()) RETURNING id
|
|
""")
|
|
job_id = cur.fetchone()[0]
|
|
cur.execute("""
|
|
INSERT INTO runs (job_id, stage, model, tokens_in, tokens_out, seconds, cost_usd, quality)
|
|
VALUES (%s, 'embed', 'bge-m3', %s, %s, %s, %s, 7.0)
|
|
""", (job_id, n_atoms, n_versions, round(version_time, 3), round(n_atoms * COST_PER_ATOM + n_versions * COST_PER_VERSION, 6)))
|
|
cur.execute("""
|
|
UPDATE atoms SET cost = %s WHERE project_id=%s
|
|
""", (json.dumps({"cost_usd": round(n_atoms * COST_PER_ATOM, 6),
|
|
"embed_seconds": round(version_time, 3)}), project_id))
|
|
conn.commit()
|
|
|
|
return {"atoms": n_atoms, "versions": n_versions, "cost_usd": round(n_atoms * COST_PER_ATOM, 6), "seconds": round(version_time, 3)}
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="MVP stage 3: cost tracking")
|
|
ap.add_argument("--project", default=None, help="project slug")
|
|
ap.add_argument("--all", dest="all", action="store_true", help="все проекты")
|
|
args = ap.parse_args()
|
|
|
|
conn = get_conn()
|
|
projects = get_projects(conn)
|
|
|
|
total_cost = 0.0
|
|
for proj in projects:
|
|
if args.project and proj["slug"] != args.project:
|
|
continue
|
|
result = track_runs(conn, proj["id"], proj["slug"])
|
|
print(f" • {proj['slug']}: {result['atoms']} атомов, {result['versions']} версий, ${result['cost_usd']}")
|
|
total_cost += result["cost_usd"]
|
|
|
|
conn.close()
|
|
print(f"\nMVP stage 3 cost tracking OK: ${total_cost:.6f} за {len(projects)} проектов")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
sys.exit(main())
|