SystemSh0cker/private-gpt
1
0
Fork 0
mirror of https://github.com/zylon-ai/private-gpt.git synced 2025-12-22 20:12:55 +01:00
Code Issues Projects Releases Packages Activity

feat: allow to autopull ollama models

Browse source
This commit is contained in:
Javier Martinez 2024-07-29 09:25:42 +02:00
parent d34b7d22f4
commit 0a9c57447b
No known key found for this signature in database
4 changed files with 56 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

25
private_gpt/components/llm/llm_component.py
View file

@ -146,8 +146,15 @@ class LLMComponent:
"repeat_penalty": ollama_settings.repeat_penalty, # ollama llama-cpp
}
self.llm = Ollama(
model=ollama_settings.llm_model,
# calculate llm model. If not provided tag, it will be use latest model
model_name = (
ollama_settings.llm_model + ':latest'
if ":" not in ollama_settings.llm_model
else ollama_settings.llm_model
)
llm = Ollama(
model=model_name,
base_url=ollama_settings.api_base,
temperature=settings.llm.temperature,
context_window=settings.llm.context_window,
@ -155,6 +162,18 @@ class LLMComponent:
request_timeout=ollama_settings.request_timeout,
)
if ollama_settings.autopull_models:
try:
installed_models = [model['name'] for model in llm.client.list().get("models", {})]
if model_name not in installed_models:
logger.info(
f"Pulling model {model_name}. Please wait..."
)
llm.client.pull(model_name)
logger.info(f"Model {model_name} pulled successfully")
except Exception as e:
logger.error(f"Failed to pull model {model_name}: {e!s}")
if (
ollama_settings.keep_alive
!= ollama_settings.model_fields["keep_alive"].default
@ -172,6 +191,8 @@ class LLMComponent:
Ollama.complete = add_keep_alive(Ollama.complete)
Ollama.stream_complete = add_keep_alive(Ollama.stream_complete)
self.llm = llm
case "azopenai":
try:
from llama_index.llms.azure_openai import ( # type: ignore