from flask import Flask, request, jsonify
import asyncio
import axios

app = Flask(__name__)

# Load Markdown files
def loadMarkdownFiles(directory):
  files = os.listdir(directory)
  markdownFiles = []

  for file in files:
    if file.endswith('.md'):
      filePath = os.path.join(directory, file)
      with open(filePath, 'r', encoding='utf-8') as f:
        content = f.read()
        markdownFiles.append({'name': file, 'content': content})

  return markdownFiles

# Initialize RAG model and tokenizer
async def initRagModel():
  from transformers import AutoTokenizer, RagTokenForGeneration
  tokenizer = await AutoTokenizer.from_pretrained('facebook/rag-token-nq')
  model = await RagTokenForGeneration.from_pretrained('facebook/rag-token-nq')

  return {'tokenizer': tokenizer, 'model': model}

# Retrieve relevant information from Markdown files using Ollama API
async def retrieveInformation(query):
  try:
    config = configparser.ConfigParser()
    config.read('ollama.ini')
    host = config.get('Ollama', 'host')
    port = config.get('Ollama', 'port')

    response = await axios.post(f'http://{host}:{port}/chat', {'query': query})
    return response.json()['response']
  except (axios.AxiosError, Exception) as error:
    print(f'Error: {error}')
    raise

# Chatbot logic
async def chatbot():
  directory = './notes'  # Directory containing Markdown files
  markdownFiles = loadMarkdownFiles(directory)
  ragModel = await initRagModel()

  print('Chatbot is ready! Ask your questions.')

  while True:
    query = input().strip()
    if query.lower() == 'exit':
      break

    try:
      response = await retrieveInformation(query)
      print(f'Chatbot: {response}')
    except Exception as error:
      print(f'Error: {error}')

if __name__ == '__main__':
  asyncio.run(chatbot())
  app.run(debug=True)