fix: declare 'SearchType' for 'search_type' arg
fix: allow 'search_type' only for text searches
This commit is contained in:
parent
b2ffc50168
commit
7953de46b9
5 changed files with 26 additions and 8 deletions
|
|
@ -19,6 +19,7 @@ from rich.progress import (
|
||||||
from rich.syntax import Syntax
|
from rich.syntax import Syntax
|
||||||
|
|
||||||
from haiku.rag.client import HaikuRAG, RebuildMode
|
from haiku.rag.client import HaikuRAG, RebuildMode
|
||||||
|
from haiku.rag.client.models import SearchType
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
from haiku.rag.mcp import create_mcp_server
|
from haiku.rag.mcp import create_mcp_server
|
||||||
from haiku.rag.monitor import FileWatcher, S3Watcher
|
from haiku.rag.monitor import FileWatcher, S3Watcher
|
||||||
|
|
@ -361,7 +362,7 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
query: str | None = None,
|
query: str | None = None,
|
||||||
limit: int | None = None,
|
limit: int | None = None,
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
search_type: str = "hybrid",
|
search_type: SearchType | None = None,
|
||||||
image: Path | None = None,
|
image: Path | None = None,
|
||||||
):
|
):
|
||||||
if query is None and image is None:
|
if query is None and image is None:
|
||||||
|
|
@ -373,6 +374,10 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
self.console.print("[red]Pass either a query or --image, not both.[/red]")
|
self.console.print("[red]Pass either a query or --image, not both.[/red]")
|
||||||
return
|
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
|
search_input: str | bytes
|
||||||
if image is not None:
|
if image is not None:
|
||||||
search_input = image.read_bytes()
|
search_input = image.read_bytes()
|
||||||
|
|
@ -387,7 +392,10 @@ class HaikuRAGApp: # pragma: no cover
|
||||||
before=self.before,
|
before=self.before,
|
||||||
) as self.client:
|
) as self.client:
|
||||||
results = await self.client.search(
|
results = await self.client.search(
|
||||||
search_input, limit=limit, filter=filter, search_type=search_type,
|
search_input,
|
||||||
|
limit=limit,
|
||||||
|
filter=filter,
|
||||||
|
search_type=search_type,
|
||||||
)
|
)
|
||||||
if not results:
|
if not results:
|
||||||
self.console.print("[yellow]No results found.[/yellow]")
|
self.console.print("[yellow]No results found.[/yellow]")
|
||||||
|
|
|
||||||
|
|
@ -15,6 +15,7 @@ from dotenv import find_dotenv, load_dotenv
|
||||||
load_dotenv(find_dotenv(usecwd=True))
|
load_dotenv(find_dotenv(usecwd=True))
|
||||||
|
|
||||||
from haiku.rag.app import HaikuRAGApp # noqa: E402
|
from haiku.rag.app import HaikuRAGApp # noqa: E402
|
||||||
|
from haiku.rag.client.models import SearchType # noqa: E402
|
||||||
from haiku.rag.config import ( # noqa: E402
|
from haiku.rag.config import ( # noqa: E402
|
||||||
AppConfig,
|
AppConfig,
|
||||||
find_config_file,
|
find_config_file,
|
||||||
|
|
@ -314,11 +315,11 @@ def search( # pragma: no cover
|
||||||
"-f",
|
"-f",
|
||||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
||||||
),
|
),
|
||||||
search_type: str | None = typer.Option(
|
search_type: SearchType = typer.Option(
|
||||||
None,
|
"hybrid",
|
||||||
"--search-type",
|
"--search-type",
|
||||||
"-s",
|
"-s",
|
||||||
help="Type of search: one of 'hybrid' (default) / 'fts' / 'vector' ",
|
help="Type of search to perform (text searches only)",
|
||||||
),
|
),
|
||||||
image: Path | None = typer.Option(
|
image: Path | None = typer.Option(
|
||||||
None,
|
None,
|
||||||
|
|
@ -334,7 +335,11 @@ def search( # pragma: no cover
|
||||||
app = create_app(db)
|
app = create_app(db)
|
||||||
asyncio.run(
|
asyncio.run(
|
||||||
app.search(
|
app.search(
|
||||||
query=query, limit=limit, filter=filter, search_type=search_type, image=image,
|
query=query,
|
||||||
|
limit=limit,
|
||||||
|
filter=filter,
|
||||||
|
search_type=search_type,
|
||||||
|
image=image,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -13,6 +13,7 @@ from urllib.parse import urlparse
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
|
from haiku.rag.client.models import SearchType
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
from haiku.rag.converters import get_converter
|
from haiku.rag.converters import get_converter
|
||||||
from haiku.rag.reranking import get_reranker
|
from haiku.rag.reranking import get_reranker
|
||||||
|
|
@ -349,7 +350,7 @@ class HaikuRAG:
|
||||||
self,
|
self,
|
||||||
query: "str | bytes | PILImage.Image",
|
query: "str | bytes | PILImage.Image",
|
||||||
limit: int | None = None,
|
limit: int | None = None,
|
||||||
search_type: str = "hybrid",
|
search_type: SearchType | None = None,
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
include_images: bool = True,
|
include_images: bool = True,
|
||||||
) -> list[SearchResult]:
|
) -> list[SearchResult]:
|
||||||
|
|
|
||||||
3
haiku_rag_slim/haiku/rag/client/models.py
Normal file
3
haiku_rag_slim/haiku/rag/client/models.py
Normal file
|
|
@ -0,0 +1,3 @@
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
SearchType = Literal["vector", "fts", "hybrid"]
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
import base64
|
import base64
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from haiku.rag.client.models import SearchType
|
||||||
from haiku.rag.reranking import get_reranker
|
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
|
||||||
|
|
||||||
|
|
@ -14,7 +15,7 @@ async def search(
|
||||||
client: "HaikuRAG",
|
client: "HaikuRAG",
|
||||||
query: "str | bytes | PILImage.Image",
|
query: "str | bytes | PILImage.Image",
|
||||||
limit: int | None = None,
|
limit: int | None = None,
|
||||||
search_type: str = "hybrid",
|
search_type: SearchType = "hybrid",
|
||||||
filter: str | None = None,
|
filter: str | None = None,
|
||||||
include_images: bool = True,
|
include_images: bool = True,
|
||||||
) -> list[SearchResult]:
|
) -> list[SearchResult]:
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue