ds_task_scp/backend/embeddings.py

import cohere
from .config import Config


class EmbeddingGenerator:
    def __init__(self):
        self.client = cohere.Client(Config.COHERE_API_KEY)

    def generate_embeddings(self, text: str):
        response = self.client.embed(
            texts=[text],
            model=Config.EMBED_MODEL,
            input_type="search_document"
        )
        return response.embeddings[0]

    def rerank_issues(self, issues: list, query: str, top_n: int = 5):
        # Handle empty issues list
        if not issues:
            return []

        response = self.client.rerank(
            query=query,
            documents=issues,
            top_n=min(top_n, len(issues)),
            model=Config.RERANK_MODEL
        )
        return [result.document for result in response.results]