From 305629b1f1e19214a03db358f395b6ff35e56d81 Mon Sep 17 00:00:00 2001
From: Amira24523 <85802878+Amira24523@users.noreply.github.com>
Date: Wed, 11 Dec 2024 23:39:46 +0100
Subject: [PATCH] BingBang

---
 lib.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/lib.py b/lib.py
index 2ca388a..8eae2bb 100644
--- a/lib.py
+++ b/lib.py
@@ -2,8 +2,12 @@ import os, pymupdf, whisper, json
 
 import cv2
 import psycopg2
+import pymongo
 import pytesseract
 import ollama
+from aiohttp.web_response import json_response
+
+import ollama_conn
 
 
 def read_files(path, output, filetypes=None):
@@ -59,5 +63,21 @@ def prompt_embedding(prompt):
 
     result = cur.fetchall()
 
+    res = get_data_from_mongo(result)
+
+    print(ollama_conn.ask_ollama(prompt, res))
+
     conn.close()
-    print(result)
\ No newline at end of file
+
+def get_data_from_mongo(dataset):
+    result_string = ""
+
+    client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
+    mongodb = client['document_table']
+    collection = mongodb['documents']
+
+    for set in dataset:
+        json_data = collection.find_one({"doc_id": set[0]})
+        result_string += json_data["chunk_content"]
+
+    return result_string
\ No newline at end of file