Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(trace): opentelemetry instrument

Browse Source
This commit is contained in:
orenzhang
2025-03-10 22:27:31 +08:00
parent c2046b9247
commit c761e4fd08
12 changed files with 337 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/loaders/tavily.py
+13 -18
View File
@@ -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
backend/open_webui/retrieval/vector/dbs/opensearch.py
+40 -51
View File
@@ -21,14 +21,14 @@ class OpenSearchClient:
verify_certs=OPENSEARCH_CERT_VERIFY,
http_auth=(OPENSEARCH_USERNAME, OPENSEARCH_PASSWORD),
)
def _get_index_name(self, collection_name: str) -> str:
return f"{self.index_prefix}_{collection_name}"
def _result_to_get_result(self, result) -> GetResult:
if not result["hits"]["hits"]:
return None
ids = []
documents = []
metadatas = []
@@ -43,7 +43,7 @@ class OpenSearchClient:
def _result_to_search_result(self, result) -> SearchResult:
if not result["hits"]["hits"]:
return None
ids = []
distances = []
documents = []
@@ -56,16 +56,15 @@ class OpenSearchClient:
metadatas.append(hit["_source"].get("metadata"))
return SearchResult(
ids=[ids], distances=[distances], documents=[documents], metadatas=[metadatas]
ids=[ids],
distances=[distances],
documents=[documents],
metadatas=[metadatas],
)
def _create_index(self, collection_name: str, dimension: int):
body = {
"settings": {
"index": {
"knn": True
}
},
"settings": {"index": {"knn": True}},
"mappings": {
"properties": {
"id": {"type": "keyword"},
@@ -81,13 +80,13 @@ class OpenSearchClient:
"parameters": {
"ef_construction": 128,
"m": 16,
}
},
},
},
"text": {"type": "text"},
"metadata": {"type": "object"},
}
}
},
}
self.client.indices.create(
index=self._get_index_name(collection_name), body=body
@@ -100,9 +99,7 @@ class OpenSearchClient:
def has_collection(self, collection_name: str) -> bool:
# has_collection here means has index.
# We are simply adapting to the norms of the other DBs.
return self.client.indices.exists(
index=self._get_index_name(collection_name)
)
return self.client.indices.exists(index=self._get_index_name(collection_name))
def delete_collection(self, collection_name: str):
# delete_collection here means delete index.
@@ -115,33 +112,30 @@ class OpenSearchClient:
try:
if not self.has_collection(collection_name):
return None
query = {
"size": limit,
"_source": ["text", "metadata"],
"query": {
"script_score": {
"query": {
"match_all": {}
},
"query": {"match_all": {}},
"script": {
"source": "cosineSimilarity(params.query_value, doc[params.field]) + 1.0",
"params": {
"field": "vector",
"query_value": vectors[0]
"field": "vector",
"query_value": vectors[0],
}, # Assuming single query vector
},
}
},
}
result = self.client.search(
index=self._get_index_name(collection_name),
body=query
index=self._get_index_name(collection_name), body=query
)
return self._result_to_search_result(result)
except Exception as e:
return None
@@ -152,20 +146,14 @@ class OpenSearchClient:
return None
query_body = {
"query": {
"bool": {
"filter": []
}
},
"query": {"bool": {"filter": []}},
"_source": ["text", "metadata"],
}
for field, value in filter.items():
query_body["query"]["bool"]["filter"].append({
"match": {
"metadata." + str(field): value
}
})
query_body["query"]["bool"]["filter"].append(
{"match": {"metadata." + str(field): value}}
)
size = limit if limit else 10
@@ -201,9 +189,9 @@ class OpenSearchClient:
for batch in self._create_batches(items):
actions = [
{
"_op_type": "index",
"_op_type": "index",
"_index": self._get_index_name(collection_name),
"_id": item["id"],
"_id": item["id"],
"_source": {
"vector": item["vector"],
"text": item["text"],
@@ -222,9 +210,9 @@ class OpenSearchClient:
for batch in self._create_batches(items):
actions = [
{
"_op_type": "update",
"_op_type": "update",
"_index": self._get_index_name(collection_name),
"_id": item["id"],
"_id": item["id"],
"doc": {
"vector": item["vector"],
"text": item["text"],
@@ -236,7 +224,12 @@ class OpenSearchClient:
]
bulk(self.client, actions)
def delete(self, collection_name: str, ids: Optional[list[str]] = None, filter: Optional[dict] = None):
def delete(
self,
collection_name: str,
ids: Optional[list[str]] = None,
filter: Optional[dict] = None,
):
if ids:
actions = [
{
@@ -249,20 +242,16 @@ class OpenSearchClient:
bulk(self.client, actions)
elif filter:
query_body = {
"query": {
"bool": {
"filter": []
}
},
"query": {"bool": {"filter": []}},
}
for field, value in filter.items():
query_body["query"]["bool"]["filter"].append({
"match": {
"metadata." + str(field): value
}
})
self.client.delete_by_query(index=self._get_index_name(collection_name), body=query_body)
query_body["query"]["bool"]["filter"].append(
{"match": {"metadata." + str(field): value}}
)
self.client.delete_by_query(
index=self._get_index_name(collection_name), body=query_body
)
def reset(self):
indices = self.client.indices.get(index=f"{self.index_prefix}_*")
for index in indices:
backend/open_webui/retrieval/web/utils.py
+8 -9
View File
@@ -136,18 +136,18 @@ class RateLimitMixin:
self.last_request_time = datetime.now()
class URLProcessingMixin:
class URLProcessingMixin:
def _verify_ssl_cert(self, url: str) -> bool:
"""Verify SSL certificate for a URL."""
return verify_ssl_cert(url)
async def _safe_process_url(self, url: str) -> bool:
"""Perform safety checks before processing a URL."""
if self.verify_ssl and not self._verify_ssl_cert(url):
raise ValueError(f"SSL certificate verification failed for {url}")
await self._wait_for_rate_limit()
return True
def _safe_process_url_sync(self, url: str) -> bool:
"""Synchronous version of safety checks."""
if self.verify_ssl and not self._verify_ssl_cert(url):
@@ -286,7 +286,7 @@ class SafeTavilyLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
proxy["server"] = env_proxy_server
else:
proxy = {"server": env_proxy_server}
# Store parameters for creating TavilyLoader instances
self.web_paths = web_paths if isinstance(web_paths, list) else [web_paths]
self.api_key = api_key
@@ -295,7 +295,7 @@ class SafeTavilyLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
self.verify_ssl = verify_ssl
self.trust_env = trust_env
self.proxy = proxy
# Add rate limiting
self.requests_per_second = requests_per_second
self.last_request_time = None
@@ -329,7 +329,7 @@ class SafeTavilyLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
log.exception(e, "Error extracting content from URLs")
else:
raise e
async def alazy_load(self) -> AsyncIterator[Document]:
"""Async version with rate limiting and SSL verification."""
valid_urls = []
@@ -341,13 +341,13 @@ class SafeTavilyLoader(BaseLoader, RateLimitMixin, URLProcessingMixin):
log.warning(f"SSL verification failed for {url}: {str(e)}")
if not self.continue_on_failure:
raise e
if not valid_urls:
if self.continue_on_failure:
log.warning("No valid URLs to process after SSL verification")
return
raise ValueError("No valid URLs to process after SSL verification")
try:
loader = TavilyLoader(
urls=valid_urls,
@@ -477,7 +477,6 @@ class SafePlaywrightURLLoader(PlaywrightURLLoader, RateLimitMixin, URLProcessing
await browser.close()
class SafeWebBaseLoader(WebBaseLoader):
"""WebBaseLoader with enhanced error handling for URLs."""