simonbb 0.0.9__tar.gz → 0.0.11__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: simonbb
3
- Version: 0.0.9
3
+ Version: 0.0.11
4
4
  Summary: A small package to facilitate teaching AI courses at business schools
5
5
  Project-URL: Homepage, https://simonbb.com
6
6
  Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
@@ -9,17 +9,27 @@ License-File: LICENSE
9
9
  Classifier: Operating System :: OS Independent
10
10
  Classifier: Programming Language :: Python :: 3
11
11
  Requires-Python: >=3.9
12
- Requires-Dist: gymnasium>=1.3.0
13
- Requires-Dist: matplotlib>=3.10.8
14
- Requires-Dist: numpy>=2.4.0
15
- Requires-Dist: opencv-python>=5.0.0.93
16
- Requires-Dist: pytorch-ignite>=0.5.4
17
- Requires-Dist: pyyaml>=6.0.3
18
- Requires-Dist: ray>=2.56.0
19
- Requires-Dist: stable-baselines3>=2.9.0
20
- Requires-Dist: torch>=2.9.1
21
- Requires-Dist: torchvision>=0.24.1
22
- Requires-Dist: transformers>=4.57.3
12
+ Requires-Dist: chromadb~=1.5.9
13
+ Requires-Dist: gymnasium~=1.3.0
14
+ Requires-Dist: langchain-chroma~=1.1.0
15
+ Requires-Dist: langchain-community~=0.4.2
16
+ Requires-Dist: langchain-core~=1.6.12
17
+ Requires-Dist: langchain-openai~=1.7.2
18
+ Requires-Dist: langchain-text-splitters~=1.1.3
19
+ Requires-Dist: langchain[mcp]~=1.4.6
20
+ Requires-Dist: langgraph~=1.2.16
21
+ Requires-Dist: matplotlib~=3.10.9
22
+ Requires-Dist: numpy~=2.4.6
23
+ Requires-Dist: opencv-python~=5.0.0.93
24
+ Requires-Dist: pydantic~=2.14.1
25
+ Requires-Dist: python-dotenv~=1.2.4
26
+ Requires-Dist: pytorch-ignite~=0.5.5
27
+ Requires-Dist: pyyaml~=6.0.3
28
+ Requires-Dist: ray~=2.56.1
29
+ Requires-Dist: stable-baselines3~=2.9.0
30
+ Requires-Dist: torchvision~=0.24.1
31
+ Requires-Dist: torch~=2.9.1
32
+ Requires-Dist: transformers~=4.57.6
23
33
  Description-Content-Type: text/markdown
24
34
 
25
35
  # Background
@@ -0,0 +1,44 @@
1
+ [build-system]
2
+ requires = ["hatchling >= 1.26"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "simonbb"
7
+ version = "0.0.11"
8
+ authors = [
9
+ { name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
10
+ ]
11
+ description = "A small package to facilitate teaching AI courses at business schools"
12
+ readme = "README.md"
13
+ requires-python = ">=3.9"
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "Operating System :: OS Independent",
17
+ ]
18
+ license = "MIT"
19
+ license-files = ["LICEN[CS]E*"]
20
+ dependencies = [
21
+ "matplotlib ~= 3.10.9",
22
+ "numpy ~= 2.4.6",
23
+ "torch ~= 2.9.1",
24
+ "torchvision ~= 0.24.1",
25
+ "transformers ~= 4.57.6",
26
+ "gymnasium ~= 1.3.0",
27
+ "pytorch-ignite ~= 0.5.5",
28
+ "stable_baselines3 ~= 2.9.0",
29
+ "ray ~= 2.56.1",
30
+ "opencv-python ~= 5.0.0.93",
31
+ "pyyaml ~= 6.0.3",
32
+ "langchain-core ~= 1.6.12",
33
+ "langchain[mcp] ~= 1.4.6",
34
+ "langgraph ~= 1.2.16",
35
+ "langchain-community ~= 0.4.2",
36
+ "langchain-chroma ~= 1.1.0",
37
+ "langchain-openai ~= 1.7.2",
38
+ "langchain_text_splitters ~= 1.1.3",
39
+ "chromadb ~= 1.5.9",
40
+ "python-dotenv ~= 1.2.4",
41
+ "pydantic ~= 2.14.1"
42
+ ]
43
+ [project.urls]
44
+ Homepage = "https://simonbb.com"
@@ -1,11 +1,11 @@
1
1
  import re
2
2
  import torch
3
3
  import torch.nn as nn
4
- from transformers import BertTokenizer, BertModel, BertConfig
4
+ from transformers import AutoTokenizer, BertModel, BertConfig
5
5
 
6
6
  #%%
7
7
 
8
- tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)
8
+ tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
9
9
 
