diff --git a/celery_app/server.py b/celery_app/server.py index e9e2743..499aff0 100644 --- a/celery_app/server.py +++ b/celery_app/server.py @@ -1,12 +1,11 @@ from celery import Celery - from utils.credentials import load_rabbitmq_credentials rabbit_creds = load_rabbitmq_credentials() -user = rabbit_creds['user'] -password = rabbit_creds['password'] -host = rabbit_creds['host'] -port = rabbit_creds['port'] +user = rabbit_creds["user"] +password = rabbit_creds["password"] +host = rabbit_creds["host"] +port = rabbit_creds["port"] app = Celery("celery_app/tasks", broker=f"pyamqp://{user}:{password}@{host}:{port}//") app.autodiscover_tasks(["celery_app"]) diff --git a/celery_app/tasks.py b/celery_app/tasks.py index b4fea92..ccd4c5e 100644 --- a/celery_app/tasks.py +++ b/celery_app/tasks.py @@ -2,16 +2,17 @@ from celery_app.server import app from utils.credentials import load_rabbitmq_credentials + # TODO: Write tasks that match our requirements @app.task def add(x, y): rabbit_creds = load_rabbitmq_credentials() - username = rabbit_creds['user'] - password = rabbit_creds['password'] - broker_url = rabbit_creds['host'] - port = rabbit_creds['port'] + username = rabbit_creds["user"] + password = rabbit_creds["password"] + broker_url = rabbit_creds["host"] + port = rabbit_creds["port"] res = x + y job_send(broker_url, port, username, password, res) diff --git a/discord_query.py b/discord_query.py index c0d9b37..ae0865a 100644 --- a/discord_query.py +++ b/discord_query.py @@ -1,8 +1,8 @@ -from llama_index import QueryBundle -from llama_index.vector_stores import ExactMatchFilter, FilterCondition, MetadataFilters from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever from bot.retrievers.process_dates import process_dates from bot.retrievers.utils.load_hyperparams import load_hyperparams +from llama_index import QueryBundle +from llama_index.vector_stores import ExactMatchFilter, FilterCondition, MetadataFilters from tc_hivemind_backend.embeddings.cohere import CohereEmbedding from tc_hivemind_backend.pg_vector_access import PGVectorAccess diff --git a/tests/unit/test_discord_summary_retriever.py b/tests/unit/test_discord_summary_retriever.py index 92448c4..d5fafa3 100644 --- a/tests/unit/test_discord_summary_retriever.py +++ b/tests/unit/test_discord_summary_retriever.py @@ -3,9 +3,9 @@ from unittest import TestCase from unittest.mock import MagicMock +from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever from dateutil import parser from llama_index import Document, MockEmbedding, ServiceContext, VectorStoreIndex -from bot.retrievers.forum_summary_retriever import ForumBasedSummaryRetriever class TestDiscordSummaryRetriever(TestCase): diff --git a/tests/unit/test_load_credentials.py b/tests/unit/test_load_credentials.py index 701b5e6..5bfa795 100644 --- a/tests/unit/test_load_credentials.py +++ b/tests/unit/test_load_credentials.py @@ -38,4 +38,4 @@ def test_rabbitmq_envs_values(self): self.assertIsInstance(rabbitmq_creds["user"], str) self.assertIsInstance(rabbitmq_creds["password"], str) self.assertIsInstance(rabbitmq_creds["host"], str) - self.assertIsInstance(rabbitmq_creds["port"], str) \ No newline at end of file + self.assertIsInstance(rabbitmq_creds["port"], str) diff --git a/tests/unit/test_summary_retriever_base.py b/tests/unit/test_summary_retriever_base.py index 3f027db..14180ac 100644 --- a/tests/unit/test_summary_retriever_base.py +++ b/tests/unit/test_summary_retriever_base.py @@ -2,9 +2,9 @@ from unittest import TestCase from unittest.mock import MagicMock +from bot.retrievers.summary_retriever_base import BaseSummarySearch from llama_index import Document, MockEmbedding, ServiceContext, VectorStoreIndex from llama_index.schema import NodeWithScore -from bot.retrievers.summary_retriever_base import BaseSummarySearch class TestSummaryRetrieverBase(TestCase): diff --git a/utils/credentials.py b/utils/credentials.py index 02b3fb7..234c82c 100644 --- a/utils/credentials.py +++ b/utils/credentials.py @@ -52,4 +52,4 @@ def load_rabbitmq_credentials() -> dict[str, str]: rabbitmq_creds["host"] = os.getenv("RABBIT_HOST", "") rabbitmq_creds["port"] = os.getenv("RABBIT_PORT", "") - return rabbitmq_creds \ No newline at end of file + return rabbitmq_creds diff --git a/worker.py b/worker.py index 1104bbd..48ddb05 100644 --- a/worker.py +++ b/worker.py @@ -2,7 +2,6 @@ from tc_messageBroker import RabbitMQ from tc_messageBroker.rabbit_mq.event import Event from tc_messageBroker.rabbit_mq.queue import Queue - from utils.credentials import load_rabbitmq_credentials @@ -32,9 +31,9 @@ def job_recieve(broker_url, port, username, password): if __name__ == "__main__": rabbit_creds = load_rabbitmq_credentials() - username = rabbit_creds['user'] - password = rabbit_creds['password'] - broker_url = rabbit_creds['host'] - port = rabbit_creds['port'] + username = rabbit_creds["user"] + password = rabbit_creds["password"] + broker_url = rabbit_creds["host"] + port = rabbit_creds["port"] job_recieve(broker_url, port, username, password)