puuuuush

lib.py

@@ -1,7 +1,9 @@
 import os, pymupdf, whisper, json
 
 import cv2
+import psycopg2
 import pytesseract
+import ollama
 
 
 def read_files(path, output, filetypes=None):
@@ -40,3 +42,20 @@ def extract_image_content(path):
     pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
     img = cv2.imread(path)
     return pytesseract.image_to_string(img)
+
+def prompt_embedding(prompt):
+    embedding = ollama.embeddings(model="mxbai-embed-large", prompt=prompt)
+
+    conn = psycopg2.connect(
+        dbname="embeddings",
+        user="python",
+        password="PasswordPassword123",
+        host="localhost",
+        port="5555"
+    )
+    cur = conn.cursor()
+
+    cur.execute(f"select id, filepath, embedding, embedding <-> {embedding} as distance from dbtable order by distance limit 1;")
+
+    result = cur.fetchall()
+    print(result)

main2.py

@@ -0,0 +1,60 @@
+import pymupdf
+import os
+import lib
+import json
+
+
+def retrieve_file_contents(path):
+    file_extraction_functions = {
+        "pdf": lambda path: lib.extract_pdf_content(path),
+        "jpg": lambda path: lib.extract_image_content(path),
+        "png": lambda path: lib.extract_image_content(path),
+        "txt": lambda path: lib.extract_pdf_content(path),
+        "mp3": lambda path: lib.extract_mp3_content(path),
+    }
+
+    lib.read_files(path, files := [])
+
+    contents = []
+
+    for file in files:
+        content = file_extraction_functions[file[0]](file[1])
+        contents.append({
+            "type": file[0],
+            "path": file[1],
+            "filename": file[2],
+            "content": content
+        })
+
+    return contents
+
+
+def reload_files():
+    path = input("Please provider path to folder: ")
+    pass
+
+
+def add_files():
+    path = input("Please provider path to folder: ")
+    pass
+
+
+def prompt_cycle():
+    while True:
+        prompt = input("Please enter prompt: ")
+        lib.prompt_embedding(prompt)
+
+
+def get_user_action():
+    user_action = input("Reload Files (r), Add Files (a), Prompt (p): ")
+
+    {
+        "r": lambda: reload_files(),
+        "a": lambda: add_files(),
+        "p": lambda: prompt_cycle()
+    }[user_action]()
+
+# contents = retrieve_file_contents("C:\\SoftwareEng")
+# print(json.dumps(contents, indent="\t"))
+
+get_user_action()