roberto
8bfce9da00
- Replaced direct instantiation of ChatOpenAI with a centralized get_llm function in CheckpointAgent, NoteAgent, ProjectAgent, and TaskAgent. - Introduced a new llm.py module to handle LLM model instantiation and API key management. - Updated settings.py to include LLM_MODEL and LLM_ROUTER_MODEL configurations. - Modified orchestrator.py to use get_router_llm for intent classification. - Updated requirements.txt to include litellm for LLM management. - Adjusted tests to mock get_llm instead of ChatOpenAI directly.
122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
"""Checkpoint agent — project milestone management (list, create, update, delete)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from langchain_core.tools import tool
|
|
|
|
from app.core.agent_registry import ChatAgent, registry
|
|
from app.core.llm import get_llm
|
|
|
|
_SYSTEM_PROMPT = (
|
|
"You are a project checkpoint assistant. Checkpoints are milestone dates that\n"
|
|
"track progress on a project — they are not calendar events.\n\n"
|
|
"Rules:\n"
|
|
" - project_id is REQUIRED for every create; confirm with the user if unknown\n"
|
|
" - date is a Unix timestamp in milliseconds; convert human-readable dates\n"
|
|
" - is_ai_suggested: 1 when proactively proposing a checkpoint, 0 otherwise\n"
|
|
" - is_approved: 0 until the user explicitly confirms; then 1\n"
|
|
" - For update_checkpoint, use -1 for integer fields you do not want to change\n"
|
|
" - Listing without a project_id returns all checkpoints across projects\n"
|
|
" - Always echo the title and formatted date in your confirmation."
|
|
)
|
|
|
|
|
|
@tool
|
|
async def list_checkpoints(project_id: str = "") -> str:
|
|
"""List checkpoints. Provide project_id to scope to a specific project."""
|
|
return json.dumps({
|
|
"action": "list",
|
|
"table": "checkpoints",
|
|
"filters": {"projectId": project_id or None},
|
|
})
|
|
|
|
|
|
@tool
|
|
async def create_checkpoint(
|
|
project_id: str,
|
|
title: str,
|
|
date: int,
|
|
is_ai_suggested: int = 0,
|
|
is_approved: int = 0,
|
|
) -> str:
|
|
"""Create a project checkpoint (milestone).
|
|
project_id: REQUIRED UUID of the parent project
|
|
title: descriptive name for the milestone
|
|
date: Unix timestamp in milliseconds
|
|
is_ai_suggested: 1 if proactively suggested, 0 if user-requested
|
|
is_approved: 0 until the user confirms
|
|
"""
|
|
return json.dumps({
|
|
"action": "create_record",
|
|
"table": "checkpoints",
|
|
"data": {
|
|
"projectId": project_id,
|
|
"title": title,
|
|
"date": date,
|
|
"isAiSuggested": is_ai_suggested,
|
|
"isApproved": is_approved,
|
|
},
|
|
})
|
|
|
|
|
|
@tool
|
|
async def update_checkpoint(
|
|
checkpoint_id: str,
|
|
title: str = "",
|
|
date: int = -1,
|
|
is_approved: int = -1,
|
|
) -> str:
|
|
"""Update a checkpoint. Only pass fields that should change.
|
|
checkpoint_id: UUID of the checkpoint (required)
|
|
date: -1 means unchanged; any other value sets the new date (ms timestamp)
|
|
is_approved: -1 means unchanged; 0 or 1 sets the approval state
|
|
"""
|
|
updates: dict[str, Any] = {}
|
|
if title:
|
|
updates["title"] = title
|
|
if date != -1:
|
|
updates["date"] = date
|
|
if is_approved != -1:
|
|
updates["isApproved"] = is_approved
|
|
return json.dumps({
|
|
"action": "update_record",
|
|
"table": "checkpoints",
|
|
"data": {"id": checkpoint_id, "updates": updates},
|
|
})
|
|
|
|
|
|
@tool
|
|
async def delete_checkpoint(checkpoint_id: str) -> str:
|
|
"""Delete a checkpoint permanently by its UUID."""
|
|
return json.dumps({
|
|
"action": "delete_record",
|
|
"table": "checkpoints",
|
|
"data": {"id": checkpoint_id},
|
|
})
|
|
|
|
|
|
@registry.register
|
|
class CheckpointAgent(ChatAgent):
|
|
def get_name(self) -> str:
|
|
return "checkpoint_agent"
|
|
|
|
def get_description(self) -> str:
|
|
return "Manages project checkpoints (milestones): list, create, update, delete"
|
|
|
|
def get_tools(self) -> list[Any]:
|
|
return [list_checkpoints, create_checkpoint, update_checkpoint, delete_checkpoint]
|
|
|
|
async def handle(self, query: str, context: dict[str, Any]) -> str:
|
|
llm = get_llm()
|
|
messages = [
|
|
SystemMessage(content=_SYSTEM_PROMPT),
|
|
HumanMessage(
|
|
content=f"User query: {query}\nContext: {json.dumps(context)[:1000]}"
|
|
),
|
|
]
|
|
return await self._tool_loop(llm, messages, self.get_tools())
|