163 lines
6.3 KiB
Python
163 lines
6.3 KiB
Python
<|channel|>final code<|message|>import os
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import httpx
|
||
from dotenv import load_dotenv
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_chroma import Chroma
|
||
from langchain_core.prompts import PromptTemplate
|
||
from langchain.agents import initialize_agent, Tool, AgentType
|
||
from langchain.schema import Document
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 1. Загрузка FAQ в ChromaDB
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
DATA_DIR = Path("data")
|
||
CHROMA_PATH = Path("./chroma_faq")
|
||
|
||
def load_faq_to_chroma() -> Chroma:
|
||
"""
|
||
Читает все *.md файлы из DATA_DIR, разбивает их на чанки и сохраняет
|
||
в локальный ChromaDB. Возвращает объект Chroma.
|
||
"""
|
||
# Создаём список документов
|
||
docs = []
|
||
for md_file in DATA_DIR.glob("*.md"):
|
||
text = md_file.read_text(encoding="utf-8")
|
||
docs.append(Document(page_content=text, metadata={"source": str(md_file)}))
|
||
|
||
# Инициализируем Chroma с Ollama embeddings
|
||
chroma = Chroma.from_documents(
|
||
documents=docs,
|
||
embedding=ChatOllama(model="nomic-embed-text"),
|
||
persist_directory=str(CHROMA_PATH),
|
||
)
|
||
return chroma
|
||
|
||
def search_course_docs(query: str, k: int = 3) -> list[Document]:
|
||
"""
|
||
Поиск в ChromaDB. Возвращает список документов с наиболее релевантными
|
||
ответами на запрос.
|
||
"""
|
||
chroma = load_faq_to_chroma()
|
||
return chroma.similarity_search_with_score(query, k=k)
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 2. MCP‑style tool (mock)
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def fetch_course_meta(query: str) -> dict:
|
||
"""
|
||
Имитация вызова внешнего MCP сервера. Читает статический JSON
|
||
из файла mock_meta.json и возвращает часть, соответствующую query.
|
||
"""
|
||
meta_file = Path("mock_meta.json")
|
||
if not meta_file.exists():
|
||
raise FileNotFoundError(f"{meta_file} не найден")
|
||
|
||
data = json.loads(meta_file.read_text(encoding="utf-8"))
|
||
# простая фильтрация по ключу
|
||
return {k: v for k, v in data.items() if query.lower() in k.lower()}
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 3. Agent + Tools
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def create_agent():
|
||
"""
|
||
Создаёт агента LangChain с двумя инструментами:
|
||
- search_course_docs (Chroma)
|
||
- fetch_course_meta (MCP‑style)
|
||
Агент использует Ollama LLM.
|
||
"""
|
||
llm = ChatOllama(model="nomic-embed-text")
|
||
|
||
tools = [
|
||
Tool(
|
||
name="search_course_docs",
|
||
func=lambda q: "\n".join([f"{d.page_content[:200]}..." for d in search_course_docs(q)]),
|
||
description=(
|
||
"Используется, если вопрос относится к материалам курса. "
|
||
"Возвращает краткие отрывки из FAQ."
|
||
),
|
||
),
|
||
Tool(
|
||
name="fetch_course_meta",
|
||
func=lambda q: json.dumps(fetch_course_meta(q), indent=2),
|
||
description=(
|
||
"Вызывается, если нужно получить расписание или метаданные курса. "
|
||
"Возвращает JSON‑объект."
|
||
),
|
||
),
|
||
]
|
||
|
||
system_prompt = (
|
||
"Ты FAQ‑бот для локального курса. "
|
||
"Если вопрос относится к материалам занятия — используй search_course_docs. "
|
||
"Если нужен график, даты или другая метаинформация — fetch_course_meta. "
|
||
"Не вызывай оба инструмента одновременно. В ответе указывай источник: "
|
||
"`source: chroma` или `source: mcp_meta`."
|
||
)
|
||
|
||
prompt = PromptTemplate(
|
||
input_variables=["input"],
|
||
template=system_prompt + "\n\nВопрос: {input}\nОтвет:",
|
||
)
|
||
|
||
agent = initialize_agent(
|
||
tools,
|
||
llm,
|
||
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
|
||
verbose=True,
|
||
max_iterations=5,
|
||
early_stopping_method="generate",
|
||
system_message=prompt.format(input="{input}"),
|
||
)
|
||
return agent
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 4. CLI
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def main():
|
||
"""
|
||
Простая консольная утилита с тремя готовыми вопросами и интерактивным режимом.
|
||
"""
|
||
print("=== FAQ‑бот для курса ===")
|
||
sample_questions = [
|
||
"Какой материал будет на следующем занятии?",
|
||
"Когда начинается следующий модуль?",
|
||
"Какие темы охватывает первый раздел?",
|
||
]
|
||
|
||
for idx, q in enumerate(sample_questions, 1):
|
||
print(f"\n{idx}. {q}")
|
||
|
||
print("\nВведите номер вопроса (1‑3) или '0' для интерактивного режима.")
|
||
choice = input("Ваш выбор: ").strip()
|
||
|
||
if choice == "0":
|
||
while True:
|
||
q = input("\nВопрос (или 'exit'): ")
|
||
if q.lower() in ("exit", "quit"):
|
||
break
|
||
agent = create_agent()
|
||
response = agent.run(q)
|
||
print(f"\nОтвет:\n{response}\n")
|
||
else:
|
||
try:
|
||
idx = int(choice) - 1
|
||
question = sample_questions[idx]
|
||
except (ValueError, IndexError):
|
||
print("Неверный выбор.")
|
||
return
|
||
|
||
agent = create_agent()
|
||
response = agent.run(question)
|
||
print(f"\nОтвет:\n{response}\n")
|
||
|
||
if __name__ == "__main__":
|
||
load_dotenv() # если понадобится переменные окружения
|
||
main() |