Spaces:
Paused
Paused
File size: 4,203 Bytes
cb1a5c9 | 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 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 | # filename: language_models.py
"""
Module for defining various language model classes and their invocation methods.
"""
import os
import logging
from abc import ABC, abstractmethod
from time import time
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from log_config import get_logger
# Configure logging
logger = get_logger('LanguageModels')
def get_api_key(service_name: str) -> str:
"""
Retrieve the API key from environment variables.
Args:
service_name (str): The name of the service.
Returns:
str: The API key for the service.
Raises:
ValueError: If the API key is not set in environment variables.
"""
api_key = os.getenv(service_name)
if not api_key:
logger.error(f"{service_name} environment variable not set.")
raise ValueError(f"API key for {service_name} is required")
return api_key
class LanguageModel(ABC):
"""
Abstract base class for language models.
Defines a common interface for different AI language models.
"""
@abstractmethod
def invoke(self, prompt: str) -> str:
"""
Send a prompt to the language model and return the response.
Args:
prompt (str): The text prompt to send to the language model.
Returns:
str: The response from the language model.
Raises:
Exception: Propagates any exceptions that occur during invocation.
"""
pass
class GPT4Model(LanguageModel):
"""
A language model class for OpenAI's GPT-4.
"""
def __init__(self):
api_key = get_api_key("OPENAI_API_KEY")
self.model = ChatOpenAI(
model=os.getenv('MODEL_NAME_GPT4', 'gpt-4o'),
temperature=0,
openai_api_key=api_key
)
def invoke(self, prompt: str) -> str:
if not prompt:
logger.error("Empty prompt provided.")
raise ValueError("Prompt cannot be empty.")
start_time = time()
try:
response = self.model.invoke([prompt])
logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
return response
except Exception as e:
logger.error(f"Failed to invoke GPT-4 with prompt: '{prompt}', error: {str(e)}")
raise
class ClaudeModel(LanguageModel):
"""
A language model class for Anthropic's Claude.
"""
def __init__(self):
api_key = get_api_key("ANTHROPIC_API_KEY")
self.model = ChatAnthropic(
model=os.getenv('MODEL_NAME_CLAUDE_OPUS', 'claude-3-opus-20240229'),
temperature=0,
anthropic_api_key=api_key
)
def invoke(self, prompt: str) -> str:
if not prompt:
logger.error("Empty prompt provided.")
raise ValueError("Prompt cannot be empty.")
start_time = time()
try:
response = self.model.invoke([prompt])
logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
return response
except Exception as e:
logger.error(f"Failed to invoke Claude with prompt: '{prompt}', error: {str(e)}")
raise
class GeminiModel(LanguageModel):
"""
A language model class for Google's Generative AI, Gemini.
"""
def __init__(self):
api_key = get_api_key("GOOGLE_API_KEY")
self.model = ChatGoogleGenerativeAI(
model=os.getenv('MODEL_NAME_GEMINI', 'gemini-1.5-pro-latest'),
api_key=api_key
)
def invoke(self, prompt: str) -> str:
if not prompt:
logger.error("Empty prompt provided.")
raise ValueError("Prompt cannot be empty.")
start_time = time()
try:
response = self.model.invoke([prompt])
logger.info(f"Invocation successful in {time() - start_time:.2f} seconds.")
return response
except Exception as e:
logger.error(f"Failed to invoke Gemini with prompt: '{prompt}', error: {str(e)}")
raise
# file: language_models.py (end)
|