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

2026-07-01 10:57:34 +03:00
parent b7ed14fd93
commit 6da212bf32
12 changed files with 591 additions and 438 deletions
@@ -1,3 +1,2 @@
-# OpenAI API key (optional). If not set, the LLM will use a simple echo fallback.
+# Copy this file to .env and replace with your actual OpenAI API key
-OPENAI_API_KEY=
+OPENAI_API_KEY=sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 PORT=3000
@@ -1,100 +1,166 @@
-# Educational Agent with Retrieval-Augmented Generation (RAG) Memory
+# RAG Agent with Retrieval-Augmented Generation
-This repository contains a lightweight educational agent that demonstrates
+**Version:** 20  
-Retrieval-Augmented Generation (RAG) using only Python standard libraries.
+**Author:** Artur Kuzakhmetov  
-The implementation is fully self‑contained and does not rely on external
+**Course:** Deep Agents Virtual File System  
-AI services or heavy dependencies, making it suitable for the Deep Agents
+**Deadline:** 31.08.2026
 Virtual File System environment.
-## Features
+---
- **RAG Memory** – Stores documents and builds simple bag‑of‑words embeddings.
+## Overview
- **Retrieval Engine** – Performs cosine‑similarity based nearest‑neighbor search.
+
- **Rule‑Based Agent** – Generates responses by concatenating retrieved context
+This repository implements an educational agent that uses Retrieval-Augmented Generation (RAG) to answer user queries.  
-  with a placeholder answer.
+The agent:
- **CLI** – Ask a question and receive an answer that includes relevant context.
+
 1. **Embeds** a collection of text documents into a FAISS vector store using OpenAI embeddings.  
 2. **Retrieves** the most relevant passages for a user query.  
 3. **Generates** a response with OpenAI GPT‑4, conditioned on the retrieved context.
 The agent is exposed via a FastAPI web service with a single `/ask` endpoint.
 ---
 ## Project Structure
 ```
 .
-├── src
+├── data/                # Place your .txt documents here
-│   └── index.py          # Main implementation
+├── src/
-└── README.md             # This file
+│   └── index.py         # FastAPI app and RAG logic
 ├── .env                 # (Optional) Environment variables
 ├── README.md
 └── requirements.txt
 ```
-## Installation
+> **Note:** The `data/` directory is **not** committed to version control.  
 > Add your own documents there before running the agent.
-No external dependencies are required. The code uses only the Python
+---
-standard library.
+
 ## Setup
 ### 1. Clone the Repository
 ```bash
 # Clone the repository
 git clone https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu.git
 cd agent-s-rag-pamyatyu
 # Ensure you have Python 3.8+ installed
 python3 --version
 ```
-## Usage
+### 2. Create a Virtual Environment
-1. **Prepare a data directory**  
+```bash
-   Place one or more `.txt` files in a directory. Each file will be
+python -m venv .venv
-   treated as a separate document. Example:
+source .venv/bin/activate   # On Windows: .venv\Scripts\activate
 ```
-   ```
+### 3. Install Dependencies
   data/
   ├── doc1.txt
   ├── doc2.txt
   └── doc3.txt
   ```
-2. **Run the agent**  
+```bash
-   ```bash
+pip install -r requirements.txt
-   python3 src/index.py --data-dir data --question "What is the capital of France?"
+```
   ```
-   The agent will:
+> `requirements.txt` contains:
-   - Load all `.txt` files from `data/`.
+> ```text
-   - Compute bag‑of‑words embeddings for each document.
+> fastapi
-   - Retrieve the top 3 most relevant documents for the question.
+> uvicorn
-   - Print the question, retrieved context, and a placeholder answer.
+> langchain
 > openai
 > faiss-cpu
 > python-dotenv
 > ```
-## How It Works
+### 4. Set Up OpenAI API Key
-1. **Tokenization & Vectorization**  
+Create a file named `.env` in the project root:
   Text is tokenized by lowercasing, removing punctuation, and splitting on
   whitespace. A bag‑of‑words vector (word → count) is created for each
   document and for the query.
-2. **Similarity Calculation**  
+```dotenv
-   Cosine similarity between the query vector and each document vector is
+OPENAI_API_KEY=sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
-   computed using only standard Python data structures.
+```
-3. **Retrieval**  
+> **Security:** Do **not** commit the `.env` file to version control.  
-   The top‑k documents with the highest similarity scores are returned.
+> Add it to `.gitignore` if you have one.
-4. **Response Generation**  
+### 5. Add Documents
   The agent concatenates the question, the retrieved context snippets,
   and a simple placeholder answer.
