from google import genai

import os

import logging

import json

import requests

from datetime import datetime

import anthropic

from dotenv import load_dotenv

# Load environment variables from .env file

load_dotenv()

class RateLimitError(Exception):

"""Raised when API rate limit is hit - signals to stop processing."""

pass

# 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, encoding='utf-8')

file_handler.setFormatter(

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

)

logger.addHandler(file_handler)

# Simple cache configuration

cache_file = "llm_cache.json"

# Stats tracking

class LLMStats:

"""Track LLM usage statistics."""

def __init__(self):

self.reset()

def reset(self):

self.total_calls = 0

self.cache_hits = 0

self.input_chars = 0

self.output_chars = 0

self.start_time = datetime.now()

def record_call(self, prompt_chars, response_chars, from_cache=False):

self.total_calls += 1

if from_cache:

self.cache_hits += 1

self.input_chars += prompt_chars

self.output_chars += response_chars

def get_summary(self):

elapsed = (datetime.now() - self.start_time).total_seconds()

# Rough token estimate (4 chars per token)

input_tokens = self.input_chars // 4

output_tokens = self.output_chars // 4

return {

"total_calls": self.total_calls,

"cache_hits": self.cache_hits,

"api_calls": self.total_calls - self.cache_hits,

"input_chars": self.input_chars,

"output_chars": self.output_chars,

"estimated_input_tokens": input_tokens,

"estimated_output_tokens": output_tokens,

"elapsed_seconds": round(elapsed, 1)

}

# Global stats instance

llm_stats = LLMStats()

def load_cache():

try:

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

return json.load(f)

except:

logger.warning(f"Failed to load cache.")

return {}

def save_cache(cache):

try:

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

json.dump(cache, f)

except:

logger.warning(f"Failed to save cache")

def get_llm_provider():

provider = os.getenv("LLM_PROVIDER")

if not provider and (os.getenv("GEMINI_PROJECT_ID") or os.getenv("GEMINI_API_KEY")):

provider = "GEMINI"

if not provider and os.getenv("ANTHROPIC_API_KEY"):

provider = "ANTHROPIC"

return provider

def _call_llm_provider(prompt: str) -> str:

"""

Call an LLM provider based on environment variables.

Environment variables:

- LLM_PROVIDER: "OLLAMA" or "XAI"

- <provider>_MODEL: Model name (e.g., OLLAMA_MODEL, XAI_MODEL)

- <provider>_BASE_URL: Base URL without endpoint (e.g., OLLAMA_BASE_URL, XAI_BASE_URL)

- <provider>_API_KEY: API key (e.g., OLLAMA_API_KEY, XAI_API_KEY; optional for providers that don't require it)

The endpoint /v1/chat/completions will be appended to the base URL.

"""

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

# Read the provider from environment variable

provider = os.environ.get("LLM_PROVIDER")

if not provider:

raise ValueError("LLM_PROVIDER environment variable is required")

# Construct the names of the other environment variables

model_var = f"{provider}_MODEL"

base_url_var = f"{provider}_BASE_URL"

api_key_var = f"{provider}_API_KEY"

# Read the provider-specific variables

model = os.environ.get(model_var)

base_url = os.environ.get(base_url_var)

api_key = os.environ.get(api_key_var, "") # API key is optional, default to empty string

# Validate required variables

if not model:

raise ValueError(f"{model_var} environment variable is required")

if not base_url:

raise ValueError(f"{base_url_var} environment variable is required")

# Append the endpoint to the base URL

url = f"{base_url.rstrip('/')}/v1/chat/completions"

# Configure headers and payload based on provider

headers = {

"Content-Type": "application/json",

}

if api_key: # Only add Authorization header if API key is provided

headers["Authorization"] = f"Bearer {api_key}"

payload = {

"model": model,

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

"temperature": 0.7,

}

try:

response = requests.post(url, headers=headers, json=payload)

response_json = response.json() # Log the response

logger.info("RESPONSE:\n%s", json.dumps(response_json, indent=2))

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

response.raise_for_status()

return response.json()["choices"][0]["message"]["content"]

except requests.exceptions.HTTPError as e:

error_message = f"HTTP error occurred: {e}"

try:

error_details = response.json().get("error", "No additional details")

error_message += f" (Details: {error_details})"

except:

pass

raise Exception(error_message)

except requests.exceptions.ConnectionError:

raise Exception(f"Failed to connect to {provider} API. Check your network connection.")

except requests.exceptions.Timeout:

raise Exception(f"Request to {provider} API timed out.")

except requests.exceptions.RequestException as e:

