Merge pull request #374 from tseaver/feat-cli-search_type

feat: add '--search-type' option to CLI 'search'
This commit is contained in:
Yiorgis Gozadinos 2026-05-18 16:39:28 +03:00 committed by GitHub
commit f208730a5b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 55 additions and 10 deletions

View file

@ -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

View file

@ -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

View file

@ -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")

View file

@ -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]:

View file

@ -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:

View file

@ -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.

View file

@ -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]]: