UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-04-08 22:55:18 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

chore: format

Browse Source
This commit is contained in:
Timothy Jaeryang Baek 2025-03-05 19:17:41 -08:00
parent aaaebfabbe
commit d4fca9dabf
1 changed files with 81 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

99
backend/open_webui/retrieval/vector/dbs/elasticsearch.py
View File

@ -12,12 +12,9 @@ from open_webui.config import (
ELASTICSEARCH_CLOUD_ID, ELASTICSEARCH_CLOUD_ID,
ELASTICSEARCH_INDEX_PREFIX, ELASTICSEARCH_INDEX_PREFIX,
SSL_ASSERT_FINGERPRINT, SSL_ASSERT_FINGERPRINT,
) )
class ElasticsearchClient: class ElasticsearchClient:
""" """
Important: Important:
@ -25,6 +22,7 @@ class ElasticsearchClient:
an index for each file but store it as a text field, while seperating to different index an index for each file but store it as a text field, while seperating to different index
baesd on the embedding length. baesd on the embedding length.
""" """
def __init__(self): def __init__(self):
self.index_prefix = ELASTICSEARCH_INDEX_PREFIX self.index_prefix = ELASTICSEARCH_INDEX_PREFIX
self.client = Elasticsearch( self.client = Elasticsearch(
@ -32,10 +30,14 @@ class ElasticsearchClient:
ca_certs=ELASTICSEARCH_CA_CERTS, ca_certs=ELASTICSEARCH_CA_CERTS,
api_key=ELASTICSEARCH_API_KEY, api_key=ELASTICSEARCH_API_KEY,
cloud_id=ELASTICSEARCH_CLOUD_ID, cloud_id=ELASTICSEARCH_CLOUD_ID,
basic_auth=(ELASTICSEARCH_USERNAME,ELASTICSEARCH_PASSWORD) if ELASTICSEARCH_USERNAME and ELASTICSEARCH_PASSWORD else None, basic_auth=(
ssl_assert_fingerprint=SSL_ASSERT_FINGERPRINT (ELASTICSEARCH_USERNAME, ELASTICSEARCH_PASSWORD)
if ELASTICSEARCH_USERNAME and ELASTICSEARCH_PASSWORD
else None
),
ssl_assert_fingerprint=SSL_ASSERT_FINGERPRINT,
) )
# Status: works # Status: works
def _get_index_name(self, dimension: int) -> str: def _get_index_name(self, dimension: int) -> str:
return f"{self.index_prefix}_d{str(dimension)}" return f"{self.index_prefix}_d{str(dimension)}"
@ -84,8 +86,12 @@ class ElasticsearchClient:
metadatas.append(hit["_source"].get("metadata")) metadatas.append(hit["_source"].get("metadata"))
return SearchResult( return SearchResult(
ids=[ids], distances=[distances], documents=[documents], metadatas=[metadatas] ids=[ids],
distances=[distances],
documents=[documents],
metadatas=[metadatas],
) )
# Status: works # Status: works
def _create_index(self, dimension: int): def _create_index(self, dimension: int):
body = { body = {
@ -94,9 +100,7 @@ class ElasticsearchClient:
{ {
"strings": { "strings": {
"match_mapping_type": "string", "match_mapping_type": "string",
"mapping": { "mapping": {"type": "keyword"},
"type": "keyword"
}
} }
} }
], ],
@ -111,10 +115,11 @@ class ElasticsearchClient:
}, },
"text": {"type": "text"}, "text": {"type": "text"},
"metadata": {"type": "object"}, "metadata": {"type": "object"},
} },
} }
} }
self.client.indices.create(index=self._get_index_name(dimension), body=body) self.client.indices.create(index=self._get_index_name(dimension), body=body)
# Status: works # Status: works
def _create_batches(self, items: list[VectorItem], batch_size=100): def _create_batches(self, items: list[VectorItem], batch_size=100):
@ -124,49 +129,32 @@ class ElasticsearchClient:
# Status: works # Status: works
def has_collection(self, collection_name) -> bool: def has_collection(self, collection_name) -> bool:
query_body = {"query": {"bool": {"filter": []}}} query_body = {"query": {"bool": {"filter": []}}}
query_body["query"]["bool"]["filter"].append({"term": {"collection": collection_name}}) query_body["query"]["bool"]["filter"].append(
{"term": {"collection": collection_name}}
)
try: try:
result = self.client.count( result = self.client.count(index=f"{self.index_prefix}*", body=query_body)
index=f"{self.index_prefix}*",
body=query_body
)
return result.body["count"] > 0 return result.body["count"] > 0
except Exception as e: except Exception as e:
return None return None
def delete_collection(self, collection_name: str): def delete_collection(self, collection_name: str):
query = { query = {"query": {"term": {"collection": collection_name}}}
"query": {
"term": {"collection": collection_name}
}
}
self.client.delete_by_query(index=f"{self.index_prefix}*", body=query) self.client.delete_by_query(index=f"{self.index_prefix}*", body=query)
# Status: works # Status: works
def search( def search(
self, collection_name: str, vectors: list[list[float]], limit: int self, collection_name: str, vectors: list[list[float]], limit: int
) -> Optional[SearchResult]: ) -> Optional[SearchResult]:
query = { query = {
"size": limit, "size": limit,
"_source": [ "_source": ["text", "metadata"],
"text",
"metadata"
],
"query": { "query": {
"script_score": { "script_score": {
"query": { "query": {
"bool": { "bool": {"filter": [{"term": {"collection": collection_name}}]}
"filter": [
{
"term": {
"collection": collection_name
}
}
]
}
}, },
"script": { "script": {
"source": "cosineSimilarity(params.vector, 'vector') + 1.0", "source": "cosineSimilarity(params.vector, 'vector') + 1.0",
@ -183,6 +171,7 @@ class ElasticsearchClient:
) )
return self._result_to_search_result(result) return self._result_to_search_result(result)
# Status: only tested halfwat # Status: only tested halfwat
def query( def query(
self, collection_name: str, filter: dict, limit: Optional[int] = None self, collection_name: str, filter: dict, limit: Optional[int] = None
@ -197,7 +186,9 @@ class ElasticsearchClient:
for field, value in filter.items(): for field, value in filter.items():
query_body["query"]["bool"]["filter"].append({"term": {field: value}}) query_body["query"]["bool"]["filter"].append({"term": {field: value}})
query_body["query"]["bool"]["filter"].append({"term": {"collection": collection_name}}) query_body["query"]["bool"]["filter"].append(
{"term": {"collection": collection_name}}
)
size = limit if limit else 10 size = limit if limit else 10
try: try:
@ -211,29 +202,24 @@ class ElasticsearchClient:
except Exception as e: except Exception as e:
return None return None
# Status: works # Status: works
def _has_index(self, dimension: int): def _has_index(self, dimension: int):
return self.client.indices.exists(index=self._get_index_name(dimension=dimension)) return self.client.indices.exists(
index=self._get_index_name(dimension=dimension)
)
def get_or_create_index(self, dimension: int): def get_or_create_index(self, dimension: int):
if not self._has_index(dimension=dimension): if not self._has_index(dimension=dimension):
self._create_index(dimension=dimension) self._create_index(dimension=dimension)
# Status: works # Status: works
def get(self, collection_name: str) -> Optional[GetResult]: def get(self, collection_name: str) -> Optional[GetResult]:
# Get all the items in the collection. # Get all the items in the collection.
query = { query = {
"query": { "query": {"bool": {"filter": [{"term": {"collection": collection_name}}]}},
"bool": { "_source": ["text", "metadata"],
"filter": [
{
"term": {
"collection": collection_name
} }
}
]
}
}, "_source": ["text", "metadata"]}
results = list(scan(self.client, index=f"{self.index_prefix}*", query=query)) results = list(scan(self.client, index=f"{self.index_prefix}*", query=query))
return self._scan_result_to_get_result(results) return self._scan_result_to_get_result(results)
@ -243,7 +229,6 @@ class ElasticsearchClient:
if not self._has_index(dimension=len(items[0]["vector"])): if not self._has_index(dimension=len(items[0]["vector"])):
self._create_index(dimension=len(items[0]["vector"])) self._create_index(dimension=len(items[0]["vector"]))
for batch in self._create_batches(items): for batch in self._create_batches(items):
actions = [ actions = [
{ {
@ -282,7 +267,6 @@ class ElasticsearchClient:
] ]
bulk(self.client, actions) bulk(self.client, actions)
# Delete specific documents from a collection by filtering on both collection and document IDs. # Delete specific documents from a collection by filtering on both collection and document IDs.
def delete( def delete(
self, self,
@ -292,21 +276,16 @@ class ElasticsearchClient:
): ):
query = { query = {
"query": { "query": {"bool": {"filter": [{"term": {"collection": collection_name}}]}}
"bool": {
"filter": [
{"term": {"collection": collection_name}}
]
}
}
} }
# logic based on chromaDB # logic based on chromaDB
if ids: if ids:
query["query"]["bool"]["filter"].append({"terms": {"_id": ids}}) query["query"]["bool"]["filter"].append({"terms": {"_id": ids}})
elif filter: elif filter:
for field, value in filter.items(): for field, value in filter.items():
query["query"]["bool"]["filter"].append({"term": {f"metadata.{field}": value}}) query["query"]["bool"]["filter"].append(
{"term": {f"metadata.{field}": value}}
)
self.client.delete_by_query(index=f"{self.index_prefix}*", body=query) self.client.delete_by_query(index=f"{self.index_prefix}*", body=query)