from __future__ import annotations import json import os from dotenv import load_dotenv from sqlalchemy import create_engine, text from sqlalchemy.engine import Engine load_dotenv() _engine: Engine | None = None def get_engine() -> Engine: global _engine if _engine is None: database_url = os.getenv("DATABASE_URL", "").strip("'\"") if not database_url: raise RuntimeError("DATABASE_URL is missing in .env") _engine = create_engine(database_url, pool_pre_ping=True) return _engine def list_users() -> list[dict]: query = text("SELECT id, username FROM auth_user ORDER BY username") with get_engine().connect() as connection: return [dict(row) for row in connection.execute(query).mappings()] def list_conversations(user_id: int) -> list[dict]: query = text(""" SELECT id, title, created_at, updated_at FROM chat_conversation WHERE user_id = :user_id ORDER BY updated_at DESC """) with get_engine().connect() as connection: return [dict(row) for row in connection.execute(query, {"user_id": user_id}).mappings()] def create_conversation(user_id: int, title: str) -> dict: query = text(""" INSERT INTO chat_conversation (user_id, title, created_at, updated_at) VALUES (:user_id, :title, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) RETURNING id, title, created_at, updated_at """) with get_engine().begin() as connection: user_exists = connection.execute( text("SELECT 1 FROM auth_user WHERE id = :user_id"), {"user_id": user_id}, ).scalar_one_or_none() if not user_exists: raise ValueError("User does not exist") return dict(connection.execute(query, {"user_id": user_id, "title": title[:255]}).mappings().one()) def ensure_conversation_owner(conversation_id: int, user_id: int) -> None: query = text(""" SELECT 1 FROM chat_conversation WHERE id = :conversation_id AND user_id = :user_id """) with get_engine().connect() as connection: if connection.execute(query, {"conversation_id": conversation_id, "user_id": user_id}).scalar_one_or_none() is None: raise ValueError("Conversation does not exist for this user") def list_messages(conversation_id: int, user_id: int, limit: int = 100) -> list[dict]: ensure_conversation_owner(conversation_id, user_id) query = text(""" SELECT id, role, content, metadata, created_at FROM ( SELECT id, role, content, metadata, created_at FROM chat_message WHERE conversation_id = :conversation_id ORDER BY created_at DESC, id DESC LIMIT :limit ) recent ORDER BY created_at, id """) with get_engine().connect() as connection: return [dict(row) for row in connection.execute( query, {"conversation_id": conversation_id, "limit": limit}, ).mappings()] def add_message(conversation_id: int, role: str, content: str, metadata: dict | None = None) -> dict: if role not in {"user", "assistant", "system", "tool"}: raise ValueError("Invalid message role") query = text(""" INSERT INTO chat_message (conversation_id, role, content, metadata, created_at) VALUES (:conversation_id, :role, :content, CAST(:metadata AS jsonb), CURRENT_TIMESTAMP) RETURNING id, role, content, metadata, created_at """) with get_engine().begin() as connection: result = connection.execute(query, { "conversation_id": conversation_id, "role": role, "content": content, "metadata": json.dumps(metadata or {}), }).mappings().one() connection.execute( text("UPDATE chat_conversation SET updated_at = CURRENT_TIMESTAMP WHERE id = :id"), {"id": conversation_id}, ) return dict(result) def format_history(messages: list[dict]) -> str: labels = {"user": "Usuario", "assistant": "Asistente", "system": "Sistema", "tool": "Herramienta"} return "\n".join(f"{labels.get(item['role'], item['role'])}: {item['content']}" for item in messages)