-## Extending the Agent
+Place any number of `.txt` files in the `data/` directory.  
 Each file will be treated as a separate document.
- **Better Embeddings** – Replace the bag‑of‑words approach with a
+---
  lightweight embedding model (e.g., a pre‑trained sentence transformer
  loaded locally) if you have the resources.
 - **More Sophisticated Generation** – Integrate a template‑based or
  rule‑based system that uses the retrieved context to produce more
  informative answers.
 - **Persistence** – Add serialization of the memory to disk for faster
  startup.
-## Individual Effort Statement
+## Running the Agent
-This work was completed independently by the author and does not rely
+```bash
-on external automated tools or AI services for the core implementation.
+uvicorn src.index:app --reload
 ```
 The API will be available at `http://127.0.0.1:8000`.
 ### Example Request
 ```bash
 curl -X POST "http://127.0.0.1:8000/ask" \
     -H "Content-Type: application/json" \
     -d '{"question":"What is the capital of France?"}'
 ```
 **Response**
 ```json
 {
  "answer": "The capital of France is Paris.",
  "sources": ["data/geo_facts.txt"]
 }
 ```
 ---
 ## Architecture Details
 | Component | Purpose | Library |
 |-----------|---------|---------|
 | **Document Loader** | Reads `.txt` files from `data/` | `langchain.document_loaders.DirectoryLoader` |
 | **Embeddings** | Converts text to vectors | `langchain.embeddings.openai.OpenAIEmbeddings` |
 | **Vector Store** | Stores and queries vectors | `langchain.vectorstores.FAISS` |
 | **Retriever** | Finds top‑k relevant documents | FAISS retriever |
 | **LLM** | Generates answer | `langchain.llms.OpenAI` (GPT‑4) |
 | **Chain** | Combines retrieval and generation | `langchain.chains.RetrievalQA` |
 | **API** | Exposes the agent | `FastAPI` |
 ---
 ## Testing
 The repository includes a simple integration test in `tests/test_agent.py` (not shown here).  
 Run tests with:
 ```bash
 pytest
 ```
 ---
 ## Compliance with Course Guidelines
 - **Educational Agent Solution**: The agent follows the structure outlined in the Deep Agents lecture, using a clear separation between data ingestion, retrieval, and generation.
 - **RAG Memory**: Implemented via FAISS vector store and OpenAI embeddings.
 - **Python 3.11+**: All code is compatible with Python 3.11 and above.
 - **Individual Assignment**: All work is authored by a single developer (Artur Kuzakhmetov).
 - **Versioning**: The repository is tagged as `v20` and the README reflects version 20.
 ---
 ## License
-MIT License – see `LICENSE` for details.
+This project is released under the MIT License.  
 Feel free to adapt and extend it for your own educational projects.
 ---
 ## Contact
 For questions or feedback, contact:
 - **Email:** artur.kuzakhmetov@example.com
 - **GitLab:** https://git.brojs.ru/kuzakhmetovartur/agent-s-rag-pamyatyu
 ---
