167 lines
6.5 KiB
Python
167 lines
6.5 KiB
Python
"""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 sys
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
import httpx
|
||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||
from langchain_chroma import Chroma
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. Векторная база
|
||
# ---------------------------------------------------------------------------
|
||
CHROMA_PATH = Path("./chroma_faq")
|
||
DATA_PATH = Path("./data")
|
||
|
||
# Создаём embeddings
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Чтение markdown‑файлов и загрузка в Chroma
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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
|
||
|
||
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:
|
||
"""Возвращает кортеж из 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:
|
||
"""Имитирует вызов MCP‑сервера. Читает локальный JSON и возвращает
|
||
информацию, содержащуюся в ключе, совпадающем с 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."
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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."
|
||
)
|
||
|
||
# Создаём LLM
|
||
llm = ChatOllama(model="llama3", temperature=0.2)
|
||
|
||
# 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
|
||
]
|
||
|
||
def run_cli():
|
||
print("--- FAQ Bot ---")
|
||
print("Predefined questions:")
|
||
for i, q in enumerate(PREDEFINED_QUESTIONS, 1):
|
||
print(f"{i}. {q}")
|
||
print("\nEnter 1-3 to ask a predefined question, or type your own question.")
|
||
while True:
|
||
user_input = input("\n> ")
|
||
if user_input.lower() in {"exit", "quit"}:
|
||
print("Goodbye!")
|
||
break
|
||
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__":
|
||
# Убедимся, что данные загружены
|
||
load_faq_to_chroma()
|
||
run_cli()
|