From fc87c12118c8cca7809e310e3f537b5ff34d1da2 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 14 Oct 2024 21:28:18 +0100 Subject: [PATCH] Fix Ollama host --- .../trustgraph/model/text_completion/ollama/llm.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index b506b3cd..17151a00 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -6,6 +6,7 @@ Input is prompt, output is response. from ollama import Client from prometheus_client import Histogram, Info +import os from .... schema import TextCompletionRequest, TextCompletionResponse, Error from .... schema import text_completion_request_queue @@ -19,8 +20,8 @@ module = ".".join(__name__.split(".")[1:-1]) default_input_queue = text_completion_request_queue default_output_queue = text_completion_response_queue default_subscriber = module -default_model = 'gemma2' -default_ollama = 'http://localhost:11434' +default_model = 'gemma2:9b' +default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434') class Processor(ConsumerProducer): @@ -152,7 +153,7 @@ class Processor(ConsumerProducer): parser.add_argument( '-m', '--model', default="gemma2", - help=f'LLM model (default: gemma2)' + help=f'LLM model (default: {default_model})' ) parser.add_argument(