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

.env.example

@@ -0,0 +1,2 @@
+# Copy this file to .env and replace the placeholder with your actual Tavily API key
+TAVILY_API_KEY=your_tavily_api_key_here

.gitignore

@@ -0,0 +1,5 @@
+node_modules/
+.env
+dist/
+build/
+*.log

README.md

@@ -0,0 +1,34 @@
+# Экзамен: RAG-агент с ChromaDB и веб-поиском
+
+Главная
+Мои задания
+Экзамен: RAG-агент с ChromaDB и веб-поиском
+5Д
+EN
+Экзамен: RAG-агент с ChromaDB и веб-поиском
+Зачёт
+Версия 1
+Дедлайн сдачи: 31.08.2026
+
+В работе
+
+Редактирование ответа
+
+Заполните ответ и отправьте работу на проверку преподавателю.
+
+Тип ответа
+Текст
+Ссылка
+Файлы
+Текст ответа
+Прикреплённые файлы
+Загрузить файл
+Отправить на проверку
+Отменить
+
+Задание
+
+Практическое задание: RAG-агент с ChromaDB и веб-поиском
+Цель
+
+Построить AI-агента с локальным RAG-хранилищем на ChromaDB и веб-п

requirements.txt

@@ -0,0 +1,7 @@
+langchain
+langchain-chroma
+langchain-tavily
+langchain-ollama
+tavily-python
+chromadb
+python-dotenv

scripts/load_data.py

@@ -0,0 +1,12 @@
+import os
+from dotenv import load_dotenv
+from src.vectorstore import create_vectorstore, load_documents
+
+def main():
+    load_dotenv()
+    vectorstore = create_vectorstore(persist_directory="./chroma_db")
+    load_documents("./documents", vectorstore)
+    print("Documents loaded into ChromaDB.")
+
+if __name__ == "__main__":
+    main()

src/__init__.py

@@ -0,0 +1 @@
+# src package initialization

src/agent.py

@@ -0,0 +1,96 @@
+import os
+from dotenv import load_dotenv
+from langchain_ollama import ChatOllama
+from langchain.tools import tool
+from langchain.agents import initialize_agent, AgentType
+from langchain.schema import Document
+from langchain.vectorstores import Chroma
+
+# Global variable to hold the vector store for tool access
+_vectorstore: Chroma | None = None
+
+def set_vectorstore(vs: Chroma) -> None:
+    global _vectorstore
+    _vectorstore = vs
+
+@tool
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    """
+    Search the local knowledge base (ChromaDB) for relevant documents.
+
+    Args:
+        query (str): The search query.
+        top_k (int): Number of top documents to return.
+
+    Returns:
+        str: Concatenated content of the top documents or a not-found message.
+    """
+    if _vectorstore is None:
+        return "Vector store not initialized."
+
+    retriever = _vectorstore.as_retriever(search_kwargs={"k": top_k})
+    docs = retriever.get_relevant_documents(query)
+
+    if not docs:
+        return "No relevant documents found in local knowledge base."
+
+    return "\n---\n".join(
+        [f"Document {i+1}:\n{doc.page_content}" for i, doc in enumerate(docs)]
+    )
+
+@tool
+def web_search(query: str) -> str:
+    """
+    Perform a web search using Tavily.
+
+    Args:
+        query (str): The search query.
+
+    Returns:
+        str: Concatenated content of the top search results or a not-found message.
+    """
+    from tavily import TavilyClient
+
+    api_key = os.getenv("TAVILY_API_KEY")
+    if not api_key:
+        return "TAVILY_API_KEY not set in environment."
+
+    client = TavilyClient(api_key=api_key)
+    try:
+        results = client.search(query, max_results=3)
+    except Exception as e:
+        return f"Web search failed: {e}"
+
+    if not results:
+        return "No results found on the web."
+
+    return "\n---\n".join(
+        [f"Result {i+1}:\n{res.get('content', '')}" for i, res in enumerate(results)]
+    )
+
+def create_agent(vectorstore: Chroma):
+    """
+    Create a LangChain agent that can route queries to either the local KB or the web.
+
+    Args:
+        vectorstore (Chroma): The vector store to use for local search.
+
+    Returns:
+        AgentExecutor: The configured agent.
+    """
+    set_vectorstore(vectorstore)
+
+    llm = ChatOllama(model="llama3")
+
+    tools = [search_local_kb, web_search]
+
+    system_prompt = """
+You are an AI assistant. When answering a question, decide whether the answer can be found in the local knowledge base or requires up-to-date information from the web. Use the tool search_local_kb if the answer is in the local knowledge base. Use web_search if the answer requires recent information. Always include the source of the information in your answer, either 'chromadb' or 'tavily'. Do not call both tools unless necessary. If you call a tool, the tool will return the content. Use that content to answer the question. Do not mention the tool usage in your answer. Just provide the answer and the source."""
+    agent = initialize_agent(
+        tools=tools,
+        llm=llm,
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        verbose=False,
+        agent_kwargs={"system_message": system_prompt},
+    )
+    return agent

src/main.py

@@ -0,0 +1,42 @@
+import os
+from dotenv import load_dotenv
+from src.vectorstore import create_vectorstore, load_documents
+from src.agent import create_agent
+
+def main():
+    load_dotenv()
+    # Initialize or load the vector store
+    vectorstore = create_vectorstore(persist_directory="./chroma_db")
+
+    # Load documents into the vector store if not already loaded
+    # (Chroma will load existing data automatically)
+    load_documents("./documents", vectorstore)
+
+    # Create the agent
+    agent = create_agent(vectorstore)
+
+    print("\n=== RAG Agent with ChromaDB and Tavily ===")
+    print("Type your question (or 'exit' to quit):")
+
+    while True:
+        try:
+            user_input = input("\nYou: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
+            break
+
+        if not user_input:
+            continue
+
+        if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
+            break
+
+        try:
+            response = agent.run(user_input)
+            print(f"\nAgent: {response}")
+        except Exception as e:
+            print(f"Error: {e}")
+
+if __name__ == "__main__":
+    main()

src/vectorstore.py

@@ -0,0 +1,47 @@
+import os
+from pathlib import Path
+from langchain_ollama import OllamaEmbeddings
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain.schema import Document
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """
+    Create or load a Chroma vector store with Ollama embeddings.
+
+    Args:
+        persist_directory (str): Directory where the Chroma DB is stored.
+
+    Returns:
+        Chroma: The Chroma vector store instance.
+    """
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+
+def load_documents(directory: str, vectorstore: Chroma) -> None:
+    """
+    Load .txt and .md files from a directory, split them into chunks,
+    and add them to the provided vector store.
+
+    Args:
+        directory (str): Path to the directory containing documents.
+        vectorstore (Chroma): The vector store to add documents to.
+    """
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+
+    for file_path in Path(directory).glob("**/*"):
+        if file_path.suffix.lower() in [".txt", ".md"]:
+            try:
+                text = file_path.read_text(encoding="utf-8")
+                chunks = splitter.split_text(text)
+                docs.extend([Document(page_content=chunk) for chunk in chunks])
+            except Exception as e:
+                print(f"Failed to read {file_path}: {e}")
+
+    if docs:
+        vectorstore.add_documents(docs)
+        vectorstore.persist()
+        print(f"Loaded {len(docs)} chunks into ChromaDB.")
+    else:
+        print("No documents found to load.")