feat: solution for 'Повторный экзамен: FAQ-бот — ChromaDB + один MCP-tool'
This commit is contained in:
@@ -0,0 +1,46 @@
|
|||||||
|
import os
|
||||||
|
from langchain_ollama import Ollama
|
||||||
|
from langchain.agents import Tool, initialize_agent, AgentType
|
||||||
|
from src.vector_store_utils import search_course_docs
|
||||||
|
from src.mcp_utils import fetch_course_meta
|
||||||
|
|
||||||
|
def chroma_search_tool(collection) -> Tool:
|
||||||
|
def _search(query: str) -> str:
|
||||||
|
results = search_course_docs(collection, query, k=3)
|
||||||
|
if not results:
|
||||||
|
return "No relevant documents found.\nSource: chroma"
|
||||||
|
return "\n\n".join(
|
||||||
|
[f"Source {i+1}:\n{res['page_content']}" for i, res in enumerate(results)]
|
||||||
|
) + "\nSource: chroma"
|
||||||
|
return Tool(
|
||||||
|
name="Chroma Search",
|
||||||
|
func=_search,
|
||||||
|
description="Search the FAQ stored in Chroma. Use this for general course questions."
|
||||||
|
)
|
||||||
|
|
||||||
|
def mcp_meta_tool() -> Tool:
|
||||||
|
def _meta(query: str) -> str:
|
||||||
|
return fetch_course_meta(query) + "\nSource: mcp_meta"
|
||||||
|
return Tool(
|
||||||
|
name="MCP Metadata",
|
||||||
|
func=_meta,
|
||||||
|
description="Fetch metadata about the course from the MCP service."
|
||||||
|
)
|
||||||
|
|
||||||
|
def create_agent(collection):
|
||||||
|
tools = [chroma_search_tool(collection), mcp_meta_tool()]
|
||||||
|
llm = Ollama(model="llama3")
|
||||||
|
system_prompt = (
|
||||||
|
"You are a helpful assistant for a course. "
|
||||||
|
"Use the 'Chroma Search' tool for general FAQ questions. "
|
||||||
|
"Use the 'MCP Metadata' tool for questions about course schedule, modules, or lessons. "
|
||||||
|
"Always include a source tag in your answer: 'Source: chroma' or 'Source: mcp_meta'."
|
||||||
|
)
|
||||||
|
agent = initialize_agent(
|
||||||
|
tools,
|
||||||
|
llm,
|
||||||
|
agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
|
||||||
|
verbose=False,
|
||||||
|
agent_kwargs={"system_message": system_prompt},
|
||||||
|
)
|
||||||
|
return agent
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
from src.vector_store_utils import load_faq_to_chroma
|
||||||
|
from src.langchain_agent import create_agent
|
||||||
|
|
||||||
|
def run_preset_questions(preset: str, agent):
|
||||||
|
questions = [q.strip() for q in preset.split(",") if q.strip()]
|
||||||
|
for q in questions:
|
||||||
|
print(f"\nQuestion: {q}")
|
||||||
|
answer = agent.run(q)
|
||||||
|
print(f"Answer:\n{answer}")
|
||||||
|
|
||||||
|
def interactive_mode(agent):
|
||||||
|
print("Enter your question (type 'exit' to quit):")
|
||||||
|
while True:
|
||||||
|
q = input("> ")
|
||||||
|
if q.lower() in ("exit", "quit"):
|
||||||
|
break
|
||||||
|
answer = agent.run(q)
|
||||||
|
print(f"Answer:\n{answer}")
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="FAQ Bot CLI")
|
||||||
|
parser.add_argument("--preset", type=str, help="Comma‑separated preset questions")
|
||||||
|
parser.add_argument("--interactive", action="store_true", help="Interactive mode")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
# Load vector store
|
||||||
|
collection = load_faq_to_chroma()
|
||||||
|
agent = create_agent(collection)
|
||||||
|
|
||||||
|
if args.preset:
|
||||||
|
run_preset_questions(args.preset, agent)
|
||||||
|
elif args.interactive:
|
||||||
|
interactive_mode(agent)
|
||||||
|
else:
|
||||||
|
print("No mode selected. Use --preset or --interactive.")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
import os
|
||||||
|
import json
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
def fetch_course_meta(query: str) -> str:
|
||||||
|
"""
|
||||||
|
Fetch metadata about the course from the MCP service.
|
||||||
|
Tries to GET from MCP_ENDPOINT; falls back to local JSON file.
|
||||||
|
"""
|
||||||
|
endpoint = os.getenv("MCP_ENDPOINT")
|
||||||
|
if endpoint:
|
||||||
|
try:
|
||||||
|
resp = httpx.get(endpoint, timeout=5.0)
|
||||||
|
resp.raise_for_status()
|
||||||
|
data = resp.json()
|
||||||
|
return f"{query}: {data.get(query, 'Not found')}"
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
# Fallback to local file
|
||||||
|
local_path = os.path.join("data", "course_meta.json")
|
||||||
|
with open(local_path, "r", encoding="utf-8") as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return f"{query}: {data.get(query, 'Not found')}"
|
||||||
@@ -0,0 +1,85 @@
|
|||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import List, Dict, Any
|
||||||
|
|
||||||
|
from chromadb import Client
|
||||||
|
from chromadb import Collection
|
||||||
|
from langchain_ollama import OllamaEmbeddings
|
||||||
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
|
||||||
|
# Constants
|
||||||
|
COLLECTION_NAME = "faq_collection"
|
||||||
|
EMBEDDING_MODEL = "nomic-embed-text"
|
||||||
|
EMBEDDING_DIM = 1024 # Adjust if the model changes
|
||||||
|
CHROMA_PATH = "./chroma_faq"
|
||||||
|
|
||||||
|
def load_faq_to_chroma() -> Collection:
|
||||||
|
"""
|
||||||
|
Load all markdown files from the data/ directory, split them into chunks,
|
||||||
|
embed them, and store them in a Chroma collection.
|
||||||
|
Returns the Chroma Collection instance.
|
||||||
|
"""
|
||||||
|
# Initialize Chroma client with persistence
|
||||||
|
client = Client(path=CHROMA_PATH)
|
||||||
|
|
||||||
|
# Create or get collection
|
||||||
|
collection = client.get_or_create_collection(
|
||||||
|
name=COLLECTION_NAME,
|
||||||
|
metadata={"hnsw:space": "cosine"},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prepare text splitter
|
||||||
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
|
|
||||||
|
# Prepare embeddings
|
||||||
|
embedder = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||||||
|
|
||||||
|
# Load markdown files
|
||||||
|
data_dir = Path("data")
|
||||||
|
docs = []
|
||||||
|
ids = []
|
||||||
|
metadatas = []
|
||||||
|
|
||||||
|
for md_file in data_dir.glob("*.md"):
|
||||||
|
text = md_file.read_text(encoding="utf-8")
|
||||||
|
chunks = splitter.split_text(text)
|
||||||
|
for idx, chunk in enumerate(chunks):
|
||||||
|
docs.append(chunk)
|
||||||
|
ids.append(f"{md_file.stem}_{idx}")
|
||||||
|
metadatas.append({"source": md_file.name})
|
||||||
|
|
||||||
|
# Embed documents
|
||||||
|
embeddings = embedder.embed_documents(docs)
|
||||||
|
|
||||||
|
# Add to collection
|
||||||
|
collection.add(
|
||||||
|
documents=docs,
|
||||||
|
ids=ids,
|
||||||
|
metadatas=metadatas,
|
||||||
|
embeddings=embeddings,
|
||||||
|
)
|
||||||
|
|
||||||
|
return collection
|
||||||
|
|
||||||
|
def search_course_docs(collection: Collection, query: str, k: int = 3) -> List[Dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Query the Chroma collection for the top k documents matching the query.
|
||||||
|
Returns a list of dicts with page_content and score.
|
||||||
|
"""
|
||||||
|
results = collection.query(
|
||||||
|
query_texts=[query],
|
||||||
|
n_results=k,
|
||||||
|
include=["documents", "distances", "metadatas"],
|
||||||
|
)
|
||||||
|
docs = []
|
||||||
|
for doc, distance, metadata in zip(
|
||||||
|
results["documents"][0], results["distances"][0], results["metadatas"][0]
|
||||||
|
):
|
||||||
|
docs.append(
|
||||||
|
{
|
||||||
|
"page_content": doc,
|
||||||
|
"score": 1 - distance, # Convert distance to similarity
|
||||||
|
"metadata": metadata,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
return docs
|
||||||
Reference in New Issue
Block a user