ML-pipeline/conversation_store.py

from datetime import datetime
import sqlite3
import os

import debug as debugModule
from config import Config


def initialize_db():
	conn = sqlite3.connect(Config.CONVERSATION_DB_PATH)
	c = conn.cursor()

	# Create conversations table
	c.execute('''CREATE TABLE IF NOT EXISTS conversations
				 (id INTEGER PRIMARY KEY AUTOINCREMENT,
				  query TEXT,
				  response TEXT,
				  links TEXT,
				  timestamp DATETIME)''')

	# Create code_executions table
	c.execute('''CREATE TABLE IF NOT EXISTS code_executions
				 (id INTEGER PRIMARY KEY AUTOINCREMENT,
				  conversation_id INTEGER,
				  code_content TEXT,
				  execution_result TEXT,
				  error_message TEXT,
				  retry_count INTEGER,
				  timestamp DATETIME,
				  FOREIGN KEY(conversation_id) REFERENCES conversations(id))''')

	conn.commit()
	conn.close()
	debugModule.log("database initialized successfully!")


def save_conversation(query, response, links):
	conn = sqlite3.connect(Config.CONVERSATION_DB_PATH)
	c = conn.cursor()
	c.execute('''INSERT INTO conversations 
				 (query, response, links, timestamp)
				 VALUES (?, ?, ?, ?)''',
           (query, response, ", ".join(links), datetime.now()))
	conn.commit()
	conn.close()
	return c.lastrowid  # Return the generated conversation ID because it's easier than uuids


def save_code_execution(code, result, error=None, retries=0, conversation_id=None):
	conn = sqlite3.connect(Config.CONVERSATION_DB_PATH)
	c = conn.cursor()

	# Convert execution result to string if needed
	if isinstance(result, dict):
		execution_result = str(result.get('err', ''))
		error_message = result.get('err', '')
	else:
		execution_result = str(result)
		error_message = error or ''

	c.execute('''INSERT INTO code_executions
				 (conversation_id, code_content, execution_result, 
				  error_message, retry_count, timestamp)
				 VALUES (?, ?, ?, ?, ?, ?)''',
           (conversation_id, code, execution_result,
            error_message, retries, datetime.now()))

	conn.commit()
	conn.close()


def save_conversation(query, response, links):
	if len(query) > Config.MAX_QUERY_LENGTH or len(response) > Config.MAX_RESPONSE_LENGTH:
		raise ValueError("Input too large")

	conn = sqlite3.connect(Config.CONVERSATION_DB_PATH)
	cursor = conn.cursor()
	links_str = ", ".join(links) if links else ""
	cursor.execute(
		"INSERT INTO conversations (query, response, links) VALUES (?, ?, ?)",
		(query, response, links_str)
	)
	conn.commit()
	conn.close()