mirror of
https://github.com/ION606/ML-pipeline.git
synced 2026-05-14 21:06:54 +00:00
85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
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.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()
|