raise Exception(f"An error occurred while making the request to {provider}: {e}")

except ValueError:

raise Exception(f"Failed to parse response as JSON from {provider}. The server might have returned an invalid response.")

# 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 = load_cache()

# Return from cache if exists

if prompt in cache:

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

llm_stats.record_call(len(prompt), len(cache[prompt]), from_cache=True)

return cache[prompt]

provider = get_llm_provider()

if provider == "GEMINI":

response_text = _call_llm_gemini(prompt)

elif provider == "ANTHROPIC":

response_text = _call_llm_anthropic(prompt)

else: # generic method using a URL that is OpenAI compatible API (Ollama, ...)

response_text = _call_llm_provider(prompt)

# Log the response

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

# Record stats

llm_stats.record_call(len(prompt), len(response_text), from_cache=False)

# Update cache if enabled

if use_cache:

# Load cache again to avoid overwrites

cache = load_cache()

# Add to cache and save

cache[prompt] = response_text

save_cache(cache)

return response_text

def _call_llm_gemini(prompt: str) -> str:

if os.getenv("GEMINI_PROJECT_ID"):

client = genai.Client(

vertexai=True,

project=os.getenv("GEMINI_PROJECT_ID"),

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

)

elif os.getenv("GEMINI_API_KEY"):

client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

else:

raise ValueError("Either GEMINI_PROJECT_ID or GEMINI_API_KEY must be set in the environment")

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

# Configure thinking budget for Gemini 2.5 models

# Set GEMINI_THINKING_BUDGET to control reasoning depth:

# 0 = disable thinking (fastest, lower quality)

# 1024 = light thinking (good balance of speed and quality)

# 8192 = default deep thinking (slowest, highest quality)

# -1 = dynamic (model decides based on complexity)

# Default: 1024 for reasonable speed with quality

thinking_budget_str = os.getenv("GEMINI_THINKING_BUDGET", "1024")

thinking_budget = int(thinking_budget_str) if thinking_budget_str else 1024

# Build config with thinking settings for 2.5 models

config = None

if "2.5" in model or "2-5" in model:

from google.genai import types

config = types.GenerateContentConfig(

thinking_config=types.ThinkingConfig(thinking_budget=thinking_budget)

)

try:

response = client.models.generate_content(

model=model,

contents=[prompt],

config=config

)

return response.text

except Exception as e:

error_str = str(e)

# Detect rate limit errors

if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str or "quota" in error_str.lower():

logger.error(f"RATE LIMIT HIT: {error_str}")

raise RateLimitError(f"Rate limit exceeded: {error_str}")

raise

def _call_llm_anthropic(prompt: str) -> str:

"""Call Anthropic's Claude API."""

api_key = os.getenv("ANTHROPIC_API_KEY")

if not api_key:

raise ValueError("ANTHROPIC_API_KEY must be set in the environment")

client = anthropic.Anthropic(api_key=api_key)

model = os.getenv("ANTHROPIC_MODEL", "claude-3-5-sonnet-20241022")

message = client.messages.create(

model=model,

max_tokens=4096,

messages=[

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

]

)

return message.content[0].text

def call_llm_parallel(prompts: list, use_cache: bool = True, max_workers: int = 5) -> list:

"""

Call LLM on multiple prompts in parallel.

Args:

prompts: List of prompts

use_cache: Whether to use cache

max_workers: Max concurrent workers

Returns:

List of responses in same order as prompts

Raises:

RateLimitError: If rate limit is hit, stops immediately and propagates

"""

import concurrent.futures

results = [None] * len(prompts)

rate_limit_hit = False

rate_limit_error = None

def call_single(idx_prompt):

idx, prompt = idx_prompt

return idx, call_llm(prompt, use_cache=use_cache)

with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:

futures = {executor.submit(call_single, (i, p)): i for i, p in enumerate(prompts)}

for future in concurrent.futures.as_completed(futures):

idx = futures[future]

try:

result_idx, result = future.result()

results[result_idx] = result

except RateLimitError as e:

# Rate limit hit - cancel remaining futures and stop

logger.error(f"RATE LIMIT HIT - stopping all parallel calls: {e}")

rate_limit_hit = True

rate_limit_error = e

results[idx] = f"ERROR: {e}"

# Cancel remaining futures

for f in futures:

f.cancel()

break

except Exception as e:

logger.error(f"Parallel LLM call failed: {e}")

results[idx] = f"ERROR: {e}"

# If rate limit was hit, raise it so the caller can handle it

if rate_limit_hit:

raise RateLimitError(str(rate_limit_error))

return results

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}")