From 02fb7240f9adccb94f844fdfbd6efab74a5ea8c9 Mon Sep 17 00:00:00 2001 From: Mohammad Amin Date: Thu, 28 Dec 2023 11:00:27 +0330 Subject: [PATCH] feat: Using the shared library! We have to update the test cases too. --- .../discord_mongo_summary_etl.py | 6 +- .../discord_mongo_vector_store_etl.py | 6 +- .../discourse_summary_etl.py | 6 +- .../discourse_vectorstore_etl.py | 6 +- dags/hivemind_etl_helpers/gdrive_etl.py | 4 +- .../src/db/gdrive/db_utils.py | 4 +- .../src/db/gdrive/delete_records.py | 2 +- .../src/utils/cohere_embedding.py | 77 ----------- .../src/utils/credentials.py | 24 ---- .../src/utils/load_llm_params.py | 33 ----- .../src/utils/pg_db_utils.py | 126 ------------------ .../src/utils/pg_vector_access.py | 6 +- .../src/utils/postgres.py | 39 ------ .../test_discord_prepare_summary.py | 2 +- .../test_gdrive_db_fetch_files_modified_at.py | 2 +- .../integration/test_gdrive_delete_records.py | 2 +- .../test_gdrive_documents_to_delete.py | 2 +- .../test_gdrive_documents_to_insert.py | 2 +- .../tests/integration/test_load_envs.py | 6 +- .../test_pg_vector_access_with_discord.py | 4 +- .../tests/unit/test_load_model_params.py | 2 +- requirements.txt | 3 +- 22 files changed, 32 insertions(+), 332 deletions(-) delete mode 100644 dags/hivemind_etl_helpers/src/utils/cohere_embedding.py delete mode 100644 dags/hivemind_etl_helpers/src/utils/load_llm_params.py delete mode 100644 dags/hivemind_etl_helpers/src/utils/pg_db_utils.py delete mode 100644 dags/hivemind_etl_helpers/src/utils/postgres.py diff --git a/dags/hivemind_etl_helpers/discord_mongo_summary_etl.py b/dags/hivemind_etl_helpers/discord_mongo_summary_etl.py index 55adb63e..e7484549 100644 --- a/dags/hivemind_etl_helpers/discord_mongo_summary_etl.py +++ b/dags/hivemind_etl_helpers/discord_mongo_summary_etl.py @@ -6,9 +6,9 @@ find_guild_id_by_community_id, ) from hivemind_etl_helpers.src.document_node_parser import configure_node_parser -from hivemind_etl_helpers.src.utils.cohere_embedding import CohereEmbedding -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams -from hivemind_etl_helpers.src.utils.pg_db_utils import setup_db +from tc_hivemind_backend.embeddings.cohere import CohereEmbedding +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams +from tc_hivemind_backend.db.pg_db_utils import setup_db from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index.response_synthesizers import get_response_synthesizer diff --git a/dags/hivemind_etl_helpers/discord_mongo_vector_store_etl.py b/dags/hivemind_etl_helpers/discord_mongo_vector_store_etl.py index f4c720c0..9c733b9c 100644 --- a/dags/hivemind_etl_helpers/discord_mongo_vector_store_etl.py +++ b/dags/hivemind_etl_helpers/discord_mongo_vector_store_etl.py @@ -9,9 +9,9 @@ find_guild_id_by_community_id, ) from hivemind_etl_helpers.src.document_node_parser import configure_node_parser -from hivemind_etl_helpers.src.utils.cohere_embedding import CohereEmbedding -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams -from hivemind_etl_helpers.src.utils.pg_db_utils import setup_db +from tc_hivemind_backend.embeddings.cohere import CohereEmbedding +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams +from tc_hivemind_backend.db.pg_db_utils import setup_db from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess diff --git a/dags/hivemind_etl_helpers/discourse_summary_etl.py b/dags/hivemind_etl_helpers/discourse_summary_etl.py index dfbd4e36..77264382 100644 --- a/dags/hivemind_etl_helpers/discourse_summary_etl.py +++ b/dags/hivemind_etl_helpers/discourse_summary_etl.py @@ -9,9 +9,9 @@ ) from hivemind_etl_helpers.src.db.discourse.utils.get_forums import get_forums from hivemind_etl_helpers.src.document_node_parser import configure_node_parser -from hivemind_etl_helpers.src.utils.cohere_embedding import CohereEmbedding -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams -from hivemind_etl_helpers.src.utils.pg_db_utils import setup_db +from tc_hivemind_backend.embeddings.cohere import CohereEmbedding +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams +from tc_hivemind_backend.db.pg_db_utils import setup_db from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index import Document from llama_index.response_synthesizers import get_response_synthesizer diff --git a/dags/hivemind_etl_helpers/discourse_vectorstore_etl.py b/dags/hivemind_etl_helpers/discourse_vectorstore_etl.py index 5f71bb96..96e5b68b 100644 --- a/dags/hivemind_etl_helpers/discourse_vectorstore_etl.py +++ b/dags/hivemind_etl_helpers/discourse_vectorstore_etl.py @@ -7,9 +7,9 @@ from hivemind_etl_helpers.src.db.discourse.utils.get_forums import get_forums from hivemind_etl_helpers.src.document_node_parser import configure_node_parser from hivemind_etl_helpers.src.utils.check_documents import check_documents -from hivemind_etl_helpers.src.utils.cohere_embedding import CohereEmbedding -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams -from hivemind_etl_helpers.src.utils.pg_db_utils import setup_db +from tc_hivemind_backend.embeddings.cohere import CohereEmbedding +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams +from tc_hivemind_backend.db.pg_db_utils import setup_db from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess diff --git a/dags/hivemind_etl_helpers/gdrive_etl.py b/dags/hivemind_etl_helpers/gdrive_etl.py index 0f5189b9..9fedd057 100644 --- a/dags/hivemind_etl_helpers/gdrive_etl.py +++ b/dags/hivemind_etl_helpers/gdrive_etl.py @@ -8,8 +8,8 @@ ) from hivemind_etl_helpers.src.document_node_parser import configure_node_parser from hivemind_etl_helpers.src.utils.check_documents import check_documents -from hivemind_etl_helpers.src.utils.cohere_embedding import CohereEmbedding -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams +from tc_hivemind_backend.embeddings.cohere import CohereEmbedding +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess diff --git a/dags/hivemind_etl_helpers/src/db/gdrive/db_utils.py b/dags/hivemind_etl_helpers/src/db/gdrive/db_utils.py index c2ac131f..16fd2bea 100644 --- a/dags/hivemind_etl_helpers/src/db/gdrive/db_utils.py +++ b/dags/hivemind_etl_helpers/src/db/gdrive/db_utils.py @@ -3,8 +3,8 @@ import psycopg2 from dateutil import parser -from hivemind_etl_helpers.src.utils.pg_db_utils import convert_tuple_str -from hivemind_etl_helpers.src.utils.postgres import PostgresSingleton +from tc_hivemind_backend.db.pg_db_utils import convert_tuple_str +from tc_hivemind_backend.db.postgresql import PostgresSingleton def setup_db(community_id: str) -> None: diff --git a/dags/hivemind_etl_helpers/src/db/gdrive/delete_records.py b/dags/hivemind_etl_helpers/src/db/gdrive/delete_records.py index 8b302170..2b7b8b6c 100644 --- a/dags/hivemind_etl_helpers/src/db/gdrive/delete_records.py +++ b/dags/hivemind_etl_helpers/src/db/gdrive/delete_records.py @@ -1,7 +1,7 @@ import logging from hivemind_etl_helpers.src.utils.pg_db_utils import convert_tuple_str -from hivemind_etl_helpers.src.utils.postgres import PostgresSingleton +from tc_hivemind_backend.db.postgresql import PostgresSingleton def delete_records(db_name: str, table_name: str, metadata_file_id: list[str]) -> None: diff --git a/dags/hivemind_etl_helpers/src/utils/cohere_embedding.py b/dags/hivemind_etl_helpers/src/utils/cohere_embedding.py deleted file mode 100644 index b522d8eb..00000000 --- a/dags/hivemind_etl_helpers/src/utils/cohere_embedding.py +++ /dev/null @@ -1,77 +0,0 @@ -import os - -import cohere -from dotenv import load_dotenv -from llama_index.embeddings import BaseEmbedding - - -class CohereEmbedding(BaseEmbedding): - def __init__(self): - super().__init__() - - def prepare_cohere( - self, - ) -> cohere.Client: - """ - setup cohere client - https://cohere.com/ - - - Returns - -------- - client : cohere.Client - the cohere client to query anything - """ - load_dotenv() - key = os.getenv("COHERE_API_KEY") - - client = cohere.Client(key) - return client - - def get_text_embedding( - self, text: str | None = None, texts: list[str] | None = None - ) -> list[float] | list[list[float]]: - co = self.prepare_cohere() - - if text is not None: - response = co.embed( - texts=[text], - model="embed-multilingual-v3.0", - input_type="classification", - truncate=None, - ) - # checking the output to be right - assert len(response.embeddings) == 1 - - return response.embeddings[0] - elif texts is not None: - response = co.embed( - texts=texts, - model="embed-multilingual-v3.0", - input_type="classification", - truncate=None, - ) - return response.embeddings - else: - raise ValueError("Both inputs cannot be None") - - def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]: - """Get text embeddings. - - By default, this is a wrapper around _get_text_embedding. - Can be overridden for batch queries. - - """ - return self.get_text_embedding(texts=texts) # type: ignore - - def _get_text_embedding(self, text: str) -> list[float]: - """Get text embedding.""" - return self.get_text_embedding(text=text) # type: ignore - - def _get_query_embedding(self, query: str) -> list[float]: - """Get query embedding.""" - return self.get_text_embedding(text=query) # type: ignore - - async def _aget_query_embedding(self, query: str) -> list[float]: - """The asynchronous version of _get_query_embedding.""" - raise NotImplementedError("Not implemented!") diff --git a/dags/hivemind_etl_helpers/src/utils/credentials.py b/dags/hivemind_etl_helpers/src/utils/credentials.py index d85cb51d..deea7c4c 100644 --- a/dags/hivemind_etl_helpers/src/utils/credentials.py +++ b/dags/hivemind_etl_helpers/src/utils/credentials.py @@ -27,27 +27,3 @@ def load_mongo_credentials() -> dict[str, str]: mongo_creds["port"] = os.getenv("MONGODB_PORT", "") return mongo_creds - - -def load_postgres_credentials() -> dict[str, str]: - """ - load postgres credentials into a dictionary - - Returns - --------- - credentials : dict[str, str] - a dictionary containing - `host`, `password`, `user`, `port`, and `db_name` as keys - and values are the values for the credentials - """ - load_dotenv() - - credentials: dict[str, str] = {} - - credentials["host"] = os.getenv("POSTGRES_HOST", "") - credentials["password"] = os.getenv("POSTGRES_PASS", "") - credentials["user"] = os.getenv("POSTGRES_USER", "") - credentials["port"] = os.getenv("POSTGRES_PORT", "") - credentials["db_name"] = os.getenv("POSTGRES_DBNAME", "") - - return credentials diff --git a/dags/hivemind_etl_helpers/src/utils/load_llm_params.py b/dags/hivemind_etl_helpers/src/utils/load_llm_params.py deleted file mode 100644 index 3a708176..00000000 --- a/dags/hivemind_etl_helpers/src/utils/load_llm_params.py +++ /dev/null @@ -1,33 +0,0 @@ -import os - -from dotenv import load_dotenv - - -def load_model_hyperparams() -> tuple[int, int]: - """ - load the llm and embedding model hyperparameters (the input parameters) - - Returns - --------- - chunk_size : int - the chunk size to chunk the data - embedding_dim : int - the embedding dimension - """ - load_dotenv() - - chunk_size_str = os.getenv("CHUNK_SIZE") - chunk_size: int - if chunk_size_str is None: - raise ValueError("Chunk size is not given in env") - else: - chunk_size = int(chunk_size_str) - - embedding_dim_str = os.getenv("EMBEDDING_DIM") - embedding_dim: int - if embedding_dim_str is None: - raise ValueError("Embedding dimension size is not given in env") - else: - embedding_dim = int(embedding_dim_str) - - return chunk_size, embedding_dim diff --git a/dags/hivemind_etl_helpers/src/utils/pg_db_utils.py b/dags/hivemind_etl_helpers/src/utils/pg_db_utils.py deleted file mode 100644 index 6026fbaf..00000000 --- a/dags/hivemind_etl_helpers/src/utils/pg_db_utils.py +++ /dev/null @@ -1,126 +0,0 @@ -import logging -from datetime import datetime - -import psycopg2 -from hivemind_etl_helpers.src.utils.postgres import PostgresSingleton - - -def setup_db(community_id: str, dbname: str, latest_date_query: str) -> datetime | None: - """ - setup the database. - create a database if not available, else get the latest message saved - - Parameters - ------------ - community_id : str - the community id for the case of logging - dbname : str - the database name to create or access its database - guild_id : str - the guild id to create a database for - latest_date_query : str - the query to get latest date of a message - - Returns - --------- - from_date : datetime | None - in case of no data available it would be None - """ - msg = f"COMMUNITYID: {community_id} " - from_date: datetime | None = None - connection: psycopg2.extensions.connection - try: - # first connecting to no database to check the database availability - postgres = PostgresSingleton(dbname=None) - connection = postgres.get_connection() - connection.autocommit = True - cursor = connection.cursor() - logging.info(f"{msg}Creating database {dbname}") - cursor.execute(f"CREATE DATABASE {dbname};") - cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;") - cursor.close() - connection.close() - except psycopg2.errors.DuplicateDatabase: - logging.info(f"{msg}Database already exist!") - logging.info(f"{msg}Checking the latest saved message!") - postgres.close_connection() - from_date = get_latest_msg(community_id, dbname, latest_date_query) - - except Exception as exp: - logging.error(f"{msg}database initialization error: {exp}") - - return from_date - - -def get_latest_msg(community_id: str, dbname: str, latest_date_query: str): - from_date: datetime | None = None - msg = f"COMMUNITYID: {community_id} " - - postgres = PostgresSingleton(dbname=dbname) - connection = postgres.get_connection() - connection.autocommit = True - with connection.cursor() as cursor: - try: - # If we had some data previously saved - # fetch the latest date we wanted to work on it - logging.info(f"{msg}Loading the latest date from previous data") - cursor.execute(latest_date_query) - data = cursor.fetchone() - if data is not None: - from_date = data[0] - logging.info(f"{msg}Latest processed message: {from_date}") - else: - logging.info(f"{msg}No processed message, starting from the first!") - except psycopg2.errors.UndefinedTable: - logging.warning(f"{msg}No data to get the latest date") - - postgres.close_connection() - return from_date - - -def delete_data(deletion_query: str, dbname) -> None: - """ - a wrapper function to add the deletion feature - - Parameters - ----------- - deletion_query : str - the query to delete or modify the database - """ - connection: psycopg2.extensions.connection - try: - # first connecting to no database to check the database availability - postgres = PostgresSingleton(dbname=dbname) - connection = postgres.get_connection() - connection.autocommit = True - with connection.cursor() as cursor: - logging.info("Deleting data from postgresql!") - cursor.execute(deletion_query) - except Exception as exp: - logging.error(f"Database deletion error: {exp}") - finally: - postgres.close_connection() - - -def convert_tuple_str(data: list[str]) -> str: - """ - convert a list of inputs to a string tuple that - can be queried within database - - Parameters - ------------ - data : list[str] - list of string items - - Returns - --------- - data_str : str - the data converted to string tuple - """ - data_str: str - if len(data) == 1: - data_str = f"('{data[0]}')" - else: - data_str = str(tuple(data)) - - return data_str diff --git a/dags/hivemind_etl_helpers/src/utils/pg_vector_access.py b/dags/hivemind_etl_helpers/src/utils/pg_vector_access.py index f8304bee..58154845 100644 --- a/dags/hivemind_etl_helpers/src/utils/pg_vector_access.py +++ b/dags/hivemind_etl_helpers/src/utils/pg_vector_access.py @@ -1,9 +1,9 @@ import logging import time -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams -from hivemind_etl_helpers.src.utils.pg_db_utils import delete_data +from tc_hivemind_backend.db.credentials import load_postgres_credentials +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams +from tc_hivemind_backend.db.utils.delete_data import delete_data from llama_index import Document, MockEmbedding, ServiceContext, StorageContext from llama_index.embeddings import BaseEmbedding, OpenAIEmbedding from llama_index.indices.vector_store import VectorStoreIndex diff --git a/dags/hivemind_etl_helpers/src/utils/postgres.py b/dags/hivemind_etl_helpers/src/utils/postgres.py deleted file mode 100644 index a06b59ad..00000000 --- a/dags/hivemind_etl_helpers/src/utils/postgres.py +++ /dev/null @@ -1,39 +0,0 @@ -import psycopg2 -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials - - -class PostgresSingleton: - _instance = None - - def __new__(cls, dbname: str | None, *args, **kwargs): - if cls._instance is None: - cls._instance = super().__new__(cls, *args, **kwargs) - cls._instance.connect(dbname) - return cls._instance - - def connect(self, dbname: str | None): - creds = load_postgres_credentials() - self.conn = psycopg2.connect( - dbname=dbname, - user=creds["user"], - password=creds["password"], - host=creds["host"], - port=creds["port"], - ) - - def get_connection(self): - """ - get a connection to database - """ - return self.conn - - def close_connection(self): - """ - will close the connection and destroy the class - """ - self.conn.close() - self.destroy_instance() - - @classmethod - def destroy_instance(cls): - cls._instance = None diff --git a/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py b/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py index 4df0315e..e1137b8a 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py @@ -2,7 +2,7 @@ from unittest import TestCase from hivemind_etl_helpers.src.db.discord.discord_summary import DiscordSummary -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams from hivemind_etl_helpers.src.utils.mongo import MongoSingleton from llama_index import Document, MockEmbedding, ServiceContext from llama_index.llms import MockLLM diff --git a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_db_fetch_files_modified_at.py b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_db_fetch_files_modified_at.py index 48b75393..24768ceb 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_db_fetch_files_modified_at.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_db_fetch_files_modified_at.py @@ -3,7 +3,7 @@ import psycopg2 from hivemind_etl_helpers.src.db.gdrive.db_utils import fetch_files_date_field, setup_db -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index import Document diff --git a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py index c47cadc4..1371ba9e 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py @@ -4,7 +4,7 @@ import psycopg2 from hivemind_etl_helpers.src.db.gdrive.db_utils import setup_db from hivemind_etl_helpers.src.db.gdrive.delete_records import delete_records -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials from hivemind_etl_helpers.src.utils.pg_db_utils import convert_tuple_str from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index import Document diff --git a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_delete.py b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_delete.py index 4609f97f..0e3f24d9 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_delete.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_delete.py @@ -4,7 +4,7 @@ import psycopg2 from hivemind_etl_helpers.src.db.gdrive.db_utils import setup_db from hivemind_etl_helpers.src.utils.check_documents import check_documents -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index import Document diff --git a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_insert.py b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_insert.py index 127c3b78..3da16091 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_insert.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_documents_to_insert.py @@ -4,7 +4,7 @@ import psycopg2 from hivemind_etl_helpers.src.db.gdrive.db_utils import setup_db from hivemind_etl_helpers.src.utils.check_documents import check_documents -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index import Document diff --git a/dags/hivemind_etl_helpers/tests/integration/test_load_envs.py b/dags/hivemind_etl_helpers/tests/integration/test_load_envs.py index c5f47627..33c7833d 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_load_envs.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_load_envs.py @@ -1,9 +1,7 @@ import unittest -from hivemind_etl_helpers.src.utils.credentials import ( - load_mongo_credentials, - load_postgres_credentials, -) +from hivemind_etl_helpers.src.utils.credentials import load_mongo_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials class TestCredentialLoadings(unittest.TestCase): diff --git a/dags/hivemind_etl_helpers/tests/integration/test_pg_vector_access_with_discord.py b/dags/hivemind_etl_helpers/tests/integration/test_pg_vector_access_with_discord.py index 884ff253..965fc714 100644 --- a/dags/hivemind_etl_helpers/tests/integration/test_pg_vector_access_with_discord.py +++ b/dags/hivemind_etl_helpers/tests/integration/test_pg_vector_access_with_discord.py @@ -5,9 +5,9 @@ from hivemind_etl_helpers.src.db.discord.discord_raw_message_to_document import ( discord_raw_to_docuemnts, ) -from hivemind_etl_helpers.src.utils.credentials import load_postgres_credentials +from tc_hivemind_backend.db.credentials import load_postgres_credentials from hivemind_etl_helpers.src.utils.mongo import MongoSingleton -from hivemind_etl_helpers.src.utils.pg_db_utils import setup_db +from tc_hivemind_backend.db.pg_db_utils import setup_db from hivemind_etl_helpers.src.utils.pg_vector_access import PGVectorAccess from llama_index.indices.vector_store import VectorStoreIndex diff --git a/dags/hivemind_etl_helpers/tests/unit/test_load_model_params.py b/dags/hivemind_etl_helpers/tests/unit/test_load_model_params.py index 7fb0544a..7ef10efa 100644 --- a/dags/hivemind_etl_helpers/tests/unit/test_load_model_params.py +++ b/dags/hivemind_etl_helpers/tests/unit/test_load_model_params.py @@ -1,7 +1,7 @@ import os import unittest -from hivemind_etl_helpers.src.utils.load_llm_params import load_model_hyperparams +from tc_hivemind_backend.db.utils.model_hyperparams import load_model_hyperparams class TestLoadModelHyperparams(unittest.TestCase): diff --git a/requirements.txt b/requirements.txt index 88b3adfc..712eeeb2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ cohere neo4j>=5.14.1, <6.0.0 coverage>=7.3.3, <8.0.0 pytest>=7.4.3, <8.0.0 -python-dotenv==1.0.0 +python-dotenv>=1.0.0, <2.0.0 +tc_hivemind_backend==1.0.0