File size: 1,849 Bytes
de850e8
 
 
edd46dc
de850e8
 
 
 
 
 
 
 
cf07068
5d46337
de850e8
 
 
 
273182d
de850e8
 
 
 
 
273182d
 
 
de850e8
 
 
 
273182d
ceaa913
de850e8
 
 
 
edd46dc
de850e8
 
 
 
 
273182d
de850e8
 
 
 
 
cf07068
 
 
 
 
 
 
de850e8
 
 
 
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
# llm.py
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_groq import ChatGroq
from langchain_community.chat_models import ChatOllama
from langchain_core.runnables import ConfigurableField
from langchain.callbacks.base import BaseCallbackHandler


class StreamCallback(BaseCallbackHandler):
    def on_llm_new_token(self, token: str, **kwargs):
        # print(token, end="", flush=True)
        pass


def get_llm(streaming=True):
    return ChatOpenAI(
        model="gpt-4o",
        temperature=0,
        streaming=streaming,
        callbacks=[StreamCallback()],
    ).configurable_alternatives(
        ConfigurableField(id="llm"),
        default_key="gpt_4o",
        claude_3_5_sonnet=ChatAnthropic(
            model="claude-3-5-sonnet-20240620",
            temperature=0,
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
        gpt_3_5_turbo=ChatOpenAI(
            model="gpt-3.5-turbo-0125",
            temperature=0,
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
        gemini_1_5_flash=ChatGoogleGenerativeAI(
            model="gemini-1.5-flash",
            temperature=0,
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
        llama3_70b=ChatGroq(
            model_name="llama3-70b-8192",
            temperature=0,
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
        eeve=ChatOllama(
            model="EEVE-Korean-10.8B",
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
        gemma2=ChatOllama(
            model="gemma2",
            streaming=streaming,
            callbacks=[StreamCallback()],
        ),
    )