This commit is contained in:
+17
-57
@@ -1,62 +1,22 @@
|
||||
import logging
|
||||
from typing import List
|
||||
from langchain_ollama import Ollama, OllamaEmbeddings
|
||||
from langchain.chains import RetrievalQA
|
||||
import config
|
||||
from src.vector_store import QdrantVectorStore
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from .knowledge_base import KnowledgeBase
|
||||
from .config import load_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RAGAgent:
|
||||
def create_agent(vector_store: QdrantVectorStore):
|
||||
"""
|
||||
Retrieval-Augmented Generation agent.
|
||||
Create a RetrievalQA agent that uses Ollama for both embeddings and LLM.
|
||||
"""
|
||||
# Embeddings for the vector store
|
||||
embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL)
|
||||
|
||||
def __init__(self, config_path: str = "src/config.yaml"):
|
||||
self.config = load_config(config_path)
|
||||
logging.basicConfig(level=self.config["logging"]["level"])
|
||||
logger.info("Initializing RAGAgent.")
|
||||
self.kb = KnowledgeBase(
|
||||
data_dir=self.config["knowledge_base"]["data_dir"],
|
||||
embedding_model=self.config["knowledge_base"]["embedding_model"],
|
||||
vector_store=self.config["knowledge_base"]["vector_store"],
|
||||
)
|
||||
self.model_name = self.config["language_model"]["model_name"]
|
||||
self.max_length = self.config["language_model"]["max_length"]
|
||||
self.top_k = self.config["retrieval"]["top_k"]
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
||||
self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
|
||||
self.model.eval()
|
||||
if torch.cuda.is_available():
|
||||
self.model.to("cuda")
|
||||
logger.info(f"Loaded language model {self.model_name}")
|
||||
# LLM for generating answers
|
||||
llm = Ollama(model=config.OLLAMA_MODEL)
|
||||
|
||||
def generate_response(self, query: str) -> str:
|
||||
"""
|
||||
Generate a response to the query using retrieved context.
|
||||
"""
|
||||
logger.info(f"Generating response for query: {query}")
|
||||
passages = self.kb.retrieve(query, top_k=self.top_k)
|
||||
context = "\n\n".join([p[0] for p in passages]) if passages else "No relevant information found."
|
||||
prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
|
||||
logger.debug(f"Prompt:\n{prompt}")
|
||||
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt")
|
||||
if torch.cuda.is_available():
|
||||
inputs = {k: v.to("cuda") for k, v in inputs.items()}
|
||||
with torch.no_grad():
|
||||
output_ids = self.model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=self.max_length,
|
||||
do_sample=True,
|
||||
top_p=0.95,
|
||||
temperature=0.7,
|
||||
)
|
||||
answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
||||
# Extract the part after "Answer:" if present
|
||||
if "Answer:" in answer:
|
||||
answer = answer.split("Answer:")[1].strip()
|
||||
logger.info(f"Generated answer: {answer}")
|
||||
return answer
|
||||
# Build the RetrievalQA chain
|
||||
qa_chain = RetrievalQA.from_chain_type(
|
||||
llm=llm,
|
||||
chain_type="stuff",
|
||||
retriever=vector_store.get_retriever(),
|
||||
)
|
||||
return qa_chain
|
||||
+35
-24
@@ -1,31 +1,42 @@
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
from .agent import RAGAgent
|
||||
import os
|
||||
from langchain.schema import Document
|
||||
from src.vector_store import QdrantVectorStore
|
||||
from src.agent import create_agent
|
||||
import config
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Educational RAG Agent CLI")
|
||||
parser.add_argument("--config", type=str, default="src/config.yaml", help="Path to config file")
|
||||
args = parser.parse_args()
|
||||
# Ensure Qdrant is reachable
|
||||
os.environ["QDRANT_URL"] = f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}"
|
||||
if config.QDRANT_API_KEY:
|
||||
os.environ["QDRANT_API_KEY"] = config.QDRANT_API_KEY
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
agent = RAGAgent(config_path=args.config)
|
||||
# Initialize embeddings and vector store
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
embeddings = OllamaEmbeddings(model=config.OLLAMA_MODEL)
|
||||
vector_store = QdrantVectorStore(embeddings)
|
||||
|
||||
print("Welcome to the Educational RAG Agent. Type 'exit' to quit.")
|
||||
while True:
|
||||
try:
|
||||
query = input("\nYour question: ").strip()
|
||||
if query.lower() in ("exit", "quit"):
|
||||
print("Goodbye!")
|
||||
break
|
||||
if not query:
|
||||
print("Please enter a non-empty question.")
|
||||
continue
|
||||
answer = agent.generate_response(query)
|
||||
print(f"\nAnswer:\n{answer}")
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted. Exiting.")
|
||||
break
|
||||
# Add sample documents (only if collection is empty)
|
||||
# In a real scenario, you would load your corpus here
|
||||
sample_docs = [
|
||||
Document(page_content="Hello world! This is a test document.", metadata={"source": "test"}),
|
||||
Document(page_content="LangChain is a powerful framework for building LLM applications.", metadata={"source": "test"}),
|
||||
]
|
||||
# Check if collection already has documents
|
||||
try:
|
||||
# Attempt to retrieve a document to see if collection is populated
|
||||
vector_store.get_retriever().get_relevant_documents("test")
|
||||
except Exception:
|
||||
# If retrieval fails, add documents
|
||||
vector_store.add_documents(sample_docs)
|
||||
|
||||
# Create the agent
|
||||
agent = create_agent(vector_store)
|
||||
|
||||
# Run a sample query
|
||||
query = "What is LangChain?"
|
||||
print(f"Query: {query}")
|
||||
result = agent.run(query)
|
||||
print(f"Answer: {result}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+18
-86
@@ -1,97 +1,29 @@
|
||||
"""
|
||||
Vector store implementation using Qdrant via langchain-qdrant.
|
||||
Provides a simple interface for adding documents and performing
|
||||
similarity search. Embeddings are generated using OpenAIEmbeddings
|
||||
by default, but can be overridden by passing a custom embedding
|
||||
function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, List, Optional
|
||||
|
||||
from langchain.embeddings import OpenAIEmbeddings
|
||||
from langchain_qdrant import Qdrant
|
||||
from langchain.vectorstores import VectorStore
|
||||
from langchain_core.documents import Document
|
||||
from langchain.schema import Document
|
||||
import config
|
||||
|
||||
from .config import (
|
||||
QDRANT_HOST,
|
||||
QDRANT_PORT,
|
||||
QDRANT_API_KEY,
|
||||
QDRANT_COLLECTION,
|
||||
)
|
||||
|
||||
|
||||
class QdrantVectorStore(VectorStore):
|
||||
class QdrantVectorStore:
|
||||
"""
|
||||
A wrapper around langchain_qdrant.Qdrant that implements the
|
||||
VectorStore interface expected by LangChain chains.
|
||||
Wrapper around langchain_qdrant.Qdrant to provide a simple interface
|
||||
for adding documents and retrieving a retriever.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embeddings: Optional[OpenAIEmbeddings] = None,
|
||||
collection_name: str = QDRANT_COLLECTION,
|
||||
):
|
||||
self.embeddings = embeddings or OpenAIEmbeddings()
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Initialize Qdrant client
|
||||
self.client = Qdrant(
|
||||
host=QDRANT_HOST,
|
||||
port=QDRANT_PORT,
|
||||
api_key=QDRANT_API_KEY,
|
||||
def __init__(self, embeddings, collection_name: str = None):
|
||||
self.collection_name = collection_name or config.QDRANT_COLLECTION
|
||||
self.qdrant = Qdrant(
|
||||
url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}",
|
||||
api_key=config.QDRANT_API_KEY,
|
||||
collection_name=self.collection_name,
|
||||
)
|
||||
|
||||
def add_documents(self, documents: Iterable[Document]) -> None:
|
||||
"""
|
||||
Add a collection of documents to the Qdrant store.
|
||||
"""
|
||||
texts = [doc.page_content for doc in documents]
|
||||
metadatas = [doc.metadata for doc in documents]
|
||||
ids = [doc.id for doc in documents if doc.id is not None]
|
||||
|
||||
# Embed the documents
|
||||
embeddings = self.embeddings.embed_documents(texts)
|
||||
|
||||
# Upsert into Qdrant
|
||||
self.client.upsert(
|
||||
embeddings=embeddings,
|
||||
documents=texts,
|
||||
metadatas=metadatas,
|
||||
ids=ids,
|
||||
)
|
||||
|
||||
def similarity_search(
|
||||
self,
|
||||
query: str,
|
||||
k: int = 5,
|
||||
filter: Optional[dict] = None,
|
||||
) -> List[Document]:
|
||||
def add_documents(self, documents: list[Document]):
|
||||
"""
|
||||
Perform a similarity search against the Qdrant store.
|
||||
Add a list of langchain Document objects to the Qdrant collection.
|
||||
"""
|
||||
query_embedding = self.embeddings.embed_query(query)
|
||||
results = self.client.search(
|
||||
query_embedding=query_embedding,
|
||||
limit=k,
|
||||
filter=filter,
|
||||
)
|
||||
# Convert results to Document objects
|
||||
return [
|
||||
Document(
|
||||
page_content=result["payload"]["text"],
|
||||
metadata=result["payload"],
|
||||
id=result["id"],
|
||||
)
|
||||
for result in results
|
||||
]
|
||||
self.qdrant.add_documents(documents)
|
||||
|
||||
# The following methods are required by the VectorStore interface
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
return self.embeddings.embed_query(query)
|
||||
|
||||
def embed_documents(self, documents: List[str]) -> List[List[float]]:
|
||||
return self.embeddings.embed_documents(documents)
|
||||
def get_retriever(self):
|
||||
"""
|
||||
Return a retriever that can be used with LangChain chains.
|
||||
"""
|
||||
return self.qdrant.as_retriever()
|
||||
Reference in New Issue
Block a user