File size: 3,247 Bytes
a2d031b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
from openai import OpenAI
import os, sys, time
from typing import List, Optional
class PDFChat:
"""
A class to interact with the OpenAI API to create an assistant for answering questions based on a PDF file.
Attributes:
client (OpenAI): Client for interacting with OpenAI API.
assistant_id (Optional[str]): ID of the created assistant. None until an assistant is created.
"""
def __init__(self) -> None:
"""
Initializes the PDFAssistant with the API key from environment variables.
"""
api_key: Optional[str] = os.getenv("OPENAI_API_KEY")
if api_key is None:
raise ValueError("API Key not found in environment variables")
self.client = OpenAI(api_key=api_key)
self.assistant_id: Optional[str] = None
def upload_file(self, filename: str) -> None:
"""
Uploads a file to the OpenAI API and creates an assistant related to that file.
Args:
filename (str): The path to the file to be uploaded.
"""
file = self.client.files.create(
file=open(filename, 'rb'),
purpose="assistants"
)
assistant = self.client.beta.assistants.create(
name="PDF Helper",
instructions="You are my assistant who can answer questions from the given pdf",
tools=[{"type": "retrieval"}],
model="gpt-3.5-turbo-0125",
file_ids=[file.id]
)
self.assistant_id = assistant.id
def get_answers(self, question: str) -> List[str]:
"""
Asks a question to the assistant and retrieves the answers.
Args:
question (str): The question to be asked to the assistant.
Returns:
List[str]: A list of answers from the assistant.
Raises:
ValueError: If the assistant has not been created yet.
"""
if self.assistant_id is None:
raise ValueError("Assistant not created. Please upload a file first.")
thread = self.client.beta.threads.create()
self.client.beta.threads.messages.create(
thread_id=thread.id,
role="user",
content=question
)
run = self.client.beta.threads.runs.create(
thread_id=thread.id,
assistant_id=self.assistant_id
)
while True:
run_status = self.client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
time.sleep(10)
if run_status.status == 'completed':
messages = self.client.beta.threads.messages.list(thread_id=thread.id)
break
else:
time.sleep(2)
return [message.content[0].text.value for message in messages.data if message.role == "assistant"]
if __name__ == "__main__":
client = PDFChat()
filename = sys.argv[1]
client.upload_file(filename)
while True:
question = input("Enter your question (or type 'exit' to quit): ")
if question.lower() in ['exit', 'quit']:
break
answers = client.get_answers(question)
for answer in answers:
print(answer) |