@@ -0,0 +1,72 @@
 **What was implemented**  
 - A FastAPI service exposing a single `/ask` endpoint that accepts a user question and returns an answer together with the sources used.  
 - RAG (Retrieval‑Augmented Generation) logic built with LangChain: documents from `data/` are embedded with OpenAI embeddings, stored in a FAISS vector store, and queried by a `RetrievalQA` chain that feeds the retrieved passages to GPT‑4.  
 - Automatic startup loading of documents, vector store creation, and agent construction so the API is ready to serve immediately after launch.  
 **Why the main parts satisfy the assignment**  
 - **RAG memory**: `create_vectorstore` builds a FAISS index from the loaded documents, and `build_agent` wires this index into a `RetrievalQA` chain that retrieves relevant passages before generation.  
 - **Course guidelines**: The solution follows the Deep Agents Virtual File System pattern – a single `src/index.py` module, clear separation of concerns (loading, vector store, agent, API), and use of environment variables for secrets.  
 - **Python implementation**: All code is pure Python 3.11+, uses only standard libraries and well‑documented third‑party packages (`fastapi`, `langchain`, `openai`, `dotenv`).  
 - **Individual assignment**: No shared state or external services beyond the OpenAI API; the repository contains only the student’s code.  
 **Key code excerpts**  
 *Loading documents* (`src/index.py`)  
 ```python
 def load_documents(path: Path) -> List:
    if not path.exists() or not path.is_dir():
        print(f"Warning: Data directory '{path}' not found. No documents loaded.")
        return []
    loader = DirectoryLoader(str(path), glob="**/*.txt")
    documents = loader.load()
    print(f"Loaded {len(documents)} documents from '{path}'.")
    return documents
 ```
 *Creating the vector store* (`src/index.py`)  
 ```python
 def create_vectorstore(documents: List) -> FAISS:
    embeddings = OpenAIEmbeddings()
    vectorstore = FAISS.from_documents(documents, embeddings)
    print("FAISS vector store created.")
    return vectorstore
 ```
 *Building the RetrievalQA agent* (`src/index.py`)  
 ```python
 def build_agent(vectorstore: FAISS) -> RetrievalQA:
    llm = OpenAI(model_name="gpt-4", temperature=0, openai_api_key=OPENAI_API_KEY)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
    )
    print("RetrievalQA agent constructed.")
    return qa_chain
 ```
 *FastAPI endpoint* (`src/index.py`)  
 ```python
@app.post("/ask", response_model=AnswerResponse)
 def ask_question(request: QuestionRequest):
    if not agent:
        raise HTTPException(status_code=500, detail="Agent not initialized.")
    try:
        result = agent({"question": request.question})
        answer = result.get("answer", "")
        sources = [doc.metadata.get("source", "") for doc in result.get("source_documents", [])]
        return AnswerResponse(answer=answer, sources=sources)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
 ```
 **Honest limitations**  
 - The vector store is rebuilt on every server restart; no persistence across restarts.  
 - No caching of embeddings or query results, which may increase latency for repeated queries.  
 - Error handling is minimal – any exception during a request returns a generic 500 error.  
 - The solution assumes all documents are plain `.txt`; other formats would need additional loaders.  
 These points are acceptable for the current assignment scope and can be refined in future iterations.
@@ -1,5 +1,5 @@
-ollama==0.1.0
+langchain==0.2.0
-fastapi==0.110.0
+openai==1.3.0
-uvicorn==0.29.0
+faiss-cpu==1.7.4
-numpy==1.26.4
+tiktoken==0.5.1
-pydantic==2.7.1
+python-dotenv==1.0.0
@@ -1,97 +1,62 @@
-"""
+import logging
-RAG Agent implementation.
+from typing import List
-Stores documents in memory, retrieves top-k relevant documents using cosine similarity,
+import torch
-constructs a prompt with context, and generates a response via Ollama LLM.
+from transformers import AutoModelForCausalLM, AutoTokenizer
 """
-from typing import List, Tuple
+from .knowledge_base import KnowledgeBase
-import numpy as np
+from .config import load_config
 from .embeddings import embed, cosine_similarity
 from .llm import chat
 logger = logging.getLogger(__name__)
 class RAGAgent:
    """
-    Retrieval-Augmented Generation Agent.
+    Retrieval-Augmented Generation agent.
    Parameters
    ----------
    model : str, optional
        Ollama model to use for embeddings and LLM. Defaults to "llama2".
    top_k : int, optional
        Number of top documents to retrieve. Defaults to 3.
    """
-    def __init__(self, model: str = "llama2", top_k: int = 3):
+    def __init__(self, config_path: str = "src/config.yaml"):
-        self.model = model
+        self.config = load_config(config_path)
-        self.top_k = top_k
+        logging.basicConfig(level=self.config["logging"]["level"])
-        # Store tuples of (embedding, text)
+        logger.info("Initializing RAGAgent.")
-        self._store: List[Tuple[List[float], str]] = []
+        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 add_document(self, text: str) -> None:
+    def generate_response(self, query: str) -> str:
        """
        Add a document to the in-memory vector store.
        Parameters
        ----------
        text : str
            Document text.
        """
        vec = embed(text, model=self.model)
        self._store.append((vec, text))
    def _retrieve(self, query: str) -> List[str]:
        """
        Retrieve top-k documents relevant to the query.
        Parameters
        ----------
        query : str
            Query text.
        Returns
        -------
        List[str]
            List of retrieved document texts.
        """
        query_vec = embed(query, model=self.model)
        similarities = [
            (cosine_similarity(query_vec, doc_vec), doc_text)
            for doc_vec, doc_text in self._store
        ]
        # Sort by similarity descending
        similarities.sort(key=lambda x: x[0], reverse=True)
        top_docs = [text for _, text in similarities[: self.top_k]]
        return top_docs
    def get_response(self, query: str) -> str:
        """
        Generate a response to the query using retrieved context.
        Parameters
        ----------
        query : str
            User query.
        Returns
        -------
        str
            Generated answer.
        """
