danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add uri to Documents

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-06-17 10:41:44 +02:00
parent e24d3cdb0a
commit 1ff17c2f50
No known key found for this signature in database
4 changed files with 92 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
src/haiku/rag/store/engine.py
View file

@ -24,6 +24,7 @@ class Store:
CREATE TABLE IF NOT EXISTS documents (
id INTEGER PRIMARY KEY AUTOINCREMENT,
content TEXT NOT NULL,
uri TEXT,
metadata TEXT DEFAULT '{}',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP

1
src/haiku/rag/store/models/document.py
View file

@ -10,6 +10,7 @@ class Document(BaseModel):
id: int | None = None
content: str
uri: str | None = None
metadata: dict = {}
created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now)

48
src/haiku/rag/store/repositories/document.py
View file

@ -30,11 +30,12 @@ class DocumentRepository(BaseRepository[Document]):
# Insert the document
cursor.execute(
"""
INSERT INTO documents (content, metadata, created_at, updated_at)
VALUES (?, ?, ?, ?)
INSERT INTO documents (content, uri, metadata, created_at, updated_at)
VALUES (?, ?, ?, ?, ?)
""",
(
entity.content,
entity.uri,
json.dumps(entity.metadata),
entity.created_at,
entity.updated_at,
@ -65,7 +66,7 @@ class DocumentRepository(BaseRepository[Document]):
cursor = self.store._connection.cursor()
cursor.execute(
"""
SELECT id, content, metadata, created_at, updated_at
SELECT id, content, uri, metadata, created_at, updated_at
FROM documents WHERE id = ?
""",
(entity_id,),
@ -75,12 +76,43 @@ class DocumentRepository(BaseRepository[Document]):
if row is None:
return None
document_id, content, metadata_json, created_at, updated_at = row
document_id, content, uri, metadata_json, created_at, updated_at = row
metadata = json.loads(metadata_json) if metadata_json else {}
return Document(
id=document_id,
content=content,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,
)
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
if self.store._connection is None:
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
cursor.execute(
"""
SELECT id, content, uri, metadata, created_at, updated_at
FROM documents WHERE uri = ?
""",
(uri,),
)
row = cursor.fetchone()
if row is None:
return None
document_id, content, uri, metadata_json, created_at, updated_at = row
metadata = json.loads(metadata_json) if metadata_json else {}
return Document(
id=document_id,
content=content,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,
@ -103,11 +135,12 @@ class DocumentRepository(BaseRepository[Document]):
cursor.execute(
"""
UPDATE documents
SET content = ?, metadata = ?, updated_at = ?
SET content = ?, uri = ?, metadata = ?, updated_at = ?
WHERE id = ?
""",
(
entity.content,
entity.uri,
json.dumps(entity.metadata),
entity.updated_at,
entity.id,
@ -150,7 +183,7 @@ class DocumentRepository(BaseRepository[Document]):
raise ValueError("Store connection is not available")
cursor = self.store._connection.cursor()
query = "SELECT id, content, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
query = "SELECT id, content, uri, metadata, created_at, updated_at FROM documents ORDER BY created_at DESC"
params = []
if limit is not None:
@ -166,12 +199,13 @@ class DocumentRepository(BaseRepository[Document]):
documents = []
for row in rows:
document_id, content, metadata_json, created_at, updated_at = row
document_id, content, uri, metadata_json, created_at, updated_at = row
metadata = json.loads(metadata_json) if metadata_json else {}
documents.append(
Document(
id=document_id,
content=content,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,

77
tests/test_document.py
View file

@ -12,52 +12,61 @@ async def test_create_document_with_chunks(qa_corpus: Dataset):
# Create an in-memory store and repository
store = Store(":memory:")
doc_repo = DocumentRepository(store)
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Create a Document instance
document = Document(
content=document_text,
metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")}
metadata={"source": "qa_corpus", "topic": first_doc.get("document_topic", "")},
)
# Create the document with chunks in the database
created_document = await doc_repo.create(document)
# Verify the document was created
assert created_document.id is not None
assert created_document.content == document_text
# Check that chunks were created in the database
if store._connection is not None:
cursor = store._connection.cursor()
cursor.execute("SELECT COUNT(*) FROM chunks WHERE document_id = ?", (created_document.id,))
cursor.execute(
"SELECT COUNT(*) FROM chunks WHERE document_id = ?", (created_document.id,)
)
chunk_count = cursor.fetchone()[0]
assert chunk_count > 0
# Check that embeddings were created
cursor.execute("""
cursor.execute(
"""
SELECT COUNT(*) FROM chunk_embeddings ce
JOIN chunks c ON c.id = ce.chunk_id
WHERE c.document_id = ?
""", (created_document.id,))
""",
(created_document.id,),
)
embedding_count = cursor.fetchone()[0]
assert embedding_count == chunk_count
# Verify chunk metadata contains order information
cursor.execute("SELECT metadata FROM chunks WHERE document_id = ? ORDER BY id", (created_document.id,))
cursor.execute(
"SELECT metadata FROM chunks WHERE document_id = ? ORDER BY id",
(created_document.id,),
)
chunk_metadata = cursor.fetchall()
for i, (metadata_json,) in enumerate(chunk_metadata):
import json
metadata = json.loads(metadata_json)
assert "order" in metadata
assert metadata["order"] == i
store.close()
@ -67,40 +76,52 @@ async def test_document_repository_crud(qa_corpus: Dataset):
# Create an in-memory store and repository
store = Store(":memory:")
doc_repo = DocumentRepository(store)
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Create a document
# Create a document with URI
test_uri = "file:///path/to/test.txt"
document = Document(
content=document_text,
metadata={"source": "test"}
content=document_text, uri=test_uri, metadata={"source": "test"}
)
created_document = await doc_repo.create(document)
# Test get_by_id
assert created_document.id is not None
retrieved_document = await doc_repo.get_by_id(created_document.id)
assert retrieved_document is not None
assert retrieved_document.content == document_text
assert retrieved_document.uri == test_uri
# Test get_by_uri
retrieved_by_uri = await doc_repo.get_by_uri(test_uri)
assert retrieved_by_uri is not None
assert retrieved_by_uri.id == created_document.id
assert retrieved_by_uri.content == document_text
assert retrieved_by_uri.uri == test_uri
# Test get_by_uri with non-existent URI
non_existent = await doc_repo.get_by_uri("file:///non/existent.txt")
assert non_existent is None
# Test update (should regenerate chunks)
retrieved_document.content = "Updated content for testing"
updated_document = await doc_repo.update(retrieved_document)
assert updated_document.content == "Updated content for testing"
# Test list_all
all_documents = await doc_repo.list_all()
assert len(all_documents) == 1
assert all_documents[0].id == created_document.id
# Test delete
deleted = await doc_repo.delete(created_document.id)
assert deleted is True
# Verify document is gone
retrieved_document = await doc_repo.get_by_id(created_document.id)
assert retrieved_document is None
store.close()
store.close()