helper.py•1.14 kB
import os
import glob
import openai
def get_files(ticker: str) -> str:
script_dir = os.path.dirname(os.path.abspath(__file__))
data_dir = os.path.join(script_dir, '..', 'data', ticker)
files = glob.glob(f'{data_dir}/*.txt')
return files
def get_sec_filing_files(ticker: str) -> list:
script_dir = os.path.dirname(os.path.abspath(__file__))
data_dir = os.path.join(script_dir, '..', 'data', ticker)
files = glob.glob(f'{data_dir}/*.txt')
result = []
for file in files:
filename = os.path.basename(file)
if filename.endswith('_summary.txt'):
continue
result.append(filename.replace('.txt', ''))
return result
def embed(query: str) -> list:
res = openai.embeddings.create(
input = query,
model = 'text-embedding-3-small',
dimensions = 512
)
return res.data[0].embedding
def build_rerank_documents(results: list) -> list:
documents = []
for result in results:
documents.append({
"id": result['id'],
"chunk_text": result['metadata']['original_text']
})
return documents