-
Notifications
You must be signed in to change notification settings - Fork 0
/
bot.py
93 lines (72 loc) · 2.98 KB
/
bot.py
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
"""Bot module."""
import os
from typing import Optional
from datetime import datetime
from dotenv import dotenv_values
import pytz
import instructor
from openai import OpenAI
from langchain_core.vectorstores import VectorStoreRetriever
from embeddings import Embeddings
from models import FunctionMessage
os.environ["TOKENIZERS_PARALLELISM"] = "false"
current_date = datetime.now().astimezone(pytz.timezone("America/New_York"))
PROMPT = f"""
You are an experienced financial analyst that can access the documentation of OpenBB Add-in for Excel
and accurately create Excel commands to retrieve the required OpenBB data.
Current date: {current_date.strftime("%Y-%m-%d")}
Current day of the week: {current_date.strftime("%A")}
From the information below, provide the relevant Excel commands for users' questions. If relevant, use the examples as guidance.
Using the information provided from the description, parameters and data field, determine the most appropriate functions to satisfy the user's request.
In case users' request include any specific query parameters, please make sure to include them in the function similar to the examples provided in the documentation.
If no specific symbol is provided, use AAPL as the default symbol.
You are able to express yourself purely through JSON, strictly and precisely adhering to the provided schemas.
"""
class Bot:
"""Bot class."""
def __init__(
self, api_key: str, model: str, retriever: Optional[VectorStoreRetriever] = None
) -> None:
self.llm = instructor.patch(OpenAI(api_key=api_key))
self.model = model
self.retriever = retriever
self.PROMPT = PROMPT
def get_context(self, question: str) -> str:
"""Get context."""
if self.retriever:
return str(self.retriever.get_relevant_documents(question))
return ""
def ask(self, question: str) -> str:
"""Ask a question."""
context = self.get_context(question)
result: FunctionMessage = self.llm.chat.completions.create(
model=self.model,
temperature=0,
response_model=FunctionMessage,
messages=[
{
"role": "system",
"content": self.PROMPT,
},
{"role": "user", "content": f"{context}"},
{"role": "user", "content": f"Question: {question}"},
],
validation_context={"text_chunk": context},
)
return result.to_xl()
@classmethod
def create(cls) -> "Bot":
config: dict = dotenv_values(".env")
e = Embeddings()
try:
e.load()
retriever = e.get_retriever()
except Exception:
e.train()
e.save()
retriever = e.get_retriever()
return cls(
api_key=config.get("OPENAI_API_KEY", ""),
model=config.get("OPENAI_MODEL", ""),
retriever=retriever,
)