Cognita: Building an Open Source, Modular, RAG applications for Production

register_dataloader("MyDataLoader", MyDataLoaderClass)

register_parser("MyParser", MyParserClass)

register_embedder("MyEmbedder", MyEmbedderClass)

from langchain_community.vectorstores.qdrant import Qdrant
from qdrant_client.http.models import VectorParams, Distance
from qdrant_client import QdrantClient, models

class QdrantVectorDB(BaseVectorDB):
    def __init__(self, config: VectorDBConfig):
        ...
        # Initialize the qdrant client
        self.qdrant_client = QdrantClient(
            url=self.url,
            port=self.port,
            prefer_grpc=self.prefer_grpc,
            prefix=self.prefix,
        )

 def create_collection(self, collection_name: str, embeddings: Embeddings):
    # Calculate embedding size
    partial_embeddings = embeddings.embed_documents(["Initial document"])
    vector_size = len(partial_embeddings[0])

    # Create collection given the embeding dimension and simalrity metric
    self.qdrant_client.create_collection(
        collection_name=collection_name,
        vectors_config=VectorParams(
            size=vector_size,  # embedding dimension
            distance=Distance.COSINE,
        ),
        replication_factor=3,
    )

    # Create metadata index for the collection
    self.qdrant_client.create_payload_index(
        collection_name=collection_name,
        field_name=f"metadata.{DATA_POINT_FQN_METADATA_KEY}",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )

def upsert_documents(
    self,
    collection_name: str,
    documents,
    embeddings: Embeddings,
    incremental: bool = True,
):
    # Get the data point fqns from the documents
    # fqns uniquely identify the documents
    # This is used to delete the outdated documents from the db
    data_point_fqns = []
    for document in documents:
        if document.metadata.get(DATA_POINT_FQN_METADATA_KEY):
            data_point_fqns.append(
                document.metadata.get(DATA_POINT_FQN_METADATA_KEY)
            )

    # Add Documents
    Qdrant(
        client=self.qdrant_client,
        collection_name=collection_name,
        embeddings=embeddings,
    ).add_documents(documents=documents)

    # Get the record ids that are already present in the db (Existing documents that are modified)
    # These docs should be removed from the db
    # Used when incremental indexing is enabled
    record_ids_to_be_upserted: List[str] = self._get_records_to_be_upserted(
        collection_name=collection_name,
        data_point_fqns=data_point_fqns,
        incremental=incremental,
    )

    # Delete Documents
    if len(record_ids_to_be_upserted):
        for i in range(0, len(record_ids_to_be_upserted), BATCH_SIZE):
            record_ids_to_be_processed = record_ids_to_be_upserted[
                i : i + BATCH_SIZE
            ]
            self.qdrant_client.delete(
                collection_name=collection_name,
                points_selector=models.PointIdsList(
                    points=record_ids_to_be_processed,
                ),
            )

# This method is used to get the records that are already present in the db
def _get_records_to_be_upserted(
    self, collection_name: str,
    data_point_fqns: List[str],
    incremental: bool
):
    if not incremental:
        return []
    # For incremental deletion, we delete the documents with the same document_id
    stop = False
    offset = None
    record_ids_to_be_upserted = []

    # we fetch the records in batches based on the data_point_fqns
    while stop is not True:
        records, next_offset = self.qdrant_client.scroll(
            collection_name=collection_name,
            scroll_filter=models.Filter(
                should=[
                    models.FieldCondition(
                        key=f"metadata.{DATA_POINT_FQN_METADATA_KEY}",
                        match=models.MatchAny(
                            any=data_point_fqns,
                        ),
                    ),
                ]
            ),
            limit=BATCH_SIZE,
            offset=offset,
            with_payload=False,
            with_vectors=False,
        )
        for record in records:
            record_ids_to_be_upserted.append(record.id)
            if len(record_ids_to_be_upserted) > MAX_SCROLL_LIMIT:
                stop = True
                break
        if next_offset is None:
            stop = True
        else:
            offset = next_offset
    return record_ids_to_be_upserted

