214 lines
7.2 KiB
Python
214 lines
7.2 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from datetime import date
|
||
|
|
from typing import Any, Optional
|
||
|
|
|
||
|
|
from config import settings
|
||
|
|
from services.distributor import (
|
||
|
|
DayEntry,
|
||
|
|
build_plan_for_dates,
|
||
|
|
build_plan_for_day_list,
|
||
|
|
entries_from_last_state,
|
||
|
|
)
|
||
|
|
from services.gitlab_client import GitLabClient, GitLabTask
|
||
|
|
from services.grok_client import GrokClient
|
||
|
|
from services.holidays_util import iter_workdays
|
||
|
|
from services.meeting_planner import merge_meetings_into_plan
|
||
|
|
from services.state_store import get_last_entries
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class Planner:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.gitlab = GitLabClient()
|
||
|
|
self.grok = GrokClient()
|
||
|
|
|
||
|
|
async def build_plan(
|
||
|
|
self,
|
||
|
|
start: date,
|
||
|
|
end: date,
|
||
|
|
*,
|
||
|
|
use_llm: Optional[bool] = None,
|
||
|
|
use_last_if_empty: bool = True,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
workdays = iter_workdays(start, end)
|
||
|
|
return await self._build_plan_for_days(
|
||
|
|
workdays,
|
||
|
|
use_llm=use_llm,
|
||
|
|
use_last_if_empty=use_last_if_empty,
|
||
|
|
rule_builder=lambda tasks, _: build_plan_for_dates(tasks, start, end),
|
||
|
|
)
|
||
|
|
|
||
|
|
async def build_plan_for_days(
|
||
|
|
self,
|
||
|
|
days: list[date],
|
||
|
|
*,
|
||
|
|
use_llm: Optional[bool] = None,
|
||
|
|
use_last_if_empty: bool = True,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
sorted_days = sorted(set(days))
|
||
|
|
return await self._build_plan_for_days(
|
||
|
|
sorted_days,
|
||
|
|
use_llm=use_llm,
|
||
|
|
use_last_if_empty=use_last_if_empty,
|
||
|
|
rule_builder=lambda tasks, workdays: build_plan_for_day_list(tasks, workdays),
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _build_plan_for_days(
|
||
|
|
self,
|
||
|
|
workdays: list[date],
|
||
|
|
*,
|
||
|
|
use_llm: Optional[bool],
|
||
|
|
use_last_if_empty: bool,
|
||
|
|
rule_builder,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
gitlab_tasks = await self.gitlab.get_assigned_tasks()
|
||
|
|
|
||
|
|
if gitlab_tasks:
|
||
|
|
use_llm = settings.use_llm if use_llm is None else use_llm
|
||
|
|
if use_llm and self.grok.enabled:
|
||
|
|
try:
|
||
|
|
plan, analysis = await self.grok.plan_pts_entries(
|
||
|
|
gitlab_tasks,
|
||
|
|
workdays,
|
||
|
|
pts_project_name=settings.pts_project_name,
|
||
|
|
pts_task_name=settings.pts_default_task_name,
|
||
|
|
max_hours_per_day=settings.max_hours_per_day,
|
||
|
|
hour_step=settings.hour_step,
|
||
|
|
)
|
||
|
|
plan, analysis = await self._apply_variation(
|
||
|
|
plan,
|
||
|
|
gitlab_tasks,
|
||
|
|
analysis,
|
||
|
|
source="llm",
|
||
|
|
)
|
||
|
|
return await self._finalize_plan(
|
||
|
|
plan,
|
||
|
|
source="llm",
|
||
|
|
gitlab_tasks=gitlab_tasks,
|
||
|
|
workdays=workdays,
|
||
|
|
analysis=analysis,
|
||
|
|
)
|
||
|
|
except Exception as exc: # noqa: BLE001
|
||
|
|
logger.warning("LLM planning failed, fallback to rules: %s", exc)
|
||
|
|
|
||
|
|
plan = rule_builder(gitlab_tasks, workdays)
|
||
|
|
plan, analysis = await self._apply_variation(
|
||
|
|
plan,
|
||
|
|
gitlab_tasks,
|
||
|
|
{"fallback_reason": "rule_based_distribution"},
|
||
|
|
source="rule_based",
|
||
|
|
)
|
||
|
|
return await self._finalize_plan(
|
||
|
|
plan,
|
||
|
|
source="rule_based",
|
||
|
|
gitlab_tasks=gitlab_tasks,
|
||
|
|
workdays=workdays,
|
||
|
|
analysis=analysis,
|
||
|
|
)
|
||
|
|
|
||
|
|
if use_last_if_empty:
|
||
|
|
plan = entries_from_last_state(get_last_entries(), workdays)
|
||
|
|
plan, analysis = await self._apply_variation(
|
||
|
|
plan,
|
||
|
|
[],
|
||
|
|
{},
|
||
|
|
source="last_entries",
|
||
|
|
)
|
||
|
|
return await self._finalize_plan(
|
||
|
|
plan,
|
||
|
|
source="last_entries",
|
||
|
|
gitlab_tasks=[],
|
||
|
|
workdays=workdays,
|
||
|
|
analysis=analysis,
|
||
|
|
)
|
||
|
|
|
||
|
|
return await self._finalize_plan(
|
||
|
|
{d.isoformat(): [] for d in workdays},
|
||
|
|
source="empty",
|
||
|
|
gitlab_tasks=[],
|
||
|
|
workdays=workdays,
|
||
|
|
analysis={},
|
||
|
|
)
|
||
|
|
|
||
|
|
def _has_duplicate_descriptions(self, plan: dict[str, list[DayEntry]]) -> bool:
|
||
|
|
descriptions = [entry.description for entries in plan.values() for entry in entries]
|
||
|
|
return len(descriptions) != len(set(descriptions))
|
||
|
|
|
||
|
|
async def _apply_variation(
|
||
|
|
self,
|
||
|
|
plan: dict[str, list[DayEntry]],
|
||
|
|
tasks: list[GitLabTask],
|
||
|
|
analysis: dict[str, Any],
|
||
|
|
*,
|
||
|
|
source: str,
|
||
|
|
) -> tuple[dict[str, list[DayEntry]], dict[str, Any]]:
|
||
|
|
if not settings.use_description_variation or not self.grok.enabled:
|
||
|
|
return plan, analysis
|
||
|
|
|
||
|
|
entry_count = sum(len(entries) for entries in plan.values())
|
||
|
|
if entry_count <= 1:
|
||
|
|
return plan, {**analysis, "variation": {"source": "variation_skipped", "reason": "single_entry"}}
|
||
|
|
|
||
|
|
if source == "llm" and not self._has_duplicate_descriptions(plan):
|
||
|
|
return plan, {**analysis, "variation": {"source": "variation_skipped", "reason": "llm_unique"}}
|
||
|
|
|
||
|
|
try:
|
||
|
|
plan, variation = await self.grok.diversify_descriptions(plan, tasks)
|
||
|
|
return plan, {**analysis, "variation": variation}
|
||
|
|
except Exception as exc: # noqa: BLE001
|
||
|
|
logger.warning("Description variation failed: %s", exc)
|
||
|
|
return plan, {**analysis, "variation": {"source": "variation_failed", "error": str(exc)}}
|
||
|
|
|
||
|
|
async def _finalize_plan(
|
||
|
|
self,
|
||
|
|
plan: dict[str, list[DayEntry]],
|
||
|
|
*,
|
||
|
|
source: str,
|
||
|
|
gitlab_tasks: list[GitLabTask],
|
||
|
|
workdays: list[date],
|
||
|
|
analysis: dict[str, Any],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
workday_keys = [d.isoformat() for d in workdays]
|
||
|
|
plan, meeting_meta = await merge_meetings_into_plan(plan, workday_keys)
|
||
|
|
if meeting_meta:
|
||
|
|
analysis = {**analysis, "meetings": meeting_meta}
|
||
|
|
return self._wrap_plan(
|
||
|
|
plan,
|
||
|
|
source=source,
|
||
|
|
gitlab_tasks=gitlab_tasks,
|
||
|
|
workdays=workdays,
|
||
|
|
analysis=analysis,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _wrap_plan(
|
||
|
|
self,
|
||
|
|
plan: dict[str, list[DayEntry]],
|
||
|
|
*,
|
||
|
|
source: str,
|
||
|
|
gitlab_tasks: list[GitLabTask],
|
||
|
|
workdays: list[date],
|
||
|
|
analysis: dict[str, Any],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
return {
|
||
|
|
"source": source,
|
||
|
|
"gitlab_task_count": len(gitlab_tasks),
|
||
|
|
"workdays": [d.isoformat() for d in workdays],
|
||
|
|
"llm_enabled": self.grok.enabled,
|
||
|
|
"analysis": analysis,
|
||
|
|
"plan": {
|
||
|
|
day: [
|
||
|
|
{
|
||
|
|
"description": entry.description,
|
||
|
|
"take_hours": entry.take_hours,
|
||
|
|
"gitlab_task_id": entry.gitlab_task_id,
|
||
|
|
"pts_task_name": entry.pts_task_name,
|
||
|
|
}
|
||
|
|
for entry in entries
|
||
|
|
]
|
||
|
|
for day, entries in plan.items()
|
||
|
|
},
|
||
|
|
"_plan_entries": plan,
|
||
|
|
}
|