From 1e9e521c3a5f7797587891ff0a0a7b3e8b82cd71 Mon Sep 17 00:00:00 2001 From: Ali Date: Wed, 11 Dec 2024 21:28:53 +0100 Subject: [PATCH] puuuuush database vektor + mongo insert of weird gay ass shit LoL --- main.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/main.py b/main.py index b63d1fc..c4c00df 100644 --- a/main.py +++ b/main.py @@ -4,8 +4,9 @@ import os import lib import json import psycopg2 +import pymongo - +count = 0 def retrieve_file_contents(path): file_extraction_functions = { @@ -33,6 +34,7 @@ def retrieve_file_contents(path): def create_embeddings(pContent): + global count conn = psycopg2.connect( dbname="embeddings", user="python", @@ -45,6 +47,7 @@ def create_embeddings(pContent): create_table_query = ''' create table if not exists dbtable ( id SERIAL PRIMARY KEY, + doc_id Serial NOT NULL, filepath TEXT NOT NULL, embedding VECTOR NOT NULL ); @@ -62,16 +65,26 @@ def create_embeddings(pContent): # 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);" + insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s) Returning doc_id;" cur.execute(insert_data, (response["embedding"],)) - + doc_id = cur.fetchone()[0] + insert_data_mongo(doc_id, "Hello World") conn.commit() conn.close() +def insert_data_mongo(id, pChunk): + client = pymongo.MongoClient('mongodb://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") + contents = retrieve_file_contents("C:\\Users\\afist\\Downloads\\SoftwareEng") create_embeddings(contents)