Create pipeline.py

pipeline.py

@@ -0,0 +1,481 @@
+from typing import TypedDict, List, Annotated
+import json
+import time
+from contextlib import contextmanager
+from pathlib import Path
+import json, re
+import tempfile
+import zipfile, tarfile
+
+from langchain_core.documents import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.llms.yandex import YandexGPT
+from langgraph.graph import StateGraph, END
+from langgraph.graph.message import add_messages
+
+from dotenv import load_dotenv
+load_dotenv()
+
+from pathlib import Path
+import os
+from pypdf import PdfReader
+
+# ====================== CONFIG ======================
+
+YANDEX_API_KEY = os.getenv("YANDEX_API_KEY")
+YANDEX_FOLDER_ID = os.getenv("YANDEX_FOLDER_ID")
+
+# --- LLMs ---
+llm_reasoning = YandexGPT(
+    api_key=YANDEX_API_KEY,
+    folder_id=YANDEX_FOLDER_ID,
+    temperature=0.3,
+    max_tokens=2000
+)
+
+llm_output = YandexGPT(
+    api_key=YANDEX_API_KEY,
+    folder_id=YANDEX_FOLDER_ID,
+    temperature=0.4,
+    max_tokens=6500
+)
+
+def llm_analyze(prompt: str) -> str:
+    return llm_reasoning.invoke(prompt)
+
+def llm_final(prompt: str) -> str:
+    return llm_output.invoke(prompt)
+
+# --- Logging ---
+@contextmanager
+def log_step(name: str):
+    print(f"\n>>> START: {name}")
+    start = time.time()
+    try:
+        yield
+    finally:
+        print(f"<<< END: {name} ({time.time() - start:.2f}s)")
+
+# ====================== STATE ======================
+
+class GraphState(TypedDict):
+    query: str
+    toc: str
+    toc_analysis: str
+    plan: dict
+    retrieval_queries: List[str]
+    contexts: List[Document]
+    result: dict
+    iteration: int
+    recurse: bool
+
+# ====================== ARCHIVE PROCESSING ======================
+
+def read_file(file_path: Path) -> str:
+    suffix = file_path.suffix.lower()
+
+    try:
+        # --- PDF ---
+        if suffix == ".pdf":
+            reader = PdfReader(str(file_path))
+            text = []
+            for page in reader.pages:
+                text.append(page.extract_text() or "")
+            return "\n".join(text)
+
+        # --- TXT / CODE ---
+        elif suffix in [".txt", ".md", ".py", ".json"]:
+            return file_path.read_text(encoding="utf-8", errors="ignore")
+
+        # --- fallback ---
+        else:
+            return file_path.read_text(encoding="utf-8", errors="ignore")
+
+    except Exception as e:
+        print(f"⚠️ Ошибка чтения {file_path}: {e}")
+        return ""
+
+
+def extract_archive_to_documents(archive_path: str) -> List[Document]:
+    documents = []
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+
+        if archive_path.endswith('.zip'):
+            with zipfile.ZipFile(archive_path) as z:
+                z.extractall(tmpdir)
+        elif archive_path.endswith(('.tar', '.tar.gz', '.tgz')):
+            with tarfile.open(archive_path) as t:
+                t.extractall(tmpdir)
+        else:
+            raise ValueError(f"Unsupported archive: {archive_path}")
+
+        for file_path in tmp_path.rglob("*"):
+            if file_path.is_file() and not file_path.name.startswith('.'):
+                try:
+                    relative_path = file_path.relative_to(tmp_path)
+                    text = read_file(file_path)
+
+                    if not text.strip():
+                        continue
+
+                    doc = Document(
+                        page_content=text,
+                        metadata={
+                            "source": str(relative_path),
+                            "file_name": file_path.name,
+                            "file_type": file_path.suffix.lower(),
+                            "size_bytes": file_path.stat().st_size,
+                        }
+                    )
+                    documents.append(doc)
+
+                except Exception as e:
+                    print(f"⚠️ Не удалось обработать {file_path}: {e}")
+
+    print(f"Извлечено {len(documents)} документов из архива")
+    return documents
+
+
+# ====================== VECTORSTORE ======================
+
+def build_vectorstore(docs: List[Document]):
+    with log_step("Building FAISS Vectorstore"):
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1200,
+            chunk_overlap=150,
+            separators=["\n\n", "\n", ". ", " ", ""]
+        )
+        
+        splits = splitter.split_documents(docs)
+        
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        
+        db = FAISS.from_documents(splits, embeddings)
+        retriever = db.as_retriever(
+            search_type="mmr",
+            search_kwargs={"k": 10, "fetch_k": 25, "lambda_mult": 0.7}
+        )
+        return db, retriever
+
+# ====================== NODES ======================
+
+def analyze_toc(state: GraphState):
+    with log_step("analyze_toc"):
+        prompt = f"Оглавление курса:\n{state['toc']}\n\nКратко опиши структуру курса, выделив основные разделы и логику."
+        toc_analysis = llm_analyze(prompt)
+        return {"toc_analysis": toc_analysis}
+
+
+def planner(state: GraphState):
+    with log_step("planner"):
+        prompt = f"""
+Оглавление:
+{state['toc']}
+
+Анализ структуры:
+{state['toc_analysis']}
+
+Запрос пользователя:
+{state['query']}
+
+Составь план retrieval'а. Верни JSON:
+{{
+  "sections": ["список ключевых тем"],
+  "queries": ["конкретные поисковые запросы для RAG"]
+}}
+"""
+        raw = llm_analyze(prompt)
+        try:
+            plan = json.loads(raw)
+        except:
+            plan = {"sections": [], "queries": [state["query"]]}
+
+        return {
+            "plan": plan,
+            "retrieval_queries": plan.get("queries", [state["query"]]),
+            "iteration": 0
+        }
+
+
+def retrieve(state: GraphState):
+    with log_step("retrieve"):
+        all_docs = []
+        for q in state["retrieval_queries"]:
+            docs = retriever.invoke(q)
+            all_docs.extend(docs)
+
+        seen = set()
+        unique_docs = []
+        for doc in all_docs:
+            content_hash = hash(doc.page_content[:200])
+            if content_hash not in seen:
+                seen.add(content_hash)
+                unique_docs.append(doc)
+
+        print(f"Retrieved {len(unique_docs)} unique documents")
+        return {
+            "contexts": state["contexts"] + unique_docs
+        }
+
+def check_completeness(state: GraphState):
+    with log_step("check_completeness"):
+        state["iteration"] += 1
+        
+        context_text = "\n\n".join(doc.page_content for doc in state["contexts"])[:6000]
+
+        prompt = f"""
+Запрос пользователя: {state['query']}
+
+Текущий контекст имеет {len(state['contexts'])} документов.
+
+Проанализируй, достаточно ли материала для создания полноценного 9-недельного интенсивного курса.
+
+Верни JSON:
+{{
+  "enough": true или false,
+  "next_query": "если не enough — один хороший поисковый запрос, иначе пустая строка"
+}}
+"""
+
+        raw = llm_analyze(prompt)
+
+        try:
+            match = re.search(r'\{[\s\S]*\}', raw)
+            data = json.loads(match.group(0)) if match else json.loads(raw)
+
+            enough = data.get("enough", False)
+            next_query = data.get("next_query", "").strip()
+
+            print(f"Enough: {enough}")
+            if next_query:
+                print(f"→ New query: {next_query}")
+
+            if not enough and next_query and state["iteration"] < 4:
+                return {
+                    "retrieval_queries": state["retrieval_queries"] + [next_query],
+                    "recurse": True
+                }
+            else:
+                return {"recurse": False}
+
+        except Exception as e:
+            print(f"Parse error in check_completeness: {e}")
+            return {"recurse": False}
+        
+
+def generate_weekly_plan(state: GraphState):
+    with log_step("generate_weekly_plan"):
+        
+        context_blocks = []
+        for doc in state["contexts"][:25]:
+            meta = doc.metadata
+            source = meta.get("source") or meta.get("file_name", "unknown")
+            preview = doc.page_content[:700].replace("\n", " ").strip()
+            context_blocks.append(f"Источник → {source}\n{preview}...\n")
+
+        context_text = "\n\n" + "-" * 80 + "\n\n".join(context_blocks)
+
+        prompt = f"""Ты — строгий профессиональный методист. 
+Отвечай **исключительно валидным JSON**, без единого слова вне объекта.
+
+Запрос пользователя: {state['query']}
+
+Оглавление курса:
+{state.get('toc', 'Не предоставлено')}
+
+Доступные материалы:
+{context_text}
+
+Создай понедельный учебный план. 
+В одной неделе может быть от 2 до 5 занятий.
+
+Если требуется усиленный курс, то надо от 3х занятий в неделю.
+
+Верни **ТОЛЬКО** этот JSON (начинай сразу с '{{'):
+
+{{
+  "course_title": "Название курса",
+  "duration_weeks": число,
+  "weekly_plan": [
+    {{
+      "week": 1,
+      "theme": "Общая тема недели",
+      "sessions": [
+        {{
+          "session_number": 1,
+          "title": "Название конкретного занятия",
+          "goal": "Подробная мотивировка и цель занятия (что студент должен понять и уметь)",
+          "main_sources": [
+            {{
+              "material": "Название учебника или лекции",
+              "chapter": "§1 Проективное пространство",
+              "section": "1.1 Соглашения об обозначениях, 1.2 Определение",
+              "file_source": "имя_файла.pdf"
+            }}
+          ],
+          "preparation_materials": ["список строк для подготовки"],
+          "practice_and_homework": ["конкретные задания, задачи, домашняя работа"],
+          "estimated_time": "3–5 часов"
+        }}
+      ]
+    }}
+  ],
+  "additional_recommendations": "Общие рекомендации"
+}}
+
+Правила, которых нужно строго придерживаться:
+- Используй реальные названия глав, параграфов и файлов из предоставленных материалов.
+- Обязательно заполняй поле "goal" — подробная мотивировка, зачем это занятие нужно.
+- Поле "file_source" должно содержать реальное имя файла из метаданных, если возможно.
+- В одной неделе можно делать несколько занятий (sessions).
+- Не пиши ничего кроме JSON.
+"""
+
+        raw = llm_final(prompt)
+
+        try:
+            cleaned = raw.strip()
+            if cleaned.startswith("```"):
+                cleaned = cleaned.split("```")[1].strip()
+            if cleaned.lower().startswith(("json", "```json")):
+                cleaned = cleaned.split("\n", 1)[-1].strip()
+
+            plan = json.loads(cleaned)
+            
+        except json.JSONDecodeError:
+            import re
+            match = re.search(r'(\{[\s\S]*\})', raw, re.DOTALL)
+            if match:
+                try:
+                    plan = json.loads(match.group(1))
+                except:
+                    plan = {"error": "JSON parse failed", "raw_preview": raw[:1000]}
+            else:
+                plan = {"error": "JSON parse failed", "raw_preview": raw[:800]}
+
+
+        return {"result": plan}
+    
+
+def make_final_output_pretty(state: GraphState):
+    import json
+
+    raw_plan = state["result"]
+
+    if isinstance(raw_plan, dict):
+        plan_str = json.dumps(raw_plan, ensure_ascii=False, indent=2)
+    else:
+        plan_str = str(raw_plan)
+
+    prompt = f"""
+Ты — отличный технический писатель и методист.
+
+У тебя есть следующий JSON с учебным планом (возможно, неидеальный):
+
+{plan_str}
+
+Твоя задача:
+преобразовать его в КРАСИВЫЙ и СТРУКТУРИРОВАННЫЙ Markdown-документ.
+
+Требования:
+- Используй заголовки (#, ##, ###)
+- Выделяй недели как отдельные секции
+- Для каждой недели:
+  - тема
+  - список занятий (bullets)
+  - краткое описание (если есть)
+- Добавь читаемую структуру курса
+- В конце добавь раздел "Рекомендации"
+- Убери технический мусор и JSON-структуру
+- Сделай текст естественным, как учебный материал
+
+Формат:
+# Название курса
+
+## Неделя 1 — Тема
+- Занятие 1
+- Занятие 2
+
+## Неделя 2 — ...
+
+## Рекомендации
+
+Верни ТОЛЬКО Markdown, без JSON и без пояснений.
+"""
+
+    pretty = llm_final(prompt)
+
+    return {"result": pretty}
+    
+
+# ====================== GRAPH ======================
+def should_continue(state: GraphState) -> str:
+    if state.get("recurse") is True:
+        return "retrieve"
+    return "generate_weekly_plan"
+    
+
+builder = StateGraph(GraphState)
+
+builder.add_node("analyze_toc", analyze_toc)
+builder.add_node("planner", planner)
+builder.add_node("retrieve", retrieve)
+builder.add_node("check_completeness", check_completeness)
+builder.add_node("generate_weekly_plan", generate_weekly_plan)
+builder.add_node("pretty_output", make_final_output_pretty)
+
+builder.set_entry_point("analyze_toc")
+
+builder.add_edge("analyze_toc", "planner")
+builder.add_edge("planner", "retrieve")
+builder.add_edge("retrieve", "check_completeness")
+
+builder.add_conditional_edges(
+    "check_completeness",
+    should_continue,
+    {
+        "retrieve": "retrieve",
+        "generate_weekly_plan": "generate_weekly_plan"
+    }
+)
+
+builder.add_edge("generate_weekly_plan", "pretty_output")
+builder.add_edge("pretty_output", END)
+
+graph = builder.compile()
+
+# ====================== GLOBAL ======================
+retriever = None
+toc_text = ""
+
+# ====================== MAIN PIPELINE ======================
+
+def run_course_builder(query: str, toc: str, archive_path: str):
+    global retriever, toc_text
+    
+    toc_text = toc
+    
+    documents = extract_archive_to_documents(archive_path)
+    
+    _, retriever = build_vectorstore(documents)
+
+    initial_state = {
+        "query": query,
+        "toc": toc_text,
+        "toc_analysis": "",
+        "plan": {},
+        "retrieval_queries": [],
+        "contexts": [],
+        "result": {},
+        "iteration": 0,
+        "recurse": False
+    }
+
+    final_state = graph.invoke(initial_state)
+    
+    return final_state["result"]
+