BingBang

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)
+
+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