+ Coverage for tinytroupe / agent / grounding.py: + 0% +
+ ++ 200 statements + + + +
++ « prev + ^ index + » next + + coverage.py v7.13.4, + created at 2026-02-28 17:48 +0000 +
+ +1from tinytroupe.utils import JsonSerializableRegistry
+2import tinytroupe.utils as utils
+ +4from tinytroupe.agent import logger
+5from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document, StorageContext, load_index_from_storage
+6from llama_index.core.vector_stores import SimpleVectorStore
+7from llama_index.readers.web import SimpleWebPageReader
+8import json
+9import tempfile
+10import os
+11import shutil
+ + +14#######################################################################################################################
+15# Grounding connectors
+16#######################################################################################################################
+ +18class GroundingConnector(JsonSerializableRegistry):
+19 """
+20 An abstract class representing a grounding connector. A grounding connector is a component that allows an agent to ground
+21 its knowledge in external sources, such as files, web pages, databases, etc.
+22 """
+ +24 serializable_attributes = ["name"]
+ +26 def __init__(self, name:str) -> None:
+27 self.name = name
+ +29 def retrieve_relevant(self, relevance_target:str, source:str, top_k=20) -> list:
+30 raise NotImplementedError("Subclasses must implement this method.")
+ +32 def retrieve_by_name(self, name:str) -> str:
+33 raise NotImplementedError("Subclasses must implement this method.")
+ +35 def list_sources(self) -> list:
+36 raise NotImplementedError("Subclasses must implement this method.")
+ + +39@utils.post_init
+40class BaseSemanticGroundingConnector(GroundingConnector):
+41 """
+42 A base class for semantic grounding connectors. A semantic grounding connector is a component that indexes and retrieves
+43 documents based on so-called "semantic search" (i.e, embeddings-based search). This specific implementation
+44 is based on the VectorStoreIndex class from the LLaMa-Index library. Here, "documents" refer to the llama-index's
+45 data structure that stores a unit of content, not necessarily a file.
+46 """
+ +48 serializable_attributes = ["documents", "index"]
+ +50 # needs custom deserialization to handle Pydantic models (Document is a Pydantic model)
+51 custom_deserializers = {"documents": lambda docs_json: [Document.from_json(doc_json) for doc_json in docs_json],
+52 "index": lambda index_json: BaseSemanticGroundingConnector._deserialize_index(index_json)}
+ +54 custom_serializers = {"documents": lambda docs: [doc.to_json() for doc in docs] if docs is not None else None,
+55 "index": lambda index: BaseSemanticGroundingConnector._serialize_index(index)}
+ +57 def __init__(self, name:str="Semantic Grounding") -> None:
+58 super().__init__(name)
+ +60 self.documents = None
+61 self.name_to_document = None
+62 self.index = None
+ +64 # @post_init ensures that _post_init is called after the __init__ method
+ +66 def _post_init(self):
+67 """
+68 This will run after __init__, since the class has the @post_init decorator.
+69 It is convenient to separate some of the initialization processes to make deserialize easier.
+70 """
+71 self.index = None
+ +73 if not hasattr(self, 'documents') or self.documents is None:
+74 self.documents = []
+ +76 if not hasattr(self, 'name_to_document') or self.name_to_document is None:
+77 self.name_to_document = {}
+ +79 if hasattr(self, 'documents') and self.documents is not None:
+80 for document in self.documents:
+81 # if the document has a semantic memory ID, we use it as the identifier
+82 name = document.metadata.get("semantic_memory_id", document.id_)
+ +84 # self.name_to_document[name] contains a list, since each source file could be split into multiple pages
+85 if name in self.name_to_document:
+86 self.name_to_document[name].append(document)
+87 else:
+88 self.name_to_document[name] = [document]
+ +90 # Rebuild index from documents if it's None or invalid
+91 if self.index is None and self.documents:
+92 logger.warning("No index found. Rebuilding index from documents.")
+93 vector_store = SimpleVectorStore()
+94 self.index = VectorStoreIndex.from_documents(
+95 self.documents,
+96 vector_store=vector_store,
+97 store_nodes_override=True
+98 )
+ +100 # TODO remove?
+101 #self.add_documents(self.documents)
+ +103 @staticmethod
+104 def _serialize_index(index):
+105 """Helper function to serialize index with proper storage context"""
+106 if index is None:
+107 return None
+ +109 try:
+110 # Create a temporary directory to store the index
+111 with tempfile.TemporaryDirectory() as temp_dir:
+112 # Persist the index to the temporary directory
+113 index.storage_context.persist(persist_dir=temp_dir)
+ +115 # Read all the persisted files and store them in a dictionary
+116 persisted_data = {}
+117 for filename in os.listdir(temp_dir):
+118 filepath = os.path.join(temp_dir, filename)
+119 if os.path.isfile(filepath):
+120 with open(filepath, 'r', encoding="utf-8", errors="replace") as f:
+121 persisted_data[filename] = f.read()
+ +123 return persisted_data
+124 except Exception as e:
+125 logger.warning(f"Failed to serialize index: {e}")
+126 return None
+ +128 @staticmethod
+129 def _deserialize_index(index_data):
+130 """Helper function to deserialize index with proper error handling"""
+131 if not index_data:
+132 return None
+ +134 try:
+135 # Create a temporary directory to restore the index
+136 with tempfile.TemporaryDirectory() as temp_dir:
+137 # Write all the persisted files to the temporary directory
+138 for filename, content in index_data.items():
+139 filepath = os.path.join(temp_dir, filename)
+140 with open(filepath, 'w', encoding="utf-8", errors="replace") as f:
+141 f.write(content)
+ +143 # Load the index from the temporary directory
+144 storage_context = StorageContext.from_defaults(persist_dir=temp_dir)
+145 index = load_index_from_storage(storage_context)
+ +147 return index
+148 except Exception as e:
+149 # If deserialization fails, return None
+150 # The index will be rebuilt from documents in _post_init
+151 logger.warning(f"Failed to deserialize index: {e}. Index will be rebuilt.")
+152 return None
+ +154 def retrieve_relevant(self, relevance_target:str, top_k=20) -> list:
+155 """
+156 Retrieves all values from memory that are relevant to a given target.
+157 """
+158 # Handle empty or None query
+159 if not relevance_target or not relevance_target.strip():
+160 return []
+ +162 if self.index is not None:
+163 retriever = self.index.as_retriever(similarity_top_k=top_k)
+164 nodes = retriever.retrieve(relevance_target)
+165 else:
+166 nodes = []
+ +168 retrieved = []
+169 for node in nodes:
+170 content = "SOURCE: " + node.metadata.get('file_name', '(unknown)')
+171 content += "\n" + "SIMILARITY SCORE:" + str(node.score)
+172 content += "\n" + "RELEVANT CONTENT:" + node.text
+173 retrieved.append(content)
+ +175 logger.debug(f"Content retrieved: {content[:200]}")
+ +177 return retrieved
+ +179 def retrieve_by_name(self, name:str) -> list:
+180 """
+181 Retrieves a content source by its name.
+182 """
+183 # TODO also optionally provide a relevance target?
+184 results = []
+185 if self.name_to_document is not None and name in self.name_to_document:
+186 docs = self.name_to_document[name]
+187 for i, doc in enumerate(docs):
+188 if doc is not None:
+189 content = f"SOURCE: {name}\n"
+190 content += f"PAGE: {i}\n"
+191 content += "CONTENT: \n" + doc.text[:10000] # TODO a more intelligent way to limit the content
+192 results.append(content)
+ +194 return results
+ + +197 def list_sources(self) -> list:
+198 """
+199 Lists the names of the available content sources.
+200 """
+201 if self.name_to_document is not None:
+202 return list(self.name_to_document.keys())
+203 else:
+204 return []
+ +206 def add_document(self, document) -> None:
+207 """
+208 Indexes a document for semantic retrieval.
+ +210 Assumes the document has a metadata field called "semantic_memory_id" that is used to identify the document within Semantic Memory.
+211 """
+212 self.add_documents([document])
+ +214 def add_documents(self, new_documents) -> list:
+215 """
+216 Indexes documents for semantic retrieval.
+217 """
+218 # index documents by name
+219 if len(new_documents) > 0:
+ +221 # process documents individually too
+222 for document in new_documents:
+223 logger.debug(f"Adding document {document} to index, text is: {document.text}")
+ +225 # out of an abundance of caution, we sanitize the text
+226 document.text = utils.sanitize_raw_string(document.text)
+ +228 logger.debug(f"Document text after sanitization: {document.text}")
+ +230 # add the new document to the list of documents after all sanitization and checks
+231 self.documents.append(document)
+ +233 if document.metadata.get("semantic_memory_id") is not None:
+234 # if the document has a semantic memory ID, we use it as the identifier
+235 name = document.metadata["semantic_memory_id"]
+ +237 # Ensure name_to_document is initialized
+238 if not hasattr(self, 'name_to_document') or self.name_to_document is None:
+239 self.name_to_document = {}
+ +241 # self.name_to_document[name] contains a list, since each source file could be split into multiple pages
+242 if name in self.name_to_document:
+243 self.name_to_document[name].append(document)
+244 else:
+245 self.name_to_document[name] = [document]
+ + +248 # index documents for semantic retrieval
+249 if self.index is None:
+250 # Create storage context with vector store
+251 vector_store = SimpleVectorStore()
+252 storage_context = StorageContext.from_defaults(vector_store=vector_store)
+ +254 self.index = VectorStoreIndex.from_documents(
+255 self.documents,
+256 storage_context=storage_context,
+257 store_nodes_override=True # This ensures nodes (with text) are stored
+258 )
+259 else:
+260 self.index.refresh(self.documents)
+ +262 @staticmethod
+263 def _set_internal_id_to_documents(documents:list, external_attribute_name:str ="file_name") -> None:
+264 """
+265 Sets the internal ID for each document in the list of documents.
+266 This is useful to ensure that each document has a unique identifier.
+267 """
+268 for doc in documents:
+269 if not hasattr(doc, 'metadata'):
+270 doc.metadata = {}
+271 doc.metadata["semantic_memory_id"] = doc.metadata.get(external_attribute_name, doc.id_)
+ +273 return documents
+ + +276@utils.post_init
+277class LocalFilesGroundingConnector(BaseSemanticGroundingConnector):
+ +279 serializable_attributes = ["folders_paths"]
+ +281 def __init__(self, name:str="Local Files", folders_paths: list=None) -> None:
+282 super().__init__(name)
+ +284 self.folders_paths = folders_paths
+ +286 # @post_init ensures that _post_init is called after the __init__ method
+ +288 def _post_init(self):
+289 """
+290 This will run after __init__, since the class has the @post_init decorator.
+291 It is convenient to separate some of the initialization processes to make deserialize easier.
+292 """
+293 self.loaded_folders_paths = []
+ +295 if not hasattr(self, 'folders_paths') or self.folders_paths is None:
+296 self.folders_paths = []
+ +298 self.add_folders(self.folders_paths)
+ +300 def add_folders(self, folders_paths:list) -> None:
+301 """
+302 Adds a path to a folder with files used for grounding.
+303 """
+ +305 if folders_paths is not None:
+306 for folder_path in folders_paths:
+307 try:
+308 logger.debug(f"Adding the following folder to grounding index: {folder_path}")
+309 self.add_folder(folder_path)
+310 except (FileNotFoundError, ValueError) as e:
+311 print(f"Error: {e}")
+312 print(f"Current working directory: {os.getcwd()}")
+313 print(f"Provided path: {folder_path}")
+314 print("Please check if the path exists and is accessible.")
+ +316 def add_folder(self, folder_path:str) -> None:
+317 """
+318 Adds a path to a folder with files used for grounding.
+319 """
+ +321 if folder_path not in self.loaded_folders_paths:
+322 self._mark_folder_as_loaded(folder_path)
+ +324 # for PDF files, please note that the document will be split into pages: https://github.com/run-llama/llama_index/issues/15903
+325 new_files = SimpleDirectoryReader(folder_path).load_data()
+326 BaseSemanticGroundingConnector._set_internal_id_to_documents(new_files, "file_name")
+ +328 self.add_documents(new_files)
+ +330 def add_file_path(self, file_path:str) -> None:
+331 """
+332 Adds a path to a file used for grounding.
+333 """
+334 # a trick to make SimpleDirectoryReader work with a single file
+335 new_files = SimpleDirectoryReader(input_files=[file_path]).load_data()
+ +337 logger.debug(f"Adding the following file to grounding index: {new_files}")
+338 BaseSemanticGroundingConnector._set_internal_id_to_documents(new_files, "file_name")
+ +340 def _mark_folder_as_loaded(self, folder_path:str) -> None:
+341 if folder_path not in self.loaded_folders_paths:
+342 self.loaded_folders_paths.append(folder_path)
+ +344 if folder_path not in self.folders_paths:
+345 self.folders_paths.append(folder_path)
+ + + + +350@utils.post_init
+351class WebPagesGroundingConnector(BaseSemanticGroundingConnector):
+ +353 serializable_attributes = ["web_urls"]
+ +355 def __init__(self, name:str="Web Pages", web_urls: list=None) -> None:
+356 super().__init__(name)
+ +358 self.web_urls = web_urls
+ +360 # @post_init ensures that _post_init is called after the __init__ method
+ +362 def _post_init(self):
+363 self.loaded_web_urls = []
+ +365 if not hasattr(self, 'web_urls') or self.web_urls is None:
+366 self.web_urls = []
+ +368 # load web urls
+369 self.add_web_urls(self.web_urls)
+ +371 def add_web_urls(self, web_urls:list) -> None:
+372 """
+373 Adds the data retrieved from the specified URLs to grounding.
+374 """
+375 filtered_web_urls = [url for url in web_urls if url not in self.loaded_web_urls]
+376 for url in filtered_web_urls:
+377 self._mark_web_url_as_loaded(url)
+ +379 if len(filtered_web_urls) > 0:
+380 new_documents = SimpleWebPageReader(html_to_text=True).load_data(filtered_web_urls)
+381 BaseSemanticGroundingConnector._set_internal_id_to_documents(new_documents, "url")
+382 self.add_documents(new_documents)
+ +384 def add_web_url(self, web_url:str) -> None:
+385 """
+386 Adds the data retrieved from the specified URL to grounding.
+387 """
+388 # we do it like this because the add_web_urls could run scrapes in parallel, so it is better
+389 # to implement this one in terms of the other
+390 self.add_web_urls([web_url])
+ +392 def _mark_web_url_as_loaded(self, web_url:str) -> None:
+393 if web_url not in self.loaded_web_urls:
+394 self.loaded_web_urls.append(web_url)
+ +396 if web_url not in self.web_urls:
+397 self.web_urls.append(web_url)
+ +