utils.py 17 KB
Newer Older
1
import os
2
import logging
3
4
import requests

5
from typing import List
6

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

12
13
from huggingface_hub import snapshot_download

14
15
from langchain_core.documents import Document
from langchain_community.retrievers import BM25Retriever
Steven Kreitzer's avatar
Steven Kreitzer committed
16
from langchain.retrievers import (
17
    ContextualCompressionRetriever,
Steven Kreitzer's avatar
Steven Kreitzer committed
18
19
20
    EnsembleRetriever,
)

21
from typing import Optional
22

23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from apps.rag.search.brave import search_brave
from apps.rag.search.google_pse import search_google_pse
from apps.rag.search.main import SearchResult
from apps.rag.search.searxng import search_searxng
from apps.rag.search.serper import search_serper
from apps.rag.search.serpstack import search_serpstack
from config import (
    SRC_LOG_LEVELS,
    CHROMA_CLIENT,
    SEARXNG_QUERY_URL,
    GOOGLE_PSE_API_KEY,
    GOOGLE_PSE_ENGINE_ID,
    BRAVE_SEARCH_API_KEY,
    SERPSTACK_API_KEY,
    SERPSTACK_HTTPS,
    SERPER_API_KEY,
)
40

41
42
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
Timothy J. Baek's avatar
Timothy J. Baek committed
43
44


Timothy J. Baek's avatar
Timothy J. Baek committed
45
def query_doc(
Steven Kreitzer's avatar
Steven Kreitzer committed
46
47
    collection_name: str,
    query: str,
Timothy J. Baek's avatar
Timothy J. Baek committed
48
    embedding_function,
49
    k: int,
Steven Kreitzer's avatar
Steven Kreitzer committed
50
):
51
    try:
Steven Kreitzer's avatar
Steven Kreitzer committed
52
        collection = CHROMA_CLIENT.get_collection(name=collection_name)
Timothy J. Baek's avatar
Timothy J. Baek committed
53
        query_embeddings = embedding_function(query)
Steven Kreitzer's avatar
Steven Kreitzer committed
54

Timothy J. Baek's avatar
Timothy J. Baek committed
55
56
57
58
        result = collection.query(
            query_embeddings=[query_embeddings],
            n_results=k,
        )
59

Timothy J. Baek's avatar
Timothy J. Baek committed
60
61
62
63
        log.info(f"query_doc:result {result}")
        return result
    except Exception as e:
        raise e
64
65


Timothy J. Baek's avatar
Timothy J. Baek committed
66
67
68
69
70
71
def query_doc_with_hybrid_search(
    collection_name: str,
    query: str,
    embedding_function,
    k: int,
    reranking_function,
tabacoWang's avatar
fix:  
tabacoWang committed
72
    r: float,
Timothy J. Baek's avatar
Timothy J. Baek committed
73
74
75
76
):
    try:
        collection = CHROMA_CLIENT.get_collection(name=collection_name)
        documents = collection.get()  # get all documents
77

Timothy J. Baek's avatar
Timothy J. Baek committed
78
79
80
81
82
        bm25_retriever = BM25Retriever.from_texts(
            texts=documents.get("documents"),
            metadatas=documents.get("metadatas"),
        )
        bm25_retriever.k = k
83

Timothy J. Baek's avatar
Timothy J. Baek committed
84
85
86
87
88
        chroma_retriever = ChromaRetriever(
            collection=collection,
            embedding_function=embedding_function,
            top_n=k,
        )
89

Timothy J. Baek's avatar
Timothy J. Baek committed
90
91
92
93
94
95
        ensemble_retriever = EnsembleRetriever(
            retrievers=[bm25_retriever, chroma_retriever], weights=[0.5, 0.5]
        )

        compressor = RerankCompressor(
            embedding_function=embedding_function,
Steven Kreitzer's avatar
Steven Kreitzer committed
96
            top_n=k,
Timothy J. Baek's avatar
Timothy J. Baek committed
97
98
99
100
101
102
103
            reranking_function=reranking_function,
            r_score=r,
        )

        compression_retriever = ContextualCompressionRetriever(
            base_compressor=compressor, base_retriever=ensemble_retriever
        )
104

Timothy J. Baek's avatar
Timothy J. Baek committed
105
106
107
108
109
110
        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]],
        }
Steven Kreitzer's avatar
Steven Kreitzer committed
111

Timothy J. Baek's avatar
Timothy J. Baek committed
112
        log.info(f"query_doc_with_hybrid_search:result {result}")
