chains.py

from langchain_google_genai import GoogleGenerativeAI
from prompts import evaluation_prompt_v1, evaluation_prompt_v2
from parsers import evaluation_parser_v1, evaluation_parser_v2
from dotenv import load_dotenv
import os
load_dotenv()


def setup_chain_v1():
    llm = GoogleGenerativeAI(model='gemini-1.5-pro',
                 temperature=0.5,
                 google_api_key=os.getenv('GOOGLE_API_KEY'))
    
    chain = evaluation_prompt_v1 | llm | evaluation_parser_v1
    return chain

def setup_chain_v2():
    llm = GoogleGenerativeAI(model='gemini-1.5-pro',
                 temperature=0.5,
                 google_api_key=os.getenv('GOOGLE_API_KEY'))
    
    chain = evaluation_prompt_v2 | llm | evaluation_parser_v2
    return chain