import os import json import zipfile import json_repair from docxtpl import DocxTemplate from openai import OpenAI from datetime import datetime import chromadb from sentence_transformers import SentenceTransformer # --- КОНФИГУРАЦИЯ --- API_KEY = os.getenv("DEEPSEEK_API_KEY") # Или вставь свой ключ жестко "sk-..." BASE_URL = "https://api.deepseek.com" TEMPLATES_DIR = "tagged_templates" DOWNLOADS_DIR = "downloads" REGISTRY_FILE = "templates_registry.json" TAGS_DB_FILE = "tags_db.json" DB_PATH = "./legal_db" ZIP_PATH = "./legal_db.zip" # -------------------- client = OpenAI(api_key=API_KEY, base_url=BASE_URL) # ============================================================================== # 1. ИНИЦИАЛИЗАЦИЯ # ============================================================================== collection = None encoder = None if not os.path.exists(DB_PATH) and os.path.exists(ZIP_PATH): try: with zipfile.ZipFile(ZIP_PATH, 'r') as z: z.extractall(".") except: pass try: encoder = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2') chroma_client = chromadb.PersistentClient(path=DB_PATH) collection = chroma_client.get_collection(name="laws") except: pass try: registry = json.load(open(REGISTRY_FILE, "r", encoding="utf-8")) tags_db = json.load(open(TAGS_DB_FILE, "r", encoding="utf-8")) clean_tags_db = {k: v for k, v in tags_db.items() if not k.startswith("_")} except: registry = [] clean_tags_db = {} # ============================================================================== # 2. ЛОГИКА # ============================================================================== async def select_best_template(user_query): """ Гибридный поиск: Сначала ИИ, потом перебор ключевых слов. """ # 1. Попытка через ИИ docs_list = "" for idx, item in enumerate(registry): # Ограничиваем длину списка, чтобы не забить контекст docs_list += f"{idx}. {item['filename']} ({item['description'][:100]}...)\n" system_prompt = f""" Ты библиотекарь. Твоя цель - найти имя файла из списка, которое лучше всего подходит к запросу. СПИСОК: {docs_list} ВЕРНИ ТОЛЬКО JSON: {{"filename": "имя_файла.docx"}} ЕСЛИ НИЧЕГО НЕ ПОДХОДИТ, ВЕРНИ: {{"filename": null}} """ try: response = client.chat.completions.create( model="deepseek-chat", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_query} ], response_format={"type": "json_object"}, temperature=0.0 ) result = json_repair.loads(response.choices[0].message.content) ai_choice = result.get("filename") if ai_choice: return ai_choice except Exception as e: print(f"⚠️ ИИ поиск не удался: {e}") # 2. ФОЛЛБЭК: Тупой поиск по словам (Если ИИ не справился) print("⚠️ ИИ не нашел шаблон. Запускаю поиск по ключевым словам...") query_words = user_query.lower().replace(",", "").split() # Убираем мусорные слова stop_words = ["хочу", "мне", "нужен", "составь", "сделай", "договор", "бланк", "образец"] keywords = [w for w in query_words if w not in stop_words and len(w) > 3] if not keywords: return None best_match = None max_hits = 0 for item in registry: hits = 0 fname = item['filename'].lower() desc = item.get('description', '').lower() for kw in keywords: # Ищем совпадения корней (очень примитивно, но работает) root = kw[:-2] if len(kw) > 5 else kw if root in fname or root in desc: hits += 1 if hits > max_hits: max_hits = hits best_match = item['filename'] return best_match async def extract_data_from_chat(user_query, filename): system_prompt = "Извлеки данные в JSON. Даты dd.mm.yyyy.\nСХЕМА:\n" for key, val in clean_tags_db.items(): system_prompt += f"- {val['tag']}: {val['description']}\n" try: response = client.chat.completions.create( model="deepseek-chat", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_query} ], response_format={"type": "json_object"}, temperature=0.1 ) return json_repair.loads(response.choices[0].message.content) except: return {} async def consult_logic(user_text): # Логика консультации (если файл не найден) context = "" if collection and encoder: try: vec = encoder.encode(user_text).tolist() res = collection.query(query_embeddings=[vec], n_results=3) if res['documents'][0]: for txt in res['documents'][0]: context += f"{txt[:1000]}\n...\n" except: pass sys_p = "Ты юрист. Отвечай, используя законы (если есть)." try: r = client.chat.completions.create( model="deepseek-chat", messages=[{"role":"system","content":sys_p}, {"role":"user","content":f"Вопрос: {user_text}\nКонтекст: {context}"}], temperature=0.3 ) return {"type": "text", "content": r.choices[0].message.content} except Exception as e: return {"type": "text", "content": f"Ошибка: {e}"} async def generate_doc_logic(user_text): # 1. Ищем шаблон best_filename = await select_best_template(user_text) # ЕСЛИ ШАБЛОН НЕ НАЙДЕН — ВОЗВРАЩАЕМ ТЕКСТ (КОНСУЛЬТАЦИЮ) if not best_filename: fallback_res = await consult_logic(f"Составь текст документа: {user_text}") # Добавляем приписку, почему так вышло fallback_res["content"] = "⚠️ **Шаблон файла не найден в базе.**\nЯ составил текст документа вручную:\n\n" + fallback_res["content"] return fallback_res # 2. Проверяем файл на диске template_path = os.path.join(TEMPLATES_DIR, best_filename) if not os.path.exists(template_path): return {"type": "text", "content": f"⚠️ Ошибка: Шаблон '{best_filename}' есть в реестре, но файл отсутствует на сервере."} # 3. Извлекаем данные и заполняем context = await extract_data_from_chat(user_text, best_filename) if "doc_date" not in context: context["doc_date"] = datetime.now().strftime("%d.%m.%Y") try: doc = DocxTemplate(template_path) doc.render(context) if not os.path.exists(DOWNLOADS_DIR): os.makedirs(DOWNLOADS_DIR) ts = datetime.now().strftime("%Y%m%d_%H%M%S") out_name = f"Ready_{best_filename.replace('.docx', '')[:20]}_{ts}.docx" out_path = os.path.join(DOWNLOADS_DIR, out_name) doc.save(out_path) return { "type": "file", "content": f"✅ Документ сформирован на основе шаблона: **{best_filename}**", "file_url": out_path } except Exception as e: return {"type": "text", "content": f"Ошибка при заполнении: {e}"}