"""

results = index.search(query)

retrieved_documents = [documents_by_id[result["ref"]] for result in results[:limit]]

"""

def cosine_similarity(a, b):

return sum(x * y for x, y in zip(a, b)) / ((sum(x * x for x in a) ** 0.5) * (sum(y * y for y in b) ** 0.5))

query_embedding = client.embeddings.create(model="text-embedding-3-small", input=query).data[0].embedding

similarities = []

for doc in documents:

doc_embedding = doc["embedding"]

similarity = cosine_similarity(query_embedding, doc_embedding)

similarities.append((doc, similarity))

similarities.sort(key=lambda x: x[1], reverse=True)

retrieved_documents = [doc for doc, _ in similarities[:limit]]

"""

scores = {}

for i, doc in enumerate(text_results):

if doc["id"] not in scores:

scores[doc["id"]] = 0

scores[doc["id"]] += 1 / (i + k)

for i, doc in enumerate(vector_results):

if doc["id"] not in scores:

scores[doc["id"]] = 0

scores[doc["id"]] += 1 / (i + k)

scored_documents = sorted(scores.items(), key=lambda x: x[1], reverse=True)

retrieved_documents = [documents_by_id[doc_id] for doc_id, _ in scored_documents]

"""

encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

scores = encoder.predict([(query, doc["text"]) for doc in retrieved_documents])

scored_documents = [v for _, v in sorted(zip(scores, retrieved_documents), reverse=True)]

"""

text_results = full_text_search(query, limit * 2)

vector_results = vector_search(query, limit * 2)

fused_results = reciprocal_rank_fusion(text_results, vector_results)

reranked_results = rerank(query, fused_results)

model=MODEL_NAME,

temperature=0.3,

messages=[

{"role": "system", "content": SYSTEM_MESSAGE},

{"role": "user", "content": f"{user_question}\nSources: {context}"},

],