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/opensearch.py
+73 -91
View File
@@ -24,7 +24,7 @@ from open_webui.config import (
class OpenSearchClient(VectorDBBase):
def __init__(self):
self.index_prefix = "open_webui"
self.index_prefix = 'open_webui'
self.client = OpenSearch(
hosts=[OPENSEARCH_URI],
use_ssl=OPENSEARCH_SSL,
@@ -33,25 +33,25 @@ class OpenSearchClient(VectorDBBase):
)
def _get_index_name(self, collection_name: str) -> str:
return f"{self.index_prefix}_{collection_name}"
return f'{self.index_prefix}_{collection_name}'
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])
def _result_to_search_result(self, result) -> SearchResult:
if not result["hits"]["hits"]:
if not result['hits']['hits']:
return None
ids = []
@@ -59,11 +59,11 @@ class OpenSearchClient(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],
@@ -74,33 +74,31 @@ class OpenSearchClient(VectorDBBase):
def _create_index(self, collection_name: str, dimension: int):
body = {
"settings": {"index": {"knn": True}},
"mappings": {
"properties": {
"id": {"type": "keyword"},
"vector": {
"type": "knn_vector",
"dimension": dimension, # Adjust based on your vector dimensions
"index": True,
"similarity": "faiss",
"method": {
"name": "hnsw",
"space_type": "innerproduct", # Use inner product to approximate cosine similarity
"engine": "faiss",
"parameters": {
"ef_construction": 128,
"m": 16,
'settings': {'index': {'knn': True}},
'mappings': {
'properties': {
'id': {'type': 'keyword'},
'vector': {
'type': 'knn_vector',
'dimension': dimension, # Adjust based on your vector dimensions
'index': True,
'similarity': 'faiss',
'method': {
'name': 'hnsw',
'space_type': 'innerproduct', # Use inner product to approximate cosine similarity
'engine': 'faiss',
'parameters': {
'ef_construction': 128,
'm': 16,
},
},
},
"text": {"type": "text"},
"metadata": {"type": "object"},
'text': {'type': 'text'},
'metadata': {'type': 'object'},
}
},
}
self.client.indices.create(
index=self._get_index_name(collection_name), body=body
)
self.client.indices.create(index=self._get_index_name(collection_name), body=body)
def _create_batches(self, items: list[VectorItem], batch_size=100):
for i in range(0, len(items), batch_size):
@@ -128,46 +126,40 @@ class OpenSearchClient(VectorDBBase):
return None
query = {
"size": limit,
"_source": ["text", "metadata"],
"query": {
"script_score": {
"query": {"match_all": {}},
"script": {
"source": "(cosineSimilarity(params.query_value, doc[params.field]) + 1.0) / 2.0",
"params": {
"field": "vector",
"query_value": vectors[0],
'size': limit,
'_source': ['text', 'metadata'],
'query': {
'script_score': {
'query': {'match_all': {}},
'script': {
'source': '(cosineSimilarity(params.query_value, doc[params.field]) + 1.0) / 2.0',
'params': {
'field': 'vector',
'query_value': vectors[0],
}, # Assuming single query vector
},
}
},
}
result = self.client.search(
index=self._get_index_name(collection_name), body=query
)
result = self.client.search(index=self._get_index_name(collection_name), body=query)
return self._result_to_search_result(result)
except Exception as e:
return None
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": {"metadata." + str(field) + ".keyword": value}}
)
query_body['query']['bool']['filter'].append({'term': {'metadata.' + str(field) + '.keyword': value}})
size = limit if limit else 10000
@@ -188,28 +180,24 @@ class OpenSearchClient(VectorDBBase):
self._create_index(collection_name, dimension)
def get(self, collection_name: str) -> Optional[GetResult]:
query = {"query": {"match_all": {}}, "_source": ["text", "metadata"]}
query = {'query': {'match_all': {}}, '_source': ['text', 'metadata']}
result = self.client.search(
index=self._get_index_name(collection_name), body=query
)
result = self.client.search(index=self._get_index_name(collection_name), body=query)
return self._result_to_get_result(result)
def insert(self, collection_name: str, items: list[VectorItem]):
self._create_index_if_not_exists(
collection_name=collection_name, dimension=len(items[0]["vector"])
)
self._create_index_if_not_exists(collection_name=collection_name, dimension=len(items[0]['vector']))
for batch in self._create_batches(items):
actions = [
{
"_op_type": "index",
"_index": self._get_index_name(collection_name),
"_id": item["id"],
"_source": {
"vector": item["vector"],
"text": item["text"],
"metadata": process_metadata(item["metadata"]),
'_op_type': 'index',
'_index': self._get_index_name(collection_name),
'_id': item['id'],
'_source': {
'vector': item['vector'],
'text': item['text'],
'metadata': process_metadata(item['metadata']),
},
}
for item in batch
@@ -218,22 +206,20 @@ class OpenSearchClient(VectorDBBase):
self.client.indices.refresh(index=self._get_index_name(collection_name))
def upsert(self, collection_name: str, items: list[VectorItem]):
self._create_index_if_not_exists(
collection_name=collection_name, dimension=len(items[0]["vector"])
)
self._create_index_if_not_exists(collection_name=collection_name, dimension=len(items[0]['vector']))
for batch in self._create_batches(items):
actions = [
{
"_op_type": "update",
"_index": self._get_index_name(collection_name),
"_id": item["id"],
"doc": {
"vector": item["vector"],
"text": item["text"],
"metadata": process_metadata(item["metadata"]),
'_op_type': 'update',
'_index': self._get_index_name(collection_name),
'_id': item['id'],
'doc': {
'vector': item['vector'],
'text': item['text'],
'metadata': process_metadata(item['metadata']),
},
"doc_as_upsert": True,
'doc_as_upsert': True,
}
for item in batch
]
@@ -249,27 +235,23 @@ class OpenSearchClient(VectorDBBase):
if ids:
actions = [
{
"_op_type": "delete",
"_index": self._get_index_name(collection_name),
"_id": id,
'_op_type': 'delete',
'_index': self._get_index_name(collection_name),
'_id': id,
}
for id in ids
]
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(
{"term": {"metadata." + str(field) + ".keyword": value}}
)
self.client.delete_by_query(
index=self._get_index_name(collection_name), body=query_body
)
query_body['query']['bool']['filter'].append({'term': {'metadata.' + str(field) + '.keyword': value}})
self.client.delete_by_query(index=self._get_index_name(collection_name), body=query_body)
self.client.indices.refresh(index=self._get_index_name(collection_name))
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)