guard / databaseengine.py
Junaidb's picture
Update databaseengine.py
76bed36 verified
raw
history blame
2.42 kB
from astrapy import DataAPIClient
import base64
import zlib
# 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="kk"
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"],
"contextual_relevance_alignment":data["contextual_relevance_alignment"],
"unit_coherence":data["unit_coherence"],
"response_specificity":data["response_specificity"]
}
},
)
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,"type":"context"})
responsedocument=collection2.find_one({"version":promptversion,"type":"response"})
actual_saved_response=responsedocument["response"]
actual_saved_context=contextdocument["context"]
actual_saved_prompt=promptdocument["prompt"]
compressed_response_data = base64.b64decode(actual_saved_response)
decompressed_data = zlib.decompress(compressed_response_data)
actual_original_response=decompressed_data.decode("utf-8")
compressed_context_data = base64.b64decode(actual_saved_context)
decompressed_context_data = zlib.decompress(compressed_context_data)
actual_original_context=decompressed_context_data.decode("utf-8")
return {
"prompt":actual_saved_prompt,
"context":actual_original_context,
"response":actual_original_response
}