Ayomide
48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
import cohere
|
|
from .config import Config
|
|
|
|
co = cohere.Client(Config.COHERE_API_KEY)
|
|
|
|
|
|
def get_embeddings(texts):
|
|
"""Generate embeddings using Cohere"""
|
|
try:
|
|
response = co.embed(
|
|
texts=texts,
|
|
model=Config.EMBEDDING_MODEL,
|
|
input_type="search_document"
|
|
)
|
|
return response.embeddings
|
|
except Exception as e:
|
|
print(f"Error generating embeddings: {str(e)}")
|
|
return None
|
|
|
|
|
|
def get_query_embedding(query):
|
|
"""Generate embedding for search query"""
|
|
try:
|
|
response = co.embed(
|
|
texts=[query],
|
|
model=Config.EMBEDDING_MODEL,
|
|
input_type="search_query"
|
|
)
|
|
return response.embeddings[0]
|
|
except Exception as e:
|
|
print(f"Error generating query embedding: {str(e)}")
|
|
return None
|
|
|
|
|
|
def rerank_results(query, documents):
|
|
"""Re-rank search results using Cohere"""
|
|
try:
|
|
response = co.rerank(
|
|
model="rerank-english-v2.0",
|
|
query=query,
|
|
documents=documents,
|
|
top_n=5
|
|
)
|
|
return response.results
|
|
except Exception as e:
|
|
print(f"Error reranking results: {str(e)}")
|
|
return []
|