OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

fixed es bugs

Browse Source
This commit is contained in:
ofek 2025-03-05 23:19:56 +02:00
parent 1639fbb544
commit a8f205213c
2 changed files with 137 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/open_webui/config.py
View File

@ -1553,7 +1553,7 @@ ELASTICSEARCH_USERNAME = os.environ.get("ELASTICSEARCH_USERNAME", None)
ELASTICSEARCH_PASSWORD = os.environ.get("ELASTICSEARCH_PASSWORD", None) ELASTICSEARCH_PASSWORD = os.environ.get("ELASTICSEARCH_PASSWORD", None)
ELASTICSEARCH_CLOUD_ID = os.environ.get("ELASTICSEARCH_CLOUD_ID", None) ELASTICSEARCH_CLOUD_ID = os.environ.get("ELASTICSEARCH_CLOUD_ID", None)
SSL_ASSERT_FINGERPRINT = os.environ.get("SSL_ASSERT_FINGERPRINT", None) SSL_ASSERT_FINGERPRINT = os.environ.get("SSL_ASSERT_FINGERPRINT", None)
ELASTICSEARCH_INDEX_PREFIX = os.environ.get("ELASTICSEARCH_INDEX_PREFIX", "open_webui_collections")
# Pgvector # Pgvector
PGVECTOR_DB_URL = os.environ.get("PGVECTOR_DB_URL", DATABASE_URL) PGVECTOR_DB_URL = os.environ.get("PGVECTOR_DB_URL", DATABASE_URL)
if VECTOR_DB == "pgvector" and not PGVECTOR_DB_URL.startswith("postgres"): if VECTOR_DB == "pgvector" and not PGVECTOR_DB_URL.startswith("postgres"):

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

