diff --git a/solution.py b/solution.py new file mode 100644 index 0000000..de9ba91 --- /dev/null +++ b/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() \ No newline at end of file