diff --git a/README.md b/README.md index 94936b1..b6ffe7b 100644 --- a/README.md +++ b/README.md @@ -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 \ No newline at end of file +MIT License +--- +Feel free to extend the agent with additional features such as custom embeddings, different LLMs, or alternative search APIs. \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 722e91d..0a3f5dd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,5 @@ -langchain -langchain-ollama -langchain-qdrant -langchain-tavily -tavily-python -chromadb -python-dotenv -qdrant-client \ No newline at end of file +chromadb==0.4.22 +openai==1.12.0 +requests==2.31.0 +beautifulsoup4==4.12.3 +python-dotenv==1.0.1 \ No newline at end of file diff --git a/src/main.py b/src/main.py index 8a2dc5f..50d5943 100644 --- a/src/main.py +++ b/src/main.py @@ -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() \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py new file mode 100644 index 0000000..42ace28 --- /dev/null +++ b/src/vector_store.py @@ -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)) \ No newline at end of file diff --git a/src/web_search.py b/src/web_search.py new file mode 100644 index 0000000..098cd83 --- /dev/null +++ b/src/web_search.py @@ -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 \ No newline at end of file