from llama_index import (
    ServiceContext,
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
)
from llama_index.vector_stores.qdrant import QdrantVectorStore
from tqdm import tqdm
import arxiv
import os
import argparse
import yaml
import qdrant_client
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index.embeddings import LangchainEmbedding

from llama_index import ServiceContext
from llama_index.llms import Ollama


class Data:
    def __init__(self, config):
        self.config = config

    def _create_data_folder(self, download_path):
        data_path = download_path
        if not os.path.exists(data_path):
            os.makedirs(self.config["data_path"])
            print("Output folder created")
        else:
            print("Output folder already exists.")

    def download_papers(self, search_query, download_path, max_results):
        self._create_data_folder(download_path)
        client = arxiv.Client()

        search = arxiv.Search(
            query=search_query,
            max_results=max_results,
            sort_by=arxiv.SortCriterion.SubmittedDate,
        )

        results = list(client.results(search))
        for paper in tqdm(results):
            if os.path.exists(download_path):
                paper_title = (paper.title).replace(" ", "_")
                paper.download_pdf(dirpath=download_path, filename=f"{paper_title}.pdf")
                print(f"{paper.title} Downloaded.")

    def ingest(self, embedder, llm):
        print("Indexing data...")
        documents = SimpleDirectoryReader(self.config["data_path"]).load_data()

        client = qdrant_client.QdrantClient(url=self.config["qdrant_url"])
        qdrant_vector_store = QdrantVectorStore(
            client=client, collection_name=self.config["collection_name"]
        )
        storage_context = StorageContext.from_defaults(vector_store=qdrant_vector_store)
        # service_context = ServiceContext.from_defaults(
        #     llm=llm, embed_model=embedder, chunk_size=self.config["chunk_size"]
        # )
        service_context = ServiceContext.from_defaults(
            llm=None, embed_model=embedder, chunk_size=self.config["chunk_size"]
        )

        index = VectorStoreIndex.from_documents(
            documents, storage_context=storage_context, service_context=service_context
        )
        print(
            f"Data indexed successfully to Qdrant. Collection: {self.config['collection_name']}"
        )
        return index


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-q", "--query",
        type=str,
        default=False,
        help="Download papers from arxiv with this query.",
    )
    # parser.add_argument(
    #     "-o", "--output", type=str, default=False, help="Download path."
    # )

    parser.add_argument(
        "-m", "--max", type=int, default=False, help="Max results to download."
    )

    parser.add_argument(
        "-i",
        "--ingest",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Ingest data to Qdrant vector Database.",
    )

    args = parser.parse_args()
    config_file = "config.yml"
    with open(config_file, "r") as conf:
        config = yaml.safe_load(conf)
    data = Data(config)
    if args.query:
        data.download_papers(
            search_query=args.query,
            download_path=config["data_path"],
            max_results=args.max,
        )
    if args.ingest:
        print("Loading Embedder...")
        embed_model = LangchainEmbedding(
            HuggingFaceEmbeddings(model_name=config["embedding_model"])
        )
        llm = Ollama(model=config["llm_name"], base_url=config["llm_url"])
        data.ingest(embedder=embed_model, llm=llm)

langchain.embeddings.huggingface.HuggingFaceEmbeddings
[((2566, 2591), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (2589, 2591), False, 'import argparse\n'), ((970, 984), 'arxiv.Client', 'arxiv.Client', ([], {}), '()\n', (982, 984), False, 'import arxiv\n'), ((1003, 1108), 'arxiv.Search', 'arxiv.Search', ([], {'query': 'search_query', 'max_results': 'max_results', 'sort_by': 'arxiv.SortCriterion.SubmittedDate'}), '(query=search_query, max_results=max_results, sort_by=arxiv.\n    SortCriterion.SubmittedDate)\n', (1015, 1108), False, 'import arxiv\n'), ((1220, 1233), 'tqdm.tqdm', 'tqdm', (['results'], {}), '(results)\n', (1224, 1233), False, 'from tqdm import tqdm\n'), ((1654, 1711), 'qdrant_client.QdrantClient', 'qdrant_client.QdrantClient', ([], {'url': "self.config['qdrant_url']"}), "(url=self.config['qdrant_url'])\n", (1680, 1711), False, 'import qdrant_client\n'), ((1742, 1827), 'llama_index.vector_stores.qdrant.QdrantVectorStore', 'QdrantVectorStore', ([], {'client': 'client', 'collection_name': "self.config['collection_name']"}), "(client=client, collection_name=self.config['collection_name']\n    )\n", (1759, 1827), False, 'from llama_index.vector_stores.qdrant import QdrantVectorStore\n'), ((1871, 1933), 'llama_index.StorageContext.from_defaults', 'StorageContext.from_defaults', ([], {'vector_store': 'qdrant_vector_store'}), '(vector_store=qdrant_vector_store)\n', (1899, 1933), False, 'from llama_index import ServiceContext, SimpleDirectoryReader, StorageContext, VectorStoreIndex\n'), ((2112, 2215), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm': 'None', 'embed_model': 'embedder', 'chunk_size': "self.config['chunk_size']"}), "(llm=None, embed_model=embedder, chunk_size=\n    self.config['chunk_size'])\n", (2140, 2215), False, 'from llama_index import ServiceContext\n'), ((2250, 2362), 'llama_index.VectorStoreIndex.from_documents', 'VectorStoreIndex.from_documents', (['documents'], {'storage_context': 'storage_context', 'service_context': 'service_context'}), '(documents, storage_context=storage_context,\n    service_context=service_context)\n', (2281, 2362), False, 'from llama_index import ServiceContext, SimpleDirectoryReader, StorageContext, VectorStoreIndex\n'), ((3283, 3303), 'yaml.safe_load', 'yaml.safe_load', (['conf'], {}), '(conf)\n', (3297, 3303), False, 'import yaml\n'), ((3700, 3760), 'llama_index.llms.Ollama', 'Ollama', ([], {'model': "config['llm_name']", 'base_url': "config['llm_url']"}), "(model=config['llm_name'], base_url=config['llm_url'])\n", (3706, 3760), False, 'from llama_index.llms import Ollama\n'), ((646, 671), 'os.path.exists', 'os.path.exists', (['data_path'], {}), '(data_path)\n', (660, 671), False, 'import os\n'), ((685, 722), 'os.makedirs', 'os.makedirs', (["self.config['data_path']"], {}), "(self.config['data_path'])\n", (696, 722), False, 'import os\n'), ((1250, 1279), 'os.path.exists', 'os.path.exists', (['download_path'], {}), '(download_path)\n', (1264, 1279), False, 'import os\n'), ((3616, 3675), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': "config['embedding_model']"}), "(model_name=config['embedding_model'])\n", (3637, 3675), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((1576, 1623), 'llama_index.SimpleDirectoryReader', 'SimpleDirectoryReader', (["self.config['data_path']"], {}), "(self.config['data_path'])\n", (1597, 1623), False, 'from llama_index import ServiceContext, SimpleDirectoryReader, StorageContext, VectorStoreIndex\n')]