Spaces:
Sleeping
Sleeping
File size: 4,316 Bytes
36bfe21 | 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 115 116 117 118 119 120 121 122 123 124 125 | import psycopg2
from psycopg2.extras import RealDictCursor
from contextlib import contextmanager
from typing import Optional, List, Dict, Any
from uuid import UUID, uuid4
from datetime import datetime
from app.config import get_settings
class PostgresDB:
"""PostgreSQL database client for conversation storage."""
def __init__(self):
self.settings = get_settings()
self.connection_string = self.settings.neon_database_url
@contextmanager
def get_connection(self):
"""Context manager for database connections."""
conn = psycopg2.connect(self.connection_string)
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
def create_conversation(self, metadata: Optional[Dict[str, Any]] = None) -> UUID:
"""Create a new conversation."""
conversation_id = uuid4()
with self.get_connection() as conn:
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO conversations (id, metadata)
VALUES (%s, %s)
RETURNING id
""",
(str(conversation_id), psycopg2.extras.Json(metadata or {}))
)
result = cur.fetchone()[0]
return UUID(result) if isinstance(result, str) else result
def add_message(
self,
conversation_id: UUID,
role: str,
content: str,
context_used: Optional[List[str]] = None,
metadata: Optional[Dict[str, Any]] = None
) -> UUID:
"""Add a message to a conversation."""
message_id = uuid4()
with self.get_connection() as conn:
with conn.cursor() as cur:
cur.execute(
"""
INSERT INTO messages (id, conversation_id, role, content, context_used, metadata)
VALUES (%s, %s, %s, %s, %s, %s)
RETURNING id
""",
(
str(message_id),
str(conversation_id),
role,
content,
context_used or [],
psycopg2.extras.Json(metadata or {})
)
)
result = cur.fetchone()[0]
return UUID(result) if isinstance(result, str) else result
def get_conversation(self, conversation_id: UUID) -> Optional[Dict[str, Any]]:
"""Get a conversation with all its messages."""
with self.get_connection() as conn:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
# Get conversation
cur.execute(
"SELECT * FROM conversations WHERE id = %s",
(str(conversation_id),)
)
conversation = cur.fetchone()
if not conversation:
return None
# Get messages
cur.execute(
"""
SELECT * FROM messages
WHERE conversation_id = %s
ORDER BY created_at ASC
""",
(str(conversation_id),)
)
messages = cur.fetchall()
return {
**dict(conversation),
'messages': [dict(msg) for msg in messages]
}
def get_conversation_history(
self,
conversation_id: UUID,
limit: int = 10
) -> List[Dict[str, Any]]:
"""Get recent messages from a conversation."""
with self.get_connection() as conn:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute(
"""
SELECT * FROM messages
WHERE conversation_id = %s
ORDER BY created_at DESC
LIMIT %s
""",
(str(conversation_id), limit)
)
messages = cur.fetchall()
return [dict(msg) for msg in reversed(messages)]
|