File size: 1,890 Bytes
fce23bc
 
a3d6903
 
 
 
 
 
 
 
 
 
 
66e4785
a3d6903
 
66e4785
 
a3d6903
 
 
fce23bc
 
80c35b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fce23bc
 
a3d6903
 
 
 
 
 
 
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
import json
import openai
import pandas as pd
import yfinance as yf

from config import Config


def convert_to_excel(dataframes: dict[pd.DataFrame], filename="updated_data", index=True) -> str:
    """
    Accepts a dictionary of dataframes and converts it to an Excel file
    """
    filename = f"{filename}.xlsx"
    print(f"Writing to {filename}")
    writer = pd.ExcelWriter(filename, engine='openpyxl')
    for ticker, df in dataframes.items():
        print(f"Sheet name: {ticker}")
        df.reset_index(drop=True, inplace=True)
        df.to_excel(writer, index=index, sheet_name=ticker)
    writer.close()
    return filename


def send_openai_request(prompt: str, max_tokens=6000, temperature=0.1, model="gpt-4") -> dict:
    try:
        response = openai.ChatCompletion.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": f"{prompt}"
                }
            ],
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=1,
            frequency_penalty=0,
            presence_penalty=0
            )
    except:
        response = openai.ChatCompletion.create(
            model=model,
            messages=[
                {
                    "role": "user",
                    "content": f"{prompt}"
                }
            ],
            temperature=temperature,
            max_tokens=max_tokens,
            top_p=1,
            frequency_penalty=0,
            presence_penalty=0
        )
    return json.loads(response.choices[0]["message"]["content"])


def get_financial_columns(data: pd.DataFrame) -> list:
    test = yf.Ticker("AAPL")
    financial_columns = [x for x in data.columns if x in (test.info.keys() or Config.FINANCE_HISTORICAL_RETURN in x)
                         and x != "Country"]
    return financial_columns