10
10
  class BertClassifier( nn.Module ):
11
11
  def __init__(self, freeze_bert=False):
@@ -33,7 +33,7 @@ def BertPreprocessor( txt ):
33
33
  for s in txt:
34
34
  s = re.sub(r'(@.*?)[\s]', ' ', s)
35
35
  s = re.sub(r'\s+', ' ', s).strip()
36
- encoded = tokenizer.encode_plus( text=s, add_special_tokens=True, max_length=128, padding='max_length', truncation=True, return_attention_mask=True )
36
+ encoded = tokenizer( text=s, add_special_tokens=True, max_length=128, padding='max_length', truncation=True, return_attention_mask=True, return_tensors="pt" )
37
37
  input_ids.append( encoded.get('input_ids') ) # Add the outputs to the lists
38
38
  attention_masks.append( encoded.get('attention_mask') )
39
39
  return torch.tensor( input_ids ), torch.tensor( attention_masks )
@@ -1,8 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
- import os
4
3
  import re
5
- from typing import Annotated, TypedDict, get_args, Any, Protocol, Sequence
4
+ from typing import Annotated, TypedDict, get_args, Any, Sequence
6
5
  from dataclasses import dataclass
7
6
  from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
8
7
 
@@ -25,6 +24,7 @@ from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLo
25
24
  from langchain_core.callbacks import CallbackManagerForRetrieverRun
26
25
 
27
26
  import hashlib
27
+ from uuid import uuid4
28
28
  import chromadb
29
29
  from langchain_core.embeddings import Embeddings
30
30
  from langchain_core.language_models import BaseChatModel
@@ -35,13 +35,10 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
35
35
 
36
36
  def load_documents(path: Path):
37
37
  """Load documents from a file or directory."""
38
-
39
- valid_extensions = {'.pdf', '.html', '.ejs'}
40
-
38
+ valid_extensions = {'.pdf', '.html', '.ejs'}
41
39
  docs = []
42
40
  success_cnt = 0
43
- page_cnt = 0
44
-
41
+ page_cnt = 0
45
42
  if path.exists():
46
43
  if path.is_dir():
47
44
  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):
77
74
 
78
75
  return docs, success_cnt, page_cnt
79
76
 
