diff --git a/README.md b/README.md index 9339b0a..177bd8b 100644 --- a/README.md +++ b/README.md @@ -1,113 +1,64 @@ -# Agent with RAG Memory +# Agent with RAG Memory using Qdrant and Ollama -This repository contains a simple **Retrieval‑Augmented Generation (RAG)** agent -implemented with LangChain, FAISS for vector storage, and OpenAI embeddings -and LLM. It also provides an `auto_check_graph` function that verifies the -generated answer against a ground‑truth mapping and returns a `verdict_row`. +This project demonstrates a simple Retrieval-Augmented Generation (RAG) agent built with LangChain, Qdrant, and Ollama. The agent uses Ollama embeddings for vector representation and Qdrant as the vector store. -> **Important** -> The auto‑check graph must return a `verdict_row`. The implementation -> below guarantees that by always including the key in the returned -> dictionary. +## Prerequisites -## Features - -- **RAG Agent** – Load documents, embed them, store in FAISS, and answer queries. -- **Auto‑Check Graph** – Run a query, generate an answer, compare it to a - ground‑truth answer, and return a verdict (`PASS`, `FAIL`, or `UNKNOWN`). -- **Unit Tests** – Verify that the agent and auto‑check graph work as - expected. +- Python 3.10+ +- Qdrant server running locally or accessible remotely +- Ollama server running locally or accessible remotely ## Installation ```bash -# Create a virtual environment (recommended) -python -m venv .venv -source .venv/bin/activate # On Windows: .venv\Scripts\activate +# Clone the repository +git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git +cd agent-s-rag-pamyatyu + +# 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 ``` -`requirements.txt` contains: +## Configuration -``` -langchain -openai -faiss-cpu -pytest -``` - -> **OpenAI API Key** -> If you want to use real embeddings and LLM, set the environment variable -> `OPENAI_API_KEY`: - -```bash -export OPENAI_API_KEY="sk-..." -``` - -If the key is not set, the agent falls back to `FakeEmbeddings` and -`FakeLLM`, which are suitable for local testing and unit tests. - -## Usage +Edit `config.py` to match your environment: ```python -from src.index import RAGAgent, auto_check_graph +# Qdrant settings +QDRANT_HOST = "localhost" +QDRANT_PORT = 6333 +QDRANT_API_KEY = None +QDRANT_COLLECTION = "rag_collection" -# Create agent -agent = RAGAgent() - -# Add documents (e.g., from a directory) -agent.add_documents([ - "The capital of France is Paris.", - "William Shakespeare wrote Hamlet." -]) - -# Define ground truth mapping -ground_truth = { - "What is the capital of France?": "Paris", - "Who wrote Hamlet?": "William Shakespeare", -} - -# Run auto‑check graph -result = auto_check_graph( - "What is the capital of France?", - agent, - ground_truth -) - -print(result) -# Output: -# { -# "verdict_row": "PASS", -# "answer": "Paris", -# "expected": "Paris" -# } +# Ollama settings +OLLAMA_MODEL = "llama3" ``` -## Running Tests +## Running the Agent ```bash -pytest +python src/main.py ``` -The tests cover: +The script will: -- Adding documents and querying. -- Auto‑check graph returning `PASS`, `FAIL`, and `UNKNOWN` verdicts. -- Handling of empty queries and missing ground‑truth. +1. Connect to Qdrant. +2. Create an Ollama embeddings instance. +3. Add sample documents to the collection if it is empty. +4. Build a RetrievalQA chain using the Ollama LLM. +5. Execute a sample query and print the answer. -## Project Structure +## Extending -``` -src/ -├── index.py # Main implementation -tests/ -├── test_agent.py # Unit tests -README.md -requirements.txt -``` +- Replace the sample documents with your own corpus. +- Adjust the `chain_type` in `src/agent.py` if you need a different retrieval strategy. +- Use environment variables or a `.env` file to store sensitive information like `QDRANT_API_KEY`. ## License -MIT License \ No newline at end of file +MIT License +--- \ No newline at end of file diff --git a/SOLUTION.md b/SOLUTION.md index 9dfedd2..7c2c489 100644 --- a/SOLUTION.md +++ b/SOLUTION.md @@ -1,70 +1,63 @@ -**Что реализовано** +**What was implemented** +- Switched the vector store from FAISS to Qdrant using the `langchain_qdrant` wrapper. +- Replaced `OpenAIEmbeddings` with `OllamaEmbeddings` from `langchain_ollama`. +- Updated the agent to use Ollama for both embeddings and the LLM. +- Added `langchain-qdrant` and `langchain-ollama` to `requirements.txt`. +- Adjusted configuration to point to a local Qdrant instance and an Ollama model. -- Добавлен класс `RAGAgent`, который умеет индексировать документы в FAISS, выполнять поиск по запросу и генерировать ответ при помощи LLM (OpenAI или `FakeLLM`). -- Реализована функция `auto_check_graph`, которая запускает агента, сравнивает полученный ответ с ожидаемым и формирует словарь‑результат с ключом `verdict_row` (`PASS`, `FAIL` или `UNKNOWN`). +**Why the main parts satisfy the requirements** +- `src/vector_store.py` now imports `langchain_qdrant.Qdrant` and passes the Ollama embeddings, fulfilling the “use langchain‑qdrant” constraint. +- `src/agent.py` constructs the RetrievalQA chain with an Ollama LLM and the Qdrant retriever, meeting the “use Ollama embeddings” and “Qdrant as RAG memory” constraints. +- `config.py` centralises Qdrant and Ollama settings, so the rest of the code stays clean and configurable. +- `requirements.txt` lists both `langchain-qdrant` and `langchain-ollama`, removing any OpenAI/FAISS dependencies. -**Почему решения удовлетворяют требованиям** - -| Требование | Как реализовано | -|------------|----------------| -| **Агент с RAG‑памятью** | `RAGAgent.add_documents` добавляет документы в FAISS, `RAGAgent.query` извлекает ближайшие документы и формирует запрос к LLM. | -| **Граф автопроверки возвращает verdict_row** | `auto_check_graph` возвращает словарь, в котором обязательно присутствует ключ `"verdict_row"`. | -| **Проверка ответа** | Сравнение выполняется сначала точным совпадением, затем (если нужно) по косинусному сходству, что покрывает как точные, так и схожие ответы. | - -**Ключевые фрагменты кода** - -*`src/index.py` – добавление документов* +**Key code excerpts** +`config.py` – Qdrant & Ollama settings ```python -def add_documents(self, documents: Iterable[str], *, ids: Optional[List[str]] = None) -> None: - docs = [ - Document(page_content=doc, metadata={"id": doc_id}) - for doc, doc_id in zip(documents, ids or [None] * len(documents)) - ] - self.vector_store.add_documents(docs) - self.vector_store.save_local(self.vector_store_path) +# Qdrant settings +QDRANT_HOST = "localhost" +QDRANT_PORT = 6333 +QDRANT_API_KEY = None +QDRANT_COLLECTION = "rag_collection" + +# Ollama settings +OLLAMA_MODEL = "llama3" ``` -*`src/index.py` – запрос и генерация ответа* - +`src/vector_store.py` – Qdrant wrapper ```python -def query(self, query: str, k: int = 4) -> str: - docs_and_scores = self.vector_store.similarity_search_with_score(query, k=k) - context = "\n\n".join( - f"Document {i+1} (score={score:.3f}):\n{doc.page_content}" - for i, (doc, score) in enumerate(docs_and_scores) +class QdrantVectorStore: + def __init__(self, embeddings, collection_name: str = None): + self.qdrant = Qdrant( + url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}", + api_key=config.QDRANT_API_KEY, + collection_name=self.collection_name, + embeddings=embeddings, + ) +``` + +`src/agent.py` – RetrievalQA with Ollama +```python +def create_agent(vector_store: QdrantVectorStore): + embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL) + llm = Ollama(model=config.OLLAMA_MODEL) + qa_chain = RetrievalQA.from_chain_type( + llm=llm, + chain_type="stuff", + retriever=vector_store.get_retriever(), ) - prompt = f"You are an assistant. Use the following documents to answer the question. ..." - answer = self.llm.invoke(prompt).content.strip() - return answer + return qa_chain ``` -*`src/index.py` – автопроверка* - +`src/main.py` – initialization and sample run ```python -def auto_check_graph(user_query: str, rag_agent: RAGAgent, ground_truth: Dict[str, str]) -> Dict[str, str]: - answer = rag_agent.query(user_query) - expected = ground_truth.get(user_query) - if expected is None: - verdict = "UNKNOWN" - else: - if answer.strip().lower() == expected.strip().lower(): - verdict = "PASS" - else: - try: - query_vec = rag_agent.embeddings.embed_query(user_query) - answer_vec = rag_agent.embeddings.embed_query(answer) - similarity = rag_agent.embeddings.cosine_similarity(query_vec, answer_vec) - verdict = "PASS" if similarity >= SIMILARITY_THRESHOLD else "FAIL" - except Exception as exc: - logger.warning(f"Similarity check failed: {exc}") - verdict = "FAIL" - return {"verdict_row": verdict, "answer": answer, "expected": expected} +vector_store = QdrantVectorStore(embeddings) +agent = create_agent(vector_store) +result = agent.run("What is LangChain?") ``` -**Ограничения** - -- При отсутствии `OPENAI_API_KEY` используется `FakeEmbeddings`, у которых нет метода `cosine_similarity`. В этом случае сравнение по сходству всегда падает в `except`, и ответ считается `FAIL`. Для корректной работы в реальном окружении нужен настоящий OpenAI‑embedding‑модель. -- Точность проверки ограничена простым сравнением строк и косинусным сходством; более сложные случаи (например, синонимы) могут не распознаваться как `PASS`. - -Таким образом, реализованный код полностью покрывает требования задания: агент с RAG‑памятью, автопроверка, и гарантированное возвращение `verdict_row`. \ No newline at end of file +**Honest limitations** +- The solution assumes a running Qdrant instance on `localhost:6333` and an Ollama model named `llama3` available locally. +- No error handling for connection failures is added; in production you’d want to wrap Qdrant/ollama calls in try/except blocks. +- The sample documents are added only if the collection is empty; this logic is simplistic but sufficient for demonstration. \ No newline at end of file diff --git a/config.py b/config.py new file mode 100644 index 0000000..3e8fbb0 --- /dev/null +++ b/config.py @@ -0,0 +1,11 @@ +# Configuration for Qdrant and Ollama +# Adjust these values according to your environment + +# Qdrant settings +QDRANT_HOST = "localhost" # Qdrant server host +QDRANT_PORT = 6333 # Qdrant server port +QDRANT_API_KEY = None # Qdrant API key if required +QDRANT_COLLECTION = "rag_collection" # Collection name for embeddings + +# Ollama settings +OLLAMA_MODEL = "llama3" # Ollama model name for embeddings and LLM \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f386e17..3459854 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -langchain==0.2.0 -openai==1.3.0 -faiss-cpu==1.7.4 -tiktoken==0.5.1 -python-dotenv==1.0.0 \ No newline at end of file +langchain>=0.1.0 +langchain-qdrant +langchain-ollama +qdrant-client +python-dotenv \ No newline at end of file diff --git a/src/agent.py b/src/agent.py index 05ea78a..34a9953 100644 --- a/src/agent.py +++ b/src/agent.py @@ -1,62 +1,22 @@ -import logging -from typing import List +from langchain_ollama import Ollama, OllamaEmbeddings +from langchain.chains import RetrievalQA +import config +from src.vector_store import QdrantVectorStore -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer - -from .knowledge_base import KnowledgeBase -from .config import load_config - -logger = logging.getLogger(__name__) - -class RAGAgent: +def create_agent(vector_store: QdrantVectorStore): """ - Retrieval-Augmented Generation agent. + Create a RetrievalQA agent that uses Ollama for both embeddings and LLM. """ + # Embeddings for the vector store + embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL) - def __init__(self, config_path: str = "src/config.yaml"): - self.config = load_config(config_path) - logging.basicConfig(level=self.config["logging"]["level"]) - logger.info("Initializing RAGAgent.") - self.kb = KnowledgeBase( - data_dir=self.config["knowledge_base"]["data_dir"], - embedding_model=self.config["knowledge_base"]["embedding_model"], - vector_store=self.config["knowledge_base"]["vector_store"], - ) - self.model_name = self.config["language_model"]["model_name"] - self.max_length = self.config["language_model"]["max_length"] - self.top_k = self.config["retrieval"]["top_k"] - self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) - self.model = AutoModelForCausalLM.from_pretrained(self.model_name) - self.model.eval() - if torch.cuda.is_available(): - self.model.to("cuda") - logger.info(f"Loaded language model {self.model_name}") + # LLM for generating answers + llm = Ollama(model=config.OLLAMA_MODEL) - def generate_response(self, query: str) -> str: - """ - Generate a response to the query using retrieved context. - """ - logger.info(f"Generating response for query: {query}") - passages = self.kb.retrieve(query, top_k=self.top_k) - context = "\n\n".join([p[0] for p in passages]) if passages else "No relevant information found." - prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:" - logger.debug(f"Prompt:\n{prompt}") - - inputs = self.tokenizer(prompt, return_tensors="pt") - if torch.cuda.is_available(): - inputs = {k: v.to("cuda") for k, v in inputs.items()} - with torch.no_grad(): - output_ids = self.model.generate( - **inputs, - max_new_tokens=self.max_length, - do_sample=True, - top_p=0.95, - temperature=0.7, - ) - answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True) - # Extract the part after "Answer:" if present - if "Answer:" in answer: - answer = answer.split("Answer:")[1].strip() - logger.info(f"Generated answer: {answer}") - return answer \ No newline at end of file + # Build the RetrievalQA chain + qa_chain = RetrievalQA.from_chain_type( + llm=llm, + chain_type="stuff", + retriever=vector_store.get_retriever(), + ) + return qa_chain \ No newline at end of file diff --git a/src/main.py b/src/main.py index b2804ac..c500f6f 100644 --- a/src/main.py +++ b/src/main.py @@ -1,31 +1,42 @@ -import argparse -import logging - -from .agent import RAGAgent +import os +from langchain.schema import Document +from src.vector_store import QdrantVectorStore +from src.agent import create_agent +import config def main(): - parser = argparse.ArgumentParser(description="Educational RAG Agent CLI") - parser.add_argument("--config", type=str, default="src/config.yaml", help="Path to config file") - args = parser.parse_args() + # Ensure Qdrant is reachable + os.environ["QDRANT_URL"] = f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}" + if config.QDRANT_API_KEY: + os.environ["QDRANT_API_KEY"] = config.QDRANT_API_KEY - logging.basicConfig(level=logging.INFO) - agent = RAGAgent(config_path=args.config) + # Initialize embeddings and vector store + from langchain_ollama import OllamaEmbeddings + embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL) + vector_store = QdrantVectorStore(embeddings) - print("Welcome to the Educational RAG Agent. Type 'exit' to quit.") - while True: - try: - query = input("\nYour question: ").strip() - if query.lower() in ("exit", "quit"): - print("Goodbye!") - break - if not query: - print("Please enter a non-empty question.") - continue - answer = agent.generate_response(query) - print(f"\nAnswer:\n{answer}") - except KeyboardInterrupt: - print("\nInterrupted. Exiting.") - break + # Add sample documents (only if collection is empty) + # In a real scenario, you would load your corpus here + sample_docs = [ + Document(page_content="Hello world! This is a test document.", metadata={"source": "test"}), + Document(page_content="LangChain is a powerful framework for building LLM applications.", metadata={"source": "test"}), + ] + # Check if collection already has documents + try: + # Attempt to retrieve a document to see if collection is populated + vector_store.get_retriever().get_relevant_documents("test") + except Exception: + # If retrieval fails, add documents + vector_store.add_documents(sample_docs) + + # Create the agent + agent = create_agent(vector_store) + + # Run a sample query + query = "What is LangChain?" + print(f"Query: {query}") + result = agent.run(query) + print(f"Answer: {result}") if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/vector_store.py b/src/vector_store.py index 6e35a38..5cd83ef 100644 --- a/src/vector_store.py +++ b/src/vector_store.py @@ -1,97 +1,29 @@ -""" -Vector store implementation using Qdrant via langchain-qdrant. -Provides a simple interface for adding documents and performing -similarity search. Embeddings are generated using OpenAIEmbeddings -by default, but can be overridden by passing a custom embedding -function. -""" - -from __future__ import annotations - -from typing import Iterable, List, Optional - -from langchain.embeddings import OpenAIEmbeddings from langchain_qdrant import Qdrant -from langchain.vectorstores import VectorStore -from langchain_core.documents import Document +from langchain.schema import Document +import config -from .config import ( - QDRANT_HOST, - QDRANT_PORT, - QDRANT_API_KEY, - QDRANT_COLLECTION, -) - - -class QdrantVectorStore(VectorStore): +class QdrantVectorStore: """ - A wrapper around langchain_qdrant.Qdrant that implements the - VectorStore interface expected by LangChain chains. + Wrapper around langchain_qdrant.Qdrant to provide a simple interface + for adding documents and retrieving a retriever. """ - - def __init__( - self, - embeddings: Optional[OpenAIEmbeddings] = None, - collection_name: str = QDRANT_COLLECTION, - ): - self.embeddings = embeddings or OpenAIEmbeddings() - self.collection_name = collection_name - - # Initialize Qdrant client - self.client = Qdrant( - host=QDRANT_HOST, - port=QDRANT_PORT, - api_key=QDRANT_API_KEY, + def __init__(self, embeddings, collection_name: str = None): + self.collection_name = collection_name or config.QDRANT_COLLECTION + self.qdrant = Qdrant( + url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}", + api_key=config.QDRANT_API_KEY, collection_name=self.collection_name, - ) - - def add_documents(self, documents: Iterable[Document]) -> None: - """ - Add a collection of documents to the Qdrant store. - """ - texts = [doc.page_content for doc in documents] - metadatas = [doc.metadata for doc in documents] - ids = [doc.id for doc in documents if doc.id is not None] - - # Embed the documents - embeddings = self.embeddings.embed_documents(texts) - - # Upsert into Qdrant - self.client.upsert( embeddings=embeddings, - documents=texts, - metadatas=metadatas, - ids=ids, ) - def similarity_search( - self, - query: str, - k: int = 5, - filter: Optional[dict] = None, - ) -> List[Document]: + def add_documents(self, documents: list[Document]): """ - Perform a similarity search against the Qdrant store. + Add a list of langchain Document objects to the Qdrant collection. """ - query_embedding = self.embeddings.embed_query(query) - results = self.client.search( - query_embedding=query_embedding, - limit=k, - filter=filter, - ) - # Convert results to Document objects - return [ - Document( - page_content=result["payload"]["text"], - metadata=result["payload"], - id=result["id"], - ) - for result in results - ] + self.qdrant.add_documents(documents) - # The following methods are required by the VectorStore interface - def embed_query(self, query: str) -> List[float]: - return self.embeddings.embed_query(query) - - def embed_documents(self, documents: List[str]) -> List[List[float]]: - return self.embeddings.embed_documents(documents) \ No newline at end of file + def get_retriever(self): + """ + Return a retriever that can be used with LangChain chains. + """ + return self.qdrant.as_retriever() \ No newline at end of file