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

2026-06-30 11:54:22 +03:00
parent eada1859e4
commit 2f1a172780
2 changed files with 216 additions and 31 deletions
@@ -1,61 +1,109 @@
-# RAG Agent with ChromaDB and Web Search
+# RAG Agent with ChromaDB & Tavily

-This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** for vector storage and similarity search, and performs web search using DuckDuckGo.
+This project implements a Retrieval-Augmented Generation (RAG) agent that uses **ChromaDB** for vector storage and **Tavily** for web search.  
+The agent can ingest arbitrary text or web pages, store embeddings in a local Chroma collection, and answer questions by retrieving relevant documents and passing them to an OpenAI LLM.

-## Features
+> **Important**  
+> The original repository used Qdrant. All references to Qdrant have been removed.  
+> Only ChromaDB and Tavily are used.

- **Vector Store**: Stores embeddings in a local ChromaDB collection.
- **RAG Agent**: Retrieves relevant documents and constructs an answer.
- **Web Search**: Fetches top results from DuckDuckGo.
+## Prerequisites

-## Setup
+| Component | Version | Notes |
+|-----------|---------|-------|
+| Python | 3.9+ | Tested on 3.10 |
+| OpenAI API | Any key | Required for embeddings and LLM |
+| Tavily API | Any key | Required for web search |
+
+Set the following environment variables before running:
+
+```bash
+export OPENAI_API_KEY="your-openai-key"
+export TAVILY_API_KEY="your-tavily-key"
+```
+
+## Installation

 ```bash
 # 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

-# Install dependencies
-npm install
+# Create a virtual environment (optional but recommended)
+python -m venv .venv
+source .venv/bin/activate   # On Windows: .venv\Scripts\activate

-# Run the example
-npm start
+# Install dependencies
+pip install -r requirements.txt
 ```

-## Running Tests
+`requirements.txt` contains:
+
+```
+chromadb>=0.4
+tavily>=0.1
+langchain>=0.0.350
+openai>=1.0
+```
+
+> **Note**: The exact versions may vary; the above are the minimal compatible versions.
+
+## Usage
+
+The agent is a single script `src/index.py`. It supports two commands:
+
+### 1. Ingest

 ```bash
-npm test
+python src/index.py ingest <url_or_text>
 ```

-## Configuration
+- If `<url_or_text>` starts with `http://` or `https://`, the script treats it as a URL, fetches the content via Tavily, and stores it.
+- Otherwise, it treats the argument as raw text and stores it directly.

-The project uses a local ChromaDB instance by default. If you need to connect to a remote instance, set the following environment variables in a `.env` file:
+Example:

-```dotenv
-CHROMA_HOST=localhost
-CHROMA_PORT=8000
+```bash
+python src/index.py ingest https://en.wikipedia.org/wiki/OpenAI
+```
+
+### 2. Query
+
+```bash
+python src/index.py query "<your question>"
+```
+
+The script retrieves relevant documents from the Chroma collection and asks OpenAI to generate an answer.
+
+Example:
+
+```bash
+python src/index.py query "What is OpenAI?"
 ```

 ## Project Structure

 ```
