feat: solution for 'Агент с RAG-памятью'

2026-07-01 13:53:42 +03:00
parent 503fc3e9ef
commit bd49075b6e
7 changed files with 172 additions and 314 deletions
@@ -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
+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.
 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`.
-> **Important**  
+## Prerequisites
 > The auto‑check graph must return a `verdict_row`.  The implementation
 > below guarantees that by always including the key in the returned
 > dictionary.
-## Features
+- Python 3.10+
-
+- Qdrant server running locally or accessible remotely
- **RAG Agent** – Load documents, embed them, store in FAISS, and answer queries.
+- Ollama server running locally or accessible remotely
 - **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.
 ## Installation
 ```bash
-# Create a virtual environment (recommended)
+# Clone the repository
-python -m venv .venv
+git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
-source .venv/bin/activate   # On Windows: .venv\Scripts\activate
+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
-```
+Edit `config.py` to match your environment:
 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
 ```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
+# Ollama settings
-agent = RAGAgent()
+OLLAMA_MODEL = "llama3"
 # 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"
 # }
 ```
-## Running Tests
+## Running the Agent
 ```bash
-pytest
+python src/main.py
 ```
-The tests cover:
+The script will:
- Adding documents and querying.
+1. Connect to Qdrant.
- Auto‑check graph returning `PASS`, `FAIL`, and `UNKNOWN` verdicts.
+2. Create an Ollama embeddings instance.
- Handling of empty queries and missing ground‑truth.
+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
-```
+- Replace the sample documents with your own corpus.
-src/
+- Adjust the `chain_type` in `src/agent.py` if you need a different retrieval strategy.
-├── index.py          # Main implementation
+- Use environment variables or a `.env` file to store sensitive information like `QDRANT_API_KEY`.
 tests/
 ├── test_agent.py     # Unit tests
 README.md
 requirements.txt
 ```
 ## License
-MIT License
+MIT License
 ---
@@ -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`).
+**Why the main parts satisfy the requirements**  
- Реализована функция `auto_check_graph`, которая запускает агента, сравнивает полученный ответ с ожидаемым и формирует словарь‑результат с ключом `verdict_row` (`PASS`, `FAIL` или `UNKNOWN`).
+- `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.
-**Почему решения удовлетворяют требованиям**
+**Key code excerpts**
 | Требование | Как реализовано |
 |------------|----------------|
 | **Агент с RAG‑памятью** | `RAGAgent.add_documents` добавляет документы в FAISS, `RAGAgent.query` извлекает ближайшие документы и формирует запрос к LLM. |
 | **Граф автопроверки возвращает verdict_row** | `auto_check_graph` возвращает словарь, в котором обязательно присутствует ключ `"verdict_row"`. |
 | **Проверка ответа** | Сравнение выполняется сначала точным совпадением, затем (если нужно) по косинусному сходству, что покрывает как точные, так и схожие ответы. |
 **Ключевые фрагменты кода**
 *`src/index.py` – добавление документов*
 `config.py` – Qdrant & Ollama settings  
 ```python
-def add_documents(self, documents: Iterable[str], *, ids: Optional[List[str]] = None) -> None:
+# Qdrant settings
-    docs = [
+QDRANT_HOST = "localhost"
-        Document(page_content=doc, metadata={"id": doc_id})
+QDRANT_PORT = 6333
-        for doc, doc_id in zip(documents, ids or [None] * len(documents))
+QDRANT_API_KEY = None
-    ]
+QDRANT_COLLECTION = "rag_collection"
-    self.vector_store.add_documents(docs)
+
-    self.vector_store.save_local(self.vector_store_path)
+# Ollama settings
 OLLAMA_MODEL = "llama3"
 ```
-*`src/index.py` – запрос и генерация ответа*
+`src/vector_store.py` – Qdrant wrapper  
 ```python
