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

2026-06-30 12:03:06 +03:00
parent ef8bb1bd2c
commit 4f5e00efdb
5 changed files with 289 additions and 65 deletions
@@ -1,51 +1,91 @@
 # RAG Agent with ChromaDB and Web Search

-This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector store and performs web search as a fallback. The agent is written in Node.js and uses only the required dependencies.
+This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** as the vector database and performs live web searches to provide up‑to‑date information.

 ## Features

- **Vector storage** with ChromaDB (in-memory by default).
- **Simple embedding** function (placeholder) – replace with a real model for production.
- **Web search** using DuckDuckGo’s HTML interface.
- **RAG agent** that retrieves relevant documents or falls back to web search.
+- **Vector store** – Documents are ingested, split into chunks, embedded with OpenAI embeddings, and stored in a persistent ChromaDB collection.
+- **Web search** – Uses DuckDuckGo scraping to fetch recent web snippets for a query.
+- **RAG pipeline** – Combines local document context and web results, then generates an answer with OpenAI GPT‑3.5‑Turbo.
+- **CLI** – Simple command line interface for ingestion and querying.

-## Installation
+## Prerequisites
+
+- Python 3.10+
+- An OpenAI API key with access to `text-embedding-ada-002` and `gpt-3.5-turbo`.
+
+## Setup

 ```bash
-npm install
+# Clone the repository
+git clone https://git.brojs.ru/kuzakhmetovartur/ekzamen-rag-agent-s-chromadb-i-veb-poisk.git
+cd ekzamen-rag-agent-s-chromadb-i-veb-poisk
+
+# Create a virtual environment (optional but recommended)
+python -m venv .venv
+source .venv/bin/activate   # On Windows use `.venv\Scripts\activate`
+
+# Install dependencies
+pip install -r requirements.txt
 ```

+## Configuration
+
+Create a `.env` file in the project root (or set environment variables directly):
+
+```
+OPENAI_API_KEY=sk-...
+CHROMA_DB_PATH=./chromadb
+CHROMA_COLLECTION_NAME=rag_collection
+```
+
+> **Note**: Do not commit your `.env` file or API key to version control.
+
 ## Usage

-```bash
-node src/index.js "Your query here"
-```
-
-If no query is provided, it defaults to `"What is ChromaDB?"`.
-
-## Running Tests
+### 1. Ingest Documents

 ```bash
-npm test
+python src/main.py ingest path/to/doc1.txt path/to/doc2.txt
 ```

+The script will read each file, split it into chunks, generate embeddings, and store them in ChromaDB.
+
+### 2. Query the Agent
+
+```bash
+python src/main.py query "What is the capital of France?"
+```
+
+The agent will:
+
+1. Retrieve relevant chunks from the local vector store.
+2. Perform a DuckDuckGo web search for the query.
+3. Combine both sources of information.
+4. Generate a response using OpenAI GPT‑3.5‑Turbo.
+
 ## Project Structure

 ```
 src/
-  index.js        # Entry point
-  agent.js        # RAG agent logic
-  vectorStore.js  # ChromaDB wrapper
-  webSearch.js    # Simple web search helper
-  test.js         # Basic test for vector store
+├── main.py          # CLI entry point
+├── vector_store.py  # ChromaDB ingestion & retrieval
+├── web_search.py    # DuckDuckGo web search
+requirements.txt
+README.md
 ```

-## Extending
+## Testing

- Replace the `embed` function in `vectorStore.js` with a real embedding model (e.g., OpenAI, HuggingFace).
- Persist the ChromaDB collection by configuring the client with a storage path.
- Add a language model to generate responses from retrieved documents.
+The project can be tested with `pytest` (tests are not included in this minimal example).  
+If you add tests, run:
+
+```bash
+pytest
+```

 ## License

