spb/uqo-research-db
Public
1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819from pathlib import Path20from typing import Any2122import pandas as pd2324from src.utils.constants import DUCKDB_PATH25from src.utils.logging import get_logger2627logger = get_logger("api.queries")282930class ResearchAPI:3132 def __init__(self, db_path: Path | str | None = None) -> None:33 self._db_path = Path(db_path) if db_path else DUCKDB_PATH34 self._db: Any = None35 self._emb: Any = None36 self._search: Any = None37 self._panel: Any = None3839 @property40 def db(self):41 if self._db is None:42 from src.storage.database import DatabaseManager43 self._db = DatabaseManager(self._db_path)44 self._db.initialize()45 return self._db4647 @property48 def embedding_manager(self):49 if self._emb is None:50 from src.embeddings.manager import EmbeddingManager51 from src.storage.parquet_manager import ParquetManager52 parquet = ParquetManager()53 self._emb = EmbeddingManager(self.db, parquet)54 return self._emb5556 @property57 def semantic_search(self):58 if self._search is None:59 from src.search.semantic import SemanticSearch60 self._search = SemanticSearch(self.db, self.embedding_manager)61 return self._search6263 @property64 def panel_builder(self):65 if self._panel is None:66 from src.analytics.panel import PanelBuilder67 self._panel = PanelBuilder(self.db)68 return self._panel6970 def get_filings(self, **kwargs) -> pd.DataFrame:71 return self.db.get_filings(**kwargs)7273 def get_sections(self, accession_number: str) -> pd.DataFrame:74 return self.db.get_sections(accession_number)7576 def get_chunks(77 self,78 accession_number: str | None = None,79 section_id: str | None = None,80 ) -> pd.DataFrame:81 return self.db.get_chunks(accession_number=accession_number, section_id=section_id)8283 def search(self, query: str, top_k: int = 20, **filters) -> pd.DataFrame:84 return self.semantic_search.search(query, top_k=top_k, **filters)8586 def build_panel(self, **kwargs) -> pd.DataFrame:87 return self.panel_builder.build_panel(**kwargs)8889 def get_company_history(self, ticker: str) -> pd.DataFrame:90 sql = """91 SELECT92 f.accession_number,93 f.form_type,94 f.filing_date,95 f.fiscal_year,96 f.fiscal_quarter,97 f.download_status,98 f.parse_status,99 COUNT(DISTINCT s.section_id) AS section_count,100 COUNT(DISTINCT ch.chunk_id) AS chunk_count,101 COALESCE(SUM(s.word_count), 0) AS total_word_count102 FROM filings f103 LEFT JOIN filing_sections s ON f.accession_number = s.accession_number104 LEFT JOIN filing_chunks ch ON f.accession_number = ch.accession_number105 WHERE f.ticker = ?106 GROUP BY f.accession_number, f.form_type, f.filing_date,107 f.fiscal_year, f.fiscal_quarter, f.download_status, f.parse_status108 ORDER BY f.filing_date DESC109 """110 return self.db.fetch_df(sql, [ticker])111112 def get_stats(self) -> dict:113 return self.db.get_stats()114115 def sql(self, query: str) -> pd.DataFrame:116 return self.db.fetch_df(query)117118 def close(self) -> None:119 if self._db is not None:120 self._db.close()121 self._db = None122 self._emb = None123 self._search = None124 self._panel = None125126 def __enter__(self) -> ResearchAPI:127 return self128129 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:130 self.close()131