chore: format

This commit is contained in:
Timothy Jaeryang Baek
2024-11-16 23:46:12 -08:00
parent c24bc60d35
commit c338f2cae1
64 changed files with 2668 additions and 1212 deletions

View File

@@ -32,7 +32,12 @@ from open_webui.config import (
)
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import ENV, SRC_LOG_LEVELS, DEVICE_TYPE, ENABLE_FORWARD_USER_INFO_HEADERS
from open_webui.env import (
ENV,
SRC_LOG_LEVELS,
DEVICE_TYPE,
ENABLE_FORWARD_USER_INFO_HEADERS,
)
from fastapi import Depends, FastAPI, File, HTTPException, Request, UploadFile, status
from fastapi.middleware.cors import CORSMiddleware
@@ -48,7 +53,11 @@ MAX_FILE_SIZE = MAX_FILE_SIZE_MB * 1024 * 1024 # Convert MB to bytes
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["AUDIO"])
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, openapi_url="/openapi.json" if ENV == "dev" else None, redoc_url=None)
app = FastAPI(
docs_url="/docs" if ENV == "dev" else None,
openapi_url="/openapi.json" if ENV == "dev" else None,
redoc_url=None,
)
app.add_middleware(
CORSMiddleware,

View File

@@ -48,7 +48,11 @@ log.setLevel(SRC_LOG_LEVELS["IMAGES"])
IMAGE_CACHE_DIR = Path(CACHE_DIR).joinpath("./image/generations/")
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, openapi_url="/openapi.json" if ENV == "dev" else None, redoc_url=None)
app = FastAPI(
docs_url="/docs" if ENV == "dev" else None,
openapi_url="/openapi.json" if ENV == "dev" else None,
redoc_url=None,
)
app.add_middleware(
CORSMiddleware,

View File

@@ -127,7 +127,11 @@ from langchain_core.documents import Document
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
app = FastAPI(docs_url="/docs" if ENV == "dev" else None, openapi_url="/openapi.json" if ENV == "dev" else None, redoc_url=None)
app = FastAPI(
docs_url="/docs" if ENV == "dev" else None,
openapi_url="/openapi.json" if ENV == "dev" else None,
redoc_url=None,
)
app.state.config = AppConfig()
@@ -537,7 +541,7 @@ async def update_rag_config(form_data: ConfigUpdateForm, user=Depends(get_admin_
if form_data.web is not None:
app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION = (
#Note: When UI "Bypass SSL verification for Websites"=True then ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION=False
# Note: When UI "Bypass SSL verification for Websites"=True then ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION=False
form_data.web.web_loader_ssl_verification
)
@@ -1230,7 +1234,9 @@ def process_web_search(form_data: SearchForm, user=Depends(get_verified_user)):
urls = [result.link for result in web_results]
loader = get_web_loader(urls, verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION)
loader = get_web_loader(
urls, verify_ssl=app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION
)
docs = loader.load()
save_docs_to_vector_db(docs, collection_name, overwrite=True)

View File

@@ -27,7 +27,9 @@ class ChromaClient:
if CHROMA_CLIENT_AUTH_PROVIDER is not None:
settings_dict["chroma_client_auth_provider"] = CHROMA_CLIENT_AUTH_PROVIDER
if CHROMA_CLIENT_AUTH_CREDENTIALS is not None:
settings_dict["chroma_client_auth_credentials"] = CHROMA_CLIENT_AUTH_CREDENTIALS
settings_dict["chroma_client_auth_credentials"] = (
CHROMA_CLIENT_AUTH_CREDENTIALS
)
if CHROMA_HTTP_HOST != "":
self.client = chromadb.HttpClient(

View File

@@ -7,9 +7,10 @@ from open_webui.config import (
OPENSEARCH_SSL,
OPENSEARCH_CERT_VERIFY,
OPENSEARCH_USERNAME,
OPENSEARCH_PASSWORD
OPENSEARCH_PASSWORD,
)
class OpenSearchClient:
def __init__(self):
self.index_prefix = "open_webui"
@@ -25,10 +26,10 @@ class OpenSearchClient:
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)
@@ -38,13 +39,15 @@ class OpenSearchClient:
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, distances=distances, documents=documents, metadatas=metadatas)
return SearchResult(
ids=ids, distances=distances, documents=documents, metadatas=metadatas
)
def _create_index(self, index_name: str, dimension: int):
body = {
@@ -52,20 +55,20 @@ class OpenSearchClient:
"properties": {
"id": {"type": "keyword"},
"vector": {
"type": "dense_vector",
"dims": dimension, # Adjust based on your vector dimensions
"index": true,
"similarity": "faiss",
"method": {
"type": "dense_vector",
"dims": dimension, # Adjust based on your vector dimensions
"index": true,
"similarity": "faiss",
"method": {
"name": "hnsw",
"space_type": "ip", # Use inner product to approximate cosine similarity
"engine": "faiss",
"ef_construction": 128,
"m": 16
}
"m": 16,
},
},
"text": {"type": "text"},
"metadata": {"type": "object"}
"metadata": {"type": "object"},
}
}
}
@@ -73,19 +76,21 @@ class OpenSearchClient:
def _create_batches(self, items: list[VectorItem], batch_size=100):
for i in range(0, len(items), batch_size):
yield items[i:i + batch_size]
yield items[i : i + batch_size]
def has_collection(self, index_name: str) -> bool:
# has_collection here means has index.
# has_collection here means has index.
# We are simply adapting to the norms of the other DBs.
return self.client.indices.exists(index=f"{self.index_prefix}_{index_name}")
def delete_colleciton(self, index_name: str):
# delete_collection here means delete index.
# delete_collection here means delete index.
# We are simply adapting to the norms of the other DBs.
self.client.indices.delete(index=f"{self.index_prefix}_{index_name}")
def search(self, index_name: str, vectors: list[list[float]], limit: int) -> Optional[SearchResult]:
def search(
self, index_name: str, vectors: list[list[float]], limit: int
) -> Optional[SearchResult]:
query = {
"size": limit,
"_source": ["text", "metadata"],
@@ -94,15 +99,16 @@ class OpenSearchClient:
"query": {"match_all": {}},
"script": {
"source": "cosineSimilarity(params.vector, 'vector') + 1.0",
"params": {"vector": vectors[0]} # Assuming single query vector
}
"params": {
"vector": vectors[0]
}, # Assuming single query vector
},
}
}
},
}
result = self.client.search(
index=f"{self.index_prefix}_{index_name}",
body=query
index=f"{self.index_prefix}_{index_name}", body=query
)
return self._result_to_search_result(result)
@@ -112,12 +118,11 @@ class OpenSearchClient:
self._create_index(index_name, dimension)
def get(self, index_name: str) -> Optional[GetResult]:
query = {
"query": {"match_all": {}},
"_source": ["text", "metadata"]
}
query = {"query": {"match_all": {}}, "_source": ["text", "metadata"]}
result = self.client.search(index=f"{self.index_prefix}_{index_name}", body=query)
result = self.client.search(
index=f"{self.index_prefix}_{index_name}", body=query
)
return self._result_to_get_result(result)
def insert(self, index_name: str, items: list[VectorItem]):
@@ -126,7 +131,16 @@ class OpenSearchClient:
for batch in self._create_batches(items):
actions = [
{"index": {"_id": item["id"], "_source": {"vector": item["vector"], "text": item["text"], "metadata": item["metadata"]}}}
{
"index": {
"_id": item["id"],
"_source": {
"vector": item["vector"],
"text": item["text"],
"metadata": item["metadata"],
},
}
}
for item in batch
]
self.client.bulk(actions)
@@ -137,13 +151,25 @@ class OpenSearchClient:
for batch in self._create_batches(items):
actions = [
{"index": {"_id": item["id"], "_source": {"vector": item["vector"], "text": item["text"], "metadata": item["metadata"]}}}
{
"index": {
"_id": item["id"],
"_source": {
"vector": item["vector"],
"text": item["text"],
"metadata": item["metadata"],
},
}
}
for item in batch
]
self.client.bulk(actions)
def delete(self, index_name: str, ids: list[str]):
actions = [{"delete": {"_index": f"{self.index_prefix}_{index_name}", "_id": id}} for id in ids]
actions = [
{"delete": {"_index": f"{self.index_prefix}_{index_name}", "_id": id}}
for id in ids
]
self.client.bulk(body=actions)
def reset(self):

