jeongsk commited on
Commit
3a47e80
1 Parent(s): 7def007

chore: Add stream method to LaasApiClient for streaming chat completions API response

Browse files
Files changed (1) hide show
  1. libs/laas.py +34 -1
libs/laas.py CHANGED
@@ -1,5 +1,5 @@
1
  import os
2
- from typing import Any, Dict, List, Optional
3
 
4
  import requests
5
 
@@ -25,6 +25,10 @@ class LaasApiClient:
25
  # For chat completions:
26
  chat_response = client.chat_completions(messages=[{"role": "user", "content": "Hello!"}])
27
  print(chat_response)
 
 
 
 
28
  """
29
 
30
  BASE_URL = "https://api-laas.wanted.co.kr"
@@ -91,3 +95,32 @@ class LaasApiClient:
91
  **kwargs,
92
  }
93
  return self._make_api_call("preset/chat/completions", payload)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
2
+ from typing import Any, Dict, List, Optional, Generator
3
 
4
  import requests
5
 
 
25
  # For chat completions:
26
  chat_response = client.chat_completions(messages=[{"role": "user", "content": "Hello!"}])
27
  print(chat_response)
28
+
29
+ # For streaming:
30
+ for chunk in client.stream(messages=[{"role": "user", "content": "Tell me a story"}]):
31
+ print(chunk, end="", flush=True)
32
  """
33
 
34
  BASE_URL = "https://api-laas.wanted.co.kr"
 
95
  **kwargs,
96
  }
97
  return self._make_api_call("preset/chat/completions", payload)
98
+
99
+ def stream(
100
+ self, messages: List[Dict[str, str]], **kwargs
101
+ ) -> Generator[str, None, None]:
102
+ """
103
+ Stream the chat completions API response.
104
+
105
+ Args:
106
+ messages (List[Dict[str, str]]): List of message dictionaries.
107
+ **kwargs: Additional keyword arguments for the API call.
108
+
109
+ Yields:
110
+ str: Chunks of the streaming response.
111
+ """
112
+ url = f"{self.base_url}/api/preset/chat/completions"
113
+ payload = {
114
+ "hash": self.hash,
115
+ "messages": messages,
116
+ "stream": True,
117
+ **kwargs,
118
+ }
119
+
120
+ with requests.post(
121
+ url, headers=self.headers, json=payload, stream=True
122
+ ) as response:
123
+ response.raise_for_status()
124
+ for line in response.iter_lines():
125
+ if line:
126
+ yield line.decode("utf-8")