Severian commited on
Commit
8b877e3
1 Parent(s): 33af6ab

Update llm_handler.py

Browse files
Files changed (1) hide show
  1. llm_handler.py +5 -32
llm_handler.py CHANGED
@@ -11,52 +11,25 @@ client = OpenAI(api_key=OPENAI_API_KEY)
11
  def send_to_chatgpt(msg_list):
12
  try:
13
  completion = client.chat.completions.create(
14
- model=OPENAI_MODEL,
15
  messages=msg_list,
16
  temperature=0.6,
17
  stream=True
18
  )
19
 
20
- chatgpt_response = ""
21
  for chunk in completion:
22
  if chunk.choices[0].delta.content is not None:
23
- chatgpt_response += chunk.choices[0].delta.content
24
 
25
- # Note: Usage information might not be available with LlamaNet
26
- chatgpt_usage = None
27
- return chatgpt_response, chatgpt_usage
28
  except Exception as e:
29
  print(f"Error in send_to_chatgpt: {str(e)}")
30
  return f"Error: {str(e)}", None
31
 
32
- def send_to_llamanet(msg_list):
33
- try:
34
- # Create a new OpenAI client for LlamaNet (no API key needed)
35
- llamanet_client = OpenAI()
36
-
37
- # Send request to LlamaNet
38
- completion = llamanet_client.chat.completions.create(
39
- model="gpt-3.5-turbo", # LlamaNet uses this as a placeholder
40
- messages=msg_list,
41
- stream=True
42
- )
43
-
44
- llamanet_response = ""
45
- for chunk in completion:
46
- if chunk.choices[0].delta.content is not None:
47
- llamanet_response += chunk.choices[0].delta.content
48
-
49
- # LlamaNet doesn't provide usage information
50
- llamanet_usage = None
51
- return llamanet_response, llamanet_usage
52
- except Exception as e:
53
- print(f"Error in send_to_llamanet: {str(e)}")
54
- return f"Error: {str(e)}", None
55
-
56
  def send_to_llm(provider, msg_list):
57
  if provider == "llamanet":
58
- return send_to_llamanet(msg_list)
59
- elif provider == "openai":
60
  return send_to_chatgpt(msg_list)
61
  else:
62
  raise ValueError(f"Unknown provider: {provider}")
 
11
  def send_to_chatgpt(msg_list):
12
  try:
13
  completion = client.chat.completions.create(
14
+ model=OPENAI_MODEL, # LlamaNet ignores this
15
  messages=msg_list,
16
  temperature=0.6,
17
  stream=True
18
  )
19
 
20
+ response = ""
21
  for chunk in completion:
22
  if chunk.choices[0].delta.content is not None:
23
+ response += chunk.choices[0].delta.content
24
 
25
+ return response, None # LlamaNet doesn't provide usage info
26
+
 
27
  except Exception as e:
28
  print(f"Error in send_to_chatgpt: {str(e)}")
29
  return f"Error: {str(e)}", None
30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
  def send_to_llm(provider, msg_list):
32
  if provider == "llamanet":
 
 
33
  return send_to_chatgpt(msg_list)
34
  else:
35
  raise ValueError(f"Unknown provider: {provider}")