retreival-api / app.py
sdYethi's picture
Update app.py
dd5f53b
from openai import OpenAI
import os
import gradio as gr
import time
# Define the dropdown options
dropdown_options = ["1_BRD_Lending_v0.1.pdf", "2_BRD_Lending_v0.1.pdf", "3_BRD_Payments_v0.1.pdf"]
query = {
dropdown_options[0] : "Generate test cases for the following process: Loan Boarding Process, Loan Repayment Management, Loan Servicing Management, Disbursement Cancellation",
dropdown_options[1] : "Generate test cases for the following process: Process Flows, Quick Data Entry, RBI Defaulters, Credit Approval Process",
dropdown_options[2] : "Generate test cases for the following process: Fund Transfer and Payments, SBM Account Payment, FCY Payment , Credit Card Payment",
}
class Assistant:
def __init__(self, path):
self.client = OpenAI()
print(str(path))
# Upload a file with an "assistants" purpose
self.file = self.client.files.create(
file=open(str(path), "rb"),
purpose='assistants'
)
# Add the file to the assistant
self.assistant = self.client.beta.assistants.create(
instructions="You are an helpful assistant to generating test scenarios and test cases for test design process automation.",
model="gpt-4-1106-preview",
tools=[{"type": "retrieval"}],
file_ids=[self.file.id]
)
self.thread = self.client.beta.threads.create()
self.query = query.get(path)
def infer(self):
message = self.client.beta.threads.messages.create(
thread_id=self.thread.id,
role="user",
content=self.query,
)
# print(message)
run = self.client.beta.threads.runs.create(
thread_id = self.thread.id,
assistant_id= self.assistant.id
)
while True:
retreive = self.client.beta.threads.runs.retrieve(
thread_id = self.thread.id,
run_id= run.id
)
print(retreive.status)
if retreive.status=="completed":
break
else:
time.sleep(1)
continue
# time.sleep(1)
# run = self.client.beta.threads.runs.retrieve(
# thread_id = thread.id,
# run_id= run.id
# )
messages = self.client.beta.threads.messages.list(
thread_id =self.thread.id
)
print(messages)
for message in messages.data:
print(message.content[0].text.value)
return message.content[0].text.value
if __name__ == "__main__":
# Define a function that takes the selected option as an argument
def create_assistant(path):
obj = Assistant(path)
return obj.infer()
demo_2 = gr.Interface(
create_assistant,
[gr.Dropdown(dropdown_options, label="path", info="Select BRDS doc"),],
"text",
)
# # Create the Gradio interface with a dropdown
# iface = gr.Interface(
# fn=create_assistant(inputs),
# inputs=gr.Dropdown(choices=dropdown_options, label="Select an option"),
# outputs="text"
# )
# gr.Dropdown(
# dropdown_options, label="Animal", info="Will add more animals later!"
# ),
# Launch the interface
print("Running main")
# demo = gr.Interface(fn=obj.infer, inputs="text", outputs="text")
demo_2.launch(show_api=False, share=True)