Async implementation of vector store population

app.py

@@ -1,17 +1,22 @@
+import asyncio
 from typing import List
 import chainlit as cl
 import json
+import os
 
+from dotenv import load_dotenv
 from langchain_experimental.text_splitter import SemanticChunker
+from langchain_openai import ChatOpenAI
 from langchain_openai.embeddings import OpenAIEmbeddings
 from langchain_core.documents import Document
 
 from langchain_qdrant import QdrantVectorStore
+from pstuts_rag.rag import RAGChainFactory, RetrieverFactory
 from qdrant_client import QdrantClient
 from qdrant_client.http.models import Distance, VectorParams
 from dataclasses import dataclass
 
-import pstuts_rag.datastore
+import pstuts_rag.rag
 
 
 @dataclass
@@ -19,6 +24,14 @@ class ApplicationParameters:
     filename = "data/test.json"
     embedding_model = "text-embedding-3-small"
     n_context_docs = 2
+    llm_model = "gpt-4.1-mini"
+
+
+def set_api_key_if_not_present(key_name, prompt_message=""):
+    if len(prompt_message) == 0:
+        prompt_message = key_name
+    if key_name not in os.environ or not os.environ[key_name]:
+        os.environ[key_name] = getpass.getpass(prompt_message)
 
 
 class ApplicationState:
@@ -26,48 +39,42 @@ class ApplicationState:
     docs: List[Document] = []
     qdrantclient: QdrantClient = None
     vectorstore: QdrantVectorStore = None
-    retriever = None
+    retriever_factory: pstuts_rag.rag.RetrieverFactory
+    rag_factory: pstuts_rag.rag.RAGChainFactory
+
+    def __init__(self) -> None:
+        load_dotenv()
+        set_api_key_if_not_present("OPENAI_API_KEY")
 
 
 state = ApplicationState()
+params = ApplicationParameters()
 
 
 @cl.on_chat_start
 async def on_chat_start():
-    params = ApplicationParameters()
-
-    await cl.Message(content=f"Loading file {params.filename}").send()
-    data = json.load(open(params.filename, "rb"))
+    state.client = QdrantClient(":memory:")
 
