rprav007/open_deep_research_agent

1import os2import chainlit as cl3from dotenv import load_dotenv4from operator import itemgetter5from langchain_huggingface import HuggingFaceEndpoint6from langchain_community.document_loaders import TextLoader7from langchain_text_splitters import RecursiveCharacterTextSplitter8from langchain_community.vectorstores import Qdrant9from langchain_huggingface import HuggingFaceEndpointEmbeddings10from langchain_core.prompts import PromptTemplate11from langchain.schema.output_parser import StrOutputParser12from langchain.schema.runnable import RunnablePassthrough13from langchain.schema.runnable.config import RunnableConfig14from tqdm.asyncio import tqdm_asyncio15import asyncio16from tqdm.asyncio import tqdm17from qdrant_client import QdrantClient18from qdrant_client.models import Distance, VectorParams, PointStruct19import json20import hashlib21from pathlib import Path22 23# GLOBAL SCOPE - ENTIRE APPLICATION HAS ACCESS TO VALUES SET IN THIS SCOPE #24# ---- ENV VARIABLES ---- # 25"""26This function will load our environment file (.env) if it is present.27 28NOTE: Make sure that .env is in your .gitignore file - it is by default, but please ensure it remains there.29"""30load_dotenv()31 32"""33We will load our environment variables here.34"""35HF_LLM_ENDPOINT = os.environ["HF_LLM_ENDPOINT"]36HF_EMBED_ENDPOINT = os.environ["HF_EMBED_ENDPOINT"]37HF_TOKEN = os.environ["HF_TOKEN"]38 39# ---- GLOBAL DECLARATIONS ---- #40 41# -- RETRIEVAL -- #42"""431. Load Documents from Text File442. Split Documents into Chunks453. Load HuggingFace Embeddings (remember to use the URL we set above)464. Index Files if they do not exist, otherwise load the vectorstore47"""48document_loader = TextLoader("./data/paul_graham_essays.txt")49documents = document_loader.load()50 51text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=30)52split_documents = text_splitter.split_documents(documents)53 54hf_embeddings = HuggingFaceEndpointEmbeddings(55    model=HF_EMBED_ENDPOINT,56    task="feature-extraction",57    huggingfacehub_api_token=HF_TOKEN,58)59 60# Add after the other global declarations61CACHE_DIR = Path("./cache")62CACHE_DIR.mkdir(exist_ok=True)63 64def get_cache_key(text):65    """Generate a cache key for a text string."""66    return hashlib.md5(text.encode()).hexdigest()67 68async def get_cached_embeddings(texts):69    """Get embeddings from cache or compute and cache them."""70    cache_file = CACHE_DIR / "embeddings_cache.json"71    cache = {}72    if cache_file.exists():73        with open(cache_file, "r") as f:74            cache = json.load(f)75    76    results = []77    texts_to_embed = []78    indices_to_embed = []79    80    for i, text in enumerate(texts):81        cache_key = get_cache_key(text)82        if cache_key in cache:83            results.append(cache[cache_key])84        else:85            texts_to_embed.append(text)86            indices_to_embed.append(i)87    88    if texts_to_embed:89        new_embeddings = await hf_embeddings.aembed_documents(texts_to_embed)90        for text, embedding, idx in zip(texts_to_embed, new_embeddings, indices_to_embed):91            cache_key = get_cache_key(text)92            cache[cache_key] = embedding93            results.insert(idx, embedding)94        95        # Save updated cache96        with open(cache_file, "w") as f:97            json.dump(cache, f)98    99    return results100 101async def process_batch(client, collection_name, batch, is_first_batch, pbar):102    texts = [doc.page_content for doc in batch]103    metadatas = [doc.metadata for doc in batch]104    105    # Use cached embeddings instead of direct computation106    embeddings = await get_cached_embeddings(texts)107    108    if is_first_batch:109        client.recreate_collection(110            collection_name=collection_name,111            vectors_config=VectorParams(size=len(embeddings[0]), distance=Distance.COSINE)112        )113    114    # Create points using PointStruct115    points = [116        PointStruct(117            id=i + pbar.n,118            vector=embedding,119            payload={120                "page_content": text,121                "metadata": metadata122            }123        )124        for i, (text, metadata, embedding) in enumerate(zip(texts, metadatas, embeddings))125    ]126    127    # Add points to the collection128    client.upsert(129        collection_name=collection_name,130        points=points131    )132    133    vectorstore = Qdrant(134        client=client,135        collection_name=collection_name,136        embeddings=hf_embeddings,137        content_payload_key="page_content",138        metadata_payload_key="metadata"139    )140    141    pbar.update(len(batch))142    return vectorstore143 144async def main():145    print("Indexing Files")146    147    collection_name = "paul_graham_essays"148    # Create Qdrant instance in memory149    client = QdrantClient(":memory:")150    vectorstore = None151    batch_size = 32152    153    batches = [split_documents[i:i+batch_size] for i in range(0, len(split_documents), batch_size)]154    155    async def process_all_batches():156        nonlocal vectorstore157        tasks = []158        pbars = []159        160        for i, batch in enumerate(batches):161            pbar = tqdm(total=len(batch), desc=f"Batch {i+1}/{len(batches)}", position=i)162            pbars.append(pbar)163            164            if i == 0:165                vectorstore = await process_batch(client, collection_name, batch, True, pbar)166            else:167                tasks.append(process_batch(client, collection_name, batch, False, pbar))168        169        if tasks:170            await asyncio.gather(*tasks)171        172        for pbar in pbars:173            pbar.close()174    175    await process_all_batches()176    177    hf_retriever = vectorstore.as_retriever()178    print("\nIndexing complete. Vectorstore is ready for use.")179    return hf_retriever180 181async def run():182    retriever = await main()183    return retriever184 185hf_retriever = asyncio.run(run())186 187# -- AUGMENTED -- #188"""1891. Define a String Template1902. Create a Prompt Template from the String Template191"""192RAG_PROMPT_TEMPLATE = """\193<|start_header_id|>system<|end_header_id|>194You are a helpful assistant. You answer user questions based on provided context. If you can't answer the question with the provided context, say you don't know.<|eot_id|>195 196<|start_header_id|>user<|end_header_id|>197User Query:198{query}199 200Context:201{context}<|eot_id|>202 203<|start_header_id|>assistant<|end_header_id|>204"""205 206rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)207 208# -- GENERATION -- #209"""2101. Create a HuggingFaceEndpoint for the LLM211"""212hf_llm = HuggingFaceEndpoint(213    endpoint_url=HF_LLM_ENDPOINT,214    max_new_tokens=512,215    top_k=10,216    top_p=0.95,217    temperature=0.3,218    repetition_penalty=1.15,219    huggingfacehub_api_token=HF_TOKEN,220)221 222@cl.author_rename223def rename(original_author: str):224    """225    This function can be used to rename the 'author' of a message. 226 227    In this case, we're overriding the 'Assistant' author to be 'Paul Graham Essay Bot'.228    """229    rename_dict = {230        "Assistant" : "Paul Graham Essay Bot"231    }232    return rename_dict.get(original_author, original_author)233 234@cl.on_chat_start235async def start_chat():236    """237    This function will be called at the start of every user session. 238 239    We will build our LCEL RAG chain here, and store it in the user session. 240 241    The user session is a dictionary that is unique to each user session, and is stored in the memory of the server.242    """243 244    lcel_rag_chain = (245        {"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}246        | rag_prompt | hf_llm247    )248 249    cl.user_session.set("lcel_rag_chain", lcel_rag_chain)250 251@cl.on_message  252async def main(message: cl.Message):253    """254    This function will be called every time a message is recieved from a session.255 256    We will use the LCEL RAG chain to generate a response to the user query.257 258    The LCEL RAG chain is stored in the user session, and is unique to each user session - this is why we can access it here.259    """260    lcel_rag_chain = cl.user_session.get("lcel_rag_chain")261 262    msg = cl.Message(content="")263 264    for chunk in await cl.make_async(lcel_rag_chain.stream)(265        {"query": message.content},266        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),267    ):268        await msg.stream_token(chunk)269 270    await msg.send()