File size: 2,783 Bytes
c6d67ac | 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 | import os
import argparse
import requests
import re
from bs4 import BeautifulSoup
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from core.sec_tools import get_latest_10k_url, HEADERS
def ingest_10k(ticker: str):
"""Downloads, cleans, and embeds a 10-K into a local Chroma Vector Database."""
ticker = ticker.upper()
persist_directory = f"./chroma_db/{ticker}_10k"
if os.path.exists(persist_directory):
print(f"[Ingest: Vector DB for {ticker} 10-K already exists at {persist_directory}. Skipping...]")
return
print(f"\n==============================================")
print(f"Starting Ingestion Pipeline for {ticker}")
print(f"==============================================")
url = get_latest_10k_url(ticker)
if url.startswith("Error") or url.startswith("No 10-K"):
print(f"[Error: SEC URL Fetch failed: {url}]")
return
print(f"[1/4] Downloading raw 10-K from SEC: {url}")
response = requests.get(url, headers=HEADERS)
response.raise_for_status()
raw_text = response.text
print(f"[2/4] Parsing HTML and isolating text payload...")
doc_match = re.search(r'<DOCUMENT>(.*?)</DOCUMENT>', raw_text, re.DOTALL | re.IGNORECASE)
if doc_match:
raw_text = doc_match.group(1)
soup = BeautifulSoup(raw_text, "html.parser")
clean_text = soup.get_text(separator=" ", strip=True)
print(f"[3/4] Chunking document...")
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200
)
docs = [Document(page_content=clean_text, metadata={"source": url, "ticker": ticker})]
chunks = text_splitter.split_documents(docs)
print(f"[4/4] Embedding {len(chunks)} chunks into Chroma DB. (This may take a minute) ...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
Chroma.from_documents(
documents=chunks,
embedding=embeddings,
persist_directory=persist_directory
)
print(f"[Success] {ticker} 10-K successfully ingested.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Ingest SEC 10-K filings into Chroma DB.")
parser.add_argument(
"--tickers",
nargs="+",
required=True,
help="List of stock tickers to ingest (e.g., --tickers AAPL MSFT TSLA)"
)
args = parser.parse_args()
os.makedirs("./chroma_db", exist_ok=True)
for t in args.tickers:
try:
ingest_10k(t)
except Exception as e:
print(f"[Error] Failed to ingest {t}: {str(e)}")
|