-    state.embeddings = OpenAIEmbeddings(model=params.embedding_model)
-    state.docs = pstuts_rag.datastore.transcripts_load(data, state.embeddings)
-    await cl.Message(
-        content=f"Loaded {len(state.docs)} chunks from file {params.filename}."
-    ).send()
-
-    state.qdrantclient = QdrantClient(":memory:")
-
-    state.vectorstore = pstuts_rag.datastore.initialize_vectorstore(
-        client=state.qdrantclient,
-        collection_name=f"{params.filename}_qdrant",
-        embeddings=state.embeddings,
+    state.retriever_factory = pstuts_rag.rag.RetrieverFactory(
+        qdrant_client=state.client, name="local_test"
     )
-
-    _ = state.vectorstore.add_documents(documents=state.docs)
-    state.retriever = state.vectorstore.as_retriever(
-        search_kwargs={"k": params.n_context_docs}
+    if state.retriever_factory.count_docs() == 0:
+        data: List[Dict[str, Any]] = json.load(open(params.filename, "rb"))
+        asyncio.run(main=state.retriever_factory.aadd_docs(raw_docs=data))
+    state.rag_factory = pstuts_rag.rag.RAGChainFactory(
+        retriever=state.retriever_factory.get_retriever()
     )
-
-    await cl.Message(content=f"Populated vector database.").send()
+    state.llm = ChatOpenAI(model=params.llm_model, temperature=0)
+    state.rag_chain = state.rag_factory.get_rag_chain(state.llm)
 
 
 @cl.on_message
 async def main(message: cl.Message):
     # Send a response back to the user
 
-    v = await state.retriever.ainvoke(message.content)
+    v = await state.rag_chain.ainvoke(message.content)
 
-    await cl.Message(content=f"Hello! {len(v)}").send()
+    await cl.Message(content=v.content).send()
 
 
 if __name__ == "__main__":

pstuts_rag/pstuts_rag/datastore.py

@@ -1,4 +1,6 @@
+import asyncio
 from typing import List, Dict, Iterator, Any
+import uuid
 
 
 from langchain_experimental.text_splitter import SemanticChunker
@@ -7,14 +9,22 @@ from langchain_core.documents import Document
 
 from .loader import VideoTranscriptBulkLoader, VideoTranscriptLoader
 
+from langchain_core.vectorstores import VectorStoreRetriever
+
 from langchain_qdrant import QdrantVectorStore
 from qdrant_client import QdrantClient
 from qdrant_client.http.models import Distance, VectorParams
+from qdrant_client.models import VectorParams, Distance, PointStruct
+
 
+def batch(iterable: List[Any], size: int = 16) -> Iterator[List[Any]]:
+    for i in range(0, len(iterable), size):
+        yield iterable[i : i + size]
 
-def transcripts_load(
+
+async def chunk_transcripts(
     json_transcripts: List[Dict[str, Any]],
-    embeddings: OpenAIEmbeddings = OpenAIEmbeddings(
+    semantic_chunker_embedding_model: OpenAIEmbeddings = OpenAIEmbeddings(
         model="text-embedding-3-small"
     ),
 ) -> List[Document]:
@@ -40,12 +50,21 @@ def transcripts_load(
         json_payload=json_transcripts
     ).load()
 
-    text_splitter = SemanticChunker(embeddings)
-
-    docs_chunks_semantic: List[Document] = text_splitter.split_documents(
-        docs_full_transcript
+    # semantically split the combined transcript
+    text_splitter = SemanticChunker(semantic_chunker_embedding_model)
+    docs_group = await asyncio.gather(
+        *[
+            text_splitter.atransform_documents(d)
+            for d in batch(docs_full_transcript, size=2)
+        ]
     )
+    # Flatten the nested list of documents
+    docs_chunks_semantic: List[Document] = []
+    for group in docs_group:
+        docs_chunks_semantic.extend(group)
 
+    # locate individual sections of the original transcript
+    # with the semantic chunks
     def is_subchunk(a: Document, ofb: Document) -> bool:
         return (a.metadata["video_id"] == ofb.metadata["video_id"]) and (
             a.page_content in ofb.page_content
@@ -83,35 +102,129 @@ def transcripts_load(
         else:
             chunk.metadata["start"], chunk.metadata["stop"] = None, None
 
+    docs_chunks_semantic[0].metadata.keys()
     return docs_chunks_semantic
 
 
-def initialize_vectorstore(
-    client: QdrantClient, collection_name: str, embeddings: OpenAIEmbeddings
-) -> QdrantVectorStore:
+class DatastoreManager:
+    """Factory class for creating and managing vector store retrievers.
+
+    This class simplifies the process of creating, populating, and managing
+    Qdrant vector stores for document retrieval.
+
+    Attributes:
+        embeddings: OpenAI embeddings model for document vectorization
+        docs: List of documents stored in the vector store
+        qdrant_client: Client for Qdrant vector database
+        name: Unique identifier for this retriever instance
+        vector_store: The Qdrant vector store instance
     """
-    Initialize a Qdrant vector store with the given client and collection name.
 
-    This function creates a new collection in Qdrant and initializes a vector
-    store with the specified embeddings model. The collection is configured
-    with appropriate vector parameters for the embedding model.
+    embeddings: OpenAIEmbeddings
+    docs: List[Document]
+    qdrant_client: QdrantClient
+    name: str
+    vector_store: QdrantVectorStore
+
+    def __init__(
+        self,
+        embeddings: OpenAIEmbeddings = OpenAIEmbeddings(
+            model="text-embedding-3-small"
+        ),
+        qdrant_client: QdrantClient = QdrantClient(location=":memory:"),
+        name: str = str(object=uuid.uuid4()),
+    ) -> None:
+        """Initialize the RetrieverFactory.
+
+        Args:
+            embeddings: OpenAI embeddings model to use
+            qdrant_client: Qdrant client for vector database operations
+            name: Unique identifier for this retriever instance
+        """
+        self.embeddings = embeddings
+        self.name = name
+        self.qdrant_client = qdrant_client
+
+        self.qdrant_client.recreate_collection(
+            collection_name=self.name,
+            vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
+        )
 
-    Args:
-        client: QdrantClient instance to use for connecting to the database
-        collection_name: Name to use for the new collection
-        embeddings: OpenAI embeddings model to use for vector encoding
+        # wrapper around the client
+        self.vector_store = QdrantVectorStore(
+            client=self.qdrant_client,
+            collection_name=self.name,
+            embedding=embeddings,
+        )
 
-    Returns:
-        Initialized QdrantVectorStore instance ready for document storage
-    """
-    client.create_collection(
-        collection_name=collection_name,
-        vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
-    )
+        self.docs = []
 
-    vector_store = QdrantVectorStore(
-        client=client,
-        collection_name=collection_name,
-        embedding=embeddings,
-    )
-    return vector_store
+    async def populate_database(self, raw_docs: List[Dict[str, Any]]):
+
+        # perform chunking
+        self.docs: List[Document] = await chunk_transcripts(
+            json_transcripts=raw_docs,
+            semantic_chunker_embedding_model=self.embeddings,
+        )
+
+        # perform embedding
+
+        vector_batches = await asyncio.gather(
+            *[
+                self.embeddings.aembed_documents(
+                    [c.page_content for c in chunk_batch]
+                )
+                for chunk_batch in batch(self.docs, 8)
+            ]
+        )
+        vectors = []
+        for vb in vector_batches:
+            vectors.extend(vb)
+        ids = list(range(len(vectors)))
+
+        points = [
+            PointStruct(
+                id=id,
+                vector=vector,
+                payload={
+                    "page_content": doc.page_content,
+                    "metadata": doc.metadata,
+                },
+            )
+            for id, vector, doc in zip(ids, vectors, self.docs)
+        ]
+
+        # upload qdrant payload
+        self.qdrant_client.upload_points(
+            collection_name=self.name,
+            points=points,
+        )
+
+    def count_docs(self) -> int:
+        try:
+            count = self.qdrant_client.get_collection(self.name).points_count
+            return count if count else 0
+        except ValueError:
+            return 0
+
+    def clear(self) -> bool:
+        """Clear all documents from the vector store.
+
+        Returns:
+            bool: True if deletion was successful, False otherwise
+        """
+        self.docs = []
+        return True if self.vector_store.delete() else False
+
+    def get_retriever(self, n_context_docs: int = 2) -> VectorStoreRetriever:
+        """Get a retriever for the vector store.
+
+        Args:
+            n_context_docs: Number of documents to retrieve for each query
+
+        Returns:
+            VectorStoreRetriever: The configured retriever
+        """
+        return self.vector_store.as_retriever(
+            search_kwargs={"k": n_context_docs}
+        )

pstuts_rag/pstuts_rag/rag.py

@@ -6,6 +6,7 @@ This module provides the core RAG functionality, including:
 """
 
 import json
+from multiprocessing import Value
 import uuid
 from operator import itemgetter
 from typing import Dict, List, Any
@@ -25,96 +26,12 @@ from langchain.prompts import ChatPromptTemplate
 from langchain_core.vectorstores import VectorStoreRetriever
 from langchain_openai import ChatOpenAI
 
-from .datastore import initialize_vectorstore, transcripts_load
 from .prompt_templates import RAG_PROMPT_TEMPLATES
 
 from langchain_core.language_models.base import BaseLanguageModel
 from langchain_core.messages import AIMessage
 
 
-class RetrieverFactory:
-    """Factory class for creating and managing vector store retrievers.
-
-    This class simplifies the process of creating, populating, and managing
-    Qdrant vector stores for document retrieval.
-
-    Attributes:
-        embeddings: OpenAI embeddings model for document vectorization
-        docs: List of documents stored in the vector store
-        qdrant_client: Client for Qdrant vector database
-        name: Unique identifier for this retriever instance
-        vector_store: The Qdrant vector store instance
-    """
-
-    embeddings: OpenAIEmbeddings
-    docs: List[Document]
-    qdrant_client: QdrantClient
-    name: str
-    vector_store: QdrantVectorStore
-
-    def __init__(
-        self,
-        embeddings: OpenAIEmbeddings = OpenAIEmbeddings(
-            model="text-embedding-3-small"
-        ),
-        qdrant_client: QdrantClient = QdrantClient(location=":memory:"),
-        name: str = str(object=uuid.uuid4()),
-    ) -> None:
-        """Initialize the RetrieverFactory.
-
-        Args:
-            embeddings: OpenAI embeddings model to use
-            qdrant_client: Qdrant client for vector database operations
-            name: Unique identifier for this retriever instance
-        """
-        self.embeddings = embeddings
-        self.name = name
-        self.qdrant_client = qdrant_client
-        self.vector_store = initialize_vectorstore(
-            client=self.qdrant_client,
-            collection_name=f"{self.name}_qdrant",
-            embeddings=self.embeddings,
-        )
-        self.docs = []
-
-    def add_docs(self, raw_docs: List[Dict[str, Any]]) -> None:
-        """Add documents to the vector store.
-
-        Takes raw document data, converts it to Document objects,
-        and adds them to the vector store.
-
-        Args:
-            raw_docs: List of raw document dictionaries
-        """
-        docs: List[Document] = transcripts_load(
-            json_transcripts=raw_docs, embeddings=self.embeddings
-        )
-        self.docs.extend(docs)
-        _ = self.vector_store.add_documents(documents=docs)
-
-    def clear(self) -> bool:
-        """Clear all documents from the vector store.
-
-        Returns:
-            bool: True if deletion was successful, False otherwise
-        """
-        self.docs = []
-        return True if self.vector_store.delete() else False
-
-    def get_retriever(self, n_context_docs: int = 2) -> VectorStoreRetriever:
-        """Get a retriever for the vector store.
-
-        Args:
-            n_context_docs: Number of documents to retrieve for each query
-
-        Returns:
-            VectorStoreRetriever: The configured retriever
-        """
-        return self.vector_store.as_retriever(
-            search_kwargs={"k": n_context_docs}
-        )
-
-
 class RAGChainFactory:
     """Factory class for creating RAG (Retrieval Augmented Generation) chains.
 

pyproject.toml

@@ -68,14 +68,16 @@ known-first-party = ["src"]
 line-length = 79
 target-version = ["py313"]
 
+# [tool.mypy]
+# python_version = "3.13"
+# warn_return_any = true
+# warn_unused_configs = true
+# disallow_untyped_defs = true
+# mypy_path           = ["pstuts_rag/pstuts_rag"]
+# namespace_packages  = true
+# explicit_package_bases = true
 [tool.mypy]
-python_version = "3.13"
-warn_return_any = true
-warn_unused_configs = true
-disallow_untyped_defs = true
-mypy_path           = ["pstuts_rag/pstuts_rag"]
-namespace_packages  = true
-explicit_package_bases = true
+ignore_errors = true
 
 [tool.flake8]
 application-import-names = "pstuts_rag"