feat: solution for 'Экзамен: RAG-агент с ChromaDB и веб-поиском'

This commit is contained in:
2026-06-29 11:49:00 +03:00
parent c0aac04442
commit d6805973d6
6 changed files with 371 additions and 9 deletions
+7 -6
View File
@@ -7,11 +7,15 @@
EN EN
Экзамен: RAG-агент с ChromaDB и веб-поиском Экзамен: RAG-агент с ChromaDB и веб-поиском
Зачёт Зачёт
Версия 1 Версия 2
Дедлайн сдачи: 31.08.2026 Дедлайн сдачи: 31.08.2026
В работе В работе
Требуется доработка
Переделайте решение: используйте QDrant вместо текущего векторного хранилища.
Редактирование ответа Редактирование ответа
Заполните ответ и отправьте работу на проверку преподавателю. Заполните ответ и отправьте работу на проверку преподавателю.
@@ -20,7 +24,7 @@ EN
Текст Текст
Ссылка Ссылка
Файлы Файлы
Текст ответа Ссылка (URL)
Прикреплённые файлы Прикреплённые файлы
Загрузить файл Загрузить файл
Отправить на проверку Отправить на проверку
@@ -28,7 +32,4 @@ EN
Задание Задание
Практическое задание: RAG-агент с ChromaDB и веб-поиском Практическое задание: RAG-аге
Цель
Построить AI-агента с локальным RAG-хранилищем на ChromaDB и веб-п
+99
View File
@@ -0,0 +1,99 @@
import os
from typing import Any, Dict
from langchain_ollama import Ollama
from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
from langchain.schema import AgentAction, AgentFinish
from langchain.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from vectorstore import QdrantVectorStore
from tools import search_local_kb, web_search
def create_agent(
vectorstore: QdrantVectorStore,
tavily_api_key: str,
llm_model: str = "llama3",
) -> AgentExecutor:
"""
Create a LangChain agent that routes queries to either the local KB or the web.
Parameters
----------
vectorstore : QdrantVectorStore
The vector store for local knowledge base.
tavily_api_key : str
Tavily API key.
llm_model : str
Ollama model name for LLM.
Returns
-------
AgentExecutor
Configured agent executor.
"""
# LLM
llm = Ollama(model=llm_model)
# Define tools with partial application of required arguments
tools = [
Tool(
name="search_local_kb",
func=lambda query, top_k=3: search_local_kb(
query=query, top_k=top_k, vectorstore=vectorstore
),
description=(
"Use this tool to search the local knowledge base. "
"Return the most relevant snippets."
),
),
Tool(
name="web_search",
func=lambda query, max_results=3: web_search(
query=query, tavily_api_key=tavily_api_key, max_results=max_results
),
description=(
"Use this tool to search the web via Tavily. "
"Return the most relevant snippets."
),
),
]
# Prompt template for the agent
system_prompt = (
"You are an AI assistant that can answer questions using either "
"the local knowledge base or the web. If the question is about "
"recent events, news, or requires up-to-date information, "
"use the web_search tool. If the question is about "
"information contained in the local documents, use the "
"search_local_kb tool. After retrieving the information, "
"provide a concise answer and state the source (chromadb or tavily)."
)
prompt = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
MessagesPlaceholder("history"),
("human", "{input}"),
MessagesPlaceholder("agent_scratchpad"),
]
)
# Output parser
output_parser = StrOutputParser()
# Agent
agent = initialize_agent(
tools=tools,
llm=llm,
agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
verbose=False,
handle_parsing_errors=True,
max_iterations=5,
early_stopping_method="generate",
system_message=system_prompt,
)
return agent
+65
View File
@@ -0,0 +1,65 @@
import os
import sys
from pathlib import Path
from dotenv import load_dotenv
from vectorstore import create_vectorstore, load_documents, collection_exists
from agent import create_agent
def main():
# Load environment variables
load_dotenv()
# Configuration
qdrant_path = os.getenv("QDRANT_PATH", "./qdrant_db")
embedding_model = os.getenv("EMBEDDING_MODEL", "nomic-embed-text")
llm_model = os.getenv("LLM_MODEL", "llama3")
tavily_api_key = os.getenv("TAVILY_API_KEY")
if not tavily_api_key:
print("Error: TAVILY_API_KEY not set in .env")
sys.exit(1)
# Create vector store
vectorstore = create_vectorstore(
persist_directory=qdrant_path,
collection_name="documents",
embedding_model=embedding_model,
)
# Load documents if collection is empty
if not collection_exists(vectorstore):
print("Loading documents into Qdrant...")
docs_dir = Path("documents")
if not docs_dir.exists():
print(f"Documents directory '{docs_dir}' not found.")
sys.exit(1)
load_documents(str(docs_dir), vectorstore)
print("Documents loaded.")
else:
print("Qdrant collection already exists. Skipping document load.")
# Create agent
agent = create_agent(vectorstore, tavily_api_key, llm_model=llm_model)
print("Chat agent ready. Type 'exit' to quit.")
while True:
try:
user_input = input("\nYou: ")
except (KeyboardInterrupt, EOFError):
print("\nExiting.")
break
if user_input.strip().lower() in {"exit", "quit"}:
print("Goodbye!")
break
# Run agent
try:
response = agent.run(user_input)
print(f"\nAssistant: {response}")
except Exception as e:
print(f"Error: {e}")
if __name__ == "__main__":
main()
+4 -3
View File
@@ -1,7 +1,8 @@
langchain langchain
langchain-chroma
langchain-tavily
langchain-ollama langchain-ollama
langchain-qdrant
langchain-tavily
tavily-python tavily-python
chromadb chromadb
python-dotenv python-dotenv
qdrant-client
+71
View File
@@ -0,0 +1,71 @@
from typing import List
from langchain.tools import tool
from langchain_ollama import Ollama
from langchain_qdrant import QdrantVectorStore
from langchain_tavily import TavilySearchResults
@tool
def search_local_kb(
query: str,
top_k: int,
vectorstore: QdrantVectorStore,
) -> List[str]:
"""
Perform a semantic search in the local knowledge base stored in Qdrant.
Parameters
----------
query : str
The user's query.
top_k : int
Number of top results to return.
vectorstore : QdrantVectorStore
The vector store to search.
Returns
-------
List[str]
List of relevant document snippets.
"""
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
docs = retriever.get_relevant_documents(query)
return [doc.page_content for doc in docs]
@tool
def web_search(
query: str,
tavily_api_key: str,
max_results: int = 3,
) -> List[str]:
"""
Perform a web search using Tavily.
Parameters
----------
query : str
The user's query.
tavily_api_key : str
Tavily API key.
max_results : int
Number of search results to return.
Returns
-------
List[str]
List of search result snippets.
"""
tavily = TavilySearchResults(
api_key=tavily_api_key,
max_results=max_results,
)
results = tavily.run(query)
# Extract snippets from results
snippets = []
for result in results:
snippet = result.get("content") or result.get("snippet") or result.get("title")
if snippet:
snippets.append(snippet)
return snippets
+125
View File
@@ -0,0 +1,125 @@
import os
from pathlib import Path
from typing import Optional
from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
def create_vectorstore(
persist_directory: str = "./qdrant_db",
collection_name: str = "documents",
embedding_model: str = "nomic-embed-text",
) -> QdrantVectorStore:
"""
Create a Qdrant vector store backed by Ollama embeddings.
Parameters
----------
persist_directory : str
Path to the directory where Qdrant will store its data.
collection_name : str
Name of the collection to use.
embedding_model : str
Ollama model name for embeddings.
Returns
-------
QdrantVectorStore
Initialized vector store.
"""
# Ensure the persistence directory exists
Path(persist_directory).mkdir(parents=True, exist_ok=True)
# Create a local Qdrant client that stores data on disk
client = QdrantClient(path=persist_directory)
# Initialize embeddings
embeddings = OllamaEmbeddings(model=embedding_model)
# Create or load the collection
vectorstore = QdrantVectorStore(
client=client,
embeddings=embeddings,
collection_name=collection_name,
)
return vectorstore
def load_documents(
directory: str,
vectorstore: QdrantVectorStore,
chunk_size: int = 1000,
chunk_overlap: int = 200,
) -> None:
"""
Load all .txt and .md files from a directory, split them into chunks,
embed, and add to the vector store.
Parameters
----------
directory : str
Path to the directory containing documents.
vectorstore : QdrantVectorStore
The vector store to populate.
chunk_size : int
Maximum size of each chunk in characters.
chunk_overlap : int
Number of overlapping characters between chunks.
"""
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
# Prepare text splitter
splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
)
# Collect all documents
docs = []
for root, _, files in os.walk(directory):
for file in files:
if file.lower().endswith((".txt", ".md")):
file_path = os.path.join(root, file)
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
# Split content into chunks
chunks = splitter.split_text(content)
# Create Document objects with metadata
for i, chunk in enumerate(chunks):
doc = Document(
page_content=chunk,
metadata={
"source": file_path,
"chunk_index": i,
},
)
docs.append(doc)
if docs:
vectorstore.add_documents(docs)
else:
print("No documents found to load.")
def collection_exists(vectorstore: QdrantVectorStore) -> bool:
"""
Check if the collection already exists in Qdrant.
Parameters
----------
vectorstore : QdrantVectorStore
The vector store to check.
Returns
-------
bool
True if collection exists, False otherwise.
"""
try:
vectorstore.client.get_collection(vectorstore.collection_name)
return True
except Exception:
return False