generator.py

from retriever import Retriever
from reranker import Reranker
from anthropic import Anthropic
from typing import List


retriever = Retriever()
reranker = Reranker()


class RAGAgent:
    def __init__(
            self,
            retriever=retriever,
            reranker=reranker,
            anthropic_api_key: str = "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",
            model: str = "claude-3-5-sonnet-20241022",
            max_tokens: int = 1024,
            temperature: float = 0.0,
    ):
        self.retriever = retriever
        self.reranker = reranker
        self.client = Anthropic(api_key=anthropic_api_key)
        self.model = model
        self.max_tokens = max_tokens
        self.temperature = temperature

    def get_context(self, query: str) -> List[str]:
        # Get initial candidates from retriever
        retrieved_docs = self.retriever.search_similar(query)

        # Rerank the candidates
        context = self.reranker.rerank(query, retrieved_docs)

        return context

    def generate_prompt(self, context: List[str]) -> str:
        context = "\n".join(context)
        prompt = f"""
                    "אתה רופא שיניים, דובר עברית בלבד. קוראים לך 'רופא השיניים העברי האלקטרוני הראשון'. ענה למטופל על השאלה שלו על סמך הקונטקס הבא: {context}. הוסף כמה שיותר פרטים, ודאג שהתחביר יהיה תקין ויפה. תעצור כשאתה מרגיש שמיצית את עצמך. אל תמציא דברים. ואל תענה בשפות שהן לא עברית.
                 """
        return prompt

    def get_response(self, question: str) -> str:
        # Get relevant context
        context = self.get_context(question)

        # Generate prompt with context
        prompt = self.generate_prompt(context)

        # Get response from Claude
        response = self.client.messages.create(
            model=self.model,
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            messages=[
                        {"role": "assistant", "content": prompt},
                        {"role": "user", "content": f"{question}"}
                    ]
        )

        return response.content[0].text