View File

@@ -44,7 +44,9 @@ class PgvectorClient:
self.session = Session
else:
engine = create_engine(PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool)
engine = create_engine(
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
)
SessionLocal = sessionmaker(
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
)

View File

@@ -15,10 +15,11 @@ class QdrantClient:
self.collection_prefix = "open-webui"
self.QDRANT_URI = QDRANT_URI
self.QDRANT_API_KEY = QDRANT_API_KEY
self.client = Qclient(
url=self.QDRANT_URI,
api_key=self.QDRANT_API_KEY
) if self.QDRANT_URI else None
self.client = (
Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
if self.QDRANT_URI
else None
)
def _result_to_get_result(self, points) -> GetResult:
ids = []

View File

@@ -1,6 +1,5 @@
import logging
import os
import os
from pprint import pprint
from typing import Optional
import requests
@@ -10,15 +9,22 @@ import argparse
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
'''
"""
Documentation: https://docs.microsoft.com/en-us/bing/search-apis/bing-web-search/overview
'''
"""
def search_bing(
subscription_key: str, endpoint: str, locale: str, query: str, count: int, filter_list: Optional[list[str]] = None
subscription_key: str,
endpoint: str,
locale: str,
query: str,
count: int,
filter_list: Optional[list[str]] = None,
) -> list[SearchResult]:
mkt = locale
params = { 'q': query, 'mkt': mkt, 'answerCount': count }
headers = { 'Ocp-Apim-Subscription-Key': subscription_key }
params = {"q": query, "mkt": mkt, "answerCount": count}
headers = {"Ocp-Apim-Subscription-Key": subscription_key}
try:
response = requests.get(endpoint, headers=headers, params=params)
@@ -38,15 +44,30 @@ def search_bing(
except Exception as ex:
log.error(f"Error: {ex}")
raise ex
def main():
parser = argparse.ArgumentParser(description="Search Bing from the command line.")
parser.add_argument("query", type=str, default="Top 10 international news today", help="The search query.")
parser.add_argument("--count", type=int, default=10, help="Number of search results to return.")
parser.add_argument("--filter", nargs='*', help="List of filters to apply to the search results.")
parser.add_argument("--locale", type=str, default="en-US", help="The locale to use for the search, maps to market in api")
parser.add_argument(
"query",
type=str,
default="Top 10 international news today",
help="The search query.",
)
parser.add_argument(
"--count", type=int, default=10, help="Number of search results to return."
)
parser.add_argument(
"--filter", nargs="*", help="List of filters to apply to the search results."
)
parser.add_argument(
"--locale",
type=str,
default="en-US",
help="The locale to use for the search, maps to market in api",
)
args = parser.parse_args()
results = search_bing(args.locale, args.query, args.count, args.filter)
pprint(results)
pprint(results)