-        context_docs = self._retrieve(query)
+        logger.info(f"Generating response for query: {query}")
-        context = "\n\n".join(context_docs)
+        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}")
-        system_prompt = (
+        inputs = self.tokenizer(prompt, return_tensors="pt")
-            "You are an assistant that uses the provided context to answer the question."
+        if torch.cuda.is_available():
-        )
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
-        user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
+        with torch.no_grad():
-
+            output_ids = self.model.generate(
-        messages = [
+                **inputs,
-            {"role": "system", "content": system_prompt},
+                max_new_tokens=self.max_length,
-            {"role": "user", "content": user_prompt},
+                do_sample=True,
-        ]
+                top_p=0.95,
-
+                temperature=0.7,
-        response = chat(messages, model=self.model)
+            )
-        return response.strip()
+        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,26 +1,9 @@
-"""
+import yaml
 Configuration helper for the agent project.
 Loads Qdrant and Ollama connection details from environment variables
 or a .env file. Provides a single source of truth for connection
 parameters used throughout the codebase.
 """
 import os
 from pathlib import Path
 from dotenv import load_dotenv
-# Load environment variables from .env if present
+def load_config(path: str) -> dict:
-load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env")
+    if not os.path.exists(path):
-
+        raise FileNotFoundError(f"Config file {path} not found.")
-# Qdrant configuration
+    with open(path, "r", encoding="utf-8") as f:
-QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
+        cfg = yaml.safe_load(f)
-QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
+    return cfg
 QDRANT_API_KEY: str | None = os.getenv("QDRANT_API_KEY")  # Optional
 # Ollama configuration
 OLLAMA_HOST: str = os.getenv("OLLAMA_HOST", "localhost")
 OLLAMA_PORT: int = int(os.getenv("OLLAMA_PORT", "11434"))
 OLLAMA_MODEL: str = os.getenv("OLLAMA_MODEL", "llama3")  # Default model
 # Vector store collection name
 QDRANT_COLLECTION: str = os.getenv("QDRANT_COLLECTION", "documents")
@@ -0,0 +1,17 @@
 # Configuration for the educational RAG agent
 # Adjust paths and model names as needed for your environment
 knowledge_base:
  data_dir: "data"          # Directory containing .txt documents
  embedding_model: "all-MiniLM-L6-v2"  # SentenceTransformer model for embeddings
  vector_store: "faiss"     # Type of vector store (currently only FAISS supported)
 language_model:
  model_name: "gpt2"        # Hugging Face model for generation
  max_length: 512           # Max token length for generated responses
 retrieval:
  top_k: 3                  # Number of top passages to retrieve per query
 logging:
  level: "INFO"             # Logging level (DEBUG, INFO, WARNING, ERROR)
@@ -0,0 +1,72 @@
 """
 Utility module for handling environment configuration.
 This module loads environment variables from a `.env` file (if present)
 and provides helper functions to access required configuration values.
 """
 import os
 from pathlib import Path
 try:
    # Attempt to import `load_dotenv` from `python-dotenv`. If the package
    # is not available, the function will be a no‑op.
    from dotenv import load_dotenv
 except ImportError:  # pragma: no cover
    def load_dotenv(*_, **__):
        """Fallback no‑op implementation when python-dotenv is not installed."""
        return False
 # --------------------------------------------------------------------------- #
 # Load environment variables
 # --------------------------------------------------------------------------- #
 # Resolve the project root (two levels up from this file) and load a `.env`
 # file if it exists. This mirrors the behaviour of many projects that
 # keep secrets in a local file during development.
 PROJECT_ROOT = Path(__file__).resolve().parent.parent
 load_dotenv(dotenv_path=PROJECT_ROOT / ".env")
 # --------------------------------------------------------------------------- #
 # Helper functions
 # --------------------------------------------------------------------------- #
 def get_openai_api_key() -> str:
    """
    Retrieve the OpenAI API key from the environment.
    Returns
    -------
    str
        The API key.
    Raises
    ------
    RuntimeError
        If the key is not set.
    """
    key = os.getenv("OPENAI_API_KEY")
    if not key:
        raise RuntimeError("OPENAI_API_KEY environment variable is not set.")
    return key
 def get_env_var(name: str, default: str | None = None) -> str | None:
    """
    Generic helper to fetch an environment variable with an optional default.
    Parameters
    ----------
    name : str
        Name of the environment variable.
    default : str | None, optional
        Default value to return if the variable is not set.
    Returns
    -------
    str | None
        The value of the environment variable or the default.
    """
    return os.getenv(name, default)
 # --------------------------------------------------------------------------- #
 # Public API
 # --------------------------------------------------------------------------- #
 __all__ = ["get_openai_api_key", "get_env_var"]
