Coverage for open_webui/utils/embeddings.py: 51%

1import logging 

2import random 

3import sys 

4 

5from fastapi import Request 

6from open_webui.env import BYPASS_MODEL_ACCESS_CONTROL, GLOBAL_LOG_LEVEL 

7from open_webui.models.models import Models 

8from open_webui.models.users import UserModel 

9from open_webui.routers.ollama import ( 

10 GenerateEmbedForm, 

11) 

12from open_webui.routers.ollama import ( 

13 embed as ollama_embed, 

14) 

15from open_webui.routers.openai import embeddings as openai_embeddings 

16from open_webui.utils.models import check_model_access 

17from open_webui.utils.payload import convert_embed_payload_openai_to_ollama 

18from open_webui.utils.response import convert_embedding_response_ollama_to_openai 

19 

20logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL) 

21log = logging.getLogger(__name__) 

22 

23 

24async def generate_embeddings( 

25 request: Request, 

26 form_data: dict, 

27 user: UserModel, 

28 bypass_filter: bool = False, 

29): 

30 """ 

31 Dispatch and handle embeddings generation based on the model type (OpenAI, Ollama). 

32 

33 Args: 

34 request (Request): The FastAPI request context. 

35 form_data (dict): The input data sent to the endpoint. 

36 user (UserModel): The authenticated user. 

37 bypass_filter (bool): If True, disables access filtering (default False). 

38 

39 Returns: 

40 dict: The embeddings response, following OpenAI API compatibility. 

41 """ 

42 if BYPASS_MODEL_ACCESS_CONTROL: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true

43 bypass_filter = True 

44 

45 # Attach extra metadata from request.state if present 

46 if hasattr(request.state, 'metadata'): 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true

47 if 'metadata' not in form_data: 

48 form_data['metadata'] = request.state.metadata 

49 else: 

50 form_data['metadata'] = { 

51 **form_data['metadata'], 

52 **request.state.metadata, 

53 } 

54 

55 # If "direct" flag present, use only that model 

56 if getattr(request.state, 'direct', False) and hasattr(request.state, 'model'): 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 models = { 

58 request.state.model['id']: request.state.model, 

59 } 

60 else: 

61 models = request.app.state.MODELS 

62 

63 model_id = form_data.get('model') 

64 if model_id not in models: 64 ↛ 66line 64 didn't jump to line 66 because the condition on line 64 was always true

65 raise Exception('Model not found') 

66 model = models[model_id] 

67 

68 # Access filtering 

69 if not getattr(request.state, 'direct', False): 

70 if not bypass_filter and user.role == 'user': 

71 await check_model_access(user, model) 

72 

73 # Ollama backend — use /api/embed which supports batch input natively 

74 if model.get('owned_by') == 'ollama': 

75 ollama_payload = convert_embed_payload_openai_to_ollama(form_data) 

76 response = await ollama_embed( 

77 request=request, 

78 form_data=GenerateEmbedForm(**ollama_payload), 

79 user=user, 

80 ) 

81 return convert_embedding_response_ollama_to_openai(response) 

82 

83 # Default: OpenAI or compatible backend 

84 return await openai_embeddings( 

85 request=request, 

86 form_data=form_data, 

87 user=user, 

88 )