jerpint commited on
Commit
5d88935
·
1 Parent(s): 062729f

move imports inline so that they're not required for running the app

Browse files
Files changed (1) hide show
  1. llms.py +24 -29
llms.py CHANGED
@@ -1,33 +1,4 @@
1
  import os
2
- import logging
3
-
4
- import google.generativeai as genai
5
- from openai import OpenAI
6
- import anthropic
7
-
8
-
9
- DEFAULT_PROVIDER = "openai"
10
- DEFAULT_MODEL = "gpt-4o-mini"
11
-
12
- # Load API keys
13
- OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
14
- AI_STUDIO_API_KEY = os.getenv("AI_STUDIO_API_KEY")
15
- ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
16
-
17
- assert (
18
- OPENAI_API_KEY
19
- ), "OPENAI_API_KEY is not set, please set it in your environment variables."
20
- assert (
21
- AI_STUDIO_API_KEY
22
- ), "AI_STUDIO_API_KEY is not set, please set it in your environment variables."
23
- assert (
24
- ANTHROPIC_API_KEY
25
- ), "ANTHROPIC_API_KEY is not set, please set it in your environment variables."
26
-
27
- # Initialize clients
28
- openai_client = OpenAI(api_key=OPENAI_API_KEY)
29
- genai.configure(api_key=AI_STUDIO_API_KEY)
30
- anthropic_client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
31
 
32
 
33
  SYSTEM_PROMPT = "You are a programming assistant. You are solving the 2024 advent of code challenge."
@@ -87,6 +58,12 @@ def get_completion(
87
  Unified function to get completions from various LLM providers.
88
  """
89
  if provider == "openai":
 
 
 
 
 
 
90
  completion = openai_client.chat.completions.create(
91
  model=model,
92
  messages=[
@@ -99,6 +76,15 @@ def get_completion(
99
  return completion.choices[0].message.content
100
 
101
  elif provider == "gemini":
 
 
 
 
 
 
 
 
 
102
  model = genai.GenerativeModel(
103
  model_name=model,
104
  system_instruction=system_prompt,
@@ -112,6 +98,15 @@ def get_completion(
112
  return response.text
113
 
114
  elif provider == "anthropic":
 
 
 
 
 
 
 
 
 
115
  response = anthropic_client.messages.create(
116
  model=model,
117
  max_tokens=2048,
 
1
  import os
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
 
4
  SYSTEM_PROMPT = "You are a programming assistant. You are solving the 2024 advent of code challenge."
 
58
  Unified function to get completions from various LLM providers.
59
  """
60
  if provider == "openai":
61
+ from openai import OpenAI
62
+ OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
63
+ assert (
64
+ OPENAI_API_KEY
65
+ ), "OPENAI_API_KEY is not set, please set it in your environment variables."
66
+ openai_client = OpenAI(api_key=OPENAI_API_KEY)
67
  completion = openai_client.chat.completions.create(
68
  model=model,
69
  messages=[
 
76
  return completion.choices[0].message.content
77
 
78
  elif provider == "gemini":
79
+
80
+ # Setup
81
+ import google.generativeai as genai
82
+ AI_STUDIO_API_KEY = os.getenv("AI_STUDIO_API_KEY")
83
+ assert (
84
+ AI_STUDIO_API_KEY
85
+ ), "AI_STUDIO_API_KEY is not set, please set it in your environment variables."
86
+ genai.configure(api_key=AI_STUDIO_API_KEY)
87
+
88
  model = genai.GenerativeModel(
89
  model_name=model,
90
  system_instruction=system_prompt,
 
98
  return response.text
99
 
100
  elif provider == "anthropic":
101
+
102
+ # Setup
103
+ import anthropic
104
+ ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
105
+ assert (
106
+ ANTHROPIC_API_KEY
107
+ ), "ANTHROPIC_API_KEY is not set, please set it in your environment variables."
108
+ anthropic_client = anthropic.Anthropic(api_key=ANTHROPIC_API_KEY)
109
+
110
  response = anthropic_client.messages.create(
111
  model=model,
112
  max_tokens=2048,