# Written by retoor@molodetz.nl

# This script connects to a Pinecone assistant service and provides a command-line interface for user interaction. It includes functionality for creating or retrieving an assistant, handling chat messages, and uploading files to the assistant with specific metadata.

# External library used: Pinecone, along with pinecone_plugins.assistant.models.chat for handling message interactions with the assistant.

# MIT License


from pinecone import Pinecone
import os
import pathlib
import sys
from pinecone_plugins.assistant.models.chat import Message

pc = Pinecone(api_key=os.environ.get('PINECONE_API_KEY'))

ASSISTANT_NAME = "liebranca2"

def get_or_create_assistant(assistant_name):
    assistant = None
    try:
        assistant = pc.assistant.create_assistant(
            assistant_name=assistant_name,
            instructions="Your data is based on someone called Liebranka. Act as Liebranka and talk like Liebranca. Respond as Liebranca would do. You also are used to be called libry.",
            region="us",
            timeout=30
        )
    except Exception as e:
        assistant = pc.assistant.Assistant(assistant_name=assistant_name)
    return assistant

assistant = get_or_create_assistant(ASSISTANT_NAME)

def cli():
    messages = []
    while True:
        messages.append(Message(role="user", content=input("> ")))
        resp = assistant.chat(messages=messages)
        print(resp['message']['content'])

def upload(path, glob_pattern):
    files = pathlib.Path(path).glob(glob_pattern)
    for file in files:
        try:
            response = assistant.upload_file(
                file_path=file.absolute(),
                metadata={"document_type": "text"},
                timeout=None
            )
            print(response)
        except Exception as e:
            print(e)

if "--test" in sys.argv or "train" in sys.argv:
    upload("/home/retoor/projects/drstats/export", "*Liebranca*")
    upload("/home/retoor/projects/drstats/export", "*mentions*")
    print("Training complete.")
    exit()

cli()