File size: 1,130 Bytes
4d7183d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from __future__ import annotations
from config import Config
import openai

CFG = Config()

openai.api_key = CFG.openai_api_key
openai.api_base = CFG.openai_api_base

from typing import Optional

def llm_response(model, 
             messages, 
             temperature: float = CFG.temperature,
             max_tokens: Optional[int] = None):
    return openai.ChatCompletion.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
        ).choices[0].message["content"]


def llm_stream_response(model, 
                        messages, 
                        temperature: float = CFG.temperature, 
                        max_tokens: Optional[int] = None):
    response = ""
    for chunk in openai.ChatCompletion.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
    ):
        content = chunk["choices"][0].get("delta", {}).get("content")
        if content is not None:
            response += content
            yield response