DS_Task_AI_News/backend/embeddings.py

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 []