import dspy

import pandas as pd

#lm = dspy.LM('ollama_chat/deepseek-r1', api_base='http://localhost:11434', api_key='')
#lm = dspy.LM('huggingface/Qwen/Qwen2.5-Coder-32B-Instruct')
lm = dspy.LM('huggingface/meta-llama/Llama-3.2-1B')
dspy.configure(lm=lm)

df = pd.read_csv("product2.csv")
df['content']=df['product']+"; "+df['purpose']+"; "+df['benefit']+"; "+df['fee']

corpus = [row['content'] for i,row in df.iterrows()]

'huggingface/BAAI/bge-small-en-v1.5'
"""
from sentence_transformers import SentenceTransformer

# Load an extremely efficient local model for retrieval
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
embedder = dspy.Embedder(model.encode)
"""
embedder = dspy.Embedder('huggingface/BAAI/bge-small-en-v1.5')
class RecommendProduct(dspy.Signature):
  """
  Recommend RBC financial product based on verbatim
  """
  context = dspy.InputField(desc="may contain relevant product information")
  verbatim = dspy.InputField()
  product = dspy.OutputField(desc="product name with benefit")

class RAG(dspy.Module):
  def __init__(self, num_passages=3):
    super().__init__()
    self.retrieve = dspy.retrievers.Embeddings(embedder=embedder, corpus=corpus, k=num_passages)
    self.recommender = dspy.ChainOfThought(RecommendProduct)

  def forward(self, verbatim):
    context = self.retrieve(verbatim).passages
    prediction = self.recommender(context=context, verbatim=verbatim)
    return dspy.Prediction(context=context, product=prediction.product)


customer="Low APR and great customer service. I would highly recommend if you’re looking for a great credit card company and looking to rebuild your credit. I have had my credit limit increased annually and the annual fee is very low."


qa = RAG(num_passages=2)

def rbc_product(customer:str):
  response = qa(verbatim=f"Which RBC personal banking product best serve the follow customer needs, pain points: {customer}")
  return response.product