74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
import os
|
|
import json
|
|
import httpx
|
|
from pathlib import Path
|
|
from langchain_ollama import ChatOllama
|
|
from langchain_chroma import Chroma
|
|
from langchain.embeddings import OllamaEmbeddings
|
|
from langchain.schema import Document
|
|
from langchain.prompts import ChatPromptTemplate
|
|
from langchain.chains import LLMChain
|
|
|
|
# 1. Load FAQ into Chroma
|
|
|
|
def load_faq_to_chroma(md_path: str, persist_dir: str = "./chroma_faq"):
|
|
from langchain.document_loaders import TextLoader
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
loader = TextLoader(md_path)
|
|
docs = loader.load()
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
texts = splitter.split_documents(docs)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
chroma = Chroma.from_documents(texts, embeddings, persist_directory=persist_dir)
|
|
chroma.persist()
|
|
return chroma
|
|
|
|
# 2. Search function
|
|
|
|
def search_course_docs(query: str, k: int = 3):
|
|
chroma = Chroma(persist_directory="./chroma_faq", embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
|
|
results = chroma.similarity_search(query, k=k)
|
|
return [doc.page_content for doc in results]
|
|
|
|
# 3. MCP-style tool
|
|
|
|
def fetch_course_meta(query: str):
|
|
# For demo, use static JSON file
|
|
meta_path = Path("meta.json")
|
|
if not meta_path.exists():
|
|
return {"error": "Meta not found"}
|
|
meta = json.loads(meta_path.read_text())
|
|
# simple search by key
|
|
return meta.get(query, {})
|
|
|
|
# 4. Agent logic
|
|
|
|
def answer_question(question: str):
|
|
# Simple heuristic: if question contains "schedule" or "метаданные" use meta
|
|
if any(word in question.lower() for word in ["schedule", "расписание", "метаданные"]):
|
|
source = "mcp_meta"
|
|
answer = fetch_course_meta(question)
|
|
else:
|
|
source = "chroma"
|
|
answer = search_course_docs(question, k=1)[0]
|
|
return {"answer": answer, "source": source}
|
|
|
|
if __name__ == "__main__":
|
|
# CLI with preset questions
|
|
preset = [
|
|
"Что такое ChromaDB?",
|
|
"Как подключить Ollama embeddings?",
|
|
"Когда будет расписание следующего занятия?"
|
|
]
|
|
for q in preset:
|
|
res = answer_question(q)
|
|
print(f"Q: {q}\nA: {res['answer']}\nSource: {res['source']}\n")
|
|
# interactive
|
|
while True:
|
|
q = input("Ask a question (or 'exit'): ")
|
|
if q.lower() == "exit":
|
|
break
|
|
res = answer_question(q)
|
|
print(f"A: {res['answer']} (source: {res['source']})")
|