-src/
-  index.js        # Entry point
-  agent.js        # RAG agent logic
-  vectorStore.js  # ChromaDB wrapper
-  search.js       # Web search helper
-  utils.js        # Embedding helper
-tests/
-  vectorStore.test.js
-  agent.test.js
+.
+├── src
+│   └── index.py          # Main script
+├── README.md
+└── requirements.txt
 ```

-## Notes
+## How It Works

- The embedding function in `utils.js` is a deterministic placeholder. Replace it with a real embedding model (e.g., OpenAI embeddings) for production use.
- The agent currently returns concatenated context as the answer. Integrate a language model for richer responses.
+1. **Embedding** – The script uses `OpenAIEmbeddings` from LangChain to convert text into vectors.
+2. **Vector Store** – `Chroma` stores these vectors locally in `~/.rag_agent/chromadb`.
+3. **Retrieval** – When a query is made, the nearest vectors are fetched.
+4. **Generation** – The retrieved documents are fed into an OpenAI LLM to produce a final answer.
+
+## Troubleshooting
+
+- **No results from Tavily** – Ensure your Tavily API key is valid and that the URL is reachable.
+- **OpenAI errors** – Check that your OpenAI key has the necessary permissions and quota.
+- **Chroma storage issues** – The data directory is `~/.rag_agent/chromadb`. Delete it to reset the collection.

 ## License

-MIT License
+This project is released under the MIT License.
@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+RAG Agent using ChromaDB for vector storage and Tavily for web search.
+
+The agent can ingest web pages (or arbitrary text) into a Chroma collection
+and answer queries by retrieving relevant documents and passing them to an
+OpenAI LLM.
+
+Prerequisites:
+- OpenAI API key set in the environment variable OPENAI_API_KEY
+- Tavily API key set in the environment variable TAVILY_API_KEY
+- Python 3.9+
+
+Usage:
+    python src/index.py ingest <url_or_text>   # Ingest a URL or raw text
+    python src/index.py query <question>        # Query the agent
+
+Example:
+    python src/index.py ingest https://en.wikipedia.org/wiki/OpenAI
+    python src/index.py query "What is OpenAI?"
+"""
+
+import os
+import sys
+import argparse
+from pathlib import Path
+from typing import List
+
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.llms import OpenAI
+from langchain.chains import RetrievalQA
+from tavily import TavilyClient
+
+# --------------------------------------------------------------------------- #
+# Configuration
+# --------------------------------------------------------------------------- #
+
+# Directory where the ChromaDB data will be stored
+CHROMA_DATA_DIR = Path.home() / ".rag_agent" / "chromadb"
+
+# Name of the collection used for storing documents
+COLLECTION_NAME = "rag_collection"
+
+# --------------------------------------------------------------------------- #
+# Helper functions
+# --------------------------------------------------------------------------- #
+
+def get_chroma_collection() -> Chroma:
+    """
+    Create or load a Chroma collection.
+    """
+    embeddings = OpenAIEmbeddings()
+    return Chroma(
+        collection_name=COLLECTION_NAME,
+        embedding_function=embeddings,
+        persist_directory=str(CHROMA_DATA_DIR),
+    )
+
+def ingest_text(text: str, collection: Chroma) -> None:
+    """
+    Add raw text to the Chroma collection.
+    """
+    collection.add_texts([text])
+
+def ingest_url(url: str, collection: Chroma, tavily_client: TavilyClient) -> None:
+    """
+    Fetch content from a URL using Tavily, embed it, and store it in Chroma.
+    """
+    # Tavily's search returns a list of results; we use the first result's content.
+    results = tavily_client.search(query=url, max_results=1)
+    if not results:
+        print(f"No results found for URL: {url}")
+        return
+    content = results[0].content
+    if not content:
+        print(f"No content extracted from URL: {url}")
+        return
+    collection.add_texts([content])
+    print(f"Ingested content from {url}")
+
+def query_agent(question: str, collection: Chroma) -> str:
+    """
+    Retrieve relevant documents from Chroma and ask OpenAI to answer.
+    """
+    llm = OpenAI(temperature=0)
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=collection.as_retriever(search_kwargs={"k": 4}),
+    )
+    return qa_chain.run(question)
+
+# --------------------------------------------------------------------------- #
+# Main entry point
+# --------------------------------------------------------------------------- #
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="RAG Agent with ChromaDB & Tavily")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    ingest_parser = subparsers.add_parser("ingest", help="Ingest a URL or raw text")
+    ingest_parser.add_argument("source", help="URL or raw text to ingest")
+
+    query_parser = subparsers.add_parser("query", help="Ask a question")
+    query_parser.add_argument("question", help="The question to ask the agent")
+
+    args = parser.parse_args()
+
+    # Ensure required environment variables are set
+    if "OPENAI_API_KEY" not in os.environ:
+        print("Error: OPENAI_API_KEY environment variable not set.")
+        sys.exit(1)
+    if "TAVILY_API_KEY" not in os.environ:
+        print("Error: TAVILY_API_KEY environment variable not set.")
+        sys.exit(1)
+
+    # Initialize Chroma collection
+    collection = get_chroma_collection()
+
+    # Initialize Tavily client
+    tavily_client = TavilyClient(api_key=os.environ["TAVILY_API_KEY"])
+
+    if args.command == "ingest":
+        source = args.source
+        if source.startswith(("http://", "https://")):
+            ingest_url(source, collection, tavily_client)
+        else:
+            ingest_text(source, collection)
+            print("Ingested raw text.")
+    elif args.command == "query":
+        answer = query_agent(args.question, collection)
+        print("\nAnswer:\n")
+        print(answer)
+
+if __name__ == "__main__":
+    main()