portfolio_manager.py

"""Portfolio rebalancing helpers."""

from __future__ import annotations

from typing import Iterable, List, Sequence, Tuple, TypedDict


class Position(TypedDict):
    """Minimal portfolio position description."""

    ticker: str
    quantity: float


class AnalysisResult(TypedDict, total=False):
    """Subset of evaluation data used for trade planning."""

    ticker: str
    price: float
    score: float
    llm_decision: str


TradePlan = Tuple[List[Tuple[str, float, float]], List[Tuple[str, int, float]]]


def _find_analysis(analyses: Iterable[AnalysisResult], ticker: str) -> AnalysisResult | None:
    """Return the analysis matching ``ticker`` if present."""

    return next((item for item in analyses if item.get("ticker") == ticker), None)


def plan_trades(
    portfolio: Sequence[Position],
    analyses: Sequence[AnalysisResult],
    base_budget: float,
) -> TradePlan:
    """Generate sell and buy plans using LLM-confirmed decisions."""

    sell_plan: List[Tuple[str, float, float]] = []
    for position in portfolio:
        analysis = _find_analysis(analyses, position["ticker"])
        if (
            analysis
            and analysis.get("llm_decision") == "SELL"
            and position["quantity"] > 0
        ):
            sell_plan.append(
                (position["ticker"], position["quantity"], analysis["price"])
            )

    sell_proceeds = sum(quantity * price for _, quantity, price in sell_plan)
    total_budget = base_budget + sell_proceeds

    buy_candidates = sorted(
        (item for item in analyses if item.get("llm_decision") == "BUY"),
        key=lambda item: item["score"],
        reverse=True,
    )

    total_score = sum(candidate["score"] for candidate in buy_candidates)
    buy_plan: List[Tuple[str, int, float]] = []

    for candidate in buy_candidates:
        if total_score == 0:
            continue
        weight = candidate["score"] / total_score
        allocated_budget = weight * total_budget
        max_quantity = int(allocated_budget // candidate["price"])
        if max_quantity <= 0:
            continue
        buy_plan.append((candidate["ticker"], max_quantity, candidate["price"]))

    return sell_plan, buy_plan