foodlinkk-command-center/tools-api/app/brain.py

"""PostgreSQL second brain — Telegram messages, graph edges, pgvector."""
from __future__ import annotations

import logging
import os
from typing import Any, Optional

import httpx

from app.db import execute, execute_returning, fetch_all, fetch_one, json_param

log = logging.getLogger("tools.brain")

OLLAMA_URL = os.getenv("OLLAMA_URL", "http://10.4.7.19:11434").rstrip("/")
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
EMBED_DIM = 768


async def embed_text(text: str) -> list[float]:
    text = (text or "").strip()[:4000]
    if not text:
        return []
    async with httpx.AsyncClient(timeout=120.0) as client:
        r = await client.post(
            f"{OLLAMA_URL}/api/embeddings",
            json={"model": EMBED_MODEL, "prompt": text},
        )
        r.raise_for_status()
        vec = r.json().get("embedding") or []
        if len(vec) != EMBED_DIM:
            raise ValueError(f"embedding dim {len(vec)} != {EMBED_DIM}")
        return vec


def _vec_param(vec: list[float]) -> str:
    return "[" + ",".join(f"{x:.8f}" for x in vec) + "]"


def upsert_conversation(
    chat_id: int,
    *,
    chat_type: str = "private",
    user_name: str | None = None,
    user_role: str | None = None,
    metadata: dict | None = None,
) -> dict[str, Any]:
    row = execute_returning(
        """
        INSERT INTO telegram_conversations (chat_id, chat_type, user_name, user_role, metadata, updated_at)
        VALUES (%s, %s, %s, %s, %s, NOW())
        ON CONFLICT (chat_id) DO UPDATE SET
            user_name = COALESCE(EXCLUDED.user_name, telegram_conversations.user_name),
            user_role = COALESCE(EXCLUDED.user_role, telegram_conversations.user_role),
            metadata = telegram_conversations.metadata || EXCLUDED.metadata,
            updated_at = NOW()
        RETURNING *
        """,
        (chat_id, chat_type, user_name, user_role, json_param(metadata or {})),
    )
    return row or {}


def store_message(
    chat_id: int,
    *,
    direction: str,
    content_text: str | None = None,
    content_type: str = "text",
    role: str = "user",
    telegram_message_id: int | None = None,
    reply_to_db_id: int | None = None,
    agent_name: str | None = None,
    content_json: dict | None = None,
    user_name: str | None = None,
    user_role: str | None = None,
    chat_type: str = "private",
) -> dict[str, Any]:
    conv = upsert_conversation(
        chat_id, chat_type=chat_type, user_name=user_name, user_role=user_role
    )
    row = execute_returning(
        """
        INSERT INTO telegram_messages (
            conversation_id, telegram_message_id, direction, role, content_type,
            content_text, content_json, reply_to_message_id, agent_name
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        RETURNING *
        """,
        (
            conv["id"],
            telegram_message_id,
            direction,
            role,
            content_type,
            (content_text or "")[:8000] or None,
            json_param(content_json or {}),
            reply_to_db_id,
            agent_name,
        ),
    )
    return row or {}


def store_edge(
    source_message_id: int,
    edge_type: str,
    *,
    target_message_id: int | None = None,
    target_entity_type: str | None = None,
    target_entity_id: int | None = None,
    weight: float = 1.0,
    metadata: dict | None = None,
) -> dict[str, Any]:
    row = execute_returning(
        """
        INSERT INTO telegram_message_edges (
            source_message_id, target_message_id, target_entity_type,
            target_entity_id, edge_type, weight, metadata
        ) VALUES (%s, %s, %s, %s, %s, %s, %s)
        RETURNING *
        """,
        (
            source_message_id,
            target_message_id,
            target_entity_type,
            target_entity_id,
            edge_type,
            weight,
            json_param(metadata or {}),
        ),
    )
    return row or {}


def store_embedding(message_id: int, vec: list[float], chunk_index: int = 0) -> None:
    execute(
        """
        INSERT INTO telegram_message_embeddings (message_id, chunk_index, embedding, model)
        VALUES (%s, %s, %s::vector, %s)
        ON CONFLICT (message_id, chunk_index) DO UPDATE SET
            embedding = EXCLUDED.embedding,
            model = EXCLUDED.model
        """,
        (message_id, chunk_index, _vec_param(vec), EMBED_MODEL),
    )


