diff --git a/.env.example b/.env.example index 231989e..a31daed 100644 --- a/.env.example +++ b/.env.example @@ -31,6 +31,9 @@ CLAUDE_MODEL=claude-opus-4-6 # --- Embedding Model --- EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2 +# Chunks embedded per batch in the embedding worker. Larger batches are +# faster on machines with plenty of RAM (e.g. 128 on a 32 GB+ machine). +EMBEDDING_BATCH_SIZE=64 # --- Query API --- API_KEYS=dev-key-1,dev-key-2 diff --git a/db/migrations/002_hybrid_search_and_facts.sql b/db/migrations/002_hybrid_search_and_facts.sql new file mode 100644 index 0000000..6e24cdd --- /dev/null +++ b/db/migrations/002_hybrid_search_and_facts.sql @@ -0,0 +1,39 @@ +-- FinDocDRAG Migration 002 — Hybrid search + structured financial facts +-- See: docs/technical-design-document.md Section 5.3 +-- +-- 1. Full-text search support on document_chunks (lexical leg of hybrid +-- retrieval — fused with vector search via Reciprocal Rank Fusion). +-- 2. financial_facts table for curated annual XBRL facts, injected into +-- query context for numeric financial questions. +-- +-- All statements are idempotent: this file may be re-applied safely. + +-- ============================================================ +-- Full-text search column + GIN index +-- ============================================================ +ALTER TABLE document_chunks + ADD COLUMN IF NOT EXISTS chunk_tsv tsvector + GENERATED ALWAYS AS (to_tsvector('english', chunk_text)) STORED; + +CREATE INDEX IF NOT EXISTS idx_chunks_tsv + ON document_chunks USING GIN (chunk_tsv); + +-- ============================================================ +-- Structured financial facts (XBRL companyfacts) +-- ============================================================ +CREATE TABLE IF NOT EXISTS financial_facts ( + ticker VARCHAR(10) NOT NULL, + cik BIGINT NOT NULL, + concept VARCHAR(120) NOT NULL, -- us-gaap tag, e.g. NetIncomeLoss + label VARCHAR(120) NOT NULL, -- human-readable label + unit VARCHAR(20) NOT NULL, -- USD | USD/shares + fiscal_year INTEGER NOT NULL, -- year of the fiscal period end + period_end DATE NOT NULL, + value NUMERIC NOT NULL, + filed DATE NOT NULL, -- filing date of the reporting 10-K + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (ticker, concept, unit, period_end) +); + +CREATE INDEX IF NOT EXISTS idx_facts_ticker_year + ON financial_facts(ticker, fiscal_year); diff --git a/docker-compose.yml b/docker-compose.yml index 4ec63e8..a7b324b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -213,7 +213,7 @@ services: deploy: resources: limits: - cpus: "8.0" + cpus: "12.0" memory: 16G depends_on: db-migrate: @@ -228,6 +228,7 @@ services: POSTGRES_PASSWORD: changeme KAFKA_BOOTSTRAP_SERVERS: kafka:9092 EMBEDDING_MODEL: "sentence-transformers/all-MiniLM-L6-v2" + EMBEDDING_BATCH_SIZE: "${EMBEDDING_BATCH_SIZE:-64}" LOG_LEVEL: "${LOG_LEVEL:-INFO}" ports: - "8002:8002" diff --git a/docs/evaluation-results.md b/docs/evaluation-results.md index 80b3d62..0324475 100644 --- a/docs/evaluation-results.md +++ b/docs/evaluation-results.md @@ -73,3 +73,61 @@ Scores range from 0 to 1; higher is better. Target thresholds: ≥ 0.70 for all | What risks does Amazon identify related to its internat… | AMZN | 0.679 | 0.625 | 1.000 | | What was Amazon's operating income for fiscal year 2024? | AMZN | 0.000 | 0.333 | 0.750 | | **Mean** | | **0.185** | **0.696** | **0.697** | + +## Evaluation Run — 2026-07-19 12:44:53 UTC + +**Samples evaluated:** 6 / 30 +**Full report:** `services/eval/results/eval_report_2026-07-19_12-44-53_UTC.json` + +**Retrieval metrics (LLM-free):** section_hit_rate = 1.000 · section_mrr = 0.833 · ticker_accuracy = 1.000 + +| Question | Ticker | Answer Relevancy | Faithfulness | Context Precision | +| --- | --- | --- | --- | --- | +| What was Apple's total net revenue for fiscal year 2024? | AAPL | 0.966 | 1.000 | — | +| What are Apple's primary supply chain risk factors disc… | AAPL | 0.908 | 1.000 | — | +| What was Apple's Services segment revenue in fiscal yea… | AAPL | 0.000 | — | — | +| How much did Apple spend on research and development in… | AAPL | 1.000 | 1.000 | — | +| Which geographic markets contribute most to Apple's net… | AAPL | 0.989 | 1.000 | — | +| What competitive risks does Apple identify in its annua… | AAPL | 0.954 | — | — | +| **Mean** | | **0.803** | **1.000** | **nan** | + +## Evaluation Run — 2026-07-19 13:25:20 UTC + +**Samples evaluated:** 30 / 30 +**Full report:** `services/eval/results/eval_report_2026-07-19_13-25-20_UTC.json` + +**Retrieval metrics (LLM-free):** section_hit_rate = 0.800 · section_mrr = 0.678 · ticker_accuracy = 1.000 + +| Question | Ticker | Answer Relevancy | Faithfulness | Context Precision | +| --- | --- | --- | --- | --- | +| What was Apple's total net revenue for fiscal year 2024? | AAPL | 1.000 | 1.000 | 0.667 | +| What are Apple's primary supply chain risk factors disc… | AAPL | 0.913 | 1.000 | 0.000 | +| What was Apple's Services segment revenue in fiscal yea… | AAPL | 0.000 | 0.750 | 0.000 | +| How much did Apple spend on research and development in… | AAPL | 1.000 | 1.000 | 0.833 | +| Which geographic markets contribute most to Apple's net… | AAPL | 0.999 | 1.000 | 0.500 | +| What competitive risks does Apple identify in its annua… | AAPL | 0.936 | 1.000 | 1.000 | +| What was Microsoft's total revenue for fiscal year 2024? | MSFT | 1.000 | 1.000 | 0.667 | +| What was the revenue from Microsoft's Intelligent Cloud… | MSFT | 1.000 | 1.000 | 0.250 | +| What are the key competition-related risk factors Micro… | MSFT | 0.924 | 0.947 | 0.000 | +| How much did Microsoft spend on research and developmen… | MSFT | 1.000 | 1.000 | 0.917 | +| How does Microsoft describe its approach to returning c… | MSFT | 0.946 | 1.000 | 0.500 | +| What are Microsoft's three main business segments as de… | MSFT | 0.774 | 1.000 | 1.000 | +| What was Alphabet's total revenue for fiscal year 2024? | GOOGL | 0.999 | 1.000 | 0.700 | +| What was Google Cloud's revenue for fiscal year 2024? | GOOGL | 1.000 | 1.000 | 0.333 | +| What regulatory and legal risks does Alphabet disclose … | GOOGL | 0.969 | 1.000 | 0.500 | +| What is Alphabet's primary source of revenue according … | GOOGL | 0.902 | — | 0.000 | +| How much did Alphabet spend on research and development… | GOOGL | 1.000 | 1.000 | 1.000 | +| How does Alphabet describe its artificial intelligence … | GOOGL | 0.823 | 0.700 | 0.000 | +| What was Amazon's total net sales for fiscal year 2024? | AMZN | 0.958 | 1.000 | 1.000 | +| What was Amazon Web Services net sales for fiscal year … | AMZN | 0.000 | 1.000 | 0.000 | +| What competition-related risks does Amazon disclose in … | AMZN | 0.967 | 0.893 | 0.000 | +| How does Amazon describe its fulfilment and logistics n… | AMZN | 0.876 | 0.583 | 0.000 | +| What risks does Amazon identify related to its internat… | AMZN | 0.978 | 1.000 | 0.500 | +| What was Amazon's operating income for fiscal year 2024? | AMZN | 1.000 | 1.000 | 1.000 | +| What was JPMorgan Chase's total net revenue for fiscal … | JPM | 1.000 | 1.000 | 1.000 | +| What was JPMorgan Chase's net interest income for fisca… | JPM | 0.000 | 1.000 | 0.000 | +| What credit risk factors does JPMorgan Chase identify i… | JPM | 0.925 | 1.000 | 0.000 | +| How does JPMorgan Chase describe its capital management… | JPM | 0.000 | 1.000 | 0.000 | +| What was JPMorgan Chase's provision for credit losses i… | JPM | 1.000 | 1.000 | 0.000 | +| What are JPMorgan Chase's four main business segments a… | JPM | 0.859 | 0.625 | 0.000 | +| **Mean** | | **0.825** | **0.948** | **0.412** | diff --git a/helm/findoc-rag/templates/configmaps.yaml b/helm/findoc-rag/templates/configmaps.yaml index 1704376..d3be118 100644 --- a/helm/findoc-rag/templates/configmaps.yaml +++ b/helm/findoc-rag/templates/configmaps.yaml @@ -65,6 +65,30 @@ data: CREATE INDEX IF NOT EXISTS idx_chunks_ticker ON document_chunks(ticker); CREATE INDEX IF NOT EXISTS idx_chunks_accession ON document_chunks(accession_number); + 002_hybrid_search_and_facts.sql: | + ALTER TABLE document_chunks + ADD COLUMN IF NOT EXISTS chunk_tsv tsvector + GENERATED ALWAYS AS (to_tsvector('english', chunk_text)) STORED; + + CREATE INDEX IF NOT EXISTS idx_chunks_tsv + ON document_chunks USING GIN (chunk_tsv); + + CREATE TABLE IF NOT EXISTS financial_facts ( + ticker VARCHAR(10) NOT NULL, + cik BIGINT NOT NULL, + concept VARCHAR(120) NOT NULL, + label VARCHAR(120) NOT NULL, + unit VARCHAR(20) NOT NULL, + fiscal_year INTEGER NOT NULL, + period_end DATE NOT NULL, + value NUMERIC NOT NULL, + filed DATE NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + PRIMARY KEY (ticker, concept, unit, period_end) + ); + + CREATE INDEX IF NOT EXISTS idx_facts_ticker_year + ON financial_facts(ticker, fiscal_year); {{- if .Values.prometheus.enabled }} --- apiVersion: v1 diff --git a/helm/findoc-rag/tests/configmaps_test.yaml b/helm/findoc-rag/tests/configmaps_test.yaml index f881a3a..8188cab 100644 --- a/helm/findoc-rag/tests/configmaps_test.yaml +++ b/helm/findoc-rag/tests/configmaps_test.yaml @@ -149,6 +149,26 @@ tests: path: data["001_initial_schema.sql"] pattern: "m = 16" + - it: db-migrations configmap contains the hybrid search migration + documentIndex: 1 + asserts: + - isNotEmpty: + path: data["002_hybrid_search_and_facts.sql"] + + - it: db-migrations 002 adds the tsvector column for full-text search + documentIndex: 1 + asserts: + - matchRegex: + path: data["002_hybrid_search_and_facts.sql"] + pattern: "chunk_tsv tsvector" + + - it: db-migrations 002 creates the financial_facts table + documentIndex: 1 + asserts: + - matchRegex: + path: data["002_hybrid_search_and_facts.sql"] + pattern: "CREATE TABLE IF NOT EXISTS financial_facts" + # ── prometheus-config configmap (documentIndex 2) ───────────────── - it: prometheus-config configmap has correct kind diff --git a/services/embedding-worker/src/chunker.py b/services/embedding-worker/src/chunker.py index 639f467..138e63d 100644 --- a/services/embedding-worker/src/chunker.py +++ b/services/embedding-worker/src/chunker.py @@ -1,9 +1,17 @@ """Section-aware chunking for 10-K filings. -Implements the three-stage chunking strategy from TDD Section 5.2.2: +Implements the chunking strategy from TDD Section 5.2.2: 1. Section split — by 10-K item headers (Item 1, 1A, 7, etc.) 2. Paragraph split — by double newlines within each section - 3. Token-based windowing — 512-token windows with 64-token overlap + 3. Paragraph packing — consecutive paragraphs are greedily packed into + chunks of up to 512 tokens, so short paragraphs (headings, single + sentences) never become their own low-signal chunks + 4. Token-based windowing — only paragraphs that alone exceed 512 tokens + are split into 512-token windows with 64-token overlap + +Each chunk also exposes ``embedding_text`` — the chunk text prefixed with +a contextual header (ticker, filing date, section) that is embedded but +not stored, which measurably improves retrieval on corpus-wide queries. References: - TDD: FR-7 (section-aware splitting with 512/64 token window) @@ -55,6 +63,20 @@ class Chunk: text: str token_count: int + @property + def embedding_text(self) -> str: + """Chunk text prefixed with a contextual header, used for embedding only. + + The header anchors the vector to the filing's identity so that + queries like "Apple supply chain risks" match AAPL chunks even when + the chunk body never repeats the company name. The stored/displayed + text (``self.text``) is unchanged. + """ + return ( + f"[{self.ticker} | 10-K | filed {self.filing_date} | {self.section_name}]\n" + f"{self.text}" + ) + # ── Tokeniser (cached) ────────────────────────────────────────── @@ -159,6 +181,44 @@ def split_by_token_window( return windows +# ── Stage 2b: Paragraph packing ───────────────────────────────── + +def pack_paragraphs(paragraphs: list[str], max_tokens: int = DEFAULT_CHUNK_SIZE) -> list[str]: + """Greedily pack consecutive paragraphs into groups of ≤ max_tokens. + + Short paragraphs (headings, one-liners, table fragments) are merged + with their neighbours instead of becoming their own low-signal chunks. + A paragraph that alone exceeds max_tokens is emitted as its own group + (the caller window-splits it). + """ + groups: list[str] = [] + current: list[str] = [] + current_tokens = 0 + + for paragraph in paragraphs: + tokens = count_tokens(paragraph) + + if tokens > max_tokens: + if current: + groups.append("\n\n".join(current)) + current, current_tokens = [], 0 + groups.append(paragraph) # oversize — window-split downstream + continue + + # +1 accounts for the "\n\n" joiner between paragraphs. + if current and current_tokens + tokens + 1 > max_tokens: + groups.append("\n\n".join(current)) + current, current_tokens = [], 0 + + current.append(paragraph) + current_tokens += tokens + 1 + + if current: + groups.append("\n\n".join(current)) + + return groups + + # ── Public API: full chunking pipeline ─────────────────────────── def chunk_filing( @@ -169,7 +229,7 @@ def chunk_filing( max_tokens: int = DEFAULT_CHUNK_SIZE, overlap: int = DEFAULT_OVERLAP, ) -> list[Chunk]: - """Run the full three-stage chunking pipeline on a filing. + """Run the full chunking pipeline on a filing. Returns a list of Chunk objects ready for embedding and storage. Implements FR-7, FR-8. @@ -181,9 +241,10 @@ def chunk_filing( for section_name, section_text in sections: paragraphs = split_into_paragraphs(section_text) + groups = pack_paragraphs(paragraphs, max_tokens) - for paragraph in paragraphs: - windows = split_by_token_window(paragraph, max_tokens, overlap) + for group in groups: + windows = split_by_token_window(group, max_tokens, overlap) for window_text in windows: token_count = count_tokens(window_text) diff --git a/services/embedding-worker/src/main.py b/services/embedding-worker/src/main.py index 4567179..25ad5c9 100644 --- a/services/embedding-worker/src/main.py +++ b/services/embedding-worker/src/main.py @@ -55,6 +55,7 @@ f"/{os.getenv('POSTGRES_DB', 'findocdrag')}" ) EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2") +EMBEDDING_BATCH_SIZE = int(os.getenv("EMBEDDING_BATCH_SIZE", "64")) LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO") TOPIC_RAW = "filings.raw" @@ -186,10 +187,11 @@ def _process_filing( for c in chunks: CHUNK_TOKENS.observe(c.token_count) - # Embed (with batch duration timing) - texts = [c.text for c in chunks] + # Embed (with batch duration timing). embedding_text prefixes each + # chunk with its filing context header — embedded, never stored. + texts = [c.embedding_text for c in chunks] t0 = time.perf_counter() - embeddings = embedder.embed(texts) + embeddings = embedder.embed(texts, batch_size=EMBEDDING_BATCH_SIZE) embed_elapsed = time.perf_counter() - t0 BATCH_DURATION.observe(embed_elapsed) diff --git a/services/embedding-worker/tests/test_chunker.py b/services/embedding-worker/tests/test_chunker.py index 5b96f7a..63d7309 100644 --- a/services/embedding-worker/tests/test_chunker.py +++ b/services/embedding-worker/tests/test_chunker.py @@ -11,6 +11,7 @@ chunk_filing, count_tokens, make_chunk_id, + pack_paragraphs, split_by_token_window, split_into_paragraphs, split_into_sections, @@ -166,3 +167,64 @@ def test_chunk_ids_are_unique(self) -> None: ) ids = [c.chunk_id for c in chunks] assert len(ids) == len(set(ids)) # all unique + + +# ── Paragraph packing ──────────────────────────────────────────── + +class TestPackParagraphs: + def test_empty_list(self) -> None: + assert pack_paragraphs([]) == [] + + def test_short_paragraphs_are_merged(self) -> None: + paragraphs = ["Item 7. MD&A", "Revenue grew 5% year over year.", "Margins expanded."] + groups = pack_paragraphs(paragraphs, max_tokens=512) + assert len(groups) == 1 + assert groups[0] == "\n\n".join(paragraphs) + + def test_budget_is_respected(self) -> None: + paragraph = "word " * 50 # ~50 tokens + groups = pack_paragraphs([paragraph.strip()] * 10, max_tokens=120) + assert len(groups) > 1 + for group in groups: + assert count_tokens(group) <= 120 + + def test_oversize_paragraph_emitted_alone(self) -> None: + small = "A short line." + huge = "token " * 600 # exceeds the budget on its own + groups = pack_paragraphs([small, huge.strip(), small], max_tokens=512) + assert len(groups) == 3 + assert groups[0] == small + assert groups[1] == huge.strip() + assert groups[2] == small + + def test_order_is_preserved(self) -> None: + paragraphs = [f"Paragraph number {i}." for i in range(20)] + groups = pack_paragraphs(paragraphs, max_tokens=30) + assert "\n\n".join(groups) == "\n\n".join(paragraphs) + + +# ── Contextual embedding text ──────────────────────────────────── + +class TestEmbeddingText: + def _chunk(self) -> Chunk: + return Chunk( + chunk_id="x" * 64, + accession_number="ACC001", + ticker="AAPL", + filing_date="2024-11-01", + section_name="Item 1A", + chunk_index=0, + text="The Company faces supply chain risks.", + token_count=7, + ) + + def test_header_contains_filing_identity(self) -> None: + chunk = self._chunk() + header, body = chunk.embedding_text.split("\n", 1) + assert header == "[AAPL | 10-K | filed 2024-11-01 | Item 1A]" + assert body == chunk.text + + def test_stored_text_is_unchanged(self) -> None: + chunk = self._chunk() + _ = chunk.embedding_text + assert chunk.text == "The Company faces supply chain risks." diff --git a/services/eval/eval_dataset.json b/services/eval/eval_dataset.json index abdad66..3f7be38 100644 --- a/services/eval/eval_dataset.json +++ b/services/eval/eval_dataset.json @@ -2,151 +2,283 @@ { "question": "What was Apple's total net revenue for fiscal year 2024?", "ground_truth": "Apple's total net revenue for fiscal year 2024 was approximately $391 billion.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What are Apple's primary supply chain risk factors disclosed in its 10-K?", "ground_truth": "Apple's 10-K identifies concentration of manufacturing in Asia, particularly China, as a primary supply chain risk, including geopolitical tensions, trade restrictions, and reliance on a limited number of suppliers for key components.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "Item 1A" + ] }, { "question": "What was Apple's Services segment revenue in fiscal year 2024?", "ground_truth": "Apple's Services segment revenue for fiscal year 2024 was approximately $96 billion, encompassing the App Store, Apple Music, iCloud, Apple TV+, and other subscription offerings.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "How much did Apple spend on research and development in fiscal year 2024?", "ground_truth": "Apple's research and development expenses for fiscal year 2024 were approximately $31 billion.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "Which geographic markets contribute most to Apple's net revenue?", "ground_truth": "The Americas is Apple's largest geographic segment, followed by Europe and Greater China. The United States represents the single largest country market.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "Item 1", + "Item 7", + "Item 8" + ] }, { "question": "What competitive risks does Apple identify in its annual filing?", "ground_truth": "Apple identifies intense competition across all product and service categories, including from Android-based smartphone manufacturers, PC makers, and streaming and digital services platforms, along with the risk that competitors may offer lower-priced alternatives.", - "ticker": "AAPL" + "ticker": "AAPL", + "expected_sections": [ + "Item 1A" + ] }, { "question": "What was Microsoft's total revenue for fiscal year 2024?", "ground_truth": "Microsoft's total revenue for fiscal year 2024 was approximately $245 billion.", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What was the revenue from Microsoft's Intelligent Cloud segment in fiscal year 2024?", "ground_truth": "Microsoft's Intelligent Cloud segment, which includes Azure, generated approximately $105 billion in revenue for fiscal year 2024.", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What are the key competition-related risk factors Microsoft identifies in its 10-K?", "ground_truth": "Microsoft identifies competition from large technology companies including Google, Amazon, Apple, and Meta across cloud, search, productivity, and gaming segments, along with the risk that open-source software and new entrants could commoditise its offerings.", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "Item 1A" + ] }, { "question": "How much did Microsoft spend on research and development in fiscal year 2024?", "ground_truth": "Microsoft's research and development expenses for fiscal year 2024 were approximately $29 billion.", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "How does Microsoft describe its approach to returning capital to shareholders?", "ground_truth": "Microsoft returns capital to shareholders through quarterly cash dividends and share repurchase programmes, with the board authorising multi-year buyback programmes and a history of consistent dividend growth.", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "Item 5", + "Item 7" + ] }, { "question": "What are Microsoft's three main business segments as described in its 10-K?", "ground_truth": "Microsoft's three reportable segments are Productivity and Business Processes (Office, LinkedIn, Dynamics), Intelligent Cloud (Azure, server products), and More Personal Computing (Windows, Xbox, Surface, Bing).", - "ticker": "MSFT" + "ticker": "MSFT", + "expected_sections": [ + "Item 1", + "Item 7" + ] }, { "question": "What was Alphabet's total revenue for fiscal year 2024?", "ground_truth": "Alphabet's total revenue for fiscal year 2024 was approximately $350 billion, driven primarily by Google advertising and Google Cloud.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What was Google Cloud's revenue for fiscal year 2024?", "ground_truth": "Google Cloud generated approximately $43 billion in revenue for fiscal year 2024, reflecting continued strong growth in cloud infrastructure and platform services.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What regulatory and legal risks does Alphabet disclose in its annual filing?", "ground_truth": "Alphabet discloses risks including antitrust investigations and litigation in the United States and European Union related to Google Search and the Play Store, GDPR and privacy regulatory actions, and potential forced structural changes to its business.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "Item 1A", + "Item 3" + ] }, { "question": "What is Alphabet's primary source of revenue according to its 10-K?", "ground_truth": "Alphabet's primary revenue source is advertising, with Google Search and Google Network Members' properties accounting for the majority of total revenue. YouTube advertising is also a significant and growing contributor.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "Item 1", + "Item 7" + ] }, { "question": "How much did Alphabet spend on research and development in fiscal year 2024?", "ground_truth": "Alphabet's research and development expenses for fiscal year 2024 were approximately $45 billion, reflecting heavy investment in AI, cloud infrastructure, and hardware.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "How does Alphabet describe its artificial intelligence strategy in its annual report?", "ground_truth": "Alphabet describes AI as foundational to all its products and services, highlighting investments in large language models such as Gemini, AI integration into Google Search and Cloud offerings, and the DeepMind research division as key pillars of its long-term AI strategy.", - "ticker": "GOOGL" + "ticker": "GOOGL", + "expected_sections": [ + "Item 1", + "Item 1A", + "Item 7" + ] }, { "question": "What was Amazon's total net sales for fiscal year 2024?", "ground_truth": "Amazon's total net sales for fiscal year 2024 were approximately $638 billion, spanning the North America, International, and Amazon Web Services segments.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What was Amazon Web Services net sales for fiscal year 2024?", "ground_truth": "Amazon Web Services generated approximately $107 billion in net sales for fiscal year 2024, making it the largest cloud infrastructure provider globally.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What competition-related risks does Amazon disclose in its 10-K?", "ground_truth": "Amazon identifies competition from large established companies in each segment, including Walmart and Target in retail, Microsoft Azure and Google Cloud in cloud services, and Meta and Google in digital advertising, noting that many competitors have greater resources in specific verticals.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "Item 1A" + ] }, { "question": "How does Amazon describe its fulfilment and logistics network in its annual filing?", "ground_truth": "Amazon describes a global network of fulfilment centres, sortation centres, delivery stations, and last-mile logistics capabilities, with continued investment in robotics and automation to reduce costs and delivery times.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "Item 1", + "Item 2", + "Item 7" + ] }, { "question": "What risks does Amazon identify related to its international operations?", "ground_truth": "Amazon identifies risks including foreign exchange fluctuations, local regulatory requirements, differences in consumer behaviour, trade restrictions, political and economic instability, and the challenge of adapting its retail and logistics model to local market conditions.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "Item 1A" + ] }, { "question": "What was Amazon's operating income for fiscal year 2024?", "ground_truth": "Amazon's consolidated operating income for fiscal year 2024 was approximately $68 billion, with AWS contributing the largest share of operating profit relative to its revenue size.", - "ticker": "AMZN" + "ticker": "AMZN", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What was JPMorgan Chase's total net revenue for fiscal year 2024?", "ground_truth": "JPMorgan Chase's total net revenue for fiscal year 2024 was approximately $177 billion, comprising net interest income and noninterest revenue across its four business segments.", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What was JPMorgan Chase's net interest income for fiscal year 2024?", "ground_truth": "JPMorgan Chase's net interest income for fiscal year 2024 was approximately $90 billion, benefiting from the higher interest rate environment.", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What credit risk factors does JPMorgan Chase identify in its annual report?", "ground_truth": "JPMorgan Chase's 10-K identifies credit risk factors including consumer and commercial loan default rates, concentration in real estate lending, exposure to leveraged finance and emerging markets, and sensitivity of credit quality to macroeconomic conditions such as unemployment and GDP growth.", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "Item 1A", + "Item 7" + ] }, { "question": "How does JPMorgan Chase describe its capital management strategy in its 10-K?", "ground_truth": "JPMorgan Chase describes maintaining capital levels above regulatory minimums under Basel III frameworks, returning excess capital through dividends and share buybacks subject to stress test results and Federal Reserve non-objection under CCAR, and targeting a standardised CET1 ratio consistent with its stated capital targets.", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "Item 7" + ] }, { "question": "What was JPMorgan Chase's provision for credit losses in fiscal year 2024?", "ground_truth": "JPMorgan Chase's provision for credit losses in fiscal year 2024 was approximately $10 billion, reflecting normalisation of credit from post-pandemic lows and reserve builds in consumer credit.", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "XBRL Financial Facts", + "Item 7", + "Item 8" + ] }, { "question": "What are JPMorgan Chase's four main business segments as described in its 10-K?", "ground_truth": "JPMorgan Chase's four reportable business segments are Consumer and Community Banking (CCB), Commercial Banking (CB), Corporate and Investment Bank (CIB), and Asset and Wealth Management (AWM).", - "ticker": "JPM" + "ticker": "JPM", + "expected_sections": [ + "Item 1", + "Item 7" + ] } -] +] \ No newline at end of file diff --git a/services/eval/requirements.txt b/services/eval/requirements.txt index 01dff1d..cea7f01 100644 --- a/services/eval/requirements.txt +++ b/services/eval/requirements.txt @@ -7,4 +7,5 @@ aiohttp==3.11.* pydantic==2.10.* langchain-openai>=0.1.0 langchain-anthropic>=0.2.0 -langchain-community>=0.2.0 \ No newline at end of file +langchain-community>=0.2.0 +langchain-ollama>=0.2.0 # local Ollama judge with reasoning=False (Qwen3 etc.) \ No newline at end of file diff --git a/services/eval/run_eval.py b/services/eval/run_eval.py index fcd85ee..03f44d4 100644 --- a/services/eval/run_eval.py +++ b/services/eval/run_eval.py @@ -102,7 +102,13 @@ async def _call_query_api( top_k: int, ) -> dict[str, Any]: """POST /v1/query and return the parsed response body.""" - payload: dict[str, Any] = {"question": question, "top_k": top_k} + payload: dict[str, Any] = { + "question": question, + "top_k": top_k, + # Full chunk text per source — ragas must score against the real + # retrieved context, not the 200-char preview. + "include_source_text": True, + } if ticker: payload["ticker_filter"] = ticker @@ -150,16 +156,18 @@ async def collect_responses( resp = await _call_query_api( session, question, ticker, api_url, api_key, top_k ) - # text_preview (200 chars) is what the API exposes per FR-17; - # it is sufficient for ragas context scoring. + # Prefer the full chunk text (include_source_text=True); + # fall back to the 200-char preview for older API versions. contexts = [ - s.get("text_preview", "") for s in resp.get("sources", []) + s.get("text") or s.get("text_preview", "") + for s in resp.get("sources", []) ] enriched.append( { "question": question, "ground_truth": sample["ground_truth"], "ticker": ticker, + "expected_sections": sample.get("expected_sections", []), "answer": resp.get("answer") or "", "contexts": contexts, "sources": resp.get("sources", []), @@ -175,6 +183,7 @@ async def collect_responses( "question": question, "ground_truth": sample["ground_truth"], "ticker": ticker, + "expected_sections": sample.get("expected_sections", []), "answer": "", "contexts": [], "sources": [], @@ -188,6 +197,60 @@ async def collect_responses( return enriched +# ── Retrieval-only metrics (no LLM judge needed) ───────────────────────────── + + +def compute_retrieval_metrics(responses: list[dict[str, Any]]) -> dict[str, float]: + """Compute LLM-free retrieval metrics from sources + expected_sections. + + - section_hit_rate: fraction of samples where at least one retrieved + chunk comes from an expected 10-K section. + - section_mrr: mean reciprocal rank of the first expected-section hit. + - ticker_accuracy: fraction of retrieved chunks whose ticker matches + the sample's ticker (catches cross-company contamination). + + These run on every sample with retrieved sources, independent of the + LLM judge, so retrieval regressions are visible even when generation + or judging fails. + """ + hits: list[float] = [] + reciprocal_ranks: list[float] = [] + ticker_matches = 0 + ticker_total = 0 + + for r in responses: + sources = r.get("sources", []) + if not sources: + continue + + expected = set(r.get("expected_sections", [])) + if expected: + rank = next( + ( + i + for i, s in enumerate(sources, start=1) + if s.get("section") in expected + ), + None, + ) + hits.append(1.0 if rank else 0.0) + reciprocal_ranks.append(1.0 / rank if rank else 0.0) + + if r.get("ticker"): + for s in sources: + ticker_total += 1 + if s.get("ticker") == r["ticker"]: + ticker_matches += 1 + + metrics: dict[str, float] = {} + if hits: + metrics["section_hit_rate"] = sum(hits) / len(hits) + metrics["section_mrr"] = sum(reciprocal_ranks) / len(reciprocal_ranks) + if ticker_total: + metrics["ticker_accuracy"] = ticker_matches / ticker_total + return metrics + + # ── ragas evaluation ────────────────────────────────────────────────────────── @@ -203,6 +266,10 @@ def _resolve_ragas_llm() -> tuple[Any, Any, str]: openai_key = os.getenv("OPENAI_API_KEY") ollama_url = os.getenv("EVAL_OLLAMA_URL", "http://localhost:11434") ollama_model = os.getenv("EVAL_OLLAMA_MODEL", "mistral:7b") + # ragas AnswerRelevancy needs an embedding model. A chat/reasoning model is a + # poor, slow embedder, so allow a dedicated embedding model (e.g. nomic-embed-text) + # while the judge LLM stays a capable chat model. Defaults to the judge model. + ollama_embed_model = os.getenv("EVAL_OLLAMA_EMBED_MODEL", ollama_model) if anthropic_key: try: @@ -226,21 +293,25 @@ def _resolve_ragas_llm() -> tuple[Any, Any, str]: # Fallback: local Ollama (heavy on RAM — stop other services first) try: - from langchain_community.chat_models import ChatOllama - from langchain_community.embeddings import OllamaEmbeddings + from langchain_ollama import ChatOllama, OllamaEmbeddings from ragas.embeddings import LangchainEmbeddingsWrapper from ragas.llms import LangchainLLMWrapper + # reasoning=False disables chain-of-thought on reasoning models + # (e.g. Qwen3): without it the judge spends ~26s "thinking" per call, + # turning a 30-sample eval into hours. Harmless for non-reasoning models. ragas_llm = LangchainLLMWrapper( - ChatOllama(model=ollama_model, base_url=ollama_url) + ChatOllama(model=ollama_model, base_url=ollama_url, reasoning=False) ) ragas_emb = LangchainEmbeddingsWrapper( - OllamaEmbeddings(model=ollama_model, base_url=ollama_url) + OllamaEmbeddings(model=ollama_embed_model, base_url=ollama_url) ) logger.warning( - "No API key found — falling back to local Ollama (%s) for ragas judge. " - "RAM usage will be high. Set ANTHROPIC_API_KEY for a lighter alternative.", + "No API key found — falling back to local Ollama for ragas judge " + "(LLM: %s, embeddings: %s). RAM usage will be high. " + "Set ANTHROPIC_API_KEY for a lighter alternative.", ollama_model, + ollama_embed_model, ) return ragas_llm, ragas_emb, f"Ollama/{ollama_model} (local)" except ImportError: @@ -325,6 +396,7 @@ def compute_ragas_metrics( """ try: from ragas import evaluate + from ragas.run_config import RunConfig except ImportError as exc: logger.error("Cannot import ragas: %s", exc) logger.error("Run: pip install -r requirements.txt") @@ -345,9 +417,21 @@ def compute_ragas_metrics( dataset, metrics, metric_names = _build_ragas_dataset(valid, ragas_llm, ragas_embeddings, judge_label) - logger.info("Running ragas evaluation (metrics: %s) …", ", ".join(metric_names)) + # A local single-GPU Ollama judge serves one request at a time, so ragas' + # default 16-way concurrency just queues calls until they time out. Serialize + # (max_workers=1) with a generous per-call timeout for the local path; API + # judges can override via EVAL_RAGAS_MAX_WORKERS. + max_workers = int(os.getenv("EVAL_RAGAS_MAX_WORKERS", "1")) + call_timeout = int(os.getenv("EVAL_RAGAS_TIMEOUT", "600")) + run_config = RunConfig(max_workers=max_workers, timeout=call_timeout) + + logger.info( + "Running ragas evaluation (metrics: %s, max_workers: %d) …", + ", ".join(metric_names), + max_workers, + ) try: - result = evaluate(dataset=dataset, metrics=metrics) + result = evaluate(dataset=dataset, metrics=metrics, run_config=run_config) except Exception as exc: logger.error("ragas evaluate() failed: %s", exc) return {}, [] @@ -367,6 +451,7 @@ def save_json_report( responses: list[dict[str, Any]], scores: dict[str, list[float]], metric_names: list[str], + retrieval_metrics: dict[str, float], run_at: str, ) -> Path: """Write the full per-sample report to results/eval_report_.json (NFR-9).""" @@ -380,6 +465,7 @@ def save_json_report( "run_at": run_at, "total_samples": len(responses), "evaluated_samples": len(valid), + "retrieval_metrics": retrieval_metrics, "aggregate_metrics": { name: { "scores": scores.get(name, []), @@ -419,6 +505,7 @@ def append_markdown_summary( scores: dict[str, list[float]], metric_names: list[str], responses: list[dict[str, Any]], + retrieval_metrics: dict[str, float], run_at: str, report_path: Path, ) -> None: @@ -431,6 +518,11 @@ def append_markdown_summary( f"**Full report:** `{report_path.relative_to(SCRIPT_DIR.parent.parent)}`\n", ] + if retrieval_metrics: + lines.append("**Retrieval metrics (LLM-free):** " + " · ".join( + f"{name} = {value:.3f}" for name, value in retrieval_metrics.items() + ) + "\n") + if metric_names and valid and scores: col_headers = ["Question", "Ticker"] + [ m.replace("_", " ").title() for m in metric_names @@ -541,20 +633,27 @@ async def _main(args: argparse.Namespace) -> int: ) return 1 - # ── 4. Compute ragas metrics ────────────────────────────────────────────── + # ── 4. Compute retrieval metrics (LLM-free) ─────────────────────────────── + retrieval_metrics = compute_retrieval_metrics(responses) + if retrieval_metrics: + print("\n── Retrieval Metrics (LLM-free) " + "─" * 28) + for name, value in retrieval_metrics.items(): + print(f" {name:<25} {value:.4f}") + + # ── 5. Compute ragas metrics ────────────────────────────────────────────── scores, metric_names = compute_ragas_metrics(responses) - # ── 5. Print aggregate scores to stdout ─────────────────────────────────── + # ── 6. Print aggregate scores to stdout ─────────────────────────────────── if metric_names and scores: print("\n── Aggregate Scores " + "─" * 40) for name in metric_names: print(f" {name:<25} {_mean(scores.get(name, [])):.4f}") print() - # ── 6. Persist results ──────────────────────────────────────────────────── + # ── 7. Persist results ──────────────────────────────────────────────────── run_at = datetime.now(UTC).strftime("%Y-%m-%d %H:%M:%S UTC") - report_path = save_json_report(responses, scores, metric_names, run_at) - append_markdown_summary(scores, metric_names, responses, run_at, report_path) + report_path = save_json_report(responses, scores, metric_names, retrieval_metrics, run_at) + append_markdown_summary(scores, metric_names, responses, retrieval_metrics, run_at, report_path) return 0 diff --git a/services/ingestion/requirements.txt b/services/ingestion/requirements.txt index fc2403e..4dbf982 100644 --- a/services/ingestion/requirements.txt +++ b/services/ingestion/requirements.txt @@ -4,6 +4,8 @@ fastapi==0.115.* uvicorn[standard]==0.34.* aiohttp==3.11.* +beautifulsoup4==4.12.* +lxml==5.3.* confluent-kafka==2.6.* psycopg2-binary==2.9.* pyyaml==6.0.* diff --git a/services/ingestion/src/db.py b/services/ingestion/src/db.py index 4d83c0b..eeb938e 100644 --- a/services/ingestion/src/db.py +++ b/services/ingestion/src/db.py @@ -18,6 +18,7 @@ from collections.abc import Generator from src.edgar_client import Filing + from src.facts import FinancialFact logger = structlog.get_logger() @@ -98,3 +99,42 @@ def record_ingestion(self, filing: Filing) -> None: accession=filing.accession_number, ticker=filing.ticker, ) + + def store_financial_facts(self, facts: list[FinancialFact]) -> int: + """Upsert XBRL facts into financial_facts; returns the row count written. + + Restated values from later filings overwrite earlier ones + (same ticker/concept/unit/period_end key). + """ + if not facts: + return 0 + with self._cursor() as cur: + for fact in facts: + cur.execute( + """ + INSERT INTO financial_facts + (ticker, cik, concept, label, unit, fiscal_year, + period_end, value, filed) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (ticker, concept, unit, period_end) + DO UPDATE SET + value = EXCLUDED.value, + label = EXCLUDED.label, + fiscal_year = EXCLUDED.fiscal_year, + filed = EXCLUDED.filed + WHERE EXCLUDED.filed >= financial_facts.filed + """, + ( + fact.ticker, + fact.cik, + fact.concept, + fact.label, + fact.unit, + fact.fiscal_year, + fact.period_end, + fact.value, + fact.filed, + ), + ) + logger.info("financial_facts_stored", ticker=facts[0].ticker, count=len(facts)) + return len(facts) diff --git a/services/ingestion/src/edgar_client.py b/services/ingestion/src/edgar_client.py index aeb6782..59bc22d 100644 --- a/services/ingestion/src/edgar_client.py +++ b/services/ingestion/src/edgar_client.py @@ -1,14 +1,22 @@ -"""SEC EDGAR EFTS client for fetching 10-K filings. - -Uses the EDGAR full-text search API to find filings, then fetches -the filing document text from the EDGAR archives. +"""SEC EDGAR client for fetching 10-K filings. + +Resolution flow (one filing): + 1. Resolve ticker → CIK via the official company_tickers.json mapping, + so we only ever ingest the company's own filings (never another + filer that merely mentions the ticker in its text). + 2. List the company's 10-K filings via the submissions API + (data.sec.gov/submissions/CIK##########.json). + 3. Fetch the filing's *primary document* (the actual 10-K HTML) rather + than the full submission .txt, which bundles exhibits, XBRL and + base64-encoded binaries. + 4. Parse the HTML to clean text (src.html_parser). References: - - TDD: FR-1 (fetch 10-K filings from EFTS API by ticker) + - TDD: FR-1 (fetch 10-K filings by ticker) - TDD: FR-5 (log and skip filings that fail to parse) - TDD: NFR-4 (respect SEC rate limit of 10 req/s) - TDD: Section 8.1.1 (EDGAR request duration histogram) - - API docs: https://efts.sec.gov/LATEST/search-index + - API docs: https://www.sec.gov/search-filings/edgar-application-programming-interfaces """ from __future__ import annotations @@ -24,12 +32,19 @@ import aiohttp import structlog +from src.html_parser import extract_text from src.metrics import EDGAR_REQUEST_DURATION, FILINGS_FETCHED_TOTAL logger = structlog.get_logger() -# EDGAR full-text search endpoint -EDGAR_SEARCH_URL = "https://efts.sec.gov/LATEST/search-index" +# Official ticker → CIK mapping maintained by the SEC +COMPANY_TICKERS_URL = "https://www.sec.gov/files/company_tickers.json" + +# Filing history per company (10 years of data, parallel arrays) +SUBMISSIONS_URL = "https://data.sec.gov/submissions/CIK{cik:010d}.json" + +# XBRL structured facts per company +COMPANY_FACTS_URL = "https://data.sec.gov/api/xbrl/companyfacts/CIK{cik:010d}.json" # Log a download-progress line every this many bytes _PROGRESS_LOG_INTERVAL = 5 * 1024 * 1024 # 5 MB @@ -50,7 +65,7 @@ class Filing: @dataclass class EdgarClient: - """Async client for SEC EDGAR EFTS full-text search API. + """Async client for the SEC EDGAR submissions and archives APIs. Rate-limited via an asyncio.Semaphore to respect SEC's 10 req/s policy (TDD: NFR-4). Retry logic uses exponential backoff with @@ -60,7 +75,10 @@ class EdgarClient: user_agent: str rate_limit_rps: int = 10 max_retries: int = 3 + filings_since: str = "2020-01-01" # earliest filing date to ingest + max_filings_per_ticker: int = 6 # newest first _semaphore: asyncio.Semaphore = field(init=False) + _cik_cache: dict[str, tuple[int, str]] = field(init=False, default_factory=dict) def __post_init__(self) -> None: self._semaphore = asyncio.Semaphore(self.rate_limit_rps) @@ -101,58 +119,129 @@ async def _with_retries( return failure_value return failure_value # unreachable; keeps mypy happy - # ── Search ─────────────────────────────────────────────────── + # ── JSON fetch helper ──────────────────────────────────────── - async def search_10k_filings( + async def _get_json( self, - ticker: str, + url: str, session: aiohttp.ClientSession, - ) -> list[dict[str, Any]]: - """Search EDGAR EFTS for 10-K filings by ticker symbol. - - Returns a list of filing metadata dicts from the EDGAR - search API. Returns an empty list on error (FR-5). - """ - params = { - "q": f'"{ticker}"', - "dateRange": "custom", - "forms": "10-K", - "startdt": "2020-01-01", - "enddt": "2026-12-31", - } + operation_name: str, + **log_ctx: Any, + ) -> dict[str, Any] | None: + """GET a JSON document with rate limiting and retries.""" headers = {"User-Agent": self.user_agent, "Accept": "application/json"} - async def _do_search() -> list[dict[str, Any]]: + async def _do_get() -> dict[str, Any]: t0 = time.perf_counter() - async with self._semaphore, session.get( - EDGAR_SEARCH_URL, params=params, headers=headers - ) as resp: + async with self._semaphore, session.get(url, headers=headers) as resp: resp.raise_for_status() data: dict[str, Any] = await resp.json(content_type=None) await asyncio.sleep(1.0 / self.rate_limit_rps) - elapsed = time.perf_counter() - t0 - EDGAR_REQUEST_DURATION.labels(ticker=ticker).observe(elapsed) - hits: list[dict[str, Any]] = data.get("hits", {}).get("hits", []) logger.info( - "edgar_search_complete", - ticker=ticker, - results_count=len(hits), - elapsed_ms=round(elapsed * 1000, 1), + f"{operation_name}_complete", + **log_ctx, + elapsed_ms=round((time.perf_counter() - t0) * 1000, 1), + ) + return data + + return await self._with_retries( # type: ignore[no-any-return] + _do_get, failure_value=None, operation_name=operation_name, **log_ctx + ) + + # ── Ticker → CIK resolution ────────────────────────────────── + + async def resolve_cik( + self, + ticker: str, + session: aiohttp.ClientSession, + ) -> tuple[int, str] | None: + """Resolve a ticker symbol to (CIK, official company name). + + Uses the SEC's company_tickers.json mapping; the full mapping is + cached in memory after the first call. Returns None when the + ticker is unknown or the mapping cannot be fetched. + """ + symbol = ticker.upper() + if not self._cik_cache: + data = await self._get_json( + COMPANY_TICKERS_URL, session, operation_name="edgar_ticker_map" ) - return hits + if data is None: + return None + for entry in data.values(): + self._cik_cache[str(entry["ticker"]).upper()] = ( + int(entry["cik_str"]), + str(entry["title"]), + ) + logger.info("edgar_ticker_map_loaded", entries=len(self._cik_cache)) - return await self._with_retries(_do_search, failure_value=[], operation_name="edgar_search", ticker=ticker) # type: ignore[return-value,no-any-return] + resolved = self._cik_cache.get(symbol) + if resolved is None: + logger.warning("edgar_unknown_ticker", ticker=symbol) + return resolved - # ── Fetch full text ────────────────────────────────────────── + # ── Filing metadata via the submissions API ────────────────── - async def fetch_filing_text( + async def list_10k_filings( self, - filing_url: str, + cik: int, session: aiohttp.ClientSession, + ) -> list[dict[str, str]]: + """List recent 10-K filings for a CIK via the submissions API. + + Returns dicts with accession_number, filing_date, primary_document + (newest first), filtered to filings on/after ``filings_since`` and + capped at ``max_filings_per_ticker``. Amendments (10-K/A) are + excluded — they would duplicate the original filing's content. + """ + data = await self._get_json( + SUBMISSIONS_URL.format(cik=cik), + session, + operation_name="edgar_submissions", + cik=cik, + ) + if data is None: + return [] + + recent = data.get("filings", {}).get("recent", {}) + forms: list[str] = recent.get("form", []) + accessions: list[str] = recent.get("accessionNumber", []) + dates: list[str] = recent.get("filingDate", []) + primary_docs: list[str] = recent.get("primaryDocument", []) + + filings: list[dict[str, str]] = [] + for form, accession, filing_date, primary_doc in zip( + forms, accessions, dates, primary_docs, strict=False + ): + if form != "10-K" or filing_date < self.filings_since: + continue + if not accession or not primary_doc: + logger.warning("edgar_incomplete_filing_entry", cik=cik, accession=accession) + continue + filings.append( + { + "accession_number": accession, + "filing_date": filing_date, + "primary_document": primary_doc, + } + ) + if len(filings) >= self.max_filings_per_ticker: + break + + logger.info("edgar_submissions_filtered", cik=cik, filings_found=len(filings)) + return filings + + # ── Fetch a filing document ────────────────────────────────── + + async def fetch_filing_document( + self, + document_url: str, + session: aiohttp.ClientSession, + ticker: str = "", ) -> str | None: - """Fetch the full text of a filing from its EDGAR archive URL. + """Fetch a filing's primary document from the EDGAR archives. - Returns the raw text content, or None on failure (FR-5). + Returns the raw document body (HTML or text), or None on failure (FR-5). """ headers = {"User-Agent": self.user_agent} @@ -162,41 +251,55 @@ async def _do_fetch() -> str: bytes_received = 0 last_log_bytes = 0 - async with self._semaphore, session.get(filing_url, headers=headers) as resp: + async with self._semaphore, session.get(document_url, headers=headers) as resp: resp.raise_for_status() content_length = resp.content_length # may be None - logger.info( - "edgar_fetch_started_http", - url=filing_url, - content_length_bytes=content_length, - content_length_mb=round(content_length / 1024 / 1024, 1) if content_length else None, - ) async for chunk in resp.content.iter_chunked(1024 * 64): # 64 KB chunks byte_chunks.append(chunk) bytes_received += len(chunk) if bytes_received - last_log_bytes >= _PROGRESS_LOG_INTERVAL: logger.info( "edgar_fetch_progress", - url=filing_url, + url=document_url, received_mb=round(bytes_received / 1024 / 1024, 1), total_mb=round(content_length / 1024 / 1024, 1) if content_length else None, - elapsed_ms=round((time.perf_counter() - t0) * 1000, 1), ) last_log_bytes = bytes_received await asyncio.sleep(1.0 / self.rate_limit_rps) raw_bytes = b"".join(byte_chunks) elapsed = time.perf_counter() - t0 + if ticker: + EDGAR_REQUEST_DURATION.labels(ticker=ticker).observe(elapsed) logger.info( "edgar_fetch_complete", - url=filing_url, + url=document_url, elapsed_ms=round(elapsed * 1000, 1), - size_mb=round(len(raw_bytes) / 1024 / 1024, 1), - throughput_mbps=round(len(raw_bytes) / 1024 / 1024 / elapsed, 2) if elapsed > 0 else None, + size_mb=round(len(raw_bytes) / 1024 / 1024, 2), ) return raw_bytes.decode("utf-8", errors="replace") - return await self._with_retries(_do_fetch, failure_value=None, operation_name="edgar_fetch", url=filing_url) # type: ignore[return-value,no-any-return] + return await self._with_retries( # type: ignore[no-any-return] + _do_fetch, failure_value=None, operation_name="edgar_fetch", url=document_url + ) + + # ── XBRL company facts ─────────────────────────────────────── + + async def fetch_company_facts( + self, + cik: int, + session: aiohttp.ClientSession, + ) -> dict[str, Any] | None: + """Fetch the XBRL companyfacts document for a CIK. + + Returns the parsed JSON, or None on failure. + """ + return await self._get_json( + COMPANY_FACTS_URL.format(cik=cik), + session, + operation_name="edgar_company_facts", + cik=cik, + ) # ── End-to-end per ticker ──────────────────────────────────── @@ -206,68 +309,67 @@ async def get_filings_for_ticker( company_name: str, session: aiohttp.ClientSession, ) -> list[Filing]: - """Search for 10-K filings and fetch their full text. + """Resolve the ticker, list its 10-Ks, fetch and parse each document. Implements: - FR-1: Fetch 10-K filings from SEC EDGAR by ticker. - FR-5: Log and skip filings that fail; continue with rest. """ t_ticker = time.perf_counter() - hits = await self.search_10k_filings(ticker, session) - filings: list[Filing] = [] - - for hit in hits: - source = hit.get("_source", {}) - # EDGAR uses "adsh" for the accession number (e.g. "0000320193-24-000123") - accession = source.get("adsh", "") - if not accession: - logger.warning("edgar_missing_accession", ticker=ticker) - FILINGS_FETCHED_TOTAL.labels(ticker=ticker, status="skipped").inc() - continue + resolved = await self.resolve_cik(ticker, session) + if resolved is None: + FILINGS_FETCHED_TOTAL.labels(ticker=ticker, status="skipped").inc() + return [] + cik, official_name = resolved - filing_date = source.get("file_date", "") - - # EDGAR doesn't return a direct filing URL — construct it from the accession number - # Format: https://www.sec.gov/Archives/edgar/data/{CIK}/{accession-no-dashes}/{accession}.txt - ciks = source.get("ciks", []) - cik = ciks[0] if ciks else "" - if not cik: - logger.warning("edgar_missing_cik", ticker=ticker, accession=accession) - FILINGS_FETCHED_TOTAL.labels(ticker=ticker, status="skipped").inc() - continue + entries = await self.list_10k_filings(cik, session) + filings: list[Filing] = [] + for entry in entries: + accession = entry["accession_number"] accession_no_dashes = accession.replace("-", "") source_url = ( - f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_no_dashes}/{accession}.txt" + f"https://www.sec.gov/Archives/edgar/data/{cik}" + f"/{accession_no_dashes}/{entry['primary_document']}" ) logger.info( "edgar_fetch_started", ticker=ticker, accession=accession, - filing_date=filing_date, + filing_date=entry["filing_date"], url=source_url, ) - raw_text = await self.fetch_filing_text(source_url, session) - if raw_text is None: - logger.warning( - "edgar_fetch_skipped", - ticker=ticker, - accession=accession, - ) + document = await self.fetch_filing_document(source_url, session, ticker=ticker) + if document is None: + logger.warning("edgar_fetch_skipped", ticker=ticker, accession=accession) FILINGS_FETCHED_TOTAL.labels(ticker=ticker, status="skipped").inc() continue + clean_text = extract_text(document) + if not clean_text: + logger.warning("edgar_parse_empty", ticker=ticker, accession=accession) + FILINGS_FETCHED_TOTAL.labels(ticker=ticker, status="skipped").inc() + continue + + logger.info( + "edgar_document_parsed", + ticker=ticker, + accession=accession, + raw_kb=round(len(document) / 1024, 1), + clean_kb=round(len(clean_text) / 1024, 1), + ) + filings.append( Filing( accession_number=accession, ticker=ticker, - company_name=company_name, - filing_date=filing_date, + company_name=official_name or company_name, + filing_date=entry["filing_date"], filing_type="10-K", source_url=source_url, - raw_text=raw_text, + raw_text=clean_text, ) ) diff --git a/services/ingestion/src/facts.py b/services/ingestion/src/facts.py new file mode 100644 index 0000000..4fbe602 --- /dev/null +++ b/services/ingestion/src/facts.py @@ -0,0 +1,138 @@ +"""Extraction of structured financial facts from XBRL companyfacts JSON. + +Financial figures ("what was total revenue in FY2024?") live in XBRL, not +in filing prose, so the ingestion service stores a small curated set of +annual us-gaap facts in the ``financial_facts`` table. The Query API +injects them as authoritative context for numeric questions. + +References: + - SEC companyfacts API: https://data.sec.gov/api/xbrl/companyfacts/CIK##########.json + - db/migrations/002_hybrid_search_and_facts.sql (financial_facts schema) +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import date, timedelta +from typing import Any + +import structlog + +logger = structlog.get_logger() + +# Curated us-gaap concepts worth storing, with human-readable labels. +# Several revenue tags exist because filers migrated tags over the years. +TRACKED_CONCEPTS: dict[str, str] = { + "RevenueFromContractWithCustomerExcludingAssessedTax": "Total revenue (net sales)", + "Revenues": "Total revenue", + "SalesRevenueNet": "Total revenue (net sales)", + "CostOfRevenue": "Cost of revenue", + "CostOfGoodsAndServicesSold": "Cost of sales", + "GrossProfit": "Gross profit", + "ResearchAndDevelopmentExpense": "Research and development expense", + "OperatingExpenses": "Operating expenses", + "OperatingIncomeLoss": "Operating income", + "NetIncomeLoss": "Net income", + "EarningsPerShareDiluted": "Diluted earnings per share", + "Assets": "Total assets", + "StockholdersEquity": "Total stockholders' equity", + "CashAndCashEquivalentsAtCarryingValue": "Cash and cash equivalents", +} + +# Accepted XBRL units per concept kind. +_ACCEPTED_UNITS = ("USD", "USD/shares") + +# A "duration" fact must cover roughly one fiscal year. +_MIN_PERIOD_DAYS = 340 +_MAX_PERIOD_DAYS = 380 + + +@dataclass +class FinancialFact: + """One annual financial fact extracted from XBRL companyfacts.""" + + ticker: str + cik: int + concept: str # us-gaap tag, e.g. "NetIncomeLoss" + label: str # human-readable, e.g. "Net income" + unit: str # "USD" or "USD/shares" + fiscal_year: int # year of the fiscal period end + period_end: str # ISO date + value: float + filed: str # ISO date the containing filing was filed + + +def _parse_date(value: str) -> date | None: + try: + return date.fromisoformat(value) + except (ValueError, TypeError): + return None + + +def _is_annual_period(item: dict[str, Any]) -> bool: + """True for instant facts, or duration facts spanning ~1 fiscal year.""" + start = item.get("start") + if start is None: + return True # instant concept (balance sheet) + start_d = _parse_date(str(start)) + end_d = _parse_date(str(item.get("end", ""))) + if start_d is None or end_d is None: + return False + return timedelta(days=_MIN_PERIOD_DAYS) <= (end_d - start_d) <= timedelta(days=_MAX_PERIOD_DAYS) + + +def extract_annual_facts( + company_facts: dict[str, Any], + ticker: str, + cik: int, +) -> list[FinancialFact]: + """Extract tracked annual facts from a companyfacts JSON document. + + Keeps facts reported in 10-K filings for full fiscal years (fp == "FY"). + The same period is re-reported in later filings (comparative columns); + the most recently filed value wins so restatements take precedence. + The fiscal year is labelled with the calendar year of the period end. + """ + facts_obj = company_facts.get("facts") + us_gaap: dict[str, Any] = facts_obj.get("us-gaap", {}) if isinstance(facts_obj, dict) else {} + if not us_gaap: + return [] + + # (concept, unit, period_end) → best fact seen so far + best: dict[tuple[str, str, str], FinancialFact] = {} + + for concept, label in TRACKED_CONCEPTS.items(): + units = us_gaap.get(concept, {}).get("units", {}) + for unit, items in units.items(): + if unit not in _ACCEPTED_UNITS: + continue + for item in items: + if item.get("form") != "10-K" or item.get("fp") != "FY": + continue + if not _is_annual_period(item): + continue + end_d = _parse_date(str(item.get("end", ""))) + value = item.get("val") + filed = item.get("filed") + if end_d is None or filed is None or not isinstance(value, (int, float)): + continue + + fact = FinancialFact( + ticker=ticker, + cik=cik, + concept=concept, + label=label, + unit=unit, + fiscal_year=end_d.year, + period_end=end_d.isoformat(), + value=float(value), + filed=str(filed), + ) + key = (concept, unit, fact.period_end) + current = best.get(key) + if current is None or fact.filed > current.filed: + best[key] = fact + + facts = sorted(best.values(), key=lambda f: (f.concept, f.period_end)) + logger.info("xbrl_facts_extracted", ticker=ticker, cik=cik, facts=len(facts)) + return facts diff --git a/services/ingestion/src/html_parser.py b/services/ingestion/src/html_parser.py new file mode 100644 index 0000000..762161b --- /dev/null +++ b/services/ingestion/src/html_parser.py @@ -0,0 +1,104 @@ +"""HTML → clean text extraction for SEC EDGAR filing documents. + +EDGAR primary documents are (inline-XBRL) HTML. Chunking and embedding +raw HTML wastes tokens on markup and destroys financial tables, so this +module converts a filing document to clean plain text: + + - "): result = await client.get_filings_for_ticker("AAPL", "Apple Inc.", MagicMock()) assert result == [] @pytest.mark.asyncio - async def test_returns_filing_on_success(self, client: EdgarClient) -> None: - """Happy path: a hit with adsh + ciks produces a Filing.""" - with patch.object(client, "search_10k_filings", return_value=[ - {"_source": { - "adsh": "0001-24-000001", - "file_date": "2024-11-01", - "ciks": ["1234567890"], - }}, - ]), patch.object(client, "fetch_filing_text", return_value="Item 1. Business..."): + async def test_returns_parsed_filing_on_success(self, client: EdgarClient) -> None: + """Happy path: HTML document is fetched and converted to clean text.""" + html = "
Item 1. Business

