Add files via upload

scripts/custom_retriever.py

@@ -0,0 +1,257 @@
+import asyncio
+import os
+import time
+import traceback
+from typing import List, Optional
+
+import logfire
+import tiktoken
+from cohere import AsyncClient
+from dotenv import load_dotenv
+from llama_index.core import Document, QueryBundle
+from llama_index.core.async_utils import run_async_tasks
+from llama_index.core.callbacks import CBEventType, EventPayload
+from llama_index.core.retrievers import (
+    BaseRetriever,
+    KeywordTableSimpleRetriever,
+    VectorIndexRetriever,
+)
+from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle, TextNode
+from llama_index.core.vector_stores import (
+    FilterCondition,
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.postprocessor.cohere_rerank import CohereRerank
+from llama_index.postprocessor.cohere_rerank.base import CohereRerank
+
+load_dotenv()
+
+
+class AsyncCohereRerank(CohereRerank):
+    def __init__(
+        self,
+        top_n: int = 5,
+        model: str = "rerank-english-v3.0",
+        api_key: Optional[str] = None,
+    ) -> None:
+        super().__init__(top_n=top_n, model=model, api_key=api_key)
+        self._api_key = api_key
+        self._model = model
+        self._top_n = top_n
+
+    async def apostprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        if query_bundle is None:
+            raise ValueError("Query bundle must be provided.")
+
+        if len(nodes) == 0:
+            return []
+
+        async_client = AsyncClient(api_key=self._api_key)
+
+        with self.callback_manager.event(
+            CBEventType.RERANKING,
+            payload={
+                EventPayload.NODES: nodes,
+                EventPayload.MODEL_NAME: self._model,
+                EventPayload.QUERY_STR: query_bundle.query_str,
+                EventPayload.TOP_K: self._top_n,
+            },
+        ) as event:
+            texts = [
+                node.node.get_content(metadata_mode=MetadataMode.EMBED)
+                for node in nodes
+            ]
+
+            results = await async_client.rerank(
+                model=self._model,
+                top_n=self._top_n,
+                query=query_bundle.query_str,
+                documents=texts,
+            )
+
+            new_nodes = []
+            for result in results.results:
+                new_node_with_score = NodeWithScore(
+                    node=nodes[result.index].node, score=result.relevance_score
+                )
+                new_nodes.append(new_node_with_score)
+            event.on_end(payload={EventPayload.NODES: new_nodes})
+
+        return new_nodes
+
+
+class CustomRetriever(BaseRetriever):
+    """Custom retriever that performs both semantic search and hybrid search."""
+
+    def __init__(
+        self,
+        vector_retriever: VectorIndexRetriever,
+        document_dict: dict,
+        keyword_retriever=None,
+        mode: str = "AND",
+    ) -> None:
+        """Init params."""
+        self._vector_retriever = vector_retriever
+        self._document_dict = document_dict
+        self._keyword_retriever = keyword_retriever
+        if mode not in ("AND", "OR"):
+            raise ValueError("Invalid mode.")
+        self._mode = mode
+        super().__init__()
+
+    async def _process_retrieval(
+        self, query_bundle: QueryBundle, is_async: bool = True
+    ) -> List[NodeWithScore]:
+        """Common processing logic for both sync and async retrieval."""
+        # Clean query string
+        query_bundle.query_str = query_bundle.query_str.replace(
+            "\ninput is ", ""
+        ).rstrip()
+        logfire.info(f"Retrieving nodes with string: '{query_bundle}'")
+
+        start = time.time()
+
+        # Get nodes from both retrievers
+        if is_async:
+            nodes = await self._vector_retriever.aretrieve(query_bundle)
+        else:
+            nodes = self._vector_retriever.retrieve(query_bundle)
+
+        keyword_nodes = []
+        if self._keyword_retriever:
+            if is_async:
+                keyword_nodes = await self._keyword_retriever.aretrieve(query_bundle)
+            else:
+                keyword_nodes = self._keyword_retriever.retrieve(query_bundle)
+
+        logfire.info(f"Number of vector nodes: {len(nodes)}")
+        logfire.info(f"Number of keyword nodes: {len(keyword_nodes)}")
+
+        # # Filter keyword nodes based on metadata filters from vector retriever
+        # if (
+        #     hasattr(self._vector_retriever, "_filters")
+        #     and self._vector_retriever._filters
+        # ):
+        #     filtered_keyword_nodes = []
+        #     for node in keyword_nodes:
+        #         node_source = node.node.metadata.get("source")
+        #         # Check if node's source matches any of the filter conditions
+        #         for filter in self._vector_retriever._filters.filters:
+        #             if (
+        #                 isinstance(filter, MetadataFilter)
+        #                 and filter.key == "source"
+        #                 and filter.operator == FilterOperator.EQ
+        #                 and filter.value == node_source
+        #             ):
+        #                 filtered_keyword_nodes.append(node)
+        #                 break
+        #     keyword_nodes = filtered_keyword_nodes
+        #     logfire.info(
+        #         f"Number of keyword nodes after filtering: {len(keyword_nodes)}"
+        #     )
+
+        # Combine results based on mode
+        vector_ids = {n.node.node_id for n in nodes}
+        keyword_ids = {n.node.node_id for n in keyword_nodes}
+        combined_dict = {n.node.node_id: n for n in nodes}
+        combined_dict.update({n.node.node_id: n for n in keyword_nodes})
+
+        # If no keyword retriever or no keyword nodes, just use vector nodes
+        if not self._keyword_retriever or not keyword_nodes:
+            retrieve_ids = vector_ids
+        else:
+            retrieve_ids = (
+                vector_ids.intersection(keyword_ids)
+                if self._mode == "AND"
+                else vector_ids.union(keyword_ids)
+            )
+
+        nodes = [combined_dict[rid] for rid in retrieve_ids]
+        logfire.info(f"Number of combined nodes: {len(nodes)}")
+
+        # Filter unique doc IDs
+        nodes = self._filter_nodes_by_unique_doc_id(nodes)
+        logfire.info(f"Number of nodes without duplicate doc IDs: {len(nodes)}")
+
+        # Process node contents
+        for node in nodes:
+            doc_id = node.node.source_node.node_id
+            if node.metadata["retrieve_doc"]:
+                doc = self._document_dict[doc_id]
+                node.node.text = doc.text
+            node.node.node_id = doc_id
+
+        # Rerank results
+        try:
+            reranker = (
+                AsyncCohereRerank(top_n=5, model="rerank-english-v3.0")
+                if is_async
+                else CohereRerank(top_n=5, model="rerank-english-v3.0")
+            )
+            nodes = (
+                await reranker.apostprocess_nodes(nodes, query_bundle)
+                if is_async
+                else reranker.postprocess_nodes(nodes, query_bundle)
+            )
+        except Exception as e:
+            error_msg = f"Error during reranking: {type(e).__name__}: {str(e)}\n"
+            error_msg += "Traceback:\n"
+            error_msg += traceback.format_exc()
+            logfire.error(error_msg)
+
+        # Filter by score and token count
+        nodes_filtered = self._filter_by_score_and_tokens(nodes)
+
+        duration = time.time() - start
+        logfire.info(f"Retrieving nodes took {duration:.2f}s")
+        logfire.info(f"Nodes sent to LLM: {nodes_filtered[:5]}")
+
+        return nodes_filtered[:5]
+
+    def _filter_nodes_by_unique_doc_id(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        """Filter nodes to keep only unique doc IDs."""
+        unique_nodes = {}
+        for node in nodes:
+            doc_id = node.node.source_node.node_id
+            if doc_id is not None and doc_id not in unique_nodes:
+                unique_nodes[doc_id] = node
+        return list(unique_nodes.values())
+
+    def _filter_by_score_and_tokens(
+        self, nodes: List[NodeWithScore]
+    ) -> List[NodeWithScore]:
+        """Filter nodes by score and token count."""
+        nodes_filtered = []
+        total_tokens = 0
+        enc = tiktoken.encoding_for_model("gpt-4o-mini")
+
+        for node in nodes:
+            if node.score < 0.10:
+                logfire.info(f"Skipping node with score {node.score}")
+                continue
+
+            node_tokens = len(enc.encode(node.node.text))
+            if total_tokens + node_tokens > 100_000:
+                logfire.info("Skipping node due to token count exceeding 100k")
+                break
+
+            total_tokens += node_tokens
+            nodes_filtered.append(node)
+
+        return nodes_filtered
+
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Async retrieve nodes given query."""
+        return await self._process_retrieval(query_bundle, is_async=True)
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Sync retrieve nodes given query."""
+        return asyncio.run(self._process_retrieval(query_bundle, is_async=False))

scripts/evaluate_rag_system.py

@@ -0,0 +1,773 @@
+import asyncio
+import html
+import json
+import logging
+import os
+import pdb
+import pickle
+import random
+import time
+from typing import Dict, List, Optional, Tuple
+
+import aiofiles
+import chromadb
+import logfire
+import pandas as pd
+from custom_retriever import CustomRetriever
+from llama_index.agent.openai import OpenAIAgent
+from llama_index.core import Document, SimpleKeywordTableIndex, VectorStoreIndex
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.bridge.pydantic import Field, SerializeAsAny
+from llama_index.core.chat_engine.types import (
+    AGENT_CHAT_RESPONSE_TYPE,
+    AgentChatResponse,
+    ChatResponseMode,
+)
+from llama_index.core.evaluation import (
+    AnswerRelevancyEvaluator,
+    BatchEvalRunner,
+    EmbeddingQAFinetuneDataset,
+    FaithfulnessEvaluator,
+    RelevancyEvaluator,
+)
+from llama_index.core.evaluation.base import EvaluationResult
+from llama_index.core.evaluation.retrieval.base import (
+    BaseRetrievalEvaluator,
+    RetrievalEvalMode,
+    RetrievalEvalResult,
+)
+from llama_index.core.indices.base_retriever import BaseRetriever
+from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
+from llama_index.core.retrievers import (
+    BaseRetriever,
+    KeywordTableSimpleRetriever,
+    VectorIndexRetriever,
+)
+from llama_index.core.schema import ImageNode, NodeWithScore, QueryBundle, TextNode
+from llama_index.core.tools import RetrieverTool, ToolMetadata
+from llama_index.core.vector_stores import (
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.gemini import Gemini
+from llama_index.llms.openai import OpenAI
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from prompts import system_message_openai_agent
+from pydantic import BaseModel, Field
+from tqdm.asyncio import tqdm_asyncio
+
+# from setup import (
+#     AVAILABLE_SOURCES,
+#     AVAILABLE_SOURCES_UI,
+#     custom_retriever_all_sources,
+#     custom_retriever_langchain,
+#     custom_retriever_llama_index,
+#     custom_retriever_openai_cookbooks,
+#     custom_retriever_peft,
+#     custom_retriever_transformers,
+#     custom_retriever_trl,
+# )
+
+
+class RotatingJSONLWriter:
+    def __init__(
+        self, base_filename: str, max_size: int = 10**6, backup_count: int = 5
+    ):
+        """
+        Initialize the rotating JSONL writer.
+
+        Args:
+            base_filename (str): The base filename for the JSONL files.
+            max_size (int): Maximum size in bytes before rotating.
+            backup_count (int): Number of backup files to keep.
+        """
+        self.base_filename = base_filename
+        self.max_size = max_size
+        self.backup_count = backup_count
+        self.current_file = base_filename
+
+    async def write(self, data: dict):
+        # Rotate if file exceeds max size
+        if (
+            os.path.exists(self.current_file)
+            and os.path.getsize(self.current_file) > self.max_size
+        ):
+            await self.rotate_files()
+
+        async with aiofiles.open(self.current_file, "a", encoding="utf-8") as f:
+            await f.write(json.dumps(data, ensure_ascii=False) + "\n")
+
+    async def rotate_files(self):
+        # Remove the oldest backup if it exists
+        oldest_backup = f"{self.base_filename}.{self.backup_count}"
+        if os.path.exists(oldest_backup):
+            os.remove(oldest_backup)
+
+        # Rotate existing backups
+        for i in range(self.backup_count - 1, 0, -1):
+            src = f"{self.base_filename}.{i}"
+            dst = f"{self.base_filename}.{i + 1}"
+            if os.path.exists(src):
+                os.rename(src, dst)
+
+        # Rename current file to backup
+        os.rename(self.current_file, f"{self.base_filename}.1")
+
+
+class AsyncKeywordTableSimpleRetriever(KeywordTableSimpleRetriever):
+    async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, self._retrieve, query_bundle)
+
+
+class SampleableEmbeddingQADataset:
+    def __init__(self, dataset: EmbeddingQAFinetuneDataset):
+        self.dataset = dataset
+
+    def sample(self, n: int) -> EmbeddingQAFinetuneDataset:
+        """
+        Sample n queries from the dataset.
+
+        Args:
+            n (int): Number of queries to sample.
+
+        Returns:
+            EmbeddingQAFinetuneDataset: A new dataset with the sampled queries.
+        """
+        if n > len(self.dataset.queries):
+            raise ValueError(
+                f"n ({n}) is greater than the number of queries ({len(self.dataset.queries)})"
+            )
+
+        sampled_query_ids = random.sample(list(self.dataset.queries.keys()), n)
+
+        sampled_queries = {qid: self.dataset.queries[qid] for qid in sampled_query_ids}
+        sampled_relevant_docs = {
+            qid: self.dataset.relevant_docs[qid] for qid in sampled_query_ids
+        }
+
+        # Collect all unique document IDs from the sampled relevant docs
+        sampled_doc_ids = set()
+        for doc_ids in sampled_relevant_docs.values():
+            sampled_doc_ids.update(doc_ids)
+
+        sampled_corpus = {
+            doc_id: self.dataset.corpus[doc_id] for doc_id in sampled_doc_ids
+        }
+
+        return EmbeddingQAFinetuneDataset(
+            queries=sampled_queries,
+            corpus=sampled_corpus,
+            relevant_docs=sampled_relevant_docs,
+            mode=self.dataset.mode,
+        )
+
+    def __getattr__(self, name):
+        return getattr(self.dataset, name)
+
+
+class RetrieverEvaluator(BaseRetrievalEvaluator):
+    """Retriever evaluator.
+
+    This module will evaluate a retriever using a set of metrics.
+
+    Args:
+        metrics (List[BaseRetrievalMetric]): Sequence of metrics to evaluate
+        retriever: Retriever to evaluate.
+        node_postprocessors (Optional[List[BaseNodePostprocessor]]): Post-processor to apply after retrieval.
+    """
+
+    retriever: BaseRetriever = Field(..., description="Retriever to evaluate")
+    node_postprocessors: Optional[List[SerializeAsAny[BaseNodePostprocessor]]] = Field(
+        default=None, description="Optional post-processor"
+    )
+
+    async def _aget_retrieved_ids_and_texts(
+        self,
+        query: str,
+        mode: RetrievalEvalMode = RetrievalEvalMode.TEXT,
+        source: str = "",
+    ) -> Tuple[List[str], List[str]]:
+        """Get retrieved ids and texts, potentially applying a post-processor."""
+        try:
+            retrieved_nodes: list[NodeWithScore] = await self.retriever.aretrieve(query)
+            logfire.info(f"Retrieved {len(retrieved_nodes)} nodes for: '{query}'")
+        except Exception as e:
+            return ["00000000-0000-0000-0000-000000000000"], [str(e)]
+
+        if len(retrieved_nodes) == 0 or retrieved_nodes is None:
+            print(f"No nodes retrieved for {query}")
+            return ["00000000-0000-0000-0000-000000000000"], ["No nodes retrieved"]
+
+        if self.node_postprocessors:
+            for node_postprocessor in self.node_postprocessors:
+                retrieved_nodes = node_postprocessor.postprocess_nodes(
+                    retrieved_nodes, query_str=query
+                )
+
+        return (
+            [node.node.node_id for node in retrieved_nodes],
+            [node.node.text for node in retrieved_nodes],  # type: ignore
+        )
+
+
+class OpenAIAgentRetrieverEvaluator(BaseRetrievalEvaluator):
+    agent: OpenAIAgent = Field(description="The OpenAI agent used for retrieval")
+
+    async def _aget_retrieved_ids_and_texts(
+        self,
+        query: str,
+        mode: RetrievalEvalMode = RetrievalEvalMode.TEXT,
+        source: str = "",
+    ) -> Tuple[List[str], List[str]]:
+
+        self.agent.memory.reset()
+
+        try:
+            logfire.info(f"Executing agent with query: {query}")
+            response: AgentChatResponse = await self.agent.achat(query)
+        except Exception as e:
+            # await self._save_response_data_async(
+            #     source, query, ["Error retrieving nodes"], "Error retrieving nodes"
+            # )
+            return ["00000000-0000-0000-0000-000000000000"], [str(e)]
+
+        retrieved_nodes: list[NodeWithScore] = get_nodes_with_score(response)
+        logfire.info(f"Retrieved {len(retrieved_nodes)} to answer: '{query}'")
+        retrieved_nodes = retrieved_nodes[:6]  # Limit to first 6 retrieved nodes
+
+        if len(retrieved_nodes) == 0 or retrieved_nodes is None:
+            # await self._save_response_data_async(
+            #     source, query, ["No retrieved nodes"], "No retrieved nodes"
+            # )
+            return ["00000000-0000-0000-0000-000000000000"], ["No nodes retrieved"]
+
+        retrieved_ids = [node.node.node_id for node in retrieved_nodes]
+        retrieved_texts = [node.node.text for node in retrieved_nodes]  # type: ignore
+
+        # Will not save context as its too long (token wise), costly and takes too much time.
+        await self._save_response_data_async(
+            source=source, query=query, context="", response=response.response
+        )
+
+        return retrieved_ids, retrieved_texts
+
+    async def _save_response_data_async(self, source, query, context, response):
+        data = {
+            "source": source,
+            "question": query,
+            # "context": context,
+            "answer": response,
+        }
+        await rotating_writer.write(data)
+
+
+def get_nodes_with_score(completion) -> list[NodeWithScore]:
+    retrieved_nodes = []
+    for source in completion.sources:  # completion.sources = list[ToolOutput]
+        if source.is_error == True:
+            continue
+        for node in source.raw_output:  # source.raw_output = list[NodeWithScore]
+            retrieved_nodes.append(node)
+    return retrieved_nodes
+
+
+def setup_basic_database(db_collection, dict_file_name, keyword_retriever):
+    db = chromadb.PersistentClient(path=f"data/{db_collection}")
+    chroma_collection = db.get_or_create_collection(db_collection)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+
+    # embed_model = OpenAIEmbedding(model="text-embedding-3-large", mode="similarity")
+    embed_model = CohereEmbedding(
+        api_key=os.environ["COHERE_API_KEY"],
+        model_name="embed-english-v3.0",
+        input_type="search_query",
+    )
+    # client = embed_model._get_client()
+    # aclient = embed_model._get_aclient()
+    # logfire.instrument_openai(client)
+    # logfire.instrument_openai(aclient)
+
+    index = VectorStoreIndex.from_vector_store(
+        vector_store=vector_store,
+        show_progress=True,
+    )
+    vector_retriever = VectorIndexRetriever(
+        index=index,
+        similarity_top_k=15,
+        embed_model=embed_model,
+    )
+    with open(f"data/{db_collection}/{dict_file_name}", "rb") as f:
+        document_dict = pickle.load(f)
+
+    return CustomRetriever(vector_retriever, document_dict, keyword_retriever, "OR")
+
+
+def update_query_engine_tools(selected_sources, custom_retriever_all_sources):
+    tools = []
+    source_mapping = {
+        # "Transformers Docs": (
+        #     custom_retriever_transformers,
+        #     "Transformers_information",
+        #     """Useful for general questions asking about the artificial intelligence (AI) field. Employ this tool to fetch information on topics such as language models (LLMs) models such as Llama3 and theory (transformer architectures), tips on prompting, quantization, etc.""",
+        # ),
+        # "PEFT Docs": (
+        #     custom_retriever_peft,
+        #     "PEFT_information",
+        #     """Useful for questions asking about efficient LLM fine-tuning. Employ this tool to fetch information on topics such as LoRA, QLoRA, etc.""",
+        # ),
+        # "TRL Docs": (
+        #     custom_retriever_trl,
+        #     "TRL_information",
+        #     """Useful for questions asking about fine-tuning LLMs with reinforcement learning (RLHF). Includes information about the Supervised Fine-tuning step (SFT), Reward Modeling step (RM), and the Proximal Policy Optimization (PPO) step.""",
+        # ),
+        # "LlamaIndex Docs": (
+        #     custom_retriever_llama_index,
+        #     "LlamaIndex_information",
+        #     """Useful for questions asking about retrieval augmented generation (RAG) with LLMs and embedding models. It is the documentation of a framework, includes info about fine-tuning embedding models, building chatbots, and agents with llms, using vector databases, embeddings, information retrieval with cosine similarity or bm25, etc.""",
+        # ),
+        # "OpenAI Cookbooks": (
+        #     custom_retriever_openai_cookbooks,
+        #     "openai_cookbooks_info",
+        #     """Useful for questions asking about accomplishing common tasks with the OpenAI API. Returns example code and guides stored in Jupyter notebooks, including info about ChatGPT GPT actions, OpenAI Assistants API,  and How to fine-tune OpenAI's GPT-4o and GPT-4o-mini models with the OpenAI API.""",
+        # ),
+        # "LangChain Docs": (
+        #     custom_retriever_langchain,
+        #     "langchain_info",
+        #     """Useful for questions asking about the LangChain framework. It is the documentation of the LangChain framework, includes info about building chains, agents, and tools, using memory, prompts, callbacks, etc.""",
+        # ),
+        "All Sources": (
+            custom_retriever_all_sources,
+            "all_sources_info",
+            """Useful for all questions, contains information about the field of AI.""",
+        ),
+    }
+
+    for source in selected_sources:
+        if source in source_mapping:
+            retriever, name, description = source_mapping[source]
+            tools.append(
+                RetrieverTool(
+                    retriever=retriever,
+                    metadata=ToolMetadata(
+                        name=name,
+                        description=description,
+                    ),
+                )
+            )
+
+    return tools
+
+
+def setup_agent(custom_retriever_all_sources) -> OpenAIAgent:
+
+    llm = OpenAI(
+        temperature=1,
+        # model="gpt-4o",
+        model="gpt-4o-mini",
+        max_tokens=5000,
+        max_retries=3,
+    )
+    client = llm._get_client()
+    logfire.instrument_openai(client)
+    aclient = llm._get_aclient()
+    logfire.instrument_openai(aclient)
+
+    tools_available = [
+        # "Transformers Docs",
+        # "PEFT Docs",
+        # "TRL Docs",
+        # "LlamaIndex Docs",
+        # "LangChain Docs",
+        # "OpenAI Cookbooks",
+        "All Sources",
+    ]
+    query_engine_tools = update_query_engine_tools(
+        tools_available, custom_retriever_all_sources
+    )
+
+    agent = OpenAIAgent.from_tools(
+        llm=llm,
+        tools=query_engine_tools,
+        system_prompt=system_message_openai_agent,
+    )
+
+    return agent
+
+
+async def evaluate_answers():
+    start_time = time.time()
+
+    # Gemini is not async here, maybe it could work with multithreading?
+    # llm = Gemini(model="models/gemini-1.5-flash-002", temperature=1, max_tokens=1000)
+    llm = OpenAI(model="gpt-4o-mini", temperature=1, max_tokens=1000)
+    relevancy_evaluator = AnswerRelevancyEvaluator(llm=llm)
+
+    # Load queries and response strings from JSONL file
+    query_response_pairs = []
+    with open("response_data.jsonl", "r") as f:
+        for line in f:
+            data = json.loads(line)
+            query_response_pairs.append(
+                (data["source"], data["query"], data["response"])
+            )
+
+    logfire.info(f"Number of queries and answers: {len(query_response_pairs)}")
+
+    semaphore = asyncio.Semaphore(90)  # Adjust this value as needed
+
+    async def evaluate_query_response(source, query, response):
+        async with semaphore:
+            try:
+                result: EvaluationResult = await relevancy_evaluator.aevaluate(
+                    query=query, response=response
+                )
+                return source, result
+            except Exception as e:
+                logfire.error(f"Error evaluating query for {source}: {str(e)}")
+                return source, None
+
+    # Use asyncio.gather to run all evaluations concurrently
+    results = await tqdm_asyncio.gather(
+        *[
+            evaluate_query_response(source, query, response)
+            for source, query, response in query_response_pairs
+        ],
+        desc="Evaluating answers",
+        total=len(query_response_pairs),
+    )
+
+    # Process results
+    eval_results = {}
+    for item in results:
+        if isinstance(item, tuple) and len(item) == 2:
+            source, result = item
+            if result is not None:
+                if source not in eval_results:
+                    eval_results[source] = []
+                eval_results[source].append(result)
+        else:
+            logfire.error(f"Unexpected result: {item}")
+
+    # Save results for each source
+    for source, results in eval_results.items():
+        with open(f"eval_answers_results_{source}.pkl", "wb") as f:
+            pickle.dump(results, f)
+
+    end_time = time.time()
+    logfire.info(f"Total evaluation time: {round(end_time - start_time, 3)} seconds")
+
+    return eval_results
+
+
+def create_docs(input_file: str) -> List[Document]:
+    with open(input_file, "r") as f:
+        documents = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+def get_sample_size(source: str, total_queries: int) -> int:
+    """Determine the number of queries to sample based on the source."""
+    # small_datasets = {"peft": 0, "trl": 0, "openai_cookbooks": 0}
+    # large_datasets = {
+    #     "transformers": 0,
+    #     "llama_index": 0,
+    #     "langchain": 1,
+    #     "tai_blog": 0,
+    # }
+    small_datasets = {"peft": 49, "trl": 34, "openai_cookbooks": 170}
+    large_datasets = {
+        "transformers": 200,
+        "llama_index": 200,
+        "langchain": 200,
+        "tai_blog": 200,
+    }
+    # small_datasets = {"peft": 49, "trl": 34, "openai_cookbooks": 100}
+    # large_datasets = {
+    #     "transformers": 100,
+    #     "llama_index": 100,
+    #     "langchain": 100,
+    #     "tai_blog": 100,
+    # }
+    # small_datasets = {"peft": 18, "trl": 12, "openai_cookbooks": 14}
+    # large_datasets = {
+    #     "transformers": 24,
+    #     "llama_index": 8,
+    #     "langchain": 6,
+    #     "tai_blog": 18,
+    # }
+    # small_datasets = {"peft": 4, "trl": 4, "openai_cookbooks": 4}
+    # large_datasets = {
+    #     "transformers": 4,
+    #     "llama_index": 4,
+    #     "langchain": 5,
+    #     "tai_blog": 5,
+    # }
+
+    if source in small_datasets:
+        return small_datasets[source]
+    elif source in large_datasets:
+        return large_datasets[source]
+    else:
+        return min(100, total_queries)  # Default to 100 or all queries if less than 100
+
+
+async def evaluate_retriever():
+    start_time = time.time()
+    with open("data/keyword_retriever_async.pkl", "rb") as f:
+        keyword_retriever = pickle.load(f)
+
+    custom_retriever_all_sources: CustomRetriever = setup_basic_database(
+        "chroma-db-all_sources", "document_dict_all_sources.pkl", keyword_retriever
+    )
+    # agent = setup_agent(custom_retriever_all_sources)
+
+    # filters = MetadataFilters(
+    #     filters=[
+    #         MetadataFilter(key="source", operator=FilterOperator.EQ, value="langchain"),
+    #     ]
+    # )
+    # custom_retriever_all_sources._vector_retriever._filters = filters
+
+    end_time = time.time()
+    logfire.info(
+        f"Time taken for setup the custom retriever: {round(end_time - start_time, 2)} seconds"
+    )
+
+    sources_to_evaluate = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+
+    # for k in [5, 7, 9, 11, 13, 15]:
+    #     custom_retriever_all_sources._vector_retriever._similarity_top_k = k
+
+    retriever_evaluator = RetrieverEvaluator.from_metric_names(
+        ["mrr", "hit_rate"], retriever=custom_retriever_all_sources
+    )
+    # retriever_evaluator = OpenAIAgentRetrieverEvaluator.from_metric_names(
+    #     metric_names=["mrr", "hit_rate"], agent=agent
+    # )
+
+    all_query_pairs = []
+    for source in sources_to_evaluate:
+        rag_eval_dataset = EmbeddingQAFinetuneDataset.from_json(
+            f"scripts/rag_eval_{source}.json"
+        )
+        sampleable_dataset = SampleableEmbeddingQADataset(rag_eval_dataset)
+        sample_size = get_sample_size(source, len(sampleable_dataset.queries))
+        sampled_dataset = sampleable_dataset.sample(n=sample_size)
+        query_expected_ids_pairs = sampled_dataset.query_docid_pairs
+        all_query_pairs.extend(
+            [(source, pair[0], pair[1]) for pair in query_expected_ids_pairs]
+        )
+
+    semaphore = asyncio.Semaphore(220)  # 250 caused a couple of errors
+    # semaphore = asyncio.Semaphore(90)  # 100 caused a couple of errors with agent
+
+    async def evaluate_query(source, query, expected_ids):
+        async with semaphore:
+            try:
+                result: RetrievalEvalResult = await retriever_evaluator.aevaluate(
+                    query=query,
+                    expected_ids=expected_ids,
+                    mode=RetrievalEvalMode.TEXT,
+                    source=source,
+                )
+                return source, result
+            except Exception as e:
+                logfire.error(f"Error evaluating query for {source}: {str(e)}")
+                return source, None
+
+    # Use asyncio.gather to run all evaluations concurrently
+    results = await tqdm_asyncio.gather(
+        *[
+            evaluate_query(source, query, expected_ids)
+            for source, query, expected_ids in all_query_pairs
+        ],
+        desc="Evaluating queries",
+        total=len(all_query_pairs),
+    )
+
+    # Process results
+    eval_results = {source: [] for source in sources_to_evaluate}
+    for item in results:
+        if isinstance(item, tuple) and len(item) == 2:
+            source, result = item
+            if result is not None:
+                eval_results[source].append(result)
+        else:
+            logfire.error(f"Unexpected result: {item}")
+
+    # Save results for each source
+    for source, results in eval_results.items():
+        with open(f"eval_results_{source}.pkl", "wb") as f:
+            pickle.dump(results, f)
+        # print(display_results_retriever(source, results))
+
+    end_time = time.time()
+    logfire.info(f"Total evaluation time: {round(end_time - start_time, 3)} seconds")
+
+
+def display_results_retriever(name, eval_results):
+    """Display results from evaluate."""
+
+    metric_dicts = []
+    for eval_result in eval_results:
+        metric_dict = eval_result.metric_vals_dict
+        metric_dicts.append(metric_dict)
+
+    full_df = pd.DataFrame(metric_dicts)
+
+    hit_rate = full_df["hit_rate"].mean()
+    mrr = full_df["mrr"].mean()
+
+    metric_df = pd.DataFrame(
+        {"Retriever Name": [name], "Hit Rate": [hit_rate], "MRR": [mrr]}
+    )
+
+    return metric_df
+
+
+def display_results():
+
+    sources = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+    # retrievers_to_evaluate = [
+    #     # "chroma-db-all_sources_400_0",
+    #     # "chroma-db-all_sources_400_200",
+    #     # "chroma-db-all_sources_500_0",
+    #     # "chroma-db-all_sources_500_250",
+    #     # "chroma-db-all_sources",
+    #     # "chroma-db-all_sources_800_400",
+    #     # "chroma-db-all_sources_1000_0",
+    #     # "chroma-db-all_sources_1000_500",
+    # ]
+
+    # topk = [5, 7, 9, 11, 13, 15]
+    # for k in topk:
+    # for db_name in retrievers_to_evaluate:
+    if True:
+        # print("-" * 20)
+        # print(f"Retriever {db_name}")
+        for source in sources:
+            with open(f"eval_results_{source}.pkl", "rb") as f:
+                eval_results = pickle.load(f)
+            print(display_results_retriever(f"{source}", eval_results))
+
+
+def display_results_answers():
+
+    sources = [
+        "transformers",
+        "peft",
+        "trl",
+        "llama_index",
+        "langchain",
+        "openai_cookbooks",
+        "tai_blog",
+    ]
+
+    for source in sources:
+        with open(f"eval_answers_results_{source}.pkl", "rb") as f:
+            eval_results = pickle.load(f)
+        print(
+            f"Score for {source}:",
+            sum(result.score for result in eval_results) / len(eval_results),
+        )
+
+
+async def main():
+    await evaluate_retriever()
+    display_results()
+    # await evaluate_answers()
+    # display_results_answers()
+    return
+
+
+if __name__ == "__main__":
+
+    logfire.configure()
+    rotating_writer = RotatingJSONLWriter(
+        "response_data.jsonl", max_size=10**7, backup_count=5
+    )
+
+    start_time = time.time()
+    asyncio.run(main())
+    end_time = time.time()
+    logfire.info(
+        f"Time taken to run script: {round((end_time - start_time), 3)} seconds"
+    )
+
+    # # Creating the keyword index and retriever
+    # logfire.info("Creating nodes from documents")
+    # documents = create_docs("data/all_sources_data.jsonl")
+    # pipeline = IngestionPipeline(
+    #     transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)]
+    # )
+    # all_nodes = pipeline.run(documents=documents, show_progress=True)
+    # # with open("data/all_nodes.pkl", "wb") as f:
+    # #     pickle.dump(all_nodes, f)
+
+    # # all_nodes = pickle.load(open("data/all_nodes.pkl", "rb"))
+    # logfire.info(f"Number of nodes: {len(all_nodes)}")
+
+    # with open("processed_chunks.pkl", "rb") as f:
+    #     all_nodes: list[TextNode] = pickle.load(f)
+
+    # keyword_index = SimpleKeywordTableIndex(
+    #     nodes=all_nodes, max_keywords_per_chunk=10, show_progress=True, use_async=False
+    # )
+    # # with open("data/keyword_index.pkl", "wb") as f:
+    # #     pickle.dump(keyword_index, f)
+    # # keyword_index = pickle.load(open("data/keyword_index.pkl", "rb"))
+
+    # logfire.info("Creating keyword retriever")
+    # keyword_retriever = AsyncKeywordTableSimpleRetriever(index=keyword_index)
+
+    # with open("data/keyword_retriever_async.pkl", "wb") as f:
+    #     pickle.dump(keyword_retriever, f)

scripts/generate_qa_dataset.ipynb

@@ -0,0 +1,287 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import nest_asyncio\n",
+    "\n",
+    "nest_asyncio.apply()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Generate synthetic dataset of questions + doc ids\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# DEFAULT_QA_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "# Context information is below.\n",
+    "\n",
+    "# ---------------------\n",
+    "# {context_str}\n",
+    "# ---------------------\n",
+    "\n",
+    "# Given the context information and not prior knowledge.\n",
+    "# generate only questions based on the below query.\n",
+    "\n",
+    "# You are a Teacher/ Professor with expertise in the field of AI. Your task is to setup \\\n",
+    "# {num_questions_per_chunk} questions for an upcoming \\\n",
+    "# quiz/examination. The questions should be diverse in nature \\\n",
+    "# across the document. Restrict the questions to the \\\n",
+    "# context information provided.\"\n",
+    "# \"\"\"\n",
+    "\n",
+    "STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL = \"\"\"\\\n",
+    "Context information is below.\n",
+    "\n",
+    "---------------------\n",
+    "{context_str}\n",
+    "---------------------\n",
+    "\n",
+    "Given the context information above, generate {num_questions_per_chunk} questions that a student might ask about AI, specifically related to the information provided in the context, \n",
+    "but without the questions mentioning the context information.\n",
+    "\n",
+    "You are simulating curious students who are learning about AI. Your task is to create questions that:\n",
+    "1. Reflect genuine curiosity about the topics covered in the context.\n",
+    "2. Vary in complexity, from basic clarifications to more advanced inquiries.\n",
+    "3. Demonstrate a student's desire to understand AI concepts better.\n",
+    "4. Can be answered using the information provided in the context.\n",
+    "\n",
+    "The questions should be diverse and cover different aspects of the content. Do not use any prior knowledge beyond what's given in the context. Ensure that each question you generate can be answered using the information provided in the context.\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.llms.openai import OpenAI\n",
+    "from llama_index.llms.gemini import Gemini\n",
+    "\n",
+    "# llm = Gemini(model=\"models/gemini-1.5-flash-latest\", temperature=1, max_tokens=1000)\n",
+    "llm = OpenAI(\n",
+    "    api_key=\"\",\n",
+    "    model=\"gpt-4o-mini\",\n",
+    "    temperature=1,\n",
+    "    max_tokens=1000,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "import re\n",
+    "import uuid\n",
+    "import warnings\n",
+    "import asyncio\n",
+    "from typing import Dict, List, Tuple\n",
+    "\n",
+    "from llama_index.core.llms.llm import LLM\n",
+    "from llama_index.core.schema import MetadataMode, TextNode\n",
+    "from tqdm.asyncio import tqdm as async_tqdm\n",
+    "from llama_index.core.llama_dataset.legacy.embedding import EmbeddingQAFinetuneDataset\n",
+    "from llama_index.core import Document\n",
+    "\n",
+    "\n",
+    "async def generate_qa_embedding_pairs(\n",
+    "    nodes: List[TextNode],\n",
+    "    llm: LLM,\n",
+    "    qa_generate_prompt_tmpl: str = STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "    num_questions_per_chunk: int = 1,\n",
+    "    max_concurrent: int = 10,\n",
+    "    delay: float = 0.5,\n",
+    ") -> EmbeddingQAFinetuneDataset:\n",
+    "    \"\"\"Generate examples given a set of nodes.\"\"\"\n",
+    "    node_dict = {\n",
+    "        node.node_id: node.get_content(metadata_mode=MetadataMode.NONE)\n",
+    "        for node in nodes\n",
+    "    }\n",
+    "\n",
+    "    queries = {}\n",
+    "    relevant_docs = {}\n",
+    "    semaphore = asyncio.Semaphore(max_concurrent)\n",
+    "\n",
+    "    async def process_node(node_id: str, text: str):\n",
+    "        async with semaphore:\n",
+    "            query = qa_generate_prompt_tmpl.format(\n",
+    "                context_str=text, num_questions_per_chunk=num_questions_per_chunk\n",
+    "            )\n",
+    "            response = await llm.acomplete(\n",
+    "                query\n",
+    "            )  # Assuming the LLM has an async method\n",
+    "\n",
+    "            result = str(response).strip().split(\"\\n\")\n",
+    "            questions = [\n",
+    "                re.sub(r\"^\\d+[\\).\\s]\", \"\", question).strip() for question in result\n",
+    "            ]\n",
+    "            questions = [question for question in questions if len(question) > 0][\n",
+    "                :num_questions_per_chunk\n",
+    "            ]\n",
+    "\n",
+    "            num_questions_generated = len(questions)\n",
+    "            if num_questions_generated < num_questions_per_chunk:\n",
+    "                warnings.warn(\n",
+    "                    f\"Fewer questions generated ({num_questions_generated}) \"\n",
+    "                    f\"than requested ({num_questions_per_chunk}).\"\n",
+    "                )\n",
+    "\n",
+    "            for question in questions:\n",
+    "                question_id = str(uuid.uuid4())\n",
+    "                queries[question_id] = question\n",
+    "                relevant_docs[question_id] = [node_id]\n",
+    "\n",
+    "            await asyncio.sleep(delay)\n",
+    "\n",
+    "    # Use asyncio.gather to process nodes concurrently\n",
+    "    await async_tqdm.gather(\n",
+    "        *[process_node(node_id, text) for node_id, text in node_dict.items()]\n",
+    "    )\n",
+    "\n",
+    "    # construct dataset\n",
+    "    return EmbeddingQAFinetuneDataset(\n",
+    "        queries=queries, corpus=node_dict, relevant_docs=relevant_docs\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import pickle\n",
+    "\n",
+    "\n",
+    "async def generate_questions(path, source_name):\n",
+    "    nodes = []\n",
+    "    with open(path, \"rb\") as f:\n",
+    "        document_dict = pickle.load(f)\n",
+    "        for doc_id in document_dict.keys():\n",
+    "            doc: Document = document_dict[doc_id]\n",
+    "            if doc.metadata[\"tokens\"] >= 100_000:\n",
+    "                print(\"skipping\", doc.metadata[\"tokens\"])\n",
+    "                continue\n",
+    "            node = TextNode(text=doc.text, metadata=doc.metadata, id_=doc_id)\n",
+    "            nodes.append(node)\n",
+    "\n",
+    "    rag_eval_dataset: EmbeddingQAFinetuneDataset = await generate_qa_embedding_pairs(\n",
+    "        nodes,\n",
+    "        llm=llm,\n",
+    "        num_questions_per_chunk=1,\n",
+    "        qa_generate_prompt_tmpl=STUDENT_AI_QUESTIONS_GENERATE_PROMPT_TMPL,\n",
+    "        max_concurrent=20,  # Adjust this to control concurrency\n",
+    "        delay=0.5,  # Adjust this to add delay between API calls\n",
+    "    )\n",
+    "    rag_eval_dataset.save_json(f\"./rag_eval_{source_name}.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-langchain/document_dict_langchain.pkl\", \"langchain\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-llama_index/document_dict_llama_index.pkl\", \"llama_index\"\n",
+    "# )\n",
+    "# await generate_questions(\n",
+    "#     \"../data/chroma-db-openai_cookbooks/document_dict_openai_cookbooks.pkl\",\n",
+    "#     \"openai_cookbooks\",\n",
+    "# )\n",
+    "# await generate_questions(\"../data/chroma-db-peft/document_dict_peft.pkl\", \"peft\")\n",
+    "# await generate_questions(\"../data/chroma-db-trl/document_dict_trl.pkl\", \"trl\")\n",
+    "\n",
+    "await generate_questions(\n",
+    "    \"../data/chroma-db-tai_blog/document_dict_tai_blog.pkl\", \"tai_blog\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # We can also load the dataset from a previously saved json file.\n",
+    "# from llama_index.core.evaluation import EmbeddingQAFinetuneDataset\n",
+    "\n",
+    "# rag_eval_dataset = EmbeddingQAFinetuneDataset.from_json(\"./rag_eval_transformers.json\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

scripts/main.py

@@ -0,0 +1,268 @@
+import pdb
+
+import gradio as gr
+import logfire
+from custom_retriever import CustomRetriever
+from llama_index.agent.openai import OpenAIAgent
+from llama_index.core.llms import MessageRole
+from llama_index.core.memory import ChatSummaryMemoryBuffer
+from llama_index.core.tools import RetrieverTool, ToolMetadata
+from llama_index.core.vector_stores import (
+    FilterCondition,
+    FilterOperator,
+    MetadataFilter,
+    MetadataFilters,
+)
+from llama_index.llms.openai import OpenAI
+from prompts import system_message_openai_agent
+from setup import (
+    AVAILABLE_SOURCES,
+    AVAILABLE_SOURCES_UI,
+    CONCURRENCY_COUNT,
+    custom_retriever_all_sources,
+)
+
+
+def update_query_engine_tools(selected_sources) -> list[RetrieverTool]:
+    tools = []
+    source_mapping: dict[str, tuple[CustomRetriever, str, str]] = {
+        "All Sources": (
+            custom_retriever_all_sources,
+            "all_sources_info",
+            """Useful tool that contains general information about the field of AI.""",
+        ),
+    }
+
+    for source in selected_sources:
+        if source in source_mapping:
+            custom_retriever, name, description = source_mapping[source]
+            tools.append(
+                RetrieverTool(
+                    retriever=custom_retriever,
+                    metadata=ToolMetadata(
+                        name=name,
+                        description=description,
+                    ),
+                )
+            )
+
+    return tools
+
+
+def generate_completion(
+    query,
+    history,
+    sources,
+    model,
+    memory,
+):
+    llm = OpenAI(temperature=1, model=model, max_tokens=None)
+    client = llm._get_client()
+    logfire.instrument_openai(client)
+
+    with logfire.span(f"Running query: {query}"):
+        logfire.info(f"User chosen sources: {sources}")
+
+        memory_chat_list = memory.get()
+
+        if len(memory_chat_list) != 0:
+            user_index_memory = [
+                i
+                for i, msg in enumerate(memory_chat_list)
+                if msg.role == MessageRole.USER
+            ]
+
+            user_index_history = [
+                i for i, msg in enumerate(history) if msg["role"] == "user"
+            ]
+
+            if len(user_index_memory) > len(user_index_history):
+                logfire.warn(f"There are more user messages in memory than in history")
+                user_index_to_remove = user_index_memory[len(user_index_history)]
+                memory_chat_list = memory_chat_list[:user_index_to_remove]
+                memory.set(memory_chat_list)
+
+        logfire.info(f"chat_history: {len(memory.get())} {memory.get()}")
+        logfire.info(f"gradio_history: {len(history)} {history}")
+
+        query_engine_tools: list[RetrieverTool] = update_query_engine_tools(
+            ["All Sources"]
+        )
+
+        filter_list = []
+        source_mapping = {
+            "Transformers Docs": "transformers",
+            "PEFT Docs": "peft",
+            "TRL Docs": "trl",
+            "LlamaIndex Docs": "llama_index",
+            "LangChain Docs": "langchain",
+            "OpenAI Cookbooks": "openai_cookbooks",
+            "Towards AI Blog": "tai_blog",
+            "8 Hour Primer": "8-hour_primer",
+            "Advanced LLM Developer": "llm_developer",
+            "Python Primer": "python_primer",
+        }
+
+        for source in sources:
+            if source in source_mapping:
+                filter_list.append(
+                    MetadataFilter(
+                        key="source",
+                        operator=FilterOperator.EQ,
+                        value=source_mapping[source],
+                    )
+                )
+
+        filters = MetadataFilters(
+            filters=filter_list,
+            condition=FilterCondition.OR,
+        )
+        logfire.info(f"Filters: {filters}")
+        query_engine_tools[0].retriever._vector_retriever._filters = filters
+
+        # pdb.set_trace()
+
+        agent = OpenAIAgent.from_tools(
+            llm=llm,
+            memory=memory,
+            tools=query_engine_tools,
+            system_prompt=system_message_openai_agent,
+        )
+
+        completion = agent.stream_chat(query)
+
+    answer_str = ""
+    for token in completion.response_gen:
+        answer_str += token
+        yield answer_str
+
+    for answer_str in add_sources(answer_str, completion):
+        yield answer_str
+
+
+def add_sources(answer_str, completion):
+    if completion is None:
+        yield answer_str
+
+    formatted_sources = format_sources(completion)
+    if formatted_sources == "":
+        yield answer_str
+
+    if formatted_sources != "":
+        answer_str += "\n\n" + formatted_sources
+
+    yield answer_str
+
+
+def format_sources(completion) -> str:
+    if len(completion.sources) == 0:
+        return ""
+
+    # logfire.info(f"Formatting sources: {completion.sources}")
+
+    display_source_to_ui = {
+        src: ui for src, ui in zip(AVAILABLE_SOURCES, AVAILABLE_SOURCES_UI)
+    }
+
+    documents_answer_template: str = (
+        "📝 Here are the sources I used to answer your question:\n{documents}"
+    )
+    document_template: str = "[🔗 {source}: {title}]({url}), relevance: {score:2.2f}"
+    all_documents = []
+    for source in completion.sources:  # looping over list[ToolOutput]
+        if isinstance(source.raw_output, Exception):
+            logfire.error(f"Error in source output: {source.raw_output}")
+            # pdb.set_trace()
+            continue
+
+        if not isinstance(source.raw_output, list):
+            logfire.warn(f"Unexpected source output type: {type(source.raw_output)}")
+            continue
+        for src in source.raw_output:  # looping over list[NodeWithScore]
+            document = document_template.format(
+                title=src.metadata["title"],
+                score=src.score,
+                source=display_source_to_ui.get(
+                    src.metadata["source"], src.metadata["source"]
+                ),
+                url=src.metadata["url"],
+            )
+            all_documents.append(document)
+
+    if len(all_documents) == 0:
+        return ""
+    else:
+        documents = "\n".join(all_documents)
+        return documents_answer_template.format(documents=documents)
+
+
+def save_completion(completion, history):
+    pass
+
+
+def vote(data: gr.LikeData):
+    pass
+
+
+accordion = gr.Accordion(label="Customize Sources (Click to expand)", open=False)
+sources = gr.CheckboxGroup(
+    AVAILABLE_SOURCES_UI,
+    label="Sources",
+    value=[
+        "Advanced LLM Developer",
+        "8 Hour Primer",
+        "Python Primer",
+        "Towards AI Blog",
+        "Transformers Docs",
+        "PEFT Docs",
+        "TRL Docs",
+        "LlamaIndex Docs",
+        "LangChain Docs",
+        "OpenAI Cookbooks",
+    ],
+    interactive=True,
+)
+model = gr.Dropdown(
+    [
+        "gpt-4o-mini",
+        #Kenny added GPT2
+        #"gpt2",
+    ],
+    label="Model",
+    value="gpt-4o-mini",
+    interactive=False,
+)
+
+with gr.Blocks(
+    title="Towards AI 🤖",
+    analytics_enabled=True,
+    fill_height=True,
+) as demo:
+
+    memory = gr.State(
+        lambda: ChatSummaryMemoryBuffer.from_defaults(
+            token_limit=120000,
+        )
+    )
+    chatbot = gr.Chatbot(
+        type="messages",
+        scale=20,
+        placeholder="<strong>Towards AI 🤖: A Question-Answering Bot for anything AI-related</strong><br>",
+        show_label=False,
+        show_copy_button=True,
+    )
+    chatbot.like(vote, None, None)
+    gr.ChatInterface(
+        fn=generate_completion,
+        type="messages",
+        chatbot=chatbot,
+        additional_inputs=[sources, model, memory],
+        additional_inputs_accordion=accordion,
+        # fill_height=True,
+        # fill_width=True,
+        analytics_enabled=True,
+    )
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=CONCURRENCY_COUNT)
+    demo.launch(debug=False, share=False)

scripts/prompts.py

@@ -0,0 +1,29 @@
+# Prompt 5
+system_message_openai_agent = """You are an AI teacher, answering questions from students of an applied AI course on Large Language Models (LLMs or llm), Retrieval Augmented Generation (RAG) for LLMs and Python programming (computer science in general).
+
+Topics covered include coding with Python, training models, fine-tuning models, giving memory to LLMs, prompting tips, hallucinations and bias, vector databases, transformer architectures, embeddings, RAG frameworks such as Langchain and LlamaIndex, making LLMs interact with tools, AI agents, reinforcement learning with human feedback (RLHF). Questions should be understood in this context.
+
+Your answers are aimed to teach students, so they should be complete, clear, and easy to understand.
+
+Use the available tools to gather insights pertinent to the field of AI.
+
+To answer student questions, always use the all_sources_info tool.
+
+Only some information returned by the tools might be relevant to the question, so ignore the irrelevant part and answer the question with what you have.
+
+Your responses are exclusively based on the output provided by the tools. Refrain from incorporating information not directly obtained from the tool's responses.
+
+When the conversation deepens or shifts focus within a topic, adapt your input to the tools to reflect these nuances. This means if a user requests further elaboration on a specific aspect of a previously discussed topic, you should reformulate your input to the tool to capture this new angle or more profound layer of inquiry.
+
+Provide comprehensive answers, ideally structured in multiple paragraphs, drawing from the tool's variety of relevant details. The depth and breadth of your responses should align with the scope and specificity of the information retrieved.
+
+Should the tools repository lack information on the queried topic, politely inform the user that the question transcends the bounds of your current knowledge base, citing the absence of relevant content in the tool's documentation.
+
+At the end of your answers, always invite the students to ask deeper questions about the topic if they have any. Make sure reformulate the question to the tool to capture this new angle or more profound layer of inquiry.
+
+Do not refer to the documentation directly, but use the information provided within it to answer questions.
+
+If code is provided in the information, share it with the students. It's important to provide complete code blocks so they can execute the code when they copy and paste them.
+
+Make sure to format your answers in Markdown format, including code blocks and snippets.
+"""

scripts/setup.py

@@ -0,0 +1,146 @@
+import asyncio
+import json
+import logging
+import os
+import pickle
+
+import chromadb
+import logfire
+from custom_retriever import CustomRetriever
+from dotenv import load_dotenv
+from llama_index.core import Document, VectorStoreIndex
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.embeddings.cohere import CohereEmbedding
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from utils import init_mongo_db
+
+load_dotenv()
+
+logfire.configure()
+
+if not os.path.exists("data/chroma-db-all_sources"):
+    # Download the vector database from the Hugging Face Hub if it doesn't exist locally
+    # https://huggingface.co/datasets/towardsai-buster/ai-tutor-vector-db/tree/main
+    logfire.warn(
+        f"Vector database does not exist at 'data/chroma-db-all_sources', downloading from Hugging Face Hub"
+    )
+    from huggingface_hub import snapshot_download
+
+    snapshot_download(
+        repo_id="towardsai-tutors/ai-tutor-vector-db",
+        local_dir="data",
+        repo_type="dataset",
+    )
+    logfire.info(f"Downloaded vector database to 'data/chroma-db-all_sources'")
+
+
+def create_docs(input_file: str) -> list[Document]:
+    with open(input_file, "r") as f:
+        documents = []
+        for line in f:
+            data = json.loads(line)
+            documents.append(
+                Document(
+                    doc_id=data["doc_id"],
+                    text=data["content"],
+                    metadata={  # type: ignore
+                        "url": data["url"],
+                        "title": data["name"],
+                        "tokens": data["tokens"],
+                        "retrieve_doc": data["retrieve_doc"],
+                        "source": data["source"],
+                    },
+                    excluded_llm_metadata_keys=[
+                        "title",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                    excluded_embed_metadata_keys=[
+                        "url",
+                        "tokens",
+                        "retrieve_doc",
+                        "source",
+                    ],
+                )
+            )
+    return documents
+
+
+def setup_database(db_collection, dict_file_name) -> CustomRetriever:
+    db = chromadb.PersistentClient(path=f"data/{db_collection}")
+    chroma_collection = db.get_or_create_collection(db_collection)
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    embed_model = CohereEmbedding(
+        api_key=os.environ["COHERE_API_KEY"],
+        model_name="embed-english-v3.0",
+        input_type="search_query",
+    )
+
+    index = VectorStoreIndex.from_vector_store(
+        vector_store=vector_store,
+        transformations=[SentenceSplitter(chunk_size=800, chunk_overlap=0)],
+        show_progress=True,
+        # use_async=True,
+    )
+    vector_retriever = VectorIndexRetriever(
+        index=index,
+        similarity_top_k=15,
+        embed_model=embed_model,
+        # use_async=True,
+    )
+    with open(f"data/{db_collection}/{dict_file_name}", "rb") as f:
+        document_dict = pickle.load(f)
+
+    return CustomRetriever(vector_retriever, document_dict)
+
+
+custom_retriever_all_sources: CustomRetriever = setup_database(
+    "chroma-db-all_sources",
+    "document_dict_all_sources.pkl",
+)
+
+
+CONCURRENCY_COUNT = int(os.getenv("CONCURRENCY_COUNT", 64))
+MONGODB_URI = os.getenv("MONGODB_URI")
+
+AVAILABLE_SOURCES_UI = [
+    "Transformers Docs",
+    "PEFT Docs",
+    "TRL Docs",
+    "LlamaIndex Docs",
+    "LangChain Docs",
+    "OpenAI Cookbooks",
+    "Towards AI Blog",
+    "8 Hour Primer",
+    "Advanced LLM Developer",
+    "Python Primer",
+]
+
+AVAILABLE_SOURCES = [
+    "transformers",
+    "peft",
+    "trl",
+    "llama_index",
+    "langchain",
+    "openai_cookbooks",
+    "tai_blog",
+    "8-hour_primer",
+    "llm_developer",
+    "python_primer",
+]
+
+mongo_db = (
+    init_mongo_db(uri=MONGODB_URI, db_name="towardsai-buster")
+    if MONGODB_URI
+    else logfire.warn("No mongodb uri found, you will not be able to save data.")
+)
+
+__all__ = [
+    "custom_retriever_all_sources",
+    "mongo_db",
+    "CONCURRENCY_COUNT",
+    "AVAILABLE_SOURCES_UI",
+    "AVAILABLE_SOURCES",
+]

scripts/utils.py

@@ -0,0 +1,16 @@
+from pymongo.mongo_client import MongoClient
+from pymongo.server_api import ServerApi
+
+
+def init_mongo_db(uri: str, db_name: str):
+    """Initialize the mongodb database."""
+
+    try:
+        assert uri is not None, "No URI passed"
+        client = MongoClient(uri, server_api=ServerApi("1"))
+        database = client[db_name]
+        print("Connected to MongoDB")
+        return database
+    except Exception as e:
+        print("Something went wrong connecting to mongodb")
+        return