@ -10,10 +10,14 @@ from open_webui.config import (
ELASTICSEARCH_USERNAME, ELASTICSEARCH_USERNAME,
ELASTICSEARCH_PASSWORD, ELASTICSEARCH_PASSWORD,
ELASTICSEARCH_CLOUD_ID, ELASTICSEARCH_CLOUD_ID,
ELASTICSEARCH_INDEX_PREFIX,
SSL_ASSERT_FINGERPRINT, SSL_ASSERT_FINGERPRINT,
) )
class ElasticsearchClient: class ElasticsearchClient:
""" """
Important: Important:
@ -21,22 +25,17 @@ 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 = "open_webui_collections" self.index_prefix = ELASTICSEARCH_INDEX_PREFIX
self.client = Elasticsearch( self.client = Elasticsearch(
hosts=[ELASTICSEARCH_URL], hosts=[ELASTICSEARCH_URL],
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=( basic_auth=(ELASTICSEARCH_USERNAME,ELASTICSEARCH_PASSWORD) if ELASTICSEARCH_USERNAME and ELASTICSEARCH_PASSWORD else None,
(ELASTICSEARCH_USERNAME, ELASTICSEARCH_PASSWORD) ssl_assert_fingerprint=SSL_ASSERT_FINGERPRINT
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)}"
@ -85,16 +84,22 @@ class ElasticsearchClient:
metadatas.append(hit["_source"].get("metadata")) metadatas.append(hit["_source"].get("metadata"))
return SearchResult( return SearchResult(
ids=[ids], ids=[ids], distances=[distances], documents=[documents], metadatas=[metadatas]
distances=[distances],
documents=[documents],
metadatas=[metadatas],
) )
#Status: works #Status: works
def _create_index(self, dimension: int): def _create_index(self, dimension: int):
body = { body = {
"mappings": { "mappings": {
"dynamic_templates": [
{
"strings": {
"match_mapping_type": "string",
"mapping": {
"type": "keyword"
}
}
}
],
"properties": { "properties": {
"collection": {"type": "keyword"}, "collection": {"type": "keyword"},
"id": {"type": "keyword"}, "id": {"type": "keyword"},
@ -110,7 +115,6 @@ class ElasticsearchClient:
} }
} }
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):
@ -120,35 +124,49 @@ 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( query_body["query"]["bool"]["filter"].append({"term": {"collection": collection_name}})
{"term": {"collection": collection_name}}
)
try: try:
result = self.client.count(index=f"{self.index_prefix}*", body=query_body) result = self.client.count(
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
# @TODO: Make this delete a collection and not an index
def delete_colleciton(self, collection_name: str):
# TODO: fix this to include the dimension or a * prefix
# delete_collection here means delete a bunch of documents for an index.
# We are simply adapting to the norms of the other DBs.
self.client.indices.delete(index=self._get_collection_name(collection_name))
def delete_collection(self, collection_name: str):
query = {
"query": {
"term": {"collection": collection_name}
}
}
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": ["text", "metadata"], "_source": [
"text",
"metadata"
],
"query": { "query": {
"script_score": { "script_score": {
"query": { "query": {
"bool": {"filter": [{"term": {"collection": collection_name}}]} "bool": {
"filter": [
{
"term": {
"collection": collection_name
}
}
]
}
}, },
"script": { "script": {
"source": "cosineSimilarity(params.vector, 'vector') + 1.0", "source": "cosineSimilarity(params.vector, 'vector') + 1.0",
@ -165,7 +183,6 @@ 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
@ -180,9 +197,7 @@ 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( query_body["query"]["bool"]["filter"].append({"term": {"collection": collection_name}})
{"term": {"collection": collection_name}}
)
size = limit if limit else 10 size = limit if limit else 10
try: try:
@ -196,24 +211,29 @@ 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( return self.client.indices.exists(index=self._get_index_name(dimension=dimension))
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": {"bool": {"filter": [{"term": {"collection": collection_name}}]}}, "query": {
"_source": ["text", "metadata"], "bool": {
"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)
@ -223,6 +243,7 @@ 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 = [
{ {
@ -239,34 +260,55 @@ class ElasticsearchClient:
] ]
bulk(self.client,actions) bulk(self.client,actions)
# Status: should work # Upsert documents using the update API with doc_as_upsert=True.
def upsert(self, collection_name: str, items: list[VectorItem]): def upsert(self, collection_name: str, items: list[VectorItem]):
if not self._has_index(dimension=len(items[0]["vector"])): if not self._has_index(dimension=len(items[0]["vector"])):
self._create_index(collection_name, 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 = [
{ {
"_index": self._get_index_name(dimension=len(items[0]["vector"])), "_op_type": "update",
"_index": self._get_index_name(dimension=len(item["vector"])),
"_id": item["id"], "_id": item["id"],
"_source": { "doc": {
"collection": collection_name,
"vector": item["vector"], "vector": item["vector"],
"text": item["text"], "text": item["text"],
"metadata": item["metadata"], "metadata": item["metadata"],
}, },
"doc_as_upsert": True,
} }
for item in batch for item in batch
] ]
self.client.bulk(actions) bulk(self.client,actions)
# TODO: This currently deletes by * which is not always supported in ElasticSearch.
# Need to read a bit before changing. Also, need to delete from a specific collection # Delete specific documents from a collection by filtering on both collection and document IDs.
def delete(self, collection_name: str, ids: list[str]): def delete(
# Assuming ID is unique across collections and indexes self,
actions = [ collection_name: str,
{"delete": {"_index": f"{self.index_prefix}*", "_id": id}} for id in ids ids: Optional[list[str]] = None,
filter: Optional[dict] = None,
):
query = {
"query": {
"bool": {
"filter": [
{"term": {"collection": collection_name}}
] ]
self.client.bulk(body=actions) }
}
}
#logic based on chromaDB
if ids:
query["query"]["bool"]["filter"].append({"terms": {"_id": ids}})
elif filter:
for field, value in filter.items():
query["query"]["bool"]["filter"].append({"term": {f"metadata.{field}": value}})
self.client.delete_by_query(index=f"{self.index_prefix}*", body=query)
def reset(self): def reset(self):
indices = self.client.indices.get(index=f"{self.index_prefix}*") indices = self.client.indices.get(index=f"{self.index_prefix}*")