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

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")
 
-# Persist directory for Chroma
-CHROMA_DIR = Path("./chroma_faq")
-
-def load_faq_to_chroma():
-    """
-    Load .md files from data/ into ChromaDB.
-    """
-    vector_store = Chroma(
+# Создание или загрузка коллекции Chroma
+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()):
+    persist_directory="./chroma_faq",
+)
+
+def load_faq_to_chroma() -> None:
+    """
+    Загружает все .md файлы из папки data/, разбивает их на чанки и сохраняет в Chroma.
+    """
+    data_dir = Path("data")
+    if not data_dir.exists():
+        print("Папка data/ не найдена. Создайте её и добавьте .md файлы.")
+        return
+
+    # Очистка коллекции перед загрузкой
+    vector_store.delete_collection()
+
+    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     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}))
+
+    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([
+# Backend для deepagents
+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."
+    ]
 )
 
+# Системный промпт, который указывает агенту использовать нужный инструмент
+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: <chroma|mcp_meta>` 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())