util.py

from numbers import Number
import json
import ollama
from collections import defaultdict

def jsonify(value):
    if isinstance(value, Number):
        return value
    elif isinstance(value, str):
        return value
    elif isinstance(value, bool):
        return value
    elif value is None:
        return value
    elif isinstance(value, list):
        return [jsonify(item) for item in value]
    elif isinstance(value, dict):
        return {str(k): jsonify(v) for k,v in value.items()}
    else:
        return str(value)

def write_jsonl(line, fp):
    line_str = json.dumps(line)
    if '\n' in line_str:
        raise ValueError(f"Newline in jsonl data: {line_str}")
    fp.write(line_str + '\n')

def load_jsonl(path, keys=None):
    items = []
    with path.open() as f:
        for line in f:
            item = json.loads(line)
            if keys:
                item = {k: item[k] for k in keys}
            items.append(item)
    return items 

def group_records_by(data, key):
    new_data = defaultdict(lambda: [])
    for item in data:
        new_data[item[key]].append(item)
    return list(new_data.values())

def find_nth(haystack: str, needle: str, n: int) -> int:
    """ https://stackoverflow.com/a/1884277 """
    start = haystack.find(needle)
    while start >= 0 and n > 1:
        start = haystack.find(needle, start+len(needle))
        n -= 1
    return start

### OLLAMA STUFF ###

def get_available_models(include_generics=True):
    models = sorted(ollama.list()['models'], key=lambda x: x['size'])
    default_model = models[0]['name']
    available_models = [m['name'] for m in models]
    if include_generics:
        available_models += list({m.split(':')[0] for m in available_models})
    return available_models

def collect_stream(stream, do_print=True):
    message = ''
    for i, chunk in enumerate(stream):
        if i == 0 and do_print:
            print(f"{chunk['message']['role']}: ", end='')
        message_chunk = chunk['message']['content']
        message += message_chunk
        if do_print:
            end = '\n' if chunk['done'] else ''
            print(message_chunk, end=end, flush=True)
    return message

def print_message(m, **kwargs):
    print(f"{m['role']}: {m['content']}", **kwargs)

### END OLLAMA STUFF ###