diff --git a/modules/vector-sets/examples/cli-tool/cli.py b/modules/vector-sets/examples/cli-tool/cli.py index 7a2eeb4f9..614a02103 100755 --- a/modules/vector-sets/examples/cli-tool/cli.py +++ b/modules/vector-sets/examples/cli-tool/cli.py @@ -8,6 +8,7 @@ # #!/usr/bin/env python3 +import argparse import redis import requests import re @@ -15,9 +16,12 @@ import shlex from prompt_toolkit import PromptSession from prompt_toolkit.history import InMemoryHistory +# Default Ollama embeddings URL (can be overridden with --ollama-url) +OLLAMA_URL = "http://localhost:11434/api/embeddings" + def get_embedding(text): """Get embedding from local Ollama API""" - url = "http://localhost:11434/api/embeddings" + url = OLLAMA_URL payload = { "model": "mxbai-embed-large", "prompt": text @@ -73,6 +77,15 @@ def format_response(response): return str(response) def main(): + global OLLAMA_URL + + parser = argparse.ArgumentParser(prog="cli.py", add_help=False) + parser.add_argument("--ollama-url", dest="ollama_url", + help="Ollama embeddings API URL (default: {OLLAMA_URL})", + default=OLLAMA_URL) + args, _ = parser.parse_known_args() + OLLAMA_URL = args.ollama_url + # Default connection to localhost:6379 r = redis.Redis(host='localhost', port=6379, decode_responses=True)