113
114
115
116
117
        return result
    except Exception as e:
        raise e


Steven Kreitzer's avatar
Steven Kreitzer committed
118
def merge_and_sort_query_results(query_results, k, reverse=False):
Timothy J. Baek's avatar
Timothy J. Baek committed
119
120
121
    # Initialize lists to store combined data
    combined_distances = []
    combined_documents = []
Steven Kreitzer's avatar
Steven Kreitzer committed
122
    combined_metadatas = []
Timothy J. Baek's avatar
Timothy J. Baek committed
123
124
125
126

    for data in query_results:
        combined_distances.extend(data["distances"][0])
        combined_documents.extend(data["documents"][0])
Steven Kreitzer's avatar
Steven Kreitzer committed
127
        combined_metadatas.extend(data["metadatas"][0])
Timothy J. Baek's avatar
Timothy J. Baek committed
128

Steven Kreitzer's avatar
Steven Kreitzer committed
129
    # Create a list of tuples (distance, document, metadata)
130
    combined = list(zip(combined_distances, combined_documents, combined_metadatas))
Timothy J. Baek's avatar
Timothy J. Baek committed
131
132

    # Sort the list based on distances
Steven Kreitzer's avatar
Steven Kreitzer committed
133
    combined.sort(key=lambda x: x[0], reverse=reverse)
Timothy J. Baek's avatar
Timothy J. Baek committed
134

135
136
137
138
139
140
141
142
    # 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)
Timothy J. Baek's avatar
Timothy J. Baek committed
143

144
145
146
147
        # 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]
Timothy J. Baek's avatar
Timothy J. Baek committed
148
149

    # Create the output dictionary
150
    result = {
Timothy J. Baek's avatar
Timothy J. Baek committed
151
152
        "distances": [sorted_distances],
        "documents": [sorted_documents],
Steven Kreitzer's avatar
Steven Kreitzer committed
153
        "metadatas": [sorted_metadatas],
Timothy J. Baek's avatar
Timothy J. Baek committed
154
155
    }

156
    return result
Timothy J. Baek's avatar
Timothy J. Baek committed
157
158


Timothy J. Baek's avatar
Timothy J. Baek committed
159
def query_collection(
Steven Kreitzer's avatar
Steven Kreitzer committed
160
161
    collection_names: List[str],
    query: str,
Timothy J. Baek's avatar
Timothy J. Baek committed
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    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,
Steven Kreitzer's avatar
Steven Kreitzer committed
184
185
    k: int,
    reranking_function,
Timothy J. Baek's avatar
Timothy J. Baek committed
186
    r: float,
Timothy J. Baek's avatar
Timothy J. Baek committed
187
):
188
189
190
    results = []
    for collection_name in collection_names:
        try:
Timothy J. Baek's avatar
Timothy J. Baek committed
191
            result = query_doc_with_hybrid_search(
192
193
                collection_name=collection_name,
                query=query,
Timothy J. Baek's avatar
Timothy J. Baek committed
194
                embedding_function=embedding_function,
195
                k=k,
Steven Kreitzer's avatar
Steven Kreitzer committed
196
                reranking_function=reranking_function,
Timothy J. Baek's avatar
Timothy J. Baek committed
197
                r=r,
198
199
200
201
            )
            results.append(result)
        except:
            pass
Timothy J. Baek's avatar
Timothy J. Baek committed
202
    return merge_and_sort_query_results(results, k=k, reverse=True)
203
204


Timothy J. Baek's avatar
Timothy J. Baek committed
205
def rag_template(template: str, context: str, query: str):
206
207
    template = template.replace("[context]", context)
    template = template.replace("[query]", query)
Timothy J. Baek's avatar
Timothy J. Baek committed
208
    return template
Timothy J. Baek's avatar
Timothy J. Baek committed
209
210


Timothy J. Baek's avatar
Timothy J. Baek committed
211
def get_embedding_function(
Steven Kreitzer's avatar
Steven Kreitzer committed
212
213
214
215
216
217
218
219
    embedding_engine,
    embedding_model,
    embedding_function,
    openai_key,
    openai_url,
):
    if embedding_engine == "":
        return lambda query: embedding_function.encode(query).tolist()
220
221
222
223
224
225
226
227
228
    elif embedding_engine in ["ollama", "openai"]:
        if embedding_engine == "ollama":
            func = lambda query: generate_ollama_embeddings(
                GenerateEmbeddingsForm(
                    **{
                        "model": embedding_model,
                        "prompt": query,
                    }
                )
Steven Kreitzer's avatar
Steven Kreitzer committed
229
            )
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
        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)
