Portkey Integration

This commit is contained in:
Josako
2024-07-18 07:08:42 +02:00
parent 908a2eaf7e
commit 88ca04136d
6 changed files with 39 additions and 8 deletions

View File

@@ -8,6 +8,7 @@ import ast
from typing import List
from openai import OpenAI
# from groq import Groq
from portkey_ai import createHeaders, PORTKEY_GATEWAY_URL
from common.models.document import EmbeddingSmallOpenAI, EmbeddingLargeOpenAI
@@ -91,16 +92,26 @@ def select_model_variables(tenant):
# Set Embedding variables
match embedding_provider:
case 'openai':
portkey_metadata = {'tenant_id': str(tenant.id)}
portkey_headers = createHeaders(api_key=current_app.config.get('PORTKEY_API_KEY'),
provider='openai',
metadata=portkey_metadata)
match embedding_model:
case 'text-embedding-3-small':
api_key = current_app.config.get('OPENAI_API_KEY')
model_variables['embedding_model'] = OpenAIEmbeddings(api_key=api_key,
model='text-embedding-3-small')
model='text-embedding-3-small',
base_url=PORTKEY_GATEWAY_URL,
default_headers=portkey_headers
)
model_variables['embedding_db_model'] = EmbeddingSmallOpenAI
case 'text-embedding-3-large':
api_key = current_app.config.get('OPENAI_API_KEY')
model_variables['embedding_model'] = OpenAIEmbeddings(api_key=api_key,
model='text-embedding-3-large')
model='text-embedding-3-large',
base_url=PORTKEY_GATEWAY_URL,
default_headers=portkey_headers
)
model_variables['embedding_db_model'] = EmbeddingLargeOpenAI
case _:
raise Exception(f'Error setting model variables for tenant {tenant.id} '
@@ -112,13 +123,22 @@ def select_model_variables(tenant):
# Set Chat model variables
match llm_provider:
case 'openai':
portkey_metadata = {'tenant_id': str(tenant.id)}
portkey_headers = createHeaders(api_key=current_app.config.get('PORTKEY_API_KEY'),
metadata=portkey_metadata,
provider='openai')
tool_calling_supported = False
api_key = current_app.config.get('OPENAI_API_KEY')
model_variables['llm'] = ChatOpenAI(api_key=api_key,
model=llm_model,
temperature=model_variables['RAG_temperature'])
temperature=model_variables['RAG_temperature'],
base_url=PORTKEY_GATEWAY_URL,
default_headers=portkey_headers)
model_variables['llm_no_rag'] = ChatOpenAI(api_key=api_key,
model=llm_model,
temperature=model_variables['no_RAG_temperature'])
temperature=model_variables['no_RAG_temperature'],
base_url=PORTKEY_GATEWAY_URL,
default_headers=portkey_headers)
tool_calling_supported = False
match llm_model:
case 'gpt-4-turbo' | 'gpt-4o':
@@ -161,9 +181,16 @@ def select_model_variables(tenant):
# model_variables['transcription_client'] = Groq(api_key=api_key)
# model_variables['transcription_model'] = 'whisper-large-v3'
# Using OpenAI
# Using OpenAI for transcriptions
portkey_metadata = {'tenant_id': str(tenant.id)}
portkey_headers = createHeaders(api_key=current_app.config.get('PORTKEY_API_KEY'),
metadata=portkey_metadata,
provider='openai'
)
api_key = current_app.config.get('OPENAI_API_KEY')
model_variables['transcription_client'] = OpenAI(api_key=api_key)
model_variables['transcription_client'] = OpenAI(api_key=api_key,
base_url=PORTKEY_GATEWAY_URL,
default_headers=portkey_headers)
model_variables['transcription_model'] = 'whisper-1'
return model_variables