114 lines
2.9 KiB
Python
114 lines
2.9 KiB
Python
import ollama
|
|
import pymupdf
|
|
import os
|
|
|
|
import chunker
|
|
import lib
|
|
import json
|
|
import psycopg2
|
|
import pymongo
|
|
|
|
count = 0
|
|
|
|
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 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, chunk)
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def insert_data_mongo(id, pChunk):
|
|
client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
|
|
mongodb = client['document_table']
|
|
collection = mongodb['documents']
|
|
dokument = {
|
|
'doc_id': id,
|
|
'chunk_content': pChunk,
|
|
}
|
|
collection.insert_one(dokument)
|
|
|
|
def reload_files():
|
|
# path = input("Please provider path to folder: ")
|
|
contents = retrieve_file_contents("C:\\SoftwareEng")
|
|
create_embeddings(contents)
|
|
|
|
|
|
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]()
|
|
|
|
get_user_action()
|
|
|
|
#print(json.dumps(contents, indent="\t")) |