Skip to content

Commit

Permalink
Add cloudflare gateway for GPT
Browse files Browse the repository at this point in the history
  • Loading branch information
kaarthik108 committed Dec 21, 2023
1 parent b1d1715 commit 5429b49
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions chain.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@ def __init__(self, config: ModelConfig):
self.model_type = config.model_type
self.secrets = config.secrets
self.callback_handler = config.callback_handler
account_tag = self.secrets['CF_ACCOUNT_TAG']
self.gateway_url = f"https://gateway.ai.cloudflare.com/v1/{account_tag}/k-1-gpt/openai"
self.setup()

def setup(self):
Expand All @@ -49,24 +51,26 @@ def setup(self):
def setup_gpt(self):
self.q_llm = OpenAI(
temperature=0.1,
openai_api_key=self.secrets["OPENAI_API_KEY"],
api_key=self.secrets["OPENAI_API_KEY"],
model_name="gpt-3.5-turbo-16k",
max_tokens=500,
base_url=self.gateway_url
)

self.llm = ChatOpenAI(
model_name="gpt-3.5-turbo-16k",
temperature=0.5,
openai_api_key=self.secrets["OPENAI_API_KEY"],
api_key=self.secrets["OPENAI_API_KEY"],
max_tokens=500,
callbacks=[self.callback_handler],
streaming=True,
base_url=self.gateway_url
)

def setup_mixtral(self):
self.q_llm = OpenAI(
temperature=0.1,
openai_api_key=self.secrets["MIXTRAL_API_KEY"],
api_key=self.secrets["MIXTRAL_API_KEY"],
model_name="mistralai/Mixtral-8x7B-Instruct-v0.1",
max_tokens=500,
base_url="https://api.together.xyz/v1",
Expand All @@ -75,7 +79,7 @@ def setup_mixtral(self):
self.llm = ChatOpenAI(
model_name="mistralai/Mixtral-8x7B-Instruct-v0.1",
temperature=0.5,
openai_api_key=self.secrets["MIXTRAL_API_KEY"],
api_key=self.secrets["MIXTRAL_API_KEY"],
max_tokens=500,
callbacks=[self.callback_handler],
streaming=True,
Expand Down

0 comments on commit 5429b49

Please sign in to comment.