def get_collections(self):
    collections = self.qdrant_client.list_collections()
    return [collection.name for collection in collections]

def delete_collection(self, collection_name: str):
    self.qdrant_client.delete_collection(collection_name=collection_name)

def get_vector_store(self, collection_name: str, embeddings: Embeddings):
    return Qdrant(
            client=self.qdrant_client,
            embeddings=embeddings,
            collection_name=collection_name,
        )

def get_vector_client(self):
    return self.qdrant_client

def list_data_point_vectors(
    self, collection_name: str,
    data_source_fqn: str,
    batch_size: int
) -> List[DataPointVector]:

    stop = False
    offset = None
    data_point_vectors: List[DataPointVector] = []

    # Get the vectors that match with the doc unique identifier
    # Data point is an object that has vector_id, fqn and hash
    # hash is calculated during parsing and chunking phase
    while stop is not True:
        records, next_offset = self.qdrant_client.scroll(
            collection_name=collection_name,
            limit=batch_size,
            with_payload=[
                f"metadata.{DATA_POINT_FQN_METADATA_KEY}",
                f"metadata.{DATA_POINT_HASH_METADATA_KEY}",
            ],
            # filter the records based on the data_point_fqn
            scroll_filter=models.Filter(
                should=[
                    models.FieldCondition(
                        key=f"metadata.{DATA_POINT_FQN_METADATA_KEY}",
                        match=models.MatchText(
                            text=data_source_fqn,
                        ),
                    ),
                ]
            ),
            with_vectors=False,
            offset=offset,
        )
        for record in records:
            metadata: dict = record.payload.get("metadata")
            if (
                metadata
                and metadata.get(DATA_POINT_FQN_METADATA_KEY)
                and metadata.get(DATA_POINT_HASH_METADATA_KEY)
            ):
                # Create a DataPointVector object
                data_point_vectors.append(
                    DataPointVector(
                        data_point_vector_id=record.id,
                        data_point_fqn=metadata.get(DATA_POINT_FQN_METADATA_KEY),
                        data_point_hash=metadata.get(DATA_POINT_HASH_METADATA_KEY),
                    )
                )
            if len(data_point_vectors) > MAX_SCROLL_LIMIT:
                stop = True
                break
        if next_offset is None:
            stop = True
        else:
            offset = next_offset
    return data_point_vectors

def delete_data_point_vectors(
    self,
    collection_name: str,
    data_point_vectors: List[DataPointVector],
    batch_size: int = BATCH_SIZE,
):

    vectors_to_be_deleted_count = len(data_point_vectors)
    deleted_vectors_count = 0

    # delete the vectors in batches
    for i in range(0, vectors_to_be_deleted_count, batch_size):
        data_point_vectors_to_be_processed = data_point_vectors[i : i + batch_size]
        # Delete the vectors from the db
        self.qdrant_client.delete(
            collection_name=collection_name,
            points_selector=models.PointIdsList(
                points=[
                    document_vector_point.data_point_vector_id
                    for document_vector_point in data_point_vectors_to_be_processed
                ],
            ),
        )
        # Increment the deleted vectors count
        deleted_vectors_count = deleted_vectors_count + len(data_point_vectors_to_be_processed)

from langchain_community.vectorstores.singlestoredb import SingleStoreDB
import singlestoredb as s2

