models.py

import os
import json
from langchain.schema import SystemMessage, HumanMessage
from langchain.prompts.chat import (
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
    ChatPromptTemplate
)
from langchain.prompts.prompt import PromptTemplate
from langchain.retrievers.multi_query import MultiQueryRetriever

from langchain_aws import BedrockEmbeddings
from langchain_aws.chat_models.bedrock_converse import ChatBedrockConverse
from langchain_cohere import ChatCohere
from langchain_fireworks.chat_models import ChatFireworks
from langchain_fireworks.embeddings import FireworksEmbeddings
from langchain_groq.chat_models import ChatGroq
from langchain_openai import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_ollama.chat_models import ChatOllama
from langchain_ollama.embeddings import OllamaEmbeddings
from langchain_cohere.embeddings import CohereEmbeddings
from langchain_cohere.chat_models import ChatCohere
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai.embeddings import GoogleGenerativeAIEmbeddings
from langchain_community.chat_models import ChatPerplexity
from langchain_together import ChatTogether
from langchain_together.embeddings import TogetherEmbeddings

def split_provider_model(provider_model):
    parts = provider_model.split(':', 1)
    provider = parts[0]
    model = parts[1] if len(parts) > 1 else None
    return provider, model

def get_model(provider_model, temperature=0.0):
    provider, model = split_provider_model(provider_model)
    match provider:
        case 'bedrock':
            if model is None:
                model = "anthropic.claude-3-sonnet-20240229-v1:0"
            chat_llm = ChatBedrockConverse(model=model, temperature=temperature)
        case 'cohere':
            if model is None:
                model = 'command-r-plus'
            chat_llm = ChatCohere(model=model, temperature=temperature)
        case 'fireworks':
            if model is None:
                model = 'accounts/fireworks/models/llama-v3p1-8b-instruct'
            chat_llm = ChatFireworks(model_name=model, temperature=temperature, max_tokens=120000)
        case 'googlegenerativeai':
            if model is None:
                model = "gemini-1.5-flash"
            chat_llm = ChatGoogleGenerativeAI(model=model, temperature=temperature, 
                                              max_tokens=None, timeout=None, max_retries=2,)
        case 'groq':
            if model is None:
                model = 'llama-3.1-8b-instant'
            chat_llm = ChatGroq(model_name=model, temperature=temperature)
        case 'ollama':
            if model is None:
                model = 'llama3.1'
            chat_llm = ChatOllama(model=model, temperature=temperature)
        case 'openai':
            if model is None:
                model = "gpt-4o-mini"
            chat_llm = ChatOpenAI(model_name=model, temperature=temperature)
        case 'perplexity':
            if model is None:
                model = 'llama-3.1-sonar-small-128k-online'
            chat_llm = ChatPerplexity(model=model, temperature=temperature)
        case 'together':
            if model is None:
                model = 'meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo'
            chat_llm = ChatTogether(model=model, temperature=temperature)
        case _:
            raise ValueError(f"Unknown LLM provider {provider}")
    
    return chat_llm


def get_embedding_model(provider_model):
    provider, model = split_provider_model(provider_model)
    match provider:
        case 'bedrock':
            if model is None:
                model = "amazon.titan-embed-text-v2:0"
            embedding_model = BedrockEmbeddings(model_id=model)
        case 'cohere':
            if model is None:
                model = "embed-english-light-v3.0"
            embedding_model = CohereEmbeddings(model=model)
        case 'fireworks':
            if model is None:
                model = 'nomic-ai/nomic-embed-text-v1.5'
            embedding_model = FireworksEmbeddings(model=model)
        case 'ollama':
            if model is None:
                model = 'nomic-embed-text:latest'
            embedding_model = OllamaEmbeddings(model=model)
        case 'openai':
            if model is None:
                model = "text-embedding-3-small"
            embedding_model = OpenAIEmbeddings(model=model)
        case 'googlegenerativeai':
            if model is None:
                model = "models/embedding-001"
            embedding_model = GoogleGenerativeAIEmbeddings(model=model)
        case 'groq':
            embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
        case 'perplexity':
            raise ValueError(f"Cannot use Perplexity for embedding model")
        case 'together':
            if model is None:
                model = 'togethercomputer/m2-bert-80M-2k-retrieval'
            embedding_model = TogetherEmbeddings(model=model)
        case _:
            raise ValueError(f"Unknown LLM provider {provider}")

    return embedding_model


import unittest
from unittest.mock import patch
from models import get_embedding_model  # Make sure this import is correct

