Add more file extensions to accepted inputs
This commit is contained in:
parent
f63dc7148a
commit
3e7f4e42f7
1364
poetry.lock
generated
1364
poetry.lock
generated
File diff suppressed because it is too large
Load Diff
@ -1,27 +1,29 @@
|
||||
[project]
|
||||
name = "code-rag"
|
||||
version = "0.1.0"
|
||||
description = "Simple RAG implementation for use with neovim"
|
||||
authors = [{ name = "Alex Selimov", email = "alex@alexselimov.com" }]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9,<4.0"
|
||||
dependencies = [
|
||||
"langchain (>=0.3.21,<0.4.0)",
|
||||
"ollama (>=0.4.7,<0.5.0)",
|
||||
"langchain-community (>=0.3.20,<0.4.0)",
|
||||
"langchain-ollama (>=0.2.3,<0.3.0)",
|
||||
"chromadb (>=0.4.0,<0.6.0)",
|
||||
"unstructured (>=0.17.2,<0.18.0)",
|
||||
"langchain-chroma (>=0.1.0,<0.2.0)"
|
||||
description = "Simple RAG implementation"
|
||||
authors = [
|
||||
{ name = "Alex Selimov", email = "alex@alexselimov.com" }
|
||||
]
|
||||
readme = "README.md"
|
||||
requires-python = "^3.11"
|
||||
|
||||
[tool.poetry]
|
||||
packages = [{ include = "code_rag", from = "src" }]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = "^3.11"
|
||||
langchain = "*"
|
||||
langchain-community = "*"
|
||||
langchain-core = "*"
|
||||
langchain-chroma = "*"
|
||||
langchain-ollama = "*"
|
||||
chromadb = "*"
|
||||
ollama = "*"
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
pytest = "^8.3.5"
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core>=2.0.0,<3.0.0"]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
|
||||
@ -1,10 +1,12 @@
|
||||
import os
|
||||
import uuid
|
||||
import glob
|
||||
from langchain_community.document_loaders import DirectoryLoader, TextLoader
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_community.vectorstores import Chroma
|
||||
from langchain_chroma import Chroma
|
||||
from langchain.chains import RetrievalQA
|
||||
from langchain.prompts import PromptTemplate
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from code_rag.doc_tracker import DocumentTracker
|
||||
from code_rag.ollama_wrapper import OllamaWrapper
|
||||
@ -61,10 +63,32 @@ class RAG:
|
||||
|
||||
return all_chunks, file_chunk_map
|
||||
|
||||
def create_vector_db(self, extension=".txt", force_refresh=False):
|
||||
def create_vector_db(self, extensions=None, force_refresh=False):
|
||||
"""
|
||||
Create or update a vector database, with complete handling of changes
|
||||
Create or update a vector database, with complete handling of changes.
|
||||
|
||||
Args:
|
||||
extensions (list[str], optional): List of file extensions to include.
|
||||
If None, defaults to common programming languages.
|
||||
force_refresh (bool): Whether to force a complete refresh of the vector database.
|
||||
"""
|
||||
# Set default extensions for common programming languages if none provided
|
||||
if extensions is None:
|
||||
extensions = [
|
||||
# Python
|
||||
'.py', '.pyi', '.pyx',
|
||||
# C/C++
|
||||
'.c', '.cpp', '.cc', '.cxx', '.h', '.hpp', '.hxx',
|
||||
# Rust
|
||||
'.rs',
|
||||
# Documentation
|
||||
'.txt', '.md',
|
||||
# Build/Config
|
||||
'.toml', '.yaml', '.json'
|
||||
]
|
||||
elif isinstance(extensions, str):
|
||||
extensions = [extensions]
|
||||
|
||||
# Text splitter configuration
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1000, chunk_overlap=200
|
||||
@ -75,121 +99,79 @@ class RAG:
|
||||
embeddings = self.ollama.embeddings
|
||||
print("after embedding")
|
||||
|
||||
if force_refresh:
|
||||
print("Force refresh: Processing all documents")
|
||||
# Load all documents
|
||||
loader = DirectoryLoader(self.docs_dir, glob=f"**/*{extension}")
|
||||
all_documents = loader.load()
|
||||
|
||||
if not all_documents:
|
||||
print("No documents found to process")
|
||||
# Create an empty vector store
|
||||
vectorstore = Chroma(
|
||||
persist_directory=self.db_dir, embedding_function=embeddings
|
||||
)
|
||||
return vectorstore
|
||||
|
||||
# Add unique IDs to each document
|
||||
for doc in all_documents:
|
||||
doc.metadata["source"] = os.path.abspath(doc.metadata["source"])
|
||||
doc.metadata["source_id"] = doc.metadata["source"]
|
||||
|
||||
# Split documents
|
||||
chunks = text_splitter.split_documents(all_documents)
|
||||
|
||||
# Add chunk IDs and update tracker
|
||||
file_chunk_map = {}
|
||||
for chunk in chunks:
|
||||
chunk_id = str(uuid.uuid4())
|
||||
chunk.metadata["chunk_id"] = chunk_id
|
||||
|
||||
source = chunk.metadata["source"]
|
||||
if source not in file_chunk_map:
|
||||
file_chunk_map[source] = []
|
||||
|
||||
file_chunk_map[source].append(chunk_id)
|
||||
|
||||
# Update tracker with chunk mappings
|
||||
for file_path, chunk_ids in file_chunk_map.items():
|
||||
self.tracker.update_chunk_mappings(file_path, chunk_ids)
|
||||
|
||||
print(
|
||||
f"Processing {len(all_documents)} documents with {len(chunks)} chunks"
|
||||
# Load or create vector store
|
||||
if os.path.exists(self.db_dir) and not force_refresh:
|
||||
print("Loading existing vector store")
|
||||
vectorstore = Chroma(
|
||||
persist_directory=self.db_dir,
|
||||
embedding_function=embeddings
|
||||
)
|
||||
else:
|
||||
print("Creating new vector store")
|
||||
vectorstore = Chroma(
|
||||
persist_directory=self.db_dir,
|
||||
embedding_function=embeddings
|
||||
)
|
||||
|
||||
# Create new vector store
|
||||
vectorstore = Chroma.from_documents(
|
||||
documents=chunks, embedding=embeddings, persist_directory=self.db_dir
|
||||
)
|
||||
# Find all files that match the extensions
|
||||
all_files = set()
|
||||
for ext in extensions:
|
||||
files = glob.glob(os.path.join(self.docs_dir, f"**/*{ext}"), recursive=True)
|
||||
all_files.update(files)
|
||||
|
||||
if not all_files:
|
||||
print("No documents found to process")
|
||||
return vectorstore
|
||||
|
||||
# Get changes since last update
|
||||
changed_files = self.tracker.get_changed_files(self.docs_dir)
|
||||
|
||||
if not any(changed_files.values()):
|
||||
print("No document changes detected")
|
||||
# Load existing vector store if available
|
||||
if os.path.exists(self.db_dir):
|
||||
return Chroma(
|
||||
persist_directory=self.db_dir,
|
||||
embedding_function=self.ollama.embeddings,
|
||||
)
|
||||
else:
|
||||
print("No vector database exists. Creating from all documents...")
|
||||
return self.create_vector_db(force_refresh=True)
|
||||
|
||||
# Process changes
|
||||
print(
|
||||
f"Changes detected - New: {len(changed_files['new'])}, Modified: {len(changed_files['modified'])}, Deleted: {len(changed_files['deleted'])}"
|
||||
)
|
||||
|
||||
# Load existing vector store if it exists
|
||||
if os.path.exists(self.db_dir):
|
||||
vectorstore = Chroma(
|
||||
persist_directory=self.db_dir, embedding_function=self.ollama.embeddings
|
||||
)
|
||||
|
||||
# 1. Handle deleted documents
|
||||
if changed_files["deleted"]:
|
||||
chunks_to_delete = self.tracker.get_chunks_to_delete(
|
||||
changed_files["deleted"]
|
||||
)
|
||||
if chunks_to_delete:
|
||||
print(
|
||||
f"Removing {len(chunks_to_delete)} chunks from deleted documents"
|
||||
)
|
||||
# Delete the chunks from vector store
|
||||
vectorstore._collection.delete(
|
||||
where={"chunk_id": {"$in": chunks_to_delete}}
|
||||
# Process all files
|
||||
all_documents = []
|
||||
for file_path in all_files:
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
doc = Document(
|
||||
page_content=content,
|
||||
metadata={
|
||||
"source": os.path.abspath(file_path),
|
||||
"source_id": os.path.abspath(file_path)
|
||||
}
|
||||
)
|
||||
all_documents.append(doc)
|
||||
print(f"Successfully loaded {file_path}")
|
||||
except Exception as e:
|
||||
print(f"Error loading file {file_path}: {str(e)}")
|
||||
|
||||
# 2. Handle modified documents (delete old chunks first)
|
||||
chunks_to_delete_modified = self.tracker.get_chunks_for_modified_files(
|
||||
changed_files["modified"]
|
||||
if not all_documents:
|
||||
print("No documents could be loaded")
|
||||
return vectorstore
|
||||
|
||||
# Split documents
|
||||
chunks = text_splitter.split_documents(all_documents)
|
||||
|
||||
if force_refresh:
|
||||
# Create new vector store from scratch
|
||||
vectorstore = Chroma.from_documents(
|
||||
documents=chunks,
|
||||
embedding=embeddings,
|
||||
persist_directory=self.db_dir
|
||||
)
|
||||
if chunks_to_delete_modified:
|
||||
print(
|
||||
f"Removing {len(chunks_to_delete_modified)} chunks from modified documents"
|
||||
)
|
||||
vectorstore._collection.delete(
|
||||
where={"chunk_id": {"$in": chunks_to_delete_modified}}
|
||||
)
|
||||
|
||||
# 3. Process new and modified documents
|
||||
files_to_process = changed_files["new"] + changed_files["modified"]
|
||||
if files_to_process:
|
||||
chunks, _ = self.process_documents(files_to_process, text_splitter)
|
||||
print(f"Adding {len(chunks)} new chunks to the vector store")
|
||||
vectorstore.add_documents(chunks)
|
||||
else:
|
||||
# If no existing DB, create from all documents
|
||||
print("No existing vector database. Creating from all documents...")
|
||||
return self.create_vector_db(force_refresh=True)
|
||||
# Update existing vector store
|
||||
# First, get all existing document IDs
|
||||
existing_docs = vectorstore._collection.get()
|
||||
existing_sources = {m["source"] for m in existing_docs["metadatas"]}
|
||||
|
||||
# Persist changes
|
||||
vectorstore.persist()
|
||||
print(f"Vector database updated at {self.db_dir}")
|
||||
# Find new documents
|
||||
new_chunks = [
|
||||
chunk for chunk in chunks
|
||||
if chunk.metadata["source"] not in existing_sources
|
||||
]
|
||||
|
||||
if new_chunks:
|
||||
print(f"Adding {len(new_chunks)} new chunks to vector store")
|
||||
vectorstore.add_documents(new_chunks)
|
||||
|
||||
print(f"Vector store updated with {len(chunks)} total chunks")
|
||||
return vectorstore
|
||||
|
||||
def setup_rag(self):
|
||||
|
||||
@ -1,11 +1,10 @@
|
||||
import os
|
||||
import pytest
|
||||
import shutil
|
||||
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
import pytest
|
||||
from langchain_chroma import Chroma
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
|
||||
from .fixtures import *
|
||||
from .utility import *
|
||||
from code_rag.rag import RAG
|
||||
from code_rag.doc_tracker import DocumentTracker
|
||||
from code_rag.ollama_wrapper import OllamaWrapper
|
||||
@ -52,18 +51,40 @@ def test_process_documents(tracker_file, docs_dir, db_dir, sample_docs, rag_pipe
|
||||
def test_create_vector_db(docs_dir, db_dir, tracker_file, sample_docs):
|
||||
"""Test creating a vector database"""
|
||||
rag_pipeline = RAG(docs_dir, db_dir, tracker_file)
|
||||
# Create initial vector database
|
||||
|
||||
# Create files with different extensions
|
||||
files = {
|
||||
"test.py": "def hello():\n print('Hello, World!')",
|
||||
"main.cpp": "#include <iostream>\nint main() { std::cout << 'Hello'; return 0; }",
|
||||
"lib.rs": "fn main() { println!('Hello from Rust!'); }",
|
||||
"config.toml": "[package]\nname = 'test'",
|
||||
"doc.md": "# Documentation\nThis is a test file."
|
||||
}
|
||||
|
||||
for filename, content in files.items():
|
||||
filepath = os.path.join(docs_dir, filename)
|
||||
with open(filepath, "w") as f:
|
||||
f.write(content)
|
||||
|
||||
# Create vector database with default extensions (should include all file types)
|
||||
vectorstore = rag_pipeline.create_vector_db(force_refresh=True)
|
||||
|
||||
# Verify it was created
|
||||
assert os.path.exists(rag_pipeline.db_dir)
|
||||
assert vectorstore is not None
|
||||
# Check the database has content
|
||||
|
||||
# Check the database has content from all file types
|
||||
loaded_db = Chroma(
|
||||
persist_directory=rag_pipeline.db_dir,
|
||||
embedding_function=rag_pipeline.ollama.embeddings,
|
||||
persist_directory=rag_pipeline.db_dir, embedding_function=rag_pipeline.ollama.embeddings
|
||||
)
|
||||
# Should have content from all files
|
||||
assert loaded_db._collection.count() > 0
|
||||
|
||||
# Verify each file type is included
|
||||
docs = loaded_db._collection.get()
|
||||
sources = {os.path.basename(m["source"]) for m in docs["metadatas"]}
|
||||
for filename in files.keys():
|
||||
assert filename in sources, f"File {filename} not found in vector store"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
@ -72,45 +93,66 @@ def test_create_vector_db(docs_dir, db_dir, tracker_file, sample_docs):
|
||||
def test_update_vector_db_with_changes(docs_dir, db_dir, tracker_file, sample_docs):
|
||||
"""Test updating a vector database with document changes"""
|
||||
rag_pipeline = RAG(docs_dir, db_dir, tracker_file)
|
||||
# Create initial vector database
|
||||
rag_pipeline.create_vector_db(force_refresh=True)
|
||||
# Create initial vector database with only Python files
|
||||
vectorstore = rag_pipeline.create_vector_db(extensions=['.py'], force_refresh=True)
|
||||
|
||||
# Get initial count
|
||||
initial_db = Chroma(
|
||||
persist_directory=rag_pipeline.db_dir,
|
||||
embedding_function=rag_pipeline.ollama.embeddings,
|
||||
persist_directory=rag_pipeline.db_dir, embedding_function=rag_pipeline.ollama.embeddings
|
||||
)
|
||||
initial_count = initial_db._collection.count()
|
||||
|
||||
# Make changes to documents
|
||||
# Add a new document
|
||||
create_test_document(
|
||||
docs_dir, "newdoc.txt", "This is a brand new document for testing."
|
||||
)
|
||||
# Add files of different types
|
||||
new_files = {
|
||||
"newdoc.cpp": "#include <iostream>\nint main() { return 0; }",
|
||||
"lib.rs": "fn main() { println!('Hello'); }",
|
||||
"config.toml": "[package]\nname = 'test'"
|
||||
}
|
||||
|
||||
for filename, content in new_files.items():
|
||||
filepath = os.path.join(docs_dir, filename)
|
||||
with open(filepath, "w") as f:
|
||||
f.write(content)
|
||||
|
||||
# Update the vector database
|
||||
rag_pipeline.create_vector_db()
|
||||
# Update the vector database to include all supported extensions
|
||||
rag_pipeline.create_vector_db() # Use default extensions
|
||||
|
||||
# Check the database has been updated
|
||||
updated_db = Chroma(
|
||||
persist_directory=rag_pipeline.db_dir,
|
||||
embedding_function=rag_pipeline.ollama.embeddings,
|
||||
persist_directory=rag_pipeline.db_dir, embedding_function=rag_pipeline.ollama.embeddings
|
||||
)
|
||||
assert updated_db._collection.count() > initial_count
|
||||
|
||||
# Verify new files are included
|
||||
docs = updated_db._collection.get()
|
||||
sources = {os.path.basename(m["source"]) for m in docs["metadatas"]}
|
||||
for filename in new_files.keys():
|
||||
assert filename in sources, f"File {filename} not found in vector store"
|
||||
|
||||
|
||||
# Final integration test - full RAG pipeline
|
||||
@pytest.mark.skipif(
|
||||
not shutil.which("ollama"), reason="Ollama not installed or not in PATH"
|
||||
)
|
||||
def test_full_rag_pipeline(docs_dir, db_dir, tracker_file, sample_docs):
|
||||
"""Test the entire RAG pipeline from document processing to querying"""
|
||||
rag_pipeline = RAG(docs_dir, db_dir, tracker_file)
|
||||
# Create a specific document with known content
|
||||
test_content = "Python is a high-level programming language known for its readability and versatility."
|
||||
create_test_document(rag_pipeline.docs_dir, "python_info.txt", test_content)
|
||||
|
||||
# Create documents with mixed content types
|
||||
test_files = {
|
||||
"python_info.py": """# Python Information
|
||||
def describe_python():
|
||||
\"\"\"Python is a high-level programming language known for its readability and versatility.\"\"\"
|
||||
pass""",
|
||||
"readme.md": "# Python\nPython is a popular programming language used in web development, data science, and automation."
|
||||
}
|
||||
|
||||
for filename, content in test_files.items():
|
||||
filepath = os.path.join(rag_pipeline.docs_dir, filename)
|
||||
with open(filepath, "w") as f:
|
||||
f.write(content)
|
||||
|
||||
# Create vector database
|
||||
# Create vector database with all default extensions
|
||||
rag_pipeline.create_vector_db(force_refresh=True)
|
||||
|
||||
# Set up RAG
|
||||
@ -124,4 +166,3 @@ def test_full_rag_pipeline(docs_dir, db_dir, tracker_file, sample_docs):
|
||||
# This is a soft test since the exact response will depend on the LLM
|
||||
assert response.strip() != ""
|
||||
assert "programming" in response.lower() or "language" in response.lower()
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user