feat: set the config.py

J-coder118 · J-coder118 · commit cd2d1c9e6458 · 2024-08-13T21:41:54.000-03:00
diff --git a/3-feature-pipeline/config.py b/3-feature-pipeline/config.py
@@ -0,0 +1,35 @@
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    # CometML config
+    COMET_API_KEY: str | None = None
+    COMET_WORKSPACE: str | None = None
+    COMET_PROJECT: str | None = None
+
+    # Embeddings config
+    EMBEDDING_MODEL_ID: str = "sentence-transformers/all-MiniLM-L6-v2"
+    EMBEDDING_MODEL_MAX_INPUT_LENGTH: int = 256
+    EMBEDDING_SIZE: int = 384
+    EMBEDDING_MODEL_DEVICE: str = "cpu"
+
+    # OpenAI
+    OPENAI_MODEL_ID: str = "gpt-4-1106-preview"
+    OPENAI_API_KEY: str | None = None
+
+    # MQ config
+    RABBITMQ_DEFAULT_USERNAME: str = "guest"
+    RABBITMQ_DEFAULT_PASSWORD: str = "guest"
+    RABBITMQ_HOST: str = "mq" # or localhost if running outside Docker
+    RABBITMQ_PORT: int = 5672
+    RABBITMQ_QUEUE_NAME: str = "default"
+
+    # QdrantDB config
+    QDRANT_DATABASE_HOST: str = "qdrant" # or localhost if running outside Docker
+    QDRANT_DATABASE_PORT: int = 6333
+    USE_QDRANT_CLOUD: bool = False # if True, fill in QDRANT_CLOUD_URL and QDRANT_APIKEY
+    QDRANT_CLOUD_URL: str | None = None
+    QDRANT_APIKEY: str | None = None
+
+
+settings = Settings()
diff --git a/3-feature-pipeline/db.py b/3-feature-pipeline/db.py
@@ -0,0 +1,82 @@
+from qdrant_client import QdrantClient, models
+from qdrant_client.http.exceptions import UnexpectedResponse
+from qdrant_client.http.models import Batch, Distance, VectorParams
+
+from utils.logging import get_logger
+from config import settings
+
+logger = get_logger(__name__)
+
+
+class QdrantDatabaseConnector:
+    _instance: QdrantClient | None = None
+
+    def __init__(self) -> None:
+        if self._instance is None:
+            try:
+                if settings.USE_QDRANT_CLOUD:
+                    self._instance = QdrantClient(
+                        url=settings.QDRANT_CLOUD_URL,
+                        api_key=settings.QDRANT_APIKEY,
+                    )
+                else:
+                    self._instance = QdrantClient(
+                        host=settings.QDRANT_DATABASE_HOST,
+                        port=settings.QDRANT_DATABASE_PORT,
+                    )
+            except UnexpectedResponse:
+                logger.exception(
+                    "Couldn't connect to Qdrant.",
+                    host=settings.QDRANT_DATABASE_HOST,
+                    port=settings.QDRANT_DATABASE_PORT,
+                    url=settings.QDRANT_CLOUD_URL,
+                )
+
+                raise
+
+    def get_collection(self, collection_name: str):
+        return self._instance.get_collection(collection_name=collection_name)
+
+    def create_non_vector_collection(self, collection_name: str):
+        self._instance.create_collection(
+            collection_name=collection_name, vectors_config={}
+        )
+
+    def create_vector_collection(self, collection_name: str):
+        self._instance.create_collection(
+            collection_name=collection_name,
+            vectors_config=VectorParams(
+                size=settings.EMBEDDING_SIZE, distance=Distance.COSINE
+            ),
+        )
+
+    def write_data(self, collection_name: str, points: Batch):
+        try:
+            self._instance.upsert(collection_name=collection_name, points=points)
+        except Exception:
+            logger.exception("An error occurred while inserting data.")
+
+            raise
+
+    def search(
+        self,
+        collection_name: str,
+        query_vector: list,
+        query_filter: models.Filter | None = None,
+        limit: int = 3,
+    ) -> list:
+        return self._instance.search(
+            collection_name=collection_name,
+            query_vector=query_vector,
+            query_filter=query_filter,
+            limit=limit,
+        )
+
+    def scroll(self, collection_name: str, limit: int):
+        return self._instance.scroll(collection_name=collection_name, limit=limit)
+
+    def close(self):
+        if self._instance:
+            self._instance.close()
+
+            logger.info("Connected to database has been closed.")
diff --git a/3-feature-pipeline/main.py b/3-feature-pipeline/main.py
@@ -0,0 +1,34 @@
+import bytewax.operators as op
+from bytewax.dataflow import Dataflow
+
+from db import QdrantDatabaseConnector
+
+from data_flow.stream_input import RabbitMQSource
+from data_flow.stream_output import QdrantOutput
+from data_logic.dispatchers import (
+    ChunkingDispatcher,
+    CleaningDispatcher,
+    EmbeddingDispatcher,
+    RawDispatcher,
+)
+
+connection = QdrantDatabaseConnector()
+
+flow = Dataflow("Streaming ingestion pipeline")
+stream = op.input("input", flow, RabbitMQSource())
+stream = op.map("raw dispatch", stream, RawDispatcher.handle_mq_message)
+stream = op.map("clean dispatch", stream, CleaningDispatcher.dispatch_cleaner)
+op.output(
+    "cleaned data insert to qdrant",
+    stream,
+    QdrantOutput(connection=connection, sink_type="clean"),
+)
+stream = op.flat_map("chunk dispatch", stream, ChunkingDispatcher.dispatch_chunker)
+stream = op.map(
+    "embedded chunk dispatch", stream, EmbeddingDispatcher.dispatch_embedder
+)
+op.output(
+    "embedded data insert to qdrant",
+    stream,
+    QdrantOutput(connection=connection, sink_type="vector"),
+)
diff --git a/3-feature-pipeline/mq.py b/3-feature-pipeline/mq.py
@@ -0,0 +1,98 @@
+import pika
+
+from utils.logging import get_logger
+from config import settings
+
+logger = get_logger(__name__)
+
+
+class RabbitMQConnection:
+    _instance = None
+
+    def __new__(
+        cls,
+        host: str | None = None,
+        port: int | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        virtual_host: str = "/",
+    ):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+
+        return cls._instance
+
+    def __init__(
+        self,
+        host: str | None = None,
+        port: int | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        virtual_host: str = "/",
+        fail_silently: bool = False,
+        **kwargs,
+    ):
+        self.host = host or settings.RABBITMQ_HOST
+        self.port = port or settings.RABBITMQ_PORT
+        self.username = username or settings.RABBITMQ_DEFAULT_USERNAME
+        self.password = password or settings.RABBITMQ_DEFAULT_PASSWORD
+        self.virtual_host = virtual_host
+        self.fail_silently = fail_silently
+        self._connection = None
+
+    def __enter__(self):
+        self.connect()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def connect(self):
+        try:
+            credentials = pika.PlainCredentials(self.username, self.password)
+            self._connection = pika.BlockingConnection(
+                pika.ConnectionParameters(
+                    host=self.host,
+                    port=self.port,
+                    virtual_host=self.virtual_host,
+                    credentials=credentials,
+                )
+            )
+        except pika.exceptions.AMQPConnectionError as e:
+            logger.exception("Failed to connect to RabbitMQ.")
+
+            if not self.fail_silently:
+                raise e
+
+    def publish_message(self, data: str, queue: str):
+        channel = self.get_channel()
+        channel.queue_declare(
+            queue=queue, durable=True, exclusive=False, auto_delete=False
+        )
+        channel.confirm_delivery()
+
+        try:
+            channel.basic_publish(
+                exchange="", routing_key=queue, body=data, mandatory=True
+            )
+            logger.info(
+                "Sent message successfully.", queue_type="RabbitMQ", queue_name=queue
+            )
+        except pika.exceptions.UnroutableError:
+            logger.info(
+                "Failed to send the message.", queue_type="RabbitMQ", queue_name=queue
+            )
+
+    def is_connected(self) -> bool:
+        return self._connection is not None and self._connection.is_open
+
+    def get_channel(self):
+        if self.is_connected():
+            return self._connection.channel()
+
+    def close(self):
+        if self.is_connected():
+            self._connection.close()
+            self._connection = None
+
+            logger.info("Closed RabbitMQ connection.")
diff --git a/3-feature-pipeline/retriever.py b/3-feature-pipeline/retriever.py
@@ -0,0 +1,22 @@
+from dotenv import load_dotenv
+from langchain.globals import set_verbose
+from rag.retriever import VectorRetriever
+
+from utils.logging import get_logger
+
+set_verbose(True)
+
+logger = get_logger(__name__)
+
+if __name__ == "__main__":
+    load_dotenv()
+    query = """
+        Could you please draft a LinkedIn post discussing RAG systems?
+        I'm particularly interested in how RAG works and how it is integrated with vector DBs and large language models (LLMs).
+        """
+    retriever = VectorRetriever(query=query)
+    hits = retriever.retrieve_top_k(k=6, to_expand_to_n_queries=5)
+
+    reranked_hits = retriever.rerank(hits=hits, keep_top_k=5)
+    for rank, hit in enumerate(reranked_hits):
+        logger.info(f"{rank}: {hit}")
