Merge pull request #374 from tseaver/feat-cli-search_type
feat: add '--search-type' option to CLI 'search'
This commit is contained in:
commit
f208730a5b
7 changed files with 55 additions and 10 deletions
12
docs/cli.md
12
docs/cli.md
|
|
@ -133,6 +133,18 @@ With options:
|
|||
haiku-rag search "python programming" --limit 10 # or -l 10
|
||||
```
|
||||
|
||||
With search type:
|
||||
```bash
|
||||
# Hybrid search (the default)
|
||||
haiku-rag search "python programming" --search-type hybrid # or -s hybrid
|
||||
|
||||
# Full-text search only
|
||||
haiku-rag search "python programming" --search-type fts # or -s fts
|
||||
|
||||
# Vector search only
|
||||
haiku-rag search "python programming" --search-type vector # or -s vector
|
||||
```
|
||||
|
||||
With filters (filter by document properties, use `--filter` or `-f`):
|
||||
```bash
|
||||
# Filter by URI pattern
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from haiku.rag.client import HaikuRAG, RebuildMode
|
|||
from haiku.rag.config import AppConfig, Config
|
||||
from haiku.rag.mcp import create_mcp_server
|
||||
from haiku.rag.monitor import FileWatcher, S3Watcher
|
||||
from haiku.rag.store.models.chunk import SearchType
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -361,6 +362,7 @@ class HaikuRAGApp: # pragma: no cover
|
|||
query: str | None = None,
|
||||
limit: int | None = None,
|
||||
filter: str | None = None,
|
||||
search_type: SearchType | None = None,
|
||||
image: Path | None = None,
|
||||
):
|
||||
if query is None and image is None:
|
||||
|
|
@ -372,6 +374,10 @@ class HaikuRAGApp: # pragma: no cover
|
|||
self.console.print("[red]Pass either a query or --image, not both.[/red]")
|
||||
return
|
||||
|
||||
if query is None and search_type is not None:
|
||||
self.console.print("[red]Pass --search-type only for text queries[/red]")
|
||||
return
|
||||
|
||||
search_input: str | bytes
|
||||
if image is not None:
|
||||
search_input = image.read_bytes()
|
||||
|
|
@ -385,7 +391,12 @@ class HaikuRAGApp: # pragma: no cover
|
|||
read_only=self.read_only,
|
||||
before=self.before,
|
||||
) as self.client:
|
||||
results = await self.client.search(search_input, limit=limit, filter=filter)
|
||||
results = await self.client.search(
|
||||
search_input,
|
||||
limit=limit,
|
||||
filter=filter,
|
||||
search_type=search_type,
|
||||
)
|
||||
if not results:
|
||||
self.console.print("[yellow]No results found.[/yellow]")
|
||||
return
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from haiku.rag.store.exceptions import ( # noqa: E402
|
|||
MigrationRequiredError,
|
||||
ReadOnlyError,
|
||||
)
|
||||
from haiku.rag.store.models.chunk import SearchType # noqa: E402
|
||||
from haiku.rag.utils import is_up_to_date # noqa: E402
|
||||
|
||||
_cli = typer.Typer(
|
||||
|
|
@ -314,6 +315,12 @@ def search( # pragma: no cover
|
|||
"-f",
|
||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
||||
),
|
||||
search_type: SearchType | None = typer.Option(
|
||||
None,
|
||||
"--search-type",
|
||||
"-s",
|
||||
help="Type of search to perform (text searches only)",
|
||||
),
|
||||
image: Path | None = typer.Option(
|
||||
None,
|
||||
"--image",
|
||||
|
|
@ -326,7 +333,15 @@ def search( # pragma: no cover
|
|||
),
|
||||
):
|
||||
app = create_app(db)
|
||||
asyncio.run(app.search(query=query, limit=limit, filter=filter, image=image))
|
||||
asyncio.run(
|
||||
app.search(
|
||||
query=query,
|
||||
limit=limit,
|
||||
filter=filter,
|
||||
search_type=search_type,
|
||||
image=image,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@_cli.command("visualize", help="Show visual grounding for a chunk")
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from haiku.rag.config import AppConfig, Config
|
|||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.reranking import get_reranker
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchResult
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.models.document_item import extract_items
|
||||
from haiku.rag.store.repositories.chunk import ChunkRepository
|
||||
|
|
@ -349,7 +349,7 @@ class HaikuRAG:
|
|||
self,
|
||||
query: "str | bytes | PILImage.Image",
|
||||
limit: int | None = None,
|
||||
search_type: str = "hybrid",
|
||||
search_type: SearchType | None = None,
|
||||
filter: str | None = None,
|
||||
include_images: bool = True,
|
||||
) -> list[SearchResult]:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import base64
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
from haiku.rag.reranking import get_reranker
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchResult
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from PIL import Image as PILImage
|
||||
|
|
@ -14,7 +14,7 @@ async def search(
|
|||
client: "HaikuRAG",
|
||||
query: "str | bytes | PILImage.Image",
|
||||
limit: int | None = None,
|
||||
search_type: str = "hybrid",
|
||||
search_type: SearchType | None = None,
|
||||
filter: str | None = None,
|
||||
include_images: bool = True,
|
||||
) -> list[SearchResult]:
|
||||
|
|
@ -25,7 +25,8 @@ async def search(
|
|||
query: Text (``str``) or image (``bytes`` / ``PIL.Image.Image``).
|
||||
Image queries require a multimodal embedder and run vector-only.
|
||||
limit: Maximum number of results to return. Defaults to config.search.limit.
|
||||
search_type: "vector", "fts", or "hybrid" (default). Text queries only.
|
||||
search_type: "vector", "fts", or "hybrid".
|
||||
Applicable only for text queries, where the default is "hybrid".
|
||||
filter: Optional SQL WHERE clause to filter documents before searching chunks.
|
||||
include_images: When True, populate ``SearchResult.image_data`` with
|
||||
base64 picture bytes for picture-labeled chunks.
|
||||
|
|
@ -37,6 +38,9 @@ async def search(
|
|||
limit = client._config.search.limit
|
||||
|
||||
if isinstance(query, str):
|
||||
if search_type is None:
|
||||
search_type = "hybrid"
|
||||
|
||||
reranker = get_reranker(config=client._config)
|
||||
|
||||
if reranker is None:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
|
||||
from pydantic import BaseModel, PrivateAttr
|
||||
|
||||
|
|
@ -113,6 +113,9 @@ class Chunk(BaseModel):
|
|||
return ChunkMetadata.model_validate(self.metadata)
|
||||
|
||||
|
||||
SearchType = Literal["vector", "fts", "hybrid"]
|
||||
|
||||
|
||||
class SearchResult(BaseModel):
|
||||
"""Search result with optional provenance information for citations.
|
||||
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ from lancedb.index import FTS
|
|||
from lancedb.rerankers import RRFReranker
|
||||
|
||||
from haiku.rag.store.engine import Store, query_to_pydantic
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -220,7 +220,7 @@ class ChunkRepository:
|
|||
self,
|
||||
query: str = "",
|
||||
limit: int = 5,
|
||||
search_type: str = "hybrid",
|
||||
search_type: SearchType = "hybrid",
|
||||
filter: str | None = None,
|
||||
query_vector: list[float] | None = None,
|
||||
) -> list[tuple[Chunk, float]]:
|
||||
|
|
|
|||
Loading…
Reference in a new issue