import openai
import dspy
import json


with open("creds.json", "r") as creds:
    api_key = json.loads(creds.read())["openai_key"]


lm = dspy.OpenAI(model='gpt-4', api_key=api_key, model_type='chat', max_tokens = 500)
dspy.settings.configure(lm=lm)


from dspy.datasets import HotPotQA

dataset = HotPotQA(train_seed=1, train_size=20, eval_seed=2023, dev_size=50, test_size=0)

trainset = [x.with_inputs('question') for x in dataset.train]
devset = [x.with_inputs('question') for x in dataset.dev]


train_example = trainset[0]
print(train_example)
print(f"Question: {train_example.question}")
print(f"Answer: {train_example.answer}")


class BasicQA(dspy.Signature):
    """Answer questions with short factoid answers."""

    question = dspy.InputField()
    answer = dspy.OutputField(desc="often between 1 and 5 words")


class BasicQABot(dspy.Module):
    def __init__(self):
        super().__init__()

        self.generate = dspy.Predict(BasicQA)

    def forward(self,question):
        prediction = self.generate(question = question)
        return dspy.Prediction(answer = prediction.answer)


qa_bot = BasicQABot()
pred = qa_bot.forward("In the 10th Century A.D. Ealhswith had a son called Æthelweard by which English king?")
pred.answer


from dspy.teleprompt import KNNFewShot
from dspy.predict.knn import KNN

knn_teleprompter = 
KNNFewShot(KNN, 7, trainset)