openllm

1
# NOTE: Make sure to install openai>1
2
import os, openai, typing as t
3
from openai.types.chat import (
4
  ChatCompletionMessageParam,
5
  ChatCompletionSystemMessageParam,
6
  ChatCompletionUserMessageParam,
7
  ChatCompletionAssistantMessageParam,
8
)
9

10
client = openai.OpenAI(base_url=os.getenv('OPENLLM_ENDPOINT', 'http://localhost:3000') + '/v1', api_key='na')
11

12
models = client.models.list()
13
print('Models:', models.model_dump_json(indent=2))
14
model = models.data[0].id
15

16
# Chat completion API
17
stream = str(os.getenv('STREAM', False)).upper() in ['TRUE', '1', 'YES', 'Y', 'ON']
18
messages: t.List[ChatCompletionMessageParam]= [
19
  ChatCompletionSystemMessageParam(role='system', content='You are acting as Ernest Hemmingway.'),
20
  ChatCompletionUserMessageParam(role='user', content='Hi there!'),
21
  ChatCompletionAssistantMessageParam(role='assistant', content='Yes?'),
22
  ChatCompletionUserMessageParam(role='user', content='What is the meaning of life?'),
23
]
24
completions = client.chat.completions.create(messages=messages, model=model, max_tokens=128, stream=stream)
25

26
print(f'Chat completion result (stream={stream}):')
27
if stream:
28
  for chunk in completions:
29
    text = chunk.choices[0].delta.content
30
    if text:
31
      print(text, flush=True, end='')
32
else:
33
  print(completions)
34