utils.py 9.8 KB
Newer Older
1
import logging
2
import requests
Steven Kreitzer's avatar
Steven Kreitzer committed
3
4
5
import operator

import sentence_transformers
6

7
from typing import List
8

9
10
11
12
from apps.ollama.main import (
    generate_ollama_embeddings,
    GenerateEmbeddingsForm,
)
Timothy J. Baek's avatar
Timothy J. Baek committed
13

Steven Kreitzer's avatar
Steven Kreitzer committed
14
15
16
17
18
from langchain.retrievers import (
    BM25Retriever,
    EnsembleRetriever,
)

19
20
from config import SRC_LOG_LEVELS, CHROMA_CLIENT

21

22
23
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
Timothy J. Baek's avatar
Timothy J. Baek committed
24
25


Steven Kreitzer's avatar
Steven Kreitzer committed
26
27
28
29
30
31
32
def query_embeddings_doc(
    collection_name: str,
    query: str,
    k: int,
    embeddings_function,
    reranking_function,
):
33
34
    try:
        # if you use docker use the model from the environment variable
35
36
        collection = CHROMA_CLIENT.get_collection(name=collection_name)

Steven Kreitzer's avatar
Steven Kreitzer committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
        # keyword search
        documents = collection.get() # get all documents
        bm25_retriever = BM25Retriever.from_texts(
            texts=documents.get("documents"),
            metadatas=documents.get("metadatas"),
        )
        bm25_retriever.k = k

        # semantic search (vector)
        chroma_retriever = ChromaRetriever(
            collection=collection,
            k=k,
            embeddings_function=embeddings_function,
        )

        # hybrid search (ensemble)
        ensemble_retriever = EnsembleRetriever(
            retrievers=[bm25_retriever, chroma_retriever],
            weights=[0.6, 0.4]
56
        )
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
57

Steven Kreitzer's avatar
Steven Kreitzer committed
58
59
60
61
62
63
64
65
66
67
68
69
70
        documents = ensemble_retriever.invoke(query)
        result = query_results_rank(
            query=query,
            documents=documents,
            k=k,
            reranking_function=reranking_function,
        )
        result = {
            "distances": [[d[1].item() for d in result]],
            "documents": [[d[0].page_content for d in result]],
            "metadatas": [[d[0].metadata for d in result]],
        }

71
72
73
74
75
        return result
    except Exception as e:
        raise e


Steven Kreitzer's avatar
Steven Kreitzer committed
76
77
78
79
80
81
82
def query_results_rank(query: str, documents, k: int, reranking_function):
    scores = reranking_function.predict([(query, doc.page_content) for doc in documents])
    docs_with_scores = list(zip(documents, scores))
    result = sorted(docs_with_scores, key=operator.itemgetter(1), reverse=True)
    return result[: k]


Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
85
86
def merge_and_sort_query_results(query_results, k):
    # Initialize lists to store combined data
    combined_distances = []
    combined_documents = []
Steven Kreitzer's avatar
Steven Kreitzer committed
87
    combined_metadatas = []
Timothy J. Baek's avatar
Timothy J. Baek committed
88
89
90
91
92

    # Combine data from each dictionary
    for data in query_results:
        combined_distances.extend(data["distances"][0])
        combined_documents.extend(data["documents"][0])
Steven Kreitzer's avatar
Steven Kreitzer committed
93
        combined_metadatas.extend(data["metadatas"][0])
Timothy J. Baek's avatar
Timothy J. Baek committed
94

Steven Kreitzer's avatar
Steven Kreitzer committed
95
    # Create a list of tuples (distance, document, metadata)
Timothy J. Baek's avatar
Timothy J. Baek committed
96
    combined = list(
Steven Kreitzer's avatar
Steven Kreitzer committed
97
        zip(combined_distances, combined_documents, combined_metadatas)
Timothy J. Baek's avatar
Timothy J. Baek committed
98
99
100
101
102
103
    )

    # Sort the list based on distances
    combined.sort(key=lambda x: x[0])

    # Unzip the sorted list
Steven Kreitzer's avatar
Steven Kreitzer committed
104
    sorted_distances, sorted_documents, sorted_metadatas = zip(*combined)
Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
107
108

    # Slicing the lists to include only k elements
    sorted_distances = list(sorted_distances)[:k]
    sorted_documents = list(sorted_documents)[:k]
Steven Kreitzer's avatar
Steven Kreitzer committed
109
    sorted_metadatas = list(sorted_metadatas)[:k]
Timothy J. Baek's avatar
Timothy J. Baek committed
110
111
112
113
114

    # Create the output dictionary
    merged_query_results = {
        "distances": [sorted_distances],
        "documents": [sorted_documents],
Steven Kreitzer's avatar
Steven Kreitzer committed
115
        "metadatas": [sorted_metadatas],
Timothy J. Baek's avatar
Timothy J. Baek committed
116
117
118
119
120
121
122
123
        "embeddings": None,
        "uris": None,
        "data": None,
    }

    return merged_query_results


