forked from open-webui/open-webui
513 lines
15 KiB
Python
513 lines
15 KiB
Python
import os
|
|
import logging
|
|
import requests
|
|
|
|
from typing import List
|
|
|
|
from apps.ollama.main import (
|
|
generate_ollama_embeddings,
|
|
GenerateEmbeddingsForm,
|
|
)
|
|
|
|
from huggingface_hub import snapshot_download
|
|
|
|
from langchain_core.documents import Document
|
|
from langchain_community.retrievers import BM25Retriever
|
|
from langchain.retrievers import (
|
|
ContextualCompressionRetriever,
|
|
EnsembleRetriever,
|
|
)
|
|
|
|
from typing import Optional
|
|
from config import SRC_LOG_LEVELS, CHROMA_CLIENT
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
log.setLevel(SRC_LOG_LEVELS["RAG"])
|
|
|
|
|
|
def query_doc(
|
|
collection_name: str,
|
|
query: str,
|
|
embedding_function,
|
|
k: int,
|
|
):
|
|
try:
|
|
collection = CHROMA_CLIENT.get_collection(name=collection_name)
|
|
query_embeddings = embedding_function(query)
|
|
|
|
result = collection.query(
|
|
query_embeddings=[query_embeddings],
|
|
n_results=k,
|
|
)
|
|
|
|
log.info(f"query_doc:result {result}")
|
|
return result
|
|
except Exception as e:
|
|
raise e
|
|
|
|
|
|
def query_doc_with_hybrid_search(
|
|
collection_name: str,
|
|
query: str,
|
|
embedding_function,
|
|
k: int,
|
|
reranking_function,
|
|
r: int,
|
|
):
|
|
try:
|
|
collection = CHROMA_CLIENT.get_collection(name=collection_name)
|
|
documents = collection.get() # get all documents
|
|
|
|
bm25_retriever = BM25Retriever.from_texts(
|
|
texts=documents.get("documents"),
|
|
metadatas=documents.get("metadatas"),
|
|
)
|
|
bm25_retriever.k = k
|
|
|
|
chroma_retriever = ChromaRetriever(
|
|
collection=collection,
|
|
embedding_function=embedding_function,
|
|
top_n=k,
|
|
)
|
|
|
|
ensemble_retriever = EnsembleRetriever(
|
|
retrievers=[bm25_retriever, chroma_retriever], weights=[0.5, 0.5]
|
|
)
|
|
|
|
compressor = RerankCompressor(
|
|
embedding_function=embedding_function,
|
|
top_n=k,
|
|
reranking_function=reranking_function,
|
|
r_score=r,
|
|
)
|
|
|
|
compression_retriever = ContextualCompressionRetriever(
|
|
base_compressor=compressor, base_retriever=ensemble_retriever
|
|
)
|
|
|
|
result = compression_retriever.invoke(query)
|
|
result = {
|
|
"distances": [[d.metadata.get("score") for d in result]],
|
|
"documents": [[d.page_content for d in result]],
|
|
"metadatas": [[d.metadata for d in result]],
|
|
}
|
|
|
|
log.info(f"query_doc_with_hybrid_search:result {result}")
|
|
return result
|
|
except Exception as e:
|
|
raise e
|
|
|
|
|
|
def merge_and_sort_query_results(query_results, k, reverse=False):
|
|
# Initialize lists to store combined data
|
|
combined_distances = []
|
|
combined_documents = []
|
|
combined_metadatas = []
|
|
|
|
for data in query_results:
|
|
combined_distances.extend(data["distances"][0])
|
|
combined_documents.extend(data["documents"][0])
|
|
combined_metadatas.extend(data["metadatas"][0])
|
|
|
|
# Create a list of tuples (distance, document, metadata)
|
|
combined = list(zip(combined_distances, combined_documents, combined_metadatas))
|
|
|
|
# Sort the list based on distances
|
|
combined.sort(key=lambda x: x[0], reverse=reverse)
|
|
|
|
# We don't have anything :-(
|
|
if not combined:
|
|
sorted_distances = []
|
|
sorted_documents = []
|
|
sorted_metadatas = []
|
|
else:
|
|
# Unzip the sorted list
|
|
sorted_distances, sorted_documents, sorted_metadatas = zip(*combined)
|
|
|
|
# Slicing the lists to include only k elements
|
|
sorted_distances = list(sorted_distances)[:k]
|
|
sorted_documents = list(sorted_documents)[:k]
|
|
sorted_metadatas = list(sorted_metadatas)[:k]
|
|
|
|
# Create the output dictionary
|
|
result = {
|
|
"distances": [sorted_distances],
|
|
"documents": [sorted_documents],
|
|
"metadatas": [sorted_metadatas],
|
|
}
|
|
|
|
return result
|
|
|
|
|
|
def query_collection(
|
|
collection_names: List[str],
|
|
query: str,
|
|
embedding_function,
|
|
k: int,
|
|
):
|
|
results = []
|
|
for collection_name in collection_names:
|
|
try:
|
|
result = query_doc(
|
|
collection_name=collection_name,
|
|
query=query,
|
|
k=k,
|
|
embedding_function=embedding_function,
|
|
)
|
|
results.append(result)
|
|
except:
|
|
pass
|
|
return merge_and_sort_query_results(results, k=k)
|
|
|
|
|
|
def query_collection_with_hybrid_search(
|
|
collection_names: List[str],
|
|
query: str,
|
|
embedding_function,
|
|
k: int,
|
|
reranking_function,
|
|
r: float,
|
|
):
|
|
results = []
|
|
for collection_name in collection_names:
|
|
try:
|
|
result = query_doc_with_hybrid_search(
|
|
collection_name=collection_name,
|
|
query=query,
|
|
embedding_function=embedding_function,
|
|
k=k,
|
|
reranking_function=reranking_function,
|
|
r=r,
|
|
)
|
|
results.append(result)
|
|
except:
|
|
pass
|
|
return merge_and_sort_query_results(results, k=k, reverse=True)
|
|
|
|
|
|
def rag_template(template: str, context: str, query: str):
|
|
template = template.replace("[context]", context)
|
|
template = template.replace("[query]", query)
|
|
return template
|
|
|
|
|
|
def get_embedding_function(
|
|
embedding_engine,
|
|
embedding_model,
|
|
embedding_function,
|
|
openai_key,
|
|
openai_url,
|
|
):
|
|
if embedding_engine == "":
|
|
return lambda query: embedding_function.encode(query).tolist()
|
|
elif embedding_engine in ["ollama", "openai"]:
|
|
if embedding_engine == "ollama":
|
|
func = lambda query: generate_ollama_embeddings(
|
|
GenerateEmbeddingsForm(
|
|
**{
|
|
"model": embedding_model,
|
|
"prompt": query,
|
|
}
|
|
)
|
|
)
|
|
elif embedding_engine == "openai":
|
|
func = lambda query: generate_openai_embeddings(
|
|
model=embedding_model,
|
|
text=query,
|
|
key=openai_key,
|
|
url=openai_url,
|
|
)
|
|
|
|
def generate_multiple(query, f):
|
|
if isinstance(query, list):
|
|
return [f(q) for q in query]
|
|
else:
|
|
return f(query)
|
|
|
|
return lambda query: generate_multiple(query, func)
|
|
|
|
|
|
def rag_messages(
|
|
docs,
|
|
messages,
|
|
template,
|
|
embedding_function,
|
|
k,
|
|
reranking_function,
|
|
r,
|
|
hybrid_search,
|
|
):
|
|
log.debug(f"docs: {docs} {messages} {embedding_function} {reranking_function}")
|
|
|
|
last_user_message_idx = None
|
|
for i in range(len(messages) - 1, -1, -1):
|
|
if messages[i]["role"] == "user":
|
|
last_user_message_idx = i
|
|
break
|
|
|
|
user_message = messages[last_user_message_idx]
|
|
|
|
if isinstance(user_message["content"], list):
|
|
# Handle list content input
|
|
content_type = "list"
|
|
query = ""
|
|
for content_item in user_message["content"]:
|
|
if content_item["type"] == "text":
|
|
query = content_item["text"]
|
|
break
|
|
elif isinstance(user_message["content"], str):
|
|
# Handle text content input
|
|
content_type = "text"
|
|
query = user_message["content"]
|
|
else:
|
|
# Fallback in case the input does not match expected types
|
|
content_type = None
|
|
query = ""
|
|
|
|
extracted_collections = []
|
|
relevant_contexts = []
|
|
|
|
for doc in docs:
|
|
context = None
|
|
|
|
collection = doc.get("collection_name")
|
|
if collection:
|
|
collection = [collection]
|
|
else:
|
|
collection = doc.get("collection_names", [])
|
|
|
|
collection = set(collection).difference(extracted_collections)
|
|
if not collection:
|
|
log.debug(f"skipping {doc} as it has already been extracted")
|
|
continue
|
|
|
|
try:
|
|
if doc["type"] == "text":
|
|
context = doc["content"]
|
|
else:
|
|
if hybrid_search:
|
|
context = query_collection_with_hybrid_search(
|
|
collection_names=(
|
|
doc["collection_names"]
|
|
if doc["type"] == "collection"
|
|
else [doc["collection_name"]]
|
|
),
|
|
query=query,
|
|
embedding_function=embedding_function,
|
|
k=k,
|
|
reranking_function=reranking_function,
|
|
r=r,
|
|
)
|
|
else:
|
|
context = query_collection(
|
|
collection_names=(
|
|
doc["collection_names"]
|
|
if doc["type"] == "collection"
|
|
else [doc["collection_name"]]
|
|
),
|
|
query=query,
|
|
embedding_function=embedding_function,
|
|
k=k,
|
|
)
|
|
except Exception as e:
|
|
log.exception(e)
|
|
context = None
|
|
|
|
if context:
|
|
relevant_contexts.append(context)
|
|
|
|
extracted_collections.extend(collection)
|
|
|
|
context_string = ""
|
|
for context in relevant_contexts:
|
|
items = context["documents"][0]
|
|
context_string += "\n\n".join(items)
|
|
context_string = context_string.strip()
|
|
|
|
ra_content = rag_template(
|
|
template=template,
|
|
context=context_string,
|
|
query=query,
|
|
)
|
|
|
|
log.debug(f"ra_content: {ra_content}")
|
|
|
|
if content_type == "list":
|
|
new_content = []
|
|
for content_item in user_message["content"]:
|
|
if content_item["type"] == "text":
|
|
# Update the text item's content with ra_content
|
|
new_content.append({"type": "text", "text": ra_content})
|
|
else:
|
|
# Keep other types of content as they are
|
|
new_content.append(content_item)
|
|
new_user_message = {**user_message, "content": new_content}
|
|
else:
|
|
new_user_message = {
|
|
**user_message,
|
|
"content": ra_content,
|
|
}
|
|
|
|
messages[last_user_message_idx] = new_user_message
|
|
|
|
return messages
|
|
|
|
|
|
def get_model_path(model: str, update_model: bool = False):
|
|
# Construct huggingface_hub kwargs with local_files_only to return the snapshot path
|
|
cache_dir = os.getenv("SENTENCE_TRANSFORMERS_HOME")
|
|
|
|
local_files_only = not update_model
|
|
|
|
snapshot_kwargs = {
|
|
"cache_dir": cache_dir,
|
|
"local_files_only": local_files_only,
|
|
}
|
|
|
|
log.debug(f"model: {model}")
|
|
log.debug(f"snapshot_kwargs: {snapshot_kwargs}")
|
|
|
|
# Inspiration from upstream sentence_transformers
|
|
if (
|
|
os.path.exists(model)
|
|
or ("\\" in model or model.count("/") > 1)
|
|
and local_files_only
|
|
):
|
|
# If fully qualified path exists, return input, else set repo_id
|
|
return model
|
|
elif "/" not in model:
|
|
# Set valid repo_id for model short-name
|
|
model = "sentence-transformers" + "/" + model
|
|
|
|
snapshot_kwargs["repo_id"] = model
|
|
|
|
# Attempt to query the huggingface_hub library to determine the local path and/or to update
|
|
try:
|
|
model_repo_path = snapshot_download(**snapshot_kwargs)
|
|
log.debug(f"model_repo_path: {model_repo_path}")
|
|
return model_repo_path
|
|
except Exception as e:
|
|
log.exception(f"Cannot determine model snapshot path: {e}")
|
|
return model
|
|
|
|
|
|
def generate_openai_embeddings(
|
|
model: str, text: str, key: str, url: str = "https://api.openai.com/v1"
|
|
):
|
|
try:
|
|
r = requests.post(
|
|
f"{url}/embeddings",
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {key}",
|
|
},
|
|
json={"input": text, "model": model},
|
|
)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
if "data" in data:
|
|
return data["data"][0]["embedding"]
|
|
else:
|
|
raise "Something went wrong :/"
|
|
except Exception as e:
|
|
print(e)
|
|
return None
|
|
|
|
|
|
from typing import Any
|
|
|
|
from langchain_core.retrievers import BaseRetriever
|
|
from langchain_core.callbacks import CallbackManagerForRetrieverRun
|
|
|
|
|
|
class ChromaRetriever(BaseRetriever):
|
|
collection: Any
|
|
embedding_function: Any
|
|
top_n: int
|
|
|
|
def _get_relevant_documents(
|
|
self,
|
|
query: str,
|
|
*,
|
|
run_manager: CallbackManagerForRetrieverRun,
|
|
) -> List[Document]:
|
|
query_embeddings = self.embedding_function(query)
|
|
|
|
results = self.collection.query(
|
|
query_embeddings=[query_embeddings],
|
|
n_results=self.top_n,
|
|
)
|
|
|
|
ids = results["ids"][0]
|
|
metadatas = results["metadatas"][0]
|
|
documents = results["documents"][0]
|
|
|
|
results = []
|
|
for idx in range(len(ids)):
|
|
results.append(
|
|
Document(
|
|
metadata=metadatas[idx],
|
|
page_content=documents[idx],
|
|
)
|
|
)
|
|
return results
|
|
|
|
|
|
import operator
|
|
|
|
from typing import Optional, Sequence
|
|
|
|
from langchain_core.documents import BaseDocumentCompressor, Document
|
|
from langchain_core.callbacks import Callbacks
|
|
from langchain_core.pydantic_v1 import Extra
|
|
|
|
from sentence_transformers import util
|
|
|
|
|
|
class RerankCompressor(BaseDocumentCompressor):
|
|
embedding_function: Any
|
|
top_n: int
|
|
reranking_function: Any
|
|
r_score: float
|
|
|
|
class Config:
|
|
extra = Extra.forbid
|
|
arbitrary_types_allowed = True
|
|
|
|
def compress_documents(
|
|
self,
|
|
documents: Sequence[Document],
|
|
query: str,
|
|
callbacks: Optional[Callbacks] = None,
|
|
) -> Sequence[Document]:
|
|
reranking = self.reranking_function is not None
|
|
|
|
if reranking:
|
|
scores = self.reranking_function.predict(
|
|
[(query, doc.page_content) for doc in documents]
|
|
)
|
|
else:
|
|
query_embedding = self.embedding_function(query)
|
|
document_embedding = self.embedding_function(
|
|
[doc.page_content for doc in documents]
|
|
)
|
|
scores = util.cos_sim(query_embedding, document_embedding)[0]
|
|
|
|
docs_with_scores = list(zip(documents, scores.tolist()))
|
|
if self.r_score:
|
|
docs_with_scores = [
|
|
(d, s) for d, s in docs_with_scores if s >= self.r_score
|
|
]
|
|
|
|
result = sorted(docs_with_scores, key=operator.itemgetter(1), reverse=True)
|
|
final_results = []
|
|
for doc, doc_score in result[: self.top_n]:
|
|
metadata = doc.metadata
|
|
metadata["score"] = doc_score
|
|
doc = Document(
|
|
page_content=doc.page_content,
|
|
metadata=metadata,
|
|
)
|
|
final_results.append(doc)
|
|
return final_results
|