128 lines
4.3 KiB
Python
128 lines
4.3 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import tempfile
|
|
import shutil
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
# Ensure the environment variable is set before importing the module
|
|
os.environ["OPENAI_API_KEY"] = "test-key"
|
|
|
|
# Import the module after setting the environment variable
|
|
import src.index as index
|
|
|
|
# Helper to create a temporary text file
|
|
def create_temp_txt(folder: Path, name: str, content: str):
|
|
file_path = folder / name
|
|
file_path.write_text(content, encoding="utf-8")
|
|
return file_path
|
|
|
|
def test_split_text_basic():
|
|
text = "Sentence one. Sentence two. Sentence three."
|
|
chunks = index._split_text(text, max_chunk_size=50)
|
|
assert len(chunks) == 3
|
|
assert chunks[0] == "Sentence one."
|
|
assert chunks[1] == "Sentence two."
|
|
assert chunks[2] == "Sentence three."
|
|
|
|
def test_split_text_long_sentence():
|
|
long_sentence = "A" * 200
|
|
text = f"{long_sentence}. Another short sentence."
|
|
chunks = index._split_text(text, max_chunk_size=100)
|
|
# The long sentence should be split into two chunks
|
|
assert len(chunks) == 2
|
|
assert chunks[0].startswith("A" * 100)
|
|
assert chunks[1].startswith("A" * 100)
|
|
|
|
def test_ingest_folder(monkeypatch):
|
|
# Create a temporary directory with a single .txt file
|
|
temp_dir = Path(tempfile.mkdtemp())
|
|
try:
|
|
content = "Hello world. This is a test."
|
|
create_temp_txt(temp_dir, "test.txt", content)
|
|
|
|
# Mock the add_documents method to capture its arguments
|
|
captured = {}
|
|
def mock_add_documents(self, documents, ids):
|
|
captured["documents"] = documents
|
|
captured["ids"] = ids
|
|
|
|
monkeypatch.setattr(index.ChromaDBWrapper, "add_documents", mock_add_documents)
|
|
|
|
# Instantiate the agent with a dummy db wrapper
|
|
dummy_db = index.ChromaDBWrapper()
|
|
agent = index.RAGAgent(dummy_db)
|
|
|
|
# Run ingestion
|
|
agent.ingest_folder(str(temp_dir))
|
|
|
|
# Verify that documents were split and added
|
|
assert "documents" in captured
|
|
assert "ids" in captured
|
|
assert len(captured["documents"]) == 2 # two sentences
|
|
assert captured["documents"][0] == "Hello world."
|
|
assert captured["documents"][1] == "This is a test."
|
|
assert len(captured["ids"]) == 2
|
|
assert captured["ids"][0].startswith("test_")
|
|
finally:
|
|
shutil.rmtree(temp_dir)
|
|
|
|
def test_answer_query_local(monkeypatch):
|
|
# Dummy database that returns a relevant document
|
|
class DummyDB:
|
|
def query(self, query_text, k=5):
|
|
return [("Relevant context about Python.", 0.8)]
|
|
|
|
dummy_db = DummyDB()
|
|
agent = index.RAGAgent(dummy_db)
|
|
|
|
# Mock the OpenAI ChatCompletion to return a predictable answer
|
|
mock_response = {
|
|
"choices": [
|
|
{"message": {"content": "Python is a programming language."}}
|
|
]
|
|
}
|
|
monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)
|
|
|
|
answer = agent.answer_query("What is Python?")
|
|
assert answer == "Python is a programming language."
|
|
|
|
def test_answer_query_fallback(monkeypatch):
|
|
# Dummy database that returns no relevant documents
|
|
class DummyDB:
|
|
def query(self, query_text, k=5):
|
|
return []
|
|
|
|
dummy_db = DummyDB()
|
|
agent = index.RAGAgent(dummy_db)
|
|
|
|
# Mock web_search to return snippets
|
|
monkeypatch.setattr(index, "web_search", lambda query, max_results=3: ["Snippet about AI.", "Another snippet."])
|
|
|
|
# Mock the OpenAI ChatCompletion to return a predictable answer
|
|
mock_response = {
|
|
"choices": [
|
|
{"message": {"content": "AI stands for Artificial Intelligence."}}
|
|
]
|
|
}
|
|
monkeypatch.setattr(index.openai.ChatCompletion, "create", lambda **kwargs: mock_response)
|
|
|
|
answer = agent.answer_query("What does AI stand for?")
|
|
assert answer == "AI stands for Artificial Intelligence."
|
|
|
|
def test_web_search_mock(monkeypatch):
|
|
# Mock ddg to return predefined results
|
|
mock_results = [
|
|
{"body": "First snippet content."},
|
|
{"body": "Second snippet content."},
|
|
]
|
|
monkeypatch.setattr(index.ddg, "__call__", lambda query, max_results=3: mock_results)
|
|
|
|
snippets = index.web_search("test query")
|
|
assert snippets == ["First snippet content.", "Second snippet content."]
|
|
|
|
if __name__ == "__main__":
|
|
# Run tests manually if executed as a script
|
|
import pytest
|
|
sys.exit(pytest.main([__file__])) |