goober/modules/markovmemory.py

import os
import json
import markovify
import pickle
from modules.globalvars import *
from modules.volta.main import _
import logging
logger = logging.getLogger("goober")

def get_file_info(file_path: str) -> dict:
    try:
        file_size = os.path.getsize(file_path)
        with open(file_path, "r") as f:
            lines = f.readlines()
        return {
            "file_size_bytes": file_size,
            "line_count": len(lines)
        }
    except Exception as e:
        return {"error": str(e)}

def load_memory() -> list:
    try:
        with open(MEMORY_FILE, "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return []

def save_memory(memory: list) -> None:
    with open(MEMORY_FILE, "w") as f:
        json.dump(memory, f, indent=4)

def train_markov_model(memory: list, additional_data: list = None):
    lines = [line for line in (memory or []) if isinstance(line, str)]

    if additional_data:
        lines.extend(line for line in additional_data if isinstance(line, str))

    if not lines:
        return None

    text = "\n".join(lines)
    return markovify.NewlineText(text, state_size=2)

def save_markov_model(model, filename: str = 'markov_model.pkl') -> None:
    with open(filename, 'wb') as f:
        pickle.dump(model, f)

def load_markov_model(filename: str = 'markov_model.pkl'):
    try:
        with open(filename, 'rb') as f:
            model = pickle.load(f)
        logger.info(f"{_('model_loaded')} {filename}.{RESET}")
        return model
    except FileNotFoundError:
        logger.error(f"{filename} {_('not_found')}{RESET}")
        return None