Загрузить файлы в «/»
This commit is contained in:
389
new_parser_test.py
Normal file
389
new_parser_test.py
Normal file
@@ -0,0 +1,389 @@
|
||||
import os
|
||||
import json
|
||||
import re
|
||||
import requests
|
||||
from typing import Dict, Any, List, Optional
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
API_KEY = "AIzaSyBXGBGH5NDY8L_jVmq2zb4i8xYEV2qN-48"
|
||||
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"]],
|
||||
"user_location": "улица Добрая, дом 1",
|
||||
"time": 120
|
||||
}
|
||||
example_user_2 = "Я хочу сходить в музей пройти по парку и где-нибудь поесть, у меня 3.5 часа, я на Коминтерна, 115"
|
||||
example_json_2 = {
|
||||
"tags": [["Музей", "Парк"], ["Ресторан", "Бистро", "Кафе", "ANY"]],
|
||||
"user_location": "улица Коминтерна, дом 115",
|
||||
"time": 120
|
||||
}
|
||||
instr = f"""
|
||||
Ты — ИИ-парсер на русском языке.
|
||||
Задача: из диалога с пользователем извлечь строго JSON c полями:
|
||||
|
||||
- tags: упорядоченный список релевантных тегов ТОЛЬКО из данного tag_list (не придумывай новых), теги должны быть сгруппированы по приоритету посещения
|
||||
- user_location: строка с текущим местоположением пользователя (примерный адрес), должен содержать название улицы/площади/... Но не может быть ориентиром на объект без конкретного адреса или чем-то абстрактным: дома, у памятника Пушкину, у станции метро, на автобусной остановке... не подходят, но я около площади Революции подходит.
|
||||
- time: целое число минут на прогулку (преобразуй выражения типа "2 часа" -> 120, "полчаса" -> 30 и т.п.). Пользователь может указать до скольки он свободен тогда используй фиксированный формат until <время в 24 формате HH:MM>. 'Я хочу погулять до 6 вечера' -> тогда в поле время 'until 18:00'.
|
||||
|
||||
Работа с тегами:
|
||||
Группируй теги по ПРИОРИТЕТУ посещения:
|
||||
- Если пользователь говорит "сначала музей, потом кофе" -> [['Музей'], ['Кофейня']]
|
||||
- Если "хочу в музей и кофе" (без явного порядка) -> [['Музей', 'Кофейня']]
|
||||
- Если неопределённые предпочтения "погулять по интересным местам", добавь ключевое слово 'ANY' к группе подходящих тегов.
|
||||
- Запрос пользователя может быть абстрактным "я хочу погулять по улице". Подбери релевантные теги на открытом воздухе (т.к. погулять можно интерпретировать как походить по улице не заходя куда-либо) например: [['Архитектура','Сквер','Площадь', 'ANY']]. Запрос может быть абстрактным, но не может быть полностью неопределенным. Примеры неопределенных ответов: "Хочу куда-нибудь сходить", "Я не знаю куда сходить", "Куда мне сходить?"
|
||||
- Указывай все теги из tag_list которые явно подходят!!! Примеры могут быть неполными!
|
||||
|
||||
ИСПОЛЬЗУЙ ТОЛЬКО ТЕГИ ИЗ ДАННОГО СПИСКА!
|
||||
|
||||
Требования:
|
||||
- Верни ТОЛЬКО валидный 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)}
|
||||
|
||||
|
||||
История диалога:
|
||||
"""
|
||||
|
||||
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}]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
out = parse()
|
||||
print(out)
|
||||
|
||||
Reference in New Issue
Block a user