38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
from dotenv import load_dotenv
|
|
import os
|
|
import openai
|
|
|
|
class GeminiModel:
|
|
|
|
def __init__(self):
|
|
load_dotenv()
|
|
self.openrouter_api_key = os.getenv("OPENROUTER_API_KEY")
|
|
self.openrouter_api_url = os.getenv("OPENROUTER_API_URL", "https://openrouter.ai/api/v1")
|
|
|
|
def call(self, prompt):
|
|
if not self.openrouter_api_key:
|
|
raise ValueError("OPENROUTER_API_KEY is not set in environment variables.")
|
|
|
|
client = openai.OpenAI(
|
|
base_url=self.openrouter_api_url,
|
|
api_key=self.openrouter_api_key,
|
|
)
|
|
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model="google/gemini-2.5-flash",
|
|
temperature=0.5,
|
|
messages=[
|
|
{"role": "user", "content": prompt}
|
|
]
|
|
)
|
|
data = response.to_dict()
|
|
return data['choices'][0]['message']['content']
|
|
except openai.APIError as e:
|
|
print(f"OpenRouter API call failed: {e}")
|
|
raise
|
|
|
|
if __name__ == "__main__":
|
|
model = GeminiModel()
|
|
print(model.call("Сколько будет 2 + 2?"))
|