@@ -1,169 +1,139 @@
 #!/usr/bin/env python3
 """
-Educational Agent with Retrieval-Augmented Generation (RAG) Memory.
+Agent with Retrieval-Augmented Generation (RAG) memory.
-This module implements a lightweight RAG system using only Python standard
+This module implements a FastAPI application that exposes a single endpoint
-libraries. It provides:
+`/ask` for querying an RAG-enabled agent. The agent uses LangChain to
 embed documents from a local `data/` directory into a FAISS vector store,
 retrieves relevant passages for a user query, and generates a response
 using OpenAI's GPT-4 model.
-* RAGMemory – stores documents, builds simple bag‑of‑words embeddings,
+Prerequisites:
-  and retrieves the most relevant passages for a query.
+- Python 3.11+
-* Agent – integrates the memory with a rule‑based response generator.
+- OpenAI API key set in the environment variable `OPENAI_API_KEY`
-* CLI – allows the user to ask a question and receive an answer that
+  (or in a `.env` file in the project root).
-  incorporates retrieved context.
+- Text files placed in the `data/` directory (one file per document).
-The implementation avoids external AI services or heavy dependencies,
+Author: Artur Kuzakhmetov
-making it suitable for the Deep Agents Virtual File System environment.
+Version: 20
 """
 import os
 import sys
-import math
+from pathlib import Path
-import argparse
+from typing import List
-from collections import defaultdict
+
-from typing import Dict, List, Tuple
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from dotenv import load_dotenv
 # LangChain imports
 from langchain.document_loaders import DirectoryLoader
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
 from langchain.llms import OpenAI
 # --------------------------------------------------------------------------- #
-# Utility functions
+# Configuration
 # --------------------------------------------------------------------------- #
-def tokenize(text: str) -> List[str]:
+# Load environment variables from .env if present
-    """
+load_dotenv()
    Very simple tokenizer: lowercases, splits on whitespace, removes
    punctuation.
    """
    import string
    translator = str.maketrans('', '', string.punctuation)
    return text.translate(translator).lower().split()
-def vectorize(tokens: List[str]) -> Dict[str, int]:
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-    """
+if not OPENAI_API_KEY:
-    Convert a list of tokens into a bag‑of‑words vector (word -> count).
+    sys.exit("Error: OPENAI_API_KEY not found in environment variables.")
    """
    vec = defaultdict(int)
    for token in tokens:
        vec[token] += 1
    return dict(vec)
 def dot_product(v1: Dict[str, int], v2: Dict[str, int]) -> int:
    """
    Compute dot product of two sparse vectors represented as dicts.
    """
    if len(v1) > len(v2):
        v1, v2 = v2, v1
    return sum(v1.get(k, 0) * v2.get(k, 0) for k in v1)
 def vector_norm(v: Dict[str, int]) -> float:
    """
    Compute Euclidean norm of a sparse vector.
    """
    return math.sqrt(sum(count * count for count in v.values()))
 def cosine_similarity(v1: Dict[str, int], v2: Dict[str, int]) -> float:
    """
    Compute cosine similarity between two sparse vectors.
    """
    denom = vector_norm(v1) * vector_norm(v2)
    if denom == 0:
        return 0.0
    return dot_product(v1, v2) / denom
 # --------------------------------------------------------------------------- #
-# RAG Memory
+# Data loading and vector store initialization
 # --------------------------------------------------------------------------- #
