diff --git a/README.md b/README.md index 50ed043..5d5b2f9 100644 --- a/README.md +++ b/README.md @@ -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 ``` -## Configuration +- If `` 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 "" +``` + +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 \ No newline at end of file +This project is released under the MIT License. \ No newline at end of file diff --git a/src/index.py b/src/index.py new file mode 100644 index 0000000..8f519ff --- /dev/null +++ b/src/index.py @@ -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 # Ingest a URL or raw text + python src/index.py query # 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() \ No newline at end of file