diff --git a/lib.py b/lib.py index 883cd5a..81cdaf7 100644 --- a/lib.py +++ b/lib.py @@ -10,6 +10,7 @@ from PIL import Image from aiohttp.web_response import json_response import ollama_conn +import redis_conn def read_files(path, output, filetypes=None): @@ -91,7 +92,9 @@ def prompt_embedding(prompt): res = get_data_from_mongo(new_result) - print(ollama_conn.ask_ollama(prompt, res)) + ollama_response = ollama_conn.ask_ollama(prompt, res) + print(ollama_response) + redis_conn.cache_prompt_to_redis(prompt.lower().strip(), ollama_response) conn.close() diff --git a/main.py b/main.py index 4dbd466..910d243 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,10 @@ import json import psycopg2 import pymongo +import mongo_conn +import pgvec_conn +import redis_conn + count = 0 def retrieve_file_contents(path): @@ -113,7 +117,13 @@ def insert_data_mongo(id, filepath,pChunk): def reload_files(): # path = input("Please provider path to folder: ") - contents = retrieve_file_contents("C:\\SoftwareEng") + 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) @@ -125,7 +135,10 @@ def add_files(): def prompt_cycle(): while True: prompt = input("Please enter prompt: ") - lib.prompt_embedding(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(): diff --git a/mongo_conn.py b/mongo_conn.py new file mode 100644 index 0000000..8e55ea1 --- /dev/null +++ b/mongo_conn.py @@ -0,0 +1,7 @@ +import pymongo + +def flush_mongo(): + client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/') + mongodb = client['document_table'] + mongodb.drop_collection("documents") + client.close() \ No newline at end of file diff --git a/pgvec_conn.py b/pgvec_conn.py new file mode 100644 index 0000000..94a3f7e --- /dev/null +++ b/pgvec_conn.py @@ -0,0 +1,16 @@ +import psycopg2 + +def flush_pg(): + conn = psycopg2.connect( + dbname="embeddings", + user="python", + password="PasswordPassword123", + host="localhost", + port="5555" + ) + cur = conn.cursor() + + cur.execute("drop table if exists dbtable;") + conn.commit() + + conn.close() \ No newline at end of file diff --git a/redis_conn.py b/redis_conn.py new file mode 100644 index 0000000..619bdf9 --- /dev/null +++ b/redis_conn.py @@ -0,0 +1,17 @@ +import redis + +def cache_prompt_to_redis(prompt, response): + r = redis.Redis(host="localhost", port=6379, decode_responses=True) + r.set(prompt, response) + r.close() + +def load_response_from_redis(prompt): + r = redis.Redis(host="localhost", port=6379, decode_responses=True) + response = r.get(prompt) + r.close() + return response + +def flush_redis(): + r = redis.Redis(host="localhost", port=6379, decode_responses=True) + r.flushall() + r.close() \ No newline at end of file