Steven Kreitzer's avatar
Steven Kreitzer committed
245
246


247
248
249
250
def rag_messages(
    docs,
    messages,
    template,
Timothy J. Baek's avatar
Timothy J. Baek committed
251
    embedding_function,
252
    k,
Timothy J. Baek's avatar
Timothy J. Baek committed
253
    reranking_function,
254
    r,
Timothy J. Baek's avatar
Timothy J. Baek committed
255
    hybrid_search,
256
):
Timothy J. Baek's avatar
Timothy J. Baek committed
257
    log.debug(f"docs: {docs} {messages} {embedding_function} {reranking_function}")
Timothy J. Baek's avatar
Timothy J. Baek committed
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283

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

284
    extracted_collections = []
Timothy J. Baek's avatar
Timothy J. Baek committed
285
286
287
288
289
    relevant_contexts = []

    for doc in docs:
        context = None

290
291
292
293
294
295
296
297
        collection_names = (
            doc["collection_names"]
            if doc["type"] == "collection"
            else [doc["collection_name"]]
        )

        collection_names = set(collection_names).difference(extracted_collections)
        if not collection_names:
298
299
            log.debug(f"skipping {doc} as it has already been extracted")
            continue
300

301
        try:
302
            if doc["type"] == "text":
303
                context = doc["content"]
Timothy J. Baek's avatar
Timothy J. Baek committed
304
            else:
Timothy J. Baek's avatar
Timothy J. Baek committed
305
306
                if hybrid_search:
                    context = query_collection_with_hybrid_search(
307
                        collection_names=collection_names,
Timothy J. Baek's avatar
Timothy J. Baek committed
308
309
310
311
312
313
314
315
                        query=query,
                        embedding_function=embedding_function,
                        k=k,
                        reranking_function=reranking_function,
                        r=r,
                    )
                else:
                    context = query_collection(
316
                        collection_names=collection_names,
Timothy J. Baek's avatar
Timothy J. Baek committed
317
318
319
320
                        query=query,
                        embedding_function=embedding_function,
                        k=k,
                    )
Timothy J. Baek's avatar
Timothy J. Baek committed
321
        except Exception as e:
322
            log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
323
324
            context = None

325
        if context:
326
            relevant_contexts.append({**context, "source": doc})
327

328
        extracted_collections.extend(collection_names)
Timothy J. Baek's avatar
Timothy J. Baek committed
329
330

    context_string = ""
Timothy J. Baek's avatar
Timothy J. Baek committed
331

332
    citations = []
Timothy J. Baek's avatar
Timothy J. Baek committed
333
    for context in relevant_contexts:
334
335
        try:
            if "documents" in context:
336
337
338
339
                context_string += "\n\n".join(
                    [text for text in context["documents"][0] if text is not None]
                )

340
341
342
                if "metadatas" in context:
                    citations.append(
                        {
343
                            "source": context["source"],
344
345
346
347
                            "document": context["documents"][0],
                            "metadata": context["metadatas"][0],
                        }
                    )
348
349
        except Exception as e:
            log.exception(e)
Timothy J. Baek's avatar
Timothy J. Baek committed
350

351
    context_string = context_string.strip()
Timothy J. Baek's avatar
Timothy J. Baek committed
352
353
354
355
356
357
358

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

359
360
    log.debug(f"ra_content: {ra_content}")

Timothy J. Baek's avatar
Timothy J. Baek committed
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
    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

379
    return messages, citations
380

Self Denial's avatar
Self Denial committed
381

382
383
384
385
386
387
388
389
390
391
392
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,
    }

Steven Kreitzer's avatar
Steven Kreitzer committed
393
    log.debug(f"model: {model}")
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
    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


