Update main.py with Ollama and LangGraph implementation

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_core.tools import BaseTool, tool
-from langchain.tools import tool
+from langchain_core.prompts import ChatPromptTemplate
-from deepagents import create_deep_agent
+from langchain_core.output_parsers import StrOutputParser
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langchain_core.runnables import Runnable
-from langchain_core.messages import HumanMessage
+from langchain.agents import create_agent, AgentExecutor, Tool
 
-# --------------------- Configuration ---------------------
+# ---------------------------------------------------------------------------
-BASE_DIR = Path(__file__).parent
+# 1. Векторная база
-DATA_DIR = BASE_DIR / "data"
+# ---------------------------------------------------------------------------
-CHROMA_DIR = BASE_DIR / "chroma_faq"
+CHROMA_PATH = Path("./chroma_faq")
-MOCK_META_FILE = BASE_DIR / "course_meta.json"
+DATA_PATH = Path("./data")
 
-# --------------------- LLM and Embeddings ---------------------
+# Создаём embeddings
-llm = ChatOpenAI(
+embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    model="openai/gpt-oss-20b:free",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
-)
 
-embeddings = OpenAIEmbeddings(
+# ---------------------------------------------------------------------------
-    model="text-embedding-3-small",
+# Чтение markdown‑файлов и загрузка в Chroma
-    base_url="https://openrouter.ai/api/v1",
+# ---------------------------------------------------------------------------
-    api_key=os.getenv("OPENAI_API_KEY"),
-)
 
-# --------------------- Chroma Vector Store ---------------------
+def load_faq_to_chroma() -> Chroma:
-vector_store = Chroma(
+    """Загружает все .md файлы из DATA_PATH в Chroma.
-    collection_name="faq_collection",
+    Если коллекция уже существует – просто возвращаем её.
+    """
+    if CHROMA_PATH.exists():
+        return Chroma(
+            collection_name="faq",
             embedding_function=embeddings,
-    persist_directory=str(CHROMA_DIR),
+            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."""
+    """Возвращает кортеж из k наиболее релевантных фрагментов.
-    docs = vector_store.similarity_search(query, k=k)
+    Формат ответа: ``source: chroma`` + текст.
-    if not docs:
+    """
-        return "No relevant information found in the course materials."
+    vector_store = load_faq_to_chroma()
-    return "\n\n---\n\n".join(f"**{doc.metadata.get('title', 'Document')}**\n{doc.page_content}" for doc in docs)
+    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.
+    """Имитирует вызов MCP‑сервера. Читает локальный JSON и возвращает
-    In production this would be an HTTP call to an MCP server.
+    информацию, содержащуюся в ключе, совпадающем с query.
-    Here we read a local JSON file for simplicity.
     """
-    import json
+    if not MCP_JSON_PATH.exists():
-    if not MOCK_META_FILE.exists():
+        return "source: mcp_meta\nMeta file not found."
-        return "Metadata source not available."
+    data = json.loads(MCP_JSON_PATH.read_text(encoding="utf-8"))
-    with open(MOCK_META_FILE, "r", encoding="utf-8") as f:
+    # простая логика: ищем ключ, содержащий query (case‑insensitive)
-        data = json.load(f)
+    for key, value in data.items():
-    # Simple keyword search in the metadata
+        if query.lower() in key.lower():
-    results = [f"{k}: {v}" for k, v in data.items() if query.lower() in k.lower() or query.lower() in str(v).lower()]
+            return f"source: mcp_meta\n{key}: {value}"
-    return "\n".join(results) if results else "No metadata matches your query."
+    return "source: mcp_meta\nNo matching metadata found."
 
-# --------------------- Backend ---------------------
+# ---------------------------------------------------------------------------
-backend = CompositeBackend([
+# 4. Создание агента
-    LocalShellBackend(workspace_dir="./workspace"),
+# ---------------------------------------------------------------------------
-    FilesystemBackend(),
+# Системный промпт, который заставляет агент выбирать нужный инструмент
+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."
+)
+
+# Создаём LLM
+llm = ChatOllama(model="llama3", temperature=0.2)
+
+# Prompt template
+prompt = ChatPromptTemplate.from_messages([
+    ("system", SYSTEM_PROMPT),
+    ("human", "{input}"),
 ])
 
-# --------------------- Agent ---------------------
+# Создаём Runnable, который будет использовать инструменты
-agent = create_deep_agent(
+agent = create_agent(
-    model=llm,
+    llm=llm,
     tools=[search_course_docs, fetch_course_meta],
-    backend=backend,
+    system_message=SYSTEM_PROMPT,
-    system_prompt=(
+    verbose=True,
-        "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'."
-    ),
 )
 
-# --------------------- Data Loading ---------------------
+# ---------------------------------------------------------------------------
-async def load_faq_to_chroma():
+# 5. CLI
-    """Load all .md files from data/ into the Chroma collection."""
+# ---------------------------------------------------------------------------
-    if not DATA_DIR.exists():
+PREDEFINED_QUESTIONS = [
-        print("Data directory not found.")
+    "What topics are covered in the first lecture?",  # Chroma
-        return
+    "Explain the concept of polymorphism in the context of the course.",  # Chroma
-    docs = []
+    "What is the schedule for the next week?",  # MCP
-    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/.")
-
-# --------------------- 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
 ]
 
-async def run_cli():
+def run_cli():
-    await load_faq_to_chroma()
+    print("--- FAQ Bot ---")
-    print("\n--- FAQ Bot CLI ---\n")
+    print("Predefined questions:")
-    for i, q in enumerate(PRESET_QUESTIONS, 1):
+    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(
+        if user_input.isdigit() and 1 <= int(user_input) <= len(PREDEFINED_QUESTIONS):
-            {"messages": [HumanMessage(content=user_input)]},
+            question = PREDEFINED_QUESTIONS[int(user_input) - 1]
-            {"configurable": {"thread_id": "session-1"}},
+        else:
-        )
+            question = user_input
-        print(result["messages"][-1].content)
+        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()