440 lines
20 KiB
Python
440 lines
20 KiB
Python
import os
|
||
import json
|
||
import re
|
||
import requests
|
||
from typing import Dict, Any, List, Optional
|
||
from datetime import datetime, timedelta
|
||
|
||
API_KEY = "AIzaSyBM4YX3AWQQOXgkg1L9WGD9_p1SZNkJMDw"
|
||
API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
|
||
|
||
# Базовый список доступных тегов
|
||
TAG_LIST = [
|
||
"Памятник", "Архитектура", "Мозаика", "Парк", "Озеро", "Ботанический сад", "Библиотека",
|
||
"Архитектура", "Площадь", "Набережная", "Канатная дорога",
|
||
"Сквер", "Музей", "Собор", "Церковь",
|
||
"Храм", "Театр", "Мост", "Вокзал", "Фонтан", "Лестница", "Кинотеатр",
|
||
"Дом Культуры", "Планетарий", "Галерея", "Ресторан", "Бар", "Бистро", "Кафе", "Кофейня"
|
||
]
|
||
|
||
|
||
def build_extraction_prompt(conversation_history: List[Dict[str, str]], tag_list: List[str]) -> str:
|
||
"""
|
||
Строим промпт для ИИ-парсера с историей диалога
|
||
"""
|
||
example_user_1 = "Я хочу культурно провести время, у меня 2 часа, я сейчас на улице Доброй около дома 1."
|
||
example_json_1 = {
|
||
"tags": [["Музей", "Галерея", "Театр", "ANY", 'MULTI']],
|
||
"user_location": "улица Добрая, дом 1",
|
||
"time": 120
|
||
}
|
||
example_user_2 = "Я хочу сходить в музей пройти по парку и где-нибудь поесть, у меня 3.5 часа, я на Коминтерна, 115"
|
||
example_json_2 = {
|
||
"tags": [["Музей", "Парк"], ["Ресторан", "Бистро", "Кафе", "ANY"]],
|
||
"user_location": "улица Коминтерна, дом 115",
|
||
"time": 210
|
||
}
|
||
example_user_3 = "Я хочу культурно провести время потом попить кофе, я свободен до 6 вечера, я сейчас на б-р 60 лет Октября, 9"
|
||
example_json_3 = {
|
||
"tags": [['Музей', 'Галерея', 'Театр', 'Планетарий', 'ANY', 'MULTI'], ["Кофейня"]],
|
||
"user_location": "бульвар 60 лет Октября, дом 9",
|
||
"time": 'until 18:00'
|
||
}
|
||
|
||
example_user_4 = "Я хочу обойти места в которых можно попить кофе и зайти в музей, у меня 4 часа, я на дальней 8"
|
||
example_json_4 = {
|
||
"tags": [['Кофейня', 'MULTI'],
|
||
["Музей"]],
|
||
"user_location": "улица Дальняя, дом 8",
|
||
"time": 240
|
||
}
|
||
|
||
example_user_5 = "Я нахожусь: казанская набережная около дома 5. Мне интересно: Хочу посмотреть на картины и прокатиться на канатной дороге, потом зайти поесть. У меня есть время: 3 часа"
|
||
example_json_5 = {
|
||
"tags": [['Галерея','Канатная дорога'], ['Ресторан', 'Бистро', 'Кафе', 'ANY']],
|
||
"user_location": "Казанская набережная, дом 5",
|
||
"time": 180
|
||
}
|
||
|
||
example_user_6 = "Я нахожусь: проспект Гагарина 42. Мне интересно: Я хочу посмотреть на интересные здания и поесть. У меня есть время: 2.5 часа"
|
||
example_json_6 = {
|
||
"tags": [['Архитектура','MULTI'], ['Ресторан', 'Бистро', 'Кафе', 'ANY']],
|
||
"user_location": "проспект Гагарина, дом 42",
|
||
"time": 150
|
||
}
|
||
|
||
instr = f"""
|
||
Ты — ИИ-парсер на русском языке.
|
||
Задача: из диалога с пользователем извлечь строго JSON c полями:
|
||
|
||
- tags: упорядоченный список релевантных тегов ТОЛЬКО из данного tag_list (не придумывай новых), теги должны быть сгруппированы по приоритету посещения
|
||
- user_location: строка с текущим местоположением пользователя (примерный адрес), должен содержать название улицы/площади/... Но не может быть ориентиром на объект без конкретного адреса или чем-то абстрактным: дома, у памятника Пушкину, у станции метро, на автобусной остановке... не подходят, но я около площади Революции подходит.
|
||
- time: целое число минут на прогулку (преобразуй выражения типа "2 часа" -> 120, "полчаса" -> 30 и т.п.). Пользователь может указать до скольки он свободен тогда используй фиксированный формат until <время в 24 формате HH:MM>. 'Я хочу погулять до 6 вечера' -> тогда в поле время 'until 18:00'.
|
||
|
||
Работа с тегами:
|
||
Группируй теги по ПРИОРИТЕТУ посещения:
|
||
- Если пользователь говорит "сначала музей, потом кофе" -> [['Музей'], ['Кофейня']]
|
||
- Если "хочу в музей и кофе" (без явного порядка) -> [['Музей', 'Кофейня']]
|
||
- Если неопределённые предпочтения "погулять по интересным местам", добавь ключевое слово 'ANY' к группе подходящих тегов.
|
||
- Если пользователь хочет посетить несколько мест одного типа, добавь ключевое слово 'MULTI' в соответсвующую группу. "Хочу походить по музеям, потом попить кофе." -> [['Музей','MULTI'], ['Кофейня']]
|
||
- Запрос пользователя может быть абстрактным "я хочу погулять по улице". Подбери релевантные теги на открытом воздухе (т.к. погулять можно интерпретировать как походить по улице не заходя куда-либо) например: [['Архитектура','Сквер','Площадь', 'ANY']]. Запрос может быть абстрактным, но не может быть полностью неопределенным. Примеры неопределенных ответов: "Хочу куда-нибудь сходить", "Я не знаю куда сходить", "Куда мне сходить?"
|
||
|
||
ИСПОЛЬЗУЙ ТОЛЬКО ТЕГИ ИЗ ДАННОГО СПИСКА!
|
||
|
||
Требования:
|
||
- Верни ТОЛЬКО валидный JSON без префиксов, комментариев и форматирования в кодовых блоках.
|
||
- Если в диалоге с пользователем нет данных для какого-либо поля, поставь null в этом поле.
|
||
- Не добавляй никаких дополнительных полей, только: tags, user_location, time.
|
||
- tags должны быть подмножеством из tag_list. Не включай нерелевантные теги.
|
||
|
||
tag_list = {json.dumps(tag_list, ensure_ascii=False, indent=0)}
|
||
|
||
Пример 1:
|
||
Пользователь: "{example_user_1}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_1, ensure_ascii=False, indent=2)}
|
||
|
||
Пример 2:
|
||
Пользователь: "{example_user_2}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_2, ensure_ascii=False, indent=2)}
|
||
|
||
Пример 3:
|
||
Пользователь: "{example_user_3}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_3, ensure_ascii=False, indent=2)}
|
||
|
||
Пример 4:
|
||
Пользователь: "{example_user_4}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_4, ensure_ascii=False, indent=2)}
|
||
|
||
Пример 5:
|
||
Пользователь: "{example_user_5}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_5, ensure_ascii=False, indent=2)}
|
||
|
||
Пример 6:
|
||
Пользователь: "{example_user_6}"
|
||
Ожидаемый JSON:
|
||
{json.dumps(example_json_6, ensure_ascii=False, indent=2)}
|
||
|
||
История диалога:
|
||
"""
|
||
|
||
for msg in conversation_history:
|
||
role = msg["role"]
|
||
content = msg["content"]
|
||
if role == "user":
|
||
instr += f"\nПользователь: \"{content}\""
|
||
elif role == "assistant":
|
||
instr += f"\nАссистент: \"{content}\""
|
||
|
||
instr += "\n\nТеперь извлеки JSON из диалога выше."
|
||
|
||
return instr.strip()
|
||
|
||
|
||
def build_conversational_prompt(conversation_history: List[Dict[str, str]], missing: List[str],
|
||
tag_list: List[str]) -> str:
|
||
"""
|
||
Строим промпт для генерации уточняющего вопроса
|
||
"""
|
||
missing_desc = {
|
||
"tags": "какие места/категории интересны",
|
||
"user_location": "текущее местоположение (адрес)",
|
||
"time": "время на прогулку"
|
||
}
|
||
|
||
missing_fields_str = ", ".join([missing_desc[f] for f in missing])
|
||
|
||
instr = f"""
|
||
Ты — вежливый ИИ-ассистент, который помогает пользователю спланировать прогулку.
|
||
Твоя задача — задать краткие вопросы, чтобы уточнить недостающую информацию.
|
||
|
||
Недостающие данные: {missing_fields_str}
|
||
|
||
Правила:
|
||
- Задай ТОЛЬКО простык, вежливый и краткий вопрос.
|
||
- НЕ давай рекомендаций, комментариев или предложений.
|
||
- НЕ предлагай варианты, кроме случаев когда нужно выбрать категории (tags).
|
||
- Если нужно уточнить категории (tags), предложи категории ИЗ ЭТОГО СПИСКА: {json.dumps(tag_list, ensure_ascii=False)}
|
||
- Верни ТОЛЬКО текст вопроса без дополнительных пояснений.
|
||
|
||
История диалога:
|
||
"""
|
||
|
||
for msg in conversation_history:
|
||
role = msg["role"]
|
||
content = msg["content"]
|
||
if role == "user":
|
||
instr += f"\nПользователь: {content}"
|
||
elif role == "assistant":
|
||
instr += f"\nАссистент: {content}"
|
||
|
||
instr += "\n\nТеперь задай уточняющий вопрос:"
|
||
|
||
return instr.strip()
|
||
|
||
|
||
def call_gemini(prompt: str) -> str:
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"x-goog-api-key": API_KEY,
|
||
}
|
||
|
||
payload = {
|
||
"contents": [
|
||
{
|
||
"role": "user",
|
||
"parts": [{"text": prompt}]
|
||
}
|
||
],
|
||
"generationConfig": {
|
||
"temperature": 0.5
|
||
}
|
||
}
|
||
|
||
resp = requests.post(API_URL, headers=headers, json=payload, timeout=30)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
# Извлекаем текст первого кандидата
|
||
text = ""
|
||
try:
|
||
text = data["candidates"][0]["content"]["parts"][0]["text"]
|
||
except Exception:
|
||
# Попытка fallback: иногда ответ может приходить иначе
|
||
text = json.dumps(data, ensure_ascii=False)
|
||
|
||
return text
|
||
|
||
|
||
def strip_code_fences(text: str) -> str:
|
||
s = text.strip()
|
||
# Убираем ограду `````` или ``````
|
||
if s.startswith("```"):
|
||
# срезаем первые ```
|
||
s = s[3:].lstrip()
|
||
# если указан язык (например, json), уберём первую строку
|
||
first_nl = s.find("\n")
|
||
if first_nl != -1:
|
||
lang = s[:first_nl].strip().lower()
|
||
# если это похоже на метку языка, отбрасываем её
|
||
if lang in ("json", "yaml", "yml", "xml", "markdown", "md", "txt"):
|
||
s = s[first_nl + 1:]
|
||
else:
|
||
# если метка не признана, всё равно продолжаем с текущим s
|
||
pass
|
||
# убираем завершающие ```
|
||
if s.rstrip().endswith("```"):
|
||
s = s.rstrip()[:-3]
|
||
|
||
return s.strip()
|
||
|
||
|
||
def extract_first_json(text: str):
|
||
# 1) снимаем ограду
|
||
s = strip_code_fences(text)
|
||
|
||
# 2) пробуем распарсить как есть
|
||
try:
|
||
return json.loads(s)
|
||
except Exception:
|
||
pass
|
||
|
||
# 3) ищем первый сбалансированный объект { ... } без рекурсивного regex
|
||
start = s.find("{")
|
||
while start != -1:
|
||
depth = 0
|
||
in_str = False
|
||
esc = False
|
||
for i in range(start, len(s)):
|
||
ch = s[i]
|
||
if in_str:
|
||
if esc:
|
||
esc = False
|
||
elif ch == "\\":
|
||
esc = True
|
||
elif ch == '"':
|
||
in_str = False
|
||
else:
|
||
if ch == '"':
|
||
in_str = True
|
||
elif ch == "{":
|
||
depth += 1
|
||
elif ch == "}":
|
||
depth -= 1
|
||
if depth == 0:
|
||
cand = s[start:i + 1]
|
||
try:
|
||
return json.loads(cand)
|
||
except Exception:
|
||
break
|
||
# ищем следующий '{', если текущий блок не распарсился
|
||
start = s.find("{", start + 1)
|
||
|
||
return None
|
||
|
||
|
||
def is_complete(result: Dict[str, Any]) -> bool:
|
||
if not isinstance(result, dict):
|
||
return False
|
||
|
||
# Проверяем наличие ключей
|
||
for key in ["tags", "user_location", "time"]:
|
||
if key not in result:
|
||
return False
|
||
|
||
# Проверяем содержимое
|
||
tags_ok = isinstance(result["tags"], list) and len(result["tags"]) > 0
|
||
loc_ok = isinstance(result["user_location"], str) and len(result["user_location"].strip()) > 0
|
||
|
||
# time может быть int или строка формата "until HH:MM"
|
||
time_val = result["time"]
|
||
if isinstance(time_val, int) and time_val > 0:
|
||
time_ok = True
|
||
elif isinstance(time_val, str) and time_val.startswith("until "):
|
||
time_ok = True
|
||
else:
|
||
time_ok = False
|
||
|
||
return tags_ok and loc_ok and time_ok
|
||
|
||
|
||
def missing_fields(result: Dict[str, Any]) -> List[str]:
|
||
missing = []
|
||
|
||
if not isinstance(result, dict):
|
||
return ["tags", "user_location", "time"]
|
||
|
||
if "tags" not in result or not isinstance(result["tags"], list) or len(result["tags"]) == 0:
|
||
missing.append("tags")
|
||
|
||
if "user_location" not in result or not isinstance(result["user_location"], str) or not result[
|
||
"user_location"].strip():
|
||
missing.append("user_location")
|
||
|
||
time_val = result.get("time")
|
||
if isinstance(time_val, int) and time_val > 0:
|
||
pass # OK
|
||
elif isinstance(time_val, str) and time_val.startswith("until "):
|
||
pass # OK
|
||
else:
|
||
missing.append("time")
|
||
|
||
return missing
|
||
|
||
|
||
def convert_time_to_minutes(time_val, current_time: datetime) -> int:
|
||
"""
|
||
Конвертирует значение времени в минуты.
|
||
Если time_val — это строка формата "until HH:MM", вычисляет разницу от текущего времени.
|
||
"""
|
||
if isinstance(time_val, int):
|
||
return time_val
|
||
|
||
if isinstance(time_val, str) and time_val.startswith("until "):
|
||
time_str = time_val.replace("until ", "").strip()
|
||
try:
|
||
target_time = datetime.strptime(time_str, "%H:%M")
|
||
# Устанавливаем дату как текущую
|
||
target_datetime = current_time.replace(
|
||
hour=target_time.hour,
|
||
minute=target_time.minute,
|
||
second=0,
|
||
microsecond=0
|
||
)
|
||
|
||
# Если целевое время уже прошло сегодня, предполагаем завтра
|
||
if target_datetime <= current_time:
|
||
target_datetime += timedelta(days=1)
|
||
|
||
delta = target_datetime - current_time
|
||
minutes = int(delta.total_seconds() / 60)
|
||
return minutes
|
||
except ValueError:
|
||
return None
|
||
|
||
return None
|
||
|
||
|
||
def extract_with_ai(conversation_history: List[Dict[str, str]]) -> Dict[str, Any]:
|
||
prompt = build_extraction_prompt(conversation_history, TAG_LIST)
|
||
model_text = call_gemini(prompt)
|
||
parsed = extract_first_json(model_text) or {}
|
||
return parsed
|
||
|
||
|
||
def ask_ai_for_clarification(conversation_history: List[Dict[str, str]], missing: List[str]) -> str:
|
||
"""
|
||
Генерируем уточняющий вопрос через ИИ
|
||
"""
|
||
prompt = build_conversational_prompt(conversation_history, missing, TAG_LIST)
|
||
question = call_gemini(prompt)
|
||
return question.strip()
|
||
|
||
|
||
def parse():
|
||
if not API_KEY or API_KEY.startswith("<"):
|
||
raise RuntimeError("Установите переменную окружения GEMINI_API_KEY с вашим ключом API.")
|
||
|
||
print(
|
||
"Опишите ваш запрос (например: 'Я хочу культурно провести время потом попить кофе, у меня 2 часа, я сейчас на улице Дальней около дома 8.')")
|
||
|
||
conversation_history = []
|
||
|
||
# Первое сообщение от пользователя
|
||
user_input = input("> ").strip()
|
||
conversation_history.append({"role": "user", "content": user_input})
|
||
|
||
# Текущее время (Москва, UTC+3)
|
||
current_time = datetime.now()
|
||
|
||
# Первая попытка извлечения
|
||
result = extract_with_ai(conversation_history)
|
||
|
||
# Итеративный диалог
|
||
max_tries = 10
|
||
tries = 0
|
||
|
||
while (not is_complete(result)) and tries < max_tries:
|
||
missing = missing_fields(result)
|
||
|
||
if not missing:
|
||
break
|
||
|
||
# Генерируем уточняющий вопрос через ИИ
|
||
ai_question = ask_ai_for_clarification(conversation_history, missing)
|
||
print(ai_question)
|
||
|
||
conversation_history.append({"role": "assistant", "content": ai_question})
|
||
|
||
# Получаем ответ пользователя
|
||
user_response = input("> ").strip()
|
||
conversation_history.append({"role": "user", "content": user_response})
|
||
|
||
# Повторно извлекаем данные с учётом новой информации
|
||
result = extract_with_ai(conversation_history)
|
||
tries += 1
|
||
|
||
# Конвертируем время если нужно
|
||
if "time" in result and isinstance(result["time"], str):
|
||
minutes = convert_time_to_minutes(result["time"], current_time)
|
||
if minutes is not None:
|
||
result["time"] = minutes
|
||
|
||
# Печатаем финальный JSON
|
||
out = {
|
||
"tags": result.get("tags", None),
|
||
"user_location": result.get("user_location", None),
|
||
"time": result.get("time", None),
|
||
}
|
||
return json.dumps(out, ensure_ascii=False, indent=2),conversation_history
|
||
|
||
|
||
if __name__ == "__main__":
|
||
out = parse()
|
||
print(out)
|
||
|