default sqlite when no path is given
This commit is contained in:
parent
186d404d11
commit
ac5c60fecc
6 changed files with 56 additions and 19 deletions
|
|
@ -7,6 +7,7 @@ from urllib.parse import urlparse
|
|||
|
||||
import httpx
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.reader import FileReader
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
|
|
@ -18,8 +19,15 @@ from haiku.rag.store.repositories.document import DocumentRepository
|
|||
class HaikuRAG:
|
||||
"""High-level haiku-rag client."""
|
||||
|
||||
def __init__(self, db_path: Path | Literal[":memory:"]):
|
||||
def __init__(
|
||||
self,
|
||||
db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
|
||||
/ "haiku.rag.sqlite",
|
||||
):
|
||||
"""Initialize the RAG client with a database path."""
|
||||
if isinstance(db_path, Path):
|
||||
if not db_path.parent.exists():
|
||||
Path.mkdir(db_path.parent, parents=True)
|
||||
self.store = Store(db_path)
|
||||
self.document_repository = DocumentRepository(self.store)
|
||||
self.chunk_repository = ChunkRepository(self.store)
|
||||
|
|
|
|||
|
|
@ -1,22 +1,28 @@
|
|||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from pydantic import BaseModel
|
||||
|
||||
from haiku.rag.utils import get_default_data_dir
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class AppConfig(BaseModel):
|
||||
ENV: str = "development"
|
||||
|
||||
OLLAMA_BASE_URL: str = "http://localhost:11434"
|
||||
DEFAULT_DATA_DIR: Path = get_default_data_dir()
|
||||
|
||||
EMBEDDING_PROVIDER: str = "ollama"
|
||||
EMBEDDING_MODEL: str = "mxbai-embed-large"
|
||||
EMBEDDING_VECTOR_DIM: int = 1024
|
||||
|
||||
CHUNK_SIZE: int = 256
|
||||
CHUNK_OVERLAP: int = 32
|
||||
|
||||
OLLAMA_BASE_URL: str = "http://localhost:11434"
|
||||
|
||||
|
||||
# Expose Config object for app to import
|
||||
Config = AppConfig.model_validate(os.environ)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from haiku.rag.config import Config
|
||||
from haiku.rag.embeddings.base import EmbedderBase
|
||||
from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder
|
||||
from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
|
||||
|
||||
|
||||
def get_embedder() -> EmbedderBase:
|
||||
|
|
@ -14,7 +13,7 @@ def get_embedder() -> EmbedderBase:
|
|||
|
||||
if Config.EMBEDDING_PROVIDER == "voyageai":
|
||||
try:
|
||||
import voyageai
|
||||
from haiku.rag.embeddings.voyageai import Embedder as VoyageAIEmbedder
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"VoyageAI embedder requires the 'voyageai' package. "
|
||||
|
|
|
|||
|
|
@ -1,14 +1,17 @@
|
|||
from voyageai.client import Client
|
||||
try:
|
||||
from voyageai.client import Client # type: ignore
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.embeddings.base import EmbedderBase
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.embeddings.base import EmbedderBase
|
||||
|
||||
class Embedder(EmbedderBase):
|
||||
_model: str = Config.EMBEDDING_MODEL
|
||||
_vector_dim: int = 1024
|
||||
|
||||
class Embedder(EmbedderBase):
|
||||
_model: str = Config.EMBEDDING_MODEL
|
||||
_vector_dim: int = 1024
|
||||
async def embed(self, text: str) -> list[float]:
|
||||
client = Client()
|
||||
res = client.embed([text], model=self._model, output_dtype="float")
|
||||
return res.embeddings[0] # type: ignore[return-value]
|
||||
|
||||
async def embed(self, text: str) -> list[float]:
|
||||
client = Client()
|
||||
res = client.embed([text], model=self._model, output_dtype="float")
|
||||
return res.embeddings[0] # type: ignore[return-value]
|
||||
except ImportError:
|
||||
pass
|
||||
|
|
|
|||
25
src/haiku/rag/utils.py
Normal file
25
src/haiku/rag/utils.py
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def get_default_data_dir() -> Path:
|
||||
"""
|
||||
Get the user data directory for the current system platform.
|
||||
|
||||
Linux: ~/.local/share/haiku.rag
|
||||
macOS: ~/Library/Application Support/haiku.rag
|
||||
Windows: C:/Users/<USER>/AppData/Roaming/haiku.rag
|
||||
|
||||
:return: User Data Path
|
||||
:rtype: Path
|
||||
"""
|
||||
home = Path.home()
|
||||
|
||||
system_paths = {
|
||||
"win32": home / "AppData/Roaming/haiku.rag",
|
||||
"linux": home / ".local/share/haiku.rag",
|
||||
"darwin": home / "Library/Application Support/haiku.rag",
|
||||
}
|
||||
|
||||
data_path = system_paths[sys.platform]
|
||||
return data_path
|
||||
6
uv.lock
6
uv.lock
|
|
@ -480,9 +480,6 @@ dependencies = [
|
|||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
all = [
|
||||
{ name = "voyageai" },
|
||||
]
|
||||
voyageai = [
|
||||
{ name = "voyageai" },
|
||||
]
|
||||
|
|
@ -507,11 +504,10 @@ requires-dist = [
|
|||
{ name = "python-dotenv", specifier = ">=1.1.0" },
|
||||
{ name = "sqlite-vec", specifier = ">=0.1.6" },
|
||||
{ name = "tiktoken", specifier = ">=0.9.0" },
|
||||
{ name = "voyageai", marker = "extra == 'all'", specifier = ">=0.3.2" },
|
||||
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" },
|
||||
{ name = "watchfiles", specifier = ">=1.1.0" },
|
||||
]
|
||||
provides-extras = ["voyageai", "all"]
|
||||
provides-extras = ["voyageai"]
|
||||
|
||||
[package.metadata.requires-dev]
|
||||
dev = [
|
||||
|
|
|
|||
Loading…
Reference in a new issue