from astrapy import DataAPIClient

# Initialize the client
client1 = DataAPIClient("AstraCS:vUPtWJNjuqSJzFwpjwGaIMCG:a763d288a24184d3d8f87e78d8f3a9bdb4e687002ea80fb3e21949b969715b7f")
db1 = client1.get_database_by_api_endpoint(
  "https://ab618096-30bc-47c4-85ac-1a97c75e746c-us-east-2.apps.astra.datastax.com"
)
coll1="pp"



client2= DataAPIClient("AstraCS:ZZUTCqocrzQiBhtvlEPfKqRq:b556beac8fc27ab35a3075db0d1fb34fde88872313673836f14ab3f5a1fb8438")
db2= client2.get_database_by_api_endpoint(
  "https://b4031b92-c3c2-46af-a2ff-bf39cd832d0a-us-east-2.apps.astra.datastax.com"
)
coll2="uq"


class DatabaseEngine():
    def __init__(self):
        pass

    
    def Update(self,data):
        
        collection=db1.get_collection(coll1)
        collection.update_one(
         {"version": data["promptversion"]},
         {"$set": {"biological_context_alignment":data["biological_context_alignment"]}},
        )

    
    def GetData(self,promptversion):
        
        collection1=db1.get_collection(coll1)
        collection2=db2.get_collection(coll2)
        
        promptdocument = collection1.find_one({"version":promptversion})
        

        contextdocument=collection2.find_one({"version":promptversion})
        responsedocument=collection2.find_one({"version":promptversion})

        
        return {
            "prompt":promptdocument,
            "context":contextdocument,
            "response":responsedocument
        }