124
def query_embeddings_collection(
Steven Kreitzer's avatar
Steven Kreitzer committed
125
126
127
128
129
    collection_names: List[str],
    query: str,
    k: int,
    embeddings_function,
    reranking_function,
Timothy J. Baek's avatar
Timothy J. Baek committed
130
131
):

132
    results = []
Timothy J. Baek's avatar
Timothy J. Baek committed
133

134
135
    for collection_name in collection_names:
        try:
136
137
138
139
            result = query_embeddings_doc(
                collection_name=collection_name,
                query=query,
                k=k,
Steven Kreitzer's avatar
Steven Kreitzer committed
140
141
                embeddings_function=embeddings_function,
                reranking_function=reranking_function,
142
143
144
145
146
147
148
149
            )
            results.append(result)
        except:
            pass

    return merge_and_sort_query_results(results, k)


Timothy J. Baek's avatar
Timothy J. Baek committed
150
def rag_template(template: str, context: str, query: str):
151
152
    template = template.replace("[context]", context)
    template = template.replace("[query]", query)
Timothy J. Baek's avatar
Timothy J. Baek committed
153
    return template
Timothy J. Baek's avatar
Timothy J. Baek committed
154
155


Steven Kreitzer's avatar
Steven Kreitzer committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def query_embeddings_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 == "ollama":
        return lambda query: generate_ollama_embeddings(
            GenerateEmbeddingsForm(
                **{
                    "model": embedding_model,
                    "prompt": query,
                }
            )
        )
    elif embedding_engine == "openai":
        return lambda query: generate_openai_embeddings(
            model=embedding_model,
            text=query,
            key=openai_key,
            url=openai_url,
        )


183
184
185
186
187
188
189
190
def rag_messages(
    docs,
    messages,
    template,
    k,
    embedding_engine,
    embedding_model,
    embedding_function,
Steven Kreitzer's avatar
Steven Kreitzer committed
191
    reranking_function,
192
193
194
    openai_key,
    openai_url,
):
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
195
    log.debug(
Steven Kreitzer's avatar
Steven Kreitzer committed
196
        f"docs: {docs} {messages} {embedding_engine} {embedding_model} {embedding_function} {reranking_function} {openai_key} {openai_url}"
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
197
    )
Timothy J. Baek's avatar
Timothy J. Baek committed
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229

    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 = ""

    relevant_contexts = []

    for doc in docs:
        context = None

        try:
230
231

            if doc["type"] == "text":
232
                context = doc["content"]
Timothy J. Baek's avatar
Timothy J. Baek committed
233
            else:
Steven Kreitzer's avatar
Steven Kreitzer committed
234
235
236
237
238
239
240
                embeddings_function = query_embeddings_function(
                    embedding_engine,
                    embedding_model,
                    embedding_function,
                    openai_key,
                    openai_url,
                )
241
242
243
244
245
246

                if doc["type"] == "collection":
                    context = query_embeddings_collection(
                        collection_names=doc["collection_names"],
                        query=query,
                        k=k,
Steven Kreitzer's avatar
Steven Kreitzer committed
247
248
                        embeddings_function=embeddings_function,
                        reranking_function=reranking_function,
249
                    )
250
                else:
251
252
253
254
                    context = query_embeddings_doc(
                        collection_name=doc["collection_name"],
                        query=query,
                        k=k,
Steven Kreitzer's avatar
Steven Kreitzer committed
255
256
                        embeddings_function=embeddings_function,
                        reranking_function=reranking_function,
257
                    )
258

Timothy J. Baek's avatar
Timothy J. Baek committed
259
        except Exception as e:
260
            log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
261
262
263
264
            context = None

        relevant_contexts.append(context)

Timothy J. Baek's avatar
Timothy J. Baek committed
265
266
    log.debug(f"relevant_contexts: {relevant_contexts}")

Timothy J. Baek's avatar
Timothy J. Baek committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
    context_string = ""
    for context in relevant_contexts:
        if context:
            context_string += " ".join(context["documents"][0]) + "\n"

    ra_content = rag_template(
        template=template,
        context=context_string,
        query=query,
    )

    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
297

Self Denial's avatar
Self Denial committed
298

299
def generate_openai_embeddings(
Timothy J. Baek's avatar
Timothy J. Baek committed
300
    model: str, text: str, key: str, url: str = "https://api.openai.com/v1"
301
302
303
):
    try:
        r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
304
            f"{url}/embeddings",
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
            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
Steven Kreitzer's avatar
Steven Kreitzer committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357


from typing import Any

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever


class ChromaRetriever(BaseRetriever):
    collection: Any
    k: int
    embeddings_function: Any

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        query_embeddings = self.embeddings_function(query)

        results = self.collection.query(
            query_embeddings=[query_embeddings],
            n_results=self.k,
        )

        ids = results["ids"][0]
        metadatas = results["metadatas"][0]
        documents = results["documents"][0]

        return [
            Document(
                metadata=metadatas[idx],
                page_content=documents[idx],
            )
            for idx in range(len(ids))
        ]