Ayomide
29 lines
792 B
Python
29 lines
792 B
Python
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]
|