|
from QdrantU import QdrantU |
|
from Processing import TextEmbedder |
|
import cohere |
|
from Helpers import generate_prompt, llama, get_docs_by_indices , query_rewriter |
|
|
|
|
|
def run_rag(query, history=None): |
|
embedding_model = TextEmbedder() |
|
uploader = QdrantU(collection_name='News_source') |
|
try: |
|
query = query_rewriter(query) |
|
print("Query after rewriting: ", query) |
|
except: |
|
print("Error in query rewriting") |
|
pass |
|
search_results = uploader.search(query, embedding_model, limit=1000) |
|
docs = list(set([result.payload['content'] for result in search_results])) |
|
|
|
apiKey = 'Q21IIAUkTtt1jk9WUgJg0XiCvaU2K73cFbq0djhM' |
|
co = cohere.Client(apiKey) |
|
rerank_docs = co.rerank( |
|
query=query, documents=docs, top_n=2, model="rerank-english-v3.0" |
|
) |
|
|
|
indices = [result.index for result in rerank_docs.results] |
|
documents = get_docs_by_indices(docs, indices) |
|
prompt = generate_prompt(documents, query, history) |
|
|
|
return prompt |