-MIT
+MIT License
+--- 
+Feel free to extend the agent with additional features such as custom embeddings, different LLMs, or alternative search APIs.
@@ -1,8 +1,5 @@
-langchain
-langchain-ollama
-langchain-qdrant
-langchain-tavily
-tavily-python
-chromadb
-python-dotenv
-qdrant-client
+chromadb==0.4.22
+openai==1.12.0
+requests==2.31.0
+beautifulsoup4==4.12.3
+python-dotenv==1.0.1
@@ -1,42 +1,79 @@
+import argparse
 import os
-from dotenv import load_dotenv
-from src.vectorstore import create_vectorstore, load_documents
-from src.agent import create_agent
+import sys
+from typing import List
+
+import openai
+
+from vector_store import ingest_documents, get_relevant_chunks
+from web_search import search_web
+
+# Load OpenAI API key
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    print("Error: OPENAI_API_KEY environment variable not set.")
+    sys.exit(1)
+openai.api_key = OPENAI_API_KEY
+
+def generate_answer(context: str, question: str) -> str:
+    """
+    Generate an answer using OpenAI ChatCompletion with the provided context.
+    """
+    system_prompt = "You are a helpful assistant. Use the provided context to answer the question."
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
+    ]
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=messages,
+            temperature=0.2,
+            max_tokens=512
+        )
+        return response["choices"][0]["message"]["content"].strip()
+    except Exception as e:
+        print(f"OpenAI request failed: {e}")
+        return ""
+
+def ingest_mode(file_paths: List[str]) -> None:
+    ingest_documents(file_paths)
+
+def query_mode(question: str) -> None:
+    # Retrieve relevant chunks from local vector store
+    local_chunks = get_relevant_chunks(question, k=5)
+    local_context = "\n\n".join([chunk for _, chunk in local_chunks])
+
+    # Perform web search for up-to-date info
+    web_snippets = search_web(question, num_results=3)
+    web_context = "\n\n".join(web_snippets)
+
+    # Combine contexts
+    combined_context = f"Local documents:\n{local_context}\n\nWeb results:\n{web_context}"
+
+    # Generate answer
+    answer = generate_answer(combined_context, question)
+    print("\nAnswer:\n")
+    print(answer)

 def main():
-    load_dotenv()
-    # Initialize or load the vector store
-    vectorstore = create_vectorstore(persist_directory="./chroma_db")
+    parser = argparse.ArgumentParser(description="RAG Agent with ChromaDB and Web Search")
+    subparsers = parser.add_subparsers(dest="command", required=True)

-    # Load documents into the vector store if not already loaded
-    # (Chroma will load existing data automatically)
-    load_documents("./documents", vectorstore)
+    ingest_parser = subparsers.add_parser("ingest", help="Ingest documents into the vector store")
+    ingest_parser.add_argument("files", nargs="+", help="Paths to text files to ingest")

-    # Create the agent
-    agent = create_agent(vectorstore)
+    query_parser = subparsers.add_parser("query", help="Ask a question to the RAG agent")
+    query_parser.add_argument("question", help="The question to ask")

-    print("\n=== RAG Agent with ChromaDB and Tavily ===")
-    print("Type your question (or 'exit' to quit):")
+    args = parser.parse_args()

-    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 args.command == "ingest":
+        ingest_mode(args.files)
+    elif args.command == "query":
+        query_mode(args.question)
+    else:
+        parser.print_help()

 if __name__ == "__main__":
    main()
