diff --git a/main.py b/main.py index f531129..1a8e1e4 100644 --- a/main.py +++ b/main.py @@ -1,124 +1,166 @@ +"""main.py + +Простой FAQ‑бот, использующий ChromaDB для локальных конспектов и один MCP‑подобный инструмент – HTTP‑запрос к статическому JSON. + +Требования: +- Python 3.10+ +- Ollama модели: `nomic-embed-text` и `llama3` +- Пакеты из requirements.txt + +Запуск: +```bash +python main.py +``` + +В интерактивном режиме можно задавать вопросы. В примере уже есть три готовых вопроса – два из Chroma, один из MCP‑тул. +""" + +import json import os -import asyncio +import sys from pathlib import Path -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from typing import List + +import httpx +from langchain_ollama import ChatOllama, OllamaEmbeddings from langchain_chroma import Chroma -from langchain_core.documents import Document -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_core.messages import HumanMessage +from langchain_core.tools import BaseTool, tool +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.output_parsers import StrOutputParser +from langchain_core.runnables import Runnable +from langchain.agents import create_agent, AgentExecutor, Tool -# --------------------- Configuration --------------------- -BASE_DIR = Path(__file__).parent -DATA_DIR = BASE_DIR / "data" -CHROMA_DIR = BASE_DIR / "chroma_faq" -MOCK_META_FILE = BASE_DIR / "course_meta.json" +# --------------------------------------------------------------------------- +# 1. Векторная база +# --------------------------------------------------------------------------- +CHROMA_PATH = Path("./chroma_faq") +DATA_PATH = Path("./data") -# --------------------- LLM and Embeddings --------------------- -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), - temperature=0.0, -) +# Создаём embeddings +embeddings = OllamaEmbeddings(model="nomic-embed-text") -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) +# --------------------------------------------------------------------------- +# Чтение markdown‑файлов и загрузка в Chroma +# --------------------------------------------------------------------------- -# --------------------- Chroma Vector Store --------------------- -vector_store = Chroma( - collection_name="faq_collection", - embedding_function=embeddings, - persist_directory=str(CHROMA_DIR), -) +def load_faq_to_chroma() -> Chroma: + """Загружает все .md файлы из DATA_PATH в Chroma. + Если коллекция уже существует – просто возвращаем её. + """ + if CHROMA_PATH.exists(): + return Chroma( + collection_name="faq", + embedding_function=embeddings, + persist_directory=str(CHROMA_PATH), + ) + # Если нет, создаём новую коллекцию + from langchain_text_splitters import MarkdownTextSplitter -# --------------------- Tools --------------------- + splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50) + docs = [] + for md_file in DATA_PATH.glob("*.md"): + text = md_file.read_text(encoding="utf-8") + docs.extend(splitter.split_text(text)) + vector_store = Chroma.from_texts( + texts=docs, + embedding=embeddings, + collection_name="faq", + persist_directory=str(CHROMA_PATH), + ) + return vector_store + +# --------------------------------------------------------------------------- +# 2. Tool: поиск по Chroma +# --------------------------------------------------------------------------- @tool def search_course_docs(query: str, k: int = 3) -> str: - """Search the local FAQ collection for relevant passages.""" - docs = vector_store.similarity_search(query, k=k) - if not docs: - return "No relevant information found in the course materials." - return "\n\n---\n\n".join(f"**{doc.metadata.get('title', 'Document')}**\n{doc.page_content}" for doc in docs) + """Возвращает кортеж из k наиболее релевантных фрагментов. + Формат ответа: ``source: chroma`` + текст. + """ + vector_store = load_faq_to_chroma() + results = vector_store.similarity_search(query, k=k) + snippets = "\n---\n".join([r.page_content for r in results]) + return f"source: chroma\n{snippets}" + +# --------------------------------------------------------------------------- +# 3. MCP‑подобный инструмент – HTTP‑запрос к статическому JSON +# --------------------------------------------------------------------------- +MCP_JSON_PATH = Path("./course_meta.json") @tool def fetch_course_meta(query: str) -> str: - """Mock MCP-style tool that returns course metadata. - In production this would be an HTTP call to an MCP server. - Here we read a local JSON file for simplicity. + """Имитирует вызов MCP‑сервера. Читает локальный JSON и возвращает + информацию, содержащуюся в ключе, совпадающем с query. """ - import json - if not MOCK_META_FILE.exists(): - return "Metadata source not available." - with open(MOCK_META_FILE, "r", encoding="utf-8") as f: - data = json.load(f) - # Simple keyword search in the metadata - results = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()] - return "\n".join(results) if results else "No metadata matches your query." + if not MCP_JSON_PATH.exists(): + return "source: mcp_meta\nMeta file not found." + data = json.loads(MCP_JSON_PATH.read_text(encoding="utf-8")) + # простая логика: ищем ключ, содержащий query (case‑insensitive) + for key, value in data.items(): + if query.lower() in key.lower(): + return f"source: mcp_meta\n{key}: {value}" + return "source: mcp_meta\nNo matching metadata found." -# --------------------- Backend --------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# --------------------- Agent --------------------- -agent = create_deep_agent( - model=llm, - tools=[search_course_docs, fetch_course_meta], - backend=backend, - system_prompt=( - "You are a helpful FAQ bot for the course.\n" - "Use the search_course_docs tool for questions about lecture materials.\n" - "Use the fetch_course_meta tool for questions about schedule or metadata.\n" - "Do not call both tools unless absolutely necessary.\n" - "In your final answer, prefix the source with 'source: chroma' or 'source: mcp_meta'." - ), +# --------------------------------------------------------------------------- +# 4. Создание агента +# --------------------------------------------------------------------------- +# Системный промпт, который заставляет агент выбирать нужный инструмент +SYSTEM_PROMPT = ( + "You are a helpful assistant for a course FAQ. Use the provided tools to answer the user. " + "If the question is about course materials, use search_course_docs. " + "If it is about schedule or metadata, use fetch_course_meta. " + "Do not use both tools unless necessary. " + "Always prefix your answer with the source: chroma or source: mcp_meta." ) -# --------------------- Data Loading --------------------- -async def load_faq_to_chroma(): - """Load all .md files from data/ into the Chroma collection.""" - if not DATA_DIR.exists(): - print("Data directory not found.") - return - docs = [] - for md_file in DATA_DIR.glob("*.md"): - text = md_file.read_text(encoding="utf-8") - docs.append(Document(page_content=text, metadata={"title": md_file.stem})) - if docs: - vector_store.add_documents(docs) - vector_store.persist() - print(f"Loaded {len(docs)} documents into Chroma.") - else: - print("No markdown files found in data/.") +# Создаём LLM +llm = ChatOllama(model="llama3", temperature=0.2) -# --------------------- CLI --------------------- -PRESET_QUESTIONS = [ - "What is the deadline for the final project?", # chroma - "Explain the concept of tokenization in NLP.", # chroma - "When is the next lecture scheduled?", # mcp_meta +# Prompt template +prompt = ChatPromptTemplate.from_messages([ + ("system", SYSTEM_PROMPT), + ("human", "{input}"), +]) + +# Создаём Runnable, который будет использовать инструменты +agent = create_agent( + llm=llm, + tools=[search_course_docs, fetch_course_meta], + system_message=SYSTEM_PROMPT, + verbose=True, +) + +# --------------------------------------------------------------------------- +# 5. CLI +# --------------------------------------------------------------------------- +PREDEFINED_QUESTIONS = [ + "What topics are covered in the first lecture?", # Chroma + "Explain the concept of polymorphism in the context of the course.", # Chroma + "What is the schedule for the next week?", # MCP ] -async def run_cli(): - await load_faq_to_chroma() - print("\n--- FAQ Bot CLI ---\n") - for i, q in enumerate(PRESET_QUESTIONS, 1): +def run_cli(): + print("--- FAQ Bot ---") + print("Predefined questions:") + for i, q in enumerate(PREDEFINED_QUESTIONS, 1): print(f"{i}. {q}") - print("\nEnter your own question (or 'exit' to quit):") + print("\nEnter 1-3 to ask a predefined question, or type your own question.") while True: - user_input = input("> ") + user_input = input("\n> ") if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") break - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print(result["messages"][-1].content) + if user_input.isdigit() and 1 <= int(user_input) <= len(PREDEFINED_QUESTIONS): + question = PREDEFINED_QUESTIONS[int(user_input) - 1] + else: + question = user_input + try: + response = agent.invoke({"input": question}) + print("\nAnswer:\n", response) + except Exception as e: + print("Error:", e) if __name__ == "__main__": - asyncio.run(run_cli()) + # Убедимся, что данные загружены + load_faq_to_chroma() + run_cli()