Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool: solution.py

solution.py

@@ -0,0 +1,163 @@
+<|channel|>final code<|message|>import os
+import json
+from pathlib import Path
+
+import httpx
+from dotenv import load_dotenv
+from langchain_ollama import ChatOllama
+from langchain_chroma import Chroma
+from langchain_core.prompts import PromptTemplate
+from langchain.agents import initialize_agent, Tool, AgentType
+from langchain.schema import Document
+
+# --------------------------------------------------------------------------- #
+# 1. Загрузка FAQ в ChromaDB
+# --------------------------------------------------------------------------- #
+
+DATA_DIR = Path("data")
+CHROMA_PATH = Path("./chroma_faq")
+
+def load_faq_to_chroma() -> Chroma:
+    """
+    Читает все *.md файлы из DATA_DIR, разбивает их на чанки и сохраняет
+    в локальный ChromaDB. Возвращает объект Chroma.
+    """
+    # Создаём список документов
+    docs = []
+    for md_file in DATA_DIR.glob("*.md"):
+        text = md_file.read_text(encoding="utf-8")
+        docs.append(Document(page_content=text, metadata={"source": str(md_file)}))
+
+    # Инициализируем Chroma с Ollama embeddings
+    chroma = Chroma.from_documents(
+        documents=docs,
+        embedding=ChatOllama(model="nomic-embed-text"),
+        persist_directory=str(CHROMA_PATH),
+    )
+    return chroma
+
+def search_course_docs(query: str, k: int = 3) -> list[Document]:
+    """
+    Поиск в ChromaDB. Возвращает список документов с наиболее релевантными
+    ответами на запрос.
+    """
+    chroma = load_faq_to_chroma()
+    return chroma.similarity_search_with_score(query, k=k)
+
+# --------------------------------------------------------------------------- #
+# 2. MCP‑style tool (mock)
+# --------------------------------------------------------------------------- #
+
+def fetch_course_meta(query: str) -> dict:
+    """
+    Имитация вызова внешнего MCP сервера. Читает статический JSON
+    из файла mock_meta.json и возвращает часть, соответствующую query.
+    """
+    meta_file = Path("mock_meta.json")
+    if not meta_file.exists():
+        raise FileNotFoundError(f"{meta_file} не найден")
+
+    data = json.loads(meta_file.read_text(encoding="utf-8"))
+    # простая фильтрация по ключу
+    return {k: v for k, v in data.items() if query.lower() in k.lower()}
+
+# --------------------------------------------------------------------------- #
+# 3. Agent + Tools
+# --------------------------------------------------------------------------- #
+
+def create_agent():
+    """
+    Создаёт агента LangChain с двумя инструментами:
+      - search_course_docs (Chroma)
+      - fetch_course_meta (MCP‑style)
+    Агент использует Ollama LLM.
+    """
+    llm = ChatOllama(model="nomic-embed-text")
+
+    tools = [
+        Tool(
+            name="search_course_docs",
+            func=lambda q: "\n".join([f"{d.page_content[:200]}..." for d in search_course_docs(q)]),
+            description=(
+                "Используется, если вопрос относится к материалам курса. "
+                "Возвращает краткие отрывки из FAQ."
+            ),
+        ),
+        Tool(
+            name="fetch_course_meta",
+            func=lambda q: json.dumps(fetch_course_meta(q), indent=2),
+            description=(
+                "Вызывается, если нужно получить расписание или метаданные курса. "
+                "Возвращает JSON‑объект."
+            ),
+        ),
+    ]
+
+    system_prompt = (
+        "Ты FAQ‑бот для локального курса. "
+        "Если вопрос относится к материалам занятия — используй search_course_docs. "
+        "Если нужен график, даты или другая метаинформация — fetch_course_meta. "
+        "Не вызывай оба инструмента одновременно. В ответе указывай источник: "
+        "`source: chroma` или `source: mcp_meta`."
+    )
+
+    prompt = PromptTemplate(
+        input_variables=["input"],
+        template=system_prompt + "\n\nВопрос: {input}\nОтвет:",
+    )
+
+    agent = initialize_agent(
+        tools,
+        llm,
+        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
+        verbose=True,
+        max_iterations=5,
+        early_stopping_method="generate",
+        system_message=prompt.format(input="{input}"),
+    )
+    return agent
+
+# --------------------------------------------------------------------------- #
+# 4. CLI
+# --------------------------------------------------------------------------- #
+
+def main():
+    """
+    Простая консольная утилита с тремя готовыми вопросами и интерактивным режимом.
+    """
+    print("=== FAQ‑бот для курса ===")
+    sample_questions = [
+        "Какой материал будет на следующем занятии?",
+        "Когда начинается следующий модуль?",
+        "Какие темы охватывает первый раздел?",
+    ]
+
+    for idx, q in enumerate(sample_questions, 1):
+        print(f"\n{idx}. {q}")
+
+    print("\nВведите номер вопроса (1‑3) или '0' для интерактивного режима.")
+    choice = input("Ваш выбор: ").strip()
+
+    if choice == "0":
+        while True:
+            q = input("\nВопрос (или 'exit'): ")
+            if q.lower() in ("exit", "quit"):
+                break
+            agent = create_agent()
+            response = agent.run(q)
+            print(f"\nОтвет:\n{response}\n")
+    else:
+        try:
+            idx = int(choice) - 1
+            question = sample_questions[idx]
+        except (ValueError, IndexError):
+            print("Неверный выбор.")
+            return
+
+        agent = create_agent()
+        response = agent.run(question)
+        print(f"\nОтвет:\n{response}\n")
+
+if __name__ == "__main__":
+    load_dotenv()  # если понадобится переменные окружения
+    main()