backend/graph/nodes/write_brain.py

import time
import json
import uuid
import datetime
from backend.graph.state import BrainState
from backend.db.supabase import get_client
from backend.llm import get_embeddings
from backend.sse import emit


async def write_brain(state: BrainState) -> dict:
    job_id = state.get("job_id")
    company_id = state.get("company_id")
    final_skills = state.get("final_skills", [])
    start_time = state.get("start_time", time.time())
    duration_ms = int((time.time() - start_time) * 1000)

    print(
        f"[{job_id}] Node write_brain: persisting {len(final_skills)} skills for {company_id}"
    )
    await emit(
        job_id,
        "stage",
        {
            "name": "WRITING_DB",
            "detail": f"Pre-embedding and persisting {len(final_skills)} skills",
        },
    )

    skill_texts = [
        f"{s.get('category', '')} {s.get('rule', '')} {s.get('rationale', '')}"
        for s in final_skills
    ]
    embeddings = get_embeddings(skill_texts)
    for skill, emb in zip(final_skills, embeddings):
        skill["embedding_vector"] = emb

    skills_file = {
        "skills": final_skills,
        "meta": {
            "company_id": company_id,
            "compiled_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
            "total_skills": len(final_skills),
            "duration_ms": duration_ms,
        },
    }

    db = get_client()
    if not db:
        await emit(job_id, "pipeline_error", {"error": "Database connection failed"})
        print(f"[{job_id}] write_brain: no DB client")
        return {
            "errors": ["DB connection failed in write_brain"],
            "skills_file": skills_file,
        }

    try:
        now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat()
        version_str = f"v_{int(time.time())}"

        source_hashes = {}
        for f in state.get("source_files", []):
            if "filename" in f and "sha256" in f:
                source_hashes[f["filename"]] = f["sha256"]

        db.table("skills_files").update({"is_current": False}).eq(
            "company_id", company_id
        ).eq("is_current", True).execute()

        sf_res = (
            db.table("skills_files")
            .insert(
                {
                    "company_id": company_id,
                    "version": version_str,
                    "brain_json": skills_file,
                    "source_hashes": source_hashes,
                    "is_current": True,
                }
            )
            .execute()
        )

        sf_id = sf_res.data[0]["id"]

        skill_rows = []
        for skill in final_skills:
            skill_copy = {k: v for k, v in skill.items() if k != "embedding_vector"}
            skill_rows.append(
                {
                    "id": skill.get("id", str(uuid.uuid4())[:8]),
                    "company_id": company_id,
                    "skills_file_id": sf_id,
                    "name": skill.get("rule", "Unknown")[:200],
                    "domain": skill.get("category", "general"),
                    "version": version_str,
                    "confidence": float(skill.get("confidence", 0.5)),
                    "skill_json": skill_copy,
                }
            )
        if skill_rows:
            db.table("skills").insert(skill_rows).execute()

        db.table("compile_runs").update(
            {
                "status": "complete",
                "completed_at": now_iso,
                "duration_ms": duration_ms,
                "result_version": version_str,
            }
        ).eq("id", job_id).execute()

    except Exception as e:
        print(f"[{job_id}] [write_brain] DB Error: {e}")
        await emit(job_id, "pipeline_error", {"error": str(e)})
        return {"errors": [f"write_brain DB error: {e}"], "skills_file": skills_file}

    await emit(
        job_id,
        "stage",
        {
            "name": "DONE",
            "detail": f"Brain {version_str} written: {len(final_skills)} skills, {len(source_hashes)} sources, {duration_ms}ms",
        },
    )
    await emit(
        job_id,
        "pipeline_complete",
        {
            "status": "success",
            "version": version_str,
            "skills_count": len(final_skills),
            "source_count": len(source_hashes),
            "duration_ms": duration_ms,
        },
    )

    print(f"[{job_id}] write_brain: done (version: {version_str})")
    return {"skills_file": skills_file, "brain_version": version_str}