Make list_documents return structured output and give feedback about total count and pages

This commit is contained in:
Yiorgis Gozadinos 2026-01-27 17:00:39 +02:00
parent cca03dbe61
commit 626f5cdd0e
No known key found for this signature in database
7 changed files with 78 additions and 18 deletions

View file

@ -4,8 +4,9 @@
### Added ### Added
- **Chat Agent Document Awareness Tools**: Two new tools for browsing and understanding the knowledge base - **Chat Agent Document Awareness Tools**: Two new tools for browsing and understanding the knowledge base
- `list_documents` — Browse available documents with title, URI, and creation date; respects session document filter - `list_documents` — Browse available documents with title, URI, and creation date; respects session document filter; paginated with total count
- `summarize_document` — Generate LLM-powered summaries of specific documents - `summarize_document` — Generate LLM-powered summaries of specific documents
- **Document Count API**: New `count_documents(filter)` method on `HaikuRAG` client for efficient document counting
- **Read-Only Initial Context**: Initial context is now locked after the first message, providing consistent session context - **Read-Only Initial Context**: Initial context is now locked after the first message, providing consistent session context
- Chat TUI: `--initial-context` CLI option sets background context for the session - Chat TUI: `--initial-context` CLI option sets background context for the session
- Context can be edited via command palette before the first message is sent - Context can be edited via command palette before the first message is sent

View file

@ -151,6 +151,15 @@ docs = await client.list_documents(
) )
``` ```
Count documents:
```python
# Count all documents
total = await client.count_documents()
# Count with filter
pdf_count = await client.count_documents(filter="uri LIKE '%.pdf'")
```
### Updating Documents ### Updating Documents
```python ```python

View file

@ -8,6 +8,8 @@ from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY, AGUI_STATE_KEY,
ChatDeps, ChatDeps,
ChatSessionState, ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse, QAResponse,
SearchDeps, SearchDeps,
SessionContext, SessionContext,
@ -20,6 +22,8 @@ __all__ = [
"SearchAgent", "SearchAgent",
"ChatDeps", "ChatDeps",
"ChatSessionState", "ChatSessionState",
"DocumentInfo",
"DocumentListResponse",
"QAResponse", "QAResponse",
"SearchDeps", "SearchDeps",
"SessionContext", "SessionContext",

View file

@ -14,6 +14,8 @@ from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY, MAX_QA_HISTORY,
ChatDeps, ChatDeps,
ChatSessionState, ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse, QAResponse,
build_document_filter, build_document_filter,
build_multi_document_filter, build_multi_document_filter,
@ -373,17 +375,18 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
@agent.tool @agent.tool
async def list_documents( async def list_documents(
ctx: RunContext[ChatDeps], ctx: RunContext[ChatDeps],
limit: int | None = None, page: int = 1,
offset: int | None = None, ) -> DocumentListResponse:
) -> str:
"""List available documents in the knowledge base. """List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available. Use this when the user wants to browse or see what documents are available.
Args: Args:
limit: Maximum number of documents to return page: Page number (default: 1, 50 documents per page)
offset: Number of documents to skip (for pagination)
""" """
page_size = 50
offset = (page - 1) * page_size
# Build session filter from document_filter # Build session filter from document_filter
doc_filter = None doc_filter = None
if ctx.deps.session_state and ctx.deps.session_state.document_filter: if ctx.deps.session_state and ctx.deps.session_state.document_filter:
@ -392,20 +395,24 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
) )
docs = await ctx.deps.client.list_documents( docs = await ctx.deps.client.list_documents(
limit=limit, offset=offset, filter=doc_filter limit=page_size, offset=offset, filter=doc_filter
) )
total = await ctx.deps.client.count_documents(filter=doc_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
if not docs: return DocumentListResponse(
return "No documents found in the knowledge base." documents=[
DocumentInfo(
lines = [f"Found {len(docs)} document(s):\n"] title=doc.title or "Untitled",
for doc in docs: uri=doc.uri or "",
title = doc.title or "Untitled" created=doc.created_at.strftime("%Y-%m-%d"),
uri = doc.uri or "N/A" )
created = doc.created_at.strftime("%Y-%m-%d") for doc in docs
lines.append(f"- **{title}** (URI: {uri}, Created: {created})") ],
page=page,
return "\n".join(lines) total_pages=total_pages,
total_documents=total,
)
async def _find_document(client: HaikuRAG, query: str): async def _find_document(client: HaikuRAG, query: str):
"""Find a document by exact URI, partial URI, or partial title match.""" """Find a document by exact URI, partial URI, or partial title match."""

View file

@ -42,6 +42,23 @@ class QAResponse(BaseModel):
) )
class DocumentInfo(BaseModel):
"""Document info for list_documents response."""
title: str
uri: str
created: str
class DocumentListResponse(BaseModel):
"""Response from list_documents tool."""
documents: list[DocumentInfo]
page: int
total_pages: int
total_documents: int
class SessionContext(BaseModel): class SessionContext(BaseModel):
"""Compressed summary of conversation history for research graph.""" """Compressed summary of conversation history for research graph."""

View file

@ -854,6 +854,17 @@ class HaikuRAG:
limit=limit, offset=offset, filter=filter limit=limit, offset=offset, filter=filter
) )
async def count_documents(self, filter: str | None = None) -> int:
"""Count documents with optional filtering.
Args:
filter: Optional SQL WHERE clause to filter documents.
Returns:
Number of documents matching the criteria.
"""
return await self.document_repository.count(filter=filter)
async def search( async def search(
self, self,
query: str, query: str,

View file

@ -167,6 +167,17 @@ class DocumentRepository:
results = list(query.to_pydantic(DocumentRecord)) results = list(query.to_pydantic(DocumentRecord))
return [self._record_to_document(doc) for doc in results] return [self._record_to_document(doc) for doc in results]
async def count(self, filter: str | None = None) -> int:
"""Count documents with optional filtering.
Args:
filter: Optional SQL WHERE clause to filter documents.
Returns:
Number of documents matching the criteria.
"""
return self.store.documents_table.count_rows(filter=filter)
async def get_by_uri(self, uri: str) -> Document | None: async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI.""" """Get a document by its URI."""
escaped_uri = _escape_sql_string(uri) escaped_uri = _escape_sql_string(uri)