class TestGetEmbeddingModel(unittest.TestCase):

    @patch('models.BedrockEmbeddings')
    def test_bedrock_embedding(self, mock_bedrock):
        result = get_embedding_model('bedrock')
        mock_bedrock.assert_called_once_with(model_id='cohere.embed-multilingual-v3')
        self.assertEqual(result, mock_bedrock.return_value)

    @patch('models.CohereEmbeddings')
    def test_cohere_embedding(self, mock_cohere):
        result = get_embedding_model('cohere')
        mock_cohere.assert_called_once_with(model='embed-english-light-v3.0')
        self.assertEqual(result, mock_cohere.return_value)

    @patch('models.FireworksEmbeddings')
    def test_fireworks_embedding(self, mock_fireworks):
        result = get_embedding_model('fireworks')
        mock_fireworks.assert_called_once_with(model='nomic-ai/nomic-embed-text-v1.5')
        self.assertEqual(result, mock_fireworks.return_value)

    @patch('models.OllamaEmbeddings')
    def test_ollama_embedding(self, mock_ollama):
        result = get_embedding_model('ollama')
        mock_ollama.assert_called_once_with(model='nomic-embed-text:latest')
        self.assertEqual(result, mock_ollama.return_value)

    @patch('models.OpenAIEmbeddings')
    def test_openai_embedding(self, mock_openai):
        result = get_embedding_model('openai')
        mock_openai.assert_called_once_with(model='text-embedding-3-small')
        self.assertEqual(result, mock_openai.return_value)

    @patch('models.GoogleGenerativeAIEmbeddings')
    def test_google_embedding(self, mock_google):
        result = get_embedding_model('googlegenerativeai')
        mock_google.assert_called_once_with(model='models/embedding-001')
        self.assertEqual(result, mock_google.return_value)

    @patch('models.TogetherEmbeddings')
    def test_together_embedding(self, mock_together):
        result = get_embedding_model('together')
        mock_together.assert_called_once_with(model='BAAI/bge-base-en-v1.5')
        self.assertEqual(result, mock_together.return_value)

    def test_invalid_provider(self):
        with self.assertRaises(ValueError):
            get_embedding_model('invalid_provider')

    def test_groq_provider(self):
        with self.assertRaises(ValueError):
            get_embedding_model('groq')

    def test_perplexity_provider(self):
        with self.assertRaises(ValueError):
            get_embedding_model('perplexity')


import unittest
from unittest.mock import patch
from models import get_model  # Make sure this import is correct

class TestGetModel(unittest.TestCase):

    @patch('models.ChatBedrockConverse')
    def test_bedrock_model(self, mock_bedrock):
        result = get_model('bedrock')
        mock_bedrock.assert_called_once_with(
            model="anthropic.claude-3-sonnet-20240229-v1:0",
            temperature=0.0
        )
        self.assertEqual(result, mock_bedrock.return_value)

    @patch('models.ChatCohere')
    def test_cohere_model(self, mock_cohere):
        result = get_model('cohere')
        mock_cohere.assert_called_once_with(model='command-r-plus', temperature=0.0)
        self.assertEqual(result, mock_cohere.return_value)

    @patch('models.ChatFireworks')
    def test_fireworks_model(self, mock_fireworks):
        result = get_model('fireworks')
        mock_fireworks.assert_called_once_with(
            model_name='accounts/fireworks/models/llama-v3p1-8b-instruct',
            temperature=0.0,
            max_tokens=120000
        )
        self.assertEqual(result, mock_fireworks.return_value)

    @patch('models.ChatGoogleGenerativeAI')
    def test_google_model(self, mock_google):
        result = get_model('googlegenerativeai')
        mock_google.assert_called_once_with(
            model="gemini-1.5-pro",
            temperature=0.0,
            max_tokens=None,
            timeout=None,
            max_retries=2
        )
        self.assertEqual(result, mock_google.return_value)

    @patch('models.ChatGroq')
    def test_groq_model(self, mock_groq):
        result = get_model('groq')
        mock_groq.assert_called_once_with(model_name='llama-3.1-8b-instant', temperature=0.0)
        self.assertEqual(result, mock_groq.return_value)

    @patch('models.ChatOllama')
    def test_ollama_model(self, mock_ollama):
        result = get_model('ollama')
        mock_ollama.assert_called_once_with(model='llama3.1', temperature=0.0)
        self.assertEqual(result, mock_ollama.return_value)

    @patch('models.ChatOpenAI')
    def test_openai_model(self, mock_openai):
        result = get_model('openai')
        mock_openai.assert_called_once_with(model_name='gpt-4o-mini', temperature=0.0)
        self.assertEqual(result, mock_openai.return_value)

    @patch('models.ChatPerplexity')
    def test_perplexity_model(self, mock_perplexity):
        result = get_model('perplexity')
        mock_perplexity.assert_called_once_with(model='llama-3.1-sonar-small-128k-online', temperature=0.0)
        self.assertEqual(result, mock_perplexity.return_value)

    @patch('models.ChatTogether')
    def test_together_model(self, mock_together):
        result = get_model('together')
        mock_together.assert_called_once_with(model='meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo', temperature=0.0)
        self.assertEqual(result, mock_together.return_value)

    def test_invalid_provider(self):
        with self.assertRaises(ValueError):
            get_model('invalid_provider')

    def test_custom_temperature(self):
        with patch('models.ChatOpenAI') as mock_openai:
            result = get_model('openai', temperature=0.5)
            mock_openai.assert_called_once_with(model_name='gpt-4o-mini', temperature=0.5)
            self.assertEqual(result, mock_openai.return_value)

    def test_custom_model(self):
        with patch('models.ChatOpenAI') as mock_openai:
            result = get_model('openai/gpt-4')
            mock_openai.assert_called_once_with(model_name='gpt-4', temperature=0.0)
            self.assertEqual(result, mock_openai.return_value)

if __name__ == '__main__':
    unittest.main()