This commit is contained in:
+73
-40
@@ -1,50 +1,83 @@
|
||||
"""
|
||||
Agent implementation that performs RAG memory retrieval and response generation.
|
||||
Agent implementation that uses ChromaDBVectorStore for RAG.
|
||||
"""
|
||||
|
||||
from langchain import PromptTemplate, LLMChain
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.memory import ConversationBufferMemory
|
||||
from langchain.llms import OpenAI
|
||||
from langchain.vectorstores import Qdrant
|
||||
from config import OPENAI_API_KEY, OPENAI_MODEL
|
||||
from vector_store import get_vector_store
|
||||
from typing import List, Optional
|
||||
|
||||
def build_agent() -> RetrievalQA:
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from langchain.schema import Document
|
||||
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
|
||||
|
||||
from .vector_store import ChromaDBVectorStore
|
||||
|
||||
|
||||
class RAGAgent:
|
||||
"""
|
||||
Builds and returns a RetrievalQA chain configured with:
|
||||
- OpenAI LLM for generation
|
||||
- Qdrant vector store for retrieval
|
||||
- ConversationBufferMemory for context
|
||||
A simple RAG agent that retrieves relevant documents from a vector store
|
||||
and generates answers using an LLM.
|
||||
"""
|
||||
# LLM for generation
|
||||
llm = OpenAI(
|
||||
temperature=0,
|
||||
openai_api_key=OPENAI_API_KEY,
|
||||
model_name=OPENAI_MODEL
|
||||
)
|
||||
|
||||
# Vector store and retriever
|
||||
vector_store: Qdrant = get_vector_store()
|
||||
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
|
||||
def __init__(
|
||||
self,
|
||||
vector_store: ChromaDBVectorStore,
|
||||
llm_model: str = "gpt-4o-mini",
|
||||
temperature: float = 0.2,
|
||||
):
|
||||
"""
|
||||
Initialize the agent.
|
||||
|
||||
# Memory to keep conversation context
|
||||
memory = ConversationBufferMemory(
|
||||
memory_key="chat_history",
|
||||
return_messages=True
|
||||
)
|
||||
Args:
|
||||
vector_store: Instance of ChromaDBVectorStore.
|
||||
llm_model: OpenAI LLM model name.
|
||||
temperature: Sampling temperature for the LLM.
|
||||
"""
|
||||
self.vector_store = vector_store
|
||||
self.llm = ChatOpenAI(model=llm_model, temperature=temperature)
|
||||
|
||||
# RetrievalQA chain
|
||||
chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=retriever,
|
||||
memory=memory
|
||||
)
|
||||
return chain
|
||||
# Prompt template
|
||||
self.prompt = ChatPromptTemplate.from_messages(
|
||||
[
|
||||
("system", "You are a helpful assistant."),
|
||||
("human", "Use the following context to answer the question."),
|
||||
MessagesPlaceholder("context"),
|
||||
("human", "Question: {question}"),
|
||||
]
|
||||
)
|
||||
|
||||
def ask_question(chain: RetrievalQA, question: str) -> str:
|
||||
"""
|
||||
Utility function to ask a question using the provided chain.
|
||||
"""
|
||||
return chain.run(question)
|
||||
def add_documents(self, documents: List[Document]) -> None:
|
||||
"""
|
||||
Add documents to the underlying vector store.
|
||||
|
||||
Args:
|
||||
documents: List of langchain.schema.Document objects.
|
||||
"""
|
||||
self.vector_store.add_documents(documents)
|
||||
|
||||
def ask(self, question: str, k: int = 4) -> str:
|
||||
"""
|
||||
Retrieve relevant documents and generate an answer.
|
||||
|
||||
Args:
|
||||
question: The user question.
|
||||
k: Number of documents to retrieve.
|
||||
|
||||
Returns:
|
||||
The LLM-generated answer as a string.
|
||||
"""
|
||||
relevant_docs = self.vector_store.similarity_search(question, k=k)
|
||||
context = "\n\n".join([doc.page_content for doc in relevant_docs])
|
||||
|
||||
# Build messages
|
||||
messages = self.prompt.format(context=context, question=question)
|
||||
|
||||
# Generate answer
|
||||
response = self.llm(messages)
|
||||
return response.content.strip()
|
||||
|
||||
def get_document_count(self) -> int:
|
||||
"""Return the number of documents stored."""
|
||||
return self.vector_store.count()
|
||||
|
||||
def clear_store(self) -> None:
|
||||
"""Clear all documents from the vector store."""
|
||||
self.vector_store.delete_all()
|
||||
+82
-20
@@ -1,32 +1,94 @@
|
||||
"""
|
||||
Entry point for the RAG agent.
|
||||
Command-line interface for the RAG agent.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
from agent import build_agent, ask_question
|
||||
from config import OPENAI_API_KEY
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
from langchain.schema import Document
|
||||
|
||||
from .agent import RAGAgent
|
||||
from .vector_store import ChromaDBVectorStore
|
||||
|
||||
|
||||
def load_text_files(directory: str) -> List[Document]:
|
||||
"""
|
||||
Load all .txt files from a directory into Document objects.
|
||||
|
||||
Args:
|
||||
directory: Path to the directory containing text files.
|
||||
|
||||
Returns:
|
||||
List of Document objects.
|
||||
"""
|
||||
docs = []
|
||||
for file_path in Path(directory).glob("*.txt"):
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
docs.append(Document(page_content=content, metadata={"source": str(file_path)}))
|
||||
return docs
|
||||
|
||||
|
||||
def main() -> None:
|
||||
# Load environment variables from .env if present
|
||||
load_dotenv()
|
||||
parser = argparse.ArgumentParser(description="RAG Agent CLI")
|
||||
parser.add_argument(
|
||||
"--docs",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to directory containing .txt documents to index.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--question",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Question to ask the agent.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--persist",
|
||||
type=str,
|
||||
default="./chromadb",
|
||||
help="Directory to persist ChromaDB data.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="gpt-4o-mini",
|
||||
help="OpenAI LLM model to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--k",
|
||||
type=int,
|
||||
default=4,
|
||||
help="Number of documents to retrieve for RAG.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Ensure OpenAI API key is available
|
||||
if not OPENAI_API_KEY:
|
||||
raise RuntimeError("OPENAI_API_KEY is not set. Please set it in environment or .env file.")
|
||||
# Initialize vector store
|
||||
vector_store = ChromaDBVectorStore(
|
||||
persist_directory=args.persist,
|
||||
collection_name="rag_collection",
|
||||
)
|
||||
|
||||
# Build the agent
|
||||
chain = build_agent()
|
||||
# If collection is empty, load documents
|
||||
if vector_store.count() == 0:
|
||||
print("Indexing documents...")
|
||||
docs = load_text_files(args.docs)
|
||||
vector_store.add_documents(docs)
|
||||
print(f"Indexed {len(docs)} documents.")
|
||||
else:
|
||||
print(f"Using existing index with {vector_store.count()} documents.")
|
||||
|
||||
# Initialize agent
|
||||
agent = RAGAgent(vector_store=vector_store, llm_model=args.model)
|
||||
|
||||
# Ask question
|
||||
answer = agent.ask(args.question, k=args.k)
|
||||
print("\n=== Answer ===")
|
||||
print(answer)
|
||||
|
||||
# Simple interactive loop
|
||||
print("RAG Agent is ready. Type 'exit' to quit.")
|
||||
while True:
|
||||
user_input = input("\nYou: ")
|
||||
if user_input.lower() in {"exit", "quit"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
response = ask_question(chain, user_input)
|
||||
print(f"Agent: {response}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+107
-33
@@ -1,42 +1,116 @@
|
||||
"""
|
||||
Vector store implementation using Qdrant.
|
||||
Vector store implementation using ChromaDB.
|
||||
"""
|
||||
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qdrant_models
|
||||
from langchain.vectorstores import Qdrant
|
||||
from config import QDRANT_HOST, QDRANT_PORT, QDRANT_COLLECTION_NAME
|
||||
from embeddings import get_ollama_embeddings
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
def get_qdrant_client() -> QdrantClient:
|
||||
"""
|
||||
Creates a Qdrant client connected to the local Qdrant instance.
|
||||
"""
|
||||
return QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
||||
from chromadb import Client
|
||||
from chromadb.config import Settings
|
||||
from chromadb.errors import ChromaError
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from langchain.schema import Document
|
||||
|
||||
def ensure_collection(client: QdrantClient, collection_name: str, vector_size: int = 768) -> None:
|
||||
|
||||
class ChromaDBVectorStore:
|
||||
"""
|
||||
Ensures that the specified collection exists in Qdrant.
|
||||
If it does not exist, it will be created with the given vector size.
|
||||
A vector store backed by ChromaDB. It handles embedding generation,
|
||||
persistence, and similarity search.
|
||||
"""
|
||||
if not client.has_collection(collection_name):
|
||||
client.recreate_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=qdrant_models.VectorParams(
|
||||
size=vector_size,
|
||||
distance="Cosine"
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
persist_directory: str = "./chromadb",
|
||||
collection_name: str = "rag_collection",
|
||||
embedding_model: str = "text-embedding-3-small",
|
||||
):
|
||||
"""
|
||||
Initialize the ChromaDB client and collection.
|
||||
|
||||
Args:
|
||||
persist_directory: Directory where ChromaDB will store data.
|
||||
collection_name: Name of the collection to use.
|
||||
embedding_model: OpenAI embedding model name.
|
||||
"""
|
||||
self.persist_directory = persist_directory
|
||||
self.collection_name = collection_name
|
||||
self.embedding_model = embedding_model
|
||||
|
||||
# Ensure persistence directory exists
|
||||
os.makedirs(self.persist_directory, exist_ok=True)
|
||||
|
||||
# Initialize Chroma client
|
||||
self.client = Client(Settings(
|
||||
chroma_db_impl="duckdb+parquet",
|
||||
persist_directory=self.persist_directory,
|
||||
))
|
||||
|
||||
# Create or get collection
|
||||
try:
|
||||
self.collection = self.client.get_collection(name=self.collection_name)
|
||||
except ChromaError:
|
||||
self.collection = self.client.create_collection(name=self.collection_name)
|
||||
|
||||
# Embedding model
|
||||
self.embedder = OpenAIEmbeddings(model=self.embedding_model)
|
||||
|
||||
def add_documents(self, documents: List[Document]) -> None:
|
||||
"""
|
||||
Add documents to the collection. Each document is embedded and stored.
|
||||
|
||||
Args:
|
||||
documents: List of langchain.schema.Document objects.
|
||||
"""
|
||||
ids = []
|
||||
metadatas = []
|
||||
embeddings = []
|
||||
|
||||
for idx, doc in enumerate(documents):
|
||||
ids.append(f"doc_{len(self.collection.get()['ids']) + idx}")
|
||||
metadatas.append({"source": doc.metadata.get("source", "")})
|
||||
embeddings.append(self.embedder.embed_query(doc.page_content))
|
||||
|
||||
self.collection.add(
|
||||
documents=[doc.page_content for doc in documents],
|
||||
embeddings=embeddings,
|
||||
ids=ids,
|
||||
metadatas=metadatas,
|
||||
)
|
||||
|
||||
def get_vector_store() -> Qdrant:
|
||||
"""
|
||||
Returns a Qdrant vector store instance ready for use with LangChain.
|
||||
"""
|
||||
client = get_qdrant_client()
|
||||
ensure_collection(client, QDRANT_COLLECTION_NAME)
|
||||
embeddings = get_ollama_embeddings()
|
||||
return Qdrant(
|
||||
client=client,
|
||||
collection_name=QDRANT_COLLECTION_NAME,
|
||||
embeddings=embeddings
|
||||
)
|
||||
def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 4,
|
||||
filter: Optional[dict] = None,
|
||||
) -> List[Document]:
|
||||
"""
|
||||
Retrieve the top-k most similar documents to the query.
|
||||
|
||||
Args:
|
||||
query: The query string.
|
||||
k: Number of results to return.
|
||||
filter: Optional metadata filter.
|
||||
|
||||
Returns:
|
||||
List of langchain.schema.Document objects.
|
||||
"""
|
||||
query_embedding = self.embedder.embed_query(query)
|
||||
results = self.collection.query(
|
||||
query_embeddings=[query_embedding],
|
||||
n_results=k,
|
||||
where=filter,
|
||||
)
|
||||
|
||||
docs = []
|
||||
for content, metadata in zip(results["documents"][0], results["metadatas"][0]):
|
||||
docs.append(Document(page_content=content, metadata=metadata))
|
||||
return docs
|
||||
|
||||
def count(self) -> int:
|
||||
"""Return the number of documents stored."""
|
||||
return len(self.collection.get()["ids"])
|
||||
|
||||
def delete_all(self) -> None:
|
||||
"""Delete all documents from the collection."""
|
||||
self.collection.delete(where={})
|
||||
self.client.persist()
|
||||
Reference in New Issue
Block a user