170 lines
5.0 KiB
Python
170 lines
5.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""News service client for news enrichment operations."""
|
|
|
|
import httpx
|
|
|
|
|
|
class NewsServiceClient:
|
|
"""Async client for the News Service API."""
|
|
|
|
def __init__(self, base_url: str = "http://localhost:8002"):
|
|
"""Initialize the client with a base URL.
|
|
|
|
Args:
|
|
base_url: Base URL for the news service API.
|
|
"""
|
|
self.base_url = base_url.rstrip("/")
|
|
self._client: httpx.AsyncClient | None = None
|
|
|
|
async def __aenter__(self) -> "NewsServiceClient":
|
|
self._client = httpx.AsyncClient(
|
|
base_url=self.base_url,
|
|
timeout=90.0,
|
|
limits=httpx.Limits(max_connections=100, max_keepalive_connections=20)
|
|
)
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
if self._client:
|
|
await self._client.aclose()
|
|
|
|
async def get_enriched_news(
|
|
self,
|
|
ticker: str,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
limit: int | None = None,
|
|
) -> dict:
|
|
"""Get enriched news for a ticker.
|
|
|
|
Args:
|
|
ticker: Stock ticker symbol.
|
|
start_date: Start date (YYYY-MM-DD).
|
|
end_date: End date (YYYY-MM-DD).
|
|
|
|
Returns:
|
|
Dictionary with enriched news data.
|
|
"""
|
|
params = {"ticker": ticker}
|
|
if start_date:
|
|
params["start_date"] = start_date
|
|
if end_date:
|
|
params["end_date"] = end_date
|
|
if limit is not None:
|
|
params["limit"] = limit
|
|
response = await self._client.get("/api/enriched-news", params=params)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_news_for_date(
|
|
self,
|
|
ticker: str,
|
|
date: str,
|
|
limit: int = 20,
|
|
) -> dict:
|
|
"""Get enriched news rows for a specific trade date."""
|
|
response = await self._client.get(
|
|
"/api/news-for-date",
|
|
params={"ticker": ticker, "date": date, "limit": limit},
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_news_timeline(
|
|
self,
|
|
ticker: str,
|
|
start_date: str,
|
|
end_date: str,
|
|
) -> dict:
|
|
"""Get aggregated news timeline for a ticker."""
|
|
response = await self._client.get(
|
|
"/api/news-timeline",
|
|
params={
|
|
"ticker": ticker,
|
|
"start_date": start_date,
|
|
"end_date": end_date,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_similar_days(
|
|
self,
|
|
ticker: str,
|
|
date: str,
|
|
n_similar: int = 5,
|
|
) -> dict:
|
|
"""Get similar trading days based on price patterns.
|
|
|
|
Args:
|
|
ticker: Stock ticker symbol.
|
|
date: Reference date (YYYY-MM-DD).
|
|
n_similar: Number of similar days to return.
|
|
|
|
Returns:
|
|
Dictionary with similar day data.
|
|
"""
|
|
params = {"ticker": ticker, "date": date, "n_similar": n_similar}
|
|
response = await self._client.get("/api/similar-days", params=params)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_story(self, ticker: str, as_of_date: str) -> dict:
|
|
"""Get or build a ticker story as of one date.
|
|
|
|
Args:
|
|
ticker: Stock ticker symbol.
|
|
as_of_date: Story date.
|
|
|
|
Returns:
|
|
Dictionary with story data.
|
|
"""
|
|
response = await self._client.get(
|
|
f"/api/stories/{ticker}",
|
|
params={"as_of_date": as_of_date},
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_categories(
|
|
self,
|
|
ticker: str,
|
|
start_date: str | None = None,
|
|
end_date: str | None = None,
|
|
limit: int = 200,
|
|
) -> dict:
|
|
"""Get categories for a ticker window.
|
|
|
|
Returns:
|
|
Dictionary with available categories.
|
|
"""
|
|
params = {"ticker": ticker, "limit": limit}
|
|
if start_date:
|
|
params["start_date"] = start_date
|
|
if end_date:
|
|
params["end_date"] = end_date
|
|
response = await self._client.get("/api/categories", params=params)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def get_range_explain(
|
|
self,
|
|
ticker: str,
|
|
start_date: str,
|
|
end_date: str,
|
|
article_ids: list[str] | None = None,
|
|
limit: int = 100,
|
|
) -> dict:
|
|
"""Get a range explanation for a ticker window."""
|
|
params: list[tuple[str, str | int]] = [
|
|
("ticker", ticker),
|
|
("start_date", start_date),
|
|
("end_date", end_date),
|
|
("limit", limit),
|
|
]
|
|
for article_id in article_ids or []:
|
|
params.append(("article_ids", article_id))
|
|
response = await self._client.get("/api/range-explain", params=params)
|
|
response.raise_for_status()
|
|
return response.json()
|