420
def generate_openai_embeddings(
Timothy J. Baek's avatar
Timothy J. Baek committed
421
    model: str, text: str, key: str, url: str = "https://api.openai.com/v1"
422
423
424
):
    try:
        r = requests.post(
Timothy J. Baek's avatar
Timothy J. Baek committed
425
            f"{url}/embeddings",
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
            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
441
442
443
444
445


from typing import Any

from langchain_core.retrievers import BaseRetriever
446
from langchain_core.callbacks import CallbackManagerForRetrieverRun
Steven Kreitzer's avatar
Steven Kreitzer committed
447
448
449
450


class ChromaRetriever(BaseRetriever):
    collection: Any
Timothy J. Baek's avatar
Timothy J. Baek committed
451
    embedding_function: Any
452
    top_n: int
Steven Kreitzer's avatar
Steven Kreitzer committed
453
454
455
456
457
458
459

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
Timothy J. Baek's avatar
Timothy J. Baek committed
460
        query_embeddings = self.embedding_function(query)
Steven Kreitzer's avatar
Steven Kreitzer committed
461
462
463

        results = self.collection.query(
            query_embeddings=[query_embeddings],
464
            n_results=self.top_n,
Steven Kreitzer's avatar
Steven Kreitzer committed
465
466
467
468
469
470
        )

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

Steven Kreitzer's avatar
Steven Kreitzer committed
471
472
473
474
475
476
477
        results = []
        for idx in range(len(ids)):
            results.append(
                Document(
                    metadata=metadatas[idx],
                    page_content=documents[idx],
                )
Steven Kreitzer's avatar
Steven Kreitzer committed
478
            )
Steven Kreitzer's avatar
Steven Kreitzer committed
479
        return results
480
481
482
483
484
485
486
487
488
489
490
491
492
493


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):
Timothy J. Baek's avatar
Timothy J. Baek committed
494
    embedding_function: Any
Steven Kreitzer's avatar
Steven Kreitzer committed
495
    top_n: int
496
497
498
499
500
501
502
503
504
505
506
507
508
    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]:
Steven Kreitzer's avatar
Steven Kreitzer committed
509
510
511
        reranking = self.reranking_function is not None

        if reranking:
512
513
514
515
            scores = self.reranking_function.predict(
                [(query, doc.page_content) for doc in documents]
            )
        else:
Timothy J. Baek's avatar
Timothy J. Baek committed
516
517
            query_embedding = self.embedding_function(query)
            document_embedding = self.embedding_function(
518
519
520
521
522
523
524
525
526
527
                [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
            ]

Steven Kreitzer's avatar
Steven Kreitzer committed
528
        result = sorted(docs_with_scores, key=operator.itemgetter(1), reverse=True)
529
530
531
532
533
534
535
536
537
538
        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
539
540


Timothy J. Baek's avatar
Timothy J. Baek committed
541
def search_web(engine: str, query: str) -> list[SearchResult]:
542
543
544
545
546
547
548
549
550
551
552
    """Search the web using a search engine and return the results as a list of SearchResult objects.
    Will look for a search engine API key in environment variables in the following order:
    - SEARXNG_QUERY_URL
    - GOOGLE_PSE_API_KEY + GOOGLE_PSE_ENGINE_ID
    - BRAVE_SEARCH_API_KEY
    - SERPSTACK_API_KEY
    - SERPER_API_KEY

    Args:
        query (str): The query to search for
    """
Timothy J. Baek's avatar
refac  
Timothy J. Baek committed
553
554

    # TODO: add playwright to search the web
Timothy J. Baek's avatar
Timothy J. Baek committed
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
    if engine == "searxng":
        if SEARXNG_QUERY_URL:
            return search_searxng(SEARXNG_QUERY_URL, query)
        else:
            raise Exception("No SEARXNG_QUERY_URL found in environment variables")
    elif engine == "google_pse":
        if GOOGLE_PSE_API_KEY and GOOGLE_PSE_ENGINE_ID:
            return search_google_pse(GOOGLE_PSE_API_KEY, GOOGLE_PSE_ENGINE_ID, query)
        else:
            raise Exception(
                "No GOOGLE_PSE_API_KEY or GOOGLE_PSE_ENGINE_ID found in environment variables"
            )
    elif engine == "brave":
        if BRAVE_SEARCH_API_KEY:
            return search_brave(BRAVE_SEARCH_API_KEY, query)
        else:
            raise Exception("No BRAVE_SEARCH_API_KEY found in environment variables")
    elif engine == "serpstack":
        if SERPSTACK_API_KEY:
            return search_serpstack(
                SERPSTACK_API_KEY, query, https_enabled=SERPSTACK_HTTPS
            )
        else:
            raise Exception("No SERPSTACK_API_KEY found in environment variables")
    elif engine == "serper":
        if SERPER_API_KEY:
            return search_serper(SERPER_API_KEY, query)
        else:
            raise Exception("No SERPER_API_KEY found in environment variables")
584
585
    else:
        raise Exception("No search engine API key found in environment variables")