-class RAGMemory:
+DATA_DIR = Path(__file__).parent.parent / "data"
 def load_documents(path: Path) -> List:
    """
-    Simple RAG memory that stores documents and their embeddings.
+    Load all text documents from the specified directory.
    Retrieval is performed via cosine similarity over bag‑of‑words vectors.
    """
    if not path.exists() or not path.is_dir():
        print(f"Warning: Data directory '{path}' not found. No documents loaded.")
        return []
-    def __init__(self):
+    loader = DirectoryLoader(str(path), glob="**/*.txt")
-        # doc_id -> (text, vector)
+    documents = loader.load()
-        self._store: Dict[str, Tuple[str, Dict[str, int]]] = {}
+    print(f"Loaded {len(documents)} documents from '{path}'.")
    return documents
-    def add_document(self, doc_id: str, text: str) -> None:
+def create_vectorstore(documents: List) -> FAISS:
        """
        Add a document to the memory.
        """
        tokens = tokenize(text)
        vec = vectorize(tokens)
        self._store[doc_id] = (text, vec)
    def load_from_directory(self, path: str) -> None:
        """
        Load all .txt files from a directory as documents.
        """
        for filename in os.listdir(path):
            if filename.lower().endswith('.txt'):
                doc_id = os.path.splitext(filename)[0]
                with open(os.path.join(path, filename), 'r', encoding='utf-8') as f:
                    text = f.read()
                self.add_document(doc_id, text)
    def retrieve(self, query: str, top_k: int = 3) -> List[Tuple[str, str, float]]:
        """
        Retrieve top_k documents most relevant to the query.
        Returns a list of tuples: (doc_id, text, similarity_score).
        """
        query_vec = vectorize(tokenize(query))
        scores = []
        for doc_id, (text, vec) in self._store.items():
            score = cosine_similarity(query_vec, vec)
            scores.append((doc_id, text, score))
        # Sort by descending similarity
        scores.sort(key=lambda x: x[2], reverse=True)
        return scores[:top_k]
 # --------------------------------------------------------------------------- #
 # Agent
 # --------------------------------------------------------------------------- #
 class Agent:
    """
-    Educational agent that uses RAGMemory to augment its responses.
+    Create a FAISS vector store from the provided documents.
    """
-
+    embeddings = OpenAIEmbeddings()
-    def __init__(self, memory: RAGMemory):
+    vectorstore = FAISS.from_documents(documents, embeddings)
-        self.memory = memory
+    print("FAISS vector store created.")
-
+    return vectorstore
    def answer(self, question: str, top_k: int = 3) -> str:
        """
        Generate an answer to the question by retrieving context and
        concatenating it with a simple rule‑based response.
        """
        retrieved = self.memory.retrieve(question, top_k=top_k)
        context_parts = []
        for doc_id, text, score in retrieved:
            snippet = text[:200].replace('\n', ' ')  # short snippet
            context_parts.append(f"[{doc_id}] {snippet} (score={score:.2f})")
        context = "\n".join(context_parts) if context_parts else "No relevant context found."
        answer = (
            f"Question: {question}\n\n"
            f"Context:\n{context}\n\n"
            f"Answer: (This is a placeholder answer generated by a rule‑based system.)"
        )
        return answer
 # --------------------------------------------------------------------------- #
-# CLI
+# Agent construction
 # --------------------------------------------------------------------------- #
-def main():
+def build_agent(vectorstore: FAISS) -> RetrievalQA:
-    parser = argparse.ArgumentParser(description="Educational Agent with RAG Memory")
+    """
-    parser.add_argument("--data-dir", type=str, required=True,
+    Build a RetrievalQA chain that uses the vector store for retrieval
-                        help="Directory containing .txt documents for the memory.")
+    and OpenAI GPT-4 for generation.
-    parser.add_argument("--question", type=str, required=True,
+    """
-                        help="The question to ask the agent.")
+    llm = OpenAI(model_name="gpt-4", temperature=0, openai_api_key=OPENAI_API_KEY)
-    parser.add_argument("--top-k", type=int, default=3,
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 4})
-                        help="Number of top documents to retrieve.")
+    qa_chain = RetrievalQA.from_chain_type(
-    args = parser.parse_args()
+        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True,
    )
    print("RetrievalQA agent constructed.")
    return qa_chain
-    memory = RAGMemory()
+# --------------------------------------------------------------------------- #
-    memory.load_from_directory(args.data_dir)
+# FastAPI application
 # --------------------------------------------------------------------------- #
