aartemka/pilot-etf-tech

1from prompts import Prompts2from config import Config3from common_functionality import send_openai_request4import openai5import os6import yfinance as yf7 8 9def get_basic_info(sector: str, n_companies: int, required_fields: list, filter_tickers=False) -> list:10    max_num_companies = Config.OPENAI_MAX_NUM_COMPANIES11    data = []12    if n_companies > max_num_companies:13        print(f"Getting {n_companies} companies")14        ignore_companies = []15        for _ in range(n_companies // max_num_companies):16            prompt = Prompts.BASIC_INFORMATION_PROMPT.format(num_companies=max_num_companies,17                                                             sector=sector,18                                                             ignore_companies=ignore_companies,19                                                             required_fields=required_fields)20            resp = send_openai_request(prompt, max_tokens=7100, temperature=0.1, model="gpt-4")21            print(f"Ignore companies: {ignore_companies}")22            ignore_companies.extend([company['name'] for company in resp])23            if type(resp) == dict:24                data.append(resp)25            else:26                data.extend(resp)27    else:28        prompt = Prompts.BASIC_INFORMATION_PROMPT_SINGLE_USE.format(num_companies=n_companies, sector=sector,29                                                                    required_fields=required_fields)30        data.extend(send_openai_request(prompt, max_tokens=7100, temperature=0.1, model="gpt-4"))31 32    if filter_tickers:33        filtered_tickers = []34        for company in data:35            try:36                ticker = yf.Ticker(company['ticker'])37                if ticker.info.get("currentPrice"):38                    filtered_tickers.append(company["ticker"])39            except:40                continue41        data = [company for company in data if company['ticker'] in filtered_tickers]42    return data43