task-6a1d75c5fd30e81cf3126ae7/main.py

import os
import asyncio
import json
import httpx
from pathlib import Path

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

# Загрузка переменных окружения
load_dotenv()
OPENAI_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",
)

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 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()
        print(f"Загружено {len(docs)} чанков из {len(list(data_dir.glob('*.md')))} файлов.")
    else:
        print("Нет документов для загрузки.")

@tool
def search_course_docs(query: str, k: int = 3) -> str:
    """
    Ищет релевантные фрагменты из локальной базы знаний.
    """
    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:
    """
    Возвращает метаданные курса, если запрос содержит ключевые слова.
    """
    if not META_JSON_PATH.exists():
        return "Metadata file not found."

    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=OPENAI_API_KEY,
    temperature=0.0,
)

# 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: <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,
)

async def ask_agent(question: str) -> str:
    """
    Отправляет вопрос агенту и возвращает последний вывод.
    """
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=question)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    # Последнее сообщение агента
    return result["messages"][-1].content

async def main() -> None:
    # Загружаем данные в Chroma
    load_faq_to_chroma()

    # Предварительно заданные вопросы
    preset_questions = [
        "Какой материал будет на следующем занятии?",
        "Кто будет вести лекцию по машинному обучению?",
        "Когда проходит экзамен по курсу?",
    ]

    print("\n=== Предварительные вопросы ===")
    for q in preset_questions:
        print(f"\nВопрос: {q}")
        answer = await ask_agent(q)
        print(f"Ответ:\n{answer}")

    # Интерактивный режим
    print("\n=== Интерактивный режим ===")
    print("Введите ваш вопрос (или 'выход' для завершения).")
    while True:
        user_input = input("\n> ").strip()
        if user_input.lower() in {"выход", "exit", "quit"}:
            print("До свидания!")
            break
        answer = await ask_agent(user_input)
        print(f"Ответ:\n{answer}")

if __name__ == "__main__":
    asyncio.run(main())