Spaces:
Running
Running
File size: 3,382 Bytes
5d7e9a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
from google import genai
import os
import logging
import json
from datetime import datetime
# 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
# Create formatter
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
# Add file handler
file_handler = logging.FileHandler(log_file, encoding='utf-8')
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
# Add stdout handler
stdout_handler = logging.StreamHandler()
stdout_handler.setFormatter(formatter)
# logger.addHandler(stdout_handler)
# Simple cache configuration - use temp directory for HF Spaces
import tempfile
cache_file = os.path.join(tempfile.gettempdir(), "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", encoding="utf-8") 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 cached: {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", ""),
)
# pro was too slow and janky
model = os.getenv("GEMINI_MODEL", "gemini-2.5-flash")
# 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", encoding="utf-8") as f:
cache = json.load(f)
except:
pass
# Add to cache and save
cache[prompt] = response_text
try:
with open(cache_file, "w", encoding="utf-8") 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}") |