Skip to content

Commit f374942

Browse files
kwlee2025cppkwlee2025cpp
authored
refactor: extract shared LLM utilities to llm_utils.py (#36)
Extract get_model_key_from_env, get_config_class, and related helpers from entrypoint.py into llm_utils.py to eliminate the circular import between entrypoint.py and prompt_pipeline/entrypoint.py. - Both consumers now import from llm_utils (no sys.path hacks) - Fix prompt_pipeline path resolution for CI (ai_tutor/ fallback) - Add llm_utils.py to Dockerfile and .dockerignore whitelist - Simplify test imports (remove importlib.util workaround) All 231 unit tests pass. Codegen integration verified across 5 LLM providers via cpf2503-20tue78/ai-tutor-ci.
1 parent 5da4634 commit f374942

10 files changed

Lines changed: 543 additions & 136 deletions

File tree

‎.dockerignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
!entrypoint.py
44
!llm_client.py
55
!llm_configs.py
6+
!llm_utils.py
67
!locale/*.json
78
!prompt.py
89
!requirements.txt

‎.github/workflows/build.yml‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,4 +151,59 @@ jobs:
151151
echo "feedback_${{ matrix.model }}=${{ steps.integration.outputs.feedback }}" >> $GITHUB_OUTPUT
152152
python3 -u -m pytest tests/integration.py
153153
154+
integration-tests-codegen:
155+
needs: build
156+
runs-on: ubuntu-latest
157+
strategy:
158+
matrix:
159+
include:
160+
- model: gemini-2.5-flash
161+
secret: GOOGLE_API_KEY
162+
- model: grok-code-fast
163+
secret: XAI_API_KEY
164+
- model: claude-sonnet-4-20250514
165+
secret: CLAUDE_API_KEY
166+
- model: google/gemma-2-9b-it
167+
secret: NVIDIA_NIM_API_KEY
168+
- model: sonar
169+
secret: PERPLEXITY_API_KEY
170+
fail-fast: false
171+
steps:
172+
- uses: actions/checkout@v6
173+
174+
- name: Install uv package manager
175+
uses: astral-sh/setup-uv@v7
176+
with:
177+
version: "latest"
178+
version-file: "pyproject.toml"
179+
python-version: 3.11
180+
enable-cache: true
181+
182+
- run: uv sync
183+
184+
- name: Generate exercise.py with ${{ matrix.model }}
185+
env:
186+
INPUT_PROMPT-FILE: tests/sample_prompt.txt
187+
CONTAINER_OUTPUT: ${{ runner.temp }}/codegen_output
188+
INPUT_CLAUDE_API_KEY: ${{ secrets.CLAUDE_API_KEY }}
189+
INPUT_GEMINI-API-KEY: ${{ secrets.GOOGLE_API_KEY }}
190+
INPUT_GROK-API-KEY: ${{ secrets.XAI_API_KEY }}
191+
INPUT_NVIDIA-API-KEY: ${{ secrets.NVIDIA_NIM_API_KEY }}
192+
INPUT_PERPLEXITY-API-KEY: ${{ secrets.PERPLEXITY_API_KEY }}
193+
INPUT_MODEL: ${{ matrix.model }}
194+
run: |
195+
. .venv/bin/activate
196+
python3 prompt_pipeline/entrypoint.py
197+
timeout-minutes: 5
198+
199+
- name: Show generated exercise.py
200+
run: cat ${{ runner.temp }}/codegen_output/exercise.py
201+
202+
- name: Verify generated code
203+
env:
204+
CONTAINER_OUTPUT: ${{ runner.temp }}/codegen_output
205+
run: |
206+
. .venv/bin/activate
207+
python3 -u -m pytest tests/integration_codegen.py -v
208+
154209
# end .github/workflows/build.yml

‎Dockerfile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ COPY requirements.txt /requirements.txt
1818
COPY prompt.py /prompt.py
1919
COPY llm_client.py /llm_client.py
2020
COPY llm_configs.py /llm_configs.py
21+
COPY llm_utils.py /llm_utils.py
2122
COPY locale/ /locale/
2223

2324
RUN python3 -m pip install --upgrade pip

‎entrypoint.py‎

Lines changed: 1 addition & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818

1919
from llm_client import LLMAPIClient
20-
from llm_configs import ClaudeConfig, GeminiConfig, GrokConfig, NvidiaNIMConfig, PerplexityConfig
20+
from llm_utils import get_config_class, get_model_key_from_env
2121

2222
import prompt
2323

@@ -144,127 +144,6 @@ def write_token_usage(
144144
logging.warning(f"Could not write token usage: {e}")
145145

146146

147-
def get_startwith(key:str, dictionary:dict) -> Any:
148-
result = None
149-
for k, v in dictionary.items():
150-
if key.startswith(k):
151-
result = v
152-
break
153-
return result
154-
155-
156-
def get_model_key_from_env() -> Tuple[str, str]:
157-
"""
158-
Extracts the LLM model and API key from environment variables with flexible selection.
159-
- Uses INPUT_API-KEY if provided, especially with a specified model.
160-
- Falls back to model-specific API keys if INPUT_API-KEY is not set.
161-
- Raises ValueError if no API keys are available.
162-
- Uses model-to-provider mapping for precise model IDs.
163-
"""
164-
api_key_dict = get_api_key_dict_from_env()
165-
valid_keys_dict = {k: v for k, v in api_key_dict.items() if v and v.strip()}
166-
167-
model = os.getenv('INPUT_MODEL', '').lower()
168-
general_api_key = os.getenv('INPUT_API-KEY', '').strip()
169-
170-
# Model-to-provider mapping for precise model IDs
171-
model_to_provider = {
172-
'google/gemma-2-9b-it': 'nvidia_nim',
173-
'sonar': 'perplexity',
174-
'gemini-2.5-flash': 'gemini',
175-
'grok-code-fast': 'grok',
176-
'claude-sonnet-4-20250514': 'claude'
177-
}
178-
179-
# Case 1: Use INPUT_API-KEY if provided
180-
if general_api_key:
181-
selected_model = model or 'gemini-2.5-flash' # Default to specific Gemini model
182-
logging.info(f"Using INPUT_API-KEY for model: {selected_model}")
183-
return selected_model, general_api_key
184-
185-
# Case 2: No INPUT_API-KEY, check model-specific keys
186-
if not valid_keys_dict:
187-
raise ValueError(
188-
"No API keys provided. Set at least one of:\n"
189-
"\tINPUT_API-KEY\n"
190-
"\tINPUT_CLAUDE_API_KEY\n"
191-
"\tINPUT_GEMINI-API-KEY\n"
192-
"\tINPUT_GROK-API-KEY\n"
193-
"\tINPUT_NVIDIA-API-KEY\n"
194-
"\tINPUT_PERPLEXITY-API-KEY\n"
195-
)
196-
197-
# Case 3: Only one API key available
198-
if len(valid_keys_dict) == 1:
199-
selected_model, api_key = next(iter(valid_keys_dict.items()))
200-
logging.info(f"Using single available model: {selected_model}")
201-
return selected_model, api_key.strip()
202-
203-
# Case 4: Use model-to-provider mapping for specified model
204-
provider = model_to_provider.get(model, None)
205-
if model and provider and provider in valid_keys_dict:
206-
logging.info(f"Using mapped model: {model} with provider: {provider}")
207-
return model, valid_keys_dict[provider].strip()
208-
209-
# Case 5: Fallback to provider-based matching
210-
if model:
211-
api_key = get_startwith(model, valid_keys_dict)
212-
if api_key:
213-
logging.info(f"Using specified model with provider matching: {model}")
214-
return model, api_key.strip()
215-
216-
# Case 6: Fallback to Gemini if available
217-
if 'gemini' in valid_keys_dict:
218-
logging.info("Falling back to Gemini model")
219-
return 'gemini-2.5-flash', valid_keys_dict['gemini'].strip()
220-
221-
# Case 7: No matching model or Gemini
222-
raise ValueError(
223-
f"No API key provided for specified model '{model}' and Gemini not available. "
224-
f"Available models: {', '.join(valid_keys_dict.keys())}"
225-
)
226-
227-
228-
def get_api_key_dict_from_env() -> Dict[str, str]:
229-
"""
230-
Retrieves API keys for different models from environment variables.
231-
Returns empty strings for unset variables.
232-
"""
233-
return {
234-
'claude': os.getenv('INPUT_CLAUDE_API_KEY', ''),
235-
'gemini': os.getenv('INPUT_GEMINI-API-KEY', ''),
236-
'grok': os.getenv('INPUT_GROK-API-KEY', ''),
237-
'nvidia_nim': os.getenv('INPUT_NVIDIA-API-KEY', ''),
238-
'perplexity': os.getenv('INPUT_PERPLEXITY-API-KEY', ''),
239-
}
240-
241-
242-
def get_config_class_dict() -> Dict[str, type]:
243-
"""
244-
Returns a dictionary mapping model names to their respective configuration classes.
245-
"""
246-
return {
247-
'claude': ClaudeConfig,
248-
'claude-sonnet-4-20250514': ClaudeConfig, # Add specific model
249-
'gemini': GeminiConfig,
250-
'gemini-2.5-flash': GeminiConfig,
251-
'grok': GrokConfig,
252-
'grok-code-fast': GrokConfig,
253-
'nvidia_nim': NvidiaNIMConfig,
254-
'google/gemma-2-9b-it': NvidiaNIMConfig, # Add specific model
255-
'perplexity': PerplexityConfig,
256-
'sonar': PerplexityConfig # Add specific model
257-
}
258-
259-
260-
def get_config_class(model: str) -> type:
261-
config_map = get_config_class_dict()
262-
config_class = get_startwith(model, config_map)
263-
if not config_class:
264-
raise ValueError(f"Unsupported LLM type: {model}. Use {', '.join(config_map.keys())}")
265-
return config_class
266-
267-
268147
def get_path_tuple(paths_str: str) -> Tuple[pathlib.Path]:
269148
"""
270149
Converts a comma-separated string of file paths to a tuple of pathlib.Path objects.

‎llm_utils.py‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
# begin llm_utils.py
2+
"""Shared LLM utilities used by both the tutor and prompt pipeline.
3+
4+
Extracted from entrypoint.py to eliminate circular imports between
5+
entrypoint.py and prompt_pipeline/entrypoint.py.
6+
"""
7+
8+
import logging
9+
import os
10+
11+
from typing import Any, Dict, Tuple
12+
13+
from llm_configs import (
14+
ClaudeConfig,
15+
GeminiConfig,
16+
GrokConfig,
17+
NvidiaNIMConfig,
18+
PerplexityConfig,
19+
)
20+
21+
22+
logging.basicConfig(level=logging.INFO)
23+
24+
25+
def get_startwith(key: str, dictionary: dict) -> Any:
26+
result = None
27+
for k, v in dictionary.items():
28+
if key.startswith(k):
29+
result = v
30+
break
31+
return result
32+
33+
34+
def get_api_key_dict_from_env() -> Dict[str, str]:
35+
"""
36+
Retrieves API keys for different models from environment variables.
37+
Returns empty strings for unset variables.
38+
"""
39+
return {
40+
'claude': os.getenv('INPUT_CLAUDE_API_KEY', ''),
41+
'gemini': os.getenv('INPUT_GEMINI-API-KEY', ''),
42+
'grok': os.getenv('INPUT_GROK-API-KEY', ''),
43+
'nvidia_nim': os.getenv('INPUT_NVIDIA-API-KEY', ''),
44+
'perplexity': os.getenv('INPUT_PERPLEXITY-API-KEY', ''),
45+
}
46+
47+
48+
def get_config_class_dict() -> Dict[str, type]:
49+
"""
50+
Returns a dictionary mapping model names to their respective configuration classes.
51+
"""
52+
return {
53+
'claude': ClaudeConfig,
54+
'claude-sonnet-4-20250514': ClaudeConfig, # Add specific model
55+
'gemini': GeminiConfig,
56+
'gemini-2.5-flash': GeminiConfig,
57+
'grok': GrokConfig,
58+
'grok-code-fast': GrokConfig,
59+
'nvidia_nim': NvidiaNIMConfig,
60+
'google/gemma-2-9b-it': NvidiaNIMConfig, # Add specific model
61+
'perplexity': PerplexityConfig,
62+
'sonar': PerplexityConfig # Add specific model
63+
}
64+
65+
66+
def get_config_class(model: str) -> type:
67+
config_map = get_config_class_dict()
68+
config_class = get_startwith(model, config_map)
69+
if not config_class:
70+
raise ValueError(f"Unsupported LLM type: {model}. Use {', '.join(config_map.keys())}")
71+
return config_class
72+
73+
74+
def get_model_key_from_env() -> Tuple[str, str]:
75+
"""
76+
Extracts the LLM model and API key from environment variables with flexible selection.
77+
- Uses INPUT_API-KEY if provided, especially with a specified model.
78+
- Falls back to model-specific API keys if INPUT_API-KEY is not set.
79+
- Raises ValueError if no API keys are available.
80+
- Uses model-to-provider mapping for precise model IDs.
81+
"""
82+
api_key_dict = get_api_key_dict_from_env()
83+
valid_keys_dict = {k: v for k, v in api_key_dict.items() if v and v.strip()}
84+
85+
model = os.getenv('INPUT_MODEL', '').lower()
86+
general_api_key = os.getenv('INPUT_API-KEY', '').strip()
87+
88+
# Model-to-provider mapping for precise model IDs
89+
model_to_provider = {
90+
'google/gemma-2-9b-it': 'nvidia_nim',
91+
'sonar': 'perplexity',
92+
'gemini-2.5-flash': 'gemini',
93+
'grok-code-fast': 'grok',
94+
'claude-sonnet-4-20250514': 'claude'
95+
}
96+
97+
# Case 1: Use INPUT_API-KEY if provided
98+
if general_api_key:
99+
selected_model = model or 'gemini-2.5-flash' # Default to specific Gemini model
100+
logging.info(f"Using INPUT_API-KEY for model: {selected_model}")
101+
return selected_model, general_api_key
102+
103+
# Case 2: No INPUT_API-KEY, check model-specific keys
104+
if not valid_keys_dict:
105+
raise ValueError(
106+
"No API keys provided. Set at least one of:\n"
107+
"\tINPUT_API-KEY\n"
108+
"\tINPUT_CLAUDE_API_KEY\n"
109+
"\tINPUT_GEMINI-API-KEY\n"
110+
"\tINPUT_GROK-API-KEY\n"
111+
"\tINPUT_NVIDIA-API-KEY\n"
112+
"\tINPUT_PERPLEXITY-API-KEY\n"
113+
)
114+
115+
# Case 3: Only one API key available
116+
if len(valid_keys_dict) == 1:
117+
selected_model, api_key = next(iter(valid_keys_dict.items()))
118+
logging.info(f"Using single available model: {selected_model}")
119+
return selected_model, api_key.strip()
120+
121+
# Case 4: Use model-to-provider mapping for specified model
122+
provider = model_to_provider.get(model, None)
123+
if model and provider and provider in valid_keys_dict:
124+
logging.info(f"Using mapped model: {model} with provider: {provider}")
125+
return model, valid_keys_dict[provider].strip()
126+
127+
# Case 5: Fallback to provider-based matching
128+
if model:
129+
api_key = get_startwith(model, valid_keys_dict)
130+
if api_key:
131+
logging.info(f"Using specified model with provider matching: {model}")
132+
return model, api_key.strip()
133+
134+
# Case 6: Fallback to Gemini if available
135+
if 'gemini' in valid_keys_dict:
136+
logging.info("Falling back to Gemini model")
137+
return 'gemini-2.5-flash', valid_keys_dict['gemini'].strip()
138+
139+
# Case 7: No matching model or Gemini
140+
raise ValueError(
141+
f"No API key provided for specified model '{model}' and Gemini not available. "
142+
f"Available models: {', '.join(valid_keys_dict.keys())}"
143+
)
144+
145+
# end llm_utils.py

0 commit comments

Comments
 (0)