diff --git a/CHANGELOG.md b/CHANGELOG.md index 2977daaa..0c4fdb72 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,9 @@ ### Added - **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 +- **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 - 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 diff --git a/docs/python.md b/docs/python.md index f04fffc7..33f336dd 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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 ```python diff --git a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py index 5ffd5116..cfecef98 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py @@ -8,6 +8,8 @@ from haiku.rag.agents.chat.state import ( AGUI_STATE_KEY, ChatDeps, ChatSessionState, + DocumentInfo, + DocumentListResponse, QAResponse, SearchDeps, SessionContext, @@ -20,6 +22,8 @@ __all__ = [ "SearchAgent", "ChatDeps", "ChatSessionState", + "DocumentInfo", + "DocumentListResponse", "QAResponse", "SearchDeps", "SessionContext", diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 929f0898..85134763 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -14,6 +14,8 @@ from haiku.rag.agents.chat.state import ( MAX_QA_HISTORY, ChatDeps, ChatSessionState, + DocumentInfo, + DocumentListResponse, QAResponse, build_document_filter, build_multi_document_filter, @@ -373,17 +375,18 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: @agent.tool async def list_documents( ctx: RunContext[ChatDeps], - limit: int | None = None, - offset: int | None = None, - ) -> str: + page: int = 1, + ) -> DocumentListResponse: """List available documents in the knowledge base. Use this when the user wants to browse or see what documents are available. Args: - limit: Maximum number of documents to return - offset: Number of documents to skip (for pagination) + page: Page number (default: 1, 50 documents per page) """ + page_size = 50 + offset = (page - 1) * page_size + # Build session filter from document_filter doc_filter = None 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( - 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 "No documents found in the knowledge base." - - lines = [f"Found {len(docs)} document(s):\n"] - for doc in docs: - title = doc.title or "Untitled" - uri = doc.uri or "N/A" - created = doc.created_at.strftime("%Y-%m-%d") - lines.append(f"- **{title}** (URI: {uri}, Created: {created})") - - return "\n".join(lines) + return DocumentListResponse( + documents=[ + DocumentInfo( + title=doc.title or "Untitled", + uri=doc.uri or "", + created=doc.created_at.strftime("%Y-%m-%d"), + ) + for doc in docs + ], + page=page, + total_pages=total_pages, + total_documents=total, + ) async def _find_document(client: HaikuRAG, query: str): """Find a document by exact URI, partial URI, or partial title match.""" diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 69800a92..5e7ccfed 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -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): """Compressed summary of conversation history for research graph.""" diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index d0f458c7..da0ffa77 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -854,6 +854,17 @@ class HaikuRAG: 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( self, query: str, diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 6daef6b0..9f45a1fb 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -167,6 +167,17 @@ class DocumentRepository: results = list(query.to_pydantic(DocumentRecord)) 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: """Get a document by its URI.""" escaped_uri = _escape_sql_string(uri)