apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 17:21:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix Ollama host

Browse source
This commit is contained in:
Cyber MacGeddon 2024-10-14 21:28:18 +01:00
parent 0eee19b540
commit fc87c12118
1 changed files with 4 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

7
trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py
View file

@ -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(