77
+ #%% RAG Chains
78
+
79
+ system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
80
+ IMPORTANT INSTRUCTIONS:
81
+ 1. Base your answer ONLY on the provided context snippets below.
82
+ 2. Synthesize information from multiple snippets if needed.
83
+ 3. If the snippets contain partial information, provide the best answer you can with what's available.
84
+ 4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
85
+ 5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
86
+
87
+ The snippets are provided in this format:
88
+ Snippet 1: <text content>
89
+ Snippet 2: <text content>
90
+ ...
91
+
92
+ {context}"""
93
+
94
+ def format_documents(documents: Sequence[Document]) -> str:
95
+ snippets = []
96
+ for i, doc in enumerate(documents, start=1):
97
+ snippets.append(f"Snippet {i}: {doc.page_content}")
98
+ context = "\n".join(snippets)
99
+ return context
100
+
101
+ def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
102
+ """Return an LCEL chain accepting a question string and producing an answer string."""
103
+ prompt = ChatPromptTemplate.from_messages(
104
+ [
105
+ ( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
106
+ ("human", "{query}"),
107
+ ]
108
+ )
109
+ inputs = {
110
+ "context": retriever | format_documents,
111
+ "query": RunnablePassthrough(),
112
+ }
113
+ return inputs | prompt | model | StrOutputParser()
114
+
115
+
116
+ def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
117
+ """Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
118
+ -- 'sources': original list of Documents from retriever
119
+ -- 'query': original query
120
+ -- 'context': merged content from relevant documents
121
+ -- 'answer': final answer """
122
+ prompt = ChatPromptTemplate.from_messages([
123
+ ("system", system_prompt),
124
+ ("human", "{query}")
125
+ ])
126
+ return (
127
+ {
128
+ "sources": retriever,
129
+ "query": RunnablePassthrough(),
130
+ }
131
+ | RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
132
+ | RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
133
+ )
134
+
135
+ build_rag = rag_main
136
+
137
+
138
+ """
139
+ Retrievers are subclasses of LangChain's BaseRetriever which inherits from Pydantic's BaseModel.
140
+ Unlike ordinary Python classes where attributes declared in the class body are class attributes shared by all instances,
141
+ Pydantic model fields declared with type annotations are instance attributes, not shared class attributes, even though they are declared in the class body.
142
+ """
80
143
  #%% MultiQueryRerankRetriever
81
144
 
82
145
  class MultiQueryRerankRetriever(BaseRetriever):
83
146
  """Chroma retriever with optional multi-query retrieval and reranking."""
84
- knowledge_base: str
85
147
  embeddings: Embeddings = Field(exclude=True)
86
148
  query_llm: BaseChatModel | None = Field(default=None, exclude=True)
87
149
  reranker: Any = Field(default=None, exclude=True)
88
150
  k: int = Field(default=3, ge=1)
89
151
  top_n: int = Field(default=5, ge=1)
90
- vstore_path: str = Field( default_factory=lambda: f"{os.path.expanduser('~')}/vstore" )
152
+ persist_directory: str | Path
153
+ collection_name: str
154
+ create_collection_if_not_exists: bool = False
155
+ store: Chroma = None
91
156
  _base_retriever: BaseRetriever = PrivateAttr()
92
157
 
158
+ """ 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."""
93
159
  def model_post_init(self, __context: Any) -> None:
94
- """Initialize the underlying Chroma retriever."""
95
- store = Chroma( persist_directory=( f"{self.vstore_path}/{self.knowledge_base}/BAAI_db/chroma_vector_store"), embedding_function=self.embeddings )
96
- self._base_retriever = store.as_retriever( search_kwargs={"k": self.k} )
160
+ """Initialize the underlying Chroma retriever. Use the last folder name as collection name"""
161
+ 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 )
162
+ self._base_retriever = self.store.as_retriever( search_kwargs={"k": self.k} )
97
163
  if self.query_llm is not None and self.reranker is None:
98
164
  raise ValueError( "A reranker is required when query_llm is provided." )
165
+
166
+ def ingest( self, documents: Sequence[Document], *, source_dir: str = 'placeholder', chunk_size: int = 1000, chunk_overlap: int = 150 ):
167
+ text_splitter = RecursiveCharacterTextSplitter( separators=["\n\n", "\n", ". ", " ", ""], chunk_size=chunk_size, chunk_overlap=chunk_overlap )
168
+ chunks = text_splitter.split_documents( documents )
169
+ self.store.add_documents( chunks )
170
+
99
171
  # run_manager manages callbacks, tracing, and monitoring during a retriever's execution, allowing integrations with LangSmith
100
172
  def _generate_queries( self, query: str, run_manager: CallbackManagerForRetrieverRun ) -> list[str]:
101
173
  """Generate five alternative formulations of the query."""
@@ -127,6 +199,7 @@ class MultiQueryRerankRetriever(BaseRetriever):
127
199
  )
128
200
  return [doc for _, doc in scored_docs[:self.top_n]]
129
201
 
202
+ # override _get_relevant_documents() method from BaseRetriever
130
203
  def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> list[Document]:
131
204
  """Retrieve documents using simple or multi-query retrieval."""
132
205
  # Simple retrieval
@@ -143,12 +216,91 @@ class MultiQueryRerankRetriever(BaseRetriever):
143
216
  # Deduplicate and rerank
144
217
  unique_docs = self._unique_union(doc_lists)
145
218
  return self._rerank(query, unique_docs)
219
+
220
+ #%%
221
+
222
+ class DummyEmbeddings(Embeddings):
223
+ """Return constant 1-dimensional embeddings without using an LLM."""
224
+ def embed_documents(self, texts: list[str]) -> list[list[float]]:
225
+ return [[0.0] for _ in texts]
226
+ def embed_query(self, text: str) -> list[float]:
227
+ return [0.0]
228
+
229
+ class ParentDocumentRetriever(BaseRetriever):
230
+ persist_directory: str | Path
231
+ parent_collection: str
232
+ child_collection: str
233
+ embeddings: Embeddings
234
+ k: int = 5
235
+ _parent_store: Chroma = PrivateAttr()
236
+ _child_store: Chroma = PrivateAttr()
237
+
238
+ def __init__(self, **kwargs):
239
+ super().__init__(**kwargs)
240
+ self._parent_store = Chroma(
241
+ embedding_function=DummyEmbeddings(),
242
+ collection_name=self.parent_collection,
243
+ persist_directory=self.persist_directory,
244
+ )
245
+
246
+ self._child_store = Chroma(
247
+ embedding_function=self.embeddings,
248
+ collection_name=self.child_collection,
249
+ persist_directory=self.persist_directory,
250
+ )
251
+
252
+ 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:
253
+ """Split documents and store parent and child chunks."""
254
+ parent_splitter = RecursiveCharacterTextSplitter( chunk_size=parent_chunk_size, chunk_overlap=parent_chunk_overlap )
255
+ child_splitter = RecursiveCharacterTextSplitter( chunk_size=child_chunk_size, chunk_overlap=child_chunk_overlap )
256
+
257
+ parent_docs = []
258
+ parent_ids = []
259
+ child_docs = []
260
+
261
+ for parent in parent_splitter.split_documents(docs):
262
+ parent_id = str(uuid4())
263
+ parent_docs.append(parent)
264
+ parent_ids.append(parent_id)
265
+ for child in child_splitter.split_documents( [parent] ):
266
+ child.metadata["parent_id"] = parent_id
267
+ child_docs.append(child)
268
+ if parent_docs:
269
+ self._parent_store.add_documents( documents=parent_docs, ids=parent_ids )
270
+ if child_docs:
271
+ self._child_store.add_documents( documents=child_docs )
272
+
273
+ # override _get_relevant_documents() method from BaseRetriever
274
+ def _get_relevant_documents(
275
+ self,
276
+ query: str,
277
+ *,
278
+ run_manager: CallbackManagerForRetrieverRun,
279
+ ) -> list[Document]:
280
+ """Retrieve parent documents using child similarity search."""
281
+
282
+ # 1. Search child chunks
283
+ child_chunks = self._child_store.similarity_search( query, k=self.k )
284
+ # 2. Extract unique parent IDs in relevance order
285
+ parent_ids = list(dict.fromkeys( # standard trick to remove duplicates from a list or iterable while preserving the original item orde
286
+ chunk.metadata["parent_id"] for chunk in child_chunks if "parent_id" in chunk.metadata
287
+ ))
288
+ if not parent_ids:
289
+ return []
290
+ # 3. Fetch parent documents
291
+ parents = self._parent_store.get_by_ids( parent_ids )
292
+ parents_by_id = { doc.id: doc for doc in parents if doc.id is not None }
293
+ # 4. Restore relevance order
294
+ return [
295
+ parents_by_id[pid]
296
+ for pid in parent_ids if pid in parents_by_id
297
+ ]
146
298
 
147
299
 
148
300
  #%% Hypothetical Question Retriever
149
301
 
150
- def _stable_id(source_id: str, index: int, text: str) -> str:
151
- digest = hashlib.sha256(f"{source_id}\0{index}\0{text}".encode()).hexdigest()
302
+ def _stable_id(source_dir: str, index: int, text: str) -> str:
303
+ digest = hashlib.sha256(f"{source_dir}\0{index}\0{text}".encode()).hexdigest()
152
304
  return f"chunk-{digest[:32]}"
153
305
 
154
306
  def _clean_questions(questions: Sequence[str]) -> list[str]:
@@ -169,9 +321,6 @@ def _chroma_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
169
321
  if value is not None and isinstance(value, (str, int, float, bool))
170
322
  }
171
323
 
172
- class QuestionGenerator(Protocol):
173
- """Anything capable of generating search questions for a chunk."""
174
- def generate(self, document: Document) -> list[str]: ...
175
324
 
176
325
  class _QuestionList(BaseModel):
177
326
  questions: list[str] = Field(description="Distinct questions answered by the text")
@@ -179,7 +328,7 @@ class _QuestionList(BaseModel):
179
328
  class HypotheticalQuestionGenerator:
180
329
  """LLM-backed question generator using provider-supported structured output."""
181
330
 
182
- def __init__(self, model: BaseChatModel, *, questions_per_chunk: int = 4) -> None:
331
+ def __init__(self, llm: BaseChatModel, *, questions_per_chunk: int = 4) -> None:
183
332
  if questions_per_chunk < 1:
184
333
  raise ValueError("questions_per_chunk must be positive")
185
334
  self.questions_per_chunk = questions_per_chunk
@@ -189,7 +338,7 @@ class HypotheticalQuestionGenerator:
189
338
  ("human", "Text:\n{text}"),
190
339
  ]
191
340
  )
192
- self._chain = prompt | model.with_structured_output(_QuestionList)
341
+ self._chain = prompt | llm.with_structured_output(_QuestionList)
193
342
 
194
343
  def generate(self, document: Document) -> list[str]:
195
344
  result = self._chain.invoke( {"count": self.questions_per_chunk, "text": document.page_content} )
@@ -203,38 +352,40 @@ class IngestionResult:
203
352
  chunk_ids: tuple[str, ...]
204
353
 
205
354
  class HypotheticalQuestionRetriever(BaseRetriever):
206
- """Search question vectors, then resolve their chunk IDs from the chunk collection."""
207
- model_config = ConfigDict(arbitrary_types_allowed=True)
208
- question_store: Chroma
209
- chunk_collection: Any
210
- generator: Any
355
+ """Search question vectors, then resolve their chunk IDs from the chunk collection."""
356
+ llm: Any = None
357
+ persist_directory: str | Path
358
+ embeddings: Embeddings
359
+ question_collection: str
360
+ chunk_collection: str
211
361
  question_k: int = 12
212
362
  chunk_k: int = 4
213
363
  metadata_filter: dict[str, Any] | None = None
364
+ create_collection_if_not_exists: bool = False
365
+ _question_store: Chroma = PrivateAttr()
366
+ _chunk_collection: Any = PrivateAttr()
367
+ _generator: Any = PrivateAttr()
214
368
 
215
- 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:
216
- client = chromadb.PersistentClient(path=str(persist_directory))
217
- question_store = Chroma( client=client, collection_name=question_collection, embedding_function=embeddings )
218
- chunks = client.get_or_create_collection( name=chunk_collection, embedding_function=None )
219
- super().__init__(
220
- question_store=question_store,
221
- chunk_collection=chunks,
222
- generator=generator,
223
- question_k=question_k,
224
- chunk_k=chunk_k,
225
- metadata_filter=metadata_filter,
226
- **kwargs,
227
- )
369
+ model_config = ConfigDict(arbitrary_types_allowed=True)
228
370
 
229
- 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:
230
- if not source_id: # source_id is a unique identifier for the source document or documents being ingested
231
- raise ValueError("source_id must not be empty")
371
+ def __init__(self, **kwargs) -> None:
372
+ super().__init__(**kwargs)
373
+ client = chromadb.PersistentClient( path= self.persist_directory )
374
+ 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 )
375
+ self._chunk_collection = client.get_or_create_collection( name=self.chunk_collection, embedding_function=None )
376
+
377
+ 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:
378
+ if not self.llm:
379
+ raise ValueError("To ingest, llm must be provided in the constructor.")
380
+ self._generator = HypotheticalQuestionGenerator( self.llm )
381
+ if not source_dir: # source_dir is a unique identifier for the source document or documents being ingested
382
+ raise ValueError("source_dir must not be empty")
232
383
  splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True )
233
384
  input_docs = splitter.split_documents( list(documents) ) if split else list(documents)
234
385
 
235
386
  if replace_source:
236
- self.question_store.delete(where={"source_id": source_id})
237
- self.chunk_collection.delete(where={"source_id": source_id})
387
+ self._question_collection.delete( where={"source_dir": source_dir} )
388
+ self._chunk_collection.delete( where={"source_dir": source_dir} )
238
389
 
239
390
  chunk_ids: list[str] = []
240
391
  question_docs: list[Document] = []
@@ -243,32 +394,31 @@ class HypotheticalQuestionRetriever(BaseRetriever):
243
394
  chunk_texts: list[str] = []
244
395
 
245
396
  for index, doc in enumerate(input_docs):
246
- chunk_id = _stable_id(source_id, index, doc.page_content)
247
- metadata = _chroma_metadata( {**doc.metadata, "source_id": source_id, "chunk_id": chunk_id} )
397
+ chunk_id = _stable_id(source_dir, index, doc.page_content)
398
+ metadata = _chroma_metadata( {**doc.metadata, "source_dir": source_dir, "chunk_id": chunk_id} )
248
399
  chunk_ids.append( chunk_id )
249
400
  chunk_metadatas.append( metadata )
250
401
  chunk_texts.append( doc.page_content )
251
-
252
- questions = _clean_questions( self.generator.generate(doc) )
402
+ questions = _clean_questions( self._generator.generate(doc) )
253
403
  for question_index, question in enumerate(questions):
254
404
  question_ids.append( f"{chunk_id}:q:{question_index}" )
255
405
  question_docs.append(
256
406
  Document( page_content=question,
257
- metadata={"chunk_id": chunk_id, "source_id": source_id, "question_index": question_index },
407
+ metadata={"chunk_id": chunk_id, "source_dir": source_dir, "question_index": question_index },
258
408
  )
259
409
  )
260
410
  if question_docs: # LangChain's Chroma.add_documents() triggers embedding
261
- self.question_store.add_documents( question_docs, ids=question_ids )
411
+ self._question_collection.add_documents( question_docs, ids=question_ids )
262
412
 
263
413
  if chunk_ids: # Chroma's underlying collection API stores text and metadata without embeddings.
264
- self.chunk_collection.upsert( ids=chunk_ids, documents=chunk_texts, metadatas=chunk_metadatas,
414
+ self._chunk_collection.upsert( ids=chunk_ids, documents=chunk_texts, metadatas=chunk_metadatas,
265
415
  embeddings=[[0.0] for _ in chunk_ids] ) # placeholder vectors since Chroma records require vectors
266
416
 
267
417
  return IngestionResult(len(chunk_ids), len(question_docs), tuple(chunk_ids))
268
418
 
269
419
  # override _get_relevant_documents() method from BaseRetriever
270
420
  def _get_relevant_documents(self, query: str, *, run_manager: Any = None) -> list[Document]:
271
- matches = self.question_store.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
421
+ matches = self._question_collection.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
272
422
  chunk_ids: list[str] = []
273
423
  seen: set[str] = set()
274
424
  for m in matches:
@@ -281,74 +431,16 @@ class HypotheticalQuestionRetriever(BaseRetriever):
281
431
  if not chunk_ids:
282
432
  return []
283
433
 
284
- records = self.chunk_collection.get(ids=chunk_ids, include=["documents", "metadatas"])
285
- by_id = {
434
+ records = self._chunk_collection.get( ids=chunk_ids, include=["documents", "metadatas"] )
435
+ records_by_id = {
286
436
  record_id: Document(page_content=text, metadata=metadata or {})
287
437
  for record_id, text, metadata in zip(
288
438
  records["ids"], records["documents"], records["metadatas"], strict=True
289
439
  )
290
440
  }
291
- return [by_id[chunk_id] for chunk_id in chunk_ids if chunk_id in by_id]
292
-
293
- #%% RAG Chains
294
-
295
- system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
296
- IMPORTANT INSTRUCTIONS:
297
- 1. Base your answer ONLY on the provided context snippets below.
298
- 2. Synthesize information from multiple snippets if needed.
299
- 3. If the snippets contain partial information, provide the best answer you can with what's available.
300
- 4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
301
- 5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
302
-
303
- The snippets are provided in this format:
304
- Snippet 1: <text content>
305
- Snippet 2: <text content>
306
- ...
307
-
308
- {context}"""
309
-
310
- def format_documents(documents: Sequence[Document]) -> str:
311
- snippets = []
312
- for i, doc in enumerate(documents, start=1):
313
- snippets.append(f"Snippet {i}: {doc.page_content}")
314
- context = "\n".join(snippets)
315
- return context
316
-
317
- def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
318
- """Return an LCEL chain accepting a question string and producing an answer string."""
319
- prompt = ChatPromptTemplate.from_messages(
320
- [
321
- ( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
322
- ("human", "{query}"),
323
- ]
324
- )
325
- inputs = {
326
- "context": retriever | format_documents,
327
- "query": RunnablePassthrough(),
328
- }
329
- return inputs | prompt | model | StrOutputParser()
441
+ return [ records_by_id[chunk_id] for chunk_id in chunk_ids if chunk_id in records_by_id ]
330
442
 
331
443
 
332
- def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
333
- """Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
334
- -- 'sources': original list of Documents from retriever
335
- -- 'query': original query
336
- -- 'context': merged content from relevant documents
337
- -- 'answer': final answer """
338
- prompt = ChatPromptTemplate.from_messages([
339
- ("system", system_prompt),
340
- ("human", "{query}")
341
- ])
342
- return (
343
- {
344
- "sources": retriever,
345
- "query": RunnablePassthrough(),
346
- }
347
- | RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
348
- | RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
349
- )
350
-
351
- build_rag = rag_main
352
444
 
353
445
  #%% RAG with Routing
354
446
 
@@ -368,7 +460,7 @@ Snippet 2: <text content>
368
460
 
369
461
  {context}"""
370
462
 
371
- def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
463
+ def multiRagBuilder( persist_directory, domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
372
464
  """
373
465
  domain_literals: the last one is a placeholder if the router finds no domain matches the query
374
466
  """
@@ -444,7 +536,7 @@ def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker,
444
536
  retrievers = {}
445
537
  retriever_nodes = {}
446
538
  for domain in domains:
447
- retrievers[domain] = MultiQueryRerankRetriever(knowledge_base=domain, embeddings=embeddings, query_llm=query_llm, reranker=reranker, k=k)
539
+ retrievers[domain] = MultiQueryRerankRetriever( persist_directory=persist_directory, collection_name=domain, embeddings=embeddings, query_llm=query_llm, reranker=reranker, k=k)
448
540
  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.
449
541
  builder.add_node(domain, retriever_nodes[domain])
450
542
  builder.add_edge(domain, "generate_answer")
@@ -1,34 +0,0 @@
1
- [build-system]
2
- requires = ["hatchling >= 1.26"]
3
- build-backend = "hatchling.build"
4
-
5
- [project]
6
- name = "simonbb"
7
- version = "0.0.9"
8
- authors = [
9
- { name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
10
- ]
11
- description = "A small package to facilitate teaching AI courses at business schools"
12
- readme = "README.md"
13
- requires-python = ">=3.9"
14
- classifiers = [
15
- "Programming Language :: Python :: 3",
16
- "Operating System :: OS Independent",
17
- ]
18
- license = "MIT"
19
- license-files = ["LICEN[CS]E*"]
20
- dependencies = [
21
- "matplotlib >= 3.10.8",
22
- "torch >= 2.9.1",
23
- "torchvision >=0.24.1",
24
- "numpy >=2.4.0",
25
- "transformers >= 4.57.3",
26
- "gymnasium >= 1.3.0",
27
- "pytorch-ignite >= 0.5.4",
28
- "stable_baselines3 >= 2.9.0",
29
- "ray >= 2.56.0",
30
- "opencv-python >= 5.0.0.93",
31
- "pyyaml >= 6.0.3"
32
- ]
33
- [project.urls]
34
- Homepage = "https://simonbb.com"
File without changes
File without changes
File without changes