Spaces:
Sleeping
Sleeping
File size: 2,844 Bytes
f810b2f |
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 |
from os import getenv
from typing import Optional
from langchain_openai import ChatOpenAI
class ModelManager:
"""
A class to manage the creation and configuration of language model instances.
Handles API keys and provides fallbacks if environment variables are not found.
"""
def __init__(
self,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
model_name: str = "google/gemini-2.5-flash-preview",
helicone_api_key: Optional[str] = None,
):
"""
Initialize the ModelManager with the specified parameters.
Args:
api_key: The API key for the language model service. If None, will try to get from environment.
api_base: The base URL for the language model service. If None, will try to get from environment.
model_name: The name of the model to use.
helicone_api_key: The Helicone API key for logging. If None, will try to get from environment.
"""
self.api_key = api_key
if not self.api_key:
print("Warning: No API key found.")
self.api_base = api_base
if not self.api_base:
print("Warning: No API base URL found.")
self.helicone_api_key = helicone_api_key
self.model_name = model_name
def create_model(self, **kwargs) -> Optional[ChatOpenAI]:
"""
Create and return a configured language model instance.
Args:
**kwargs: Additional keyword arguments to pass to the model constructor.
Returns:
A configured ChatOpenAI instance, or None if required parameters are missing.
"""
# Check if required parameters are available
if not self.api_key or not self.api_base:
print("Error: Cannot create model. Missing required API key or base URL.")
return None
# Prepare model kwargs
model_kwargs = kwargs.get("model_kwargs", {})
# Add Helicone headers if available
if self.helicone_api_key:
extra_headers = model_kwargs.get("extra_headers", {})
extra_headers["Helicone-Auth"] = f"Bearer {self.helicone_api_key}"
model_kwargs["extra_headers"] = extra_headers
# Update kwargs with new model_kwargs
kwargs["model_kwargs"] = model_kwargs
# Create and return the model
try:
return ChatOpenAI(
openai_api_key=self.api_key,
openai_api_base=self.api_base,
model_name=self.model_name,
**kwargs
)
except Exception as e:
print(f"Error creating model: {str(e)}")
return None |