Kairo-Brain / storage.py
ekjotsingh's picture
Update storage.py
7675694 verified
import sqlite3
import pandas as pd
import time
import threading
DB_FILE = "hunter_data.db"
db_lock = threading.Lock()
def init_db():
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
c = conn.cursor()
c.execute('''CREATE TABLE IF NOT EXISTS results
(symbol TEXT PRIMARY KEY, score INTEGER, decision TEXT,
trend TEXT, metrics TEXT, reason TEXT, timestamp REAL)''')
c.execute('''CREATE TABLE IF NOT EXISTS state
(key TEXT PRIMARY KEY, value TEXT)''')
conn.commit()
conn.close()
def save_result(data):
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
c = conn.cursor()
c.execute('''INSERT OR REPLACE INTO results VALUES (?, ?, ?, ?, ?, ?, ?)''',
(data['Ticker'], data['Score'], data['Decision'],
data['Trend'], data['Key Metrics'], data['AI Reasoning'], time.time()))
conn.commit()
conn.close()
def fetch_results():
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
try:
df = pd.read_sql("SELECT * FROM results ORDER BY timestamp DESC", conn)
except:
df = pd.DataFrame()
conn.close()
return df
def save_progress(index):
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
c = conn.cursor()
c.execute("INSERT OR REPLACE INTO state (key, value) VALUES ('last_index', ?)", (str(index),))
conn.commit()
conn.close()
def load_progress():
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
c = conn.cursor()
c.execute("SELECT value FROM state WHERE key='last_index'")
result = c.fetchone()
conn.close()
return int(result[0]) if result else 0
def reset_database():
with db_lock:
conn = sqlite3.connect(DB_FILE, check_same_thread=False)
c = conn.cursor()
c.execute("DELETE FROM results")
c.execute("DELETE FROM state")
conn.commit()
conn.close()