feat: implement Phase 3 -- LLM enrichment, labeler, review queue, dataset export
- libs/llm/: OllamaClient (httpx async, retry), LLMCacheStore (SHA-256 DB cache), prompt registry (event_classifier_v1), LLMParser (cache→prompt→validate→repair) - libs/parser/merger.py: rule+LLM canonical merge with provenance tracking, conflict detection (both confident + disagree), should_queue_for_review() - libs/db/models.py: LLMCallCache, ReviewItem, EventLabel 3개 ORM 모델 추가 - libs/db/migrations/versions/0002_phase3_tables.py: Phase 3 Alembic migration - libs/labeler/: filing_time_bucket→reaction_date, 1D/3D/5D fwd return, MFE/MAE - libs/review/queue.py: create(dedup)/resolve/list ReviewItem - libs/export/snapshot_export.py: temporal split + Parquet + manifest.json - apps/: label_generator, dataset_export, review CLI, gold set evaluator - 42개 신규 테스트 추가 (unit 32 + integration 4 + replay 1) — 152/152 통과 - libs/common/config.py: OLLAMA_URL/MODEL/TIMEOUT 설정 추가 - libs/common/logging.py: bugfix — add_logger_name incompatible with PrintLoggerFactory Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>main
parent
7fed850fbf
commit
bbe0e31150
@ -0,0 +1,69 @@
|
||||
"""Dataset Export: snapshot features + labels to Parquet."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from libs.common.config import get_settings
|
||||
from libs.common.logging import configure_logging, get_logger
|
||||
from libs.db.session import get_session
|
||||
from libs.export.snapshot_export import export_dataset_snapshot
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def run_dataset_export(
|
||||
snapshot_id: str | None,
|
||||
split_policy: str,
|
||||
output_dir: str,
|
||||
) -> dict:
|
||||
async with get_session() as session:
|
||||
manifest = await export_dataset_snapshot(
|
||||
session=session,
|
||||
snapshot_id=snapshot_id,
|
||||
split_policy=split_policy,
|
||||
output_dir=output_dir,
|
||||
)
|
||||
return manifest
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Dataset Export")
|
||||
parser.add_argument("--snapshot-id", default=None, help="Snapshot ID (UUID, auto-generated if omitted)")
|
||||
parser.add_argument(
|
||||
"--split-policy",
|
||||
default="temporal_70_15_15",
|
||||
help="Split policy string (default: temporal_70_15_15)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="./data/datasets/snapshots",
|
||||
help="Output directory for Parquet files",
|
||||
)
|
||||
parser.add_argument("--json", action="store_true", help="Print manifest JSON to stdout")
|
||||
args = parser.parse_args()
|
||||
|
||||
settings = get_settings()
|
||||
configure_logging(settings.log_level)
|
||||
|
||||
manifest = asyncio.run(
|
||||
run_dataset_export(
|
||||
snapshot_id=args.snapshot_id,
|
||||
split_policy=args.split_policy,
|
||||
output_dir=args.output_dir,
|
||||
)
|
||||
)
|
||||
|
||||
if args.json:
|
||||
print(json.dumps(manifest, indent=2))
|
||||
else:
|
||||
print(f"Snapshot exported: {manifest['snapshot_id']}")
|
||||
print(f" Output: {manifest['output_dir']}")
|
||||
print(f" Total rows: {manifest['total_rows']}")
|
||||
for split, count in manifest["row_counts"].items():
|
||||
print(f" {split}: {count} rows")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,129 @@
|
||||
"""Label Generator: compute forward-return labels for all valid events."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import datetime as dt
|
||||
import uuid
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from libs.common.config import get_settings
|
||||
from libs.common.ids import new_job_run_id
|
||||
from libs.common.logging import bind_job_run_id, configure_logging, get_logger
|
||||
from libs.db.models import Event, EventLabel, JobRun, SymbolMaster
|
||||
from libs.db.session import get_session
|
||||
from libs.labeler.label_generator import LABEL_VERSION, generate_labels
|
||||
from libs.oracle_client.client import make_oracle_client
|
||||
from libs.oracle_client.price import PriceService
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def run_label_generator(
|
||||
run_id: str,
|
||||
entry_convention: str = "next_open_after_reaction_close",
|
||||
event_id_filter: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
stats = {"seen": 0, "labeled": 0, "skipped": 0, "errors": 0}
|
||||
|
||||
async with get_session() as session, make_oracle_client() as oracle:
|
||||
price_svc = PriceService(oracle)
|
||||
|
||||
job = JobRun(
|
||||
job_run_id=uuid.UUID(run_id),
|
||||
job_name="label_generator",
|
||||
source_name="oracle",
|
||||
run_date=dt.date.today(),
|
||||
status="running",
|
||||
)
|
||||
session.add(job)
|
||||
await session.flush()
|
||||
|
||||
# Query: events with status=valid (or specific event_id)
|
||||
stmt = select(Event, SymbolMaster).join(
|
||||
SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id, isouter=True
|
||||
).where(Event.status == "valid")
|
||||
|
||||
if event_id_filter:
|
||||
stmt = stmt.where(Event.event_id == event_id_filter)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
rows = result.all()
|
||||
stats["seen"] = len(rows)
|
||||
|
||||
for event, symbol in rows:
|
||||
if symbol is None:
|
||||
logger.warning("label_no_symbol", event_id=event.event_id)
|
||||
stats["skipped"] += 1
|
||||
continue
|
||||
|
||||
# Skip if label already exists
|
||||
existing = await session.execute(
|
||||
select(EventLabel).where(
|
||||
EventLabel.event_id == event.event_id,
|
||||
EventLabel.entry_convention == entry_convention,
|
||||
EventLabel.label_version == LABEL_VERSION,
|
||||
)
|
||||
)
|
||||
if existing.scalar_one_or_none() is not None:
|
||||
stats["skipped"] += 1
|
||||
continue
|
||||
|
||||
try:
|
||||
label = await generate_labels(
|
||||
session=session,
|
||||
event=event,
|
||||
price_svc=price_svc,
|
||||
ticker=symbol.ticker,
|
||||
entry_convention=entry_convention,
|
||||
)
|
||||
session.add(label)
|
||||
await session.flush()
|
||||
stats["labeled"] += 1
|
||||
logger.info(
|
||||
"label_created",
|
||||
event_id=event.event_id,
|
||||
label_status=label.label_status,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("label_error", event_id=event.event_id, error=str(exc))
|
||||
stats["errors"] += 1
|
||||
|
||||
job.status = "succeeded" if stats["errors"] == 0 else "partial"
|
||||
job.finished_at_utc = dt.datetime.now(tz=dt.UTC)
|
||||
job.records_seen = stats["seen"]
|
||||
job.records_written = stats["labeled"]
|
||||
job.error_count = stats["errors"]
|
||||
|
||||
logger.info("label_generator_done", **stats)
|
||||
return stats
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Event Label Generator")
|
||||
parser.add_argument("--run-id", default=new_job_run_id())
|
||||
parser.add_argument(
|
||||
"--entry-convention",
|
||||
default="next_open_after_reaction_close",
|
||||
choices=["next_open_after_reaction_close", "reaction_close"],
|
||||
help="Entry price convention",
|
||||
)
|
||||
parser.add_argument("--event-id", default=None, help="Process a single event by ID")
|
||||
args = parser.parse_args()
|
||||
|
||||
settings = get_settings()
|
||||
configure_logging(settings.log_level)
|
||||
bind_job_run_id(args.run_id)
|
||||
|
||||
asyncio.run(
|
||||
run_label_generator(
|
||||
run_id=args.run_id,
|
||||
entry_convention=args.entry_convention,
|
||||
event_id_filter=args.event_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,198 @@
|
||||
"""QA Tool: evaluate parser output against a gold set of manually labeled documents."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from libs.common.config import get_settings
|
||||
from libs.common.logging import configure_logging, get_logger
|
||||
from libs.parser.rule_parser import RuleBasedParser
|
||||
from libs.parser.text_normalizer import normalize_text
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_GOLD_SET_DIR = Path("data/gold_set")
|
||||
|
||||
_DIRECTION_FIELDS = ("event_type", "event_direction")
|
||||
_ONEOFF_FIELD = "oneoff_item"
|
||||
|
||||
|
||||
def _load_gold_set(gold_dir: Path) -> list[dict[str, Any]]:
|
||||
"""Load all gold set JSON files from directory."""
|
||||
entries = []
|
||||
for path in sorted(gold_dir.glob("*.json")):
|
||||
data = json.loads(path.read_text())
|
||||
if isinstance(data, list):
|
||||
entries.extend(data)
|
||||
else:
|
||||
entries.append(data)
|
||||
return entries
|
||||
|
||||
|
||||
def _precision_recall_f1(tp: int, fp: int, fn: int) -> dict[str, float]:
|
||||
precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
|
||||
recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
|
||||
f1 = (
|
||||
2 * precision * recall / (precision + recall)
|
||||
if (precision + recall) > 0
|
||||
else 0.0
|
||||
)
|
||||
return {"precision": precision, "recall": recall, "f1": f1}
|
||||
|
||||
|
||||
def evaluate(gold_entries: list[dict[str, Any]], parser: RuleBasedParser) -> dict[str, Any]:
|
||||
"""Run parser on gold set and compute accuracy metrics.
|
||||
|
||||
Each gold entry must have:
|
||||
- document_id: str
|
||||
- form_type: str
|
||||
- text: str (raw exhibit text)
|
||||
- gold_event_type: str
|
||||
- gold_event_direction: str
|
||||
- gold_guidance_status: str (optional)
|
||||
- gold_oneoff: bool (optional)
|
||||
- gold_quality_assessment: str (optional)
|
||||
"""
|
||||
total = len(gold_entries)
|
||||
if total == 0:
|
||||
return {"error": "Empty gold set"}
|
||||
|
||||
# Accuracy trackers
|
||||
event_type_correct = 0
|
||||
direction_correct = 0
|
||||
guidance_correct = 0
|
||||
guidance_total = 0
|
||||
|
||||
# Oneoff precision/recall
|
||||
oneoff_tp = oneoff_fp = oneoff_fn = 0
|
||||
|
||||
# Evidence presence
|
||||
evidence_present = 0
|
||||
|
||||
per_doc: list[dict[str, Any]] = []
|
||||
|
||||
for entry in gold_entries:
|
||||
text = normalize_text(entry.get("text", ""))
|
||||
meta = {
|
||||
"filing_date": entry.get("filing_date", "1900-01-01"),
|
||||
"form_type": entry.get("form_type", "8-K"),
|
||||
}
|
||||
output = parser.parse(
|
||||
document_id=entry.get("document_id", "unknown"),
|
||||
form_type=entry.get("form_type", "8-K"),
|
||||
text=text,
|
||||
metadata=meta,
|
||||
)
|
||||
out = output.model_dump()
|
||||
|
||||
gold_et = entry.get("gold_event_type", "")
|
||||
gold_dir = entry.get("gold_event_direction", "")
|
||||
gold_guidance = entry.get("gold_guidance_status", None)
|
||||
gold_oneoff = entry.get("gold_oneoff", None)
|
||||
|
||||
et_match = out["event_type"] == gold_et
|
||||
dir_match = out["event_direction"] == gold_dir
|
||||
event_type_correct += int(et_match)
|
||||
direction_correct += int(dir_match)
|
||||
|
||||
if gold_guidance is not None:
|
||||
guidance_correct += int(out["guidance"]["status"] == gold_guidance)
|
||||
guidance_total += 1
|
||||
|
||||
if gold_oneoff is not None:
|
||||
pred_oneoff = out["risk_flags"]["oneoff_item"]
|
||||
if pred_oneoff and gold_oneoff:
|
||||
oneoff_tp += 1
|
||||
elif pred_oneoff and not gold_oneoff:
|
||||
oneoff_fp += 1
|
||||
elif not pred_oneoff and gold_oneoff:
|
||||
oneoff_fn += 1
|
||||
|
||||
has_evidence = len(out.get("evidence", [])) > 0
|
||||
evidence_present += int(has_evidence)
|
||||
|
||||
per_doc.append({
|
||||
"document_id": entry.get("document_id"),
|
||||
"et_match": et_match,
|
||||
"dir_match": dir_match,
|
||||
"pred_event_type": out["event_type"],
|
||||
"gold_event_type": gold_et,
|
||||
"pred_direction": out["event_direction"],
|
||||
"gold_direction": gold_dir,
|
||||
"confidence": out["confidence"]["overall"],
|
||||
})
|
||||
|
||||
oneoff_metrics = _precision_recall_f1(oneoff_tp, oneoff_fp, oneoff_fn)
|
||||
summary: dict[str, Any] = {
|
||||
"total": total,
|
||||
"event_type_accuracy": event_type_correct / total,
|
||||
"direction_accuracy": direction_correct / total,
|
||||
"guidance_accuracy": guidance_correct / guidance_total if guidance_total > 0 else None,
|
||||
"oneoff_precision": oneoff_metrics["precision"],
|
||||
"oneoff_recall": oneoff_metrics["recall"],
|
||||
"oneoff_f1": oneoff_metrics["f1"],
|
||||
"evidence_presence_ratio": evidence_present / total,
|
||||
"per_document": per_doc,
|
||||
}
|
||||
return summary
|
||||
|
||||
|
||||
def main() -> None:
|
||||
arg_parser = argparse.ArgumentParser(description="Evaluate parser against gold set")
|
||||
arg_parser.add_argument(
|
||||
"--gold-dir",
|
||||
default=str(_GOLD_SET_DIR),
|
||||
help="Directory containing gold set JSON files",
|
||||
)
|
||||
arg_parser.add_argument("--output", default=None, help="Write JSON report to this file")
|
||||
arg_parser.add_argument("--verbose", action="store_true", help="Print per-doc results")
|
||||
args = arg_parser.parse_args()
|
||||
|
||||
settings = get_settings()
|
||||
configure_logging(settings.log_level)
|
||||
|
||||
gold_dir = Path(args.gold_dir)
|
||||
if not gold_dir.exists():
|
||||
print(f"Gold set directory not found: {gold_dir}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
gold_entries = _load_gold_set(gold_dir)
|
||||
if not gold_entries:
|
||||
print(f"No gold set entries found in {gold_dir}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
parser = RuleBasedParser()
|
||||
report = evaluate(gold_entries, parser)
|
||||
|
||||
print(f"\n=== Gold Set Evaluation Report ({report['total']} documents) ===")
|
||||
print(f" Event type accuracy : {report['event_type_accuracy']:.1%}")
|
||||
print(f" Direction accuracy : {report['direction_accuracy']:.1%}")
|
||||
if report["guidance_accuracy"] is not None:
|
||||
print(f" Guidance accuracy : {report['guidance_accuracy']:.1%}")
|
||||
print(f" One-off F1 : {report['oneoff_f1']:.3f}")
|
||||
print(f" Precision : {report['oneoff_precision']:.3f}")
|
||||
print(f" Recall : {report['oneoff_recall']:.3f}")
|
||||
print(f" Evidence presence : {report['evidence_presence_ratio']:.1%}")
|
||||
|
||||
if args.verbose and report.get("per_document"):
|
||||
print("\n--- Per-document results ---")
|
||||
for doc in report["per_document"]:
|
||||
status = "OK" if doc["et_match"] and doc["dir_match"] else "MISMATCH"
|
||||
print(
|
||||
f" [{status}] {doc['document_id']}: "
|
||||
f"type={doc['pred_event_type']}({doc['gold_event_type']}) "
|
||||
f"dir={doc['pred_direction']}({doc['gold_direction']}) "
|
||||
f"conf={doc['confidence']:.2f}"
|
||||
)
|
||||
|
||||
if args.output:
|
||||
out_path = Path(args.output)
|
||||
out_path.write_text(json.dumps(report, indent=2))
|
||||
print(f"\nReport written to {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,145 @@
|
||||
"""Review Queue CLI: list, show, and resolve review items."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
|
||||
from libs.common.config import get_settings
|
||||
from libs.common.logging import configure_logging, get_logger
|
||||
from libs.db.session import get_session
|
||||
from libs.review.queue import list_review_items, resolve_review_item
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _fmt_item(item: object) -> str:
|
||||
"""Format a ReviewItem for display."""
|
||||
lines = [
|
||||
f" review_id : {item.review_id}", # type: ignore[attr-defined]
|
||||
f" entity_type : {item.entity_type}", # type: ignore[attr-defined]
|
||||
f" entity_id : {item.entity_id}", # type: ignore[attr-defined]
|
||||
f" priority : {item.priority}", # type: ignore[attr-defined]
|
||||
f" status : {item.status}", # type: ignore[attr-defined]
|
||||
f" reasons : {item.reason_codes}", # type: ignore[attr-defined]
|
||||
f" created_at : {item.created_at_utc}", # type: ignore[attr-defined]
|
||||
]
|
||||
reviewer_id = getattr(item, "reviewer_id", None)
|
||||
if reviewer_id:
|
||||
lines.append(f" reviewer : {reviewer_id}")
|
||||
if item.resolution_type: # type: ignore[attr-defined]
|
||||
lines.append(f" resolution : {item.resolution_type}") # type: ignore[attr-defined]
|
||||
lines.append(f" root_cause : {item.root_cause}") # type: ignore[attr-defined]
|
||||
lines.append(f" notes : {item.notes}") # type: ignore[attr-defined]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def cmd_list(args: argparse.Namespace) -> None:
|
||||
async with get_session() as session:
|
||||
items = await list_review_items(
|
||||
session,
|
||||
status=args.status,
|
||||
priority=getattr(args, "priority", None),
|
||||
)
|
||||
if not items:
|
||||
print("No review items found.")
|
||||
return
|
||||
print(f"Found {len(items)} review item(s):\n")
|
||||
for item in items:
|
||||
print(f"[{item.priority}] {item.review_id}") # type: ignore[attr-defined]
|
||||
print(f" entity: {item.entity_type}/{item.entity_id}") # type: ignore[attr-defined]
|
||||
print(f" status: {item.status} | reasons: {item.reason_codes}") # type: ignore[attr-defined]
|
||||
print()
|
||||
|
||||
|
||||
async def cmd_show(args: argparse.Namespace) -> None:
|
||||
from sqlalchemy import select
|
||||
|
||||
from libs.db.models import ReviewItem
|
||||
|
||||
async with get_session() as session:
|
||||
import uuid
|
||||
|
||||
result = await session.execute(
|
||||
select(ReviewItem).where(ReviewItem.review_id == uuid.UUID(args.review_id))
|
||||
)
|
||||
item = result.scalar_one_or_none()
|
||||
|
||||
if item is None:
|
||||
print(f"Review item {args.review_id} not found.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(_fmt_item(item))
|
||||
if item.snapshot_refs: # type: ignore[attr-defined]
|
||||
print(f" snapshot_refs: {json.dumps(item.snapshot_refs, indent=4)}") # type: ignore[attr-defined]
|
||||
if item.suggested_overrides: # type: ignore[attr-defined]
|
||||
print(f" suggested_overrides: {json.dumps(item.suggested_overrides, indent=4)}") # type: ignore[attr-defined]
|
||||
|
||||
|
||||
async def cmd_resolve(args: argparse.Namespace) -> None:
|
||||
async with get_session() as session:
|
||||
item = await resolve_review_item(
|
||||
session=session,
|
||||
review_id=args.review_id,
|
||||
reviewer_id=args.reviewer,
|
||||
resolution_type=args.resolution_type,
|
||||
root_cause=args.root_cause or "",
|
||||
notes=args.notes or "",
|
||||
field_overrides=json.loads(args.overrides) if args.overrides else None,
|
||||
)
|
||||
await session.commit()
|
||||
print(f"Resolved review item {item.review_id} as '{item.resolution_type}'.") # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description="ACE-F Review Queue CLI")
|
||||
sub = p.add_subparsers(dest="command", required=True)
|
||||
|
||||
# list
|
||||
p_list = sub.add_parser("list", help="List review items")
|
||||
p_list.add_argument("--status", default="open", help="Filter by status (default: open)")
|
||||
p_list.add_argument("--priority", default=None, help="Filter by priority (P0/P1/P2)")
|
||||
|
||||
# show
|
||||
p_show = sub.add_parser("show", help="Show review item details")
|
||||
p_show.add_argument("review_id", help="UUID of the review item")
|
||||
|
||||
# resolve
|
||||
p_resolve = sub.add_parser("resolve", help="Resolve a review item")
|
||||
p_resolve.add_argument("review_id", help="UUID of the review item")
|
||||
p_resolve.add_argument("--reviewer", required=True, help="Reviewer ID/name")
|
||||
p_resolve.add_argument(
|
||||
"--resolution-type", required=True,
|
||||
choices=["accepted", "corrected", "wont_fix", "escalated"],
|
||||
)
|
||||
p_resolve.add_argument("--root-cause", default="", help="Root cause description")
|
||||
p_resolve.add_argument("--notes", default="", help="Free-text notes")
|
||||
p_resolve.add_argument("--overrides", default=None, help="JSON string of field overrides")
|
||||
|
||||
return p
|
||||
|
||||
|
||||
def main() -> None:
|
||||
settings = get_settings()
|
||||
configure_logging(settings.log_level)
|
||||
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
handlers = {
|
||||
"list": cmd_list,
|
||||
"show": cmd_show,
|
||||
"resolve": cmd_resolve,
|
||||
}
|
||||
|
||||
handler = handlers.get(args.command)
|
||||
if handler is None:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
asyncio.run(handler(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,81 +1,151 @@
|
||||
# Phase 3 Testing Checklist
|
||||
|
||||
> Phase 3 implementation status: **COMPLETE** (2026-03-12)
|
||||
|
||||
## 1. 단위 테스트
|
||||
|
||||
### text normalizer
|
||||
- HTML 문서가 안정적으로 plain text로 변환된다.
|
||||
- 동일 문서에 대해 해시가 안정적으로 재생산된다.
|
||||
- disclaimer strip 옵션이 본문을 과도하게 삭제하지 않는다.
|
||||
- [x] HTML 문서가 안정적으로 plain text로 변환된다. → `test_text_normalizer.py`
|
||||
- [x] 동일 문서에 대해 해시가 안정적으로 재생산된다. → `test_text_normalizer.py`
|
||||
- [x] disclaimer strip 옵션이 본문을 과도하게 삭제하지 않는다. → `test_text_normalizer.py`
|
||||
|
||||
### rule parser
|
||||
- guidance 키워드가 올바르게 분류된다.
|
||||
- one-off 키워드가 검출된다.
|
||||
- demand/pricing/margin 키워드가 검출된다.
|
||||
- section parser가 없는 문서에서도 안전하게 실패한다.
|
||||
- [x] guidance 키워드가 올바르게 분류된다. → `test_rule_parser.py`
|
||||
- [x] one-off 키워드가 검출된다. → `test_rule_parser.py`
|
||||
- [x] demand/pricing/margin 키워드가 검출된다. → `test_rule_parser.py`
|
||||
- [x] section parser가 없는 문서에서도 안전하게 실패한다. → `test_rule_parser.py`
|
||||
|
||||
### span mapper
|
||||
- evidence char offsets가 원문 구간과 일치한다.
|
||||
- normalization 후에도 span reference가 추적 가능하다.
|
||||
- [ ] evidence char offsets가 원문 구간과 일치한다.
|
||||
- [ ] normalization 후에도 span reference가 추적 가능하다.
|
||||
|
||||
### llm wrapper (Ollama)
|
||||
- [x] 캐시 히트 시 외부 호출이 발생하지 않는다. → `test_llm_client.py::TestOllamaClientChat::test_cache_hit_skips_llm_call`
|
||||
- [x] 타임아웃 시 LLMTimeoutError 발생 → `test_llm_client.py::TestOllamaClientChat::test_chat_timeout_raises_llm_timeout_error`
|
||||
- [x] 5xx 응답 시 RetryableError 발생 → `test_llm_client.py::TestOllamaClientChat::test_chat_5xx_raises_retryable_error`
|
||||
- [x] 성공 응답 파싱 → `test_llm_client.py::TestOllamaClientChat::test_chat_success_returns_parsed_json`
|
||||
|
||||
### llm wrapper
|
||||
- 캐시 히트 시 외부 호출이 발생하지 않는다.
|
||||
- 타임아웃/레이트리밋 시 재시도 정책이 지켜진다.
|
||||
- schema invalid 응답이 repair path를 탄다.
|
||||
### LLM cache
|
||||
- [x] 캐시 미스 시 None 반환 → `test_llm_cache.py::TestLLMCacheStore::test_get_returns_none_on_miss`
|
||||
- [x] 캐시 put이 DB 행 생성 → `test_llm_cache.py::TestLLMCacheStore::test_put_stores_entry`
|
||||
|
||||
### canonical merge
|
||||
- rule/llm 충돌 시 보수적 merge가 적용된다.
|
||||
- provenance 필드가 누락되지 않는다.
|
||||
- review queue trigger가 올바르게 작동한다.
|
||||
- [x] rule 우선 (both confident) → `test_merger.py::TestMerge::test_rule_wins_when_both_confident`
|
||||
- [x] llm이 unknown signal 채움 → `test_merger.py::TestMerge::test_llm_fills_unknown_signals`
|
||||
- [x] 충돌 시 보수적 merge + warning → `test_merger.py::TestMerge::test_conflict_flags_rule_llm_conflict`
|
||||
- [x] llm=None → all provenance=rule → `test_merger.py::TestMerge::test_llm_none_all_provenance_is_rule`
|
||||
- [x] risk_flags OR 결합 → `test_merger.py::TestMerge::test_risk_flags_are_ored`
|
||||
- [x] review queue trigger (low confidence) → `test_merger.py::TestShouldQueueForReview::test_low_confidence_triggers_review`
|
||||
- [x] review queue trigger (conflict) → `test_merger.py::TestShouldQueueForReview::test_conflict_triggers_review`
|
||||
- [x] review queue trigger (oneoff) → `test_merger.py::TestShouldQueueForReview::test_oneoff_triggers_review`
|
||||
|
||||
### feature builder
|
||||
- reaction_close_location 계산이 정확하다.
|
||||
- rolling window가 미래 데이터를 보지 않는다.
|
||||
- null feature가 정책대로 처리된다.
|
||||
- available_ts가 entry convention과 맞는다.
|
||||
- [x] reaction_close_location 계산이 정확하다. → `test_event_features.py`
|
||||
- [x] rolling window가 미래 데이터를 보지 않는다. → `test_market_features.py`
|
||||
- [x] null feature가 정책대로 처리된다. → `test_event_features.py`
|
||||
|
||||
### labeler
|
||||
- reaction date 계산이 장전/장중/장후에 맞게 동작한다.
|
||||
- 1D/3D/5D forward return이 정확하다.
|
||||
- MFE/MAE가 고저가 경로로 정확히 계산된다.
|
||||
- 비거래일/휴일 처리에 오류가 없다.
|
||||
- [x] pre_market/regular_hours → same-day reaction → `test_labeler.py::TestComputeReactionDate::test_pre_market_on_trading_day_returns_same_day`
|
||||
- [x] post_market → next trading day → `test_labeler.py::TestComputeReactionDate::test_post_market_returns_next_trading_day`
|
||||
- [x] unknown → next trading day → `test_labeler.py::TestComputeReactionDate::test_unknown_returns_next_trading_day`
|
||||
- [x] 주말(비거래일) pre_market → next trading day → `test_labeler.py::TestComputeReactionDate::test_pre_market_on_weekend_returns_next_trading_day`
|
||||
- [x] 금요일 post_market → 월요일 → `test_labeler.py::TestComputeReactionDate::test_post_market_on_friday_returns_monday`
|
||||
- [x] 1D forward return 정확 → `test_labeler.py::TestComputeLabelsFromBars::test_1d_return_calculation`
|
||||
- [x] MFE = max(high-entry)/entry → `test_labeler.py::TestComputeLabelsFromBars::test_mfe_is_max_high_minus_entry`
|
||||
- [x] MAE = min(low-entry)/entry → `test_labeler.py::TestComputeLabelsFromBars::test_mae_is_min_low_minus_entry`
|
||||
- [x] hit_pos_1r True/False → `test_labeler.py::TestComputeLabelsFromBars::test_hit_pos_1r_true/false`
|
||||
- [x] close_up_after_3d logic → `test_labeler.py::TestComputeLabelsFromBars::test_close_up_after_3d_true_when_final_close_above_entry`
|
||||
- [x] 빈 bars → empty dict → `test_labeler.py::TestComputeLabelsFromBars::test_empty_bars_returns_empty_dict`
|
||||
- [x] label=ok when price data available → `test_labeler.py::TestGenerateLabels::test_generate_labels_with_valid_prices`
|
||||
- [x] label=unavailable when Oracle fails → `test_labeler.py::TestGenerateLabels::test_generate_labels_unavailable_when_no_price_data`
|
||||
|
||||
### review queue
|
||||
- [x] 신규 ReviewItem 생성 → `test_review_queue.py::TestCreateReviewItem::test_create_new_review_item`
|
||||
- [x] 중복 open item → 업데이트 (priority escalation) → `test_review_queue.py::TestCreateReviewItem::test_deduplicate_open_items`
|
||||
- [x] resolve → status=resolved → `test_review_queue.py::TestCreateReviewItem::test_resolve_review_item`
|
||||
- [x] list_review_items status 필터 → `test_review_queue.py::TestCreateReviewItem::test_list_review_items_with_status_filter`
|
||||
|
||||
### snapshot export
|
||||
- [x] manifest.json 생성 (snapshot_id, created_at, row_counts) → `test_snapshot_export.py::TestExportDatasetSnapshot::test_manifest_is_written`
|
||||
- [x] train/valid/test Parquet 파일 생성 → `test_snapshot_export.py::TestExportDatasetSnapshot::test_parquet_files_created`
|
||||
- [x] temporal split proportions → `test_snapshot_export.py::TestTemporalSplit::test_split_proportions`
|
||||
- [x] temporal split order preserved → `test_snapshot_export.py::TestTemporalSplit::test_split_preserves_temporal_order`
|
||||
|
||||
## 2. 통합 테스트
|
||||
|
||||
- SEC raw 문서 하나가 parser output까지 도달한다.
|
||||
- parser output이 feature builder로 연결된다.
|
||||
- feature + price data가 labeler로 연결된다.
|
||||
- review queue item이 실제로 생성된다.
|
||||
- snapshot export가 manifest 포함해 생성된다.
|
||||
- [x] SEC raw 문서 하나가 parser output까지 도달한다. → `test_filing_pipeline.py`
|
||||
- [x] parser output이 feature builder로 연결된다. → `test_feature_pipeline.py`
|
||||
- [x] feature + price data가 labeler로 연결된다. → `test_label_pipeline.py::test_label_pipeline_end_to_end`
|
||||
- [x] Oracle 실패 시 unavailable label 생성 → `test_label_pipeline.py::test_label_pipeline_handles_missing_price_data`
|
||||
- [x] review queue item이 실제로 생성된다. → `test_review_queue_integration.py::test_low_confidence_merge_creates_review_item`
|
||||
- [x] conflict 시 P0 review item 생성 → `test_review_queue_integration.py::test_conflict_merge_creates_p0_review_item`
|
||||
- [ ] snapshot export가 manifest 포함해 생성된다. (needs real DB)
|
||||
|
||||
## 3. Replay 테스트
|
||||
|
||||
- 동일 문서 재처리 시 canonical output이 동일하다.
|
||||
- 동일 문서 + 동일 prompt_version에서 캐시 결과가 재현된다.
|
||||
- parser_version 변경 시 이전 결과와 diff report 생성 가능하다.
|
||||
- historical day replay가 live path와 같은 코드 경로를 탄다.
|
||||
- [x] 동일 문서 + 동일 prompt_version → LLM 캐시 히트 → `test_llm_cache_replay.py::test_same_document_hits_cache_on_replay`
|
||||
- [ ] 동일 문서 재처리 시 canonical output이 동일하다. → `test_determinism.py`
|
||||
- [ ] parser_version 변경 시 diff report 생성 가능하다.
|
||||
- [ ] historical day replay가 live path와 같은 코드 경로를 탄다.
|
||||
|
||||
## 4. Gold set 테스트
|
||||
|
||||
- event_type accuracy baseline 이상
|
||||
- guidance_direction accuracy baseline 이상
|
||||
- oneoff precision/recall baseline 이상
|
||||
- confidence calibration sanity check
|
||||
- evidence presence ratio 기준 이상
|
||||
- [ ] event_type accuracy baseline 이상 → `apps/qa/evaluate_gold_set.py`
|
||||
- [ ] guidance_direction accuracy baseline 이상 → `apps/qa/evaluate_gold_set.py`
|
||||
- [ ] oneoff precision/recall baseline 이상 → `apps/qa/evaluate_gold_set.py`
|
||||
- [ ] confidence calibration sanity check
|
||||
- [ ] evidence presence ratio 기준 이상 → `apps/qa/evaluate_gold_set.py`
|
||||
|
||||
## 5. Leakage 테스트
|
||||
|
||||
- next_open entry dataset에 entry day 장중 정보가 포함되지 않는다.
|
||||
- FINRA post-close data가 당일 아침 feature로 들어가지 않는다.
|
||||
- forward returns를 만드는 price bars가 feature 계산에 재사용되지 않는다.
|
||||
- snapshot split이 시간 순서를 위반하지 않는다.
|
||||
- [ ] next_open entry dataset에 entry day 장중 정보가 포함되지 않는다.
|
||||
- [ ] FINRA post-close data가 당일 아침 feature로 들어가지 않는다.
|
||||
- [ ] forward returns를 만드는 price bars가 feature 계산에 재사용되지 않는다.
|
||||
- [x] snapshot split이 시간 순서를 위반하지 않는다. → `test_snapshot_export.py::TestTemporalSplit::test_split_preserves_temporal_order`
|
||||
|
||||
## 6. 운영 전 체크리스트
|
||||
|
||||
- parser schema version 고정
|
||||
- prompt version 고정
|
||||
- gold set 리포트 생성 완료
|
||||
- review backlog acceptable
|
||||
- null rate report 검토 완료
|
||||
- label distribution sanity check 완료
|
||||
- dataset manifest에 commit hash 포함
|
||||
- raw prompt/response 보관 정책 확인
|
||||
- [x] parser schema version 고정 → `libs/parser/rule_parser.py::PARSER_VERSION`
|
||||
- [x] prompt version 고정 → `libs/llm/prompts.py::PROMPT_VERSION`
|
||||
- [ ] gold set 리포트 생성 완료
|
||||
- [ ] review backlog acceptable
|
||||
- [ ] null rate report 검토 완료
|
||||
- [ ] label distribution sanity check 완료
|
||||
- [x] dataset manifest에 commit hash 포함 → `libs/export/snapshot_export.py::_get_git_commit_hash`
|
||||
- [x] raw prompt/response 보관 정책 확인 → `libs/db/models.py::LLMCallCache`
|
||||
|
||||
## 7. 신규 파일 목록 (Phase 3)
|
||||
|
||||
| 경로 | 상태 |
|
||||
|---|---|
|
||||
| `libs/llm/__init__.py` | ✅ |
|
||||
| `libs/llm/exceptions.py` | ✅ |
|
||||
| `libs/llm/client.py` | ✅ |
|
||||
| `libs/llm/cache.py` | ✅ |
|
||||
| `libs/llm/prompts.py` | ✅ |
|
||||
| `libs/llm/parser.py` | ✅ |
|
||||
| `libs/parser/merger.py` | ✅ |
|
||||
| `libs/labeler/__init__.py` | ✅ |
|
||||
| `libs/labeler/reaction_date.py` | ✅ |
|
||||
| `libs/labeler/label_generator.py` | ✅ |
|
||||
| `libs/review/__init__.py` | ✅ |
|
||||
| `libs/review/queue.py` | ✅ |
|
||||
| `libs/export/__init__.py` | ✅ |
|
||||
| `libs/export/snapshot_export.py` | ✅ |
|
||||
| `libs/db/models.py` (+ 3 models) | ✅ |
|
||||
| `libs/db/migrations/versions/0002_phase3_tables.py` | ✅ |
|
||||
| `libs/common/config.py` (+ Ollama settings) | ✅ |
|
||||
| `apps/pipeline/label_generator/main.py` | ✅ |
|
||||
| `apps/pipeline/dataset_export/main.py` | ✅ |
|
||||
| `apps/review/cli.py` | ✅ |
|
||||
| `apps/qa/evaluate_gold_set.py` | ✅ |
|
||||
| `tests/unit/test_llm_client.py` | ✅ |
|
||||
| `tests/unit/test_llm_cache.py` | ✅ |
|
||||
| `tests/unit/test_merger.py` | ✅ |
|
||||
| `tests/unit/test_labeler.py` | ✅ |
|
||||
| `tests/unit/test_review_queue.py` | ✅ |
|
||||
| `tests/unit/test_snapshot_export.py` | ✅ |
|
||||
| `tests/integration/test_label_pipeline.py` | ✅ |
|
||||
| `tests/integration/test_review_queue_integration.py` | ✅ |
|
||||
| `tests/replay/test_llm_cache_replay.py` | ✅ |
|
||||
| `.env.example` (+ Ollama vars) | ✅ |
|
||||
|
||||
@ -0,0 +1,153 @@
|
||||
"""Phase 3 tables: llm_call_cache, review_items, event_labels.
|
||||
|
||||
Revision ID: 0002
|
||||
Revises: 0001
|
||||
Create Date: 2026-03-12
|
||||
|
||||
Tables created:
|
||||
- llm_call_cache: DB-backed cache for LLM API calls
|
||||
- review_items: human-in-the-loop review queue
|
||||
- event_labels: forward-return labels for parsed events
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
||||
|
||||
revision: str = "0002"
|
||||
down_revision: str | None = "0001"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"llm_call_cache",
|
||||
sa.Column("cache_key", sa.Text, primary_key=True),
|
||||
sa.Column("document_id", sa.Text, nullable=True),
|
||||
sa.Column("model_name", sa.Text, nullable=False),
|
||||
sa.Column("prompt_version", sa.Text, nullable=False),
|
||||
sa.Column("schema_version", sa.Text, nullable=False),
|
||||
sa.Column("raw_prompt", sa.Text, nullable=False),
|
||||
sa.Column("raw_response", sa.Text, nullable=False),
|
||||
sa.Column("normalized_json", JSONB, nullable=False),
|
||||
sa.Column("token_usage_json", JSONB, nullable=False),
|
||||
sa.Column("elapsed_ms", sa.Integer, nullable=False),
|
||||
sa.Column(
|
||||
"created_at_utc",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.text("now()"),
|
||||
),
|
||||
)
|
||||
|
||||
op.create_table(
|
||||
"review_items",
|
||||
sa.Column("review_id", UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("entity_type", sa.Text, nullable=False),
|
||||
sa.Column("entity_id", sa.Text, nullable=False),
|
||||
sa.Column("priority", sa.Text, nullable=False),
|
||||
sa.Column("reason_codes", JSONB, nullable=False),
|
||||
sa.Column("status", sa.Text, nullable=False, server_default=sa.text("'open'")),
|
||||
sa.Column("assigned_to", sa.Text, nullable=True),
|
||||
sa.Column("snapshot_refs", JSONB, nullable=False, server_default=sa.text("'{}'")),
|
||||
sa.Column("suggested_overrides", JSONB, nullable=True),
|
||||
sa.Column("reviewer_id", sa.Text, nullable=True),
|
||||
sa.Column("resolution_type", sa.Text, nullable=True),
|
||||
sa.Column("field_overrides", JSONB, nullable=True),
|
||||
sa.Column("root_cause", sa.Text, nullable=True),
|
||||
sa.Column("notes", sa.Text, nullable=True),
|
||||
sa.Column("resolved_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column(
|
||||
"created_at_utc",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.text("now()"),
|
||||
),
|
||||
sa.Column(
|
||||
"updated_at_utc",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.text("now()"),
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_review_items_entity_open", "review_items", ["entity_type", "entity_id", "status"]
|
||||
)
|
||||
op.create_index(
|
||||
"ix_review_items_status_priority", "review_items", ["status", "priority"]
|
||||
)
|
||||
# Partial unique index: only one open review per entity
|
||||
op.create_index(
|
||||
"uix_review_items_entity_open",
|
||||
"review_items",
|
||||
["entity_type", "entity_id"],
|
||||
unique=True,
|
||||
postgresql_where=sa.text("status = 'open'"),
|
||||
)
|
||||
|
||||
op.create_table(
|
||||
"event_labels",
|
||||
sa.Column("label_id", sa.BigInteger, primary_key=True, autoincrement=True),
|
||||
sa.Column(
|
||||
"event_id",
|
||||
sa.Text,
|
||||
sa.ForeignKey("events.event_id"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("entry_convention", sa.Text, nullable=False),
|
||||
sa.Column("reaction_date", sa.Date, nullable=True),
|
||||
sa.Column("entry_date", sa.Date, nullable=True),
|
||||
sa.Column("entry_price", sa.Numeric, nullable=True),
|
||||
sa.Column("fwd_return_1d", sa.Numeric, nullable=True),
|
||||
sa.Column("fwd_return_3d", sa.Numeric, nullable=True),
|
||||
sa.Column("fwd_return_5d", sa.Numeric, nullable=True),
|
||||
sa.Column("hit_pos_1r_within_3d", sa.Boolean, nullable=True),
|
||||
sa.Column("hit_neg_1r_within_3d", sa.Boolean, nullable=True),
|
||||
sa.Column("close_up_after_3d", sa.Boolean, nullable=True),
|
||||
sa.Column("close_up_after_5d", sa.Boolean, nullable=True),
|
||||
sa.Column("mfe_3d", sa.Numeric, nullable=True),
|
||||
sa.Column("mae_3d", sa.Numeric, nullable=True),
|
||||
sa.Column("mfe_5d", sa.Numeric, nullable=True),
|
||||
sa.Column("mae_5d", sa.Numeric, nullable=True),
|
||||
sa.Column("bars_to_mfe_3d", sa.Integer, nullable=True),
|
||||
sa.Column("bars_to_mae_3d", sa.Integer, nullable=True),
|
||||
sa.Column("days_to_peak_close_5d", sa.Integer, nullable=True),
|
||||
sa.Column("label_status", sa.Text, nullable=False),
|
||||
sa.Column(
|
||||
"invalid_event_for_labeling",
|
||||
sa.Boolean,
|
||||
nullable=False,
|
||||
server_default=sa.text("false"),
|
||||
),
|
||||
sa.Column("risk_model_name", sa.Text, nullable=True),
|
||||
sa.Column("initial_stop_price", sa.Numeric, nullable=True),
|
||||
sa.Column("initial_r_value", sa.Numeric, nullable=True),
|
||||
sa.Column("label_version", sa.Text, nullable=False),
|
||||
sa.Column(
|
||||
"created_at_utc",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.text("now()"),
|
||||
),
|
||||
sa.UniqueConstraint(
|
||||
"event_id",
|
||||
"entry_convention",
|
||||
"label_version",
|
||||
name="uq_event_labels_event_convention_version",
|
||||
),
|
||||
)
|
||||
op.create_index("ix_event_labels_event_id", "event_labels", ["event_id"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_event_labels_event_id", table_name="event_labels")
|
||||
op.drop_table("event_labels")
|
||||
op.drop_index("uix_review_items_entity_open", table_name="review_items")
|
||||
op.drop_index("ix_review_items_status_priority", table_name="review_items")
|
||||
op.drop_index("ix_review_items_entity_open", table_name="review_items")
|
||||
op.drop_table("review_items")
|
||||
op.drop_table("llm_call_cache")
|
||||
@ -0,0 +1 @@
|
||||
"""Dataset export: snapshot features + labels to Parquet."""
|
||||
@ -0,0 +1,185 @@
|
||||
"""Export feature snapshots + labels to Parquet with train/valid/test split."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.common.time_utils import utc_now
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
MANIFEST_FILENAME = "manifest.json"
|
||||
|
||||
|
||||
def _get_git_commit_hash() -> str:
|
||||
"""Return the current git commit hash (short), or 'unknown'."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-parse", "--short", "HEAD"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
return result.stdout.strip() or "unknown"
|
||||
except Exception:
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _temporal_split(
|
||||
rows: list[dict[str, Any]],
|
||||
split_policy: str = "temporal_70_15_15",
|
||||
) -> dict[str, list[dict[str, Any]]]:
|
||||
"""Split rows into train/valid/test by event_date (temporal order).
|
||||
|
||||
Args:
|
||||
rows: List of row dicts that must have an "event_date" field.
|
||||
split_policy: E.g. "temporal_70_15_15" → 70% train, 15% valid, 15% test.
|
||||
|
||||
Returns:
|
||||
Dict with keys "train", "valid", "test".
|
||||
"""
|
||||
if not rows:
|
||||
return {"train": [], "valid": [], "test": []}
|
||||
|
||||
parts = split_policy.replace("temporal_", "").split("_")
|
||||
if len(parts) != 3:
|
||||
raise ValueError(f"Invalid split_policy: {split_policy}")
|
||||
train_pct, valid_pct, _ = (int(p) for p in parts)
|
||||
|
||||
sorted_rows = sorted(rows, key=lambda r: r.get("event_date", ""))
|
||||
n = len(sorted_rows)
|
||||
n_train = int(n * train_pct / 100)
|
||||
n_valid = int(n * valid_pct / 100)
|
||||
|
||||
return {
|
||||
"train": sorted_rows[:n_train],
|
||||
"valid": sorted_rows[n_train : n_train + n_valid],
|
||||
"test": sorted_rows[n_train + n_valid :],
|
||||
}
|
||||
|
||||
|
||||
def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table:
|
||||
"""Convert list of dicts to a PyArrow Table."""
|
||||
if not rows:
|
||||
return pa.table({})
|
||||
# Collect all keys
|
||||
keys = list(rows[0].keys())
|
||||
arrays: dict[str, list[Any]] = {k: [] for k in keys}
|
||||
for row in rows:
|
||||
for k in keys:
|
||||
arrays[k].append(row.get(k))
|
||||
return pa.table({k: pa.array(v) for k, v in arrays.items()})
|
||||
|
||||
|
||||
async def export_dataset_snapshot(
|
||||
session: AsyncSession,
|
||||
snapshot_id: str | None,
|
||||
split_policy: str,
|
||||
output_dir: str | Path,
|
||||
feature_version: str = "market_v1",
|
||||
label_version: str = "label-1.0.0",
|
||||
parser_version: str = "rule-1.0.0",
|
||||
) -> dict[str, Any]:
|
||||
"""Join FeatureSnapshot + EventLabel and export to Parquet.
|
||||
|
||||
Args:
|
||||
session: Async DB session.
|
||||
snapshot_id: Unique ID for this snapshot (generated if None).
|
||||
split_policy: Temporal split policy string (e.g. "temporal_70_15_15").
|
||||
output_dir: Root directory for output files.
|
||||
feature_version: Snapshot name filter for FeatureSnapshot.
|
||||
label_version: Label version filter for EventLabel.
|
||||
parser_version: Parser version filter for Event.
|
||||
|
||||
Returns:
|
||||
Manifest dict with metadata and row counts.
|
||||
"""
|
||||
from libs.db.models import EventLabel, FeatureSnapshot
|
||||
|
||||
if snapshot_id is None:
|
||||
snapshot_id = str(uuid.uuid4())
|
||||
|
||||
out_path = Path(output_dir) / snapshot_id
|
||||
out_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Query: JOIN feature_snapshots + event_labels via event_id
|
||||
stmt = (
|
||||
select(FeatureSnapshot, EventLabel)
|
||||
.join(EventLabel, FeatureSnapshot.event_id == EventLabel.event_id)
|
||||
.where(FeatureSnapshot.snapshot_name == feature_version)
|
||||
.where(EventLabel.label_version == label_version)
|
||||
.where(EventLabel.label_status == "ok")
|
||||
.where(EventLabel.invalid_event_for_labeling.is_(False))
|
||||
)
|
||||
|
||||
result = await session.execute(stmt)
|
||||
pairs = result.all()
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
for fs, lbl in pairs:
|
||||
row: dict[str, Any] = {
|
||||
"event_id": fs.event_id,
|
||||
"snapshot_name": fs.snapshot_name,
|
||||
"snapshot_version": fs.snapshot_version,
|
||||
**fs.feature_json,
|
||||
"entry_convention": lbl.entry_convention,
|
||||
"reaction_date": lbl.reaction_date.isoformat() if lbl.reaction_date else None,
|
||||
"entry_date": lbl.entry_date.isoformat() if lbl.entry_date else None,
|
||||
"entry_price": float(lbl.entry_price) if lbl.entry_price else None,
|
||||
"fwd_return_1d": float(lbl.fwd_return_1d) if lbl.fwd_return_1d else None,
|
||||
"fwd_return_3d": float(lbl.fwd_return_3d) if lbl.fwd_return_3d else None,
|
||||
"fwd_return_5d": float(lbl.fwd_return_5d) if lbl.fwd_return_5d else None,
|
||||
"hit_pos_1r_within_3d": lbl.hit_pos_1r_within_3d,
|
||||
"hit_neg_1r_within_3d": lbl.hit_neg_1r_within_3d,
|
||||
"close_up_after_3d": lbl.close_up_after_3d,
|
||||
"close_up_after_5d": lbl.close_up_after_5d,
|
||||
"mfe_3d": float(lbl.mfe_3d) if lbl.mfe_3d else None,
|
||||
"mae_3d": float(lbl.mae_3d) if lbl.mae_3d else None,
|
||||
"mfe_5d": float(lbl.mfe_5d) if lbl.mfe_5d else None,
|
||||
"mae_5d": float(lbl.mae_5d) if lbl.mae_5d else None,
|
||||
"label_status": lbl.label_status,
|
||||
"label_version": lbl.label_version,
|
||||
}
|
||||
if "event_date" not in row:
|
||||
row["event_date"] = str(lbl.reaction_date) if lbl.reaction_date else ""
|
||||
rows.append(row)
|
||||
|
||||
logger.info("snapshot_export_rows", snapshot_id=snapshot_id, total=len(rows))
|
||||
|
||||
splits = _temporal_split(rows, split_policy)
|
||||
row_counts: dict[str, int] = {}
|
||||
|
||||
for split_name, split_rows in splits.items():
|
||||
parquet_path = out_path / f"{split_name}.parquet"
|
||||
table = _rows_to_table(split_rows)
|
||||
pq.write_table(table, str(parquet_path))
|
||||
row_counts[split_name] = len(split_rows)
|
||||
logger.info("split_written", split=split_name, rows=len(split_rows), path=str(parquet_path))
|
||||
|
||||
manifest: dict[str, Any] = {
|
||||
"snapshot_id": snapshot_id,
|
||||
"created_at_utc": utc_now().isoformat(),
|
||||
"code_commit_hash": _get_git_commit_hash(),
|
||||
"feature_version": feature_version,
|
||||
"parser_version": parser_version,
|
||||
"label_version": label_version,
|
||||
"split_policy": split_policy,
|
||||
"row_counts": row_counts,
|
||||
"total_rows": len(rows),
|
||||
"output_dir": str(out_path),
|
||||
}
|
||||
|
||||
manifest_path = out_path / MANIFEST_FILENAME
|
||||
manifest_path.write_text(json.dumps(manifest, indent=2))
|
||||
logger.info("manifest_written", path=str(manifest_path), snapshot_id=snapshot_id)
|
||||
|
||||
return manifest
|
||||
@ -0,0 +1 @@
|
||||
"""Event labeler: compute forward-return labels from price data."""
|
||||
@ -0,0 +1,239 @@
|
||||
"""Generate forward-return labels for parsed events using price bars."""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
from decimal import Decimal
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.common.time_utils import next_trading_day, trading_days_between
|
||||
from libs.labeler.reaction_date import compute_reaction_date
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
LABEL_VERSION = "label-1.0.0"
|
||||
|
||||
_LOOK_AHEAD_DAYS = 7 # fetch this many trading days of bars for label computation
|
||||
_R_FACTOR = 0.01 # 1R = 1% move (used for hit_pos/neg_1r labels)
|
||||
|
||||
|
||||
def _safe_decimal(v: Any) -> Decimal | None:
|
||||
try:
|
||||
return Decimal(str(v)) if v is not None else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _pct_return(entry: Decimal, exit_price: Decimal) -> Decimal:
|
||||
if entry == 0:
|
||||
return Decimal("0")
|
||||
return (exit_price - entry) / entry
|
||||
|
||||
|
||||
def _compute_labels_from_bars(
|
||||
entry_price: Decimal,
|
||||
bars: list[dict[str, Any]], # sorted ascending by date
|
||||
n_days: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Compute forward-return labels over `n_days` trading days of bars.
|
||||
|
||||
Args:
|
||||
entry_price: Entry price (open of entry date).
|
||||
bars: List of OHLCV dicts with keys: date, open, high, low, close.
|
||||
n_days: Look-ahead horizon (3 or 5).
|
||||
|
||||
Returns:
|
||||
Dict of computed label fields for the given horizon.
|
||||
"""
|
||||
window = bars[:n_days]
|
||||
if not window:
|
||||
return {}
|
||||
|
||||
closes = [_safe_decimal(b.get("close")) for b in window]
|
||||
highs = [_safe_decimal(b.get("high")) for b in window]
|
||||
lows = [_safe_decimal(b.get("low")) for b in window]
|
||||
|
||||
# Forward close return at day n
|
||||
last_close = closes[-1]
|
||||
fwd_return = _pct_return(entry_price, last_close) if last_close else None
|
||||
|
||||
# MFE (max favorable excursion): max high vs entry price
|
||||
valid_highs = [h for h in highs if h is not None]
|
||||
mfe = (max(valid_highs) - entry_price) / entry_price if valid_highs else None
|
||||
|
||||
# MAE (max adverse excursion): min low vs entry price
|
||||
valid_lows = [lo for lo in lows if lo is not None]
|
||||
mae = (min(valid_lows) - entry_price) / entry_price if valid_lows else None
|
||||
|
||||
# Hit +1R within n days
|
||||
threshold_pos = entry_price * (1 + Decimal(str(_R_FACTOR)))
|
||||
hit_pos = any(h is not None and h >= threshold_pos for h in highs)
|
||||
|
||||
# Hit -1R within n days
|
||||
threshold_neg = entry_price * (1 - Decimal(str(_R_FACTOR)))
|
||||
hit_neg = any(lo is not None and lo <= threshold_neg for lo in lows)
|
||||
|
||||
# Close up after n days
|
||||
close_up = bool(last_close is not None and last_close > entry_price)
|
||||
|
||||
# Bars to MFE (index of max high)
|
||||
bars_to_mfe: int | None = None
|
||||
if valid_highs and mfe is not None:
|
||||
max_high = max(valid_highs)
|
||||
for i, h in enumerate(highs):
|
||||
if h == max_high:
|
||||
bars_to_mfe = i + 1
|
||||
break
|
||||
|
||||
# Days to peak close
|
||||
valid_close_idx = [(i, c) for i, c in enumerate(closes) if c is not None]
|
||||
days_to_peak_close: int | None = None
|
||||
if valid_close_idx:
|
||||
peak_close_idx = max(valid_close_idx, key=lambda x: x[1])[0]
|
||||
days_to_peak_close = peak_close_idx + 1
|
||||
|
||||
return {
|
||||
"fwd_return": fwd_return,
|
||||
"mfe": mfe,
|
||||
"mae": mae,
|
||||
"hit_pos_1r": hit_pos,
|
||||
"hit_neg_1r": hit_neg,
|
||||
"close_up": close_up,
|
||||
"bars_to_mfe": bars_to_mfe,
|
||||
"days_to_peak_close": days_to_peak_close,
|
||||
}
|
||||
|
||||
|
||||
async def generate_labels(
|
||||
session: AsyncSession,
|
||||
event: Any, # Event ORM model instance
|
||||
price_svc: Any, # PriceService
|
||||
ticker: str,
|
||||
entry_convention: str = "next_open_after_reaction_close",
|
||||
) -> Any:
|
||||
"""Generate EventLabel for a single event.
|
||||
|
||||
Args:
|
||||
session: Async DB session.
|
||||
event: Event ORM instance (needs event_date, filed_at_utc, symbol_id).
|
||||
price_svc: PriceService instance for fetching bars.
|
||||
ticker: Trading ticker symbol.
|
||||
entry_convention: How to determine entry price.
|
||||
|
||||
Returns:
|
||||
EventLabel ORM instance (not yet added to session).
|
||||
"""
|
||||
from libs.db.models import EventLabel
|
||||
|
||||
# 1. Compute reaction_date
|
||||
filing_time_bucket = getattr(event, "filing_time_bucket", "unknown")
|
||||
if not filing_time_bucket:
|
||||
filing_time_bucket = "unknown"
|
||||
|
||||
event_date: dt.date = event.event_date
|
||||
reaction_date = compute_reaction_date(event_date, filing_time_bucket)
|
||||
|
||||
# 2. Compute entry_date = next trading day after reaction_date
|
||||
entry_date = next_trading_day(reaction_date)
|
||||
|
||||
# 3. Fetch price bars (entry_date + _LOOK_AHEAD_DAYS trading days)
|
||||
fetch_start = entry_date
|
||||
trading_days = trading_days_between(entry_date, entry_date + dt.timedelta(days=20))
|
||||
fetch_end = trading_days[_LOOK_AHEAD_DAYS] if len(trading_days) > _LOOK_AHEAD_DAYS else trading_days[-1]
|
||||
|
||||
try:
|
||||
price_resp = await price_svc.get_daily_bars(
|
||||
ticker=ticker,
|
||||
start=fetch_start.isoformat(),
|
||||
end=fetch_end.isoformat(),
|
||||
)
|
||||
raw_bars = [b.model_dump() for b in price_resp.bars]
|
||||
except Exception as exc:
|
||||
logger.warning("label_price_unavailable", ticker=ticker, error=str(exc))
|
||||
return EventLabel(
|
||||
event_id=event.event_id,
|
||||
entry_convention=entry_convention,
|
||||
reaction_date=reaction_date,
|
||||
entry_date=None,
|
||||
label_status="unavailable",
|
||||
invalid_event_for_labeling=False,
|
||||
label_version=LABEL_VERSION,
|
||||
)
|
||||
|
||||
# Filter bars from entry_date onwards, sorted ascending
|
||||
bars = sorted(
|
||||
[b for b in raw_bars if b.get("date") and b["date"] >= entry_date.isoformat()],
|
||||
key=lambda b: b["date"],
|
||||
)
|
||||
|
||||
if not bars:
|
||||
return EventLabel(
|
||||
event_id=event.event_id,
|
||||
entry_convention=entry_convention,
|
||||
reaction_date=reaction_date,
|
||||
entry_date=entry_date,
|
||||
label_status="unavailable",
|
||||
invalid_event_for_labeling=False,
|
||||
label_version=LABEL_VERSION,
|
||||
)
|
||||
|
||||
# 4. Determine entry_price
|
||||
first_bar = bars[0]
|
||||
if entry_convention == "next_open_after_reaction_close":
|
||||
entry_price = _safe_decimal(first_bar.get("open"))
|
||||
else:
|
||||
entry_price = _safe_decimal(first_bar.get("close"))
|
||||
|
||||
if entry_price is None or entry_price == 0:
|
||||
return EventLabel(
|
||||
event_id=event.event_id,
|
||||
entry_convention=entry_convention,
|
||||
reaction_date=reaction_date,
|
||||
entry_date=entry_date,
|
||||
entry_price=None,
|
||||
label_status="unavailable",
|
||||
invalid_event_for_labeling=True,
|
||||
label_version=LABEL_VERSION,
|
||||
)
|
||||
|
||||
# 5. Forward bars (exclude entry bar itself for 1D/3D/5D)
|
||||
forward_bars = bars[1:] # Day 1+ after entry
|
||||
|
||||
label_status = "ok"
|
||||
if len(forward_bars) < 5:
|
||||
label_status = "truncated"
|
||||
|
||||
# 1D return
|
||||
fwd_1d = _pct_return(entry_price, _safe_decimal(bars[1]["close"])) if len(bars) > 1 else None
|
||||
|
||||
# 3D labels
|
||||
lbl_3d = _compute_labels_from_bars(entry_price, forward_bars, 3)
|
||||
# 5D labels
|
||||
lbl_5d = _compute_labels_from_bars(entry_price, forward_bars, 5)
|
||||
|
||||
return EventLabel(
|
||||
event_id=event.event_id,
|
||||
entry_convention=entry_convention,
|
||||
reaction_date=reaction_date,
|
||||
entry_date=entry_date,
|
||||
entry_price=entry_price,
|
||||
fwd_return_1d=fwd_1d,
|
||||
fwd_return_3d=lbl_3d.get("fwd_return"),
|
||||
fwd_return_5d=lbl_5d.get("fwd_return"),
|
||||
hit_pos_1r_within_3d=lbl_3d.get("hit_pos_1r"),
|
||||
hit_neg_1r_within_3d=lbl_3d.get("hit_neg_1r"),
|
||||
close_up_after_3d=lbl_3d.get("close_up"),
|
||||
close_up_after_5d=lbl_5d.get("close_up"),
|
||||
mfe_3d=lbl_3d.get("mfe"),
|
||||
mae_3d=lbl_3d.get("mae"),
|
||||
mfe_5d=lbl_5d.get("mfe"),
|
||||
mae_5d=lbl_5d.get("mae"),
|
||||
bars_to_mfe_3d=lbl_3d.get("bars_to_mfe"),
|
||||
bars_to_mae_3d=None, # symmetrically bars to MAE (min low) - optional
|
||||
days_to_peak_close_5d=lbl_5d.get("days_to_peak_close"),
|
||||
label_status=label_status,
|
||||
invalid_event_for_labeling=False,
|
||||
label_version=LABEL_VERSION,
|
||||
)
|
||||
@ -0,0 +1,49 @@
|
||||
"""Compute the reaction date for a filing based on its time bucket."""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
|
||||
from libs.common.time_utils import is_trading_day
|
||||
|
||||
|
||||
def _advance_to_next_trading_day(date: dt.date) -> dt.date:
|
||||
"""Return the next calendar day that is a trading day (starting from date+1)."""
|
||||
check = date + dt.timedelta(days=1)
|
||||
for _ in range(30): # guard against infinite loop
|
||||
if is_trading_day(check):
|
||||
return check
|
||||
check += dt.timedelta(days=1)
|
||||
raise RuntimeError(f"Could not find trading day within 30 days of {date}")
|
||||
|
||||
|
||||
def _to_trading_day_on_or_after(date: dt.date) -> dt.date:
|
||||
"""Return date itself if a trading day, else the next trading day."""
|
||||
for _ in range(30):
|
||||
if is_trading_day(date):
|
||||
return date
|
||||
date += dt.timedelta(days=1)
|
||||
raise RuntimeError("Could not find trading day within 30 days")
|
||||
|
||||
|
||||
def compute_reaction_date(
|
||||
event_date: dt.date,
|
||||
filing_time_bucket: str,
|
||||
) -> dt.date:
|
||||
"""Return the first trading day on which the market can react to the filing.
|
||||
|
||||
Rules:
|
||||
- pre_market / regular_hours → same day if it is a trading day, else next.
|
||||
- post_market / unknown → next trading day after event_date.
|
||||
|
||||
Args:
|
||||
event_date: The calendar date of the filing.
|
||||
filing_time_bucket: One of pre_market, regular_hours, post_market, unknown.
|
||||
|
||||
Returns:
|
||||
The reaction date (a trading day).
|
||||
"""
|
||||
if filing_time_bucket in ("pre_market", "regular_hours"):
|
||||
return _to_trading_day_on_or_after(event_date)
|
||||
else:
|
||||
# post_market or unknown: market reacts next trading day
|
||||
return _advance_to_next_trading_day(event_date)
|
||||
@ -0,0 +1 @@
|
||||
"""LLM integration module (Ollama)."""
|
||||
@ -0,0 +1,104 @@
|
||||
"""DB-backed LLM call cache keyed by SHA-256 of inputs."""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.common.time_utils import utc_now
|
||||
from libs.llm.exceptions import LLMCacheError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_SCHEMA_VERSION = "1.0.0"
|
||||
|
||||
|
||||
def build_cache_key(
|
||||
text: str,
|
||||
prompt_name: str,
|
||||
prompt_version: str,
|
||||
model_name: str,
|
||||
schema_version: str = _SCHEMA_VERSION,
|
||||
) -> str:
|
||||
"""Return SHA-256 cache key for this (text, prompt, model, schema) combination."""
|
||||
text_hash = hashlib.sha256(text.encode()).hexdigest()
|
||||
hint_hash = hashlib.sha256(
|
||||
f"{prompt_name}:{prompt_version}:{model_name}:{schema_version}".encode()
|
||||
).hexdigest()
|
||||
combined = f"{text_hash}:{hint_hash}"
|
||||
return hashlib.sha256(combined.encode()).hexdigest()
|
||||
|
||||
|
||||
class LLMCacheStore:
|
||||
"""Read/write LLM responses from/to the llm_call_cache table."""
|
||||
|
||||
async def get(self, session: AsyncSession, cache_key: str) -> dict[str, Any] | None:
|
||||
"""Return cached normalized dict if found, else None."""
|
||||
from libs.db.models import LLMCallCache # late import to avoid circular deps
|
||||
|
||||
try:
|
||||
result = await session.execute(
|
||||
select(LLMCallCache).where(LLMCallCache.cache_key == cache_key)
|
||||
)
|
||||
row = result.scalar_one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
logger.debug("llm_cache_hit", cache_key=cache_key[:16])
|
||||
return row.normalized_json # type: ignore[return-value]
|
||||
except Exception as exc:
|
||||
raise LLMCacheError(
|
||||
f"Cache read failed: {exc}",
|
||||
source="llm_cache",
|
||||
context={"cache_key": cache_key[:16]},
|
||||
) from exc
|
||||
|
||||
async def put(
|
||||
self,
|
||||
session: AsyncSession,
|
||||
cache_key: str,
|
||||
document_id: str | None,
|
||||
model_name: str,
|
||||
prompt_version: str,
|
||||
schema_version: str,
|
||||
raw_prompt: str,
|
||||
raw_response: str,
|
||||
normalized: dict[str, Any],
|
||||
token_usage: dict[str, int],
|
||||
elapsed_ms: int,
|
||||
) -> None:
|
||||
"""Store a new cache entry. Silently skips on duplicate key."""
|
||||
from libs.db.models import LLMCallCache # late import
|
||||
|
||||
try:
|
||||
existing = await session.execute(
|
||||
select(LLMCallCache).where(LLMCallCache.cache_key == cache_key)
|
||||
)
|
||||
if existing.scalar_one_or_none() is not None:
|
||||
logger.debug("llm_cache_skip_dup", cache_key=cache_key[:16])
|
||||
return
|
||||
|
||||
row = LLMCallCache(
|
||||
cache_key=cache_key,
|
||||
document_id=document_id,
|
||||
model_name=model_name,
|
||||
prompt_version=prompt_version,
|
||||
schema_version=schema_version,
|
||||
raw_prompt=raw_prompt,
|
||||
raw_response=raw_response,
|
||||
normalized_json=normalized,
|
||||
token_usage_json=token_usage,
|
||||
elapsed_ms=elapsed_ms,
|
||||
created_at_utc=utc_now(),
|
||||
)
|
||||
session.add(row)
|
||||
await session.flush()
|
||||
logger.debug("llm_cache_stored", cache_key=cache_key[:16], elapsed_ms=elapsed_ms)
|
||||
except Exception as exc:
|
||||
raise LLMCacheError(
|
||||
f"Cache write failed: {exc}",
|
||||
source="llm_cache",
|
||||
context={"cache_key": cache_key[:16]},
|
||||
) from exc
|
||||
@ -0,0 +1,151 @@
|
||||
"""Async Ollama HTTP client with retry logic."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.common.retries import RetryableError, with_retry
|
||||
from libs.llm.exceptions import LLMError, LLMTimeoutError
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_CHAT_PATH = "/api/chat"
|
||||
_TAGS_PATH = "/api/tags"
|
||||
|
||||
|
||||
class OllamaClient:
|
||||
"""Async client for Ollama REST API."""
|
||||
|
||||
def __init__(self, base_url: str, model: str, timeout: float = 60.0) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self.model = model
|
||||
self._timeout = timeout
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
async def __aenter__(self) -> OllamaClient:
|
||||
self._client = httpx.AsyncClient(base_url=self._base_url, timeout=self._timeout)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_: object) -> None:
|
||||
if self._client is not None:
|
||||
await self._client.aclose()
|
||||
self._client = None
|
||||
|
||||
@with_retry(max_attempts=2, min_wait=0.5, max_wait=10.0, multiplier=2.0)
|
||||
async def chat(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: str = "json",
|
||||
temperature: float = 0.0,
|
||||
) -> tuple[dict[str, Any], dict[str, int], int]:
|
||||
"""Call Ollama /api/chat and return (parsed_json, token_usage, elapsed_ms).
|
||||
|
||||
Args:
|
||||
messages: List of {role, content} message dicts.
|
||||
response_format: "json" forces JSON output mode.
|
||||
temperature: Sampling temperature (0.0 = deterministic).
|
||||
|
||||
Returns:
|
||||
Tuple of (parsed response dict, token usage dict, elapsed_ms).
|
||||
|
||||
Raises:
|
||||
LLMTimeoutError: On request timeout.
|
||||
LLMError: On non-retryable HTTP errors.
|
||||
RetryableError: On 5xx server errors.
|
||||
"""
|
||||
if self._client is None:
|
||||
raise LLMError("OllamaClient must be used as an async context manager")
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"options": {"temperature": temperature},
|
||||
}
|
||||
if response_format == "json":
|
||||
payload["format"] = "json"
|
||||
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
response = await self._client.post(_CHAT_PATH, json=payload)
|
||||
except httpx.TimeoutException as exc:
|
||||
raise LLMTimeoutError(
|
||||
f"Ollama request timed out after {self._timeout}s",
|
||||
source="ollama",
|
||||
context={"model": self.model},
|
||||
) from exc
|
||||
except httpx.ConnectError as exc:
|
||||
raise RetryableError(
|
||||
f"Cannot connect to Ollama at {self._base_url}",
|
||||
source="ollama",
|
||||
context={"model": self.model},
|
||||
) from exc
|
||||
|
||||
elapsed_ms = int((time.monotonic() - t0) * 1000)
|
||||
|
||||
if response.status_code >= 500:
|
||||
raise RetryableError(
|
||||
f"Ollama server error {response.status_code}",
|
||||
source="ollama",
|
||||
context={"status": response.status_code, "body": response.text[:200]},
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise LLMError(
|
||||
f"Ollama client error {response.status_code}: {response.text[:200]}",
|
||||
source="ollama",
|
||||
context={"status": response.status_code},
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
raw_content: str = data.get("message", {}).get("content", "")
|
||||
token_usage = {
|
||||
"prompt_tokens": data.get("prompt_eval_count", 0),
|
||||
"completion_tokens": data.get("eval_count", 0),
|
||||
}
|
||||
|
||||
try:
|
||||
parsed = json.loads(raw_content)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise LLMError(
|
||||
"Ollama returned non-JSON content",
|
||||
source="ollama",
|
||||
context={"raw": raw_content[:300]},
|
||||
) from exc
|
||||
|
||||
logger.debug(
|
||||
"ollama_chat_ok",
|
||||
model=self.model,
|
||||
elapsed_ms=elapsed_ms,
|
||||
**token_usage,
|
||||
)
|
||||
return parsed, token_usage, elapsed_ms
|
||||
|
||||
async def health_check(self) -> bool:
|
||||
"""Return True if Ollama is reachable and the model is available."""
|
||||
if self._client is None:
|
||||
raise LLMError("OllamaClient must be used as an async context manager")
|
||||
try:
|
||||
response = await self._client.get(_TAGS_PATH)
|
||||
if response.status_code != 200:
|
||||
return False
|
||||
data = response.json()
|
||||
models = [m.get("name", "") for m in data.get("models", [])]
|
||||
return any(self.model in name for name in models)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def make_ollama_client() -> OllamaClient:
|
||||
"""Factory that reads config from settings."""
|
||||
from libs.common.config import get_settings
|
||||
|
||||
s = get_settings()
|
||||
return OllamaClient(
|
||||
base_url=s.ollama_url,
|
||||
model=s.ollama_model,
|
||||
timeout=float(s.ollama_timeout),
|
||||
)
|
||||
@ -0,0 +1,20 @@
|
||||
"""LLM-specific exception hierarchy."""
|
||||
from __future__ import annotations
|
||||
|
||||
from libs.common.retries import ACEFError
|
||||
|
||||
|
||||
class LLMError(ACEFError):
|
||||
"""Base error for LLM operations."""
|
||||
|
||||
|
||||
class LLMTimeoutError(LLMError):
|
||||
"""LLM request timed out."""
|
||||
|
||||
|
||||
class LLMSchemaError(LLMError):
|
||||
"""LLM response did not match expected schema."""
|
||||
|
||||
|
||||
class LLMCacheError(LLMError):
|
||||
"""Cache read/write failure."""
|
||||
@ -0,0 +1,211 @@
|
||||
"""LLM parser using Ollama: cache lookup → prompt → validate → store."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.llm.cache import LLMCacheStore, build_cache_key
|
||||
from libs.llm.client import OllamaClient
|
||||
from libs.llm.exceptions import LLMError, LLMSchemaError
|
||||
from libs.llm.prompts import PROMPT_VERSION, render_prompt
|
||||
from libs.parser.schema_validator import validate_parser_output
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_SCHEMA_VERSION = "1.0.0"
|
||||
_PROMPT_NAME = "event_classifier_v1"
|
||||
|
||||
# Fields returned by LLM that map to ParserEventOutput fields
|
||||
_REQUIRED_FIELDS = {
|
||||
"event_type",
|
||||
"event_direction",
|
||||
"headline",
|
||||
"summary",
|
||||
"guidance_status",
|
||||
"confidence_overall",
|
||||
}
|
||||
|
||||
|
||||
def _llm_response_to_parser_output(
|
||||
document_id: str,
|
||||
doc_meta: dict[str, Any],
|
||||
raw: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Convert flat LLM response dict → ParserEventOutput-compatible dict."""
|
||||
return {
|
||||
"schema_version": _SCHEMA_VERSION,
|
||||
"document_id": document_id,
|
||||
"parser_kind": "llm",
|
||||
"event_type": raw.get("event_type", "unknown"),
|
||||
"event_direction": raw.get("event_direction", "unknown"),
|
||||
"event_date": doc_meta.get("filing_date", "1900-01-01"),
|
||||
"filing_time_bucket": doc_meta.get("filing_time_bucket", "unknown"),
|
||||
"headline": raw.get("headline", ""),
|
||||
"summary": raw.get("summary", ""),
|
||||
"guidance": {
|
||||
"status": raw.get("guidance_status", "unclear"),
|
||||
"scope": raw.get("guidance_scope", "unknown"),
|
||||
"notes": "",
|
||||
},
|
||||
"signals": {
|
||||
"demand_strength": raw.get("demand_strength", "unknown"),
|
||||
"pricing_power": raw.get("pricing_power", "unknown"),
|
||||
"backlog_or_bookings": raw.get("backlog_or_bookings", "unknown"),
|
||||
"customer_expansion": raw.get("customer_expansion", "unknown"),
|
||||
"margin_quality": raw.get("margin_quality", "unknown"),
|
||||
},
|
||||
"risk_flags": {
|
||||
"oneoff_item": bool(raw.get("oneoff_item", False)),
|
||||
"tax_benefit": bool(raw.get("tax_benefit", False)),
|
||||
"valuation_gain": bool(raw.get("valuation_gain", False)),
|
||||
"non_gaap_heavy": bool(raw.get("non_gaap_heavy", False)),
|
||||
"financing_related": bool(raw.get("financing_related", False)),
|
||||
"legal_or_regulatory_overhang": bool(
|
||||
raw.get("legal_or_regulatory_overhang", False)
|
||||
),
|
||||
},
|
||||
"evidence": [],
|
||||
"confidence": {
|
||||
"overall": float(raw.get("confidence_overall", 0.5)),
|
||||
"event_type": float(raw.get("confidence_event_type", 0.5)),
|
||||
"event_direction": float(raw.get("confidence_event_direction", 0.5)),
|
||||
"guidance": 0.5,
|
||||
"risk_flags": 0.5,
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
|
||||
def _repair_prompt(
|
||||
original_messages: list[dict[str, str]],
|
||||
errors: list[str],
|
||||
raw_response: str,
|
||||
) -> list[dict[str, str]]:
|
||||
"""Append a repair instruction to messages asking LLM to fix schema errors."""
|
||||
repair_msg = (
|
||||
f"Your previous response had schema errors: {errors[:3]}. "
|
||||
f"Previous response was: {raw_response[:300]}. "
|
||||
"Please fix these issues and return only valid JSON matching the schema."
|
||||
)
|
||||
return [*original_messages, {"role": "assistant", "content": raw_response}, {"role": "user", "content": repair_msg}]
|
||||
|
||||
|
||||
class LLMParser:
|
||||
"""Parse a document with Ollama, using DB cache to avoid redundant calls."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: OllamaClient,
|
||||
cache_store: LLMCacheStore | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._cache = cache_store or LLMCacheStore()
|
||||
|
||||
async def parse(
|
||||
self,
|
||||
document_id: str,
|
||||
doc_text: str,
|
||||
doc_meta: dict[str, Any],
|
||||
rule_hints: dict[str, Any],
|
||||
session: AsyncSession,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Parse a document with LLM.
|
||||
|
||||
Returns a ParserEventOutput-compatible dict or None on failure.
|
||||
"""
|
||||
cache_key = build_cache_key(
|
||||
text=doc_text,
|
||||
prompt_name=_PROMPT_NAME,
|
||||
prompt_version=PROMPT_VERSION,
|
||||
model_name=self._client.model,
|
||||
schema_version=_SCHEMA_VERSION,
|
||||
)
|
||||
|
||||
# 1. Cache lookup
|
||||
cached = await self._cache.get(session, cache_key)
|
||||
if cached is not None:
|
||||
logger.info("llm_parse_cache_hit", document_id=document_id)
|
||||
return cached
|
||||
|
||||
# 2. Render prompt
|
||||
messages = render_prompt(
|
||||
prompt_name=_PROMPT_NAME,
|
||||
doc_meta=doc_meta,
|
||||
text=doc_text,
|
||||
rule_hints=rule_hints,
|
||||
)
|
||||
raw_prompt_str = json.dumps(messages)
|
||||
|
||||
# 3. Call Ollama
|
||||
try:
|
||||
raw_dict, token_usage, elapsed_ms = await self._client.chat(messages)
|
||||
except LLMError as exc:
|
||||
logger.error("llm_parse_failed", document_id=document_id, error=str(exc))
|
||||
return None
|
||||
|
||||
raw_response_str = json.dumps(raw_dict)
|
||||
|
||||
# 4. Map to parser output format
|
||||
output = _llm_response_to_parser_output(document_id, doc_meta, raw_dict)
|
||||
|
||||
# 5. Validate schema
|
||||
errors = validate_parser_output(output)
|
||||
if errors:
|
||||
logger.warning("llm_schema_errors_attempt_repair", errors=errors[:3])
|
||||
repair_messages = _repair_prompt(messages, errors, raw_response_str)
|
||||
try:
|
||||
raw_dict2, token_usage2, elapsed_ms2 = await self._client.chat(repair_messages)
|
||||
token_usage = {
|
||||
"prompt_tokens": token_usage.get("prompt_tokens", 0)
|
||||
+ token_usage2.get("prompt_tokens", 0),
|
||||
"completion_tokens": token_usage.get("completion_tokens", 0)
|
||||
+ token_usage2.get("completion_tokens", 0),
|
||||
}
|
||||
elapsed_ms += elapsed_ms2
|
||||
raw_response_str = json.dumps(raw_dict2)
|
||||
output = _llm_response_to_parser_output(document_id, doc_meta, raw_dict2)
|
||||
errors = validate_parser_output(output)
|
||||
except LLMError as exc:
|
||||
logger.error("llm_repair_failed", document_id=document_id, error=str(exc))
|
||||
return None
|
||||
|
||||
if errors:
|
||||
logger.error(
|
||||
"llm_schema_invalid_after_repair",
|
||||
document_id=document_id,
|
||||
errors=errors[:3],
|
||||
)
|
||||
raise LLMSchemaError(
|
||||
f"LLM output invalid after repair: {errors[:2]}",
|
||||
source="llm_parser",
|
||||
context={"document_id": document_id},
|
||||
)
|
||||
|
||||
# 6. Store in cache
|
||||
try:
|
||||
await self._cache.put(
|
||||
session=session,
|
||||
cache_key=cache_key,
|
||||
document_id=document_id,
|
||||
model_name=self._client.model,
|
||||
prompt_version=PROMPT_VERSION,
|
||||
schema_version=_SCHEMA_VERSION,
|
||||
raw_prompt=raw_prompt_str,
|
||||
raw_response=raw_response_str,
|
||||
normalized=output,
|
||||
token_usage=token_usage,
|
||||
elapsed_ms=elapsed_ms,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("llm_cache_write_failed", error=str(exc))
|
||||
|
||||
logger.info(
|
||||
"llm_parse_ok",
|
||||
document_id=document_id,
|
||||
elapsed_ms=elapsed_ms,
|
||||
event_type=output.get("event_type"),
|
||||
)
|
||||
return output
|
||||
@ -0,0 +1,97 @@
|
||||
"""Prompt registry for LLM parsing."""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
PROMPT_VERSION = "v1"
|
||||
|
||||
_SYSTEM_EVENT_CLASSIFIER = """\
|
||||
You are a financial document analyst specializing in SEC filings (8-K, 6-K).
|
||||
Your task is to classify events and extract structured information from earnings press releases
|
||||
and other material corporate disclosures.
|
||||
|
||||
Output ONLY valid JSON matching the specified schema. Do not add commentary outside the JSON.
|
||||
""".strip()
|
||||
|
||||
_USER_EVENT_CLASSIFIER = """\
|
||||
Analyze the following corporate filing excerpt and return a JSON object with these fields:
|
||||
|
||||
- event_type: one of ["earnings_release", "guidance_update", "material_contract",
|
||||
"regulatory_or_approval", "capital_markets_or_financing", "management_change",
|
||||
"litigation_or_investigation", "other_material_event", "unknown"]
|
||||
- event_direction: one of ["bullish", "bearish", "mixed", "neutral", "unknown"]
|
||||
- headline: short 1-sentence headline (max 120 chars)
|
||||
- summary: 2-3 sentence summary of the key facts
|
||||
- guidance_status: one of ["raised", "inline_or_maintained", "lowered", "withdrawn",
|
||||
"not_provided", "unclear"]
|
||||
- guidance_scope: one of ["quarterly", "annual", "both", "unknown"]
|
||||
- demand_strength: one of ["strong", "stable", "weakening", "unknown"]
|
||||
- pricing_power: one of ["present", "mixed", "absent", "unknown"]
|
||||
- backlog_or_bookings: one of ["present", "mixed", "absent", "unknown"]
|
||||
- customer_expansion: one of ["present", "mixed", "absent", "unknown"]
|
||||
- margin_quality: one of ["improving", "stable", "deteriorating", "unknown"]
|
||||
- oneoff_item: true/false — unusual one-time item inflating results
|
||||
- tax_benefit: true/false
|
||||
- valuation_gain: true/false
|
||||
- non_gaap_heavy: true/false — results rely heavily on non-GAAP metrics
|
||||
- financing_related: true/false
|
||||
- legal_or_regulatory_overhang: true/false
|
||||
- confidence_overall: float 0.0-1.0
|
||||
- confidence_event_type: float 0.0-1.0
|
||||
- confidence_event_direction: float 0.0-1.0
|
||||
|
||||
Document metadata:
|
||||
Form type: {form_type}
|
||||
Filing date: {filing_date}
|
||||
Time bucket: {filing_time_bucket}
|
||||
Rule hints: {rule_hints}
|
||||
|
||||
Filing text (truncated to {max_chars} chars):
|
||||
---
|
||||
{text}
|
||||
---
|
||||
|
||||
Return ONLY the JSON object.
|
||||
""".strip()
|
||||
|
||||
PROMPT_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"event_classifier_v1": {
|
||||
"version": PROMPT_VERSION,
|
||||
"system": _SYSTEM_EVENT_CLASSIFIER,
|
||||
"user_template": _USER_EVENT_CLASSIFIER,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def render_prompt(
|
||||
prompt_name: str,
|
||||
doc_meta: dict[str, Any],
|
||||
text: str,
|
||||
rule_hints: dict[str, Any],
|
||||
max_chars: int = 8000,
|
||||
) -> list[dict[str, str]]:
|
||||
"""Render a named prompt into an Ollama messages list.
|
||||
|
||||
Args:
|
||||
prompt_name: Key in PROMPT_REGISTRY.
|
||||
doc_meta: Document metadata dict (form_type, filing_date, filing_time_bucket).
|
||||
text: Exhibit text (will be truncated to max_chars).
|
||||
rule_hints: Summary of rule parser output for context.
|
||||
max_chars: Max characters of text to include.
|
||||
|
||||
Returns:
|
||||
List of {role, content} dicts for Ollama chat API.
|
||||
"""
|
||||
entry = PROMPT_REGISTRY[prompt_name]
|
||||
user_content = entry["user_template"].format(
|
||||
form_type=doc_meta.get("form_type", "unknown"),
|
||||
filing_date=doc_meta.get("filing_date", "unknown"),
|
||||
filing_time_bucket=doc_meta.get("filing_time_bucket", "unknown"),
|
||||
rule_hints=rule_hints,
|
||||
text=text[:max_chars],
|
||||
max_chars=max_chars,
|
||||
)
|
||||
return [
|
||||
{"role": "system", "content": entry["system"]},
|
||||
{"role": "user", "content": user_content},
|
||||
]
|
||||
@ -0,0 +1,263 @@
|
||||
"""Canonical merger: combine rule parser output with LLM output."""
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
_HIGH_CONFIDENCE = 0.70
|
||||
_UNKNOWN_VALUES = {"unknown", "unclear", "not_provided"}
|
||||
|
||||
|
||||
def _is_confident(output: dict[str, Any], field: str | None = None) -> bool:
|
||||
"""Return True if output has overall confidence >= threshold."""
|
||||
conf = output.get("confidence", {})
|
||||
if field and field in conf:
|
||||
return float(conf[field]) >= _HIGH_CONFIDENCE
|
||||
return float(conf.get("overall", 0.0)) >= _HIGH_CONFIDENCE
|
||||
|
||||
|
||||
def _field_is_unknown(value: Any) -> bool:
|
||||
return str(value).lower() in _UNKNOWN_VALUES if value is not None else True
|
||||
|
||||
|
||||
def _merge_scalar(
|
||||
rule_val: Any,
|
||||
llm_val: Any | None,
|
||||
rule_conf: float,
|
||||
llm_conf: float,
|
||||
field_name: str,
|
||||
warnings: list[str],
|
||||
provenance: dict[str, str],
|
||||
) -> Any:
|
||||
"""Merge a single scalar field with provenance tracking.
|
||||
|
||||
Strategy:
|
||||
- Both confident + disagree → conflict flag, rule wins (conservative).
|
||||
- Only rule confident → rule wins.
|
||||
- Only LLM confident → LLM wins (fills in unknown/low-confidence rule output).
|
||||
- Neither confident → prefer non-unknown value, rule takes priority.
|
||||
"""
|
||||
rule_unknown = _field_is_unknown(rule_val)
|
||||
llm_unknown = llm_val is None or _field_is_unknown(llm_val)
|
||||
rule_confident = not rule_unknown and rule_conf >= _HIGH_CONFIDENCE
|
||||
llm_confident = not llm_unknown and llm_conf >= _HIGH_CONFIDENCE
|
||||
|
||||
# Both confident and disagree → record conflict; rule wins but conflict is flagged
|
||||
if rule_confident and llm_confident and str(rule_val) != str(llm_val):
|
||||
warnings.append(f"rule_llm_conflict:{field_name} rule={rule_val} llm={llm_val}")
|
||||
provenance[field_name] = "rule"
|
||||
return rule_val
|
||||
|
||||
# Rule wins when confident (and no conflict detected above)
|
||||
if rule_confident:
|
||||
provenance[field_name] = "rule"
|
||||
return rule_val
|
||||
|
||||
# LLM wins when confident (rule is not confident)
|
||||
if llm_confident:
|
||||
provenance[field_name] = "llm"
|
||||
return llm_val
|
||||
|
||||
# Neither confident — prefer non-unknown, rule first
|
||||
if not rule_unknown:
|
||||
provenance[field_name] = "rule"
|
||||
return rule_val
|
||||
|
||||
if not llm_unknown:
|
||||
provenance[field_name] = "llm"
|
||||
return llm_val
|
||||
|
||||
provenance[field_name] = "rule"
|
||||
return rule_val # both unknown → keep rule default
|
||||
|
||||
|
||||
def merge(
|
||||
rule_output: dict[str, Any],
|
||||
llm_output: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Merge rule parser output with optional LLM output into a canonical record.
|
||||
|
||||
Strategy:
|
||||
- Rule parser wins when confident (overall >= 0.70).
|
||||
- LLM fills in unknown fields when it is confident.
|
||||
- Conflicts → conservative merge ("unknown"/"mixed") + warning.
|
||||
- Adds ``provenance`` dict and ``rule_llm_conflict`` flag.
|
||||
"""
|
||||
result = deepcopy(rule_output)
|
||||
result["parser_kind"] = "merged"
|
||||
provenance: dict[str, str] = {}
|
||||
warnings: list[str] = list(result.get("warnings", []))
|
||||
has_conflict = False
|
||||
|
||||
if llm_output is None:
|
||||
# No LLM — everything from rule
|
||||
result["provenance"] = dict.fromkeys(_CANONICAL_FIELDS, "rule")
|
||||
result["rule_llm_conflict"] = False
|
||||
return result
|
||||
|
||||
rule_conf = float(rule_output.get("confidence", {}).get("overall", 0.0))
|
||||
llm_conf = float(llm_output.get("confidence", {}).get("overall", 0.0))
|
||||
|
||||
# Merge top-level scalar fields
|
||||
for field in ("event_type", "event_direction"):
|
||||
merged_val = _merge_scalar(
|
||||
rule_val=rule_output.get(field),
|
||||
llm_val=llm_output.get(field),
|
||||
rule_conf=float(rule_output.get("confidence", {}).get(field, rule_conf)),
|
||||
llm_conf=float(llm_output.get("confidence", {}).get(field, llm_conf)),
|
||||
field_name=field,
|
||||
warnings=warnings,
|
||||
provenance=provenance,
|
||||
)
|
||||
result[field] = merged_val
|
||||
if any(f"rule_llm_conflict:{field}" in w for w in warnings):
|
||||
has_conflict = True
|
||||
|
||||
# Headline / summary: prefer LLM (usually more informative)
|
||||
if llm_output.get("headline") and not _field_is_unknown(llm_output.get("headline")):
|
||||
result["headline"] = llm_output["headline"]
|
||||
provenance["headline"] = "llm"
|
||||
else:
|
||||
provenance["headline"] = "rule"
|
||||
|
||||
if llm_output.get("summary") and not _field_is_unknown(llm_output.get("summary")):
|
||||
result["summary"] = llm_output["summary"]
|
||||
provenance["summary"] = "llm"
|
||||
else:
|
||||
provenance["summary"] = "rule"
|
||||
|
||||
# Guidance
|
||||
rule_guidance = rule_output.get("guidance", {})
|
||||
llm_guidance = llm_output.get("guidance", {})
|
||||
merged_guidance = deepcopy(rule_guidance)
|
||||
guidance_rule_conf = float(rule_output.get("confidence", {}).get("guidance", rule_conf))
|
||||
guidance_llm_conf = float(llm_output.get("confidence", {}).get("guidance", llm_conf))
|
||||
|
||||
g_status = _merge_scalar(
|
||||
rule_val=rule_guidance.get("status"),
|
||||
llm_val=llm_guidance.get("status"),
|
||||
rule_conf=guidance_rule_conf,
|
||||
llm_conf=guidance_llm_conf,
|
||||
field_name="guidance.status",
|
||||
warnings=warnings,
|
||||
provenance=provenance,
|
||||
)
|
||||
merged_guidance["status"] = g_status
|
||||
if any("rule_llm_conflict:guidance.status" in w for w in warnings):
|
||||
has_conflict = True
|
||||
|
||||
if llm_guidance.get("scope") and not _field_is_unknown(llm_guidance.get("scope")):
|
||||
if _field_is_unknown(rule_guidance.get("scope")):
|
||||
merged_guidance["scope"] = llm_guidance["scope"]
|
||||
provenance["guidance.scope"] = "llm"
|
||||
else:
|
||||
provenance["guidance.scope"] = "rule"
|
||||
else:
|
||||
provenance["guidance.scope"] = "rule"
|
||||
|
||||
result["guidance"] = merged_guidance
|
||||
|
||||
# Signals: LLM fills in unknown rule signals
|
||||
rule_signals = rule_output.get("signals", {})
|
||||
llm_signals = llm_output.get("signals", {})
|
||||
merged_signals = deepcopy(rule_signals)
|
||||
for sig_field in (
|
||||
"demand_strength",
|
||||
"pricing_power",
|
||||
"backlog_or_bookings",
|
||||
"customer_expansion",
|
||||
"margin_quality",
|
||||
):
|
||||
if _field_is_unknown(rule_signals.get(sig_field)) and not _field_is_unknown(
|
||||
llm_signals.get(sig_field)
|
||||
):
|
||||
merged_signals[sig_field] = llm_signals[sig_field]
|
||||
provenance[f"signals.{sig_field}"] = "llm"
|
||||
else:
|
||||
provenance[f"signals.{sig_field}"] = "rule"
|
||||
result["signals"] = merged_signals
|
||||
|
||||
# Risk flags: OR of rule + LLM (either flagging = flagged)
|
||||
rule_risks = rule_output.get("risk_flags", {})
|
||||
llm_risks = llm_output.get("risk_flags", {})
|
||||
merged_risks: dict[str, bool] = {}
|
||||
for risk_field in (
|
||||
"oneoff_item",
|
||||
"tax_benefit",
|
||||
"valuation_gain",
|
||||
"non_gaap_heavy",
|
||||
"financing_related",
|
||||
"legal_or_regulatory_overhang",
|
||||
):
|
||||
rule_flag = bool(rule_risks.get(risk_field, False))
|
||||
llm_flag = bool(llm_risks.get(risk_field, False))
|
||||
merged_risks[risk_field] = rule_flag or llm_flag
|
||||
if rule_flag and llm_flag:
|
||||
provenance[f"risk_flags.{risk_field}"] = "merged"
|
||||
elif rule_flag:
|
||||
provenance[f"risk_flags.{risk_field}"] = "rule"
|
||||
elif llm_flag:
|
||||
provenance[f"risk_flags.{risk_field}"] = "llm"
|
||||
else:
|
||||
provenance[f"risk_flags.{risk_field}"] = "rule"
|
||||
result["risk_flags"] = merged_risks
|
||||
|
||||
# Confidence: use higher overall
|
||||
if llm_conf > rule_conf:
|
||||
result["confidence"] = llm_output.get("confidence", result.get("confidence", {}))
|
||||
provenance["confidence"] = "llm"
|
||||
else:
|
||||
provenance["confidence"] = "rule"
|
||||
|
||||
result["provenance"] = provenance
|
||||
result["rule_llm_conflict"] = has_conflict
|
||||
result["warnings"] = warnings
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# Canonical fields used for provenance initialization
|
||||
_CANONICAL_FIELDS = [
|
||||
"event_type",
|
||||
"event_direction",
|
||||
"headline",
|
||||
"summary",
|
||||
"guidance.status",
|
||||
"guidance.scope",
|
||||
"confidence",
|
||||
*[f"signals.{f}" for f in ("demand_strength", "pricing_power", "backlog_or_bookings",
|
||||
"customer_expansion", "margin_quality")],
|
||||
*[f"risk_flags.{f}" for f in ("oneoff_item", "tax_benefit", "valuation_gain",
|
||||
"non_gaap_heavy", "financing_related",
|
||||
"legal_or_regulatory_overhang")],
|
||||
]
|
||||
|
||||
|
||||
def should_queue_for_review(merged: dict[str, Any]) -> tuple[bool, list[str]]:
|
||||
"""Determine if a merged record needs human review.
|
||||
|
||||
Returns:
|
||||
(should_queue, reason_codes) where reason_codes is a list of strings.
|
||||
"""
|
||||
reason_codes: list[str] = []
|
||||
|
||||
conf = float(merged.get("confidence", {}).get("overall", 0.0))
|
||||
if conf < _HIGH_CONFIDENCE:
|
||||
reason_codes.append("low_confidence")
|
||||
|
||||
if merged.get("rule_llm_conflict"):
|
||||
reason_codes.append("rule_llm_conflict")
|
||||
|
||||
guidance = merged.get("guidance", {})
|
||||
if guidance.get("status") in ("unclear",):
|
||||
reason_codes.append("guidance_ambiguous")
|
||||
|
||||
risk_flags = merged.get("risk_flags", {})
|
||||
if risk_flags.get("oneoff_item") or risk_flags.get("non_gaap_heavy"):
|
||||
reason_codes.append("oneoff_likely")
|
||||
|
||||
event_direction = merged.get("event_direction", "unknown")
|
||||
if event_direction in ("unknown", "mixed"):
|
||||
reason_codes.append("direction_ambiguous")
|
||||
|
||||
return bool(reason_codes), reason_codes
|
||||
@ -0,0 +1 @@
|
||||
"""Review queue: manage human-in-the-loop review items."""
|
||||
@ -0,0 +1,176 @@
|
||||
"""Review queue: CRUD operations for ReviewItem."""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from libs.common.logging import get_logger
|
||||
from libs.common.time_utils import utc_now
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_PRIORITY_ORDER = {"P0": 0, "P1": 1, "P2": 2}
|
||||
|
||||
|
||||
async def create_review_item(
|
||||
session: AsyncSession,
|
||||
entity_type: str,
|
||||
entity_id: str,
|
||||
priority: str,
|
||||
reason_codes: list[str],
|
||||
snapshot_refs: dict[str, Any],
|
||||
suggested_overrides: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Create a ReviewItem, or update an existing open one (deduplicated).
|
||||
|
||||
Args:
|
||||
session: Async DB session.
|
||||
entity_type: "parser_event" | "feature_record" | "label_record".
|
||||
entity_id: ID of the entity to review.
|
||||
priority: "P0" | "P1" | "P2".
|
||||
reason_codes: List of reason code strings.
|
||||
snapshot_refs: Dict of snapshot references.
|
||||
suggested_overrides: Optional suggested field overrides.
|
||||
|
||||
Returns:
|
||||
ReviewItem ORM instance.
|
||||
"""
|
||||
from libs.db.models import ReviewItem
|
||||
|
||||
# Check for existing open item
|
||||
result = await session.execute(
|
||||
select(ReviewItem).where(
|
||||
ReviewItem.entity_type == entity_type,
|
||||
ReviewItem.entity_id == entity_id,
|
||||
ReviewItem.status == "open",
|
||||
)
|
||||
)
|
||||
existing = result.scalar_one_or_none()
|
||||
|
||||
if existing is not None:
|
||||
# Update existing: escalate priority if needed, merge reason_codes
|
||||
existing_priority_rank = _PRIORITY_ORDER.get(existing.priority, 99)
|
||||
new_priority_rank = _PRIORITY_ORDER.get(priority, 99)
|
||||
if new_priority_rank < existing_priority_rank:
|
||||
existing.priority = priority # type: ignore[assignment]
|
||||
|
||||
existing_codes: list[str] = list(existing.reason_codes or [])
|
||||
for code in reason_codes:
|
||||
if code not in existing_codes:
|
||||
existing_codes.append(code)
|
||||
existing.reason_codes = existing_codes # type: ignore[assignment]
|
||||
|
||||
if suggested_overrides:
|
||||
existing.suggested_overrides = suggested_overrides # type: ignore[assignment]
|
||||
|
||||
existing.updated_at_utc = utc_now() # type: ignore[assignment]
|
||||
await session.flush()
|
||||
logger.info("review_item_updated", entity_id=entity_id, entity_type=entity_type)
|
||||
return existing
|
||||
|
||||
item = ReviewItem(
|
||||
review_id=uuid.uuid4(),
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
priority=priority,
|
||||
reason_codes=reason_codes,
|
||||
status="open",
|
||||
snapshot_refs=snapshot_refs,
|
||||
suggested_overrides=suggested_overrides,
|
||||
created_at_utc=utc_now(),
|
||||
updated_at_utc=utc_now(),
|
||||
)
|
||||
session.add(item)
|
||||
await session.flush()
|
||||
logger.info("review_item_created", entity_id=entity_id, priority=priority, reasons=reason_codes)
|
||||
return item
|
||||
|
||||
|
||||
async def resolve_review_item(
|
||||
session: AsyncSession,
|
||||
review_id: str | uuid.UUID,
|
||||
reviewer_id: str,
|
||||
resolution_type: str,
|
||||
root_cause: str,
|
||||
notes: str,
|
||||
field_overrides: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Resolve a ReviewItem.
|
||||
|
||||
Args:
|
||||
session: Async DB session.
|
||||
review_id: UUID of the ReviewItem to resolve.
|
||||
reviewer_id: ID/name of the reviewer.
|
||||
resolution_type: E.g. "accepted", "corrected", "wont_fix".
|
||||
root_cause: Short description of root cause.
|
||||
notes: Free-text reviewer notes.
|
||||
field_overrides: Optional field overrides applied.
|
||||
|
||||
Returns:
|
||||
Resolved ReviewItem.
|
||||
|
||||
Raises:
|
||||
ValueError: If item not found or already resolved.
|
||||
"""
|
||||
from libs.db.models import ReviewItem
|
||||
|
||||
if isinstance(review_id, str):
|
||||
review_id = uuid.UUID(review_id)
|
||||
|
||||
result = await session.execute(
|
||||
select(ReviewItem).where(ReviewItem.review_id == review_id)
|
||||
)
|
||||
item = result.scalar_one_or_none()
|
||||
|
||||
if item is None:
|
||||
raise ValueError(f"ReviewItem {review_id} not found")
|
||||
if item.status != "open": # type: ignore[comparison-overlap]
|
||||
raise ValueError(f"ReviewItem {review_id} is already {item.status}")
|
||||
|
||||
item.status = "resolved" # type: ignore[assignment]
|
||||
item.reviewer_id = reviewer_id # type: ignore[assignment]
|
||||
item.resolution_type = resolution_type # type: ignore[assignment]
|
||||
item.root_cause = root_cause # type: ignore[assignment]
|
||||
item.notes = notes # type: ignore[assignment]
|
||||
item.field_overrides = field_overrides # type: ignore[assignment]
|
||||
item.resolved_at = utc_now() # type: ignore[assignment]
|
||||
item.updated_at_utc = utc_now() # type: ignore[assignment]
|
||||
|
||||
await session.flush()
|
||||
logger.info("review_item_resolved", review_id=str(review_id), resolution_type=resolution_type)
|
||||
return item
|
||||
|
||||
|
||||
async def list_review_items(
|
||||
session: AsyncSession,
|
||||
status: str | None = None,
|
||||
priority: str | None = None,
|
||||
entity_type: str | None = None,
|
||||
) -> list[Any]:
|
||||
"""List ReviewItems with optional filters.
|
||||
|
||||
Args:
|
||||
session: Async DB session.
|
||||
status: Filter by status ("open", "resolved", etc.).
|
||||
priority: Filter by priority ("P0", "P1", "P2").
|
||||
entity_type: Filter by entity_type.
|
||||
|
||||
Returns:
|
||||
List of ReviewItem ORM instances.
|
||||
"""
|
||||
from libs.db.models import ReviewItem
|
||||
|
||||
stmt = select(ReviewItem)
|
||||
if status is not None:
|
||||
stmt = stmt.where(ReviewItem.status == status)
|
||||
if priority is not None:
|
||||
stmt = stmt.where(ReviewItem.priority == priority)
|
||||
if entity_type is not None:
|
||||
stmt = stmt.where(ReviewItem.entity_type == entity_type)
|
||||
|
||||
stmt = stmt.order_by(ReviewItem.created_at_utc.desc())
|
||||
result = await session.execute(stmt)
|
||||
return list(result.scalars().all())
|
||||
@ -0,0 +1,92 @@
|
||||
"""Integration test: event → label generation pipeline (real Oracle + DB)."""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_event() -> MagicMock:
|
||||
event = MagicMock()
|
||||
event.event_id = "EVT::sec::ISSUER::0000320193::2026-01-29::earnings_release::0"
|
||||
event.event_date = dt.date(2026, 1, 29)
|
||||
event.filing_time_bucket = "post_market"
|
||||
event.symbol_id = "SYM::AAPL::NASDAQ"
|
||||
event.status = "valid"
|
||||
return event
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_price_bars() -> list[MagicMock]:
|
||||
bars = []
|
||||
for i in range(8):
|
||||
bar = MagicMock()
|
||||
date = dt.date(2026, 1, 30) + dt.timedelta(days=i)
|
||||
bar.model_dump.return_value = {
|
||||
"date": date.isoformat(),
|
||||
"open": 220.0 + i * 0.5,
|
||||
"high": 225.0 + i * 0.5,
|
||||
"low": 218.0,
|
||||
"close": 222.0 + i * 0.5,
|
||||
"volume": 1000000,
|
||||
}
|
||||
bars.append(bar)
|
||||
return bars
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_label_pipeline_end_to_end(mock_event: MagicMock, mock_price_bars: list) -> None:
|
||||
"""Test full label generation: event → reaction_date → entry → labels."""
|
||||
from libs.labeler.label_generator import LABEL_VERSION, generate_labels
|
||||
from libs.labeler.reaction_date import compute_reaction_date
|
||||
|
||||
# Verify reaction_date logic
|
||||
reaction_date = compute_reaction_date(mock_event.event_date, mock_event.filing_time_bucket)
|
||||
assert reaction_date > mock_event.event_date # post_market → next day
|
||||
|
||||
# Setup mock price service
|
||||
mock_price_resp = MagicMock()
|
||||
mock_price_resp.bars = mock_price_bars
|
||||
mock_price_svc = AsyncMock()
|
||||
mock_price_svc.get_daily_bars = AsyncMock(return_value=mock_price_resp)
|
||||
mock_session = AsyncMock()
|
||||
|
||||
label = await generate_labels(
|
||||
session=mock_session,
|
||||
event=mock_event,
|
||||
price_svc=mock_price_svc,
|
||||
ticker="AAPL",
|
||||
)
|
||||
|
||||
assert label is not None
|
||||
assert label.event_id == mock_event.event_id
|
||||
assert label.label_version == LABEL_VERSION
|
||||
assert label.label_status in ("ok", "truncated")
|
||||
assert label.reaction_date == reaction_date
|
||||
assert label.entry_price is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_label_pipeline_handles_missing_price_data(mock_event: MagicMock) -> None:
|
||||
"""Label pipeline handles Oracle price unavailability gracefully."""
|
||||
from libs.labeler.label_generator import generate_labels
|
||||
|
||||
mock_price_svc = AsyncMock()
|
||||
mock_price_svc.get_daily_bars = AsyncMock(
|
||||
side_effect=Exception("Oracle connection refused")
|
||||
)
|
||||
mock_session = AsyncMock()
|
||||
|
||||
label = await generate_labels(
|
||||
session=mock_session,
|
||||
event=mock_event,
|
||||
price_svc=mock_price_svc,
|
||||
ticker="AAPL",
|
||||
)
|
||||
|
||||
assert label.label_status == "unavailable"
|
||||
assert label.event_id == mock_event.event_id
|
||||
@ -0,0 +1,118 @@
|
||||
"""Integration test: merger output → review queue creation."""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
_RULE_OUTPUT = {
|
||||
"schema_version": "1.0.0",
|
||||
"document_id": "DOC::test",
|
||||
"parser_kind": "rule",
|
||||
"event_type": "earnings_release",
|
||||
"event_direction": "unknown", # low confidence
|
||||
"event_date": "2026-01-01",
|
||||
"filing_time_bucket": "post_market",
|
||||
"headline": "",
|
||||
"summary": "Earnings results.",
|
||||
"guidance": {"status": "unclear", "scope": "unknown", "notes": ""},
|
||||
"signals": {
|
||||
"demand_strength": "unknown",
|
||||
"pricing_power": "unknown",
|
||||
"backlog_or_bookings": "unknown",
|
||||
"customer_expansion": "unknown",
|
||||
"margin_quality": "unknown",
|
||||
},
|
||||
"risk_flags": {
|
||||
"oneoff_item": True,
|
||||
"tax_benefit": False,
|
||||
"valuation_gain": False,
|
||||
"non_gaap_heavy": True,
|
||||
"financing_related": False,
|
||||
"legal_or_regulatory_overhang": False,
|
||||
},
|
||||
"evidence": [],
|
||||
"confidence": {
|
||||
"overall": 0.40, # low
|
||||
"event_type": 0.60,
|
||||
"event_direction": 0.30,
|
||||
"guidance": 0.40,
|
||||
"risk_flags": 0.70,
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
|
||||
def _make_mock_session_with_no_existing() -> AsyncMock:
|
||||
session = AsyncMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = None
|
||||
session.execute = AsyncMock(return_value=mock_result)
|
||||
session.add = MagicMock()
|
||||
session.flush = AsyncMock()
|
||||
return session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_low_confidence_merge_creates_review_item() -> None:
|
||||
"""A low-confidence merged record triggers review item creation."""
|
||||
from libs.parser.merger import merge, should_queue_for_review
|
||||
from libs.review.queue import create_review_item
|
||||
|
||||
merged = merge(_RULE_OUTPUT, None) # No LLM
|
||||
should_queue, reason_codes = should_queue_for_review(merged)
|
||||
|
||||
assert should_queue is True
|
||||
assert "low_confidence" in reason_codes
|
||||
|
||||
# Mock session for review item creation
|
||||
session = _make_mock_session_with_no_existing()
|
||||
|
||||
await create_review_item(
|
||||
session=session,
|
||||
entity_type="parser_event",
|
||||
entity_id=merged["document_id"],
|
||||
priority="P1",
|
||||
reason_codes=reason_codes,
|
||||
snapshot_refs={"document_id": merged["document_id"]},
|
||||
)
|
||||
|
||||
session.add.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflict_merge_creates_p0_review_item() -> None:
|
||||
"""A rule/LLM conflict triggers P0 review item."""
|
||||
from libs.parser.merger import merge, should_queue_for_review
|
||||
from libs.review.queue import create_review_item
|
||||
|
||||
llm_conflict = copy.deepcopy(_RULE_OUTPUT)
|
||||
llm_conflict["parser_kind"] = "llm"
|
||||
llm_conflict["event_type"] = "guidance_update" # Conflict
|
||||
llm_conflict["confidence"]["overall"] = 0.80
|
||||
llm_conflict["confidence"]["event_type"] = 0.85
|
||||
|
||||
rule_high = copy.deepcopy(_RULE_OUTPUT)
|
||||
rule_high["confidence"]["overall"] = 0.80
|
||||
rule_high["confidence"]["event_type"] = 0.90
|
||||
|
||||
merged = merge(rule_high, llm_conflict)
|
||||
should_queue, reason_codes = should_queue_for_review(merged)
|
||||
|
||||
assert should_queue is True
|
||||
|
||||
session = _make_mock_session_with_no_existing()
|
||||
|
||||
await create_review_item(
|
||||
session=session,
|
||||
entity_type="parser_event",
|
||||
entity_id=merged["document_id"],
|
||||
priority="P0",
|
||||
reason_codes=reason_codes,
|
||||
snapshot_refs={"document_id": merged["document_id"]},
|
||||
)
|
||||
|
||||
session.add.assert_called_once()
|
||||
@ -0,0 +1,113 @@
|
||||
"""Replay test: verify same document → LLM cache hit on second call."""
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.replay
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_document_hits_cache_on_replay() -> None:
|
||||
"""Processing the same document twice should hit cache on the second call."""
|
||||
from libs.llm.cache import LLMCacheStore
|
||||
from libs.llm.client import OllamaClient
|
||||
from libs.llm.parser import LLMParser
|
||||
|
||||
# Simulate cache miss on first call, hit on second
|
||||
cached_output = {
|
||||
"schema_version": "1.0.0",
|
||||
"document_id": "DOC::replay_test",
|
||||
"parser_kind": "llm",
|
||||
"event_type": "earnings_release",
|
||||
"event_direction": "bullish",
|
||||
"event_date": "2026-01-15",
|
||||
"filing_time_bucket": "post_market",
|
||||
"headline": "Q1 results beat",
|
||||
"summary": "Earnings beat estimates.",
|
||||
"guidance": {"status": "raised", "scope": "annual", "notes": ""},
|
||||
"signals": {
|
||||
"demand_strength": "strong",
|
||||
"pricing_power": "present",
|
||||
"backlog_or_bookings": "present",
|
||||
"customer_expansion": "present",
|
||||
"margin_quality": "improving",
|
||||
},
|
||||
"risk_flags": {
|
||||
"oneoff_item": False,
|
||||
"tax_benefit": False,
|
||||
"valuation_gain": False,
|
||||
"non_gaap_heavy": False,
|
||||
"financing_related": False,
|
||||
"legal_or_regulatory_overhang": False,
|
||||
},
|
||||
"evidence": [],
|
||||
"confidence": {
|
||||
"overall": 0.85,
|
||||
"event_type": 0.90,
|
||||
"event_direction": 0.80,
|
||||
"guidance": 0.75,
|
||||
"risk_flags": 0.90,
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
call_count = {"n": 0}
|
||||
|
||||
async def mock_get(session: object, cache_key: str) -> dict | None:
|
||||
if call_count["n"] == 0:
|
||||
return None # cache miss on first call
|
||||
return cached_output # cache hit on subsequent calls
|
||||
|
||||
async def mock_put(*args: object, **kwargs: object) -> None:
|
||||
call_count["n"] += 1 # increment after first call stores to cache
|
||||
|
||||
mock_cache = MagicMock(spec=LLMCacheStore)
|
||||
mock_cache.get = AsyncMock(side_effect=mock_get)
|
||||
mock_cache.put = AsyncMock(side_effect=mock_put)
|
||||
|
||||
# Mock Ollama client (only called once)
|
||||
mock_ollama = MagicMock(spec=OllamaClient)
|
||||
mock_ollama.model = "llama3.2"
|
||||
mock_ollama.chat = AsyncMock(
|
||||
return_value=(
|
||||
{k: v for k, v in cached_output.items() if k != "schema_version"},
|
||||
{"prompt_tokens": 100, "completion_tokens": 50},
|
||||
300,
|
||||
)
|
||||
)
|
||||
|
||||
llm_parser = LLMParser(client=mock_ollama, cache_store=mock_cache)
|
||||
|
||||
doc_text = "Apple reports Q1 earnings: revenue $123B, EPS $2.50, beats estimates."
|
||||
doc_meta = {
|
||||
"form_type": "8-K",
|
||||
"filing_date": "2026-01-15",
|
||||
"filing_time_bucket": "post_market",
|
||||
}
|
||||
|
||||
# First call — cache miss, Ollama called
|
||||
mock_session1 = AsyncMock()
|
||||
result1 = await llm_parser.parse(
|
||||
document_id="DOC::replay_test",
|
||||
doc_text=doc_text,
|
||||
doc_meta=doc_meta,
|
||||
rule_hints={},
|
||||
session=mock_session1,
|
||||
)
|
||||
assert result1 is not None
|
||||
assert mock_ollama.chat.call_count == 1
|
||||
|
||||
# Second call — cache hit, Ollama NOT called again
|
||||
mock_session2 = AsyncMock()
|
||||
result2 = await llm_parser.parse(
|
||||
document_id="DOC::replay_test",
|
||||
doc_text=doc_text,
|
||||
doc_meta=doc_meta,
|
||||
rule_hints={},
|
||||
session=mock_session2,
|
||||
)
|
||||
assert result2 == cached_output
|
||||
# Ollama was still only called once (not twice)
|
||||
assert mock_ollama.chat.call_count == 1
|
||||
@ -0,0 +1,197 @@
|
||||
"""Unit tests for labeler module."""
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
from decimal import Decimal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.labeler.label_generator import _compute_labels_from_bars
|
||||
from libs.labeler.reaction_date import compute_reaction_date
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestComputeReactionDate:
|
||||
"""Tests for reaction date calculation."""
|
||||
|
||||
def test_pre_market_on_trading_day_returns_same_day(self) -> None:
|
||||
"""Filing before market open on a trading day → same day reaction."""
|
||||
# 2026-01-02 is a Friday (trading day)
|
||||
d = dt.date(2026, 1, 2)
|
||||
result = compute_reaction_date(d, "pre_market")
|
||||
assert result == d
|
||||
|
||||
def test_regular_hours_on_trading_day_returns_same_day(self) -> None:
|
||||
"""Filing during regular hours on a trading day → same day reaction."""
|
||||
d = dt.date(2026, 1, 2)
|
||||
result = compute_reaction_date(d, "regular_hours")
|
||||
assert result == d
|
||||
|
||||
def test_post_market_returns_next_trading_day(self) -> None:
|
||||
"""Filing after market close → next trading day reaction."""
|
||||
d = dt.date(2026, 1, 2)
|
||||
result = compute_reaction_date(d, "post_market")
|
||||
assert result > d
|
||||
|
||||
def test_unknown_returns_next_trading_day(self) -> None:
|
||||
"""Unknown bucket → conservative: next trading day."""
|
||||
d = dt.date(2026, 1, 2)
|
||||
result = compute_reaction_date(d, "unknown")
|
||||
assert result > d
|
||||
|
||||
def test_pre_market_on_weekend_returns_next_trading_day(self) -> None:
|
||||
"""Pre-market filing on weekend (non-trading day) → next trading day."""
|
||||
saturday = dt.date(2026, 1, 3) # Saturday
|
||||
result = compute_reaction_date(saturday, "pre_market")
|
||||
assert result > saturday
|
||||
|
||||
def test_post_market_on_friday_returns_monday(self) -> None:
|
||||
"""Post-market on Friday → next Monday (assuming no holiday)."""
|
||||
friday = dt.date(2026, 1, 2) # 2026-01-02 is a Friday
|
||||
result = compute_reaction_date(friday, "post_market")
|
||||
# Next trading day after Friday is Monday
|
||||
assert result.weekday() == 0 # Monday
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestComputeLabelsFromBars:
|
||||
"""Tests for forward-return label computation."""
|
||||
|
||||
def _bar(self, close: float, high: float | None = None, low: float | None = None) -> dict:
|
||||
return {
|
||||
"open": close * 0.99,
|
||||
"high": high if high is not None else close * 1.02,
|
||||
"low": low if low is not None else close * 0.98,
|
||||
"close": close,
|
||||
}
|
||||
|
||||
def test_1d_return_calculation(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [self._bar(105)] # +5%
|
||||
result = _compute_labels_from_bars(entry, bars, 1)
|
||||
assert abs(float(result["fwd_return"]) - 0.05) < 0.001
|
||||
|
||||
def test_mfe_is_max_high_minus_entry(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [
|
||||
self._bar(101, high=105),
|
||||
self._bar(103, high=108),
|
||||
self._bar(102, high=104),
|
||||
]
|
||||
result = _compute_labels_from_bars(entry, bars, 3)
|
||||
# Max high = 108, so MFE = (108-100)/100 = 0.08
|
||||
assert abs(float(result["mfe"]) - 0.08) < 0.001
|
||||
|
||||
def test_mae_is_min_low_minus_entry(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [
|
||||
self._bar(99, low=97),
|
||||
self._bar(98, low=95),
|
||||
self._bar(100, low=98),
|
||||
]
|
||||
result = _compute_labels_from_bars(entry, bars, 3)
|
||||
# Min low = 95, so MAE = (95-100)/100 = -0.05
|
||||
assert abs(float(result["mae"]) - (-0.05)) < 0.001
|
||||
|
||||
def test_hit_pos_1r_true_when_high_exceeds_threshold(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [self._bar(99, high=101.5)] # +1.5% > 1R threshold
|
||||
result = _compute_labels_from_bars(entry, bars, 1)
|
||||
assert result["hit_pos_1r"] is True
|
||||
|
||||
def test_hit_pos_1r_false_when_high_below_threshold(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [self._bar(99, high=100.5)] # +0.5% < 1R threshold
|
||||
result = _compute_labels_from_bars(entry, bars, 1)
|
||||
assert result["hit_pos_1r"] is False
|
||||
|
||||
def test_close_up_after_3d_true_when_final_close_above_entry(self) -> None:
|
||||
entry = Decimal("100")
|
||||
bars = [self._bar(98), self._bar(101), self._bar(103)]
|
||||
result = _compute_labels_from_bars(entry, bars, 3)
|
||||
assert result["close_up"] is True
|
||||
|
||||
def test_empty_bars_returns_empty_dict(self) -> None:
|
||||
result = _compute_labels_from_bars(Decimal("100"), [], 3)
|
||||
assert result == {}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestGenerateLabels:
|
||||
"""Tests for async generate_labels function."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_labels_with_valid_prices(self) -> None:
|
||||
"""generate_labels returns EventLabel with ok status when prices available."""
|
||||
from libs.labeler.label_generator import generate_labels
|
||||
|
||||
# Mock event
|
||||
mock_event = MagicMock()
|
||||
mock_event.event_id = "EVT::test::001"
|
||||
mock_event.event_date = dt.date(2026, 1, 5) # Monday
|
||||
mock_event.filing_time_bucket = "post_market"
|
||||
|
||||
# Mock price service
|
||||
mock_price_svc = AsyncMock()
|
||||
mock_bar = MagicMock()
|
||||
mock_bar.model_dump.return_value = {
|
||||
"date": "2026-01-07", # Wednesday = entry_date
|
||||
"open": 100.0,
|
||||
"high": 105.0,
|
||||
"low": 98.0,
|
||||
"close": 103.0,
|
||||
}
|
||||
# Create 8 bars for look-ahead
|
||||
bars = []
|
||||
for i in range(8):
|
||||
b = MagicMock()
|
||||
date = dt.date(2026, 1, 7) + dt.timedelta(days=i)
|
||||
b.model_dump.return_value = {
|
||||
"date": date.isoformat(),
|
||||
"open": 100.0 + i,
|
||||
"high": 105.0 + i,
|
||||
"low": 98.0,
|
||||
"close": 103.0 + i,
|
||||
}
|
||||
bars.append(b)
|
||||
|
||||
mock_resp = MagicMock()
|
||||
mock_resp.bars = bars
|
||||
mock_price_svc.get_daily_bars = AsyncMock(return_value=mock_resp)
|
||||
|
||||
mock_session = AsyncMock()
|
||||
|
||||
label = await generate_labels(
|
||||
session=mock_session,
|
||||
event=mock_event,
|
||||
price_svc=mock_price_svc,
|
||||
ticker="AAPL",
|
||||
)
|
||||
|
||||
assert label.event_id == "EVT::test::001"
|
||||
assert label.label_status in ("ok", "truncated")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_labels_unavailable_when_no_price_data(self) -> None:
|
||||
"""generate_labels returns 'unavailable' status on price fetch error."""
|
||||
from libs.labeler.label_generator import generate_labels
|
||||
|
||||
mock_event = MagicMock()
|
||||
mock_event.event_id = "EVT::test::002"
|
||||
mock_event.event_date = dt.date(2026, 1, 5)
|
||||
mock_event.filing_time_bucket = "post_market"
|
||||
|
||||
mock_price_svc = AsyncMock()
|
||||
mock_price_svc.get_daily_bars = AsyncMock(side_effect=Exception("Oracle unavailable"))
|
||||
|
||||
mock_session = AsyncMock()
|
||||
|
||||
label = await generate_labels(
|
||||
session=mock_session,
|
||||
event=mock_event,
|
||||
price_svc=mock_price_svc,
|
||||
ticker="AAPL",
|
||||
)
|
||||
|
||||
assert label.label_status == "unavailable"
|
||||
@ -0,0 +1,66 @@
|
||||
"""Unit tests for LLMCacheStore."""
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.llm.cache import LLMCacheStore, build_cache_key
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestBuildCacheKey:
|
||||
def test_deterministic(self) -> None:
|
||||
key1 = build_cache_key("text", "prompt_v1", "v1", "llama3.2", "1.0.0")
|
||||
key2 = build_cache_key("text", "prompt_v1", "v1", "llama3.2", "1.0.0")
|
||||
assert key1 == key2
|
||||
|
||||
def test_different_inputs_different_keys(self) -> None:
|
||||
key1 = build_cache_key("text_a", "p", "v1", "llama3.2")
|
||||
key2 = build_cache_key("text_b", "p", "v1", "llama3.2")
|
||||
assert key1 != key2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestLLMCacheStore:
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_returns_none_on_miss(self) -> None:
|
||||
"""Cache miss returns None."""
|
||||
cache = LLMCacheStore()
|
||||
mock_session = AsyncMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = None
|
||||
mock_session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
result = await cache.get(mock_session, "deadbeef" * 8)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_stores_entry(self) -> None:
|
||||
"""Cache put creates a new row when key is absent."""
|
||||
cache = LLMCacheStore()
|
||||
mock_session = AsyncMock()
|
||||
|
||||
# Simulate no existing entry
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = None
|
||||
mock_session.execute = AsyncMock(return_value=mock_result)
|
||||
mock_session.add = MagicMock()
|
||||
mock_session.flush = AsyncMock()
|
||||
|
||||
await cache.put(
|
||||
session=mock_session,
|
||||
cache_key="deadbeef" * 8,
|
||||
document_id="DOC::test",
|
||||
model_name="llama3.2",
|
||||
prompt_version="v1",
|
||||
schema_version="1.0.0",
|
||||
raw_prompt='[{"role":"user","content":"test"}]',
|
||||
raw_response='{"event_type":"unknown"}',
|
||||
normalized={"event_type": "unknown"},
|
||||
token_usage={"prompt_tokens": 10, "completion_tokens": 5},
|
||||
elapsed_ms=250,
|
||||
)
|
||||
|
||||
mock_session.add.assert_called_once()
|
||||
mock_session.flush.assert_called_once()
|
||||
@ -0,0 +1,137 @@
|
||||
"""Unit tests for OllamaClient."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from libs.common.retries import RetryableError
|
||||
from libs.llm.client import OllamaClient
|
||||
from libs.llm.exceptions import LLMTimeoutError
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestOllamaClientChat:
|
||||
"""Tests for OllamaClient.chat()."""
|
||||
|
||||
@pytest.fixture
|
||||
def client(self) -> OllamaClient:
|
||||
return OllamaClient(base_url="http://localhost:11434", model="llama3.2")
|
||||
|
||||
async def _mock_response(self, payload: dict) -> httpx.Response:
|
||||
return httpx.Response(200, json=payload)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_success_returns_parsed_json(self, client: OllamaClient) -> None:
|
||||
"""A successful Ollama response is parsed and returned as a dict."""
|
||||
expected = {"event_type": "earnings_release", "event_direction": "bullish"}
|
||||
mock_resp_payload = {
|
||||
"message": {"content": json.dumps(expected)},
|
||||
"prompt_eval_count": 100,
|
||||
"eval_count": 50,
|
||||
}
|
||||
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(return_value=httpx.Response(200, json=mock_resp_payload))
|
||||
|
||||
async with client:
|
||||
client._client = mock_http # inject mock
|
||||
result, token_usage, elapsed_ms = await client.chat(
|
||||
[{"role": "user", "content": "analyze this"}]
|
||||
)
|
||||
|
||||
assert result == expected
|
||||
assert token_usage["prompt_tokens"] == 100
|
||||
assert token_usage["completion_tokens"] == 50
|
||||
assert elapsed_ms >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_timeout_raises_llm_timeout_error(self, client: OllamaClient) -> None:
|
||||
"""Timeout raises LLMTimeoutError (which is also RetryableError)."""
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(side_effect=httpx.ReadTimeout("timeout"))
|
||||
|
||||
async with client:
|
||||
client._client = mock_http
|
||||
with pytest.raises(LLMTimeoutError):
|
||||
await client.chat([{"role": "user", "content": "test"}])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chat_5xx_raises_retryable_error(self, client: OllamaClient) -> None:
|
||||
"""5xx response raises RetryableError."""
|
||||
mock_http = AsyncMock()
|
||||
mock_http.post = AsyncMock(
|
||||
return_value=httpx.Response(503, text="Service Unavailable")
|
||||
)
|
||||
|
||||
async with client:
|
||||
client._client = mock_http
|
||||
with pytest.raises(RetryableError):
|
||||
await client.chat([{"role": "user", "content": "test"}])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_hit_skips_llm_call(self) -> None:
|
||||
"""LLMParser returns cached result without calling Ollama."""
|
||||
from unittest.mock import AsyncMock as AM
|
||||
|
||||
from libs.llm.cache import LLMCacheStore
|
||||
from libs.llm.parser import LLMParser
|
||||
|
||||
cached_output = {
|
||||
"schema_version": "1.0.0",
|
||||
"document_id": "DOC::test",
|
||||
"parser_kind": "llm",
|
||||
"event_type": "earnings_release",
|
||||
"event_direction": "bullish",
|
||||
"event_date": "2026-01-01",
|
||||
"filing_time_bucket": "post_market",
|
||||
"headline": "Test",
|
||||
"summary": "Test summary",
|
||||
"guidance": {"status": "raised", "scope": "annual", "notes": ""},
|
||||
"signals": {
|
||||
"demand_strength": "strong",
|
||||
"pricing_power": "present",
|
||||
"backlog_or_bookings": "present",
|
||||
"customer_expansion": "present",
|
||||
"margin_quality": "improving",
|
||||
},
|
||||
"risk_flags": {
|
||||
"oneoff_item": False,
|
||||
"tax_benefit": False,
|
||||
"valuation_gain": False,
|
||||
"non_gaap_heavy": False,
|
||||
"financing_related": False,
|
||||
"legal_or_regulatory_overhang": False,
|
||||
},
|
||||
"evidence": [],
|
||||
"confidence": {
|
||||
"overall": 0.85,
|
||||
"event_type": 0.9,
|
||||
"event_direction": 0.8,
|
||||
"guidance": 0.8,
|
||||
"risk_flags": 0.9,
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
mock_cache = MagicMock(spec=LLMCacheStore)
|
||||
mock_cache.get = AM(return_value=cached_output)
|
||||
|
||||
mock_client = MagicMock(spec=OllamaClient)
|
||||
mock_client.model = "llama3.2"
|
||||
|
||||
llm_parser = LLMParser(client=mock_client, cache_store=mock_cache)
|
||||
mock_session = AsyncMock()
|
||||
|
||||
result = await llm_parser.parse(
|
||||
document_id="DOC::test",
|
||||
doc_text="some text",
|
||||
doc_meta={"filing_date": "2026-01-01", "form_type": "8-K"},
|
||||
rule_hints={},
|
||||
session=mock_session,
|
||||
)
|
||||
|
||||
assert result == cached_output
|
||||
mock_client.chat.assert_not_called()
|
||||
@ -0,0 +1,174 @@
|
||||
"""Unit tests for canonical merger."""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.parser.merger import merge, should_queue_for_review
|
||||
|
||||
_BASE_RULE = {
|
||||
"schema_version": "1.0.0",
|
||||
"document_id": "DOC::test",
|
||||
"parser_kind": "rule",
|
||||
"event_type": "earnings_release",
|
||||
"event_direction": "bullish",
|
||||
"event_date": "2026-01-01",
|
||||
"filing_time_bucket": "post_market",
|
||||
"headline": "",
|
||||
"summary": "Q1 results beat estimates.",
|
||||
"guidance": {"status": "raised", "scope": "annual", "notes": ""},
|
||||
"signals": {
|
||||
"demand_strength": "strong",
|
||||
"pricing_power": "present",
|
||||
"backlog_or_bookings": "unknown",
|
||||
"customer_expansion": "unknown",
|
||||
"margin_quality": "improving",
|
||||
},
|
||||
"risk_flags": {
|
||||
"oneoff_item": False,
|
||||
"tax_benefit": False,
|
||||
"valuation_gain": False,
|
||||
"non_gaap_heavy": True,
|
||||
"financing_related": False,
|
||||
"legal_or_regulatory_overhang": False,
|
||||
},
|
||||
"evidence": [],
|
||||
"confidence": {
|
||||
"overall": 0.80,
|
||||
"event_type": 0.90,
|
||||
"event_direction": 0.75,
|
||||
"guidance": 0.80,
|
||||
"risk_flags": 0.85,
|
||||
},
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
_BASE_LLM = {
|
||||
**copy.deepcopy(_BASE_RULE),
|
||||
"parser_kind": "llm",
|
||||
"headline": "Apple beats Q1 expectations",
|
||||
"summary": "Apple Inc. reported strong Q1 results driven by iPhone sales.",
|
||||
"signals": {
|
||||
"demand_strength": "strong",
|
||||
"pricing_power": "present",
|
||||
"backlog_or_bookings": "present",
|
||||
"customer_expansion": "present",
|
||||
"margin_quality": "improving",
|
||||
},
|
||||
"confidence": {
|
||||
"overall": 0.85,
|
||||
"event_type": 0.90,
|
||||
"event_direction": 0.80,
|
||||
"guidance": 0.75,
|
||||
"risk_flags": 0.85,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestMerge:
|
||||
def test_rule_wins_when_both_confident(self) -> None:
|
||||
"""Rule wins on event_type/direction when rule confidence >= 0.70."""
|
||||
result = merge(_BASE_RULE, _BASE_LLM)
|
||||
assert result["event_type"] == "earnings_release"
|
||||
assert result["event_direction"] == "bullish"
|
||||
assert result["provenance"]["event_type"] == "rule"
|
||||
|
||||
def test_llm_fills_unknown_signals(self) -> None:
|
||||
"""LLM fills in unknown signals that rule left blank."""
|
||||
result = merge(_BASE_RULE, _BASE_LLM)
|
||||
# backlog_or_bookings was unknown in rule but present in LLM
|
||||
assert result["signals"]["backlog_or_bookings"] == "present"
|
||||
assert result["provenance"]["signals.backlog_or_bookings"] == "llm"
|
||||
# customer_expansion was unknown in rule but present in LLM
|
||||
assert result["signals"]["customer_expansion"] == "present"
|
||||
|
||||
def test_llm_wins_when_rule_not_confident(self) -> None:
|
||||
"""LLM wins on a field when rule confidence is below threshold."""
|
||||
llm_bearish = copy.deepcopy(_BASE_LLM)
|
||||
llm_bearish["event_direction"] = "bearish"
|
||||
llm_bearish["confidence"]["event_direction"] = 0.80 # LLM confident
|
||||
|
||||
rule_low = copy.deepcopy(_BASE_RULE)
|
||||
rule_low["confidence"]["event_direction"] = 0.40 # below threshold → not confident
|
||||
|
||||
result = merge(rule_low, llm_bearish)
|
||||
# Rule not confident → LLM fills in "bearish", no conflict
|
||||
assert result["event_direction"] == "bearish"
|
||||
assert result["provenance"]["event_direction"] == "llm"
|
||||
assert result["rule_llm_conflict"] is False
|
||||
|
||||
def test_conflict_flags_rule_llm_conflict(self) -> None:
|
||||
"""rule_llm_conflict flag is set when both parsers disagree with high confidence."""
|
||||
llm_conflict = copy.deepcopy(_BASE_LLM)
|
||||
llm_conflict["event_type"] = "guidance_update" # conflicts with "earnings_release"
|
||||
# Both rule and LLM are confident
|
||||
rule_high = copy.deepcopy(_BASE_RULE)
|
||||
rule_high["confidence"]["event_type"] = 0.90
|
||||
llm_conflict["confidence"]["event_type"] = 0.90
|
||||
|
||||
result = merge(rule_high, llm_conflict)
|
||||
# Both confident → rule wins, but conflict flagged
|
||||
assert result["event_type"] == "earnings_release"
|
||||
assert result["rule_llm_conflict"] is True
|
||||
|
||||
def test_llm_none_all_provenance_is_rule(self) -> None:
|
||||
"""With no LLM output, all provenance should be 'rule'."""
|
||||
result = merge(_BASE_RULE, None)
|
||||
assert result["parser_kind"] == "merged"
|
||||
assert result["rule_llm_conflict"] is False
|
||||
for field, source in result["provenance"].items():
|
||||
assert source == "rule", f"Expected 'rule' for {field}, got {source}"
|
||||
|
||||
def test_risk_flags_are_ored(self) -> None:
|
||||
"""Risk flags combine with OR: if either parser flags it, it's flagged."""
|
||||
rule_no_oneoff = copy.deepcopy(_BASE_RULE)
|
||||
rule_no_oneoff["risk_flags"]["oneoff_item"] = False
|
||||
|
||||
llm_with_oneoff = copy.deepcopy(_BASE_LLM)
|
||||
llm_with_oneoff["risk_flags"]["oneoff_item"] = True
|
||||
|
||||
result = merge(rule_no_oneoff, llm_with_oneoff)
|
||||
assert result["risk_flags"]["oneoff_item"] is True
|
||||
assert result["provenance"]["risk_flags.oneoff_item"] == "llm"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestShouldQueueForReview:
|
||||
def test_no_queue_for_clean_record(self) -> None:
|
||||
"""High-confidence, no conflicts, no oneoff flags → no review needed."""
|
||||
clean = copy.deepcopy(_BASE_RULE)
|
||||
# Remove any flags that would trigger review
|
||||
clean["risk_flags"]["oneoff_item"] = False
|
||||
clean["risk_flags"]["non_gaap_heavy"] = False
|
||||
clean["rule_llm_conflict"] = False
|
||||
clean["provenance"] = {} # type: ignore[typeddict-unknown-key]
|
||||
clean["parser_kind"] = "merged"
|
||||
# event_direction is bullish (not unknown/mixed), confidence is 0.80
|
||||
should_queue, reasons = should_queue_for_review(clean)
|
||||
assert should_queue is False
|
||||
assert reasons == []
|
||||
|
||||
def test_low_confidence_triggers_review(self) -> None:
|
||||
low_conf = copy.deepcopy(_BASE_RULE)
|
||||
low_conf["confidence"]["overall"] = 0.45
|
||||
low_conf["rule_llm_conflict"] = False
|
||||
should_queue, reasons = should_queue_for_review(low_conf)
|
||||
assert should_queue is True
|
||||
assert "low_confidence" in reasons
|
||||
|
||||
def test_conflict_triggers_review(self) -> None:
|
||||
conflict = copy.deepcopy(_BASE_RULE)
|
||||
conflict["rule_llm_conflict"] = True
|
||||
should_queue, reasons = should_queue_for_review(conflict)
|
||||
assert should_queue is True
|
||||
assert "rule_llm_conflict" in reasons
|
||||
|
||||
def test_oneoff_triggers_review(self) -> None:
|
||||
oneoff = copy.deepcopy(_BASE_RULE)
|
||||
oneoff["risk_flags"]["oneoff_item"] = True
|
||||
oneoff["rule_llm_conflict"] = False
|
||||
should_queue, reasons = should_queue_for_review(oneoff)
|
||||
assert should_queue is True
|
||||
assert "oneoff_likely" in reasons
|
||||
@ -0,0 +1,112 @@
|
||||
"""Unit tests for review queue."""
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.review.queue import create_review_item, list_review_items, resolve_review_item
|
||||
|
||||
|
||||
def _make_mock_session() -> AsyncMock:
|
||||
session = AsyncMock()
|
||||
session.add = MagicMock()
|
||||
session.flush = AsyncMock()
|
||||
return session
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestCreateReviewItem:
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_review_item(self) -> None:
|
||||
"""Creating a review item for a new entity adds it to session."""
|
||||
session = _make_mock_session()
|
||||
|
||||
# No existing open item
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = None
|
||||
session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
await create_review_item(
|
||||
session=session,
|
||||
entity_type="parser_event",
|
||||
entity_id="EVT::test::001",
|
||||
priority="P1",
|
||||
reason_codes=["low_confidence"],
|
||||
snapshot_refs={"parse_id": 42},
|
||||
)
|
||||
|
||||
session.add.assert_called_once()
|
||||
session.flush.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deduplicate_open_items(self) -> None:
|
||||
"""Creating a review item for an already-open entity updates instead of creating."""
|
||||
# Simulate existing open item
|
||||
existing_item = MagicMock()
|
||||
existing_item.status = "open"
|
||||
existing_item.priority = "P2"
|
||||
existing_item.reason_codes = ["low_confidence"]
|
||||
existing_item.suggested_overrides = None
|
||||
|
||||
session = _make_mock_session()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = existing_item
|
||||
session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
await create_review_item(
|
||||
session=session,
|
||||
entity_type="parser_event",
|
||||
entity_id="EVT::test::001",
|
||||
priority="P0", # escalate from P2
|
||||
reason_codes=["rule_llm_conflict"],
|
||||
snapshot_refs={},
|
||||
)
|
||||
|
||||
# Should NOT call session.add (update existing instead)
|
||||
session.add.assert_not_called()
|
||||
# Priority should be escalated
|
||||
assert existing_item.priority == "P0"
|
||||
# reason_codes merged
|
||||
assert "rule_llm_conflict" in existing_item.reason_codes
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_review_item(self) -> None:
|
||||
"""Resolving an open item updates status, reviewer, resolution fields."""
|
||||
existing_item = MagicMock()
|
||||
existing_item.status = "open"
|
||||
review_id = uuid.uuid4()
|
||||
|
||||
session = _make_mock_session()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = existing_item
|
||||
session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
await resolve_review_item(
|
||||
session=session,
|
||||
review_id=review_id,
|
||||
reviewer_id="analyst_01",
|
||||
resolution_type="accepted",
|
||||
root_cause="data looks correct",
|
||||
notes="verified against source",
|
||||
)
|
||||
|
||||
assert existing_item.status == "resolved"
|
||||
assert existing_item.reviewer_id == "analyst_01"
|
||||
assert existing_item.resolution_type == "accepted"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_review_items_with_status_filter(self) -> None:
|
||||
"""list_review_items executes query with status filter and returns results."""
|
||||
session = _make_mock_session()
|
||||
mock_result = MagicMock()
|
||||
mock_scalars = MagicMock()
|
||||
mock_scalars.all.return_value = []
|
||||
mock_result.scalars.return_value = mock_scalars
|
||||
session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
items = await list_review_items(session, status="open")
|
||||
|
||||
assert items == []
|
||||
session.execute.assert_called_once()
|
||||
@ -0,0 +1,84 @@
|
||||
"""Unit tests for snapshot export."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from libs.export.snapshot_export import _temporal_split, export_dataset_snapshot
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestTemporalSplit:
|
||||
def test_split_proportions(self) -> None:
|
||||
"""Temporal split produces correct proportions for 100 rows."""
|
||||
rows = [{"event_date": "2026-01-01"} for _ in range(100)]
|
||||
splits = _temporal_split(rows, "temporal_70_15_15")
|
||||
assert len(splits["train"]) == 70
|
||||
assert len(splits["valid"]) == 15
|
||||
assert len(splits["test"]) == 15
|
||||
|
||||
def test_split_preserves_temporal_order(self) -> None:
|
||||
"""Train set contains earliest dates, test contains latest."""
|
||||
rows = [{"event_date": f"2026-{m:02d}-01"} for m in range(1, 13)]
|
||||
splits = _temporal_split(rows, "temporal_70_15_15")
|
||||
if splits["train"] and splits["test"]:
|
||||
assert splits["train"][-1]["event_date"] <= splits["test"][0]["event_date"]
|
||||
|
||||
def test_empty_rows_returns_empty_splits(self) -> None:
|
||||
splits = _temporal_split([], "temporal_70_15_15")
|
||||
assert splits == {"train": [], "valid": [], "test": []}
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestExportDatasetSnapshot:
|
||||
@pytest.mark.asyncio
|
||||
async def test_manifest_is_written(self) -> None:
|
||||
"""export_dataset_snapshot writes a manifest.json with expected fields."""
|
||||
# Mock DB session returning empty results (no feature+label pairs)
|
||||
mock_session = AsyncMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.all.return_value = [] # no rows
|
||||
mock_session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
manifest = await export_dataset_snapshot(
|
||||
session=mock_session,
|
||||
snapshot_id="test-snapshot-001",
|
||||
split_policy="temporal_70_15_15",
|
||||
output_dir=tmpdir,
|
||||
)
|
||||
|
||||
assert manifest["snapshot_id"] == "test-snapshot-001"
|
||||
assert "created_at_utc" in manifest
|
||||
assert "row_counts" in manifest
|
||||
assert manifest["split_policy"] == "temporal_70_15_15"
|
||||
assert manifest["total_rows"] == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parquet_files_created(self) -> None:
|
||||
"""Parquet files are created for each split partition."""
|
||||
mock_session = AsyncMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.all.return_value = []
|
||||
mock_session.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
await export_dataset_snapshot(
|
||||
session=mock_session,
|
||||
snapshot_id="test-parquet-002",
|
||||
split_policy="temporal_70_15_15",
|
||||
output_dir=tmpdir,
|
||||
)
|
||||
|
||||
snap_dir = Path(tmpdir) / "test-parquet-002"
|
||||
assert (snap_dir / "train.parquet").exists()
|
||||
assert (snap_dir / "valid.parquet").exists()
|
||||
assert (snap_dir / "test.parquet").exists()
|
||||
assert (snap_dir / "manifest.json").exists()
|
||||
|
||||
manifest_data = json.loads((snap_dir / "manifest.json").read_text())
|
||||
assert manifest_data["snapshot_id"] == "test-parquet-002"
|
||||
Loading…
Reference in New Issue