From 68e1d68a26454caedadfdd8f3daa9977481cf445 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 09:17:51 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=9F=D0=BE=D0=B2?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BD=D1=8B=D0=B9=20=D1=8D=D0=BA=D0=B7=D0=B0?= =?UTF-8?q?=D0=BC=D0=B5=D0=BD:=20FAQ-=D0=B1=D0=BE=D1=82=20=E2=80=94=20Chro?= =?UTF-8?q?maDB=20+=20=D0=BE=D0=B4=D0=B8=D0=BD=20MCP-tool?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 189 ++++++++++++++++++++++++++++++++++---------------------- 1 file changed, 115 insertions(+), 74 deletions(-) diff --git a/main.py b/main.py index 499ce0b..f0564d3 100644 --- a/main.py +++ b/main.py @@ -1,130 +1,171 @@ import os import asyncio import json +import httpx from pathlib import Path -from langchain_openai import ChatOpenAI, OpenAIEmbeddings + +from dotenv import load_dotenv +from langchain_openai import ChatOpenAI from langchain_chroma import Chroma +from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_core.documents import Document +from langchain_core.messages import HumanMessage from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage -# DESIGN DECISION: Use ChromaDB for local vector store -# NECESSITY: The assignment explicitly requires ChromaDB + Ollama embeddings. -# OPTIMALITY: ChromaDB is lightweight, file-based, and integrates directly with LangChain. -# ALTERNATIVES CONSIDERED: QDrant would need a separate server process and more setup. +# Загрузка переменных окружения +load_dotenv() +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -# Embeddings via OpenRouter -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), +# Инициализация эмбеддингов Ollama +embeddings = OllamaEmbeddings(model="nomic-embed-text") + +# Создание или загрузка коллекции Chroma +vector_store = Chroma( + collection_name="faq", + embedding_function=embeddings, + persist_directory="./chroma_faq", ) -# Persist directory for Chroma -CHROMA_DIR = Path("./chroma_faq") +def load_faq_to_chroma() -> None: + """ + Загружает все .md файлы из папки data/, разбивает их на чанки и сохраняет в Chroma. + """ + data_dir = Path("data") + if not data_dir.exists(): + print("Папка data/ не найдена. Создайте её и добавьте .md файлы.") + return -def load_faq_to_chroma(): - """ - Load .md files from data/ into ChromaDB. - """ - vector_store = Chroma( - collection_name="faq", - embedding_function=embeddings, - persist_directory=str(CHROMA_DIR), - ) - if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()): - docs = [] - for md_file in Path("data").glob("*.md"): - content = md_file.read_text(encoding="utf-8") - docs.append(Document(page_content=content, metadata={"source": md_file.name})) + # Очистка коллекции перед загрузкой + vector_store.delete_collection() + + splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + docs = [] + + for md_file in data_dir.glob("*.md"): + text = md_file.read_text(encoding="utf-8") + chunks = splitter.split_text(text) + for i, chunk in enumerate(chunks): + docs.append(Document(page_content=chunk, metadata={"source": md_file.name, "chunk": i})) + + if docs: vector_store.add_documents(docs) vector_store.persist() - return vector_store - -vector_store = load_faq_to_chroma() + print(f"Загружено {len(docs)} чанков из {len(list(data_dir.glob('*.md')))} файлов.") + else: + print("Нет документов для загрузки.") @tool -def search_course_docs(query: str) -> str: +def search_course_docs(query: str, k: int = 3) -> str: """ - Search the knowledge base for relevant information. + Ищет релевантные фрагменты из локальной базы знаний. """ - docs = vector_store.similarity_search(query, k=3) - return "\n\n".join(d.page_content for d in docs) if docs else "No results found." + docs = vector_store.similarity_search(query, k=k) + if not docs: + return "No results found in the knowledge base." + return "\n\n".join( + f"[{doc.metadata.get('source', 'unknown')} - chunk {doc.metadata.get('chunk', 0)}]\n{doc.page_content}" + for doc in docs + ) + +# Статический JSON с метаданными курса +META_JSON_PATH = Path("meta.json") @tool def fetch_course_meta(query: str) -> str: """ - Fetch course metadata from a static JSON file. + Возвращает метаданные курса, если запрос содержит ключевые слова. """ - meta_path = Path("meta.json") - if not meta_path.exists(): + if not META_JSON_PATH.exists(): return "Metadata file not found." - data = json.loads(meta_path.read_text(encoding="utf-8")) - # Simple case-insensitive search in keys and values - matches = [] - for key, value in data.items(): - if isinstance(value, dict): - for subkey, subvalue in value.items(): - if query.lower() in subkey.lower() or query.lower() in str(subvalue).lower(): - matches.append(f"{subkey}: {subvalue}") - else: - if query.lower() in key.lower() or query.lower() in str(value).lower(): - matches.append(f"{key}: {value}") - return "\n".join(matches) if matches else "No metadata matches your query." -# LLM via OpenRouter + with META_JSON_PATH.open("r", encoding="utf-8") as f: + meta = json.load(f) + + # Простая логика поиска: если запрос содержит слово "расписание", возвращаем расписание + if "расписание" in query.lower(): + return json.dumps(meta.get("schedule", {}), indent=2) + # Если запрос содержит слово "преподаватель", возвращаем список преподавателей + if "преподаватель" in query.lower(): + return json.dumps(meta.get("instructors", []), indent=2) + # Иначе возвращаем всю мета + return json.dumps(meta, indent=2) + +# Инициализация LLM через OpenRouter llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), + api_key=OPENAI_API_KEY, temperature=0.0, ) -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -system_prompt = ( - "You are a helpful FAQ bot. Use search_course_docs for questions about course materials. " - "Use fetch_course_meta for questions about schedule or metadata. " - "Do not use both tools unless necessary. " - "Indicate source in your answer: source: chroma | mcp_meta." +# Backend для deepagents +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] ) +# Системный промпт, который указывает агенту использовать нужный инструмент +SYSTEM_PROMPT = """ +You are a helpful FAQ assistant for a course. +When answering a question, first decide whether the answer can be found in the local knowledge base or requires metadata. +If the answer is in the knowledge base, use the tool `search_course_docs`. +If the answer requires metadata (e.g., schedule, instructors), use the tool `fetch_course_meta`. +Do not call both tools unless absolutely necessary. +In your final answer, prepend the line `source: ` to indicate which tool was used. +If no tool is needed, just provide the answer and set source to `none`. +Be concise and accurate. +""" + agent = create_deep_agent( model=llm, tools=[search_course_docs, fetch_course_meta], backend=backend, - system_prompt=system_prompt, + system_prompt=SYSTEM_PROMPT, ) -async def ask_agent(question: str, thread_id: str = "session-1") -> str: +async def ask_agent(question: str) -> str: + """ + Отправляет вопрос агенту и возвращает последний вывод. + """ result = await agent.ainvoke( {"messages": [HumanMessage(content=question)]}, - {"configurable": {"thread_id": thread_id}}, + {"configurable": {"thread_id": "session-1"}}, ) + # Последнее сообщение агента return result["messages"][-1].content -async def main(): +async def main() -> None: + # Загружаем данные в Chroma + load_faq_to_chroma() + + # Предварительно заданные вопросы preset_questions = [ - "What is covered in Lecture 1?", - "Explain supervised learning.", - "When is Lecture 2 scheduled?", + "Какой материал будет на следующем занятии?", + "Кто будет вести лекцию по машинному обучению?", + "Когда проходит экзамен по курсу?", ] - print("=== Preset questions ===") + + print("\n=== Предварительные вопросы ===") for q in preset_questions: + print(f"\nВопрос: {q}") answer = await ask_agent(q) - print(f"\nQ: {q}\nA: {answer}\n") - print("=== Interactive mode (type 'exit' to quit) ===") + print(f"Ответ:\n{answer}") + + # Интерактивный режим + print("\n=== Интерактивный режим ===") + print("Введите ваш вопрос (или 'выход' для завершения).") while True: - user_input = input("\nYour question: ") - if user_input.lower() in {"exit", "quit"}: + user_input = input("\n> ").strip() + if user_input.lower() in {"выход", "exit", "quit"}: + print("До свидания!") break answer = await ask_agent(user_input) - print(f"\nAnswer: {answer}") + print(f"Ответ:\n{answer}") if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file