simonbb 0.0.8__tar.gz → 0.0.10__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: simonbb
-Version: 0.0.8
+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>
@@ -9,18 +9,17 @@ License-File: LICENSE
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.9
-Requires-Dist: gymnasium~=1.3.0
-Requires-Dist: matplotlib~=3.10.8
-Requires-Dist: numpy~=2.4.0
-Requires-Dist: opencv-python~=5.0.0.93
-Requires-Dist: pytorch-ignite~=0.5.4
-Requires-Dist: pyyaml~=6.0.3
-Requires-Dist: ray~=2.56.0
-Requires-Dist: sentence-transformers~=5.3.0
-Requires-Dist: stable-baselines3~=2.9.0
-Requires-Dist: torchvision~=0.24.1
-Requires-Dist: torch~=2.9.1
-Requires-Dist: transformers~=4.57.3
+Requires-Dist: gymnasium>=1.3.0
+Requires-Dist: matplotlib>=3.10.8
+Requires-Dist: numpy>=2.4.0
+Requires-Dist: opencv-python>=5.0.0.93
+Requires-Dist: pytorch-ignite>=0.5.4
+Requires-Dist: pyyaml>=6.0.3
+Requires-Dist: ray>=2.56.0
+Requires-Dist: stable-baselines3>=2.9.0
+Requires-Dist: torch>=2.9.1
+Requires-Dist: torchvision>=0.24.1
+Requires-Dist: transformers>=4.57.3
 Description-Content-Type: text/markdown
 
 # Background

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "simonbb"
-version = "0.0.8"
+version = "0.0.10"
 authors = [
 	{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
 ]
@@ -18,18 +18,17 @@ classifiers = [
 license = "MIT"
 license-files = ["LICEN[CS]E*"]
 dependencies = [
-	"matplotlib~=3.10.8",
-	"torch~=2.9.1",
-	"torchvision~=0.24.1",
-	"numpy~=2.4.0",
-	"transformers~=4.57.3",
-	"sentence-transformers~=5.3.0",
-	"gymnasium~=1.3.0",
-	"pytorch-ignite~=0.5.4",
-	"stable_baselines3~=2.9.0",
-	"ray~=2.56.0",
-	"opencv-python~=5.0.0.93",
-	"pyyaml~=6.0.3"
+	"matplotlib >= 3.10.8",
+	"torch >= 2.9.1",
+	"torchvision >=0.24.1",
+	"numpy >=2.4.0",
+	"transformers >= 4.57.3",
+	"gymnasium >= 1.3.0",
+	"pytorch-ignite >= 0.5.4",
+	"stable_baselines3 >= 2.9.0",
+	"ray >= 2.56.0",
+	"opencv-python >= 5.0.0.93",
+	"pyyaml >= 6.0.3"
 ]
 [project.urls]
 Homepage = "https://simonbb.com"

simonbb-0.0.10/src/simonbb/cis438/rag.py

@@ -0,0 +1,545 @@
+from __future__ import annotations
+
+import re
+from typing import Annotated, TypedDict, get_args, Any, Protocol, Sequence
+from dataclasses import dataclass
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+
+from pathlib import Path
+
+from functools import partial
+
+from langchain_core.documents import Document
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import Runnable, RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
+
+from langgraph.graph import END, START, StateGraph
+from langgraph.graph.message import add_messages
+
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLoader
+
+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
+from langchain_core.retrievers import BaseRetriever
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+#%%
+
+def load_documents(path: Path):
+    """Load documents from a file or directory."""    
+    valid_extensions = {'.pdf', '.html', '.ejs'}   
+    docs = []
+    success_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 ]:
+                print(file.name)                
+                if file.suffix == '.pdf':
+                    documents = PyPDFLoader(file).load()
+                elif file.suffix in {'.html', '.ejs'}:
+                    documents = UnstructuredHTMLLoader(file).load()
+                else:
+                    continue
+                success_cnt += 1
+                page_cnt += len(documents)
+                docs.extend(documents)
+        else:            
+            print(path.name)
+            if path.suffix == '.pdf':
+                documents = PyPDFLoader(path).load()
+            elif path.suffix in {'.html', '.ejs'}:
+                documents = UnstructuredHTMLLoader(path).load()
+            else:
+                print(f"Unsupported file type: {path}")
+                return None, 0, 0
+            success_cnt = 1
+            page_cnt = len(documents)
+            docs.extend(documents)
+    else:
+        print(f"Error: '{path}' is not a valid file or directory")
+        return None, 0, 0
+
+    if success_cnt == 0:
+        print("No valid documents found!")
+        return None, 0, 0
+
+    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."""
+    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)
+    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. 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."""
+        prompt = ChatPromptTemplate.from_template( "You are an AI language model assistant. Generate 5 different versions of the following question. Return exactly 5 alternative questions, one per line. Do not include explanations or numbering. Original question: {question}" )
+        chain = prompt | self.query_llm | StrOutputParser()
+        result = chain.invoke( {"question": query}, config={"callbacks": run_manager.get_child()} )
+        return [q.strip() for q in result.splitlines() if q.strip()]
+
+    @staticmethod
+    def _unique_union( document_lists: list[list[Document]] ) -> list[Document]:
+        """Flatten document lists and remove duplicates."""
+        unique_docs = {
+            doc.page_content: doc
+            for docs in document_lists
+            for doc in docs
+        }
+        return list(unique_docs.values())
+
+    def _rerank( self, query: str, docs: list[Document] ) -> list[Document]:
+        """Rerank retrieved documents using a cross-encoder."""
+        if not docs:
+            return []
+        pairs = [(query, doc.page_content) for doc in docs]
+        scores = self.reranker.predict(pairs)
+        scored_docs = sorted(
+            zip(scores, docs),
+            key=lambda x: float(x[0]),
+            reverse=True,
+        )
+        return [doc for _, doc in scored_docs[:self.top_n]]
+
+    def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> list[Document]:
+        """Retrieve documents using simple or multi-query retrieval."""
+        # Simple retrieval
+        if self.query_llm is None:
+            return self._base_retriever.invoke( query, config={"callbacks": run_manager.get_child()} )
+        # Generate alternative queries
+        alt_queries = self._generate_queries(query, run_manager)
+        all_queries = list( dict.fromkeys([query, *alt_queries]) ) # standard, most efficient way to remove duplicates from a list (or any iterable) while preserving the original insertion order
+        # Retrieve documents for all queries
+        doc_lists = self._base_retriever.batch(
+            all_queries,
+            config={"callbacks": run_manager.get_child()},
+        )
+        # 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_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]:
+    clean: list[str] = []
+    seen: set[str] = set()
+    for value in questions:
+        question = " ".join(str(value).split()).strip()
+        key = question.casefold()
+        if question and key not in seen:
+            seen.add(key)
+            clean.append(question)
+    return clean
+
+def _chroma_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+    """Keep Chroma-compatible scalar metadata and drop None/compound values."""
+    return {
+        str(key): value for key, value in metadata.items()
+        if value is not None and isinstance(value, (str, int, float, bool))
+    }
+
+
+class _QuestionList(BaseModel):
+    questions: list[str] = Field(description="Distinct questions answered by the text")
+
+class HypotheticalQuestionGenerator:
+    """LLM-backed question generator using provider-supported structured output."""
+
+    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
+        prompt = ChatPromptTemplate.from_messages(
+            [
+                ( "system", "Generate exactly {count} distinct, natural questions that can be answered from the supplied text. Do not answer them and do not use outside knowledge."),
+                ("human", "Text:\n{text}"),
+            ]
+        )
+        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} )
+        questions = result.questions if isinstance(result, _QuestionList) else result["questions"]
+        return _clean_questions(questions)[: self.questions_per_chunk]
+
+@dataclass(frozen=True)
+class IngestionResult:
+    chunks: int
+    hypothetical_questions: int
+    chunk_ids: tuple[str, ...]
+
+class HypotheticalQuestionRetriever(BaseRetriever):
+    """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()
+    
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    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_collection.delete( where={"source_dir": source_dir} )
+            self._chunk_collection.delete( where={"source_dir": source_dir} )
+
+        chunk_ids: list[str] = []
+        question_docs: list[Document] = []
+        question_ids: list[str] = []
+        chunk_metadatas: list[dict[str, Any]] = []
+        chunk_texts: list[str] = []
+
+        for index, doc in enumerate(input_docs):
+            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) )
+            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_dir": source_dir, "question_index": question_index },
+                    )
+                )
+        if question_docs: # LangChain's Chroma.add_documents() triggers embedding
+            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,
+                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_collection.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
+        chunk_ids: list[str] = []
+        seen: set[str] = set()
+        for m in matches:
+            chunk_id = str(m.metadata["chunk_id"])
+            if chunk_id not in seen:
+                seen.add(chunk_id)
+                chunk_ids.append(chunk_id)
+                if len(chunk_ids) == self.chunk_k:
+                    break
+        if not chunk_ids:
+            return []
+
+        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(
+                records["ids"], records["documents"], records["metadatas"], strict=True
+            )
+        }
+        return [by_id[chunk_id] for chunk_id in chunk_ids if chunk_id in by_id]
+
+
+
+#%% RAG with Routing
+
+domain_prompt = """You are a helpful assistant answering questions based on provided document snippets relating to {domain}.
+
+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 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
+    """
+    domains = get_args(domain_literals)[:-1] # exclude the placeholder
+    nodomain = get_args(domain_literals)[-1]
+    # use the functional form of TypedDict instead of the class form to accommodate Python 3.14
+    State = TypedDict(
+        "State",
+        {
+            "messages": Annotated[list, add_messages],
+            "query": str,
+            "domain": domain_literals,
+            "documents": list[Document],
+            "answer": str,
+        }
+    )
+
+    class Input(TypedDict):
+        query: str
+
+    Output = TypedDict( 
+        "Output",
+        {
+            "domain": domain_literals,
+            "documents": list[Document],
+            "answer": str
+        }
+    )
+
+    def router_node(state) -> State:
+        user_message = HumanMessage(state["query"])
+        messages = [ SystemMessage(router_prompt), *state["messages"], user_message ]
+        res = llm.invoke(messages)
+        return {
+            "domain": res.content.strip(),
+            "messages": [user_message, res], # update conversation history
+        }
+
+    def pick_retriever( state ):
+        for domain in domains:
+            if re.search( domain, state['domain'] ):
+                return domain
+        return 'generate_answer'
+    
+    def generate_answer(state) -> State:    
+        snippets = []    
+        for i, doc in enumerate(state.get('documents', []), start=1):
+            snippets.append(f"Snippet {i}: {doc.page_content}")       
+        context = "\n".join(snippets)
+        if state.get('domain', '')!=nodomain:
+            prompt = ChatPromptTemplate.from_messages( [ ("system", domain_prompt),  ("human", "{query}")  ])
+            messages = prompt.format_messages( domain=state.get('domain', ''), context=context, query=state["query"] )
+        else:
+            prompt = ChatPromptTemplate.from_messages( [ ("system", 'You are a helpful assistant.'),  ("human", "{query}")  ])
+            messages = prompt.format_messages( query=state["query"] )
+        res = llm.invoke(messages)
+        return {
+            "answer": res.content,
+            "messages": res,  # update conversation history
+        }
+    
+    def retriever_node( domain, state):        
+        query = re.sub(domain, 'this class', state["query"] ) # using the course code hurts RAG since the course code is rarely used in documents.
+        #print( f"domain={domain}; query={query}")
+        documents = retrievers[domain].invoke( query )
+        return { "documents": documents, 'query': query }
+    
+    builder = StateGraph(State, input_schema=Input, output_schema=Output)
+    builder.add_node("router", router_node)
+    builder.add_edge(START, "router")
+    builder.add_conditional_edges("router", pick_retriever)
+
+    retrievers = {}
+    retriever_nodes = {}
+    for domain in domains:    
+        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")
+        
+    builder.add_node("generate_answer", generate_answer)
+    builder.add_edge("generate_answer", END)
+    
+    return builder

simonbb-0.0.8/src/simonbb/cis438/rag.py

@@ -1,258 +0,0 @@
-# module dependencies: torch torchvision langchain_core langchain langgraph langchain_community langchain_chroma langchain_huggingface sentence_transformers
-# additional packages for class projects: langchain_ollama langchain_text_splitters unstructured streamlit_javascript
-
-import os
-import re
-from typing import Annotated, Literal, TypedDict, get_args
-from pathlib import Path
-
-from functools import partial
-
-import torch
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-from langchain_core.documents import Document
-from langchain_core.messages import HumanMessage, SystemMessage
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnableLambda
-
-from langgraph.graph import END, START, StateGraph
-from langgraph.graph.message import add_messages
-
-from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLoader
-from sentence_transformers import CrossEncoder
-
-embedding_model = HuggingFaceEmbeddings(
-    model_name="BAAI/bge-large-en-v1.5",
-    model_kwargs={'device': device},
-    encode_kwargs={'normalize_embeddings': False}
-)   
-
-#%%
-
-def load_documents(path: Path):
-    """Load documents from a file or directory."""    
-
-    valid_extensions = {'.pdf', '.html', '.ejs'}
-    
-    docs = []
-    success_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 ]:
-                print(file.name)                
-                if file.suffix == '.pdf':
-                    documents = PyPDFLoader(file).load()
-                elif file.suffix in {'.html', '.ejs'}:
-                    documents = UnstructuredHTMLLoader(file).load()
-                else:
-                    continue
-                success_cnt += 1
-                page_cnt += len(documents)
-                docs.extend(documents)
-        else:            
-            print(path.name)
-            if path.suffix == '.pdf':
-                documents = PyPDFLoader(path).load()
-            elif path.suffix in {'.html', '.ejs'}:
-                documents = UnstructuredHTMLLoader(path).load()
-            else:
-                print(f"Unsupported file type: {path}")
-                return None, 0, 0
-            success_cnt = 1
-            page_cnt = len(documents)
-            docs.extend(documents)
-    else:
-        print(f"Error: '{path}' is not a valid file or directory")
-        return None, 0, 0
-
-    if success_cnt == 0:
-        print("No valid documents found!")
-        return None, 0, 0
-
-    return docs, success_cnt, page_cnt
-
-reranker_model = CrossEncoder("BAAI/bge-reranker-large")
-def _getRetriever( knowledge_base, query_llm=None, k=3, vstore_path = f'{os.getenv("HOME")}/vstore' ):
-    store = Chroma( persist_directory=f'{vstore_path}/{knowledge_base}/BAAI_db/chroma_vector_store', embedding_function=embedding_model )
-    base_retriever = store.as_retriever(search_kwargs={"k": k}) # a vector store becomes a retriever through the as_retriever() method       
-    if query_llm==None:
-        return base_retriever
-
-    perspectives_prompt = ChatPromptTemplate.from_template(
-    """You are an AI language model assistant. Your task is to generate 5 different versions of the given original question
-    These should be in 5 lines and provide ONLY these alternative questions separated by newlines, no explaination is needed.
-    Original question: {question}""")
-    def parse_queries_output(message):
-        #print( message.content )
-        return message.content.split('\n')
-    
-    query_gen = perspectives_prompt | query_llm | parse_queries_output
-    
-    def get_unique_union(document_lists): # Flatten list of lists, and deduplicate them
-        deduped_docs = { doc.page_content: doc for sublist in document_lists for doc in sublist}
-        unique_docs = list(deduped_docs.values()) # a flat list of unique docs        
-        return unique_docs
-   
-    def rerank_docs(query, docs, top_n=3):
-        """Rerank documents using cross-encoder and return top_n."""
-        if not docs:
-            return []
-        pairs = [(query, doc.page_content) for doc in docs]
-        scores = reranker_model.predict(pairs)
-        scored_docs = sorted(zip(scores, docs), key=lambda x: float(x[0]), reverse=True)
-        return [doc for _, doc in scored_docs[:top_n]]
-    
-    def retrieve_and_rerank(query):
-        # Generate alternative queries + include original
-        alt_queries = query_gen.invoke({"question": query})
-        all_queries = [query] + alt_queries
-        # Retrieve chunks for each query
-        doc_lists = base_retriever.batch(all_queries)
-        unique_docs = get_unique_union(doc_lists)
-        # Rerank with cross-encoder and return top 3
-        return rerank_docs(query, unique_docs, top_n=5)
-
-    return RunnableLambda(retrieve_and_rerank)
-    
-
-"""
-RAG Chain
-"""
-
-def getRagModel(knowledge_base, llm, query_llm=None, k=3):
-    """
-    Create and return a RAG model that returns a dictionary with 'result' and 'source_documents'.
-    """
-    retriever = _getRetriever(knowledge_base, query_llm, k)
-
-    def _rag_chain(input_dict):
-        query = input_dict["query"]
-        docs = retriever.invoke(query)
-        snippets = []
-        for i, doc in enumerate(docs, start=1):            
-            snippets.append(f"Snippet {i}: {doc.page_content}")
-        context = "\n".join(snippets)        
-        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}"""
-        prompt = ChatPromptTemplate.from_messages( [ ("system", system_prompt),  ("human", "{query}")  ])
-        messages = prompt.format_messages( context=context, query=query )
-        answer = llm.invoke(messages).content
-        return {"answer": answer, "sources": docs}
-
-    return RunnableLambda(_rag_chain)
-
-
-"""
-multi-rag system with router
-"""
-
-domain_prompt = """You are a helpful assistant answering questions based on provided document snippets relating to {domain}.
-
-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 multiRagBuilder( domain_literals, router_prompt, chat_router, chat, query_llm=None, k=3):
-    """
-    domain_literals: the last one is a placeholder if the router finds no domain matches the query
-    """
-    domains = get_args(domain_literals)[:-1] # exclude the placeholder
-    nodomain = get_args(domain_literals)[-1]
-    
-    class State(TypedDict):    
-        messages: Annotated[list, add_messages] # annotated with the reducer function to append new messages to the existing list, rather than overwriting
-        query: str # input
-        domain: domain_literals # output
-        documents: list[Document]
-        answer: str
-
-    class Input(TypedDict):
-        query: str
-
-    class Output(TypedDict):
-        domain: domain_literals
-        documents: list[Document]
-        answer: str
-    
-    def router_node(state: State) -> State:
-        user_message = HumanMessage(state["query"])
-        messages = [ SystemMessage(router_prompt), *state["messages"], user_message ]
-        res = chat_router.invoke(messages)
-        return {
-            "domain": res.content.strip(),
-            "messages": [user_message, res], # update conversation history
-        }
-
-    def pick_retriever( state: State) -> Literal[*domains, "generate_answer"]:
-        for domain in domains:
-            if re.search( domain, state['domain'] ):
-                return domain
-        return 'generate_answer'
-    
-    def generate_answer(state: State) -> State:    
-        snippets = []    
-        for i, doc in enumerate(state.get('documents', []), start=1):
-            snippets.append(f"Snippet {i}: {doc.page_content}")       
-        context = "\n".join(snippets)
-        if state.get('domain', '')!=nodomain:
-            prompt = ChatPromptTemplate.from_messages( [ ("system", domain_prompt),  ("human", "{query}")  ])
-            messages = prompt.format_messages( domain=state.get('domain', ''), context=context, query=state["query"] )
-        else:
-            prompt = ChatPromptTemplate.from_messages( [ ("system", 'You are a helpful assistant.'),  ("human", "{query}")  ])
-            messages = prompt.format_messages( query=state["query"] )
-        res = chat.invoke(messages)
-        return {
-            "answer": res.content,
-            "messages": res,  # update conversation history
-        }
-    
-    def retriever_node( domain: domains, state: State):        
-        query = re.sub(domain, 'this class', state["query"] ) # using the course code hurts RAG since the course code is rarely used in documents.
-        #print( f"domain={domain}; query={query}")
-        documents = retrievers[domain].invoke( query )
-        return { "documents": documents, 'query': query }
-    
-    builder = StateGraph(State, input_schema=Input, output_schema=Output)
-    builder.add_node("router", router_node)
-    builder.add_edge(START, "router")
-    builder.add_conditional_edges("router", pick_retriever)
-
-    retrievers = {}
-    retriever_nodes = {}
-    for domain in domains:
-        retrievers[domain] = _getRetriever(domain, query_llm, 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")
-        
-    builder.add_node("generate_answer", generate_answer)
-    builder.add_edge("generate_answer", END)
-    
-    return builder