jdwh08s/Autodoc-Lifter
0
1#####################################################2### DOCUMENT PROCESSOR [FULLDOC]3#####################################################4### Jonathan Wang5 6# ABOUT:7# This creates an app to chat with PDFs.8 9# This is the FULLDOC10# which is a class that associates documents11# with their critical information12# and their tools. (keywords, summary, queryengine, etc.)13#####################################################14### TODO Board:15# Automatically determine which reader to use for each document based on the file type.16 17#####################################################18### PROGRAM SETTINGS19 20#####################################################21### PROGRAM IMPORTS22from __future__ import annotations23 24import asyncio25from pathlib import Path26from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, TypeVar27from uuid import UUID, uuid428 29from llama_index.core import StorageContext, VectorStoreIndex30from llama_index.core.query_engine import SubQuestionQueryEngine31from llama_index.core.schema import BaseNode, TransformComponent32from llama_index.core.settings import Settings33from llama_index.core.tools import QueryEngineTool, ToolMetadata34from streamlit import session_state as ss35 36if TYPE_CHECKING:37 from llama_index.core.base.base_query_engine import BaseQueryEngine38 from llama_index.core.callbacks import CallbackManager39 from llama_index.core.node_parser import NodeParser40 from llama_index.core.readers.base import BaseReader41 from llama_index.core.response_synthesizers import BaseSynthesizer42 from llama_index.core.retrievers import BaseRetriever43 44# Own Modules45from engine import get_engine46from keywords import KeywordMetadataAdder47from retriever import get_retriever48from storage import get_docstore, get_vector_store49from summary import DEFAULT_ONELINE_SUMMARY_TEMPLATE, DEFAULT_TREE_SUMMARY_TEMPLATE50 51#####################################################52### SCRIPT53 54GenericNode = TypeVar("GenericNode", bound=BaseNode)55 56class FullDocument:57 """Bundles all the information about a document together.58 59 Args:60 name (str): The name of the document.61 file_path (Path): The path to the document.62 summary (str): The summary of the document.63 keywords (List[str]): The keywords of the document.64 entities (List[str]): The entities of the document.65 vector_store (BaseDocumentStore): The vector store of the document.66 """67 68 # Identifiers69 id: UUID70 name: str71 file_path: Path72 file_name: str73 74 # Basic Contents75 summary: str76 summary_oneline: str # A one line summary of the document.77 keywords: set[str] # List of keywords in document.78 # entities: Set[str] # list of entities in document ## TODO: Add entities79 metadata: dict[str, Any] | None80 # NOTE: other metdata that might be useful:81 # Document Creation / Last Date (e.g., recency important for legal/medical questions)82 # Document Source and Trustworthiness83 # Document Access Level (though this isn't important for us here.)84 # Document Citations?85 # Document Format? (text/spreadsheet/presentation/image/etc.)86 87 # RAG Components88 nodes: list[BaseNode]89 storage_context: StorageContext # NOTE: current setup has single storage context per document.90 vector_store_index: VectorStoreIndex91 retriever: BaseRetriever # TODO(Jonathan Wang): Consider multiple retrievers for keywords vs semantic.92 engine: BaseQueryEngine # TODO(Jonathan Wang): Consider mulitple engines. 93 subquestion_engine: SubQuestionQueryEngine94 95 def __init__(96 self,97 name: str,98 file_path: Path | str,99 metadata: dict[str, Any] | None = None100 ) -> None:101 self.id = uuid4()102 self.name = name103 104 if (isinstance(file_path, str)):105 file_path = Path(file_path)106 self.file_path = file_path107 self.file_name = file_path.name108 109 self.metadata = metadata110 111 112 @classmethod113 def class_name(cls) -> str:114 return "FullDocument"115 116 def add_name_to_nodes(self, nodes: list[GenericNode]) -> list[GenericNode]:117 """Add the name of the document to the nodes.118 119 Args:120 nodes (List[GenericNode]): The nodes to add the name to.121 122 Returns:123 List[GenericNode]: The nodes with the name added.124 """125 for node in nodes:126 node.metadata["name"] = self.name127 return nodes128 129 def file_to_nodes(130 self,131 reader: BaseReader,132 postreaders: list[Callable[[list[GenericNode]], list[GenericNode]] | TransformComponent] | None=None, # NOTE: these should be used in order. and probably all TransformComponent instead.133 node_parser: NodeParser | None=None,134 postparsers: list[Callable[[list[GenericNode]], list[GenericNode]] | TransformComponent] | None=None, # Stuff like chunking, adding Embeddings, etc.135 ) -> None:136 """Read in the file path and get the nodes.137 138 Args:139 file_path (Optional[Path], optional): The path to the file. Defaults to file_path from init.140 reader (Optional[BaseReader], optional): The reader to use. Defaults to reader from init.141 """142 # Use the provided reader to read in the file.143 print("NEWPDF: Reading input file...")144 nodes = reader.load_data(file_path=self.file_path)145 146 # Use node postreaders to post process the nodes.147 if (postreaders is not None):148 for node_postreader in postreaders:149 nodes = node_postreader(nodes) # type: ignore (TransformComponent allows a list of nodes)150 151 # Use node parser to parse the nodes.152 if (node_parser is None):153 node_parser = Settings.node_parser154 nodes = node_parser(nodes) # type: ignore (Document is a child of BaseNode)155 156 # Use node postreaders to post process the nodes. (also add the common name to the nodes)157 if (postparsers is None):158 postparsers = [self.add_name_to_nodes]159 else:160 postparsers.append(self.add_name_to_nodes)161 162 for node_postparser in postparsers:163 nodes = node_postparser(nodes) # type: ignore (TransformComponent allows a list of nodes)164 165 # Save nodes166 self.nodes = nodes # type: ignore167 168 def nodes_to_summary(169 self,170 summarizer: BaseSynthesizer, # NOTE: this is typically going to be a TreeSummarizer / SimpleSummarize for our use case171 query_str: str = DEFAULT_TREE_SUMMARY_TEMPLATE,172 ) -> None:173 """Summarize the nodes.174 175 Args:176 summarizer (BaseSynthesizer): The summarizer to use. Takes in nodes and returns summary.177 """178 if (not hasattr(self, "nodes")):179 msg = "Nodes must be extracted from document using `file_to_nodes` before calling `nodes_to_summary`."180 raise ValueError(msg)181 182 text_chunks = [getattr(node, "text", "") for node in self.nodes if hasattr(node, "text")]183 summary_responses = summarizer.aget_response(query_str=query_str, text_chunks=text_chunks)184 185 loop = asyncio.get_event_loop()186 summary = loop.run_until_complete(summary_responses)187 188 if (not isinstance(summary, str)):189 # TODO(Jonathan Wang): ... this should always give us a string, right? we're not doing anything fancy with TokenGen/TokenAsyncGen/Pydantic BaseModel...190 msg = f"Summarizer must return a string summary. Actual type: {type(summary)}, with value {summary}."191 raise TypeError(msg)192 193 self.summary = summary194 195 def summary_to_oneline(196 self,197 summarizer: BaseSynthesizer, # NOTE: this is typically going to be a SimpleSummarize / TreeSummarizer for our use case198 query_str: str = DEFAULT_ONELINE_SUMMARY_TEMPLATE,199 ) -> None:200 201 if (not hasattr(self, "summary")):202 msg = "Summary must be extracted from document using `nodes_to_summary` before calling `summary_to_oneline`."203 raise ValueError(msg)204 205 oneline = summarizer.get_response(query_str=query_str, text_chunks=[self.summary]) # There's only one chunk.206 self.summary_oneline = oneline # type: ignore | shouldn't have fancy TokenGenerators / TokenAsyncGenerators / Pydantic BaseModels207 208 def nodes_to_document_keywords(self, keyword_extractor: Optional[KeywordMetadataAdder] = None) -> None:209 """Save the keywords from the nodes into the document.210 211 Args:212 keyword_extractor (Optional[BaseKeywordExtractor], optional): The keyword extractor to use. Defaults to None.213 """214 if (not hasattr(self, "nodes")):215 msg = "Nodes must be extracted from document using `file_to_nodes` before calling `nodes_to_keywords`."216 raise ValueError(msg)217 218 if (keyword_extractor is None):219 keyword_extractor = KeywordMetadataAdder()220 221 # Add keywords to nodes using KeywordMetadataAdder222 keyword_extractor.process_nodes(self.nodes)223 224 # Save keywords225 keywords: list[str] = []226 for node in self.nodes:227 node_keywords = node.metadata.get("keyword_metadata", "").split(", ") # NOTE: KeywordMetadataAdder concatinates b/c required string output228 keywords = keywords + node_keywords229 230 # TODO(Jonathan Wang): handle dedupling keywords which are similar to each other (fuzzy?)231 self.keywords = set(keywords)232 233 def nodes_to_storage(self, create_new_storage: bool = True) -> None:234 """Save the nodes to storage."""235 if (not hasattr(self, "nodes")):236 msg = "Nodes must be extracted from document using `file_to_nodes` before calling `nodes_to_storage`."237 raise ValueError(msg)238 239 if (create_new_storage):240 docstore = get_docstore(documents=self.nodes)241 self.docstore = docstore242 243 vector_store = get_vector_store()244 245 storage_context = StorageContext.from_defaults(246 docstore=docstore,247 vector_store=vector_store248 )249 self.storage_context = storage_context250 251 vector_store_index = VectorStoreIndex(252 self.nodes, storage_context=storage_context253 )254 self.vector_store_index = vector_store_index255 256 else:257 ### TODO(Jonathan Wang): use an existing storage instead of creating a new one.258 msg = "Currently creates new storage for every document."259 raise NotImplementedError(msg)260 261 # TODO(Jonathan Wang): Create multiple different retrievers based on the question type(?)262 # E.g., if the question is focused on specific keywords or phrases, use a retriever oriented towards sparse scores.263 def storage_to_retriever(264 self,265 semantic_nodes: int = 6,266 sparse_nodes: int = 3,267 fusion_nodes: int = 3,268 semantic_weight: float = 0.6,269 merge_up_thresh: float = 0.5,270 callback_manager: CallbackManager | None=None271 ) -> None:272 """Create retriever from storage."""273 if (not hasattr(self, "vector_store_index")):274 msg = "Vector store must be extracted from document using `nodes_to_storage` before calling `storage_to_retriever`."275 raise ValueError(msg)276 277 retriever = get_retriever(278 _vector_store_index=self.vector_store_index,279 semantic_top_k=semantic_nodes,280 sparse_top_k=sparse_nodes,281 fusion_similarity_top_k=fusion_nodes,282 semantic_weight_fraction=semantic_weight,283 merge_up_thresh=merge_up_thresh,284 verbose=True,285 _callback_manager=callback_manager or ss.callback_manager286 )287 self.retriever = retriever288 289 def retriever_to_engine(290 self,291 response_synthesizer: BaseSynthesizer,292 callback_manager: CallbackManager | None=None293 ) -> None:294 """Create query engine from retriever."""295 if (not hasattr(self, "retriever")):296 msg = "Retriever must be extracted from document using `storage_to_retriever` before calling `retriver_to_engine`."297 raise ValueError(msg)298 299 engine = get_engine(300 retriever=self.retriever,301 response_synthesizer=response_synthesizer,302 callback_manager=callback_manager or ss.callback_manager303 )304 self.engine = engine305 306 # TODO(Jonathan Wang): Create Summarization Index and Engine.307 def engine_to_sub_question_engine(self) -> None:308 """Convert a basic query engine into a sub-question query engine for handling complex, multi-step questions.309 310 Args:311 query_engine (BaseQueryEngine): The Base Query Engine to convert.312 """313 if (not hasattr(self, "summary_oneline")):314 msg = "One Line Summary must be created for the document before calling `engine_to_sub_query_engine`"315 raise ValueError(msg)316 elif (not hasattr(self, "engine")):317 msg = "Basic Query Engine must be created before calling `engine_to_sub_query_engine`"318 raise ValueError(msg)319 320 sqe_tools = [321 QueryEngineTool(322 query_engine=self.engine, # TODO(Jonathan Wang): handle mulitple engines?323 metadata=ToolMetadata(324 name=(self.name + "simple query answerer"),325 description=f"""A tool that answers simple questions about the following document: {self.summary_oneline}"""326 )327 )328 # TODO(Jonathan Wang): add more tools329 ]330 331 subquestion_engine = SubQuestionQueryEngine.from_defaults(332 query_engine_tools=sqe_tools, 333 verbose=True,334 use_async=True335 )336 self.subquestion_engine = subquestion_engine337