File size: 626 Bytes
457352d
 
 
 
 
 
 
 
 
 
369fdbd
457352d
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import os

import google.generativeai as genai
from dotenv import load_dotenv

load_dotenv()


class BaseAIClient:
    def __init__(self,system,model_response):
        genai.configure(api_key=os.getenv('GOOGLE_API_KEY'))
        self.system = system
        self.model = model_response

    def create(self, conversation):
        model = genai.GenerativeModel('gemini-pro')
        new_conversation = [
            {"role": 'user', "parts": [self.system]},
            {"role": 'model', "parts": [self.model]},
        ]
        new_conversation.extend(conversation)
        return model.start_chat(history=new_conversation)