We sell devices.

" + with patch.object(client, "resolve_cik", return_value=(320193, "Apple Inc.")), \ + patch.object(client, "list_10k_filings", return_value=[ + {"accession_number": "0001-24-000001", "filing_date": "2024-11-01", + "primary_document": "aapl-2024.htm"}, + ]), \ + patch.object(client, "fetch_filing_document", return_value=html): result = await client.get_filings_for_ticker("AAPL", "Apple Inc.", MagicMock()) + assert len(result) == 1 - assert result[0].ticker == "AAPL" - assert result[0].accession_number == "0001-24-000001" - assert result[0].raw_text == "Item 1. Business..." - assert result[0].company_name == "Apple Inc." - assert result[0].filing_type == "10-K" - assert "sec.gov" in result[0].source_url + filing = result[0] + assert filing.ticker == "AAPL" + assert filing.accession_number == "0001-24-000001" + assert filing.company_name == "Apple Inc." + assert filing.filing_type == "10-K" + assert "Item 1. Business" in filing.raw_text + assert "We sell devices." in filing.raw_text + assert "<" not in filing.raw_text # no HTML survives @pytest.mark.asyncio - async def test_constructs_correct_url(self, client: EdgarClient) -> None: - """Verify the SEC archive URL is constructed correctly from adsh + CIK.""" - with patch.object(client, "search_10k_filings", return_value=[ - {"_source": { - "adsh": "0001-24-000001", - "file_date": "2024-11-01", - "ciks": ["1234567890"], - }}, - ]), patch.object(client, "fetch_filing_text", return_value="text") as mock_fetch: + async def test_constructs_primary_document_url(self, client: EdgarClient) -> None: + with patch.object(client, "resolve_cik", return_value=(320193, "Apple Inc.")), \ + patch.object(client, "list_10k_filings", return_value=[ + {"accession_number": "0001-24-000001", "filing_date": "2024-11-01", + "primary_document": "aapl-2024.htm"}, + ]), \ + patch.object(client, "fetch_filing_document", return_value="Item 1. text") as mock_fetch: await client.get_filings_for_ticker("AAPL", "Apple Inc.", MagicMock()) - # Verify the URL passed to fetch_filing_text - call_args = mock_fetch.call_args - url = call_args[0][0] - assert "1234567890" in url - assert "000124000001" in url # dashes removed from accession - assert url.endswith("0001-24-000001.txt") - - @pytest.mark.asyncio - async def test_returns_empty_when_search_returns_empty(self, client: EdgarClient) -> None: - """If search returns no hits, get_filings_for_ticker returns [].""" - with patch.object(client, "search_10k_filings", return_value=[]): - result = await client.get_filings_for_ticker("AAPL", "Apple Inc.", MagicMock()) - assert result == [] + url = mock_fetch.call_args[0][0] + assert url == ( + "https://www.sec.gov/Archives/edgar/data/320193" + "/000124000001/aapl-2024.htm" + ) @pytest.mark.asyncio async def test_multiple_filings(self, client: EdgarClient) -> None: - """Multiple hits with valid data produce multiple Filings.""" - with patch.object(client, "search_10k_filings", return_value=[ - {"_source": {"adsh": "0001-24-000001", "file_date": "2024-11-01", "ciks": ["123"]}}, - {"_source": {"adsh": "0001-24-000002", "file_date": "2024-11-02", "ciks": ["123"]}}, - ]), patch.object(client, "fetch_filing_text", return_value="filing text"): + with patch.object(client, "resolve_cik", return_value=(320193, "Apple Inc.")), \ + patch.object(client, "list_10k_filings", return_value=[ + {"accession_number": "0001-24-000001", "filing_date": "2024-11-01", + "primary_document": "a.htm"}, + {"accession_number": "0001-23-000001", "filing_date": "2023-11-03", + "primary_document": "b.htm"}, + ]), \ + patch.object(client, "fetch_filing_document", return_value="Item 1. text"): result = await client.get_filings_for_ticker("AAPL", "Apple Inc.", MagicMock()) - assert len(result) == 2 - assert result[0].accession_number == "0001-24-000001" - assert result[1].accession_number == "0001-24-000002" + assert [f.accession_number for f in result] == ["0001-24-000001", "0001-23-000001"] # ── Filing dataclass ───────────────────────────────────────────── @@ -607,6 +689,15 @@ def test_delivery_callback_on_success(self) -> None: # ── FastAPI App (main.py) ──────────────────────────────────────── +def _make_edgar_mock(*, filings: list[Filing] | None = None) -> MagicMock: + """Edgar client mock whose facts flow no-ops (resolve_cik → None).""" + mock_edgar = MagicMock() + mock_edgar.get_filings_for_ticker = AsyncMock(return_value=filings or []) + mock_edgar.resolve_cik = AsyncMock(return_value=None) + mock_edgar.fetch_company_facts = AsyncMock(return_value=None) + return mock_edgar + + class TestFastAPIApp: """Tests for the FastAPI endpoints in main.py. @@ -724,7 +815,7 @@ def test_ingest_no_tickers_and_no_config(self) -> None: original_edgar = main_module._edgar_client original_kafka = main_module._kafka_producer original_db = main_module._db - main_module._edgar_client = MagicMock() + main_module._edgar_client = _make_edgar_mock() main_module._kafka_producer = MagicMock() main_module._db = MagicMock() @@ -744,10 +835,6 @@ def test_ingest_success_with_tickers(self) -> None: import src.main as main_module - mock_edgar = MagicMock() - mock_kafka = MagicMock() - mock_db = MagicMock() - filing = Filing( accession_number="0001-24-000001", ticker="AAPL", @@ -758,7 +845,10 @@ def test_ingest_success_with_tickers(self) -> None: raw_text="Item 1...", ) - mock_edgar.get_filings_for_ticker = AsyncMock(return_value=[filing]) + mock_edgar = _make_edgar_mock(filings=[filing]) + mock_kafka = MagicMock() + mock_db = MagicMock() + mock_kafka.publish_filing = MagicMock() mock_kafka.flush = MagicMock() mock_db.is_already_ingested.return_value = False @@ -779,6 +869,7 @@ def test_ingest_success_with_tickers(self) -> None: assert data["tickers_processed"] == ["AAPL"] assert data["filings_published"] == 1 assert data["filings_skipped"] == 0 + assert data["facts_stored"] == 0 assert data["errors"] == [] mock_kafka.publish_filing.assert_called_once_with(filing) mock_db.record_ingestion.assert_called_once_with(filing) @@ -787,18 +878,77 @@ def test_ingest_success_with_tickers(self) -> None: main_module._kafka_producer = original_kafka main_module._db = original_db - def test_ingest_handles_ticker_error(self) -> None: - """POST /v1/ingest captures per-ticker errors without crashing.""" + def test_ingest_stores_company_facts(self) -> None: + """POST /v1/ingest fetches XBRL companyfacts and stores annual facts.""" from fastapi.testclient import TestClient import src.main as main_module + company_facts = { + "facts": { + "us-gaap": { + "NetIncomeLoss": { + "units": { + "USD": [ + { + "start": "2023-10-01", + "end": "2024-09-28", + "val": 93_736_000_000, + "form": "10-K", + "fp": "FY", + "fy": 2024, + "filed": "2024-11-01", + } + ] + } + } + } + } + } + mock_edgar = MagicMock() + mock_edgar.get_filings_for_ticker = AsyncMock(return_value=[]) + mock_edgar.resolve_cik = AsyncMock(return_value=(320193, "Apple Inc.")) + mock_edgar.fetch_company_facts = AsyncMock(return_value=company_facts) + mock_kafka = MagicMock() + mock_kafka.flush = MagicMock() mock_db = MagicMock() + mock_db.store_financial_facts.return_value = 1 + + original_edgar = main_module._edgar_client + original_kafka = main_module._kafka_producer + original_db = main_module._db + main_module._edgar_client = mock_edgar + main_module._kafka_producer = mock_kafka + main_module._db = mock_db + try: + client = TestClient(app=main_module.app, raise_server_exceptions=False) + response = client.post("/v1/ingest", json={"tickers": ["AAPL"]}) + assert response.status_code == 200 + data = response.json() + assert data["facts_stored"] == 1 + stored_facts = mock_db.store_financial_facts.call_args[0][0] + assert len(stored_facts) == 1 + assert stored_facts[0].concept == "NetIncomeLoss" + assert stored_facts[0].ticker == "AAPL" + finally: + main_module._edgar_client = original_edgar + main_module._kafka_producer = original_kafka + main_module._db = original_db + + def test_ingest_handles_ticker_error(self) -> None: + """POST /v1/ingest captures per-ticker errors without crashing.""" + from fastapi.testclient import TestClient + + import src.main as main_module + + mock_edgar = _make_edgar_mock() mock_edgar.get_filings_for_ticker = AsyncMock(side_effect=RuntimeError("EDGAR down")) + mock_kafka = MagicMock() mock_kafka.flush = MagicMock() + mock_db = MagicMock() original_edgar = main_module._edgar_client original_kafka = main_module._kafka_producer @@ -827,11 +977,10 @@ def test_ingest_with_config_file_tickers(self) -> None: import src.main as main_module - mock_edgar = MagicMock() + mock_edgar = _make_edgar_mock() mock_kafka = MagicMock() - mock_db = MagicMock() - mock_edgar.get_filings_for_ticker = AsyncMock(return_value=[]) mock_kafka.flush = MagicMock() + mock_db = MagicMock() original_edgar = main_module._edgar_client original_kafka = main_module._kafka_producer @@ -860,10 +1009,6 @@ def test_ingest_skips_already_ingested_filing(self) -> None: import src.main as main_module - mock_edgar = MagicMock() - mock_kafka = MagicMock() - mock_db = MagicMock() - filing = Filing( accession_number="0001-24-000001", ticker="AAPL", @@ -874,8 +1019,10 @@ def test_ingest_skips_already_ingested_filing(self) -> None: raw_text="Item 1...", ) - mock_edgar.get_filings_for_ticker = AsyncMock(return_value=[filing]) + mock_edgar = _make_edgar_mock(filings=[filing]) + mock_kafka = MagicMock() mock_kafka.flush = MagicMock() + mock_db = MagicMock() # Simulate filing already present in ingestion_log mock_db.is_already_ingested.return_value = True @@ -907,10 +1054,6 @@ def test_ingest_mixed_new_and_duplicate_filings(self) -> None: import src.main as main_module - mock_edgar = MagicMock() - mock_kafka = MagicMock() - mock_db = MagicMock() - new_filing = Filing( accession_number="0001-24-000001", ticker="AAPL", @@ -930,8 +1073,10 @@ def test_ingest_mixed_new_and_duplicate_filings(self) -> None: raw_text="Item 1 old...", ) - mock_edgar.get_filings_for_ticker = AsyncMock(return_value=[new_filing, dup_filing]) + mock_edgar = _make_edgar_mock(filings=[new_filing, dup_filing]) + mock_kafka = MagicMock() mock_kafka.flush = MagicMock() + mock_db = MagicMock() # new_filing is new, dup_filing is already ingested mock_db.is_already_ingested.side_effect = lambda acc: acc == dup_filing.accession_number diff --git a/services/ingestion/tests/test_facts.py b/services/ingestion/tests/test_facts.py new file mode 100644 index 0000000..6ccad8a --- /dev/null +++ b/services/ingestion/tests/test_facts.py @@ -0,0 +1,162 @@ +"""Unit tests for XBRL companyfacts extraction.""" + +from __future__ import annotations + +from typing import Any + +from src.facts import FinancialFact, extract_annual_facts + + +def _companyfacts(concept_items: dict[str, list[dict[str, Any]]], unit: str = "USD") -> dict[str, Any]: + """Build a minimal companyfacts JSON document.""" + return { + "cik": 320193, + "entityName": "Apple Inc.", + "facts": { + "us-gaap": { + concept: {"units": {unit: items}} + for concept, items in concept_items.items() + } + }, + } + + +def _annual_item( + end: str, + val: float, + filed: str, + start: str | None = None, + form: str = "10-K", + fp: str = "FY", +) -> dict[str, Any]: + item: dict[str, Any] = {"end": end, "val": val, "filed": filed, "form": form, "fp": fp, "fy": 2024} + if start is not None: + item["start"] = start + return item + + +class TestExtractAnnualFacts: + def test_extracts_annual_duration_fact(self) -> None: + doc = _companyfacts({ + "NetIncomeLoss": [ + _annual_item(start="2023-10-01", end="2024-09-28", val=93_736_000_000, filed="2024-11-01"), + ], + }) + facts = extract_annual_facts(doc, ticker="AAPL", cik=320193) + assert len(facts) == 1 + fact = facts[0] + assert fact.concept == "NetIncomeLoss" + assert fact.fiscal_year == 2024 + assert fact.period_end == "2024-09-28" + assert fact.value == 93_736_000_000 + assert fact.ticker == "AAPL" + assert fact.cik == 320193 + + def test_skips_quarterly_periods(self) -> None: + doc = _companyfacts({ + "NetIncomeLoss": [ + # ~3 month duration — not an annual figure + _annual_item(start="2024-06-30", end="2024-09-28", val=14_736_000_000, filed="2024-11-01"), + ], + }) + assert extract_annual_facts(doc, ticker="AAPL", cik=320193) == [] + + def test_skips_non_10k_forms(self) -> None: + doc = _companyfacts({ + "NetIncomeLoss": [ + _annual_item( + start="2023-10-01", end="2024-09-28", val=1.0, + filed="2024-08-01", form="10-Q", fp="Q3", + ), + ], + }) + assert extract_annual_facts(doc, ticker="AAPL", cik=320193) == [] + + def test_instant_facts_need_no_duration(self) -> None: + doc = _companyfacts({ + "Assets": [_annual_item(end="2024-09-28", val=364_980_000_000, filed="2024-11-01")], + }) + facts = extract_annual_facts(doc, ticker="AAPL", cik=320193) + assert len(facts) == 1 + assert facts[0].concept == "Assets" + + def test_latest_filed_value_wins_for_same_period(self) -> None: + """Comparative columns re-report the same period; restatements must win.""" + doc = _companyfacts({ + "NetIncomeLoss": [ + _annual_item(start="2022-09-25", end="2023-09-30", val=96_995_000_000, filed="2023-11-03"), + _annual_item(start="2022-09-25", end="2023-09-30", val=97_000_000_000, filed="2024-11-01"), + ], + }) + facts = extract_annual_facts(doc, ticker="AAPL", cik=320193) + assert len(facts) == 1 + assert facts[0].value == 97_000_000_000 + assert facts[0].filed == "2024-11-01" + + def test_untracked_concepts_are_ignored(self) -> None: + doc = _companyfacts({ + "SomeObscureConcept": [ + _annual_item(start="2023-10-01", end="2024-09-28", val=1.0, filed="2024-11-01"), + ], + }) + assert extract_annual_facts(doc, ticker="AAPL", cik=320193) == [] + + def test_unaccepted_units_are_ignored(self) -> None: + doc = _companyfacts( + {"NetIncomeLoss": [ + _annual_item(start="2023-10-01", end="2024-09-28", val=1.0, filed="2024-11-01"), + ]}, + unit="EUR", + ) + assert extract_annual_facts(doc, ticker="AAPL", cik=320193) == [] + + def test_eps_in_usd_per_share(self) -> None: + doc = _companyfacts( + {"EarningsPerShareDiluted": [ + _annual_item(start="2023-10-01", end="2024-09-28", val=6.08, filed="2024-11-01"), + ]}, + unit="USD/shares", + ) + facts = extract_annual_facts(doc, ticker="AAPL", cik=320193) + assert len(facts) == 1 + assert facts[0].unit == "USD/shares" + assert facts[0].value == 6.08 + + def test_missing_fields_are_skipped(self) -> None: + doc = _companyfacts({ + "NetIncomeLoss": [ + {"end": "2024-09-28", "form": "10-K", "fp": "FY"}, # no val/filed + {"val": 1.0, "form": "10-K", "fp": "FY", "filed": "2024-11-01"}, # no end + ], + }) + assert extract_annual_facts(doc, ticker="AAPL", cik=320193) == [] + + def test_empty_document(self) -> None: + assert extract_annual_facts({}, ticker="AAPL", cik=320193) == [] + + def test_multiple_years_sorted(self) -> None: + doc = _companyfacts({ + "Revenues": [ + _annual_item(start="2023-10-01", end="2024-09-28", val=391.0, filed="2024-11-01"), + _annual_item(start="2022-09-25", end="2023-09-30", val=383.0, filed="2023-11-03"), + ], + }) + facts = extract_annual_facts(doc, ticker="AAPL", cik=320193) + assert [f.fiscal_year for f in facts] == [2023, 2024] + + +class TestFinancialFactDataclass: + def test_stores_all_fields(self) -> None: + fact = FinancialFact( + ticker="AAPL", + cik=320193, + concept="NetIncomeLoss", + label="Net income", + unit="USD", + fiscal_year=2024, + period_end="2024-09-28", + value=93_736_000_000.0, + filed="2024-11-01", + ) + assert fact.label == "Net income" + assert fact.fiscal_year == 2024 diff --git a/services/ingestion/tests/test_html_parser.py b/services/ingestion/tests/test_html_parser.py new file mode 100644 index 0000000..4035ec1 --- /dev/null +++ b/services/ingestion/tests/test_html_parser.py @@ -0,0 +1,139 @@ +"""Unit tests for the EDGAR HTML → clean text parser.""" + +from __future__ import annotations + +from src.html_parser import extract_text + +# ── Plain-text passthrough ─────────────────────────────────────── + + +class TestPlainTextPassthrough: + def test_empty_input(self) -> None: + assert extract_text("") == "" + + def test_plain_text_is_normalised_not_parsed(self) -> None: + text = "Item 1. Business\n\n\n\nWe are a company." + result = extract_text(text) + assert result == "Item 1. Business\n\nWe are a company." + + def test_plain_text_preserves_paragraphs(self) -> None: + text = "First paragraph.\n\nSecond paragraph." + assert extract_text(text) == text + + +# ── HTML parsing ───────────────────────────────────────────────── + + +class TestHtmlParsing: + def test_strips_tags(self) -> None: + html = "

