155 lines
4.1 KiB
Python
155 lines
4.1 KiB
Python
import ollama
|
|
import pymupdf
|
|
import os
|
|
from minio import Minio
|
|
from minio.error import S3Error
|
|
import chunker
|
|
import lib
|
|
import json
|
|
import psycopg2
|
|
import pymongo
|
|
|
|
import mongo_conn
|
|
import pgvec_conn
|
|
import redis_conn
|
|
|
|
count = 0
|
|
|
|
def retrieve_file_contents(path):
|
|
|
|
#Minio Connection
|
|
minioClient = Minio(
|
|
"localhost:9000",
|
|
access_key="PVFOeJbx87rQyi0WXF1X",
|
|
secret_key="Am8Cd9auYGbEGuEXfJtnWEPsMwJCx9N58NCNHCgs",
|
|
secure=False,
|
|
)
|
|
|
|
file_extraction_functions = {
|
|
"pdf": lambda path: lib.extract_text_and_pictures(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])
|
|
with open(f"./{file[2]}.txt", "w") as future_s3_file:
|
|
future_s3_file.writelines(file[1] + "\n" + content)
|
|
future_s3_file.flush()
|
|
future_s3_file.close()
|
|
|
|
minioClient.fput_object(
|
|
bucket_name="datafiles",
|
|
object_name=f"{file[2]}.txt",
|
|
file_path=file[1],
|
|
)
|
|
|
|
os.remove(f"./{file[2]}.txt")
|
|
|
|
|
|
contents.append({
|
|
"type": file[0],
|
|
"path": file[1],
|
|
"filename": file[2],
|
|
"content": content
|
|
})
|
|
|
|
return contents
|
|
|
|
|
|
def create_embeddings(pContent):
|
|
global count
|
|
conn = psycopg2.connect(
|
|
dbname="embeddings",
|
|
user="python",
|
|
password="PasswordPassword123",
|
|
host="localhost",
|
|
port="5555"
|
|
)
|
|
cur = conn.cursor()
|
|
|
|
create_table_query = '''
|
|
create table if not exists dbtable (
|
|
id SERIAL PRIMARY KEY,
|
|
filepath TEXT NOT NULL,
|
|
embedding VECTOR NOT NULL
|
|
);
|
|
'''
|
|
|
|
cur.execute('CREATE EXTENSION IF NOT EXISTS vector;')
|
|
|
|
cur.execute(create_table_query)
|
|
|
|
conn.commit()
|
|
|
|
for content in pContent:
|
|
for chunk in chunker.generate_chunks(content["content"]):
|
|
merged_info = "Dateiname: " + content["filename"] + " Dateiinhalt: " + chunk
|
|
# print(merged_info)
|
|
response = ollama.embeddings(model="mxbai-embed-large", prompt=merged_info)
|
|
#embedding_list.append(response["embedding"])
|
|
insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s) Returning id;"
|
|
cur.execute(insert_data, (response["embedding"],))
|
|
doc_id = cur.fetchone()[0]
|
|
insert_data_mongo(doc_id, content['path'], chunk)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def insert_data_mongo(id, filepath,pChunk):
|
|
client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
|
|
mongodb = client['document_table']
|
|
collection = mongodb['documents']
|
|
dokument = {
|
|
'doc_id': id,
|
|
'filepath': filepath,
|
|
'chunk_content': pChunk,
|
|
}
|
|
result = collection.insert_one(dokument)
|
|
|
|
if result.acknowledged:
|
|
collection.create_index("doc_id")
|
|
|
|
def reload_files():
|
|
# path = input("Please provider path to folder: ")
|
|
pgvec_conn.flush_pg()
|
|
redis_conn.flush_redis()
|
|
mongo_conn.flush_mongo()
|
|
|
|
provided_path = input("Please Provide Full Qualified Path: ")
|
|
|
|
contents = retrieve_file_contents(provided_path)
|
|
create_embeddings(contents)
|
|
|
|
|
|
def add_files():
|
|
path = input("Please provider path to folder: ")
|
|
pass
|
|
|
|
|
|
def prompt_cycle():
|
|
while True:
|
|
prompt = input("Please enter prompt: ")
|
|
|
|
response = redis_conn.load_response_from_redis(prompt.lower().strip())
|
|
|
|
lib.prompt_embedding(prompt) if response is None else print("Cached Response:", response)
|
|
|
|
|
|
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]()
|
|
|
|
get_user_action()
|
|
|
|
#print(json.dumps(contents, indent="\t")) |