This commit is contained in:
+49
-84
@@ -1,97 +1,62 @@
|
||||
"""
|
||||
RAG Agent implementation.
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
Stores documents in memory, retrieves top-k relevant documents using cosine similarity,
|
||||
constructs a prompt with context, and generates a response via Ollama LLM.
|
||||
"""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from typing import List, Tuple
|
||||
import numpy as np
|
||||
|
||||
from .embeddings import embed, cosine_similarity
|
||||
from .llm import chat
|
||||
from .knowledge_base import KnowledgeBase
|
||||
from .config import load_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
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.
|
||||
Retrieval-Augmented Generation agent.
|
||||
"""
|
||||
|
||||
def __init__(self, model: str = "llama2", top_k: int = 3):
|
||||
self.model = model
|
||||
self.top_k = top_k
|
||||
# Store tuples of (embedding, text)
|
||||
self._store: List[Tuple[List[float], str]] = []
|
||||
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}")
|
||||
|
||||
def add_document(self, text: str) -> None:
|
||||
"""
|
||||
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:
|
||||
def generate_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)
|
||||
context = "\n\n".join(context_docs)
|
||||
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}")
|
||||
|
||||
system_prompt = (
|
||||
"You are an assistant that uses the provided context to answer the question."
|
||||
)
|
||||
user_prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
]
|
||||
|
||||
response = chat(messages, model=self.model)
|
||||
return response.strip()
|
||||
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
|
||||
+7
-24
@@ -1,26 +1,9 @@
|
||||
"""
|
||||
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 yaml
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env")
|
||||
|
||||
# Qdrant configuration
|
||||
QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
|
||||
QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
|
||||
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")
|
||||
def load_config(path: str) -> dict:
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"Config file {path} not found.")
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
return cfg
|
||||
@@ -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"]
|
||||
+109
-139
@@ -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
|
||||
libraries. It provides:
|
||||
This module implements a FastAPI application that exposes a single endpoint
|
||||
`/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,
|
||||
and retrieves the most relevant passages for a query.
|
||||
* Agent – integrates the memory with a rule‑based response generator.
|
||||
* CLI – allows the user to ask a question and receive an answer that
|
||||
incorporates retrieved context.
|
||||
Prerequisites:
|
||||
- Python 3.11+
|
||||
- OpenAI API key set in the environment variable `OPENAI_API_KEY`
|
||||
(or in a `.env` file in the project root).
|
||||
- Text files placed in the `data/` directory (one file per document).
|
||||
|
||||
The implementation avoids external AI services or heavy dependencies,
|
||||
making it suitable for the Deep Agents Virtual File System environment.
|
||||
Author: Artur Kuzakhmetov
|
||||
Version: 20
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import math
|
||||
import argparse
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
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]:
|
||||
"""
|
||||
Very simple tokenizer: lowercases, splits on whitespace, removes
|
||||
punctuation.
|
||||
"""
|
||||
import string
|
||||
translator = str.maketrans('', '', string.punctuation)
|
||||
return text.translate(translator).lower().split()
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv()
|
||||
|
||||
def vectorize(tokens: List[str]) -> Dict[str, int]:
|
||||
"""
|
||||
Convert a list of tokens into a bag‑of‑words vector (word -> count).
|
||||
"""
|
||||
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
|
||||
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
if not OPENAI_API_KEY:
|
||||
sys.exit("Error: OPENAI_API_KEY not found in environment variables.")
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 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.
|
||||
Retrieval is performed via cosine similarity over bag‑of‑words vectors.
|
||||
Load all text documents from the specified directory.
|
||||
"""
|
||||
if not path.exists() or not path.is_dir():
|
||||
print(f"Warning: Data directory '{path}' not found. No documents loaded.")
|
||||
return []
|
||||
|
||||
def __init__(self):
|
||||
# doc_id -> (text, vector)
|
||||
self._store: Dict[str, Tuple[str, Dict[str, int]]] = {}
|
||||
loader = DirectoryLoader(str(path), glob="**/*.txt")
|
||||
documents = loader.load()
|
||||
print(f"Loaded {len(documents)} documents from '{path}'.")
|
||||
return documents
|
||||
|
||||
def add_document(self, doc_id: str, text: str) -> None:
|
||||
"""
|
||||
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:
|
||||
def create_vectorstore(documents: List) -> FAISS:
|
||||
"""
|
||||
Educational agent that uses RAGMemory to augment its responses.
|
||||
Create a FAISS vector store from the provided documents.
|
||||
"""
|
||||
|
||||
def __init__(self, memory: RAGMemory):
|
||||
self.memory = memory
|
||||
|
||||
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
|
||||
embeddings = OpenAIEmbeddings()
|
||||
vectorstore = FAISS.from_documents(documents, embeddings)
|
||||
print("FAISS vector store created.")
|
||||
return vectorstore
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# CLI
|
||||
# Agent construction
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Educational Agent with RAG Memory")
|
||||
parser.add_argument("--data-dir", type=str, required=True,
|
||||
help="Directory containing .txt documents for the memory.")
|
||||
parser.add_argument("--question", type=str, required=True,
|
||||
help="The question to ask the agent.")
|
||||
parser.add_argument("--top-k", type=int, default=3,
|
||||
help="Number of top documents to retrieve.")
|
||||
args = parser.parse_args()
|
||||
def build_agent(vectorstore: FAISS) -> RetrievalQA:
|
||||
"""
|
||||
Build a RetrievalQA chain that uses the vector store for retrieval
|
||||
and OpenAI GPT-4 for generation.
|
||||
"""
|
||||
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
|
||||
|
||||
memory = RAGMemory()
|
||||
memory.load_from_directory(args.data_dir)
|
||||
# --------------------------------------------------------------------------- #
|
||||
# FastAPI application
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
agent = Agent(memory)
|
||||
response = agent.answer(args.question, top_k=args.top_k)
|
||||
print(response)
|
||||
app = FastAPI(title="RAG Agent API", version="1.0.0")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
class QuestionRequest(BaseModel):
|
||||
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
|
||||
# --------------------------------------------------------------------------- #
|
||||
+67
-78
@@ -1,89 +1,78 @@
|
||||
"""
|
||||
Knowledge base implementation using Ollama embeddings.
|
||||
Provides tools for searching and adding documents.
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
from typing import List
|
||||
from langchain.docstore.document import Document
|
||||
from embeddings import get_embedding_model
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
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):
|
||||
self.embedding_model = get_embedding_model()
|
||||
self.documents: List[Document] = []
|
||||
self.embeddings: np.ndarray = np.empty((0, self.embedding_model.get_sentence_embedding_dimension()))
|
||||
def __init__(self, data_dir: str, embedding_model: str, vector_store: str = "faiss"):
|
||||
self.data_dir = data_dir
|
||||
self.embedding_model_name = embedding_model
|
||||
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:
|
||||
"""
|
||||
Adds a new document to the knowledge base.
|
||||
|
||||
Args:
|
||||
content (str): The text content to add.
|
||||
|
||||
Returns:
|
||||
str: Confirmation message.
|
||||
"""
|
||||
doc = Document(page_content=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.
|
||||
"""
|
||||
def _load_documents(self):
|
||||
"""Load all .txt files from the data directory."""
|
||||
logger.info(f"Loading documents from {self.data_dir}")
|
||||
if not os.path.isdir(self.data_dir):
|
||||
raise FileNotFoundError(f"Data directory {self.data_dir} does not exist.")
|
||||
for filename in os.listdir(self.data_dir):
|
||||
if filename.lower().endswith(".txt"):
|
||||
path = os.path.join(self.data_dir, filename)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.documents.append(content)
|
||||
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)
|
||||
query_embedding = np.array(query_embedding).reshape(1, -1)
|
||||
|
||||
similarities = np.dot(self.embeddings, query_embedding.T).flatten()
|
||||
top_indices = similarities.argsort()[-k:][::-1]
|
||||
return [self.documents[i] for i in top_indices]
|
||||
|
||||
# Global knowledge base instance
|
||||
kb = KnowledgeBase()
|
||||
|
||||
def add_to_knowledge_base(content: str) -> str:
|
||||
"""
|
||||
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)
|
||||
logger.debug(f"Retrieving top {top_k} passages for query: {query}")
|
||||
model = SentenceTransformer(self.embedding_model_name)
|
||||
query_vec = model.encode([query], convert_to_numpy=True, normalize_embeddings=True)
|
||||
distances, indices = self.index.search(query_vec, top_k)
|
||||
results = []
|
||||
for idx, score in zip(indices[0], distances[0]):
|
||||
if idx < 0:
|
||||
continue
|
||||
results.append((self.documents[idx], float(score)))
|
||||
logger.debug(f"Retrieved {len(results)} passages.")
|
||||
return results
|
||||
+25
-38
@@ -1,44 +1,31 @@
|
||||
"""
|
||||
FastAPI application exposing the RAG agent as a REST endpoint.
|
||||
|
||||
This file is optional but useful for running the agent in a container.
|
||||
"""
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
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
|
||||
agent = RAGAgent(model="llama2", top_k=3)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
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):
|
||||
text: str
|
||||
|
||||
|
||||
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}
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user