Item 1. Business

We sell devices.

" + result = extract_text(html) + assert "

" not in result + assert "Item 1. Business" in result + assert "We sell devices." in result + + def test_paragraphs_become_blank_line_separated(self) -> None: + html = "

First.

Second.

" + result = extract_text(html) + assert "First.\n\nSecond." in result + + def test_drops_script_and_style(self) -> None: + html = ( + "" + "

Visible text.

" + ) + result = extract_text(html) + assert "Visible text." in result + assert "alert" not in result + assert "color" not in result + + def test_drops_inline_xbrl_hidden_header(self) -> None: + html = ( + "" + "MACHINE-ONLY-FACT" + "

Human readable.

" + ) + result = extract_text(html) + assert "Human readable." in result + assert "MACHINE-ONLY-FACT" not in result + + def test_inline_spans_do_not_break_words(self) -> None: + # EDGAR filings wrap runs of text in adjacent spans. + html = "
Net sales
" + result = extract_text(html) + assert "Net sales" in result + + def test_nbsp_normalised_to_space(self) -> None: + html = "

Total revenue

" + assert "Total revenue" in extract_text(html) + + def test_no_triple_blank_lines(self) -> None: + html = ( + "

A

" + "

B

" + ) + result = extract_text(html) + assert "\n\n\n" not in result + + +# ── Table rendering ────────────────────────────────────────────── + + +class TestTableRendering: + def test_table_rows_become_pipe_separated(self) -> None: + html = ( + "" + "" + "" + "
Net sales391,035383,285
Operating income123,216114,301
" + ) + result = extract_text(html) + assert "Net sales | 391,035 | 383,285" in result + assert "Operating income | 123,216 | 114,301" in result + + def test_empty_cells_are_dropped(self) -> None: + html = ( + "" + "" + "
20242023
" + ) + assert "2024 | 2023" in extract_text(html) + + def test_nested_markup_in_cells_is_flattened(self) -> None: + html = ( + "" + "" + "
R&D expense
31,370
" + ) + assert "R&D expense | 31,370" in extract_text(html) + + def test_header_cells_included(self) -> None: + html = ( + "" + "" + "" + "
SegmentRevenue
iPhone201,183
" + ) + result = extract_text(html) + assert "Segment | Revenue" in result + assert "iPhone | 201,183" in result + + +# ── Realistic 10-K shape ───────────────────────────────────────── + + +class TestSectionHeadersSurviveParsing: + def test_item_headers_stay_on_their_own_lines(self) -> None: + """The embedding worker's section splitter needs 'Item N.' at line starts.""" + html = ( + "" + "
Item 1A. Risk Factors
" + "

