Merge pull request #15 from ggozad/feat/settings
Store settings to db & check for incompatibilities.
This commit is contained in:
commit
96945e2c98
13 changed files with 614 additions and 423 deletions
|
|
@ -2,6 +2,9 @@
|
||||||
|
|
||||||
Configuration is done through the use of environment variables.
|
Configuration is done through the use of environment variables.
|
||||||
|
|
||||||
|
!!! note
|
||||||
|
If you create a db with certain settings and later change them, `haiku.rag` will detect incompatibilities (for example, if you change embedding provider) and will exit. You can **rebuild** the database to apply the new settings, see [Rebuild Database](./cli.md#rebuild-database).
|
||||||
|
|
||||||
## File Monitoring
|
## File Monitoring
|
||||||
|
|
||||||
Set directories to monitor for automatic indexing:
|
Set directories to monitor for automatic indexing:
|
||||||
|
|
|
||||||
|
|
@ -56,7 +56,7 @@ dev = [
|
||||||
"mkdocs>=1.6.1",
|
"mkdocs>=1.6.1",
|
||||||
"mkdocs-material>=9.6.14",
|
"mkdocs-material>=9.6.14",
|
||||||
"pre-commit>=4.2.0",
|
"pre-commit>=4.2.0",
|
||||||
"pyright>=1.1.402",
|
"pyright>=1.1.403",
|
||||||
"pytest>=8.4.0",
|
"pytest>=8.4.0",
|
||||||
"pytest-asyncio>=1.0.0",
|
"pytest-asyncio>=1.0.0",
|
||||||
"pytest-cov>=6.2.1",
|
"pytest-cov>=6.2.1",
|
||||||
|
|
|
||||||
|
|
@ -74,7 +74,7 @@ class HaikuRAGApp:
|
||||||
self.console.print(f"[red]Error: {e}[/red]")
|
self.console.print(f"[red]Error: {e}[/red]")
|
||||||
|
|
||||||
async def rebuild(self):
|
async def rebuild(self):
|
||||||
async with HaikuRAG(db_path=self.db_path) as client:
|
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
|
||||||
try:
|
try:
|
||||||
documents = await client.list_documents()
|
documents = await client.list_documents()
|
||||||
total_docs = len(documents)
|
total_docs = len(documents)
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,7 @@ import typer
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
|
|
||||||
from haiku.rag.app import HaikuRAGApp
|
from haiku.rag.app import HaikuRAGApp
|
||||||
from haiku.rag.utils import get_default_data_dir
|
from haiku.rag.utils import get_default_data_dir, is_up_to_date
|
||||||
|
|
||||||
cli = typer.Typer(
|
cli = typer.Typer(
|
||||||
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
|
context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
|
||||||
|
|
@ -15,6 +15,23 @@ console = Console()
|
||||||
event_loop = asyncio.get_event_loop()
|
event_loop = asyncio.get_event_loop()
|
||||||
|
|
||||||
|
|
||||||
|
async def check_version():
|
||||||
|
"""Check if haiku.rag is up to date and show warning if not."""
|
||||||
|
up_to_date, current_version, latest_version = await is_up_to_date()
|
||||||
|
if not up_to_date:
|
||||||
|
console.print(
|
||||||
|
f"[yellow]Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}[/yellow]"
|
||||||
|
)
|
||||||
|
console.print("[yellow]Please update.[/yellow]")
|
||||||
|
|
||||||
|
|
||||||
|
@cli.callback()
|
||||||
|
def main():
|
||||||
|
"""haiku.rag CLI - SQLite-based RAG system"""
|
||||||
|
# Run version check before any command
|
||||||
|
event_loop.run_until_complete(check_version())
|
||||||
|
|
||||||
|
|
||||||
@cli.command("list", help="List all stored documents")
|
@cli.command("list", help="List all stored documents")
|
||||||
def list_documents(
|
def list_documents(
|
||||||
db: Path = typer.Option(
|
db: Path = typer.Option(
|
||||||
|
|
|
||||||
|
|
@ -24,12 +24,13 @@ class HaikuRAG:
|
||||||
self,
|
self,
|
||||||
db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
|
db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
|
||||||
/ "haiku.rag.sqlite",
|
/ "haiku.rag.sqlite",
|
||||||
|
skip_validation: bool = False,
|
||||||
):
|
):
|
||||||
"""Initialize the RAG client with a database path."""
|
"""Initialize the RAG client with a database path."""
|
||||||
if isinstance(db_path, Path):
|
if isinstance(db_path, Path):
|
||||||
if not db_path.parent.exists():
|
if not db_path.parent.exists():
|
||||||
Path.mkdir(db_path.parent, parents=True)
|
Path.mkdir(db_path.parent, parents=True)
|
||||||
self.store = Store(db_path)
|
self.store = Store(db_path, skip_validation=skip_validation)
|
||||||
self.document_repository = DocumentRepository(self.store)
|
self.document_repository = DocumentRepository(self.store)
|
||||||
self.chunk_repository = ChunkRepository(self.store)
|
self.chunk_repository = ChunkRepository(self.store)
|
||||||
|
|
||||||
|
|
@ -165,29 +166,26 @@ class HaikuRAG:
|
||||||
|
|
||||||
# Create a temporary file with the appropriate extension
|
# Create a temporary file with the appropriate extension
|
||||||
with tempfile.NamedTemporaryFile(
|
with tempfile.NamedTemporaryFile(
|
||||||
mode="wb", suffix=file_extension, delete=False
|
mode="wb", suffix=file_extension
|
||||||
) as temp_file:
|
) as temp_file:
|
||||||
temp_file.write(response.content)
|
temp_file.write(response.content)
|
||||||
|
temp_file.flush() # Ensure content is written to disk
|
||||||
temp_path = Path(temp_file.name)
|
temp_path = Path(temp_file.name)
|
||||||
|
|
||||||
try:
|
|
||||||
# Parse the content using FileReader
|
# Parse the content using FileReader
|
||||||
content = FileReader.parse_file(temp_path)
|
content = FileReader.parse_file(temp_path)
|
||||||
|
|
||||||
# Merge metadata with contentType and md5
|
# Merge metadata with contentType and md5
|
||||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||||
|
|
||||||
if existing_doc:
|
if existing_doc:
|
||||||
existing_doc.content = content
|
existing_doc.content = content
|
||||||
existing_doc.metadata = metadata
|
existing_doc.metadata = metadata
|
||||||
return await self.update_document(existing_doc)
|
return await self.update_document(existing_doc)
|
||||||
else:
|
else:
|
||||||
return await self.create_document(
|
return await self.create_document(
|
||||||
content=content, uri=url, metadata=metadata
|
content=content, uri=url, metadata=metadata
|
||||||
)
|
)
|
||||||
finally:
|
|
||||||
# Clean up temporary file
|
|
||||||
temp_path.unlink(missing_ok=True)
|
|
||||||
|
|
||||||
def _get_extension_from_content_type_or_url(
|
def _get_extension_from_content_type_or_url(
|
||||||
self, url: str, content_type: str
|
self, url: str, content_type: str
|
||||||
|
|
@ -277,12 +275,16 @@ class HaikuRAG:
|
||||||
Yields:
|
Yields:
|
||||||
int: The ID of the document currently being processed
|
int: The ID of the document currently being processed
|
||||||
"""
|
"""
|
||||||
documents = await self.list_documents()
|
|
||||||
|
|
||||||
if not documents:
|
|
||||||
return
|
|
||||||
|
|
||||||
await self.chunk_repository.delete_all()
|
await self.chunk_repository.delete_all()
|
||||||
|
self.store.recreate_embeddings_table()
|
||||||
|
|
||||||
|
# Update settings to current config
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
settings_repo = SettingsRepository(self.store)
|
||||||
|
settings_repo.save()
|
||||||
|
|
||||||
|
documents = await self.list_documents()
|
||||||
|
|
||||||
for doc in documents:
|
for doc in documents:
|
||||||
if doc.id is not None:
|
if doc.id is not None:
|
||||||
|
|
|
||||||
|
|
@ -9,10 +9,19 @@ from haiku.rag.embeddings import get_embedder
|
||||||
|
|
||||||
|
|
||||||
class Store:
|
class Store:
|
||||||
def __init__(self, db_path: Path | Literal[":memory:"]):
|
def __init__(
|
||||||
|
self, db_path: Path | Literal[":memory:"], skip_validation: bool = False
|
||||||
|
):
|
||||||
self.db_path: Path | Literal[":memory:"] = db_path
|
self.db_path: Path | Literal[":memory:"] = db_path
|
||||||
self._connection = self.create_db()
|
self._connection = self.create_db()
|
||||||
|
|
||||||
|
# Validate config compatibility after connection is established
|
||||||
|
if not skip_validation:
|
||||||
|
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||||
|
|
||||||
|
settings_repo = SettingsRepository(self)
|
||||||
|
settings_repo.validate_config_compatibility()
|
||||||
|
|
||||||
def create_db(self) -> sqlite3.Connection:
|
def create_db(self) -> sqlite3.Connection:
|
||||||
"""Create the database and tables with sqlite-vec support for embeddings."""
|
"""Create the database and tables with sqlite-vec support for embeddings."""
|
||||||
db = sqlite3.connect(self.db_path)
|
db = sqlite3.connect(self.db_path)
|
||||||
|
|
@ -60,14 +69,50 @@ class Store:
|
||||||
)
|
)
|
||||||
""")
|
""")
|
||||||
|
|
||||||
|
# Create settings table for storing current configuration
|
||||||
|
db.execute("""
|
||||||
|
CREATE TABLE IF NOT EXISTS settings (
|
||||||
|
id INTEGER PRIMARY KEY DEFAULT 1,
|
||||||
|
settings TEXT NOT NULL DEFAULT '{}'
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
|
||||||
# Create indexes for better performance
|
# Create indexes for better performance
|
||||||
db.execute(
|
db.execute(
|
||||||
"CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
|
"CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Save current settings to the new database
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
settings_json = Config.model_dump_json()
|
||||||
|
db.execute(
|
||||||
|
"INSERT OR IGNORE INTO settings (id, settings) VALUES (1, ?)",
|
||||||
|
(settings_json,),
|
||||||
|
)
|
||||||
|
|
||||||
db.commit()
|
db.commit()
|
||||||
return db
|
return db
|
||||||
|
|
||||||
|
def recreate_embeddings_table(self) -> None:
|
||||||
|
"""Recreate the embeddings table with current vector dimensions."""
|
||||||
|
if self._connection is None:
|
||||||
|
raise ValueError("Store connection is not available")
|
||||||
|
|
||||||
|
# Drop existing embeddings table
|
||||||
|
self._connection.execute("DROP TABLE IF EXISTS chunk_embeddings")
|
||||||
|
|
||||||
|
# Recreate with current dimensions
|
||||||
|
embedder = get_embedder()
|
||||||
|
self._connection.execute(f"""
|
||||||
|
CREATE VIRTUAL TABLE chunk_embeddings USING vec0(
|
||||||
|
chunk_id INTEGER PRIMARY KEY,
|
||||||
|
embedding FLOAT[{embedder._vector_dim}]
|
||||||
|
)
|
||||||
|
""")
|
||||||
|
|
||||||
|
self._connection.commit()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def serialize_embedding(embedding: list[float]) -> bytes:
|
def serialize_embedding(embedding: list[float]) -> bytes:
|
||||||
"""Serialize a list of floats to bytes for sqlite-vec storage."""
|
"""Serialize a list of floats to bytes for sqlite-vec storage."""
|
||||||
|
|
|
||||||
78
src/haiku/rag/store/repositories/settings.py
Normal file
78
src/haiku/rag/store/repositories/settings.py
Normal file
|
|
@ -0,0 +1,78 @@
|
||||||
|
import json
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
|
||||||
|
|
||||||
|
class ConfigMismatchError(Exception):
|
||||||
|
"""Raised when current config doesn't match stored settings."""
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SettingsRepository:
|
||||||
|
def __init__(self, store: Store):
|
||||||
|
self.store = store
|
||||||
|
|
||||||
|
def get(self) -> dict[str, Any]:
|
||||||
|
"""Get all settings from the database."""
|
||||||
|
if self.store._connection is None:
|
||||||
|
raise ValueError("Store connection is not available")
|
||||||
|
|
||||||
|
cursor = self.store._connection.execute("SELECT settings FROM settings LIMIT 1")
|
||||||
|
row = cursor.fetchone()
|
||||||
|
if row:
|
||||||
|
return json.loads(row[0])
|
||||||
|
return {}
|
||||||
|
|
||||||
|
def save(self) -> None:
|
||||||
|
"""Sync settings from the current AppConfig to database."""
|
||||||
|
if self.store._connection is None:
|
||||||
|
raise ValueError("Store connection is not available")
|
||||||
|
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
settings_json = Config.model_dump_json()
|
||||||
|
|
||||||
|
self.store._connection.execute(
|
||||||
|
"INSERT INTO settings (id, settings) VALUES (1, ?) ON CONFLICT(id) DO UPDATE SET settings = excluded.settings",
|
||||||
|
(settings_json,),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.store._connection.commit()
|
||||||
|
|
||||||
|
def validate_config_compatibility(self) -> None:
|
||||||
|
"""Check if current config is compatible with stored settings.
|
||||||
|
|
||||||
|
Raises ConfigMismatchError if there are incompatible differences.
|
||||||
|
If no settings exist, saves current config.
|
||||||
|
"""
|
||||||
|
db_settings = self.get()
|
||||||
|
if not db_settings:
|
||||||
|
# No settings in DB, save current config
|
||||||
|
self.save()
|
||||||
|
return
|
||||||
|
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
|
||||||
|
current_config = Config.model_dump(mode="json")
|
||||||
|
|
||||||
|
# Critical settings that must match
|
||||||
|
critical_settings = [
|
||||||
|
"EMBEDDINGS_PROVIDER",
|
||||||
|
"EMBEDDINGS_MODEL",
|
||||||
|
"EMBEDDINGS_VECTOR_DIM",
|
||||||
|
"CHUNK_SIZE",
|
||||||
|
"CHUNK_OVERLAP",
|
||||||
|
]
|
||||||
|
|
||||||
|
errors = []
|
||||||
|
for setting in critical_settings:
|
||||||
|
if db_settings.get(setting) != current_config.get(setting):
|
||||||
|
errors.append(
|
||||||
|
f"{setting}: current={current_config.get(setting)}, stored={db_settings.get(setting)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
if errors:
|
||||||
|
error_msg = f"Config mismatch detected: {'; '.join(errors)}. Consider rebuilding the database with the current configuration."
|
||||||
|
raise ConfigMismatchError(error_msg)
|
||||||
|
|
@ -1,6 +1,10 @@
|
||||||
import sys
|
import sys
|
||||||
|
from importlib import metadata
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from packaging.version import Version, parse
|
||||||
|
|
||||||
|
|
||||||
def get_default_data_dir() -> Path:
|
def get_default_data_dir() -> Path:
|
||||||
"""
|
"""
|
||||||
|
|
@ -23,3 +27,23 @@ def get_default_data_dir() -> Path:
|
||||||
|
|
||||||
data_path = system_paths[sys.platform]
|
data_path = system_paths[sys.platform]
|
||||||
return data_path
|
return data_path
|
||||||
|
|
||||||
|
|
||||||
|
async def is_up_to_date() -> tuple[bool, Version, Version]:
|
||||||
|
"""
|
||||||
|
Checks whether haiku.rag is current.
|
||||||
|
|
||||||
|
:return: A tuple containing a boolean indicating whether haiku.rag is current, the running version and the latest version
|
||||||
|
:rtype: tuple[bool, Version, Version]
|
||||||
|
"""
|
||||||
|
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
running_version = parse(metadata.version("haiku.rag"))
|
||||||
|
try:
|
||||||
|
response = await client.get("https://pypi.org/pypi/haiku.rag/json")
|
||||||
|
data = response.json()
|
||||||
|
pypi_version = parse(data["info"]["version"])
|
||||||
|
except Exception:
|
||||||
|
# If no network connection, do not raise alarms.
|
||||||
|
pypi_version = running_version
|
||||||
|
return running_version >= pypi_version, running_version, pypi_version
|
||||||
|
|
|
||||||
|
|
@ -12,300 +12,270 @@ from haiku.rag.client import HaikuRAG
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_document_crud(qa_corpus: Dataset):
|
async def test_client_document_crud(qa_corpus: Dataset):
|
||||||
"""Test HaikuRAG CRUD operations for documents."""
|
"""Test HaikuRAG CRUD operations for documents."""
|
||||||
# Create client with in-memory database
|
async with HaikuRAG(":memory:") as client:
|
||||||
client = HaikuRAG(":memory:")
|
# Get test data
|
||||||
|
first_doc = qa_corpus[0]
|
||||||
|
document_text = first_doc["document_extracted"]
|
||||||
|
test_uri = "file:///path/to/test.txt"
|
||||||
|
test_metadata = {"source": "test", "topic": "testing"}
|
||||||
|
|
||||||
# Get test data
|
# Test create_document
|
||||||
first_doc = qa_corpus[0]
|
created_doc = await client.create_document(
|
||||||
document_text = first_doc["document_extracted"]
|
content=document_text, uri=test_uri, metadata=test_metadata
|
||||||
test_uri = "file:///path/to/test.txt"
|
)
|
||||||
test_metadata = {"source": "test", "topic": "testing"}
|
|
||||||
|
|
||||||
# Test create_document
|
assert created_doc.id is not None
|
||||||
created_doc = await client.create_document(
|
assert created_doc.content == document_text
|
||||||
content=document_text, uri=test_uri, metadata=test_metadata
|
assert created_doc.uri == test_uri
|
||||||
)
|
assert created_doc.metadata == test_metadata
|
||||||
|
|
||||||
assert created_doc.id is not None
|
# Test get_document_by_id
|
||||||
assert created_doc.content == document_text
|
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
||||||
assert created_doc.uri == test_uri
|
assert retrieved_doc is not None
|
||||||
assert created_doc.metadata == test_metadata
|
assert retrieved_doc.id == created_doc.id
|
||||||
|
assert retrieved_doc.content == document_text
|
||||||
|
assert retrieved_doc.uri == test_uri
|
||||||
|
|
||||||
# Test get_document_by_id
|
# Test get_document_by_uri
|
||||||
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
retrieved_by_uri = await client.get_document_by_uri(test_uri)
|
||||||
assert retrieved_doc is not None
|
assert retrieved_by_uri is not None
|
||||||
assert retrieved_doc.id == created_doc.id
|
assert retrieved_by_uri.id == created_doc.id
|
||||||
assert retrieved_doc.content == document_text
|
assert retrieved_by_uri.content == document_text
|
||||||
assert retrieved_doc.uri == test_uri
|
|
||||||
|
|
||||||
# Test get_document_by_uri
|
# Test get_document_by_uri with non-existent URI
|
||||||
retrieved_by_uri = await client.get_document_by_uri(test_uri)
|
non_existent = await client.get_document_by_uri("file:///non/existent.txt")
|
||||||
assert retrieved_by_uri is not None
|
assert non_existent is None
|
||||||
assert retrieved_by_uri.id == created_doc.id
|
|
||||||
assert retrieved_by_uri.content == document_text
|
|
||||||
|
|
||||||
# Test get_document_by_uri with non-existent URI
|
# Test update_document
|
||||||
non_existent = await client.get_document_by_uri("file:///non/existent.txt")
|
retrieved_doc.content = "Updated content"
|
||||||
assert non_existent is None
|
retrieved_doc.uri = "file:///updated/path.txt"
|
||||||
|
updated_doc = await client.update_document(retrieved_doc)
|
||||||
|
assert updated_doc.content == "Updated content"
|
||||||
|
assert updated_doc.uri == "file:///updated/path.txt"
|
||||||
|
|
||||||
# Test update_document
|
# Test list_documents
|
||||||
retrieved_doc.content = "Updated content"
|
all_docs = await client.list_documents()
|
||||||
retrieved_doc.uri = "file:///updated/path.txt"
|
assert len(all_docs) == 1
|
||||||
updated_doc = await client.update_document(retrieved_doc)
|
assert all_docs[0].id == created_doc.id
|
||||||
assert updated_doc.content == "Updated content"
|
|
||||||
assert updated_doc.uri == "file:///updated/path.txt"
|
|
||||||
|
|
||||||
# Test list_documents
|
# Test list_documents with pagination
|
||||||
all_docs = await client.list_documents()
|
limited_docs = await client.list_documents(limit=10, offset=0)
|
||||||
assert len(all_docs) == 1
|
assert len(limited_docs) == 1
|
||||||
assert all_docs[0].id == created_doc.id
|
|
||||||
|
|
||||||
# Test list_documents with pagination
|
# Test delete_document
|
||||||
limited_docs = await client.list_documents(limit=10, offset=0)
|
deleted = await client.delete_document(created_doc.id)
|
||||||
assert len(limited_docs) == 1
|
assert deleted is True
|
||||||
|
|
||||||
# Test delete_document
|
# Verify document is gone
|
||||||
deleted = await client.delete_document(created_doc.id)
|
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
||||||
assert deleted is True
|
assert retrieved_doc is None
|
||||||
|
|
||||||
# Verify document is gone
|
# Test delete non-existent document
|
||||||
retrieved_doc = await client.get_document_by_id(created_doc.id)
|
deleted_again = await client.delete_document(created_doc.id)
|
||||||
assert retrieved_doc is None
|
assert deleted_again is False
|
||||||
|
|
||||||
# Test delete non-existent document
|
|
||||||
deleted_again = await client.delete_document(created_doc.id)
|
|
||||||
assert deleted_again is False
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_source():
|
async def test_client_create_document_from_source():
|
||||||
"""Test creating a document from a file source."""
|
"""Test creating a document from a file source."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
|
test_content = "This is test content from a file."
|
||||||
|
temp_path = Path(temp_dir) / "test.txt"
|
||||||
|
temp_path.write_text(test_content)
|
||||||
|
|
||||||
# Create a temporary text file
|
# Test create_document_from_source with Path
|
||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
doc = await client.create_document_from_source(source=temp_path)
|
||||||
test_content = "This is test content from a file."
|
|
||||||
f.write(test_content)
|
|
||||||
temp_path = Path(f.name)
|
|
||||||
|
|
||||||
try:
|
assert doc.id is not None
|
||||||
# Test create_document_from_source with Path
|
assert doc.content == test_content
|
||||||
doc = await client.create_document_from_source(
|
assert doc.uri == temp_path.as_uri()
|
||||||
source=temp_path, metadata={"source_type": "file"}
|
assert "contentType" in doc.metadata
|
||||||
)
|
assert "md5" in doc.metadata
|
||||||
|
assert doc.metadata["contentType"] == "text/plain"
|
||||||
|
|
||||||
assert doc.id is not None
|
# Test create_document_from_source with string path
|
||||||
assert doc.content == test_content
|
doc2 = await client.create_document_from_source(source=str(temp_path))
|
||||||
assert doc.uri == temp_path.as_uri()
|
|
||||||
assert doc.metadata["source_type"] == "file"
|
|
||||||
assert "contentType" in doc.metadata
|
|
||||||
assert "md5" in doc.metadata
|
|
||||||
assert doc.metadata["contentType"] == "text/plain"
|
|
||||||
|
|
||||||
# Test create_document_from_source with string path
|
assert doc2.id is not None
|
||||||
doc2 = await client.create_document_from_source(source=str(temp_path))
|
assert doc2.content == test_content
|
||||||
|
assert doc2.uri == temp_path.as_uri()
|
||||||
assert doc2.id is not None
|
assert "contentType" in doc2.metadata
|
||||||
assert doc2.content == test_content
|
assert "md5" in doc2.metadata
|
||||||
assert doc2.uri == temp_path.as_uri()
|
|
||||||
assert "contentType" in doc2.metadata
|
|
||||||
assert "md5" in doc2.metadata
|
|
||||||
|
|
||||||
finally:
|
|
||||||
# Clean up
|
|
||||||
temp_path.unlink()
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_source_unsupported():
|
async def test_client_create_document_from_source_unsupported():
|
||||||
"""Test creating a document from an unsupported file type."""
|
"""Test creating a document from an unsupported file type."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Create a temporary file with unsupported extension
|
||||||
|
with tempfile.NamedTemporaryFile(
|
||||||
|
mode="w", suffix=".unsupported", delete=False
|
||||||
|
) as f:
|
||||||
|
f.write("content")
|
||||||
|
temp_path = Path(f.name)
|
||||||
|
|
||||||
# Create a temporary file with unsupported extension
|
# Should raise ValueError for unsupported extension
|
||||||
with tempfile.NamedTemporaryFile(
|
with pytest.raises(ValueError, match="Unsupported file extension"):
|
||||||
mode="w", suffix=".unsupported", delete=False
|
await client.create_document_from_source(temp_path)
|
||||||
) as f:
|
|
||||||
f.write("content")
|
|
||||||
temp_path = Path(f.name)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# Should raise ValueError for unsupported extension
|
|
||||||
with pytest.raises(ValueError, match="Unsupported file extension"):
|
|
||||||
await client.create_document_from_source(temp_path)
|
|
||||||
|
|
||||||
finally:
|
|
||||||
temp_path.unlink()
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_source_nonexistent():
|
async def test_client_create_document_from_source_nonexistent():
|
||||||
"""Test creating a document from a non-existent file."""
|
"""Test creating a document from a non-existent file."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
non_existent_path = Path("/non/existent/file.txt")
|
||||||
|
|
||||||
non_existent_path = Path("/non/existent/file.txt")
|
# Should raise ValueError when file doesn't exist
|
||||||
|
with pytest.raises(ValueError, match="File does not exist"):
|
||||||
# Should raise ValueError when file doesn't exist
|
await client.create_document_from_source(non_existent_path)
|
||||||
with pytest.raises(ValueError, match="File does not exist"):
|
|
||||||
await client.create_document_from_source(non_existent_path)
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_url():
|
async def test_client_create_document_from_url():
|
||||||
"""Test creating a document from a URL."""
|
"""Test creating a document from a URL."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Mock the HTTP response
|
||||||
|
mock_response = AsyncMock()
|
||||||
|
mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
|
||||||
|
mock_response.headers = {"content-type": "text/html"}
|
||||||
|
mock_response.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
# Mock the HTTP response
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
||||||
mock_response = AsyncMock()
|
doc = await client.create_document_from_source(
|
||||||
mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
|
source="https://example.com/test.html", metadata={"source_type": "web"}
|
||||||
mock_response.headers = {"content-type": "text/html"}
|
)
|
||||||
mock_response.raise_for_status = AsyncMock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
assert doc.id is not None
|
||||||
doc = await client.create_document_from_source(
|
assert "Test Page" in doc.content
|
||||||
source="https://example.com/test.html", metadata={"source_type": "web"}
|
assert "test content" in doc.content
|
||||||
)
|
assert doc.uri == "https://example.com/test.html"
|
||||||
|
assert doc.metadata["source_type"] == "web"
|
||||||
assert doc.id is not None
|
assert "contentType" in doc.metadata
|
||||||
assert "Test Page" in doc.content
|
assert "md5" in doc.metadata
|
||||||
assert "test content" in doc.content
|
assert doc.metadata["contentType"] == "text/html"
|
||||||
assert doc.uri == "https://example.com/test.html"
|
|
||||||
assert doc.metadata["source_type"] == "web"
|
|
||||||
assert "contentType" in doc.metadata
|
|
||||||
assert "md5" in doc.metadata
|
|
||||||
assert doc.metadata["contentType"] == "text/html"
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_url_with_different_content_types():
|
async def test_client_create_document_from_url_with_different_content_types():
|
||||||
"""Test creating documents from URLs with different content types."""
|
"""Test creating documents from URLs with different content types."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Test JSON content
|
||||||
# Test JSON content
|
mock_json_response = AsyncMock()
|
||||||
mock_json_response = AsyncMock()
|
mock_json_response.content = (
|
||||||
mock_json_response.content = (
|
b'{"title": "Test JSON", "content": "This is JSON content"}'
|
||||||
b'{"title": "Test JSON", "content": "This is JSON content"}'
|
|
||||||
)
|
|
||||||
mock_json_response.headers = {"content-type": "application/json"}
|
|
||||||
mock_json_response.raise_for_status = AsyncMock()
|
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_json_response):
|
|
||||||
doc = await client.create_document_from_source(
|
|
||||||
"https://api.example.com/data.json"
|
|
||||||
)
|
)
|
||||||
|
mock_json_response.headers = {"content-type": "application/json"}
|
||||||
|
mock_json_response.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
assert doc.id is not None
|
with patch("httpx.AsyncClient.get", return_value=mock_json_response):
|
||||||
assert "Test JSON" in doc.content
|
doc = await client.create_document_from_source(
|
||||||
assert doc.uri == "https://api.example.com/data.json"
|
"https://api.example.com/data.json"
|
||||||
assert "contentType" in doc.metadata
|
)
|
||||||
assert "md5" in doc.metadata
|
|
||||||
assert doc.metadata["contentType"] == "application/json"
|
|
||||||
|
|
||||||
# Test plain text content
|
assert doc.id is not None
|
||||||
mock_text_response = AsyncMock()
|
assert "Test JSON" in doc.content
|
||||||
mock_text_response.content = b"This is plain text content from a URL."
|
assert doc.uri == "https://api.example.com/data.json"
|
||||||
mock_text_response.headers = {"content-type": "text/plain"}
|
assert "contentType" in doc.metadata
|
||||||
mock_text_response.raise_for_status = AsyncMock()
|
assert "md5" in doc.metadata
|
||||||
|
assert doc.metadata["contentType"] == "application/json"
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_text_response):
|
# Test plain text content
|
||||||
doc = await client.create_document_from_source("https://example.com/readme.txt")
|
mock_text_response = AsyncMock()
|
||||||
|
mock_text_response.content = b"This is plain text content from a URL."
|
||||||
|
mock_text_response.headers = {"content-type": "text/plain"}
|
||||||
|
mock_text_response.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
assert doc.id is not None
|
with patch("httpx.AsyncClient.get", return_value=mock_text_response):
|
||||||
assert doc.content == "This is plain text content from a URL."
|
doc = await client.create_document_from_source(
|
||||||
assert doc.uri == "https://example.com/readme.txt"
|
"https://example.com/readme.txt"
|
||||||
assert "contentType" in doc.metadata
|
)
|
||||||
assert "md5" in doc.metadata
|
|
||||||
assert doc.metadata["contentType"] == "text/plain"
|
|
||||||
|
|
||||||
client.close()
|
assert doc.id is not None
|
||||||
|
assert doc.content == "This is plain text content from a URL."
|
||||||
|
assert doc.uri == "https://example.com/readme.txt"
|
||||||
|
assert "contentType" in doc.metadata
|
||||||
|
assert "md5" in doc.metadata
|
||||||
|
assert doc.metadata["contentType"] == "text/plain"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_url_unsupported_content():
|
async def test_client_create_document_from_url_unsupported_content():
|
||||||
"""Test creating a document from URL with unsupported content type."""
|
"""Test creating a document from URL with unsupported content type."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Mock response with unsupported content type
|
||||||
|
mock_response = AsyncMock()
|
||||||
|
mock_response.content = b"binary content"
|
||||||
|
mock_response.headers = {"content-type": "application/octet-stream"}
|
||||||
|
mock_response.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
# Mock response with unsupported content type
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
||||||
mock_response = AsyncMock()
|
with pytest.raises(ValueError, match="Unsupported content type"):
|
||||||
mock_response.content = b"binary content"
|
await client.create_document_from_source(
|
||||||
mock_response.headers = {"content-type": "application/octet-stream"}
|
"https://example.com/binary.bin"
|
||||||
mock_response.raise_for_status = AsyncMock()
|
)
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
|
||||||
with pytest.raises(ValueError, match="Unsupported content type"):
|
|
||||||
await client.create_document_from_source("https://example.com/binary.bin")
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_document_from_url_http_error():
|
async def test_client_create_document_from_url_http_error():
|
||||||
"""Test handling HTTP errors when creating document from URL."""
|
"""Test handling HTTP errors when creating document from URL."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
with patch("httpx.AsyncClient.get") as mock_get:
|
||||||
with patch("httpx.AsyncClient.get") as mock_get:
|
mock_get.side_effect = httpx.HTTPStatusError(
|
||||||
mock_get.side_effect = httpx.HTTPStatusError(
|
"404 Not Found",
|
||||||
"404 Not Found",
|
request=httpx.Request("GET", "https://example.com/notfound.html"),
|
||||||
request=httpx.Request("GET", "https://example.com/notfound.html"),
|
response=httpx.Response(404),
|
||||||
response=httpx.Response(404),
|
|
||||||
)
|
|
||||||
|
|
||||||
with pytest.raises(httpx.HTTPStatusError):
|
|
||||||
await client.create_document_from_source(
|
|
||||||
"https://example.com/notfound.html"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
client.close()
|
with pytest.raises(httpx.HTTPStatusError):
|
||||||
|
await client.create_document_from_source(
|
||||||
|
"https://example.com/notfound.html"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_get_extension_from_content_type_or_url():
|
async def test_get_extension_from_content_type_or_url():
|
||||||
"""Test the helper method for determining file extensions."""
|
"""Test the helper method for determining file extensions."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Test content type mappings
|
||||||
# Test content type mappings
|
assert (
|
||||||
assert client._get_extension_from_content_type_or_url("", "text/html") == ".html"
|
client._get_extension_from_content_type_or_url("", "text/html") == ".html"
|
||||||
assert (
|
|
||||||
client._get_extension_from_content_type_or_url("", "application/pdf") == ".pdf"
|
|
||||||
)
|
|
||||||
assert client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
|
|
||||||
|
|
||||||
# Test URL extension detection
|
|
||||||
assert (
|
|
||||||
client._get_extension_from_content_type_or_url(
|
|
||||||
"https://example.com/doc.pdf", ""
|
|
||||||
)
|
)
|
||||||
== ".pdf"
|
assert (
|
||||||
)
|
client._get_extension_from_content_type_or_url("", "application/pdf")
|
||||||
assert (
|
== ".pdf"
|
||||||
client._get_extension_from_content_type_or_url(
|
|
||||||
"https://example.com/data.json", ""
|
|
||||||
)
|
)
|
||||||
== ".json"
|
assert (
|
||||||
)
|
client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
|
||||||
|
|
||||||
# Test default fallback
|
|
||||||
assert (
|
|
||||||
client._get_extension_from_content_type_or_url("https://example.com/", "")
|
|
||||||
== ".html"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Test content type priority over URL extension
|
|
||||||
assert (
|
|
||||||
client._get_extension_from_content_type_or_url(
|
|
||||||
"https://example.com/file.txt", "application/pdf"
|
|
||||||
)
|
)
|
||||||
== ".pdf"
|
|
||||||
)
|
|
||||||
|
|
||||||
client.close()
|
# Test URL extension detection
|
||||||
|
assert (
|
||||||
|
client._get_extension_from_content_type_or_url(
|
||||||
|
"https://example.com/doc.pdf", ""
|
||||||
|
)
|
||||||
|
== ".pdf"
|
||||||
|
)
|
||||||
|
assert (
|
||||||
|
client._get_extension_from_content_type_or_url(
|
||||||
|
"https://example.com/data.json", ""
|
||||||
|
)
|
||||||
|
== ".json"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Test default fallback
|
||||||
|
assert (
|
||||||
|
client._get_extension_from_content_type_or_url("https://example.com/", "")
|
||||||
|
== ".html"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Test content type priority over URL extension
|
||||||
|
assert (
|
||||||
|
client._get_extension_from_content_type_or_url(
|
||||||
|
"https://example.com/file.txt", "application/pdf"
|
||||||
|
)
|
||||||
|
== ".pdf"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
@ -313,165 +283,147 @@ async def test_client_metadata_content_type_and_md5():
|
||||||
"""Test that contentType and md5 metadata are correctly set."""
|
"""Test that contentType and md5 metadata are correctly set."""
|
||||||
import hashlib
|
import hashlib
|
||||||
|
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Create a temporary file with known content
|
||||||
|
test_content = "Test content for MD5 calculation."
|
||||||
|
expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
|
||||||
|
|
||||||
# Create a temporary file with known content
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
test_content = "Test content for MD5 calculation."
|
temp_path = Path(temp_dir) / "test.txt"
|
||||||
expected_md5 = hashlib.md5(test_content.encode()).hexdigest()
|
temp_path.write_text(test_content)
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
doc = await client.create_document_from_source(temp_path)
|
||||||
f.write(test_content)
|
|
||||||
temp_path = Path(f.name)
|
|
||||||
|
|
||||||
try:
|
assert doc.metadata["contentType"] == "text/plain"
|
||||||
doc = await client.create_document_from_source(temp_path)
|
assert doc.metadata["md5"] == expected_md5
|
||||||
|
|
||||||
assert doc.metadata["contentType"] == "text/plain"
|
mock_response = AsyncMock()
|
||||||
assert doc.metadata["md5"] == expected_md5
|
mock_response.content = test_content.encode()
|
||||||
|
mock_response.headers = {"content-type": "text/plain"}
|
||||||
|
mock_response.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
mock_response = AsyncMock()
|
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
||||||
mock_response.content = test_content.encode()
|
url_doc = await client.create_document_from_source(
|
||||||
mock_response.headers = {"content-type": "text/plain"}
|
"https://example.com/test.txt"
|
||||||
mock_response.raise_for_status = AsyncMock()
|
)
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_response):
|
assert url_doc.metadata["contentType"] == "text/plain"
|
||||||
url_doc = await client.create_document_from_source(
|
assert url_doc.metadata["md5"] == expected_md5
|
||||||
"https://example.com/test.txt"
|
|
||||||
)
|
|
||||||
|
|
||||||
assert url_doc.metadata["contentType"] == "text/plain"
|
|
||||||
assert url_doc.metadata["md5"] == expected_md5
|
|
||||||
|
|
||||||
finally:
|
|
||||||
temp_path.unlink()
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_create_update_no_op_behavior():
|
async def test_client_create_update_no_op_behavior():
|
||||||
"""Test create/update/no-op behavior based on MD5 changes."""
|
"""Test create/update/no-op behavior based on MD5 changes."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Create a temporary file
|
||||||
|
test_content = "Original content for testing."
|
||||||
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
|
temp_path = Path(temp_dir) / "test.txt"
|
||||||
|
temp_path.write_text(test_content)
|
||||||
|
|
||||||
# Create a temporary file
|
# First call - should create new document
|
||||||
test_content = "Original content for testing."
|
doc1 = await client.create_document_from_source(temp_path)
|
||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
assert doc1.id is not None
|
||||||
f.write(test_content)
|
assert doc1.content == test_content
|
||||||
temp_path = Path(f.name)
|
original_id = doc1.id
|
||||||
|
|
||||||
try:
|
# Second call with same content - should return existing document (no-op)
|
||||||
# First call - should create new document
|
doc2 = await client.create_document_from_source(temp_path)
|
||||||
doc1 = await client.create_document_from_source(temp_path)
|
assert doc2.id == original_id # Same document
|
||||||
assert doc1.id is not None
|
assert doc2.content == test_content
|
||||||
assert doc1.content == test_content
|
|
||||||
original_id = doc1.id
|
|
||||||
|
|
||||||
# Second call with same content - should return existing document (no-op)
|
# Modify file content
|
||||||
doc2 = await client.create_document_from_source(temp_path)
|
updated_content = "Updated content for testing."
|
||||||
assert doc2.id == original_id # Same document
|
temp_path.write_text(updated_content)
|
||||||
assert doc2.content == test_content
|
|
||||||
|
|
||||||
# Modify file content
|
# Third call with changed content - should update existing document
|
||||||
updated_content = "Updated content for testing."
|
doc3 = await client.create_document_from_source(temp_path)
|
||||||
temp_path.write_text(updated_content)
|
assert doc3.id == original_id # Same document ID
|
||||||
|
assert doc3.content == updated_content # Updated content
|
||||||
|
|
||||||
# Third call with changed content - should update existing document
|
# Verify the document was actually updated in database
|
||||||
doc3 = await client.create_document_from_source(temp_path)
|
retrieved_doc = await client.get_document_by_id(original_id)
|
||||||
assert doc3.id == original_id # Same document ID
|
assert retrieved_doc is not None
|
||||||
assert doc3.content == updated_content # Updated content
|
assert retrieved_doc.content == updated_content
|
||||||
|
|
||||||
# Verify the document was actually updated in database
|
|
||||||
retrieved_doc = await client.get_document_by_id(original_id)
|
|
||||||
assert retrieved_doc is not None
|
|
||||||
assert retrieved_doc.content == updated_content
|
|
||||||
|
|
||||||
finally:
|
|
||||||
temp_path.unlink()
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_url_create_update_no_op_behavior():
|
async def test_client_url_create_update_no_op_behavior():
|
||||||
"""Test create/update/no-op behavior for URLs based on MD5 changes."""
|
"""Test create/update/no-op behavior for URLs based on MD5 changes."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
url = "https://example.com/test.txt"
|
||||||
|
original_content = b"Original URL content"
|
||||||
|
updated_content = b"Updated URL content"
|
||||||
|
|
||||||
url = "https://example.com/test.txt"
|
# Mock first response
|
||||||
original_content = b"Original URL content"
|
mock_response1 = AsyncMock()
|
||||||
updated_content = b"Updated URL content"
|
mock_response1.content = original_content
|
||||||
|
mock_response1.headers = {"content-type": "text/plain"}
|
||||||
|
mock_response1.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
# Mock first response
|
with patch("httpx.AsyncClient.get", return_value=mock_response1):
|
||||||
mock_response1 = AsyncMock()
|
# First call - should create new document
|
||||||
mock_response1.content = original_content
|
doc1 = await client.create_document_from_source(url)
|
||||||
mock_response1.headers = {"content-type": "text/plain"}
|
assert doc1.id is not None
|
||||||
mock_response1.raise_for_status = AsyncMock()
|
original_id = doc1.id
|
||||||
|
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_response1):
|
# Second call with same content - should return existing document (no-op)
|
||||||
# First call - should create new document
|
doc2 = await client.create_document_from_source(url)
|
||||||
doc1 = await client.create_document_from_source(url)
|
assert doc2.id == original_id # Same document
|
||||||
assert doc1.id is not None
|
|
||||||
original_id = doc1.id
|
|
||||||
|
|
||||||
# Second call with same content - should return existing document (no-op)
|
mock_response2 = AsyncMock()
|
||||||
doc2 = await client.create_document_from_source(url)
|
mock_response2.content = updated_content
|
||||||
assert doc2.id == original_id # Same document
|
mock_response2.headers = {"content-type": "text/plain"}
|
||||||
|
mock_response2.raise_for_status = AsyncMock()
|
||||||
|
|
||||||
mock_response2 = AsyncMock()
|
with patch("httpx.AsyncClient.get", return_value=mock_response2):
|
||||||
mock_response2.content = updated_content
|
# Third call with changed content - should update existing document
|
||||||
mock_response2.headers = {"content-type": "text/plain"}
|
doc3 = await client.create_document_from_source(url)
|
||||||
mock_response2.raise_for_status = AsyncMock()
|
assert doc3.id == original_id # Same document ID
|
||||||
|
assert doc3.content == updated_content.decode() # Updated content
|
||||||
with patch("httpx.AsyncClient.get", return_value=mock_response2):
|
|
||||||
# Third call with changed content - should update existing document
|
|
||||||
doc3 = await client.create_document_from_source(url)
|
|
||||||
assert doc3.id == original_id # Same document ID
|
|
||||||
assert doc3.content == updated_content.decode() # Updated content
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_client_search():
|
async def test_client_search():
|
||||||
"""Test HaikuRAG search functionality."""
|
"""Test HaikuRAG search functionality."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
# Add multiple documents to search from
|
||||||
|
doc1_text = "Python is a high-level programming language known for its simplicity and readability."
|
||||||
|
doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
|
||||||
|
doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."
|
||||||
|
|
||||||
# Add multiple documents to search from
|
# Create documents
|
||||||
doc1_text = "Python is a high-level programming language known for its simplicity and readability."
|
doc1 = await client.create_document(
|
||||||
doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
|
content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
|
||||||
doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."
|
)
|
||||||
|
doc2 = await client.create_document(
|
||||||
|
content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"}
|
||||||
|
)
|
||||||
|
await client.create_document(
|
||||||
|
content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"}
|
||||||
|
)
|
||||||
|
|
||||||
# Create documents
|
# Test search with keyword that should match doc1
|
||||||
doc1 = await client.create_document(
|
results = await client.search("Python programming", limit=3)
|
||||||
content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
|
|
||||||
)
|
|
||||||
doc2 = await client.create_document(
|
|
||||||
content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"}
|
|
||||||
)
|
|
||||||
await client.create_document(
|
|
||||||
content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"}
|
|
||||||
)
|
|
||||||
|
|
||||||
# Test search with keyword that should match doc1
|
assert len(results) > 0
|
||||||
results = await client.search("Python programming", limit=3)
|
assert all(len(result) == 2 for result in results)
|
||||||
|
|
||||||
assert len(results) > 0
|
# Verify first result is from the Python document (doc1)
|
||||||
assert all(len(result) == 2 for result in results)
|
first_chunk, _ = results[0]
|
||||||
|
assert first_chunk.document_id == doc1.id
|
||||||
|
|
||||||
# Verify first result is from the Python document (doc1)
|
# Test search with different query
|
||||||
first_chunk, _ = results[0]
|
ml_results = await client.search("machine learning data", limit=2)
|
||||||
assert first_chunk.document_id == doc1.id
|
assert len(ml_results) > 0
|
||||||
|
|
||||||
# Test search with different query
|
# Verify first result is from the machine learning document (doc2)
|
||||||
ml_results = await client.search("machine learning data", limit=2)
|
first_ml_chunk, _ = ml_results[0]
|
||||||
assert len(ml_results) > 0
|
assert first_ml_chunk.document_id == doc2.id
|
||||||
|
|
||||||
# Verify first result is from the machine learning document (doc2)
|
# Test search with limit parameter
|
||||||
first_ml_chunk, _ = ml_results[0]
|
limited_results = await client.search("programming", limit=1)
|
||||||
assert first_ml_chunk.document_id == doc2.id
|
assert len(limited_results) <= 1
|
||||||
|
|
||||||
# Test search with limit parameter
|
|
||||||
limited_results = await client.search("programming", limit=1)
|
|
||||||
assert len(limited_results) <= 1
|
|
||||||
|
|
||||||
client.close()
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
|
|
@ -13,11 +13,10 @@ from haiku.rag.store.models.document import Document
|
||||||
async def test_file_watcher_upsert_document():
|
async def test_file_watcher_upsert_document():
|
||||||
"""Test FileWatcher._upsert_document method."""
|
"""Test FileWatcher._upsert_document method."""
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
f.write("Test content for file watcher")
|
temp_path = Path(temp_dir) / "test.txt"
|
||||||
temp_path = Path(f.name)
|
temp_path.write_text("Test content for file watcher")
|
||||||
|
|
||||||
try:
|
|
||||||
mock_client = AsyncMock(spec=HaikuRAG)
|
mock_client = AsyncMock(spec=HaikuRAG)
|
||||||
mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri())
|
mock_doc = Document(id=1, content="Test content", uri=temp_path.as_uri())
|
||||||
mock_client.create_document_from_source.return_value = mock_doc
|
mock_client.create_document_from_source.return_value = mock_doc
|
||||||
|
|
@ -32,19 +31,15 @@ async def test_file_watcher_upsert_document():
|
||||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||||
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
||||||
|
|
||||||
finally:
|
|
||||||
temp_path.unlink(missing_ok=True)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_file_watcher_upsert_existing_document():
|
async def test_file_watcher_upsert_existing_document():
|
||||||
"""Test FileWatcher._upsert_document with existing document."""
|
"""Test FileWatcher._upsert_document with existing document."""
|
||||||
|
|
||||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
|
with tempfile.TemporaryDirectory() as temp_dir:
|
||||||
f.write("Test content for file watcher")
|
temp_path = Path(temp_dir) / "test.txt"
|
||||||
temp_path = Path(f.name)
|
temp_path.write_text("Test content for file watcher")
|
||||||
|
|
||||||
try:
|
|
||||||
mock_client = AsyncMock(spec=HaikuRAG)
|
mock_client = AsyncMock(spec=HaikuRAG)
|
||||||
existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri())
|
existing_doc = Document(id=1, content="Old content", uri=temp_path.as_uri())
|
||||||
updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri())
|
updated_doc = Document(id=1, content="Updated content", uri=temp_path.as_uri())
|
||||||
|
|
@ -61,9 +56,6 @@ async def test_file_watcher_upsert_existing_document():
|
||||||
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
mock_client.get_document_by_uri.assert_called_once_with(temp_path.as_uri())
|
||||||
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
mock_client.create_document_from_source.assert_called_once_with(str(temp_path))
|
||||||
|
|
||||||
finally:
|
|
||||||
temp_path.unlink(missing_ok=True)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_file_watcher_delete_document():
|
async def test_file_watcher_delete_document():
|
||||||
|
|
|
||||||
|
|
@ -8,45 +8,42 @@ from haiku.rag.store.models.document import Document
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_rebuild_database(qa_corpus: Dataset):
|
async def test_rebuild_database(qa_corpus: Dataset):
|
||||||
"""Test rebuild functionality with existing documents."""
|
"""Test rebuild functionality with existing documents."""
|
||||||
client = HaikuRAG(":memory:")
|
async with HaikuRAG(":memory:") as client:
|
||||||
|
created_docs: list[Document] = []
|
||||||
|
for content in qa_corpus["document_extracted"][:3]:
|
||||||
|
doc = await client.create_document(
|
||||||
|
content=content,
|
||||||
|
)
|
||||||
|
created_docs.append(doc)
|
||||||
|
|
||||||
created_docs: list[Document] = []
|
documents_before = await client.list_documents()
|
||||||
for content in qa_corpus["document_extracted"][:3]:
|
assert len(documents_before) == 3
|
||||||
doc = await client.create_document(
|
|
||||||
content=content,
|
|
||||||
)
|
|
||||||
created_docs.append(doc)
|
|
||||||
|
|
||||||
documents_before = await client.list_documents()
|
chunks_before = []
|
||||||
assert len(documents_before) == 3
|
for doc in created_docs:
|
||||||
|
assert doc.id is not None
|
||||||
chunks_before = []
|
|
||||||
for doc in created_docs:
|
|
||||||
assert doc.id is not None
|
|
||||||
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
|
||||||
chunks_before.extend(doc_chunks)
|
|
||||||
|
|
||||||
assert len(chunks_before) > 0
|
|
||||||
|
|
||||||
# Perform rebuild
|
|
||||||
processed_doc_ids = []
|
|
||||||
async for doc_id in client.rebuild_database():
|
|
||||||
processed_doc_ids.append(doc_id)
|
|
||||||
|
|
||||||
# Verify all documents were processed
|
|
||||||
expected_doc_ids = [doc.id for doc in created_docs]
|
|
||||||
assert set(processed_doc_ids) == set(expected_doc_ids)
|
|
||||||
|
|
||||||
documents_after = await client.list_documents()
|
|
||||||
assert len(documents_after) == 3
|
|
||||||
|
|
||||||
# Verify chunks were recreated
|
|
||||||
chunks_after = []
|
|
||||||
for doc in documents_after:
|
|
||||||
if doc.id is not None:
|
|
||||||
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||||
chunks_after.extend(doc_chunks)
|
chunks_before.extend(doc_chunks)
|
||||||
|
|
||||||
assert len(chunks_after) > 0
|
assert len(chunks_before) > 0
|
||||||
|
|
||||||
client.close()
|
# Perform rebuild
|
||||||
|
processed_doc_ids = []
|
||||||
|
async for doc_id in client.rebuild_database():
|
||||||
|
processed_doc_ids.append(doc_id)
|
||||||
|
|
||||||
|
# Verify all documents were processed
|
||||||
|
expected_doc_ids = [doc.id for doc in created_docs]
|
||||||
|
assert set(processed_doc_ids) == set(expected_doc_ids)
|
||||||
|
|
||||||
|
documents_after = await client.list_documents()
|
||||||
|
assert len(documents_after) == 3
|
||||||
|
|
||||||
|
# Verify chunks were recreated
|
||||||
|
chunks_after = []
|
||||||
|
for doc in documents_after:
|
||||||
|
if doc.id is not None:
|
||||||
|
doc_chunks = await client.chunk_repository.get_by_document_id(doc.id)
|
||||||
|
chunks_after.extend(doc_chunks)
|
||||||
|
|
||||||
|
assert len(chunks_after) > 0
|
||||||
|
|
|
||||||
81
tests/test_settings.py
Normal file
81
tests/test_settings.py
Normal file
|
|
@ -0,0 +1,81 @@
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
from haiku.rag.config import Config
|
||||||
|
from haiku.rag.store.engine import Store
|
||||||
|
from haiku.rag.store.repositories.settings import (
|
||||||
|
ConfigMismatchError,
|
||||||
|
SettingsRepository,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_settings_table_populated_on_store_init():
|
||||||
|
"""Test that settings table is populated with current config when store is initialized."""
|
||||||
|
|
||||||
|
store = Store(":memory:")
|
||||||
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
|
db_settings = settings_repo.get()
|
||||||
|
config_dict = Config.model_dump(mode="json")
|
||||||
|
|
||||||
|
assert db_settings == config_dict
|
||||||
|
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_settings_save_and_retrieve():
|
||||||
|
"""Test saving and retrieving settings after config change."""
|
||||||
|
store = Store(":memory:")
|
||||||
|
settings_repo = SettingsRepository(store)
|
||||||
|
|
||||||
|
original_chunk_size = Config.CHUNK_SIZE
|
||||||
|
Config.CHUNK_SIZE = 2 * original_chunk_size
|
||||||
|
|
||||||
|
settings_repo.save()
|
||||||
|
retrieved_settings = settings_repo.get()
|
||||||
|
assert retrieved_settings["CHUNK_SIZE"] == 2 * original_chunk_size
|
||||||
|
|
||||||
|
Config.CHUNK_SIZE = original_chunk_size
|
||||||
|
store.close()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_config_validation_on_db_load():
|
||||||
|
"""Test that config validation fails when loading db with mismatched settings."""
|
||||||
|
# Create a temporary database file
|
||||||
|
with tempfile.NamedTemporaryFile(suffix=".sqlite") as tmp:
|
||||||
|
db_path = Path(tmp.name)
|
||||||
|
|
||||||
|
# Create store and save settings
|
||||||
|
store1 = Store(db_path)
|
||||||
|
SettingsRepository(store1)
|
||||||
|
store1.close()
|
||||||
|
|
||||||
|
# Change config
|
||||||
|
original_chunk_size = Config.CHUNK_SIZE
|
||||||
|
Config.CHUNK_SIZE = 999
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Loading the database should raise ConfigMismatchError
|
||||||
|
with pytest.raises(ConfigMismatchError) as exc_info:
|
||||||
|
Store(db_path)
|
||||||
|
|
||||||
|
assert "CHUNK_SIZE" in str(exc_info.value)
|
||||||
|
assert "Consider rebuilding" in str(exc_info.value)
|
||||||
|
|
||||||
|
# Rebuild
|
||||||
|
async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
|
||||||
|
async for _ in client.rebuild_database():
|
||||||
|
pass # Process all documents
|
||||||
|
|
||||||
|
# Verify we can now load the database without exception (settings were updated)
|
||||||
|
store2 = Store(db_path)
|
||||||
|
settings_repo2 = SettingsRepository(store2)
|
||||||
|
db_settings = settings_repo2.get()
|
||||||
|
assert db_settings["CHUNK_SIZE"] == 999
|
||||||
|
store2.close()
|
||||||
|
|
||||||
|
finally:
|
||||||
|
Config.CHUNK_SIZE = original_chunk_size
|
||||||
8
uv.lock
8
uv.lock
|
|
@ -881,7 +881,7 @@ dev = [
|
||||||
{ name = "mkdocs", specifier = ">=1.6.1" },
|
{ name = "mkdocs", specifier = ">=1.6.1" },
|
||||||
{ name = "mkdocs-material", specifier = ">=9.6.14" },
|
{ name = "mkdocs-material", specifier = ">=9.6.14" },
|
||||||
{ name = "pre-commit", specifier = ">=4.2.0" },
|
{ name = "pre-commit", specifier = ">=4.2.0" },
|
||||||
{ name = "pyright", specifier = ">=1.1.402" },
|
{ name = "pyright", specifier = ">=1.1.403" },
|
||||||
{ name = "pytest", specifier = ">=8.4.0" },
|
{ name = "pytest", specifier = ">=8.4.0" },
|
||||||
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
|
{ name = "pytest-asyncio", specifier = ">=1.0.0" },
|
||||||
{ name = "pytest-cov", specifier = ">=6.2.1" },
|
{ name = "pytest-cov", specifier = ">=6.2.1" },
|
||||||
|
|
@ -2305,15 +2305,15 @@ wheels = [
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "pyright"
|
name = "pyright"
|
||||||
version = "1.1.402"
|
version = "1.1.403"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "nodeenv" },
|
{ name = "nodeenv" },
|
||||||
{ name = "typing-extensions" },
|
{ name = "typing-extensions" },
|
||||||
]
|
]
|
||||||
sdist = { url = "https://files.pythonhosted.org/packages/aa/04/ce0c132d00e20f2d2fb3b3e7c125264ca8b909e693841210534b1ea1752f/pyright-1.1.402.tar.gz", hash = "sha256:85a33c2d40cd4439c66aa946fd4ce71ab2f3f5b8c22ce36a623f59ac22937683", size = 3888207, upload-time = "2025-06-11T08:48:35.759Z" }
|
sdist = { url = "https://files.pythonhosted.org/packages/fe/f6/35f885264ff08c960b23d1542038d8da86971c5d8c955cfab195a4f672d7/pyright-1.1.403.tar.gz", hash = "sha256:3ab69b9f41c67fb5bbb4d7a36243256f0d549ed3608678d381d5f51863921104", size = 3913526, upload-time = "2025-07-09T07:15:52.882Z" }
|
||||||
wheels = [
|
wheels = [
|
||||||
{ url = "https://files.pythonhosted.org/packages/fe/37/1a1c62d955e82adae588be8e374c7f77b165b6cb4203f7d581269959abbc/pyright-1.1.402-py3-none-any.whl", hash = "sha256:2c721f11869baac1884e846232800fe021c33f1b4acb3929cff321f7ea4e2982", size = 5624004, upload-time = "2025-06-11T08:48:33.998Z" },
|
{ url = "https://files.pythonhosted.org/packages/49/b6/b04e5c2f41a5ccad74a1a4759da41adb20b4bc9d59a5e08d29ba60084d07/pyright-1.1.403-py3-none-any.whl", hash = "sha256:c0eeca5aa76cbef3fcc271259bbd785753c7ad7bcac99a9162b4c4c7daed23b3", size = 5684504, upload-time = "2025-07-09T07:15:50.958Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue