40 lines
1.9 KiB
Python
40 lines
1.9 KiB
Python
# planner.py
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_core.messages import SystemMessage, HumanMessage
|
||
import json
|
||
from typing import List, Dict
|
||
from config import MODEL_NAME
|
||
|
||
class PlannerAgent:
|
||
def __init__(self, llm=None):
|
||
self.llm = llm or ChatOllama(
|
||
model=MODEL_NAME,
|
||
temperature=0.0
|
||
)
|
||
|
||
def create_plan(self, task: str, tools_description: List[Dict]) -> List[Dict]:
|
||
system_prompt = """Ты - планировщик. Твоя задача - разбить запрос пользователя на последовательность шагов.
|
||
Каждый шаг - это вызов одного из доступных инструментов с конкретными аргументами.
|
||
Верни JSON массив шагов, где каждый шаг содержит:
|
||
- "tool": имя инструмента (строго из списка)
|
||
- "args": словарь аргументов для этого инструмента
|
||
- "description": краткое пояснение, зачем этот шаг
|
||
|
||
Инструменты и их аргументы:
|
||
{tools_description}
|
||
|
||
План должен быть линейным (шаг за шагом). Если нужны условия или циклы - разворачивай их в последовательность.
|
||
После выполнения всех шагов задача должна быть решена.
|
||
Верни ТОЛЬКО JSON массив, без лишнего текста."""
|
||
|
||
user_prompt = f"Задача: {task}"
|
||
tools_str = json.dumps(tools_description, indent=2)
|
||
|
||
messages = [
|
||
SystemMessage(content=system_prompt.format(tools_description=tools_str)),
|
||
HumanMessage(content=user_prompt)
|
||
]
|
||
response = self.llm.invoke(messages)
|
||
# парсим JSON
|
||
plan = json.loads(response.content)
|
||
return plan |