class SSDB(SingleStoreDB):
    def _create_table(self: SingleStoreDB) -> None:
        """Create table if it doesn't exist."""
        conn = self.connection_pool.connect()
        try:
            cur = conn.cursor()
            # Overriding the default table creation behaviour this adds id as autoinc primary key
            try:
                if self.use_vector_index:
                    index_options = ""
                    if self.vector_index_options and len(self.vector_index_options) > 0:
                        index_options = "INDEX_OPTIONS '{}'".format(
                            json.dumps(self.vector_index_options)
                        )
                    cur.execute(
                        """CREATE TABLE IF NOT EXISTS {}
                        (id BIGINT AUTO_INCREMENT PRIMARY KEY, {} TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci,
                        {} VECTOR({}, F32) NOT NULL, {} JSON,
                        VECTOR INDEX {} ({}) {});""".format(
                            self.table_name,
                            self.content_field,
                            self.vector_field,
                            self.vector_size,
                            self.metadata_field,
                            self.vector_index_name,
                            self.vector_field,
                            index_options,
                        ),
                    )
                else:
                    cur.execute(
                        """CREATE TABLE IF NOT EXISTS {}
                        (id BIGINT AUTO_INCREMENT PRIMARY KEY, {} TEXT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci,
                        {} BLOB, {} JSON);""".format(
                            self.table_name,
                            self.content_field,
                            self.vector_field,
                            self.metadata_field,
                        ),
                    )
            finally:
                cur.close()
        finally:
            conn.close()


    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        embeddings: Optional[List[List[float]]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Add more texts to the vectorstore.

        Args:
            texts (Iterable[str]): Iterable of strings/text to add to the vectorstore.
            metadatas (Optional[List[dict]], optional): Optional list of metadatas.
                Defaults to None.
            embeddings (Optional[List[List[float]]], optional): Optional pre-generated
                embeddings. Defaults to None.

        Returns:
            List[str]: empty list
        """
        conn = self.connection_pool.connect()
        try:
            cur = conn.cursor()
            try:
                # Write data to singlestore db
                for i, text in enumerate(texts):
                    # Use provided values by default or fallback
                    metadata = metadatas[i] if metadatas else {}
                    embedding = (
                        embeddings[i]
                        if embeddings
                        else self.embedding.embed_documents([text])[0]
                    )
                    # Overriding insert statement to handle autoincrement id
                    cur.execute(
                        "INSERT INTO {} (content, vector, metadata) VALUES (%s, JSON_ARRAY_PACK(%s), %s)".format(
                            self.table_name
                        ),
                        (
                            text,
                            "[{}]".format(",".join(map(str, embedding))),
                            json.dumps(metadata),
                        ),
                    )
                if self.use_vector_index:
                    cur.execute("OPTIMIZE TABLE {} FLUSH;".format(self.table_name))
            finally:
                cur.close()
        finally:
            conn.close()
        return []

class SingleStoreVectorDB(BaseVectorDB):
    def __init__(self, config: VectorDBConfig):
        # url: mysql://{user}:{password}@{host}:{port}/{db}
        self.host = config.url

def create_collection(self, collection_name: str, embeddings: Embeddings):
    # Calculate embedding size
    partial_embeddings = embeddings.embed_documents(["Initial document"])
    vector_size = len(partial_embeddings[0])

    # Create collection
    # we keep vector_filed, content_field as default values
    SSDB(
        embedding=embeddings,
        host=self.host,
        table_name=collection_name,
        vector_size=vector_size,
        use_vector_index=True,
    )

def upsert_documents(
    self,
    collection_name: str,
    documents: List[Document],
    embeddings: Embeddings,
    incremental: bool = True,
):
    # Get the data point fqns from the documents
    data_point_fqns = []
    for document in documents:
        if document.metadata.get(DATA_POINT_FQN_METADATA_KEY):
            data_point_fqns.append(
                document.metadata.get(DATA_POINT_FQN_METADATA_KEY)
            )

    try:
        # Add documents to the table
        SSDB.from_documents(
            embedding=embeddings,
            documents=documents,
            table_name=collection_name,
            host=self.host,
        )
        logger.debug(
            f"[SingleStore] Added {len(documents)} documents to collection {collection_name}"
        )
    except Exception as e:
        logger.error(
            f"[SingleStore] Failed to add documents to collection {collection_name}: {e}"
        )

def get_collections(self) -> List[str]:
    # Create connection to SingleStore
    conn = s2.connect(self.host)
    try:
        cur = conn.cursor()
        try:
            # Get all the tables in the db
            cur.execute("SHOW TABLES")
            return [row[0] for row in cur.fetchall()]
        except Exception as e:
            logger.error(
                f"[SingleStore] Failed to get collections: {e}"
            )
        finally:
            cur.close()
    except Exception as e:
        logger.error(
            f"[SingleStore] Failed to get collections: {e}"
        )
    finally:
        conn.close()

def delete_collection(self, collection_name: str):
    # Create connection to SingleStore
    conn = s2.connect(self.host)
    try:
        cur = conn.cursor()
        try:
            # Drop the table
            cur.execute(f"DROP TABLE {collection_name}")
            logger.debug(f"[SingleStore] Deleted collection {collection_name}")
        except Exception as e:
            logger.error(
                f"[SingleStore] Failed to delete collection {collection_name}: {e}"
            )
        finally:
            cur.close()
    except Exception as e:
        logger.error(
            f"[SingleStore] Failed to delete collection {collection_name}: {e}"
        )
    finally:
        conn.close()

def get_vector_store(self, collection_name: str, embeddings: Embeddings):
    return SSDB(
        embedding=embeddings,
        host=self.host,
        table_name=collection_name,
    )

def get_vector_client(self):
    return s2.connect(self.host)

def list_data_point_vectors(
    self,
    collection_name: str,
    data_source_fqn: str,
) -> List[DataPointVector]:
    data_point_vectors: List[DataPointVector] = []
    logger.debug(f"data_source_fqn: {data_source_fqn}")

    # Get all data point vectors from table upto MAX_SCROLL_LIMIT
    conn = s2.connect(self.host)
    try:
        curr = conn.cursor()

        # Remove all data point vectors with the same data_source_fqn
        curr.execute(f"SELECT * FROM {collection_name} WHERE JSON_EXTRACT_JSON(metadata, '{DATA_POINT_FQN_METADATA_KEY}') LIKE '%{data_source_fqn}%' LIMIT {MAX_SCROLL_LIMIT}")

        for record in curr:
            # id, content, vector, metadata
            id, _, _, metadata = record
            if (
                metadata
                and metadata.get(DATA_POINT_FQN_METADATA_KEY)
                and metadata.get(DATA_POINT_HASH_METADATA_KEY)
            ):
                # Create a DataPointVector object
                data_point_vectors.append(
                    DataPointVector(
                        data_point_vector_id=id,
                        data_point_fqn=metadata.get(DATA_POINT_FQN_METADATA_KEY),
                        data_point_hash=metadata.get(DATA_POINT_HASH_METADATA_KEY),
                    )
                )
    except Exception as e:
        logger.error(
            f"[SingleStore] Failed to list data point vectors: {e}"
        )
    finally:
        conn.close()

    logger.debug(
        f"[SingleStore] Listing {len(data_point_vectors)} data point vectors for collection {collection_name}"
    )
    return data_point_vectors

def delete_data_point_vectors(
    self,
    collection_name: str,
    data_point_vectors: List[DataPointVector],
    batch_size: int =BATCH_SIZE,
):
    if len(data_point_vectors) > 0:
        # Create connection to SingleStore
        conn = s2.connect(self.host)

        try:
            vectors_to_be_deleted_count = len(data_point_vectors)
            curr = conn.cursor()

            # Delete the vectors as per the data_point_vector_id
            curr.execute(
                f"DELETE FROM {collection_name} WHERE id in ({', '.join(data_point_vector.data_point_vector_id for data_point_vector in data_point_vectors)})"
            )
            logger.debug(
                f"[SingleStore] Deleted {vectors_to_be_deleted_count} data point vectors"
            )
        except Exception as e:
            logger.error(
                f"[SingleStore] Failed to delete data point vectors: {e}"
            )
        finally:
            conn.close()

from backend.server.decorator import post
@query_controller("/app-2")
class MyCustomController():
    ...
    @post("/answer")
    def answer(query: str):
        # Write code to express your logic for answer
        # This API will be exposed as POST /my-controller/answer
        ...

from backend.modules.query_controllers.sample_controller.controller import MyCustomController