async def store_message_with_embedding(
    chat_id: int,
    *,
    direction: str,
    content_text: str | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    msg = store_message(chat_id, direction=direction, content_text=content_text, **kwargs)
    text = (content_text or "").strip()
    if text and len(text) >= 8:
        try:
            vec = await embed_text(text)
            if vec:
                store_embedding(msg["id"], vec)
        except Exception as exc:
            log.warning("embed failed msg=%s: %s", msg.get("id"), exc)
    return msg


def get_graph(chat_id: int, limit: int = 50) -> dict[str, Any]:
    conv = fetch_one("SELECT * FROM telegram_conversations WHERE chat_id = %s", (chat_id,))
    if not conv:
        return {"chat_id": chat_id, "nodes": [], "edges": []}

    messages = fetch_all(
        """
        SELECT id, direction, role, content_type, content_text, agent_name,
               reply_to_message_id, created_at
        FROM telegram_messages
        WHERE conversation_id = %s
        ORDER BY created_at DESC
        LIMIT %s
        """,
        (conv["id"], limit),
    )
    msg_ids = [m["id"] for m in messages]
    edges: list[dict] = []
    if msg_ids:
        edges = fetch_all(
            """
            SELECT e.*, sm.content_text AS source_preview, tm.content_text AS target_preview
            FROM telegram_message_edges e
            JOIN telegram_messages sm ON sm.id = e.source_message_id
            LEFT JOIN telegram_messages tm ON tm.id = e.target_message_id
            WHERE e.source_message_id = ANY(%s) OR e.target_message_id = ANY(%s)
            ORDER BY e.created_at DESC
            LIMIT %s
            """,
            (msg_ids, msg_ids, limit * 2),
        )

    nodes = [
        {
            "id": m["id"],
            "label": (m.get("content_text") or "")[:80],
            "direction": m.get("direction"),
            "role": m.get("role"),
            "agent": m.get("agent_name"),
            "type": m.get("content_type"),
            "created_at": m.get("created_at"),
        }
        for m in reversed(messages)
    ]
    edge_list = [
        {
            "id": e["id"],
            "from": e["source_message_id"],
            "to": e.get("target_message_id"),
            "type": e["edge_type"],
            "entity_type": e.get("target_entity_type"),
            "entity_id": e.get("target_entity_id"),
            "weight": e.get("weight"),
        }
        for e in edges
    ]
    return {
        "chat_id": chat_id,
        "conversation_id": conv["id"],
        "nodes": nodes,
        "edges": edge_list,
        "stats": {"messages": len(nodes), "edges": len(edge_list)},
    }


async def search_memory(
    query: str,
    *,
    chat_id: int | None = None,
    limit: int = 8,
) -> list[dict[str, Any]]:
    query = query.strip()
    if not query:
        return []

    results: list[dict[str, Any]] = []

    try:
        vec = await embed_text(query)
        if vec:
            v = _vec_param(vec)
            if chat_id is not None:
                rows = fetch_all(
                    """
                    SELECT m.id, m.content_text, m.direction, m.agent_name, m.created_at,
                           c.chat_id,
                           1 - (e.embedding <=> %s::vector) AS score
                    FROM telegram_message_embeddings e
                    JOIN telegram_messages m ON m.id = e.message_id
                    JOIN telegram_conversations c ON c.id = m.conversation_id
                    WHERE m.content_text IS NOT NULL AND c.chat_id = %s
                    ORDER BY e.embedding <=> %s::vector
                    LIMIT %s
                    """,
                    (v, chat_id, v, limit),
                )
            else:
                rows = fetch_all(
                    """
                    SELECT m.id, m.content_text, m.direction, m.agent_name, m.created_at,
                           c.chat_id,
                           1 - (e.embedding <=> %s::vector) AS score
                    FROM telegram_message_embeddings e
                    JOIN telegram_messages m ON m.id = e.message_id
                    JOIN telegram_conversations c ON c.id = m.conversation_id
                    WHERE m.content_text IS NOT NULL
                    ORDER BY e.embedding <=> %s::vector
                    LIMIT %s
                    """,
                    (v, v, limit),
                )
            results.extend(rows)
    except Exception as exc:
        log.warning("vector search failed: %s", exc)

    if len(results) < limit:
        params2: list[Any] = [query, limit - len(results)]
        chat_clause = ""
        if chat_id is not None:
            chat_clause = "AND c.chat_id = %s"
            params2.append(chat_id)
        fts = fetch_all(
            f"""
            SELECT m.id, m.content_text, m.direction, m.agent_name, m.created_at, c.chat_id,
                   ts_rank(to_tsvector('simple', coalesce(m.content_text, '')),
                           plainto_tsquery('simple', %s)) AS score
            FROM telegram_messages m
            JOIN telegram_conversations c ON c.id = m.conversation_id
            WHERE to_tsvector('simple', coalesce(m.content_text, '')) @@ plainto_tsquery('simple', %s)
            {chat_clause}
            ORDER BY score DESC
            LIMIT %s
            """,
            tuple([query, query, *([chat_id] if chat_id else []), limit - len(results)]),
        )
        seen = {r["id"] for r in results}
        for row in fts:
            if row["id"] not in seen:
                results.append(row)

    return results[:limit]

# Append to tools-api/app/brain.py

def list_conversations(limit: int = 50) -> list[dict[str, Any]]:
    return fetch_all(
        """
        SELECT c.*,
               (SELECT COUNT(*) FROM telegram_messages m WHERE m.conversation_id = c.id) AS message_count,
               (SELECT MAX(m2.created_at) FROM telegram_messages m2 WHERE m2.conversation_id = c.id) AS last_message_at
        FROM telegram_conversations c
        ORDER BY (SELECT MAX(m3.created_at) FROM telegram_messages m3 WHERE m3.conversation_id = c.id) DESC NULLS LAST
        LIMIT %s
        """,
        (limit,),
    )


def list_feed(*, chat_id: int | None = None, limit: int = 80, offset: int = 0) -> list[dict[str, Any]]:
    if chat_id is not None:
        return fetch_all(
            """
            SELECT m.*, c.chat_id, c.user_name, c.user_role,
                   EXISTS(SELECT 1 FROM telegram_message_embeddings e WHERE e.message_id = m.id) AS has_embedding
            FROM telegram_messages m
            JOIN telegram_conversations c ON c.id = m.conversation_id
            WHERE c.chat_id = %s
            ORDER BY m.created_at DESC
            LIMIT %s OFFSET %s
            """,
            (chat_id, limit, offset),
        )
    return fetch_all(
        """
        SELECT m.*, c.chat_id, c.user_name, c.user_role,
               EXISTS(SELECT 1 FROM telegram_message_embeddings e WHERE e.message_id = m.id) AS has_embedding
        FROM telegram_messages m
        JOIN telegram_conversations c ON c.id = m.conversation_id
        ORDER BY m.created_at DESC
        LIMIT %s OFFSET %s
        """,
        (limit, offset),
    )


def get_dashboard_stats() -> dict[str, Any]:
    stats = fetch_one(
        """
        SELECT
            (SELECT COUNT(*) FROM telegram_conversations) AS conversations,
            (SELECT COUNT(*) FROM telegram_messages) AS messages,
            (SELECT COUNT(*) FROM telegram_messages WHERE direction = 'in') AS inbound,
            (SELECT COUNT(*) FROM telegram_messages WHERE direction = 'out') AS outbound,
            (SELECT COUNT(*) FROM telegram_message_edges) AS edges,
            (SELECT COUNT(*) FROM telegram_message_embeddings) AS embeddings
        """
    ) or {}
    recent = list_feed(limit=15)
    events = fetch_all(
        """
        SELECT id, agent_name, event_type, title, body, status, channel, created_at
        FROM agent_events
        WHERE channel = 'telegram' OR agent_name IN ('hermes', 'herman', 'personal_pa')
        ORDER BY created_at DESC
        LIMIT 20
        """
    )
    return {"stats": stats, "recent_messages": recent, "agent_events": events}


def get_global_graph(limit: int = 100) -> dict[str, Any]:
    messages = fetch_all(
        """
        SELECT m.id, m.direction, m.role, m.content_type, m.content_text, m.agent_name,
               m.created_at, c.chat_id, c.user_name
        FROM telegram_messages m
        JOIN telegram_conversations c ON c.id = m.conversation_id
        ORDER BY m.created_at DESC
        LIMIT %s
        """,
        (limit,),
    )
    if not messages:
        return {"nodes": [], "edges": []}
    ids = [m["id"] for m in messages]
    edges = fetch_all(
        """
        SELECT e.*
        FROM telegram_message_edges e
        WHERE e.source_message_id = ANY(%s) OR e.target_message_id = ANY(%s)
        """,
        (ids, ids),
    )
    nodes = [
        {
            "id": m["id"],
            "label": (m.get("content_text") or "")[:100],
            "direction": m.get("direction"),
            "role": m.get("role"),
            "agent": m.get("agent_name"),
            "chat_id": m.get("chat_id"),
            "user_name": m.get("user_name"),
            "type": m.get("content_type"),
            "created_at": m.get("created_at"),
        }
        for m in reversed(messages)
    ]
    edge_list = [
        {
            "id": e["id"],
            "from": e["source_message_id"],
            "to": e.get("target_message_id"),
            "type": e["edge_type"],
            "entity_type": e.get("target_entity_type"),
            "entity_id": e.get("target_entity_id"),
        }
        for e in edges
    ]
    return {"nodes": nodes, "edges": edge_list, "stats": {"nodes": len(nodes), "edges": len(edge_list)}}