From 02ec9c1759e509e8d870f56a2aa69857611ed44d Mon Sep 17 00:00:00 2001 From: Pushkinmazila2 Date: Mon, 4 May 2026 23:27:31 +0300 Subject: [PATCH 1/3] add api for local LLM --- src/obsidian_rag/cli.py | 76 +++++++++++++++++++++++++++++++------ src/obsidian_rag/config.py | 21 ++++++++++ src/obsidian_rag/indexer.py | 21 ++++++++-- src/obsidian_rag/server.py | 8 +++- src/obsidian_rag/watcher.py | 15 +++++++- 5 files changed, 121 insertions(+), 20 deletions(-) diff --git a/src/obsidian_rag/cli.py b/src/obsidian_rag/cli.py index bcb5b3a..d2f2781 100644 --- a/src/obsidian_rag/cli.py +++ b/src/obsidian_rag/cli.py @@ -37,9 +37,13 @@ help="Ollama API URL (only used with --provider ollama)") @click.option("--lmstudio-url", default=None, help="LM Studio API URL (only used with --provider lmstudio)") +@click.option("--ollama-api-key", default=None, + help="Bearer token for Ollama (only used with --provider ollama)") +@click.option("--lmstudio-api-key", default=None, + help="Bearer token for LM Studio (only used with --provider lmstudio)") @click.option("--model", default=None, help="Override embedding model name") @click.pass_context -def main(ctx, vault, data, provider, ollama_url, lmstudio_url, model): +def main(ctx, vault, data, provider, ollama_url, lmstudio_url, ollama_api_key, lmstudio_api_key, model): """Obsidian RAG - Semantic search for your Obsidian vault.""" ctx.ensure_object(dict) @@ -51,6 +55,8 @@ def main(ctx, vault, data, provider, ollama_url, lmstudio_url, model): ctx.obj["provider"] = provider or config.provider ctx.obj["ollama_url"] = ollama_url or config.ollama_url ctx.obj["lmstudio_url"] = lmstudio_url or config.lmstudio_url + ctx.obj["ollama_api_key"] = ollama_api_key or config.get_ollama_api_key() + ctx.obj["lmstudio_api_key"] = lmstudio_api_key or config.get_lmstudio_api_key() ctx.obj["model"] = model # None means use provider default ctx.obj["config"] = config @@ -105,6 +111,10 @@ def setup(): ) config.ollama_url = ollama_url + # Optional Bearer token + if click.confirm("\nDoes your Ollama instance require a Bearer token?", default=False): + config.ollama_api_key = click.prompt("Enter Bearer token", hide_input=True) + # Verify connection and get available models click.echo("Checking Ollama server...", nl=False) server_running = is_ollama_running(ollama_url) @@ -149,6 +159,10 @@ def setup(): ) config.lmstudio_url = lmstudio_url + # Optional Bearer token + if click.confirm("\nDoes your LM Studio instance require a Bearer token?", default=False): + config.lmstudio_api_key = click.prompt("Enter Bearer token", hide_input=True) + # Verify connection and get available models click.echo("Checking LM Studio server...", nl=False) server_running = is_lmstudio_running(lmstudio_url) @@ -226,13 +240,15 @@ def setup(): embedder = create_embedder( provider="ollama", model=config.ollama_model, - base_url=config.ollama_url + base_url=config.ollama_url, + api_key=config.get_ollama_api_key(), ) else: # lmstudio embedder = create_embedder( provider="lmstudio", model=config.lmstudio_model, - base_url=config.lmstudio_url + base_url=config.lmstudio_url, + api_key=config.get_lmstudio_api_key(), ) store = VectorStore(data_path=config.get_data_path()) @@ -291,7 +307,9 @@ def setup(): config.data_path or str(get_data_dir()), config.provider, config.ollama_url, - None # model + None, # model + config.get_ollama_api_key(), + config.get_lmstudio_api_key(), ) plist_path.write_text(plist_content) @@ -327,6 +345,8 @@ def index(ctx, clear, path_filter): provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] lmstudio_url = ctx.obj["lmstudio_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] config = ctx.obj["config"] # Get model from CLI override or config file based on provider @@ -344,16 +364,19 @@ def index(ctx, clear, path_filter): click.echo(f"Provider: {provider}") click.echo(f"Model: {model}") - # Determine the correct base_url based on provider + # Determine the correct base_url and api_key based on provider if provider == "ollama": base_url = ollama_url + api_key = ollama_api_key elif provider == "lmstudio": base_url = lmstudio_url + api_key = lmstudio_api_key else: base_url = None + api_key = None # Initialize components - embedder = create_embedder(provider=provider, model=model, base_url=base_url) + embedder = create_embedder(provider=provider, model=model, base_url=base_url, api_key=api_key) store = VectorStore(data_path=data_path) indexer = VaultIndexer(vault_path=vault_path, embedder=embedder, config=config.indexer) @@ -413,6 +436,8 @@ def search(ctx, query, limit, note_type): provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] lmstudio_url = ctx.obj["lmstudio_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] config = ctx.obj["config"] # Get model from CLI override or config file based on provider @@ -425,16 +450,19 @@ def search(ctx, query, limit, note_type): elif provider == "lmstudio": model = config.lmstudio_model - # Determine the correct base_url based on provider + # Determine the correct base_url and api_key based on provider if provider == "ollama": base_url = ollama_url + api_key = ollama_api_key elif provider == "lmstudio": base_url = lmstudio_url + api_key = lmstudio_api_key else: base_url = None + api_key = None # Initialize components - embedder = create_embedder(provider=provider, model=model, base_url=base_url) + embedder = create_embedder(provider=provider, model=model, base_url=base_url, api_key=api_key) store = VectorStore(data_path=data_path) # Generate query embedding @@ -486,6 +514,8 @@ def similar(ctx, note_path, limit): provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] lmstudio_url = ctx.obj["lmstudio_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] config = ctx.obj["config"] model = ctx.obj["model"] @@ -499,12 +529,15 @@ def similar(ctx, note_path, limit): if provider == "ollama": base_url = ollama_url + api_key = ollama_api_key elif provider == "lmstudio": base_url = lmstudio_url + api_key = lmstudio_api_key else: base_url = None + api_key = None - embedder = create_embedder(provider=provider, model=model, base_url=base_url) + embedder = create_embedder(provider=provider, model=model, base_url=base_url, api_key=api_key) store = VectorStore(data_path=data_path) click.echo(f"Finding notes similar to: {note_path}\n") @@ -552,6 +585,8 @@ def context(ctx, note_path, limit): provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] lmstudio_url = ctx.obj["lmstudio_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] config = ctx.obj["config"] model = ctx.obj["model"] @@ -565,12 +600,15 @@ def context(ctx, note_path, limit): if provider == "ollama": base_url = ollama_url + api_key = ollama_api_key elif provider == "lmstudio": base_url = lmstudio_url + api_key = lmstudio_api_key else: base_url = None + api_key = None - embedder = create_embedder(provider=provider, model=model, base_url=base_url) + embedder = create_embedder(provider=provider, model=model, base_url=base_url, api_key=api_key) store = VectorStore(data_path=data_path) click.echo(f"Getting context for: {note_path}\n") @@ -634,6 +672,8 @@ def watch(ctx, debounce): provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] lmstudio_url = ctx.obj["lmstudio_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] model = ctx.obj["model"] click.echo(f"Watching vault: {vault_path}") @@ -648,6 +688,8 @@ def watch(ctx, debounce): provider=provider, ollama_url=ollama_url, lmstudio_url=lmstudio_url, + ollama_api_key=ollama_api_key, + lmstudio_api_key=lmstudio_api_key, model=model, debounce_delay=debounce, ) @@ -698,7 +740,7 @@ def _uninstall_wrapper_script(): wrapper_path.unlink() -def _get_plist_content(vault_path: str, data_path: str, provider: str, ollama_url: str, model: str | None) -> str: +def _get_plist_content(vault_path: str, data_path: str, provider: str, ollama_url: str, model: str | None, ollama_api_key: str | None = None, lmstudio_api_key: str | None = None) -> str: """Generate launchd plist content.""" # Use wrapper script for better System Settings appearance wrapper_path = WRAPPER_SCRIPT_DIR / WRAPPER_SCRIPT_NAME @@ -715,6 +757,14 @@ def _get_plist_content(vault_path: str, data_path: str, provider: str, ollama_ur env_vars += f""" OBSIDIAN_RAG_OLLAMA_URL {ollama_url}""" + if ollama_api_key: + env_vars += f""" + OBSIDIAN_RAG_OLLAMA_API_KEY + {ollama_api_key}""" + elif provider == "lmstudio" and lmstudio_api_key: + env_vars += f""" + OBSIDIAN_RAG_LMSTUDIO_API_KEY + {lmstudio_api_key}""" if model: env_vars += f""" @@ -765,6 +815,8 @@ def install_service(ctx): data_path = ctx.obj["data"] provider = ctx.obj["provider"] ollama_url = ctx.obj["ollama_url"] + ollama_api_key = ctx.obj["ollama_api_key"] + lmstudio_api_key = ctx.obj["lmstudio_api_key"] model = ctx.obj["model"] plist_path = LAUNCH_AGENTS_DIR / PLIST_NAME @@ -783,7 +835,7 @@ def install_service(ctx): click.echo(f"Created: {wrapper_path}") # Write plist - plist_content = _get_plist_content(vault_path, data_path, provider, ollama_url, model) + plist_content = _get_plist_content(vault_path, data_path, provider, ollama_url, model, ollama_api_key, lmstudio_api_key) plist_path.write_text(plist_content) click.echo(f"Created: {plist_path}") diff --git a/src/obsidian_rag/config.py b/src/obsidian_rag/config.py index 1de8756..231865c 100644 --- a/src/obsidian_rag/config.py +++ b/src/obsidian_rag/config.py @@ -80,10 +80,12 @@ class Config: # Ollama settings ollama_url: str = "http://localhost:11434" ollama_model: str = "nomic-embed-text" + ollama_api_key: Optional[str] = None # Bearer token for protected Ollama instances # LM Studio settings lmstudio_url: str = "http://localhost:1234" lmstudio_model: str = "text-embedding-nomic-embed-text-v1.5" + lmstudio_api_key: Optional[str] = None # Bearer token for protected LM Studio instances # OpenAI model (optional override) openai_model: str = "text-embedding-3-small" @@ -97,6 +99,14 @@ def get_openai_api_key(self) -> Optional[str]: """Get OpenAI API key from config or environment.""" return self.openai_api_key or os.environ.get("OPENAI_API_KEY") + def get_ollama_api_key(self) -> Optional[str]: + """Get Ollama Bearer token from config or environment.""" + return self.ollama_api_key or os.environ.get("OBSIDIAN_RAG_OLLAMA_API_KEY") + + def get_lmstudio_api_key(self) -> Optional[str]: + """Get LM Studio Bearer token from config or environment.""" + return self.lmstudio_api_key or os.environ.get("OBSIDIAN_RAG_LMSTUDIO_API_KEY") + def load_config() -> Config: """Load configuration from file with environment variable overrides. @@ -130,11 +140,13 @@ def load_config() -> Config: if "ollama" in data: config.ollama_url = data["ollama"].get("url", config.ollama_url) config.ollama_model = data["ollama"].get("model", config.ollama_model) + config.ollama_api_key = data["ollama"].get("api_key", config.ollama_api_key) # LM Studio settings if "lmstudio" in data: config.lmstudio_url = data["lmstudio"].get("url", config.lmstudio_url) config.lmstudio_model = data["lmstudio"].get("model", config.lmstudio_model) + config.lmstudio_api_key = data["lmstudio"].get("api_key", config.lmstudio_api_key) # Indexer settings if "indexer" in data: @@ -154,6 +166,11 @@ def load_config() -> Config: config.ollama_url = os.environ["OBSIDIAN_RAG_OLLAMA_URL"] if os.environ.get("OBSIDIAN_RAG_LMSTUDIO_URL"): config.lmstudio_url = os.environ["OBSIDIAN_RAG_LMSTUDIO_URL"] + # Bearer token overrides for local providers + if os.environ.get("OBSIDIAN_RAG_OLLAMA_API_KEY"): + config.ollama_api_key = os.environ["OBSIDIAN_RAG_OLLAMA_API_KEY"] + if os.environ.get("OBSIDIAN_RAG_LMSTUDIO_API_KEY"): + config.lmstudio_api_key = os.environ["OBSIDIAN_RAG_LMSTUDIO_API_KEY"] if os.environ.get("OBSIDIAN_RAG_MODEL"): if config.provider == "ollama": config.ollama_model = os.environ["OBSIDIAN_RAG_MODEL"] @@ -200,6 +217,8 @@ def save_config(config: Config) -> Path: ollama_section["url"] = config.ollama_url if config.ollama_model != "nomic-embed-text": ollama_section["model"] = config.ollama_model + if config.ollama_api_key: + ollama_section["api_key"] = config.ollama_api_key if ollama_section: data["ollama"] = ollama_section @@ -210,6 +229,8 @@ def save_config(config: Config) -> Path: lmstudio_section["url"] = config.lmstudio_url if config.lmstudio_model != "text-embedding-nomic-embed-text-v1.5": lmstudio_section["model"] = config.lmstudio_model + if config.lmstudio_api_key: + lmstudio_section["api_key"] = config.lmstudio_api_key if lmstudio_section: data["lmstudio"] = lmstudio_section diff --git a/src/obsidian_rag/indexer.py b/src/obsidian_rag/indexer.py index a4c01f0..532a037 100644 --- a/src/obsidian_rag/indexer.py +++ b/src/obsidian_rag/indexer.py @@ -352,11 +352,15 @@ class OllamaEmbedder: def __init__( self, base_url: str = "http://localhost:11434", - model: str = "nomic-embed-text" + model: str = "nomic-embed-text", + api_key: Optional[str] = None, ): self.base_url = base_url self.model = model - self.client = httpx.Client(timeout=60.0) + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self.client = httpx.Client(timeout=60.0, headers=headers) def _get_prefix(self, task_type: str) -> str: """Get task-specific prefix for models that support them.""" @@ -393,11 +397,15 @@ class LMStudioEmbedder: def __init__( self, base_url: str = "http://localhost:1234", - model: str = "text-embedding-nomic-embed-text-v1.5" + model: str = "text-embedding-nomic-embed-text-v1.5", + api_key: Optional[str] = None, ): self.base_url = base_url.rstrip("/") self.model = model - self.client = httpx.Client(timeout=60.0) + headers = {} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self.client = httpx.Client(timeout=60.0, headers=headers) def _get_prefix(self, task_type: str) -> str: model = self.model.lower() @@ -500,6 +508,7 @@ def create_embedder( provider: str = "openai", model: Optional[str] = None, base_url: Optional[str] = None, + api_key: Optional[str] = None, ) -> Embedder: """Create an embedder instance for the specified provider.""" if provider == "openai": @@ -513,6 +522,8 @@ def create_embedder( kwargs["model"] = model if base_url: kwargs["base_url"] = base_url + if api_key: + kwargs["api_key"] = api_key return OllamaEmbedder(**kwargs) elif provider == "lmstudio": kwargs = {} @@ -520,6 +531,8 @@ def create_embedder( kwargs["model"] = model if base_url: kwargs["base_url"] = base_url + if api_key: + kwargs["api_key"] = api_key return LMStudioEmbedder(**kwargs) else: raise ValueError(f"Unknown provider: {provider}. Use 'openai', 'ollama', or 'lmstudio'.") diff --git a/src/obsidian_rag/server.py b/src/obsidian_rag/server.py index eb7c168..ee3494d 100644 --- a/src/obsidian_rag/server.py +++ b/src/obsidian_rag/server.py @@ -37,21 +37,25 @@ def get_embedder() -> Embedder: if config.provider == "openai" and config.openai_api_key: os.environ["OPENAI_API_KEY"] = config.openai_api_key - # Determine model and base_url based on provider + # Determine model, base_url and api_key based on provider if config.provider == "openai": model = config.openai_model base_url = None + api_key = config.get_openai_api_key() elif config.provider == "ollama": model = config.ollama_model base_url = config.ollama_url + api_key = config.get_ollama_api_key() else: # lmstudio model = config.lmstudio_model base_url = config.lmstudio_url - + api_key = config.get_lmstudio_api_key() + _embedder = create_embedder( provider=config.provider, model=model, base_url=base_url, + api_key=api_key, ) return _embedder diff --git a/src/obsidian_rag/watcher.py b/src/obsidian_rag/watcher.py index a44f9f8..1799cfc 100644 --- a/src/obsidian_rag/watcher.py +++ b/src/obsidian_rag/watcher.py @@ -44,6 +44,8 @@ DEFAULT_PROVIDER = _config.provider DEFAULT_OLLAMA_URL = _config.ollama_url DEFAULT_LMSTUDIO_URL = _config.lmstudio_url +DEFAULT_OLLAMA_API_KEY: Optional[str] = _config.get_ollama_api_key() +DEFAULT_LMSTUDIO_API_KEY: Optional[str] = _config.get_lmstudio_api_key() DEFAULT_MODEL: Optional[str] = None # Use provider default DEFAULT_DEBOUNCE = float(os.environ.get("OBSIDIAN_RAG_DEBOUNCE", "2.0")) @@ -326,6 +328,8 @@ def __init__( provider: str = DEFAULT_PROVIDER, ollama_url: str = DEFAULT_OLLAMA_URL, lmstudio_url: str = DEFAULT_LMSTUDIO_URL, + ollama_api_key: Optional[str] = DEFAULT_OLLAMA_API_KEY, + lmstudio_api_key: Optional[str] = DEFAULT_LMSTUDIO_API_KEY, model: Optional[str] = DEFAULT_MODEL, debounce_delay: float = DEFAULT_DEBOUNCE, ): @@ -337,9 +341,10 @@ def __init__( if provider == "openai" and _config.openai_api_key: os.environ["OPENAI_API_KEY"] = _config.openai_api_key - # Determine correct base_url based on provider + # Determine correct base_url and api_key based on provider if provider == "ollama": base_url = ollama_url + api_key = ollama_api_key # Health check for Ollama before starting if not check_ollama_health(ollama_url): logger.warning("Ollama is not running! Waiting for it to start...") @@ -347,10 +352,12 @@ def __init__( self._wait_for_ollama(ollama_url) elif provider == "lmstudio": base_url = lmstudio_url + api_key = lmstudio_api_key else: base_url = None + api_key = None - self.embedder = create_embedder(provider=provider, model=model, base_url=base_url) + self.embedder = create_embedder(provider=provider, model=model, base_url=base_url, api_key=api_key) self.store = VectorStore(data_path=data_path) self.debounce_delay = debounce_delay self.retry_queue = RetryQueue() @@ -508,6 +515,8 @@ def run_watcher( provider: str = DEFAULT_PROVIDER, ollama_url: str = DEFAULT_OLLAMA_URL, lmstudio_url: str = DEFAULT_LMSTUDIO_URL, + ollama_api_key: Optional[str] = DEFAULT_OLLAMA_API_KEY, + lmstudio_api_key: Optional[str] = DEFAULT_LMSTUDIO_API_KEY, model: Optional[str] = DEFAULT_MODEL, debounce: float = DEFAULT_DEBOUNCE, ): @@ -524,6 +533,8 @@ def run_watcher( provider=provider, ollama_url=ollama_url, lmstudio_url=lmstudio_url, + ollama_api_key=ollama_api_key, + lmstudio_api_key=lmstudio_api_key, model=model, debounce_delay=debounce, ) From e229e613f7414cb10bf30a2f92737a9b73409970 Mon Sep 17 00:00:00 2001 From: Pushkinmazila2 Date: Mon, 4 May 2026 23:55:13 +0300 Subject: [PATCH 2/3] Update cli.py fix eddit token --- src/obsidian_rag/cli.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/obsidian_rag/cli.py b/src/obsidian_rag/cli.py index d2f2781..caa2ec4 100644 --- a/src/obsidian_rag/cli.py +++ b/src/obsidian_rag/cli.py @@ -113,7 +113,7 @@ def setup(): # Optional Bearer token if click.confirm("\nDoes your Ollama instance require a Bearer token?", default=False): - config.ollama_api_key = click.prompt("Enter Bearer token", hide_input=True) + config.ollama_api_key = click.prompt("Enter Bearer token", hide_input=False) # Verify connection and get available models click.echo("Checking Ollama server...", nl=False) @@ -161,7 +161,7 @@ def setup(): # Optional Bearer token if click.confirm("\nDoes your LM Studio instance require a Bearer token?", default=False): - config.lmstudio_api_key = click.prompt("Enter Bearer token", hide_input=True) + config.lmstudio_api_key = click.prompt("Enter Bearer token", hide_input=False) # Verify connection and get available models click.echo("Checking LM Studio server...", nl=False) From 922d2e5e3007a6185ca2d4ee98aec7e17b16b6ed Mon Sep 17 00:00:00 2001 From: Pushkinmazila2 Date: Tue, 5 May 2026 00:14:04 +0300 Subject: [PATCH 3/3] fix api_key=config.ollama_api_key + hide_input --- src/obsidian_rag/cli.py | 8 ++++---- src/obsidian_rag/indexer.py | 24 ++++++++++++++---------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/src/obsidian_rag/cli.py b/src/obsidian_rag/cli.py index caa2ec4..1723fdc 100644 --- a/src/obsidian_rag/cli.py +++ b/src/obsidian_rag/cli.py @@ -117,14 +117,14 @@ def setup(): # Verify connection and get available models click.echo("Checking Ollama server...", nl=False) - server_running = is_ollama_running(ollama_url) + server_running = is_ollama_running(ollama_url, api_key=config.ollama_api_key) if server_running: click.echo(" ✓ connected") # Fetch available embedding models click.echo("Fetching available embedding models...", nl=False) - available_models = get_ollama_models(ollama_url) + available_models = get_ollama_models(ollama_url, api_key=config.ollama_api_key) if available_models: click.echo(f" found {len(available_models)}") @@ -165,14 +165,14 @@ def setup(): # Verify connection and get available models click.echo("Checking LM Studio server...", nl=False) - server_running = is_lmstudio_running(lmstudio_url) + server_running = is_lmstudio_running(lmstudio_url, api_key=config.lmstudio_api_key) if server_running: click.echo(" ✓ connected") # Fetch available embedding models click.echo("Fetching available embedding models...", nl=False) - available_models = get_lmstudio_models(lmstudio_url) + available_models = get_lmstudio_models(lmstudio_url, api_key=config.lmstudio_api_key) if available_models: click.echo(f" found {len(available_models)}") diff --git a/src/obsidian_rag/indexer.py b/src/obsidian_rag/indexer.py index 532a037..1783d8e 100644 --- a/src/obsidian_rag/indexer.py +++ b/src/obsidian_rag/indexer.py @@ -443,31 +443,34 @@ def close(self): self.client.close() -def is_lmstudio_running(base_url: str = "http://localhost:1234") -> bool: +def is_lmstudio_running(base_url: str = "http://localhost:1234", api_key: Optional[str] = None) -> bool: """Check if LM Studio server is running.""" + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: - with httpx.Client(timeout=2.0) as client: + with httpx.Client(timeout=2.0, headers=headers) as client: response = client.get(f"{base_url.rstrip('/')}/v1/models") - return response.status_code == 200 + return response.status_code in (200, 401) # 401 = server up but wrong key except (httpx.RequestError, httpx.TimeoutException): return False -def is_ollama_running(base_url: str = "http://localhost:11434") -> bool: +def is_ollama_running(base_url: str = "http://localhost:11434", api_key: Optional[str] = None) -> bool: """Check if Ollama server is running.""" + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: - with httpx.Client(timeout=2.0) as client: + with httpx.Client(timeout=2.0, headers=headers) as client: response = client.get(f"{base_url.rstrip('/')}/api/tags") - return response.status_code == 200 + return response.status_code in (200, 401) except (httpx.RequestError, httpx.TimeoutException): return False -def get_lmstudio_models(base_url: str = "http://localhost:1234") -> List[str]: +def get_lmstudio_models(base_url: str = "http://localhost:1234", api_key: Optional[str] = None) -> List[str]: """Get list of available embedding models from LM Studio.""" embedding_keywords = ['embed', 'bge', 'minilm', 'e5', 'gte', 'instructor'] + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: - with httpx.Client(timeout=5.0) as client: + with httpx.Client(timeout=5.0, headers=headers) as client: response = client.get(f"{base_url.rstrip('/')}/v1/models") if response.status_code != 200: return [] @@ -482,11 +485,12 @@ def get_lmstudio_models(base_url: str = "http://localhost:1234") -> List[str]: return [] -def get_ollama_models(base_url: str = "http://localhost:11434") -> List[str]: +def get_ollama_models(base_url: str = "http://localhost:11434", api_key: Optional[str] = None) -> List[str]: """Get list of available embedding models from Ollama.""" embedding_keywords = ['embed', 'bge', 'minilm', 'e5', 'gte', 'instructor', 'nomic'] + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: - with httpx.Client(timeout=5.0) as client: + with httpx.Client(timeout=5.0, headers=headers) as client: response = client.get(f"{base_url.rstrip('/')}/api/tags") if response.status_code != 200: return []