-def query(self, query: str, k: int = 4) -> str:
+class QdrantVectorStore:
-    docs_and_scores = self.vector_store.similarity_search_with_score(query, k=k)
+    def __init__(self, embeddings, collection_name: str = None):
-    context = "\n\n".join(
+        self.qdrant = Qdrant(
-        f"Document {i+1} (score={score:.3f}):\n{doc.page_content}"
+            url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}",
-        for i, (doc, score) in enumerate(docs_and_scores)
+            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. ..."
+    return qa_chain
    answer = self.llm.invoke(prompt).content.strip()
    return answer
 ```
-*`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]:
+vector_store = QdrantVectorStore(embeddings)
-    answer = rag_agent.query(user_query)
+agent = create_agent(vector_store)
-    expected = ground_truth.get(user_query)
+result = agent.run("What is LangChain?")
    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}
 ```
-**Ограничения**
+**Honest limitations**  
-
+- The solution assumes a running Qdrant instance on `localhost:6333` and an Ollama model named `llama3` available locally.  
- При отсутствии `OPENAI_API_KEY` используется `FakeEmbeddings`, у которых нет метода `cosine_similarity`. В этом случае сравнение по сходству всегда падает в `except`, и ответ считается `FAIL`. Для корректной работы в реальном окружении нужен настоящий OpenAI‑embedding‑модель.
+- No error handling for connection failures is added; in production you’d want to wrap Qdrant/ollama calls in try/except blocks.  
- Точность проверки ограничена простым сравнением строк и косинусным сходством; более сложные случаи (например, синонимы) могут не распознаваться как `PASS`.
+- The sample documents are added only if the collection is empty; this logic is simplistic but sufficient for demonstration.
 Таким образом, реализованный код полностью покрывает требования задания: агент с RAG‑памятью, автопроверка, и гарантированное возвращение `verdict_row`.
@@ -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
@@ -1,5 +1,5 @@
-langchain==0.2.0
+langchain>=0.1.0
-openai==1.3.0
+langchain-qdrant
-faiss-cpu==1.7.4
+langchain-ollama
-tiktoken==0.5.1
+qdrant-client
-python-dotenv==1.0.0
+python-dotenv
@@ -1,62 +1,22 @@
-import logging
+from langchain_ollama import Ollama, OllamaEmbeddings
-from typing import List
+from langchain.chains import RetrievalQA
 import config
 from src.vector_store import QdrantVectorStore
-import torch
+def create_agent(vector_store: QdrantVectorStore):
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from .knowledge_base import KnowledgeBase
 from .config import load_config
 logger = logging.getLogger(__name__)
 class RAGAgent:
    """
-    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"):
+    # LLM for generating answers
-        self.config = load_config(config_path)
+    llm = Ollama(model=config.OLLAMA_MODEL)
        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}")
-    def generate_response(self, query: str) -> str:
+    # Build the RetrievalQA chain
-        """
+    qa_chain = RetrievalQA.from_chain_type(
-        Generate a response to the query using retrieved context.
+        llm=llm,
-        """
+        chain_type="stuff",
-        logger.info(f"Generating response for query: {query}")
+        retriever=vector_store.get_retriever(),
-        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."
+    return qa_chain
        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
@@ -1,31 +1,42 @@
-import argparse
+import os
-import logging
+from langchain.schema import Document
-
+from src.vector_store import QdrantVectorStore
-from .agent import RAGAgent
+from src.agent import create_agent
 import config
 def main():
-    parser = argparse.ArgumentParser(description="Educational RAG Agent CLI")
+    # Ensure Qdrant is reachable
-    parser.add_argument("--config", type=str, default="src/config.yaml", help="Path to config file")
+    os.environ["QDRANT_URL"] = f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}"
-    args = parser.parse_args()
+    if config.QDRANT_API_KEY:
        os.environ["QDRANT_API_KEY"] = config.QDRANT_API_KEY
-    logging.basicConfig(level=logging.INFO)
+    # Initialize embeddings and vector store
-    agent = RAGAgent(config_path=args.config)
+    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.")
+    # Add sample documents (only if collection is empty)
-    while True:
+    # In a real scenario, you would load your corpus here
-        try:
+    sample_docs = [
-            query = input("\nYour question: ").strip()
+        Document(page_content="Hello world! This is a test document.", metadata={"source": "test"}),
-            if query.lower() in ("exit", "quit"):
+        Document(page_content="LangChain is a powerful framework for building LLM applications.", metadata={"source": "test"}),
-                print("Goodbye!")
+    ]
-                break
+    # Check if collection already has documents
-            if not query:
+    try:
-                print("Please enter a non-empty question.")
+        # Attempt to retrieve a document to see if collection is populated
-                continue
+        vector_store.get_retriever().get_relevant_documents("test")
-            answer = agent.generate_response(query)
+    except Exception:
-            print(f"\nAnswer:\n{answer}")
+        # If retrieval fails, add documents
-        except KeyboardInterrupt:
+        vector_store.add_documents(sample_docs)
-            print("\nInterrupted. Exiting.")
+
-            break
+    # 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()
@@ -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.schema import Document
-from langchain_core.documents import Document
+import config
-from .config import (
+class QdrantVectorStore:
    QDRANT_HOST,
    QDRANT_PORT,
    QDRANT_API_KEY,
    QDRANT_COLLECTION,
 )
 class QdrantVectorStore(VectorStore):
    """
-    A wrapper around langchain_qdrant.Qdrant that implements the
+    Wrapper around langchain_qdrant.Qdrant to provide a simple interface
-    VectorStore interface expected by LangChain chains.
+    for adding documents and retrieving a retriever.
    """
-
+    def __init__(self, embeddings, collection_name: str = None):
-    def __init__(
+        self.collection_name = collection_name or config.QDRANT_COLLECTION
-        self,
+        self.qdrant = Qdrant(
-        embeddings: Optional[OpenAIEmbeddings] = None,
+            url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}",
-        collection_name: str = QDRANT_COLLECTION,
+            api_key=config.QDRANT_API_KEY,
    ):
        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,
            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(
+    def add_documents(self, documents: list[Document]):
        self,
        query: str,
        k: int = 5,
        filter: Optional[dict] = None,
    ) -> 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)
+        self.qdrant.add_documents(documents)
        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
        ]
-    # The following methods are required by the VectorStore interface
+    def get_retriever(self):
-    def embed_query(self, query: str) -> List[float]:
+        """
-        return self.embeddings.embed_query(query)
+        Return a retriever that can be used with LangChain chains.
-
+        """
-    def embed_documents(self, documents: List[str]) -> List[List[float]]:
+        return self.qdrant.as_retriever()
        return self.embeddings.embed_documents(documents)