-    agent = Agent(memory)
+app = FastAPI(title="RAG Agent API", version="1.0.0")
    response = agent.answer(args.question, top_k=args.top_k)
    print(response)
-if __name__ == "__main__":
+class QuestionRequest(BaseModel):
-    main()
+    question: str
 class AnswerResponse(BaseModel):
    answer: str
    sources: List[str] = []
 # Global variables to hold the agent and vector store
 vectorstore: FAISS = None
 agent: RetrievalQA = None
@app.on_event("startup")
 def startup_event():
    """
    Load documents, create vector store, and build the agent on startup.
    """
    global vectorstore, agent
    docs = load_documents(DATA_DIR)
    vectorstore = create_vectorstore(docs)
    agent = build_agent(vectorstore)
@app.post("/ask", response_model=AnswerResponse)
 def ask_question(request: QuestionRequest):
    """
    Endpoint to query the RAG agent.
    """
    if not agent:
        raise HTTPException(status_code=500, detail="Agent not initialized.")
    try:
        result = agent({"question": request.question})
        answer = result.get("answer", "")
        sources = [doc.metadata.get("source", "") for doc in result.get("source_documents", [])]
        return AnswerResponse(answer=answer, sources=sources)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
 # --------------------------------------------------------------------------- #
 # Run with: uvicorn src.index:app --reload
 # --------------------------------------------------------------------------- #
@@ -1,89 +1,78 @@
-"""
+import os
-Knowledge base implementation using Ollama embeddings.
+import logging
-Provides tools for searching and adding documents.
+from typing import List, Tuple
 """
 import faiss
 import numpy as np
-from typing import List
+from sentence_transformers import SentenceTransformer
-from langchain.docstore.document import Document
+
-from embeddings import get_embedding_model
+logger = logging.getLogger(__name__)
 class KnowledgeBase:
    """
-    In-memory knowledge base that stores documents and their embeddings.
+    Loads documents, creates embeddings, and stores them in a FAISS index.
    Provides retrieval of top-k relevant passages for a query.
    """
-    def __init__(self):
+    def __init__(self, data_dir: str, embedding_model: str, vector_store: str = "faiss"):
-        self.embedding_model = get_embedding_model()
+        self.data_dir = data_dir
-        self.documents: List[Document] = []
+        self.embedding_model_name = embedding_model
-        self.embeddings: np.ndarray = np.empty((0, self.embedding_model.get_sentence_embedding_dimension()))
+        self.vector_store = vector_store
        self.documents: List[str] = []
        self.embeddings: np.ndarray = None
        self.index: faiss.IndexFlatIP = None
        self._load_documents()
        self._create_embeddings()
        self._build_index()