The Company faces risks.

" + "
Item 7. Management's Discussion
" + "

Revenue grew.

" + "" + ) + result = extract_text(html) + lines = result.split("\n") + assert any(line.startswith("Item 1A.") for line in lines) + assert any(line.startswith("Item 7.") for line in lines) diff --git a/services/query-api/src/llm/ollama_backend.py b/services/query-api/src/llm/ollama_backend.py index 2b0787c..381491d 100644 --- a/services/query-api/src/llm/ollama_backend.py +++ b/services/query-api/src/llm/ollama_backend.py @@ -45,6 +45,11 @@ async def generate(self, prompt: str, max_tokens: int = 1024) -> LLMResponse: "model": self._model, "prompt": prompt, "stream": False, + # Suppress chain-of-thought on reasoning models (e.g. Qwen3): we want + # the grounded answer in `response`, not thinking tokens that would + # otherwise consume the num_predict budget and leave `response` empty. + # Ignored by non-reasoning models (e.g. mistral). + "think": False, "options": { "num_predict": max_tokens, }, diff --git a/services/query-api/src/main.py b/services/query-api/src/main.py index 424fb5c..32b5890 100644 --- a/services/query-api/src/main.py +++ b/services/query-api/src/main.py @@ -43,6 +43,7 @@ QueryRequest, QueryResponse, ) +from src.rag.facts import FactsRepository from src.rag.generator import RAGGenerator from src.rag.retriever import Retriever @@ -120,12 +121,13 @@ def _get_api_key(request: Request) -> str: _retriever: Retriever | None = None _generator: RAGGenerator | None = None +_facts_repo: FactsRepository | None = None @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: """Startup / shutdown lifecycle.""" - global _retriever, _generator # noqa: PLW0603 + global _retriever, _generator, _facts_repo # noqa: PLW0603 logger.info("query_api_starting", llm_backend=LLM_BACKEND) @@ -134,6 +136,10 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: _retriever.connect() _retriever.verify_embedding_model_consistency() + # Structured XBRL facts (injected into numeric financial questions) + _facts_repo = FactsRepository(dsn=POSTGRES_DSN) + _facts_repo.connect() + # LLM backend (FR-18) — select via LLM_BACKEND env var if LLM_BACKEND == "openai": llm = OpenAIBackend(model=OPENAI_MODEL) # type: ignore[assignment] @@ -142,7 +148,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: else: llm = OllamaBackend(base_url=OLLAMA_URL, model=OLLAMA_MODEL) # type: ignore[assignment] - _generator = RAGGenerator(retriever=_retriever, llm=llm) + _generator = RAGGenerator(retriever=_retriever, llm=llm, facts=_facts_repo) logger.info("query_api_started") @@ -150,6 +156,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: if _retriever is not None: _retriever.close() + if _facts_repo is not None: + _facts_repo.close() logger.info("query_api_stopped") @@ -247,6 +255,9 @@ async def query( question=body.question, top_k=body.top_k, ticker_filter=body.ticker_filter, + filing_date_from=body.filing_date_from, + filing_date_to=body.filing_date_to, + include_source_text=body.include_source_text, ) # Record Prometheus metrics (TDD Section 8.1.3) diff --git a/services/query-api/src/models.py b/services/query-api/src/models.py index 758d904..c9a833d 100644 --- a/services/query-api/src/models.py +++ b/services/query-api/src/models.py @@ -10,6 +10,8 @@ from __future__ import annotations +import datetime + from pydantic import BaseModel, Field, field_validator # ── Query endpoint ─────────────────────────────────────────────── @@ -19,6 +21,13 @@ class QueryRequest(BaseModel): question: str ticker_filter: str | None = None + # Optional filing-date window (ISO dates). Useful when several years of + # 10-Ks are ingested and the question targets a specific period. + filing_date_from: str | None = None + filing_date_to: str | None = None + # When true, each source includes the full chunk text alongside the + # 200-char preview (used by the evaluation harness for ragas scoring). + include_source_text: bool = False # Upper bound of 20 is a context-window budget: at ~512 tokens per chunk, # 20 chunks consume ~10 K tokens, leaving headroom for the system prompt, # question, and answer inside a 32 K-token context window. Raise the cap @@ -32,6 +41,14 @@ def question_must_not_be_blank(cls, v: str) -> str: raise ValueError("question must not be empty or whitespace") return v.strip() + @field_validator("filing_date_from", "filing_date_to") + @classmethod + def date_must_be_iso(cls, v: str | None) -> str | None: + if v is None: + return v + datetime.date.fromisoformat(v) # raises ValueError if malformed + return v + class SourceChunk(BaseModel): """A single source chunk returned alongside the answer (FR-17).""" @@ -42,6 +59,7 @@ class SourceChunk(BaseModel): section: str relevance_score: float text_preview: str # first 200 characters + text: str | None = None # full chunk text, only when include_source_text class TimingInfo(BaseModel): diff --git a/services/query-api/src/rag/facts.py b/services/query-api/src/rag/facts.py new file mode 100644 index 0000000..146b05f --- /dev/null +++ b/services/query-api/src/rag/facts.py @@ -0,0 +1,232 @@ +"""Structured XBRL facts — lookup and injection for numeric questions. + +Financial figures ("what was total revenue in FY2024?") are answered far +more reliably from XBRL companyfacts than from prose retrieval, so the +ingestion service stores curated annual facts in the ``financial_facts`` +table and this module injects the relevant ones into the RAG context. + +Flow (per query): detect financial-metric intent from the question text, +look up matching facts for the requested ticker and fiscal years, and +return them as high-relevance RetrievedChunk entries that the generator +prepends to the retrieved context. Questions without a ticker filter or +without metric keywords are unaffected. + +References: + - db/migrations/002_hybrid_search_and_facts.sql (financial_facts schema) + - services/ingestion/src/facts.py (extraction side) +""" + +from __future__ import annotations + +import contextlib +import re +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Generator + +import psycopg2 +import structlog + +from src.rag.prompts import RetrievedChunk + +logger = structlog.get_logger() + +# Keyword → prioritised us-gaap concept list. The first concept in a group +# that has data for the requested year wins (filers migrated revenue tags +# over the years, so several aliases may exist). +# Order matters: more specific phrases must precede generic ones +# ("earnings per share" before "earnings"/"net income"). +_METRIC_KEYWORDS: list[tuple[tuple[str, ...], list[str]]] = [ + ( + ("earnings per share", "eps", "diluted earnings"), + ["EarningsPerShareDiluted"], + ), + ( + ("research and development", "r&d"), + ["ResearchAndDevelopmentExpense"], + ), + ( + ("operating income", "income from operations", "operating profit"), + ["OperatingIncomeLoss"], + ), + ( + ("operating expense",), + ["OperatingExpenses"], + ), + ( + ("gross profit", "gross margin"), + ["GrossProfit"], + ), + ( + ("cost of revenue", "cost of sales", "cost of goods"), + ["CostOfRevenue", "CostOfGoodsAndServicesSold"], + ), + ( + ("net income", "net profit", "net earnings", "bottom line"), + ["NetIncomeLoss"], + ), + ( + ("total assets",), + ["Assets"], + ), + ( + ("stockholders' equity", "stockholders equity", "shareholders' equity", "shareholders equity"), + ["StockholdersEquity"], + ), + ( + ("cash and cash equivalents",), + ["CashAndCashEquivalentsAtCarryingValue"], + ), + ( + ("revenue", "net sales", "total sales", "turnover"), + [ + "RevenueFromContractWithCustomerExcludingAssessedTax", + "Revenues", + "SalesRevenueNet", + ], + ), +] + +_YEAR_PATTERN = re.compile(r"\b(?:19|20)\d{2}\b") + +# Cap the number of injected facts so they cannot crowd out prose context. +_MAX_FACTS = 8 + + +def detect_metric_concepts(question: str) -> list[list[str]]: + """Return prioritised concept groups whose keywords appear in the question.""" + lowered = question.lower() + return [ + concepts + for keywords, concepts in _METRIC_KEYWORDS + if any(kw in lowered for kw in keywords) + ] + + +def extract_years(question: str) -> list[int]: + """Extract four-digit years (fiscal year references) from the question.""" + return [int(y) for y in _YEAR_PATTERN.findall(question)] + + +def format_fact_value(value: float, unit: str) -> str: + """Format an XBRL value for the prompt ("$391,035,000,000" / "$6.08 per share").""" + if unit == "USD/shares": + return f"${value:,.2f} per share" + if value == int(value): + return f"${int(value):,}" + return f"${value:,.2f}" + + +class FactsRepository: + """Reads curated annual XBRL facts from the financial_facts table.""" + + def __init__(self, dsn: str) -> None: + self._dsn = dsn + self._conn: psycopg2.extensions.connection | None = None + + def connect(self) -> None: + self._conn = psycopg2.connect(self._dsn) + self._conn.autocommit = True + logger.info("facts_repository_connected") + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + self._conn = None + + def _get_conn(self) -> psycopg2.extensions.connection: + if self._conn is None or self._conn.closed: + self.connect() + if self._conn is None: + raise RuntimeError("Failed to establish database connection") + return self._conn + + @contextlib.contextmanager + def _cursor(self) -> Generator[psycopg2.extensions.cursor, None, None]: + cur = self._get_conn().cursor() + try: + yield cur + finally: + cur.close() + + def _lookup_group( + self, + ticker: str, + concepts: list[str], + years: list[int], + ) -> list[tuple[Any, ...]]: + """Fetch facts for one concept group; per fiscal year, the highest-priority + concept with data wins.""" + sql = """ + SELECT concept, label, unit, fiscal_year, period_end, value + FROM financial_facts + WHERE ticker = %s AND concept = ANY(%s) + ORDER BY fiscal_year DESC + """ + with self._cursor() as cur: + cur.execute(sql, (ticker, concepts)) + rows = cur.fetchall() + + by_year: dict[int, tuple[Any, ...]] = {} + for row in rows: + year = int(row[3]) + current = by_year.get(year) + if current is None or concepts.index(row[0]) < concepts.index(current[0]): + by_year[year] = row + + if years: + selected_years = [y for y in years if y in by_year] + elif by_year: + selected_years = [max(by_year)] # no year asked → most recent + else: + selected_years = [] + + return [by_year[y] for y in selected_years] + + def facts_for_question( + self, + question: str, + ticker: str | None, + ) -> list[RetrievedChunk]: + """Return XBRL facts relevant to the question as context chunks. + + Empty unless a ticker filter is set and the question names a + tracked financial metric. Lookup failures degrade to no facts — + never to a failed query. + """ + if not ticker: + return [] + concept_groups = detect_metric_concepts(question) + if not concept_groups: + return [] + years = extract_years(question) + + chunks: list[RetrievedChunk] = [] + try: + for concepts in concept_groups: + for concept, label, unit, fiscal_year, period_end, value in self._lookup_group( + ticker, concepts, years + ): + text = ( + f"{ticker} {label}, fiscal year {fiscal_year} " + f"(period ending {period_end}): {format_fact_value(float(value), unit)}. " + f"Source: XBRL us-gaap:{concept}, SEC companyfacts (authoritative)." + ) + chunks.append( + RetrievedChunk( + chunk_id=f"xbrl:{ticker}:{concept}:{fiscal_year}", + ticker=ticker, + filing_date=str(period_end), + section="XBRL Financial Facts", + relevance_score=1.0, + text=text, + ) + ) + except Exception as exc: + logger.warning("facts_lookup_failed", ticker=ticker, error=str(exc)) + return [] + + if chunks: + logger.info("facts_injected", ticker=ticker, facts=len(chunks)) + return chunks[:_MAX_FACTS] diff --git a/services/query-api/src/rag/generator.py b/services/query-api/src/rag/generator.py index ece88ed..a4c0864 100644 --- a/services/query-api/src/rag/generator.py +++ b/services/query-api/src/rag/generator.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from src.llm.backend import LLMBackend, LLMResponse + from src.rag.facts import FactsRepository from src.rag.retriever import Retriever from src.metrics import ( @@ -32,11 +33,17 @@ class RAGGenerator: - """Orchestrates the full RAG pipeline: retrieve → prompt → generate.""" + """Orchestrates the full RAG pipeline: retrieve → facts → prompt → generate.""" - def __init__(self, retriever: Retriever, llm: LLMBackend) -> None: + def __init__( + self, + retriever: Retriever, + llm: LLMBackend, + facts: FactsRepository | None = None, + ) -> None: self._retriever = retriever self._llm = llm + self._facts = facts def _record_llm_metrics( self, @@ -58,6 +65,9 @@ async def answer( question: str, top_k: int = 5, ticker_filter: str | None = None, + filing_date_from: str | None = None, + filing_date_to: str | None = None, + include_source_text: bool = False, ) -> QueryResponse: """Run the full RAG pipeline and return a QueryResponse. @@ -75,6 +85,8 @@ async def answer( question=question, top_k=top_k, ticker_filter=ticker_filter, + filing_date_from=filing_date_from, + filing_date_to=filing_date_to, ) retrieval_ms = (time.perf_counter() - t0_retrieve) * 1000 - embedding_ms @@ -86,6 +98,15 @@ async def answer( if chunks: RETRIEVAL_SCORE.observe(chunks[0].relevance_score) + # ── Structured XBRL facts ──────────────────────────────── + # For numeric financial questions with a ticker filter, prepend + # authoritative XBRL facts so figures come from structured data, + # not prose retrieval. + if self._facts is not None: + fact_chunks = self._facts.facts_for_question(question, ticker_filter) + if fact_chunks: + chunks = fact_chunks + chunks + # Build sources for the response (FR-17) sources = [ SourceChunk( @@ -95,6 +116,7 @@ async def answer( section=c.section, relevance_score=round(c.relevance_score, 4), text_preview=c.text[:200], + text=c.text if include_source_text else None, ) for c in chunks ] diff --git a/services/query-api/src/rag/prompts.py b/services/query-api/src/rag/prompts.py index 6d06a33..7481743 100644 --- a/services/query-api/src/rag/prompts.py +++ b/services/query-api/src/rag/prompts.py @@ -13,6 +13,10 @@ the provided context from SEC filings. If the context does not contain enough information to answer, say "I don't have enough information to answer this." +Context chunks from the "XBRL Financial Facts" section are authoritative +structured figures from SEC XBRL data — prefer them over prose when they +answer a numeric question. + For every claim in your answer, cite the source using [Source N] notation, where N corresponds to the context chunk number.""" diff --git a/services/query-api/src/rag/retriever.py b/services/query-api/src/rag/retriever.py index 79d5aeb..e510ea3 100644 --- a/services/query-api/src/rag/retriever.py +++ b/services/query-api/src/rag/retriever.py @@ -1,10 +1,25 @@ -"""Vector retriever — embeds the query and searches pgvector. +"""Hybrid retriever — vector search fused with Postgres full-text search. + +Retrieval runs two legs over document_chunks and fuses them with +Reciprocal Rank Fusion (RRF): + + 1. Vector leg — query embedding vs pgvector HNSW (cosine distance), + strong on paraphrase and semantic similarity. + 2. Lexical leg — websearch_to_tsquery over the chunk_tsv tsvector + column, strong on exact terms ("Intelligent Cloud", + "fiscal 2024") that embeddings blur. + +The fused candidate pool is then reranked with Maximal Marginal +Relevance (MMR) to balance relevance and diversity. If the lexical +column is missing (migration 002 not applied), retrieval degrades +gracefully to vector-only. References: - - TDD: FR-13 (embed query, retrieve top-k via cosine distance, - optional ticker filter) + - TDD: FR-13 (embed query, retrieve top-k, optional ticker filter) - TDD: FR-19, FR-20 (list ingested filings) - TDD: NFR-1 (retrieval within 200ms at p99) + - Cormack, Clarke & Buettcher (2009), "Reciprocal Rank Fusion + outperforms Condorcet and individual rank learning methods", SIGIR. """ from __future__ import annotations @@ -28,12 +43,19 @@ DEFAULT_MODEL = "sentence-transformers/all-MiniLM-L6-v2" -# How many candidates to fetch from pgvector before MMR reranking. +# How many candidates each leg fetches before fusion and MMR reranking. # 4× top_k gives the algorithm enough diversity headroom without a # meaningful latency cost (HNSW lookup is O(log n) regardless of LIMIT). _CANDIDATE_MULTIPLIER = 4 _MAX_CANDIDATES = 100 +# RRF constant from Cormack et al. (2009); dampens the weight of top ranks +# so one leg cannot dominate the fusion. +_RRF_K = 60 + +# Columns shared by both retrieval legs (score is appended per leg). +_CHUNK_COLUMNS = "chunk_id, ticker, filing_date, section_name, chunk_text, embedding" + def _apply_mmr( candidates: list[tuple[RetrievedChunk, np.ndarray]], @@ -86,6 +108,46 @@ def _apply_mmr( return selected +def _fuse_rrf( + ranked_lists: list[list[tuple[Any, ...]]], + query_embedding: np.ndarray, + k: int = _RRF_K, +) -> list[tuple[RetrievedChunk, np.ndarray]]: + """Fuse ranked result lists with Reciprocal Rank Fusion. + + Each input row is (chunk_id, ticker, filing_date, section_name, + chunk_text, embedding, leg_score). A chunk appearing in several lists + accumulates 1/(k + rank) per appearance; the fused pool is ordered by + that sum, descending. + + The returned chunks carry cosine similarity to the query as + relevance_score (comparable across legs, unlike leg-native scores), + which is what MMR and the API response report. + """ + rrf_scores: dict[str, float] = {} + by_id: dict[str, tuple[RetrievedChunk, np.ndarray]] = {} + + for rows in ranked_lists: + for rank, row in enumerate(rows, start=1): + chunk_id = row[0] + rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + 1.0 / (k + rank) + if chunk_id not in by_id: + emb = np.asarray(row[5], dtype=float) + chunk = RetrievedChunk( + chunk_id=chunk_id, + ticker=row[1], + filing_date=str(row[2]), + section=row[3], + # Embeddings are L2-normalised → dot == cosine similarity. + relevance_score=float(np.dot(emb, query_embedding)), + text=row[4], + ) + by_id[chunk_id] = (chunk, emb) + + fused_ids = sorted(rrf_scores, key=lambda cid: rrf_scores[cid], reverse=True) + return [by_id[cid] for cid in fused_ids] + + class Retriever: """Embeds a query and retrieves the top-k most similar chunks.""" @@ -165,13 +227,90 @@ def verify_embedding_model_consistency(self) -> None: else: logger.info("embedding_model_consistent", dim=model_dim) + @staticmethod + def _build_filters( + ticker_filter: str | None, + filing_date_from: str | None, + filing_date_to: str | None, + ) -> tuple[list[str], list[Any]]: + """Build optional WHERE clauses shared by both retrieval legs.""" + clauses: list[str] = [] + params: list[Any] = [] + if ticker_filter: + clauses.append("ticker = %s") + params.append(ticker_filter) + if filing_date_from: + clauses.append("filing_date >= %s") + params.append(filing_date_from) + if filing_date_to: + clauses.append("filing_date <= %s") + params.append(filing_date_to) + return clauses, params + + def _vector_search( + self, + query_embedding: list[float], + candidate_k: int, + filter_clauses: list[str], + filter_params: list[Any], + ) -> list[tuple[Any, ...]]: + """Vector leg: top candidates by cosine distance (HNSW).""" + where = f"WHERE {' AND '.join(filter_clauses)}" if filter_clauses else "" + sql = f""" + SELECT {_CHUNK_COLUMNS}, + 1 - (embedding <=> %s::vector) AS score + FROM document_chunks + {where} + ORDER BY embedding <=> %s::vector + LIMIT %s + """ + params = [query_embedding, *filter_params, query_embedding, candidate_k] + with self._cursor() as cur: + cur.execute(sql, params) + rows: list[tuple[Any, ...]] = cur.fetchall() + return rows + + def _lexical_search( + self, + question: str, + candidate_k: int, + filter_clauses: list[str], + filter_params: list[Any], + ) -> list[tuple[Any, ...]]: + """Lexical leg: top candidates by full-text rank. + + Returns [] when the tsvector column is missing (migration 002 not + applied) or the question yields an empty tsquery — retrieval then + degrades to vector-only. + """ + clauses = ["chunk_tsv @@ websearch_to_tsquery('english', %s)", *filter_clauses] + sql = f""" + SELECT {_CHUNK_COLUMNS}, + ts_rank_cd(chunk_tsv, websearch_to_tsquery('english', %s)) AS score + FROM document_chunks + WHERE {" AND ".join(clauses)} + ORDER BY score DESC + LIMIT %s + """ + params = [question, question, *filter_params, candidate_k] + try: + with self._cursor() as cur: + cur.execute(sql, params) + rows: list[tuple[Any, ...]] = cur.fetchall() + return rows + except Exception as exc: + logger.warning("lexical_search_unavailable", error=str(exc)) + return [] + def retrieve( self, question: str, top_k: int = 5, ticker_filter: str | None = None, + filing_date_from: str | None = None, + filing_date_to: str | None = None, ) -> tuple[list[RetrievedChunk], list[float], float]: - """Embed the question and retrieve the top-k chunks. + """Embed the question and retrieve the top-k chunks (hybrid + MMR). Returns: (chunks, query_embedding, embedding_time_ms) @@ -181,60 +320,38 @@ def retrieve( query_embedding = self.embed_query(question) embedding_ms = (time.perf_counter() - t0) * 1000 - # Step 2: Query pgvector — fetch more candidates than requested so - # MMR has enough material to trade off relevance against diversity. + # Step 2: Run both legs — fetch more candidates than requested so + # fusion + MMR have enough material to work with. candidate_k = min(top_k * _CANDIDATE_MULTIPLIER, _MAX_CANDIDATES) + filter_clauses, filter_params = self._build_filters( + ticker_filter, filing_date_from, filing_date_to + ) - if ticker_filter: - sql = """ - SELECT chunk_id, ticker, filing_date, section_name, - chunk_text, embedding, - 1 - (embedding <=> %s::vector) AS relevance_score - FROM document_chunks - WHERE ticker = %s - ORDER BY embedding <=> %s::vector - LIMIT %s - """ - params: tuple[Any, ...] = (query_embedding, ticker_filter, query_embedding, candidate_k) - else: - sql = """ - SELECT chunk_id, ticker, filing_date, section_name, - chunk_text, embedding, - 1 - (embedding <=> %s::vector) AS relevance_score - FROM document_chunks - ORDER BY embedding <=> %s::vector - LIMIT %s - """ - params = (query_embedding, query_embedding, candidate_k) - - with self._cursor() as cur: - cur.execute(sql, params) - rows = cur.fetchall() + vector_rows = self._vector_search( + query_embedding, candidate_k, filter_clauses, filter_params + ) + lexical_rows = self._lexical_search( + question, candidate_k, filter_clauses, filter_params + ) - # Build (chunk, embedding_vector) pairs for MMR. - # row[5] is the pgvector embedding (numpy array after register_vector). - # row[6] is the relevance score. - candidates: list[tuple[RetrievedChunk, np.ndarray]] = [] - for row in rows: - chunk = RetrievedChunk( - chunk_id=row[0], - ticker=row[1], - filing_date=str(row[2]), - section=row[3], - relevance_score=float(row[6]), - text=row[4], - ) - candidates.append((chunk, np.asarray(row[5]))) + # Step 3: Fuse with RRF, keep the top candidate_k of the fused pool. + candidates = _fuse_rrf( + [vector_rows, lexical_rows], np.asarray(query_embedding) + )[:candidate_k] - # Step 3: Apply MMR to select top_k diverse chunks. + # Step 4: Apply MMR to select top_k diverse chunks. chunks = _apply_mmr(candidates, top_k) logger.info( "retrieval_complete", top_k=top_k, - candidates_fetched=len(candidates), + vector_candidates=len(vector_rows), + lexical_candidates=len(lexical_rows), + fused_candidates=len(candidates), results=len(chunks), ticker_filter=ticker_filter, + filing_date_from=filing_date_from, + filing_date_to=filing_date_to, embedding_ms=round(embedding_ms, 1), ) diff --git a/services/query-api/tests/test_facts.py b/services/query-api/tests/test_facts.py new file mode 100644 index 0000000..0fda46f --- /dev/null +++ b/services/query-api/tests/test_facts.py @@ -0,0 +1,181 @@ +"""Unit tests for XBRL facts lookup and injection (src/rag/facts.py).""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +from src.rag.facts import ( + FactsRepository, + detect_metric_concepts, + extract_years, + format_fact_value, +) + +# ── Metric intent detection ────────────────────────────────────── + + +class TestDetectMetricConcepts: + def test_revenue_keywords(self) -> None: + groups = detect_metric_concepts("What was Apple's total revenue in 2024?") + assert ["RevenueFromContractWithCustomerExcludingAssessedTax", "Revenues", "SalesRevenueNet"] in groups + + def test_net_sales_maps_to_revenue(self) -> None: + groups = detect_metric_concepts("Report net sales for fiscal 2023") + assert any("Revenues" in g for g in groups) + + def test_eps_beats_net_income(self) -> None: + """'diluted earnings per share' must match EPS, and the EPS group + must come first (more specific keyword group).""" + groups = detect_metric_concepts("What was diluted earnings per share?") + assert groups[0] == ["EarningsPerShareDiluted"] + + def test_rnd_abbreviation(self) -> None: + groups = detect_metric_concepts("How much did MSFT spend on R&D?") + assert ["ResearchAndDevelopmentExpense"] in groups + + def test_non_financial_question_matches_nothing(self) -> None: + assert detect_metric_concepts("What are the main supply chain risks?") == [] + + def test_case_insensitive(self) -> None: + assert detect_metric_concepts("NET INCOME please") == [["NetIncomeLoss"]] + + +class TestExtractYears: + def test_single_year(self) -> None: + assert extract_years("revenue in 2024") == [2024] + + def test_multiple_years(self) -> None: + assert extract_years("compare 2022 and 2024") == [2022, 2024] + + def test_no_years(self) -> None: + assert extract_years("latest revenue") == [] + + def test_ignores_non_year_numbers(self) -> None: + assert extract_years("top 100 customers, 391035 million") == [] + + +class TestFormatFactValue: + def test_large_usd_value(self) -> None: + assert format_fact_value(391_035_000_000.0, "USD") == "$391,035,000,000" + + def test_per_share_value(self) -> None: + assert format_fact_value(6.08, "USD/shares") == "$6.08 per share" + + def test_fractional_usd(self) -> None: + assert format_fact_value(1234.5, "USD") == "$1,234.50" + + +# ── FactsRepository ────────────────────────────────────────────── + + +def _repo_with_rows(rows: list[tuple[Any, ...]]) -> tuple[FactsRepository, MagicMock]: + repo = FactsRepository(dsn="postgresql://fake") + mock_cur = MagicMock() + mock_cur.fetchall.return_value = rows + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cur + mock_conn.closed = False + repo._conn = mock_conn + return repo, mock_cur + + +# Row shape: (concept, label, unit, fiscal_year, period_end, value) +_REVENUE_2024 = ( + "RevenueFromContractWithCustomerExcludingAssessedTax", + "Total revenue", + "USD", + 2024, + "2024-09-28", + 391_035_000_000.0, +) +_REVENUE_2023 = ( + "RevenueFromContractWithCustomerExcludingAssessedTax", + "Total revenue", + "USD", + 2023, + "2023-09-30", + 383_285_000_000.0, +) + + +class TestFactsForQuestion: + def test_no_ticker_returns_empty(self) -> None: + repo, mock_cur = _repo_with_rows([_REVENUE_2024]) + assert repo.facts_for_question("What was total revenue in 2024?", None) == [] + mock_cur.execute.assert_not_called() + + def test_no_metric_keywords_returns_empty(self) -> None: + repo, mock_cur = _repo_with_rows([_REVENUE_2024]) + assert repo.facts_for_question("Summarise the risk factors", "AAPL") == [] + mock_cur.execute.assert_not_called() + + def test_injects_fact_chunk(self) -> None: + repo, _ = _repo_with_rows([_REVENUE_2024]) + chunks = repo.facts_for_question("What was Apple's revenue in 2024?", "AAPL") + assert len(chunks) == 1 + chunk = chunks[0] + assert chunk.chunk_id == "xbrl:AAPL:RevenueFromContractWithCustomerExcludingAssessedTax:2024" + assert chunk.section == "XBRL Financial Facts" + assert chunk.relevance_score == 1.0 + assert "$391,035,000,000" in chunk.text + assert "fiscal year 2024" in chunk.text + + def test_no_year_asked_returns_most_recent(self) -> None: + repo, _ = _repo_with_rows([_REVENUE_2024, _REVENUE_2023]) + chunks = repo.facts_for_question("What is Apple's latest revenue?", "AAPL") + assert len(chunks) == 1 + assert "2024" in chunks[0].chunk_id + + def test_specific_years_selected(self) -> None: + repo, _ = _repo_with_rows([_REVENUE_2024, _REVENUE_2023]) + chunks = repo.facts_for_question("Compare revenue in 2023 and 2024", "AAPL") + years = {c.chunk_id.rsplit(":", 1)[1] for c in chunks} + assert years == {"2023", "2024"} + + def test_higher_priority_concept_wins_per_year(self) -> None: + """When both revenue aliases exist for a year, the first concept in + the group's priority order must win.""" + fallback = ("Revenues", "Total revenue", "USD", 2024, "2024-09-28", 1.0) + repo, _ = _repo_with_rows([fallback, _REVENUE_2024]) + chunks = repo.facts_for_question("Revenue in 2024?", "AAPL") + assert len(chunks) == 1 + assert "RevenueFromContractWithCustomerExcludingAssessedTax" in chunks[0].chunk_id + + def test_db_error_degrades_to_no_facts(self) -> None: + repo, mock_cur = _repo_with_rows([]) + mock_cur.execute.side_effect = RuntimeError("relation financial_facts does not exist") + assert repo.facts_for_question("Revenue in 2024?", "AAPL") == [] + + def test_facts_capped_at_max(self) -> None: + rows = [ + ("RevenueFromContractWithCustomerExcludingAssessedTax", "Total revenue", "USD", 2015 + i, + f"{2015 + i}-09-30", float(i)) + for i in range(12) + ] + repo, _ = _repo_with_rows(rows) + question = "Revenue in " + ", ".join(str(2015 + i) for i in range(12)) + chunks = repo.facts_for_question(question, "AAPL") + assert len(chunks) == 8 + + def test_query_filters_by_ticker(self) -> None: + repo, mock_cur = _repo_with_rows([_REVENUE_2024]) + repo.facts_for_question("Revenue in 2024?", "AAPL") + sql, params = mock_cur.execute.call_args[0] + assert "ticker = %s" in sql + assert params[0] == "AAPL" + + +class TestRepositoryLifecycle: + def test_close_without_connect_is_noop(self) -> None: + repo = FactsRepository(dsn="postgresql://fake") + repo.close() # must not raise + assert repo._conn is None + + def test_close_closes_connection(self) -> None: + repo, _ = _repo_with_rows([]) + conn = repo._conn + repo.close() + assert conn is not None + conn.close.assert_called_once() + assert repo._conn is None diff --git a/services/query-api/tests/test_llm_backends.py b/services/query-api/tests/test_llm_backends.py index af0c040..0245091 100644 --- a/services/query-api/tests/test_llm_backends.py +++ b/services/query-api/tests/test_llm_backends.py @@ -105,6 +105,35 @@ async def test_generate_missing_optional_fields(self) -> None: assert result.prompt_tokens == 0 assert result.completion_tokens == 0 + @pytest.mark.asyncio + async def test_generate_disables_thinking(self) -> None: + """Payload must set think=False so reasoning models (Qwen3) return the + grounded answer in `response` instead of spending tokens on thinking.""" + from src.llm.ollama_backend import OllamaBackend + + mock_resp = AsyncMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.json = AsyncMock(return_value={"response": "4"}) + + mock_post_ctx = MagicMock() + mock_post_ctx.__aenter__ = AsyncMock(return_value=mock_resp) + mock_post_ctx.__aexit__ = AsyncMock(return_value=False) + + mock_session = MagicMock() + mock_session.post = MagicMock(return_value=mock_post_ctx) + + mock_session_ctx = MagicMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session) + mock_session_ctx.__aexit__ = AsyncMock(return_value=False) + + with patch("src.llm.ollama_backend.aiohttp.ClientSession", return_value=mock_session_ctx): + backend = OllamaBackend(base_url="http://localhost:11434", model="qwen3.5-122b-ctx") + await backend.generate("What is 2+2?", max_tokens=64) + + payload = mock_session.post.call_args.kwargs["json"] + assert payload["think"] is False + assert payload["options"]["num_predict"] == 64 + # ── OpenAIBackend ──────────────────────────────────────────────── diff --git a/services/query-api/tests/test_main.py b/services/query-api/tests/test_main.py index 6024393..f6d0632 100644 --- a/services/query-api/tests/test_main.py +++ b/services/query-api/tests/test_main.py @@ -330,10 +330,17 @@ def test_query_passes_ticker_filter_and_top_k(self) -> None: captured: dict = {} async def fake_answer( - question: str, top_k: int = 5, ticker_filter: str | None = None + question: str, + top_k: int = 5, + ticker_filter: str | None = None, + filing_date_from: str | None = None, + filing_date_to: str | None = None, + include_source_text: bool = False, ) -> QueryResponse: captured["top_k"] = top_k captured["ticker_filter"] = ticker_filter + captured["filing_date_from"] = filing_date_from + captured["include_source_text"] = include_source_text return _make_query_response() mock_generator = MagicMock() diff --git a/services/query-api/tests/test_retriever.py b/services/query-api/tests/test_retriever.py index 603a7d9..6ea5ec2 100644 --- a/services/query-api/tests/test_retriever.py +++ b/services/query-api/tests/test_retriever.py @@ -11,47 +11,53 @@ from src.rag.prompts import RetrievedChunk +# Unit basis vectors → dot product == cosine similarity, scores stay exact. +_DIM = 384 -class TestRetriever: - """Tests for the Retriever class.""" - @patch("src.rag.retriever.SentenceTransformer") - def test_embed_query_returns_list(self, mock_st_class: MagicMock) -> None: - """embed_query should return a list of floats.""" - from src.rag.retriever import Retriever +def _unit(axis: int) -> np.ndarray: + vec = np.zeros(_DIM) + vec[axis] = 1.0 + return vec - mock_model = MagicMock() - mock_model.encode.return_value = np.array([0.1] * 384) - mock_st_class.return_value = mock_model - with patch("src.rag.retriever.psycopg2"): - retriever = Retriever(dsn="postgresql://fake", model_name="test") +def _row(chunk_id: str, embedding: np.ndarray, score: float = 0.9) -> tuple: + return (chunk_id, "AAPL", "2024-11-01", "Item 1A", f"Text of {chunk_id}", embedding, score) - result = retriever.embed_query("What is Apple's revenue?") - assert isinstance(result, list) - assert len(result) == 384 - @patch("src.rag.retriever.SentenceTransformer") - def test_retrieve_with_ticker_filter(self, mock_st_class: MagicMock) -> None: - """retrieve should include a WHERE clause when ticker_filter is set.""" - from src.rag.retriever import Retriever +def _make_retriever(query_embedding: np.ndarray) -> tuple: + """Build a Retriever with mocked model and DB connection.""" + from src.rag.retriever import Retriever + with patch("src.rag.retriever.SentenceTransformer") as mock_st_class: mock_model = MagicMock() - mock_model.encode.return_value = np.array([0.1] * 384) + mock_model.encode.return_value = query_embedding mock_st_class.return_value = mock_model - with patch("src.rag.retriever.psycopg2"): retriever = Retriever(dsn="postgresql://fake", model_name="test") - # Mock the connection and cursor - mock_cur = MagicMock() - mock_cur.fetchall.return_value = [ - ("chunk1", "AAPL", "2024-11-01", "Item 1A", "Risk text...", np.array([0.1] * 384), 0.87), - ] - mock_conn = MagicMock() - mock_conn.cursor.return_value = mock_cur - mock_conn.closed = False - retriever._conn = mock_conn + mock_cur = MagicMock() + mock_conn = MagicMock() + mock_conn.cursor.return_value = mock_cur + mock_conn.closed = False + retriever._conn = mock_conn + return retriever, mock_cur + + +class TestRetriever: + """Tests for the Retriever class.""" + + def test_embed_query_returns_list(self) -> None: + retriever, _ = _make_retriever(_unit(0)) + result = retriever.embed_query("What is Apple's revenue?") + assert isinstance(result, list) + assert len(result) == _DIM + + def test_retrieve_with_ticker_filter(self) -> None: + """Both legs must include the ticker WHERE clause; the fused chunk + carries cosine similarity to the query as relevance_score.""" + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] chunks, embedding, emb_ms = retriever.retrieve( "What are Apple's risks?", top_k=5, ticker_filter="AAPL" @@ -59,35 +65,102 @@ def test_retrieve_with_ticker_filter(self, mock_st_class: MagicMock) -> None: assert len(chunks) == 1 assert chunks[0].ticker == "AAPL" - assert chunks[0].relevance_score == 0.87 - # Verify the SQL included ticker filter - executed_sql = mock_cur.execute.call_args[0][0] - assert "ticker = %s" in executed_sql - - @patch("src.rag.retriever.SentenceTransformer") - def test_retrieve_without_ticker_filter(self, mock_st_class: MagicMock) -> None: - """retrieve without filter should not have WHERE ticker clause.""" - from src.rag.retriever import Retriever + # chunk embedding == query embedding → cosine similarity 1.0 + assert chunks[0].relevance_score == 1.0 + assert len(embedding) == _DIM + # Both legs (vector + lexical) ran, each with the ticker filter. + executed = [call[0][0] for call in mock_cur.execute.call_args_list] + assert len(executed) == 2 + for sql in executed: + assert "ticker = %s" in sql + + def test_retrieve_without_ticker_filter(self) -> None: + retriever, mock_cur = _make_retriever(_unit(0)) + mock_cur.fetchall.return_value = [] - mock_model = MagicMock() - mock_model.encode.return_value = np.array([0.1] * 384) - mock_st_class.return_value = mock_model + chunks, _, _ = retriever.retrieve("General question", top_k=3) - with patch("src.rag.retriever.psycopg2"): - retriever = Retriever(dsn="postgresql://fake", model_name="test") + assert chunks == [] + for call in mock_cur.execute.call_args_list: + assert "ticker = %s" not in call[0][0] - mock_cur = MagicMock() + def test_retrieve_with_date_filters(self) -> None: + retriever, mock_cur = _make_retriever(_unit(0)) mock_cur.fetchall.return_value = [] - mock_conn = MagicMock() - mock_conn.cursor.return_value = mock_cur - mock_conn.closed = False - retriever._conn = mock_conn - chunks, _, _ = retriever.retrieve("General question", top_k=3) + retriever.retrieve( + "Revenue trends", + top_k=3, + filing_date_from="2023-01-01", + filing_date_to="2024-12-31", + ) - assert chunks == [] - executed_sql = mock_cur.execute.call_args[0][0] - assert "ticker = %s" not in executed_sql + for call in mock_cur.execute.call_args_list: + sql = call[0][0] + assert "filing_date >= %s" in sql + assert "filing_date <= %s" in sql + + def test_lexical_failure_degrades_to_vector_only(self) -> None: + """If migration 002 is missing, the tsvector query fails — retrieval + must still return vector-leg results instead of raising.""" + retriever, mock_cur = _make_retriever(_unit(0)) + + def execute(sql: str, params: object = None) -> None: + if "ts_rank_cd" in sql: + raise RuntimeError('column "chunk_tsv" does not exist') + + mock_cur.execute.side_effect = execute + mock_cur.fetchall.return_value = [_row("chunk1", _unit(0))] + + chunks, _, _ = retriever.retrieve("What are Apple's risks?", top_k=5) + + assert len(chunks) == 1 + assert chunks[0].chunk_id == "chunk1" + + +class TestFuseRRF: + """Tests for Reciprocal Rank Fusion of the two retrieval legs.""" + + def test_empty_lists(self) -> None: + from src.rag.retriever import _fuse_rrf + assert _fuse_rrf([[], []], _unit(0)) == [] + + def test_chunk_in_both_legs_outranks_single_leg_chunk(self) -> None: + from src.rag.retriever import _fuse_rrf + both = _row("in-both", _unit(1)) + vector_only = _row("vector-only", _unit(2)) + fused = _fuse_rrf([[vector_only, both], [both]], _unit(0)) + assert [c.chunk_id for c, _ in fused] == ["in-both", "vector-only"] + + def test_duplicate_chunk_returned_once(self) -> None: + from src.rag.retriever import _fuse_rrf + row = _row("chunk1", _unit(1)) + fused = _fuse_rrf([[row], [row]], _unit(0)) + assert len(fused) == 1 + + def test_relevance_is_cosine_with_query(self) -> None: + from src.rag.retriever import _fuse_rrf + query = _unit(0) + aligned = _row("aligned", query.copy()) + orthogonal = _row("orthogonal", _unit(1)) + fused = _fuse_rrf([[aligned, orthogonal]], query) + scores = {c.chunk_id: c.relevance_score for c, _ in fused} + assert scores["aligned"] == 1.0 + assert scores["orthogonal"] == 0.0 + + +class TestBuildFilters: + def test_no_filters(self) -> None: + from src.rag.retriever import Retriever + clauses, params = Retriever._build_filters(None, None, None) + assert clauses == [] + assert params == [] + + def test_all_filters(self) -> None: + from src.rag.retriever import Retriever + clauses, params = Retriever._build_filters("AAPL", "2023-01-01", "2024-12-31") + assert clauses == ["ticker = %s", "filing_date >= %s", "filing_date <= %s"] + assert params == ["AAPL", "2023-01-01", "2024-12-31"] class TestEmbeddingModelConsistency: