From 49d59c3dfa05a73bf207960090c51838b7a70e5d Mon Sep 17 00:00:00 2001 From: rohan-uiuc Date: Mon, 30 Sep 2024 15:54:26 -0500 Subject: [PATCH] Changes for redis migration --- ai_ta_backend/database/redis.py | 12 +++++++ ai_ta_backend/main.py | 3 ++ ai_ta_backend/service/project_service.py | 43 +++++++++++++----------- requirements.txt | 1 + 4 files changed, 40 insertions(+), 19 deletions(-) create mode 100644 ai_ta_backend/database/redis.py diff --git a/ai_ta_backend/database/redis.py b/ai_ta_backend/database/redis.py new file mode 100644 index 00000000..eecbfbfb --- /dev/null +++ b/ai_ta_backend/database/redis.py @@ -0,0 +1,12 @@ +import os + +import redis + + +class RedisDatabase: + + def __init__(self): + self.redis_client = redis.Redis(host=os.environ['REDIS_URL'], password=os.environ['REDIS_PASSWORD'], db=0) + + def hset(self, key: str, value: str): + return self.redis_client.hset(key, value) diff --git a/ai_ta_backend/main.py b/ai_ta_backend/main.py index aea7937b..001ddb1d 100644 --- a/ai_ta_backend/main.py +++ b/ai_ta_backend/main.py @@ -20,6 +20,7 @@ from ai_ta_backend.beam.nomic_logging import create_document_map from ai_ta_backend.database.aws import AWSStorage +from ai_ta_backend.database.redis import RedisDatabase from ai_ta_backend.database.sql import SQLDatabase from ai_ta_backend.database.vector import VectorDatabase from ai_ta_backend.executors.flask_executor import ( @@ -106,6 +107,7 @@ def getTopContexts(service: RetrievalService) -> Response: Testing how exceptions are handled. """ data = request.get_json() + print("data: ", data) search_query: str = data.get('search_query', '') course_name: str = data.get('course_name', '') token_limit: int = data.get('token_limit', 3000) @@ -588,6 +590,7 @@ def configure(binder: Binder) -> None: binder.bind(VectorDatabase, to=VectorDatabase, scope=SingletonScope) binder.bind(SQLDatabase, to=SQLDatabase, scope=SingletonScope) binder.bind(AWSStorage, to=AWSStorage, scope=SingletonScope) + binder.bind(RedisDatabase, to=RedisDatabase, scope=SingletonScope) binder.bind(ExecutorInterface, to=FlaskExecutorAdapter(executor), scope=SingletonScope) diff --git a/ai_ta_backend/service/project_service.py b/ai_ta_backend/service/project_service.py index d336d2d2..eae80a35 100644 --- a/ai_ta_backend/service/project_service.py +++ b/ai_ta_backend/service/project_service.py @@ -1,9 +1,8 @@ import json -import os -import requests from injector import inject +from ai_ta_backend.database.redis import RedisDatabase from ai_ta_backend.database.sql import SQLDatabase from ai_ta_backend.service.posthog_service import PosthogService from ai_ta_backend.service.sentry_service import SentryService @@ -18,8 +17,10 @@ class ProjectService: """ @inject - def __init__(self, sql_db: SQLDatabase, posthog_service: PosthogService, sentry_service: SentryService): + def __init__(self, sql_db: SQLDatabase, posthog_service: PosthogService, sentry_service: SentryService, + redis_db: RedisDatabase): self.sqlDb = sql_db + self.redisDb = redis_db self.posthog = posthog_service self.sentry = sentry_service @@ -44,11 +45,11 @@ def create_project(self, project_name: str, project_description: str | None, pro """ try: # Insert project into Redis - headers = { - "Authorization": - f"Bearer {os.environ['KV_REST_API_TOKEN']}", # Ensure you use the appropriate write-access API key - "Content-Type": "application/json" - } + # headers = { + # "Authorization": + # f"Bearer {os.environ['KV_REST_API_TOKEN']}", # Ensure you use the appropriate write-access API key + # "Content-Type": "application/json" + # } # Define the key-value pair you want to insert key = project_name # Replace with your key value = { @@ -66,16 +67,21 @@ def create_project(self, project_name: str, project_description: str | None, pro } # Construct the URL for the HSET request - hset_url = str(os.environ['KV_REST_API_URL']) + f"/hset/course_metadatas/{key}" + # hset_url = str(os.environ['KV_REST_API_URL']) + f"/hset/course_metadatas/{key}" # Make the POST request to insert the key-value pair - response = requests.post(hset_url, headers=headers, data=json.dumps(value)) + # response = requests.post(hset_url, headers=headers, data=json.dumps(value)) + response = self.redisDb.hset(project_name, json.dumps(value)) - # Check the response status - if response.status_code == 200: + if response == 1: print("Course metadata inserted successfully.") else: - print(f"Failed to insert course metadata. Status code: {response.status_code}, Response: {response.text}") + print(f"Failed to insert course metadata. Status code: {response}") + # Check the response status + # if response.status_code == 200: + # print("Course metadata inserted successfully.") + # else: + # print(f"Failed to insert course metadata. Status code: {response.status_code}, Response: {response.text}") # check if the project owner has pre-assigned API keys if project_owner_email: @@ -93,15 +99,14 @@ def create_project(self, project_name: str, project_description: str | None, pro llm_val[row['providerName']] = row['providerBodyNoModels'] # Insert the pre-assigned API keys into Redis - set_llm_url = str(os.environ['KV_REST_API_URL']) + f"/set/{redis_key}" - set_response = requests.post(set_llm_url, headers=headers, data=json.dumps(llm_val)) - + # set_llm_url = str(os.environ['KV_REST_API_URL']) + f"/set/{redis_key}" + # set_response = requests.post(set_llm_url, headers=headers, data=json.dumps(llm_val)) + set_response = self.redisDb.hset(redis_key, json.dumps(llm_val)) # Check the response status - if set_response.status_code == 200: + if set_response == 1: print("LLM key-value pair inserted successfully.") else: - print( - f"Failed to insert LLM key-value pair. Status code: {response.status_code}, Response: {response.text}") + print(f"Failed to insert LLM key-value pair. Status code: {set_response}") return "success" except Exception as e: diff --git a/requirements.txt b/requirements.txt index cffc3c97..4c502fde 100644 --- a/requirements.txt +++ b/requirements.txt @@ -37,6 +37,7 @@ langchainhub==0.1.14 boto3==1.28.79 qdrant-client==1.7.3 supabase==2.5.3 +redis==5.1.0 # Logging posthog==3.1.0