Refakrotiserat så att den externa funktionen fetch_models_from_endpoints() gjorts om till klassmetoden fetch_models()

This commit is contained in:
2024-08-04 11:48:38 +02:00
parent 9e2cebd6cc
commit 8e983919e5
+39 -40
View File
@@ -4,45 +4,6 @@
import logging import logging
import json import json
import requests import requests
#from backend import GlobalState # Assuming GlobalState is defined there
def fetch_models_from_endpoints(endpoints, global_state):
"""
Fetch models from endpoints and update the endpoint dictionaries.
Args:
endpoints (list): List of endpoint dictionaries.
global_state: The global state object.
Returns:
None
"""
logger = logging.getLogger(__name__)
for endpoint in endpoints:
if endpoint["provider"] == "ollama":
headers = {
"Content-Type": "application/json",
}
if "requestOptions" in endpoint: # Check if authentication is needed
headers.update({
"Authorization": endpoint["requestOptions"]["headers"]["Authorization"]
})
try:
models_response = requests.get(endpoint["url"] + "/api/tags", headers=headers)
models_response.raise_for_status() # Raise an exception for HTTP errors
models = models_response.json()
except requests.exceptions.RequestException as e:
logger.error("Error fetching models from backend: %s", str(e))
continue
if isinstance(models, dict) and 'error' in models:
logger.error('Error fetching models from backend: %s', models['error'])
else:
endpoint["models"] = models.get("models", []) # Get the list of models directly
# logger.debug("models = \n{}".format(json.dumps(models, indent=4)))
class GlobalState: class GlobalState:
""" """
@@ -140,9 +101,47 @@ class GlobalState:
self.endpoints = endpoints self.endpoints = endpoints
def get_endpoints(self): def get_endpoints(self):
"""Return the list of endpoints""" """
Get the list of endpoints.
Returns:
List of endpoints
"""
return self.endpoints return self.endpoints
def fetch_models(self):
"""
Fetch models from endpoints and update the endpoint dictionaries.
Returns:
None
"""
logger = logging.getLogger(__name__)
for endpoint in self.endpoints:
if endpoint["provider"] == "ollama":
headers = {
"Content-Type": "application/json",
}
if "requestOptions" in endpoint: # Check if authentication is needed
headers.update({
"Authorization": endpoint["requestOptions"]["headers"]["Authorization"]
})
try:
models_response = requests.get(endpoint["url"] + "/api/tags", headers=headers)
models_response.raise_for_status() # Raise an exception for HTTP errors
models = models_response.json()
except requests.exceptions.RequestException as e:
logger.error("Error fetching models from backend: %s", str(e))
continue
if isinstance(models, dict) and 'error' in models:
logger.error('Error fetching models from backend: %s', models['error'])
else:
endpoint["models"] = models.get("models", []) # Get the list of models directly
def get_list_of_available_llms(self, endpoint=None): def get_list_of_available_llms(self, endpoint=None):
"""Return a sorted list of LLMs available at endpoint""" """Return a sorted list of LLMs available at endpoint"""
llm_list = None llm_list = None