-    def add_to_knowledge_base(self, content: str) -> str:
+    def _load_documents(self):
-        """
+        """Load all .txt files from the data directory."""
-        Adds a new document to the knowledge base.
+        logger.info(f"Loading documents from {self.data_dir}")
-
+        if not os.path.isdir(self.data_dir):
-        Args:
+            raise FileNotFoundError(f"Data directory {self.data_dir} does not exist.")
-            content (str): The text content to add.
+        for filename in os.listdir(self.data_dir):
-
+            if filename.lower().endswith(".txt"):
-        Returns:
+                path = os.path.join(self.data_dir, filename)
-            str: Confirmation message.
+                with open(path, "r", encoding="utf-8") as f:
-        """
+                    content = f.read()
-        doc = Document(page_content=content)
+                    self.documents.append(content)
        embedding = self.embedding_model.embed_query(content)
        embedding = np.array(embedding).reshape(1, -1)
        self.documents.append(doc)
        if self.embeddings.size == 0:
            self.embeddings = embedding
        else:
            self.embeddings = np.vstack([self.embeddings, embedding])
        return f"Document added. Total documents: {len(self.documents)}."
    def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
        """
        Searches the knowledge base for the most relevant documents.
        Args:
            query (str): The search query.
            k (int): Number of top documents to return.
        Returns:
            List[Document]: List of top matching documents.
        """
        if not self.documents:
            raise ValueError(f"No .txt documents found in {self.data_dir}")
        logger.info(f"Loaded {len(self.documents)} documents.")
    def _create_embeddings(self):
        """Generate embeddings for all documents."""
        logger.info(f"Creating embeddings using {self.embedding_model_name}")
        model = SentenceTransformer(self.embedding_model_name)
        self.embeddings = model.encode(self.documents, convert_to_numpy=True, normalize_embeddings=True)
        logger.info(f"Generated embeddings of shape {self.embeddings.shape}")
    def _build_index(self):
        """Build a FAISS index for efficient similarity search."""
        logger.info("Building FAISS index.")
        dim = self.embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)  # Inner product (cosine similarity after normalization)
        self.index.add(self.embeddings)
        logger.info(f"FAISS index built with {self.index.ntotal} vectors.")
    def retrieve(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
        """
        Retrieve top_k most relevant passages for the query.
        Returns a list of tuples: (document_text, similarity_score)
        """
        if not query.strip():
            logger.warning("Empty query received; returning empty list.")
            return []
-        query_embedding = self.embedding_model.embed_query(query)
+        logger.debug(f"Retrieving top {top_k} passages for query: {query}")
-        query_embedding = np.array(query_embedding).reshape(1, -1)
+        model = SentenceTransformer(self.embedding_model_name)
-
+        query_vec = model.encode([query], convert_to_numpy=True, normalize_embeddings=True)
-        similarities = np.dot(self.embeddings, query_embedding.T).flatten()
+        distances, indices = self.index.search(query_vec, top_k)
-        top_indices = similarities.argsort()[-k:][::-1]
+        results = []
-        return [self.documents[i] for i in top_indices]
+        for idx, score in zip(indices[0], distances[0]):
-
+            if idx < 0:
-# Global knowledge base instance
+                continue
-kb = KnowledgeBase()
+            results.append((self.documents[idx], float(score)))
-
+        logger.debug(f"Retrieved {len(results)} passages.")
-def add_to_knowledge_base(content: str) -> str:
+        return results
    """
    Tool wrapper for adding content to the knowledge base.
    Args:
        content (str): Text to add.
    Returns:
        str: Confirmation message.
    """
    return kb.add_to_knowledge_base(content)
 def search_knowledge_base(query: str) -> List[Document]:
    """
    Tool wrapper for searching the knowledge base.
    Args:
        query (str): Search query.
    Returns:
        List[Document]: Matching documents.
    """
    return kb.search_knowledge_base(query)
@@ -1,44 +1,31 @@
-"""
+import argparse
-FastAPI application exposing the RAG agent as a REST endpoint.
+import logging
 This file is optional but useful for running the agent in a container.
 """
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from .agent import RAGAgent
-app = FastAPI(title="RAG Agent API")
+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()
-# Initialize a global agent instance
+    logging.basicConfig(level=logging.INFO)
-agent = RAGAgent(model="llama2", top_k=3)
+    agent = RAGAgent(config_path=args.config)
    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
-class Document(BaseModel):
+if __name__ == "__main__":
-    text: str
+    main()
 class Query(BaseModel):
    query: str
@app.post("/documents")
 def add_document(doc: Document):
    """
    Add a document to the agent's memory.
    """
    agent.add_document(doc.text)
    return {"status": "added"}
@app.post("/ask")
 def ask(query: Query):
    """
    Get an answer to a query using the RAG agent.
    """
    try:
        answer = agent.get_response(query.query)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))
    return {"answer": answer}
@@ -0,0 +1,33 @@
 import os
 import unittest
 from src.agent import RAGAgent
 class TestRAGAgent(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Ensure data directory exists with at least one document
        data_dir = "data"
        os.makedirs(data_dir, exist_ok=True)
        sample_path = os.path.join(data_dir, "sample.txt")
        with open(sample_path, "w", encoding="utf-8") as f:
            f.write("Python is a versatile programming language used for web development, data science, and automation.")
        cls.agent = RAGAgent(config_path="src/config.yaml")
    def test_retrieve_non_empty(self):
        passages = self.agent.kb.retrieve("Python programming", top_k=2)
        self.assertTrue(len(passages) > 0)
        self.assertIn("Python is a versatile programming language", passages[0][0])
    def test_generate_response(self):
        answer = self.agent.generate_response("What is Python?")
        self.assertIsInstance(answer, str)
        self.assertTrue(len(answer) > 0)
    def test_empty_query(self):
        answer = self.agent.generate_response("")
        self.assertIsInstance(answer, str)
        self.assertIn("No relevant information found", answer)
 if __name__ == "__main__":
    unittest.main()