@@ -0,0 +1,85 @@
+import os
+from typing import List, Tuple
+
+import chromadb
+from chromadb import PersistentClient
+from chromadb.config import Settings
+import openai
+
+# Load OpenAI API key from environment
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    raise RuntimeError("OPENAI_API_KEY environment variable not set.")
+openai.api_key = OPENAI_API_KEY
+
+# ChromaDB persistent client settings
+CHROMA_DB_PATH = os.getenv("CHROMA_DB_PATH", "./chromadb")
+CHROMA_COLLECTION_NAME = os.getenv("CHROMA_COLLECTION_NAME", "rag_collection")
+
+# Initialize Chroma client
+client = PersistentClient(path=CHROMA_DB_PATH, settings=Settings(chroma_api_impl="chromadb.api.fastapi.FastAPI"))
+collection = client.get_or_create_collection(name=CHROMA_COLLECTION_NAME)
+
+
+def _split_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
+    """
+    Split text into chunks of approximately chunk_size characters with overlap.
+    """
+    chunks = []
+    start = 0
+    text_length = len(text)
+    while start < text_length:
+        end = min(start + chunk_size, text_length)
+        chunk = text[start:end]
+        chunks.append(chunk)
+        start += chunk_size - overlap
+    return chunks
+
+
+def _embed_text(text: str) -> List[float]:
+    """
+    Generate embedding for a single text string using OpenAI embeddings.
+    """
+    response = openai.Embedding.create(
+        model="text-embedding-ada-002",
+        input=text
+    )
+    return response["data"][0]["embedding"]
+
+
+def ingest_documents(file_paths: List[str]) -> None:
+    """
+    Ingest a list of file paths into the Chroma collection.
+    Each file is read, split into chunks, embedded, and stored.
+    """
+    for file_path in file_paths:
+        if not os.path.isfile(file_path):
+            print(f"Skipping non-existent file: {file_path}")
+            continue
+        with open(file_path, "r", encoding="utf-8") as f:
+            content = f.read()
+        chunks = _split_text(content)
+        embeddings = [_embed_text(chunk) for chunk in chunks]
+        ids = [f"{os.path.basename(file_path)}_{i}" for i in range(len(chunks))]
+        collection.add(
+            ids=ids,
+            documents=chunks,
+            embeddings=embeddings
+        )
+        print(f"Ingested {len(chunks)} chunks from {file_path}.")
+
+
+def get_relevant_chunks(query: str, k: int = 5) -> List[Tuple[str, str]]:
+    """
+    Retrieve top-k relevant chunks for a query.
+    Returns a list of tuples (chunk_id, chunk_text).
+    """
+    query_embedding = _embed_text(query)
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=k,
+        include=["documents", "ids"]
+    )
+    ids = results["ids"][0]
+    docs = results["documents"][0]
+    return list(zip(ids, docs))
@@ -0,0 +1,65 @@
+import os
+import re
+import requests
+from bs4 import BeautifulSoup
+from typing import List
+
+# DuckDuckGo search URL
+DDG_SEARCH_URL = "https://duckduckgo.com/html/"
+
+def _extract_text_from_html(html: str) -> str:
+    """
+    Extract visible text from HTML, removing scripts and styles.
+    """
+    soup = BeautifulSoup(html, "html.parser")
+    for script in soup(["script", "style"]):
+        script.decompose()
+    text = soup.get_text(separator="\n")
+    lines = (line.strip() for line in text.splitlines())
+    chunks = [phrase.strip() for phrase in lines if phrase.strip()]
+    return "\n".join(chunks)
+
+def search_web(query: str, num_results: int = 3) -> List[str]:
+    """
+    Perform a web search using DuckDuckGo and return the top num_results snippets.
+    """
+    params = {
+        "q": query,
+        "s": "0",
+        "dc": "0",
+        "kl": "us-en",
+        "kp": "-2",
+        "kp": "-2",
+        "kp": "-2",
+        "kp": "-2",
+    }
+    headers = {
+        "User-Agent": "Mozilla/5.0 (compatible; RAG-Agent/1.0; +https://example.com/bot)"
+    }
+    try:
+        response = requests.get(DDG_SEARCH_URL, params=params, headers=headers, timeout=10)
+        response.raise_for_status()
+    except requests.RequestException as e:
+        print(f"Web search request failed: {e}")
+        return []
+
+    soup = BeautifulSoup(response.text, "html.parser")
+    results = []
+    for a in soup.select("a.result__a"):
+        href = a.get("href")
+        if href:
+            results.append(href)
+        if len(results) >= num_results:
+            break
+
+    snippets = []
+    for url in results:
+        try:
+            page_resp = requests.get(url, headers=headers, timeout=10)
+            page_resp.raise_for_status()
+            snippet = _extract_text_from_html(page_resp.text)[:500]  # limit snippet size
+            snippets.append(snippet)
+        except requests.RequestException:
+            continue
+
+    return snippets