diff --git a/docs/cli.md b/docs/cli.md index d580e288..86cfd69a 100644 --- a/docs/cli.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index a6ecd55f..45aa4220 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index c55d6a7f..985404bf 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -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") diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 3a26a533..537c5e18 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -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]: diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 6d19ef4f..0cfaa9b1 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 76e2a932..5e86ec80 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index f32121c8..96a8bbf2 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -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]]: