common_functionality.py

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