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

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/vector/dbs/elasticsearch.py
+79 -98
View File
@@ -51,7 +51,7 @@ class ElasticsearchClient(VectorDBBase):
# Status: works
def _get_index_name(self, dimension: int) -> str:
return f"{self.index_prefix}_d{str(dimension)}"
return f'{self.index_prefix}_d{str(dimension)}'
# Status: works
def _scan_result_to_get_result(self, result) -> GetResult:
@@ -62,24 +62,24 @@ class ElasticsearchClient(VectorDBBase):
metadatas = []
for hit in result:
ids.append(hit["_id"])
documents.append(hit["_source"].get("text"))
metadatas.append(hit["_source"].get("metadata"))
ids.append(hit['_id'])
documents.append(hit['_source'].get('text'))
metadatas.append(hit['_source'].get('metadata'))
return GetResult(ids=[ids], documents=[documents], metadatas=[metadatas])
# Status: works
def _result_to_get_result(self, result) -> GetResult:
if not result["hits"]["hits"]:
if not result['hits']['hits']:
return None
ids = []
documents = []
metadatas = []
for hit in result["hits"]["hits"]:
ids.append(hit["_id"])
documents.append(hit["_source"].get("text"))
metadatas.append(hit["_source"].get("metadata"))
for hit in result['hits']['hits']:
ids.append(hit['_id'])
documents.append(hit['_source'].get('text'))
metadatas.append(hit['_source'].get('metadata'))
return GetResult(ids=[ids], documents=[documents], metadatas=[metadatas])
@@ -90,11 +90,11 @@ class ElasticsearchClient(VectorDBBase):
documents = []
metadatas = []
for hit in result["hits"]["hits"]:
ids.append(hit["_id"])
distances.append(hit["_score"])
documents.append(hit["_source"].get("text"))
metadatas.append(hit["_source"].get("metadata"))
for hit in result['hits']['hits']:
ids.append(hit['_id'])
distances.append(hit['_score'])
documents.append(hit['_source'].get('text'))
metadatas.append(hit['_source'].get('metadata'))
return SearchResult(
ids=[ids],
@@ -106,26 +106,26 @@ class ElasticsearchClient(VectorDBBase):
# Status: works
def _create_index(self, dimension: int):
body = {
"mappings": {
"dynamic_templates": [
'mappings': {
'dynamic_templates': [
{
"strings": {
"match_mapping_type": "string",
"mapping": {"type": "keyword"},
'strings': {
'match_mapping_type': 'string',
'mapping': {'type': 'keyword'},
}
}
],
"properties": {
"collection": {"type": "keyword"},
"id": {"type": "keyword"},
"vector": {
"type": "dense_vector",
"dims": dimension, # Adjust based on your vector dimensions
"index": True,
"similarity": "cosine",
'properties': {
'collection': {'type': 'keyword'},
'id': {'type': 'keyword'},
'vector': {
'type': 'dense_vector',
'dims': dimension, # Adjust based on your vector dimensions
'index': True,
'similarity': 'cosine',
},
"text": {"type": "text"},
"metadata": {"type": "object"},
'text': {'type': 'text'},
'metadata': {'type': 'object'},
},
}
}
@@ -139,21 +139,19 @@ class ElasticsearchClient(VectorDBBase):
# Status: works
def has_collection(self, collection_name) -> bool:
query_body = {"query": {"bool": {"filter": []}}}
query_body["query"]["bool"]["filter"].append(
{"term": {"collection": collection_name}}
)
query_body = {'query': {'bool': {'filter': []}}}
query_body['query']['bool']['filter'].append({'term': {'collection': collection_name}})
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:
return None
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)
query = {'query': {'term': {'collection': collection_name}}}
self.client.delete_by_query(index=f'{self.index_prefix}*', body=query)
# Status: works
def search(
@@ -164,51 +162,41 @@ class ElasticsearchClient(VectorDBBase):
limit: int = 10,
) -> Optional[SearchResult]:
query = {
"size": limit,
"_source": ["text", "metadata"],
"query": {
"script_score": {
"query": {
"bool": {"filter": [{"term": {"collection": collection_name}}]}
},
"script": {
"source": "cosineSimilarity(params.vector, 'vector') + 1.0",
"params": {
"vector": vectors[0]
}, # Assuming single query vector
'size': limit,
'_source': ['text', 'metadata'],
'query': {
'script_score': {
'query': {'bool': {'filter': [{'term': {'collection': collection_name}}]}},
'script': {
'source': "cosineSimilarity(params.vector, 'vector') + 1.0",
'params': {'vector': vectors[0]}, # Assuming single query vector
},
}
},
}
result = self.client.search(
index=self._get_index_name(len(vectors[0])), body=query
)
result = self.client.search(index=self._get_index_name(len(vectors[0])), body=query)
return self._result_to_search_result(result)
# Status: only tested halfwat
def query(
self, collection_name: str, filter: dict, limit: Optional[int] = None
) -> Optional[GetResult]:
def query(self, collection_name: str, filter: dict, limit: Optional[int] = None) -> Optional[GetResult]:
if not self.has_collection(collection_name):
return None
query_body = {
"query": {"bool": {"filter": []}},
"_source": ["text", "metadata"],
'query': {'bool': {'filter': []}},
'_source': ['text', 'metadata'],
}
for field, value in filter.items():
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': {field: value}})
query_body['query']['bool']['filter'].append({'term': {'collection': collection_name}})
size = limit if limit else 10
try:
result = self.client.search(
index=f"{self.index_prefix}*",
index=f'{self.index_prefix}*',
body=query_body,
size=size,
)
@@ -220,9 +208,7 @@ class ElasticsearchClient(VectorDBBase):
# Status: works
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):
if not self._has_index(dimension=dimension):
@@ -232,28 +218,28 @@ class ElasticsearchClient(VectorDBBase):
def get(self, collection_name: str) -> Optional[GetResult]:
# Get all the items in the collection.
query = {
"query": {"bool": {"filter": [{"term": {"collection": collection_name}}]}},
"_source": ["text", "metadata"],
'query': {'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)
# Status: works
def insert(self, collection_name: str, items: list[VectorItem]):
if not self._has_index(dimension=len(items[0]["vector"])):
self._create_index(dimension=len(items[0]["vector"]))
if not self._has_index(dimension=len(items[0]['vector'])):
self._create_index(dimension=len(items[0]['vector']))
for batch in self._create_batches(items):
actions = [
{
"_index": self._get_index_name(dimension=len(items[0]["vector"])),
"_id": item["id"],
"_source": {
"collection": collection_name,
"vector": item["vector"],
"text": item["text"],
"metadata": process_metadata(item["metadata"]),
'_index': self._get_index_name(dimension=len(items[0]['vector'])),
'_id': item['id'],
'_source': {
'collection': collection_name,
'vector': item['vector'],
'text': item['text'],
'metadata': process_metadata(item['metadata']),
},
}
for item in batch
@@ -262,21 +248,21 @@ class ElasticsearchClient(VectorDBBase):
# Upsert documents using the update API with doc_as_upsert=True.
def upsert(self, collection_name: str, items: list[VectorItem]):
if not self._has_index(dimension=len(items[0]["vector"])):
self._create_index(dimension=len(items[0]["vector"]))
if not self._has_index(dimension=len(items[0]['vector'])):
self._create_index(dimension=len(items[0]['vector']))
for batch in self._create_batches(items):
actions = [
{
"_op_type": "update",
"_index": self._get_index_name(dimension=len(item["vector"])),
"_id": item["id"],
"doc": {
"collection": collection_name,
"vector": item["vector"],
"text": item["text"],
"metadata": process_metadata(item["metadata"]),
'_op_type': 'update',
'_index': self._get_index_name(dimension=len(item['vector'])),
'_id': item['id'],
'doc': {
'collection': collection_name,
'vector': item['vector'],
'text': item['text'],
'metadata': process_metadata(item['metadata']),
},
"doc_as_upsert": True,
'doc_as_upsert': True,
}
for item in batch
]
@@ -289,22 +275,17 @@ class ElasticsearchClient(VectorDBBase):
ids: Optional[list[str]] = None,
filter: Optional[dict] = None,
):
query = {
"query": {"bool": {"filter": [{"term": {"collection": collection_name}}]}}
}
query = {'query': {'bool': {'filter': [{'term': {'collection': collection_name}}]}}}
# logic based on chromaDB
if ids:
query["query"]["bool"]["filter"].append({"terms": {"_id": 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}}
)
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)
def reset(self):
indices = self.client.indices.get(index=f"{self.index_prefix}*")
indices = self.client.indices.get(index=f'{self.index_prefix}*')
for index in indices:
self.client.indices.delete(index=index)