Spaces:
Running
Running
File size: 4,228 Bytes
9a1014e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | 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)
|