sunholo 0.114.2__tar.gz → 0.115.1__tar.gz

{sunholo-0.114.2 → sunholo-0.115.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: sunholo
-Version: 0.114.2
+Version: 0.115.1
 Summary: Large Language Model DevOps - a package to help deploy LLMs to the Cloud.
 Home-page: https://github.com/sunholo-data/sunholo-py
-Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.114.2.tar.gz
+Download-URL: https://github.com/sunholo-data/sunholo-py/archive/refs/tags/v0.115.1.tar.gz
 Author: Holosun ApS
 Author-email: multivac@sunholo.com
 License: Apache License, Version 2.0
@@ -18,13 +18,14 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
+Requires-Dist: aiohttp
 Requires-Dist: google-auth
+Requires-Dist: pydantic
+Requires-Dist: requests
 Requires-Dist: ruamel.yaml
-Requires-Dist: langchain==0.2.16
-Requires-Dist: langchain_experimental==0.0.65
-Requires-Dist: langchain-community==0.2.17
-Requires-Dist: langsmith==0.1.143
+Requires-Dist: tenacity
 Provides-Extra: all
+Requires-Dist: aiohttp; extra == "all"
 Requires-Dist: anthropic[vertex]; extra == "all"
 Requires-Dist: asyncpg; extra == "all"
 Requires-Dist: azure-identity; extra == "all"
@@ -68,10 +69,12 @@ Requires-Dist: pillow; extra == "all"
 Requires-Dist: playwright; extra == "all"
 Requires-Dist: psutil; extra == "all"
 Requires-Dist: psycopg2-binary; extra == "all"
+Requires-Dist: pydantic; extra == "all"
 Requires-Dist: pypdf; extra == "all"
 Requires-Dist: python-hcl2; extra == "all"
 Requires-Dist: python-socketio; extra == "all"
 Requires-Dist: pytesseract; extra == "all"
+Requires-Dist: requests; extra == "all"
 Requires-Dist: rich; extra == "all"
 Requires-Dist: sounddevice; extra == "all"
 Requires-Dist: supabase; extra == "all"
@@ -82,10 +85,10 @@ Requires-Dist: tiktoken; extra == "all"
 Requires-Dist: unstructured[all-docs,local-inference]; extra == "all"
 Requires-Dist: xlwings; extra == "all"
 Provides-Extra: langchain
-Requires-Dist: langchain==0.2.16; extra == "langchain"
-Requires-Dist: langchain_experimental==0.0.65; extra == "langchain"
-Requires-Dist: langchain-community==0.2.17; extra == "langchain"
-Requires-Dist: langsmith==0.1.143; extra == "langchain"
+Requires-Dist: langchain; extra == "langchain"
+Requires-Dist: langchain_experimental; extra == "langchain"
+Requires-Dist: langchain-community; extra == "langchain"
+Requires-Dist: langsmith; extra == "langchain"
 Provides-Extra: azure
 Requires-Dist: azure-identity; extra == "azure"
 Requires-Dist: azure-storage-blob; extra == "azure"
@@ -104,6 +107,7 @@ Requires-Dist: tantivy; extra == "database"
 Provides-Extra: pipeline
 Requires-Dist: GitPython; extra == "pipeline"
 Requires-Dist: lark; extra == "pipeline"
+Requires-Dist: langchain>=0.2.16; extra == "pipeline"
 Requires-Dist: langchain-unstructured; extra == "pipeline"
 Requires-Dist: psutil; extra == "pipeline"
 Requires-Dist: pypdf; extra == "pipeline"
@@ -125,8 +129,8 @@ Requires-Dist: google-cloud-logging; extra == "gcp"
 Requires-Dist: google-cloud-pubsub; extra == "gcp"
 Requires-Dist: google-cloud-discoveryengine; extra == "gcp"
 Requires-Dist: google-cloud-texttospeech; extra == "gcp"
-Requires-Dist: google-generativeai>=0.7.1; extra == "gcp"
-Requires-Dist: langchain-google-genai==1.0.10; extra == "gcp"
+Requires-Dist: google-generativeai>=0.8.3; extra == "gcp"
+Requires-Dist: langchain-google-genai>=2.0.0; extra == "gcp"
 Requires-Dist: langchain_google_alloydb_pg>=0.2.2; extra == "gcp"
 Requires-Dist: langchain-google-vertexai; extra == "gcp"
 Requires-Dist: pillow; extra == "gcp"

{sunholo-0.114.2 → sunholo-0.115.1}/setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-version = '0.114.2'
+version = '0.115.1'
 
 setup(
     name='sunholo',
@@ -27,16 +27,17 @@ setup(
     },
     install_requires=[
         # Base dependencies
+        "aiohttp",
         "google-auth", # to check if on gcp
+        "pydantic",
+        "requests",
         "ruamel.yaml",
-        "langchain==0.2.16",
-        "langchain_experimental==0.0.65",
-        "langchain-community==0.2.17",
-        "langsmith==0.1.143",
+        "tenacity"
     ],
     extras_require={
         # Define optional dependencies with feature names
         'all': [
+            "aiohttp",
             "anthropic[vertex]",
             "asyncpg",
             "azure-identity",
@@ -80,10 +81,12 @@ setup(
             "playwright",
             "psutil",
             "psycopg2-binary",
+            "pydantic",
             "pypdf",
             "python-hcl2",
             "python-socketio",
             "pytesseract",
+            "requests",
             "rich",
             "sounddevice",
             "supabase",
@@ -95,10 +98,10 @@ setup(
             "xlwings"
         ],
         'langchain': [
-            "langchain==0.2.16",
-            "langchain_experimental==0.0.65",
-            "langchain-community==0.2.17",
-            "langsmith==0.1.143",
+            "langchain",
+            "langchain_experimental",
+            "langchain-community",
+            "langsmith",
         ],
         'azure': [
             "azure-identity",
@@ -121,6 +124,7 @@ setup(
         'pipeline': [
             "GitPython",
             "lark",
+            "langchain>=0.2.16",
             "langchain-unstructured",
             "psutil",
             "pypdf",
@@ -143,8 +147,8 @@ setup(
             "google-cloud-pubsub",
             "google-cloud-discoveryengine",
             "google-cloud-texttospeech",
-            "google-generativeai>=0.7.1",
-            "langchain-google-genai==1.0.10",
+            "google-generativeai>=0.8.3",
+            "langchain-google-genai>=2.0.0",
             "langchain_google_alloydb_pg>=0.2.2",
             "langchain-google-vertexai",
             "pillow",
@@ -189,10 +193,10 @@ setup(
     },
     classifiers=[
         'Development Status :: 3 - Alpha',      # Chose either "3 - Alpha", "4 - Beta" or "5 - Production/Stable" as the current state of your package
-        'Intended Audience :: Developers',      # Define that your audience are developers
+        'Intended Audience :: Developers',     
         'Topic :: Software Development :: Build Tools',
         'License :: OSI Approved :: Apache Software License', 
-        'Programming Language :: Python :: 3',      #Specify which python versions that you want to support
+        'Programming Language :: Python :: 3', 
         'Programming Language :: Python :: 3.10',
         'Programming Language :: Python :: 3.11',
         'Programming Language :: Python :: 3.12'

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/__init__.py

@@ -15,7 +15,6 @@ from . import invoke
 from . import langfuse
 from . import llamaindex
 from . import lookup
-from . import patches
 from . import pubsub
 from . import qna
 from . import senses
@@ -44,7 +43,6 @@ __all__ = ['agents',
            'langfuse', 
            'llamaindex',
            'lookup', 
-           'patches', 
            'pubsub',
            'qna', 
            'senses',

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/chunker/message_data.py

@@ -29,10 +29,7 @@ try:
 except ImportError:
     BlobServiceClient = None
 
-try:
-    from langchain.schema import Document
-except ImportError:
-    Document = None
+from ..types import Document
 
 from .splitter import chunk_doc_to_docs
 from .pdfs import split_pdf_to_pages

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/chunker/pdfs.py

@@ -14,6 +14,7 @@
 import os
 import pathlib
 from ..custom_logging import log
+from ..types import Document
 
 def split_pdf_to_pages(pdf_path, temp_dir):
 
@@ -51,7 +52,6 @@ def split_pdf_to_pages(pdf_path, temp_dir):
     return page_files
 
 def read_pdf_file(pdf_path, metadata):
-    from langchain.schema import Document
     from pypdf import PdfReader
     log.info(f"Trying local PDF parsing.  Reading PDF {pdf_path}...")
 

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/chunker/publish.py

@@ -2,11 +2,7 @@ from ..custom_logging import log
 from ..pubsub import PubSubManager
 from ..utils.parsers import contains_url, extract_urls
 from ..utils.gcp_project import get_gcp_project
-
-try:
-    from langchain.schema import Document
-except ImportError:
-    Document=None
+from ..types import Document
     
 def publish_if_urls(the_content, vector_name):
     """

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/chunker/splitter.py

@@ -28,6 +28,9 @@ except ImportError:
 def chunk_doc_to_docs(documents: list, extension: str = ".md", min_size: int = 800, vector_name=None, **kwargs):
     """Turns a Document object into a list of many Document chunks.
        If a document or chunk is smaller than min_size, it will be merged with adjacent documents or chunks."""
+    
+    if Document is None:
+        raise ImportError("Document chunker needs langchain installed via sunholo[pipeline]")
 
     if len(documents)==0:
         log.warning("No documents found to chunk in chunk_doc_to_docs")
@@ -137,6 +140,9 @@ def choose_splitter(extension: str, chunk_size: int=1024, chunk_overlap:int=200,
 
                     return semantic_splitter
 
+    if not text_splitter:
+        raise ImportError("text_splitter needs langchain installed via sunholo[pipeline]")
+
 
     if extension == ".py":
         return text_splitter.PythonCodeTextSplitter()

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/cli/cli_init.py

@@ -1,9 +1,11 @@
 import os
-import yaml
+from ruamel.yaml import YAML
 import shutil
 from ..utils.config import get_module_filepath
 from ..utils.parsers import sanitize_cloudrun_name
 
+yaml = YAML(typ='safe')
+
 def init_project(args):
     """
     Initializes a new sunholo project with a basic configuration file and directory structure.

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/components/llm.py

@@ -110,7 +110,7 @@ def llm_str_to_llm(llm_str, model=None, vector_name=None, config=None):
         return VertexAI(model_name = model, temperature=0, max_output_tokens=1024)
 
     elif llm_str == 'model_garden':
-        from ..patches.langchain.vertexai import VertexAIModelGarden
+        from langchain_google_vertexai import VertexAIModelGarden
         model_garden_config = config.vacConfig("gcp_config")
         if model_garden_config is None:
             raise ValueError("llm='model_garden' requires a gcp_config entry in config yaml file")

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/components/vectorstore.py

@@ -102,7 +102,7 @@ def pick_vectorstore(vs_str: str, vector_name: str=None, embeddings=None, config
         return vectorstore
         
     elif vs_str == "lancedb":
-        from ..patches.langchain.lancedb import LanceDB
+        from langchain_community.vectorstores import LanceDB
         import lancedb
 
         LANCEDB_BUCKET = os.environ.get("LANCEDB_BUCKET")

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/embedder/embed_chunk.py

@@ -34,6 +34,9 @@ def embed_pubsub_chunk(data: dict):
          data JSON
     """
 
+    if Document is None:
+        raise ImportError("Embeddin requires langchain installed via sunholo[pipeline]")
+    
     message_data = base64.b64decode(data['message']['data']).decode('utf-8')
     messageId = data['message'].get('messageId')
     publishTime = data['message'].get('publishTime')

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/senses/stream_voice.py

@@ -1,3 +1,12 @@
+from typing import Optional, TYPE_CHECKING, Union, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+    ArrayType = NDArray[np.int16]
+else:
+    ArrayType = Any  # Fallback type when numpy isn't available
+
 try:
     from google.cloud import texttospeech
 except ImportError:
@@ -29,9 +38,6 @@ import io
 import wave
 
 import argparse
-import json
-from typing import Optional
-from pathlib import Path
 import sys
 
 class StreamingTTS:
@@ -195,8 +201,11 @@ class StreamingTTS:
             log.error(f"Error initializing audio device: {e}")
             raise
 
-    def _make_fade(self, length: int, fade_type: str='l') -> np.ndarray:
+    def _make_fade(self, length: int, fade_type: str='l') -> ArrayType:
         """Generate a fade curve of specified length and type."""
+        if np is None:  # Runtime check
+            raise ImportError("numpy is required. Install with pip install sunholo[tts]")
+    
         fade = np.arange(length, dtype=np.float32) / length
         
         if fade_type == 't':  # triangle
@@ -214,8 +223,11 @@ class StreamingTTS:
         
         return fade
     
-    def _apply_fade(self, audio: np.ndarray, fade_duration: float, fade_in: bool = True, fade_out: bool = True) -> np.ndarray:
+    def _apply_fade(self, audio: ArrayType, fade_duration: float, fade_in: bool = True, fade_out: bool = True) -> ArrayType:
         """Apply fade in/out to audio with specified duration."""
+        if np is None:  # Runtime check
+            raise ImportError("numpy is required. Install with pip install sunholo[tts]")
+    
         if audio.ndim != 1:
             raise ValueError("Audio must be 1-dimensional")
         
@@ -233,8 +245,11 @@ class StreamingTTS:
         return audio.astype(np.int16)
 
     
-    def _play_audio_chunk(self, audio_chunk: np.ndarray, is_final_chunk: bool = False):
+    def _play_audio_chunk(self, audio_chunk: ArrayType, is_final_chunk: bool = False):
         """Play a single audio chunk with proper device handling."""
+        if np is None:  # Runtime check
+            raise ImportError("numpy is required. Install with pip install sunholo[tts]")
+    
         try:
             # Add longer padding for the final chunk
             padding_duration = 0.1 if is_final_chunk else 0.02
@@ -415,7 +430,7 @@ def tts_command(args):
                 Panel((
                     f"Saying: {args.text}"
                     ), 
-                    title=f"Text to Speech",
+                    title="Text to Speech",
                     subtitle=f"{tts.voice_name} is talking"),
                     )
             tts.process_text_stream(

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/streaming/content_buffer.py

@@ -13,13 +13,6 @@
 #   limitations under the License.
 from typing import Any, Dict, List, Union
 
-try:
-    from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-    from langchain.schema import LLMResult
-except ImportError:
-    StreamingStdOutCallbackHandler = None
-    LLMResult = None
-
 import threading
 import asyncio
 import re
@@ -114,7 +107,7 @@ class ContentBuffer:
             self.content_available.clear()
 
 
-class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
+class BufferStreamingStdOutCallbackHandler:
     """
     A callback handler for streaming LLM output to a content buffer.
 
@@ -212,12 +205,12 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
                 self.content_buffer.write(self.buffer)
                 self.buffer = ""
 
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+    def on_llm_end(self, response, **kwargs: Any) -> None:
         """
         Handles the end of LLM streaming.
 
         Args:
-            response (LLMResult): The result returned by the LLM.
+            response: The result returned by the LLM.
             **kwargs: Additional keyword arguments.
 
         Writes any remaining buffer content to the content buffer, and sets a signal indicating
@@ -233,7 +226,7 @@ class BufferStreamingStdOutCallbackHandler(StreamingStdOutCallbackHandler):
 
 
 
-class BufferStreamingStdOutCallbackHandlerAsync(StreamingStdOutCallbackHandler):
+class BufferStreamingStdOutCallbackHandlerAsync:
     """
     An async callback handler for streaming LLM output to a content buffer.
 
@@ -315,12 +308,12 @@ class BufferStreamingStdOutCallbackHandlerAsync(StreamingStdOutCallbackHandler):
                 await self.content_buffer.async_write(self.buffer)
                 self.buffer = ""
 
-    async def async_on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+    async def async_on_llm_end(self, response, **kwargs: Any) -> None:
         """
         Asynchronously handles the end of LLM streaming.
 
         Args:
-            response (LLMResult): The result returned by the LLM.
+            response: The result returned by the LLM.
             **kwargs: Additional keyword arguments.
         """
         if self.buffer:

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/summarise/summarise.py

@@ -13,8 +13,6 @@
 #   limitations under the License.
 from ..custom_logging import log
 
-
-
 from ..components import get_llm
 from ..chunker.splitter import chunk_doc_to_docs
 
@@ -36,6 +34,8 @@ except ImportError:
     VertexAI=None
     load_summarize_chain=None
     Document=None
+import time
+import random
 
 prompt_template = """Write a summary for below, including key concepts, people and distinct information but do not add anything that is not in the original text:
 
@@ -45,10 +45,10 @@ SUMMARY:"""
 MAP_PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
 
 
-import time
-import random
-
 def summarise_docs(docs, vector_name, skip_if_less=10000):
+    if Document is None:
+        raise ImportError("summarise_docs requires langchain installed via sunholo[pipeline]")
+    
     llm  = get_llm(vector_name)
   
     if isinstance(llm, ChatOpenAI) or isinstance(llm, OpenAI):

sunholo-0.115.1/sunholo/types.py

@@ -0,0 +1,52 @@
+from typing import Any, Dict, Optional, TYPE_CHECKING, Union
+from dataclasses import dataclass, asdict
+import json
+
+if TYPE_CHECKING:
+    from langchain.schema import Document as LangchainDocument
+
+
+@dataclass
+class Document:
+    """A simple document class with content and metadata.
+    
+    Used for storing text content and associated metadata when not using LangChain.
+    Maintains the same basic interface (page_content and metadata) for compatibility.
+    
+    Using @dataclass makes it automatically serializable and provides nice defaults.
+    """
+    page_content: str
+    metadata: Dict[str, Any] = None
+
+    def __post_init__(self) -> None:
+        """Initialize metadata if None."""
+        if self.metadata is None:
+            self.metadata = {}
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Document":
+        """Create from dictionary."""
+        return cls(**data)
+
+    def json(self) -> str:
+        """Convert to JSON string - for compatibility with LangChain's Document."""
+        return json.dumps(self.to_dict())
+
+def convert_to_langchain_doc(doc: Document) -> Union[Any, "LangchainDocument"]:
+    """Convert our Document to a LangChain Document.
+    
+    Returns Any when LangChain isn't available to avoid type errors.
+    Only imports LangChain when the function is actually called.
+    """
+    try:
+        from langchain.schema import Document as LangchainDocument
+        return LangchainDocument(
+            page_content=doc.page_content,
+            metadata=doc.metadata
+        )
+    except ImportError:
+        raise ImportError("LangChain is required for this conversion. Please install langchain.")

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/utils/config.py

@@ -14,11 +14,12 @@
 
 import os
 import json
-import yaml
 from datetime import datetime, timedelta
 from collections import defaultdict
 from .timedelta import format_timedelta
 
+from ruamel.yaml import YAML
+yaml = YAML(typ='safe')
 
 def get_module_filepath(filepath: str):
     """
@@ -107,7 +108,7 @@ def reload_config_file(config_file, filename):
         if filename.endswith('.json'):
             config = json.load(file)
         else:
-            config = yaml.safe_load(file)
+            config = yaml.load(file)
     
     config_cache[filename] = (config, datetime.now())
     log.debug(f"Loaded and cached {config_file}")
@@ -166,7 +167,7 @@ def load_config(filename: str=None) -> tuple[dict, str]:
         if filename.endswith(".json"):
             config = json.load(f)
         elif filename.endswith(".yaml") or filename.endswith(".yml"):
-            config = yaml.safe_load(f)
+            config = yaml.load(f)
         else:
             raise ValueError(f"Unsupported config file format: {config_file}. The supported formats are JSON and YAML.")
 

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/utils/config_class.py

@@ -1,9 +1,7 @@
 import os
 import json
-import yaml
 from datetime import datetime, timedelta
 from collections import defaultdict
-from yaml.constructor import SafeConstructor, ConstructorError
 
 from .timedelta import format_timedelta
 
@@ -122,30 +120,32 @@ class ConfigManager:
         Returns:
             dict: The loaded configuration.
         """
-        from ..custom_logging import log
+    def _reload_config_file(self, config_file, filename, is_local=False):
+        """
+        Helper function to load a config file and update the cache.
 
-        class NoDuplicateKeyConstructor(SafeConstructor):
-            def construct_mapping(self, node, deep=False):
-                mapping = {}
-                for key_node, value_node in node.value:
-                    key = self.construct_object(key_node, deep=deep)
-                    if key in mapping:
-                        raise ConstructorError(f"Duplicate key found: {key_node.start_mark}")
-                    value = self.construct_object(value_node, deep=deep)
-                    mapping[key] = value
-                return mapping
-
-        NoDuplicateKeyLoader = yaml.Loader
-        NoDuplicateKeyLoader.add_constructor(
-            yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
-            NoDuplicateKeyConstructor.construct_mapping
-        )
+        Args:
+            config_file (str): The path to the configuration file.
+            filename (str): The name of the configuration file.
+            is_local (bool): Indicates if the config file is from the local folder.
+
+        Returns:
+            dict: The loaded configuration.
+        """
+        from ..custom_logging import log
+        from ruamel.yaml import YAML, DuplicateKeyError
 
         with open(config_file, 'r') as file:
             if filename.endswith('.json'):
                 config = json.load(file)
             else:
-                config = yaml.load(file, Loader=NoDuplicateKeyLoader)
+                # Create YAML parser that forbids duplicates
+                yaml = YAML(typ='safe')
+                yaml.allow_duplicate_keys = False
+                try:
+                    config = yaml.load(file)
+                except DuplicateKeyError as e:
+                    raise ValueError(f"Duplicate key found in {filename}: {str(e)}")
 
         self.config_cache[filename] = (config, datetime.now())
         if is_local:

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/utils/gcp.py

@@ -14,9 +14,6 @@
 import os
 import requests
 import socket
-# can't install due to circular import sunholo.logging
-import logging
-
 
 def is_running_on_cloudrun():
     """

{sunholo-0.114.2 → sunholo-0.115.1}/sunholo/vertex/extensions_class.py

@@ -13,6 +13,8 @@ import json
 from io import StringIO
 import os
 import re
+from ruamel.yaml import YAML
+yaml = YAML(typ='safe')
 
 class VertexAIExtensions:
     """
@@ -115,13 +117,12 @@ class VertexAIExtensions:
     def upload_openapi_file(self, filename: str, extension_name:str, vac:str=None):
         if vac:
             from ..agents.route import route_vac
-            import yaml
 
             new_url = route_vac(vac)
 
             log.info(f'Overwriting extension URL with VAC url for {vac=} - {new_url=}')
             
-            openapi = yaml.safe_load(filename)
+            openapi = yaml.load(filename)
 
             openapi['servers'][0]['url'] = new_url
             with open(filename, 'w') as file:
@@ -145,10 +146,9 @@ class VertexAIExtensions:
         return self.current_extension.api_spec()
         
     def load_tool_use_examples(self, filename: str):
-        import yaml
 
         with open(filename, 'r') as file:
-            self.tool_use_examples = yaml.safe_load(file)
+            self.tool_use_examples = yaml.load(file)
 
         # google.cloud.aiplatform_v1beta1.types.ToolUseExample
         return self.tool_use_examples