simonbb 0.0.9__tar.gz → 0.0.10__tar.gz

{simonbb-0.0.9 → simonbb-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: simonbb
-Version: 0.0.9
+Version: 0.0.10
 Summary: A small package to facilitate teaching AI courses at business schools
 Project-URL: Homepage, https://simonbb.com
 Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>

{simonbb-0.0.9 → simonbb-0.0.10}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "simonbb"
-version = "0.0.9"
+version = "0.0.10"
 authors = [
 	{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
 ]

{simonbb-0.0.9 → simonbb-0.0.10}/src/simonbb/cis438/rag.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 import re
 from typing import Annotated, TypedDict, get_args, Any, Protocol, Sequence
 from dataclasses import dataclass
@@ -25,6 +24,7 @@ from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLo
 from langchain_core.callbacks import CallbackManagerForRetrieverRun
 
 import hashlib
+from uuid import uuid4
 import chromadb
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import BaseChatModel
@@ -35,13 +35,10 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 
 def load_documents(path: Path):
     """Load documents from a file or directory."""    
-
-    valid_extensions = {'.pdf', '.html', '.ejs'}
-    
+    valid_extensions = {'.pdf', '.html', '.ejs'}   
     docs = []
     success_cnt = 0
-    page_cnt = 0
-    
+    page_cnt = 0    
     if path.exists():
         if path.is_dir():
             for file in [ item for item in path.rglob('*') if item.is_file() and item.suffix in valid_extensions ]:
@@ -77,25 +74,100 @@ def load_documents(path: Path):
 
     return docs, success_cnt, page_cnt
 
+#%% RAG Chains
+
+system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
+IMPORTANT INSTRUCTIONS:
+1. Base your answer ONLY on the provided context snippets below.
+2. Synthesize information from multiple snippets if needed.
+3. If the snippets contain partial information, provide the best answer you can with what's available.
+4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
+5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
+
+The snippets are provided in this format:
+Snippet 1: <text content>
+Snippet 2: <text content>
+...
+
+{context}"""
+
+def format_documents(documents: Sequence[Document]) -> str:
+    snippets = []
+    for i, doc in enumerate(documents, start=1):            
+        snippets.append(f"Snippet {i}: {doc.page_content}")
+    context = "\n".join(snippets)     
+    return context
+
+def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
+    """Return an LCEL chain accepting a question string and producing an answer string."""
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
+            ("human", "{query}"),
+        ]
+    )
+    inputs = {
+        "context": retriever | format_documents,
+        "query": RunnablePassthrough(),
+    }
+    return inputs | prompt | model | StrOutputParser()
+
+
+def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
+    """Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
+        -- 'sources': original list of Documents from retriever
+        -- 'query': original query
+        -- 'context': merged content from relevant documents
+        -- 'answer': final answer """
+    prompt = ChatPromptTemplate.from_messages([
+        ("system", system_prompt),
+        ("human", "{query}")
+    ])
+    return (
+        {
+            "sources": retriever,
+            "query": RunnablePassthrough(),
+        }
+        | RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
+        | RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
+    )
+
+build_rag = rag_main
+
+
+"""
+Retrievers are subclasses of LangChain's BaseRetriever which inherits from Pydantic's BaseModel.
+Unlike ordinary Python classes where attributes declared in the class body are class attributes shared by all instances,
+Pydantic model fields declared with type annotations are instance attributes, not shared class attributes, even though they are declared in the class body.
+"""
 #%% MultiQueryRerankRetriever
 
 class MultiQueryRerankRetriever(BaseRetriever):
     """Chroma retriever with optional multi-query retrieval and reranking."""
-    knowledge_base: str
     embeddings: Embeddings = Field(exclude=True)
     query_llm: BaseChatModel | None = Field(default=None, exclude=True)
     reranker: Any = Field(default=None, exclude=True)
     k: int = Field(default=3, ge=1)
     top_n: int = Field(default=5, ge=1)
-    vstore_path: str = Field( default_factory=lambda: f"{os.path.expanduser('~')}/vstore" )
+    persist_directory: str | Path
+    collection_name: str
+    create_collection_if_not_exists: bool = False
+    store: Chroma = None
     _base_retriever: BaseRetriever = PrivateAttr()
 
+    """ model_post_init() is a special method in Pydantic v2 that runs automatically after a model instance has been initialized and its fields have been validated."""
     def model_post_init(self, __context: Any) -> None:
-        """Initialize the underlying Chroma retriever."""
-        store = Chroma( persist_directory=( f"{self.vstore_path}/{self.knowledge_base}/BAAI_db/chroma_vector_store"),  embedding_function=self.embeddings )
-        self._base_retriever = store.as_retriever( search_kwargs={"k": self.k} )
+        """Initialize the underlying Chroma retriever. Use the last folder name as collection name"""
+        self.store = Chroma( persist_directory=self.persist_directory, embedding_function=self.embeddings, collection_name=self.collection_name, create_collection_if_not_exists=self.create_collection_if_not_exists )
+        self._base_retriever = self.store.as_retriever( search_kwargs={"k": self.k} )
         if self.query_llm is not None and self.reranker is None:
             raise ValueError( "A reranker is required when query_llm is provided." )
+    
+    def ingest( self, documents: Sequence[Document], *,  source_dir: str = 'placeholder', chunk_size: int = 1000, chunk_overlap: int = 150 ):        
+        text_splitter = RecursiveCharacterTextSplitter( separators=["\n\n", "\n", ". ", " ", ""], chunk_size=chunk_size, chunk_overlap=chunk_overlap )
+        chunks = text_splitter.split_documents( documents )
+        self.store.add_documents( chunks )
+    
     # run_manager manages callbacks, tracing, and monitoring during a retriever's execution, allowing integrations with LangSmith
     def _generate_queries( self, query: str, run_manager: CallbackManagerForRetrieverRun ) -> list[str]:
         """Generate five alternative formulations of the query."""
@@ -143,12 +215,90 @@ class MultiQueryRerankRetriever(BaseRetriever):
         # Deduplicate and rerank
         unique_docs = self._unique_union(doc_lists)
         return self._rerank(query, unique_docs)
+
+#%% 
+
+class DummyEmbeddings(Embeddings):
+    """Return constant 1-dimensional embeddings without using an LLM."""
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        return [[0.0] for _ in texts]
+    def embed_query(self, text: str) -> list[float]:
+        return [0.0]
+    
+class ParentDocumentRetriever(BaseRetriever):
+    persist_directory: str | Path
+    parent_collection: str
+    child_collection: str
+    embeddings: Embeddings
+    k: int = 5
+    _parent_store: Chroma = PrivateAttr()
+    _child_store: Chroma = PrivateAttr()
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._parent_store = Chroma(
+            embedding_function=DummyEmbeddings(),
+            collection_name=self.parent_collection,
+            persist_directory=self.persist_directory,
+        )
+
+        self._child_store = Chroma(
+            embedding_function=self.embeddings,
+            collection_name=self.child_collection,
+            persist_directory=self.persist_directory,
+        )
+        
+    def ingest(self, docs: list[Document], *, source_dir: str = 'placeholder', parent_chunk_size: int = 600, parent_chunk_overlap: int = 100, child_chunk_size: int = 200, child_chunk_overlap: int = 50) -> None:
+        """Split documents and store parent and child chunks."""
+        parent_splitter = RecursiveCharacterTextSplitter( chunk_size=parent_chunk_size, chunk_overlap=parent_chunk_overlap )
+        child_splitter = RecursiveCharacterTextSplitter( chunk_size=child_chunk_size, chunk_overlap=child_chunk_overlap )
+
+        parent_docs = []
+        parent_ids = []
+        child_docs = []
+
+        for parent in parent_splitter.split_documents(docs):
+            parent_id = str(uuid4())
+            parent_docs.append(parent)
+            parent_ids.append(parent_id)
+            for child in child_splitter.split_documents( [parent] ):
+                child.metadata["parent_id"] = parent_id
+                child_docs.append(child)
+        if parent_docs:
+            self._parent_store.add_documents( documents=parent_docs, ids=parent_ids )
+        if child_docs:
+            self._child_store.add_documents( documents=child_docs )
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+    ) -> list[Document]:
+        """Retrieve parent documents using child similarity search."""
+
+        # 1. Search child chunks
+        child_chunks = self._child_store.similarity_search( query, k=self.k )
+        # 2. Extract unique parent IDs in relevance order
+        parent_ids = list(dict.fromkeys(  # standard trick to remove duplicates from a list or iterable while preserving the original item orde
+            chunk.metadata["parent_id"] for chunk in child_chunks if "parent_id" in chunk.metadata 
+        ))
+        if not parent_ids:
+            return []
+        # 3. Fetch parent documents
+        parents = self._parent_store.get_by_ids( parent_ids )        
+        parents_by_id = { doc.id: doc for doc in parents if doc.id is not None }
+        # 4. Restore relevance order
+        return [
+            parents_by_id[pid]
+            for pid in parent_ids if pid in parents_by_id
+        ]
     
 
 #%% Hypothetical Question Retriever
 
-def _stable_id(source_id: str, index: int, text: str) -> str:
-    digest = hashlib.sha256(f"{source_id}\0{index}\0{text}".encode()).hexdigest()
+def _stable_id(source_dir: str, index: int, text: str) -> str:
+    digest = hashlib.sha256(f"{source_dir}\0{index}\0{text}".encode()).hexdigest()
     return f"chunk-{digest[:32]}"
 
 def _clean_questions(questions: Sequence[str]) -> list[str]:
@@ -169,9 +319,6 @@ def _chroma_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
         if value is not None and isinstance(value, (str, int, float, bool))
     }
 
-class QuestionGenerator(Protocol):
-    """Anything capable of generating search questions for a chunk."""
-    def generate(self, document: Document) -> list[str]: ...
 
 class _QuestionList(BaseModel):
     questions: list[str] = Field(description="Distinct questions answered by the text")
@@ -179,7 +326,7 @@ class _QuestionList(BaseModel):
 class HypotheticalQuestionGenerator:
     """LLM-backed question generator using provider-supported structured output."""
 
-    def __init__(self, model: BaseChatModel, *, questions_per_chunk: int = 4) -> None:
+    def __init__(self, llm: BaseChatModel, *, questions_per_chunk: int = 4) -> None:
         if questions_per_chunk < 1:
             raise ValueError("questions_per_chunk must be positive")
         self.questions_per_chunk = questions_per_chunk
@@ -189,7 +336,7 @@ class HypotheticalQuestionGenerator:
                 ("human", "Text:\n{text}"),
             ]
         )
-        self._chain = prompt | model.with_structured_output(_QuestionList)
+        self._chain = prompt | llm.with_structured_output(_QuestionList)
 
     def generate(self, document: Document) -> list[str]:
         result = self._chain.invoke( {"count": self.questions_per_chunk, "text": document.page_content} )
@@ -203,38 +350,40 @@ class IngestionResult:
     chunk_ids: tuple[str, ...]
 
 class HypotheticalQuestionRetriever(BaseRetriever):
-    """Search question vectors, then resolve their chunk IDs from the chunk collection."""
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    question_store: Chroma
-    chunk_collection: Any
-    generator: Any
+    """Search question vectors, then resolve their chunk IDs from the chunk collection."""    
+    llm: Any = None
+    persist_directory: str | Path
+    embeddings: Embeddings      
+    question_collection: str
+    chunk_collection: str
     question_k: int = 12
     chunk_k: int = 4
     metadata_filter: dict[str, Any] | None = None
+    create_collection_if_not_exists: bool = False
+    _question_store: Chroma = PrivateAttr()
+    _chunk_collection: Any = PrivateAttr()
+    _generator: Any = PrivateAttr()
     
-    def __init__(self, *, persist_directory: str | Path, embeddings: Embeddings, generator: QuestionGenerator, question_collection: str = "hypothetical_questions", chunk_collection: str = "original_chunks", question_k: int = 12, chunk_k: int = 4, metadata_filter: dict[str, Any] | None = None, **kwargs: Any) -> None:
-        client = chromadb.PersistentClient(path=str(persist_directory))
-        question_store = Chroma( client=client, collection_name=question_collection, embedding_function=embeddings )
-        chunks = client.get_or_create_collection( name=chunk_collection, embedding_function=None )
-        super().__init__(
-            question_store=question_store,
-            chunk_collection=chunks,
-            generator=generator,
-            question_k=question_k,
-            chunk_k=chunk_k,
-            metadata_filter=metadata_filter,
-            **kwargs,
-        )     
+    model_config = ConfigDict(arbitrary_types_allowed=True)
     
-    def ingest( self, documents: Sequence[Document], *, source_id: str, split: bool = True, chunk_size: int = 1000, chunk_overlap: int = 150, replace_source: bool = True ) -> IngestionResult:
-        if not source_id: # source_id is a unique identifier for the source document or documents being ingested
-            raise ValueError("source_id must not be empty")
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        client = chromadb.PersistentClient( path= self.persist_directory )
+        self._question_collection = Chroma( client=client, collection_name=self.question_collection, embedding_function=self.embeddings, create_collection_if_not_exists=self.create_collection_if_not_exists )
+        self._chunk_collection = client.get_or_create_collection( name=self.chunk_collection, embedding_function=None )        
+        
+    def ingest( self, documents: Sequence[Document], *, source_dir: str, split: bool = True, chunk_size: int = 500, chunk_overlap: int = 150, replace_source: bool = True ) -> IngestionResult:
+        if not self.llm:
+            raise ValueError("To ingest, llm must be provided in the constructor.")
+        self._generator = HypotheticalQuestionGenerator( self.llm )
+        if not source_dir: # source_dir is a unique identifier for the source document or documents being ingested
+            raise ValueError("source_dir must not be empty")
         splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True )
         input_docs = splitter.split_documents( list(documents) ) if split else list(documents)
         
         if replace_source:
-            self.question_store.delete(where={"source_id": source_id})
-            self.chunk_collection.delete(where={"source_id": source_id})
+            self._question_collection.delete( where={"source_dir": source_dir} )
+            self._chunk_collection.delete( where={"source_dir": source_dir} )
 
         chunk_ids: list[str] = []
         question_docs: list[Document] = []
@@ -243,32 +392,31 @@ class HypotheticalQuestionRetriever(BaseRetriever):
         chunk_texts: list[str] = []
 
         for index, doc in enumerate(input_docs):
-            chunk_id = _stable_id(source_id, index, doc.page_content)
-            metadata = _chroma_metadata( {**doc.metadata, "source_id": source_id, "chunk_id": chunk_id} )
+            chunk_id = _stable_id(source_dir, index, doc.page_content)
+            metadata = _chroma_metadata( {**doc.metadata, "source_dir": source_dir, "chunk_id": chunk_id} )
             chunk_ids.append( chunk_id )
             chunk_metadatas.append( metadata )
             chunk_texts.append( doc.page_content )
-
-            questions = _clean_questions( self.generator.generate(doc) )
+            questions = _clean_questions( self._generator.generate(doc) )
             for question_index, question in enumerate(questions):
                 question_ids.append( f"{chunk_id}:q:{question_index}" )
                 question_docs.append(
                     Document( page_content=question,
-                        metadata={"chunk_id": chunk_id, "source_id": source_id, "question_index": question_index },
+                        metadata={"chunk_id": chunk_id, "source_dir": source_dir, "question_index": question_index },
                     )
                 )
         if question_docs: # LangChain's Chroma.add_documents() triggers embedding
-            self.question_store.add_documents( question_docs, ids=question_ids )
+            self._question_collection.add_documents( question_docs, ids=question_ids )
             
         if chunk_ids: # Chroma's underlying collection API stores text and metadata without embeddings.
-            self.chunk_collection.upsert( ids=chunk_ids, documents=chunk_texts, metadatas=chunk_metadatas,
+            self._chunk_collection.upsert( ids=chunk_ids, documents=chunk_texts, metadatas=chunk_metadatas,
                 embeddings=[[0.0] for _ in chunk_ids] ) # placeholder vectors since Chroma records require vectors
             
         return IngestionResult(len(chunk_ids), len(question_docs), tuple(chunk_ids))
 
     # override _get_relevant_documents() method from BaseRetriever
     def _get_relevant_documents(self, query: str, *, run_manager: Any = None) -> list[Document]:
-        matches = self.question_store.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
+        matches = self._question_collection.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
         chunk_ids: list[str] = []
         seen: set[str] = set()
         for m in matches:
@@ -281,7 +429,7 @@ class HypotheticalQuestionRetriever(BaseRetriever):
         if not chunk_ids:
             return []
 
-        records = self.chunk_collection.get(ids=chunk_ids, include=["documents", "metadatas"])
+        records = self._chunk_collection.get(ids=chunk_ids, include=["documents", "metadatas"])
         by_id = {
             record_id: Document(page_content=text, metadata=metadata or {})
             for record_id, text, metadata in zip(
@@ -290,65 +438,7 @@ class HypotheticalQuestionRetriever(BaseRetriever):
         }
         return [by_id[chunk_id] for chunk_id in chunk_ids if chunk_id in by_id]
 
-#%% RAG Chains
-
-system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
-IMPORTANT INSTRUCTIONS:
-1. Base your answer ONLY on the provided context snippets below.
-2. Synthesize information from multiple snippets if needed.
-3. If the snippets contain partial information, provide the best answer you can with what's available.
-4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
-5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
-
-The snippets are provided in this format:
-Snippet 1: <text content>
-Snippet 2: <text content>
-...
-
-{context}"""
-
-def format_documents(documents: Sequence[Document]) -> str:
-    snippets = []
-    for i, doc in enumerate(documents, start=1):            
-        snippets.append(f"Snippet {i}: {doc.page_content}")
-    context = "\n".join(snippets)     
-    return context
-
-def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
-    """Return an LCEL chain accepting a question string and producing an answer string."""
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
-            ("human", "{query}"),
-        ]
-    )
-    inputs = {
-        "context": retriever | format_documents,
-        "query": RunnablePassthrough(),
-    }
-    return inputs | prompt | model | StrOutputParser()
-
-
-def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
-    """Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
-        -- 'sources': original list of Documents from retriever
-        -- 'query': original query
-        -- 'context': merged content from relevant documents
-        -- 'answer': final answer """
-    prompt = ChatPromptTemplate.from_messages([
-        ("system", system_prompt),
-        ("human", "{query}")
-    ])
-    return (
-        {
-            "sources": retriever,
-            "query": RunnablePassthrough(),
-        }
-        | RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
-        | RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
-    )
 
-build_rag = rag_main
 
 #%% RAG with Routing
 
@@ -368,7 +458,7 @@ Snippet 2: <text content>
 
 {context}"""
 
-def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
+def multiRagBuilder( persist_directory, domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
     """
     domain_literals: the last one is a placeholder if the router finds no domain matches the query
     """
@@ -444,7 +534,7 @@ def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker,
     retrievers = {}
     retriever_nodes = {}
     for domain in domains:    
-        retrievers[domain] = MultiQueryRerankRetriever(knowledge_base=domain, embeddings=embeddings, query_llm=query_llm, reranker=reranker, k=k)
+        retrievers[domain] = MultiQueryRerankRetriever( persist_directory=persist_directory, collection_name=domain, embeddings=embeddings, query_llm=query_llm, reranker=reranker, k=k)
         retriever_nodes[domain] = partial( retriever_node, domain ) # the 1st argument is the partial function. the rest are frozen arguments passed positionally to fill the leftmost parameters of the original function.
         builder.add_node(domain, retriever_nodes[domain])
         builder.add_edge(domain, "generate_answer")