GeezerVison is showing github.com — PocketFlow-Tutorial-Codebase-Knowledge/utils/call_llm.py at main · tylerdgenius/PocketFlow-Tutorial-Codebase-Knowledge

Expand file tree

call_llm.py

Copy path

More file actions

Latest commit

History

207 lines (175 loc) · 6.29 KB

call_llm.py

Copy path

File metadata and controls

207 lines (175 loc) · 6.29 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

from google import genai

import os

import logging

import json

from datetime import datetime

import requests

# Configure logging

log_directory = os.getenv("LOG_DIR", "logs")

os.makedirs(log_directory, exist_ok=True)

log_file = os.path.join(

log_directory, f"llm_calls_{datetime.now().strftime('%Y%m%d')}.log"

)

# Set up logger

logger = logging.getLogger("llm_logger")

logger.setLevel(logging.INFO)

logger.propagate = False # Prevent propagation to root logger

file_handler = logging.FileHandler(log_file)

file_handler.setFormatter(

logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")

)

logger.addHandler(file_handler)

# Simple cache configuration

cache_file = "llm_cache.json"

# By default, we Google Gemini 2.5 pro, as it shows great performance for code understanding

def call_llm(prompt: str, use_cache: bool = True) -> str:

# Log the prompt

logger.info(f"PROMPT: {prompt}")

# Check cache if enabled

if use_cache:

# Load cache from disk

cache = {}

if os.path.exists(cache_file):

try:

with open(cache_file, "r") as f:

cache = json.load(f)

except:

logger.warning(f"Failed to load cache, starting with empty cache")

# Return from cache if exists

if prompt in cache:

logger.info(f"RESPONSE: {cache[prompt]}")

return cache[prompt]

# # Call the LLM if not in cache or cache disabled

# client = genai.Client(

# vertexai=True,

# # TODO: change to your own project id and location

# project=os.getenv("GEMINI_PROJECT_ID", "your-project-id"),

# location=os.getenv("GEMINI_LOCATION", "us-central1")

# )

# You can comment the previous line and use the AI Studio key instead:

client = genai.Client(

api_key=os.getenv("GEMINI_API_KEY", ""),

)

model = os.getenv("GEMINI_MODEL", "gemini-2.5-pro-exp-03-25")

# model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash-preview-04-17")

response = client.models.generate_content(model=model, contents=[prompt])

response_text = response.text

# Log the response

logger.info(f"RESPONSE: {response_text}")

# Update cache if enabled

if use_cache:

# Load cache again to avoid overwrites

cache = {}

if os.path.exists(cache_file):

try:

with open(cache_file, "r") as f:

cache = json.load(f)

except:

pass

# Add to cache and save

cache[prompt] = response_text

try:

with open(cache_file, "w") as f:

json.dump(cache, f)

except Exception as e:

logger.error(f"Failed to save cache: {e}")

return response_text

# # Use Anthropic Claude 3.7 Sonnet Extended Thinking

# def call_llm(prompt, use_cache: bool = True):

# from anthropic import Anthropic

# client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", "your-api-key"))

# response = client.messages.create(

# model="claude-3-7-sonnet-20250219",

# max_tokens=21000,

# thinking={

# "type": "enabled",

# "budget_tokens": 20000

# },

# messages=[

# {"role": "user", "content": prompt}

# ]

# )

# return response.content[1].text

# # Use OpenAI o1

# def call_llm(prompt, use_cache: bool = True):

# from openai import OpenAI

# client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "your-api-key"))

# r = client.chat.completions.create(

# model="o1",

# messages=[{"role": "user", "content": prompt}],

# response_format={

# "type": "text"

# },

# reasoning_effort="medium",

# store=False

# )

# return r.choices[0].message.content

# Use OpenRouter API

# def call_llm(prompt: str, use_cache: bool = True) -> str:

# # Log the prompt

# logger.info(f"PROMPT: {prompt}")

# # Check cache if enabled

# if use_cache:

# # Load cache from disk

# cache = {}

# if os.path.exists(cache_file):

# try:

# with open(cache_file, "r") as f:

# cache = json.load(f)

# except:

# logger.warning(f"Failed to load cache, starting with empty cache")

# # Return from cache if exists

# if prompt in cache:

# logger.info(f"RESPONSE: {cache[prompt]}")

# return cache[prompt]

# # OpenRouter API configuration

# api_key = os.getenv("OPENROUTER_API_KEY", "")

# model = os.getenv("OPENROUTER_MODEL", "google/gemini-2.0-flash-exp:free")

# headers = {

# "Authorization": f"Bearer {api_key}",

# }

# data = {

# "model": model,

# "messages": [{"role": "user", "content": prompt}]

# }

# response = requests.post(

# "https://openrouter.ai/api/v1/chat/completions",

# headers=headers,

# json=data

# )

# if response.status_code != 200:

# error_msg = f"OpenRouter API call failed with status {response.status_code}: {response.text}"

# logger.error(error_msg)

# raise Exception(error_msg)

# try:

# response_text = response.json()["choices"][0]["message"]["content"]

# except Exception as e:

# error_msg = f"Failed to parse OpenRouter response: {e}; Response: {response.text}"

# logger.error(error_msg)

# raise Exception(error_msg)

# # Log the response

# logger.info(f"RESPONSE: {response_text}")

# # Update cache if enabled

# if use_cache:

# # Load cache again to avoid overwrites

# cache = {}

# if os.path.exists(cache_file):

# try:

# with open(cache_file, "r") as f:

# cache = json.load(f)

# except:

# pass

# # Add to cache and save

# cache[prompt] = response_text

# try:

# with open(cache_file, "w") as f:

# json.dump(cache, f)

# except Exception as e:

# logger.error(f"Failed to save cache: {e}")

# return response_text

if __name__ == "__main__":

test_prompt = "Hello, how are you?"

# First call - should hit the API

print("Making call...")

response1 = call_llm(test_prompt, use_cache=False)

print(f"Response: {response1}")