feat(trace): opentelemetry instrument

backend/open_webui/retrieval/loaders/tavily.py

@@ -9,18 +9,20 @@ from open_webui.env import SRC_LOG_LEVELS
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
 
+
 class TavilyLoader(BaseLoader):
     """Extract web page content from URLs using Tavily Extract API.
-    
+
     This is a LangChain document loader that uses Tavily's Extract API to
     retrieve content from web pages and return it as Document objects.
-    
+
     Args:
         urls: URL or list of URLs to extract content from.
         api_key: The Tavily API key.
         extract_depth: Depth of extraction, either "basic" or "advanced".
         continue_on_failure: Whether to continue if extraction of a URL fails.
     """
+
     def __init__(
         self,
         urls: Union[str, List[str]],
@@ -29,13 +31,13 @@ class TavilyLoader(BaseLoader):
         continue_on_failure: bool = True,
     ) -> None:
         """Initialize Tavily Extract client.
-        
+
         Args:
             urls: URL or list of URLs to extract content from.
             api_key: The Tavily API key.
             include_images: Whether to include images in the extraction.
             extract_depth: Depth of extraction, either "basic" or "advanced".
-                advanced extraction retrieves more data, including tables and 
+                advanced extraction retrieves more data, including tables and
                 embedded content, with higher success but may increase latency.
                 basic costs 1 credit per 5 successful URL extractions,
                 advanced costs 2 credits per 5 successful URL extractions.
@@ -43,35 +45,28 @@ class TavilyLoader(BaseLoader):
         """
         if not urls:
             raise ValueError("At least one URL must be provided.")
-            
+
         self.api_key = api_key
         self.urls = urls if isinstance(urls, list) else [urls]
         self.extract_depth = extract_depth
         self.continue_on_failure = continue_on_failure
         self.api_url = "https://api.tavily.com/extract"
-        
+
     def lazy_load(self) -> Iterator[Document]:
         """Extract and yield documents from the URLs using Tavily Extract API."""
         batch_size = 20
         for i in range(0, len(self.urls), batch_size):
-            batch_urls = self.urls[i:i + batch_size]
+            batch_urls = self.urls[i : i + batch_size]
             try:
                 headers = {
                     "Content-Type": "application/json",
-                    "Authorization": f"Bearer {self.api_key}"
+                    "Authorization": f"Bearer {self.api_key}",
                 }
                 # Use string for single URL, array for multiple URLs
                 urls_param = batch_urls[0] if len(batch_urls) == 1 else batch_urls
-                payload = {
-                    "urls": urls_param,
-                    "extract_depth": self.extract_depth
-                }
+                payload = {"urls": urls_param, "extract_depth": self.extract_depth}
                 # Make the API call
-                response = requests.post(
-                    self.api_url,
-                    headers=headers,
-                    json=payload
-                )
+                response = requests.post(self.api_url, headers=headers, json=payload)
                 response.raise_for_status()
                 response_data = response.json()
                 # Process successful results
@@ -95,4 +90,4 @@ class TavilyLoader(BaseLoader):
                 if self.continue_on_failure:
                     log.error(f"Error extracting content from batch {batch_urls}: {e}")
                 else:
-                    raise e
+                    raise e