File size: 3,640 Bytes
35676b4 7880373 d1e793b 35676b4 7880373 d1e793b 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 7880373 35676b4 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | from storage import vector_store
def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
words = text.split()
if len(words) <= chunk_size:
return [text]
chunks, start = [], 0
while start < len(words):
end = min(start + chunk_size, len(words))
chunks.append(" ".join(words[start:end]))
start += chunk_size - overlap
return chunks
def clear_ticker_data(ticker: str) -> None:
"""Delete all existing chunks for a ticker from both Chroma collections."""
vector_store.delete_by_ticker("filings", ticker)
vector_store.delete_by_ticker("transcripts", ticker)
def embed_and_store_filing(
ticker: str,
company_name: str,
mda_text: str,
risk_text: str,
filing_date: str,
period: str,
form_type: str,
accession: str = "",
source_url: str = "",
document_id: str = "",
business_text: str = "",
segments_geography_text: str = "",
) -> None:
ticker = ticker.upper()
filing_document_id = document_id or f"sec:{ticker}:{accession or f'{form_type}:{period}:{filing_date}'}"
sections = [
("MD&A", "mda", mda_text),
("Risk Factors", "risk_factors", risk_text),
("Business", "business", business_text),
("Segments & Geography", "segments_geography", segments_geography_text),
]
for section, section_slug, text in sections:
if not text.strip():
print(f"WARNING: empty section '{section}' for {ticker} {period} ({form_type}) — Chroma chunk not written")
continue
chunks = chunk_text(text)
total = len(chunks)
chunk_ids = [f"{section_slug}:{i}" for i in range(total)]
metadatas = [
{
"ticker": ticker,
"company_name": company_name,
"source": form_type,
"filing_date": filing_date,
"period": period,
"section": section,
"accession": accession,
"source_url": source_url,
"document_id": filing_document_id,
"chunk_id": chunk_ids[i],
"chunk_context": f"{company_name} | {form_type} {period} | {section} | Chunk {i + 1}/{total}",
}
for i in range(total)
]
ids = [f"{ticker}-{form_type}-{period}-{section.replace(' ', '_')}-{i}" for i in range(total)]
vector_store.add_chunks("filings", chunks, metadatas, ids)
def embed_and_store_transcript(
ticker: str,
company_name: str,
transcript_text: str,
transcript_date: str,
period: str,
source_url: str = "",
document_id: str = "",
provider: str = "alphavantage",
) -> None:
ticker = ticker.upper()
if not transcript_text.strip():
return
chunks = chunk_text(transcript_text)
total = len(chunks)
transcript_document_id = document_id or f"transcript:{ticker}:{period}:{transcript_date}"
metadatas = [
{
"ticker": ticker,
"company_name": company_name,
"source": "transcript",
"date": transcript_date,
"period": period,
"speaker": "mixed",
"source_url": source_url,
"provider": provider,
"document_id": transcript_document_id,
"chunk_id": f"transcript:{i}",
"chunk_context": f"{company_name} | Earnings Call {period} | Chunk {i + 1}/{total}",
}
for i in range(total)
]
ids = [f"{ticker}-transcript-{period}-{i}" for i in range(total)]
vector_store.add_chunks("transcripts", chunks, metadatas, ids)
|