Make list_documents return structured output and give feedback about total count and pages
This commit is contained in:
parent
cca03dbe61
commit
626f5cdd0e
7 changed files with 78 additions and 18 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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",
|
||||||
|
|
|
||||||
|
|
@ -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."""
|
||||||
|
|
|
||||||
|
|
@ -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."""
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue