-
Notifications
You must be signed in to change notification settings - Fork 242
Expand file tree
/
Copy pathfeathr_chat.py
More file actions
39 lines (33 loc) · 1.49 KB
/
feathr_chat.py
File metadata and controls
39 lines (33 loc) · 1.49 KB
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
import os
from typing import Optional
from revChatGPT.V3 import Chatbot
from feathr.chat.notebook_utils import *
class FeathrChat(object):
def __init__(self):
key = self.get_api_key()
self.chat_bot = None
self.feathr_client = None
if key:
self.chat_bot = Chatbot(key)
def get_api_key(self):
return os.getenv("CHATGPT_API_KEY", None)
def ask_llm_in_notebook(self, question_with_prompt: str):
"""Ask LLM model a question within the notebook"""
if not self.chat_bot:
key = self.get_api_key()
if not key:
raise RuntimeError(
"Please set environment variable CHATGPT_API_KEY before using Feathr Chat. You can get your API key for ChatGPT at https://platform.openai.com/account/api-keys. For example, run: os.environ['CHATGPT_API_KEY'] = 'your api key' and retry."
)
self.chat_bot = Chatbot(key)
content = self.chat_bot.ask(question_with_prompt)
if "```" in content and self.is_a_code_gen_question(question_with_prompt):
# for debugging
# res = content
# print(re.sub(r"```.*```", "", res, flags=re.DOTALL))
code = extract_code_from_string(content)
create_new_cell(code)
else:
print(content)
def is_a_code_gen_question(self, question_with_prompt):
return "explain" not in question_with_prompt.lower() and "what" not in question_with_prompt.lower()